From 9885dc89621697e235fc65e0e86e147da87398c1 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 22:52:42 +0000 Subject: [PATCH] test(e2e): add azure, bedrock converse and vertex wires to the cost suite Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/conftest.py | 54 +++- tests/e2e/cost_calculation/cost_matrix.py | 144 ++++++++- .../e2e/cost_calculation/scripted_provider.py | 284 +++++++++++++++++- .../cost_calculation/test_wire_formats_e2e.py | 64 ++++ tests/e2e/cost_map.json | 158 ++++++++++ tests/e2e/models.py | 1 + 6 files changed, 681 insertions(+), 24 deletions(-) diff --git a/tests/e2e/cost_calculation/conftest.py b/tests/e2e/cost_calculation/conftest.py index 345ca26f7e3..8c6db7c0010 100644 --- a/tests/e2e/cost_calculation/conftest.py +++ b/tests/e2e/cost_calculation/conftest.py @@ -12,6 +12,7 @@ Deselected unless E2E_COST_MAP_STACK is set (marker `cost_map_stack`). from __future__ import annotations import importlib.util +import json import sys from collections.abc import Callable, Mapping from dataclasses import dataclass @@ -22,7 +23,7 @@ from typing import Final, Protocol, cast import pytest from cost_matrix import Case, FrontierModel -from e2e_config import COST_MAP_PROXY_URL +from e2e_config import COST_MAP_PROXY_URL, SCRIPTED_PROVIDER_PROXY_BASE from lifecycle import ResourceManager from models import LiteLLMParamsBody, ModelInfoBody, ModelNewBody from proxy_client import ProxyClient, build_proxy_client @@ -111,6 +112,41 @@ def client() -> CostCalcClient: return CostCalcClient(proxy=proxy) +_vertex_key_pem: str | None = None + + +def _vertex_service_account_json() -> str: + """A service-account credential JSON whose token_uri is the sidecar's + /_oauth/token route: the proxy's google-auth refresh then gets a scripted + access token without touching Google. One generated RSA key per process.""" + global _vertex_key_pem # mutable-ok: session-scoped key generation cached for reuse + if _vertex_key_pem is None: + from cryptography.hazmat.primitives import serialization + from cryptography.hazmat.primitives.asymmetric import rsa + + _vertex_key_pem = ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + .decode() + ) + return json.dumps( + { + "type": "service_account", + "project_id": "cc-scripted-project", + "private_key_id": "scripted", + "private_key": _vertex_key_pem, + "client_email": "scripted@cc-scripted-project.iam.gserviceaccount.com", + "client_id": "0", + "auth_uri": f"{SCRIPTED_PROVIDER_PROXY_BASE}/_oauth/authorize", + "token_uri": f"{SCRIPTED_PROVIDER_PROXY_BASE}/_oauth/token", + } + ) + + def register_scenario_deployment( client: CostCalcClient, resources: ResourceManager, @@ -126,15 +162,21 @@ def register_scenario_deployment( handle: Final = register_scenario(scenario) resources.defer(lambda: delete_scenario(handle)) model_name: Final = f"{model.model_name}-{marker}" + extra_params: Final[dict[str, str]] = dict(model.litellm_params) + if model.wire == "vertex_generate": + extra_params["vertex_credentials"] = _vertex_service_account_json() model_id: Final = client.proxy.register_model( ModelNewBody( model_name=model_name, - litellm_params=LiteLLMParamsBody( - model=model.litellm_model, - api_key=model.api_key, - api_base=handle.api_base(), + litellm_params=LiteLLMParamsBody.model_validate( + { + "model": model.litellm_model, + "api_key": model.api_key, + "api_base": handle.api_base(), + **extra_params, + } ), - model_info=ModelInfoBody(), + model_info=ModelInfoBody(base_model=model.base_model), ) ) resources.defer(lambda: client.proxy.delete_model(model_id)) diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index 5f634712778..37495f37b0b 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -88,7 +88,14 @@ class FrontierModel: litellm_model: str wire: Wire map_key: str - override_model: str + override_model: str | None = None + override_map_key: str | None = None + # Registered as model_info.base_model; when set, the provider-reported + # model loses to it and every case bills at this deployment's own rates. + base_model: str | None = None + # Extra litellm_params merged into the /model/new registration (api_version, + # aws_* credentials, vertex_* auth). + litellm_params: Mapping[str, str] = MappingProxyType({}) @property def rates(self) -> CostMapEntry: @@ -96,11 +103,16 @@ class FrontierModel: @property def override_rates(self) -> CostMapEntry: + if self.base_model is not None or self.override_map_key is None: + return self.rates return _COST_MAP[self.override_map_key] @property - def override_map_key(self) -> str: - return _OVERRIDE_MAP_KEYS[self.override_model] + def provider_model(self) -> str: + """The bare provider-facing model name: litellm_model minus the provider + prefix and any routing segment (converse/, responses/).""" + tail: Final = self.litellm_model.split("/")[1:] + return "/".join(tail[1:] if tail and tail[0] in ("converse", "responses") else tail) @property def provider(self) -> str: @@ -166,6 +178,92 @@ _FRONTIER_SPECS: Final[tuple[tuple[str, str, Wire], ...]] = ( ) +@dataclass(frozen=True, slots=True) +class _ExtendedSpec: + """A frontier entry whose override target, model_info.base_model or extra + litellm_params can't be derived from the map key alone.""" + + map_key: str + litellm_model: str + wire: Wire + override_model: str | None = None + override_map_key: str | None = None + base_model: str | None = None + litellm_params: Mapping[str, str] = MappingProxyType({}) + + +_AZURE_PARAMS: Final[Mapping[str, str]] = MappingProxyType({"api_version": "2025-04-01-preview"}) +_BEDROCK_PARAMS: Final[Mapping[str, str]] = MappingProxyType( + { + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + } +) +_VERTEX_PARAMS: Final[Mapping[str, str]] = MappingProxyType( + { + "vertex_project": "cc-scripted-project", + "vertex_location": "us-central1", + } +) + +_EXTENDED_SPECS: Final[tuple[_ExtendedSpec, ...]] = ( + _ExtendedSpec( + map_key="azure/gpt-5.6", + litellm_model="azure/gpt-5.6", + wire="azure_chat", + override_model="gpt-5.4-mini", + override_map_key="azure/gpt-5.4-mini", + litellm_params=_AZURE_PARAMS, + ), + _ExtendedSpec( + # Deployment name is not a model; base_model pins billing so the + # response's model field loses, proving base_model wins. + map_key="azure/gpt-5.4-mini", + litellm_model="azure/cc-pinned-deployment", + wire="azure_chat", + override_model="gpt-5.6", + override_map_key="azure/gpt-5.6", + base_model="azure/gpt-5.4-mini", + litellm_params=_AZURE_PARAMS, + ), + _ExtendedSpec( + map_key="anthropic.claude-sonnet-5-v1:0", + litellm_model="bedrock/converse/anthropic.claude-sonnet-5-v1:0", + wire="bedrock_converse", + litellm_params=_BEDROCK_PARAMS, + ), + _ExtendedSpec( + map_key="us.anthropic.claude-opus-5-v1:0", + litellm_model="bedrock/converse/us.anthropic.claude-opus-5-v1:0", + wire="bedrock_converse", + litellm_params=_BEDROCK_PARAMS, + ), + _ExtendedSpec( + map_key="meta.llama4-maverick-17b-instruct-v1:0", + litellm_model="bedrock/converse/meta.llama4-maverick-17b-instruct-v1:0", + wire="bedrock_converse", + litellm_params=_BEDROCK_PARAMS, + ), + _ExtendedSpec( + map_key="gemini-3.8-flash", + litellm_model="vertex_ai/gemini-3.8-flash", + wire="vertex_generate", + override_model="gemini-3.1-pro-preview", + override_map_key="gemini-3.1-pro-preview", + litellm_params=_VERTEX_PARAMS, + ), + _ExtendedSpec( + map_key="gemini-3.1-pro-preview", + litellm_model="vertex_ai/gemini-3.1-pro-preview", + wire="vertex_generate", + override_model="gemini-3.8-flash", + override_map_key="gemini-3.8-flash", + litellm_params=_VERTEX_PARAMS, + ), +) + + def _frontier() -> tuple[FrontierModel, ...]: return tuple( FrontierModel( @@ -174,8 +272,21 @@ def _frontier() -> tuple[FrontierModel, ...]: wire=wire, map_key=map_key, override_model=_OVERRIDE_MODELS[map_key], + override_map_key=_OVERRIDE_MAP_KEYS[_OVERRIDE_MODELS[map_key]], ) for map_key, litellm_model, wire in _FRONTIER_SPECS + ) + tuple( + FrontierModel( + model_name=f"cc-{spec.map_key.replace('/', '-').replace(':', '-').replace('.', '-').lower()}", + litellm_model=spec.litellm_model, + wire=spec.wire, + map_key=spec.map_key, + override_model=spec.override_model, + override_map_key=spec.override_map_key, + base_model=spec.base_model, + litellm_params=spec.litellm_params, + ) + for spec in _EXTENDED_SPECS ) @@ -219,6 +330,24 @@ _WIRE_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType({ "web_search", "response_model", "absent_usage", "tool_call", "image_input", } ), + "azure_chat": frozenset( + { + "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", + "web_search", "response_model", "absent_usage", "tool_call", "image_input", + } + ), + "bedrock_converse": frozenset( + { + "cache_read", "cache_write_5m", "cache_write_1h", "absent_usage", + "tool_call", "image_input", + } + ), + "vertex_generate": frozenset( + { + "cache_read", "reasoning", "audio", "web_search", "response_model", + "absent_usage", "tool_call", "image_input", "prompt_blocked", + } + ), }) CaseName: TypeAlias = Literal[ @@ -270,6 +399,7 @@ class Case: scenario_id=scenario_id, wire=model.wire, usage=self.usage, + model=model.provider_model, output=ScriptedOutput( text=text, response_model=model.override_model if self.response_model_override else None, @@ -296,7 +426,9 @@ _PROMPT_BLOCKED_USAGE: Final = ScriptedUsage(fresh_input_tokens=1000, output_tok def _web_search_case(model: FrontierModel) -> Case: - counts_exactly: Final = model.wire in ("openai_responses", "anthropic_messages", "gemini_generate") + counts_exactly: Final = model.wire in ( + "openai_responses", "anthropic_messages", "gemini_generate", "vertex_generate" + ) return Case( name="web_search", usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, web_search_calls=3), @@ -611,12 +743,12 @@ def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: wire's normalization: Anthropic folds cache read/write into prompt_tokens, everyone else reports the totals the wire emitted.""" u: Final = case.usage - if model.wire == "anthropic_messages": + if model.wire in ("anthropic_messages", "bedrock_converse"): return ( u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens, u.output_tokens, ) - if model.wire == "gemini_generate": + if model.wire in ("gemini_generate", "vertex_generate"): return ( u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens, u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py index e1a6c430307..00230fabeba 100644 --- a/tests/e2e/cost_calculation/scripted_provider.py +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -15,10 +15,15 @@ Layout on one port: - ``GET /health`` liveness - ``POST /_scenarios`` register a Scenario JSON, returns its id - ``DELETE /_scenarios/`` remove it +- ``POST /_oauth/token`` fake Google OAuth token endpoint for the + Vertex service-account credential's refresh call - ``POST ///`` provider wire; mount is one of - ``openai``, ``anthropic``, ``gemini``, ``together``, ``fireworks`` and the - remainder is whatever path the provider client appends (``chat/completions``, - ``responses``, ``v1/messages``, ``models/:generateContent`` ...) + ``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`` A request carrying ``"stream": true`` (or the ``:streamGenerateContent`` Gemini verb) gets an SSE answer; ``stream_usage`` on the Scenario decides whether the @@ -28,15 +33,17 @@ final stream chunk carries usage or the provider reports none. from __future__ import annotations import json +import struct import sys import threading import time +import zlib from collections.abc import Mapping from dataclasses import dataclass from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from types import MappingProxyType from typing import Final, Literal, TypeAlias -from urllib.parse import urlsplit +from urllib.parse import unquote, urlsplit from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, model_validator @@ -47,6 +54,9 @@ Wire: TypeAlias = Literal[ "gemini_generate", "together_chat", "fireworks_chat", + "azure_chat", + "bedrock_converse", + "vertex_generate", ] WIRE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( @@ -57,6 +67,9 @@ WIRE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( "gemini_generate": "gemini", "together_chat": "together", "fireworks_chat": "fireworks", + "azure_chat": "azure", + "bedrock_converse": "bedrock", + "vertex_generate": "vertex", } ) @@ -69,6 +82,7 @@ _TERMINAL_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType( { "openai_responses": frozenset({"incomplete", "unvalidated"}), "gemini_generate": frozenset({"prompt_blocked"}), + "vertex_generate": frozenset({"prompt_blocked"}), } ) @@ -131,6 +145,10 @@ class Scenario(BaseModel): wire: Wire usage: ScriptedUsage output: ScriptedOutput + # The bare provider-facing model name the renderer echoes when the request + # carries no model of its own (Vertex and Bedrock name the model in the URL + # path, not the body). + model: str stream_usage: StreamUsage = "final_chunk" service_tier: ServiceTier | None = None @@ -904,7 +922,208 @@ def _responses_sse(scenario: Scenario, requested_model: str) -> bytes: ) -def _render(scenario: Scenario, *, stream: bool, requested_model: str) -> RenderedResponse: +def _bedrock_usage(u: ScriptedUsage) -> Mapping[str, object]: + # Converse reports uncached input in inputTokens and rides cache reads and + # writes on top-level fields; totalTokens covers every input kind + output. + cache_writes: Final = u.cache_write_5m_tokens + u.cache_write_1h_tokens + return _jobj_opt( + ("inputTokens", u.fresh_input_tokens), + ("outputTokens", u.output_tokens), + ( + "totalTokens", + u.fresh_input_tokens + u.cache_read_tokens + cache_writes + u.output_tokens, + ), + ("cacheReadInputTokens", u.cache_read_tokens) if u.cache_read_tokens else None, + ("cacheWriteInputTokens", cache_writes) if cache_writes else None, + ( + ( + "cacheDetails", + tuple( + _jobj(("inputTokens", count), ("ttl", ttl)) + for count, ttl in ( + (u.cache_write_5m_tokens, "5m"), + (u.cache_write_1h_tokens, "1h"), + ) + if count + ), + ) + if cache_writes + else None + ), + ) + + +def _bedrock_stop_reason(scenario: Scenario) -> str: + if scenario.output.tool_call is not None: + return "tool_use" + return "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason + + +def _bedrock_content(scenario: Scenario) -> tuple[Mapping[str, object], ...]: + tool_call: Final = scenario.output.tool_call + if tool_call is not None: + return ( + _jobj( + ( + "toolUse", + _jobj( + ("toolUseId", f"tooluse_{scenario.scenario_id}"), + ("name", tool_call.name), + ("input", json.loads(tool_call.arguments)), + ), + ), + ), + ) + return (_jobj(("text", scenario.output.text)),) + + +def _bedrock_body(scenario: Scenario) -> Mapping[str, object]: + return _jobj( + ( + "output", + _jobj( + ( + "message", + _jobj( + ("role", "assistant"), + ("content", _bedrock_content(scenario)), + ), + ), + ), + ), + ("stopReason", _bedrock_stop_reason(scenario)), + ("usage", _bedrock_usage(scenario.usage)), + ("metrics", _jobj(("latencyMs", 42))), + ) + + +def _aws_event_frame(event_type: str, payload: Mapping[str, object]) -> bytes: + """One application/vnd.amazon.eventstream frame: prelude + prelude CRC32 + + headers + JSON payload + message CRC32, matching botocore EventStreamBuffer.""" + try: + from botocore.eventstream import crc32 as _crc32 + except ImportError: + _crc32 = zlib.crc32 + + def _str_header(name: str, value: str) -> bytes: + name_b: Final = name.encode() + value_b: Final = value.encode() + return ( + struct.pack("!B", len(name_b)) + + name_b + + struct.pack("!B", 7) + + struct.pack("!H", len(value_b)) + + value_b + ) + + payload_bytes: Final = json.dumps(payload, default=dict, separators=(",", ":")).encode() + headers_bytes: Final = ( + _str_header(":event-type", event_type) + + _str_header(":content-type", "application/json") + + _str_header(":message-type", "event") + ) + total_length: Final = 12 + len(headers_bytes) + len(payload_bytes) + 4 + prelude: Final = struct.pack("!II", total_length, len(headers_bytes)) + prelude_crc: Final = struct.pack("!I", _crc32(prelude) & 0xFFFFFFFF) + message: Final = prelude + prelude_crc + headers_bytes + payload_bytes + return message + struct.pack("!I", _crc32(message, 0) & 0xFFFFFFFF) + + +def _bedrock_eventstream(scenario: Scenario) -> bytes: + tool_call: Final = scenario.output.tool_call + block_start: Final[tuple[bytes, ...]] = ( + ( + _aws_event_frame( + "contentBlockStart", + _jobj( + ( + "start", + _jobj( + ( + "toolUse", + _jobj( + ("toolUseId", f"tooluse_{scenario.scenario_id}"), + ("name", tool_call.name), + ), + ), + ), + ), + ("contentBlockIndex", 0), + ), + ), + ) + if tool_call is not None + else () + ) + deltas: Final[tuple[bytes, ...]] = ( + tuple( + _aws_event_frame( + "contentBlockDelta", + _jobj( + ("delta", _jobj(("toolUse", _jobj(("input", arguments_slice))))), + ("contentBlockIndex", 0), + ), + ) + for arguments_slice in _split_arguments(tool_call.arguments) + ) + if tool_call is not None + else ( + _aws_event_frame( + "contentBlockDelta", + _jobj( + ("delta", _jobj(("text", scenario.output.text))), + ("contentBlockIndex", 0), + ), + ), + ) + ) + return b"".join( + ( + _aws_event_frame("messageStart", _jobj(("role", "assistant"))), + *block_start, + *deltas, + _aws_event_frame("contentBlockStop", _jobj(("contentBlockIndex", 0))), + _aws_event_frame("messageStop", _jobj(("stopReason", _bedrock_stop_reason(scenario)))), + *( + ( + _aws_event_frame( + "metadata", + _jobj( + ("usage", _bedrock_usage(scenario.usage)), + ("metrics", _jobj(("latencyMs", 42))), + ), + ), + ) + if scenario.stream_usage == "final_chunk" + else () + ), + ) + ) + + +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"): + if stream: + return RenderedResponse( + 200, "text/event-stream", _responses_sse(scenario, requested_model) + ) + return RenderedResponse( + 200, "application/json", _json_bytes(_responses_body(scenario, requested_model)) + ) + if scenario.wire == "bedrock_converse": + if stream: + return RenderedResponse( + 200, "application/vnd.amazon.eventstream", _bedrock_eventstream(scenario) + ) + return RenderedResponse(200, "application/json", _json_bytes(_bedrock_body(scenario))) + if scenario.wire == "vertex_generate": + if stream: + return RenderedResponse(200, "text/event-stream", _gemini_sse(scenario, requested_model)) + return RenderedResponse(200, "application/json", _json_bytes(_gemini_body(scenario, requested_model))) if scenario.wire == "anthropic_messages": if stream: return RenderedResponse(200, "text/event-stream", _anthropic_sse(scenario, requested_model)) @@ -917,7 +1136,8 @@ def _render(scenario: Scenario, *, stream: bool, requested_model: str) -> Render if stream: return RenderedResponse(200, "text/event-stream", _responses_sse(scenario, requested_model)) return RenderedResponse(200, "application/json", _json_bytes(_responses_body(scenario, requested_model))) - # openai_chat, together_chat, fireworks_chat share the OpenAI chat shape. + # openai_chat, together_chat, fireworks_chat and azure_chat share the + # OpenAI chat shape. if stream: return RenderedResponse(200, "text/event-stream", _openai_chat_sse(scenario, requested_model)) return RenderedResponse(200, "application/json", _json_bytes(_openai_chat_body(scenario, requested_model))) @@ -954,17 +1174,28 @@ def _request_body(body: bytes) -> Mapping[str, object]: return MappingProxyType({}) -def _request_wants_stream(path_tail: str, body: bytes) -> bool: - if ":streamGenerateContent" in path_tail: +def _request_wants_stream(mount_endpoint: str | None, path_tail: str, body: bytes) -> bool: + if mount_endpoint == "streamGenerateContent" or ":streamGenerateContent" in path_tail: + return True + if path_tail.endswith("converse-stream"): return True if not body: return False return _request_body(body).get("stream") is True -def _request_model(body: bytes) -> str: +def _request_model(body: bytes, path_tail: str, scenario: Scenario) -> str: model: Final = _request_body(body).get("model") - return model if isinstance(model, str) else "unknown" + if isinstance(model, str): + return model + # Bedrock Converse names the model in the path: model//converse[-stream]. + if path_tail.startswith("model/"): + 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. + return scenario.model def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: bytes) -> RenderedResponse: @@ -972,6 +1203,22 @@ def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: byte segments: Final = tuple(segment for segment in path.split("/") if segment) if method == "GET" and segments == ("health",): return RenderedResponse(200, "application/json", _json_bytes(_jobj(("status", "ok")))) + if segments and segments[0] == "_oauth": + if method == "POST" and segments == ("_oauth", "token"): + return RenderedResponse( + 200, + "application/json", + _json_bytes( + _jobj( + ("access_token", "scripted-token"), + ("token_type", "Bearer"), + ("expires_in", 3600), + ) + ), + ) + return RenderedResponse( + 404, "application/json", _json_bytes(_jobj(("error", "unknown control route"))) + ) if segments and segments[0] == "_scenarios": if method == "POST" and len(segments) == 1: try: @@ -998,7 +1245,15 @@ def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: byte return RenderedResponse( 404, "application/json", _json_bytes(_jobj(("error", f"no route for {method} {path}"))) ) - scenario_id, mount = segments[0], segments[1] + 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) + ) found: Final = store.get(scenario_id) if found is None: return RenderedResponse( @@ -1013,7 +1268,12 @@ def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: byte ), ) tail: Final = "/".join(segments[2:]) - return _render(found, stream=_request_wants_stream(tail, body), requested_model=_request_model(body)) + return _render( + found, + stream=_request_wants_stream(mount_endpoint, tail, body), + requested_model=_request_model(body, tail, found), + path_tail=tail, + ) class _ScriptedHandler(BaseHTTPRequestHandler): diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py index 4da7b31a6ef..a36bb1a8662 100644 --- a/tests/e2e/cost_calculation/test_wire_formats_e2e.py +++ b/tests/e2e/cost_calculation/test_wire_formats_e2e.py @@ -94,6 +94,40 @@ _WIRE_USAGE: Final[Mapping[str, tuple[str, ScriptedUsage]]] = MappingProxyType({ "fireworks_ai/kimi-k3", ScriptedUsage(fresh_input_tokens=80, cache_read_tokens=40, output_tokens=25), ), + "azure_chat": ( + "azure/gpt-5.6", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + cache_write_5m_tokens=20, + cache_write_1h_tokens=10, + output_tokens=25, + reasoning_tokens=15, + audio_input_tokens=5, + audio_output_tokens=3, + ), + ), + "bedrock_converse": ( + "anthropic.claude-sonnet-5-v1:0", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + cache_write_5m_tokens=20, + cache_write_1h_tokens=10, + output_tokens=25, + ), + ), + "vertex_generate": ( + "gemini-3.8-flash", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + output_tokens=25, + reasoning_tokens=15, + audio_input_tokens=5, + audio_output_tokens=3, + ), + ), }) _SHAPE_USAGE: Final = ScriptedUsage(fresh_input_tokens=80, output_tokens=25) @@ -141,6 +175,36 @@ _SHAPES: Final[tuple[tuple[str, str, Case], ...]] = ( response_model_override=True, ), ), + ( + "vertex_prompt_blocked", + "vertex_generate", + Case( + name="prompt_blocked", + usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), + terminal="prompt_blocked", + response_model_override=True, + ), + ), + ( + "vertex_prompt_blocked_stream", + "vertex_generate", + Case( + name="stream_prompt_blocked", + usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), + stream=True, + terminal="prompt_blocked", + response_model_override=True, + ), + ), + ( + "azure_served_model_override", + "azure_chat", + Case( + name="response_model_override", + usage=_SHAPE_USAGE, + response_model_override=True, + ), + ), ) diff --git a/tests/e2e/cost_map.json b/tests/e2e/cost_map.json index 68d840870c9..4fba337b701 100644 --- a/tests/e2e/cost_map.json +++ b/tests/e2e/cost_map.json @@ -304,6 +304,149 @@ "supports_reasoning": true, "supports_web_search": true }, + "anthropic.claude-sonnet-5-v1:0": { + "cache_creation_input_token_cost": 0.00051, + "cache_creation_input_token_cost_above_1hr": 0.00068, + "cache_read_input_token_cost": 1.7e-05, + "input_cost_per_token": 0.00017, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00034, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true + }, + "azure/gpt-5.4-mini": { + "cache_creation_input_token_cost": 0.00048, + "cache_creation_input_token_cost_above_1hr": 0.00064, + "cache_read_input_token_cost": 1.6e-05, + "input_cost_per_audio_token": 0.00096, + "input_cost_per_token": 0.00016, + "input_cost_per_token_above_200k_tokens": 0.00128, + "input_cost_per_token_flex": 0.00024, + "input_cost_per_token_priority": 0.000272, + "litellm_provider": "azure", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00112, + "output_cost_per_reasoning_token": 0.0008, + "output_cost_per_token": 0.00032, + "output_cost_per_token_above_200k_tokens": 0.00144, + "output_cost_per_token_flex": 0.0004, + "output_cost_per_token_priority": 0.000432, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "azure/gpt-5.6": { + "cache_creation_input_token_cost": 0.00044999999999999996, + "cache_creation_input_token_cost_above_1hr": 0.0006000000000000001, + "cache_read_input_token_cost": 1.5e-05, + "input_cost_per_audio_token": 0.0009000000000000001, + "input_cost_per_token": 0.00015000000000000001, + "input_cost_per_token_above_200k_tokens": 0.0012000000000000001, + "input_cost_per_token_flex": 0.000225, + "input_cost_per_token_priority": 0.000255, + "litellm_provider": "azure", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.0010500000000000002, + "output_cost_per_reasoning_token": 0.00075, + "output_cost_per_token": 0.00030000000000000003, + "output_cost_per_token_above_200k_tokens": 0.00135, + "output_cost_per_token_flex": 0.000375, + "output_cost_per_token_priority": 0.00040499999999999996, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true + }, + "gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 2.1e-05, + "input_cost_per_audio_token": 0.00126, + "input_cost_per_token": 0.00021, + "input_cost_per_token_above_200k_tokens": 0.00168, + "input_cost_per_token_flex": 0.000315, + "input_cost_per_token_priority": 0.000357, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00147, + "output_cost_per_reasoning_token": 0.0010500000000000002, + "output_cost_per_token": 0.00042, + "output_cost_per_token_above_200k_tokens": 0.0018900000000000001, + "output_cost_per_token_flex": 0.000525, + "output_cost_per_token_priority": 0.000567, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "gemini-3.8-flash": { + "cache_read_input_token_cost": 2e-05, + "input_cost_per_audio_token": 0.0012, + "input_cost_per_token": 0.0002, + "input_cost_per_token_above_200k_tokens": 0.0016, + "input_cost_per_token_flex": 0.0003, + "input_cost_per_token_priority": 0.00034, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.0014000000000000002, + "output_cost_per_reasoning_token": 0.001, + "output_cost_per_token": 0.0004, + "output_cost_per_token_above_200k_tokens": 0.0018000000000000001, + "output_cost_per_token_flex": 0.0005, + "output_cost_per_token_priority": 0.00054, + "search_context_cost_per_query": { + "search_context_size_high": 0.03, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02 + }, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "meta.llama4-maverick-17b-instruct-v1:0": { + "input_cost_per_token": 0.00019, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00038, + "supports_function_calling": true + }, "together_ai/moonshotai/Kimi-K3": { "cache_creation_input_token_cost": 0.00030000000000000003, "cache_creation_input_token_cost_above_1hr": 0.0004, @@ -363,5 +506,20 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_web_search": true + }, + "us.anthropic.claude-opus-5-v1:0": { + "cache_creation_input_token_cost": 0.0005400000000000001, + "cache_creation_input_token_cost_above_1hr": 0.00072, + "cache_read_input_token_cost": 1.8e-05, + "input_cost_per_token": 0.00018, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00036000000000000004, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true } } diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 7101438c5f8..d96478de1c4 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1000,6 +1000,7 @@ class ModelInfoBody(BaseModel): access_groups: list[str] | None = None team_id: str | None = None allowed_fails_policy: dict[str, int] | None = None + base_model: str | None = None class ModelNewBody(BaseModel):