from __future__ import annotations from dataclasses import dataclass from typing import Any, Protocol import logging import time from openai import APIError, AzureOpenAI, OpenAI from data_agent_baseline.langgraph_agent.llm_cache import get_cached, set_cached from data_agent_baseline.langgraph_agent.deterministic_seed import DETERMINISTIC_SEED, SEED_ENABLED _log = logging.getLogger(__name__) _RETRY_ERRORS = (500, 502, 503, 504, 429) _MAX_RETRIES = 2 _RETRY_BASE_DELAY = 1.0 # seconds (1s, 2s backoff) # Per-LLM-call HTTP timeout (seconds). Without this, a hung TCP socket can hold # a worker thread indefinitely and starve the task-level deadline / pool. # Override via env var LLM_HTTP_TIMEOUT for debugging. import os as _os LLM_HTTP_TIMEOUT = float(_os.environ.get("LLM_HTTP_TIMEOUT", "90")) class LLMApiError(Exception): """Non-retryable LLM API error (e.g., 403 Forbidden, 401 Unauthorized).""" pass @dataclass(frozen=True, slots=True) class ModelMessage: role: str content: str @dataclass(frozen=True, slots=True) class ModelStep: thought: str action: str action_input: dict[str, Any] raw_response: str class ModelAdapter(Protocol): def complete(self, messages: list[ModelMessage]) -> str: raise NotImplementedError #TODO: Add retry logic to handle transient API errors, and consider adding support for streaming responses if needed in the future. #TODO: Consider adding support for additional model providers in the future, such as Hugging Face or Cohere, by implementing additional adapter classes that conform to the ModelAdapter protocol. class OpenAIModelAdapter: def __init__( self, *, model: str, api_base: str, api_key: str, temperature: float, ) -> None: self.model = model self.api_base = api_base.rstrip("/") self.api_key = api_key self.temperature = temperature def complete(self, messages: list[ModelMessage]) -> str: if not self.api_key: raise RuntimeError("Missing model API key in config.agent.api_key.") client = OpenAI( api_key=self.api_key, base_url=self.api_base, timeout=LLM_HTTP_TIMEOUT, ) try: response = None last_exc = None for attempt in range(_MAX_RETRIES): try: _create_kwargs = dict( model=self.model, messages=[{"role": message.role, "content": message.content} for message in messages], temperature=self.temperature, ) if SEED_ENABLED: _create_kwargs["seed"] = DETERMINISTIC_SEED response = client.chat.completions.create(**_create_kwargs) break except APIError as exc: last_exc = exc status = getattr(exc, "status_code", None) if status and status not in _RETRY_ERRORS: raise LLMApiError( f"Non-retryable API error (HTTP {status}): {exc}" ) from exc if status in _RETRY_ERRORS and attempt < _MAX_RETRIES - 1: delay = _RETRY_BASE_DELAY * (2 ** attempt) _log.warning(f"API error {status}, retrying in {delay}s (attempt {attempt+1}/{_MAX_RETRIES})") time.sleep(delay) else: raise if response is None: raise RuntimeError(f"Model request failed after {_MAX_RETRIES} retries: {last_exc}") from last_exc except LLMApiError: raise except APIError as exc: raise RuntimeError(f"Model request failed: {exc}") from exc choices = response.choices or [] if not choices: raise RuntimeError("Model response missing choices.") content = choices[0].message.content if not isinstance(content, str): raise RuntimeError("Model response missing text content.") return content class AzureOpenAIModelAdapter: """Adapter for Azure OpenAI API with support for custom headers.""" def __init__( self, *, azure_endpoint: str, azure_api_key: str, azure_api_version: str, azure_deployment: str, temperature: float, azure_custom_headers: dict[str, str] | None = None, ) -> None: self.azure_endpoint = azure_endpoint self.azure_api_key = azure_api_key self.azure_api_version = azure_api_version self.azure_deployment = azure_deployment self.temperature = temperature self.azure_custom_headers = azure_custom_headers or {} def complete(self, messages: list[ModelMessage]) -> str: if not self.azure_api_key: raise RuntimeError("Missing Azure API key in config.agent.azure_api_key.") if not self.azure_endpoint: raise RuntimeError("Missing Azure endpoint in config.agent.azure_endpoint.") client = AzureOpenAI( api_key=self.azure_api_key, api_version=self.azure_api_version, azure_endpoint=self.azure_endpoint, default_headers=self.azure_custom_headers if self.azure_custom_headers else None, timeout=LLM_HTTP_TIMEOUT, ) try: response = None last_exc = None for attempt in range(_MAX_RETRIES): try: _create_kwargs = dict( model=self.azure_deployment, messages=[{"role": message.role, "content": message.content} for message in messages], temperature=self.temperature, ) if SEED_ENABLED: _create_kwargs["seed"] = DETERMINISTIC_SEED response = client.chat.completions.create(**_create_kwargs) break except APIError as exc: last_exc = exc status = getattr(exc, "status_code", None) if status and status not in _RETRY_ERRORS: raise LLMApiError( f"Non-retryable Azure API error (HTTP {status}): {exc}" ) from exc if status in _RETRY_ERRORS and attempt < _MAX_RETRIES - 1: delay = _RETRY_BASE_DELAY * (2 ** attempt) _log.warning(f"Azure API error {status}, retrying in {delay}s (attempt {attempt+1}/{_MAX_RETRIES})") time.sleep(delay) else: raise if response is None: raise RuntimeError(f"Azure model request failed after {_MAX_RETRIES} retries: {last_exc}") from last_exc except LLMApiError: raise except APIError as exc: raise RuntimeError(f"Azure model request failed: {exc}") from exc choices = response.choices or [] if not choices: raise RuntimeError("Model response missing choices.") content = choices[0].message.content if not isinstance(content, str): raise RuntimeError("Model response missing text content.") return content class ScriptedModelAdapter: def __init__(self, responses: list[str]) -> None: self._responses = list(responses) def complete(self, messages: list[ModelMessage]) -> str: del messages if not self._responses: raise RuntimeError("No scripted model responses remaining.") return self._responses.pop(0) class CachedModelAdapter: """Wraps any ModelAdapter with a disk-based response cache.""" def __init__(self, inner: "ModelAdapter", model_name: str = "") -> None: self._inner = inner self._model_name = model_name or getattr(inner, "model", "") or getattr(inner, "azure_deployment", "unknown") def complete(self, messages: list[ModelMessage]) -> str: cached = get_cached(self._model_name, messages) if cached is not None: return cached response = self._inner.complete(messages) set_cached(self._model_name, messages, response) return response