# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. """ LLM client wrapper with per-caller token counting (cortex/CLAUDE.md APIs). This is the shared LLM surface for harnesses outside ``server/``: ``inference.py``, ``baselines/*``, future ``training/train_router.py``, and Cortex subagents (sessions 9+). Per the Q1 decision in root ``CLAUDE.md``, ``r_budget`` is harness-tracked from this module's token counters, never from env state. Design notes: - Backed by the OpenAI Python SDK against an OpenAI-compatible endpoint (HF Router default: ``https://router.huggingface.co/v1``). The HF Router accepts the same chat-completions schema as openai.com. - Token counting reads ``response.usage.{prompt_tokens, completion_tokens}`` only — no local tokenizer fallback. If a provider omits ``usage``, the counter increments by 0 and a one-line warning hits stderr; the caller still gets the response content. - Caller IDs are short colon-separated strings ("inference:t3", "b1:t3", "cortex:epi:planner:t3"), passed explicitly per call. Not thread-local — robust to async / concurrent rollouts. - ``reset_counters`` is harness-driven: harnesses call it at episode boundaries. The client never auto-resets. """ from __future__ import annotations import os import sys from dataclasses import dataclass from typing import Dict, List, Literal, Optional # OpenAI SDK is approved as a prod dep (Session 7a). The HF Router and # OpenAI's own API both speak this protocol; switching providers is a # base_url + api_key change, not a code change. try: from openai import OpenAI as _OpenAI except ImportError: # pragma: no cover - dep listed in pyproject.toml _OpenAI = None # type: ignore[assignment] __all__ = ["LLMClient", "ChatMessage", "ChatResponse"] # ============================================================================ # Defaults — match Session 7b inference.py spec # ============================================================================ DEFAULT_API_BASE_URL = "https://router.huggingface.co/v1" DEFAULT_MODEL = "Qwen/Qwen2.5-72B-Instruct" DEFAULT_TEMPERATURE = 0.0 DEFAULT_MAX_TOKENS = 512 # ============================================================================ # Typed message and response shapes # ============================================================================ @dataclass(frozen=True) class ChatMessage: """One chat-completions message. ``role`` is the OpenAI chat role.""" role: Literal["system", "user", "assistant"] content: str @dataclass class ChatResponse: """Decoded LLM response with token-usage fields surfaced. ``finish_reason`` mirrors the SDK's value (typically ``"stop"`` / ``"length"`` / ``"content_filter"``). Token fields default to 0 if the provider didn't include a ``usage`` block. """ content: str finish_reason: str = "stop" prompt_tokens: int = 0 completion_tokens: int = 0 # ============================================================================ # Client # ============================================================================ class LLMClient: """Per-caller token-counting wrapper around OpenAI chat-completions. Args: api_base_url: Endpoint URL. Falls back to ``$API_BASE_URL`` then ``DEFAULT_API_BASE_URL``. api_key: API key. Falls back to ``$HF_TOKEN`` then ``$OPENAI_API_KEY``. model: Model identifier. Falls back to ``$MODEL_NAME`` then ``DEFAULT_MODEL``. temperature: Sampling temperature. 0.0 for reproducibility. max_tokens: Per-call output cap. client: Pre-built SDK client. Tests inject a stub here; production leaves it ``None`` and the OpenAI SDK is constructed from ``api_base_url`` + ``api_key``. """ def __init__( self, api_base_url: Optional[str] = None, api_key: Optional[str] = None, model: Optional[str] = None, temperature: float = DEFAULT_TEMPERATURE, max_tokens: int = DEFAULT_MAX_TOKENS, client: Optional[object] = None, ) -> None: self.api_base_url = api_base_url or os.getenv( "API_BASE_URL", DEFAULT_API_BASE_URL, ) self.api_key = api_key or os.getenv("HF_TOKEN") or os.getenv("OPENAI_API_KEY") self.model = model or os.getenv("MODEL_NAME", DEFAULT_MODEL) self.temperature = temperature self.max_tokens = max_tokens if client is not None: self._client = client else: if _OpenAI is None: raise RuntimeError( "openai SDK is not installed but no test client was passed. " "Install with `uv sync` (openai>=1.0 is in pyproject.toml)." ) if not self.api_key: raise ValueError( "LLMClient requires an api_key. Set HF_TOKEN or OPENAI_API_KEY, " "or pass api_key=... explicitly." ) self._client = _OpenAI(base_url=self.api_base_url, api_key=self.api_key) self._token_counters: Dict[str, int] = {} # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ def chat( self, caller_id: str, messages: List[ChatMessage], max_tokens: Optional[int] = None, temperature: Optional[float] = None, ) -> ChatResponse: """Call chat-completions; bill prompt+completion tokens to ``caller_id``. Per-call ``max_tokens`` and ``temperature`` overrides are accepted for harnesses that want finer control without constructing a new client. """ completion = self._client.chat.completions.create( model=self.model, messages=[{"role": m.role, "content": m.content} for m in messages], temperature=temperature if temperature is not None else self.temperature, max_tokens=max_tokens if max_tokens is not None else self.max_tokens, stream=False, ) # Defensive extraction — SDK returns rich objects, but tests use # dataclasses with the same attribute shape. choice = completion.choices[0] content = (choice.message.content or "").strip() finish_reason = getattr(choice, "finish_reason", "stop") or "stop" usage = getattr(completion, "usage", None) if usage is None: prompt_tokens = 0 completion_tokens = 0 print( f"[WARN] llm_client: response missing .usage for caller_id={caller_id!r}", file=sys.stderr, flush=True, ) else: prompt_tokens = int(getattr(usage, "prompt_tokens", 0) or 0) completion_tokens = int(getattr(usage, "completion_tokens", 0) or 0) # Cumulative — defaults to 0 for new caller_ids. self._token_counters[caller_id] = ( self._token_counters.get(caller_id, 0) + prompt_tokens + completion_tokens ) return ChatResponse( content=content, finish_reason=finish_reason, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, ) def tokens_used_for(self, caller_id: str) -> int: """Cumulative prompt+completion tokens billed to ``caller_id``. Unknown caller_ids read as 0 (not a KeyError). Harnesses use this to compose ``r_budget`` per design §14.3. """ return self._token_counters.get(caller_id, 0) def reset_counters(self, caller_id_prefix: Optional[str] = None) -> None: """Zero counters whose key starts with ``caller_id_prefix``. With no prefix, clears all counters. Harnesses call this at episode boundaries (B1, inference.py, future training loops). The client never auto-resets — counters are sticky until cleared explicitly. """ if caller_id_prefix is None: self._token_counters.clear() return for key in list(self._token_counters.keys()): if key.startswith(caller_id_prefix): self._token_counters[key] = 0