Spaces:
Sleeping
Sleeping
File size: 8,447 Bytes
d3d0e0e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 | 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
|