diff --git a/tests/integration/README.md b/tests/integration/README.md index 7e3cf67cb08..dcdf0e9fa96 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -2,7 +2,7 @@ These tests exercise a running gateway, PostgreSQL and Redis with an owned local upstream. CircleCI owns this suite. Tests are grouped by behavior, with no automatic test retries or fallback to paid provider calls -The `cost` group runs the scripted-wire cost matrix through the shared integration upstream, which serves the test-owned cost map over loopback through `LITELLM_MODEL_COST_MAP_URL`; cost goldens are checked into the integration suite and must not be copied into the E2E coverage registry. A provider speaking an existing response shape is a `wires.json` row, a `cases.json` `providers` row and cost-map entries; a new response shape needs a renderer in `scripted_wires.py` +The `cost` group runs the scripted-shape cost matrix through the shared integration upstream, which serves the test-owned cost map over loopback through `LITELLM_MODEL_COST_MAP_URL`; cost goldens are checked into the integration suite and must not be copied into the E2E coverage registry. The upstream renders a scenario in the shape LiteLLM's own provider config resolves to for the deployment, so a provider LiteLLM already parses with one of the five rendered families is a cost-map entry plus a `cases.json` `providers` row with its deployment parameters; a provider whose config class is none of those families fails at collection until `scripted_shapes.py` gains a renderer Use `tests/integration/run.py management`, `accounting`, `database`, `providers`, `extensions`, `sdk` or `cost` to run a selected group. Set `INTEGRATION_PROXY_URL`, `INTEGRATION_UPSTREAM_URL`, `INTEGRATION_MASTER_KEY` and `DATABASE_URL` to an isolated test deployment. The runner selects the new domain directories explicitly; the legacy OCI and sandbox selections remain separate diff --git a/tests/integration/_support/scripted_wires.py b/tests/integration/_support/scripted_shapes.py similarity index 91% rename from tests/integration/_support/scripted_wires.py rename to tests/integration/_support/scripted_shapes.py index 8da2c57c9a0..61bfe7c24f1 100644 --- a/tests/integration/_support/scripted_wires.py +++ b/tests/integration/_support/scripted_shapes.py @@ -1,24 +1,20 @@ -"""Scripted provider wires for the cost-calculation integration suite. +"""Scripted response shapes for the cost-calculation integration suite. -The shared integration upstream registers a Scenario over a small control API; -the provider wire routes answer the proxy's upstream calls with the scripted usage figures, in the exact wire shape -the real provider would emit (OpenAI chat completions, OpenAI Responses, -Anthropic Messages, Gemini generateContent, or the OpenAI-compatible Together / -Fireworks surfaces). Because the usage is scripted, expected spend is literal -arithmetic on the test cost map's rates, with no dependency on what a real -provider would report. +This module owns the Scenario schema, the five renderers, one per LiteLLM +parser family, and the dispatcher. Because the usage is scripted, expected +spend is literal arithmetic on the test cost map's rates, with no dependency +on what a real provider would report. The upstream exposes: - ``POST /__scenarios`` register a Scenario JSON, returns its id - ``DELETE /__scenarios/`` remove it -- ``POST ///`` provider wire; mount is one of - ``openai``, ``anthropic``, ``gemini``, ``together``, ``fireworks``, ``azure``, - ``bedrock``, ``vertex`` and the remainder is whatever path the provider - client appends (``chat/completions``, ``responses``, ``v1/messages``, - ``models/:generateContent`` ...). Vertex appends ``:generateContent`` / - ``:streamGenerateContent`` to the mount segment itself, and Bedrock Converse - targets ``model//converse`` / ``converse-stream`` +- ``POST //`` provider response; the remainder is whatever + path the provider client appends (``chat/completions``, ``responses``, + ``v1/messages``, ``models/:generateContent`` ...). Vertex appends + ``:generateContent`` / ``:streamGenerateContent`` to the scenario segment, + and Bedrock Converse targets ``model//converse`` / + ``converse-stream`` A request carrying ``"stream": true`` (or the ``:streamGenerateContent`` Gemini verb) gets an SSE answer; ``stream_usage`` on the Scenario decides whether the @@ -34,14 +30,12 @@ import time import zlib from collections.abc import Mapping from dataclasses import dataclass -from pathlib import Path from types import MappingProxyType from typing import Final, Literal, TypeAlias, assert_never from urllib.parse import unquote, urlsplit from pydantic import BaseModel, ConfigDict, TypeAdapter, model_validator -Wire: TypeAlias = str Shape: TypeAlias = Literal[ "openai_chat", "openai_responses", @@ -49,6 +43,77 @@ Shape: TypeAlias = Literal[ "gemini_generate", "bedrock_converse", ] + + +@dataclass(frozen=True, slots=True) +class ShapeSpec: + usage: frozenset[str] + terminals: frozenset[str] + + +SHAPES: Final[Mapping[Shape, ShapeSpec]] = MappingProxyType( + { + "openai_chat": ShapeSpec( + usage=frozenset( + { + "cache_read_tokens", + "reasoning_tokens", + "audio_input_tokens", + "audio_output_tokens", + "web_search_calls", + } + ), + terminals=frozenset(), + ), + "openai_responses": ShapeSpec( + usage=frozenset( + { + "cache_read_tokens", + "reasoning_tokens", + "web_search_calls", + "file_search_calls", + } + ), + terminals=frozenset({"incomplete", "unvalidated"}), + ), + "anthropic_messages": ShapeSpec( + usage=frozenset( + { + "cache_read_tokens", + "web_search_calls", + "cache_write_5m_tokens", + "cache_write_1h_tokens", + } + ), + terminals=frozenset(), + ), + "gemini_generate": ShapeSpec( + usage=frozenset( + { + "cache_read_tokens", + "reasoning_tokens", + "audio_input_tokens", + "audio_output_tokens", + "image_input_tokens", + "video_input_tokens", + "web_search_calls", + "google_maps_calls", + } + ), + terminals=frozenset({"prompt_blocked"}), + ), + "bedrock_converse": ShapeSpec( + usage=frozenset( + { + "cache_read_tokens", + "cache_write_5m_tokens", + "cache_write_1h_tokens", + } + ), + terminals=frozenset(), + ), + } +) StreamUsage: TypeAlias = Literal["final_chunk", "absent"] ServiceTier: TypeAlias = Literal["flex", "priority"] TerminalKind: TypeAlias = Literal["completed", "incomplete", "unvalidated", "prompt_blocked"] @@ -58,7 +123,7 @@ _BASE_USAGE_FIELDS: Final = frozenset({"fresh_input_tokens", "output_tokens"}) class ScriptedToolCall(BaseModel): """A single function call the scripted output emits instead of text. - ``arguments`` is the wire's JSON string (~250 chars), sliced into deltas + ``arguments`` is the shape's JSON string (~250 chars), sliced into deltas for streams.""" model_config = ConfigDict(frozen=True) @@ -71,7 +136,7 @@ class ScriptedUsage(BaseModel): """Physical token counts the scripted response reports. ``fresh_input_tokens`` is the uncached, never-written, non-audio input count; ``output_tokens`` is the non-reasoning, non-audio output count. Renderers add the cached, written, - audio, and reasoning counts into the wire's total fields the way the real + audio, and reasoning counts into the shape's total fields the way the real provider does (inside prompt_tokens for OpenAI/Gemini, as uncached-only input_tokens for Anthropic).""" @@ -92,32 +157,6 @@ class ScriptedUsage(BaseModel): file_search_calls: int = 0 -class WireSpec(BaseModel): - model_config = ConfigDict(frozen=True) - - shape: Shape - mount: str - usage: frozenset[str] - terminals: frozenset[TerminalKind] - - -def _load_wires() -> Mapping[str, WireSpec]: - adapter: Final = TypeAdapter(dict[str, WireSpec]) - loaded: Final = adapter.validate_json((Path(__file__).resolve().with_name("wires.json")).read_bytes()) - known_usage_fields: Final = frozenset(ScriptedUsage.model_fields) - _BASE_USAGE_FIELDS - unknown: Final = { - wire: sorted(spec.usage - known_usage_fields) - for wire, spec in loaded.items() - if spec.usage - known_usage_fields - } - if unknown: - raise ValueError(f"wires.json has unknown usage fields: {unknown}") - return MappingProxyType(loaded) - - -WIRES: Final[Mapping[str, WireSpec]] = _load_wires() - - class ScriptedOutput(BaseModel): model_config = ConfigDict(frozen=True) @@ -127,9 +166,9 @@ class ScriptedOutput(BaseModel): # prove the biller prices the provider-reported model. response_model: str | None = None # OpenAI-compatible providers can report a provider-computed cost; emitted as - # the top-level "cost" field on the together/fireworks wire. + # the top-level "cost" field on the together/fireworks response. provider_cost: float | None = None - # When set, the response is a tool call only: no text content on any wire. + # When set, the response is a tool call only: no text content on any response. tool_call: ScriptedToolCall | None = None # Terminal shape: "unvalidated" makes the Responses terminal response fail # pydantic validation so the proxy takes its model_construct dict path; @@ -141,7 +180,7 @@ class Scenario(BaseModel): model_config = ConfigDict(frozen=True) scenario_id: str - wire: Wire + shape: Shape usage: ScriptedUsage output: ScriptedOutput # The bare provider-facing model name the renderer echoes when the request @@ -157,17 +196,13 @@ class Scenario(BaseModel): @model_validator(mode="after") def _check_terminal_supported(self) -> Scenario: - spec: Final = WIRES.get(self.wire) - if spec is None: - raise ValueError( - f"unknown wire {self.wire}; known wires: {', '.join(sorted(WIRES))}" - ) + spec: Final = SHAPES[self.shape] if ( self.output.terminal != "completed" and self.output.terminal not in spec.terminals ): raise ValueError( - f"wire {self.wire} cannot emit terminal={self.output.terminal}" + f"shape {self.shape} cannot emit terminal={self.output.terminal}" ) unsupported: Final = frozenset( field @@ -177,18 +212,14 @@ class Scenario(BaseModel): ) if unsupported: raise ValueError( - f"wire {self.wire} cannot express usage fields {sorted(unsupported)}" + f"shape {self.shape} cannot express usage fields {sorted(unsupported)}" ) - if (self.speed or self.inference_geo) and self.wire != "anthropic_messages": + if (self.speed or self.inference_geo) and self.shape != "anthropic_messages": raise ValueError( - f"wire {self.wire} cannot emit speed/inference_geo (anthropic usage fields)" + f"shape {self.shape} cannot emit speed/inference_geo (anthropic usage fields)" ) return self - @property - def mount(self) -> str: - return WIRES[self.wire].mount - class ScenarioRegistered(BaseModel): scenario_id: str @@ -233,7 +264,7 @@ def _sse(events: tuple[tuple[str | None, Mapping[str, object] | str], ...]) -> b return "".join(_sse_frame(event_name, data) for event_name, data in events).encode("utf-8") -# ---------- per-wire usage shapes ---------- + # ---------- per-shape usage shapes ---------- def _openai_usage(u: ScriptedUsage) -> Mapping[str, object]: @@ -400,7 +431,7 @@ def _responses_usage(u: ScriptedUsage) -> Mapping[str, object]: ) -# ---------- per-wire responses ---------- + # ---------- per-shape responses ---------- def _split_arguments(arguments: str) -> tuple[str, ...]: @@ -1214,8 +1245,8 @@ def _render( scenario: Scenario, *, stream: bool, requested_model: str, path_tail: str ) -> RenderedResponse: # Azure bridges gpt-5.4+ chat requests carrying function tools onto the - # Responses API, which lands on the same mount at openai/responses. - if scenario.wire == "azure_chat" and path_tail.endswith("openai/responses"): + # Responses API, which lands on the same shape at openai/responses. + if scenario.shape == "openai_chat" and path_tail.endswith("openai/responses"): if stream: return RenderedResponse( 200, "text/event-stream", _responses_sse(scenario, requested_model) @@ -1223,7 +1254,7 @@ def _render( return RenderedResponse( 200, "application/json", _json_bytes(_responses_body(scenario, requested_model)) ) - shape: Final = WIRES[scenario.wire].shape + shape: Final = scenario.shape match shape: case "bedrock_converse": if stream: @@ -1282,8 +1313,8 @@ def _request_body(body: bytes) -> Mapping[str, object]: return MappingProxyType({}) -def _request_wants_stream(mount_endpoint: str | None, path_tail: str, body: bytes) -> bool: - if mount_endpoint == "streamGenerateContent" or ":streamGenerateContent" in path_tail: +def _request_wants_stream(endpoint: str | None, path_tail: str, body: bytes) -> bool: + if endpoint == "streamGenerateContent" or ":streamGenerateContent" in path_tail: return True if path_tail.endswith("converse-stream"): return True @@ -1301,44 +1332,33 @@ def _request_model(body: bytes, path_tail: str, scenario: Scenario) -> str: path_model: Final = path_tail.split("/", 2)[1] if path_tail.count("/") >= 2 else "" if path_model: return unquote(path_model) - # Vertex names it in the URL too, but the mount segment swallowed it when - # the api_base carried a path; fall back to the scenario's declared model. + # Vertex names it in the URL too, but the path may carry only the endpoint; + # fall back to the scenario's declared model. return scenario.model def render(store: ScenarioStore, method: str, raw_path: str, body: bytes) -> RenderedResponse: path: Final = urlsplit(raw_path).path segments: Final = tuple(segment for segment in path.split("/") if segment) - if len(segments) < 2 or method != "POST": + if len(segments) < 1 or method != "POST": return RenderedResponse( 404, "application/json", _json_bytes(_jobj(("error", f"no route for {method} {path}"))) ) - scenario_id: Final = segments[0] - # Vertex builds {api_base}:{endpoint}, so the mount segment can carry a - # :generateContent / :streamGenerateContent suffix. - mount_segment: Final = segments[1] - mount, mount_endpoint = ( - mount_segment.split(":", 1) - if ":" in mount_segment - else (mount_segment, None) + scenario_segment: Final = segments[0] + scenario_id, endpoint = ( + scenario_segment.split(":", 1) + if ":" in scenario_segment + else (scenario_segment, None) ) found: Final = store.get(scenario_id) if found is None: return RenderedResponse( 404, "application/json", _json_bytes(_jobj(("error", f"unknown scenario {scenario_id}"))) ) - if found.mount != mount: - return RenderedResponse( - 400, - "application/json", - _json_bytes( - _jobj(("error", f"scenario {scenario_id} is wire {found.wire}, not mount {mount}")) - ), - ) - tail: Final = "/".join(segments[2:]) + tail: Final = "/".join(segments[1:]) return _render( found, - stream=_request_wants_stream(mount_endpoint, tail, body), + stream=_request_wants_stream(endpoint, tail, body), requested_model=_request_model(body, tail, found), path_tail=tail, ) diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index c24212c489c..5374d420b6a 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -18,14 +18,12 @@ from starlette.responses import JSONResponse, Response from starlette.routing import Route from _fake_openai_endpoint_server import chat_completions, completions, embeddings, health, moderations -from integration._support.scripted_wires import ( +from integration._support.scripted_shapes import ( RenderedResponse, Scenario, ScenarioDeleted, ScenarioRegistered, ScenarioStore, - WIRES, - Wire, render, ) @@ -211,14 +209,10 @@ CONTROL_URL: Final = os.environ.get("INTEGRATION_UPSTREAM_URL", "http://127.0.0. @dataclass(frozen=True, slots=True) class ScenarioHandle: scenario_id: str - wire: Wire control_url: str def api_base(self) -> str: - return f"{self.control_url}/{self.scenario_id}/{self._mount()}" - - def _mount(self) -> str: - return WIRES[self.wire].mount + return f"{self.control_url}/{self.scenario_id}" def register_scenario(scenario: Scenario) -> ScenarioHandle: @@ -232,7 +226,6 @@ def register_scenario(scenario: Scenario) -> ScenarioHandle: result: Final = ScenarioRegistered.model_validate_json(response.content) return ScenarioHandle( scenario_id=result.scenario_id, - wire=scenario.wire, control_url=CONTROL_URL, ) diff --git a/tests/integration/_support/wires.json b/tests/integration/_support/wires.json deleted file mode 100644 index b298ccd33aa..00000000000 --- a/tests/integration/_support/wires.json +++ /dev/null @@ -1,119 +0,0 @@ -{ - "openai_chat": { - "shape": "openai_chat", - "mount": "openai", - "usage": [ - "cache_read_tokens", - "reasoning_tokens", - "audio_input_tokens", - "audio_output_tokens", - "web_search_calls" - ], - "terminals": [] - }, - "openai_responses": { - "shape": "openai_responses", - "mount": "openai", - "usage": [ - "cache_read_tokens", - "reasoning_tokens", - "web_search_calls", - "file_search_calls" - ], - "terminals": [ - "incomplete", - "unvalidated" - ] - }, - "anthropic_messages": { - "shape": "anthropic_messages", - "mount": "anthropic", - "usage": [ - "cache_read_tokens", - "web_search_calls", - "cache_write_5m_tokens", - "cache_write_1h_tokens" - ], - "terminals": [] - }, - "gemini_generate": { - "shape": "gemini_generate", - "mount": "gemini", - "usage": [ - "cache_read_tokens", - "reasoning_tokens", - "audio_input_tokens", - "audio_output_tokens", - "image_input_tokens", - "video_input_tokens", - "web_search_calls", - "google_maps_calls" - ], - "terminals": [ - "prompt_blocked" - ] - }, - "together_chat": { - "shape": "openai_chat", - "mount": "together", - "usage": [ - "cache_read_tokens", - "reasoning_tokens", - "audio_input_tokens", - "audio_output_tokens", - "web_search_calls" - ], - "terminals": [] - }, - "fireworks_chat": { - "shape": "openai_chat", - "mount": "fireworks", - "usage": [ - "cache_read_tokens", - "reasoning_tokens", - "audio_input_tokens", - "audio_output_tokens", - "web_search_calls" - ], - "terminals": [] - }, - "azure_chat": { - "shape": "openai_chat", - "mount": "azure", - "usage": [ - "cache_read_tokens", - "reasoning_tokens", - "audio_input_tokens", - "audio_output_tokens", - "web_search_calls" - ], - "terminals": [] - }, - "bedrock_converse": { - "shape": "bedrock_converse", - "mount": "bedrock", - "usage": [ - "cache_read_tokens", - "cache_write_5m_tokens", - "cache_write_1h_tokens" - ], - "terminals": [] - }, - "vertex_generate": { - "shape": "gemini_generate", - "mount": "vertex", - "usage": [ - "cache_read_tokens", - "reasoning_tokens", - "audio_input_tokens", - "audio_output_tokens", - "image_input_tokens", - "video_input_tokens", - "web_search_calls", - "google_maps_calls" - ], - "terminals": [ - "prompt_blocked" - ] - } -} diff --git a/tests/integration/cost_calculation/cases.json b/tests/integration/cost_calculation/cases.json index 8ff6783ae6c..478aa069f1e 100644 --- a/tests/integration/cost_calculation/cases.json +++ b/tests/integration/cost_calculation/cases.json @@ -3,49 +3,42 @@ { "litellm_provider": "openai", "mode": "chat", - "wire": "openai_chat", "model_prefix": "openai", "litellm_params": {} }, { "litellm_provider": "openai", "mode": "responses", - "wire": "openai_responses", "model_prefix": "openai/responses", "litellm_params": {} }, { "litellm_provider": "anthropic", "mode": "chat", - "wire": "anthropic_messages", "model_prefix": "anthropic", "litellm_params": {} }, { "litellm_provider": "gemini", "mode": "chat", - "wire": "gemini_generate", "model_prefix": null, "litellm_params": {} }, { "litellm_provider": "together_ai", "mode": "chat", - "wire": "together_chat", "model_prefix": null, "litellm_params": {} }, { "litellm_provider": "fireworks_ai", "mode": "chat", - "wire": "fireworks_chat", "model_prefix": null, "litellm_params": {} }, { "litellm_provider": "azure", "mode": "chat", - "wire": "azure_chat", "model_prefix": null, "litellm_params": { "api_version": "2025-04-01-preview" @@ -54,7 +47,6 @@ { "litellm_provider": "bedrock_converse", "mode": "chat", - "wire": "bedrock_converse", "model_prefix": "bedrock/converse", "litellm_params": { "aws_access_key_id": "AKIASCRIPTEDPROVIDER", @@ -65,7 +57,6 @@ { "litellm_provider": "vertex_ai-language-models", "mode": "chat", - "wire": "vertex_generate", "model_prefix": "vertex_ai", "litellm_params": { "vertex_project": "cc-scripted-project", diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index 0cbc837c184..9229bb47817 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -136,7 +136,7 @@ def register_scenario_deployment( **model.litellm_params, **( {"vertex_credentials": _vertex_service_account_json(control_url)} - if model.wire == "vertex_generate" + if model.llm_provider == "vertex_ai" else {} ), } diff --git a/tests/integration/cost_calculation/cost_matrix.py b/tests/integration/cost_calculation/cost_matrix.py index db054edd321..b261deb68b2 100644 --- a/tests/integration/cost_calculation/cost_matrix.py +++ b/tests/integration/cost_calculation/cost_matrix.py @@ -26,14 +26,22 @@ from pathlib import Path from types import MappingProxyType from typing import Final, Literal +from litellm import get_llm_provider +from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.azure.chat.gpt_transformation import AzureOpenAIConfig +from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig +from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig +from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager from pydantic import BaseModel, ConfigDict, Field, TypeAdapter -from integration._support.scripted_wires import ( - WIRES, +from integration._support.scripted_shapes import ( Scenario, + Shape, ScriptedOutput, ScriptedToolCall, ScriptedUsage, - Wire, ) COST_MAP_PATH: Final = Path(__file__).resolve().parent / "cost_map.json" @@ -151,8 +159,8 @@ def _entry_has_rate_key(entry: CostMapEntry, rate_key: str) -> bool: return value is not None -SERVICE_TIER_REQUEST_WIRES: Final = frozenset( - {"openai_chat", "azure_chat", "openai_responses", "bedrock_converse"} +SERVICE_TIER_REQUEST_SHAPES: Final = frozenset( + {"openai_chat", "openai_responses", "bedrock_converse"} ) @@ -240,7 +248,7 @@ class Case(BaseModel): def scenario(self, scenario_id: str, model: FrontierModel, text: str) -> Scenario: return Scenario( scenario_id=scenario_id, - wire=model.wire, + shape=model.shape, usage=self.usage_for(model.map_key), model=model.provider_model, output=ScriptedOutput( @@ -263,7 +271,6 @@ class _ProviderWiringRow(BaseModel): litellm_provider: str mode: str - wire: str model_prefix: str | None litellm_params: Mapping[str, str] @@ -284,26 +291,19 @@ _DEPLOYMENTS: Final[Mapping[str, DeploymentSpec]] = MappingProxyType( @dataclass(frozen=True, slots=True) -class _ProviderWiring: - """How a (litellm_provider, mode) pair maps to a provider wire, the provider - prefix on the registered litellm model string, and extra litellm_params.""" +class _DeploymentDefaults: + """How a (litellm_provider, mode) pair maps to deployment defaults.""" - wire: Wire model_prefix: str | None litellm_params: Mapping[str, str] -def _provider_wiring(rows: tuple[_ProviderWiringRow, ...]) -> Mapping[tuple[str, str], _ProviderWiring]: - unknown_wires: Final = sorted({row.wire for row in rows if row.wire not in WIRES}) - if unknown_wires: - raise ValueError( - f"cases.json providers has unknown wires: {unknown_wires}; " - f"known wires are {sorted(WIRES)}" - ) +def _deployment_defaults( + rows: tuple[_ProviderWiringRow, ...], +) -> Mapping[tuple[str, str], _DeploymentDefaults]: return MappingProxyType( { - (row.litellm_provider, row.mode): _ProviderWiring( - row.wire, + (row.litellm_provider, row.mode): _DeploymentDefaults( row.model_prefix, MappingProxyType(dict(row.litellm_params)), ) @@ -312,19 +312,22 @@ def _provider_wiring(rows: tuple[_ProviderWiringRow, ...]) -> Mapping[tuple[str, ) -_PROVIDER_WIRING: Final[Mapping[tuple[str, str], _ProviderWiring]] = _provider_wiring(CASES_FILE.providers) +_DEPLOYMENT_DEFAULTS: Final[Mapping[tuple[str, str], _DeploymentDefaults]] = _deployment_defaults( + CASES_FILE.providers +) @dataclass(frozen=True, slots=True) class FrontierModel: """One deployment under test, derived from a cost-map entry: the model_name - the suite registers, the provider-prefixed litellm model string, the wire - the scripted upstream speaks, and the sibling map model the response_model - override case reports.""" + the suite registers, the provider-prefixed litellm model string, the + response shape the scripted upstream speaks, and the sibling map model the + response_model override case reports.""" model_name: str litellm_model: str - wire: Wire + shape: Shape + llm_provider: str map_key: str override_model: str | None = None override_map_key: str | None = None @@ -343,7 +346,7 @@ class FrontierModel: # override can never repoint pricing there, same as a base_model pin. if ( self.base_model is not None - or self.wire == "bedrock_converse" + or self.shape == "bedrock_converse" or self.override_map_key is None ): return self.rates @@ -371,12 +374,35 @@ def _provider_model(litellm_model: str) -> str: return "/".join(tail[1:] if tail and tail[0] in ("converse", "responses") else tail) -def _litellm_model_for(map_key: str, wiring: _ProviderWiring) -> str: - if wiring.model_prefix is None: +def _litellm_model_for(map_key: str, defaults: _DeploymentDefaults) -> str: + if defaults.model_prefix is None: return map_key - if map_key.startswith(f"{wiring.model_prefix}/"): + if map_key.startswith(f"{defaults.model_prefix}/"): return map_key - return f"{wiring.model_prefix}/{map_key}" + return f"{defaults.model_prefix}/{map_key}" + + +def _resolve(litellm_model: str, mode: str) -> tuple[str, Shape]: + model, provider, _, _ = get_llm_provider(model=litellm_model) + llm_provider: Final = LlmProviders(provider) + if mode == "responses": + responses_config: Final = ProviderConfigManager.get_provider_responses_api_config( + model=model, + provider=llm_provider, + ) + if isinstance(responses_config, OpenAIResponsesAPIConfig): + return provider, "openai_responses" + raise ValueError(f"no scripted renderer for {type(responses_config).__name__} ({litellm_model})") + config: Final = ProviderConfigManager.get_provider_chat_config(model=model, provider=llm_provider) + if isinstance(config, AmazonConverseConfig): + return provider, "bedrock_converse" + if isinstance(config, VertexGeminiConfig): + return provider, "gemini_generate" + if isinstance(config, AnthropicConfig): + return provider, "anthropic_messages" + if isinstance(config, (AzureOpenAIConfig, OpenAIGPTConfig)): + return provider, "openai_chat" + raise ValueError(f"no scripted renderer for {type(config).__name__} ({litellm_model})") def _frontier() -> tuple[FrontierModel, ...]: @@ -390,26 +416,29 @@ def _frontier() -> tuple[FrontierModel, ...]: for map_key in sorted(COST_MAP): entry = COST_MAP[map_key] pair = (entry.litellm_provider, entry.mode) - wiring = _PROVIDER_WIRING.get(pair) - if wiring is None: + defaults = _DEPLOYMENT_DEFAULTS.get(pair) + if defaults is None: continue siblings = groups[pair] override_key = ( siblings[(siblings.index(map_key) + 1) % len(siblings)] if len(siblings) > 1 else None ) override_litellm = ( - _litellm_model_for(override_key, wiring) if override_key is not None else None + _litellm_model_for(override_key, defaults) if override_key is not None else None ) deployment = _DEPLOYMENTS.get(map_key) + litellm_model = ( + deployment.litellm_model + if deployment is not None and deployment.litellm_model is not None + else _litellm_model_for(map_key, defaults) + ) + llm_provider, shape = _resolve(litellm_model, entry.mode) models.append( FrontierModel( model_name=f"cc-{map_key.replace('/', '-').replace(':', '-').replace('.', '-').lower()}", - litellm_model=( - deployment.litellm_model - if deployment is not None and deployment.litellm_model is not None - else _litellm_model_for(map_key, wiring) - ), - wire=wiring.wire, + litellm_model=litellm_model, + shape=shape, + llm_provider=llm_provider, map_key=map_key, override_model=( _provider_model(override_litellm) @@ -418,7 +447,7 @@ def _frontier() -> tuple[FrontierModel, ...]: ), override_map_key=override_key, base_model=deployment.base_model if deployment is not None else None, - litellm_params=wiring.litellm_params, + litellm_params=defaults.litellm_params, ) ) return tuple(models) @@ -471,7 +500,7 @@ def audio_input_data_url() -> str: def video_input_data_url() -> str: """A deterministic mp4-looking blob (ftyp box plus a fixed mdat payload) - as a data URL; only the media type and bytes matter to the wire.""" + as a data URL; only the media type and bytes matter to the response.""" ftyp: Final = struct.pack(">I4s4sI4s4s", 24, b"ftyp", b"isom", 0x200, b"isom", b"iso6") mdat_payload: Final = bytes((i * 7 + 13) % 256 for i in range(4096)) mdat: Final = struct.pack(">I4s", 8 + len(mdat_payload), b"mdat") + mdat_payload @@ -570,7 +599,7 @@ def matrix_data_errors() -> tuple[str, ...]: f"(litellm_provider={entry.litellm_provider}, mode={entry.mode}); " f"add a providers row in cases.json" for map_key, entry in COST_MAP.items() - if (entry.litellm_provider, entry.mode) not in _PROVIDER_WIRING + if (entry.litellm_provider, entry.mode) not in _DEPLOYMENT_DEFAULTS ) input_rates: Final = tuple(entry.input_cost_per_token for entry in COST_MAP.values()) findings: Final = ( diff --git a/tests/integration/cost_calculation/test_token_pricing.py b/tests/integration/cost_calculation/test_token_pricing.py index 0b4e9948dfa..cc48da2b819 100644 --- a/tests/integration/cost_calculation/test_token_pricing.py +++ b/tests/integration/cost_calculation/test_token_pricing.py @@ -1,4 +1,4 @@ -"""Token pricing coverage for the integration scripted-wire cost shard.""" +"""Token pricing coverage for the integration scripted-shape cost shard.""" from __future__ import annotations @@ -9,7 +9,7 @@ import pytest from pydantic import JsonValue from integration._support.client import JSON_OBJECT, Gateway -from integration._support.scripted_wires import WIRES, ScriptedUsage, Wire +from integration._support.scripted_shapes import ScriptedUsage, Shape from integration.cost_calculation.conftest import ( approx_equal, assert_total_is_sum_of_components, @@ -20,7 +20,7 @@ from integration.cost_calculation.cost_matrix import ( AUDIO_INPUT_DATA_URL, FRONTIER_MODELS, IMAGE_INPUT_DATA_URL, - SERVICE_TIER_REQUEST_WIRES, + SERVICE_TIER_REQUEST_SHAPES, VIDEO_INPUT_DATA_URL, Case, FrontierModel, @@ -54,8 +54,8 @@ _CACHE_SHAPES: Final = frozenset({"anthropic_messages", "bedrock_converse"}) _WEB_SEARCH_OPTION_SHAPES: Final = frozenset({"openai_chat", "openai_responses"}) -def _cache_control(usage: ScriptedUsage, wire: Wire) -> dict[str, JsonValue] | None: - if WIRES[wire].shape not in _CACHE_SHAPES: +def _cache_control(usage: ScriptedUsage, shape: Shape) -> dict[str, JsonValue] | None: + if shape not in _CACHE_SHAPES: return None if not (usage.cache_read_tokens or usage.cache_write_5m_tokens or usage.cache_write_1h_tokens): return None @@ -107,18 +107,18 @@ def _chat_body(model: FrontierModel, case: Case, model_name: str, marker: str) - ), *( [{"type": "web_search_20250305", "name": "web_search", "max_uses": 5}] - if case.web_search is not None and model.wire == "anthropic_messages" + if case.web_search is not None and model.shape == "anthropic_messages" else [] ), *( [{"googleSearch": {}}] - if case.web_search is not None and model.wire in ("gemini_generate", "vertex_generate") + if case.web_search is not None and model.shape == "gemini_generate" else [] ), *([{"googleMaps": {}}] if case.google_maps else []), *([{"type": "file_search", "vector_store_ids": ["vs_cost_calc_fixture"]}] if case.file_search else []), ] - cache_control: Final = _cache_control(usage, model.wire) + cache_control: Final = _cache_control(usage, model.shape) message: Final = { "role": "system", "content": [ @@ -136,7 +136,7 @@ def _chat_body(model: FrontierModel, case: Case, model_name: str, marker: str) - **({"stream_options": {"include_usage": True}} if case.stream else {}), **( {"service_tier": case.service_tier} - if case.service_tier is not None and model.wire in SERVICE_TIER_REQUEST_WIRES + if case.service_tier is not None and model.shape in SERVICE_TIER_REQUEST_SHAPES else {} ), **({"reasoning_effort": "medium"} if case.reasoning else {}), @@ -148,15 +148,15 @@ def _chat_body(model: FrontierModel, case: Case, model_name: str, marker: str) - **({"audio": {"voice": "alloy", "format": "pcm16"}} if case.audio_output else {}), **( {"web_search_options": {"search_context_size": case.web_search}} - if case.web_search is not None and WIRES[model.wire].shape in _WEB_SEARCH_OPTION_SHAPES + if case.web_search is not None and model.shape in _WEB_SEARCH_OPTION_SHAPES else {} ), **({"tools": tools} if tools else {}), - **({"tool_choice": "auto"} if case.tool_call and model.wire != "bedrock_converse" else {}), + **({"tool_choice": "auto"} if case.tool_call and model.shape != "bedrock_converse" else {}), "allowed_openai_params": [ name for name, sent in ( - ("tool_choice", case.tool_call and model.wire != "bedrock_converse"), + ("tool_choice", case.tool_call and model.shape != "bedrock_converse"), ("modalities", case.audio_input or case.audio_output), ("audio", case.audio_output), ("web_search_options", case.web_search is not None),