HTML 77.2%
TypeScript 10.5%
Python 9.6%
JavaScript 2.5%
1"""LLM gateway: stage cascade (small → medium → large), OpenAI-compatible engine, strict JSON output validated by pydantic2schemas, full cost accounting in `llm_jobs`. Deterministic processing always runs before any call here."""3from __future__ import annotations45import asyncio6import json7import logging8import re9import socket10import time11from dataclasses import dataclass12from typing import Any1314import httpx15from pydantic import BaseModel, ValidationError1617from aiatlas.config import settings18from aiatlas.db import execute, jsonb, transaction19from aiatlas.ids import new_id2021log = logging.getLogger(__name__)2223STAGES = ("small", "medium", "large")242526class LLMUnavailable(Exception):27 pass282930@dataclass31class LLMResult:32 ok: bool33 data: dict[str, Any] | None34 raw: str35 model: str36 stage: str37 input_tokens: int | None38 output_tokens: int | None39 duration_ms: int40 error: str | None = None41 llm_job_id: str | None = None424344class Engine:45 name = "null"4647 async def complete(self, *, model: str, system: str, user: str, max_tokens: int, temperature: float, json_mode: bool) -> tuple[str, dict[str, Any]]:48 raise LLMUnavailable("no LLM engine configured")4950 async def embed(self, *, model: str, texts: list[str]) -> list[list[float]]:51 raise LLMUnavailable("no embedding engine configured")5253 async def health(self) -> bool:54 return False555657class OpenAICompatEngine(Engine):58 """Works with MacLustr llm-api.io, vLLM, llama.cpp server, Ollama (/v1), LM Studio…"""5960 name = "openai_compat"6162 def __init__(self, base_url: str, api_key: str, timeout_s: float):63 self.base_url = base_url.rstrip("/")64 self.api_key = api_key65 self.timeout_s = timeout_s6667 def _client(self) -> httpx.AsyncClient:68 return httpx.AsyncClient(base_url=self.base_url, headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},69 timeout=httpx.Timeout(self.timeout_s, connect=30))7071 async def _post(self, client: httpx.AsyncClient, path: str, body: dict[str, Any]) -> httpx.Response:72 """Local inference servers answer 503/429 while they load or swap a model: wait and retry instead of failing the job."""73 delay = 15.074 for attempt in range(6):75 r = await client.post(path, json=body)76 if r.status_code in (503, 429, 502, 504) and attempt < 5:77 retry_after = r.headers.get("retry-after")78 wait = min(120.0, float(retry_after)) if retry_after and retry_after.replace(".", "", 1).isdigit() else delay79 log.info("llm server busy, retrying", extra={"status": r.status_code, "wait": wait, "path": path})80 await asyncio.sleep(wait)81 delay = min(120.0, delay * 1.7)82 continue83 return r84 return r8586 async def complete(self, *, model: str, system: str, user: str, max_tokens: int, temperature: float, json_mode: bool) -> tuple[str, dict[str, Any]]:87 body: dict[str, Any] = {"model": model, "messages": [{"role": "system", "content": system}, {"role": "user", "content": user}],88 "max_tokens": max_tokens, "temperature": temperature, "stream": False}89 if json_mode:90 body["response_format"] = {"type": "json_object"}91 async with self._client() as client:92 r = await self._post(client, "/chat/completions", body)93 if r.status_code == 400 and json_mode:94 body.pop("response_format")95 r = await self._post(client, "/chat/completions", body)96 r.raise_for_status()97 data = r.json()98 choice = (data.get("choices") or [{}])[0]99 content = (choice.get("message") or {}).get("content") or ""100 return content, data.get("usage") or {}101102 async def embed(self, *, model: str, texts: list[str]) -> list[list[float]]:103 async with self._client() as client:104 r = await self._post(client, "/embeddings", {"model": model, "input": texts})105 r.raise_for_status()106 data = r.json()107 items = sorted(data.get("data") or [], key=lambda d: d.get("index", 0))108 return [d["embedding"] for d in items]109110 async def health(self) -> bool:111 try:112 async with self._client() as client:113 r = await client.get("/models", timeout=20)114 return r.status_code == 200115 except Exception: # noqa: BLE001116 return False117118119SYSTEM_PROMPT = """You are the extraction engine of AI Atlas, a structured database of the AI ecosystem.120Extract ONLY facts explicitly stated in the document. Never guess, never fill gaps from memory, never invent numbers or dates.121If a field is not stated, use null. Quote units exactly as written. Prefer official statements over marketing phrasing.122Return a single JSON object matching the requested schema and nothing else."""123124125class LLMGateway:126 def __init__(self, engine: Engine | None = None):127 self.engine = engine or self._default_engine()128 self.models = {"small": settings.llm_small_model, "medium": settings.llm_medium_model, "large": settings.llm_large_model}129130 @staticmethod131 def _default_engine() -> Engine:132 if settings.llm_available:133 return OpenAICompatEngine(settings.llm_base_url, settings.llm_api_key, settings.llm_timeout_s)134 return Engine()135136 @property137 def available(self) -> bool:138 return self.engine.name != "null"139140 async def health(self) -> dict[str, Any]:141 return {"engine": self.engine.name, "available": self.available, "reachable": await self.engine.health() if self.available else False,142 "models": self.models, "embedding_model": settings.embedding_model}143144 async def extract(self, *, task_type: str, document: str, schema: type[BaseModel], stage: str = "medium", instructions: str = "",145 snapshot_id: str | None = None, entity_id: str | None = None, job_id: str | None = None, max_tokens: int = 2000,146 escalate_on_failure: bool = True, max_chars: int = 24000) -> LLMResult:147 if not self.available:148 raise LLMUnavailable("LLM engine not configured (AIA_LLM_BASE_URL / AIA_LLM_API_KEY)")149 if stage not in STAGES:150 raise ValueError(f"stage must be one of {STAGES}")151 doc = document if len(document) <= max_chars else document[:max_chars] + "\n…[truncated]"152 schema_json = json.dumps(schema.model_json_schema(), ensure_ascii=False)153 user = (f"TASK: {task_type}\n{instructions}\n\nJSON SCHEMA:\n{schema_json}\n\nDOCUMENT:\n<<<\n{doc}\n>>>\n\n"154 f"Return only the JSON object.")155 stages = STAGES[STAGES.index(stage):] if escalate_on_failure else (stage,)156 last: LLMResult | None = None157 for st in stages:158 model = self.models[st]159 t0 = time.perf_counter()160 error = None161 data: dict[str, Any] | None = None162 raw = ""163 usage: dict[str, Any] = {}164 status = "ok"165 try:166 raw, usage = await self.engine.complete(model=model, system=SYSTEM_PROMPT, user=user, max_tokens=max_tokens, temperature=0.0, json_mode=True)167 parsed = _parse_json(raw)168 if parsed is None:169 status, error = "invalid_json", "model did not return JSON"170 else:171 try:172 data = schema.model_validate(parsed).model_dump(mode="json")173 except ValidationError as ve:174 status, error = "schema_error", str(ve)[:2000]175 data = parsed # keep for debugging176 except Exception as exc: # noqa: BLE001177 status, error = "failed", f"{exc.__class__.__name__}: {exc}"[:2000]178 duration = int((time.perf_counter() - t0) * 1000)179 llm_job_id = await self._account(task_type=task_type, stage=st, model=model, schema_name=schema.__name__, snapshot_id=snapshot_id,180 entity_id=entity_id, job_id=job_id, usage=usage, duration=duration, status=status,181 output=data if status in ("ok", "schema_error") else None, error=error)182 last = LLMResult(ok=status == "ok", data=data if status == "ok" else None, raw=raw, model=model, stage=st, input_tokens=usage.get("prompt_tokens"),183 output_tokens=usage.get("completion_tokens"), duration_ms=duration, error=error, llm_job_id=llm_job_id)184 if last.ok or status == "failed":185 break # transport failures: do not escalate blindly (probably the server), schema errors: try the next stage186 assert last is not None187 return last188189 async def classify(self, *, text: str, labels: list[str], task_type: str = "classify", snapshot_id: str | None = None) -> str | None:190 from pydantic import Field, create_model191192 Model = create_model("Classification", label=(str, Field(description=f"one of: {', '.join(labels)}")), confidence=(float, Field(ge=0, le=1)))193 res = await self.extract(task_type=task_type, document=text[:6000], schema=Model, stage="small", snapshot_id=snapshot_id,194 instructions=f"Choose exactly one label among: {', '.join(labels)}.", max_tokens=100, escalate_on_failure=False)195 if res.ok and res.data and res.data.get("label") in labels:196 return str(res.data["label"])197 return None198199 async def embed(self, texts: list[str]) -> list[list[float]]:200 if not self.available:201 raise LLMUnavailable("embedding engine not configured")202 t0 = time.perf_counter()203 vectors = await self.engine.embed(model=settings.embedding_model, texts=texts)204 await self._account(task_type="embed", stage="small", model=settings.embedding_model, schema_name=None, snapshot_id=None, entity_id=None,205 job_id=None, usage={"prompt_tokens": sum(len(t) // 4 for t in texts)}, duration=int((time.perf_counter() - t0) * 1000),206 status="ok", output=None, error=None)207 return vectors208209 async def _account(self, *, task_type: str, stage: str, model: str, schema_name: str | None, snapshot_id: str | None, entity_id: str | None,210 job_id: str | None, usage: dict[str, Any], duration: int, status: str, output: dict[str, Any] | None, error: str | None) -> str:211 lid = new_id("llm_job")212 try:213 async with transaction() as conn:214 await execute(conn, """insert into llm_jobs (id, job_id, task_type, stage, engine, model, node, schema_name, snapshot_id, entity_id, input_tokens, output_tokens,215 duration_ms, status, output, error) values (:id, :j, :t, :st, :eng, :m, :node, :schema, :snap, :e, :it, :ot, :d, :status, cast(:o as jsonb), :err)""",216 id=lid, j=job_id, t=task_type, st=stage, eng=self.engine.name, m=model, node=socket.gethostname(), schema=schema_name, snap=snapshot_id,217 e=entity_id, it=usage.get("prompt_tokens"), ot=usage.get("completion_tokens"), d=duration, status=status,218 o=jsonb(output) if output is not None else None, err=error)219 except Exception as exc: # noqa: BLE001220 log.warning("llm accounting failed", extra={"error": str(exc)})221 return lid222223224def _parse_json(raw: str) -> dict[str, Any] | None:225 s = raw.strip()226 s = re.sub(r"<think>.*?</think>", "", s, flags=re.DOTALL).strip()227 if s.startswith("```"):228 s = re.sub(r"^```(?:json)?\s*", "", s)229 s = re.sub(r"\s*```$", "", s)230 try:231 v = json.loads(s)232 return v if isinstance(v, dict) else None233 except json.JSONDecodeError:234 pass235 m = re.search(r"\{.*\}", s, flags=re.DOTALL)236 if m:237 try:238 v = json.loads(m.group(0))239 return v if isinstance(v, dict) else None240 except json.JSONDecodeError:241 return None242 return None243244245gateway = LLMGateway()246247__all__ = ["Engine", "LLMGateway", "LLMResult", "LLMUnavailable", "OpenAICompatEngine", "gateway"]248