ehs-llm-client 0.1.0__py3-none-any.whl

ehs_llm_client/__init__.py

@@ -0,0 +1,3 @@
+from .client import LLM
+
+__all__ = ["LLM"]

ehs_llm_client/client.py

@@ -0,0 +1,412 @@
+import os
+import io
+import json
+import asyncio
+from typing import Any
+
+import openai
+from google import genai
+from dotenv import load_dotenv
+
+from .config import load_llm_config
+from .utils import to_timestamp
+from .exceptions import LLMProviderError
+
+load_dotenv()
+
+LLM_TIMEOUT = os.getenv('LLM_TIMEOUT', 180)
+LLM_TIMEOUT = int(LLM_TIMEOUT)
+RETRIES = os.getenv('RETRIES', 3)
+RETRIES = int(RETRIES)
+
+class LLM:
+    def __init__(
+        self,
+        config_section_name: str,
+        config_file_path: str | None = None,
+        config: dict | None = None,
+    ):
+        self.config_section_name = config_section_name
+
+        cfg = load_llm_config(
+            config_section_name,
+            config_file_path=config_file_path,
+            config_dict=config,
+        )
+
+        self.provider = cfg.get("provider")
+        self.model = cfg.get("model")
+        self.model_batch = cfg.get("model_batch")
+        self.base_url = cfg.get("base_url")
+        self.api_version = cfg.get("api_version")
+
+        if not self.provider or not self.model:
+            raise ValueError("provider and model are required")
+
+        key_env = f"{self.provider}APIkey"
+        self.api_key = os.getenv(key_env)
+
+        if not self.api_key:
+            raise ValueError(f"Missing API key env var: {key_env}")
+
+        self.client = self._init_client()
+
+    def _init_client(self):
+        if self.provider == "openai":
+            return openai.AsyncOpenAI(api_key=self.api_key)
+
+        if self.provider == "azure_openai":
+            return openai.AsyncAzureOpenAI(
+                api_key=self.api_key,
+                api_version=self.api_version,
+                azure_endpoint=self.base_url,
+            )
+
+        if self.provider == "google":
+            return genai.Client(api_key=self.api_key)
+
+        raise LLMProviderError(f"Unsupported provider: {self.provider}")
+
+    def __repr__(self):
+        return (
+            f"LLM(provider={self.provider}, model={self.model}, "
+            f"api_key_set={bool(self.api_key)})"
+        )
+
+    async def call_async(self, messages: list[dict[str, str]], schema: dict | str = None) -> Any:
+        # 1. Handle Google Provider
+        if self.provider == 'google':
+            # If schema is the special string, we enable JSON mode without a schema
+            if schema == "return json":
+                gemini_schema = None
+            else:
+                gemini_schema = schema
+                # Clean up schema if it's a dictionary
+                if isinstance(gemini_schema, dict) and "additionalProperties" in gemini_schema:
+                    del gemini_schema["additionalProperties"]
+
+            system_instr = None
+            contents = []
+            for m in messages:
+                if m["role"] == "system":
+                    system_instr = m["content"]
+                else:
+                    role = "model" if m["role"] == "assistant" else "user"
+                    contents.append({"role": role, "parts": [{"text": m["content"]}]})
+
+            return await self.client.aio.models.generate_content(
+                model=self.model,
+                contents=contents,
+                config={
+                    "system_instruction": system_instr,
+                    "response_mime_type": "application/json",
+                    "response_schema": gemini_schema,
+                    "thinking_config": {"include_thoughts": False}
+                }
+            )
+
+        # 2. Handle OpenAI / Azure Provider
+        elif self.provider in ['openai', 'azure_openai']:
+            if schema == "return json":
+                # Enable basic JSON Mode (no strict schema adherence)
+                text = {"format": {"type": "json_object"}}
+                # OpenAI requires "JSON" to be mentioned in the prompt for json_object mode
+                if not any("json" in m["content"].lower() for m in messages):
+                    messages.append({"role": "system", "content": "Respond in valid JSON format."})
+            
+            elif schema:
+                # Enable Structured Outputs with a specific schema
+                text = {
+                    "format": {
+                        "type": "json_schema",
+                        "name": "schema",
+                        "schema": schema 
+                    }
+                }
+            else:
+                text = None
+
+            api_kwargs = {
+                "model": self.model,
+                "input": messages,
+                # "reasoning": {"effort": "medium"},
+                "text": text
+            }
+
+            return await self.client.responses.create(**api_kwargs)
+    
+    async def get_response_async(self, messages: list[dict[str, str]], schema: dict | str = None, max_retries: int = 3, timeout: int = LLM_TIMEOUT):
+        retries = 1
+        while retries <= max_retries:
+            try:
+                response = await asyncio.wait_for(self.call_async(messages, schema), timeout=timeout)
+
+                if self.provider == "google":
+                    # --- GOOGLE PARSING ---
+                    finish_reason = response.candidates[0].finish_reason
+                    if finish_reason in ["SAFETY", "RECITATION"]:
+                        retries += 1
+                        continue
+                    
+                    output_text = "".join(
+                        part.text for part in response.candidates[0].content.parts if part.text
+                    )
+                    in_tokens = response.usage_metadata.prompt_token_count
+                    out_tokens = response.usage_metadata.candidates_token_count
+
+                elif self.provider == "openai" or self.provider == "azure_openai":
+                    # --- OPENAI PARSING ---
+                    if getattr(response, 'incomplete_details', None) == "content_filter":
+                        retries += 1
+                        continue
+                    
+                    output_text = response.output_text
+                    in_tokens = response.usage.input_tokens
+                    out_tokens = response.usage.output_tokens
+
+                # Both providers now return valid JSON strings due to the strict schema
+                if isinstance(output_text, dict):
+                    print('Returning dict directly')
+                    return output_text, in_tokens, out_tokens
+                
+                return output_text, in_tokens, out_tokens
+                # # Otherwise (e.g., for Gemini), try to parse the string
+                # try:
+                #     return json.loads(output_text), in_tokens, out_tokens
+                # except (json.JSONDecodeError, TypeError):
+                #     # If parsing fails or it's an unexpected type, return as-is
+                #     return output_text, in_tokens, out_tokens
+                
+            except Exception as e:
+                error_message = str(e)
+                # Handle exponential backoff for rate limits (429)
+                if "429" in error_message or "ResourceExhausted" in error_message:
+                    retry_seconds = int(2 ** retries)
+                else:
+                    retry_seconds = 1
+
+                if retries == max_retries:
+                    raise RuntimeError(f"LLM call failed after {max_retries} attempts: {e}")
+
+                await asyncio.sleep(retry_seconds)
+                retries += 1
+
+    #New code, need testing
+    def create_batch_request(self, custom_id: str, messages: list[dict[str, str]], schema: dict):
+        """
+        Generates a single line for a Batch API JSONL file.
+        Matches the logic of call_async by supporting multi-turn messages and system instructions.
+        """
+        provider = self.provider.lower()
+        
+        if schema == "return json":
+                # Enable basic JSON Mode (no strict schema adherence)
+                text = {"format": {"type": "json_object"}}
+                # OpenAI requires "JSON" to be mentioned in the prompt for json_object mode
+                if not any("json" in m["content"].lower() for m in messages):
+                    messages.append({"role": "system", "content": "Respond in valid JSON format."})
+            
+        elif schema:
+            text = {
+                "format": {
+                    "type": "json_schema",
+                    "name": "schema",
+                    "schema": schema 
+                    }
+                }
+        else:
+            text = None
+
+        if provider == "openai" or provider == "azure_openai":
+            return {
+                "custom_id": custom_id,
+                "method": "POST",
+                "url": "/v1/responses",
+                "body": {
+                    "model": self.model_batch,
+                    "input": messages, 
+                    # "reasoning": {"effort": "medium"},
+                    "text": text
+                }
+            }
+
+        elif provider == "google":
+            if not isinstance(schema, str): 
+                # 1. Clean the schema for Gemini compatibility
+                gemini_schema = schema.copy() if schema else None
+                if gemini_schema and "additionalProperties" in gemini_schema:
+                    del gemini_schema["additionalProperties"]
+            else: 
+                gemini_schema = None  
+
+            # 2. Process messages into Gemini format (contents + system_instruction)
+            system_instr = None
+            contents = []
+            for m in messages:
+                if m["role"] == "system":
+                    # In Gemini Batch, system_instruction is typically a Content object
+                    system_instr = {"parts": [{"text": m["content"]}]}
+                else:
+                    role = "model" if m["role"] == "assistant" else "user"
+                    contents.append({
+                        "role": role, 
+                        "parts": [{"text": m["content"]}]
+                    })
+
+            return {
+                "key": custom_id, # Gemini uses 'key' instead of 'custom_id'
+                "request": {
+                    "model": self.model_batch,
+                    "contents": contents,
+                    "system_instruction": system_instr,
+                    "generation_config": {
+                        "response_mime_type": "application/json",
+                        "response_schema": gemini_schema,
+                        "thinking_config": {"include_thoughts": False}
+                    }
+                }
+            }
+        
+        else:
+            raise ValueError(f"Unsupported provider: {provider}")
+
+    async def run_batch_process(self,
+        batch_lines: list[str],
+        submission_id: str
+    ) -> str:
+        """
+        Unified function to upload and start a batch job.
+        Returns the batch_id (OpenAI) or batch_name (Google).
+        """
+        if not batch_lines:
+            return None
+
+        file_content_str = "\n".join(batch_lines)
+        provider = self.provider.lower()
+
+        if provider == "openai":
+            # 1. Upload File to OpenAI
+            batch_file = await self.client.files.create(
+                file=file_content_str.encode('utf-8'),
+                purpose="batch"
+            )
+            file_id = batch_file.id
+
+            # 2. Create OpenAI Batch Job
+            batch_job = await self.client.batches.create(
+                input_file_id=file_id,
+                endpoint="/v1/responses",
+                completion_window="24h"
+            )
+            return batch_job.id
+
+        elif provider == "google":
+            # 1. Upload File to Google
+            # Google SDK prefers a file path or a file-like object
+            file_io = io.BytesIO(file_content_str.encode('utf-8'))
+            
+            # Note: Google's upload is currently synchronous in the SDK
+            # We use a unique name to avoid collisions
+            uploaded_file = self.client.files.upload(
+                file=file_io,
+                config={'display_name': f"batch_{submission_id}", 'mime_type': 'application/jsonl'}
+            )
+
+            # 2. Create Google Batch Job
+            # Note: Gemini 3.0 Flash uses 'models/gemini-3-flash-preview'
+            batch_job = self.client.batches.create(
+                model=self.model,
+                src=uploaded_file.name,
+            )
+            
+            return batch_job.name
+
+        else:
+            raise ValueError(f"Unsupported provider: {provider}")
+    
+    
+ 
+    async def cancel_batch(self,batch_id):
+        try:
+            await self.client.batches.cancel(batch_id)
+            return True
+        except Exception as e:
+            print(f"Error cancelling batch {batch_id}: {e}")
+            return False
+        
+
+    async def download_batch_results(self, output_file_id: str) -> str:
+        """
+        Downloads the result file content as a string.
+        Works for OpenAI and Google Gemini.
+        """
+        provider = self.provider.lower()
+
+        if provider == "openai":
+            # OpenAI returns a response object; we must read and decode it
+            result_content = await self.client.files.content(output_file_id)
+            # result_content.read() returns bytes
+            return result_content.read().decode('utf-8')
+
+        elif provider == "google":
+            # Google's download method returns the full bytes directly.
+            # Note: The current Google Gen AI SDK download is synchronous.
+            # We wrap it in to_thread to keep the function non-blocking.
+            result_bytes = await asyncio.to_thread(
+                self.client.files.download, 
+                file=output_file_id
+            )
+            return result_bytes.decode('utf-8')
+
+        else:
+            raise ValueError(f"Unsupported provider: {provider}")
+        
+    async def get_batch_status(self, batch_id: str):
+        provider = self.provider.lower()
+
+        if provider == "openai":
+            job = await self.client.batches.retrieve(batch_id)
+            # OpenAI property: job.request_counts
+            total = job.request_counts.total if job.request_counts else 0
+            completed = job.request_counts.completed if job.request_counts else 0
+            failed = job.request_counts.failed if job.request_counts else 0
+            
+            return {
+                "status": job.status,
+                "is_completed": job.status == "completed",
+                "is_terminal": job.status in ["failed", "expired", "cancelled"],
+                "start_time": job.in_progress_at,
+                "output_file_id": job.output_file_id,
+                "stats": {
+                    "total": total,
+                    "completed": completed,
+                    "failed": failed
+                }
+            }
+
+        elif provider == "google":
+            job = await asyncio.to_thread(self.client.batches.get, name=batch_id)
+
+            # Gemini batches do not provide per-request stats
+            total = 0
+            completed = 0
+            failed = 0
+
+            # Job state
+            state = job.state.name if hasattr(job.state, 'name') else job.state
+
+            # Output file is in job.dest.file_name
+            output_file_id = job.dest.file_name if hasattr(job, 'dest') and job.dest else None
+
+            return {
+                "status": state,
+                "is_completed": state == "JOB_STATE_SUCCEEDED",
+                "is_terminal": state in ["JOB_STATE_FAILED", "JOB_STATE_CANCELLED", "JOB_STATE_EXPIRED"],
+                "start_time": to_timestamp(job.create_time) if hasattr(job, "create_time") else None,
+                "output_file_id": output_file_id,
+                "stats": {
+                    "total": total,
+                    "completed": completed,
+                    "failed": failed
+                }
+            }

