hamtaa-texttools 1.1.8__py3-none-any.whl → 1.1.10__py3-none-any.whl

{hamtaa_texttools-1.1.8.dist-info → hamtaa_texttools-1.1.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hamtaa-texttools
-Version: 1.1.8
+Version: 1.1.10
 Summary: A high-level NLP toolkit built on top of modern LLMs.
 Author-email: Tohidi <the.mohammad.tohidi@gmail.com>, Montazer <montazerh82@gmail.com>, Givechi <mohamad.m.givechi@gmail.com>, MoosaviNejad <erfanmoosavi84@gmail.com>
 License: MIT License
@@ -29,6 +29,7 @@ Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: openai==1.97.1
+Requires-Dist: pydantic>=2.0.0
 Requires-Dist: pyyaml>=6.0
 Dynamic: license-file
 
@@ -40,26 +41,26 @@ Dynamic: license-file
 
 It provides both **sync (`TheTool`)** and **async (`AsyncTheTool`)** APIs for maximum flexibility.
 
-It provides ready-to-use utilities for **translation, question detection, keyword extraction, categorization, NER extraction, and more** — designed to help you integrate AI-powered text processing into your applications with minimal effort.
+It provides ready-to-use utilities for **translation, question detection, keyword extraction, categorization, NER extraction, and more** - designed to help you integrate AI-powered text processing into your applications with minimal effort.
 
 ---
 
 ## ✨ Features
 
-TextTools provides a rich collection of high-level NLP utilities built on top of LLMs.  
+TextTools provides a rich collection of high-level NLP utilities,
 Each tool is designed to work with structured outputs (JSON / Pydantic).
 
 - **`categorize()`** - Classifies text into Islamic studies categories 
-- **`is_question()`** - Binary detection of whether input is a question
 - **`extract_keywords()`** - Extracts keywords from text
 - **`extract_entities()`** - Named Entity Recognition (NER) system
-- **`summarize()`** - Text summarization
+- **`is_question()`** - Binary detection of whether input is a question
 - **`text_to_question()`** - Generates questions from text
 - **`merge_questions()`** - Merges multiple questions with different modes
 - **`rewrite()`** - Rewrites text with different wording/meaning
 - **`subject_to_question()`** - Generates questions about a specific subject
+- **`summarize()`** - Text summarization
 - **`translate()`** - Text translation between languages
-- **`run_custom()`** - Allows users to define a custom tool with arbitrary BaseModel
+- **`run_custom()`** - Allows users to define a custom tool with an arbitrary BaseModel
 
 ---
 
@@ -67,18 +68,18 @@ Each tool is designed to work with structured outputs (JSON / Pydantic).
 
 TextTools provides several optional flags to customize LLM behavior:
 
-- **`with_analysis=True`** → Adds a reasoning step before generating the final output. Useful for debugging, improving prompts, or understanding model behavior.  
-Note: This doubles token usage per call because it triggers an additional LLM request.
+- **`with_analysis (bool)`** → Adds a reasoning step before generating the final output.
+**Note:** This doubles token usage per call because it triggers an additional LLM request.
 
-- **`logprobs=True`** → Returns token-level probabilities for the generated output. You can also specify `top_logprobs=<N>` to get the top N alternative tokens and their probabilities.  
+- **`logprobs (bool)`** → Returns token-level probabilities for the generated output. You can also specify `top_logprobs=<N>` to get the top N alternative tokens and their probabilities.  
 
-- **`output_lang="en"`** → Forces the model to respond in a specific language. The model will ignore other instructions about language and respond strictly in the requested language.
+- **`output_lang (str)`** → Forces the model to respond in a specific language. The model will ignore other instructions about language and respond strictly in the requested language.
 
-- **`user_prompt="..."`** → Allows you to inject a custom instruction or prompt into the model alongside the main template. This gives you fine-grained control over how the model interprets or modifies the input text.
+- **`user_prompt (str)`** → Allows you to inject a custom instruction or prompt into the model alongside the main template. This gives you fine-grained control over how the model interprets or modifies the input text.
 
-- **`temperature=0.0`** → Determines how creative the model should respond. Takes a float number from `0.0` to `1.0`.
+- **`temperature (float)`** → Determines how creative the model should respond. Takes a float number from `0.0` to `2.0`.
 
-- **`validator=validation_function`** → Forces TheTool to validate the output result based on your custom validator. Validator should return bool (True if there were no problem, False if the validation failed.) If validator failed, TheTool will retry to get another output by modifying `temperature`.
+- **`validator (Callable)`** → Forces TheTool to validate the output result based on your custom validator. Validator should return bool (True if there were no problem, False if the validation failed.) If validator failed, TheTool will retry to get another output by modifying `temperature`.
 
 All these parameters can be used individually or together to tailor the behavior of any tool in **TextTools**.
 
@@ -89,10 +90,10 @@ All these parameters can be used individually or together to tailor the behavior
 ## 🧩 ToolOutput
 
 Every tool of `TextTools` returns a `ToolOutput` object which is a BaseModel with attributes:
-- **`result`** → The output of LLM (`type=Any`)
-- **`analysis`** → The reasoning step before generating the final output (`type=str`)
-- **`logprobs`** → Token-level probabilities for the generated output (`type=list`)
-- **`errors`** → Any error that have occured during calling LLM (`type=str`)
+- **`result (Any)`** → The output of LLM
+- **`analysis (str)`** → The reasoning step before generating the final output
+- **`logprobs (list)`** → Token-level probabilities for the generated output 
+- **`errors (list[str])`** → Any error that have occured during calling LLM
 
 **None:** You can use `repr(ToolOutput)` to see details of an output.
 
@@ -108,7 +109,7 @@ pip install -U hamtaa-texttools
 
 ---
 
-## Sync vs Async
+## 🧨 Sync vs Async
 | Tool         | Style   | Use case                                    |
 |--------------|---------|---------------------------------------------|
 | `TheTool`    | Sync    | Simple scripts, sequential workflows        |
@@ -187,13 +188,34 @@ Use **TextTools** when you need to:
 
 ---
 
+## 🔍 Logging
+
+TextTools uses Python's standard `logging` module. The library's default logger level is `WARNING`, so if you want to modify it, follow instructions:
+
+
+```python
+import logging
+
+# Default: warnings and errors only
+logging.basicConfig(level=logging.WARNING)
+
+# Debug everything (verbose)
+logging.basicConfig(level=logging.DEBUG)
+
+# Complete silence
+logging.basicConfig(level=logging.CRITICAL)
+```
+
+---
+
 ## 📚 Batch Processing
 
 Process large datasets efficiently using OpenAI's batch API.
 
-## Quick Start
+## ⚡ Quick Start (Batch)
 
 ```python
+from pydantic import BaseModel
 from texttools import BatchJobRunner, BatchConfig
 
 # Configure your batch job
@@ -223,6 +245,6 @@ Feel free to **open issues, suggest new features, or submit pull requests**.
 
 ---
 
-## License
+## 🌿 License
 
 This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

hamtaa_texttools-1.1.10.dist-info/RECORD

@@ -0,0 +1,30 @@
+hamtaa_texttools-1.1.10.dist-info/licenses/LICENSE,sha256=Hb2YOBKy2MJQLnyLrX37B4ZVuac8eaIcE71SvVIMOLg,1082
+texttools/__init__.py,sha256=EZPPNPafVGvBaxjG9anP0piqH3gAC0DdjdAckQeAgNU,251
+texttools/batch/batch_config.py,sha256=FCDXy9TfH7xjd1PHvn_CtdwEQSq-YO5sktiaMZEId58,740
+texttools/batch/batch_runner.py,sha256=zzzVIXedmaq-8fqsFtGRR64F7CtYRLlhQeBu8uMwJQg,9385
+texttools/batch/internals/batch_manager.py,sha256=UoBe76vmFG72qrSaGKDZf4HzkykFBkkkbL9TLfV8TuQ,8730
+texttools/batch/internals/utils.py,sha256=F1_7YlVFKhjUROAFX4m0SaP8KiZVZyHRMIIB87VUGQc,373
+texttools/prompts/README.md,sha256=-5YO93CN93QLifqZpUeUnCOCBbDiOTV-cFQeJ7Gg0I4,1377
+texttools/prompts/categorizer.yaml,sha256=GMqIIzQFhgnlpkgU1qi3FAD3mD4A2jiWD5TilQ2XnnE,1204
+texttools/prompts/extract_entities.yaml,sha256=KiKjeDpHaeh3JVtZ6q1pa3k4DYucUIU9WnEcRTCA-SE,651
+texttools/prompts/extract_keywords.yaml,sha256=0O7ypL_OsEOxtvlQ2CZjnsv9637DJwAKprZsf9Vo2_s,769
+texttools/prompts/is_question.yaml,sha256=d0-vKRbXWkxvO64ikvxRjEmpAXGpCYIPGhgexvPPjws,471
+texttools/prompts/merge_questions.yaml,sha256=0J85GvTirZB4ELwH3sk8ub_WcqqpYf6PrMKr3djlZeo,1792
+texttools/prompts/rewrite.yaml,sha256=LO7He_IA3MZKz8a-LxH9DHJpOjpYwaYN1pbjp1Y0tFo,5392
+texttools/prompts/run_custom.yaml,sha256=38OkCoVITbuuS9c08UZSP1jZW4WjSmRIi8fR0RAiPu4,108
+texttools/prompts/subject_to_question.yaml,sha256=C7x7rNNm6U_ZG9HOn6zuzYOtvJUZ2skuWbL1-aYdd3E,1147
+texttools/prompts/summarize.yaml,sha256=o6rxGPfWtZd61Duvm8NVvCJqfq73b-wAuMSKR6UYUqY,459
+texttools/prompts/text_to_question.yaml,sha256=UheKYpDn6iyKI8NxunHZtFpNyfCLZZe5cvkuXpurUJY,783
+texttools/prompts/translate.yaml,sha256=mGT2uBCei6uucWqVbs4silk-UV060v3G0jnt0P6sr50,634
+texttools/tools/async_tools.py,sha256=yEj4dM2bdW_12hxvimhxPOGfGhl1PqFsHM3Z4toCTaM,14813
+texttools/tools/sync_tools.py,sha256=wHY0O8R9HipUz0P268zk1w3SlFxEIffm5EjX4tcWxNM,14579
+texttools/tools/internals/async_operator.py,sha256=5sVc5K5-Vuulsxly0IfrLmzd8W7ySI4cY09myyOGL0I,7022
+texttools/tools/internals/base_operator.py,sha256=l2Mg59MGIf396yPx1CAgcplKclOptQWeTR0UIz9VTdk,2255
+texttools/tools/internals/formatters.py,sha256=tACNLP6PeoqaRpNudVxBaHA25zyWqWYPZQuYysIu88g,941
+texttools/tools/internals/operator.py,sha256=W0DxTGB3cbtDfzLqwMCM8x5xiVWgN0vZWX8PzJwAQKE,6795
+texttools/tools/internals/output_models.py,sha256=ekpbyocmXj_dee7ieOT1zOkMo9cPHT7xcUFCZoUaXA0,1886
+texttools/tools/internals/prompt_loader.py,sha256=4g6-U8kqrGN7VpNaRcrBcnF-h03PXjUDBP0lL0_4EZY,1953
+hamtaa_texttools-1.1.10.dist-info/METADATA,sha256=6X7Qd8nOAIn7reZ5l-66CyK5fczMADtlIfdIbypJIUE,9101
+hamtaa_texttools-1.1.10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+hamtaa_texttools-1.1.10.dist-info/top_level.txt,sha256=5Mh0jIxxZ5rOXHGJ6Mp-JPKviywwN0MYuH0xk5bEWqE,10
+hamtaa_texttools-1.1.10.dist-info/RECORD,,

texttools/__init__.py

@@ -1,4 +1,6 @@
-from .batch import BatchJobRunner, BatchConfig
-from .tools import AsyncTheTool, TheTool
+from .batch.batch_runner import BatchJobRunner
+from .batch.batch_config import BatchConfig
+from .tools.sync_tools import TheTool
+from .tools.async_tools import AsyncTheTool
 
 __all__ = ["TheTool", "AsyncTheTool", "BatchJobRunner", "BatchConfig"]

texttools/batch/batch_config.py

@@ -0,0 +1,26 @@
+from dataclasses import dataclass
+from typing import Callable
+
+from texttools.batch.internals.utils import import_data, export_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/batch_runner.py

@@ -1,16 +1,16 @@
 import json
 import os
 import time
-from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Callable, Type, TypeVar
+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 BatchManager
+from texttools.batch.internals.batch_manager import BatchManager
+from texttools.batch.batch_config import BatchConfig
 from texttools.tools.internals.output_models import StrOutput
 
 # Base Model type for output models
@@ -19,43 +19,6 @@ T = TypeVar("T", bound=BaseModel)
 logger = logging.getLogger("texttools.batch_runner")
 
 
-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
-
-
 class BatchJobRunner:
     """
     Handles running batch jobs using a batch manager and configuration.
@@ -64,22 +27,22 @@ class BatchJobRunner:
     def __init__(
         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._partition_data()
-        Path(self.config.BASE_OUTPUT_DIR).mkdir(parents=True, exist_ok=True)
+        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] = {}
+        self._part_idx_to_job_name: dict[int, str] = {}
         # Track retry attempts per part
-        self.part_attempts: dict[int, int] = {}
+        self._part_attempts: dict[int, int] = {}
+        self._partition_data()
+        Path(self._config.BASE_OUTPUT_DIR).mkdir(parents=True, exist_ok=True)
 
     def _init_manager(self) -> BatchManager:
         load_dotenv()
@@ -87,15 +50,15 @@ class BatchJobRunner:
         client = OpenAI(api_key=api_key)
         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):
@@ -112,50 +75,50 @@ class BatchJobRunner:
         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
+        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)
             ]
-        logger.info(f"Data split into {len(self.parts)} part(s)")
+        logger.info(f"Data split into {len(self._parts)} part(s)")
 
     def _submit_all_jobs(self) -> None:
-        for idx, part in enumerate(self.parts):
+        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
+                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)
+            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)
+                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}"
+                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)
+            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)
@@ -166,10 +129,10 @@ class BatchJobRunner:
         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:
             logger.info("No output data to save. Skipping this part.")
@@ -179,17 +142,17 @@ class BatchJobRunner:
                 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_filename).stem}{part_suffix}.json"
+            Path(self._config.BASE_OUTPUT_DIR)
+            / f"{Path(self._output_data_filename).stem}{part_suffix}.json"
         )
         return result_path.exists()
 
@@ -201,41 +164,41 @@ class BatchJobRunner:
         """
         # Submit all jobs up-front for concurrent execution
         self._submit_all_jobs()
-        pending_parts: set[int] = set(self.part_idx_to_job_name.keys())
+        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)
+                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(
+                    output_data, log = self._manager.fetch_results(
                         job_name=job_name, remove_cache=False
                     )
-                    output_data = self.config.import_function(output_data)
+                    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:
+                    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)
+                        self._manager._clear_state(job_name)
                         time.sleep(10)
-                        payload = self._to_manager_payload(self.parts[part_idx])
+                        payload = self._to_manager_payload(self._parts[part_idx])
                         new_job_name = (
-                            f"{self.job_name}_part_{part_idx + 1}_retry_{attempt}"
+                            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
+                        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."
@@ -249,6 +212,6 @@ class BatchJobRunner:
                 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)}"
+                    f"Waiting {self._config.poll_interval_seconds}s before next status check for parts: {sorted(pending_parts)}"
                 )
-                time.sleep(self.config.poll_interval_seconds)
+                time.sleep(self._config.poll_interval_seconds)

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

@@ -33,15 +33,15 @@ class BatchManager:
         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.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
+        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
@@ -49,7 +49,7 @@ class BatchManager:
             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]]:
         """
@@ -83,17 +83,17 @@ class BatchManager:
         """
         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,
                 },
             }
@@ -105,11 +105,11 @@ class BatchManager:
             "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,
             },
         }
 
@@ -130,7 +130,7 @@ class BatchManager:
                 "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")
@@ -145,8 +145,8 @@ class BatchManager:
             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",
@@ -162,7 +162,7 @@ class BatchManager:
         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])
         logger.info("Batch job status: %s", job)
@@ -180,18 +180,18 @@ class BatchManager:
             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")
                 )
                 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 = []
@@ -202,7 +202,7 @@ class BatchManager:
                 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"}

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