hamtaa-texttools 1.0.1__py3-none-any.whl → 1.1.7__py3-none-any.whl

texttools/tools/internals/base_operator.py

@@ -0,0 +1,100 @@
+from typing import TypeVar, Type, Any
+import json
+import re
+import math
+import logging
+import random
+
+from pydantic import BaseModel
+from openai import OpenAI, AsyncOpenAI
+
+# Base Model type for output models
+T = TypeVar("T", bound=BaseModel)
+
+logger = logging.getLogger("texttools.base_operator")
+
+
+class BaseOperator:
+    def __init__(self, client: OpenAI | AsyncOpenAI, model: str):
+        self.client = client
+        self.model = model
+
+    def _build_user_message(self, prompt: str) -> dict[str, str]:
+        return {"role": "user", "content": prompt}
+
+    def _clean_json_response(self, response: str) -> str:
+        """
+        Clean JSON response by removing code block markers and whitespace.
+        Handles cases like:
+        - ```json{"result": "value"}```
+        """
+        stripped = response.strip()
+        cleaned = re.sub(r"^```(?:json)?\s*", "", stripped)
+        cleaned = re.sub(r"\s*```$", "", cleaned)
+
+        return cleaned.strip()
+
+    def _convert_to_output_model(
+        self, response_string: str, output_model: Type[T]
+    ) -> Type[T]:
+        """
+        Convert a JSON response string to output model.
+        """
+        # Clean the response string
+        cleaned_json = self._clean_json_response(response_string)
+
+        # Fix Python-style booleans
+        cleaned_json = cleaned_json.replace("False", "false").replace("True", "true")
+
+        # Convert string to Python dictionary
+        response_dict = json.loads(cleaned_json)
+
+        # Convert dictionary to output model
+        return output_model(**response_dict)
+
+    def _extract_logprobs(self, completion: dict) -> list[dict[str, Any]]:
+        """
+        Extracts and filters token probabilities from completion logprobs.
+        Skips punctuation and structural tokens, returns cleaned probability data.
+        """
+        logprobs_data = []
+
+        ignore_pattern = re.compile(r'^(result|[\s\[\]\{\}",:]+)$')
+
+        for choice in completion.choices:
+            if not getattr(choice, "logprobs", None):
+                logger.error("logprobs is not avalible in the chosen model.")
+                return []
+
+            for logprob_item in choice.logprobs.content:
+                if ignore_pattern.match(logprob_item.token):
+                    continue
+                token_entry = {
+                    "token": logprob_item.token,
+                    "prob": round(math.exp(logprob_item.logprob), 8),
+                    "top_alternatives": [],
+                }
+                for alt in logprob_item.top_logprobs:
+                    if ignore_pattern.match(alt.token):
+                        continue
+                    token_entry["top_alternatives"].append(
+                        {
+                            "token": alt.token,
+                            "prob": round(math.exp(alt.logprob), 8),
+                        }
+                    )
+                logprobs_data.append(token_entry)
+
+        return logprobs_data
+
+    def _get_retry_temp(self, base_temp: float) -> float:
+        """
+        Calculate temperature for retry attempts.
+        """
+        delta_temp = random.choice([-1, 1]) * random.uniform(0.1, 0.9)
+        new_temp = base_temp + delta_temp
+        print(f"Base Temp: {base_temp}")
+        print(f"Delta Temp: {delta_temp}")
+        print(f"New Temp: {new_temp}")
+
+        return max(0.0, min(new_temp, 1.5))

texttools/tools/internals/formatters.py

@@ -0,0 +1,24 @@
+class Formatter:
+    @staticmethod
+    def user_merge_format(messages: list[dict[str, str]]) -> list[dict[str, str]]:
+        """
+        Merges consecutive user messages into a single message, separated by newlines.
+
+        This is useful for condensing a multi-turn user input into a single
+        message for the LLM. Assistant and system messages are left unchanged and
+        act as separators between user message groups.
+        """
+        merged: list[dict[str, str]] = []
+
+        for message in messages:
+            role, content = message["role"], message["content"].strip()
+
+            # Merge with previous user turn
+            if merged and role == "user" and merged[-1]["role"] == "user":
+                merged[-1]["content"] += "\n" + content
+
+            # Otherwise, start a new turn
+            else:
+                merged.append({"role": role, "content": content})
+
+        return merged

