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

texttools/batch/batch_runner.py

@@ -1,212 +1,217 @@
 import json
 import os
 import time
-from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Callable
+from typing import Any, Type, TypeVar
+import logging
 
+from dotenv import load_dotenv
 from openai import OpenAI
 from pydantic import BaseModel
 
-from texttools.batch.batch_manager import SimpleBatchManager
+from texttools.batch.internals.batch_manager import BatchManager
+from texttools.batch.batch_config import BatchConfig
+from texttools.tools.internals.models import StrOutput
 
+# Base Model type for output models
+T = TypeVar("T", bound=BaseModel)
 
-class Output(BaseModel):
-    output: str
-
-
-def export_data(data):
-    """
-    Produces a structure of the following form from an initial data structure:
-    [
-        {"id": str, "content": str},...
-    ]
-    """
-    return data
-
-
-def import_data(data):
-    """
-    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 = 2000000
-    CHARS_PER_TOKEN: float = 2.7
-    PROMPT_TOKEN_MULTIPLIER: int = 1000
-    BASE_OUTPUT_DIR: str = "Data/batch_entity_result"
-    import_function: Callable = import_data
-    export_function: Callable = export_data
+logger = logging.getLogger("texttools.batch_runner")
 
 
 class BatchJobRunner:
     """
-    Orchestrates the execution of batched LLM processing jobs.
-
-    Handles data loading, partitioning, job execution via SimpleBatchManager,
-    and result saving. Manages the complete workflow from input data to processed outputs,
-    including retries and progress tracking across multiple batch parts.
+    Handles running batch jobs using a batch manager and configuration.
     """
 
     def __init__(
-        self, config: BatchConfig = BatchConfig(), output_model: type = Output
+        self, config: BatchConfig = BatchConfig(), output_model: Type[T] = StrOutput
     ):
-        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]]] = []
+        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)
+        Path(self._config.BASE_OUTPUT_DIR).mkdir(parents=True, exist_ok=True)
 
-    def _init_manager(self) -> SimpleBatchManager:
+    def _init_manager(self) -> BatchManager:
+        load_dotenv()
         api_key = os.getenv("OPENAI_API_KEY")
         client = OpenAI(api_key=api_key)
-        return SimpleBatchManager(
+        return BatchManager(
             client=client,
-            model=self.model,
-            prompt_template=self.system_prompt,
-            output_model=self.output_model,
+            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:
+        with open(self._input_data_path, "r", encoding="utf-8") as f:
             data = json.load(f)
-        data = self.config.export_function(data)
+        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 in this form:  [ {"id": str, "content": str},...]'
+                "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(
-                    "Each item must be a dict with 'id' and 'content' keys."
+                    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
-        print(
+        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]
+        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)
+            self._parts = [
+                self._data[i : i + self._config.MAX_BATCH_SIZE]
+                for i in range(0, len(self._data), self._config.MAX_BATCH_SIZE)
             ]
-        print(f"Data split into {len(self.parts)} part(s)")
+        logger.info(f"Data split into {len(self._parts)} part(s)")
 
-    def run(self):
-        for idx, part in enumerate(self.parts):
+    def _submit_all_jobs(self) -> None:
+        for idx, part in enumerate(self._parts):
             if self._result_exists(idx):
-                print(f"Skipping part {idx + 1}: result already exists.")
+                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
+                f"{self._job_name}_part_{idx + 1}"
+                if len(self._parts) > 1
+                else self._job_name
             )
-            print(
-                f"\n--- Processing part {idx + 1}/{len(self.parts)}: {part_job_name} ---"
-            )
-            self._process_part(part, part_job_name, idx)
+            # 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
 
-    def _process_part(
-        self, part: list[dict[str, Any]], part_job_name: str, part_idx: int
-    ):
-        while True:
-            print(f"Starting job for part: {part_job_name}")
-            self.manager.start(part, job_name=part_job_name)
-            print("Started batch job. Checking status...")
-            while True:
-                status = self.manager.check_status(job_name=part_job_name)
-                print(f"Status: {status}")
-                if status == "completed":
-                    print("Job completed. Fetching results...")
-                    output_data, log = self.manager.fetch_results(
-                        job_name=part_job_name, remove_cache=False
-                    )
-                    output_data = self.config.import_function(output_data)
-                    self._save_results(output_data, log, part_idx)
-                    print("Fetched and saved results for this part.")
-                    return
-                elif status == "failed":
-                    print("Job failed. Clearing state, waiting, and retrying...")
-                    self.manager._clear_state(part_job_name)
-                    # Wait before retrying
-                    time.sleep(10)
-                    # Break inner loop to restart the job
-                    break
-                else:
-                    # Wait before checking again
-                    time.sleep(5)
+            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]], log: list[Any], part_idx: int
+        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 ""
+        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"
+            Path(self._config.BASE_OUTPUT_DIR)
+            / f"{Path(self._output_data_filename).stem}{part_suffix}.json"
         )
         if not output_data:
-            print("No output data to save. Skipping this part.")
+            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"
+                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 ""
+        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_path).stem}{part_suffix}.json"
+            Path(self._config.BASE_OUTPUT_DIR)
+            / f"{Path(self._output_data_filename).stem}{part_suffix}.json"
         )
         return result_path.exists()
 
-
-if __name__ == "__main__":
-    print("=== Batch Job Runner ===")
-    config = BatchConfig(
-        system_prompt="",
-        job_name="job_name",
-        input_data_path="Data.json",
-        output_data_filename="output",
-    )
-    runner = BatchJobRunner(config)
-    runner.run()
+    def run(self):
+        """
+        Execute the batch job processing pipeline.
+
+        Submits jobs, monitors progress, handles retries, and saves results.
+        """
+        # 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)

texttools/batch/{batch_manager.py → internals/batch_manager.py}

@@ -1,14 +1,20 @@
 import json
 import uuid
 from pathlib import Path
-from typing import Any, Type
+from typing import Any, Type, TypeVar
+import logging
 
 from pydantic import BaseModel
 from openai import OpenAI
 from openai.lib._pydantic import to_strict_json_schema
 
+# Base Model type for output models
+T = TypeVar("T", bound=BaseModel)
 
-class SimpleBatchManager:
+logger = logging.getLogger("texttools.batch_manager")
+
+
+class BatchManager:
     """
     Manages batch processing jobs for OpenAI's chat completions with structured outputs.
 
