Spaces:
Sleeping
Sleeping
| 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 | |
| class ModelMessage: | |
| role: str | |
| content: str | |
| 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 | |