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

texttools/tools/sync_tools.py

@@ -1,4 +1,3 @@
-import sys
 from collections.abc import Callable
 from time import perf_counter
 from typing import Any, Literal
@@ -66,7 +65,7 @@ class TheTool:
             ToolOutput
 
         """
-        tool_name = sys._getframe().f_code.co_name
+        tool_name = "categorize"
         start = perf_counter()
 
         try:
@@ -198,7 +197,7 @@ class TheTool:
         Returns:
             ToolOutput
         """
-        tool_name = sys._getframe().f_code.co_name
+        tool_name = "extract_keywords"
         start = perf_counter()
 
         try:
@@ -211,7 +210,6 @@ class TheTool:
                 temperature=temperature,
                 logprobs=logprobs,
                 top_logprobs=top_logprobs,
-                mode=mode,
                 number_of_keywords=number_of_keywords,
                 validator=validator,
                 max_validation_retries=max_validation_retries,
@@ -219,6 +217,7 @@ class TheTool:
                 # Internal parameters
                 tool_name=tool_name,
                 output_model=ListStr,
+                mode=mode,
             )
 
             metadata = ToolOutputMetadata(
@@ -272,7 +271,7 @@ class TheTool:
         Returns:
             ToolOutput
         """
-        tool_name = sys._getframe().f_code.co_name
+        tool_name = "extract_entities"
         start = perf_counter()
 
         try:
@@ -343,7 +342,7 @@ class TheTool:
         Returns:
             ToolOutput
         """
-        tool_name = sys._getframe().f_code.co_name
+        tool_name = "is_question"
         start = perf_counter()
 
         try:
@@ -416,7 +415,7 @@ class TheTool:
         Returns:
             ToolOutput
         """
-        tool_name = sys._getframe().f_code.co_name
+        tool_name = "text_to_question"
         start = perf_counter()
 
         try:
@@ -489,7 +488,7 @@ class TheTool:
         Returns:
             ToolOutput
         """
-        tool_name = sys._getframe().f_code.co_name
+        tool_name = "merge_questions"
         start = perf_counter()
 
         try:
@@ -562,7 +561,7 @@ class TheTool:
         Returns:
             ToolOutput
         """
-        tool_name = sys._getframe().f_code.co_name
+        tool_name = "rewrite"
         start = perf_counter()
 
         try:
@@ -635,7 +634,7 @@ class TheTool:
         Returns:
             ToolOutput
         """
-        tool_name = sys._getframe().f_code.co_name
+        tool_name = "subject_to_question"
         start = perf_counter()
 
         try:
@@ -707,7 +706,7 @@ class TheTool:
         Returns:
             ToolOutput
         """
-        tool_name = sys._getframe().f_code.co_name
+        tool_name = "summarize"
         start = perf_counter()
 
         try:
@@ -782,7 +781,7 @@ class TheTool:
         Returns:
             ToolOutput
         """
-        tool_name = sys._getframe().f_code.co_name
+        tool_name = "translate"
         start = perf_counter()
 
         try:
@@ -900,7 +899,7 @@ class TheTool:
         Returns:
             ToolOutput
         """
-        tool_name = sys._getframe().f_code.co_name
+        tool_name = "propositionize"
         start = perf_counter()
 
         try:
@@ -961,7 +960,7 @@ class TheTool:
 
         Arguments:
             text: The input text
-            source_text: the source text that we want to check relation of text to it
+            source_text: The source text that we want to check relation of text to it
             with_analysis: Whether to include detailed reasoning analysis
             output_lang: Language for the output
             user_prompt: Additional instructions
@@ -975,13 +974,14 @@ class TheTool:
         Returns:
             ToolOutput
         """
-        tool_name = sys._getframe().f_code.co_name
+        tool_name = "check_fact"
         start = perf_counter()
 
         try:
             operator_output = self._operator.run(
                 # User parameters
                 text=text,
+                source_text=source_text,
                 with_analysis=with_analysis,
                 output_lang=output_lang,
                 user_prompt=user_prompt,
@@ -995,7 +995,6 @@ class TheTool:
                 tool_name=tool_name,
                 output_model=Bool,
                 mode=None,
-                source_text=source_text,
             )
 
             metadata = ToolOutputMetadata(
@@ -1049,7 +1048,7 @@ class TheTool:
         Returns:
             ToolOutput
         """
-        tool_name = sys._getframe().f_code.co_name
+        tool_name = "run_custom"
         start = perf_counter()
 
         try:

texttools/batch/config.py

@@ -1,40 +0,0 @@
-from collections.abc import Callable
-from dataclasses import dataclass
-from typing import Any
-
-
-def export_data(data) -> list[dict[str, str]]:
-    """
-    Produces a structure of the following form from an initial data structure:
-    [{"id": str, "text": str},...]
-    """
-    return data
-
-
-def import_data(data) -> Any:
-    """
-    Takes the output and adds and aggregates it to the original structure.
-    """
-    return data
-
-
-@dataclass
-class BatchConfig:
-    """
-    Configuration for batch job runner.
-    """
-
-    system_prompt: str = ""
-    job_name: str = ""
-    input_data_path: str = ""
-    output_data_filename: str = ""
-    model: str = "gpt-4.1-mini"
-    MAX_BATCH_SIZE: int = 100
-    MAX_TOTAL_TOKENS: int = 2_000_000
-    CHARS_PER_TOKEN: float = 2.7
-    PROMPT_TOKEN_MULTIPLIER: int = 1_000
-    BASE_OUTPUT_DIR: str = "Data/batch_entity_result"
-    import_function: Callable = import_data
-    export_function: Callable = export_data
-    poll_interval_seconds: int = 30
-    max_retries: int = 3

texttools/batch/manager.py

@@ -1,228 +0,0 @@
-import json
-import logging
-import uuid
-from pathlib import Path
-from typing import Any, Type, TypeVar
-
-from openai import OpenAI
-from openai.lib._pydantic import to_strict_json_schema
-from pydantic import BaseModel
-
-# Base Model type for output models
-T = TypeVar("T", bound=BaseModel)
-
-logger = logging.getLogger("texttools.batch_manager")
-
-
-class BatchManager:
-    """
-    Manages batch processing jobs for OpenAI's chat completions with structured outputs.
-
-    Handles the full lifecycle of a batch job: creating tasks from input texts,
-    starting the job, monitoring status, and fetching results. Results are automatically
-    parsed into the specified Pydantic output model. Job state is persisted to disk.
-    """
-
-    def __init__(
-        self,
-        client: OpenAI,
-        model: str,
-        output_model: Type[T],
-        prompt_template: str,
-        state_dir: Path = Path(".batch_jobs"),
-        custom_json_schema_obj_str: dict | None = None,
-        **client_kwargs: Any,
-    ):
-        self._client = client
-        self._model = model
-        self._output_model = output_model
-        self._prompt_template = prompt_template
-        self._state_dir = state_dir
-        self._custom_json_schema_obj_str = custom_json_schema_obj_str
-        self._client_kwargs = client_kwargs
-        self._dict_input = False
-        self._state_dir.mkdir(parents=True, exist_ok=True)
-
-        if custom_json_schema_obj_str and not isinstance(
-            custom_json_schema_obj_str, dict
-        ):
-            raise ValueError("Schema should be a dict")
-
-    def _state_file(self, job_name: str) -> Path:
-        return self._state_dir / f"{job_name}.json"
-
-    def _load_state(self, job_name: str) -> list[dict[str, Any]]:
-        """
-        Loads the state (job information) from the state file for the given job name.
-        Returns an empty list if the state file does not exist.
-        """
-        path = self._state_file(job_name)
-        if path.exists():
-            with open(path, "r", encoding="utf-8") as f:
-                return json.load(f)
-        return []
-
-    def _save_state(self, job_name: str, jobs: list[dict[str, Any]]) -> None:
-        """
-        Saves the job state to the state file for the given job name.
-        """
-        with open(self._state_file(job_name), "w", encoding="utf-8") as f:
-            json.dump(jobs, f)
-
-    def _clear_state(self, job_name: str) -> None:
-        """
-        Deletes the state file for the given job name if it exists.
-        """
-        path = self._state_file(job_name)
-        if path.exists():
-            path.unlink()
-
-    def _build_task(self, text: str, idx: str) -> dict[str, Any]:
-        """
-        Builds a single task dictionary for the batch job, including the prompt, model, and response format configuration.
-        """
-        response_format_config: dict[str, Any]
-
-        if self._custom_json_schema_obj_str:
-            response_format_config = {
-                "type": "json_schema",
-                "json_schema": self._custom_json_schema_obj_str,
-            }
-        else:
-            raw_schema = to_strict_json_schema(self._output_model)
-            response_format_config = {
-                "type": "json_schema",
-                "json_schema": {
-                    "name": self._output_model.__name__,
-                    "schema": raw_schema,
-                },
-            }
-
-        return {
-            "custom_id": str(idx),
-            "method": "POST",
-            "url": "/v1/chat/completions",
-            "body": {
-                "model": self.model,
-                "messages": [
-                    {"role": "system", "content": self._prompt_template},
-                    {"role": "user", "content": text},
-                ],
-                "response_format": response_format_config,
-                **self._client_kwargs,
-            },
-        }
-
-    def _prepare_file(self, payload: list[str] | list[dict[str, str]]) -> Path:
-        """
-        Prepares a JSONL file containing all tasks for the batch job, based on the input payload.
-        Returns the path to the created file.
-        """
-        if not payload:
-            raise ValueError("Payload must not be empty")
-        if isinstance(payload[0], str):
-            tasks = [self._build_task(text, uuid.uuid4().hex) for text in payload]
-        elif isinstance(payload[0], dict):
-            tasks = [self._build_task(dic["text"], dic["id"]) for dic in payload]
-
-        else:
-            raise TypeError(
-                "The input must be either a list of texts or a dictionary in the form {'id': str, 'text': str}"
-            )
-
-        file_path = self._state_dir / f"batch_{uuid.uuid4().hex}.jsonl"
-        with open(file_path, "w", encoding="utf-8") as f:
-            for task in tasks:
-                f.write(json.dumps(task) + "\n")
-        return file_path
-
-    def start(self, payload: list[str | dict[str, str]], job_name: str):
-        """
-        Starts a new batch job by uploading the prepared file and creating a batch job on the server.
-        If a job with the same name already exists, it does nothing.
-        """
-        if self._load_state(job_name):
-            return
-
-        path = self._prepare_file(payload)
-        upload = self._client.files.create(file=open(path, "rb"), purpose="batch")
-        job = self._client.batches.create(
-            input_file_id=upload.id,
-            endpoint="/v1/chat/completions",
-            completion_window="24h",
-        ).to_dict()
-        self._save_state(job_name, [job])
-
-    def check_status(self, job_name: str) -> str:
-        """
-        Checks and returns the current status of the batch job with the given job name.
-        Updates the job state with the latest information from the server.
-        """
-        job = self._load_state(job_name)[0]
-        if not job:
-            return "completed"
-
-        info = self._client.batches.retrieve(job["id"])
-        job = info.to_dict()
-        self._save_state(job_name, [job])
-        logger.info("Batch job status: %s", job)
-        return job["status"]
-
-    def fetch_results(
-        self, job_name: str, remove_cache: bool = True
-    ) -> tuple[dict[str, str], list]:
-        """
-        Fetches the results of a completed batch job. Optionally saves the results to a file and/or removes the job cache.
-        Returns a tuple containing the parsed results and a log of errors (if any).
-        """
-        job = self._load_state(job_name)[0]
-        if not job:
-            return {}
-        batch_id = job["id"]
-
-        info = self._client.batches.retrieve(batch_id)
-        out_file_id = info.output_file_id
-        if not out_file_id:
-            error_file_id = info.error_file_id
-            if error_file_id:
-                err_content = (
-                    self._client.files.content(error_file_id).read().decode("utf-8")
-                )
-                logger.error("Error file content:", err_content)
-            return {}
-
-        content = self._client.files.content(out_file_id).read().decode("utf-8")
-        lines = content.splitlines()
-        results = {}
-        log = []
-        for line in lines:
-            result = json.loads(line)
-            custom_id = result["custom_id"]
-            if result["response"]["status_code"] == 200:
-                content = result["response"]["body"]["choices"][0]["message"]["content"]
-                try:
-                    parsed_content = json.loads(content)
-                    model_instance = self._output_model(**parsed_content)
-                    results[custom_id] = model_instance.model_dump(mode="json")
-                except json.JSONDecodeError:
-                    results[custom_id] = {"error": "Failed to parse content as JSON"}
-                    error_d = {custom_id: results[custom_id]}
-                    log.append(error_d)
-                except Exception as e:
-                    results[custom_id] = {"error": str(e)}
-                    error_d = {custom_id: results[custom_id]}
-                    log.append(error_d)
-            else:
-                error_message = (
-                    result["response"]["body"]
-                    .get("error", {})
-                    .get("message", "Unknown error")
-                )
-                results[custom_id] = {"error": error_message}
-                error_d = {custom_id: results[custom_id]}
-                log.append(error_d)
-
-        if remove_cache:
-            self._clear_state(job_name)
-
-        return results, log

texttools/batch/runner.py

@@ -1,228 +0,0 @@
-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}")