@@ -21,30 +27,29 @@ class SimpleBatchManager:
         self,
         client: OpenAI,
         model: str,
-        output_model: Type[BaseModel],
+        output_model: Type[T],
         prompt_template: str,
-        handlers: list[Any] | None = None,
         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.handlers = handlers or []
-        self.state_dir = state_dir
-        self.state_dir.mkdir(parents=True, exist_ok=True)
-        self.custom_json_schema_obj_str = custom_json_schema_obj_str
-        self.client_kwargs = client_kwargs
-        self.dict_input = False
-
-        if self.custom_json_schema_obj_str:
-            if self.custom_json_schema_obj_str is not dict:
-                raise ValueError("schema should be a dict")
+        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"
+        return self._state_dir / f"{job_name}.json"
 
     def _load_state(self, job_name: str) -> list[dict[str, Any]]:
         """
@@ -78,17 +83,17 @@ class SimpleBatchManager:
         """
         response_format_config: dict[str, Any]
 
-        if self.custom_json_schema_obj_str:
+        if self._custom_json_schema_obj_str:
             response_format_config = {
                 "type": "json_schema",
-                "json_schema": self.custom_json_schema_obj_str,
+                "json_schema": self._custom_json_schema_obj_str,
             }
         else:
-            raw_schema = to_strict_json_schema(self.output_model)
+            raw_schema = to_strict_json_schema(self._output_model)
             response_format_config = {
                 "type": "json_schema",
                 "json_schema": {
-                    "name": self.output_model.__name__,
+                    "name": self._output_model.__name__,
                     "schema": raw_schema,
                 },
             }
