hamtaa-texttools 1.1.16__py3-none-any.whl → 1.1.18__py3-none-any.whl

texttools/{tools/internals → internals}/models.py

@@ -6,19 +6,24 @@ from pydantic import BaseModel, Field, create_model
 
 class ToolOutput(BaseModel):
     result: Any = None
-    analysis: str = ""
     logprobs: list[dict[str, Any]] = []
-    process: str = ""
+    analysis: str = ""
+    process: str | None = None
     processed_at: datetime = datetime.now()
-    execution_time: float = -1.0
+    execution_time: float | None = None
     errors: list[str] = []
 
     def __repr__(self) -> str:
-        return f"ToolOutput(process='{self.process}', result_type='{type(self.result)}', result='{self.result}', analysis='{self.analysis}', logprobs='{self.logprobs}', errors='{self.errors}', processed_at='{self.processed_at}', execution_time='{self.execution_time}'"
+        return f"""
+        ToolOutput(process='{self.process}', result_type='{type(self.result)}', 
+        result='{self.result}', analysis='{self.analysis}', 
+        logprobs='{self.logprobs}', errors='{self.errors}', 
+        processed_at='{self.processed_at}', execution_time='{self.execution_time}'
+        """
 
 
 class StrOutput(BaseModel):
-    result: str = Field(..., description="The output string")
+    result: str = Field(..., description="The output string", example="text")
 
 
 class BoolOutput(BaseModel):
@@ -37,13 +42,15 @@ class ListDictStrStrOutput(BaseModel):
     result: list[dict[str, str]] = Field(
         ...,
         description="List of dictionaries containing string key-value pairs",
-        example=[{"text": "Mohammad", "type": "PER"}],
+        example=[{"text": "Mohammad", "type": "PER"}, {"text": "Iran", "type": "LOC"}],
     )
 
 
 class ReasonListStrOutput(BaseModel):
     reason: str = Field(..., description="Thinking process that led to the output")
-    result: list[str] = Field(..., description="The output list of strings")
+    result: list[str] = Field(
+        ..., description="The output list of strings", example=["text_1", "text_2"]
+    )
 
 
 class Node(BaseModel):
@@ -51,14 +58,44 @@ class Node(BaseModel):
     name: str
     level: int
     parent_id: int | None
-    description: str = "No description provided"
+    description: str
 
 
 class CategoryTree:
     def __init__(self, tree_name):
-        self.root = Node(node_id=0, name=tree_name, level=0, parent_id=None)
-        self.all_nodes: list[Node] = [self.root]
-        self.new_id = 1
+        self._root = Node(
+            node_id=0, name=tree_name, level=0, parent_id=None, description="Root node"
+        )
+        self._all_nodes: list[Node] = [self._root]
+        self._new_id = 1
+
+    def get_all_nodes(self) -> list[Node]:
+        return self._all_nodes
+
+    def get_level_count(self) -> int:
+        return max([item.level for item in self._all_nodes])
+
+    def get_node(self, identifier: int | str) -> Node | None:
+        if isinstance(identifier, str):
+            for node in self.get_all_nodes():
+                if node.name == identifier:
+                    return node
+            return None
+        elif isinstance(identifier, int):
+            for node in self.get_all_nodes():
+                if node.node_id == identifier:
+                    return node
+            return None
+        else:
+            return None
+
+    def get_children(self, parent_node: Node) -> list[Node] | None:
+        children = [
+            node
+            for node in self.get_all_nodes()
+            if parent_node.node_id == node.parent_id
+        ]
+        return children if children else None
 
     def add_node(
         self,
@@ -66,12 +103,12 @@ class CategoryTree:
         parent_name: str | None = None,
         description: str | None = None,
     ) -> None:
-        if self.find_node(node_name):
+        if self.get_node(node_name):
             raise ValueError(f"{node_name} has been chosen for another category before")
 
         if parent_name:
-            parent_node = self.find_node(parent_name)
-            if parent_node is None:
+            parent_node = self.get_node(parent_name)
+            if not parent_node:
                 raise ValueError(f"Parent category '{parent_name}' not found")
             parent_id = parent_node.node_id
             level = parent_node.level + 1
@@ -80,61 +117,31 @@ class CategoryTree:
             parent_id = 0
 
         node_data = {
-            "node_id": self.new_id,
+            "node_id": self._new_id,
             "name": node_name,
             "level": level,
             "parent_id": parent_id,
+            "description": description if description else "No description provided",
         }
 
