"""Provider-agnostic LLM client wrapping the OpenAI SDK (Section 4). All three providers (OpenAI, OpenRouter, custom) are OpenAI-Chat-Completions compatible, so a single implementation handles them via a custom ``base_url``. Every call is routed by *tier name*, retried with backoff on transient errors, and logged to the token ledger. """ from __future__ import annotations import json import re import time from dataclasses import dataclass from typing import Any from openai import APIConnectionError, APIStatusError, APITimeoutError, OpenAI, RateLimitError from ..config import Config from ..models import UsageRecord from .usage import UsageLedger _FENCE_RE = re.compile(r"^\s*```(?:json)?\s*|\s*```\s*$", re.IGNORECASE) _TRANSIENT = (APIConnectionError, APITimeoutError, RateLimitError) class LLMError(RuntimeError): """Raised on non-recoverable LLM failures (auth/config/exhausted retries).""" class JSONParseError(LLMError): """Raised when a response that must be JSON cannot be parsed.""" @dataclass class LLMResponse: text: str prompt_tokens: int completion_tokens: int total_tokens: int retries: int class LLMClient: def __init__( self, config: Config, usage: UsageLedger | None = None, *, world_id: str = "", session_id: str = "", ) -> None: self.config = config self.usage = usage self.world_id = world_id self.session_id = session_id self._clients: dict[str, OpenAI] = {} def bind(self, *, world_id: str = "", session_id: str = "", usage: UsageLedger | None = None) -> LLMClient: """Return a shallow copy with updated logging context.""" c = LLMClient( self.config, usage or self.usage, world_id=world_id or self.world_id, session_id=session_id or self.session_id, ) c._clients = self._clients # reuse pooled SDK clients return c def _client_for(self, provider: str) -> OpenAI: if provider not in self._clients: if provider == "local": # Self-contained, no-API mode: run a small model in-process via # llama.cpp behind an OpenAI-shaped adapter. from .local import LocalLlamaClient self._clients[provider] = LocalLlamaClient() # type: ignore[assignment] return self._clients[provider] pcfg = self.config.providers[provider] self._clients[provider] = OpenAI( base_url=pcfg.base_url, api_key=pcfg.api_key(), default_headers=pcfg.default_headers or None, timeout=self.config.engine.request_timeout, max_retries=0, # we manage retries ourselves for logging ) return self._clients[provider] # -- core call ---------------------------------------------------------- def complete( self, *, tier: str, task: str, system: str | None = None, user: str, messages: list[dict[str, str]] | None = None, json_mode: bool = False, max_tokens: int | None = None, ) -> LLMResponse: """Call chat completions for ``tier``, logging usage under ``task``.""" tcfg, pcfg = self.config.resolve_tier(tier) client = self._client_for(tcfg.provider) msgs: list[dict[str, str]] = [] if messages is not None: msgs = list(messages) else: if system: msgs.append({"role": "system", "content": system}) msgs.append({"role": "user", "content": user}) kwargs: dict[str, Any] = { "model": tcfg.model, "messages": msgs, "temperature": tcfg.temperature, } if tcfg.top_p is not None: kwargs["top_p"] = tcfg.top_p eff_max = max_tokens or tcfg.max_tokens if eff_max is not None: kwargs["max_tokens"] = eff_max if json_mode: kwargs["response_format"] = {"type": "json_object"} retries = 0 last_exc: Exception | None = None max_retries = self.config.engine.max_retries while retries <= max_retries: try: resp = client.chat.completions.create(**kwargs) text = resp.choices[0].message.content or "" usage = resp.usage pt = getattr(usage, "prompt_tokens", 0) or 0 ct = getattr(usage, "completion_tokens", 0) or 0 tt = getattr(usage, "total_tokens", 0) or (pt + ct) self._log(task, tier, tcfg.provider, tcfg.model, pt, ct, tt, ok=True, retries=retries) return LLMResponse(text, pt, ct, tt, retries) except _TRANSIENT as exc: # transient -> backoff + retry last_exc = exc if retries >= max_retries: break time.sleep(min(2 ** retries, 8) + 0.1) retries += 1 except APIStatusError as exc: # 4xx/5xx -> fail loudly self._log(task, tier, tcfg.provider, tcfg.model, 0, 0, 0, ok=False, retries=retries) raise LLMError( f"{task}: API error {exc.status_code} from {tcfg.provider}: " f"{getattr(exc, 'message', exc)}" ) from exc self._log(task, tier, tcfg.provider, tcfg.model, 0, 0, 0, ok=False, retries=retries) raise LLMError(f"{task}: exhausted retries against {tcfg.provider}: {last_exc}") # -- JSON helper -------------------------------------------------------- def complete_json( self, *, tier: str, task: str, system: str | None = None, user: str, max_tokens: int | None = None, ) -> tuple[Any, LLMResponse]: """Call and parse a JSON response, stripping ``` fences defensively.""" resp = self.complete( tier=tier, task=task, system=system, user=user, json_mode=True, max_tokens=max_tokens, ) try: return parse_json(resp.text), resp except JSONParseError: # Some endpoints ignore json_mode; retry once without it then parse. resp = self.complete( tier=tier, task=task, system=system, user=user, json_mode=False, max_tokens=max_tokens, ) return parse_json(resp.text), resp def _log(self, task: str, tier: str, provider: str, model: str, pt: int, ct: int, tt: int, *, ok: bool, retries: int) -> None: if self.usage is None: return self.usage.record(UsageRecord( world_id=self.world_id, session_id=self.session_id, task=task, tier=tier, provider=provider, model=model, prompt_tokens=pt, completion_tokens=ct, total_tokens=tt, ok=ok, retries=retries, )) def strip_fences(text: str) -> str: s = text.strip() if s.startswith("```"): # remove leading and trailing fence lines lines = s.splitlines() if lines and lines[0].lstrip().startswith("```"): lines = lines[1:] if lines and lines[-1].strip().startswith("```"): lines = lines[:-1] s = "\n".join(lines) return s.strip() def parse_json(text: str) -> Any: """Parse JSON from a model response, tolerating fences and surrounding prose.""" candidate = strip_fences(text) try: return json.loads(candidate) except json.JSONDecodeError: pass # Fall back: grab the first balanced {...} or [...] span. for opener, closer in (("{", "}"), ("[", "]")): start = candidate.find(opener) end = candidate.rfind(closer) if start != -1 and end > start: try: return json.loads(candidate[start : end + 1]) except json.JSONDecodeError: continue raise JSONParseError(f"could not parse JSON from response: {text[:200]!r}")