File size: 15,229 Bytes
750b88d | 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 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 | """LLM and embedding clients.
ChatClient answers through the OpenAI SDK; `base_url` points it at any OpenAI-compatible
endpoint. OpenAIChatClient is the free-form chat client used by the methods' internal LLM
calls. Every call is logged with its phase (memorize, query, judge). Each endpoint's API key is read from
the environment variable that backbone.yaml names for it (`api_key_env`, OPENAI_API_KEY by default), so a key
is sent only to its own endpoint.
"""
from __future__ import annotations
import collections
import json
import os
import re
import sys
import threading
import time
from pathlib import Path
def extract_json_object(text: str) -> dict:
"""Parse the first balanced top-level JSON object out of `text`."""
text = text.strip()
if text.startswith("```"):
text = re.sub(r"^```[a-zA-Z]*\n?|\n?```$", "", text).strip()
decoder = json.JSONDecoder()
start = text.find("{")
while start != -1:
try:
obj, _end = decoder.raw_decode(text, start)
return obj
except json.JSONDecodeError:
start = text.find("{", start + 1)
raise ValueError("no parseable JSON object in model output")
class ContextOverflow(Exception):
"""The prompt exceeded the model's input limit."""
_REPAIR = ("\n\nYour previous reply was not valid JSON. Reply with ONLY a single valid JSON "
"object, no prose, no code fences.")
HTTP_TIMEOUT_S = 600.0
HTTP_CONNECT_TIMEOUT_S = 10.0
SDK_MAX_RETRIES = 0
MAX_ATTEMPTS = 4
def http_timeout():
"""Per-request timeout of every client: HTTP_TIMEOUT_S, with HTTP_CONNECT_TIMEOUT_S to connect."""
import httpx
return httpx.Timeout(HTTP_TIMEOUT_S, connect=HTTP_CONNECT_TIMEOUT_S)
def _http_kwargs() -> dict:
"""Timeout + retry settings shared by every openai.OpenAI() built in this module."""
return {"timeout": http_timeout(), "max_retries": SDK_MAX_RETRIES}
JUDGE_EFFORT, JUDGE_MAX_TOKENS = "medium", 800
REASONING_MODELS: set[str] = set()
RATE_LIMIT_HITS: "collections.Counter[str]" = collections.Counter()
_KEY_ENV: dict = {}
ENDPOINTS: dict = {"judge": None, "embedding": None}
def set_endpoints(roles: dict) -> None:
"""Record each role's endpoint and key variable from backbone.yaml ({role: {base_url, api_key_env}}); an
endpoint used by two roles must name the same key variable. A memory role without a model uses the answer
endpoint and is skipped."""
_KEY_ENV.clear()
for role, sec in roles.items():
if role == "memory" and not (sec or {}).get("model"):
continue
url = (sec or {}).get("base_url") or None
env = (sec or {}).get("api_key_env") or "OPENAI_API_KEY"
if _KEY_ENV.get(url, env) != env:
raise SystemExit(f"backbone.yaml: {url or 'the default endpoint'} is given two api_key_env values")
_KEY_ENV[url] = env
for role in ENDPOINTS:
ENDPOINTS[role] = (roles.get(role) or {}).get("base_url") or None
def api_key_for(base_url: str | None) -> str:
"""The API key for an endpoint (None = the default endpoint)."""
env = _KEY_ENV.get(base_url or None, "OPENAI_API_KEY")
key = os.environ.get(env)
if not key:
raise RuntimeError(f"{env} is not set (the API key for {base_url or 'the default endpoint'})")
return key
def set_reasoning_models(models) -> None:
"""Record the configured models marked `reasoning: true` in backbone.yaml."""
REASONING_MODELS.clear()
REASONING_MODELS.update(m for m in models if m)
def _is_reasoning_model(model: str) -> bool:
"""Return True for a reasoning model: it takes `max_completion_tokens` and accepts `reasoning_effort`."""
return model in REASONING_MODELS
def _note_rate_limit(who: str) -> None:
RATE_LIMIT_HITS[who] += 1
n = RATE_LIMIT_HITS[who]
if n <= 5 or n % 25 == 0:
print(f"[rate-limit] {who} 429/throttle #{n}", file=sys.stderr, flush=True)
class UsageTracker:
def __init__(self):
self._lock = threading.Lock()
self.path: Path | None = None
self.phase = "?"
self.totals: dict = {}
def bind(self, path: Path) -> None:
"""Direct the usage log to a run directory and reset its totals."""
self.path = Path(path)
with self._lock:
self.totals = {}
def set_phase(self, phase: str) -> None:
self.phase = phase
def log(self, model: str, kind: str, prompt: int, completion: int,
reasoning: int = 0, cached: int = 0) -> None:
rec = {"model": model, "kind": kind, "phase": self.phase, "prompt": prompt,
"completion": completion, "reasoning": reasoning, "cached": cached,
"total": prompt + completion}
with self._lock:
t = self.totals.setdefault(f"{self.phase}/{kind}", {"prompt": 0, "completion": 0,
"reasoning": 0, "cached": 0, "calls": 0})
t["prompt"] += prompt
t["completion"] += completion
t["reasoning"] += reasoning
t["cached"] += cached
t["calls"] += 1
if self.path:
try:
with open(self.path, "a", encoding="utf-8") as fh:
fh.write(json.dumps(rec, ensure_ascii=False) + "\n")
except Exception:
pass
USAGE = UsageTracker()
class ChatClient:
def __init__(self, model: str, reasoning_effort: str | None, max_completion_tokens: int,
max_retries: int = MAX_ATTEMPTS, temperature: float | None = None,
base_url: str | None = None):
import openai
key = api_key_for(base_url)
self._client = (openai.OpenAI(api_key=key, base_url=base_url, **_http_kwargs()) if base_url
else openai.OpenAI(api_key=key, **_http_kwargs()))
self.model = model
self.reasoning_effort = reasoning_effort
self.max_completion_tokens = max_completion_tokens
self.max_retries = max_retries
self.temperature = temperature
def complete_json(self, system: str, user: str) -> dict:
messages = [{"role": "system", "content": system}, {"role": "user", "content": user}]
effort_ok = format_ok = True
last_err = None
for attempt in range(self.max_retries):
try:
kw = dict(model=self.model, messages=messages, max_completion_tokens=self.max_completion_tokens)
if format_ok:
kw["response_format"] = {"type": "json_object"}
if self.temperature is not None:
kw["temperature"] = self.temperature
if self.reasoning_effort and effort_ok:
kw["reasoning_effort"] = self.reasoning_effort
resp = self._client.chat.completions.create(**kw)
u = getattr(resp, "usage", None)
if u is not None:
cd = getattr(u, "completion_tokens_details", None)
pdt = getattr(u, "prompt_tokens_details", None)
USAGE.log(self.model, "chat",
int(getattr(u, "prompt_tokens", 0) or 0),
int(getattr(u, "completion_tokens", 0) or 0),
int(getattr(cd, "reasoning_tokens", 0) or 0) if cd else 0,
int(getattr(pdt, "cached_tokens", 0) or 0) if pdt else 0)
return extract_json_object(resp.choices[0].message.content or "")
except Exception as e:
last_err = e
msg = str(e).lower()
if any(k in msg for k in ("context", "exceed the configured limit", "too long",
"maximum context length")):
raise ContextOverflow(str(e)) from e
if "429" in msg or "rate limit" in msg or "rate_limit" in msg or "overloaded" in msg:
_note_rate_limit("chat")
if "reasoning_effort" in msg:
effort_ok = False
if "response_format" in msg:
format_ok = False
if attempt == self.max_retries - 2:
messages = messages + [{"role": "user", "content": _REPAIR}]
time.sleep(1.5 * (attempt + 1))
raise RuntimeError(f"ChatClient failed after {self.max_retries} tries: {last_err}")
class LLMResponse:
"""Result of a free-form chat call: content and token usage."""
def __init__(self, content: str, input_tokens: int = 0, output_tokens: int = 0,
model: str = "", finish_reason: str | None = None):
self.content = content
self.input_tokens = input_tokens
self.output_tokens = output_tokens
self.model = model
self.finish_reason = finish_reason
class BaseLLMClient:
"""Interface of a free-form chat client."""
def chat(self, messages, temperature=None, max_tokens=None, **kwargs) -> LLMResponse:
raise NotImplementedError
class OpenAIChatClient(BaseLLMClient):
"""Free-form chat client; token usage is logged under the current phase."""
def __init__(self, model: str, temperature: float | None, max_tokens: int,
api_key: str | None = None, base_url: str | None = None,
max_retries: int = MAX_ATTEMPTS):
import openai
base_url = base_url or os.environ.get("OPENAI_BASE_URL")
self._client = openai.OpenAI(api_key=api_key or api_key_for(base_url), base_url=base_url,
**_http_kwargs())
self.model = model
self.temperature = temperature
self.max_tokens = max_tokens
self.max_retries = max_retries
def _use_max_completion_tokens(self) -> bool:
return _is_reasoning_model(self.model)
def chat(self, messages, temperature=None, max_tokens=None, **kwargs) -> LLMResponse:
"""One completion. A client built with temperature None never sends a temperature (for models that accept
only their default), also when a call passes one."""
params = {"model": self.model, "messages": messages}
if self.temperature is not None:
params["temperature"] = self.temperature if temperature is None else temperature
params["max_completion_tokens" if self._use_max_completion_tokens()
else "max_tokens"] = max_tokens or self.max_tokens
params.update(kwargs)
last_err = None
for attempt in range(self.max_retries):
try:
resp = self._client.chat.completions.create(**params)
u = getattr(resp, "usage", None)
p_tok = int(getattr(u, "prompt_tokens", 0) or 0) if u else 0
c_tok = int(getattr(u, "completion_tokens", 0) or 0) if u else 0
if u is not None:
cd = getattr(u, "completion_tokens_details", None)
USAGE.log(self.model, "chat", p_tok, c_tok,
int(getattr(cd, "reasoning_tokens", 0) or 0) if cd else 0)
_ch = resp.choices[0]
return LLMResponse(_ch.message.content or "", p_tok, c_tok,
self.model, getattr(_ch, "finish_reason", None))
except Exception as e:
last_err = e
if "429" in str(e).lower() or "rate limit" in str(e).lower():
_note_rate_limit("chat")
if attempt < self.max_retries - 1:
time.sleep(1.5 * (attempt + 1))
raise RuntimeError(f"OpenAIChatClient failed after {self.max_retries} tries: {last_err}")
def create_llm_client(*, max_tokens: int, provider: str = "openai", model: str | None = None,
temperature: float | None = 0.3, api_key: str | None = None,
base_url: str | None = None) -> BaseLLMClient:
"""Return a free-form chat client for an OpenAI or OpenAI-compatible endpoint."""
if provider not in ("openai", "api"):
raise ValueError(f"unsupported provider '{provider}' for the memory-system methods")
return OpenAIChatClient(model=model, temperature=temperature, max_tokens=max_tokens,
api_key=api_key, base_url=base_url)
class OpenAIEmbedder:
def __init__(self, model: str, batch_size: int | None = None, max_retries: int = MAX_ATTEMPTS,
base_url: str | None = None):
"""Embedding client at `base_url`, else embedding.base_url of backbone.yaml, else the default endpoint."""
import openai
base_url = base_url or ENDPOINTS["embedding"]
self._client = (openai.OpenAI(api_key=api_key_for(base_url), base_url=base_url, **_http_kwargs())
if base_url else openai.OpenAI(api_key=api_key_for(None), **_http_kwargs()))
self.model = model
self.batch_size = batch_size
self.max_retries = max_retries
def embed(self, texts: list[str]) -> list[list[float]]:
if not self.batch_size:
raise RuntimeError("set embedding.batch_size in backbone.yaml")
out: list[list[float]] = []
for i in range(0, len(texts), int(self.batch_size)):
batch = texts[i:i + int(self.batch_size)]
last_err = None
for attempt in range(self.max_retries):
try:
resp = self._client.embeddings.create(model=self.model, input=batch)
u = getattr(resp, "usage", None)
if u is not None:
USAGE.log(self.model, "embed", int(getattr(u, "prompt_tokens", 0) or 0), 0)
out.extend(d.embedding for d in resp.data)
last_err = None
break
except Exception as e:
last_err = e
if "Error code: 400" in str(e) or "invalid_request_error" in str(e):
break
if "429" in str(e).lower():
_note_rate_limit("embed")
time.sleep(1.5 * (attempt + 1))
if last_err is not None:
raise RuntimeError(f"OpenAIEmbedder failed: {last_err}")
return out
def make_answerer(*, max_tokens: int, model: str | None = None, temperature: float | None = None,
base_url: str | None = None):
"""Answering client: temperature and max tokens from the config, JSON response format, no reasoning_effort."""
return ChatClient(model=model, reasoning_effort=None, max_completion_tokens=max_tokens,
temperature=temperature, base_url=base_url)
def make_judge(model: str | None = None):
"""Judge client: JUDGE_MAX_TOKENS output tokens; JUDGE_EFFORT for a reasoning model."""
return ChatClient(model=model, reasoning_effort=JUDGE_EFFORT if _is_reasoning_model(model) else None,
max_completion_tokens=JUDGE_MAX_TOKENS, base_url=ENDPOINTS["judge"])
def make_embedder(model: str | None = None, batch_size: int | None = None):
return OpenAIEmbedder(model=model, batch_size=batch_size)
|