-        if description is not None:
-            node_data["description"] = description
-
-        self.all_nodes.append(Node(**node_data))
-        self.new_id += 1
-
-    def get_nodes(self) -> list[Node]:
-        return self.all_nodes
-
-    def get_level_count(self) -> int:
-        return max([item.level for item in self.all_nodes])
-
-    def find_node(self, identifier: int | str) -> Node | None:
-        if isinstance(identifier, str):
-            for node in self.get_nodes():
-                if node.name == identifier:
-                    return node
-            return None
-        elif isinstance(identifier, int):
-            for node in self.get_nodes():
-                if node.node_id == identifier:
-                    return node
-            return None
-        else:
-            return None
-
-    def find_children(self, parent_node: Node) -> list[Node] | None:
-        children = [
-            node for node in self.get_nodes() if parent_node.node_id == node.parent_id
-        ]
-        return children if children else None
+        self._all_nodes.append(Node(**node_data))
+        self._new_id += 1
 
     def remove_node(self, identifier: int | str) -> None:
-        node = self.find_node(identifier)
+        node = self.get_node(identifier)
 
-        if node is not None:
+        if node:
             # Remove node's children recursively
-            children = self.find_children(node)
+            children = self.get_children(node)
 
-            # Ending condition
-            if children is None:
-                self.all_nodes.remove(node)
+            if not children:
+                self._all_nodes.remove(node)
                 return
 
             for child in children:
                 self.remove_node(child.name)
 
-            # Remove the node from tree
-            self.all_nodes.remove(node)
+            self._all_nodes.remove(node)
         else:
             raise ValueError(f"Node with identifier: '{identifier}' not found.")
 
@@ -142,7 +149,7 @@ class CategoryTree:
         def build_dict(node: Node) -> dict:
             children = [
                 build_dict(child)
-                for child in self.all_nodes
+                for child in self._all_nodes
                 if child.parent_id == node.node_id
             ]
             return {
@@ -153,7 +160,7 @@ class CategoryTree:
                 "children": children,
             }
 
-        return {"category_tree": build_dict(self.root)["children"]}
+        return {"category_tree": build_dict(self._root)["children"]}
 
 
 # This function is needed to create CategorizerOutput with dynamic categories

texttools/internals/prompt_loader.py

@@ -0,0 +1,80 @@
+from functools import lru_cache
+from pathlib import Path
+import yaml
+
+from texttools.internals.exceptions import PromptError
+
+
+class PromptLoader:
+    """
+    Utility for loading and formatting YAML prompt templates.
+
+    Responsibilities:
+    - Load and parse YAML prompt definitions.
+    - Select the right template (by mode, if applicable).
+    - Inject variables (`{input}`, plus any extra kwargs) into the templates.
+    """
+
+    MAIN_TEMPLATE = "main_template"
+    ANALYZE_TEMPLATE = "analyze_template"
+
+    @staticmethod
+    def _build_format_args(text: str, **extra_kwargs) -> dict[str, str]:
+        # Base formatting args
+        format_args = {"input": text}
+        # Merge extras
+        format_args.update(extra_kwargs)
+        return format_args
+
+    # Use lru_cache to load each file once
+    @lru_cache(maxsize=32)
+    def _load_templates(self, prompt_file: str, mode: str | None) -> dict[str, str]:
+        """
+        Loads prompt templates from YAML file with optional mode selection.
+        """
+        try:
+            base_dir = Path(__file__).parent.parent / Path("prompts")
+            prompt_path = base_dir / prompt_file
+
+            if not prompt_path.exists():
+                raise PromptError(f"Prompt file not found: {prompt_file}")
+
+            data = yaml.safe_load(prompt_path.read_text(encoding="utf-8"))
+
+            if self.MAIN_TEMPLATE not in data:
+                raise PromptError(f"Missing 'main_template' in {prompt_file}")
+
+            if mode and mode not in data.get(self.MAIN_TEMPLATE, {}):
+                raise PromptError(f"Mode '{mode}' not found in {prompt_file}")
+
+            return {
+                self.MAIN_TEMPLATE: data[self.MAIN_TEMPLATE][mode]
+                if mode and isinstance(data[self.MAIN_TEMPLATE], dict)
+                else data[self.MAIN_TEMPLATE],
+                self.ANALYZE_TEMPLATE: data.get(self.ANALYZE_TEMPLATE, {}).get(mode)
+                if mode and isinstance(data.get(self.ANALYZE_TEMPLATE), dict)
+                else data.get(self.ANALYZE_TEMPLATE, ""),
+            }
+
+        except yaml.YAMLError as e:
+            raise PromptError(f"Invalid YAML in {prompt_file}: {e}")
+        except Exception as e:
+            raise PromptError(f"Failed to load prompt {prompt_file}: {e}")
+
+    def load(
+        self, prompt_file: str, text: str, mode: str, **extra_kwargs
+    ) -> dict[str, str]:
+        try:
+            template_configs = self._load_templates(prompt_file, mode)
+            format_args = self._build_format_args(text, **extra_kwargs)
+
+            # Inject variables inside each template
+            for key in template_configs.keys():
+                template_configs[key] = template_configs[key].format(**format_args)
+
+            return template_configs
+
+        except KeyError as e:
+            raise PromptError(f"Missing template variable: {e}")
+        except Exception as e:
+            raise PromptError(f"Failed to format prompt: {e}")

