"""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)