"""LLM client — async wrapper over the OpenAI-compat router on :14011. Design: - One LlmClient instance per long-running process (reusable httpx.AsyncClient). - Roles (REASONING / FAST / VISION) decide model + backend automatically. - Hard timeouts and retries are enforced; no LLM call hangs forever. - Helpers for the two patterns we use most: structured JSON and single-token NLI. """ from __future__ import annotations import json import re from collections.abc import AsyncIterator from contextlib import asynccontextmanager from typing import Any import httpx from tenacity import ( AsyncRetrying, retry_if_exception_type, stop_after_attempt, wait_exponential, ) from shared.config import LlmRole, settings from shared.logging import get_logger log = get_logger(__name__) class LlmError(Exception): """Wraps any failure from the LLM stack with enough context to debug.""" def __init__(self, message: str, *, status: int | None = None, body: str | None = None): super().__init__(message) self.status = status self.body = body class RoleNotAvailable(LlmError): """The requested role has no enabled model behind it.""" class LlmClient: """Async OpenAI-compat client targeting the unified router.""" def __init__( self, *, base_url: str | None = None, timeout: float = 120.0, max_attempts: int = 3, ): self._base_url = (base_url or settings.llm_router_url).rstrip("/") self._timeout = timeout self._max_attempts = max_attempts headers = {"Content-Type": "application/json"} if settings.llm_router_api_key: headers["Authorization"] = f"Bearer {settings.llm_router_api_key}" self._http = httpx.AsyncClient( base_url=self._base_url, headers=headers, timeout=httpx.Timeout(timeout, connect=10.0), ) async def aclose(self) -> None: await self._http.aclose() async def __aenter__(self) -> LlmClient: return self async def __aexit__(self, *args: Any) -> None: await self.aclose() # ------------------------------------------------------------------ core async def chat( self, *, role: LlmRole = LlmRole.REASONING, messages: list[dict[str, str]], temperature: float = 0.1, max_tokens: int = 1024, stop: list[str] | None = None, extra: dict[str, Any] | None = None, ) -> dict[str, Any]: """Raw chat completion. Returns the parsed response dict. The response shape is OpenAI-compatible: {"choices": [{"message": {"content": "..."}}], "usage": {...}, "backend": "..."} """ model_info = settings.model_for(role) if not model_info: raise RoleNotAvailable(f"Role {role.value} has no enabled model") model_id, backend = model_info payload: dict[str, Any] = { "model": model_id, "backend": backend, "messages": messages, "temperature": temperature, "max_tokens": max_tokens, } if stop: payload["stop"] = stop if extra: payload.update(extra) async for attempt in AsyncRetrying( stop=stop_after_attempt(self._max_attempts), wait=wait_exponential(multiplier=1, min=1, max=10), retry=retry_if_exception_type((httpx.HTTPError, httpx.TimeoutException)), reraise=True, ): with attempt: resp = await self._http.post("/v1/chat/completions", json=payload) if resp.status_code >= 400: raise LlmError( f"HTTP {resp.status_code} from LLM router", status=resp.status_code, body=resp.text[:1000], ) data = resp.json() if "choices" not in data: raise LlmError( "LLM response missing 'choices'", body=json.dumps(data)[:1000], ) return data raise LlmError("retry loop exited without result") # unreachable # ------------------------------------------------------------------ helpers async def chat_text( self, *, role: LlmRole = LlmRole.REASONING, system: str | None = None, user: str, temperature: float = 0.1, max_tokens: int = 1024, stop: list[str] | None = None, ) -> tuple[str, dict[str, Any]]: """Convenience: send (system?, user) → return (text, usage_dict).""" msgs: list[dict[str, str]] = [] if system: msgs.append({"role": "system", "content": system}) msgs.append({"role": "user", "content": user}) data = await self.chat( role=role, messages=msgs, temperature=temperature, max_tokens=max_tokens, stop=stop, ) text = data["choices"][0]["message"]["content"] usage = data.get("usage", {}) | {"backend": data.get("backend", "")} return text, usage async def chat_json( self, *, role: LlmRole = LlmRole.REASONING, system: str | None = None, user: str, max_tokens: int = 2048, temperature: float = 0.0, ) -> tuple[dict[str, Any] | list[Any], dict[str, Any]]: """Send a request expected to return JSON. Strips fences if present. Raises LlmError if the response is not parseable as JSON. """ text, usage = await self.chat_text( role=role, system=system, user=user, temperature=temperature, max_tokens=max_tokens, ) cleaned = _extract_json(text) try: return json.loads(cleaned), usage except json.JSONDecodeError as e: raise LlmError( f"LLM did not return valid JSON: {e}", body=text[:1000], ) from e async def chat_label( self, *, role: LlmRole = LlmRole.REASONING, system: str, user: str, allowed: list[str], max_tokens: int = 16, ) -> tuple[str, dict[str, Any]]: """Single-label classification. Returns the matched label uppercased. Useful for NLI (SUPPORT/CONTRADICT/NEUTRAL), credibility tiers, etc. Raises LlmError if no allowed label is found in the response. """ text, usage = await self.chat_text( role=role, system=system, user=user, temperature=0.0, max_tokens=max_tokens, ) upper = text.upper() for label in allowed: if label.upper() in upper: return label.upper(), usage raise LlmError( f"LLM response did not contain any allowed label {allowed}", body=text[:500], ) # ------------------------------------------------------------------ health async def list_models(self) -> list[dict[str, Any]]: resp = await self._http.get("/v1/models") resp.raise_for_status() data = resp.json() return data.get("data", []) async def list_backends(self) -> list[str]: try: resp = await self._http.get("/v1/backends") resp.raise_for_status() return resp.json().get("backends", []) except httpx.HTTPError: return [] # --------------------------------------------------------------------- utilities _FENCE_RE = re.compile(r"```(?:json)?\s*(.*?)\s*```", re.DOTALL) def _extract_json(text: str) -> str: """Strip markdown code fences and isolate the JSON object/array if needed.""" text = text.strip() m = _FENCE_RE.search(text) if m: return m.group(1).strip() # Find first { or [ and last matching close starts = [text.find("{"), text.find("[")] starts = [s for s in starts if s >= 0] if not starts: return text start = min(starts) return text[start:].strip() @asynccontextmanager async def llm_client() -> AsyncIterator[LlmClient]: """`async with llm_client() as llm:` for short-lived scripts.""" client = LlmClient() try: yield client finally: await client.aclose()