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