scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 73

test(integration): derive scripted shapes from litellm provider configs

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-09-19 03:16:16 +00:00
parent 380ec1a004
commit 57d2fefa8d
8 changed files with 196 additions and 282 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
tests/integration/README.md
View file

@ -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

198
tests/integration/_support/scripted_wires.py → tests/integration/_support/scripted_shapes.py
View file

@ -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/<id>`` remove it
- ``POST /<id>/<mount>/<provider path>`` 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/<m>:generateContent`` ...). Vertex appends ``:generateContent`` /
``:streamGenerateContent`` to the mount segment itself, and Bedrock Converse
targets ``model/<modelId>/converse`` / ``converse-stream``
- ``POST /<id>/<provider path>`` provider response; the remainder is whatever
path the provider client appends (``chat/completions``, ``responses``,
``v1/messages``, ``models/<m>:generateContent`` ...). Vertex appends
``:generateContent`` / ``:streamGenerateContent`` to the scenario segment,
and Bedrock Converse targets ``model/<modelId>/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,
)

11
tests/integration/_support/upstream.py
View file

@ -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,
)

119
tests/integration/_support/wires.json
View file

@ -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"
]
}
}

9
tests/integration/cost_calculation/cases.json
View file

@ -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",

2
tests/integration/cost_calculation/conftest.py
View file

@ -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 {}
),
}

113
tests/integration/cost_calculation/cost_matrix.py
View file

@ -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 = (

24
tests/integration/cost_calculation/test_token_pricing.py
View file

@ -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),