From bd92ec1c1947be0e3969579c13404b33eb70c123 Mon Sep 17 00:00:00 2001 From: Naftali Goldstein Date: Tue, 25 Aug 2026 12:21:04 +0300 Subject: [PATCH 1/7] support openrouter --- composer/input/files.py | 10 +- composer/llm/anthropic.py | 18 +- composer/llm/openai.py | 22 +- composer/llm/openrouter.py | 464 +++++++++++++++++++++++++++++++++++++ composer/llm/provider.py | 35 ++- pyproject.toml | 1 + 6 files changed, 523 insertions(+), 27 deletions(-) create mode 100644 composer/llm/openrouter.py diff --git a/composer/input/files.py b/composer/input/files.py index ce79c486..ef1384e8 100644 --- a/composer/input/files.py +++ b/composer/input/files.py @@ -34,7 +34,11 @@ class ContentRenderer(Protocol): def text_block(self, text: str, *, cache_level: CacheLevel = CacheLevel.NONE) -> dict: ... - def file_block(self, file_id: str, *, cache_level: CacheLevel = CacheLevel.NONE) -> dict: ... + # ``filename`` is redundant for the Files-API providers, which reference an + # upload by id alone, but is required by a renderer that inlines the bytes + # as a data URL (OpenRouter): the content block carries no id to look a name + # up by, and OpenAI-compatible inline file blocks require ``filename``. + def file_block(self, file_id: str, *, filename: str, cache_level: CacheLevel = CacheLevel.NONE) -> dict: ... # --------------------------------------------------------------------------- # Protocols (the public surface) @@ -224,7 +228,9 @@ class UploadedFile: renderer: ContentRenderer def to_dict(self, cache_level: CacheLevel = CacheLevel.NONE) -> dict: - return self.renderer.file_block(file_id=self.file_id, cache_level=cache_level) + return self.renderer.file_block( + file_id=self.file_id, filename=self.basename, cache_level=cache_level + ) def to_digest(self) -> str: return self.digest diff --git a/composer/llm/anthropic.py b/composer/llm/anthropic.py index b1101fcf..4ab192df 100644 --- a/composer/llm/anthropic.py +++ b/composer/llm/anthropic.py @@ -12,7 +12,7 @@ from composer.input.files import UploaderBase, ContentRenderer from composer.input.types import ModelConfiguration from composer.llm.provider import ( - ProviderServiceBase, ProviderSpec, compaction_threshold + ProviderServiceBase, ProviderSpec, compaction_threshold, standard_callbacks ) from composer.llm.pricing import PriceProvider, price_provider_for from .types import CacheLevel @@ -137,7 +137,10 @@ def text_block(self, text: str, *, cache_level: CacheLevel = CacheLevel.NONE) -> } return to_ret - def file_block(self, file_id: str, *, cache_level: CacheLevel = CacheLevel.NONE) -> dict: + def file_block( + self, file_id: str, *, filename: str, cache_level: CacheLevel = CacheLevel.NONE + ) -> dict: + # filename unused: the Files API upload already carries it. to_ret : dict[str, Any] = { "type": "document", "source": { @@ -263,8 +266,6 @@ def builder_for( self, *, cache_level: CacheLevel = CacheLevel.NONE, disable_thinking: bool = False ) -> "BaseChatModel": from langchain_anthropic import ChatAnthropic - from composer.diagnostics.usage_callback import UsageCallback - from composer.diagnostics.cost_callback import CostAccumulator opts = self.options thinking: dict[str, Any] | None @@ -299,12 +300,9 @@ def builder_for( betas=betas, thinking=thinking, model_kwargs=model_kwargs, - callbacks=[ - UsageCallback(), - CostAccumulator( - self.price_provider, long_cache=cache_level == CacheLevel.LONG - ), - ], + callbacks=standard_callbacks( + self.price_provider, long_cache=cache_level == CacheLevel.LONG + ), ) ANTHROPIC_SPEC = ProviderSpec( diff --git a/composer/llm/openai.py b/composer/llm/openai.py index ae7619f3..b7d3766d 100644 --- a/composer/llm/openai.py +++ b/composer/llm/openai.py @@ -22,7 +22,8 @@ from composer.input.files import UploaderBase, ContentRenderer from composer.input.types import ModelConfiguration from .provider import ( - ProviderServiceBase, ProviderSpec, compaction_threshold + ProviderServiceBase, ProviderSpec, compaction_threshold, reasoning_effort, + standard_callbacks ) from .pricing import PriceProvider, price_provider_for from .types import CacheLevel @@ -135,14 +136,6 @@ def _context_window(features: OpenAIModelFeatures) -> int: return _assumed_context_window -def _reasoning_effort(thinking_tokens: int) -> Literal["low", "medium", "high"]: - """Map a thinking-token budget onto OpenAI's three-step effort knob.""" - if thinking_tokens <= 2048: - return "low" - if thinking_tokens <= 8192: - return "medium" - return "high" - class OpenAIService(ProviderServiceBase): def __init__(self): from graphcore.tools.memory import openai_async_memory_tool @@ -168,7 +161,10 @@ class OpenAIRenderer: def text_block(self, text: str, *, cache_level: CacheLevel = CacheLevel.NONE) -> dict: to_ret : dict[str, Any] = {"type": "text", "text": text} return to_ret - def file_block(self, file_id: str, *, cache_level: CacheLevel = CacheLevel.NONE) -> dict: + def file_block( + self, file_id: str, *, filename: str, cache_level: CacheLevel = CacheLevel.NONE + ) -> dict: + # filename unused: the Files API upload already carries it. return { "type": "file", "file": { @@ -250,8 +246,6 @@ def builder_for( self, *, cache_level: CacheLevel = CacheLevel.NONE, disable_thinking: bool = False ) -> "BaseChatModel": from langchain_openai import ChatOpenAI - from composer.diagnostics.usage_callback import UsageCallback - from composer.diagnostics.cost_callback import CostAccumulator opts = self.options kwargs: dict[str, Any] = { @@ -262,7 +256,7 @@ def builder_for( if opts.thinking_tokens is not None and not disable_thinking and self.features.reasoning: kwargs["reasoning"] = { - "effort": _reasoning_effort(opts.thinking_tokens), + "effort": reasoning_effort(opts.thinking_tokens), "summary": "auto" } @@ -273,7 +267,7 @@ def builder_for( max_retries=2, # OpenAI has no cache-TTL knob, so long_cache stays False; cache_write_1h # mirrors cache_write in the table anyway. - callbacks=[UsageCallback(), CostAccumulator(self.price_provider)], + callbacks=standard_callbacks(self.price_provider), **kwargs, ) diff --git a/composer/llm/openrouter.py b/composer/llm/openrouter.py new file mode 100644 index 00000000..39a7bb6f --- /dev/null +++ b/composer/llm/openrouter.py @@ -0,0 +1,464 @@ +"""OpenRouter LLM backend: metadata probing, an inlining "uploader", and the +``ModelProvider`` that mints ``ChatOpenAI`` instances pointed at OpenRouter's +OpenAI-compatible API. Its own backend rather than a flag on ``openai.py`` because +little carries over: vendor-qualified ids, no Files API, a different request shape. + +Three things differ from the OpenAI backend: + +* **Probing is live, not parsed.** OpenRouter fronts 400+ models whose names encode + nothing and whose roster turns over weekly, so ``GET /api/v1/models`` supplies the + window, output cap, reasoning support and price card instead of a name parser and + a hand-maintained table. See :func:`_catalog`. +* **Files are inlined.** There is no Files API — see :class:`InlineFileUploader`. +* **Every route goes through the Responses API**, not Chat Completions — see + ``builder_for``. + +Requires ``OPENROUTER_API_KEY``. +""" +from typing import Any, TYPE_CHECKING, Mapping, AsyncIterator, override +from dataclasses import dataclass, field +from functools import cache +import asyncio +import base64 +import json +import logging +import os +import urllib.request + +import httpx +import openai + +from composer.input.files import UploaderBase, ContentRenderer +from composer.input.types import ModelConfiguration +from .provider import ( + ProviderServiceBase, ProviderSpec, compaction_threshold, reasoning_effort, + standard_callbacks +) +from .openai import OpenAIRenderer +from .pricing import PriceProvider, PriceTier, price_provider_for +from .types import CacheLevel + +if TYPE_CHECKING: + from langchain_core.language_models.chat_models import BaseChatModel + from langchain_core.outputs import ChatGenerationChunk + from langchain_openai import ChatOpenAI + +logger = logging.getLogger(__name__) + + +BASE_URL = "https://openrouter.ai/api/v1" + +_MODELS_URL = f"{BASE_URL}/models" +_API_KEY_ENV = "OPENROUTER_API_KEY" + +_PROBE_TIMEOUT_SECONDS = 15 + +# The only fields read off a roster record. +_PROBED_KEYS = ("context_length", "top_provider", "supported_parameters", "pricing") + +# Transient stream failures and how many times to re-issue the request. The SDK's +# own `max_retries` covers *establishing* a request; these surface while the SSE +# body is being consumed, by which point the request has already succeeded, so +# nothing below us retries them and they take the whole run down. `TimeoutError` +# covers langchain's `StreamChunkTimeoutError` (a subclass) for a content stall; +# `httpx.TransportError` covers the dropped-connection family. +_RETRYABLE_STREAM_ERRORS = (httpx.TransportError, TimeoutError, openai.APIConnectionError) +_STREAM_ATTEMPTS = 3 +_STREAM_RETRY_BACKOFF_SECONDS = 2.0 + + +def matches(model: str) -> bool: + """OpenRouter ids are vendor-qualified (``moonshotai/kimi-k2.5``, + ``openrouter/auto``), and the slash is what separates them from a native name. No + overlap with the other predicates: those split on "-", so + ``anthropic/claude-sonnet-5`` heads at ``anthropic/claude``, never ``claude``.""" + return "/" in model + + +# --- live model metadata --------------------------------------------------- + +# What to assume when the roster fetch failed. The window understates +# nearly every current model, which costs earlier compaction rather than a failed +# request; overstating it would hard-fail mid-run on a context-length error. +_FALLBACK_CONTEXT_WINDOW = 128_000 + + +@dataclass(frozen=True) +class OpenRouterModelFeatures: + """What the request shape needs to know about one route.""" + + context_window: int + # Output-token ceiling the route advertises, or None if it publishes none. + max_output_tokens: int | None + # Route accepts the ``reasoning`` knob at all. Assumed True when the fetch + # failed: OpenRouter drops a parameter the route doesn't support, so guessing + # "reasoning" wrongly costs nothing, while guessing "no reasoning" wrongly + # disables thinking silently. + reasoning: bool + + +@cache +def _catalog() -> dict[str, Mapping[str, Any]]: + """OpenRouter's model roster, by id: one blocking unauthenticated GET, no retry, + from :meth:`OpenRouterModelProvider.create` at startup. Any failure yields an + empty catalog and a run on conservative defaults.""" + try: + with urllib.request.urlopen(_MODELS_URL, timeout=_PROBE_TIMEOUT_SECONDS) as resp: + payload = json.load(resp) + except (OSError, json.JSONDecodeError) as exc: + logger.warning( + "Could not fetch OpenRouter model metadata from %s (%s); falling back to " + "a %d-token context window and no price card.", + _MODELS_URL, exc, _FALLBACK_CONTEXT_WINDOW, + ) + return {} + records = payload.get("data") if isinstance(payload, dict) else None + if not isinstance(records, list): + logger.warning("Unexpected OpenRouter /models payload shape; ignoring it.") + return {} + # Projected to what's actually read: the full roster is ~690KB on the wire and + # ~2.3MB of retained objects, nearly all of it prose this module never touches. + return { + rec["id"]: {k: rec[k] for k in _PROBED_KEYS if k in rec} + for rec in records + if isinstance(rec, dict) and isinstance(rec.get("id"), str) + } + + +def _record_for(model_name: str) -> Mapping[str, Any] | None: + catalog = _catalog() + if (exact := catalog.get(model_name)) is not None: + return exact + # Variant suffixes (`:free`, `:nitro`, `:floor`) select a routing policy, not a + # different model; most are absent from the roster under their suffixed id. + base, _, variant = model_name.partition(":") + return catalog.get(base) if variant else None + + +def _positive_int(value: Any) -> int | None: + return value if isinstance(value, int) and not isinstance(value, bool) and value > 0 else None + + +def _features_from(record: Mapping[str, Any] | None) -> OpenRouterModelFeatures: + if record is None: + # Only reachable when the fetch itself failed — an id absent from a roster + # that did load is rejected in `create`. + return OpenRouterModelFeatures( + context_window=_FALLBACK_CONTEXT_WINDOW, + max_output_tokens=None, + reasoning=True, + ) + top = record.get("top_provider") + top = top if isinstance(top, Mapping) else {} + supported = record.get("supported_parameters") + supported = set(supported) if isinstance(supported, list) else set() + return OpenRouterModelFeatures( + # Two windows are published: the model's own and the serving provider's. The + # smaller is the one a request actually has to fit in. + context_window=min( + ( + w for w in ( + _positive_int(record.get("context_length")), + _positive_int(top.get("context_length")), + ) if w is not None + ), + default=_FALLBACK_CONTEXT_WINDOW, + ), + max_output_tokens=_positive_int(top.get("max_completion_tokens")), + reasoning="reasoning" in supported, + ) + + +# --- pricing --------------------------------------------------------------- + +# OpenRouter quotes USD per token; PriceTier is USD per million. +_TOKENS_PER_MTOK = 1_000_000 + + +def _per_mtok(raw: Any) -> float | None: + """One price field, converted to per-MTok. A "0" is a real zero (a free route); + an absent field is a missing key, which comes back as None.""" + if not isinstance(raw, (str, int, float)) or isinstance(raw, bool): + return None + try: + return float(raw) * _TOKENS_PER_MTOK + except ValueError: + return None + + +def _price_tier(prices: Mapping[str, Any]) -> PriceTier | None: + prompt = _per_mtok(prices.get("prompt")) + completion = _per_mtok(prices.get("completion")) + if prompt is None or completion is None: + return None + # A route that publishes no cache bucket bills those tokens as fresh input, and + # one with no separate 1h rate charges the 5m one. `or` would not do here: a + # real 0.0 is a free route. + if (cache_read := _per_mtok(prices.get("input_cache_read"))) is None: + cache_read = prompt + if (cache_write := _per_mtok(prices.get("input_cache_write"))) is None: + cache_write = prompt + if (cache_write_1h := _per_mtok(prices.get("input_cache_write_1h"))) is None: + cache_write_1h = cache_write + return PriceTier( + input=prompt, + output=completion, + cache_read=cache_read, + cache_write=cache_write, + cache_write_1h=cache_write_1h, + ) + + +def _bare_model_name(model_name: str) -> str: + """The vendor's own name for a route, which is what the static price table in + ``composer.llm.pricing`` is keyed by: ``openai/gpt-5.5:floor`` -> ``gpt-5.5``.""" + _, _, rest = model_name.partition("/") + return rest.partition(":")[0] + + +def _price_provider_from( + record: Mapping[str, Any] | None, model_name: str +) -> PriceProvider: + """The route's pricing curve, live from the catalog where possible, else the + static table on the vendor's bare model name — which covers ``openai/*`` and + ``anthropic/*``, and yields None (an uncosted run, not a wrong one) elsewhere.""" + prices = (record or {}).get("pricing") + if not isinstance(prices, Mapping) or (short := _price_tier(prices)) is None: + return price_provider_for(_bare_model_name(model_name)) + + # Long-context surcharges arrive as overrides keyed by a prompt-size floor, each + # restating only the fields it changes (see openai/gpt-5.5's >272K tier). + raw_overrides = prices.get("overrides") + tiers: list[tuple[int, PriceTier]] = [] + if isinstance(raw_overrides, list): + for override in raw_overrides: + if not isinstance(override, Mapping): + continue + floor = _positive_int(override.get("min_prompt_tokens")) + tier = _price_tier({**prices, **override}) + if floor is not None and tier is not None: + tiers.append((floor, tier)) + tiers.sort(key=lambda t: t[0], reverse=True) + + def provider(input_tokens: int) -> PriceTier | None: + for floor, tier in tiers: + if input_tokens > floor: + return tier + return short + + return provider + + +# --- request shape --------------------------------------------------------- + +def _api_key() -> str: + if not (key := os.environ.get(_API_KEY_ENV)): + raise ValueError( + f"{_API_KEY_ENV} is not set, and an OpenRouter model was requested. " + f"Get a key at https://openrouter.ai/keys." + ) + return key + + +@dataclass +class OpenRouterRenderer(OpenAIRenderer): + """Content blocks in the Chat-Completions shape, which is what langchain converts + from: ``_convert_chat_completions_blocks_to_responses`` turns ``file`` into + ``input_file`` and ``image_url`` into ``input_image`` on the way out. Only the + file block differs from OpenAI's; the text block is inherited.""" + + @override + def file_block( + self, file_id: str, *, filename: str, cache_level: CacheLevel = CacheLevel.NONE + ) -> dict: + # `file_id` is a `data:` URL rather than a remote id: OpenRouter has no + # Files API, so InlineFileUploader carries the bytes here instead. + if file_id.startswith("data:image/"): + return {"type": "image_url", "image_url": {"url": file_id}} + return {"type": "file", "file": {"filename": filename, "file_data": file_id}} + + +@dataclass +class InlineFileUploader(UploaderBase): + """``FileUploader`` impl for a provider with no Files API: the "upload" is a + ``data:`` URL built in memory, which the renderer inlines into the request. So a + large binary is re-sent with every request carrying it, and a PDF no route reads + natively goes through OpenRouter's ``file-parser`` plugin, which falls back to + ``mistral-ocr`` at $2/1K pages (pin an engine via ``plugins`` to avoid that).""" + + renderer: ContentRenderer = field(default_factory=OpenRouterRenderer) + + async def _upload_bytes( + self, crc_basename: str, file_data: bytes, mime: str + ) -> str: + # No dedup cache: the "id" *is* the content, so there is nothing to reuse. + # Off-thread because the encode is ~1.5ms/MB of blocked loop, matching how + # `composer.input.files` already offloads the read. + encoded = await asyncio.to_thread(base64.b64encode, file_data) + return f"data:{mime};base64,{encoded.decode('ascii')}" + + +class OpenRouterService(ProviderServiceBase): + """Provider services for OpenRouter. ``cache_marker`` stays the base no-op: + OpenRouter caches prompts itself on the Responses API, so an explicit + ``cache_control`` breakpoint would add nothing.""" + + def __init__(self): + from graphcore.tools.memory import openai_async_memory_tool + super().__init__( + # The OpenAI-flavored memory tool is a plain client-side function tool + # differing only in its args schema, so it is portable to any route. + openai_async_memory_tool, + InlineFileUploader, + ) + + +@cache +def _openrouter_service(): + return OpenRouterService() + + +# --- chat model ------------------------------------------------------------ + +@cache +def _chat_model_cls() -> type["ChatOpenAI"]: + """``ChatOpenAI`` that re-issues a streamed request when the stream breaks. + + Defined behind a cached factory so ``langchain_openai`` stays a lazy import, + as it is in the sibling backends.""" + from langchain_openai import ChatOpenAI + + class RetryingChatOpenAI(ChatOpenAI): + @override + async def _astream( + self, *args: Any, **kwargs: Any + ) -> AsyncIterator["ChatGenerationChunk"]: + for attempt in range(1, _STREAM_ATTEMPTS + 1): + # Buffered, not forwarded as they arrive: a retry must not emit a + # partial response twice. Callers aggregate through `ainvoke` + # anyway, so this costs nothing today — an incremental consumer + # would lose its incrementality. + chunks: list["ChatGenerationChunk"] = [] + try: + async for chunk in super()._astream(*args, **kwargs): + chunks.append(chunk) + except _RETRYABLE_STREAM_ERRORS as exc: + if attempt == _STREAM_ATTEMPTS: + raise + delay = _STREAM_RETRY_BACKOFF_SECONDS * attempt + logger.warning( + "OpenRouter stream failed after %d chunk(s) (%s: %s); " + "re-issuing in %.0fs (attempt %d/%d).", + len(chunks), type(exc).__name__, exc, delay, + attempt + 1, _STREAM_ATTEMPTS, + ) + await asyncio.sleep(delay) + continue + for chunk in chunks: + yield chunk + return + + return RetryingChatOpenAI + + +# --- ModelProvider --------------------------------------------------------- + +@dataclass +class OpenRouterModelProvider: + """``ModelProvider`` for OpenRouter. Probes the route's metadata once at + construction and shapes the request from it. ``cache_level`` picks the + cache-write rate for costing only — OpenRouter has no cache-TTL knob.""" + + model_name: str + options: ModelConfiguration + features: OpenRouterModelFeatures + price_provider: PriceProvider + api_key: str + provider: OpenRouterService = field(default_factory=_openrouter_service) + + @staticmethod + def create(model_name: str, options: ModelConfiguration) -> "OpenRouterModelProvider": + # Before the probe: a run with no key can't start, so it shouldn't pay for a + # round-trip first. Reading it here rather than in `builder_for` is what + # makes that failure a startup one instead of a first-LLM-call one. + api_key = _api_key() + record = _record_for(model_name) + if record is None and _catalog(): + # The roster loaded and this route isn't on it. + raise ValueError( + f"{model_name!r} is not an OpenRouter model; see " + f"https://openrouter.ai/models for the roster." + ) + return OpenRouterModelProvider( + model_name=model_name, + options=options, + features=_features_from(record), + price_provider=_price_provider_from(record, model_name), + api_key=api_key, + ) + + @property + def max_prompt_tokens(self) -> int: + return compaction_threshold(self.features.context_window) + + def _output_token_cap(self) -> int: + """The response budget to ask for, clamped to what the route allows — an + ``opts.tokens`` above the route's ceiling is a 400, not a truncation.""" + requested = self.options.tokens + ceiling = self.features.max_output_tokens + return requested if ceiling is None else min(requested, ceiling) + + def builder_for( + self, *, cache_level: CacheLevel = CacheLevel.NONE, disable_thinking: bool = False + ) -> "BaseChatModel": + from pydantic import SecretStr + + opts = self.options + kwargs: dict[str, Any] = {} + + if opts.thinking_tokens is not None and not disable_thinking and self.features.reasoning: + # OpenRouter's unified reasoning knob, which it translates per route: + # passed straight through to the families whose native knob is also an + # effort level, converted to a token budget for the rest. + kwargs["reasoning"] = { + # OpenRouter converts effort to a token budget for the vendors + # whose native knob is one. + "effort": reasoning_effort(opts.thinking_tokens), + "summary": "auto", + } + # Ask for the encrypted chain of thought, which is what langchain echoes + # back with the next tool result so the model can resume it. + kwargs["include"] = ["reasoning.encrypted_content"] + + return _chat_model_cls()( + model=self.model_name, + base_url=BASE_URL, + api_key=SecretStr(self.api_key), + # Load-bearing, not a default: the Responses API is the only surface on + # which reasoning survives a tool round-trip, because langchain echoes a + # prior turn's reasoning items back into the next request. Chat + # Completions drops them, so a long tool loop re-derives its reasoning + # every round at the output token rate. + use_responses_api=True, + # Unstreamed, OpenRouter holds the whole generation and an upstream + # pause trips its gateway idle timeout, failing the request. + streaming=True, + # OpenRouter is stateless: store=True is a 400, not a no-op. + store=False, + # langchain renames this to `max_output_tokens` for the Responses API. + max_completion_tokens=self._output_token_cap(), + timeout=None, + max_retries=2, + # Names the run in OpenRouter's activity dashboard. + default_headers={"X-Title": "AutoProver"}, + callbacks=standard_callbacks( + self.price_provider, long_cache=cache_level == CacheLevel.LONG + ), + **kwargs, + ) + + +OPENROUTER_SPEC = ProviderSpec( + matches=matches, + build=OpenRouterModelProvider.create +) diff --git a/composer/llm/provider.py b/composer/llm/provider.py index a77a110c..87913426 100644 --- a/composer/llm/provider.py +++ b/composer/llm/provider.py @@ -9,16 +9,18 @@ the per-provider modules can import it without an import cycle. """ -from typing import Protocol, TYPE_CHECKING, Callable +from typing import Protocol, TYPE_CHECKING, Callable, Literal from dataclasses import dataclass from functools import cached_property from composer.input.files import FileUploader from composer.input.types import ModelConfiguration +from .pricing import PriceProvider from .types import CacheLevel from abc import ABC, abstractmethod if TYPE_CHECKING: + from langchain_core.callbacks import BaseCallbackHandler from langchain_core.language_models.chat_models import BaseChatModel from graphcore.tools.memory import AsyncPostgresBackend from graphcore.graph import RawMessageType @@ -110,3 +112,34 @@ class ProviderSpec: def compaction_threshold(context_window: int) -> int: """The prompt-token budget to allow a model with a ``context_window``-token window.""" return int(context_window * _PROMPT_TOKEN_SHARE) + + +# Thresholds for the three-step effort knob the OpenAI-shaped providers take in +# place of a thinking-token count. Shared so a retune can't reach one backend and +# miss another — the same budget has to mean the same effort on every route. +_MEDIUM_EFFORT_FROM_TOKENS = 2048 +_HIGH_EFFORT_FROM_TOKENS = 8192 + + +def reasoning_effort(thinking_tokens: int) -> Literal["low", "medium", "high"]: + """Map a thinking-token budget onto the three-step effort knob.""" + if thinking_tokens <= _MEDIUM_EFFORT_FROM_TOKENS: + return "low" + if thinking_tokens <= _HIGH_EFFORT_FROM_TOKENS: + return "medium" + return "high" + + +def standard_callbacks( + price_provider: PriceProvider, *, long_cache: bool = False +) -> list["BaseCallbackHandler"]: + """The usage + cost instrumentation to attach at model construction. + + Shared because a provider that forgets either one produces an unmetered or + uncosted run rather than an error. ``long_cache`` selects the 1-hour + cache-write rate; providers with no cache-TTL knob leave it False. Imported + lazily to keep this module a dependency-free leaf.""" + from composer.diagnostics.usage_callback import UsageCallback + from composer.diagnostics.cost_callback import CostAccumulator + + return [UsageCallback(), CostAccumulator(price_provider, long_cache=long_cache)] diff --git a/pyproject.toml b/pyproject.toml index 4fab6c6b..83449bc7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -244,3 +244,4 @@ certora_autosetup = [ [project.entry-points."certora.autoprove.llm_provider"] openai = "composer.llm.openai:OPEN_AI_SPEC" anthropic = "composer.llm.anthropic:ANTHROPIC_SPEC" +openrouter = "composer.llm.openrouter:OPENROUTER_SPEC" From a25de33da0f2a7928cad776cc2798b481d0fd114 Mon Sep 17 00:00:00 2001 From: Naftali Goldstein Date: Sun, 30 Aug 2026 11:54:31 +0300 Subject: [PATCH 2/7] John's CR --- composer/input/files.py | 141 ++++++++------ composer/llm/anthropic.py | 11 +- composer/llm/openai.py | 11 +- composer/llm/openrouter.py | 367 ++++++++++++++++++++----------------- 4 files changed, 306 insertions(+), 224 deletions(-) diff --git a/composer/input/files.py b/composer/input/files.py index ef1384e8..8e7e14f6 100644 --- a/composer/input/files.py +++ b/composer/input/files.py @@ -34,11 +34,21 @@ class ContentRenderer(Protocol): def text_block(self, text: str, *, cache_level: CacheLevel = CacheLevel.NONE) -> dict: ... - # ``filename`` is redundant for the Files-API providers, which reference an - # upload by id alone, but is required by a renderer that inlines the bytes - # as a data URL (OpenRouter): the content block carries no id to look a name - # up by, and OpenAI-compatible inline file blocks require ``filename``. - def file_block(self, file_id: str, *, filename: str, cache_level: CacheLevel = CacheLevel.NONE) -> dict: ... + + def file_block(self, file_id: str, *, cache_level: CacheLevel = CacheLevel.NONE) -> dict: + """Reference a Files-API upload by id. Providers without a Files API + raise ``NotImplementedError`` and produce + :class:`InMemoryBytesFile` instead.""" + ... + + def inline_file_block( + self, basename: str, contents: bytes, mime: str, + *, cache_level: CacheLevel = CacheLevel.NONE + ) -> dict: + """Carry the bytes in the request itself, for providers with no Files + API. The mirror of :meth:`file_block`; a provider implements one or the + other, matching what its uploader produces.""" + ... # --------------------------------------------------------------------------- # Protocols (the public surface) @@ -211,6 +221,38 @@ def to_digest(self) -> str: return _bytes_digest(self.bytes_contents) +@dataclass(frozen=True) +class InMemoryBytesFile: + """Binary content carried inline in the request, for a provider with no + Files API. The binary analogue of :class:`InMemoryTextFile`: the bytes ride + along in every request that carries the document rather than being uploaded + once and referenced by id.""" + + basename: str + contents: bytes + mime: str + renderer: ContentRenderer + + def to_dict(self, cache_level: CacheLevel = CacheLevel.NONE) -> dict: + return self.renderer.inline_file_block( + self.basename, self.contents, self.mime, cache_level=cache_level + ) + + def to_digest(self) -> str: + return _bytes_digest(self.contents) + + @property + def bytes_contents(self) -> bytes: + return self.contents + + @property + def string_contents(self) -> str | None: + try: + return self.contents.decode("utf-8") + except UnicodeDecodeError: + return None + + @dataclass(frozen=True) class UploadedFile: """A (potentially-binary) file uploaded to the Files API. Bytes are @@ -228,9 +270,7 @@ class UploadedFile: renderer: ContentRenderer def to_dict(self, cache_level: CacheLevel = CacheLevel.NONE) -> dict: - return self.renderer.file_block( - file_id=self.file_id, filename=self.basename, cache_level=cache_level - ) + return self.renderer.file_block(file_id=self.file_id, cache_level=cache_level) def to_digest(self) -> str: return self.digest @@ -265,7 +305,7 @@ def string_contents(self) -> str: # --------------------------------------------------------------------------- @dataclass -class _FileData: +class FileData: basename: str raw_data: bytes is_binary: bool @@ -277,28 +317,28 @@ class _FileData: async def _file_data( *, path: str | pathlib.Path -) -> _FileData: +) -> FileData: ... @overload async def _file_data( *, basename: str, data: bytes -) -> _FileData: +) -> FileData: ... async def _file_data( path: str | pathlib.Path | None = None, basename: str | None = None, data: bytes | None = None -) -> _FileData: +) -> FileData: return await asyncio.to_thread(_file_data_impl, path, basename, data) def _file_data_impl( path: str | pathlib.Path | None, basename: str | None, data: bytes | None -) -> _FileData: +) -> FileData: if path is not None: if isinstance(path, str): path = pathlib.Path(path) @@ -321,18 +361,18 @@ def _file_data_impl( mime = "application/octet-stream" if is_binary else "text/plain" crc = hex(zlib.crc32(data)) digest = _bytes_digest(data) - return _FileData(raw_data=data, is_binary=is_binary, mime=mime, crc_basename=f"{crc}_{basename}", digest=digest, basename=basename) + return FileData(raw_data=data, is_binary=is_binary, mime=mime, crc_basename=f"{crc}_{basename}", digest=digest, basename=basename) class FileUploader(Protocol): """Upload+dedup contract. Obtain via ``ModelProvider.uploader()`` (``composer.llm``).""" async def upload_file_if_needed( self, file_path: str | pathlib.Path - ) -> UploadedFile: ... + ) -> Document: ... async def upload_text_file_if_needed( self, file_path: str | pathlib.Path - ) -> UploadedTextFile: ... + ) -> TextDocument: ... async def get_document( self, path: str | pathlib.Path @@ -356,7 +396,11 @@ class UploaderBase(ABC): The dedup cache lives in ``self.uploaded`` (CRC-prefixed filename → remote file id) and is seeded by each subclass's ``fresh`` factory so we don't reupload a file whose bytes the account has already - seen.""" + seen. + + A provider with no Files API overrides :meth:`_binary_document` and + :meth:`_text_upload_document` to return inline shapes instead, and never + implements ``_upload_bytes``.""" renderer: ContentRenderer @@ -366,14 +410,9 @@ async def _upload_bytes( ) -> str: ... - async def upload_file_if_needed( - self, file_path: str | pathlib.Path - ) -> UploadedFile: - """Upload ``file_path`` (or reuse cached upload). Intended for - binary inputs — callers that know they have text should prefer - :meth:`get_document` (default text-inline) or - :meth:`upload_text_file_if_needed` (explicit upload of text).""" - data = await _file_data(path=file_path) + async def _binary_document(self, data: FileData) -> Document: + """How this provider represents binary content. Uploads by default; + override to inline the bytes instead.""" file_id = await self._upload_bytes(data.crc_basename, data.raw_data, data.mime) return UploadedFile( file_id=file_id, @@ -383,15 +422,9 @@ async def upload_file_if_needed( renderer=self.renderer, ) - async def upload_text_file_if_needed( - self, file_path: str | pathlib.Path - ) -> UploadedTextFile: - """Upload ``file_path`` and tag the result as text. Use for - very-large text inputs that would otherwise blow the prompt - budget if inlined; ordinary text should go through - :meth:`get_document`, which keeps the content in-prompt for - transcript debuggability.""" - data = await _file_data(path=file_path) + async def _text_upload_document(self, data: FileData) -> TextDocument: + """How this provider represents text explicitly destined for upload. + Uploads by default; override to keep it in the prompt instead.""" file_id = await self._upload_bytes(data.crc_basename, data.raw_data, data.mime) return UploadedTextFile( file_id=file_id, @@ -401,6 +434,25 @@ async def upload_text_file_if_needed( renderer=self.renderer, ) + async def upload_file_if_needed( + self, file_path: str | pathlib.Path + ) -> Document: + """Upload ``file_path`` (or reuse cached upload). Intended for + binary inputs — callers that know they have text should prefer + :meth:`get_document` (default text-inline) or + :meth:`upload_text_file_if_needed` (explicit upload of text).""" + return await self._binary_document(await _file_data(path=file_path)) + + async def upload_text_file_if_needed( + self, file_path: str | pathlib.Path + ) -> TextDocument: + """Upload ``file_path`` and tag the result as text. Use for + very-large text inputs that would otherwise blow the prompt + budget if inlined; ordinary text should go through + :meth:`get_document`, which keeps the content in-prompt for + transcript debuggability.""" + return await self._text_upload_document(await _file_data(path=file_path)) + async def get_document( self, path: str | pathlib.Path ) -> Document | None: @@ -419,14 +471,7 @@ async def get_document( return None data = await _file_data(path=p) if data.is_binary: - file_id = await self._upload_bytes(data.crc_basename, data.raw_data, data.mime) - return UploadedFile( - file_id=file_id, - basename=data.basename, - contents=data.raw_data, - digest=data.digest, - renderer=self.renderer - ) + return await self._binary_document(data) return InMemoryTextFile( basename=p.name, string_contents=data.raw_data.decode("utf-8"), @@ -435,18 +480,12 @@ async def get_document( async def upload_bytes_if_needed( self, basename: str, raw: bytes - ) -> UploadedFile: + ) -> Document: """Upload in-memory ``raw`` bytes (e.g. an audit-restored binary document) to the Files API, reusing a cached upload by CRC. The bytes-sourced analogue of :meth:`upload_file_if_needed`.""" - data = await _file_data(basename=basename, data=raw) - file_id = await self._upload_bytes(data.crc_basename, data.raw_data, data.mime) - return UploadedFile( - file_id=file_id, - basename=data.basename, - contents=data.raw_data, - digest=data.digest, - renderer=self.renderer + return await self._binary_document( + await _file_data(basename=basename, data=raw) ) def text_document_from(self, src: TextUploadable) -> TextDocument: diff --git a/composer/llm/anthropic.py b/composer/llm/anthropic.py index 4ab192df..f77f1f82 100644 --- a/composer/llm/anthropic.py +++ b/composer/llm/anthropic.py @@ -138,9 +138,8 @@ def text_block(self, text: str, *, cache_level: CacheLevel = CacheLevel.NONE) -> return to_ret def file_block( - self, file_id: str, *, filename: str, cache_level: CacheLevel = CacheLevel.NONE + self, file_id: str, *, cache_level: CacheLevel = CacheLevel.NONE ) -> dict: - # filename unused: the Files API upload already carries it. to_ret : dict[str, Any] = { "type": "document", "source": { @@ -155,6 +154,14 @@ def file_block( } return to_ret + def inline_file_block( + self, basename: str, contents: bytes, mime: str, + *, cache_level: CacheLevel = CacheLevel.NONE + ) -> dict: + raise NotImplementedError( + "Anthropic content is uploaded to the Files API, not inlined." + ) + @cache def _get_service(): return AnthropicService() diff --git a/composer/llm/openai.py b/composer/llm/openai.py index b7d3766d..2e776c0c 100644 --- a/composer/llm/openai.py +++ b/composer/llm/openai.py @@ -162,9 +162,8 @@ def text_block(self, text: str, *, cache_level: CacheLevel = CacheLevel.NONE) -> to_ret : dict[str, Any] = {"type": "text", "text": text} return to_ret def file_block( - self, file_id: str, *, filename: str, cache_level: CacheLevel = CacheLevel.NONE + self, file_id: str, *, cache_level: CacheLevel = CacheLevel.NONE ) -> dict: - # filename unused: the Files API upload already carries it. return { "type": "file", "file": { @@ -172,6 +171,14 @@ def file_block( }, } + def inline_file_block( + self, basename: str, contents: bytes, mime: str, + *, cache_level: CacheLevel = CacheLevel.NONE + ) -> dict: + raise NotImplementedError( + "OpenAI content is uploaded to the Files API, not inlined." + ) + # --- Files API uploader ---------------------------------------------------- @dataclass diff --git a/composer/llm/openrouter.py b/composer/llm/openrouter.py index 39a7bb6f..b0e68aca 100644 --- a/composer/llm/openrouter.py +++ b/composer/llm/openrouter.py @@ -15,20 +15,21 @@ Requires ``OPENROUTER_API_KEY``. """ -from typing import Any, TYPE_CHECKING, Mapping, AsyncIterator, override +from typing import Any, TYPE_CHECKING, override from dataclasses import dataclass, field from functools import cache -import asyncio import base64 -import json import logging import os -import urllib.request import httpx import openai +from pydantic import BaseModel, Field, PositiveInt, SecretStr, ValidationError -from composer.input.files import UploaderBase, ContentRenderer +from composer.input.files import ( + ContentRenderer, Document, FileData, InMemoryBytesFile, InMemoryTextFile, + TextDocument, UploaderBase +) from composer.input.types import ModelConfiguration from .provider import ( ProviderServiceBase, ProviderSpec, compaction_threshold, reasoning_effort, @@ -40,8 +41,6 @@ if TYPE_CHECKING: from langchain_core.language_models.chat_models import BaseChatModel - from langchain_core.outputs import ChatGenerationChunk - from langchain_openai import ChatOpenAI logger = logging.getLogger(__name__) @@ -53,19 +52,6 @@ _PROBE_TIMEOUT_SECONDS = 15 -# The only fields read off a roster record. -_PROBED_KEYS = ("context_length", "top_provider", "supported_parameters", "pricing") - -# Transient stream failures and how many times to re-issue the request. The SDK's -# own `max_retries` covers *establishing* a request; these surface while the SSE -# body is being consumed, by which point the request has already succeeded, so -# nothing below us retries them and they take the whole run down. `TimeoutError` -# covers langchain's `StreamChunkTimeoutError` (a subclass) for a content stall; -# `httpx.TransportError` covers the dropped-connection family. -_RETRYABLE_STREAM_ERRORS = (httpx.TransportError, TimeoutError, openai.APIConnectionError) -_STREAM_ATTEMPTS = 3 -_STREAM_RETRY_BACKOFF_SECONDS = 2.0 - def matches(model: str) -> bool: """OpenRouter ids are vendor-qualified (``moonshotai/kimi-k2.5``, @@ -97,35 +83,98 @@ class OpenRouterModelFeatures: reasoning: bool +# --- the `GET /api/v1/models` schema --------------------------------------- +# +# https://openrouter.ai/docs/api-reference/list-available-models. Only the fields +# this module reads are modelled; `extra="ignore"` (pydantic's default) drops the +# rest, which is most of the ~690KB payload. Prices arrive as decimal *strings*, +# which pydantic coerces to float on the way in. + +class _TopProvider(BaseModel): + context_length: PositiveInt | None = None + max_completion_tokens: PositiveInt | None = None + + +class _PriceCard(BaseModel): + """Per-token USD prices. A missing bucket is not a zero — it means the route + publishes no separate rate for it (see :func:`_price_tier`).""" + + prompt: float | None = None + completion: float | None = None + input_cache_read: float | None = None + input_cache_write: float | None = None + input_cache_write_1h: float | None = None + + +class _PriceOverride(_PriceCard): + """A conditional price, restating only the fields it changes. Only prompt-size + floors are modelled; OpenRouter also publishes time-of-day windows + (``utc_start``/``utc_end``), which a per-call price curve cannot express, so + those arrive with ``min_prompt_tokens`` unset and are skipped.""" + + min_prompt_tokens: PositiveInt | None = None + + +class _Pricing(_PriceCard): + overrides: list[_PriceOverride] = Field(default_factory=list) + + +class _ModelRecord(BaseModel): + id: str + context_length: PositiveInt | None = None + top_provider: _TopProvider = Field(default_factory=_TopProvider) + supported_parameters: set[str] = Field(default_factory=set) + pricing: _Pricing | None = None + + +class _ModelsEnvelope(BaseModel): + """Records stay raw here so one unreadable model can't take the roster with it + — 400+ vendors publish into this feed, and a single odd record blanking the + catalog would silently downgrade every route to the fallback window.""" + + data: list[Any] + + @cache -def _catalog() -> dict[str, Mapping[str, Any]]: - """OpenRouter's model roster, by id: one blocking unauthenticated GET, no retry, - from :meth:`OpenRouterModelProvider.create` at startup. Any failure yields an - empty catalog and a run on conservative defaults.""" +def _catalog() -> dict[str, _ModelRecord]: + """OpenRouter's model roster, by id: one blocking unauthenticated GET per + process, no retry, from :meth:`OpenRouterModelProvider.create` at startup. Any + failure yields an empty catalog and a run on conservative defaults.""" + return _fetch_catalog() + + +def _fetch_catalog() -> dict[str, _ModelRecord]: try: - with urllib.request.urlopen(_MODELS_URL, timeout=_PROBE_TIMEOUT_SECONDS) as resp: - payload = json.load(resp) - except (OSError, json.JSONDecodeError) as exc: + with httpx.Client(timeout=_PROBE_TIMEOUT_SECONDS) as client: + response = client.get(_MODELS_URL) + response.raise_for_status() + envelope = _ModelsEnvelope.model_validate_json(response.content) + except (httpx.HTTPError, ValidationError) as exc: logger.warning( "Could not fetch OpenRouter model metadata from %s (%s); falling back to " "a %d-token context window and no price card.", _MODELS_URL, exc, _FALLBACK_CONTEXT_WINDOW, ) return {} - records = payload.get("data") if isinstance(payload, dict) else None - if not isinstance(records, list): - logger.warning("Unexpected OpenRouter /models payload shape; ignoring it.") - return {} - # Projected to what's actually read: the full roster is ~690KB on the wire and - # ~2.3MB of retained objects, nearly all of it prose this module never touches. - return { - rec["id"]: {k: rec[k] for k in _PROBED_KEYS if k in rec} - for rec in records - if isinstance(rec, dict) and isinstance(rec.get("id"), str) - } + + catalog: dict[str, _ModelRecord] = {} + unreadable = 0 + for raw in envelope.data: + try: + record = _ModelRecord.model_validate(raw) + except ValidationError: + unreadable += 1 + continue + catalog[record.id] = record + if unreadable: + logger.warning( + "Skipped %d OpenRouter roster record(s) that did not match the expected " + "schema; those routes fall back to defaults.", unreadable, + ) + return catalog -def _record_for(model_name: str) -> Mapping[str, Any] | None: +def _record_for(model_name: str) -> _ModelRecord | None: catalog = _catalog() if (exact := catalog.get(model_name)) is not None: return exact @@ -135,11 +184,7 @@ def _record_for(model_name: str) -> Mapping[str, Any] | None: return catalog.get(base) if variant else None -def _positive_int(value: Any) -> int | None: - return value if isinstance(value, int) and not isinstance(value, bool) and value > 0 else None - - -def _features_from(record: Mapping[str, Any] | None) -> OpenRouterModelFeatures: +def _features_from(record: _ModelRecord | None) -> OpenRouterModelFeatures: if record is None: # Only reachable when the fetch itself failed — an id absent from a roster # that did load is rejected in `create`. @@ -148,64 +193,53 @@ def _features_from(record: Mapping[str, Any] | None) -> OpenRouterModelFeatures: max_output_tokens=None, reasoning=True, ) - top = record.get("top_provider") - top = top if isinstance(top, Mapping) else {} - supported = record.get("supported_parameters") - supported = set(supported) if isinstance(supported, list) else set() return OpenRouterModelFeatures( # Two windows are published: the model's own and the serving provider's. The # smaller is the one a request actually has to fit in. context_window=min( ( - w for w in ( - _positive_int(record.get("context_length")), - _positive_int(top.get("context_length")), - ) if w is not None + w for w in (record.context_length, record.top_provider.context_length) + if w is not None ), default=_FALLBACK_CONTEXT_WINDOW, ), - max_output_tokens=_positive_int(top.get("max_completion_tokens")), - reasoning="reasoning" in supported, + max_output_tokens=record.top_provider.max_completion_tokens, + reasoning="reasoning" in record.supported_parameters, ) # --- pricing --------------------------------------------------------------- -# OpenRouter quotes USD per token; PriceTier is USD per million. -_TOKENS_PER_MTOK = 1_000_000 - +# OpenRouter quotes USD per token; PriceTier is USD per million tokens. +_TOKENS_PER_MILLION = 1_000_000 -def _per_mtok(raw: Any) -> float | None: - """One price field, converted to per-MTok. A "0" is a real zero (a free route); - an absent field is a missing key, which comes back as None.""" - if not isinstance(raw, (str, int, float)) or isinstance(raw, bool): - return None - try: - return float(raw) * _TOKENS_PER_MTOK - except ValueError: - return None +def _price_tier(card: _PriceCard) -> PriceTier | None: + """One published price card as a :class:`PriceTier`, or None if the route + publishes no prompt/completion rate at all (leaving it uncosted). -def _price_tier(prices: Mapping[str, Any]) -> PriceTier | None: - prompt = _per_mtok(prices.get("prompt")) - completion = _per_mtok(prices.get("completion")) - if prompt is None or completion is None: + An unpublished *cache* bucket is not a guess: on OpenRouter it means the route + offers no separate rate for those tokens, so they bill at the ordinary prompt + rate — which is what falls through here. The only real inference is + ``cache_write_1h``, where a route with no 1-hour rate is assumed to charge its + 5-minute one; no route this backend has seen publishes the second without the + first.""" + if card.prompt is None or card.completion is None: return None - # A route that publishes no cache bucket bills those tokens as fresh input, and - # one with no separate 1h rate charges the 5m one. `or` would not do here: a - # real 0.0 is a free route. - if (cache_read := _per_mtok(prices.get("input_cache_read"))) is None: - cache_read = prompt - if (cache_write := _per_mtok(prices.get("input_cache_write"))) is None: - cache_write = prompt - if (cache_write_1h := _per_mtok(prices.get("input_cache_write_1h"))) is None: - cache_write_1h = cache_write + per_million = _TOKENS_PER_MILLION + cache_write = card.input_cache_write if card.input_cache_write is not None else card.prompt return PriceTier( - input=prompt, - output=completion, - cache_read=cache_read, - cache_write=cache_write, - cache_write_1h=cache_write_1h, + input=card.prompt * per_million, + output=card.completion * per_million, + cache_read=( + card.input_cache_read if card.input_cache_read is not None else card.prompt + ) * per_million, + cache_write=cache_write * per_million, + cache_write_1h=( + card.input_cache_write_1h + if card.input_cache_write_1h is not None + else cache_write + ) * per_million, ) @@ -217,27 +251,27 @@ def _bare_model_name(model_name: str) -> str: def _price_provider_from( - record: Mapping[str, Any] | None, model_name: str + record: _ModelRecord | None, model_name: str ) -> PriceProvider: """The route's pricing curve, live from the catalog where possible, else the static table on the vendor's bare model name — which covers ``openai/*`` and ``anthropic/*``, and yields None (an uncosted run, not a wrong one) elsewhere.""" - prices = (record or {}).get("pricing") - if not isinstance(prices, Mapping) or (short := _price_tier(prices)) is None: + pricing = record.pricing if record is not None else None + if pricing is None or (short := _price_tier(pricing)) is None: return price_provider_for(_bare_model_name(model_name)) - # Long-context surcharges arrive as overrides keyed by a prompt-size floor, each - # restating only the fields it changes (see openai/gpt-5.5's >272K tier). - raw_overrides = prices.get("overrides") - tiers: list[tuple[int, PriceTier]] = [] - if isinstance(raw_overrides, list): - for override in raw_overrides: - if not isinstance(override, Mapping): - continue - floor = _positive_int(override.get("min_prompt_tokens")) - tier = _price_tier({**prices, **override}) - if floor is not None and tier is not None: - tiers.append((floor, tier)) + # Long-context surcharges are keyed by a prompt-size floor and restate only the + # fields they change (see openai/gpt-5.5's >272K tier), so each one is merged + # over the base card before becoming a tier. Highest floor first, so the first + # match wins. + tiers = [ + (override.min_prompt_tokens, tier) + for override in pricing.overrides + if override.min_prompt_tokens is not None + and (tier := _price_tier( + pricing.model_copy(update=override.model_dump(exclude_none=True)) + )) is not None + ] tiers.sort(key=lambda t: t[0], reverse=True) def provider(input_tokens: int) -> PriceTier | None: @@ -264,38 +298,66 @@ def _api_key() -> str: class OpenRouterRenderer(OpenAIRenderer): """Content blocks in the Chat-Completions shape, which is what langchain converts from: ``_convert_chat_completions_blocks_to_responses`` turns ``file`` into - ``input_file`` and ``image_url`` into ``input_image`` on the way out. Only the - file block differs from OpenAI's; the text block is inherited.""" + ``input_file`` and ``image_url`` into ``input_image`` on the way out. The text + block is inherited from OpenAI's renderer.""" @override def file_block( - self, file_id: str, *, filename: str, cache_level: CacheLevel = CacheLevel.NONE + self, file_id: str, *, cache_level: CacheLevel = CacheLevel.NONE ) -> dict: - # `file_id` is a `data:` URL rather than a remote id: OpenRouter has no - # Files API, so InlineFileUploader carries the bytes here instead. - if file_id.startswith("data:image/"): - return {"type": "image_url", "image_url": {"url": file_id}} - return {"type": "file", "file": {"filename": filename, "file_data": file_id}} + raise NotImplementedError( + "OpenRouter has no Files API; binary content is inlined by " + "InlineFileUploader as an InMemoryBytesFile." + ) + + @override + def inline_file_block( + self, basename: str, contents: bytes, mime: str, + *, cache_level: CacheLevel = CacheLevel.NONE + ) -> dict: + url = f"data:{mime};base64,{base64.b64encode(contents).decode('ascii')}" + if mime.startswith("image/"): + return {"type": "image_url", "image_url": {"url": url}} + return {"type": "file", "file": {"filename": basename, "file_data": url}} @dataclass class InlineFileUploader(UploaderBase): - """``FileUploader`` impl for a provider with no Files API: the "upload" is a - ``data:`` URL built in memory, which the renderer inlines into the request. So a - large binary is re-sent with every request carrying it, and a PDF no route reads - natively goes through OpenRouter's ``file-parser`` plugin, which falls back to - ``mistral-ocr`` at $2/1K pages (pin an engine via ``plugins`` to avoid that).""" + """``FileUploader`` for a provider with no Files API: nothing is uploaded, so + binary content becomes an :class:`InMemoryBytesFile` and text destined for + upload simply stays in the prompt. + + A large binary therefore rides along in every request that carries it, and a + PDF no route reads natively goes through OpenRouter's ``file-parser`` plugin, + which falls back to ``mistral-ocr`` at $2/1K pages (pin an engine via + ``plugins`` to avoid that).""" renderer: ContentRenderer = field(default_factory=OpenRouterRenderer) + @override async def _upload_bytes( self, crc_basename: str, file_data: bytes, mime: str ) -> str: - # No dedup cache: the "id" *is* the content, so there is nothing to reuse. - # Off-thread because the encode is ~1.5ms/MB of blocked loop, matching how - # `composer.input.files` already offloads the read. - encoded = await asyncio.to_thread(base64.b64encode, file_data) - return f"data:{mime};base64,{encoded.decode('ascii')}" + raise NotImplementedError("OpenRouter has no Files API.") + + @override + async def _binary_document(self, data: FileData) -> Document: + return InMemoryBytesFile( + basename=data.basename, + contents=data.raw_data, + mime=data.mime, + renderer=self.renderer, + ) + + @override + async def _text_upload_document(self, data: FileData) -> TextDocument: + # There is no upload to make it smaller than the prompt, so the + # very-large-text case this exists for collapses into the ordinary one. + return InMemoryTextFile( + basename=data.basename, + string_contents=data.raw_data.decode("utf-8"), + renderer=self.renderer, + ) class OpenRouterService(ProviderServiceBase): @@ -306,61 +368,31 @@ class OpenRouterService(ProviderServiceBase): def __init__(self): from graphcore.tools.memory import openai_async_memory_tool super().__init__( - # The OpenAI-flavored memory tool is a plain client-side function tool - # differing only in its args schema, so it is portable to any route. + # The memory tool differs from the Anthropic one only in its args + # schema, so it is portable to any route. openai_async_memory_tool, InlineFileUploader, ) + @override + def should_retry(self, exc: Exception) -> bool: + """The OpenAI taxonomy, plus the two failures a *streamed* request adds. + Both surface while the SSE body is being read, when the request has already + succeeded, so the SDK's own ``max_retries`` never sees them: a dropped + connection (``httpx.TransportError``) and a content stall (langchain's + ``StreamChunkTimeoutError``, a ``TimeoutError`` subclass).""" + if isinstance(exc, (httpx.TransportError, TimeoutError, openai.APIConnectionError)): + return True + if isinstance(exc, openai.APIStatusError): + return exc.status_code in (408, 409, 429) or exc.status_code >= 500 + return False + @cache def _openrouter_service(): return OpenRouterService() -# --- chat model ------------------------------------------------------------ - -@cache -def _chat_model_cls() -> type["ChatOpenAI"]: - """``ChatOpenAI`` that re-issues a streamed request when the stream breaks. - - Defined behind a cached factory so ``langchain_openai`` stays a lazy import, - as it is in the sibling backends.""" - from langchain_openai import ChatOpenAI - - class RetryingChatOpenAI(ChatOpenAI): - @override - async def _astream( - self, *args: Any, **kwargs: Any - ) -> AsyncIterator["ChatGenerationChunk"]: - for attempt in range(1, _STREAM_ATTEMPTS + 1): - # Buffered, not forwarded as they arrive: a retry must not emit a - # partial response twice. Callers aggregate through `ainvoke` - # anyway, so this costs nothing today — an incremental consumer - # would lose its incrementality. - chunks: list["ChatGenerationChunk"] = [] - try: - async for chunk in super()._astream(*args, **kwargs): - chunks.append(chunk) - except _RETRYABLE_STREAM_ERRORS as exc: - if attempt == _STREAM_ATTEMPTS: - raise - delay = _STREAM_RETRY_BACKOFF_SECONDS * attempt - logger.warning( - "OpenRouter stream failed after %d chunk(s) (%s: %s); " - "re-issuing in %.0fs (attempt %d/%d).", - len(chunks), type(exc).__name__, exc, delay, - attempt + 1, _STREAM_ATTEMPTS, - ) - await asyncio.sleep(delay) - continue - for chunk in chunks: - yield chunk - return - - return RetryingChatOpenAI - - # --- ModelProvider --------------------------------------------------------- @dataclass @@ -411,7 +443,7 @@ def _output_token_cap(self) -> int: def builder_for( self, *, cache_level: CacheLevel = CacheLevel.NONE, disable_thinking: bool = False ) -> "BaseChatModel": - from pydantic import SecretStr + from langchain_openai import ChatOpenAI opts = self.options kwargs: dict[str, Any] = {} @@ -430,15 +462,12 @@ def builder_for( # back with the next tool result so the model can resume it. kwargs["include"] = ["reasoning.encrypted_content"] - return _chat_model_cls()( + return ChatOpenAI( model=self.model_name, base_url=BASE_URL, api_key=SecretStr(self.api_key), - # Load-bearing, not a default: the Responses API is the only surface on - # which reasoning survives a tool round-trip, because langchain echoes a - # prior turn's reasoning items back into the next request. Chat - # Completions drops them, so a long tool loop re-derives its reasoning - # every round at the output token rate. + # Only surface on which langchain echoes a prior turn's reasoning back, + # so a tool loop doesn't re-derive it every round. use_responses_api=True, # Unstreamed, OpenRouter holds the whole generation and an upstream # pause trips its gateway idle timeout, failing the request. From d682bcfd60e75153f9b5f2c94771e18ae5163e25 Mon Sep 17 00:00:00 2001 From: Naftali Goldstein Date: Sun, 30 Aug 2026 18:13:27 +0300 Subject: [PATCH 3/7] fix timeouts --- composer/llm/openrouter.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/composer/llm/openrouter.py b/composer/llm/openrouter.py index b0e68aca..3ce91c67 100644 --- a/composer/llm/openrouter.py +++ b/composer/llm/openrouter.py @@ -52,6 +52,13 @@ _PROBE_TIMEOUT_SECONDS = 15 +# How long a stream may go quiet before it is treated as dead. langchain's guard +# measures the gap between *content* chunks and defaults to 120s, which assumes the +# provider streams its reasoning; OpenRouter does not forward reasoning deltas for +# every route, so a thinking model goes silent for the whole reasoning phase — over +# 200s on a kimi-k3 burst — and a working request looks stalled. +_STREAM_QUIET_SECONDS = 900.0 + def matches(model: str) -> bool: """OpenRouter ids are vendor-qualified (``moonshotai/kimi-k2.5``, @@ -472,6 +479,7 @@ def builder_for( # Unstreamed, OpenRouter holds the whole generation and an upstream # pause trips its gateway idle timeout, failing the request. streaming=True, + stream_chunk_timeout=_STREAM_QUIET_SECONDS, # OpenRouter is stateless: store=True is a 400, not a no-op. store=False, # langchain renames this to `max_output_tokens` for the Responses API. From cd70edc5b5e6d2c0cf0f4e04ce54c70d5f50fe60 Mon Sep 17 00:00:00 2001 From: Naftali Goldstein Date: Sun, 30 Aug 2026 18:39:52 +0300 Subject: [PATCH 4/7] reword a comment --- composer/llm/openrouter.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/composer/llm/openrouter.py b/composer/llm/openrouter.py index 3ce91c67..3b95faab 100644 --- a/composer/llm/openrouter.py +++ b/composer/llm/openrouter.py @@ -53,10 +53,11 @@ _PROBE_TIMEOUT_SECONDS = 15 # How long a stream may go quiet before it is treated as dead. langchain's guard -# measures the gap between *content* chunks and defaults to 120s, which assumes the -# provider streams its reasoning; OpenRouter does not forward reasoning deltas for -# every route, so a thinking model goes silent for the whole reasoning phase — over -# 200s on a kimi-k3 burst — and a working request looks stalled. +# measures the gap between chunks *it emits* and defaults to 120s. OpenRouter streams +# reasoning as `response.reasoning_text.delta`, where OpenAI sends the summary event +# `response.reasoning_summary_text.delta` — the only reasoning event langchain turns +# into a chunk. So a thinking route's reasoning phase yields no chunks at all (over +# 200s of it on a kimi-k3 burst) and a healthy request looks stalled. _STREAM_QUIET_SECONDS = 900.0 From f96d290e3775b19872e6e62292b86621715519ff Mon Sep 17 00:00:00 2001 From: Naftali Goldstein Date: Sun, 30 Aug 2026 19:09:03 +0300 Subject: [PATCH 5/7] update tests and graphcore --- graphcore | 2 +- pyproject.toml | 2 +- tests/test_token_usage.py | 16 ++++++++-------- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/graphcore b/graphcore index 54a6852e..d11a2169 160000 --- a/graphcore +++ b/graphcore @@ -1 +1 @@ -Subproject commit 54a6852eac72896d17d3bc4f3068a775f9200d1f +Subproject commit d11a2169f9ddfa034345762d85612f6c1ef5e7a4 diff --git a/pyproject.toml b/pyproject.toml index 9f7d26ff..146e100c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ authors = [ ] dependencies = [ - "graphcore @ git+ssh://git@github.com/Certora/graphcore.git@54a6852eac72896d17d3bc4f3068a775f9200d1f", + "graphcore @ git+ssh://git@github.com/Certora/graphcore.git@d11a2169f9ddfa034345762d85612f6c1ef5e7a4", "aiohttp>=3.13", "attrs>=26.1", "Jinja2>=3.1", diff --git a/tests/test_token_usage.py b/tests/test_token_usage.py index 04596223..ef64d12c 100644 --- a/tests/test_token_usage.py +++ b/tests/test_token_usage.py @@ -132,14 +132,14 @@ async def test_token_usage_persisted_to_run_meta_tags(): def _fake_model(callbacks): resp = AIMessage( content="ok", - response_metadata={ - "model_name": "claude-test", - "usage": { - "input_tokens": 100, - "output_tokens": 10, - "cache_read_input_tokens": 5, - "cache_creation_input_tokens": 2, - }, + response_metadata={"model_name": "claude-test"}, + # langchain's normalized shape. `input_tokens` is the inclusive total + # (100 fresh + 5 read + 2 written); the buckets come apart downstream. + usage_metadata={ + "input_tokens": 107, + "output_tokens": 10, + "total_tokens": 117, + "input_token_details": {"cache_read": 5, "cache_creation": 2}, }, ) return FakeMessagesListChatModel(responses=[resp, resp], callbacks=callbacks) From 25a4e0be3546ef62ad10d60554f7dba6815938d1 Mon Sep 17 00:00:00 2001 From: Naftali Goldstein Date: Tue, 1 Sep 2026 13:37:37 +0300 Subject: [PATCH 6/7] fix context length detection --- composer/llm/openrouter.py | 120 +++++++++++++++++++++++++++++++------ 1 file changed, 102 insertions(+), 18 deletions(-) diff --git a/composer/llm/openrouter.py b/composer/llm/openrouter.py index 3b95faab..b0fd9d12 100644 --- a/composer/llm/openrouter.py +++ b/composer/llm/openrouter.py @@ -98,11 +98,6 @@ class OpenRouterModelFeatures: # rest, which is most of the ~690KB payload. Prices arrive as decimal *strings*, # which pydantic coerces to float on the way in. -class _TopProvider(BaseModel): - context_length: PositiveInt | None = None - max_completion_tokens: PositiveInt | None = None - - class _PriceCard(BaseModel): """Per-token USD prices. A missing bucket is not a zero — it means the route publishes no separate rate for it (see :func:`_price_tier`).""" @@ -129,8 +124,9 @@ class _Pricing(_PriceCard): class _ModelRecord(BaseModel): id: str + # A fallback window for when the per-endpoint fetch fails; the pool is the + # better source when it is reachable. context_length: PositiveInt | None = None - top_provider: _TopProvider = Field(default_factory=_TopProvider) supported_parameters: set[str] = Field(default_factory=set) pricing: _Pricing | None = None @@ -182,6 +178,82 @@ def _fetch_catalog() -> dict[str, _ModelRecord]: return catalog +class _Endpoint(BaseModel): + # `tag` carries the provider slug `provider.only` wants, sometimes with a + # quantization suffix: `deepinfra/bf16` -> `deepinfra`. + tag: str + context_length: PositiveInt | None = None + max_completion_tokens: PositiveInt | None = None + + @property + def provider(self) -> str: + return self.tag.partition("/")[0] + + +class _EndpointsEnvelope(BaseModel): + class _Data(BaseModel): + endpoints: list[_Endpoint] + + data: _Data + + +@dataclass(frozen=True) +class _Pool: + """What the providers serving one route can do, per provider. The roster only + summarises its best one, so the limits a request actually meets live here.""" + + # Highest cap in the pool: asking beyond it can't be served by anyone. + output_cap: int | None = None + # Smallest window among the providers routing is confined to. + context_window: int | None = None + # Provider slugs to confine routing to, or None for OpenRouter's own choice. + routable: list[str] | None = None + + +def _pool_for(model_name: str, requested_output: int) -> _Pool: + """The serving pool's real limits, and who to route to. + + A model is served by a pool OpenRouter load-balances across, whose caps can + differ by orders of magnitude, so the roster's per-model figures are the best + case rather than what a given request gets. Confining routing to the providers + that can serve the ask keeps the balancing, minus the ones a request would have + failed on — and the window then follows from that narrowed set.""" + try: + with httpx.Client(timeout=_PROBE_TIMEOUT_SECONDS) as client: + response = client.get(f"{_MODELS_URL}/{model_name}/endpoints") + response.raise_for_status() + endpoints = _EndpointsEnvelope.model_validate_json( + response.content + ).data.endpoints + except (httpx.HTTPError, ValidationError) as exc: + logger.warning( + "Could not fetch OpenRouter endpoints for %s (%s); falling back to the " + "roster's own limits and leaving provider routing unrestricted.", + model_name, exc, + ) + return _Pool() + + caps = [e.max_completion_tokens for e in endpoints if e.max_completion_tokens] + cap = min(requested_output, max(caps)) if caps else requested_output + serving = [e for e in endpoints if (e.max_completion_tokens or 0) >= cap] + if not serving: + logger.warning( + "No OpenRouter provider for %s publishes an output cap of %d tokens; " + "leaving routing unrestricted so the rejection comes from them.", + model_name, cap, + ) + serving = endpoints + + able = {e.provider for e in serving} + windows = [e.context_length for e in serving if e.context_length] + return _Pool( + output_cap=max(caps) if caps else None, + context_window=min(windows) if windows else None, + # Nothing to exclude: say nothing rather than pin the whole pool. + routable=sorted(able) if able < {e.provider for e in endpoints} else None, + ) + + def _record_for(model_name: str) -> _ModelRecord | None: catalog = _catalog() if (exact := catalog.get(model_name)) is not None: @@ -192,7 +264,9 @@ def _record_for(model_name: str) -> _ModelRecord | None: return catalog.get(base) if variant else None -def _features_from(record: _ModelRecord | None) -> OpenRouterModelFeatures: +def _features_from( + record: _ModelRecord | None, pool: _Pool = _Pool() +) -> OpenRouterModelFeatures: if record is None: # Only reachable when the fetch itself failed — an id absent from a roster # that did load is rejected in `create`. @@ -202,16 +276,10 @@ def _features_from(record: _ModelRecord | None) -> OpenRouterModelFeatures: reasoning=True, ) return OpenRouterModelFeatures( - # Two windows are published: the model's own and the serving provider's. The - # smaller is the one a request actually has to fit in. - context_window=min( - ( - w for w in (record.context_length, record.top_provider.context_length) - if w is not None - ), - default=_FALLBACK_CONTEXT_WINDOW, + context_window=( + pool.context_window or record.context_length or _FALLBACK_CONTEXT_WINDOW ), - max_output_tokens=record.top_provider.max_completion_tokens, + max_output_tokens=pool.output_cap, reasoning="reasoning" in record.supported_parameters, ) @@ -263,7 +331,15 @@ def _price_provider_from( ) -> PriceProvider: """The route's pricing curve, live from the catalog where possible, else the static table on the vendor's bare model name — which covers ``openai/*`` and - ``anthropic/*``, and yields None (an uncosted run, not a wrong one) elsewhere.""" + ``anthropic/*``, and yields None (an uncosted run, not a wrong one) elsewhere. + + TODO: this is the roster's *model-level* card, which is the modal price across the + serving pool rather than a bound — kimi-k3's providers span $2.55-$6.00 in and + $12.75-$22.50 out per MTok, so a run can be off by ~2x in either direction. Two + ways out: price from the max among ``_Pool.routable`` (already fetched, makes the + figure an upper bound), or read the exact cost OpenRouter returns under + ``usage: {"include": true}`` — which langchain currently drops, since its + Responses ``response_metadata`` whitelist has no ``usage`` key.""" pricing = record.pricing if record is not None else None if pricing is None or (short := _price_tier(pricing)) is None: return price_provider_for(_bare_model_name(model_name)) @@ -415,6 +491,8 @@ class OpenRouterModelProvider: price_provider: PriceProvider api_key: str provider: OpenRouterService = field(default_factory=_openrouter_service) + # Provider slugs to confine routing to, or None for OpenRouter's own choice. + routable_providers: list[str] | None = None @staticmethod def create(model_name: str, options: ModelConfiguration) -> "OpenRouterModelProvider": @@ -429,12 +507,14 @@ def create(model_name: str, options: ModelConfiguration) -> "OpenRouterModelProv f"{model_name!r} is not an OpenRouter model; see " f"https://openrouter.ai/models for the roster." ) + pool = _pool_for(model_name, options.tokens) return OpenRouterModelProvider( model_name=model_name, options=options, - features=_features_from(record), + features=_features_from(record, pool), price_provider=_price_provider_from(record, model_name), api_key=api_key, + routable_providers=pool.routable, ) @property @@ -470,6 +550,10 @@ def builder_for( # back with the next tool result so the model can resume it. kwargs["include"] = ["reasoning.encrypted_content"] + if self.routable_providers is not None: + # OpenRouter's own routing field; the SDK has no parameter for it. + kwargs["extra_body"] = {"provider": {"only": self.routable_providers}} + return ChatOpenAI( model=self.model_name, base_url=BASE_URL, From 7121a5fe47426d6f619f82f812db16f3213b214c Mon Sep 17 00:00:00 2001 From: Naftali Goldstein Date: Tue, 1 Sep 2026 13:44:49 +0300 Subject: [PATCH 7/7] reword a comment --- composer/llm/openrouter.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/composer/llm/openrouter.py b/composer/llm/openrouter.py index b0fd9d12..44228e6c 100644 --- a/composer/llm/openrouter.py +++ b/composer/llm/openrouter.py @@ -333,9 +333,9 @@ def _price_provider_from( static table on the vendor's bare model name — which covers ``openai/*`` and ``anthropic/*``, and yields None (an uncosted run, not a wrong one) elsewhere. - TODO: this is the roster's *model-level* card, which is the modal price across the - serving pool rather than a bound — kimi-k3's providers span $2.55-$6.00 in and - $12.75-$22.50 out per MTok, so a run can be off by ~2x in either direction. Two + TODO: this is the roster's *model-level* card, which is one price point inside the + serving pool rather than a bound on it — kimi-k3's providers span $2.55-$6.00 in + and $12.75-$22.50 out per MTok, so a run can be off by ~2x either way. Two ways out: price from the max among ``_Pool.routable`` (already fetched, makes the figure an upper bound), or read the exact cost OpenRouter returns under ``usage: {"include": true}`` — which langchain currently drops, since its