Spaces:
Paused
Paused
| """Chat client for any OpenAI-compatible API (OpenRouter by default) and a scripted fake for tests.""" | |
| from __future__ import annotations | |
| import time | |
| from collections.abc import Sequence | |
| from dataclasses import dataclass, field | |
| from typing import Protocol | |
| import httpx | |
| from .config import Settings | |
| class LLMError(RuntimeError): | |
| pass | |
| class Message: | |
| role: str # "system" | "user" | "assistant" | |
| content: str | |
| class Completion: | |
| text: str | |
| model: str | |
| latency_ms: int | |
| usage: dict = field(default_factory=dict) | |
| class LLM(Protocol): | |
| model: str | |
| def complete( | |
| self, messages: Sequence[Message], temperature: float = 0.2, model: str | None = None | |
| ) -> Completion: ... | |
| class ChatClient: | |
| """POST {base_url}/chat/completions. The same key and endpoint are used for every tester.""" | |
| def __init__(self, settings: Settings, transport: httpx.BaseTransport | None = None, model: str | None = None): | |
| if not settings.llm_api_key: | |
| raise LLMError("LLM_API_KEY is not set") | |
| self.model = model or settings.llm_model | |
| self._openrouter = "openrouter.ai" in settings.llm_base_url | |
| self._client = httpx.Client( | |
| base_url=settings.llm_base_url, | |
| # X-Title labels the requests in the OpenRouter dashboard; other providers ignore it. | |
| headers={"Authorization": f"Bearer {settings.llm_api_key}", "X-Title": "RosDiff"}, | |
| timeout=settings.llm_timeout, | |
| transport=transport, | |
| ) | |
| def complete(self, messages: Sequence[Message], temperature: float = 0.2, model: str | None = None) -> Completion: | |
| """`model` overrides the client's default for this call (the router picks it per attempt).""" | |
| payload = { | |
| "model": model or self.model, | |
| "messages": [{"role": m.role, "content": m.content} for m in messages], | |
| "temperature": temperature, | |
| "max_tokens": 8192, | |
| } | |
| if self._openrouter: | |
| payload["usage"] = {"include": True} # OpenRouter then reports what the call cost (usage.cost) | |
| start = time.monotonic() | |
| last_error = "" | |
| for attempt in range(3): | |
| try: | |
| r = self._client.post("/chat/completions", json=payload) | |
| except httpx.HTTPError as e: | |
| last_error = f"{type(e).__name__}: {e}" | |
| else: | |
| if r.status_code == 200: | |
| try: | |
| data = r.json() | |
| except Exception: | |
| last_error = f"non-JSON response (HTTP 200): {r.text[:200]}" | |
| break | |
| try: | |
| text = data["choices"][0]["message"]["content"] or "" | |
| except (KeyError, IndexError) as e: | |
| raise LLMError(f"unexpected response shape: {str(data)[:300]}") from e | |
| return Completion( | |
| text=text, | |
| model=data.get("model", model or self.model), | |
| latency_ms=int((time.monotonic() - start) * 1000), | |
| usage=data.get("usage", {}), | |
| ) | |
| last_error = f"HTTP {r.status_code}: {r.text[:300]}" | |
| if r.status_code not in (429, 500, 502, 503, 504): | |
| break | |
| time.sleep(2**attempt) | |
| raise LLMError(f"LLM request failed: {last_error}") | |
| class ScriptedLLM: | |
| """Returns canned responses in order; records every prompt. For tests and offline dry runs.""" | |
| def __init__(self, responses: Sequence[str], model: str = "scripted"): | |
| self.responses = list(responses) | |
| self.model = model | |
| self.calls: list[list[Message]] = [] | |
| self.models: list[str] = [] | |
| def complete(self, messages: Sequence[Message], temperature: float = 0.2, model: str | None = None) -> Completion: | |
| self.calls.append(list(messages)) | |
| self.models.append(model or self.model) | |
| if not self.responses: | |
| raise LLMError("ScriptedLLM ran out of responses") | |
| usage = {"prompt_tokens": 1000, "completion_tokens": 500} | |
| return Completion(text=self.responses.pop(0), model=model or self.model, latency_ms=0, usage=usage) | |