hamtaa-texttools 0.1.43__py3-none-any.whl

texttools/base/base_task_performer.py

@@ -0,0 +1,53 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+
+class BaseTaskPerformer(ABC):
+    """
+    Base class for common functionalities of LLM-based task performers.
+    This includes features like text preprocessing and dispatching results
+    to registered handlers.
+    """
+
+    def __init__(self, handlers: Optional[list[Any]] = None):
+        """
+        Initializes the BaseTaskPerformer with optional result handlers.
+
+        :param handlers: An optional list of handlers to process the component's results.
+        """
+        self.handlers = handlers or []
+
+    def _preprocess(self, text: str) -> str:
+        """
+        Preprocesses input text by stripping leading/trailing whitespace.
+        This can be extended for more complex preprocessing if needed.
+
+        :param text: The raw input text.
+        :return: The preprocessed text.
+        """
+        return text.strip()
+
+    @abstractmethod
+    def perform(self, *args, **kwargs) -> Any:
+        """
+        Abstract method to be implemented by concrete task performers.
+        This method will execute the primary task of the class (e.g., scoring, sorting).
+        The signature of args and kwargs will vary based on the specific task.
+        """
+        pass
+
+    def _dispatch(self, result_data: dict[str, Any]) -> None:
+        """
+        Dispatches the component's results to any registered result handlers.
+        Each handler receives a dictionary of result data.
+
+        :param result_data: A dictionary containing the results specific to the component.
+        """
+        for handler in self.handlers:
+            try:
+                handler.handle(result_data)
+            except Exception as e:
+                logging.error(
+                    f"Handler {handler.__class__.__name__} failed: {e}", exc_info=True
+                )

texttools/base/base_translator.py

@@ -0,0 +1,38 @@
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+
+class BaseTranslator(ABC):
+    """
+    Base class for all translators that  output a translated string.
+    """
+
+    def __init__(
+        self,
+        handlers: Optional[list[Any]] = None,
+    ):
+        self.handlers = handlers or []
+
+    @abstractmethod
+    def translate(
+        self, text: str, target_language: str, source_language: Optional[str] = None
+    ) -> str:
+        """
+        Translate the input text from the source language to the target language.
+        Should return the translated string.
+        The source_language can be optional if the LLM can detect it automatically.
+        """
+        pass
+
+    def preprocess(self, text: str) -> str:
+        """
+        Optional text preprocessing step.
+        """
+        return text.strip()
+
+    def _dispatch(self, result: dict) -> None:
+        """
+        Dispatch the result to handlers.
+        """
+        for handler in self.handlers:
+            handler.handle(result)

texttools/batch_manager/__init__.py

@@ -0,0 +1,2 @@
+from .batch_manager import SimpleBatchManager
+from .batch_runner import BatchJobRunner

texttools/batch_manager/batch_manager.py

