Spaces:
Paused
Paused
File size: 4,308 Bytes
da839c2 f850954 5b063c2 f850954 da839c2 f850954 4508afe da839c2 4508afe 5b063c2 f850954 da839c2 f850954 5b063c2 f850954 5b063c2 f850954 5b063c2 f850954 8501103 f850954 5b063c2 f850954 da839c2 f850954 5b063c2 f850954 5b063c2 f850954 5b063c2 f850954 5b063c2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """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
@dataclass
class Message:
role: str # "system" | "user" | "assistant"
content: str
@dataclass
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)
|