texttools/{tools/internals → internals}/sync_operator.py

@@ -5,15 +5,21 @@ import logging
 from openai import OpenAI
 from pydantic import BaseModel
 
-from texttools.tools.internals.models import ToolOutput
-from texttools.tools.internals.operator_utils import OperatorUtils
-from texttools.tools.internals.formatters import Formatter
-from texttools.tools.internals.prompt_loader import PromptLoader
+from texttools.internals.models import ToolOutput
+from texttools.internals.operator_utils import OperatorUtils
+from texttools.internals.formatters import Formatter
+from texttools.internals.prompt_loader import PromptLoader
+from texttools.internals.exceptions import (
+    TextToolsError,
+    LLMError,
+    ValidationError,
+    PromptError,
+)
 
 # Base Model type for output models
 T = TypeVar("T", bound=BaseModel)
 
-logger = logging.getLogger("texttools.operator")
+logger = logging.getLogger("texttools.sync_operator")
 
 
 class Operator:
@@ -35,15 +41,33 @@ class Operator:
         Calls OpenAI API for analysis using the configured prompt template.
         Returns the analyzed content as a string.
         """
-        analyze_prompt = prompt_configs["analyze_template"]
-        analyze_message = [OperatorUtils.build_user_message(analyze_prompt)]
-        completion = self._client.chat.completions.create(
-            model=self._model,
-            messages=analyze_message,
-            temperature=temperature,
-        )
-        analysis = completion.choices[0].message.content.strip()
-        return analysis
+        try:
+            analyze_prompt = prompt_configs["analyze_template"]
+
+            if not analyze_prompt:
+                raise PromptError("Analyze template is empty")
+
+            analyze_message = [OperatorUtils.build_user_message(analyze_prompt)]
+            completion = self._client.chat.completions.create(
+                model=self._model,
+                messages=analyze_message,
+                temperature=temperature,
+            )
+
+            if not completion.choices:
+                raise LLMError("No choices returned from LLM")
+
+            analysis = completion.choices[0].message.content.strip()
+
+            if not analysis:
+                raise LLMError("Empty analysis response")
+
+            return analysis.strip()
+
+        except Exception as e:
+            if isinstance(e, (PromptError, LLMError)):
+                raise
+            raise LLMError(f"Analysis failed: {e}")
 
     def _parse_completion(
         self,
@@ -58,23 +82,35 @@ class Operator:
         Parses a chat completion using OpenAI's structured output format.
         Returns both the parsed object and the raw completion for logprobs.
         """
-        request_kwargs = {
-            "model": self._model,
-            "messages": message,
-            "response_format": output_model,
-            "temperature": temperature,
-        }
+        try:
+            request_kwargs = {
+                "model": self._model,
+                "messages": message,
+                "response_format": output_model,
+                "temperature": temperature,
+            }
+
+            if logprobs:
+                request_kwargs["logprobs"] = True
+                request_kwargs["top_logprobs"] = top_logprobs
+            if priority:
+                request_kwargs["extra_body"] = {"priority": priority}
+            completion = self._client.beta.chat.completions.parse(**request_kwargs)
+
+            if not completion.choices:
+                raise LLMError("No choices returned from LLM")
+
+            parsed = completion.choices[0].message.parsed
 
-        if logprobs:
-            request_kwargs["logprobs"] = True
-            request_kwargs["top_logprobs"] = top_logprobs
+            if not parsed:
+                raise LLMError("Failed to parse LLM response")
 