ehs_llm_client/config.py

@@ -0,0 +1,59 @@
+import os
+import configparser
+from typing import Any
+from .exceptions import LLMConfigError
+
+
+def load_llm_config(
+    section_name: str,
+    config_file_path: str | None = None,
+    config_dict: dict | None = None,
+) -> dict[str, Any]:
+    """
+    Load LLM configuration from:
+      1. Explicit config dict
+      2. Explicit config file path
+      3. Environment variables fallback
+    """
+
+    parser = configparser.ConfigParser()
+
+    if config_dict:
+        parser.read_dict(config_dict)
+
+    elif config_file_path:
+        if not os.path.exists(config_file_path):
+            raise LLMConfigError(f"Config file not found: {config_file_path}")
+        parser.read(config_file_path)
+
+    else:
+        # ENV fallback (minimal but safe)
+        provider = os.getenv("LLM_PROVIDER")
+        model = os.getenv("LLM_MODEL")
+
+        if not provider or not model:
+            raise LLMConfigError(
+                "No config provided. Set config_file_path, config dict, "
+                "or env vars LLM_PROVIDER and LLM_MODEL."
+            )
+
+        return {
+            "provider": provider,
+            "model": model,
+            "model_batch": os.getenv("LLM_MODEL_BATCH"),
+            "base_url": os.getenv("LLM_BASE_URL"),
+            "api_version": os.getenv("LLM_API_VERSION"),
+        }
+
+    if not parser.has_section(section_name):
+        raise LLMConfigError(f"Config section '{section_name}' not found")
+
+    defaults = (
+        dict(parser.items("default_settings"))
+        if parser.has_section("default_settings")
+        else {}
+    )
+
+    section = dict(parser.items(section_name))
+
+    return {**defaults, **section}

