hamtaa-texttools 1.1.1__py3-none-any.whl → 1.2.0__py3-none-any.whl

texttools/batch/runner.py

@@ -0,0 +1,228 @@
+import json
+import logging
+import os
+import time
+from pathlib import Path
+from typing import Any, Type, TypeVar
+
+from dotenv import load_dotenv
+from openai import OpenAI
+from pydantic import BaseModel
+
+from ..core.exceptions import TextToolsError
+from ..core.internal_models import Str
+from .config import BatchConfig
+from .manager import BatchManager
+
+# Base Model type for output models
+T = TypeVar("T", bound=BaseModel)
+
+logger = logging.getLogger("texttools.batch_runner")
+
+
+class BatchRunner:
+    """
+    Handles running batch jobs using a batch manager and configuration.
+    """
+
+    def __init__(
+        self, config: BatchConfig = BatchConfig(), output_model: Type[T] = Str
+    ):
+        try:
+            self._config = config
+            self._system_prompt = config.system_prompt
+            self._job_name = config.job_name
+            self._input_data_path = config.input_data_path
+            self._output_data_filename = config.output_data_filename
+            self._model = config.model
+            self._output_model = output_model
+            self._manager = self._init_manager()
+            self._data = self._load_data()
+            self._parts: list[list[dict[str, Any]]] = []
+            # Map part index to job name
+            self._part_idx_to_job_name: dict[int, str] = {}
+            # Track retry attempts per part
+            self._part_attempts: dict[int, int] = {}
+            self._partition_data()
+            Path(self._config.BASE_OUTPUT_DIR).mkdir(parents=True, exist_ok=True)
+
+        except Exception as e:
+            raise TextToolsError(f"Batch runner initialization failed: {e}")
+
+    def _init_manager(self) -> BatchManager:
+        load_dotenv()
+        api_key = os.getenv("OPENAI_API_KEY")
+        client = OpenAI(api_key=api_key)
+        return BatchManager(
+            client=client,
+            model=self._model,
+            prompt_template=self._system_prompt,
+            output_model=self._output_model,
+        )
+
+    def _load_data(self):
+        with open(self._input_data_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        data = self._config.export_function(data)
+
+        # Ensure data is a list of dicts with 'id' and 'content' as strings
+        if not isinstance(data, list):
+            raise ValueError(
+                "Exported data must be a list of dicts with 'id' and 'content' keys"
+            )
+        for item in data:
+            if not (isinstance(item, dict) and "id" in item and "content" in item):
+                raise ValueError(
+                    f"Item must be a dict with 'id' and 'content' keys. Got: {type(item)}"
+                )
+            if not (isinstance(item["id"], str) and isinstance(item["content"], str)):
+                raise ValueError("'id' and 'content' must be strings.")
+        return data
+
+    def _partition_data(self):
+        total_length = sum(len(item["content"]) for item in self._data)
+        prompt_length = len(self._system_prompt)
+        total = total_length + (prompt_length * len(self._data))
+        calculation = total / self._config.CHARS_PER_TOKEN
+        logger.info(
+            f"Total chars: {total_length}, Prompt chars: {prompt_length}, Total: {total}, Tokens: {calculation}"
+        )
+        if calculation < self._config.MAX_TOTAL_TOKENS:
+            self._parts = [self._data]
+        else:
+            # Partition into chunks of MAX_BATCH_SIZE
+            self._parts = [
+                self._data[i : i + self._config.MAX_BATCH_SIZE]
+                for i in range(0, len(self._data), self._config.MAX_BATCH_SIZE)
+            ]
+        logger.info(f"Data split into {len(self._parts)} part(s)")
+
+    def _submit_all_jobs(self) -> None:
+        for idx, part in enumerate(self._parts):
+            if self._result_exists(idx):
+                logger.info(f"Skipping part {idx + 1}: result already exists.")
+                continue
+            part_job_name = (
+                f"{self._job_name}_part_{idx + 1}"
+                if len(self._parts) > 1
+                else self._job_name
+            )
+            # If a job with this name already exists, register and skip submitting
+            existing_job = self._manager._load_state(part_job_name)
+            if existing_job:
+                logger.info(
+                    f"Skipping part {idx + 1}: job already exists ({part_job_name})."
+                )
+                self._part_idx_to_job_name[idx] = part_job_name
+                self._part_attempts.setdefault(idx, 0)
+                continue
+
+            payload = part
+            logger.info(
+                f"Submitting job for part {idx + 1}/{len(self._parts)}: {part_job_name}"
+            )
+            self._manager.start(payload, job_name=part_job_name)
+            self._part_idx_to_job_name[idx] = part_job_name
+            self._part_attempts.setdefault(idx, 0)
+            # This is added for letting file get uploaded, before starting the next part.
+            logger.info("Uploading...")
+            time.sleep(30)
+
+    def _save_results(
+        self,
+        output_data: list[dict[str, Any]] | dict[str, Any],
+        log: list[Any],
+        part_idx: int,
+    ):
+        part_suffix = f"_part_{part_idx + 1}" if len(self._parts) > 1 else ""
+        result_path = (
+            Path(self._config.BASE_OUTPUT_DIR)
+            / f"{Path(self._output_data_filename).stem}{part_suffix}.json"
+        )
+        if not output_data:
+            logger.info("No output data to save. Skipping this part.")
+            return
+        else:
+            with open(result_path, "w", encoding="utf-8") as f:
+                json.dump(output_data, f, ensure_ascii=False, indent=4)
+        if log:
+            log_path = (
+                Path(self._config.BASE_OUTPUT_DIR)
+                / f"{Path(self._output_data_filename).stem}{part_suffix}_log.json"
+            )
+            with open(log_path, "w", encoding="utf-8") as f:
+                json.dump(log, f, ensure_ascii=False, indent=4)
+
+    def _result_exists(self, part_idx: int) -> bool:
+        part_suffix = f"_part_{part_idx + 1}" if len(self._parts) > 1 else ""
+        result_path = (
+            Path(self._config.BASE_OUTPUT_DIR)
+            / f"{Path(self._output_data_filename).stem}{part_suffix}.json"
+        )
+        return result_path.exists()
+
+    def run(self):
+        """
+        Execute the batch job processing pipeline.
+
+        Submits jobs, monitors progress, handles retries, and saves results.
+        """
+        try:
+            # Submit all jobs up-front for concurrent execution
+            self._submit_all_jobs()
+            pending_parts: set[int] = set(self._part_idx_to_job_name.keys())
+            logger.info(f"Pending parts: {sorted(pending_parts)}")
+            # Polling loop
+            while pending_parts:
+                finished_this_round: list[int] = []
+                for part_idx in list(pending_parts):
+                    job_name = self._part_idx_to_job_name[part_idx]
+                    status = self._manager.check_status(job_name=job_name)
+                    logger.info(f"Status for {job_name}: {status}")
+                    if status == "completed":
+                        logger.info(
+                            f"Job completed. Fetching results for part {part_idx + 1}..."
+                        )
+                        output_data, log = self._manager.fetch_results(
+                            job_name=job_name, remove_cache=False
+                        )
+                        output_data = self._config.import_function(output_data)
+                        self._save_results(output_data, log, part_idx)
+                        logger.info(
+                            f"Fetched and saved results for part {part_idx + 1}."
+                        )
+                        finished_this_round.append(part_idx)
+                    elif status == "failed":
+                        attempt = self._part_attempts.get(part_idx, 0) + 1
+                        self._part_attempts[part_idx] = attempt
+                        if attempt <= self._config.max_retries:
+                            logger.info(
+                                f"Job {job_name} failed (attempt {attempt}). Retrying after short backoff..."
+                            )
+                            self._manager._clear_state(job_name)
+                            time.sleep(10)
+                            payload = self._to_manager_payload(self._parts[part_idx])
+                            new_job_name = (
+                                f"{self._job_name}_part_{part_idx + 1}_retry_{attempt}"
+                            )
+                            self._manager.start(payload, job_name=new_job_name)
+                            self._part_idx_to_job_name[part_idx] = new_job_name
+                        else:
+                            logger.info(
+                                f"Job {job_name} failed after {attempt - 1} retries. Marking as failed."
+                            )
+                            finished_this_round.append(part_idx)
+                    else:
+                        # Still running or queued
+                        continue
+                # Remove finished parts
+                for part_idx in finished_this_round:
+                    pending_parts.discard(part_idx)
+                if pending_parts:
+                    logger.info(
+                        f"Waiting {self._config.poll_interval_seconds}s before next status check for parts: {sorted(pending_parts)}"
+                    )
+                    time.sleep(self._config.poll_interval_seconds)
+
+        except Exception as e:
+            raise TextToolsError(f"Batch job execution failed: {e}")

texttools/core/engine.py

@@ -0,0 +1,254 @@
+import math
+import random
+import re
+from functools import lru_cache
+from pathlib import Path
+
+import yaml
+
+from .exceptions import PromptError
+
+
+class PromptLoader:
+    """
+    Utility for loading and formatting YAML prompt templates.
+    """
+
+    MAIN_TEMPLATE = "main_template"
+    ANALYZE_TEMPLATE = "analyze_template"
+
+    @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 self.ANALYZE_TEMPLATE not in data:
+                raise PromptError(f"Missing 'analyze_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}")
+
+            main_template = (
+                data[self.MAIN_TEMPLATE][mode]
+                if mode and isinstance(data[self.MAIN_TEMPLATE], dict)
+                else data[self.MAIN_TEMPLATE]
+            )
+
+            analyze_template = (
+                data[self.ANALYZE_TEMPLATE][mode]
+                if mode and isinstance(data[self.ANALYZE_TEMPLATE], dict)
+                else data[self.ANALYZE_TEMPLATE]
+            )
+
+            if not main_template or not main_template.strip():
+                raise PromptError(
+                    f"Empty main_template in {prompt_file}"
+                    + (f" for mode '{mode}'" if mode else "")
+                )
+
+            return {
+                self.MAIN_TEMPLATE: main_template,
+                self.ANALYZE_TEMPLATE: 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 = {"text": text}
+            format_args.update(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}")
+
+
+class OperatorUtils:
+    @staticmethod
+    def build_main_prompt(
+        main_template: str,
+        analysis: str | None,
+        output_lang: str | None,
+        user_prompt: str | None,
+    ) -> str:
+        main_prompt = ""
+
+        if analysis:
+            main_prompt += f"Based on this analysis:\n{analysis}\n"
+
+        if output_lang:
+            main_prompt += f"Respond only in the {output_lang} language.\n"
+
+        if user_prompt:
+            main_prompt += f"Consider this instruction {user_prompt}\n"
+
+        main_prompt += main_template
+
+        return main_prompt
+
+    @staticmethod
+    def build_message(prompt: str) -> list[dict[str, str]]:
+        return [{"role": "user", "content": prompt}]
+
+    @staticmethod
+    def extract_logprobs(completion: dict) -> list[dict]:
+        """
+        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):
+                raise ValueError("Your model does not support logprobs")
+
+            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
+
+    @staticmethod
+    def get_retry_temp(base_temp: float) -> float:
+        delta_temp = random.choice([-1, 1]) * random.uniform(0.1, 0.9)
+        new_temp = base_temp + delta_temp
+
+        return max(0.0, min(new_temp, 1.5))
+
+
+def text_to_chunks(text: str, size: int, overlap: int) -> list[str]:
+    separators = ["\n\n", "\n", " ", ""]
+    is_separator_regex = False
+    keep_separator = True  # Equivalent to 'start'
+    length_function = len
+    strip_whitespace = True
+    chunk_size = size
+    chunk_overlap = overlap
+
+    def _split_text_with_regex(
+        text: str, separator: str, keep_separator: bool
+    ) -> list[str]:
+        if not separator:
+            return [text]
+        if not keep_separator:
+            return re.split(separator, text)
+        _splits = re.split(f"({separator})", text)
+        splits = [_splits[i] + _splits[i + 1] for i in range(1, len(_splits), 2)]
+        if len(_splits) % 2 == 0:
+            splits += [_splits[-1]]
+        return [_splits[0]] + splits if _splits[0] else splits
+
+    def _join_docs(docs: list[str], separator: str) -> str | None:
+        text = separator.join(docs)
+        if strip_whitespace:
+            text = text.strip()
+        return text if text else None
+
+    def _merge_splits(splits: list[str], separator: str) -> list[str]:
+        separator_len = length_function(separator)
+        docs = []
+        current_doc = []
+        total = 0
+        for d in splits:
+            len_ = length_function(d)
+            if total + len_ + (separator_len if current_doc else 0) > chunk_size:
+                if total > chunk_size:
+                    pass
+                if current_doc:
+                    doc = _join_docs(current_doc, separator)
+                    if doc is not None:
+                        docs.append(doc)
+                    while total > chunk_overlap or (
+                        total + len_ + (separator_len if current_doc else 0)
+                        > chunk_size
+                        and total > 0
+                    ):
+                        total -= length_function(current_doc[0]) + (
+                            separator_len if len(current_doc) > 1 else 0
+                        )
+                        current_doc = current_doc[1:]
+            current_doc.append(d)
+            total += len_ + (separator_len if len(current_doc) > 1 else 0)
+        doc = _join_docs(current_doc, separator)
+        if doc is not None:
+            docs.append(doc)
+        return docs
+
+    def _split_text(text: str, separators: list[str]) -> list[str]:
+        final_chunks = []
+        separator = separators[-1]
+        new_separators = []
+        for i, _s in enumerate(separators):
+            separator_ = _s if is_separator_regex else re.escape(_s)
+            if not _s:
+                separator = _s
+                break
+            if re.search(separator_, text):
+                separator = _s
+                new_separators = separators[i + 1 :]
+                break
+        separator_ = separator if is_separator_regex else re.escape(separator)
+        splits = _split_text_with_regex(text, separator_, keep_separator)
+        _separator = "" if keep_separator else separator
+        good_splits = []
+        for s in splits:
+            if length_function(s) < chunk_size:
+                good_splits.append(s)
+            else:
+                if good_splits:
+                    merged_text = _merge_splits(good_splits, _separator)
+                    final_chunks.extend(merged_text)
+                    good_splits = []
+                if not new_separators:
+                    final_chunks.append(s)
+                else:
+                    other_info = _split_text(s, new_separators)
+                    final_chunks.extend(other_info)
+        if good_splits:
+            merged_text = _merge_splits(good_splits, _separator)
+            final_chunks.extend(merged_text)
+        return final_chunks
+
+    return _split_text(text, separators)

texttools/core/exceptions.py

@@ -0,0 +1,22 @@
+class TextToolsError(Exception):
+    """Base exception for all TextTools errors."""
+
+    pass
+
+
+class PromptError(TextToolsError):
+    """Errors related to prompt loading and formatting."""
+
+    pass
+
+
+class LLMError(TextToolsError):
+    """Errors from LLM API calls."""
+
+    pass
+
+
+class ValidationError(TextToolsError):
+    """Errors from output validation."""
+
+    pass

texttools/core/internal_models.py

@@ -0,0 +1,58 @@
+from typing import Any, Literal, Type
+
+from pydantic import BaseModel, Field, create_model
+
+
+class OperatorOutput(BaseModel):
+    result: Any
+    analysis: str | None
+    logprobs: list[dict[str, Any]] | None
+
+
+class Str(BaseModel):
+    result: str = Field(..., description="The output string", example="text")
+
+
+class Bool(BaseModel):
+    result: bool = Field(
+        ..., description="Boolean indicating the output state", example=True
+    )
+
+
+class ListStr(BaseModel):
+    result: list[str] = Field(
+        ..., description="The output list of strings", example=["text_1", "text_2"]
+    )
+
+
+class ListDictStrStr(BaseModel):
+    result: list[dict[str, str]] = Field(
+        ...,
+        description="List of dictionaries containing string key-value pairs",
+        example=[{"text": "Mohammad", "type": "PER"}, {"text": "Iran", "type": "LOC"}],
+    )
+
+
+class ReasonListStr(BaseModel):
+    reason: str = Field(..., description="Thinking process that led to the output")
+    result: list[str] = Field(
+        ..., description="The output list of strings", example=["text_1", "text_2"]
+    )
+
+
+# This function is needed to create CategorizerOutput with dynamic categories
+def create_dynamic_model(allowed_values: list[str]) -> Type[BaseModel]:
+    literal_type = Literal[*allowed_values]
+
+    CategorizerOutput = create_model(
+        "CategorizerOutput",
+        reason=(
+            str,
+            Field(
+                ..., description="Explanation of why the input belongs to the category"
+            ),
+        ),
+        result=(literal_type, Field(..., description="Predicted category label")),
+    )
+
+    return CategorizerOutput