-        if priority:
-            request_kwargs["extra_body"] = {"priority": priority}
+            return parsed, completion
 
-        completion = self._client.beta.chat.completions.parse(**request_kwargs)
-        parsed = completion.choices[0].message.parsed
-        return parsed, completion
+        except Exception as e:
+            if isinstance(e, LLMError):
+                raise
+            raise LLMError(f"Completion failed: {e}")
 
     def run(
         self,
@@ -96,12 +132,13 @@ class Operator:
         **extra_kwargs,
     ) -> ToolOutput:
         """
-        Execute the LLM pipeline with the given input text.
+        Execute the LLM pipeline with the given input text. (Sync)
         """
-        prompt_loader = PromptLoader()
-        formatter = Formatter()
-        output = ToolOutput()
         try:
+            prompt_loader = PromptLoader()
+            formatter = Formatter()
+            output = ToolOutput()
+
             # Prompt configs contain two keys: main_template and analyze template, both are string
             prompt_configs = prompt_loader.load(
                 prompt_file=prompt_file,
@@ -140,6 +177,9 @@ class Operator:
 
             messages = formatter.user_merge_format(messages)
 
+            if logprobs and (not isinstance(top_logprobs, int) or top_logprobs < 2):
+                raise ValueError("top_logprobs should be an integer greater than 1")
+
             parsed, completion = self._parse_completion(
                 messages, output_model, temperature, logprobs, top_logprobs, priority
             )
@@ -148,6 +188,15 @@ class Operator:
 
             # Retry logic if validation fails
             if validator and not validator(output.result):
+                if (
+                    not isinstance(max_validation_retries, int)
+                    or max_validation_retries < 1
+                ):
+                    raise ValueError(
+                        "max_validation_retries should be a positive integer"
+                    )
+
+                succeeded = False
                 for attempt in range(max_validation_retries):
                     logger.warning(
                         f"Validation failed, retrying for the {attempt + 1} time."
@@ -155,6 +204,7 @@ class Operator:
 
                     # Generate new temperature for retry
                     retry_temperature = OperatorUtils.get_retry_temp(temperature)
+
                     try:
                         parsed, completion = self._parse_completion(
                             messages,
@@ -162,28 +212,23 @@ class Operator:
                             retry_temperature,
                             logprobs,
                             top_logprobs,
+                            priority=priority,
                         )
 
                         output.result = parsed.result
 
                         # Check if retry was successful
                         if validator(output.result):
-                            logger.info(
-                                f"Validation passed on retry attempt {attempt + 1}"
-                            )
+                            succeeded = True
                             break
-                        else:
-                            logger.warning(
-                                f"Validation still failing after retry attempt {attempt + 1}"
-                            )
 
-                    except Exception as e:
+                    except LLMError as e:
                         logger.error(f"Retry attempt {attempt + 1} failed: {e}")
-                        # Continue to next retry attempt if this one fails
 
-            # Final check after all retries
-            if validator and not validator(output.result):
-                output.errors.append("Validation failed after all retry attempts")
+                if not succeeded:
+                    raise ValidationError(
+                        f"Validation failed after {max_validation_retries} retries"
+                    )
 
             if logprobs:
                 output.logprobs = OperatorUtils.extract_logprobs(completion)
@@ -195,7 +240,7 @@ class Operator:
 
             return output
 
+        except (PromptError, LLMError, ValidationError):
+            raise
         except Exception as e:
-            logger.error(f"TheTool failed: {e}")
-            output.errors.append(str(e))
-            return output
+            raise TextToolsError(f"Unexpected error in operator: {e}")

texttools/prompts/propositionize.yaml

@@ -0,0 +1,15 @@
+main_template: |
+  You are an expert in breaking down text into atomic propositions in that language.
+  An atomic proposition is a single, self-contained fact that is concise, verifiable,
+  and does not rely on external context.
+  Each proposition must stand alone.
+  Rewrite sentences if needed to keep the context saved in each sentence.
+  Extract the atomic propositions of this text:
+  {input}
+
+analyze_template: |
+  We want to analyze this text snippet and think about where we can split sentence to atomic meaningful propositions.
+  An atomic proposition is a single, self-contained fact that is concise,
+  verifiable, and does not rely on external context.
+  You just have to think around the possible propositions in the text and how a proposition can be made.
+  {input}