@@ -0,0 +1,241 @@
+import json
+import uuid
+from pathlib import Path
+from typing import Any, Optional, Type
+
+from pydantic import BaseModel
+from openai import OpenAI
+from openai.lib._pydantic import to_strict_json_schema
+# from openai.lib._parsing._completions import type_to_response_format_param
+
+
+class SimpleBatchManager:
+    def __init__(
+        self,
+        client: OpenAI,
+        model: str,
+        output_model: Type[BaseModel],
+        prompt_template: str,
+        handlers: Optional[list[Any]] = None,
+        state_dir: Path = Path(".batch_jobs"),
+        custom_json_schema_obj_str: Optional[dict] = 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")
+
+    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:
+            # try:
+            # parsed_custom_schema = json.loads(self.custom_json_schema_obj_str)
+            response_format_config = {
+                "type": "json_schema",
+                "json_schema": self.custom_json_schema_obj_str,
+            }
+        # except json.JSONDecodeError as e:
+        #     raise ValueError(
+        #         "Failed to parse custom_json_schema_obj_str. "
+        #         "Please ensure it's a valid JSON string."
+        #     ) from e
+        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])
+        print("HERE is the job", 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 = []
+        # errors = []
+        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
+        # return modified_result , errors
+
+    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")
+                )
+                print("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)
+
+        for handler in self.handlers:
+            handler.handle(results)
+        if remove_cache:
+            self._clear_state(job_name)
+        # results = {"results": results, "log": log}
+        # return results
+        return results, log

texttools/batch_manager/batch_runner.py

@@ -0,0 +1,207 @@
+import json
+import os
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable
+
+# from dotenv import load_dotenv
+from openai import OpenAI
+from pydantic import BaseModel
+
+from texttools.batch_manager import SimpleBatchManager
+
+
+class OutputModel(BaseModel):
+    desired_output: str
+
+
+def exporting_data(data):
+    """
+    Produces a structure of the following form from an initial data structure:
+    [
+        {"id": str, "content": str},...
+    ]
+    """
+    return data
+
+
+def importing_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 = importing_data
+    export_function: Callable = exporting_data
+
+
+class BatchJobRunner:
+    """
+    Handles running batch jobs using a batch manager and configuration.
+    """
+
+    def __init__(
+        self, config: BatchConfig = BatchConfig(), output_model: type = OutputModel
+    ):
+        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)
+
+    def _init_manager(self) -> SimpleBatchManager:
+        # load_dotenv()
+        api_key = os.getenv("OPENAI_API_KEY")
+        client = OpenAI(api_key=api_key)
+        return SimpleBatchManager(
+            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)
+
+        # Validation: 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},...]'
+            )
+        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."
+                )
+            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(
+            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)
+            ]
+        print(f"Data split into {len(self.parts)} part(s)")
+
+    def run(self):
+        for idx, part in enumerate(self.parts):
+            if self._result_exists(idx):
+                print(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
+            )
+            print(
+                f"\n--- Processing part {idx + 1}/{len(self.parts)}: {part_job_name} ---"
+            )
+            self._process_part(part, part_job_name, idx)
+
+    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)
+                    time.sleep(10)  # Wait before retrying
+                    break  # Break inner loop to restart the job
+                else:
+                    time.sleep(5)  # Wait before checking again
+
+    def _save_results(
+        self, output_data: list[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:
+            print("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_path).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()

texttools/formatter/__init__.py

@@ -0,0 +1 @@
+from .gemma3_formatter import Gemma3Formatter

texttools/formatter/base.py

@@ -0,0 +1,26 @@
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+
+class ChatFormatter(ABC):
+    """
+    Given (raw_text, reason, maybe other hints), produce whatever payload
+    A) single string prompt (for providers that don t support multiple messages), or
+    B) list of {role, content} dicts, or
+    C) whatever shape the provider needs.
+    """
+
+    @abstractmethod
+    def format(
+        self,
+        text: str,
+        reason: Optional[str],
+        schema_instr: str,
+        prompt_template: Optional[str],
+    ) -> Any:
+        """
+        - For an OpenAI style API, this might return list[{"role": "user"/"assistant", "content": "…"}].
+        - For a one shot “text only” API, this might return a single string combining everything.
+        - For some niche service, it might return JSON: {"inputs": […], "parameters": {…}}.
+        """
+        pass

texttools/formatter/gemma3_formatter.py

@@ -0,0 +1,51 @@
+from typing import Literal
+
+from texttools.formatter.base import ChatFormatter
+
+
+class Gemma3Formatter(ChatFormatter):
+    """
+    Formatter that merges consecutive user messages (strings) with '\n'
+    and leaves assistant messages alone. No image‐handling, no extra tokens.
+    """
+
+    ROLE = "role"
+    USER_ROLE = "user"
+    ASSISTANT_ROLE = "assistant"
+    CONTENT = "content"
+    VALID_ROLES = {USER_ROLE, ASSISTANT_ROLE}
+
+    def format(
+        self, messages: list[dict[Literal["role", "content"], str]]
+    ) -> list[dict[str, str]]:
+        """
+        :param messages: list of {"role": ..., "content": ...}, where role is "user", "assistant", or "system"
+        :return: a new list where consecutive "user" messages are merged into single entries
+        """
+
+        merged: list[dict[str, str]] = []
+
+        for msg in messages:
+            role, content = msg[self.ROLE], msg[self.CONTENT].strip()
+
+            # Replace "system" role with "user" role
+            if role == "system":
+                role = self.USER_ROLE
+
+            # Raise value error if msg["role"] wan't a valid role
+            if role not in self.VALID_ROLES:
+                raise ValueError(f"Unexpected role: {role}")
+
+            # Merge with previous user turn
+            if (
+                merged
+                and role == self.USER_ROLE
+                and merged[-1][self.ROLE] == self.USER_ROLE
+            ):
+                merged[-1][self.CONTENT] += "\n" + content
+
+            # Otherwise, start a new turn
+            else:
+                merged.append({self.ROLE: role, self.CONTENT: content})
+
+        return merged

texttools/handlers/__init__.py

@@ -0,0 +1,6 @@
+from .handlers import (
+    NoOpResultHandler,
+    PrintResultHandler,
+    ResultHandler,
+    SaveToFileResultHandler,
+)

texttools/handlers/categorizer/__init__.py

@@ -0,0 +1,6 @@
+from .categorizer import (
+    ResultHandler,
+    NoOpResultHandler,
+    PrintResultHandler,
+    SaveToElasticResultHandler,
+)