texttools/tools/internals/operator.py

@@ -0,0 +1,242 @@
+from typing import Any, TypeVar, Type, Literal, Callable
+import logging
+
+from openai import OpenAI
+from pydantic import BaseModel
+
+from texttools.tools.internals.output_models import ToolOutput
+from texttools.tools.internals.base_operator import BaseOperator
+from texttools.tools.internals.formatters import Formatter
+from texttools.tools.internals.prompt_loader import PromptLoader
+
+# Base Model type for output models
+T = TypeVar("T", bound=BaseModel)
+
+logger = logging.getLogger("texttools.operator")
+
+
+class Operator(BaseOperator):
+    """
+    Core engine for running text-processing operations with an LLM (Sync).
+
+    It wires together:
+    - `PromptLoader` → loads YAML prompt templates.
+    - `UserMergeFormatter` → applies formatting to messages (e.g., merging).
+    - OpenAI client → executes completions/parsed completions.
+    """
+
+    def __init__(self, client: OpenAI, model: str):
+        self.client = client
+        self.model = model
+
+    def _analyze(self, prompt_configs: dict[str, str], temperature: float) -> str:
+        """
+        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 = [self._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
+
+    def _parse_completion(
+        self,
+        message: list[dict[str, str]],
+        output_model: Type[T],
+        temperature: float,
+        logprobs: bool = False,
+        top_logprobs: int = 3,
+    ) -> tuple[Type[T], Any]:
+        """
+        Parses a chat completion using OpenAI's structured output format.
+        Returns both the parsed object and the raw completion for logging.
+        """
+        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
+
+        completion = self.client.beta.chat.completions.parse(**request_kwargs)
+        parsed = completion.choices[0].message.parsed
+        return parsed, completion
+
+    def _vllm_completion(
+        self,
+        message: list[dict[str, str]],
+        output_model: Type[T],
+        temperature: float,
+        logprobs: bool = False,
+        top_logprobs: int = 3,
+    ) -> tuple[Type[T], Any]:
+        """
+        Generates a completion using vLLM with JSON schema guidance.
+        Returns the parsed output model and raw completion.
+        """
+        json_schema = output_model.model_json_schema()
+
+        # Build kwargs dynamically
+        request_kwargs = {
+            "model": self.model,
+            "messages": message,
+            "extra_body": {"guided_json": json_schema},
+            "temperature": temperature,
+        }
+
+        if logprobs:
+            request_kwargs["logprobs"] = True
+            request_kwargs["top_logprobs"] = top_logprobs
+
+        completion = self.client.chat.completions.create(**request_kwargs)
+        response = completion.choices[0].message.content
+
+        # Convert the string response to output model
+        parsed = self._convert_to_output_model(response, output_model)
+        return parsed, completion
+
+    def run(
+        self,
+        # User parameters
+        text: str,
+        with_analysis: bool,
+        output_lang: str | None,
+        user_prompt: str | None,
+        temperature: float,
+        logprobs: bool,
+        top_logprobs: int | None,
+        validator: Callable[[Any], bool] | None,
+        # Internal parameters
+        prompt_file: str,
+        output_model: Type[T],
+        resp_format: Literal["vllm", "parse"],
+        mode: str | None,
+        **extra_kwargs,
+    ) -> ToolOutput:
+        """
+        Execute the LLM pipeline with the given input text.
+        """
+        prompt_loader = PromptLoader()
+        formatter = Formatter()
+        output = ToolOutput()
+
+        try:
+            # Prompt configs contain two keys: main_template and analyze template, both are string
+            prompt_configs = prompt_loader.load(
+                prompt_file=prompt_file,
+                text=text.strip(),
+                mode=mode,
+                **extra_kwargs,
+            )
+
+            messages: list[dict[str, str]] = []
+
+            if with_analysis:
+                analysis = self._analyze(prompt_configs, temperature)
+                messages.append(
+                    self._build_user_message(f"Based on this analysis: {analysis}")
+                )
+
+            if output_lang:
+                messages.append(
+                    self._build_user_message(
+                        f"Respond only in the {output_lang} language."
+                    )
+                )
+
+            if user_prompt:
+                messages.append(
+                    self._build_user_message(f"Consider this instruction {user_prompt}")
+                )
+
+            messages.append(self._build_user_message(prompt_configs["main_template"]))
+            messages = formatter.user_merge_format(messages)
+
+            if resp_format == "vllm":
+                parsed, completion = self._vllm_completion(
+                    messages, output_model, temperature, logprobs, top_logprobs
+                )
+            elif resp_format == "parse":
+                parsed, completion = self._parse_completion(
+                    messages, output_model, temperature, logprobs, top_logprobs
+                )
+
+            # Ensure output_model has a `result` field
+            if not hasattr(parsed, "result"):
+                error = "The provided output_model must define a field named 'result'"
+                logger.error(error)
+                output.errors.append(error)
+                return output
+
+            output.result = parsed.result
+
+            # Retry logic if validation fails
+            if validator and not validator(output.result):
+                max_retries = 3
+                for attempt in range(max_retries):
+                    logger.warning(
+                        f"Validation failed, retrying for the {attempt + 1} time."
+                    )
+
+                    # Generate new temperature for retry
+                    retry_temperature = self._get_retry_temp(temperature)
+                    try:
+                        if resp_format == "vllm":
+                            parsed, completion = self._vllm_completion(
+                                messages,
+                                output_model,
+                                retry_temperature,
+                                logprobs,
+                                top_logprobs,
+                            )
+                        elif resp_format == "parse":
+                            parsed, completion = self._parse_completion(
+                                messages,
+                                output_model,
+                                retry_temperature,
+                                logprobs,
+                                top_logprobs,
+                            )
+
+                        output.result = parsed.result
+
+                        # Check if retry was successful
+                        if validator(output.result):
+                            logger.info(
+                                f"Validation passed on retry attempt {attempt + 1}"
+                            )
+                            break
+                        else:
+                            logger.warning(
+                                f"Validation still failing after retry attempt {attempt + 1}"
+                            )
+
+                    except Exception 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 logprobs:
+                output.logprobs = self._extract_logprobs(completion)
+
+            if with_analysis:
+                output.analysis = analysis
+
+            return output
+
+        except Exception as e:
+            logger.error(f"TheTool failed: {e}")
+            output.errors.append(str(e))
+            return output

texttools/tools/internals/output_models.py

@@ -0,0 +1,62 @@
+from typing import Literal, Any
+
+from pydantic import BaseModel, Field
+
+
+class ToolOutput(BaseModel):
+    result: Any = None
+    analysis: str = ""
+    logprobs: list[dict[str, Any]] = []
+    errors: list[str] = []
+
+    def __repr__(self) -> str:
+        return f"ToolOutput(result_type='{type(self.result)}', result='{self.result}', analysis='{self.analysis}', logprobs='{self.logprobs}', errors='{self.errors}'"
+
+
+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

@@ -0,0 +1,60 @@
+from functools import lru_cache
+from pathlib import Path
+import yaml
+
+
+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.
+    - Return a dict with:
+        {
+            "main_template": "...",
+            "analyze_template": "..." | None
+        }
+    """
+
+    MAIN_TEMPLATE: str = "main_template"
+    ANALYZE_TEMPLATE: str = "analyze_template"
+
+    # 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.
+        """
+        base_dir = Path(__file__).parent.parent.parent / Path("prompts")
+        prompt_path = base_dir / prompt_file
+        data = yaml.safe_load(prompt_path.read_text(encoding="utf-8"))
+
+        return {
+            self.MAIN_TEMPLATE: data[self.MAIN_TEMPLATE][mode]
+            if mode
+            else data[self.MAIN_TEMPLATE],
+            self.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, **extra_kwargs
+    ) -> dict[str, str]:
+        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