spb/company-atlas

1"""LLM provider abstraction (spec §25): OpenAI-compatible chat completions with JSON mode, strict pydantic validation, one repair2retry, backoff on 429/503 (the MacLustr llm-api.io server loads models on demand and answers 503/429 while swapping), token3accounting and health. Keep ONE model per job stream (sticky) to avoid evictions — see `enrich.py`.45    provider = get_provider()6    result = await provider.complete_json("small", system, user_json, ChangeClassification, max_tokens=600)7    result.data  → validated pydantic model · result.request_tokens / response_tokens / latency_ms / model / attempts8"""9from __future__ import annotations1011import asyncio12import json13import logging14import re15import time16from abc import ABC, abstractmethod17from collections.abc import Awaitable, Callable18from dataclasses import dataclass, field19from typing import Any2021import httpx22from pydantic import BaseModel, ValidationError2324from companyatlas.config import settings2526log = logging.getLogger(__name__)2728Tier = str                                              # "small" | "medium" | "large"29RETRY_STATUS = {408, 425, 429, 500, 502, 503, 504}30_JSON_BLOCK = re.compile(r"```(?:json)?\s*(\{.*?\})\s*```", re.DOTALL)31_THINK = re.compile(r"<think>.*?</think>", re.DOTALL)323334class LLMError(Exception):35    def __init__(self, message: str, *, status: int | None = None, retryable: bool = False):36        super().__init__(message)37        self.status = status38        self.retryable = retryable394041class LLMNotConfigured(LLMError):42    pass434445class LLMValidationError(LLMError):46    """Model output could not be parsed/validated even after the repair retry."""474849@dataclass(slots=True)50class LLMResult[T: BaseModel]:51    data: T52    model: str53    request_tokens: int = 054    response_tokens: int = 055    latency_ms: int = 056    attempts: int = 157    repaired: bool = False58    raw: str = ""596061@dataclass(slots=True)62class Health:63    ok: bool64    base_url: str65    models: list[str] = field(default_factory=list)66    latency_ms: int = 067    error: str | None = None686970class LLMProvider(ABC):71    """Vendor-neutral interface. Implementations must be safe to share across tasks in one process."""7273    @abstractmethod74    def model_for(self, tier: Tier) -> str: ...7576    @abstractmethod77    async def complete_json[T: BaseModel](self, model_tier: Tier, system: str, user: str, schema: type[T], *, max_tokens: int = 800,78                            temperature: float = 0.1) -> LLMResult[T]: ...7980    @abstractmethod81    async def complete_text(self, model_tier: Tier, system: str, user: str, *, max_tokens: int = 400, temperature: float = 0.2) -> LLMResult[Any]: ...8283    @abstractmethod84    async def embed(self, texts: list[str]) -> list[list[float]]: ...8586    @abstractmethod87    async def health(self) -> Health: ...8889    async def close(self) -> None:90        return None919293class OpenAICompatibleProvider(LLMProvider):94    def __init__(self, *, base_url: str | None = None, api_key: str | None = None, timeout_s: float | None = None,95                 max_tries: int | None = None, backoff_initial_s: float | None = None, backoff_max_s: float | None = None,96                 sleep: Callable[[float], Awaitable[None]] | None = None, transport: httpx.AsyncBaseTransport | None = None):97        self.base_url = (base_url if base_url is not None else settings.llm_base_url).rstrip("/")98        self.api_key = api_key if api_key is not None else settings.llm_api_key99        self.timeout_s = timeout_s or settings.llm_timeout_s100        self.max_tries = max_tries or settings.llm_max_tries101        self.backoff_initial_s = backoff_initial_s if backoff_initial_s is not None else settings.llm_backoff_initial_s102        self.backoff_max_s = backoff_max_s if backoff_max_s is not None else settings.llm_backoff_max_s103        self._sleep = sleep or asyncio.sleep104        self._transport = transport105        self._client: httpx.AsyncClient | None = None106        self._models = {"small": settings.llm_small_model, "medium": settings.llm_medium_model, "large": settings.llm_large_model,107                        "embedding": settings.llm_embedding_model}108109    # ------------------------------------------------------------------ plumbing110    @property111    def configured(self) -> bool:112        return bool(self.base_url)113114    def model_for(self, tier: Tier) -> str:115        return self._models.get(tier) or tier                   # a literal model name is accepted too116117    def _headers(self) -> dict[str, str]:118        h = {"Content-Type": "application/json", "Accept": "application/json", "User-Agent": "CompanyAtlas-LLM/0.1"}119        if self.api_key:120            h["Authorization"] = f"Bearer {self.api_key}"121        return h122123    def client(self) -> httpx.AsyncClient:124        if self._client is None:125            self._client = httpx.AsyncClient(base_url=self.base_url, headers=self._headers(), timeout=httpx.Timeout(self.timeout_s, connect=20),126                                             transport=self._transport)127        return self._client128129    async def close(self) -> None:130        if self._client is not None:131            await self._client.aclose()132            self._client = None133134    def _backoff(self, attempt: int) -> float:135        return min(self.backoff_max_s, self.backoff_initial_s * (2 ** attempt))136137    async def _post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:138        if not self.configured:139            raise LLMNotConfigured("CA_LLM_BASE_URL is not set")140        last: LLMError | None = None141        for attempt in range(self.max_tries):142            try:143                r = await self.client().post(path, json=payload)144            except httpx.TimeoutException as exc:145                last = LLMError(f"timeout: {exc}", retryable=True)146            except httpx.HTTPError as exc:147                last = LLMError(f"transport: {exc}", retryable=True)148            else:149                if r.status_code < 300:150                    try:151                        return r.json()152                    except ValueError as exc:153                        raise LLMError(f"non-JSON response: {exc}", status=r.status_code) from exc154                body = r.text[:300]155                if r.status_code in RETRY_STATUS:156                    last = LLMError(f"HTTP {r.status_code}: {body}", status=r.status_code, retryable=True)157                    retry_after = r.headers.get("Retry-After")158                    if retry_after and retry_after.isdigit():159                        await self._sleep(min(self.backoff_max_s, float(retry_after)))160                        continue161                else:162                    raise LLMError(f"HTTP {r.status_code}: {body}", status=r.status_code)163            if attempt + 1 < self.max_tries:164                delay = self._backoff(attempt)165                log.warning("llm retry", extra={"attempt": attempt + 1, "delay_s": delay, "error": str(last)})166                await self._sleep(delay)167        raise last or LLMError("unknown LLM failure", retryable=True)168169    # ------------------------------------------------------------------ completions170    async def _chat(self, model: str, messages: list[dict[str, str]], *, max_tokens: int, temperature: float, json_mode: bool) -> tuple[str, dict[str, int], int]:171        payload: dict[str, Any] = {"model": model, "messages": messages, "max_tokens": max_tokens, "temperature": temperature, "stream": False}172        if json_mode:173            payload["response_format"] = {"type": "json_object"}174        t0 = time.monotonic()175        data = await self._post("/chat/completions", payload)176        latency = int((time.monotonic() - t0) * 1000)177        try:178            content = data["choices"][0]["message"]["content"] or ""179        except (KeyError, IndexError, TypeError) as exc:180            raise LLMError(f"malformed completion: {json.dumps(data)[:200]}") from exc181        usage = data.get("usage") or {}182        tokens = {"request": int(usage.get("prompt_tokens") or 0), "response": int(usage.get("completion_tokens") or 0)}183        return str(content), tokens, latency184185    async def complete_text(self, model_tier: Tier, system: str, user: str, *, max_tokens: int = 400, temperature: float = 0.2) -> LLMResult[Any]:186        model = self.model_for(model_tier)187        content, tokens, latency = await self._chat(model, [{"role": "system", "content": system}, {"role": "user", "content": user}],188                                                    max_tokens=max_tokens, temperature=temperature, json_mode=False)189        text = _THINK.sub("", content).strip()190        return LLMResult(data=text, model=model, request_tokens=tokens["request"], response_tokens=tokens["response"], latency_ms=latency, raw=content)191192    async def complete_json[T: BaseModel](self, model_tier: Tier, system: str, user: str, schema: type[T], *, max_tokens: int = 800,193                            temperature: float = 0.1) -> LLMResult[T]:194        model = self.model_for(model_tier)195        sys_prompt = system.rstrip() + "\n\nRespond with a single JSON object only. JSON schema:\n" + json.dumps(_compact_schema(schema), ensure_ascii=False)196        messages = [{"role": "system", "content": sys_prompt}, {"role": "user", "content": user}]197        content, tokens, latency = await self._chat(model, messages, max_tokens=max_tokens, temperature=temperature, json_mode=True)198        req, resp = tokens["request"], tokens["response"]199        try:200            return LLMResult(data=parse_json_object(content, schema), model=model, request_tokens=req, response_tokens=resp, latency_ms=latency, raw=content)201        except (ValueError, ValidationError) as exc:202            error = str(exc)[:800]203            log.info("llm output invalid, repairing", extra={"model": model, "schema": schema.__name__, "error": error[:160]})204        # one repair retry: show the model its own output and the validation error205        messages += [{"role": "assistant", "content": content[:4000]},206                     {"role": "user", "content": "Your previous answer was not valid for the schema. Error:\n" + error +207                      "\nReturn ONLY the corrected JSON object with the same meaning. No prose, no markdown."}]208        content2, tokens2, latency2 = await self._chat(model, messages, max_tokens=max_tokens, temperature=0.0, json_mode=True)209        req += tokens2["request"]210        resp += tokens2["response"]211        try:212            data = parse_json_object(content2, schema)213        except (ValueError, ValidationError) as exc:214            raise LLMValidationError(f"invalid after repair: {str(exc)[:300]}") from exc215        return LLMResult(data=data, model=model, request_tokens=req, response_tokens=resp, latency_ms=latency + latency2, attempts=2, repaired=True, raw=content2)216217    async def embed(self, texts: list[str]) -> list[list[float]]:218        if not texts:219            return []220        data = await self._post("/embeddings", {"model": self.model_for("embedding"), "input": texts})221        items = sorted(data.get("data") or [], key=lambda d: d.get("index", 0))222        return [list(map(float, d["embedding"])) for d in items]223224    async def health(self) -> Health:225        if not self.configured:226            return Health(ok=False, base_url="", error="not configured")227        t0 = time.monotonic()228        try:229            r = await self.client().get("/models")230            latency = int((time.monotonic() - t0) * 1000)231            if r.status_code >= 300:232                return Health(ok=False, base_url=self.base_url, latency_ms=latency, error=f"HTTP {r.status_code}")233            body = r.json()234            models = [m.get("id") for m in (body.get("data") or []) if isinstance(m, dict) and m.get("id")]235            return Health(ok=True, base_url=self.base_url, models=models, latency_ms=latency)236        except (httpx.HTTPError, ValueError) as exc:237            return Health(ok=False, base_url=self.base_url, latency_ms=int((time.monotonic() - t0) * 1000), error=str(exc)[:200])238239240# ---------------------------------------------------------------------------------------------------------------- parsing241242243def extract_json_text(content: str) -> str:244    """Tolerate reasoning tags, code fences and prose around the object; return the outermost {...} span."""245    text = _THINK.sub("", content or "").strip()246    m = _JSON_BLOCK.search(text)247    if m:248        return m.group(1)249    start = text.find("{")250    end = text.rfind("}")251    if start == -1 or end == -1 or end < start:252        raise ValueError("no JSON object in model output")253    return text[start:end + 1]254255256def parse_json_object[T: BaseModel](content: str, schema: type[T]) -> T:257    raw = extract_json_text(content)258    try:259        obj = json.loads(raw)260    except json.JSONDecodeError as exc:261        raise ValueError(f"invalid JSON: {exc.msg} at {exc.pos}") from exc262    if isinstance(obj, dict):263        return schema.model_validate(obj)264    raise ValueError("top-level JSON value must be an object")265266267def _compact_schema(schema: type[BaseModel]) -> dict[str, Any]:268    full = schema.model_json_schema()269    props = {}270    for name, spec in (full.get("properties") or {}).items():271        t = spec.get("type") or ("/".join(x.get("type", "?") for x in spec.get("anyOf", [])) if spec.get("anyOf") else "any")272        entry: dict[str, Any] = {"type": t}273        if spec.get("description"):274            entry["description"] = spec["description"]275        for k in ("maxLength", "minimum", "maximum", "maxItems"):276            if k in spec:277                entry[k] = spec[k]278        props[name] = entry279    return {"type": "object", "properties": props, "required": full.get("required", [])}280281282# ---------------------------------------------------------------------------------------------------------------- singleton283284_provider: LLMProvider | None = None285286287def get_provider() -> LLMProvider:288    global _provider289    if _provider is None:290        _provider = OpenAICompatibleProvider()291    return _provider292293294def set_provider(provider: LLMProvider | None) -> None:295    """Tests / alternative vendors."""296    global _provider297    _provider = provider298299300__all__ = ["Health", "LLMError", "LLMNotConfigured", "LLMProvider", "LLMResult", "LLMValidationError", "OpenAICompatibleProvider", "extract_json_text",301           "get_provider", "parse_json_object", "set_provider"]302