263 lines
9.5 KiB
Python
263 lines
9.5 KiB
Python
"""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 Callable, 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)
|
|
request_guard: Callable[[], None] | None = field(default=None, repr=False)
|
|
|
|
@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):
|
|
if self.request_guard is not None:
|
|
self.request_guard()
|
|
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())
|