@@ -100,11 +105,11 @@ class SimpleBatchManager:
             "body": {
                 "model": self.model,
                 "messages": [
-                    {"role": "system", "content": self.prompt_template},
+                    {"role": "system", "content": self._prompt_template},
                     {"role": "user", "content": text},
                 ],
                 "response_format": response_format_config,
-                **self.client_kwargs,
+                **self._client_kwargs,
             },
         }
 
@@ -122,10 +127,10 @@ class SimpleBatchManager:
 
         else:
             raise TypeError(
-                "The input must be either a list of texts or a dictionary in the form {'id': str, 'text': str}."
+                "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"
+        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")
@@ -138,9 +143,10 @@ class SimpleBatchManager:
         """
         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(
+        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",
@@ -156,28 +162,12 @@ class SimpleBatchManager:
         if not job:
             return "completed"
 
-        info = self.client.batches.retrieve(job["id"])
+        info = self._client.batches.retrieve(job["id"])
         job = info.to_dict()
         self._save_state(job_name, [job])
-        print("HERE is the job", job)
+        logger.info("Batch job status: %s", job)
         return job["status"]
 
-    def _parsed(self, result: dict) -> list:
-        """
-        Parses the result dictionary, extracting the desired output or error for each item.
-        Returns a list of dictionaries with 'id' and 'output' keys.
-        """
-        modified_result = []
-
-        for key, d in result.items():
-            if "desired_output" in d:
-                new_dict = {"id": key, "output": d["desired_output"]}
-                modified_result.append(new_dict)
-            else:
-                new_dict = {"id": key, "output": d["error"]}
-                modified_result.append(new_dict)
-        return modified_result
-
     def fetch_results(
         self, job_name: str, remove_cache: bool = True
     ) -> tuple[dict[str, str], list]:
@@ -190,18 +180,18 @@ class SimpleBatchManager:
             return {}
         batch_id = job["id"]
 
-        info = self.client.batches.retrieve(batch_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")
+                    self._client.files.content(error_file_id).read().decode("utf-8")
                 )
-                print("Error file content:", err_content)
+                logger.error("Error file content:", err_content)
             return {}
 
-        content = self.client.files.content(out_file_id).read().decode("utf-8")
+        content = self._client.files.content(out_file_id).read().decode("utf-8")
         lines = content.splitlines()
         results = {}
         log = []
@@ -212,7 +202,7 @@ class SimpleBatchManager:
                 content = result["response"]["body"]["choices"][0]["message"]["content"]
                 try:
                     parsed_content = json.loads(content)
-                    model_instance = self.output_model(**parsed_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"}
@@ -232,8 +222,6 @@ class SimpleBatchManager:
                 error_d = {custom_id: results[custom_id]}
                 log.append(error_d)
 
-        for handler in self.handlers:
-            handler.handle(results)
         if remove_cache:
             self._clear_state(job_name)
 

texttools/batch/internals/utils.py

@@ -0,0 +1,16 @@
+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

texttools/prompts/README.md

@@ -3,6 +3,8 @@
 ## Overview
 This folder contains YAML files for all prompts used in the project. Each file represents a separate prompt template, which can be loaded by tools or scripts that require structured prompts for AI models.
 
+---
+
 ## Structure
 - **prompt_file.yaml**: Each YAML file represents a single prompt template.
 - **main_template**: The main instruction template for the model.
@@ -12,18 +14,20 @@ This folder contains YAML files for all prompts used in the project. Each file r
 ### Example YAML Structure
 ```yaml
 main_template:
-  default: |
+  mode_1: |
     Your main instructions here with placeholders like {input}.
-  reason: |
+  mode_2: |
     Optional reasoning instructions here.
 
 analyze_template:
-  default: |
+  mode_1: |
     Analyze and summarize the input.
-  reason: |
+  mode_2: |
     Optional detailed analysis template.
 ```
 
+---
+
 ## Guidelines
 1. **Naming**: Use descriptive names for each YAML file corresponding to the tool or task it serves.
 2. **Placeholders**: Use `{input}` or other relevant placeholders to dynamically inject data.