"""The LLM connector: step 1 of the pipeline, and the thing everything else sits on. Deliberately dependency-free (stdlib ``urllib`` only) and provider-neutral: any OpenAI-compatible ``/chat/completions`` endpoint works -- DashScope/Qwen by default, but also vLLM, Ollama, OpenRouter, DeepSeek, or OpenAI itself. Swap it out entirely by passing any object with a ``complete(messages, **kwargs) -> str`` method. """ from __future__ import annotations import json import os import random import re import time import urllib.error import urllib.request from collections.abc import Iterable from dataclasses import dataclass, field from typing import Any, Protocol from .config import LLMSettings RETRY_STATUS = {408, 409, 425, 429, 500, 502, 503, 504} _FENCE = re.compile(r"```(?:json|lean4|lean)?\s*(.*?)```", re.DOTALL) class LLMError(RuntimeError): """Any failure to obtain a usable answer from the provider.""" class ChatClient(Protocol): """The whole contract the rest of the pipeline depends on.""" def complete(self, messages: list[dict[str, str]], **kwargs: Any) -> str: ... @dataclass class Usage: calls: int = 0 prompt_tokens: int = 0 completion_tokens: int = 0 def add(self, data: dict[str, Any]) -> None: self.calls += 1 usage = data.get("usage") or {} self.prompt_tokens += int(usage.get("prompt_tokens") or 0) self.completion_tokens += int(usage.get("completion_tokens") or 0) def as_dict(self) -> dict[str, int]: return { "calls": self.calls, "prompt_tokens": self.prompt_tokens, "completion_tokens": self.completion_tokens, } @dataclass class OpenAICompatibleClient: """POST to ``{base_url}/chat/completions``, with retries and JSON helpers.""" model: str base_url: str api_key_env: str = "DASHSCOPE_API_KEY" timeout: float = 300.0 max_retries: int = 4 temperature: float = 0.0 usage: Usage = field(default_factory=Usage) @classmethod def from_settings(cls, settings: LLMSettings) -> OpenAICompatibleClient: return cls( model=settings.model, base_url=settings.base_url, api_key_env=settings.api_key_env, timeout=settings.timeout, max_retries=settings.max_retries, temperature=settings.temperature, ) # -- transport --------------------------------------------------------- def _api_key(self) -> str: key = os.environ.get(self.api_key_env, "").strip() if not key: raise LLMError( f"No API key: set {self.api_key_env} in the environment or in .env" ) return key def _post(self, payload: dict[str, Any]) -> dict[str, Any]: url = self.base_url.rstrip("/") + "/chat/completions" body = json.dumps(payload).encode("utf-8") headers = { "Content-Type": "application/json", "Authorization": "Bearer " + self._api_key(), } last: Exception | None = None for attempt in range(self.max_retries + 1): request = urllib.request.Request(url, data=body, headers=headers, method="POST") try: with urllib.request.urlopen(request, timeout=self.timeout) as response: return json.loads(response.read().decode("utf-8")) except urllib.error.HTTPError as error: detail = error.read().decode("utf-8", "replace")[:600] last = LLMError(f"HTTP {error.code} from {url}: {detail}") if error.code not in RETRY_STATUS: raise last from error except (urllib.error.URLError, TimeoutError, OSError) as error: last = LLMError(f"Network failure talking to {url}: {error}") if attempt < self.max_retries: time.sleep(min(2 ** attempt, 30) + random.random()) raise last or LLMError("Request failed for an unknown reason") # -- api --------------------------------------------------------------- def complete( self, messages: list[dict[str, str]], *, temperature: float | None = None, json_object: bool = False, max_tokens: int | None = None, **extra: Any, ) -> str: payload: dict[str, Any] = { "model": self.model, "messages": messages, "temperature": self.temperature if temperature is None else temperature, } if max_tokens: payload["max_tokens"] = max_tokens if json_object: payload["response_format"] = {"type": "json_object"} payload.update(extra) try: data = self._post(payload) except LLMError as error: # Not every OpenAI-compatible server implements response_format. if json_object and "response_format" in str(error): payload.pop("response_format", None) data = self._post(payload) else: raise self.usage.add(data) try: content = data["choices"][0]["message"]["content"] except (KeyError, IndexError, TypeError) as error: raise LLMError(f"Unexpected response shape: {json.dumps(data)[:600]}") from error if isinstance(content, list): # some servers return content parts content = "".join(part.get("text", "") for part in content if isinstance(part, dict)) if not isinstance(content, str) or not content.strip(): raise LLMError("Provider returned empty content") return content def complete_json( self, messages: list[dict[str, str]], *, attempts: int = 3, temperature: float | None = None, **extra: Any, ) -> Any: """Ask for JSON and keep asking until it parses (or give up loudly).""" conversation = list(messages) error_text = "" for attempt in range(attempts): raw = self.complete( conversation, temperature=temperature, json_object=attempt == 0, **extra, ) try: return parse_json(raw) except ValueError as error: error_text = str(error) conversation = list(messages) + [ {"role": "assistant", "content": raw[:4000]}, { "role": "user", "content": ( f"That was not valid JSON ({error_text}). Reply with the JSON " "object only: no prose, no markdown fences." ), }, ] raise LLMError(f"Model never produced valid JSON: {error_text}") def parse_json(raw: str) -> Any: """Parse JSON out of a model reply, tolerating fences and surrounding prose.""" text = raw.strip() try: return json.loads(text) except json.JSONDecodeError: pass fenced = _FENCE.search(text) if fenced: try: return json.loads(fenced.group(1).strip()) except json.JSONDecodeError: text = fenced.group(1).strip() candidates = [i for i in (text.find("{"), text.find("[")) if i != -1] if candidates: start = min(candidates) end = max(text.rfind("}"), text.rfind("]")) if end > start: try: return json.loads(text[start : end + 1]) except json.JSONDecodeError as error: raise ValueError(str(error)) from error raise ValueError("no JSON found in reply") def extract_code(raw: str, language: str = "lean") -> str: """Pull a fenced code block out of a reply, or return the reply unchanged.""" fenced = re.search(r"```(?:" + language + r"4?)?[ \t]*\n(.*?)```", raw, re.DOTALL) return (fenced.group(1) if fenced else raw).strip() class RecordingClient: """Wraps a client and keeps every exchange -- handy for debugging a bad run.""" def __init__(self, inner: ChatClient) -> None: self.inner = inner self.log: list[dict[str, Any]] = [] def complete(self, messages: list[dict[str, str]], **kwargs: Any) -> str: reply = self.inner.complete(messages, **kwargs) self.log.append({"messages": messages, "reply": reply}) return reply def complete_json(self, messages: list[dict[str, str]], **kwargs: Any) -> Any: result = self.inner.complete_json(messages, **kwargs) # type: ignore[attr-defined] self.log.append({"messages": messages, "reply": result}) return result class ScriptedClient: """A canned-reply client for tests; replies are consumed in order.""" def __init__(self, replies: Iterable[str]) -> None: self.replies = list(replies) self.calls: list[list[dict[str, str]]] = [] def complete(self, messages: list[dict[str, str]], **_: Any) -> str: self.calls.append(messages) if not self.replies: raise LLMError("ScriptedClient ran out of replies") return self.replies.pop(0) def complete_json(self, messages: list[dict[str, str]], **_: Any) -> Any: return parse_json(self.complete(messages)) def build_client(settings: LLMSettings | None = None) -> OpenAICompatibleClient: return OpenAICompatibleClient.from_settings(settings or LLMSettings.from_env())