diff --git a/composer/input/files.py b/composer/input/files.py index ce79c486..8e7e14f6 100644 --- a/composer/input/files.py +++ b/composer/input/files.py @@ -34,7 +34,21 @@ 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: ... + + 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) @@ -207,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 @@ -259,7 +305,7 @@ def string_contents(self) -> str: # --------------------------------------------------------------------------- @dataclass -class _FileData: +class FileData: basename: str raw_data: bytes is_binary: bool @@ -271,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) @@ -315,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 @@ -350,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 @@ -360,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, @@ -377,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, @@ -395,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: @@ -413,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"), @@ -429,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 3baa8aba..3a6b2814 100644 --- a/composer/llm/anthropic.py +++ b/composer/llm/anthropic.py @@ -14,7 +14,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 @@ -139,7 +139,9 @@ 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, *, cache_level: CacheLevel = CacheLevel.NONE + ) -> dict: to_ret : dict[str, Any] = { "type": "document", "source": { @@ -154,6 +156,14 @@ def file_block(self, file_id: str, *, cache_level: CacheLevel = CacheLevel.NONE) } 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() @@ -302,8 +312,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 @@ -347,12 +355,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..2e776c0c 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,9 @@ 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, *, cache_level: CacheLevel = CacheLevel.NONE + ) -> dict: return { "type": "file", "file": { @@ -176,6 +171,14 @@ def file_block(self, file_id: str, *, cache_level: CacheLevel = CacheLevel.NONE) }, } + 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 @@ -250,8 +253,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 +263,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 +274,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..44228e6c --- /dev/null +++ b/composer/llm/openrouter.py @@ -0,0 +1,586 @@ +"""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, override +from dataclasses import dataclass, field +from functools import cache +import base64 +import logging +import os + +import httpx +import openai +from pydantic import BaseModel, Field, PositiveInt, SecretStr, ValidationError + +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, + 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 + +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 + +# How long a stream may go quiet before it is treated as dead. langchain's guard +# 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 + + +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 + + +# --- 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 _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 + # 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 + 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, _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 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 {} + + 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 + + +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: + 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 _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`. + return OpenRouterModelFeatures( + context_window=_FALLBACK_CONTEXT_WINDOW, + max_output_tokens=None, + reasoning=True, + ) + return OpenRouterModelFeatures( + context_window=( + pool.context_window or record.context_length or _FALLBACK_CONTEXT_WINDOW + ), + max_output_tokens=pool.output_cap, + reasoning="reasoning" in record.supported_parameters, + ) + + +# --- pricing --------------------------------------------------------------- + +# OpenRouter quotes USD per token; PriceTier is USD per million tokens. +_TOKENS_PER_MILLION = 1_000_000 + + +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). + + 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 + 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=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, + ) + + +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: _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. + + 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 + 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)) + + # 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: + 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. The text + block is inherited from OpenAI's renderer.""" + + @override + def file_block( + self, file_id: str, *, cache_level: CacheLevel = CacheLevel.NONE + ) -> dict: + 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`` 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: + 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): + """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 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() + + +# --- 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) + # 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": + # 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." + ) + pool = _pool_for(model_name, options.tokens) + return OpenRouterModelProvider( + model_name=model_name, + options=options, + features=_features_from(record, pool), + price_provider=_price_provider_from(record, model_name), + api_key=api_key, + routable_providers=pool.routable, + ) + + @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 langchain_openai import ChatOpenAI + + 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"] + + 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, + api_key=SecretStr(self.api_key), + # 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. + 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. + 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 81a38037..e535654f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -247,3 +247,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"