ehs_llm_client/exceptions.py

@@ -0,0 +1,5 @@
+class LLMConfigError(Exception):
+    pass
+
+class LLMProviderError(Exception):
+    pass

ehs_llm_client/utils.py

@@ -0,0 +1,12 @@
+from datetime import datetime
+from dateutil.parser import isoparse
+
+
+def to_timestamp(value):
+    if value is None:
+        return None
+    if isinstance(value, datetime):
+        return value.timestamp()
+    if isinstance(value, str):
+        return isoparse(value).timestamp()
+    raise TypeError(f"Unsupported time type: {type(value)}")

ehs_llm_client-0.1.0.dist-info/METADATA

@@ -0,0 +1,8 @@
+Metadata-Version: 2.4
+Name: ehs-llm-client
+Version: 0.1.0
+Summary: Unified LLM client. Currently supports Openai, Azure Openai and Google Gemini
+Requires-Dist: openai
+Requires-Dist: google-genai
+Requires-Dist: python-dotenv
+Requires-Dist: python-dateutil

ehs_llm_client-0.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+ehs_llm_client/__init__.py,sha256=AzmovmTtvj21YseUb4RVYBt4pq2JT6RpfZb41dyLK6c,46
+ehs_llm_client/client.py,sha256=fbyMOPTw_NayG2GCbiIQbwZHyBrU76s9Lirh5uBpv3o,16201
+ehs_llm_client/config.py,sha256=1sYOvX5biBUoLYMa22CiJ_MElRlBfBgQsatfCTs_kN0,1698
+ehs_llm_client/exceptions.py,sha256=UEihJBOZufiMPjkfStEuiE7rhJTUJXib6puGN65FykA,92
+ehs_llm_client/utils.py,sha256=zVLF8obG1cKJ-9cyHulzoTZBT0LQtGkx4cdiG2UaKuk,352
+ehs_llm_client-0.1.0.dist-info/METADATA,sha256=4Zte70VYMJw3Fe3rz5UgM_x9P8Q4Xx-SgfJxFIUMWqU,263
+ehs_llm_client-0.1.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+ehs_llm_client-0.1.0.dist-info/top_level.txt,sha256=8xnm7U82x1dZYaanP_IxTIw7e4PKG-GEpmaR0jU6FSE,15
+ehs_llm_client-0.1.0.dist-info/RECORD,,

ehs_llm_client-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.10.2)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

ehs_llm_client-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+ehs_llm_client