"""Concrete OpenAI-compatible streaming client.
The only engine-aware code in the harness, and it stays thin: it adapts the
OpenAI SDK's streaming chunks to the engine-agnostic ``StreamEvent`` the runner
consumes. Works unchanged against LM Studio, mlx_lm.server and mlx-openai-server.
"""
from __future__ import annotations
from collections.abc import Iterator
from dataclasses import dataclass, field
from touchstone.runner import StreamEvent
from touchstone.toolbench import ToolStreamEvent
@dataclass
class BuildMetadata:
engine_version: str | None = None
runtime: str | None = None
quant_by_model: dict[str, str] = field(default_factory=dict)
# Only the models whose state is "loaded", with their quant/context — the evidence of what
# actually answered. quant_by_model covers every listed model, including not-loaded ones.
loaded: list[dict[str, object]] = field(default_factory=list)
def parse_lmstudio_models(payload: dict[str, object]) -> BuildMetadata:
"""Extract runtime + per-model quant from an LM Studio /api/v0/models payload.
Best-effort: unknown shape → empty metadata (callers fall back to config)."""
data = payload.get("data") if isinstance(payload, dict) else None
if not isinstance(data, list):
return BuildMetadata()
quant_by_model: dict[str, str] = {}
runtimes: list[str] = []
loaded: list[dict[str, object]] = []
for m in data:
if not isinstance(m, dict):
continue
mid = m.get("id")
q = m.get("quantization")
if m.get("state") == "loaded":
loaded.append(
{k: m.get(k) for k in ("id", "quantization", "loaded_context_length") if k in m}
)
rt = m.get("compatibility_type")
if isinstance(mid, str) and isinstance(q, str):
quant_by_model[mid] = q
if isinstance(rt, str):
runtimes.append(rt)
return BuildMetadata(
runtime=runtimes[0] if runtimes else None, quant_by_model=quant_by_model, loaded=loaded
)
class OpenAIStreamClient:
engine: str
engine_version: str
def __init__(
self,
base_url: str,
api_key: str = "not-needed",
*,
engine: str = "openai-compat",
engine_version: str = "unknown",
timeout: float | None = None,
max_retries: int | None = None,
) -> None:
from openai import NOT_GIVEN, OpenAI
# Pass timeout=NOT_GIVEN (the SDK's "use default" sentinel) rather than None — None would
# mean an *infinite* wait and silently break every existing eval/judge caller that relies
# on the SDK's 600s read timeout. max_retries is a plain int (no sentinel), so omit it
# entirely unless set; discovery passes max_retries=0 so the 3s timeout is the true bound.
_timeout = NOT_GIVEN if timeout is None else timeout
if max_retries is None:
self._client = OpenAI(base_url=base_url, api_key=api_key, timeout=_timeout)
else:
self._client = OpenAI(
base_url=base_url, api_key=api_key, timeout=_timeout, max_retries=max_retries
)
self.base_url = base_url
self.engine = engine
self.engine_version = engine_version
def probe_build_metadata(self) -> BuildMetadata:
"""Best-effort build/quant probe. LM Studio exposes /api/v0/models; others don't.
Never raises into the eval loop — any failure yields empty metadata. The import belongs
inside the guard: httpx used to arrive with openai, but openai >= 3 depends on httpx2,
so on a fresh install the import itself is what raises — after the run, taking the
report and the manifest with it (measured 2026-09-21)."""
base = self.base_url.rstrip("/")
host = base[:-3] if base.endswith("/v1") else base # strip OpenAI suffix
try:
import httpx
r = httpx.get(f"{host}/api/v0/models", timeout=3.0)
if r.status_code == 200:
return parse_lmstudio_models(r.json())
except Exception:
pass
return BuildMetadata()
def list_models(self) -> list[str]:
"""Model ids the endpoint advertises at /v1/models (for the GUI picker dropdown)."""
return [m.id for m in self._client.models.list().data]
def stream(
self,
*,
messages: list[dict[str, object]],
model: str,
max_tokens: int | None,
temperature: float,
seed: int,
extra_body: dict[str, object] | None = None,
) -> Iterator[StreamEvent]:
# Our engine-agnostic message dicts don't match the SDK's TypedDict overloads;
# the server-side schema is what actually validates them.
# max_tokens=None → omit the cap so the server answers freely (eval default); a positive int
# caps it (latency runner). extra_body forwarded only when set — keeps the call clean.
extra: dict[str, object] = {}
if max_tokens is not None:
extra["max_tokens"] = max_tokens
if extra_body:
extra["extra_body"] = extra_body
stream = self._client.chat.completions.create( # type: ignore[call-overload]
model=model,
messages=messages,
temperature=temperature,
seed=seed,
stream=True,
stream_options={"include_usage": True},
**extra,
)
for chunk in stream:
choices = getattr(chunk, "choices", None) or []
if choices:
delta = getattr(choices[0], "delta", None)
content = getattr(delta, "content", None) if delta else None
# "Thinking" models (LM Studio / mlx / DeepSeek-style) put their reasoning
# on a separate field — capture it instead of dropping it on the floor.
reasoning = None
if delta is not None:
reasoning = getattr(delta, "reasoning_content", None) or getattr(
delta, "reasoning", None
)
if content:
yield StreamEvent(delta_text=content)
if reasoning:
yield StreamEvent(reasoning_text=reasoning)
usage = getattr(chunk, "usage", None)
if usage is not None:
# Prefix-cache hits (vLLM, Splash, LM Studio …) — lets a bundle prove which
# request was served from cache instead of leaving it to inference.
details = getattr(usage, "prompt_tokens_details", None)
cached = getattr(details, "cached_tokens", None) if details is not None else None
yield StreamEvent(
prompt_tokens=getattr(usage, "prompt_tokens", None),
completion_tokens=getattr(usage, "completion_tokens", None),
cached_tokens=cached if isinstance(cached, int) else None,
)
def stream_tools(
self,
*,
messages: list[dict[str, object]],
model: str,
tools: list[dict[str, object]],
max_tokens: int | None,
temperature: float | None,
seed: int | None,
extra_body: dict[str, object] | None = None,
) -> Iterator[ToolStreamEvent]:
"""Stream a tool-enabled turn as raw fragments (``toolbench.collect_turn`` assembles).
Streams like opencode does, so a server that ships a truncated call with empty
``arguments`` at the budget edge shows up here exactly as it does there.
``temperature``/``seed`` of ``None`` are omitted from the request entirely rather than
sent as 0.0/some int — needed to replicate a client that never sets them (e.g. opencode
without `temperature: true` on the model entry), so the server's own default sampling
applies instead of our forced values."""
extra: dict[str, object] = {}
if max_tokens is not None:
extra["max_tokens"] = max_tokens
if temperature is not None:
extra["temperature"] = temperature
if seed is not None:
extra["seed"] = seed
if extra_body:
extra["extra_body"] = extra_body
stream = self._client.chat.completions.create( # type: ignore[call-overload]
model=model,
messages=messages,
tools=tools,
stream=True,
stream_options={"include_usage": True},
**extra,
)
for chunk in stream:
choices = getattr(chunk, "choices", None) or []
if choices:
ch = choices[0]
delta = getattr(ch, "delta", None)
if delta is not None:
content = getattr(delta, "content", None)
reasoning = getattr(delta, "reasoning_content", None) or getattr(
delta, "reasoning", None
)
if content or reasoning:
yield ToolStreamEvent(content=content or None, reasoning=reasoning or None)
for tc in getattr(delta, "tool_calls", None) or []:
fn = getattr(tc, "function", None)
yield ToolStreamEvent(
is_tool=True,
tc_index=getattr(tc, "index", None),
tc_id=getattr(tc, "id", None),
tc_name=getattr(fn, "name", None) if fn else None,
tc_args=getattr(fn, "arguments", None) if fn else None,
)
fr = getattr(ch, "finish_reason", None)
if fr:
yield ToolStreamEvent(finish_reason=fr)
usage = getattr(chunk, "usage", None)
if usage is not None:
details = getattr(usage, "completion_tokens_details", None)
yield ToolStreamEvent(
prompt_tokens=getattr(usage, "prompt_tokens", None),
completion_tokens=getattr(usage, "completion_tokens", None),
reasoning_tokens=getattr(details, "reasoning_tokens", None)
if details
else None,
)
def embed(self, *, model: str, inputs: list[str]) -> int:
"""Embed a batch; returns the number of vectors produced."""
resp = self._client.embeddings.create(model=model, input=inputs)
return len(resp.data)