hamtaa-texttools 1.0.4__py3-none-any.whl → 1.0.5__py3-none-any.whl

texttools/tools/internals/output_models.py

@@ -1,62 +1,52 @@
-from typing import Literal
-
-from pydantic import BaseModel
-
-
-class StrOutput(BaseModel):
-    """
-    Output model for a single string result.
-    """
-
-    result: str
-
-
-class BoolOutput(BaseModel):
-    """
-    Output model for a single boolean result.
-    """
-
-    result: bool
-
-
-class ListStrOutput(BaseModel):
-    """
-    Output model for a list of strings result.
-    """
-
-    result: list[str]
-
-
-class ListDictStrStrOutput(BaseModel):
-    """
-    Output model for a list of dictionaries with string key-value pairs.
-    """
-
-    result: list[dict[str, str]]
-
-
-class ReasonListStrOutput(BaseModel):
-    """
-    Output model containing a reasoning string followed by a list of strings.
-    """
-
-    reason: str
-    result: list[str]
-
-
-class CategorizerOutput(BaseModel):
-    """
-    Output model for categorization with reasoning and a predefined category result.
-    """
-
-    reason: str
-    result: Literal[
-        "باورهای دینی",
-        "اخلاق اسلامی",
-        "احکام و فقه",
-        "تاریخ اسلام و شخصیت ها",
-        "منابع دینی",
-        "دین و جامعه/سیاست",
-        "عرفان و معنویت",
-        "هیچکدام",
-    ]
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class StrOutput(BaseModel):
+    result: str = Field(..., description="The output string")
+
+
+class BoolOutput(BaseModel):
+    result: bool = Field(
+        ..., description="Boolean indicating the output state", example=True
+    )
+
+
+class ListStrOutput(BaseModel):
+    result: list[str] = Field(
+        ..., description="The output list of strings", example=["text_1", "text_2"]
+    )
+
+
+class ListDictStrStrOutput(BaseModel):
+    result: list[dict[str, str]] = Field(
+        ...,
+        description="List of dictionaries containing string key-value pairs",
+        example=[{"text": "Mohammad", "type": "PER"}],
+    )
+
+
+class ReasonListStrOutput(BaseModel):
+    reason: str = Field(..., description="Thinking process that led to the output")
+    result: list[str] = Field(..., description="The output list of strings")
+
+
+class CategorizerOutput(BaseModel):
+    reason: str = Field(
+        ..., description="Explanation of why the input belongs to the category"
+    )
+    result: Literal[
+        "باورهای دینی",
+        "اخلاق اسلامی",
+        "احکام و فقه",
+        "تاریخ اسلام و شخصیت ها",
+        "منابع دینی",
+        "دین و جامعه/سیاست",
+        "عرفان و معنویت",
+        "هیچکدام",
+    ] = Field(
+        ...,
+        description="Predicted category label",
+        example="اخلاق اسلامی",
+    )

texttools/tools/internals/prompt_loader.py

@@ -1,82 +1,76 @@
-from pathlib import Path
-import yaml
-
-
-class PromptLoader:
-    """
-    Utility for loading and formatting YAML prompt templates.
-
-    Each YAML file under `prompts/` must define at least a `main_template`,
-    and optionally an `analyze_template`. These can either be a single string
-    or a dictionary keyed by mode names (if `use_modes=True`).
-
-    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.
-    - Return a dict with:
-        {
-            "main_template": "...",
-            "analyze_template": "..." | None
-        }
-    """
-
-    MAIN_TEMPLATE: str = "main_template"
-    ANALYZE_TEMPLATE: str = "analyze_template"
-
-    def _get_prompt_path(self, prompt_file: str, prompts_dir: str) -> Path:
-        return Path(__file__).parent.parent.parent / prompts_dir / prompt_file
-
-    def _load_templates(
-        self,
-        prompts_dir: str,
-        prompt_file: str,
-        use_modes: bool,
-        mode: str,
-    ) -> dict[str, str]:
-        prompt_path = self._get_prompt_path(prompt_file, prompts_dir)
-
-        if not prompt_path.exists():
-            raise FileNotFoundError(f"Prompt file not found: {prompt_path}")
-
-        try:
-            # Load the data
-            data = yaml.safe_load(prompt_path.read_text(encoding="utf-8"))
-        except yaml.YAMLError as e:
-            raise ValueError(f"Invalid YAML in {prompt_path}: {e}")
-
-        return {
-            "main_template": data["main_template"][mode]
-            if use_modes
-            else data["main_template"],
-            "analyze_template": data.get("analyze_template")[mode]
-            if use_modes
-            else data.get("analyze_template"),
-        }
-
-    def _build_format_args(self, input_text: str, **extra_kwargs) -> dict[str, str]:
-        # Base formatting args
-        format_args = {"input": input_text}
-        # Merge extras
-        format_args.update(extra_kwargs)
-        return format_args
-
-    def load_prompts(
-        self,
-        prompt_file: str,
-        use_modes: bool,
-        mode: str,
-        input_text: str,
-        prompts_dir: str = "prompts",
-        **extra_kwargs,
-    ) -> dict[str, str]:
-        template_configs = self._load_templates(
-            prompts_dir, prompt_file, use_modes, mode
-        )
-        format_args = self._build_format_args(input_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
+from typing import Optional
+from pathlib import Path
+import yaml
+
+
+class PromptLoader:
+    """
+    Utility for loading and formatting YAML prompt templates.
+
+    Each YAML file under `prompts/` must define at least a `main_template`,
+    and optionally an `analyze_template`. These can either be a single string
+    or a dictionary keyed by mode names (if `use_modes=True`).
+
+    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.
+    - Return a dict with:
+        {
+            "main_template": "...",
+            "analyze_template": "..." | None
+        }
+    """
+
+    MAIN_TEMPLATE: str = "main_template"
+    ANALYZE_TEMPLATE: str = "analyze_template"
+
+    def _load_templates(
+        self,
+        prompts_dir: str,
+        prompt_file: str,
+        mode: str | None,
+    ) -> dict[str, str]:
+        prompt_path = Path(__file__).parent.parent.parent / prompts_dir / prompt_file
+
+        if not prompt_path.exists():
+            raise FileNotFoundError(f"Prompt file not found: {prompt_path}")
+
+        try:
+            # Load the data
+            data = yaml.safe_load(prompt_path.read_text(encoding="utf-8"))
+        except yaml.YAMLError as e:
+            raise ValueError(f"Invalid YAML in {prompt_path}: {e}")
+
+        return {
+            "main_template": data[self.MAIN_TEMPLATE][mode]
+            if mode
+            else data[self.MAIN_TEMPLATE],
+            "analyze_template": data.get(self.ANALYZE_TEMPLATE)[mode]
+            if mode
+            else data.get(self.ANALYZE_TEMPLATE),
+        }
+
+    def _build_format_args(self, text: str, **extra_kwargs) -> dict[str, str]:
+        # Base formatting args
+        format_args = {"input": text}
+        # Merge extras
+        format_args.update(extra_kwargs)
+        return format_args
+
+    def load(
+        self,
+        prompt_file: str,
+        text: str,
+        mode: str,
+        prompts_dir: str = "prompts",
+        **extra_kwargs,
+    ) -> dict[str, str]:
+        template_configs = self._load_templates(prompts_dir, 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