From 91761d984d55b2e4729e38c36ef7c7510c9ffc76 Mon Sep 17 00:00:00 2001 From: abhirup7 Date: Tue, 8 Sep 2026 00:30:00 +0530 Subject: [PATCH 001/251] fix(azure): send the resolved Entra ID token on image generation requests Azure image generation calls initialize_azure_sdk_client like the chat path does, but then sends the request through httpx with the headers it was given, so a credential resolved from litellm_params or the environment (Entra ID client credentials, managed or workload identity, OIDC, a static azure_ad_token) never reached the wire and Azure answered 401. Only an explicitly passed azure_ad_token_provider was applied Add get_azure_request_auth_headers, which turns the credential in azure_client_params into an Authorization: Bearer header (or api-key, following the SDK's precedence) while keeping any auth header the caller already set, and use it for both the sync and async image requests. The pre-call logging metadata receives a redacted copy of those headers so the token never reaches logging callbacks Fixes #16422 --- litellm/llms/azure/azure.py | 31 ++- litellm/llms/azure/common_utils.py | 35 +++ .../test_azure_image_generation_init.py | 249 ++++++++++++++++++ 3 files changed, 301 insertions(+), 14 deletions(-) diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 46a9dd1a531..bfbaffd972c 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -46,7 +46,9 @@ from .common_utils import ( AzureOpenAIError, BaseAzureLLM, get_azure_ad_token_from_oidc, + get_azure_request_auth_headers, process_azure_headers, + redact_azure_auth_headers, select_azure_base_url_or_endpoint, ) from .image_generation import ( @@ -1142,7 +1144,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): api_key: str, input: list, logging_obj: LiteLLMLoggingObj, - headers: dict, + headers: dict[str, str], client=None, timeout=None, model: str | None = None, @@ -1167,7 +1169,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): additional_args={ "complete_input_dict": data, "api_base": img_gen_api_base, - "headers": headers, + "headers": redact_azure_auth_headers(headers), }, ) httpx_response: Final[httpx.Response] = await self.make_async_azure_httpx_request( @@ -1226,7 +1228,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): timeout: float, optional_params: dict, logging_obj: LiteLLMLoggingObj, - headers: dict, + headers: dict[str, str], model: str | None = None, api_key: str | None = None, api_base: str | None = None, @@ -1261,21 +1263,22 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): if not isinstance(max_retries, int): raise AzureOpenAIError(status_code=422, message="max retries must be an int") - if api_key is None and azure_ad_token_provider is not None: - azure_ad_token = azure_ad_token_provider() - if azure_ad_token: - headers.pop("api-key", None) - headers["Authorization"] = f"Bearer {azure_ad_token}" - - # init AzureOpenAI Client + auth_params: Final[dict[str, object]] = {**(litellm_params or {})} # mutable-ok: SDK init takes a dict + if azure_ad_token is not None: + auth_params["azure_ad_token"] = azure_ad_token + if azure_ad_token_provider is not None: + auth_params["azure_ad_token_provider"] = azure_ad_token_provider azure_client_params: Final[dict[str, object]] = self.initialize_azure_sdk_client( - litellm_params=litellm_params or {}, + litellm_params=auth_params, api_key=api_key, model_name=model or "", api_version=api_version, api_base=api_base, is_async=False, ) + request_headers: Final = dict( # mutable-ok: the httpx request helpers take a dict + get_azure_request_auth_headers(headers=headers, azure_client_params=azure_client_params) + ) if aimg_generation is True: return self.aimage_generation( data=data, @@ -1286,7 +1289,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): client=client, azure_client_params=azure_client_params, timeout=timeout, - headers=headers, + headers=request_headers, model=model, ) @@ -1303,7 +1306,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): additional_args={ "complete_input_dict": data, "api_base": img_gen_api_base, - "headers": headers, + "headers": redact_azure_auth_headers(request_headers), }, ) httpx_response: Final[httpx.Response] = self.make_sync_azure_httpx_request( @@ -1313,7 +1316,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): api_version=api_version or "", api_key=api_key or "", data=data, - headers=headers, + headers=request_headers, deployment_name=model, ) provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2")) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 6cb7d09cec4..f276d8b18d1 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -405,6 +405,41 @@ def get_azure_ad_token( return azure_ad_token +_AZURE_AUTH_HEADER_NAMES: Final = frozenset(("api-key", "authorization")) +_REDACTED_AZURE_HEADER_VALUE: Final = "***REDACTED***" + + +def _resolve_azure_ad_token(azure_client_params: Mapping[str, object]) -> str | None: + azure_ad_token: Final = azure_client_params.get("azure_ad_token") + if isinstance(azure_ad_token, str) and azure_ad_token: + return azure_ad_token + token_provider: Final = azure_client_params.get("azure_ad_token_provider") + provided_token: Final = token_provider() if callable(token_provider) else None + return provided_token if isinstance(provided_token, str) and provided_token else None + + +def get_azure_request_auth_headers( + headers: Mapping[str, str], + azure_client_params: Mapping[str, object], +) -> Mapping[str, str]: + if any(name.lower() in _AZURE_AUTH_HEADER_NAMES for name in headers): + return headers + azure_ad_token: Final = _resolve_azure_ad_token(azure_client_params) + if azure_ad_token is not None: + return MappingProxyType({**headers, "Authorization": f"Bearer {azure_ad_token}"}) + api_key: Final = azure_client_params.get("api_key") + if isinstance(api_key, str) and api_key: + return MappingProxyType({**headers, "api-key": api_key}) + return headers + + +def redact_azure_auth_headers(headers: Mapping[str, str]) -> Mapping[str, str]: + return { # mutable-ok: logging callbacks JSON-serialize this copy + name: (_REDACTED_AZURE_HEADER_VALUE if name.lower() in _AZURE_AUTH_HEADER_NAMES else value) + for name, value in headers.items() + } + + class BaseAzureLLM(BaseOpenAILLM): @staticmethod def _try_get_default_azure_credential_provider( diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py index 70b5eab5c37..a30aa277f3d 100644 --- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py @@ -10,6 +10,11 @@ import respx import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.azure.azure import AzureChatCompletion +from litellm.llms.azure.common_utils import ( + _cached_entra_id_token_provider, + get_azure_request_auth_headers, + redact_azure_auth_headers, +) from litellm.llms.azure.image_generation.http_utils import ( azure_deployment_image_generation_json_body, ) @@ -587,3 +592,247 @@ def test_azure_image_generation_v1_route_base_model_vs_deployment_name(respx_moc sent_body = json.loads(request.content) assert sent_body["model"] == model assert sent_body["prompt"] == prompt + + +@pytest.fixture +def fake_entra_id(monkeypatch: pytest.MonkeyPatch): + built_credentials = [] + + class FakeClientSecretCredential: + def __init__(self, tenant_id: str, client_id: str, client_secret: str) -> None: + built_credentials.append((tenant_id, client_id, client_secret)) + + monkeypatch.setattr("azure.identity.ClientSecretCredential", FakeClientSecretCredential) + monkeypatch.setattr("azure.identity.get_bearer_token_provider", lambda credential, scope: lambda: "entra-id-token") + _cached_entra_id_token_provider.cache_clear() + yield built_credentials + _cached_entra_id_token_provider.cache_clear() + + +def _mock_image_generation_route(respx_mock: respx.MockRouter, api_base: str, model: str) -> respx.Route: + return respx_mock.post(f"{api_base}/openai/deployments/{model}/images/generations").mock( + return_value=httpx.Response(200, json={"created": 1234567890, "data": [{"b64_json": "aaaa"}]}) + ) + + +@pytest.mark.parametrize("credentials_in_litellm_params", [False, True]) +def test_azure_image_generation_keyless_entra_id_sends_bearer_token( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + fake_entra_id: list, + credentials_in_litellm_params: bool, +): + for name in ("AZURE_TENANT_ID", "AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET"): + monkeypatch.delenv(name, raising=False) + api_base = "https://my-resource.openai.azure.com" + api_version = "2025-04-01-preview" + litellm_params = {"api_base": api_base, "api_version": api_version} + if credentials_in_litellm_params: + litellm_params.update( + tenant_id="tenant-from-params", client_id="client-from-params", client_secret="secret-from-params" + ) + expected_credential = ("tenant-from-params", "client-from-params", "secret-from-params") + else: + monkeypatch.setenv("AZURE_TENANT_ID", "tenant-from-env") + monkeypatch.setenv("AZURE_CLIENT_ID", "client-from-env") + monkeypatch.setenv("AZURE_CLIENT_SECRET", "secret-from-env") + expected_credential = ("tenant-from-env", "client-from-env", "secret-from-env") + route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1") + logging_obj = MagicMock() + + response = AzureChatCompletion().image_generation( + prompt="a cat", + timeout=60.0, + optional_params={"n": 1, "size": "1024x1024"}, + logging_obj=logging_obj, + headers={"Content-Type": "application/json"}, + model="gpt-image-1", + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params=litellm_params, + ) + + request = route.calls.last.request + assert request.headers["Authorization"] == "Bearer entra-id-token" + assert "api-key" not in request.headers + assert fake_entra_id == [expected_credential] + assert response.data[0].b64_json == "aaaa" + logged_headers = logging_obj.pre_call.call_args.kwargs["additional_args"]["headers"] + assert logged_headers == {"Content-Type": "application/json", "Authorization": "***REDACTED***"} + assert "entra-id-token" not in str(logging_obj.pre_call.call_args) + + +@pytest.mark.asyncio +async def test_azure_aimage_generation_keyless_entra_id_sends_bearer_token( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, fake_entra_id: list +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + api_base = "https://my-resource.openai.azure.com" + api_version = "2025-04-01-preview" + route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1") + logging_obj = MagicMock() + + response = await AzureChatCompletion().image_generation( + prompt="a cat", + timeout=60.0, + optional_params={"n": 1, "size": "1024x1024"}, + logging_obj=logging_obj, + headers={"Content-Type": "application/json"}, + model="gpt-image-1", + api_key=None, + api_base=api_base, + api_version=api_version, + aimg_generation=True, + litellm_params={ + "api_base": api_base, + "api_version": api_version, + "tenant_id": "tenant-from-params", + "client_id": "client-from-params", + "client_secret": "secret-from-params", + }, + ) + + request = route.calls.last.request + assert request.headers["Authorization"] == "Bearer entra-id-token" + assert "api-key" not in request.headers + assert fake_entra_id == [("tenant-from-params", "client-from-params", "secret-from-params")] + assert response.data[0].b64_json == "aaaa" + logged_headers = logging_obj.pre_call.call_args.kwargs["additional_args"]["headers"] + assert logged_headers == {"Content-Type": "application/json", "Authorization": "***REDACTED***"} + assert "entra-id-token" not in str(logging_obj.pre_call.call_args) + + +@pytest.mark.parametrize( + "credential_kwargs, expected_authorization", + [ + ({"azure_ad_token": "static-ad-token"}, "Bearer static-ad-token"), + ({"azure_ad_token_provider": lambda: "provider-token"}, "Bearer provider-token"), + ], +) +def test_azure_image_generation_explicit_azure_ad_credential_sends_bearer_token( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + credential_kwargs: dict, + expected_authorization: str, +): + for name in ("AZURE_TENANT_ID", "AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET"): + monkeypatch.delenv(name, raising=False) + api_base = "https://my-resource.openai.azure.com" + api_version = "2025-04-01-preview" + route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1") + + response = AzureChatCompletion().image_generation( + prompt="a cat", + timeout=60.0, + optional_params={"n": 1, "size": "1024x1024"}, + logging_obj=MagicMock(), + headers={"Content-Type": "application/json"}, + model="gpt-image-1", + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"api_base": api_base, "api_version": api_version}, + **credential_kwargs, + ) + + request = route.calls.last.request + assert request.headers["Authorization"] == expected_authorization + assert "api-key" not in request.headers + assert response.data[0].b64_json == "aaaa" + + +def test_azure_image_generation_with_api_key_keeps_api_key_header( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, fake_entra_id: list +): + monkeypatch.setenv("AZURE_TENANT_ID", "tenant-from-env") + monkeypatch.setenv("AZURE_CLIENT_ID", "client-from-env") + monkeypatch.setenv("AZURE_CLIENT_SECRET", "secret-from-env") + api_base = "https://my-resource.openai.azure.com" + api_version = "2025-04-01-preview" + route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1") + logging_obj = MagicMock() + + response = AzureChatCompletion().image_generation( + prompt="a cat", + timeout=60.0, + optional_params={"n": 1, "size": "1024x1024"}, + logging_obj=logging_obj, + headers={"Content-Type": "application/json", "api-key": "sk-test"}, + model="gpt-image-1", + api_key="sk-test", + api_base=api_base, + api_version=api_version, + litellm_params={"api_base": api_base, "api_version": api_version}, + ) + + request = route.calls.last.request + assert request.headers["api-key"] == "sk-test" + assert "Authorization" not in request.headers + assert fake_entra_id == [] + assert response.data[0].b64_json == "aaaa" + assert logging_obj.pre_call.call_args.kwargs["additional_args"]["headers"]["api-key"] == "***REDACTED***" + + +@pytest.mark.parametrize( + "caller_auth_header", + [{"api-key": "caller-key"}, {"Authorization": "Bearer caller-token"}, {"authorization": "Bearer caller-token"}], +) +def test_get_azure_request_auth_headers_keeps_caller_auth_header(caller_auth_header: dict): + headers = {"Content-Type": "application/json", **caller_auth_header} + azure_client_params = { + "api_key": "sk-resolved", + "azure_ad_token": "resolved-token", + "azure_ad_token_provider": lambda: "provider-token", + } + assert get_azure_request_auth_headers(headers=headers, azure_client_params=azure_client_params) is headers + + +def test_get_azure_request_auth_headers_prefers_azure_ad_token_over_provider_and_api_key(): + headers = {"Content-Type": "application/json"} + azure_client_params = { + "api_key": "sk-resolved", + "azure_ad_token": "static-token", + "azure_ad_token_provider": lambda: "provider-token", + } + out = get_azure_request_auth_headers(headers=headers, azure_client_params=azure_client_params) + assert dict(out) == {"Content-Type": "application/json", "Authorization": "Bearer static-token"} + assert headers == {"Content-Type": "application/json"} + + +def test_get_azure_request_auth_headers_uses_token_provider_over_api_key(): + azure_client_params = {"api_key": "sk-resolved", "azure_ad_token": None, "azure_ad_token_provider": lambda: "pt"} + out = get_azure_request_auth_headers(headers={}, azure_client_params=azure_client_params) + assert dict(out) == {"Authorization": "Bearer pt"} + + +def test_get_azure_request_auth_headers_falls_back_to_api_key(): + azure_client_params = {"api_key": "sk-resolved", "azure_ad_token": None, "azure_ad_token_provider": None} + out = get_azure_request_auth_headers(headers={"Content-Type": "application/json"}, azure_client_params=azure_client_params) + assert dict(out) == {"Content-Type": "application/json", "api-key": "sk-resolved"} + + +@pytest.mark.parametrize( + "azure_client_params", + [ + {}, + {"api_key": "", "azure_ad_token": "", "azure_ad_token_provider": None}, + {"azure_ad_token_provider": lambda: None}, + {"azure_ad_token_provider": lambda: ""}, + ], +) +def test_get_azure_request_auth_headers_without_credential_leaves_headers_unchanged(azure_client_params: dict): + headers = {"Content-Type": "application/json"} + assert get_azure_request_auth_headers(headers=headers, azure_client_params=azure_client_params) is headers + + +def test_redact_azure_auth_headers_masks_only_credential_values(): + headers = {"Content-Type": "application/json", "api-key": "sk-secret", "authorization": "Bearer secret"} + assert redact_azure_auth_headers(headers) == { + "Content-Type": "application/json", + "api-key": "***REDACTED***", + "authorization": "***REDACTED***", + } + assert headers["api-key"] == "sk-secret" + assert headers["authorization"] == "Bearer secret" From 66ebc722d6751a7a8d0aef751e8d2e5f6a2efe49 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 19:33:09 -0700 Subject: [PATCH 002/251] perf(azure): reuse the token refresh credential across image generation requests With enable_azure_ad_token_refresh, every keyless image request built a new DefaultAzureCredential and fetched a token. Cache the provider per scope like the Entra ID one. --- litellm/llms/azure/common_utils.py | 9 ++-- .../test_azure_image_generation_init.py | 47 +++++++++++++++++++ .../llms/azure/test_azure_common_utils.py | 3 ++ 3 files changed, 56 insertions(+), 3 deletions(-) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index f276d8b18d1..9f70761514b 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -95,6 +95,11 @@ def _cached_entra_id_token_provider( return get_bearer_token_provider(ClientSecretCredential(tenant_id, client_id, client_secret), scope) +@lru_cache(maxsize=128) +def _cached_azure_ad_token_refresh_provider(scope: str) -> Callable[[], str]: + return get_azure_ad_token_provider(azure_scope=scope) + + def get_azure_ad_token_from_entra_id( tenant_id: str, client_id: str, @@ -649,9 +654,7 @@ class BaseAzureLLM(BaseOpenAILLM): "Using Azure AD token provider based on Service Principal with Secret workflow for Azure Auth" ) try: - azure_ad_token_provider = get_azure_ad_token_provider( - azure_scope=scope, - ) + azure_ad_token_provider = _cached_azure_ad_token_refresh_provider(scope) except ValueError: verbose_logger.debug("Azure AD Token Provider could not be used.") if api_version is None: diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py index a30aa277f3d..cfde1760389 100644 --- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py @@ -11,6 +11,7 @@ import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.azure.azure import AzureChatCompletion from litellm.llms.azure.common_utils import ( + _cached_azure_ad_token_refresh_provider, _cached_entra_id_token_provider, get_azure_request_auth_headers, redact_azure_auth_headers, @@ -775,6 +776,52 @@ def test_azure_image_generation_with_api_key_keeps_api_key_header( assert logging_obj.pre_call.call_args.kwargs["additional_args"]["headers"]["api-key"] == "***REDACTED***" +@pytest.fixture +def fake_default_azure_credential(monkeypatch: pytest.MonkeyPatch): + built_credentials = [] + + class FakeDefaultAzureCredential: + def __init__(self) -> None: + built_credentials.append(self) + + for name in ("AZURE_TENANT_ID", "AZURE_CLIENT_ID", "AZURE_CLIENT_SECRET", "AZURE_CREDENTIAL", "AZURE_AD_TOKEN"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr("azure.identity.DefaultAzureCredential", FakeDefaultAzureCredential) + monkeypatch.setattr( + "azure.identity.get_bearer_token_provider", lambda credential, scope: lambda: "default-credential-token" + ) + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", True) + _cached_azure_ad_token_refresh_provider.cache_clear() + yield built_credentials + _cached_azure_ad_token_refresh_provider.cache_clear() + + +def test_azure_image_generation_token_refresh_reuses_credential_across_requests( + respx_mock: respx.MockRouter, fake_default_azure_credential: list +): + api_base = "https://my-resource.openai.azure.com" + api_version = "2025-04-01-preview" + route = _mock_image_generation_route(respx_mock, api_base, "gpt-image-1") + + for _ in range(3): + AzureChatCompletion().image_generation( + prompt="a cat", + timeout=60.0, + optional_params={"n": 1, "size": "1024x1024"}, + logging_obj=MagicMock(), + headers={"Content-Type": "application/json"}, + model="gpt-image-1", + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"api_base": api_base, "api_version": api_version}, + ) + + assert route.call_count == 3 + assert all(call.request.headers["Authorization"] == "Bearer default-credential-token" for call in route.calls) + assert len(fake_default_azure_credential) == 1 + + @pytest.mark.parametrize( "caller_auth_header", [{"api-key": "caller-key"}, {"Authorization": "Bearer caller-token"}, {"authorization": "Bearer caller-token"}], diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index f000abb4c9a..7189be7c052 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -9,6 +9,7 @@ import pytest import litellm from litellm.llms.azure.common_utils import ( BaseAzureLLM, + _cached_azure_ad_token_refresh_provider, _cached_entra_id_token_provider, get_azure_ad_token, get_azure_ad_token_from_entra_id, @@ -34,6 +35,7 @@ def setup_mocks(monkeypatch): monkeypatch.delenv("AZURE_TENANT_ID", raising=False) monkeypatch.delenv("AZURE_SCOPE", raising=False) monkeypatch.delenv("AZURE_AD_TOKEN", raising=False) + _cached_azure_ad_token_refresh_provider.cache_clear() with ( patch( @@ -78,6 +80,7 @@ def setup_mocks(monkeypatch): "logger": mock_logger, "select_url": mock_select_url, } + _cached_azure_ad_token_refresh_provider.cache_clear() def test_initialize_with_api_key(setup_mocks): From 233337628f4b6f1c9ec0527d5d442cfe512ac08b Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 18:06:25 -0400 Subject: [PATCH 003/251] feat(batches): support Mistral files/batches and per-page OCR batch cost tracking Adds MistralFilesConfig and MistralBatchesConfig so Mistral can be used as a Files and Batches provider through the shared BaseLLMHTTPHandler path, the same way Bedrock plugs in. /v1/ocr is now an accepted batch endpoint, and completed OCR batches are billed per page (ocr_cost_per_page_batches, half the synchronous rate) instead of per token. Resolves #29914 --- litellm/batches/batch_utils.py | 41 ++- litellm/batches/main.py | 20 +- litellm/cost_calculator.py | 61 ++++ litellm/files/main.py | 9 +- litellm/files/types.py | 2 +- litellm/llms/mistral/batches/__init__.py | 0 .../llms/mistral/batches/transformation.py | 186 +++++++++++++ litellm/llms/mistral/common_utils.py | 36 +++ litellm/llms/mistral/files/__init__.py | 0 litellm/llms/mistral/files/transformation.py | 226 +++++++++++++++ ...odel_prices_and_context_window_backup.json | 40 ++- litellm/types/llms/openai.py | 2 +- litellm/types/utils.py | 4 + litellm/utils.py | 10 + model_prices_and_context_window.json | 40 ++- .../test_litellm/batches/test_batch_utils.py | 83 ++++++ tests/test_litellm/batches/test_main.py | 42 +++ .../llms/mistral/batches/__init__.py | 0 .../test_mistral_batches_transformation.py | 260 ++++++++++++++++++ .../llms/mistral/files/__init__.py | 0 .../test_mistral_files_transformation.py | 189 +++++++++++++ .../llms/mistral/ocr/test_mistral_ocr_cost.py | 4 +- tests/test_litellm/test_utils.py | 2 + 23 files changed, 1216 insertions(+), 41 deletions(-) create mode 100644 litellm/llms/mistral/batches/__init__.py create mode 100644 litellm/llms/mistral/batches/transformation.py create mode 100644 litellm/llms/mistral/common_utils.py create mode 100644 litellm/llms/mistral/files/__init__.py create mode 100644 litellm/llms/mistral/files/transformation.py create mode 100644 tests/test_litellm/llms/mistral/batches/__init__.py create mode 100644 tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py create mode 100644 tests/test_litellm/llms/mistral/files/__init__.py create mode 100644 tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 87f8fd3946e..077d6e72fd7 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -9,6 +9,7 @@ import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details +from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo from litellm.types.llms.openai import Batch from litellm.types.utils import ModelInfo, Usage from litellm.utils import token_counter @@ -50,7 +51,7 @@ def batch_cost_is_final(batch: Batch) -> bool: async def calculate_batch_cost_and_usage( file_content_dictionary: list[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], model_name: str | None = None, model_info: ModelInfo | None = None, ) -> BatchCostUsageResult: @@ -80,7 +81,7 @@ async def calculate_batch_cost_and_usage( async def _handle_completed_batch( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], model_name: str | None = None, litellm_params: dict | None = None, model_info: ModelInfo | None = None, @@ -166,7 +167,7 @@ class _BatchOutputLineStats: def _classify_output_line_stats( entries: Iterable[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, model_info: ModelInfo | None, ) -> Iterator[_BatchOutputLineStats | _LineOutcome]: @@ -185,7 +186,7 @@ def _classify_output_line_stats( def _safe_output_line_stats( entry: Mapping[str, object], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, model_info: ModelInfo | None, ) -> _BatchOutputLineStats | None: @@ -207,7 +208,7 @@ def _safe_output_line_stats( def _compute_output_line_stats( entry: Mapping[str, object], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, model_info: ModelInfo | None, ) -> _BatchOutputLineStats: @@ -218,6 +219,7 @@ def _compute_output_line_stats( response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None completion_details: Final = usage.completion_tokens_details line_prompt_cost, line_completion_cost = _output_line_cost( + response_body=response_body, usage=usage, custom_llm_provider=custom_llm_provider, model_name=model_name, @@ -237,19 +239,36 @@ def _compute_output_line_stats( ) +def _ocr_usage_info_from_response_body(response_body: Mapping[str, object]) -> OCRUsageInfo | None: + """OCR results report ``usage_info`` (pages) instead of ``usage`` (tokens); None for non-OCR lines.""" + raw_usage_info: Final = response_body.get("usage_info") + if not isinstance(raw_usage_info, Mapping): + return None + return OCRUsageInfo.model_validate(raw_usage_info) + + def _output_line_cost( + response_body: Mapping[str, object], usage: Usage, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, response_model: str | None, model_info: ModelInfo | None, ) -> tuple[float, float]: """(prompt_cost, completion_cost) for one output line, priced at batch rates.""" - from litellm.cost_calculator import batch_cost_calculator + from litellm.cost_calculator import batch_cost_calculator, ocr_batch_cost cost_model: Final = ( model_name if custom_llm_provider == "bedrock" and model_name else response_model or model_name or "" ) + ocr_usage: Final = _ocr_usage_info_from_response_body(response_body) + if ocr_usage is not None: + return ocr_batch_cost( + model=cost_model, + custom_llm_provider=custom_llm_provider, + usage_info=ocr_usage, + model_info=model_info, + ) return batch_cost_calculator( usage=usage, model=cost_model, @@ -260,7 +279,7 @@ def _output_line_cost( def _aggregate_batch_cost_usage_models( entries: Iterable[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None = None, model_info: ModelInfo | None = None, ) -> BatchCostUsageResult: @@ -427,7 +446,7 @@ def _provider_output_file_id(output_file_id: str) -> str: async def _fetch_batch_managed_file_content( file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai", litellm_params: dict | None = None, ) -> bytes: """ @@ -457,7 +476,7 @@ async def _fetch_batch_managed_file_content( async def _fetch_batch_output_file_content( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai", litellm_params: dict | None = None, ) -> bytes: """ @@ -479,7 +498,7 @@ async def _fetch_batch_output_file_content( async def count_error_file_failed_requests( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], litellm_params: dict | None, ) -> int: """Count failed requests reported only in the batch's separate error file. diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 77a4fdebf16..76b6c73b375 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -105,9 +105,11 @@ def _resolve_timeout( @client async def acreate_batch( completion_window: Literal["24h"], - endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"], + endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"], input_file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "mistral" + ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, extra_body: dict[str, str] | None = None, @@ -155,9 +157,11 @@ async def acreate_batch( @client def create_batch( completion_window: Literal["24h"], - endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"], + endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"], input_file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "mistral" + ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, extra_body: dict[str, str] | None = None, @@ -341,7 +345,7 @@ def create_batch( async def aretrieve_batch( batch_id: str, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral" ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, @@ -389,7 +393,7 @@ def _handle_retrieve_batch_providers_without_provider_config( _retrieve_batch_request: RetrieveBatchRequest, _is_async: bool, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral" ] = "openai", logging_obj: LiteLLMLoggingObj | None = None, ): @@ -497,7 +501,7 @@ def _handle_retrieve_batch_providers_without_provider_config( message=( f"LiteLLM doesn't support custom_llm_provider={custom_llm_provider} for 'retrieve_batch' without a `model` kwarg. " "Supported via this path: 'openai', 'azure', 'vertex_ai', 'anthropic'. " - "'bedrock' is supported but requires `model` to be passed so the provider config can be loaded." + "'bedrock' and 'mistral' are supported but require `model` to be passed so the provider config can be loaded." ), model="n/a", llm_provider=custom_llm_provider, @@ -514,7 +518,7 @@ def _handle_retrieve_batch_providers_without_provider_config( def retrieve_batch( batch_id: str, custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "mistral" ] = "openai", metadata: dict[str, str] | None = None, extra_headers: dict[str, str] | None = None, diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 814eaaf76f7..0c0c6b05df8 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -139,6 +139,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import ( Logging as LitellmLoggingObject, ) + from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo else: LitellmLoggingObject = Any @@ -1982,6 +1983,66 @@ def ocr_cost( return ocr_pages_cost + annotation_pages_cost, 0.0 +_OCR_PRICING_KEYS: Final = ( + "ocr_cost_per_page", + "ocr_cost_per_page_batches", + "annotation_cost_per_page", + "annotation_cost_per_page_batches", +) + + +def ocr_batch_cost( + model: str, + custom_llm_provider: str | None, + usage_info: "OCRUsageInfo", + model_info: ModelInfo | None = None, +) -> tuple[float, float]: + """Per-page cost of one OCR result inside a batch output file. + + Batch OCR is billed per page at the ``*_batches`` rate, falling back to the + synchronous per-page rate when a model has no batch price recorded, the same + fallback ``batch_cost_calculator`` applies to per-token batch pricing. Returns + ``(prompt_cost, completion_cost)`` with the whole cost in the first slot, like + ``ocr_cost``. + """ + has_ocr_pricing: Final = model_info is not None and any(model_info.get(k) is not None for k in _OCR_PRICING_KEYS) + if has_ocr_pricing: + resolved_info: ModelInfo | None = model_info + else: + try: + resolved_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: + resolved_info = None + if resolved_info is None: + verbose_logger.warning( + "OCR batch cost: model=%s custom_llm_provider=%s has no pricing entry; returning 0.0 cost.", + model, + custom_llm_provider, + ) + return 0.0, 0.0 + + page_rate: Final = _first_price(resolved_info, "ocr_cost_per_page_batches", "ocr_cost_per_page") + annotation_rate: Final = _first_price( + resolved_info, "annotation_cost_per_page_batches", "annotation_cost_per_page" + ) + pages_processed: Final = usage_info.pages_processed or 0 + annotation_pages: Final = usage_info.pages_processed_annotation or 0 + if page_rate is None and pages_processed > 0: + verbose_logger.warning( + "OCR batch cost: model=%s custom_llm_provider=%s reported pages_processed=%s but no " + "ocr_cost_per_page is configured; returning 0.0 cost for those pages.", + model, + custom_llm_provider, + pages_processed, + ) + effective_annotation_rate: Final = annotation_rate if annotation_rate is not None else page_rate + return (page_rate or 0.0) * pages_processed + (effective_annotation_rate or 0.0) * annotation_pages, 0.0 + + +def _first_price(model_info: ModelInfo, *keys: str) -> float | None: + return next((price for price in (model_info.get(k) for k in keys) if isinstance(price, (int, float))), None) + + def vector_store_search_cost( model: str | None, custom_llm_provider: str, diff --git a/litellm/files/main.py b/litellm/files/main.py index 218518eb3cd..3d90bf4f299 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -27,12 +27,15 @@ FileCreateProvider = Literal[ "litellm_proxy", "manus", "anthropic", + "mistral", ] FileRetrieveProvider = Literal[ - "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic" + "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic", "mistral" ] -FileDeleteProvider = Literal["openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic"] -FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic"] +FileDeleteProvider = Literal[ + "openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic", "mistral" +] +FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic", "mistral"] import litellm from litellm import get_secret_str from litellm.files.streaming import FileContentStreamingResponse diff --git a/litellm/files/types.py b/litellm/files/types.py index b4ec9996f37..01c7970144b 100644 --- a/litellm/files/types.py +++ b/litellm/files/types.py @@ -2,7 +2,7 @@ from collections.abc import AsyncIterator, Iterator from typing import Literal, NamedTuple FileContentProvider = Literal[ - "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "manus" + "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic", "manus", "mistral" ] diff --git a/litellm/llms/mistral/batches/__init__.py b/litellm/llms/mistral/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/mistral/batches/transformation.py b/litellm/llms/mistral/batches/transformation.py new file mode 100644 index 00000000000..399319e590b --- /dev/null +++ b/litellm/llms/mistral/batches/transformation.py @@ -0,0 +1,186 @@ +""" +Mistral Batch API. Reference: https://docs.mistral.ai/api/#tag/batch + +Mistral runs one model per job (set on the job, not per input line) and accepts +``/v1/ocr`` as a batch endpoint, which is how OCR gets its 50% batch discount. +Output and error files are OpenAI-shaped JSONL (``{custom_id, response: {status_code, body}}``), +so the shared batch cost accounting reads them without a provider branch. +""" + +from types import MappingProxyType +from typing import Final, Literal + +import httpx +from openai.types.batch import BatchRequestCounts +from openai.types.batch import Errors as BatchErrors +from openai.types.batch_error import BatchError +from pydantic import BaseModel, ConfigDict + +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.llms.openai import AllMessageValues, CreateBatchRequest +from litellm.types.utils import LiteLLMBatch, LlmProviders + +from ..common_utils import get_mistral_api_base, get_mistral_auth_headers, mistral_error + +MistralBatchStatus = Literal[ + "QUEUED", "RUNNING", "SUCCESS", "FAILED", "TIMEOUT_EXCEEDED", "CANCELLATION_REQUESTED", "CANCELLED" +] +OpenAIBatchStatus = Literal[ + "validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled" +] + +_STATUS_MAP: Final[MappingProxyType[MistralBatchStatus, OpenAIBatchStatus]] = MappingProxyType( + { + "QUEUED": "validating", + "RUNNING": "in_progress", + "SUCCESS": "completed", + "FAILED": "failed", + "TIMEOUT_EXCEEDED": "expired", + "CANCELLATION_REQUESTED": "cancelling", + "CANCELLED": "cancelled", + } +) + + +class MistralBatchError(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + message: str + count: int = 1 + + +class MistralBatchJob(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + id: str + input_files: tuple[str, ...] = () + endpoint: str + model: str | None = None + status: MistralBatchStatus + created_at: int + started_at: int | None = None + completed_at: int | None = None + total_requests: int = 0 + completed_requests: int = 0 + succeeded_requests: int = 0 + failed_requests: int = 0 + output_file: str | None = None + error_file: str | None = None + errors: tuple[MistralBatchError, ...] = () + metadata: dict[str, str] | None = None + + +def _to_litellm_batch(job: MistralBatchJob) -> LiteLLMBatch: + status: Final = _STATUS_MAP[job.status] + terminal_at: Final = job.completed_at + return LiteLLMBatch( + id=job.id, + object="batch", + endpoint=job.endpoint, + input_file_id=job.input_files[0] if job.input_files else "", + completion_window="24h", + status=status, + created_at=job.created_at, + in_progress_at=job.started_at, + completed_at=terminal_at if status == "completed" else None, + failed_at=terminal_at if status == "failed" else None, + expired_at=terminal_at if status == "expired" else None, + cancelled_at=terminal_at if status == "cancelled" else None, + output_file_id=job.output_file, + error_file_id=job.error_file, + errors=( + BatchErrors( + object="list", + data=[BatchError(message=f"{e.message} (x{e.count})" if e.count > 1 else e.message) for e in job.errors], + ) + if job.errors + else None + ), + request_counts=BatchRequestCounts( + total=job.total_requests, + completed=job.succeeded_requests, + failed=job.failed_requests, + ), + metadata=job.metadata, + ) + + +class MistralBatchesConfig(BaseBatchesConfig): + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.MISTRAL + + def validate_environment( + self, + headers: dict, + model: str, + messages: list[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: + return get_mistral_auth_headers(headers, api_key) + + def get_complete_batch_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: dict, + litellm_params: dict, + data: CreateBatchRequest, + ) -> str: + return f"{get_mistral_api_base(api_base)}/v1/batch/jobs" + + def transform_create_batch_request( + self, + model: str, + create_batch_data: CreateBatchRequest, + optional_params: dict, + litellm_params: dict, + ) -> dict[str, object]: + metadata: Final = create_batch_data.get("metadata") + return { + "input_files": [create_batch_data["input_file_id"]], + "endpoint": create_batch_data["endpoint"], + "model": model, + **({"metadata": metadata} if metadata else {}), + **(create_batch_data.get("extra_body") or {}), + } + + def transform_create_batch_response( + self, + model: str | None, + raw_response: httpx.Response, + logging_obj: object, + litellm_params: dict, + ) -> LiteLLMBatch: + return _to_litellm_batch(MistralBatchJob.model_validate(raw_response.json())) + + def transform_retrieve_batch_request( + self, + batch_id: str, + optional_params: dict, + litellm_params: dict, + ) -> dict[str, object]: + encoded_batch_id: Final = encode_url_path_segment(batch_id, field_name="batch_id") + return { + "method": "GET", + "url": f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/batch/jobs/{encoded_batch_id}", + "headers": get_mistral_auth_headers({}, litellm_params.get("api_key")), + } + + def transform_retrieve_batch_response( + self, + model: str | None, + raw_response: httpx.Response, + logging_obj: object, + litellm_params: dict, + ) -> LiteLLMBatch: + return _to_litellm_batch(MistralBatchJob.model_validate(raw_response.json())) + + def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: + return mistral_error(error_message, status_code, headers) diff --git a/litellm/llms/mistral/common_utils.py b/litellm/llms/mistral/common_utils.py new file mode 100644 index 00000000000..9ea501c860d --- /dev/null +++ b/litellm/llms/mistral/common_utils.py @@ -0,0 +1,36 @@ +from typing import Final + +import httpx + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.secret_managers.main import get_secret_str + +MISTRAL_API_BASE: Final = "https://api.mistral.ai" +MISTRAL_API_KEY_ENV_VAR: Final = "MISTRAL_API_KEY" + + +class MistralError(BaseLLMException): + pass + + +def get_mistral_api_base(api_base: str | None) -> str: + """Return the Mistral origin without a trailing ``/v1``, so callers can append ``/v1/``.""" + resolved: Final = (api_base or get_secret_str("MISTRAL_API_BASE") or MISTRAL_API_BASE).rstrip("/") + return resolved.removesuffix("/v1") + + +def get_mistral_auth_headers(headers: dict, api_key: str | None) -> dict: + resolved_key: Final = api_key or get_secret_str(MISTRAL_API_KEY_ENV_VAR) + if resolved_key is None: + raise ValueError( + "Missing Mistral API Key - A call is being made to Mistral but no key is set either in the environment variables or via params" + ) + return {**headers, "Authorization": f"Bearer {resolved_key}"} + + +def mistral_error(error_message: str, status_code: int, headers: dict | httpx.Headers) -> MistralError: + return MistralError( + status_code=status_code, + message=error_message, + headers=headers if isinstance(headers, httpx.Headers) else httpx.Headers(headers), + ) diff --git a/litellm/llms/mistral/files/__init__.py b/litellm/llms/mistral/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/mistral/files/transformation.py b/litellm/llms/mistral/files/transformation.py new file mode 100644 index 00000000000..071b6f58569 --- /dev/null +++ b/litellm/llms/mistral/files/transformation.py @@ -0,0 +1,226 @@ +""" +Mistral Files API. Reference: https://docs.mistral.ai/api/#tag/files + +Mistral's file objects already carry the OpenAI field names (id, bytes, created_at, +filename, purpose), so this config is URL routing, auth, and a purpose mapping: +Mistral only accepts ``fine-tune``, ``batch`` and ``ocr`` as upload purposes. +""" + +import time +from typing import Final, Literal + +import httpx +from openai.types.file_deleted import FileDeleted +from pydantic import BaseModel, ConfigDict + +from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.files.transformation import BaseFilesConfig, LiteLLMLoggingObj +from litellm.types.llms.openai import ( + CreateFileRequest, + FileContentRequest, + HttpxBinaryResponseContent, + OpenAICreateFileRequestOptionalParams, + OpenAIFileObject, + OpenAIFilesPurpose, +) +from litellm.types.utils import LlmProviders + +from ..common_utils import get_mistral_api_base, get_mistral_auth_headers, mistral_error + +MistralFilePurpose = Literal["fine-tune", "batch", "ocr"] + + +class MistralFile(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + id: str + bytes: int = 0 + created_at: int | None = None + filename: str = "" + purpose: MistralFilePurpose = "batch" + expires_at: int | None = None + + +class MistralFileList(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + data: tuple[MistralFile, ...] = () + + +class MistralFileDeleted(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + id: str + deleted: bool = True + + +def _to_openai_file_object(file: MistralFile) -> OpenAIFileObject: + return OpenAIFileObject( + id=file.id, + bytes=file.bytes, + created_at=file.created_at if file.created_at is not None else int(time.time()), + filename=file.filename, + object="file", + purpose=_to_openai_purpose(file.purpose), + status="uploaded", + expires_at=file.expires_at, + ) + + +def _to_openai_purpose(purpose: MistralFilePurpose) -> OpenAIFilesPurpose: + match purpose: + case "fine-tune" | "batch": + return purpose + case "ocr": + return "user_data" + + +def _to_mistral_purpose(purpose: str) -> MistralFilePurpose: + match purpose: + case "fine-tune" | "ocr": + return purpose + case _: + return "batch" + + +class MistralFilesConfig(BaseFilesConfig): + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.MISTRAL + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: dict, + litellm_params: dict, + stream: bool | None = None, + ) -> str: + return f"{get_mistral_api_base(api_base)}/v1/files" + + def _file_url(self, file_id: str, litellm_params: dict, suffix: str = "") -> str: + encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") + return f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/files/{encoded_file_id}{suffix}" + + def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: + return mistral_error(error_message, status_code, headers) + + def validate_environment( + self, + headers: dict, + model: str, + messages: list, + optional_params: dict, + litellm_params: dict, + api_key: str | None = None, + api_base: str | None = None, + ) -> dict: + return get_mistral_auth_headers(headers, api_key) + + def get_supported_openai_params(self, model: str) -> list[OpenAICreateFileRequestOptionalParams]: + return ["purpose"] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + return optional_params + + def transform_create_file_request( + self, + model: str, + create_file_data: CreateFileRequest, + optional_params: dict, + litellm_params: dict, + ) -> dict: + file_data: Final = create_file_data.get("file") + if file_data is None: + raise ValueError("File data is required") + extracted: Final = extract_file_data(file_data) + filename: Final = extracted["filename"] or f"file_{int(time.time())}.jsonl" + content_type: Final = extracted.get("content_type") or "application/octet-stream" + return { + "file": (filename, extracted["content"], content_type), + "purpose": (None, _to_mistral_purpose(create_file_data.get("purpose", "batch"))), + } + + def transform_create_file_response( + self, + model: str | None, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> OpenAIFileObject: + return _to_openai_file_object(MistralFile.model_validate(raw_response.json())) + + def transform_retrieve_file_request( + self, + file_id: str, + optional_params: dict, + litellm_params: dict, + ) -> tuple[str, dict]: + return self._file_url(file_id, litellm_params), {} + + def transform_retrieve_file_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> OpenAIFileObject: + return _to_openai_file_object(MistralFile.model_validate(raw_response.json())) + + def transform_delete_file_request( + self, + file_id: str, + optional_params: dict, + litellm_params: dict, + ) -> tuple[str, dict]: + return self._file_url(file_id, litellm_params), {} + + def transform_delete_file_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> FileDeleted: + deleted: Final = MistralFileDeleted.model_validate(raw_response.json()) + return FileDeleted(id=deleted.id, deleted=deleted.deleted, object="file") + + def transform_list_files_request( + self, + purpose: str | None, + optional_params: dict, + litellm_params: dict, + ) -> tuple[str, dict]: + params: Final = {"purpose": _to_mistral_purpose(purpose)} if purpose else {} + return f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/files", params + + def transform_list_files_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> list[OpenAIFileObject]: + return [_to_openai_file_object(f) for f in MistralFileList.model_validate(raw_response.json()).data] + + def transform_file_content_request( + self, + file_content_request: FileContentRequest, + optional_params: dict, + litellm_params: dict, + ) -> tuple[str, dict]: + return self._file_url(file_content_request["file_id"], litellm_params, suffix="/content"), {} + + def transform_file_content_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> HttpxBinaryResponseContent: + return HttpxBinaryResponseContent(response=raw_response) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 54ebdc85be9..5a934301edd 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -35262,51 +35262,66 @@ "mistral/mistral-ocr-latest": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4-0": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4-1": { "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "litellm_provider": "mistral", "mode": "ocr", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "source": "https://docs.mistral.ai/models/model-cards/ocr-4-1", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ] }, "mistral/mistral-ocr-2505-completion": { "deprecation_date": "2026-05-31", "litellm_provider": "mistral", "ocr_cost_per_page": 0.001, + "ocr_cost_per_page_batches": 0.0005, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-2512": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, @@ -59822,31 +59837,40 @@ "mistral/mistral-ocr-3": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-3-0": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4": { "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "litellm_provider": "mistral", "mode": "ocr", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "source": "https://docs.mistral.ai/models/model-cards/ocr-4-1", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ] }, "mistral/voxtral-mini-latest": { diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index b6da9490e01..defd59f2be8 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -498,7 +498,7 @@ class CreateBatchRequest(TypedDict, total=False): """ completion_window: Literal["24h"] - endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"] + endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses", "/v1/ocr"] input_file_id: str metadata: dict[str, str] | None output_expires_after: FileExpiresAfter diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d62f00f3676..ea23e00d2bb 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -320,8 +320,10 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_second_480p: ReadOnly[float | None] output_cost_per_second_4k: ReadOnly[float | None] ocr_cost_per_page: float | None # for OCR models + ocr_cost_per_page_batches: ReadOnly[float | None] ocr_cost_per_credit: float | None # for OCR models priced by credit annotation_cost_per_page: float | None # for OCR models + annotation_cost_per_page_batches: ReadOnly[float | None] search_context_cost_per_query: SearchContextCostPerQuery | None # Cost for using web search tool web_search_billing_unit: ( Literal["per_query", "per_prompt"] | None @@ -3598,8 +3600,10 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_token_above_512k_tokens: float | None = None output_vector_size: int | None = None ocr_cost_per_page: float | None = None + ocr_cost_per_page_batches: float | None = None ocr_cost_per_credit: float | None = None annotation_cost_per_page: float | None = None + annotation_cost_per_page_batches: float | None = None regional_processing_uplift_multiplier_eu: float | None = None regional_processing_uplift_multiplier_us: float | None = None regional_endpoint_uplift_multiplier: float | None = None diff --git a/litellm/utils.py b/litellm/utils.py index 8df28870544..3ca81605a95 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5963,8 +5963,10 @@ def _get_model_info_helper( tpm=_model_info.get("tpm", None), rpm=_model_info.get("rpm", None), ocr_cost_per_page=_model_info.get("ocr_cost_per_page", None), + ocr_cost_per_page_batches=_model_info.get("ocr_cost_per_page_batches", None), ocr_cost_per_credit=_model_info.get("ocr_cost_per_credit", None), annotation_cost_per_page=_model_info.get("annotation_cost_per_page", None), + annotation_cost_per_page_batches=_model_info.get("annotation_cost_per_page_batches", None), provider_specific_entry=_model_info.get("provider_specific_entry", None), uses_embed_content=_model_info.get("uses_embed_content", None), supports_image_size=_model_info.get("supports_image_size", None), @@ -8909,6 +8911,10 @@ class ProviderConfigManager: from litellm.llms.anthropic.files.transformation import AnthropicFilesConfig return AnthropicFilesConfig() + elif LlmProviders.MISTRAL == provider: + from litellm.llms.mistral.files.transformation import MistralFilesConfig + + return MistralFilesConfig() return None @staticmethod @@ -8920,6 +8926,10 @@ class ProviderConfigManager: from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig return BedrockBatchesConfig() + elif LlmProviders.MISTRAL == provider: + from litellm.llms.mistral.batches.transformation import MistralBatchesConfig + + return MistralBatchesConfig() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 54ebdc85be9..5a934301edd 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -35262,51 +35262,66 @@ "mistral/mistral-ocr-latest": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4-0": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4-1": { "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "litellm_provider": "mistral", "mode": "ocr", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "source": "https://docs.mistral.ai/models/model-cards/ocr-4-1", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ] }, "mistral/mistral-ocr-2505-completion": { "deprecation_date": "2026-05-31", "litellm_provider": "mistral", "ocr_cost_per_page": 0.001, + "ocr_cost_per_page_batches": 0.0005, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-2512": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, @@ -59822,31 +59837,40 @@ "mistral/mistral-ocr-3": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-3-0": { "litellm_provider": "mistral", "ocr_cost_per_page": 0.002, + "ocr_cost_per_page_batches": 0.001, "annotation_cost_per_page": 0.003, + "annotation_cost_per_page_batches": 0.0015, "mode": "ocr", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ], "source": "https://mistral.ai/pricing#api-pricing" }, "mistral/mistral-ocr-4": { "annotation_cost_per_page": 0.005, + "annotation_cost_per_page_batches": 0.0025, "litellm_provider": "mistral", "mode": "ocr", "ocr_cost_per_page": 0.004, + "ocr_cost_per_page_batches": 0.002, "source": "https://docs.mistral.ai/models/model-cards/ocr-4-1", "supported_endpoints": [ - "/v1/ocr" + "/v1/ocr", + "/v1/batch" ] }, "mistral/voxtral-mini-latest": { diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 768ea332677..85f267c2db0 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1787,3 +1787,86 @@ class TestBatchCostIsFinal: @pytest.mark.parametrize("status", ["failed", "expired", "cancelled"]) def test_other_terminal_statuses_are_final(self, status): assert bu.batch_cost_is_final(_retrieved_batch(status)) is True + + +# =========================================================================== # +# OCR batch output lines (Mistral /v1/ocr batches) are billed per page, not per token +# =========================================================================== # + + +def _ocr_row(pages_processed, annotation_pages=None, model="mistral-ocr-latest"): + usage_info = {"pages_processed": pages_processed, "doc_size_bytes": 4096} + if annotation_pages is not None: + usage_info["pages_processed_annotation"] = annotation_pages + return _success_row(model=model, pages=[{"index": i, "markdown": "x"} for i in range(pages_processed)], usage_info=usage_info) + + +def test_ocr_rows_are_priced_per_page_at_batch_rate(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.004, "ocr_cost_per_page_batches": 0.002}, + ) + result = bu._aggregate_batch_cost_usage_models( + entries=[_ocr_row(3), _ocr_row(5), _failed_row(model="mistral-ocr-latest")], + custom_llm_provider="mistral", + model_name="mistral/mistral-ocr-latest", + ) + assert result.cost == pytest.approx(8 * 0.002) + assert result.prompt_cost == pytest.approx(8 * 0.002) + assert result.completion_cost == 0.0 + assert (result.successful_requests, result.failed_requests) == (2, 1) + assert result.usage.total_tokens == 0 + assert result.models == ["mistral/mistral-ocr-latest"] + + +def test_ocr_rows_fall_back_to_sync_page_rate_without_batch_price(monkeypatch): + monkeypatch.setattr(litellm, "get_model_info", lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.004}) + result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(2)], custom_llm_provider="mistral") + assert result.cost == pytest.approx(2 * 0.004) + + +def test_ocr_rows_bill_annotation_pages_separately(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: { + "ocr_cost_per_page_batches": 0.002, + "annotation_cost_per_page_batches": 0.0025, + }, + ) + result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(4, annotation_pages=4)], custom_llm_provider="mistral") + assert result.cost == pytest.approx(4 * 0.002 + 4 * 0.0025) + + +def test_ocr_rows_use_deployment_model_info_pricing_over_cost_map(monkeypatch): + monkeypatch.setattr( + litellm, "get_model_info", lambda model, custom_llm_provider=None: pytest.fail("cost map must not be consulted") + ) + result = bu._aggregate_batch_cost_usage_models( + entries=[_ocr_row(10)], + custom_llm_provider="mistral", + model_info={"ocr_cost_per_page_batches": 0.001}, + ) + assert result.cost == pytest.approx(0.01) + + +def test_ocr_rows_without_pricing_bill_zero_but_count_as_successful(monkeypatch): + monkeypatch.setattr(litellm, "get_model_info", lambda model, custom_llm_provider=None: {"mode": "ocr"}) + result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(3)], custom_llm_provider="mistral") + assert result.cost == 0.0 + assert (result.successful_requests, result.failed_requests) == (1, 0) + + +def test_chat_rows_from_mistral_still_use_token_pricing(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}, + ) + result = bu._aggregate_batch_cost_usage_models( + entries=[_success_row(model="mistral-small-latest", usage=_usage(10, 5))], + custom_llm_provider="mistral", + ) + assert result.cost == pytest.approx((10 * 0.001 + 5 * 0.002) / 2) + assert result.usage.total_tokens == 15 diff --git a/tests/test_litellm/batches/test_main.py b/tests/test_litellm/batches/test_main.py index b87f9489250..c5a33dd6508 100644 --- a/tests/test_litellm/batches/test_main.py +++ b/tests/test_litellm/batches/test_main.py @@ -778,3 +778,45 @@ def test_retrieve__omits_trusted_model_credentials_when_not_supplied(seams): litellm_params = logging_obj.update_from_kwargs.call_args.kwargs["litellm_params"] assert "_litellm_internal_model_credentials" not in litellm_params + + +# =========================================================================== # +# mistral - a provider-config provider, like bedrock, so it requires `model` +# =========================================================================== # + + +def test_create__mistral_ocr_routes_to_base_http_handler_with_mistral_config(seams): + with patch.object(bm.ProviderConfigManager, "get_provider_batches_config", wraps=bm.ProviderConfigManager.get_provider_batches_config) as get_cfg: + result = bm.create_batch( + completion_window="24h", + endpoint="/v1/ocr", + input_file_id="file-abc", + custom_llm_provider="mistral", + model="mistral/mistral-ocr-latest", + ) + + assert result is seams.base_http.create_batch.return_value + _assert_only(seams.base_http.create_batch, seams, "create_batch") + get_cfg.assert_called_once() + forwarded = seams.base_http.create_batch.call_args.kwargs + assert type(forwarded["provider_config"]).__name__ == "MistralBatchesConfig" + assert forwarded["model"] == "mistral-ocr-latest" + assert forwarded["create_batch_data"]["endpoint"] == "/v1/ocr" + + +def test_create__mistral_without_model_raises_badrequest(seams): + with pytest.raises(litellm.exceptions.BadRequestError): + bm.create_batch(**CREATE_KW, custom_llm_provider="mistral") + + for m in _all_seam_methods(seams, "create_batch"): + m.assert_not_called() + + +def test_retrieve__mistral_routes_to_base_http_handler_with_mistral_config(seams): + result = bm.retrieve_batch(batch_id="job-1", custom_llm_provider="mistral", model="mistral/mistral-ocr-latest") + + assert result is seams.base_http.retrieve_batch.return_value + _assert_only(seams.base_http.retrieve_batch, seams, "retrieve_batch") + forwarded = seams.base_http.retrieve_batch.call_args.kwargs + assert type(forwarded["provider_config"]).__name__ == "MistralBatchesConfig" + assert forwarded["batch_id"] == "job-1" diff --git a/tests/test_litellm/llms/mistral/batches/__init__.py b/tests/test_litellm/llms/mistral/batches/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py new file mode 100644 index 00000000000..03e9c351a30 --- /dev/null +++ b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py @@ -0,0 +1,260 @@ +""" +Regression tests for ``MistralBatchesConfig``, the BaseBatchesConfig implementation +behind ``custom_llm_provider="mistral"`` on /v1/batches. + +Locks the request shape Mistral's ``POST /v1/batch/jobs`` accepts (input_files list, +model set on the job, endpoint passed through untouched so ``/v1/ocr`` batches work), +the Mistral -> OpenAI status mapping, request-count and file-id mapping, and auth. +Everything runs for real against canned httpx responses; only the API key env var is +set. +""" + +import json + +import httpx +import pytest + +from litellm.llms.mistral.batches.transformation import MistralBatchesConfig +from litellm.llms.mistral.common_utils import MistralError +from litellm.types.llms.openai import CreateBatchRequest +from litellm.types.utils import LiteLLMBatch, LlmProviders + +STATUS_MAP = { + "QUEUED": "validating", + "RUNNING": "in_progress", + "SUCCESS": "completed", + "FAILED": "failed", + "TIMEOUT_EXCEEDED": "expired", + "CANCELLATION_REQUESTED": "cancelling", + "CANCELLED": "cancelled", +} + + +def _job(**overrides): + base = { + "id": "8ff5e0d1-6bc2-4c3a-9f7d-0d1c2e3f4a5b", + "object": "batch", + "input_files": ["c1a2b3d4-0000-4000-8000-000000000001"], + "endpoint": "/v1/ocr", + "model": "mistral-ocr-latest", + "status": "SUCCESS", + "created_at": 1_757_400_000, + "started_at": 1_757_400_010, + "completed_at": 1_757_400_500, + "total_requests": 3, + "completed_requests": 3, + "succeeded_requests": 2, + "failed_requests": 1, + "output_file": "out-0000-4000-8000-000000000002", + "error_file": "err-0000-4000-8000-000000000003", + "errors": [], + "metadata": {"job_type": "testing"}, + } + return {**base, **overrides} + + +def _response(payload: dict, status_code: int = 200) -> httpx.Response: + return httpx.Response( + status_code=status_code, + content=json.dumps(payload).encode(), + request=httpx.Request("GET", "https://api.mistral.ai/v1/batch/jobs/x"), + ) + + +@pytest.fixture +def config() -> MistralBatchesConfig: + return MistralBatchesConfig() + + +@pytest.fixture +def api_key(monkeypatch) -> str: + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + return "sk-mistral-test" + + +def test_custom_llm_provider(config): + assert config.custom_llm_provider == LlmProviders.MISTRAL + + +# --------------------------------------------------------------------------- # +# create +# --------------------------------------------------------------------------- # + + +def test_create_request_maps_openai_fields_onto_mistral_job(config): + data = CreateBatchRequest( + completion_window="24h", + endpoint="/v1/ocr", + input_file_id="file-123", + metadata={"team": "docs"}, + ) + body = config.transform_create_batch_request( + model="mistral-ocr-latest", create_batch_data=data, optional_params={}, litellm_params={} + ) + assert body == { + "input_files": ["file-123"], + "endpoint": "/v1/ocr", + "model": "mistral-ocr-latest", + "metadata": {"team": "docs"}, + } + + +def test_create_request_omits_empty_metadata_and_forwards_extra_body(config): + data = CreateBatchRequest( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-123", + metadata=None, + extra_body={"timeout_hours": 48}, + ) + body = config.transform_create_batch_request( + model="mistral-small-latest", create_batch_data=data, optional_params={}, litellm_params={} + ) + assert "metadata" not in body + assert body["timeout_hours"] == 48 + + +@pytest.mark.parametrize( + "api_base,expected", + [ + (None, "https://api.mistral.ai/v1/batch/jobs"), + ("https://api.mistral.ai/v1", "https://api.mistral.ai/v1/batch/jobs"), + ("https://proxy.example.com/", "https://proxy.example.com/v1/batch/jobs"), + ], +) +def test_create_url(config, api_base, expected): + url = config.get_complete_batch_url( + api_base=api_base, api_key="k", model="m", optional_params={}, litellm_params={}, data={} + ) + assert url == expected + + +def test_validate_environment_uses_bearer_auth(config, api_key): + headers = config.validate_environment( + headers={"x-extra": "1"}, model="m", messages=[], optional_params={}, litellm_params={} + ) + assert headers == {"x-extra": "1", "Authorization": f"Bearer {api_key}"} + + +def test_validate_environment_explicit_key_wins(config, api_key): + headers = config.validate_environment( + headers={}, model="m", messages=[], optional_params={}, litellm_params={}, api_key="sk-explicit" + ) + assert headers["Authorization"] == "Bearer sk-explicit" + + +def test_validate_environment_without_key_raises(config, monkeypatch): + monkeypatch.delenv("MISTRAL_API_KEY", raising=False) + with pytest.raises(ValueError, match="Missing Mistral API Key"): + config.validate_environment(headers={}, model="m", messages=[], optional_params={}, litellm_params={}) + + +def test_create_response_maps_job_onto_openai_batch(config): + batch = config.transform_create_batch_response( + model="mistral-ocr-latest", + raw_response=_response(_job(status="QUEUED", started_at=None, completed_at=None)), + logging_obj=None, + litellm_params={}, + ) + assert isinstance(batch, LiteLLMBatch) + assert batch.id == "8ff5e0d1-6bc2-4c3a-9f7d-0d1c2e3f4a5b" + assert batch.endpoint == "/v1/ocr" + assert batch.input_file_id == "c1a2b3d4-0000-4000-8000-000000000001" + assert batch.status == "validating" + assert batch.created_at == 1_757_400_000 + assert batch.in_progress_at is None + assert batch.completed_at is None + assert batch.metadata == {"job_type": "testing"} + + +# --------------------------------------------------------------------------- # +# retrieve +# --------------------------------------------------------------------------- # + + +def test_retrieve_request_is_presigned_get_with_auth(config, api_key): + req = config.transform_retrieve_batch_request( + batch_id="job/with slash", optional_params={}, litellm_params={"api_base": "https://api.mistral.ai"} + ) + assert req["method"] == "GET" + assert req["url"] == "https://api.mistral.ai/v1/batch/jobs/job%2Fwith%20slash" + assert req["headers"] == {"Authorization": f"Bearer {api_key}"} + + +def test_retrieve_request_prefers_litellm_params_api_key(config, api_key): + req = config.transform_retrieve_batch_request( + batch_id="job-1", optional_params={}, litellm_params={"api_key": "sk-from-deployment"} + ) + assert req["headers"]["Authorization"] == "Bearer sk-from-deployment" + + +@pytest.mark.parametrize("mistral_status,openai_status", sorted(STATUS_MAP.items())) +def test_retrieve_response_status_mapping(config, mistral_status, openai_status): + batch = config.transform_retrieve_batch_response( + model=None, raw_response=_response(_job(status=mistral_status)), logging_obj=None, litellm_params={} + ) + assert batch.status == openai_status + + +@pytest.mark.parametrize( + "mistral_status,populated_field", + [ + ("SUCCESS", "completed_at"), + ("FAILED", "failed_at"), + ("TIMEOUT_EXCEEDED", "expired_at"), + ("CANCELLED", "cancelled_at"), + ], +) +def test_retrieve_response_terminal_timestamp_lands_on_matching_field(config, mistral_status, populated_field): + batch = config.transform_retrieve_batch_response( + model=None, raw_response=_response(_job(status=mistral_status)), logging_obj=None, litellm_params={} + ) + terminal_fields = {"completed_at", "failed_at", "expired_at", "cancelled_at"} + assert getattr(batch, populated_field) == 1_757_400_500 + for other in terminal_fields - {populated_field}: + assert getattr(batch, other) is None + assert batch.in_progress_at == 1_757_400_010 + + +def test_retrieve_response_maps_counts_and_files(config): + batch = config.transform_retrieve_batch_response( + model=None, raw_response=_response(_job()), logging_obj=None, litellm_params={} + ) + assert batch.request_counts.total == 3 + assert batch.request_counts.completed == 2 + assert batch.request_counts.failed == 1 + assert batch.output_file_id == "out-0000-4000-8000-000000000002" + assert batch.error_file_id == "err-0000-4000-8000-000000000003" + assert batch.errors is None + + +def test_retrieve_response_surfaces_job_errors(config): + batch = config.transform_retrieve_batch_response( + model=None, + raw_response=_response( + _job(status="FAILED", errors=[{"message": "invalid document", "count": 2}, {"message": "timeout"}]) + ), + logging_obj=None, + litellm_params={}, + ) + assert [e.message for e in batch.errors.data] == ["invalid document (x2)", "timeout"] + + +def test_retrieve_response_without_files_or_input(config): + batch = config.transform_retrieve_batch_response( + model=None, + raw_response=_response(_job(input_files=[], output_file=None, error_file=None, metadata=None)), + logging_obj=None, + litellm_params={}, + ) + assert batch.input_file_id == "" + assert batch.output_file_id is None + assert batch.error_file_id is None + assert batch.metadata is None + + +def test_get_error_class(config): + err = config.get_error_class("nope", 401, {"x-request-id": "r1"}) + assert isinstance(err, MistralError) + assert err.status_code == 401 + assert err.message == "nope" diff --git a/tests/test_litellm/llms/mistral/files/__init__.py b/tests/test_litellm/llms/mistral/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py new file mode 100644 index 00000000000..d6ad8a34b35 --- /dev/null +++ b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py @@ -0,0 +1,189 @@ +""" +Regression tests for ``MistralFilesConfig``, the BaseFilesConfig implementation behind +``custom_llm_provider="mistral"`` on /v1/files. + +Locks the URL routing for each file operation, the multipart upload shape Mistral's +``POST /v1/files`` accepts (purpose restricted to fine-tune/batch/ocr), and the +Mistral -> OpenAI file object mapping. Runs against canned httpx responses. +""" + +import json + +import httpx +import pytest +from openai.types.file_deleted import FileDeleted + +from litellm.llms.mistral.files.transformation import MistralFilesConfig +from litellm.types.llms.openai import CreateFileRequest, FileContentRequest, OpenAIFileObject +from litellm.types.utils import LlmProviders + +FILE_ID = "497f6eca-6276-4993-bfeb-53cbbbba6f09" + + +def _file(**overrides): + base = { + "id": FILE_ID, + "object": "file", + "bytes": 13000, + "created_at": 1_716_963_433, + "filename": "batch_input.jsonl", + "purpose": "batch", + "sample_type": "batch_request", + "num_lines": 3, + "source": "upload", + } + return {**base, **overrides} + + +def _response(payload) -> httpx.Response: + return httpx.Response( + status_code=200, + content=json.dumps(payload).encode(), + request=httpx.Request("GET", "https://api.mistral.ai/v1/files"), + ) + + +@pytest.fixture +def config() -> MistralFilesConfig: + return MistralFilesConfig() + + +@pytest.fixture +def api_key(monkeypatch) -> str: + monkeypatch.setenv("MISTRAL_API_KEY", "sk-mistral-test") + return "sk-mistral-test" + + +def test_custom_llm_provider(config): + assert config.custom_llm_provider == LlmProviders.MISTRAL + + +@pytest.mark.parametrize( + "api_base,expected", + [ + (None, "https://api.mistral.ai/v1/files"), + ("https://api.mistral.ai/v1/", "https://api.mistral.ai/v1/files"), + ("https://proxy.example.com", "https://proxy.example.com/v1/files"), + ], +) +def test_upload_url(config, api_base, expected): + url = config.get_complete_url(api_base=api_base, api_key="k", model="", optional_params={}, litellm_params={}) + assert url == expected + + +def test_validate_environment_uses_bearer_auth(config, api_key): + headers = config.validate_environment(headers={}, model="", messages=[], optional_params={}, litellm_params={}) + assert headers == {"Authorization": f"Bearer {api_key}"} + + +def test_upload_request_is_multipart_with_batch_purpose(config): + body = config.transform_create_file_request( + model="", + create_file_data=CreateFileRequest(file=("in.jsonl", b'{"custom_id":"0"}\n', "application/jsonl"), purpose="batch"), + optional_params={}, + litellm_params={}, + ) + assert body == { + "file": ("in.jsonl", b'{"custom_id":"0"}\n', "application/jsonl"), + "purpose": (None, "batch"), + } + + +@pytest.mark.parametrize( + "openai_purpose,mistral_purpose", + [("batch", "batch"), ("fine-tune", "fine-tune"), ("ocr", "ocr"), ("assistants", "batch"), ("user_data", "batch")], +) +def test_upload_request_maps_purpose_onto_mistral_enum(config, openai_purpose, mistral_purpose): + body = config.transform_create_file_request( + model="", + create_file_data=CreateFileRequest(file=("f.bin", b"x"), purpose=openai_purpose), + optional_params={}, + litellm_params={}, + ) + assert body["purpose"] == (None, mistral_purpose) + + +def test_upload_request_requires_file(config): + with pytest.raises(ValueError, match="File data is required"): + config.transform_create_file_request( + model="", create_file_data=CreateFileRequest(purpose="batch"), optional_params={}, litellm_params={} + ) + + +def test_upload_response_maps_onto_openai_file_object(config): + obj = config.transform_create_file_response( + model=None, raw_response=_response(_file()), logging_obj=None, litellm_params={} + ) + assert obj == OpenAIFileObject( + id=FILE_ID, + bytes=13000, + created_at=1_716_963_433, + filename="batch_input.jsonl", + object="file", + purpose="batch", + status="uploaded", + ) + + +def test_file_response_with_ocr_purpose_maps_onto_user_data(config): + obj = config.transform_retrieve_file_response( + raw_response=_response(_file(purpose="ocr", expires_at=1_800_000_000)), logging_obj=None, litellm_params={} + ) + assert obj.purpose == "user_data" + assert obj.expires_at == 1_800_000_000 + + +@pytest.mark.parametrize( + "method,suffix", + [ + ("transform_retrieve_file_request", ""), + ("transform_delete_file_request", ""), + ], +) +def test_single_file_urls_encode_id_and_honor_api_base(config, method, suffix): + url, params = getattr(config, method)( + file_id="id/with slash", optional_params={}, litellm_params={"api_base": "https://mistral.internal/v1"} + ) + assert url == f"https://mistral.internal/v1/files/id%2Fwith%20slash{suffix}" + assert params == {} + + +def test_file_content_url(config): + url, params = config.transform_file_content_request( + file_content_request=FileContentRequest(file_id=FILE_ID), optional_params={}, litellm_params={} + ) + assert url == f"https://api.mistral.ai/v1/files/{FILE_ID}/content" + assert params == {} + + +def test_file_content_response_is_binary_passthrough(config): + raw = httpx.Response( + 200, content=b'{"custom_id":"0","response":{"status_code":200}}\n', request=httpx.Request("GET", "https://x") + ) + out = config.transform_file_content_response(raw_response=raw, logging_obj=None, litellm_params={}) + assert out.content == b'{"custom_id":"0","response":{"status_code":200}}\n' + + +def test_delete_response(config): + out = config.transform_delete_file_response( + raw_response=_response({"id": FILE_ID, "object": "file", "deleted": True}), logging_obj=None, litellm_params={} + ) + assert out == FileDeleted(id=FILE_ID, deleted=True, object="file") + + +def test_list_request_filters_by_mapped_purpose(config): + url, params = config.transform_list_files_request(purpose="batch", optional_params={}, litellm_params={}) + assert url == "https://api.mistral.ai/v1/files" + assert params == {"purpose": "batch"} + _, no_params = config.transform_list_files_request(purpose=None, optional_params={}, litellm_params={}) + assert no_params == {} + + +def test_list_response(config): + out = config.transform_list_files_response( + raw_response=_response({"data": [_file(), _file(id="second", filename="b.jsonl")], "object": "list", "total": 2}), + logging_obj=None, + litellm_params={}, + ) + assert [f.id for f in out] == [FILE_ID, "second"] + assert out[1].filename == "b.jsonl" diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py index 40e54f71eeb..9fe6f38003f 100644 --- a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py +++ b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py @@ -72,9 +72,11 @@ def test_ocr3_pricing_entry(cost_map_path: Path) -> None: assert info is not None, f"{OCR3_MODEL} missing from {cost_map_path.name}" assert info["litellm_provider"] == "mistral" assert info["mode"] == "ocr" - assert info["supported_endpoints"] == ["/v1/ocr"] + assert info["supported_endpoints"] == ["/v1/ocr", "/v1/batch"] assert info["ocr_cost_per_page"] == OCR3_COST_PER_PAGE assert info["annotation_cost_per_page"] == OCR3_ANNOTATION_COST_PER_PAGE + assert info["ocr_cost_per_page_batches"] == OCR3_COST_PER_PAGE / 2 + assert info["annotation_cost_per_page_batches"] == OCR3_ANNOTATION_COST_PER_PAGE / 2 def test_ocr3_model_info_price(local_model_cost_map) -> None: diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index e42608c9904..61739846706 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -992,7 +992,9 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_cost_per_video_per_second_above_128k_tokens": {"type": "number"}, "input_dbu_cost_per_token": {"type": "number"}, "annotation_cost_per_page": {"type": "number"}, + "annotation_cost_per_page_batches": {"type": "number"}, "ocr_cost_per_page": {"type": "number"}, + "ocr_cost_per_page_batches": {"type": "number"}, "ocr_cost_per_credit": {"type": "number"}, "code_interpreter_cost_per_session": {"type": "number"}, "inference_geo": {"type": "string"}, From 2e5f5a95c813b940dc7f654c10b2ce2036c6fb2a Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 18:14:16 -0400 Subject: [PATCH 004/251] fix(proxy): retrieve model-routed file ids from the deployment's provider GET /v1/files/{id} for an id encoded with a non-OpenAI deployment forwarded the deployment credentials but let custom_llm_provider default to openai, so a Mistral file was fetched from api.openai.com with the Mistral key and 401'd. Delete and content already passed the provider through; retrieve now does too. --- .../openai_files_endpoints/files_endpoints.py | 5 +- .../test_files_endpoint.py | 59 +++++++++++++++++++ 2 files changed, 63 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index c315d30b8f3..49a01495d65 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -1139,7 +1139,10 @@ async def get_file( include_internal_credentials=True, ) - response = await litellm.afile_retrieve(**data) + response = await litellm.afile_retrieve( + custom_llm_provider=credentials["custom_llm_provider"], + **data, + ) # Keep the encoded ID in response if it was originally encoded if original_file_id and response and hasattr(response, "id") and response.id: diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 5f1e7e1fe0c..5faae166fca 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -4819,3 +4819,62 @@ def test_get_file_content_keeps_the_status_of_a_rejection_raised_inside_the_rout error = response.json()["error"] assert error["message"].startswith("Storage backend error") assert (error["type"], error["param"], error["code"]) == ("invalid_request_error", "file_id", "400") + + +def test_get_file_model_routed_id_forwards_deployment_provider(mocker: MockerFixture, monkeypatch): + """ + Regression: a file id encoded with a non-OpenAI deployment (here Mistral) must be + retrieved from that deployment's provider. Before the fix the retrieve path only + forwarded the credentials and let ``custom_llm_provider`` default to openai, so a + Mistral file id was sent to api.openai.com with the Mistral key and 401'd. + """ + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + router = Router( + model_list=[ + { + "model_name": "mistral-ocr", + "litellm_params": {"model": "mistral/mistral-ocr-latest", "api_key": "mistral-key"}, + "model_info": {"id": "mistral-ocr-id"}, + } + ] + ) + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + + captured_kwargs: dict = {} + + async def _mock_afile_retrieve(**kwargs): + captured_kwargs.update(kwargs) + return OpenAIFileObject( + id="7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", + object="file", + bytes=2, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(litellm, "afile_retrieve", _mock_afile_retrieve) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user" + ) + encoded_id = encode_file_id_with_model("7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", "mistral-ocr") + + try: + response = client.get(f"/v1/files/{encoded_id}", headers={"Authorization": "Bearer test-key"}) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + assert captured_kwargs["custom_llm_provider"] == "mistral" + assert captured_kwargs["api_key"] == "mistral-key" + assert captured_kwargs["file_id"] == "7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df" + assert response.json()["id"] == encoded_id From c246f75e3ec93192f95e0cb2fc3f50485bf13ebc Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 18:37:01 -0400 Subject: [PATCH 005/251] refactor(mistral): satisfy type-discipline and basedpyright gates for files/batches configs --- litellm/cost_calculator.py | 23 ++-- litellm/files/main.py | 4 +- .../llms/mistral/batches/transformation.py | 114 ++++++++++------ litellm/llms/mistral/common_utils.py | 13 +- litellm/llms/mistral/files/transformation.py | 129 +++++++++++------- .../openai_files_endpoints/files_endpoints.py | 2 +- .../test_litellm/batches/test_batch_utils.py | 46 +++++-- tests/test_litellm/batches/test_main.py | 76 +++-------- .../test_mistral_batches_transformation.py | 16 ++- .../test_mistral_files_transformation.py | 8 +- 10 files changed, 251 insertions(+), 180 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 0c0c6b05df8..7c95941d77d 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2006,13 +2006,11 @@ def ocr_batch_cost( ``ocr_cost``. """ has_ocr_pricing: Final = model_info is not None and any(model_info.get(k) is not None for k in _OCR_PRICING_KEYS) - if has_ocr_pricing: - resolved_info: ModelInfo | None = model_info - else: - try: - resolved_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) - except Exception: - resolved_info = None + resolved_info: Final = ( + model_info + if has_ocr_pricing + else _lookup_model_info_or_none(model=model, custom_llm_provider=custom_llm_provider) + ) if resolved_info is None: verbose_logger.warning( "OCR batch cost: model=%s custom_llm_provider=%s has no pricing entry; returning 0.0 cost.", @@ -2022,9 +2020,7 @@ def ocr_batch_cost( return 0.0, 0.0 page_rate: Final = _first_price(resolved_info, "ocr_cost_per_page_batches", "ocr_cost_per_page") - annotation_rate: Final = _first_price( - resolved_info, "annotation_cost_per_page_batches", "annotation_cost_per_page" - ) + annotation_rate: Final = _first_price(resolved_info, "annotation_cost_per_page_batches", "annotation_cost_per_page") pages_processed: Final = usage_info.pages_processed or 0 annotation_pages: Final = usage_info.pages_processed_annotation or 0 if page_rate is None and pages_processed > 0: @@ -2039,6 +2035,13 @@ def ocr_batch_cost( return (page_rate or 0.0) * pages_processed + (effective_annotation_rate or 0.0) * annotation_pages, 0.0 +def _lookup_model_info_or_none(model: str, custom_llm_provider: str | None) -> ModelInfo | None: + try: + return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: + return None + + def _first_price(model_info: ModelInfo, *keys: str) -> float | None: return next((price for price in (model_info.get(k) for k in keys) if isinstance(price, (int, float))), None) diff --git a/litellm/files/main.py b/litellm/files/main.py index 3d90bf4f299..5cf0f8e576a 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -32,9 +32,7 @@ FileCreateProvider = Literal[ FileRetrieveProvider = Literal[ "openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic", "mistral" ] -FileDeleteProvider = Literal[ - "openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic", "mistral" -] +FileDeleteProvider = Literal["openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic", "mistral"] FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic", "mistral"] import litellm from litellm import get_secret_str diff --git a/litellm/llms/mistral/batches/transformation.py b/litellm/llms/mistral/batches/transformation.py index 399319e590b..ef9ee5ff503 100644 --- a/litellm/llms/mistral/batches/transformation.py +++ b/litellm/llms/mistral/batches/transformation.py @@ -7,14 +7,16 @@ Output and error files are OpenAI-shaped JSONL (``{custom_id, response: {status_ so the shared batch cost accounting reads them without a provider branch. """ +from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Final, Literal +from typing import Final, Literal, TypeAlias import httpx from openai.types.batch import BatchRequestCounts from openai.types.batch import Errors as BatchErrors from openai.types.batch_error import BatchError from pydantic import BaseModel, ConfigDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig @@ -24,13 +26,14 @@ from litellm.types.utils import LiteLLMBatch, LlmProviders from ..common_utils import get_mistral_api_base, get_mistral_auth_headers, mistral_error -MistralBatchStatus = Literal[ +MistralBatchStatus: TypeAlias = Literal[ "QUEUED", "RUNNING", "SUCCESS", "FAILED", "TIMEOUT_EXCEEDED", "CANCELLATION_REQUESTED", "CANCELLED" ] -OpenAIBatchStatus = Literal[ +OpenAIBatchStatus: TypeAlias = Literal[ "validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled" ] +_NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) # mutable-ok: frozen at module scope _STATUS_MAP: Final[MappingProxyType[MistralBatchStatus, OpenAIBatchStatus]] = MappingProxyType( { "QUEUED": "validating", @@ -44,6 +47,23 @@ _STATUS_MAP: Final[MappingProxyType[MistralBatchStatus, OpenAIBatchStatus]] = Ma ) +class MistralCreateBatchJobRequest(TypedDict): + """Body of ``POST /v1/batch/jobs``.""" + + input_files: ReadOnly[tuple[str, ...]] + endpoint: ReadOnly[str] + model: ReadOnly[str] + metadata: NotRequired[ReadOnly[Mapping[str, str]]] + + +class MistralPresignedRequest(TypedDict): + """A fully-formed request the shared HTTP handler sends as-is (its ``method`` branch).""" + + method: ReadOnly[Literal["GET"]] + url: ReadOnly[str] + headers: ReadOnly[Mapping[str, str]] + + class MistralBatchError(BaseModel): model_config = ConfigDict(frozen=True, extra="ignore") @@ -69,7 +89,18 @@ class MistralBatchJob(BaseModel): output_file: str | None = None error_file: str | None = None errors: tuple[MistralBatchError, ...] = () - metadata: dict[str, str] | None = None + metadata: dict[str, str] | None = None # mutable-ok: LiteLLMBatch.metadata is typed as dict + + +def _to_batch_errors(errors: Sequence[MistralBatchError]) -> BatchErrors | None: + if not errors: + return None + return BatchErrors( + object="list", + data=[ # mutable-ok: openai Batch.Errors.data is typed as list + BatchError(message=f"{e.message} (x{e.count})" if e.count > 1 else e.message) for e in errors + ], + ) def _to_litellm_batch(job: MistralBatchJob) -> LiteLLMBatch: @@ -90,14 +121,7 @@ def _to_litellm_batch(job: MistralBatchJob) -> LiteLLMBatch: cancelled_at=terminal_at if status == "cancelled" else None, output_file_id=job.output_file, error_file_id=job.error_file, - errors=( - BatchErrors( - object="list", - data=[BatchError(message=f"{e.message} (x{e.count})" if e.count > 1 else e.message) for e in job.errors], - ) - if job.errors - else None - ), + errors=_to_batch_errors(job.errors), request_counts=BatchRequestCounts( total=job.total_requests, completed=job.succeeded_requests, @@ -114,14 +138,14 @@ class MistralBatchesConfig(BaseBatchesConfig): def validate_environment( self, - headers: dict, + headers: Mapping[str, str], model: str, - messages: list[AllMessageValues], - optional_params: dict, - litellm_params: dict, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], api_key: str | None = None, api_base: str | None = None, - ) -> dict: + ) -> dict[str, str]: # mutable-ok: BaseBatchesConfig signature return get_mistral_auth_headers(headers, api_key) def get_complete_batch_url( @@ -129,8 +153,8 @@ class MistralBatchesConfig(BaseBatchesConfig): api_base: str | None, api_key: str | None, model: str, - optional_params: dict, - litellm_params: dict, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], data: CreateBatchRequest, ) -> str: return f"{get_mistral_api_base(api_base)}/v1/batch/jobs" @@ -139,48 +163,58 @@ class MistralBatchesConfig(BaseBatchesConfig): self, model: str, create_batch_data: CreateBatchRequest, - optional_params: dict, - litellm_params: dict, - ) -> dict[str, object]: + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> dict[str, object]: # mutable-ok: BaseBatchesConfig signature + input_file_id: Final = create_batch_data.get("input_file_id") + endpoint: Final = create_batch_data.get("endpoint") + if input_file_id is None or endpoint is None: + raise ValueError("input_file_id and endpoint are required to create a Mistral batch job") metadata: Final = create_batch_data.get("metadata") - return { - "input_files": [create_batch_data["input_file_id"]], - "endpoint": create_batch_data["endpoint"], - "model": model, - **({"metadata": metadata} if metadata else {}), - **(create_batch_data.get("extra_body") or {}), - } + body: Final = ( + MistralCreateBatchJobRequest( + input_files=(input_file_id,), endpoint=endpoint, model=model, metadata=metadata + ) + if metadata + else MistralCreateBatchJobRequest(input_files=(input_file_id,), endpoint=endpoint, model=model) + ) + return dict(body) # mutable-ok: BaseBatchesConfig signature def transform_create_batch_response( self, model: str | None, raw_response: httpx.Response, logging_obj: object, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> LiteLLMBatch: return _to_litellm_batch(MistralBatchJob.model_validate(raw_response.json())) def transform_retrieve_batch_request( self, batch_id: str, - optional_params: dict, - litellm_params: dict, - ) -> dict[str, object]: + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> dict[str, object]: # mutable-ok: BaseBatchesConfig signature encoded_batch_id: Final = encode_url_path_segment(batch_id, field_name="batch_id") - return { - "method": "GET", - "url": f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/batch/jobs/{encoded_batch_id}", - "headers": get_mistral_auth_headers({}, litellm_params.get("api_key")), - } + api_base: Final = litellm_params.get("api_base") + api_key: Final = litellm_params.get("api_key") + request: Final = MistralPresignedRequest( + method="GET", + url=f"{get_mistral_api_base(api_base if isinstance(api_base, str) else None)}/v1/batch/jobs/{encoded_batch_id}", + headers=get_mistral_auth_headers(_NO_HEADERS, api_key if isinstance(api_key, str) else None), + ) + return dict(request) # mutable-ok: BaseBatchesConfig signature def transform_retrieve_batch_response( self, model: str | None, raw_response: httpx.Response, logging_obj: object, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> LiteLLMBatch: return _to_litellm_batch(MistralBatchJob.model_validate(raw_response.json())) - def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: + def get_error_class( + self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers + ) -> BaseLLMException: return mistral_error(error_message, status_code, headers) diff --git a/litellm/llms/mistral/common_utils.py b/litellm/llms/mistral/common_utils.py index 9ea501c860d..2f14328afdf 100644 --- a/litellm/llms/mistral/common_utils.py +++ b/litellm/llms/mistral/common_utils.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Final import httpx @@ -19,18 +20,22 @@ def get_mistral_api_base(api_base: str | None) -> str: return resolved.removesuffix("/v1") -def get_mistral_auth_headers(headers: dict, api_key: str | None) -> dict: +def get_mistral_auth_headers( + headers: Mapping[str, str], api_key: str | None +) -> dict[str, str]: # mutable-ok: BaseConfig.validate_environment contract returns dict resolved_key: Final = api_key or get_secret_str(MISTRAL_API_KEY_ENV_VAR) if resolved_key is None: raise ValueError( "Missing Mistral API Key - A call is being made to Mistral but no key is set either in the environment variables or via params" ) - return {**headers, "Authorization": f"Bearer {resolved_key}"} + return dict(headers, Authorization=f"Bearer {resolved_key}") # mutable-ok: BaseConfig contract returns dict -def mistral_error(error_message: str, status_code: int, headers: dict | httpx.Headers) -> MistralError: +def mistral_error(error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers) -> MistralError: return MistralError( status_code=status_code, message=error_message, - headers=headers if isinstance(headers, httpx.Headers) else httpx.Headers(headers), + headers=headers + if isinstance(headers, httpx.Headers) + else httpx.Headers(dict(headers)), # mutable-ok: httpx.Headers takes a dict ) diff --git a/litellm/llms/mistral/files/transformation.py b/litellm/llms/mistral/files/transformation.py index 071b6f58569..6d58311813c 100644 --- a/litellm/llms/mistral/files/transformation.py +++ b/litellm/llms/mistral/files/transformation.py @@ -7,11 +7,13 @@ Mistral only accepts ``fine-tune``, ``batch`` and ``ocr`` as upload purposes. """ import time -from typing import Final, Literal +from collections.abc import Mapping, Sequence +from typing import Final, Literal, TypeAlias import httpx from openai.types.file_deleted import FileDeleted from pydantic import BaseModel, ConfigDict +from typing_extensions import ReadOnly, TypedDict from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data from litellm.litellm_core_utils.url_utils import encode_url_path_segment @@ -29,7 +31,16 @@ from litellm.types.utils import LlmProviders from ..common_utils import get_mistral_api_base, get_mistral_auth_headers, mistral_error -MistralFilePurpose = Literal["fine-tune", "batch", "ocr"] +MistralFilePurpose: TypeAlias = Literal["fine-tune", "batch", "ocr"] + +_NO_QUERY_PARAMS: Final[dict[str, str]] = {} # mutable-ok: BaseFilesConfig request transforms return tuple[str, dict] + + +class MistralMultipartUpload(TypedDict): + """``files=`` payload for ``POST /v1/files``: each value is an httpx multipart tuple.""" + + file: ReadOnly[tuple[str, object, str]] + purpose: ReadOnly[tuple[None, MistralFilePurpose]] class MistralFile(BaseModel): @@ -85,6 +96,11 @@ def _to_mistral_purpose(purpose: str) -> MistralFilePurpose: return "batch" +def _api_base_from(litellm_params: Mapping[str, object]) -> str: + api_base: Final = litellm_params.get("api_base") + return get_mistral_api_base(api_base if isinstance(api_base, str) else None) + + class MistralFilesConfig(BaseFilesConfig): @property def custom_llm_provider(self) -> LlmProviders: @@ -95,99 +111,103 @@ class MistralFilesConfig(BaseFilesConfig): api_base: str | None, api_key: str | None, model: str, - optional_params: dict, - litellm_params: dict, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], stream: bool | None = None, ) -> str: return f"{get_mistral_api_base(api_base)}/v1/files" - def _file_url(self, file_id: str, litellm_params: dict, suffix: str = "") -> str: + def _file_url(self, file_id: str, litellm_params: Mapping[str, object], suffix: str = "") -> str: encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") - return f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/files/{encoded_file_id}{suffix}" + return f"{_api_base_from(litellm_params)}/v1/files/{encoded_file_id}{suffix}" - def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: + def get_error_class( + self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers + ) -> BaseLLMException: return mistral_error(error_message, status_code, headers) def validate_environment( self, - headers: dict, + headers: Mapping[str, str], model: str, - messages: list, - optional_params: dict, - litellm_params: dict, + messages: Sequence[object], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], api_key: str | None = None, api_base: str | None = None, - ) -> dict: + ) -> dict[str, str]: # mutable-ok: BaseFilesConfig signature return get_mistral_auth_headers(headers, api_key) - def get_supported_openai_params(self, model: str) -> list[OpenAICreateFileRequestOptionalParams]: - return ["purpose"] + def get_supported_openai_params( + self, model: str + ) -> list[OpenAICreateFileRequestOptionalParams]: # mutable-ok: BaseFilesConfig signature + return ["purpose"] # mutable-ok: BaseFilesConfig signature def map_openai_params( self, - non_default_params: dict, - optional_params: dict, + non_default_params: Mapping[str, object], + optional_params: dict[str, object], # mutable-ok: BaseConfig signature, returned as-is model: str, drop_params: bool, - ) -> dict: + ) -> dict[str, object]: # mutable-ok: BaseConfig signature return optional_params def transform_create_file_request( self, model: str, create_file_data: CreateFileRequest, - optional_params: dict, - litellm_params: dict, - ) -> dict: - file_data: Final = create_file_data.get("file") - if file_data is None: + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> dict[str, object]: # mutable-ok: BaseFilesConfig signature + if "file" not in create_file_data: raise ValueError("File data is required") - extracted: Final = extract_file_data(file_data) + extracted: Final = extract_file_data(create_file_data["file"]) filename: Final = extracted["filename"] or f"file_{int(time.time())}.jsonl" content_type: Final = extracted.get("content_type") or "application/octet-stream" - return { - "file": (filename, extracted["content"], content_type), - "purpose": (None, _to_mistral_purpose(create_file_data.get("purpose", "batch"))), - } + upload: Final = MistralMultipartUpload( + file=(filename, extracted["content"], content_type), + purpose=(None, _to_mistral_purpose(create_file_data.get("purpose", "batch"))), + ) + return dict(upload) # mutable-ok: BaseFilesConfig signature def transform_create_file_response( self, model: str | None, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> OpenAIFileObject: return _to_openai_file_object(MistralFile.model_validate(raw_response.json())) def transform_retrieve_file_request( self, file_id: str, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - return self._file_url(file_id, litellm_params), {} + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + return self._file_url(file_id, litellm_params), _NO_QUERY_PARAMS def transform_retrieve_file_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> OpenAIFileObject: return _to_openai_file_object(MistralFile.model_validate(raw_response.json())) def transform_delete_file_request( self, file_id: str, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - return self._file_url(file_id, litellm_params), {} + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + return self._file_url(file_id, litellm_params), _NO_QUERY_PARAMS def transform_delete_file_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> FileDeleted: deleted: Final = MistralFileDeleted.model_validate(raw_response.json()) return FileDeleted(id=deleted.id, deleted=deleted.deleted, object="file") @@ -195,32 +215,39 @@ class MistralFilesConfig(BaseFilesConfig): def transform_list_files_request( self, purpose: str | None, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - params: Final = {"purpose": _to_mistral_purpose(purpose)} if purpose else {} - return f"{get_mistral_api_base(litellm_params.get('api_base'))}/v1/files", params + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + url: Final = f"{_api_base_from(litellm_params)}/v1/files" + if not purpose: + return url, _NO_QUERY_PARAMS + return url, {"purpose": _to_mistral_purpose(purpose)} # mutable-ok: BaseFilesConfig signature returns dict def transform_list_files_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, - ) -> list[OpenAIFileObject]: - return [_to_openai_file_object(f) for f in MistralFileList.model_validate(raw_response.json()).data] + litellm_params: Mapping[str, object], + ) -> list[OpenAIFileObject]: # mutable-ok: BaseFilesConfig signature + return [ # mutable-ok: BaseFilesConfig signature + _to_openai_file_object(f) for f in MistralFileList.model_validate(raw_response.json()).data + ] def transform_file_content_request( self, file_content_request: FileContentRequest, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - return self._file_url(file_content_request["file_id"], litellm_params, suffix="/content"), {} + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature + file_id: Final = file_content_request.get("file_id") + if file_id is None: + raise ValueError("file_id is required to download file content") + return self._file_url(file_id, litellm_params, suffix="/content"), _NO_QUERY_PARAMS def transform_file_content_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> HttpxBinaryResponseContent: return HttpxBinaryResponseContent(response=raw_response) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 49a01495d65..b3bb1fa9a01 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -1130,7 +1130,7 @@ async def get_file( check_file_id_encoding=True, ) - if should_route: + if should_route and credentials is not None: # Use model-based routing with credentials from config prepare_data_with_credentials( data=data, diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 85f267c2db0..56cd3298db6 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -645,9 +645,7 @@ async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch): lambda content, model: pytest.fail("raw vertex path should not run"), ) - result = await bu.calculate_batch_cost_and_usage( - file_content_dictionary=[], custom_llm_provider="vertex_ai" - ) + result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=[], custom_llm_provider="vertex_ai") assert result.cost == 0.0 assert result.usage.total_tokens == 0 assert result.models == [] @@ -1250,6 +1248,7 @@ async def test_handle_completed_batch_no_output_file_is_zero(monkeypatch): result set - zero cost, zero usage, no models - instead of letting the file fetch raise "Output file id is None" on every aretrieve_batch logging poll. """ + # The output-file fetch must not even be attempted when there is no output file. async def _must_not_fetch(*args, **kwargs): pytest.fail("_fetch_batch_output_file_content should not be called") @@ -1376,7 +1375,10 @@ def test_anthropic_response_body_is_result_message(): def test_anthropic_usage_conversion_includes_cache_tokens(): - body = {"model": "claude-sonnet-4-5-20250929", "usage": _anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000)} + body = { + "model": "claude-sonnet-4-5-20250929", + "usage": _anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000), + } usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="anthropic") assert usage.prompt_tokens == 11000 assert usage.completion_tokens == 200 @@ -1391,7 +1393,9 @@ def test_bedrock_model_output_line_success_check(): "modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}}, } assert bu._batch_response_was_successful(row, custom_llm_provider="bedrock") is True - assert bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6" + assert ( + bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6" + ) def test_bedrock_cost_uses_deployment_model_name(): @@ -1445,7 +1449,13 @@ def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch): rows = [ { "custom_id": "req-1", - "response": {"status_code": 200, "body": {"model": "gpt-5.2", "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}}, + "response": { + "status_code": 200, + "body": { + "model": "gpt-5.2", + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + }, } ] result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") @@ -1487,7 +1497,9 @@ def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatc lambda **kw: pytest.fail("anthropic rows must not go through completion_cost"), ) - result = bu._aggregate_batch_cost_usage_models(entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic") + result = bu._aggregate_batch_cost_usage_models( + entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic" + ) assert result.cost == pytest.approx(0.3) assert seen[0]["model"] == "claude-sonnet-4-5-20250929" @@ -1522,7 +1534,11 @@ async def test_calculate_batch_cost_and_usage_anthropic_end_to_end(): ) assert result.cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2) - assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (11000, 200, 11200) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( + 11000, + 200, + 11200, + ) assert result.models == ["claude-sonnet-4-5"] @@ -1689,7 +1705,10 @@ async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) -> def test_bedrock_converse_shaped_batch_usage_is_parsed(): - body = {"model": "us.amazon.nova-lite-v1:0", "usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742}} + body = { + "model": "us.amazon.nova-lite-v1:0", + "usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742}, + } usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock") assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (2202, 540, 2742) @@ -1739,6 +1758,7 @@ def test_unparsable_bedrock_batch_usage_warns(caplog): # batch_cost_is_final # --------------------------------------------------------------------------- # + def _retrieved_batch( status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None ) -> LiteLLMBatch: @@ -1798,7 +1818,9 @@ def _ocr_row(pages_processed, annotation_pages=None, model="mistral-ocr-latest") usage_info = {"pages_processed": pages_processed, "doc_size_bytes": 4096} if annotation_pages is not None: usage_info["pages_processed_annotation"] = annotation_pages - return _success_row(model=model, pages=[{"index": i, "markdown": "x"} for i in range(pages_processed)], usage_info=usage_info) + return _success_row( + model=model, pages=[{"index": i, "markdown": "x"} for i in range(pages_processed)], usage_info=usage_info + ) def test_ocr_rows_are_priced_per_page_at_batch_rate(monkeypatch): @@ -1835,7 +1857,9 @@ def test_ocr_rows_bill_annotation_pages_separately(monkeypatch): "annotation_cost_per_page_batches": 0.0025, }, ) - result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(4, annotation_pages=4)], custom_llm_provider="mistral") + result = bu._aggregate_batch_cost_usage_models( + entries=[_ocr_row(4, annotation_pages=4)], custom_llm_provider="mistral" + ) assert result.cost == pytest.approx(4 * 0.002 + 4 * 0.0025) diff --git a/tests/test_litellm/batches/test_main.py b/tests/test_litellm/batches/test_main.py index c5a33dd6508..26dc4083b0b 100644 --- a/tests/test_litellm/batches/test_main.py +++ b/tests/test_litellm/batches/test_main.py @@ -66,9 +66,7 @@ def seams(): stack.enter_context(patch.object(bm, "openai_batches_instance", openai_i)) stack.enter_context(patch.object(bm, "azure_batches_instance", azure_i)) stack.enter_context(patch.object(bm, "vertex_ai_batches_instance", vertex_i)) - stack.enter_context( - patch.object(bm, "anthropic_batches_instance", anthropic_i) - ) + stack.enter_context(patch.object(bm, "anthropic_batches_instance", anthropic_i)) stack.enter_context(patch.object(bm, "base_llm_http_handler", base_http)) stack.enter_context(patch.object(bm, "BedrockBatchesHandler", bedrock_arn)) yield Seams( @@ -174,9 +172,7 @@ def test_create__provider_config_routes_to_base_http_handler(seams): "get_provider_batches_config", return_value=MagicMock(name="provider_config"), ): - result = bm.create_batch( - **CREATE_KW, custom_llm_provider="bedrock", model="bedrock/my-batch-model" - ) + result = bm.create_batch(**CREATE_KW, custom_llm_provider="bedrock", model="bedrock/my-batch-model") assert result is seams.base_http.create_batch.return_value _assert_only(seams.base_http.create_batch, seams, "create_batch") @@ -281,9 +277,7 @@ def test_retrieve__bedrock_model_invocation_job_arn(seams): result = bm.retrieve_batch(batch_id=arn, custom_llm_provider="bedrock") seams.bedrock_arn._handle_model_invocation_job_status.assert_called_once() - assert ( - result is seams.bedrock_arn._handle_model_invocation_job_status.return_value - ) + assert result is seams.bedrock_arn._handle_model_invocation_job_status.return_value seams.bedrock_arn._handle_async_invoke_status.assert_not_called() @@ -385,9 +379,7 @@ def test_cancel__unsupported_provider_raises_badrequest(seams): def test_cancel__async_flag_propagates_is_async(seams): - bm.cancel_batch( - batch_id="batch-1", custom_llm_provider="openai", acancel_batch=True - ) + bm.cancel_batch(batch_id="batch-1", custom_llm_provider="openai", acancel_batch=True) assert seams.openai.cancel_batch.call_args.kwargs["_is_async"] is True @@ -415,9 +407,7 @@ async def test_acreate_batch_delegates_to_create_batch(): @pytest.mark.asyncio async def test_aretrieve_batch_delegates_to_retrieve_batch(): with patch.object(bm, "retrieve_batch", MagicMock(return_value="SENTINEL")) as m: - result = await bm.aretrieve_batch( - batch_id="batch-1", custom_llm_provider="azure" - ) + result = await bm.aretrieve_batch(batch_id="batch-1", custom_llm_provider="azure") assert result == "SENTINEL" assert m.call_count == 1 @@ -429,9 +419,7 @@ async def test_aretrieve_batch_delegates_to_retrieve_batch(): @pytest.mark.asyncio async def test_alist_batches_delegates_to_list_batches(): with patch.object(bm, "list_batches", MagicMock(return_value="SENTINEL")) as m: - result = await bm.alist_batches( - after="cur", limit=3, custom_llm_provider="vertex_ai" - ) + result = await bm.alist_batches(after="cur", limit=3, custom_llm_provider="vertex_ai") assert result == "SENTINEL" assert m.call_count == 1 @@ -444,9 +432,7 @@ async def test_alist_batches_delegates_to_list_batches(): @pytest.mark.asyncio async def test_acancel_batch_delegates_to_cancel_batch(): with patch.object(bm, "cancel_batch", MagicMock(return_value="SENTINEL")) as m: - result = await bm.acancel_batch( - batch_id="batch-1", custom_llm_provider="openai" - ) + result = await bm.acancel_batch(batch_id="batch-1", custom_llm_provider="openai") assert result == "SENTINEL" assert m.call_count == 1 @@ -499,9 +485,7 @@ def _sent(mock_method, *keys): def test_create__openai_credentials_passthrough(seams): bm.create_batch(**CREATE_KW, custom_llm_provider="openai", **OPENAI_CREDS) - assert _sent( - seams.openai.create_batch, "api_key", "api_base", "organization", "max_retries" - ) == { + assert _sent(seams.openai.create_batch, "api_key", "api_base", "organization", "max_retries") == { "api_key": "sk-user-openai", "api_base": "https://openai.user.test", "organization": "org-user-123", @@ -512,9 +496,7 @@ def test_create__openai_credentials_passthrough(seams): def test_create__azure_credentials_passthrough(seams): bm.create_batch(**CREATE_KW, custom_llm_provider="azure", **AZURE_CREDS) - assert _sent( - seams.azure.create_batch, "api_key", "api_base", "api_version" - ) == { + assert _sent(seams.azure.create_batch, "api_key", "api_base", "api_version") == { "api_key": "sk-user-azure", "api_base": "https://azure.user.test", "api_version": "2024-12-99", @@ -564,9 +546,7 @@ def test_create__provider_config_credentials_passthrough(seams): def test_retrieve__openai_credentials_passthrough(seams): bm.retrieve_batch(batch_id="b1", custom_llm_provider="openai", **OPENAI_CREDS) - assert _sent( - seams.openai.retrieve_batch, "api_key", "api_base", "organization" - ) == { + assert _sent(seams.openai.retrieve_batch, "api_key", "api_base", "organization") == { "api_key": "sk-user-openai", "api_base": "https://openai.user.test", "organization": "org-user-123", @@ -576,9 +556,7 @@ def test_retrieve__openai_credentials_passthrough(seams): def test_retrieve__azure_credentials_passthrough(seams): bm.retrieve_batch(batch_id="b1", custom_llm_provider="azure", **AZURE_CREDS) - assert _sent( - seams.azure.retrieve_batch, "api_key", "api_base", "api_version" - ) == { + assert _sent(seams.azure.retrieve_batch, "api_key", "api_base", "api_version") == { "api_key": "sk-user-azure", "api_base": "https://azure.user.test", "api_version": "2024-12-99", @@ -640,9 +618,7 @@ def test_retrieve__provider_config_credentials_passthrough(seams): def test_list__openai_credentials_passthrough(seams): bm.list_batches(custom_llm_provider="openai", **OPENAI_CREDS) - assert _sent( - seams.openai.list_batches, "api_key", "api_base", "organization" - ) == { + assert _sent(seams.openai.list_batches, "api_key", "api_base", "organization") == { "api_key": "sk-user-openai", "api_base": "https://openai.user.test", "organization": "org-user-123", @@ -652,9 +628,7 @@ def test_list__openai_credentials_passthrough(seams): def test_list__azure_credentials_passthrough(seams): bm.list_batches(custom_llm_provider="azure", **AZURE_CREDS) - assert _sent( - seams.azure.list_batches, "api_key", "api_base", "api_version" - ) == { + assert _sent(seams.azure.list_batches, "api_key", "api_base", "api_version") == { "api_key": "sk-user-azure", "api_base": "https://azure.user.test", "api_version": "2024-12-99", @@ -682,9 +656,7 @@ def test_list__vertex_credentials_passthrough(seams): def test_cancel__openai_credentials_passthrough(seams): bm.cancel_batch(batch_id="b1", custom_llm_provider="openai", **OPENAI_CREDS) - assert _sent( - seams.openai.cancel_batch, "api_key", "api_base", "organization" - ) == { + assert _sent(seams.openai.cancel_batch, "api_key", "api_base", "organization") == { "api_key": "sk-user-openai", "api_base": "https://openai.user.test", "organization": "org-user-123", @@ -694,9 +666,7 @@ def test_cancel__openai_credentials_passthrough(seams): def test_cancel__azure_credentials_passthrough(seams): bm.cancel_batch(batch_id="b1", custom_llm_provider="azure", **AZURE_CREDS) - assert _sent( - seams.azure.cancel_batch, "api_key", "api_base", "api_version" - ) == { + assert _sent(seams.azure.cancel_batch, "api_key", "api_base", "api_version") == { "api_key": "sk-user-azure", "api_base": "https://azure.user.test", "api_version": "2024-12-99", @@ -786,18 +756,16 @@ def test_retrieve__omits_trusted_model_credentials_when_not_supplied(seams): def test_create__mistral_ocr_routes_to_base_http_handler_with_mistral_config(seams): - with patch.object(bm.ProviderConfigManager, "get_provider_batches_config", wraps=bm.ProviderConfigManager.get_provider_batches_config) as get_cfg: - result = bm.create_batch( - completion_window="24h", - endpoint="/v1/ocr", - input_file_id="file-abc", - custom_llm_provider="mistral", - model="mistral/mistral-ocr-latest", - ) + result = bm.create_batch( + completion_window="24h", + endpoint="/v1/ocr", + input_file_id="file-abc", + custom_llm_provider="mistral", + model="mistral/mistral-ocr-latest", + ) assert result is seams.base_http.create_batch.return_value _assert_only(seams.base_http.create_batch, seams, "create_batch") - get_cfg.assert_called_once() forwarded = seams.base_http.create_batch.call_args.kwargs assert type(forwarded["provider_config"]).__name__ == "MistralBatchesConfig" assert forwarded["model"] == "mistral-ocr-latest" diff --git a/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py index 03e9c351a30..03cfeedece2 100644 --- a/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py +++ b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py @@ -92,26 +92,34 @@ def test_create_request_maps_openai_fields_onto_mistral_job(config): model="mistral-ocr-latest", create_batch_data=data, optional_params={}, litellm_params={} ) assert body == { - "input_files": ["file-123"], + "input_files": ("file-123",), "endpoint": "/v1/ocr", "model": "mistral-ocr-latest", "metadata": {"team": "docs"}, } -def test_create_request_omits_empty_metadata_and_forwards_extra_body(config): +def test_create_request_omits_empty_metadata(config): data = CreateBatchRequest( completion_window="24h", endpoint="/v1/chat/completions", input_file_id="file-123", metadata=None, - extra_body={"timeout_hours": 48}, ) body = config.transform_create_batch_request( model="mistral-small-latest", create_batch_data=data, optional_params={}, litellm_params={} ) assert "metadata" not in body - assert body["timeout_hours"] == 48 + + +def test_create_request_requires_input_file_and_endpoint(config): + with pytest.raises(ValueError, match="input_file_id and endpoint are required"): + config.transform_create_batch_request( + model="m", + create_batch_data=CreateBatchRequest(completion_window="24h"), + optional_params={}, + litellm_params={}, + ) @pytest.mark.parametrize( diff --git a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py index d6ad8a34b35..f62645be7ee 100644 --- a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py +++ b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py @@ -79,7 +79,9 @@ def test_validate_environment_uses_bearer_auth(config, api_key): def test_upload_request_is_multipart_with_batch_purpose(config): body = config.transform_create_file_request( model="", - create_file_data=CreateFileRequest(file=("in.jsonl", b'{"custom_id":"0"}\n', "application/jsonl"), purpose="batch"), + create_file_data=CreateFileRequest( + file=("in.jsonl", b'{"custom_id":"0"}\n', "application/jsonl"), purpose="batch" + ), optional_params={}, litellm_params={}, ) @@ -181,7 +183,9 @@ def test_list_request_filters_by_mapped_purpose(config): def test_list_response(config): out = config.transform_list_files_response( - raw_response=_response({"data": [_file(), _file(id="second", filename="b.jsonl")], "object": "list", "total": 2}), + raw_response=_response( + {"data": [_file(), _file(id="second", filename="b.jsonl")], "object": "list", "total": 2} + ), logging_obj=None, litellm_params={}, ) From 91e7df3d8fe83c37268b61080a37e48ef0e63e3f Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 10 Sep 2026 13:35:35 -0400 Subject: [PATCH 006/251] chore: regenerate cost-map schema and UI API types, drop test banner comments --- model_prices_and_context_window.schema.json | 8 ++++++++ tests/test_litellm/batches/test_batch_utils.py | 5 ----- .../batches/test_mistral_batches_transformation.py | 10 ---------- .../llms/mistral/ocr/test_mistral_ocr_cost.py | 1 - ui/litellm-dashboard/src/lib/http/schema.d.ts | 8 ++++++++ 5 files changed, 16 insertions(+), 16 deletions(-) diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 47a1934a703..09385db728b 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -53,6 +53,10 @@ "type": "number", "minimum": 0 }, + "annotation_cost_per_page_batches": { + "type": "number", + "minimum": 0 + }, "audio_transcription_config": { "type": "string" }, @@ -432,6 +436,10 @@ "type": "number", "minimum": 0 }, + "ocr_cost_per_page_batches": { + "type": "number", + "minimum": 0 + }, "output_cost_per_audio_token": { "type": "number", "minimum": 0 diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 56cd3298db6..0d66c0eb5ec 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1809,11 +1809,6 @@ class TestBatchCostIsFinal: assert bu.batch_cost_is_final(_retrieved_batch(status)) is True -# =========================================================================== # -# OCR batch output lines (Mistral /v1/ocr batches) are billed per page, not per token -# =========================================================================== # - - def _ocr_row(pages_processed, annotation_pages=None, model="mistral-ocr-latest"): usage_info = {"pages_processed": pages_processed, "doc_size_bytes": 4096} if annotation_pages is not None: diff --git a/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py index 03cfeedece2..4073879e3b8 100644 --- a/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py +++ b/tests/test_litellm/llms/mistral/batches/test_mistral_batches_transformation.py @@ -76,11 +76,6 @@ def test_custom_llm_provider(config): assert config.custom_llm_provider == LlmProviders.MISTRAL -# --------------------------------------------------------------------------- # -# create -# --------------------------------------------------------------------------- # - - def test_create_request_maps_openai_fields_onto_mistral_job(config): data = CreateBatchRequest( completion_window="24h", @@ -175,11 +170,6 @@ def test_create_response_maps_job_onto_openai_batch(config): assert batch.metadata == {"job_type": "testing"} -# --------------------------------------------------------------------------- # -# retrieve -# --------------------------------------------------------------------------- # - - def test_retrieve_request_is_presigned_get_with_auth(config, api_key): req = config.transform_retrieve_batch_request( batch_id="job/with slash", optional_params={}, litellm_params={"api_base": "https://api.mistral.ai"} diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py index 9fe6f38003f..d72e866949f 100644 --- a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py +++ b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py @@ -63,7 +63,6 @@ def test_ocr4_cost_scales_with_pages(model: str, pages_processed: int) -> None: assert cost == pytest.approx(OCR4_COST_PER_PAGE * pages_processed) - @pytest.mark.parametrize("cost_map_path", [MAIN_COST_MAP, BACKUP_COST_MAP]) def test_ocr3_pricing_entry(cost_map_path: Path) -> None: with open(cost_map_path) as f: diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0f0c1fc9af4..5d1ed79098e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29493,6 +29493,8 @@ export interface components { allow_client_keepalive_override: boolean | null; /** Annotation Cost Per Page */ annotation_cost_per_page?: number | null; + /** Annotation Cost Per Page Batches */ + annotation_cost_per_page_batches?: number | null; /** Api Base */ api_base?: string | null; /** Api Key */ @@ -29704,6 +29706,8 @@ export interface components { ocr_cost_per_credit?: number | null; /** Ocr Cost Per Page */ ocr_cost_per_page?: number | null; + /** Ocr Cost Per Page Batches */ + ocr_cost_per_page_batches?: number | null; /** Organization */ organization?: string | null; /** Otpm */ @@ -39684,6 +39688,8 @@ export interface components { allow_client_keepalive_override: boolean | null; /** Annotation Cost Per Page */ annotation_cost_per_page?: number | null; + /** Annotation Cost Per Page Batches */ + annotation_cost_per_page_batches?: number | null; /** Api Base */ api_base?: string | null; /** Api Key */ @@ -39895,6 +39901,8 @@ export interface components { ocr_cost_per_credit?: number | null; /** Ocr Cost Per Page */ ocr_cost_per_page?: number | null; + /** Ocr Cost Per Page Batches */ + ocr_cost_per_page_batches?: number | null; /** Organization */ organization?: string | null; /** Otpm */ From edd5727f3c9077f449eec37367ace43945e649fa Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 10 Sep 2026 17:28:27 -0400 Subject: [PATCH 007/251] fix(proxy): enforce key/team/org/project model grants on model-routed file and batch credentials Files and batches routes take their model from a header, query param or a model-encoded resource id, which the auth layer never sees, so any key could name any deployment and act on that provider account with its server-side key. Every caller-supplied model now goes through can_key_call_resolved_model before deployment credentials are resolved, covering file create/retrieve/content/ delete/list, batch create/retrieve/list/cancel, and vector store files. --- litellm/proxy/batches_endpoints/endpoints.py | 17 +- .../openai_files_endpoints/common_utils.py | 61 +++++- .../openai_files_endpoints/files_endpoints.py | 20 +- .../vector_store_files_endpoints/endpoints.py | 6 +- .../proxy/batches_endpoints/test_endpoints.py | 58 +++++- .../test_files_endpoint.py | 184 ++++++++++++++++-- .../test_batch_x_litellm_model_encoding.py | 53 ++--- 7 files changed, 322 insertions(+), 77 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 5c4bacd757c..c99f66d032e 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -34,9 +34,9 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( encode_batch_response_ids, encode_file_id_with_model, ensure_batch_response_managed_file_ids, + get_authorized_credentials_for_model, get_batch_from_database, get_batch_id_from_unified_batch_id, - get_credentials_for_model, get_model_id_from_unified_batch_id, get_models_from_unified_file_id, get_original_file_id, @@ -218,9 +218,10 @@ async def create_batch( # SCENARIO 1: File ID is encoded with model info if model_from_file_id is not None and input_file_id: - credentials = get_credentials_for_model( + credentials = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_file_id, + user_api_key_dict=user_api_key_dict, operation_context="batch creation (file created with model)", ) @@ -310,9 +311,10 @@ async def create_batch( # SCENARIO 2 & 3: Model from header/query OR custom_llm_provider fallback if model_param: # SCENARIO 2: Use model-based routing from header/query/body - credentials = get_credentials_for_model( + credentials = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_param, + user_api_key_dict=user_api_key_dict, operation_context="batch creation", ) @@ -540,9 +542,10 @@ async def retrieve_batch( # Retrieve from provider (for non-terminal states or if DB lookup failed) # SCENARIO 1: Batch ID is encoded with model info if model_from_id is not None: - credentials: Final = get_credentials_for_model( + credentials: Final = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_id, + user_api_key_dict=user_api_key_dict, operation_context="batch retrieval (batch created with model)", ) @@ -764,9 +767,10 @@ async def list_batches( data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") ): # SCENARIO 2: Use model-based routing from header/query/body - credentials: Final = get_credentials_for_model( + credentials: Final = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_param, + user_api_key_dict=user_api_key_dict, operation_context="batch listing", ) @@ -952,9 +956,10 @@ async def cancel_batch( # SCENARIO 1: Batch ID is encoded with model info if model_from_id is not None: - credentials: Final = get_credentials_for_model( + credentials: Final = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_id, + user_api_key_dict=user_api_key_dict, operation_context="batch cancellation (batch created with model)", ) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index b1f282a0978..4202a6d1689 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -351,6 +351,10 @@ def get_credentials_for_model( """ Retrieve API credentials for a model from the LLM Router. + Does not check whether the caller may use ``model_id``; use + ``get_authorized_credentials_for_model`` for anything driven by a caller-supplied + model name (request body, header, query param, or a model-encoded resource id). + Args: llm_router: LiteLLM Router instance model_id: Model name or deployment ID @@ -381,6 +385,48 @@ def get_credentials_for_model( return credentials +async def authorize_model_for_key( + model_id: str, + llm_router: Optional["Router"], + user_api_key_dict: "UserAPIKeyAuth", +) -> None: + """ + Enforce the caller's model grants on a model name the auth layer never saw. + + The files and batches routes carry their model in a header, query param, or a + model-encoded resource id rather than the request body, so ``user_api_key_auth`` + cannot check it. Run the same key, team (incl. team-member and access-group + fallbacks), org and project allowlist checks a chat request would get, so a + restricted key cannot borrow another deployment's server-side credentials. + + Raises: + ProxyException (403): the caller is not allowed to use ``model_id`` + """ + from litellm.proxy.auth.auth_checks import can_key_call_resolved_model + + await can_key_call_resolved_model( + model=model_id, + llm_model_list=None, + valid_token=user_api_key_dict, + llm_router=llm_router, + ) + + +async def get_authorized_credentials_for_model( + llm_router: Optional["Router"], + model_id: str, + user_api_key_dict: "UserAPIKeyAuth", + operation_context: str = "file operation", +) -> dict: # mutable-ok: same contract as get_credentials_for_model, callers merge it into request data + """``get_credentials_for_model`` gated by ``authorize_model_for_key``.""" + await authorize_model_for_key(model_id=model_id, llm_router=llm_router, user_api_key_dict=user_api_key_dict) + return get_credentials_for_model( + llm_router=llm_router, + model_id=model_id, + operation_context=operation_context, + ) + + def get_team_provider_credentials( llm_router: Optional["Router"], user_api_key_dict: "UserAPIKeyAuth", @@ -573,21 +619,27 @@ def prepare_data_with_credentials( data["file_id"] = file_id -def handle_model_based_routing( +async def handle_model_based_routing( file_id: str, request, # FastAPI Request object llm_router, # Router instance data: dict, + user_api_key_dict: "UserAPIKeyAuth", check_file_id_encoding: bool = True, ) -> tuple[bool, str | None, str | None, dict | None]: """ Orchestrate model-based credential routing for file operations. + The model name comes from the caller (embedded in the file id, or a header, query + param or body field), so it is authorized against the caller's key, team, org and + project grants before any deployment credentials are resolved. + Args: file_id: File ID (may contain embedded model info) request: FastAPI request object llm_router: LiteLLM Router instance data: Request data dictionary + user_api_key_dict: The authenticated caller check_file_id_encoding: Whether to check for embedded model in file_id Returns: @@ -599,6 +651,7 @@ def handle_model_based_routing( Raises: HTTPException: If router unavailable or model not found + ProxyException: If the caller is not allowed to use the model """ model_from_id, model_from_param = extract_model_from_sources( file_id=file_id, @@ -608,9 +661,10 @@ def handle_model_based_routing( # Priority 1: Model embedded in file_id if check_file_id_encoding and model_from_id is not None: - credentials = get_credentials_for_model( + credentials = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_id, + user_api_key_dict=user_api_key_dict, operation_context=f"file operation (file created with model '{model_from_id}')", ) original_file_id: Final = get_original_file_id(file_id) @@ -618,9 +672,10 @@ def handle_model_based_routing( # Priority 2: Model from header/query/body elif model_from_param is not None: - credentials = get_credentials_for_model( + credentials = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model_from_param, + user_api_key_dict=user_api_key_dict, operation_context="file operation", ) return True, model_from_param, None, credentials diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index b3bb1fa9a01..0efd618e171 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -68,7 +68,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( apply_team_provider_credentials, encode_file_id_with_model, extract_file_creation_params, - get_credentials_for_model, + get_authorized_credentials_for_model, handle_model_based_routing, prepare_data_with_credentials, validate_file_list_limit, @@ -267,9 +267,10 @@ async def route_create_file( # NEW: Handle model-based routing (no DB required) if model is not None: # Get credentials from model_list via router - credentials: Final = get_credentials_for_model( + credentials: Final = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=model, + user_api_key_dict=user_api_key_dict, operation_context="file upload", ) @@ -907,11 +908,12 @@ async def get_file_content( model_used, original_file_id, credentials, - ) = handle_model_based_routing( + ) = await handle_model_based_routing( file_id=file_id, request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=True, ) @@ -1122,11 +1124,12 @@ async def get_file( model_used, original_file_id, credentials, - ) = handle_model_based_routing( + ) = await handle_model_based_routing( file_id=file_id, request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=True, ) @@ -1330,11 +1333,12 @@ async def delete_file( model_used, original_file_id, credentials, - ) = handle_model_based_routing( + ) = await handle_model_based_routing( file_id=file_id, request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=True, ) @@ -1517,11 +1521,12 @@ async def list_files( response: Any | None = None # Check for model-based credential routing (no file_id encoding check for list) - should_route, model_used, _, credentials = handle_model_based_routing( + should_route, model_used, _, credentials = await handle_model_based_routing( file_id="", # No file_id for list endpoint request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=False, ) @@ -1548,9 +1553,10 @@ async def list_files( status_code=500, detail="LLM Router not initialized. Ensure models added to proxy.", ) - credentials = get_credentials_for_model( + credentials = await get_authorized_credentials_for_model( llm_router=llm_router, model_id=target_model_names_list[0], + user_api_key_dict=user_api_key_dict, operation_context="file list", ) prepare_data_with_credentials(data=data, credentials=credentials) diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index 957ed9fd0b9..11ef8efb598 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -144,11 +144,12 @@ async def _update_request_data_with_managed_file_id( model_used, original_file_id, credentials, - ) = handle_model_based_routing( + ) = await handle_model_based_routing( file_id=file_id, request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=True, ) @@ -273,11 +274,12 @@ async def _update_request_data_with_model_routing_hint( _model_used, _original_file_id, credentials, - ) = handle_model_based_routing( + ) = await handle_model_based_routing( file_id="", request=request, llm_router=llm_router, data=data, + user_api_key_dict=user_api_key_dict, check_file_id_encoding=False, ) diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index a37c8ff2bb4..57e42e79a42 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -177,6 +177,8 @@ def harness(): logging.get_proxy_hook = MagicMock(return_value=None) router = MagicMock(spec=Router) + router.model_group_alias = {} + router.get_model_access_groups = MagicMock(return_value={}) router.acreate_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -1161,6 +1163,8 @@ def retrieve_harness(): logging.get_proxy_hook = MagicMock(return_value=None) router = MagicMock(spec=Router) + router.model_group_alias = {} + router.get_model_access_groups = MagicMock(return_value={}) router.aretrieve_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -1616,6 +1620,8 @@ def list_harness(): logging.get_proxy_hook = MagicMock(return_value=None) router = MagicMock(spec=Router) + router.model_group_alias = {} + router.get_model_access_groups = MagicMock(return_value={}) router.alist_batches = AsyncMock(return_value=FakeListPage([])) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -2012,6 +2018,8 @@ def cancel_harness(): logging.get_proxy_hook = MagicMock(return_value=None) router = MagicMock(spec=Router) + router.model_group_alias = {} + router.get_model_access_groups = MagicMock(return_value={}) router.acancel_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -2733,8 +2741,6 @@ async def test_cancel__unified_batch_id_allowed_when_managed_files_required(canc assert cancel_harness.router_acancel.call_count == 1 - - @pytest.mark.asyncio async def test_retrieve__managed_batch_defers_cost_to_the_poller_when_it_is_running(retrieve_harness): with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=True)): @@ -2762,3 +2768,51 @@ async def test_retrieve__raw_batch_id_is_untouched_by_the_poller_handoff(retriev metadata = retrieve_harness.litellm_aretrieve.await_args.kwargs.get("litellm_metadata") or {} assert metadata.get("batch_ignore_default_logging") is None + + +def _key_restricted_to(*models: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-restricted", team_id="team-a", team_models=list(models), models=list(models)) + + +@pytest.mark.asyncio +async def test_create__header_model_rejects_key_without_model_grant(harness): + """A key not granted the model named in x-litellm-model must not receive that deployment's credentials.""" + set_body(harness, {"input_file_id": "file-plain", "endpoint": "/v1/chat/completions", "completion_window": "24h"}) + + with pytest.raises(ProxyException) as exc_info: + await call_create(harness, user=_key_restricted_to("azure/gpt-4o"), headers={"x-litellm-model": "vertex-model"}) + + assert exc_info.value.code == "403" + harness.creds_resolver.assert_not_called() + harness.litellm_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_create__header_model_allows_key_with_model_grant(harness): + set_body(harness, {"input_file_id": "file-plain", "endpoint": "/v1/chat/completions", "completion_window": "24h"}) + + await call_create(harness, user=_key_restricted_to("vertex-model"), headers={"x-litellm-model": "vertex-model"}) + + harness.creds_resolver.assert_called_once_with(model_id="vertex-model") + assert harness.acreate_kwargs()["custom_llm_provider"] == "vertex_ai" + + +@pytest.mark.asyncio +async def test_retrieve__model_encoded_id_rejects_key_without_model_grant(retrieve_harness): + """The model embedded in a batch id is caller-controlled, so it is checked against the key's grants too.""" + with pytest.raises(ProxyException) as exc_info: + await call_retrieve(retrieve_harness, AZURE_BATCH_ID, user=_key_restricted_to("vertex-model")) + + assert exc_info.value.code == "403" + retrieve_harness.creds_resolver.assert_not_called() + retrieve_harness.litellm_aretrieve.assert_not_called() + + +@pytest.mark.asyncio +async def test_cancel__model_encoded_id_rejects_key_without_model_grant(cancel_harness): + with pytest.raises(ProxyException) as exc_info: + await call_cancel(cancel_harness, AZURE_BATCH_ID, user=_key_restricted_to("vertex-model")) + + assert exc_info.value.code == "403" + cancel_harness.creds_resolver.assert_not_called() + cancel_harness.litellm_acancel.assert_not_called() diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 5faae166fca..548c0eb0d91 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -1877,7 +1877,7 @@ def test_get_file_content_streams_openai_direct_path( monkeypatch.setattr(litellm, "afile_content", _mock_afile_content) monkeypatch.setattr( "litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing", - lambda **kwargs: (False, None, None, None), + AsyncMock(return_value=(False, None, None, None)), ) app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( @@ -1942,15 +1942,17 @@ def test_get_file_content_routed_provider_skips_streaming_when_resolved_provider ) monkeypatch.setattr( "litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing", - lambda **kwargs: ( - True, - "azure-gpt-3-5-turbo", - "file-original-123", - { - "custom_llm_provider": "azure", - "api_key": "azure-key", - "api_base": "https://azure.example.com", - }, + AsyncMock( + return_value=( + True, + "azure-gpt-3-5-turbo", + "file-original-123", + { + "custom_llm_provider": "azure", + "api_key": "azure-key", + "api_base": "https://azure.example.com", + }, + ) ), ) @@ -2015,7 +2017,7 @@ def test_get_file_content_non_openai_provider_skips_streaming_handler( ) monkeypatch.setattr( "litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing", - lambda **kwargs: (False, None, None, None), + AsyncMock(return_value=(False, None, None, None)), ) app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( @@ -2463,14 +2465,16 @@ def test_list_files_model_routing_does_not_forward_custom_llm_provider_twice( monkeypatch.setattr(litellm, "afile_list", _mock_afile_list) monkeypatch.setattr( "litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing", - lambda **kwargs: ( - True, - "azure-gpt-4o", - None, - { - "custom_llm_provider": "azure", - "api_key": "azure-key", - }, + AsyncMock( + return_value=( + True, + "azure-gpt-4o", + None, + { + "custom_llm_provider": "azure", + "api_key": "azure-key", + }, + ) ), ) @@ -4878,3 +4882,145 @@ def test_get_file_model_routed_id_forwards_deployment_provider(mocker: MockerFix assert captured_kwargs["api_key"] == "mistral-key" assert captured_kwargs["file_id"] == "7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df" assert response.json()["id"] == encoded_id + + +def _mistral_plus_anthropic_router() -> Router: + return Router( + model_list=[ + { + "model_name": "mistral-ocr", + "litellm_params": {"model": "mistral/mistral-ocr-latest", "api_key": "mistral-key"}, + "model_info": {"id": "mistral-ocr-id"}, + }, + { + "model_name": "claude-opus-4-6", + "litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "anthropic-key"}, + "model_info": {"id": "claude-id"}, + }, + ] + ) + + +def _restricted_key(key_models: list[str]) -> UserAPIKeyAuth: + from litellm.proxy._types import LitellmUserRoles + + return UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="test-user", + team_id="team-a", + team_models=["claude-opus-4-6", "mistral-ocr"], + models=key_models, + ) + + +@pytest.mark.parametrize( + "http_method, path_suffix, litellm_fn", + [ + ("get", "", "afile_retrieve"), + ("get", "/content", "afile_content"), + ("delete", "", "afile_delete"), + ], +) +def test_model_routed_file_ops_reject_key_without_model_grant( + mocker: MockerFixture, monkeypatch, http_method: str, path_suffix: str, litellm_fn: str +): + """ + Regression: a key whose allowlist does not include the deployment named in a + model-encoded file id must be refused before that deployment's server-side + credentials are resolved. Previously any key could name any deployment via the + id (or the x-litellm-model header) and act on that provider account's files. + """ + import litellm.proxy.proxy_server as ps + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + router = _mistral_plus_anthropic_router() + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + + upstream = mocker.AsyncMock(side_effect=AssertionError("provider must not be called")) + monkeypatch.setattr(litellm, litellm_fn, upstream) + app.dependency_overrides[ps.user_api_key_auth] = lambda: _restricted_key(["claude-opus-4-6"]) + encoded_id = encode_file_id_with_model("7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", "mistral-ocr") + + try: + response = getattr(client, http_method)( + f"/v1/files/{encoded_id}{path_suffix}", headers={"Authorization": "Bearer test-key"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 403, response.text + assert "not allowed to access model" in response.text + upstream.assert_not_called() + + +def test_list_files_header_model_rejects_key_without_model_grant(mocker: MockerFixture, monkeypatch): + import litellm.proxy.proxy_server as ps + + router = _mistral_plus_anthropic_router() + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + + upstream = mocker.AsyncMock(side_effect=AssertionError("provider must not be called")) + monkeypatch.setattr(litellm, "afile_list", upstream) + app.dependency_overrides[ps.user_api_key_auth] = lambda: _restricted_key(["claude-opus-4-6"]) + + try: + response = client.get( + "/v1/files", headers={"Authorization": "Bearer test-key", "x-litellm-model": "mistral-ocr"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 403, response.text + upstream.assert_not_called() + + +def test_model_routed_file_retrieve_allows_key_with_model_grant(mocker: MockerFixture, monkeypatch): + """The grant check must not break the happy path: a key allowed the deployment still resolves its credentials.""" + import litellm.proxy.proxy_server as ps + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + router = _mistral_plus_anthropic_router() + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + + captured_kwargs: dict = {} + + async def _mock_afile_retrieve(**kwargs): + captured_kwargs.update(kwargs) + return OpenAIFileObject( + id="7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", + object="file", + bytes=2, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(litellm, "afile_retrieve", _mock_afile_retrieve) + app.dependency_overrides[ps.user_api_key_auth] = lambda: _restricted_key(["mistral-ocr"]) + encoded_id = encode_file_id_with_model("7a13fa8e-fcf8-42c5-aa61-c93c10e2c7df", "mistral-ocr") + + try: + response = client.get(f"/v1/files/{encoded_id}", headers={"Authorization": "Bearer test-key"}) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + assert captured_kwargs["api_key"] == "mistral-key" + assert captured_kwargs["custom_llm_provider"] == "mistral" diff --git a/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py index 3d1831bb4cd..fe4903b547c 100644 --- a/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py +++ b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py @@ -58,10 +58,7 @@ def _make_batch_response( def test_get_batch_id_from_unified_batch_id_handles_appended_fields(): - decoded_id = ( - "litellm_proxy;model_id:deployment-123;" - "llm_batch_id:batch_openai_123;llm_output_file_id:file-output" - ) + decoded_id = "litellm_proxy;model_id:deployment-123;llm_batch_id:batch_openai_123;llm_output_file_id:file-output" assert get_batch_id_from_unified_batch_id(decoded_id) == "batch_openai_123" @@ -107,12 +104,10 @@ async def test_create_batch_with_x_litellm_model_encodes_batch_id(): } ), ), + patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls, patch( - "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_processor_cls, - patch( - "litellm.proxy.batches_endpoints.endpoints.get_credentials_for_model", - return_value=mock_credentials, + "litellm.proxy.batches_endpoints.endpoints.get_authorized_credentials_for_model", + new=AsyncMock(return_value=mock_credentials), ), patch( "litellm.proxy.batches_endpoints.endpoints.prepare_data_with_credentials", @@ -165,23 +160,15 @@ async def test_create_batch_with_x_litellm_model_encodes_batch_id(): ) # The batch_id should be encoded with model info - assert ( - response.id != raw_batch_id - ), f"Expected batch_id to be encoded, but got raw ID: {response.id}" - assert response.id.startswith( - "batch_" - ), f"Encoded batch_id should keep batch_ prefix, got: {response.id}" + assert response.id != raw_batch_id, f"Expected batch_id to be encoded, but got raw ID: {response.id}" + assert response.id.startswith("batch_"), f"Encoded batch_id should keep batch_ prefix, got: {response.id}" # Should be decodable back to the original decoded_model = decode_model_from_file_id(response.id) - assert ( - decoded_model == model_name - ), f"Expected model '{model_name}' from decoded batch_id, got: {decoded_model}" + assert decoded_model == model_name, f"Expected model '{model_name}' from decoded batch_id, got: {decoded_model}" original_id = get_original_file_id(response.id) - assert ( - original_id == raw_batch_id - ), f"Expected original ID '{raw_batch_id}', got: {original_id}" + assert original_id == raw_batch_id, f"Expected original ID '{raw_batch_id}', got: {original_id}" assert mock_create_batch.call_args.kwargs["metadata"] == {"customer_id": "cust-123"} @@ -227,12 +214,10 @@ async def test_create_batch_with_x_litellm_model_encodes_output_and_error_file_i } ), ), + patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls, patch( - "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_processor_cls, - patch( - "litellm.proxy.batches_endpoints.endpoints.get_credentials_for_model", - return_value=mock_credentials, + "litellm.proxy.batches_endpoints.endpoints.get_authorized_credentials_for_model", + new=AsyncMock(return_value=mock_credentials), ), patch( "litellm.proxy.batches_endpoints.endpoints.prepare_data_with_credentials", @@ -316,9 +301,7 @@ async def test_create_batch_without_x_litellm_model_returns_raw_ids(monkeypatch) } ), ), - patch( - "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_processor_cls, + patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls, patch( "litellm.acreate_batch", new=AsyncMock(return_value=mock_response), @@ -383,9 +366,7 @@ class TestBatchIdRoundTripWithRetrieve: raw_batch_id = "batch_vllm_12345" # What create_batch does: - encoded_id = encode_file_id_with_model( - file_id=raw_batch_id, model=model_name, id_type="batch" - ) + encoded_id = encode_file_id_with_model(file_id=raw_batch_id, model=model_name, id_type="batch") # What retrieve_batch does: decoded_model = decode_model_from_file_id(encoded_id) @@ -410,9 +391,7 @@ class TestBatchIdRoundTripWithRetrieve: ] for raw_id, model in test_cases: - encoded = encode_file_id_with_model( - file_id=raw_id, model=model, id_type="batch" - ) + encoded = encode_file_id_with_model(file_id=raw_id, model=model, id_type="batch") assert encoded.startswith("batch_") assert decode_model_from_file_id(encoded) == model assert get_original_file_id(encoded) == raw_id @@ -440,9 +419,7 @@ async def test_cancel_batch_with_unified_id_routes_with_decoded_model_and_batch_ mock_user_api_key_dict.team_metadata = {} with ( - patch( - "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" - ) as mock_processor_cls, + patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls, patch( "litellm.proxy.batches_endpoints.endpoints.update_batch_in_database", new=AsyncMock(), From bae731ddfc394b23d3c6f44a85cfaa58472897e7 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 10 Sep 2026 17:54:57 -0400 Subject: [PATCH 008/251] fix(proxy): apply model grants to unified file and batch ids on batch routes Unified ids carry the deployment model inside the id, so a restricted key could create, retrieve or cancel a batch on a deployment it is not granted. The model parsed from a unified id now goes through the same grant check as header, query and model-encoded id sources before the router is called. --- litellm/proxy/batches_endpoints/endpoints.py | 16 ++++- .../proxy/batches_endpoints/test_endpoints.py | 58 +++++++++++++++++++ .../test_batch_x_litellm_model_encoding.py | 7 +-- 3 files changed, 75 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index c99f66d032e..c2489ce52ea 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -29,6 +29,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, add_internal_model_credentials, apply_team_provider_credentials, + authorize_model_for_key, batch_cost_poller_is_active, decode_model_from_file_id, encode_batch_response_ids, @@ -286,6 +287,7 @@ async def create_batch( detail={"error": f"Expected 1 model, got {len(target_model_names)}"}, ) model: Final = target_model_names[0] + await authorize_model_for_key(model_id=model, llm_router=llm_router, user_api_key_dict=user_api_key_dict) _create_batch_data["model"] = model resolved_storage_url: Final = await _resolve_managed_input_file_storage_url(input_file_id) @@ -582,10 +584,17 @@ async def retrieve_batch( ) if unified_batch_id: + unified_model_id: Final = get_model_id_from_unified_batch_id(unified_batch_id) + if unified_model_id is not None: + await authorize_model_for_key( + model_id=llm_router.resolve_model_name_from_model_id(unified_model_id) or unified_model_id, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + ) add_internal_model_credentials( data=data, llm_router=llm_router, - model_id=get_model_id_from_unified_batch_id(unified_batch_id), + model_id=unified_model_id, ) response = await llm_router.aretrieve_batch(**data) @@ -998,6 +1007,11 @@ async def cancel_batch( status_code=400, detail={"error": "Invalid LiteLLM managed batch ID. Missing model_id."}, ) + await authorize_model_for_key( + model_id=llm_router.resolve_model_name_from_model_id(model_id_from_batch) or model_id_from_batch, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + ) data["model"] = model_id_from_batch data["batch_id"] = get_batch_id_from_unified_batch_id(unified_batch_id) response = await llm_router.acancel_batch(**data) diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 57e42e79a42..5be64ca0847 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -179,6 +179,8 @@ def harness(): router = MagicMock(spec=Router) router.model_group_alias = {} router.get_model_access_groups = MagicMock(return_value={}) + router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id) + router.model_list = [] router.acreate_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -1165,6 +1167,8 @@ def retrieve_harness(): router = MagicMock(spec=Router) router.model_group_alias = {} router.get_model_access_groups = MagicMock(return_value={}) + router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id) + router.model_list = [] router.aretrieve_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -1622,6 +1626,8 @@ def list_harness(): router = MagicMock(spec=Router) router.model_group_alias = {} router.get_model_access_groups = MagicMock(return_value={}) + router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id) + router.model_list = [] router.alist_batches = AsyncMock(return_value=FakeListPage([])) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -2020,6 +2026,8 @@ def cancel_harness(): router = MagicMock(spec=Router) router.model_group_alias = {} router.get_model_access_groups = MagicMock(return_value={}) + router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id) + router.model_list = [] router.acancel_batch = AsyncMock(return_value=make_batch()) router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup) @@ -2816,3 +2824,53 @@ async def test_cancel__model_encoded_id_rejects_key_without_model_grant(cancel_h assert exc_info.value.code == "403" cancel_harness.creds_resolver.assert_not_called() cancel_harness.litellm_acancel.assert_not_called() + + +def _b64_unified_id(decoded: str) -> str: + return base64.urlsafe_b64encode(decoded.encode()).decode().rstrip("=") + + +UNIFIED_FILE_ID_FOR_GPT4O_MINI = _b64_unified_id( + "litellm_proxy:application/octet-stream;unified_id,c4843482-b176-4901-8292-7523fd0f2c6e;" + "target_model_names,gpt-4o-mini;llm_output_file_id,file-provider;llm_output_file_model_id,dep-1" +) +UNIFIED_BATCH_ID_FOR_GPT4O_MINI = _b64_unified_id(UNIFIED_BATCH_ID) + + +@pytest.mark.asyncio +async def test_create__unified_file_id_rejects_key_without_model_grant(harness): + """The model carried inside a unified file id is caller-controlled too, so it is checked against the key's grants.""" + set_body( + harness, + { + "input_file_id": UNIFIED_FILE_ID_FOR_GPT4O_MINI, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + with pytest.raises(ProxyException) as exc_info: + await call_create(harness, user=_key_restricted_to("vertex-model")) + + assert exc_info.value.code == "403" + harness.router_acreate.assert_not_called() + harness.litellm_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_retrieve__unified_batch_id_rejects_key_without_model_grant(retrieve_harness): + with pytest.raises(ProxyException) as exc_info: + await call_retrieve(retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to("vertex-model")) + + assert exc_info.value.code == "403" + retrieve_harness.router_aretrieve.assert_not_called() + retrieve_harness.creds_resolver.assert_not_called() + + +@pytest.mark.asyncio +async def test_cancel__unified_batch_id_rejects_key_without_model_grant(cancel_harness): + with pytest.raises(ProxyException) as exc_info: + await call_cancel(cancel_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to("vertex-model")) + + assert exc_info.value.code == "403" + cancel_harness.router_acancel.assert_not_called() diff --git a/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py index fe4903b547c..3161fe99e68 100644 --- a/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py +++ b/tests/test_litellm/proxy/test_batch_x_litellm_model_encoding.py @@ -11,6 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.openai_files_endpoints.common_utils import ( decode_model_from_file_id, get_batch_id_from_unified_batch_id, @@ -412,11 +413,7 @@ async def test_cancel_batch_with_unified_id_routes_with_decoded_model_and_batch_ mock_request.url.path = f"/v1/batches/{unified_batch_id}/cancel" mock_fastapi_response = MagicMock() mock_fastapi_response.headers = {} - mock_user_api_key_dict = MagicMock() - mock_user_api_key_dict.parent_otel_span = None - mock_user_api_key_dict.user_id = "test_user" - mock_user_api_key_dict.allowed_model_region = None - mock_user_api_key_dict.team_metadata = {} + mock_user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="test_user", team_metadata={}) with ( patch("litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing") as mock_processor_cls, From 95fdefa390af6586affb5ff825955b0ccb3bce17 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 10 Sep 2026 19:14:45 -0400 Subject: [PATCH 009/251] fix(logging): tolerate a missing api_base in pre_call for presigned batch retrieves Provider batch configs that build their own request URL (Mistral, Bedrock) hand pre_call api_base=None, and mask_api_base_credentials raised TypeError on it, so every such retrieve logged a non-blocking LoggingError and lost its pre-call logging. --- litellm/litellm_core_utils/litellm_logging.py | 4 +- .../test_litellm_logging.py | 63 ++++++++++--------- 2 files changed, 35 insertions(+), 32 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index cb9209be267..38e88493986 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1199,8 +1199,8 @@ class Logging(LiteLLMLoggingBaseClass): return {"error": f"Unable to parse raw request body. Got - {data}"} return data - def _get_masked_api_base(self, api_base: str) -> str: - return str(mask_api_base_credentials(api_base)) + def _get_masked_api_base(self, api_base: str | None) -> str: + return str(mask_api_base_credentials(api_base or "")) def _pre_call(self, input, api_key, model=None, additional_args={}): """ diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 6aa77745e3d..570f4339cd4 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -58,6 +58,16 @@ def test_get_masked_api_base(logging_obj): assert type(masked_api_base) == str +def test_pre_call_tolerates_missing_api_base(logging_obj): + """Presigned batch retrieves (Mistral, Bedrock) build their own URL and pass api_base=None + to pre_call; masking must not raise or the request's pre-call logging is silently lost.""" + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + + logging_obj.pre_call(input="", api_key="", additional_args={"api_base": None, "headers": {}}) + + assert logging_obj.model_call_details["litellm_params"]["api_base"] == "" + + def test_post_call_serializes_dict_with_datetime(logging_obj): import datetime @@ -3976,9 +3986,7 @@ def test_get_standard_logging_object_payload_carries_matched_access_groups(loggi "model": "gpt-4o", "messages": [], "litellm_params": { - "metadata": { - "user_api_key_matched_model_access_groups": ["premium-pool", "shared-pool"] - }, + "metadata": {"user_api_key_matched_model_access_groups": ["premium-pool", "shared-pool"]}, "proxy_server_request": {"body": {}}, }, }, @@ -4062,9 +4070,7 @@ def _model_router_response(selected_model: str, stamp: bool): from litellm.types.utils import ModelResponse response = ModelResponse(model=selected_model) - response._hidden_params = ( - {AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model} if stamp else {} - ) + response._hidden_params = {AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model} if stamp else {} return response @@ -4088,9 +4094,7 @@ def test_standard_logging_payload_uses_stamped_model_router_model(logging_obj): "messages": [], "litellm_params": {"metadata": {}}, }, - init_response_obj=_model_router_response( - "azure_ai/grok-4-1-fast-reasoning", stamp=True - ), + init_response_obj=_model_router_response("azure_ai/grok-4-1-fast-reasoning", stamp=True), start_time=now, end_time=now, logging_obj=logging_obj, @@ -4122,9 +4126,7 @@ def test_standard_logging_payload_keeps_requested_model_without_router_stamp( "messages": [], "litellm_params": {"metadata": {}}, }, - init_response_obj=_model_router_response( - "azure_ai/grok-4-1-fast-reasoning", stamp=False - ), + init_response_obj=_model_router_response("azure_ai/grok-4-1-fast-reasoning", stamp=False), start_time=now, end_time=now, logging_obj=logging_obj, @@ -5536,9 +5538,7 @@ class TestNonInferenceCallTypesAreNotBilled: init_response_obj=self._retrieved_response(), start_time=now, end_time=now, - logging_obj=self._logging_obj( - "aget_responses", litellm_metadata=self.BACKGROUND_POLL_METADATA - ), + logging_obj=self._logging_obj("aget_responses", litellm_metadata=self.BACKGROUND_POLL_METADATA), status="success", ) @@ -5784,9 +5784,7 @@ async def test_streaming_success_callbacks_survive_cost_calculation_failure(): releasing.async_log_success_event = AsyncMock() patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing]) - with patcher, patch.object( - logging_obj, "_response_cost_calculator", side_effect=ValueError("bad usage block") - ): + with patcher, patch.object(logging_obj, "_response_cost_calculator", side_effect=ValueError("bad usage block")): await logging_obj.async_success_handler(result=_assembled_stream_result()) assert logging_obj.model_call_details["response_cost"] is None @@ -5799,8 +5797,9 @@ async def test_streaming_success_callbacks_survive_standard_logging_payload_fail releasing.async_log_success_event = AsyncMock() patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing]) - with patcher, patch.object( - logging_obj, "_build_standard_logging_payload", side_effect=ValueError("incomplete stream") + with ( + patcher, + patch.object(logging_obj, "_build_standard_logging_payload", side_effect=ValueError("incomplete stream")), ): await logging_obj.async_success_handler(result=_assembled_stream_result()) @@ -6073,6 +6072,8 @@ def test_prompt_hooks_skip_prompt_managers_when_no_prompt_id(logging_obj, tmp_pa ) for hook in [cb for cb in litellm.callbacks if isinstance(cb, VectorStorePreCallHook)]: litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, hook) + + def test_newrelic_dispatch_prefers_otel_v2_when_flag_on(monkeypatch): """With LITELLM_OTEL_V2 on and operator credentials present, the "newrelic" callback builds the OTel v2 logger (per-team credential routing); with the @@ -6228,7 +6229,9 @@ def test_get_error_information_skips_traceback_for_budget_rejection_with_provide from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup assert litellm.log_client_error_tracebacks is False - over_budget = _raise_and_catch(litellm.BudgetExceededError(current_cost=0.01, max_budget=0.0, llm_provider="anthropic")) + over_budget = _raise_and_catch( + litellm.BudgetExceededError(current_cost=0.01, max_budget=0.0, llm_provider="anthropic") + ) result = StandardLoggingPayloadSetup.get_error_information(over_budget) assert result["error_code"] == "429" assert result["llm_provider"] == "anthropic" @@ -6745,9 +6748,7 @@ def test_passthrough_embeddings_result_swapped_for_callbacks(): ], "model": "EmbeddingsGigaR", }, - request=httpx.Request( - "POST", "https://gigachat.devices.sberbank.ru/api/v1/embeddings" - ), + request=httpx.Request("POST", "https://gigachat.devices.sberbank.ru/api/v1/embeddings"), ) _, _, swapped_result = logging_obj._success_handler_helper_fn( @@ -6766,12 +6767,14 @@ def test_get_status_fields_ranks_guardrail_flagged_between_success_and_intervene request-level guardrail_status but never mask an intervention.""" flagged = {"guardrail_status": "guardrail_flagged"} - assert _get_status_fields( - "success", [{"guardrail_status": "success"}, flagged], None - )["guardrail_status"] == "guardrail_flagged" - assert _get_status_fields( - "success", [flagged, {"guardrail_status": "guardrail_intervened"}], None - )["guardrail_status"] == "guardrail_intervened" + assert ( + _get_status_fields("success", [{"guardrail_status": "success"}, flagged], None)["guardrail_status"] + == "guardrail_flagged" + ) + assert ( + _get_status_fields("success", [flagged, {"guardrail_status": "guardrail_intervened"}], None)["guardrail_status"] + == "guardrail_intervened" + ) def test_get_error_information_redacts_provider_key_from_upstream_url(): From e6bc4e47c7a63e8540a856f1b2f66710a4b47142 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 10 Sep 2026 19:46:55 -0400 Subject: [PATCH 010/251] fix(mistral): reject file purposes Mistral lacks instead of mapping them to batch The proxy runs batch-file validation and guardrails only for purpose=batch, so a purpose such as assistants that was silently rewritten to batch on the way to Mistral let an upload skip both. Only batch, fine-tune and ocr pass through now; anything else is a 400. --- litellm/llms/mistral/files/transformation.py | 9 ++++-- .../test_mistral_files_transformation.py | 29 ++++++++++++++----- 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/litellm/llms/mistral/files/transformation.py b/litellm/llms/mistral/files/transformation.py index 6d58311813c..bf1ef7cb69f 100644 --- a/litellm/llms/mistral/files/transformation.py +++ b/litellm/llms/mistral/files/transformation.py @@ -89,11 +89,14 @@ def _to_openai_purpose(purpose: MistralFilePurpose) -> OpenAIFilesPurpose: def _to_mistral_purpose(purpose: str) -> MistralFilePurpose: + """Only Mistral's own purposes pass through. Silently mapping anything else to ``batch`` + would let an upload skip the proxy's batch-file validation and guardrails, which only + run when the caller says ``purpose=batch``.""" match purpose: - case "fine-tune" | "ocr": + case "batch" | "fine-tune" | "ocr": return purpose case _: - return "batch" + raise ValueError(f"Mistral does not support purpose={purpose!r}. Use one of: batch, fine-tune, ocr") def _api_base_from(litellm_params: Mapping[str, object]) -> str: @@ -166,7 +169,7 @@ class MistralFilesConfig(BaseFilesConfig): content_type: Final = extracted.get("content_type") or "application/octet-stream" upload: Final = MistralMultipartUpload( file=(filename, extracted["content"], content_type), - purpose=(None, _to_mistral_purpose(create_file_data.get("purpose", "batch"))), + purpose=(None, _to_mistral_purpose(create_file_data.get("purpose") or "batch")), ) return dict(upload) # mutable-ok: BaseFilesConfig signature diff --git a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py index f62645be7ee..b81740c0429 100644 --- a/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py +++ b/tests/test_litellm/llms/mistral/files/test_mistral_files_transformation.py @@ -91,18 +91,28 @@ def test_upload_request_is_multipart_with_batch_purpose(config): } -@pytest.mark.parametrize( - "openai_purpose,mistral_purpose", - [("batch", "batch"), ("fine-tune", "fine-tune"), ("ocr", "ocr"), ("assistants", "batch"), ("user_data", "batch")], -) -def test_upload_request_maps_purpose_onto_mistral_enum(config, openai_purpose, mistral_purpose): +@pytest.mark.parametrize("purpose", ["batch", "fine-tune", "ocr"]) +def test_upload_request_passes_mistral_purposes_through(config, purpose): body = config.transform_create_file_request( model="", - create_file_data=CreateFileRequest(file=("f.bin", b"x"), purpose=openai_purpose), + create_file_data=CreateFileRequest(file=("f.bin", b"x"), purpose=purpose), optional_params={}, litellm_params={}, ) - assert body["purpose"] == (None, mistral_purpose) + assert body["purpose"] == (None, purpose) + + +@pytest.mark.parametrize("purpose", ["assistants", "user_data", "vision", "evals"]) +def test_upload_request_rejects_purposes_mistral_lacks(config, purpose): + """Regression: these used to be silently rewritten to ``batch``, so an upload that skipped the + proxy's batch-only validation and guardrails still landed on Mistral as a batch input file.""" + with pytest.raises(ValueError, match=f"purpose={purpose!r}"): + config.transform_create_file_request( + model="", + create_file_data=CreateFileRequest(file=("f.bin", b"x"), purpose=purpose), + optional_params={}, + litellm_params={}, + ) def test_upload_request_requires_file(config): @@ -181,6 +191,11 @@ def test_list_request_filters_by_mapped_purpose(config): assert no_params == {} +def test_list_request_rejects_purposes_mistral_lacks(config): + with pytest.raises(ValueError, match="purpose='assistants'"): + config.transform_list_files_request(purpose="assistants", optional_params={}, litellm_params={}) + + def test_list_response(config): out = config.transform_list_files_response( raw_response=_response( From 0f3c4ccfbba894c0a66d92ed44ca18f312bf75f0 Mon Sep 17 00:00:00 2001 From: jesus Date: Fri, 11 Sep 2026 00:00:17 +0000 Subject: [PATCH 011/251] fix(auth): inherit organization_alias from the org for JWT and team-linked keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 42 ++++++- .../proxy/auth/test_user_api_key_auth.py | 104 +++++++++++++++++- 2 files changed, 143 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 20ab9904f46..110c524ecdf 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -54,6 +54,7 @@ from litellm.proxy.auth.auth_checks import ( get_end_user_object, get_jwt_key_mapping_object, get_object_permission, + get_org_object, get_project_object, get_team_object, get_user_object, @@ -2398,6 +2399,37 @@ def _token_can_vouch_for_team(valid_token: UserAPIKeyAuth, lookup_error: BaseExc return PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() +async def _inherit_org_identity( + user_api_key_auth_obj: UserAPIKeyAuth, + team_object: LiteLLM_TeamTableCachedObj | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging | None, +) -> None: + if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: + user_api_key_auth_obj.org_id = team_object.organization_id + if ( + user_api_key_auth_obj.org_id is None + or user_api_key_auth_obj.organization_alias is not None + or prisma_client is None + ): + return + try: + org_object: Final = await get_org_object( + org_id=user_api_key_auth_obj.org_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception: + verbose_proxy_logger.debug("org alias lookup failed for org_id=%s", user_api_key_auth_obj.org_id, exc_info=True) + return + if org_object is not None: + user_api_key_auth_obj.organization_alias = org_object.organization_alias + + @tracer.wrap() async def _run_centralized_common_checks( user_api_key_auth_obj: UserAPIKeyAuth, @@ -2622,8 +2654,14 @@ async def _run_centralized_common_checks( ) global_proxy_spend: float | None = None if isinstance(global_spend_result, BaseException) else global_spend_result - if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: - user_api_key_auth_obj.org_id = team_object.organization_id + await _inherit_org_identity( + user_api_key_auth_obj=user_api_key_auth_obj, + team_object=cast(LiteLLM_TeamTableCachedObj | None, team_object), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) # common_checks identifies admin via user_object, not the token # (non_proxy_admin_allowed_routes_check). JWT admin shortcut and diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 6cce6d0316b..03efbfa7185 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -23,6 +23,7 @@ from litellm.proxy._types import ( LiteLLM_JWTAuth, LiteLLM_BudgetTable, LiteLLM_EndUserTable, + LiteLLM_OrganizationTable, LiteLLM_UserTable, LitellmUserRoles, ProxyErrorTypes, @@ -31,7 +32,7 @@ from litellm.proxy._types import ( JWTRoutingOverride, ) from litellm.proxy.auth.handle_jwt import JWTHandler -from litellm.proxy.auth.auth_checks import get_key_object, _cache_key_object +from litellm.proxy.auth.auth_checks import OrganizationNotFoundError, get_key_object, _cache_key_object from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( _check_key_model_budget_with_fallback, @@ -5293,6 +5294,107 @@ async def test_centralized_common_checks_backfills_org_id_from_team(key_org_id, setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_org_id,team_id,team_org_id,existing_alias,lookup_mode,expected_org_id,expected_alias", + [ + (None, "t1", "org-from-team", None, "success", "org-from-team", "acme-org"), + ("org-jwt", None, None, None, "success", "org-jwt", "acme-org"), + ("org-pinned", None, None, "preset", "success", "org-pinned", "preset"), + ("org-missing", None, None, None, "missing", "org-missing", None), + ], +) +async def test_centralized_common_checks_inherits_org_alias( + key_org_id, + team_id, + team_org_id, + existing_alias, + lookup_mode, + expected_org_id, + expected_alias, +): + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + + token = UserAPIKeyAuth( + api_key="sk-test", + user_id="u", + team_id=team_id, + org_id=key_org_id, + organization_alias=existing_alias, + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + fetched_team = ( + LiteLLM_TeamTableCachedObj(team_id="t1", organization_id=team_org_id) if team_id is not None else None + ) + organization = LiteLLM_OrganizationTable( + organization_id=expected_org_id, + organization_alias="acme-org", + budget_id="budget-id", + models=[], + created_by="test", + updated_by="test", + ) + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + attrs["prisma_client"] = MagicMock() + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + identity_seen_by_common_checks = [] + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + return_value=fetched_team, + ) as mock_get_team_object, + patch( + "litellm.proxy.auth.user_api_key_auth.get_org_object", + new_callable=AsyncMock, + return_value=organization, + ) as mock_get_org_object, + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + side_effect=lambda **kw: identity_seen_by_common_checks.append( + (kw["valid_token"].org_id, kw["valid_token"].organization_alias) + ), + ) as mock_checks, + ): + if lookup_mode == "missing": + mock_get_org_object.side_effect = OrganizationNotFoundError("x") + + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + + mock_checks.assert_awaited_once() + assert token.org_id == expected_org_id + assert token.organization_alias == expected_alias + assert identity_seen_by_common_checks == [(expected_org_id, expected_alias)] + if team_id is None: + mock_get_team_object.assert_not_awaited() + else: + mock_get_team_object.assert_awaited_once() + if existing_alias is not None: + mock_get_org_object.assert_not_awaited() + else: + mock_get_org_object.assert_awaited_once() + assert mock_get_org_object.await_args.kwargs["org_id"] == expected_org_id + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + @pytest.mark.asyncio async def test_cli_session_token_org_backfilled_from_team(monkeypatch): """LIT-4688 root cause: CLI session tokens (from /sso/cli/poll) are minted From 785c6cffc4826ef44c73a797981e28a294a45668 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Thu, 13 Aug 2026 23:34:40 -0400 Subject: [PATCH 012/251] fix(cost): carry image and video input tokens through the Responses usage bridge Realtime cost is computed from *_tokens_details after the usage round-trips through the Responses shape, and the input half of that shape carried audio only, so image and video prompt tokens stopped being billable as themselves. Vertex splits prompt tokens by modality, so a session sending camera frames arrives with image_tokens set. Those were folded into text_tokens and lost their attribution. The amount happens not to move today, because the calculator falls back to input_cost_per_token when no per-modality rate is set, but the tokens have to survive before any such rate can ever apply. InputTokensDetails now declares image_tokens and video_tokens instead of leaning on pydantic extras, the repeated per-field copying is a loop over the modality names so adding a modality no longer adds a branch, and the read-back in ResponseAPILoggingUtils picks up video_tokens, which PromptTokensDetailsWrapper already declared. The output half of the original change is dropped: 449c091391 landed the same OutputTokensDetails.audio_tokens fix upstream, with its own coverage in test_gemini_realtime_transformation.py, and it always sets output_tokens_details rather than only when non-empty. That structure is kept as upstream wrote it. --- .../transformation.py | 2 ++ litellm/responses/utils.py | 1 + litellm/types/llms/openai.py | 2 ++ .../test_litellm_completion_responses.py | 33 +++++++++++++++++++ 4 files changed, 38 insertions(+) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 01fb6cb483d..ffd7ce491b1 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2851,6 +2851,8 @@ class LiteLLMCompletionResponsesConfig: cached_tokens=prompt_details.cached_tokens if prompt_details.cached_tokens is not None else 0, text_tokens=prompt_details.text_tokens, audio_tokens=prompt_details.audio_tokens, + image_tokens=prompt_details.image_tokens, + video_tokens=prompt_details.video_tokens, cached_tokens_details=( cached_tokens_details if isinstance(cached_tokens_details, CachedTokensDetails) else None ), diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 41a3ded7022..f50c17aff85 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1182,6 +1182,7 @@ class ResponseAPILoggingUtils: cached_tokens_details=getattr( response_api_usage.input_tokens_details, "cached_tokens_details", None ), + video_tokens=getattr(response_api_usage.input_tokens_details, "video_tokens", None), cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None), web_search_requests=getattr(response_api_usage.input_tokens_details, "web_search_requests", None), google_maps_grounding_requests=getattr( diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index e3eac9b9205..98548705979 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1291,7 +1291,9 @@ class InputTokensDetails(BaseLiteLLMOpenAIResponseObject): audio_tokens: int | None = None cached_tokens: int = 0 cached_tokens_details: CachedTokensDetails | None = None + image_tokens: int | None = None text_tokens: int | None = None + video_tokens: int | None = None model_config = {"extra": "allow"} diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 0ed101952be..2f9f7adfcf1 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -2885,6 +2885,39 @@ class TestUsageTransformation: assert getattr(response_usage.input_tokens_details, "cache_write_tokens", None) == 800 assert response_usage.input_tokens_details.model_dump()["cache_write_tokens"] == 800 + def test_transform_usage_preserves_input_modality_tokens(self): + """Regression: the bridge dropped image and video input tokens. + + Vertex reports prompt tokens split by modality, so a Live session that sends + camera frames arrives with image_tokens set. InputTokensDetails declared only + audio/cached/text, so those tokens were folded into text and lost their + attribution, and any per-modality rate could never apply to them. + """ + usage = Usage( + prompt_tokens=300, + completion_tokens=10, + total_tokens=310, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=20, audio_tokens=80, image_tokens=150, video_tokens=50, cached_tokens=0 + ), + completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=10), + ) + + response_usage = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage( + chat_completion_response=usage + ) + details = response_usage.input_tokens_details + assert details is not None + assert getattr(details, "image_tokens", None) == 150 + assert getattr(details, "video_tokens", None) == 50 + assert getattr(details, "audio_tokens", None) == 80 + + from litellm.responses.utils import ResponseAPILoggingUtils + + back = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response_usage.model_dump()) + assert back.prompt_tokens_details.image_tokens == 150 + assert back.prompt_tokens_details.video_tokens == 50 + def test_transform_usage_with_reasoning_tokens_gemini(self): """Test that reasoning_tokens from Gemini are properly transformed to output_tokens_details""" # Setup: Simulate Gemini usage with thoughtsTokenCount From 76488beaf8d1a44e0f07b6a4a66c06b9b4390222 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 05:55:25 -0700 Subject: [PATCH 013/251] fix(utils): reject an untranslatable tool_choice with a 400 instead of a 500 --- litellm/main.py | 2 +- litellm/utils.py | 16 +++- tests/litellm_utils_tests/test_utils.py | 6 +- .../test_validate_tool_choice.py | 74 ++++++++++--------- .../test_litellm_completion_responses.py | 15 ++++ tests/test_litellm/test_main.py | 14 ++++ 6 files changed, 85 insertions(+), 42 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index f6f4ec1bf63..9eea6779abf 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5102,7 +5102,7 @@ def completion( messages = validate_and_fix_openai_messages(messages=messages) tools = validate_and_fix_openai_tools(tools=tools) # validate tool_choice - tool_choice = validate_chat_completion_tool_choice(tool_choice=tool_choice) + tool_choice = validate_chat_completion_tool_choice(tool_choice=tool_choice, model=model) # validate optional params stop = validate_openai_optional_params(stop=stop) thinking = validate_and_fix_thinking_param(thinking=thinking) diff --git a/litellm/utils.py b/litellm/utils.py index 18df5e2abf7..17088f0475b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8038,6 +8038,7 @@ def validate_chat_completion_user_messages(messages: list[AllMessageValues]): def validate_chat_completion_tool_choice( tool_choice: dict | str | None, + model: str, ) -> dict | str | None: """ Confirm the tool choice is passed in the OpenAI format. @@ -8053,12 +8054,19 @@ def validate_chat_completion_tool_choice( # Standard OpenAI format: {"type": "function", "function": {...}} if tool_choice.get("type") is None or tool_choice.get("function") is None: - raise Exception( - f"Invalid tool choice, tool_choice={tool_choice}. Please ensure tool_choice follows the OpenAI spec" + raise BadRequestError( + message=f"Invalid tool choice, tool_choice={tool_choice}. Please ensure tool_choice follows the OpenAI spec", + model=model, + llm_provider="", ) return tool_choice - raise Exception( - f"Invalid tool choice, tool_choice={tool_choice}. Got={type(tool_choice)}. Expecting str, or dict. Please ensure tool_choice follows the OpenAI tool_choice spec" + raise BadRequestError( + message=( + f"Invalid tool choice, tool_choice={tool_choice}. Got={type(tool_choice)}. Expecting str, or dict. " + "Please ensure tool_choice follows the OpenAI tool_choice spec" + ), + model=model, + llm_provider="", ) diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index 0ccfae55290..11d089719c6 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -1334,10 +1334,10 @@ def test_validate_chat_completion_tool_choice(tool_choice, expected_bool): from litellm.utils import validate_chat_completion_tool_choice if expected_bool: - validate_chat_completion_tool_choice(tool_choice=tool_choice) + validate_chat_completion_tool_choice(tool_choice=tool_choice, model="gpt-5.6-sol") else: - with pytest.raises(Exception, match="Invalid tool choice"): - validate_chat_completion_tool_choice(tool_choice=tool_choice) + with pytest.raises(litellm.BadRequestError, match="Invalid tool choice"): + validate_chat_completion_tool_choice(tool_choice=tool_choice, model="gpt-5.6-sol") def test_models_by_provider(): diff --git a/tests/litellm_utils_tests/test_validate_tool_choice.py b/tests/litellm_utils_tests/test_validate_tool_choice.py index b8246fe0deb..b4272af7b90 100644 --- a/tests/litellm_utils_tests/test_validate_tool_choice.py +++ b/tests/litellm_utils_tests/test_validate_tool_choice.py @@ -1,60 +1,66 @@ +import re +from typing import Final + import pytest - +import litellm from litellm.utils import validate_chat_completion_tool_choice +MODEL: Final = "anthropic/claude-haiku-4-5" + def test_validate_tool_choice_none(): """Test that None is returned as-is.""" - result = validate_chat_completion_tool_choice(None) + result = validate_chat_completion_tool_choice(None, model=MODEL) assert result is None def test_validate_tool_choice_string(): """Test that string values are returned as-is.""" - assert validate_chat_completion_tool_choice("auto") == "auto" - assert validate_chat_completion_tool_choice("none") == "none" - assert validate_chat_completion_tool_choice("required") == "required" + assert validate_chat_completion_tool_choice("auto", model=MODEL) == "auto" + assert validate_chat_completion_tool_choice("none", model=MODEL) == "none" + assert validate_chat_completion_tool_choice("required", model=MODEL) == "required" def test_validate_tool_choice_standard_dict(): """Test standard OpenAI format with function.""" tool_choice = {"type": "function", "function": {"name": "my_function"}} - result = validate_chat_completion_tool_choice(tool_choice) + result = validate_chat_completion_tool_choice(tool_choice, model=MODEL) assert result == tool_choice def test_validate_tool_choice_cursor_format(): """Cursor IDE format {"type": "auto"} is unwrapped to the bare string.""" - assert validate_chat_completion_tool_choice({"type": "auto"}) == "auto" - assert validate_chat_completion_tool_choice({"type": "none"}) == "none" - assert validate_chat_completion_tool_choice({"type": "required"}) == "required" + assert validate_chat_completion_tool_choice({"type": "auto"}, model=MODEL) == "auto" + assert validate_chat_completion_tool_choice({"type": "none"}, model=MODEL) == "none" + assert validate_chat_completion_tool_choice({"type": "required"}, model=MODEL) == "required" -def test_validate_tool_choice_invalid_dict(): - """Test that invalid dict formats raise exceptions.""" - # Missing both type and function - with pytest.raises(Exception, match='Invalid tool choice, tool_choice=\\{\\}\\. Please ensure') as exc_info: - validate_chat_completion_tool_choice({}) - assert "Invalid tool choice" in str(exc_info.value) - - # Invalid type value - with pytest.raises(Exception, match="Invalid tool choice, tool_choice=\\{'type': 'invalid'\\}\\.") as exc_info: - validate_chat_completion_tool_choice({"type": "invalid"}) - assert "Invalid tool choice" in str(exc_info.value) - - # Has type but missing function when type is "function" - with pytest.raises(Exception, match="Invalid tool choice, tool_choice=\\{'type': 'function'\\}\\.") as exc_info: - validate_chat_completion_tool_choice({"type": "function"}) - assert "Invalid tool choice" in str(exc_info.value) +@pytest.mark.parametrize( + "tool_choice", + [ + {}, + {"type": "invalid"}, + {"type": "function"}, + {"name": "lookup_fruit"}, + {"type": "file_search"}, + ], +) +def test_validate_tool_choice_invalid_dict_is_a_400(tool_choice): + """A dict shape chat completions cannot carry is the caller's mistake: a 400 that names the field, never a 500.""" + with pytest.raises( + litellm.BadRequestError, match=f"Invalid tool choice, tool_choice={re.escape(str(tool_choice))}\\. Please ensure" + ) as exc_info: + validate_chat_completion_tool_choice(tool_choice, model=MODEL) + assert exc_info.value.status_code == 400 + assert exc_info.value.model == MODEL -def test_validate_tool_choice_invalid_type(): - """Test that invalid types raise exceptions.""" - with pytest.raises(Exception, match="\\. Expecting str, or dict\\. Please ensure") as exc_info: - validate_chat_completion_tool_choice(123) - assert "Got=" in str(exc_info.value) - - with pytest.raises(Exception, match="Invalid tool choice, tool_choice=\\[\\]\\. Got=\\.") as exc_info: - validate_chat_completion_tool_choice([]) - assert "Got=" in str(exc_info.value) +@pytest.mark.parametrize("tool_choice", [123, []]) +def test_validate_tool_choice_invalid_type_is_a_400(tool_choice): + """A non-str, non-dict tool_choice is rejected as a 400 that names the type it got.""" + with pytest.raises( + litellm.BadRequestError, match=f"Got={re.escape(str(type(tool_choice)))}\\. Expecting str, or dict\\." + ) as exc_info: + validate_chat_completion_tool_choice(tool_choice, model=MODEL) + assert exc_info.value.status_code == 400 diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 0ed101952be..3950fd549ef 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -4906,3 +4906,18 @@ class TestStreamingSnapshotItemIds: reasoning_items = _bridged_output_items(completed_event.response, "reasoning") assert len(reasoning_items) == 1 assert reasoning_items[0].id == streamed_event.item_id + + +@pytest.mark.parametrize("stream", [True, False]) +async def test_bridge_rejects_untranslatable_tool_choice_with_a_400(stream: bool): + with pytest.raises(litellm.BadRequestError) as exc_info: + await litellm.aresponses( + model="anthropic/claude-haiku-4-5", + input="Which fruit is red?", + tools=[{"type": "function", "name": "lookup_fruit", "parameters": {"type": "object"}}], + tool_choice={"type": "file_search"}, + stream=stream, + api_key="sk-unused", + ) + assert exc_info.value.status_code == 400 + assert "tool_choice={'type': 'file_search'}" in str(exc_info.value) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 3dccb2b35bf..81ad161772e 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -3850,3 +3850,17 @@ def test_bridged_responses_with_openai_http_handler_keeps_forwarded_headers_out_ assert "extra_headers" not in body assert body["model"] == "gpt-5.4" assert {k: request.headers[k] for k in FORWARDED_CLIENT_HEADERS} == FORWARDED_CLIENT_HEADERS + + +@pytest.mark.parametrize("tool_choice", [{"type": "bogus"}, {"name": "lookup_fruit"}, {"type": "file_search"}]) +def test_completion_rejects_untranslatable_tool_choice_with_a_400(tool_choice): + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.completion( + model="anthropic/claude-haiku-4-5", + messages=[{"role": "user", "content": "Which fruit is red?"}], + tools=[{"type": "function", "function": {"name": "lookup_fruit", "parameters": {"type": "object"}}}], + tool_choice=tool_choice, + api_key="sk-unused", + ) + assert exc_info.value.status_code == 400 + assert f"tool_choice={tool_choice}" in str(exc_info.value) From 2bbf34c6520ef3266a62121510470868f73e499e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 15 Sep 2026 07:00:33 -0700 Subject: [PATCH 014/251] fix(utils): keep the tool_choice validator's model argument optional --- litellm/utils.py | 2 +- tests/litellm_utils_tests/test_validate_tool_choice.py | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 17088f0475b..c2dbbda2b68 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8038,7 +8038,7 @@ def validate_chat_completion_user_messages(messages: list[AllMessageValues]): def validate_chat_completion_tool_choice( tool_choice: dict | str | None, - model: str, + model: str = "", ) -> dict | str | None: """ Confirm the tool choice is passed in the OpenAI format. diff --git a/tests/litellm_utils_tests/test_validate_tool_choice.py b/tests/litellm_utils_tests/test_validate_tool_choice.py index b4272af7b90..a9dacf9fa15 100644 --- a/tests/litellm_utils_tests/test_validate_tool_choice.py +++ b/tests/litellm_utils_tests/test_validate_tool_choice.py @@ -64,3 +64,11 @@ def test_validate_tool_choice_invalid_type_is_a_400(tool_choice): ) as exc_info: validate_chat_completion_tool_choice(tool_choice, model=MODEL) assert exc_info.value.status_code == 400 + + +def test_validate_tool_choice_without_model_is_still_a_400(): + """Callers that predate the model argument keep getting a 400, with an empty model on the error.""" + with pytest.raises(litellm.BadRequestError, match="Invalid tool choice") as exc_info: + validate_chat_completion_tool_choice({"type": "bogus"}) + assert exc_info.value.status_code == 400 + assert exc_info.value.model == "" From 61ac4f57394e022382459235baa598ac89ce00d1 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 23:03:56 +0000 Subject: [PATCH 015/251] test(e2e): add scripted-provider cost calculation suite Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/CLAUDE.md | 3 +- tests/e2e/conftest.py | 9 +- tests/e2e/cost_calculation/conftest.py | 139 ++++ tests/e2e/cost_calculation/cost_matrix.py | 458 +++++++++++++ tests/e2e/cost_calculation/scripted_client.py | 70 ++ .../e2e/cost_calculation/scripted_provider.py | 631 ++++++++++++++++++ .../test_token_pricing_e2e.py | 115 ++++ .../cost_calculation/test_wire_formats_e2e.py | 186 ++++++ tests/e2e/cost_map.json | 352 ++++++++++ .../coverage_registry/quota_management.yaml | 2 + tests/e2e/e2e_config.py | 16 + tests/e2e/pytest.ini | 1 + 12 files changed, 1980 insertions(+), 2 deletions(-) create mode 100644 tests/e2e/cost_calculation/conftest.py create mode 100644 tests/e2e/cost_calculation/cost_matrix.py create mode 100644 tests/e2e/cost_calculation/scripted_client.py create mode 100644 tests/e2e/cost_calculation/scripted_provider.py create mode 100644 tests/e2e/cost_calculation/test_token_pricing_e2e.py create mode 100644 tests/e2e/cost_calculation/test_wire_formats_e2e.py create mode 100644 tests/e2e/cost_map.json diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 0541ce25d4b..b6c3840f626 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -21,6 +21,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see); needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in), and markerless harness unit tests for the locust, process-usage, and session-anomaly aggregation logic - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests +- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; asserts literal rate arithmetic on scripted usage across every provider wire and pricing component, deselected unless `E2E_COST_MAP_STACK` is set - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` @@ -221,7 +222,7 @@ other... ``` ## Hard Rules -- no unit tests of any kind under `tests/e2e`. a product feature is proven end to end against a live proxy, never with a unit test, and the harness itself is not unit-tested here either. no monkeypatching or mock tests. if a contributor asks you to write an end to end test, do NOT stage a unit test with it; if you find a product gap, call it out in the PR description +- no unit tests of any kind under `tests/e2e`. a product feature is proven end to end against a live proxy, never with a unit test, and the harness itself is not unit-tested here either. no monkeypatching or mock tests; the one carve-out is a scripted upstream served through a real HTTP sidecar (the cost_calculation suite's scripted provider), allowed because provider-response-shape coverage needs a controlled usage payload and every hop from the proxy's upstream call to the spend row still executes for real. if a contributor asks you to write an end to end test, do NOT stage a unit test with it; if you find a product gap, call it out in the PR description - use model management endpoints to create new models for a test. this could be in a conftest / inline for each test. ask the user what they want. diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index b1a75d5f862..7ab41b8ff68 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -22,9 +22,9 @@ from typing import Final import pytest import requests - from e2e_config import ( CONTROL_PLANE_BASE_URL, + COST_MAP_OPT_IN_ENV, FIXTURE_DIR, FIXTURE_MODE_RAW, MANAGED_FILES_OPT_IN_ENV, @@ -53,6 +53,7 @@ OPT_IN_MARKERS: Final = MappingProxyType( "managed_files": MANAGED_FILES_OPT_IN_ENV, "prompt_caching_stack": PROMPT_CACHING_OPT_IN_ENV, "redis_chaos": REDIS_CHAOS_OPT_IN_ENV, + "cost_map_stack": COST_MAP_OPT_IN_ENV, } ) @@ -120,6 +121,12 @@ def pytest_configure(config: pytest.Config) -> None: "redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from " "gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set", ) + config.addinivalue_line( + "markers", + "cost_map_stack: needs a proxy whose whole cost map is tests/e2e/cost_map.json " + "(LITELLM_MODEL_COST_MAP_URL) plus a scripted-provider sidecar; deselected unless " + "E2E_COST_MAP_STACK is set", + ) def pytest_sessionstart(session: pytest.Session) -> None: diff --git a/tests/e2e/cost_calculation/conftest.py b/tests/e2e/cost_calculation/conftest.py new file mode 100644 index 00000000000..1bba3d50e1d --- /dev/null +++ b/tests/e2e/cost_calculation/conftest.py @@ -0,0 +1,139 @@ +"""Cost-calculation suite fixtures. + +Runs against a dedicated proxy whose whole model cost map is the test-owned +``tests/e2e/cost_map.json`` (LITELLM_MODEL_COST_MAP_URL), so every deployment +bills at rates the test asserts literal arithmetic on. Provider calls are +answered by the scripted-provider sidecar (``scripted_provider.py``), registered +per scenario over its control API. + +Deselected unless E2E_COST_MAP_STACK is set (marker `cost_map_stack`). +""" + +from __future__ import annotations + +import importlib.util +import sys +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType +from typing import Final, Protocol, cast + +import pytest + +from cost_matrix import Case, FrontierModel +from e2e_config import COST_MAP_PROXY_URL +from lifecycle import ResourceManager +from models import LiteLLMParamsBody, ModelInfoBody, ModelNewBody +from proxy_client import ProxyClient, build_proxy_client +from scripted_client import ScenarioHandle, delete_scenario, register_scenario +from scripted_provider import Scenario + + +def _load_cost_rows() -> ModuleType: + """Load quota_management/spend_tracking/cost_rows.py by path (the e2e tree + has no package layout), the same trick the mcp suite uses for + logging/datadog_reader.py.""" + path = ( + Path(__file__).resolve().parent.parent + / "quota_management" + / "spend_tracking" + / "cost_rows.py" + ) + name = "e2e_spend_tracking_cost_rows" + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +class SpendCostBreakdown(Protocol): + input_cost: float | None + output_cost: float | None + cache_read_cost: float | None + cache_creation_cost: float | None + reasoning_cost: float | None + tool_usage_cost: float | None + total_cost: float | None + service_tier: str | None + + def model_dump(self) -> dict[str, object]: ... + + +class SpendRowMetadata(Protocol): + cost_breakdown: SpendCostBreakdown | None + + +class SpendCostRow(Protocol): + """The slice of spend_tracking.cost_rows.CostRow this suite reads.""" + + spend: float | None + prompt_tokens: int | None + completion_tokens: int | None + metadata: SpendRowMetadata | None + + @property + def breakdown(self) -> SpendCostBreakdown: ... + + +class CostRowsModule(Protocol): + """cost_rows.py loaded by path has no importable name for basedpyright, so + its surface is declared here and reached through a single cast.""" + + approx_equal: Callable[[float, float], bool] + assert_total_is_sum_of_components: Callable[[SpendCostRow], None] + poll_cost_row_where: Callable[ + [ProxyClient, str, Callable[[SpendCostRow], bool]], SpendCostRow | None + ] + + +cost_rows: Final[CostRowsModule] = cast(CostRowsModule, _load_cost_rows()) + + +@dataclass(frozen=True, slots=True) +class CostCalcClient: + """The suite's client: a ProxyClient pointed at the cost-map proxy pod.""" + + proxy: ProxyClient + + +@pytest.fixture(scope="session") +def client() -> CostCalcClient: + proxy = build_proxy_client( + base_url=COST_MAP_PROXY_URL, + control_plane_base_url=COST_MAP_PROXY_URL, + replica_urls=(COST_MAP_PROXY_URL,), + ) + return CostCalcClient(proxy=proxy) + + +def register_scenario_deployment( + client: CostCalcClient, + resources: ResourceManager, + model: FrontierModel, + case: Case, + marker: str, +) -> tuple[str, ScenarioHandle]: + """Register the case's scenario on the sidecar plus a deployment pointed at + it; both are torn down by ``resources``. Returns the callable model_name.""" + scenario: Scenario = case.scenario( + scenario_id=f"sc-{marker}", model=model, text=f"scripted answer {marker}" + ) + handle = register_scenario(scenario) + resources.defer(lambda: delete_scenario(handle)) + model_name = f"{model.model_name}-{marker}" + model_id = client.proxy.register_model( + ModelNewBody( + model_name=model_name, + litellm_params=LiteLLMParamsBody( + model=model.litellm_model, + api_key="sk-scripted-provider", + api_base=handle.api_base(), + ), + model_info=ModelInfoBody(), + ) + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return model_name, handle diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py new file mode 100644 index 00000000000..bc466d7d823 --- /dev/null +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -0,0 +1,458 @@ +"""The cost-calculation matrix: frontier model set, the pricing-component cases +each model runs, and the expected-cost arithmetic. + +Rates come from ``tests/e2e/cost_map.json``, which the proxy under test loads as +its ENTIRE model cost map (LITELLM_MODEL_COST_MAP_URL), so an entry's rates are +exactly what the proxy bills and nothing in the suite depends on the bundled +map. Each model's rates are a distinct multiple of a shared base set, so a +component billed at the wrong model's rate (or the wrong case's rate) can never +coincidentally match. + +Case applicability is pricing-field-gated AND wire-gated: a case runs for a +model only when the entry carries the rate the case exercises and the wire can +report the token kind that rate prices. When the wire cannot report a kind +(e.g. Anthropic has no reasoning-token field, Responses reports no cache +creation), the case is absent from the matrix rather than silently zero. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Final, Literal + +from pydantic import BaseModel, ConfigDict, TypeAdapter + +from scripted_provider import Scenario, ScriptedOutput, ScriptedUsage, Wire + +COST_MAP_PATH: Final = Path(__file__).resolve().parent.parent / "cost_map.json" + + +class SearchContextCostPerQuery(BaseModel): + model_config = ConfigDict(frozen=True) + + search_context_size_low: float | None = None + search_context_size_medium: float | None = None + search_context_size_high: float | None = None + + +class CostMapEntry(BaseModel): + """The pricing fields of a cost-map entry the matrix reads. Shaped like a + ``model_prices_and_context_window.json`` entry; unmodelled keys are ignored.""" + + model_config = ConfigDict(frozen=True, extra="ignore") + + litellm_provider: str + mode: str + input_cost_per_token: float | None = None + output_cost_per_token: float | None = None + cache_read_input_token_cost: float | None = None + cache_creation_input_token_cost: float | None = None + cache_creation_input_token_cost_above_1hr: float | None = None + output_cost_per_reasoning_token: float | None = None + input_cost_per_audio_token: float | None = None + output_cost_per_audio_token: float | None = None + input_cost_per_token_above_200k_tokens: float | None = None + output_cost_per_token_above_200k_tokens: float | None = None + input_cost_per_token_flex: float | None = None + output_cost_per_token_flex: float | None = None + input_cost_per_token_priority: float | None = None + output_cost_per_token_priority: float | None = None + search_context_cost_per_query: SearchContextCostPerQuery | None = None + web_search_billing_unit: str | None = None + + +_COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, CostMapEntry]) +_COST_MAP: Final[dict[str, CostMapEntry]] = _COST_MAP_ADAPTER.validate_python( + json.loads(COST_MAP_PATH.read_text()) +) + +TIER_THRESHOLD_TOKENS: Final = 200_000 + + +@dataclass(frozen=True, slots=True) +class FrontierModel: + """One deployment under test: the model_name the suite registers, the + provider-prefixed litellm model string, the wire the scripted upstream + speaks, its cost-map key, and the sibling map model the response_model + override case reports.""" + + model_name: str + litellm_model: str + wire: Wire + map_key: str + override_model: str + + @property + def rates(self) -> CostMapEntry: + return _COST_MAP[self.map_key] + + @property + def override_rates(self) -> CostMapEntry: + return _COST_MAP[self.override_map_key] + + @property + def override_map_key(self) -> str: + return _OVERRIDE_MAP_KEYS[self.override_model] + + @property + def provider(self) -> str: + return self.rates.litellm_provider + + @property + def api_key(self) -> str: + # The scripted upstream ignores auth; a fixed bogus key proves the suite + # spends zero real provider calls. + return "sk-scripted-provider" + + +# Response-model override targets: emit a sibling's bare provider-facing name so +# the biller's provider-prefixed lookup lands on that sibling's map key. +_OVERRIDE_MODELS: Final[dict[str, str]] = { + "gpt-5.6": "gpt-5.4-mini", + "gpt-5.5-pro": "gpt-5.3-codex", + "gpt-5.3-codex": "gpt-5.5-pro", + "gpt-5.4-mini": "gpt-5.6", + "claude-opus-5": "claude-sonnet-5", + "claude-sonnet-5": "claude-opus-5", + "claude-haiku-4-5": "claude-sonnet-5", + "gemini/gemini-3.8-flash": "gemini-3.1-pro-preview", + "gemini/gemini-3.1-pro-preview": "gemini-3.8-flash", + "together_ai/moonshotai/Kimi-K3": "zai-org/GLM-5.3", + "together_ai/zai-org/GLM-5.3": "moonshotai/Kimi-K3", + "fireworks_ai/kimi-k3": "qwen3p8-max", + "fireworks_ai/qwen3p8-max": "kimi-k3", + "fireworks_ai/deepseek-v4p1-flash": "kimi-k3", +} + +_OVERRIDE_MAP_KEYS: Final[dict[str, str]] = { + "gpt-5.4-mini": "gpt-5.4-mini", + "gpt-5.6": "gpt-5.6", + "gpt-5.3-codex": "gpt-5.3-codex", + "gpt-5.5-pro": "gpt-5.5-pro", + "claude-sonnet-5": "claude-sonnet-5", + "claude-opus-5": "claude-opus-5", + "gemini-3.1-pro-preview": "gemini/gemini-3.1-pro-preview", + "gemini-3.8-flash": "gemini/gemini-3.8-flash", + "zai-org/GLM-5.3": "together_ai/zai-org/GLM-5.3", + "moonshotai/Kimi-K3": "together_ai/moonshotai/Kimi-K3", + "qwen3p8-max": "fireworks_ai/qwen3p8-max", + "kimi-k3": "fireworks_ai/kimi-k3", +} + + +_FRONTIER_SPECS: Final[tuple[tuple[str, str, Wire], ...]] = ( + ("gpt-5.6", "openai/gpt-5.6", "openai_chat"), + ("gpt-5.5-pro", "openai/gpt-5.5-pro", "openai_responses"), + ("gpt-5.3-codex", "openai/gpt-5.3-codex", "openai_responses"), + ("gpt-5.4-mini", "openai/gpt-5.4-mini", "openai_chat"), + ("claude-opus-5", "anthropic/claude-opus-5", "anthropic_messages"), + ("claude-sonnet-5", "anthropic/claude-sonnet-5", "anthropic_messages"), + ("claude-haiku-4-5", "anthropic/claude-haiku-4-5", "anthropic_messages"), + ("gemini/gemini-3.8-flash", "gemini/gemini-3.8-flash", "gemini_generate"), + ("gemini/gemini-3.1-pro-preview", "gemini/gemini-3.1-pro-preview", "gemini_generate"), + ("together_ai/moonshotai/Kimi-K3", "together_ai/moonshotai/Kimi-K3", "together_chat"), + ("together_ai/zai-org/GLM-5.3", "together_ai/zai-org/GLM-5.3", "together_chat"), + ("fireworks_ai/kimi-k3", "fireworks_ai/kimi-k3", "fireworks_chat"), + ("fireworks_ai/qwen3p8-max", "fireworks_ai/qwen3p8-max", "fireworks_chat"), + ("fireworks_ai/deepseek-v4p1-flash", "fireworks_ai/deepseek-v4p1-flash", "fireworks_chat"), +) + + +def _frontier() -> tuple[FrontierModel, ...]: + return tuple( + FrontierModel( + model_name=f"cc-{map_key.replace('/', '-').lower()}", + litellm_model=litellm_model, + wire=wire, + map_key=map_key, + override_model=_OVERRIDE_MODELS[map_key], + ) + for map_key, litellm_model, wire in _FRONTIER_SPECS + ) + + +FRONTIER_MODELS: Final[tuple[FrontierModel, ...]] = _frontier() + +# Token kinds each wire can report, gating which pricing cases apply. +_WIRE_CAPS: Final[dict[str, frozenset[str]]] = { + "openai_chat": frozenset( + { + "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", + "web_search", "response_model", "absent_usage", + } + ), + "openai_responses": frozenset({"cache_read", "reasoning", "web_search", "response_model", "absent_usage"}), + # Product gap: litellm hard-indexes message_delta["usage"] in + # anthropic/chat/handler.py, so a usage-absent anthropic stream raises + # KeyError; the real wire always carries it, so the case cannot be + # represented. + "anthropic_messages": frozenset({"cache_read", "cache_write_5m", "cache_write_1h", "web_search", "response_model"}), + # Product gap: the gemini transform sets ModelResponse.model from the + # request and drops the provider's modelVersion, so a response-model + # override can never be priced on this wire. + "gemini_generate": frozenset({"cache_read", "reasoning", "audio", "web_search", "absent_usage"}), + "together_chat": frozenset( + { + "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", + "web_search", "response_model", "absent_usage", + } + ), + "fireworks_chat": frozenset( + { + "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", + "web_search", "response_model", "absent_usage", + } + ), +} + +CaseName = Literal[ + "basic", + "cache_read", + "cache_write_5m", + "cache_write_1h", + "reasoning", + "audio", + "tiered", + "service_tier_flex", + "service_tier_priority", + "web_search", + "stream", + "stream_no_usage", + "response_model_override", +] + + +@dataclass(frozen=True, slots=True) +class Case: + name: CaseName + usage: ScriptedUsage + stream: bool = False + stream_usage: Literal["final_chunk", "absent"] = "final_chunk" + service_tier: Literal["flex", "priority"] | None = None + # For web_search the wire's reported call count is not always what gets + # billed: chat-completions surfaces only expose url_citation annotations, so + # the biller floors to one call; responses/messages/gemini report a real + # count. + billed_web_search_calls: int = 0 + response_model_override: bool = False + exact_spend: bool = True + # stream_usage=absent on a wire with no proxy-side token recount means the + # bill is exactly zero; asserted as such rather than skipped. + expect_zero_bill: bool = False + + def scenario(self, scenario_id: str, model: FrontierModel, text: str) -> Scenario: + return Scenario( + scenario_id=scenario_id, + wire=model.wire, + usage=self.usage, + output=ScriptedOutput( + text=text, + response_model=model.override_model if self.response_model_override else None, + ), + stream_usage=self.stream_usage, + service_tier=self.service_tier, + ) + + +_BASIC_USAGE: Final = ScriptedUsage(fresh_input_tokens=120, output_tokens=40) + + +def _web_search_case(model: FrontierModel) -> Case: + counts_exactly = model.wire in ("openai_responses", "anthropic_messages", "gemini_generate") + return Case( + name="web_search", + usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, web_search_calls=3), + billed_web_search_calls=3 if counts_exactly else 1, + ) + + +def cases_for(model: FrontierModel) -> tuple[Case, ...]: + rates = model.rates + caps = _WIRE_CAPS[model.wire] + cases: list[Case] = [Case(name="basic", usage=_BASIC_USAGE)] + if rates.cache_read_input_token_cost is not None and "cache_read" in caps: + cases.append( + Case(name="cache_read", usage=ScriptedUsage(fresh_input_tokens=100, cache_read_tokens=50, output_tokens=30)) + ) + if rates.cache_creation_input_token_cost is not None and "cache_write_5m" in caps: + cases.append( + Case( + name="cache_write_5m", + usage=ScriptedUsage(fresh_input_tokens=90, cache_write_5m_tokens=60, output_tokens=30), + ) + ) + if ( + rates.cache_creation_input_token_cost_above_1hr is not None + and rates.cache_creation_input_token_cost is not None + and "cache_write_1h" in caps + ): + cases.append( + Case( + name="cache_write_1h", + usage=ScriptedUsage( + fresh_input_tokens=90, + cache_write_5m_tokens=20, + cache_write_1h_tokens=40, + output_tokens=30, + ), + ) + ) + if rates.output_cost_per_reasoning_token is not None and "reasoning" in caps: + cases.append( + Case( + name="reasoning", + usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, reasoning_tokens=70), + ) + ) + if ( + rates.input_cost_per_audio_token is not None + and rates.output_cost_per_audio_token is not None + and "audio" in caps + ): + cases.append( + Case( + name="audio", + usage=ScriptedUsage( + fresh_input_tokens=100, audio_input_tokens=25, output_tokens=30, audio_output_tokens=15 + ), + ) + ) + if ( + rates.input_cost_per_token_above_200k_tokens is not None + and rates.output_cost_per_token_above_200k_tokens is not None + ): + cases.append( + Case( + name="tiered", + usage=ScriptedUsage( + fresh_input_tokens=TIER_THRESHOLD_TOKENS + 1, output_tokens=30 + ), + ) + ) + if rates.input_cost_per_token_flex is not None and rates.output_cost_per_token_flex is not None: + cases.append( + Case(name="service_tier_flex", usage=_BASIC_USAGE, service_tier="flex") + ) + if rates.input_cost_per_token_priority is not None and rates.output_cost_per_token_priority is not None: + cases.append( + Case(name="service_tier_priority", usage=_BASIC_USAGE, service_tier="priority") + ) + if rates.search_context_cost_per_query is not None and "web_search" in caps: + cases.append(_web_search_case(model)) + cases.append(Case(name="stream", usage=_BASIC_USAGE, stream=True)) + if "absent_usage" in caps: + cases.append( + Case( + name="stream_no_usage", + usage=_BASIC_USAGE, + stream=True, + stream_usage="absent", + exact_spend=False, + # The responses surface bills only provider-reported usage; + # with no usage in the stream the spend row is zero. Other + # wires recount tokens proxy-side and bill a nonzero amount. + expect_zero_bill=model.wire == "openai_responses", + ) + ) + if "response_model" in caps: + cases.append(Case(name="response_model_override", usage=_BASIC_USAGE, response_model_override=True)) + return tuple(cases) + + +@dataclass(frozen=True, slots=True) +class ExpectedCost: + """The expected bill split the way the spend row's cost_breakdown reports + it: the gross input component (cache reads/writes folded in), the output + component, and the tool-usage component.""" + + input_cost: float + output_cost: float + tool_cost: float + + @property + def total(self) -> float: + return self.input_cost + self.output_cost + self.tool_cost + + +def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: + """Literal arithmetic on the test-map rates over the scripted token counts. + + Input = fresh*in + read*read + 5m*create + 1h*create_1h + audio_in*audio_in; + output = text*out + reasoning*reasoning + audio_out*audio_out; plus the + billed web-search calls at the medium search-context rate. Above-threshold + swaps every input/output rate to its ``_above_200k_tokens`` variant when + total prompt tokens exceed the threshold; a service tier swaps input/output + to the tier's variants, falling back to the base rate when a variant is + unset -- mirroring _get_token_base_cost in litellm's cost calculator. + """ + rates = model.override_rates if case.response_model_override else model.rates + u = case.usage + prompt_tokens = ( + u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + + u.cache_write_1h_tokens + u.audio_input_tokens + ) + tiered = prompt_tokens > TIER_THRESHOLD_TOKENS + in_rate = rates.input_cost_per_token or 0.0 + out_rate = rates.output_cost_per_token or 0.0 + if case.service_tier == "flex": + in_rate = rates.input_cost_per_token_flex or in_rate + out_rate = rates.output_cost_per_token_flex or out_rate + if case.service_tier == "priority": + in_rate = rates.input_cost_per_token_priority or in_rate + out_rate = rates.output_cost_per_token_priority or out_rate + if tiered: + in_rate = rates.input_cost_per_token_above_200k_tokens or in_rate + out_rate = rates.output_cost_per_token_above_200k_tokens or out_rate + input_cost = ( + u.fresh_input_tokens * in_rate + + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) + + u.cache_write_5m_tokens * (rates.cache_creation_input_token_cost or 0.0) + + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or 0.0) + + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) + ) + output_cost = ( + u.output_tokens * out_rate + + u.reasoning_tokens * (rates.output_cost_per_reasoning_token or out_rate) + + u.audio_output_tokens * (rates.output_cost_per_audio_token or out_rate) + ) + search = rates.search_context_cost_per_query + tool_cost = case.billed_web_search_calls * ( + search.search_context_size_medium if search and search.search_context_size_medium else 0.0 + ) + return ExpectedCost(input_cost=input_cost, output_cost=output_cost, tool_cost=tool_cost) + + +def expected_cost(model: FrontierModel, case: Case) -> float: + return expected_breakdown(model, case).total + + +def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: + """(prompt_tokens, completion_tokens) the spend row should carry, per the + wire's normalization: Anthropic folds cache read/write into prompt_tokens, + everyone else reports the totals the wire emitted.""" + u = case.usage + if model.wire == "anthropic_messages": + 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": + return ( + u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens, + u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, + ) + if model.wire == "openai_responses": + return ( + u.fresh_input_tokens + u.cache_read_tokens, + u.output_tokens + u.reasoning_tokens, + ) + return ( + u.fresh_input_tokens + + u.cache_read_tokens + + u.cache_write_5m_tokens + + u.cache_write_1h_tokens + + u.audio_input_tokens, + u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, + ) diff --git a/tests/e2e/cost_calculation/scripted_client.py b/tests/e2e/cost_calculation/scripted_client.py new file mode 100644 index 00000000000..dceec02630a --- /dev/null +++ b/tests/e2e/cost_calculation/scripted_client.py @@ -0,0 +1,70 @@ +"""Client side of the scripted-provider sidecar: register scenarios over its +control API through the shared transport helpers and get back a handle whose +``api_base`` is what a /model/new deployment should register for the proxy to +reach the scripted wire.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final + +from e2e_config import SCRIPTED_PROVIDER_CONTROL_URL, SCRIPTED_PROVIDER_PROXY_BASE +from e2e_http import URL, NoBody, unwrap, post +from e2e_http import delete as http_delete +from scripted_provider import ( + Scenario, + ScenarioDeleted, + ScenarioRegistered, + Wire, +) + + +@dataclass(frozen=True, slots=True) +class ScenarioHandle: + scenario_id: str + wire: Wire + proxy_base: str + + def api_base(self) -> str: + return f"{self.proxy_base}/{self.scenario_id}/{self._mount()}" + + def _mount(self) -> str: + return { + "openai_chat": "openai", + "openai_responses": "openai", + "anthropic_messages": "anthropic", + "gemini_generate": "gemini", + "together_chat": "together", + "fireworks_chat": "fireworks", + }[self.wire] + + +def register_scenario(scenario: Scenario) -> ScenarioHandle: + """POST the scenario to the sidecar's control API and return its handle.""" + result = unwrap( + post( + URL(f"{SCRIPTED_PROVIDER_CONTROL_URL}/_scenarios"), + headers=NoBody(), + json=scenario, + response_type=ScenarioRegistered, + ) + ) + return ScenarioHandle( + scenario_id=result.scenario_id, + wire=scenario.wire, + proxy_base=SCRIPTED_PROVIDER_PROXY_BASE, + ) + + +def delete_scenario(handle: ScenarioHandle) -> None: + unwrap( + http_delete( + URL(f"{SCRIPTED_PROVIDER_CONTROL_URL}/_scenarios/{handle.scenario_id}"), + headers=NoBody(), + json=NoBody(), + response_type=ScenarioDeleted, + ) + ) + + +CONTROL_URL: Final = SCRIPTED_PROVIDER_CONTROL_URL diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py new file mode 100644 index 00000000000..93a6f49ec25 --- /dev/null +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -0,0 +1,631 @@ +"""Scripted provider sidecar for the cost-calculation e2e suite. + +A standalone process (``python -m cost_calculation.scripted_provider``) that +pretends to be an LLM provider for the proxy under test. The suite registers a +Scenario over a small control API; the provider wire routes then 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. + +Layout on one port: + +- ``GET /health`` liveness +- ``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`` and the + remainder is whatever path the provider client appends (``chat/completions``, + ``responses``, ``v1/messages``, ``models/:generateContent`` ...) + +A request carrying ``"stream": true`` (or the ``:streamGenerateContent`` Gemini +verb) gets an SSE answer; ``stream_usage`` on the Scenario decides whether the +final stream chunk carries usage or the provider reports none. +""" + +from __future__ import annotations + +import json +import sys +import threading +import time +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final, Literal +from urllib.parse import urlsplit + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +Wire = Literal[ + "openai_chat", + "openai_responses", + "anthropic_messages", + "gemini_generate", + "together_chat", + "fireworks_chat", +] + +_WIRE_MOUNTS: Final[dict[str, str]] = { + "openai_chat": "openai", + "openai_responses": "openai", + "anthropic_messages": "anthropic", + "gemini_generate": "gemini", + "together_chat": "together", + "fireworks_chat": "fireworks", +} + +StreamUsage = Literal["final_chunk", "absent"] +ServiceTier = Literal["flex", "priority"] + + +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 + provider does (inside prompt_tokens for OpenAI/Gemini, as uncached-only + input_tokens for Anthropic).""" + + model_config = ConfigDict(frozen=True) + + fresh_input_tokens: int = 0 + output_tokens: int = 0 + cache_read_tokens: int = 0 + cache_write_5m_tokens: int = 0 + cache_write_1h_tokens: int = 0 + reasoning_tokens: int = 0 + audio_input_tokens: int = 0 + audio_output_tokens: int = 0 + web_search_calls: int = 0 + + +class ScriptedOutput(BaseModel): + model_config = ConfigDict(frozen=True) + + text: str + finish_reason: str = "stop" + # When set, emitted verbatim as the response's model field, letting a test + # 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. + provider_cost: float | None = None + + +class Scenario(BaseModel): + model_config = ConfigDict(frozen=True) + + scenario_id: str + wire: Wire + usage: ScriptedUsage + output: ScriptedOutput + stream_usage: StreamUsage = "final_chunk" + service_tier: ServiceTier | None = None + + @property + def mount(self) -> str: + return _WIRE_MOUNTS[self.wire] + + +class ScenarioRegistered(BaseModel): + scenario_id: str + + +class ScenarioDeleted(BaseModel): + deleted: bool + + +class HealthStatus(BaseModel): + status: str + + +@dataclass(frozen=True, slots=True) +class RenderedResponse: + status_code: int + content_type: str + body: bytes + + +def _json_bytes(payload: dict[str, object]) -> bytes: + return json.dumps(payload).encode("utf-8") + + +def _sse(events: tuple[tuple[str | None, dict[str, object] | str], ...]) -> bytes: + frames: list[str] = [] + for event_name, data in events: + head = f"event: {event_name}\n" if event_name is not None else "" + payload = data if isinstance(data, str) else json.dumps(data) + frames.append(f"{head}data: {payload}\n\n") + return "".join(frames).encode("utf-8") + + +# ---------- per-wire usage shapes ---------- + + +def _openai_usage(u: ScriptedUsage) -> dict[str, object]: + prompt_tokens = ( + u.fresh_input_tokens + + u.cache_read_tokens + + u.cache_write_5m_tokens + + u.cache_write_1h_tokens + + u.audio_input_tokens + ) + completion_tokens = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + prompt_details: dict[str, object] = {} + if u.cache_read_tokens: + prompt_details["cached_tokens"] = u.cache_read_tokens + if u.cache_write_5m_tokens or u.cache_write_1h_tokens: + prompt_details["cache_write_tokens"] = u.cache_write_5m_tokens + u.cache_write_1h_tokens + prompt_details["cache_creation_token_details"] = { + "ephemeral_5m_input_tokens": u.cache_write_5m_tokens, + "ephemeral_1h_input_tokens": u.cache_write_1h_tokens, + } + if u.audio_input_tokens: + prompt_details["audio_tokens"] = u.audio_input_tokens + completion_details: dict[str, object] = {} + if u.reasoning_tokens: + completion_details["reasoning_tokens"] = u.reasoning_tokens + if u.audio_output_tokens: + completion_details["audio_tokens"] = u.audio_output_tokens + usage: dict[str, object] = { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + } + if prompt_details: + usage["prompt_tokens_details"] = prompt_details + if completion_details: + usage["completion_tokens_details"] = completion_details + return usage + + +def _anthropic_usage(u: ScriptedUsage) -> dict[str, object]: + # Anthropic reports uncached-only input_tokens; cache reads and writes ride + # top-level fields, with the 5m/1h write split under cache_creation. + usage: dict[str, object] = { + "input_tokens": u.fresh_input_tokens, + "output_tokens": u.output_tokens, + } + if u.cache_read_tokens: + usage["cache_read_input_tokens"] = u.cache_read_tokens + if u.cache_write_5m_tokens or u.cache_write_1h_tokens: + usage["cache_creation_input_tokens"] = u.cache_write_5m_tokens + u.cache_write_1h_tokens + usage["cache_creation"] = { + "ephemeral_5m_input_tokens": u.cache_write_5m_tokens, + "ephemeral_1h_input_tokens": u.cache_write_1h_tokens, + } + if u.web_search_calls: + usage["server_tool_use"] = {"web_search_requests": u.web_search_calls} + return usage + + +def _gemini_usage(u: ScriptedUsage) -> dict[str, object]: + # promptTokenCount carries the cached count inside it; TEXT modality is the + # cached-inclusive text count so litellm's implicit-caching subtraction lands + # on the fresh figure. candidatesTokenCount includes reasoning + audio. + prompt_tokens = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens + candidates = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + usage: dict[str, object] = { + "promptTokenCount": prompt_tokens, + "candidatesTokenCount": candidates, + "totalTokenCount": prompt_tokens + candidates, + } + if u.cache_read_tokens: + usage["cachedContentTokenCount"] = u.cache_read_tokens + if u.reasoning_tokens: + usage["thoughtsTokenCount"] = u.reasoning_tokens + prompt_details = [{"modality": "TEXT", "tokenCount": u.fresh_input_tokens + u.cache_read_tokens}] + if u.audio_input_tokens: + prompt_details.append({"modality": "AUDIO", "tokenCount": u.audio_input_tokens}) + usage["promptTokensDetails"] = prompt_details + if u.audio_output_tokens: + usage["candidatesTokensDetails"] = [ + {"modality": "TEXT", "tokenCount": u.output_tokens + u.reasoning_tokens}, + {"modality": "AUDIO", "tokenCount": u.audio_output_tokens}, + ] + return usage + + +def _responses_usage(u: ScriptedUsage) -> dict[str, object]: + input_tokens = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens + output_tokens = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + usage: dict[str, object] = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + } + input_details: dict[str, object] = {} + if u.cache_read_tokens: + input_details["cached_tokens"] = u.cache_read_tokens + if input_details: + usage["input_tokens_details"] = input_details + if u.reasoning_tokens: + usage["output_tokens_details"] = {"reasoning_tokens": u.reasoning_tokens} + return usage + + +# ---------- per-wire responses ---------- + + +def _openai_message(scenario: Scenario) -> dict[str, object]: + message: dict[str, object] = {"role": "assistant", "content": scenario.output.text} + if scenario.usage.web_search_calls: + message["annotations"] = [ + { + "type": "url_citation", + "url_citation": { + "url": "https://scripted.example/source", + "title": "scripted source", + "start_index": 0, + "end_index": 1, + }, + } + for _ in range(scenario.usage.web_search_calls) + ] + return message + + +def _openai_chat_body(scenario: Scenario, requested_model: str) -> dict[str, object]: + body: dict[str, object] = { + "id": f"chatcmpl-{scenario.scenario_id}", + "object": "chat.completion", + "created": int(time.time()), + "model": scenario.output.response_model or requested_model, + "choices": [ + { + "index": 0, + "message": _openai_message(scenario), + "finish_reason": scenario.output.finish_reason, + } + ], + "usage": _openai_usage(scenario.usage), + } + if scenario.service_tier is not None: + body["service_tier"] = scenario.service_tier + if scenario.output.provider_cost is not None: + body["cost"] = scenario.output.provider_cost + return body + + +def _openai_chunk(scenario: Scenario, requested_model: str, **kw: object) -> dict[str, object]: + chunk: dict[str, object] = { + "id": f"chatcmpl-{scenario.scenario_id}", + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": scenario.output.response_model or requested_model, + } + chunk.update(kw) + return chunk + + +def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: + _EMPTY_DELTA: Final[dict[str, object]] = {} + delta: dict[str, object] = {"role": "assistant", "content": scenario.output.text} + if scenario.usage.web_search_calls: + delta["annotations"] = _openai_message(scenario)["annotations"] + events: list[tuple[str | None, dict[str, object] | str]] = [ + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=[{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}], + ), + ), + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=[{"index": 0, "delta": delta, "finish_reason": None}], + ), + ), + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=[ + { + "index": 0, + "delta": _EMPTY_DELTA, + "finish_reason": scenario.output.finish_reason, + } + ], + ), + ), + ] + if scenario.stream_usage == "final_chunk": + events.append( + (None, _openai_chunk(scenario, requested_model, choices=(), usage=_openai_usage(scenario.usage))) + ) + events.append((None, "[DONE]")) + return _sse(tuple(events)) + + +def _anthropic_body(scenario: Scenario, requested_model: str) -> dict[str, object]: + return { + "id": f"msg_{scenario.scenario_id}", + "type": "message", + "role": "assistant", + "model": scenario.output.response_model or requested_model, + "content": [{"type": "text", "text": scenario.output.text}], + "stop_reason": "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, + "usage": _anthropic_usage(scenario.usage), + } + + +def _anthropic_sse(scenario: Scenario, requested_model: str) -> bytes: + emit_usage = scenario.stream_usage == "final_chunk" + input_usage = {k: v for k, v in _anthropic_usage(scenario.usage).items() if k != "output_tokens"} + message_start: dict[str, object] = { + "type": "message_start", + "message": { + "id": f"msg_{scenario.scenario_id}", + "type": "message", + "role": "assistant", + "model": scenario.output.response_model or requested_model, + "content": [], + "stop_reason": None, + **({"usage": input_usage} if emit_usage else {}), + }, + } + message_delta: dict[str, object] = { + "type": "message_delta", + "delta": { + "stop_reason": "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason + }, + **({"usage": {"output_tokens": scenario.usage.output_tokens}} if emit_usage else {}), + } + return _sse( + ( + ("message_start", message_start), + ( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ), + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": scenario.output.text}, + }, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", message_delta), + ("message_stop", {"type": "message_stop"}), + ) + ) + + +def _gemini_body(scenario: Scenario, requested_model: str) -> dict[str, object]: + candidate: dict[str, object] = { + "content": {"parts": [{"text": scenario.output.text}], "role": "model"}, + "finishReason": "STOP" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason.upper(), + "index": 0, + } + if scenario.usage.web_search_calls: + candidate["groundingMetadata"] = { + "webSearchQueries": [f"query {i}" for i in range(scenario.usage.web_search_calls)] + } + return { + "candidates": [candidate], + "usageMetadata": _gemini_usage(scenario.usage), + "modelVersion": scenario.output.response_model or requested_model, + } + + +def _gemini_sse(scenario: Scenario, requested_model: str) -> bytes: + first = _gemini_body(scenario, requested_model) + if scenario.stream_usage == "absent": + first = {k: v for k, v in first.items() if k != "usageMetadata"} + events: list[tuple[str | None, dict[str, object] | str]] = [(None, first)] + if scenario.stream_usage == "final_chunk": + events.append( + ( + None, + { + "candidates": [], + "usageMetadata": _gemini_usage(scenario.usage), + "modelVersion": scenario.output.response_model or requested_model, + }, + ) + ) + return _sse(tuple(events)) + + +def _responses_body(scenario: Scenario, requested_model: str) -> dict[str, object]: + output: list[dict[str, object]] = [ + {"type": "web_search_call", "id": f"ws_{i}", "status": "completed"} + for i in range(scenario.usage.web_search_calls) + ] + output.append( + { + "type": "message", + "id": f"msg_{scenario.scenario_id}", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": scenario.output.text, + "annotations": [], + } + ], + } + ) + return { + "id": f"resp_{scenario.scenario_id}", + "object": "response", + "created_at": int(time.time()), + "status": "completed", + "model": scenario.output.response_model or requested_model, + "output": output, + "usage": _responses_usage(scenario.usage), + } + + +def _responses_sse(scenario: Scenario, requested_model: str) -> bytes: + completed = _responses_body(scenario, requested_model) + if scenario.stream_usage == "absent": + completed = {k: v for k, v in completed.items() if k != "usage"} + created = {**completed, "status": "in_progress", "usage": None} + return _sse( + ( + ("response.created", {"type": "response.created", "response": created}), + ( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": f"msg_{scenario.scenario_id}", + "output_index": scenario.usage.web_search_calls, + "content_index": 0, + "delta": scenario.output.text, + }, + ), + ("response.completed", {"type": "response.completed", "response": completed}), + ) + ) + + +def _render(scenario: Scenario, *, stream: bool, requested_model: str) -> RenderedResponse: + if scenario.wire == "anthropic_messages": + if stream: + return RenderedResponse(200, "text/event-stream", _anthropic_sse(scenario, requested_model)) + return RenderedResponse(200, "application/json", _json_bytes(_anthropic_body(scenario, requested_model))) + if scenario.wire == "gemini_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 == "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))) + # openai_chat, together_chat, fireworks_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))) + + +# ---------- registry + request routing ---------- + + +class _ScenarioStore: + def __init__(self) -> None: + self._lock: Final = threading.Lock() + self._scenarios: dict[str, Scenario] = {} # mutable-ok: server state, guarded by _lock + + def put(self, scenario: Scenario) -> None: + with self._lock: + self._scenarios[scenario.scenario_id] = scenario + + def drop(self, scenario_id: str) -> bool: + with self._lock: + return self._scenarios.pop(scenario_id, None) is not None + + def get(self, scenario_id: str) -> Scenario | None: + with self._lock: + return self._scenarios.get(scenario_id) + + +_REQUEST_BODY: Final = TypeAdapter(dict[str, object]) + + +def _request_body(body: bytes) -> dict[str, object]: + try: + return _REQUEST_BODY.validate_json(body) + except ValueError: + return {} + + +def _request_wants_stream(path_tail: str, body: bytes) -> bool: + if ":streamGenerateContent" in path_tail: + return True + if not body: + return False + return _request_body(body).get("stream") is True + + +def _request_model(body: bytes) -> str: + model = _request_body(body).get("model") + return model if isinstance(model, str) else "unknown" + + +def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: bytes) -> RenderedResponse: + path = urlsplit(raw_path).path + segments = [segment for segment in path.split("/") if segment] + if method == "GET" and segments == ["health"]: + return RenderedResponse(200, "application/json", _json_bytes({"status": "ok"})) + if segments and segments[0] == "_scenarios": + if method == "POST" and len(segments) == 1: + try: + scenario = Scenario.model_validate_json(body) + except ValidationError as exc: + return RenderedResponse(400, "application/json", _json_bytes({"error": str(exc)})) + store.put(scenario) + return RenderedResponse(200, "application/json", _json_bytes({"scenario_id": scenario.scenario_id})) + if method == "DELETE" and len(segments) == 2: + deleted = store.drop(segments[1]) + return RenderedResponse( + 200 if deleted else 404, "application/json", _json_bytes({"deleted": deleted}) + ) + return RenderedResponse(404, "application/json", _json_bytes({"error": "unknown control route"})) + if len(segments) < 2 or method != "POST": + return RenderedResponse(404, "application/json", _json_bytes({"error": f"no route for {method} {path}"})) + scenario_id, mount = segments[0], segments[1] + scenario = store.get(scenario_id) + if scenario is None: + return RenderedResponse(404, "application/json", _json_bytes({"error": f"unknown scenario {scenario_id}"})) + if scenario.mount != mount: + return RenderedResponse( + 400, + "application/json", + _json_bytes({"error": f"scenario {scenario_id} is wire {scenario.wire}, not mount {mount}"}), + ) + tail = "/".join(segments[2:]) + return _render(scenario, stream=_request_wants_stream(tail, body), requested_model=_request_model(body)) + + +class _ScriptedHandler(BaseHTTPRequestHandler): + store: Final[_ScenarioStore] = _ScenarioStore() + + def _dispatch(self, method: str) -> None: + length = int(self.headers.get("content-length") or 0) + body = self.rfile.read(length) if length else b"" + rendered = handle_request(self.store, method, self.path, body) + self.send_response(rendered.status_code) + self.send_header("content-type", rendered.content_type) + self.send_header("content-length", str(len(rendered.body))) + self.end_headers() + self.wfile.write(rendered.body) + + def do_GET(self) -> None: + self._dispatch("GET") + + def do_POST(self) -> None: + self._dispatch("POST") + + def do_DELETE(self) -> None: + self._dispatch("DELETE") + + + +DEFAULT_PORT: Final = 9100 + + +def serve(port: int = DEFAULT_PORT, bind_host: str = "127.0.0.1") -> None: + server = ThreadingHTTPServer((bind_host, port), _ScriptedHandler) + sys.stderr.write(f"scripted-provider listening on http://{bind_host}:{port}\n") + server.serve_forever() + + +if __name__ == "__main__": + port_arg = int(sys.argv[1]) if len(sys.argv) > 1 else DEFAULT_PORT + serve(port=port_arg) diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py new file mode 100644 index 00000000000..8d7678cf9ca --- /dev/null +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -0,0 +1,115 @@ +"""Token-pricing e2e: every (frontier model, pricing-component case) cell runs a +scripted-usage call through a deployment registered on the cost-map proxy, and +the spend row plus response-cost header must equal literal arithmetic on the +test map's rates. + +Nothing here touches a real provider or the bundled cost map: the proxy's +upstream is the scripted-provider sidecar and its entire cost map is +tests/e2e/cost_map.json. +""" + +from __future__ import annotations + +import pytest + +from conftest import CostCalcClient, cost_rows, register_scenario_deployment +from cost_matrix import ( + FRONTIER_MODELS, + Case, + FrontierModel, + cases_for, + expected_cost, + expected_token_columns, +) +from e2e_config import unique_marker +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, ChatStreamOptions + +pytestmark = [pytest.mark.e2e, pytest.mark.cost_map_stack] + +_MATRIX: list[tuple[FrontierModel, Case]] = [ + (model, case) for model in FRONTIER_MODELS for case in cases_for(model) +] + + +def _case_id(param: tuple[FrontierModel, Case]) -> str: + model, case = param + return f"{model.map_key.replace('/', '-')}-{case.name}" + + +def _chat_body(model_name: str, marker: str, case: Case) -> ChatBody: + return ChatBody( + model=model_name, + messages=[ChatMessage(role="user", content=f"{marker} scripted pricing call")], + stream=case.stream, + stream_options=ChatStreamOptions(include_usage=True) if case.stream else None, + service_tier=case.service_tier, + ) + + +class TestTokenPricing: + @pytest.mark.parametrize("model_case", _MATRIX, ids=_case_id) + @pytest.mark.covers("quota_management.spend_tracking.cost_matrix.logs_cost") + def test_scripted_usage_bills_at_map_rates( + self, + client: CostCalcClient, + resources: ResourceManager, + scoped_key: str, + model_case: tuple[FrontierModel, Case], + ) -> None: + model, case = model_case + marker = unique_marker() + model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) + response = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(scoped_key), + json=_chat_body(model_name, marker, case), + stream=case.stream, + ) + assert response.ok, ( + f"{model.map_key}/{case.name}: proxy returned {response.status_code}: {response.body[:400]}" + ) + assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" + + expected = expected_cost(model, case) + if case.exact_spend and not case.stream: + # Streamed responses commit headers before the bill is computed, so + # the x-litellm-response-cost header is asserted only on non-stream + # calls. + assert response.response_cost is not None and cost_rows.approx_equal( + response.response_cost, expected + ), ( + f"x-litellm-response-cost {response.response_cost} != expected {expected}" + ) + + row = cost_rows.poll_cost_row_where( + client.proxy, + scoped_key, + lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, + ) + assert row is not None, f"no spend row with a cost breakdown landed for {model.map_key}/{case.name}" + + if not case.exact_spend and case.expect_zero_bill: + # The provider reported no usage and this wire has no proxy-side + # recount, so the bill is exactly zero. + assert row.spend is not None and row.spend == 0, f"no-usage stream billed {row.spend}: {row}" + return + if not case.exact_spend: + # stream_usage=absent: the provider reported no usage, so the row's + # token counts are the proxy's own recount; only assert a bill landed. + assert row.spend is not None and row.spend > 0, f"no-usage stream billed nothing: {row}" + return + + assert row.spend is not None and cost_rows.approx_equal(row.spend, expected), ( + f"{model.map_key}/{case.name}: spend {row.spend} != expected {expected} " + f"(breakdown {row.breakdown.model_dump()})" + ) + + prompt_tokens, completion_tokens = expected_token_columns(model, case) + assert row.prompt_tokens == prompt_tokens, ( + f"prompt_tokens {row.prompt_tokens} != {prompt_tokens}" + ) + assert row.completion_tokens == completion_tokens, ( + f"completion_tokens {row.completion_tokens} != {completion_tokens}" + ) + cost_rows.assert_total_is_sum_of_components(row) diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py new file mode 100644 index 00000000000..b1ef675d9ef --- /dev/null +++ b/tests/e2e/cost_calculation/test_wire_formats_e2e.py @@ -0,0 +1,186 @@ +"""Wire-format e2e: one scripted upstream per provider wire, answering with a +usage payload where every token kind the wire can report is nonzero. The spend +row's gross input cost must equal fresh tokens at the input rate plus each cache +and audio component at its own rate -- proving the wire's usage shape landed the +cached tokens inside the total (OpenAI/Gemini) or as separate fields +(Anthropic), and that the biller subtracted them before billing fresh tokens. + +Also covers the Responses API wire (an openai/gpt-5.5-pro deployment bridged by +the proxy to POST /responses) and a streamed Anthropic-messages case. +""" + +from __future__ import annotations + +import pytest + +from conftest import CostCalcClient, cost_rows, register_scenario_deployment +from cost_matrix import ( + FRONTIER_MODELS, + Case, + FrontierModel, + expected_breakdown, + expected_token_columns, +) +from e2e_config import unique_marker +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, ChatStreamOptions +from scripted_provider import ScriptedUsage + +pytestmark = [pytest.mark.e2e, pytest.mark.cost_map_stack] + +_MODELS: dict[str, FrontierModel] = {model.map_key: model for model in FRONTIER_MODELS} + +# One scripted usage per wire, every reportable token kind nonzero. +_WIRE_USAGE: dict[str, tuple[str, ScriptedUsage]] = { + "openai_chat": ( + "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, + ), + ), + "openai_responses": ( + "gpt-5.5-pro", + ScriptedUsage( + fresh_input_tokens=80, cache_read_tokens=40, output_tokens=25, reasoning_tokens=15 + ), + ), + "anthropic_messages": ( + "claude-sonnet-5", + ScriptedUsage( + fresh_input_tokens=80, + cache_read_tokens=40, + cache_write_5m_tokens=20, + cache_write_1h_tokens=10, + output_tokens=25, + ), + ), + "gemini_generate": ( + "gemini/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, + ), + ), + "together_chat": ( + "together_ai/moonshotai/Kimi-K3", + 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, + ), + ), + "fireworks_chat": ( + "fireworks_ai/kimi-k3", + ScriptedUsage(fresh_input_tokens=80, cache_read_tokens=40, output_tokens=25), + ), +} + + +class TestWireFormats: + @pytest.mark.parametrize("wire", tuple(_WIRE_USAGE)) + @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") + def test_wire_usage_shape_bills_each_component( + self, + client: CostCalcClient, + resources: ResourceManager, + scoped_key: str, + wire: str, + ) -> None: + map_key, usage = _WIRE_USAGE[wire] + model = _MODELS[map_key] + case = Case(name="basic", usage=usage) + marker = unique_marker() + model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) + response = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(scoped_key), + json=ChatBody( + model=model_name, + messages=[ChatMessage(role="user", content=f"{marker} scripted wire call")], + ), + ) + assert response.ok, f"{wire}: proxy returned {response.status_code}: {response.body[:400]}" + + expected = expected_breakdown(model, case) + row = cost_rows.poll_cost_row_where( + client.proxy, + scoped_key, + lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, + ) + assert row is not None, f"{wire}: no spend row landed" + assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( + f"{wire}: spend {row.spend} != expected {expected.total} " + f"(breakdown {row.breakdown.model_dump()})" + ) + breakdown = row.breakdown + assert breakdown.input_cost is not None and cost_rows.approx_equal( + breakdown.input_cost, expected.input_cost + ), ( + f"{wire}: gross input_cost {breakdown.input_cost} != expected {expected.input_cost}; " + "cached/written tokens billed at the input rate" + ) + assert breakdown.output_cost is not None and cost_rows.approx_equal( + breakdown.output_cost, expected.output_cost + ), f"{wire}: output_cost {breakdown.output_cost} != expected {expected.output_cost}" + + prompt_tokens, completion_tokens = expected_token_columns(model, case) + assert row.prompt_tokens == prompt_tokens, ( + f"{wire}: prompt_tokens {row.prompt_tokens} != {prompt_tokens}" + ) + assert row.completion_tokens == completion_tokens, ( + f"{wire}: completion_tokens {row.completion_tokens} != {completion_tokens}" + ) + cost_rows.assert_total_is_sum_of_components(row) + + @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") + def test_anthropic_streamed_usage_bills_each_component( + self, client: CostCalcClient, resources: ResourceManager, scoped_key: str + ) -> None: + map_key, usage = _WIRE_USAGE["anthropic_messages"] + model = _MODELS[map_key] + case = Case(name="stream", usage=usage, stream=True) + marker = unique_marker() + model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) + response = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(scoped_key), + json=ChatBody( + model=model_name, + messages=[ChatMessage(role="user", content=f"{marker} scripted anthropic stream")], + stream=True, + stream_options=ChatStreamOptions(include_usage=True), + ), + stream=True, + ) + assert response.ok, f"anthropic stream: proxy returned {response.status_code}: {response.body[:400]}" + assert response.stream_done, "anthropic stream did not reach its terminal event" + assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" + + expected = expected_breakdown(model, case) + row = cost_rows.poll_cost_row_where( + client.proxy, + scoped_key, + lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, + ) + assert row is not None, "anthropic stream: no spend row landed" + assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( + f"anthropic stream: spend {row.spend} != expected {expected.total} " + f"(breakdown {row.breakdown.model_dump()})" + ) + cost_rows.assert_total_is_sum_of_components(row) diff --git a/tests/e2e/cost_map.json b/tests/e2e/cost_map.json new file mode 100644 index 00000000000..b761710bae3 --- /dev/null +++ b/tests/e2e/cost_map.json @@ -0,0 +1,352 @@ +{ + "claude-haiku-4-5": { + "cache_creation_input_token_cost": 0.00021, + "cache_creation_input_token_cost_above_1hr": 0.00028000000000000003, + "cache_read_input_token_cost": 7e-06, + "input_cost_per_token": 7.000000000000001e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00014000000000000001, + "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 + }, + "claude-opus-5": { + "cache_creation_input_token_cost": 0.00015000000000000001, + "cache_creation_input_token_cost_above_1hr": 0.0002, + "cache_read_input_token_cost": 4.9999999999999996e-06, + "input_cost_per_token": 5e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.0001, + "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 + }, + "claude-sonnet-5": { + "cache_creation_input_token_cost": 0.00018, + "cache_creation_input_token_cost_above_1hr": 0.00024000000000000003, + "cache_read_input_token_cost": 6e-06, + "input_cost_per_token": 6.000000000000001e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00012000000000000002, + "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 + }, + "fireworks_ai/deepseek-v4p1-flash": { + "cache_read_input_token_cost": 1.4e-05, + "input_cost_per_token": 0.00014000000000000001, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00028000000000000003, + "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 + }, + "fireworks_ai/kimi-k3": { + "cache_read_input_token_cost": 1.2e-05, + "input_cost_per_token": 0.00012000000000000002, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00024000000000000003, + "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 + }, + "fireworks_ai/qwen3p8-max": { + "cache_read_input_token_cost": 1.3e-05, + "input_cost_per_token": 0.00013000000000000002, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 0.00026000000000000003, + "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/gemini-3.1-pro-preview": { + "cache_read_input_token_cost": 9e-06, + "input_cost_per_audio_token": 0.00054, + "input_cost_per_token": 9e-05, + "input_cost_per_token_above_200k_tokens": 0.00072, + "input_cost_per_token_flex": 0.000135, + "input_cost_per_token_priority": 0.000153, + "litellm_provider": "gemini", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.0006299999999999999, + "output_cost_per_reasoning_token": 0.00045000000000000004, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_200k_tokens": 0.0008100000000000001, + "output_cost_per_token_flex": 0.00022500000000000002, + "output_cost_per_token_priority": 0.000243, + "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/gemini-3.8-flash": { + "cache_read_input_token_cost": 8e-06, + "input_cost_per_audio_token": 0.00048, + "input_cost_per_token": 8e-05, + "input_cost_per_token_above_200k_tokens": 0.00064, + "input_cost_per_token_flex": 0.00012, + "input_cost_per_token_priority": 0.000136, + "litellm_provider": "gemini", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00056, + "output_cost_per_reasoning_token": 0.0004, + "output_cost_per_token": 0.00016, + "output_cost_per_token_above_200k_tokens": 0.00072, + "output_cost_per_token_flex": 0.0002, + "output_cost_per_token_priority": 0.000216, + "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" + }, + "gpt-5.3-codex": { + "cache_read_input_token_cost": 3e-06, + "input_cost_per_token": 3.0000000000000004e-05, + "input_cost_per_token_above_200k_tokens": 0.00024000000000000003, + "input_cost_per_token_flex": 4.5e-05, + "input_cost_per_token_priority": 5.1e-05, + "litellm_provider": "openai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_reasoning_token": 0.00015000000000000001, + "output_cost_per_token": 6.000000000000001e-05, + "output_cost_per_token_above_200k_tokens": 0.00027, + "output_cost_per_token_flex": 7.500000000000001e-05, + "output_cost_per_token_priority": 8.099999999999999e-05, + "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 + }, + "gpt-5.4-mini": { + "cache_creation_input_token_cost": 0.00012, + "cache_creation_input_token_cost_above_1hr": 0.00016, + "cache_read_input_token_cost": 4e-06, + "input_cost_per_audio_token": 0.00024, + "input_cost_per_token": 4e-05, + "input_cost_per_token_above_200k_tokens": 0.00032, + "input_cost_per_token_flex": 6e-05, + "input_cost_per_token_priority": 6.8e-05, + "litellm_provider": "openai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00028, + "output_cost_per_reasoning_token": 0.0002, + "output_cost_per_token": 8e-05, + "output_cost_per_token_above_200k_tokens": 0.00036, + "output_cost_per_token_flex": 0.0001, + "output_cost_per_token_priority": 0.000108, + "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 + }, + "gpt-5.5-pro": { + "cache_read_input_token_cost": 2e-06, + "input_cost_per_token": 2e-05, + "input_cost_per_token_above_200k_tokens": 0.00016, + "input_cost_per_token_flex": 3e-05, + "input_cost_per_token_priority": 3.4e-05, + "litellm_provider": "openai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "output_cost_per_reasoning_token": 0.0001, + "output_cost_per_token": 4e-05, + "output_cost_per_token_above_200k_tokens": 0.00018, + "output_cost_per_token_flex": 5e-05, + "output_cost_per_token_priority": 5.4e-05, + "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 + }, + "gpt-5.6": { + "cache_creation_input_token_cost": 3e-05, + "cache_creation_input_token_cost_above_1hr": 4e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_audio_token": 6e-05, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_200k_tokens": 8e-05, + "input_cost_per_token_flex": 1.5e-05, + "input_cost_per_token_priority": 1.7e-05, + "litellm_provider": "openai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 7e-05, + "output_cost_per_reasoning_token": 5e-05, + "output_cost_per_token": 2e-05, + "output_cost_per_token_above_200k_tokens": 9e-05, + "output_cost_per_token_flex": 2.5e-05, + "output_cost_per_token_priority": 2.7e-05, + "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 + }, + "together_ai/moonshotai/Kimi-K3": { + "cache_creation_input_token_cost": 0.00030000000000000003, + "cache_creation_input_token_cost_above_1hr": 0.0004, + "cache_read_input_token_cost": 9.999999999999999e-06, + "input_cost_per_audio_token": 0.0006000000000000001, + "input_cost_per_token": 0.0001, + "input_cost_per_token_above_200k_tokens": 0.0008, + "input_cost_per_token_flex": 0.00015000000000000001, + "input_cost_per_token_priority": 0.00017, + "litellm_provider": "together_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.0006999999999999999, + "output_cost_per_reasoning_token": 0.0005, + "output_cost_per_token": 0.0002, + "output_cost_per_token_above_200k_tokens": 0.0009000000000000001, + "output_cost_per_token_flex": 0.00025, + "output_cost_per_token_priority": 0.00027, + "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 + }, + "together_ai/zai-org/GLM-5.3": { + "cache_creation_input_token_cost": 0.00033, + "cache_creation_input_token_cost_above_1hr": 0.00044, + "cache_read_input_token_cost": 1.1e-05, + "input_cost_per_audio_token": 0.00066, + "input_cost_per_token": 0.00011, + "input_cost_per_token_above_200k_tokens": 0.00088, + "input_cost_per_token_flex": 0.000165, + "input_cost_per_token_priority": 0.000187, + "litellm_provider": "together_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_audio_token": 0.00077, + "output_cost_per_reasoning_token": 0.00055, + "output_cost_per_token": 0.00022, + "output_cost_per_token_above_200k_tokens": 0.00099, + "output_cost_per_token_flex": 0.000275, + "output_cost_per_token_priority": 0.000297, + "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 + } +} diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index ad0914d455b..6b40e70125c 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -63,3 +63,5 @@ - {id: quota_management.spend_tracking.key_attribution.health_rows_keep_service_account, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [health_rows_keep_service_account], exercised_on: [chat_completions], source: "proxy/health_check.py", rationale: "A /health probe's spend row stays keyed by the literal litellm-internal-health-check service account rather than a hash of it, so health spend never appears as an unattributed key"} - {id: quota_management.spend_tracking.key_attribution.retrieve_batch_cost_joins_retrieving_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [retrieve_batch_cost_joins_retrieving_key], exercised_on: [batches], source: "proxy/batches_endpoints/endpoints.py", rationale: "The retrieve that first sees a batch in a terminal state prices it inline and writes its {provider_batch_id}_batch_cost row against the retrieving key, so the batch each run creates is one OpenAI fails at validation within seconds and the test retrieves it by its raw provider id with the same key until it is failed; a raw id is never owned by the CheckBatchCost poller, and the row must carry that key's token hash and alias"} - {id: quota_management.spend_tracking.key_attribution.poller_batch_cost_joins_creating_key, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_attribution, assertions: [poller_batch_cost_joins_creating_key], exercised_on: [batches], source: "enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py", rationale: "The CheckBatchCost poller bills a completed, positive-cost batch created through a unified id against the key that created it, a different writer from the inline retrieve. No test claims this cell yet: OpenAI's completion window is 24h and both e2e stacks boot a fresh Postgres per build, so a completed batch is out of one run's reach and the managed list never shows an earlier run's batch; the cell stays visible as a gap until a run can hand a completed batch to the poller"} +- {id: quota_management.spend_tracking.cost_matrix.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_matrix, assertions: [logs_cost], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/utils.py", rationale: "A scripted-usage call through the cost-map proxy bills every reported token kind at the deployment's test-map rate (input, output, cache read, 5m/1h cache write, reasoning, audio, above-threshold tiers, flex/priority service tiers, web search, response-model override) and lands on the row's cost_breakdown, streamed or not"} +- {id: quota_management.spend_tracking.scripted_wire.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: scripted_wire, assertions: [logs_cost], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/utils.py", rationale: "Each provider wire shape (openai chat, responses, anthropic messages, gemini generateContent, together, fireworks) parses usage into the same spend components: the gross input cost is fresh tokens at the input rate plus each cache/audio component at its own rate, streamed anthropic included"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 896cb3e7efe..a891d9dcba2 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -145,6 +145,22 @@ WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY" MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK" PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK" REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS" +# The cost_calculation suite needs a proxy booted with LITELLM_MODEL_COST_MAP_URL +# pointing at tests/e2e/cost_map.json (its whole map is test-owned rates) plus a +# scripted-provider sidecar; deselected unless the opt-in env var is set. +COST_MAP_OPT_IN_ENV = "E2E_COST_MAP_STACK" +# Base URL of the proxy running the test cost map. Defaults to the shared proxy +# so a local run only has to set the opt-in and boot the proxy accordingly. +COST_MAP_PROXY_URL = os.environ.get("E2E_COST_MAP_PROXY_URL", PROXY_BASE_URL).rstrip("/") +# Where the test runner reaches the scripted-provider sidecar's control API. +SCRIPTED_PROVIDER_CONTROL_URL = os.environ.get( + "E2E_SCRIPTED_PROVIDER_CONTROL_URL", "http://127.0.0.1:9100" +).rstrip("/") +# The api_base root deployments register with: how the proxy (possibly in +# another container) reaches the sidecar's provider wire. +SCRIPTED_PROVIDER_PROXY_BASE = os.environ.get( + "E2E_SCRIPTED_PROVIDER_PROXY_BASE", SCRIPTED_PROVIDER_CONTROL_URL +).rstrip("/") ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index 1fdd3bd28ad..7d37bcc6d3e 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -11,3 +11,4 @@ markers = managed_files: needs a proxy running with require_managed_files enabled; deselected unless E2E_MANAGED_FILES_STACK is set prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set + cost_map_stack: needs a proxy whose whole cost map is tests/e2e/cost_map.json (LITELLM_MODEL_COST_MAP_URL) plus a scripted-provider sidecar; deselected unless E2E_COST_MAP_STACK is set From 269afbe382df06d33780571a40a55e527afea2b7 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 23:28:44 +0000 Subject: [PATCH 016/251] test(e2e): apply review nits to cost calculation suite Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/conftest.py | 28 +- tests/e2e/cost_calculation/cost_matrix.py | 174 ++-- tests/e2e/cost_calculation/scripted_client.py | 12 +- .../e2e/cost_calculation/scripted_provider.py | 803 +++++++++++------- .../test_token_pricing_e2e.py | 17 +- .../cost_calculation/test_wire_formats_e2e.py | 43 +- 6 files changed, 620 insertions(+), 457 deletions(-) diff --git a/tests/e2e/cost_calculation/conftest.py b/tests/e2e/cost_calculation/conftest.py index 1bba3d50e1d..345ca26f7e3 100644 --- a/tests/e2e/cost_calculation/conftest.py +++ b/tests/e2e/cost_calculation/conftest.py @@ -13,7 +13,7 @@ from __future__ import annotations import importlib.util import sys -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass from pathlib import Path from types import ModuleType @@ -34,16 +34,16 @@ def _load_cost_rows() -> ModuleType: """Load quota_management/spend_tracking/cost_rows.py by path (the e2e tree has no package layout), the same trick the mcp suite uses for logging/datadog_reader.py.""" - path = ( + path: Final = ( Path(__file__).resolve().parent.parent / "quota_management" / "spend_tracking" / "cost_rows.py" ) - name = "e2e_spend_tracking_cost_rows" - spec = importlib.util.spec_from_file_location(name, path) + name: Final = "e2e_spend_tracking_cost_rows" + spec: Final = importlib.util.spec_from_file_location(name, path) assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) + module: Final = importlib.util.module_from_spec(spec) sys.modules[name] = module spec.loader.exec_module(module) return module @@ -59,7 +59,7 @@ class SpendCostBreakdown(Protocol): total_cost: float | None service_tier: str | None - def model_dump(self) -> dict[str, object]: ... + def model_dump(self) -> Mapping[str, object]: ... class SpendRowMetadata(Protocol): @@ -89,7 +89,9 @@ class CostRowsModule(Protocol): ] -cost_rows: Final[CostRowsModule] = cast(CostRowsModule, _load_cost_rows()) +cost_rows: Final[CostRowsModule] = cast( # cast-ok: cost_rows.py is loaded by path, so basedpyright has no importable name for it; its surface is declared in CostRowsModule + CostRowsModule, _load_cost_rows() +) @dataclass(frozen=True, slots=True) @@ -101,7 +103,7 @@ class CostCalcClient: @pytest.fixture(scope="session") def client() -> CostCalcClient: - proxy = build_proxy_client( + proxy: Final = build_proxy_client( base_url=COST_MAP_PROXY_URL, control_plane_base_url=COST_MAP_PROXY_URL, replica_urls=(COST_MAP_PROXY_URL,), @@ -118,18 +120,18 @@ def register_scenario_deployment( ) -> tuple[str, ScenarioHandle]: """Register the case's scenario on the sidecar plus a deployment pointed at it; both are torn down by ``resources``. Returns the callable model_name.""" - scenario: Scenario = case.scenario( + scenario: Final[Scenario] = case.scenario( scenario_id=f"sc-{marker}", model=model, text=f"scripted answer {marker}" ) - handle = register_scenario(scenario) + handle: Final = register_scenario(scenario) resources.defer(lambda: delete_scenario(handle)) - model_name = f"{model.model_name}-{marker}" - model_id = client.proxy.register_model( + model_name: Final = f"{model.model_name}-{marker}" + model_id: Final = client.proxy.register_model( ModelNewBody( model_name=model_name, litellm_params=LiteLLMParamsBody( model=model.litellm_model, - api_key="sk-scripted-provider", + api_key=model.api_key, api_base=handle.api_base(), ), model_info=ModelInfoBody(), diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index bc466d7d823..e8b1d249559 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -18,9 +18,11 @@ creation), the case is absent from the matrix rather than silently zero. from __future__ import annotations import json +from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path -from typing import Final, Literal +from types import MappingProxyType +from typing import Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict, TypeAdapter @@ -64,8 +66,8 @@ class CostMapEntry(BaseModel): _COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, CostMapEntry]) -_COST_MAP: Final[dict[str, CostMapEntry]] = _COST_MAP_ADAPTER.validate_python( - json.loads(COST_MAP_PATH.read_text()) +_COST_MAP: Final[Mapping[str, CostMapEntry]] = MappingProxyType( + _COST_MAP_ADAPTER.validate_python(json.loads(COST_MAP_PATH.read_text())) ) TIER_THRESHOLD_TOKENS: Final = 200_000 @@ -109,7 +111,7 @@ class FrontierModel: # Response-model override targets: emit a sibling's bare provider-facing name so # the biller's provider-prefixed lookup lands on that sibling's map key. -_OVERRIDE_MODELS: Final[dict[str, str]] = { +_OVERRIDE_MODELS: Final[Mapping[str, str]] = MappingProxyType({ "gpt-5.6": "gpt-5.4-mini", "gpt-5.5-pro": "gpt-5.3-codex", "gpt-5.3-codex": "gpt-5.5-pro", @@ -124,9 +126,9 @@ _OVERRIDE_MODELS: Final[dict[str, str]] = { "fireworks_ai/kimi-k3": "qwen3p8-max", "fireworks_ai/qwen3p8-max": "kimi-k3", "fireworks_ai/deepseek-v4p1-flash": "kimi-k3", -} +}) -_OVERRIDE_MAP_KEYS: Final[dict[str, str]] = { +_OVERRIDE_MAP_KEYS: Final[Mapping[str, str]] = MappingProxyType({ "gpt-5.4-mini": "gpt-5.4-mini", "gpt-5.6": "gpt-5.6", "gpt-5.3-codex": "gpt-5.3-codex", @@ -139,7 +141,7 @@ _OVERRIDE_MAP_KEYS: Final[dict[str, str]] = { "moonshotai/Kimi-K3": "together_ai/moonshotai/Kimi-K3", "qwen3p8-max": "fireworks_ai/qwen3p8-max", "kimi-k3": "fireworks_ai/kimi-k3", -} +}) _FRONTIER_SPECS: Final[tuple[tuple[str, str, Wire], ...]] = ( @@ -176,7 +178,7 @@ def _frontier() -> tuple[FrontierModel, ...]: FRONTIER_MODELS: Final[tuple[FrontierModel, ...]] = _frontier() # Token kinds each wire can report, gating which pricing cases apply. -_WIRE_CAPS: Final[dict[str, frozenset[str]]] = { +_WIRE_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType({ "openai_chat": frozenset( { "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", @@ -205,9 +207,9 @@ _WIRE_CAPS: Final[dict[str, frozenset[str]]] = { "web_search", "response_model", "absent_usage", } ), -} +}) -CaseName = Literal[ +CaseName: TypeAlias = Literal[ "basic", "cache_read", "cache_write_5m", @@ -260,7 +262,7 @@ _BASIC_USAGE: Final = ScriptedUsage(fresh_input_tokens=120, output_tokens=40) def _web_search_case(model: FrontierModel) -> Case: - counts_exactly = model.wire in ("openai_responses", "anthropic_messages", "gemini_generate") + counts_exactly: Final = model.wire in ("openai_responses", "anthropic_messages", "gemini_generate") return Case( name="web_search", usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, web_search_calls=3), @@ -269,26 +271,24 @@ def _web_search_case(model: FrontierModel) -> Case: def cases_for(model: FrontierModel) -> tuple[Case, ...]: - rates = model.rates - caps = _WIRE_CAPS[model.wire] - cases: list[Case] = [Case(name="basic", usage=_BASIC_USAGE)] - if rates.cache_read_input_token_cost is not None and "cache_read" in caps: - cases.append( + rates: Final = model.rates + caps: Final = _WIRE_CAPS[model.wire] + candidates: Final[tuple[Case | None, ...]] = ( + Case(name="basic", usage=_BASIC_USAGE), + ( Case(name="cache_read", usage=ScriptedUsage(fresh_input_tokens=100, cache_read_tokens=50, output_tokens=30)) - ) - if rates.cache_creation_input_token_cost is not None and "cache_write_5m" in caps: - cases.append( + if rates.cache_read_input_token_cost is not None and "cache_read" in caps + else None + ), + ( Case( name="cache_write_5m", usage=ScriptedUsage(fresh_input_tokens=90, cache_write_5m_tokens=60, output_tokens=30), ) - ) - if ( - rates.cache_creation_input_token_cost_above_1hr is not None - and rates.cache_creation_input_token_cost is not None - and "cache_write_1h" in caps - ): - cases.append( + if rates.cache_creation_input_token_cost is not None and "cache_write_5m" in caps + else None + ), + ( Case( name="cache_write_1h", usage=ScriptedUsage( @@ -298,52 +298,61 @@ def cases_for(model: FrontierModel) -> tuple[Case, ...]: output_tokens=30, ), ) - ) - if rates.output_cost_per_reasoning_token is not None and "reasoning" in caps: - cases.append( + if ( + rates.cache_creation_input_token_cost_above_1hr is not None + and rates.cache_creation_input_token_cost is not None + and "cache_write_1h" in caps + ) + else None + ), + ( Case( name="reasoning", usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, reasoning_tokens=70), ) - ) - if ( - rates.input_cost_per_audio_token is not None - and rates.output_cost_per_audio_token is not None - and "audio" in caps - ): - cases.append( + if rates.output_cost_per_reasoning_token is not None and "reasoning" in caps + else None + ), + ( Case( name="audio", usage=ScriptedUsage( fresh_input_tokens=100, audio_input_tokens=25, output_tokens=30, audio_output_tokens=15 ), ) - ) - if ( - rates.input_cost_per_token_above_200k_tokens is not None - and rates.output_cost_per_token_above_200k_tokens is not None - ): - cases.append( + if ( + rates.input_cost_per_audio_token is not None + and rates.output_cost_per_audio_token is not None + and "audio" in caps + ) + else None + ), + ( Case( name="tiered", usage=ScriptedUsage( fresh_input_tokens=TIER_THRESHOLD_TOKENS + 1, output_tokens=30 ), ) - ) - if rates.input_cost_per_token_flex is not None and rates.output_cost_per_token_flex is not None: - cases.append( + if ( + rates.input_cost_per_token_above_200k_tokens is not None + and rates.output_cost_per_token_above_200k_tokens is not None + ) + else None + ), + ( Case(name="service_tier_flex", usage=_BASIC_USAGE, service_tier="flex") - ) - if rates.input_cost_per_token_priority is not None and rates.output_cost_per_token_priority is not None: - cases.append( + if rates.input_cost_per_token_flex is not None and rates.output_cost_per_token_flex is not None + else None + ), + ( Case(name="service_tier_priority", usage=_BASIC_USAGE, service_tier="priority") - ) - if rates.search_context_cost_per_query is not None and "web_search" in caps: - cases.append(_web_search_case(model)) - cases.append(Case(name="stream", usage=_BASIC_USAGE, stream=True)) - if "absent_usage" in caps: - cases.append( + if rates.input_cost_per_token_priority is not None and rates.output_cost_per_token_priority is not None + else None + ), + _web_search_case(model) if rates.search_context_cost_per_query is not None and "web_search" in caps else None, + Case(name="stream", usage=_BASIC_USAGE, stream=True), + ( Case( name="stream_no_usage", usage=_BASIC_USAGE, @@ -355,10 +364,16 @@ def cases_for(model: FrontierModel) -> tuple[Case, ...]: # wires recount tokens proxy-side and bill a nonzero amount. expect_zero_bill=model.wire == "openai_responses", ) - ) - if "response_model" in caps: - cases.append(Case(name="response_model_override", usage=_BASIC_USAGE, response_model_override=True)) - return tuple(cases) + if "absent_usage" in caps + else None + ), + ( + Case(name="response_model_override", usage=_BASIC_USAGE, response_model_override=True) + if "response_model" in caps + else None + ), + ) + return tuple(case for case in candidates if case is not None) @dataclass(frozen=True, slots=True) @@ -387,38 +402,41 @@ def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: to the tier's variants, falling back to the base rate when a variant is unset -- mirroring _get_token_base_cost in litellm's cost calculator. """ - rates = model.override_rates if case.response_model_override else model.rates - u = case.usage - prompt_tokens = ( + rates: Final = model.override_rates if case.response_model_override else model.rates + u: Final = case.usage + prompt_tokens: Final = ( u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens + u.audio_input_tokens ) - tiered = prompt_tokens > TIER_THRESHOLD_TOKENS - in_rate = rates.input_cost_per_token or 0.0 - out_rate = rates.output_cost_per_token or 0.0 - if case.service_tier == "flex": - in_rate = rates.input_cost_per_token_flex or in_rate - out_rate = rates.output_cost_per_token_flex or out_rate - if case.service_tier == "priority": - in_rate = rates.input_cost_per_token_priority or in_rate - out_rate = rates.output_cost_per_token_priority or out_rate - if tiered: - in_rate = rates.input_cost_per_token_above_200k_tokens or in_rate - out_rate = rates.output_cost_per_token_above_200k_tokens or out_rate - input_cost = ( + tiered: Final = prompt_tokens > TIER_THRESHOLD_TOKENS + in_rate: Final = ( + (rates.input_cost_per_token_above_200k_tokens if tiered else None) + or (rates.input_cost_per_token_priority if case.service_tier == "priority" else None) + or (rates.input_cost_per_token_flex if case.service_tier == "flex" else None) + or rates.input_cost_per_token + or 0.0 + ) + out_rate: Final = ( + (rates.output_cost_per_token_above_200k_tokens if tiered else None) + or (rates.output_cost_per_token_priority if case.service_tier == "priority" else None) + or (rates.output_cost_per_token_flex if case.service_tier == "flex" else None) + or rates.output_cost_per_token + or 0.0 + ) + input_cost: Final = ( u.fresh_input_tokens * in_rate + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) + u.cache_write_5m_tokens * (rates.cache_creation_input_token_cost or 0.0) + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or 0.0) + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) ) - output_cost = ( + output_cost: Final = ( u.output_tokens * out_rate + u.reasoning_tokens * (rates.output_cost_per_reasoning_token or out_rate) + u.audio_output_tokens * (rates.output_cost_per_audio_token or out_rate) ) - search = rates.search_context_cost_per_query - tool_cost = case.billed_web_search_calls * ( + search: Final = rates.search_context_cost_per_query + tool_cost: Final = case.billed_web_search_calls * ( search.search_context_size_medium if search and search.search_context_size_medium else 0.0 ) return ExpectedCost(input_cost=input_cost, output_cost=output_cost, tool_cost=tool_cost) @@ -432,7 +450,7 @@ def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: """(prompt_tokens, completion_tokens) the spend row should carry, per the wire's normalization: Anthropic folds cache read/write into prompt_tokens, everyone else reports the totals the wire emitted.""" - u = case.usage + u: Final = case.usage if model.wire == "anthropic_messages": return ( u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens, diff --git a/tests/e2e/cost_calculation/scripted_client.py b/tests/e2e/cost_calculation/scripted_client.py index dceec02630a..9dbf9c98986 100644 --- a/tests/e2e/cost_calculation/scripted_client.py +++ b/tests/e2e/cost_calculation/scripted_client.py @@ -12,6 +12,7 @@ from e2e_config import SCRIPTED_PROVIDER_CONTROL_URL, SCRIPTED_PROVIDER_PROXY_BA from e2e_http import URL, NoBody, unwrap, post from e2e_http import delete as http_delete from scripted_provider import ( + WIRE_MOUNTS, Scenario, ScenarioDeleted, ScenarioRegistered, @@ -29,19 +30,12 @@ class ScenarioHandle: return f"{self.proxy_base}/{self.scenario_id}/{self._mount()}" def _mount(self) -> str: - return { - "openai_chat": "openai", - "openai_responses": "openai", - "anthropic_messages": "anthropic", - "gemini_generate": "gemini", - "together_chat": "together", - "fireworks_chat": "fireworks", - }[self.wire] + return WIRE_MOUNTS[self.wire] def register_scenario(scenario: Scenario) -> ScenarioHandle: """POST the scenario to the sidecar's control API and return its handle.""" - result = unwrap( + result: Final = unwrap( post( URL(f"{SCRIPTED_PROVIDER_CONTROL_URL}/_scenarios"), headers=NoBody(), diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py index 93a6f49ec25..f1deafd1bc5 100644 --- a/tests/e2e/cost_calculation/scripted_provider.py +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -31,14 +31,16 @@ import json import sys import threading import time +from collections.abc import Mapping from dataclasses import dataclass from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from typing import Final, Literal +from types import MappingProxyType +from typing import Final, Literal, TypeAlias from urllib.parse import urlsplit from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError -Wire = Literal[ +Wire: TypeAlias = Literal[ "openai_chat", "openai_responses", "anthropic_messages", @@ -47,17 +49,19 @@ Wire = Literal[ "fireworks_chat", ] -_WIRE_MOUNTS: Final[dict[str, str]] = { - "openai_chat": "openai", - "openai_responses": "openai", - "anthropic_messages": "anthropic", - "gemini_generate": "gemini", - "together_chat": "together", - "fireworks_chat": "fireworks", -} +WIRE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( + { + "openai_chat": "openai", + "openai_responses": "openai", + "anthropic_messages": "anthropic", + "gemini_generate": "gemini", + "together_chat": "together", + "fireworks_chat": "fireworks", + } +) -StreamUsage = Literal["final_chunk", "absent"] -ServiceTier = Literal["flex", "priority"] +StreamUsage: TypeAlias = Literal["final_chunk", "absent"] +ServiceTier: TypeAlias = Literal["flex", "priority"] class ScriptedUsage(BaseModel): @@ -106,7 +110,7 @@ class Scenario(BaseModel): @property def mount(self) -> str: - return _WIRE_MOUNTS[self.wire] + return WIRE_MOUNTS[self.wire] class ScenarioRegistered(BaseModel): @@ -128,369 +132,494 @@ class RenderedResponse: body: bytes -def _json_bytes(payload: dict[str, object]) -> bytes: - return json.dumps(payload).encode("utf-8") +def _jobj(*pairs: tuple[str, object]) -> Mapping[str, object]: + """A JSON object payload built in one shot and frozen.""" + return MappingProxyType(dict(pairs)) -def _sse(events: tuple[tuple[str | None, dict[str, object] | str], ...]) -> bytes: - frames: list[str] = [] - for event_name, data in events: - head = f"event: {event_name}\n" if event_name is not None else "" - payload = data if isinstance(data, str) else json.dumps(data) - frames.append(f"{head}data: {payload}\n\n") - return "".join(frames).encode("utf-8") +def _jobj_opt(*pairs: tuple[str, object] | None) -> Mapping[str, object]: + """``_jobj`` where a ``None`` pair means the field is absent.""" + return MappingProxyType(dict(pair for pair in pairs if pair is not None)) + + +def _json_bytes(payload: Mapping[str, object]) -> bytes: + return json.dumps(payload, default=dict).encode("utf-8") + + +def _sse_frame(event_name: str | None, data: Mapping[str, object] | str) -> str: + head: Final = f"event: {event_name}\n" if event_name is not None else "" + payload: Final = data if isinstance(data, str) else json.dumps(data, default=dict) + return f"{head}data: {payload}\n\n" + + +def _sse(events: tuple[tuple[str | None, Mapping[str, object] | str], ...]) -> bytes: + return "".join(_sse_frame(event_name, data) for event_name, data in events).encode("utf-8") # ---------- per-wire usage shapes ---------- -def _openai_usage(u: ScriptedUsage) -> dict[str, object]: - prompt_tokens = ( +def _openai_usage(u: ScriptedUsage) -> Mapping[str, object]: + prompt_tokens: Final = ( u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + u.cache_write_1h_tokens + u.audio_input_tokens ) - completion_tokens = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens - prompt_details: dict[str, object] = {} - if u.cache_read_tokens: - prompt_details["cached_tokens"] = u.cache_read_tokens - if u.cache_write_5m_tokens or u.cache_write_1h_tokens: - prompt_details["cache_write_tokens"] = u.cache_write_5m_tokens + u.cache_write_1h_tokens - prompt_details["cache_creation_token_details"] = { - "ephemeral_5m_input_tokens": u.cache_write_5m_tokens, - "ephemeral_1h_input_tokens": u.cache_write_1h_tokens, - } - if u.audio_input_tokens: - prompt_details["audio_tokens"] = u.audio_input_tokens - completion_details: dict[str, object] = {} - if u.reasoning_tokens: - completion_details["reasoning_tokens"] = u.reasoning_tokens - if u.audio_output_tokens: - completion_details["audio_tokens"] = u.audio_output_tokens - usage: dict[str, object] = { - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - "total_tokens": prompt_tokens + completion_tokens, - } - if prompt_details: - usage["prompt_tokens_details"] = prompt_details - if completion_details: - usage["completion_tokens_details"] = completion_details - return usage + completion_tokens: Final = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + prompt_details: Final = _jobj_opt( + ("cached_tokens", u.cache_read_tokens) if u.cache_read_tokens else None, + ( + ("cache_write_tokens", u.cache_write_5m_tokens + u.cache_write_1h_tokens) + if u.cache_write_5m_tokens or u.cache_write_1h_tokens + else None + ), + ( + ( + "cache_creation_token_details", + _jobj( + ("ephemeral_5m_input_tokens", u.cache_write_5m_tokens), + ("ephemeral_1h_input_tokens", u.cache_write_1h_tokens), + ), + ) + if u.cache_write_5m_tokens or u.cache_write_1h_tokens + else None + ), + ("audio_tokens", u.audio_input_tokens) if u.audio_input_tokens else None, + ) + completion_details: Final = _jobj_opt( + ("reasoning_tokens", u.reasoning_tokens) if u.reasoning_tokens else None, + ("audio_tokens", u.audio_output_tokens) if u.audio_output_tokens else None, + ) + return _jobj_opt( + ("prompt_tokens", prompt_tokens), + ("completion_tokens", completion_tokens), + ("total_tokens", prompt_tokens + completion_tokens), + ("prompt_tokens_details", prompt_details) if prompt_details else None, + ("completion_tokens_details", completion_details) if completion_details else None, + ) -def _anthropic_usage(u: ScriptedUsage) -> dict[str, object]: +def _anthropic_usage(u: ScriptedUsage) -> Mapping[str, object]: # Anthropic reports uncached-only input_tokens; cache reads and writes ride # top-level fields, with the 5m/1h write split under cache_creation. - usage: dict[str, object] = { - "input_tokens": u.fresh_input_tokens, - "output_tokens": u.output_tokens, - } - if u.cache_read_tokens: - usage["cache_read_input_tokens"] = u.cache_read_tokens - if u.cache_write_5m_tokens or u.cache_write_1h_tokens: - usage["cache_creation_input_tokens"] = u.cache_write_5m_tokens + u.cache_write_1h_tokens - usage["cache_creation"] = { - "ephemeral_5m_input_tokens": u.cache_write_5m_tokens, - "ephemeral_1h_input_tokens": u.cache_write_1h_tokens, - } - if u.web_search_calls: - usage["server_tool_use"] = {"web_search_requests": u.web_search_calls} - return usage + return _jobj_opt( + ("input_tokens", u.fresh_input_tokens), + ("output_tokens", u.output_tokens), + ("cache_read_input_tokens", u.cache_read_tokens) if u.cache_read_tokens else None, + ( + ("cache_creation_input_tokens", u.cache_write_5m_tokens + u.cache_write_1h_tokens) + if u.cache_write_5m_tokens or u.cache_write_1h_tokens + else None + ), + ( + ( + "cache_creation", + _jobj( + ("ephemeral_5m_input_tokens", u.cache_write_5m_tokens), + ("ephemeral_1h_input_tokens", u.cache_write_1h_tokens), + ), + ) + if u.cache_write_5m_tokens or u.cache_write_1h_tokens + else None + ), + ( + ("server_tool_use", _jobj(("web_search_requests", u.web_search_calls))) + if u.web_search_calls + else None + ), + ) -def _gemini_usage(u: ScriptedUsage) -> dict[str, object]: +def _gemini_usage(u: ScriptedUsage) -> Mapping[str, object]: # promptTokenCount carries the cached count inside it; TEXT modality is the # cached-inclusive text count so litellm's implicit-caching subtraction lands # on the fresh figure. candidatesTokenCount includes reasoning + audio. - prompt_tokens = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens - candidates = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens - usage: dict[str, object] = { - "promptTokenCount": prompt_tokens, - "candidatesTokenCount": candidates, - "totalTokenCount": prompt_tokens + candidates, - } - if u.cache_read_tokens: - usage["cachedContentTokenCount"] = u.cache_read_tokens - if u.reasoning_tokens: - usage["thoughtsTokenCount"] = u.reasoning_tokens - prompt_details = [{"modality": "TEXT", "tokenCount": u.fresh_input_tokens + u.cache_read_tokens}] - if u.audio_input_tokens: - prompt_details.append({"modality": "AUDIO", "tokenCount": u.audio_input_tokens}) - usage["promptTokensDetails"] = prompt_details - if u.audio_output_tokens: - usage["candidatesTokensDetails"] = [ - {"modality": "TEXT", "tokenCount": u.output_tokens + u.reasoning_tokens}, - {"modality": "AUDIO", "tokenCount": u.audio_output_tokens}, - ] - return usage + prompt_tokens: Final = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens + candidates: Final = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + return _jobj_opt( + ("promptTokenCount", prompt_tokens), + ("candidatesTokenCount", candidates), + ("totalTokenCount", prompt_tokens + candidates), + ("cachedContentTokenCount", u.cache_read_tokens) if u.cache_read_tokens else None, + ("thoughtsTokenCount", u.reasoning_tokens) if u.reasoning_tokens else None, + ( + "promptTokensDetails", + ( + _jobj(("modality", "TEXT"), ("tokenCount", u.fresh_input_tokens + u.cache_read_tokens)), + *( + (_jobj(("modality", "AUDIO"), ("tokenCount", u.audio_input_tokens)),) + if u.audio_input_tokens + else () + ), + ), + ), + ( + ( + "candidatesTokensDetails", + ( + _jobj(("modality", "TEXT"), ("tokenCount", u.output_tokens + u.reasoning_tokens)), + _jobj(("modality", "AUDIO"), ("tokenCount", u.audio_output_tokens)), + ), + ) + if u.audio_output_tokens + else None + ), + ) -def _responses_usage(u: ScriptedUsage) -> dict[str, object]: - input_tokens = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens - output_tokens = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens - usage: dict[str, object] = { - "input_tokens": input_tokens, - "output_tokens": output_tokens, - "total_tokens": input_tokens + output_tokens, - } - input_details: dict[str, object] = {} - if u.cache_read_tokens: - input_details["cached_tokens"] = u.cache_read_tokens - if input_details: - usage["input_tokens_details"] = input_details - if u.reasoning_tokens: - usage["output_tokens_details"] = {"reasoning_tokens": u.reasoning_tokens} - return usage +def _responses_usage(u: ScriptedUsage) -> Mapping[str, object]: + input_tokens: Final = u.fresh_input_tokens + u.cache_read_tokens + u.audio_input_tokens + output_tokens: Final = u.output_tokens + u.reasoning_tokens + u.audio_output_tokens + input_details: Final = _jobj_opt( + ("cached_tokens", u.cache_read_tokens) if u.cache_read_tokens else None, + ) + return _jobj_opt( + ("input_tokens", input_tokens), + ("output_tokens", output_tokens), + ("total_tokens", input_tokens + output_tokens), + ("input_tokens_details", input_details) if input_details else None, + ( + ("output_tokens_details", _jobj(("reasoning_tokens", u.reasoning_tokens))) + if u.reasoning_tokens + else None + ), + ) # ---------- per-wire responses ---------- -def _openai_message(scenario: Scenario) -> dict[str, object]: - message: dict[str, object] = {"role": "assistant", "content": scenario.output.text} - if scenario.usage.web_search_calls: - message["annotations"] = [ - { - "type": "url_citation", - "url_citation": { - "url": "https://scripted.example/source", - "title": "scripted source", - "start_index": 0, - "end_index": 1, - }, - } - for _ in range(scenario.usage.web_search_calls) - ] - return message +def _openai_message(scenario: Scenario) -> Mapping[str, object]: + return _jobj_opt( + ("role", "assistant"), + ("content", scenario.output.text), + ( + ( + "annotations", + tuple( + _jobj( + ("type", "url_citation"), + ( + "url_citation", + _jobj( + ("url", "https://scripted.example/source"), + ("title", "scripted source"), + ("start_index", 0), + ("end_index", 1), + ), + ), + ) + for _ in range(scenario.usage.web_search_calls) + ), + ) + if scenario.usage.web_search_calls + else None + ), + ) -def _openai_chat_body(scenario: Scenario, requested_model: str) -> dict[str, object]: - body: dict[str, object] = { - "id": f"chatcmpl-{scenario.scenario_id}", - "object": "chat.completion", - "created": int(time.time()), - "model": scenario.output.response_model or requested_model, - "choices": [ - { - "index": 0, - "message": _openai_message(scenario), - "finish_reason": scenario.output.finish_reason, - } - ], - "usage": _openai_usage(scenario.usage), - } - if scenario.service_tier is not None: - body["service_tier"] = scenario.service_tier - if scenario.output.provider_cost is not None: - body["cost"] = scenario.output.provider_cost - return body +def _openai_chat_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + return _jobj_opt( + ("id", f"chatcmpl-{scenario.scenario_id}"), + ("object", "chat.completion"), + ("created", int(time.time())), + ("model", scenario.output.response_model or requested_model), + ( + "choices", + ( + _jobj( + ("index", 0), + ("message", _openai_message(scenario)), + ("finish_reason", scenario.output.finish_reason), + ), + ), + ), + ("usage", _openai_usage(scenario.usage)), + ("service_tier", scenario.service_tier) if scenario.service_tier is not None else None, + ("cost", scenario.output.provider_cost) if scenario.output.provider_cost is not None else None, + ) -def _openai_chunk(scenario: Scenario, requested_model: str, **kw: object) -> dict[str, object]: - chunk: dict[str, object] = { - "id": f"chatcmpl-{scenario.scenario_id}", - "object": "chat.completion.chunk", - "created": int(time.time()), - "model": scenario.output.response_model or requested_model, - } - chunk.update(kw) - return chunk +def _openai_chunk( + scenario: Scenario, + requested_model: str, + choices: tuple[Mapping[str, object], ...] = (), + usage: Mapping[str, object] | None = None, +) -> Mapping[str, object]: + return _jobj_opt( + ("id", f"chatcmpl-{scenario.scenario_id}"), + ("object", "chat.completion.chunk"), + ("created", int(time.time())), + ("model", scenario.output.response_model or requested_model), + ("choices", choices), + ("usage", usage), + ) def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: - _EMPTY_DELTA: Final[dict[str, object]] = {} - delta: dict[str, object] = {"role": "assistant", "content": scenario.output.text} - if scenario.usage.web_search_calls: - delta["annotations"] = _openai_message(scenario)["annotations"] - events: list[tuple[str | None, dict[str, object] | str]] = [ + delta: Final = _jobj_opt( + ("role", "assistant"), + ("content", scenario.output.text), ( - None, - _openai_chunk( - scenario, - requested_model, - choices=[{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}], - ), + ("annotations", _openai_message(scenario)["annotations"]) + if scenario.usage.web_search_calls + else None ), + ) + return _sse( ( - None, - _openai_chunk( - scenario, - requested_model, - choices=[{"index": 0, "delta": delta, "finish_reason": None}], + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=(_jobj(("index", 0), ("delta", _jobj(("role", "assistant"))), ("finish_reason", None)),), + ), ), - ), - ( - None, - _openai_chunk( - scenario, - requested_model, - choices=[ - { - "index": 0, - "delta": _EMPTY_DELTA, - "finish_reason": scenario.output.finish_reason, - } - ], + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=(_jobj(("index", 0), ("delta", delta), ("finish_reason", None)),), + ), ), - ), - ] - if scenario.stream_usage == "final_chunk": - events.append( - (None, _openai_chunk(scenario, requested_model, choices=(), usage=_openai_usage(scenario.usage))) + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=( + _jobj( + ("index", 0), + ("delta", _jobj()), + ("finish_reason", scenario.output.finish_reason), + ), + ), + ), + ), + *( + ((None, _openai_chunk(scenario, requested_model, usage=_openai_usage(scenario.usage))),) + if scenario.stream_usage == "final_chunk" + else () + ), + (None, "[DONE]"), ) - events.append((None, "[DONE]")) - return _sse(tuple(events)) + ) -def _anthropic_body(scenario: Scenario, requested_model: str) -> dict[str, object]: - return { - "id": f"msg_{scenario.scenario_id}", - "type": "message", - "role": "assistant", - "model": scenario.output.response_model or requested_model, - "content": [{"type": "text", "text": scenario.output.text}], - "stop_reason": "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, - "usage": _anthropic_usage(scenario.usage), - } +def _anthropic_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + return _jobj( + ("id", f"msg_{scenario.scenario_id}"), + ("type", "message"), + ("role", "assistant"), + ("model", scenario.output.response_model or requested_model), + ("content", (_jobj(("type", "text"), ("text", scenario.output.text)),)), + ( + "stop_reason", + "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, + ), + ("usage", _anthropic_usage(scenario.usage)), + ) def _anthropic_sse(scenario: Scenario, requested_model: str) -> bytes: - emit_usage = scenario.stream_usage == "final_chunk" - input_usage = {k: v for k, v in _anthropic_usage(scenario.usage).items() if k != "output_tokens"} - message_start: dict[str, object] = { - "type": "message_start", - "message": { - "id": f"msg_{scenario.scenario_id}", - "type": "message", - "role": "assistant", - "model": scenario.output.response_model or requested_model, - "content": [], - "stop_reason": None, - **({"usage": input_usage} if emit_usage else {}), - }, - } - message_delta: dict[str, object] = { - "type": "message_delta", - "delta": { - "stop_reason": "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason - }, - **({"usage": {"output_tokens": scenario.usage.output_tokens}} if emit_usage else {}), - } + emit_usage: Final = scenario.stream_usage == "final_chunk" + input_usage: Final = _jobj( + *( + (key, value) + for key, value in _anthropic_usage(scenario.usage).items() + if key != "output_tokens" + ) + ) + message_start: Final = _jobj( + ("type", "message_start"), + ( + "message", + _jobj_opt( + ("id", f"msg_{scenario.scenario_id}"), + ("type", "message"), + ("role", "assistant"), + ("model", scenario.output.response_model or requested_model), + ("content", ()), + ("stop_reason", None), + ("usage", input_usage) if emit_usage else None, + ), + ), + ) + message_delta: Final = _jobj_opt( + ("type", "message_delta"), + ( + "delta", + _jobj( + ( + "stop_reason", + "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, + ) + ), + ), + ( + ("usage", _jobj(("output_tokens", scenario.usage.output_tokens))) + if emit_usage + else None + ), + ) return _sse( ( ("message_start", message_start), ( "content_block_start", - { - "type": "content_block_start", - "index": 0, - "content_block": {"type": "text", "text": ""}, - }, + _jobj( + ("type", "content_block_start"), + ("index", 0), + ("content_block", _jobj(("type", "text"), ("text", ""))), + ), ), ( "content_block_delta", - { - "type": "content_block_delta", - "index": 0, - "delta": {"type": "text_delta", "text": scenario.output.text}, - }, + _jobj( + ("type", "content_block_delta"), + ("index", 0), + ("delta", _jobj(("type", "text_delta"), ("text", scenario.output.text))), + ), ), - ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("content_block_stop", _jobj(("type", "content_block_stop"), ("index", 0))), ("message_delta", message_delta), - ("message_stop", {"type": "message_stop"}), + ("message_stop", _jobj(("type", "message_stop"))), ) ) -def _gemini_body(scenario: Scenario, requested_model: str) -> dict[str, object]: - candidate: dict[str, object] = { - "content": {"parts": [{"text": scenario.output.text}], "role": "model"}, - "finishReason": "STOP" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason.upper(), - "index": 0, - } - if scenario.usage.web_search_calls: - candidate["groundingMetadata"] = { - "webSearchQueries": [f"query {i}" for i in range(scenario.usage.web_search_calls)] - } - return { - "candidates": [candidate], - "usageMetadata": _gemini_usage(scenario.usage), - "modelVersion": scenario.output.response_model or requested_model, - } +def _gemini_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + return _jobj( + ( + "candidates", + ( + _jobj_opt( + ( + "content", + _jobj( + ("parts", (_jobj(("text", scenario.output.text)),)), + ("role", "model"), + ), + ), + ( + "finishReason", + "STOP" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason.upper(), + ), + ("index", 0), + ( + ( + "groundingMetadata", + _jobj( + ( + "webSearchQueries", + tuple(f"query {i}" for i in range(scenario.usage.web_search_calls)), + ) + ), + ) + if scenario.usage.web_search_calls + else None + ), + ), + ), + ), + ("usageMetadata", _gemini_usage(scenario.usage)), + ("modelVersion", scenario.output.response_model or requested_model), + ) def _gemini_sse(scenario: Scenario, requested_model: str) -> bytes: - first = _gemini_body(scenario, requested_model) - if scenario.stream_usage == "absent": - first = {k: v for k, v in first.items() if k != "usageMetadata"} - events: list[tuple[str | None, dict[str, object] | str]] = [(None, first)] - if scenario.stream_usage == "final_chunk": - events.append( - ( - None, - { - "candidates": [], - "usageMetadata": _gemini_usage(scenario.usage), - "modelVersion": scenario.output.response_model or requested_model, - }, - ) - ) - return _sse(tuple(events)) - - -def _responses_body(scenario: Scenario, requested_model: str) -> dict[str, object]: - output: list[dict[str, object]] = [ - {"type": "web_search_call", "id": f"ws_{i}", "status": "completed"} - for i in range(scenario.usage.web_search_calls) - ] - output.append( - { - "type": "message", - "id": f"msg_{scenario.scenario_id}", - "status": "completed", - "role": "assistant", - "content": [ - { - "type": "output_text", - "text": scenario.output.text, - "annotations": [], - } - ], - } + emit_usage: Final = scenario.stream_usage == "final_chunk" + first: Final = ( + _jobj(*((key, value) for key, value in _gemini_body(scenario, requested_model).items() if key != "usageMetadata")) + if scenario.stream_usage == "absent" + else _gemini_body(scenario, requested_model) + ) + return _sse( + ( + (None, first), + *( + ( + ( + None, + _jobj( + ("candidates", ()), + ("usageMetadata", _gemini_usage(scenario.usage)), + ("modelVersion", scenario.output.response_model or requested_model), + ), + ), + ) + if emit_usage + else () + ), + ) + ) + + +def _responses_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + return _jobj( + ("id", f"resp_{scenario.scenario_id}"), + ("object", "response"), + ("created_at", int(time.time())), + ("status", "completed"), + ("model", scenario.output.response_model or requested_model), + ( + "output", + ( + *( + _jobj(("type", "web_search_call"), ("id", f"ws_{i}"), ("status", "completed")) + for i in range(scenario.usage.web_search_calls) + ), + _jobj( + ("type", "message"), + ("id", f"msg_{scenario.scenario_id}"), + ("status", "completed"), + ("role", "assistant"), + ( + "content", + ( + _jobj( + ("type", "output_text"), + ("text", scenario.output.text), + ("annotations", ()), + ), + ), + ), + ), + ), + ), + ("usage", _responses_usage(scenario.usage)), ) - return { - "id": f"resp_{scenario.scenario_id}", - "object": "response", - "created_at": int(time.time()), - "status": "completed", - "model": scenario.output.response_model or requested_model, - "output": output, - "usage": _responses_usage(scenario.usage), - } def _responses_sse(scenario: Scenario, requested_model: str) -> bytes: - completed = _responses_body(scenario, requested_model) - if scenario.stream_usage == "absent": - completed = {k: v for k, v in completed.items() if k != "usage"} - created = {**completed, "status": "in_progress", "usage": None} + completed: Final = ( + _jobj(*((key, value) for key, value in _responses_body(scenario, requested_model).items() if key != "usage")) + if scenario.stream_usage == "absent" + else _responses_body(scenario, requested_model) + ) + created: Final = _jobj( + *((key, value) for key, value in completed.items() if key not in ("status", "usage")), + ("status", "in_progress"), + ("usage", None), + ) return _sse( ( - ("response.created", {"type": "response.created", "response": created}), + ("response.created", _jobj(("type", "response.created"), ("response", created))), ( "response.output_text.delta", - { - "type": "response.output_text.delta", - "item_id": f"msg_{scenario.scenario_id}", - "output_index": scenario.usage.web_search_calls, - "content_index": 0, - "delta": scenario.output.text, - }, + _jobj( + ("type", "response.output_text.delta"), + ("item_id", f"msg_{scenario.scenario_id}"), + ("output_index", scenario.usage.web_search_calls), + ("content_index", 0), + ("delta", scenario.output.text), + ), ), - ("response.completed", {"type": "response.completed", "response": completed}), + ("response.completed", _jobj(("type", "response.completed"), ("response", completed))), ) ) @@ -538,11 +667,11 @@ class _ScenarioStore: _REQUEST_BODY: Final = TypeAdapter(dict[str, object]) -def _request_body(body: bytes) -> dict[str, object]: +def _request_body(body: bytes) -> Mapping[str, object]: try: return _REQUEST_BODY.validate_json(body) except ValueError: - return {} + return MappingProxyType({}) def _request_wants_stream(path_tail: str, body: bytes) -> bool: @@ -554,52 +683,66 @@ def _request_wants_stream(path_tail: str, body: bytes) -> bool: def _request_model(body: bytes) -> str: - model = _request_body(body).get("model") + model: Final = _request_body(body).get("model") return model if isinstance(model, str) else "unknown" def handle_request(store: _ScenarioStore, method: str, raw_path: str, body: bytes) -> RenderedResponse: - path = urlsplit(raw_path).path - segments = [segment for segment in path.split("/") if segment] - if method == "GET" and segments == ["health"]: - return RenderedResponse(200, "application/json", _json_bytes({"status": "ok"})) + path: Final = urlsplit(raw_path).path + 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] == "_scenarios": if method == "POST" and len(segments) == 1: try: - scenario = Scenario.model_validate_json(body) + scenario: Final = Scenario.model_validate_json(body) except ValidationError as exc: - return RenderedResponse(400, "application/json", _json_bytes({"error": str(exc)})) + return RenderedResponse( + 400, "application/json", _json_bytes(_jobj(("error", str(exc)))) + ) store.put(scenario) - return RenderedResponse(200, "application/json", _json_bytes({"scenario_id": scenario.scenario_id})) - if method == "DELETE" and len(segments) == 2: - deleted = store.drop(segments[1]) return RenderedResponse( - 200 if deleted else 404, "application/json", _json_bytes({"deleted": deleted}) + 200, "application/json", _json_bytes(_jobj(("scenario_id", scenario.scenario_id))) ) - return RenderedResponse(404, "application/json", _json_bytes({"error": "unknown control route"})) + if method == "DELETE" and len(segments) == 2: + deleted: Final = store.drop(segments[1]) + return RenderedResponse( + 200 if deleted else 404, + "application/json", + _json_bytes(_jobj(("deleted", deleted))), + ) + return RenderedResponse( + 404, "application/json", _json_bytes(_jobj(("error", "unknown control route"))) + ) if len(segments) < 2 or method != "POST": - return RenderedResponse(404, "application/json", _json_bytes({"error": f"no route for {method} {path}"})) + return RenderedResponse( + 404, "application/json", _json_bytes(_jobj(("error", f"no route for {method} {path}"))) + ) scenario_id, mount = segments[0], segments[1] - scenario = store.get(scenario_id) - if scenario is None: - return RenderedResponse(404, "application/json", _json_bytes({"error": f"unknown scenario {scenario_id}"})) - if scenario.mount != mount: + 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({"error": f"scenario {scenario_id} is wire {scenario.wire}, not mount {mount}"}), + _json_bytes( + _jobj(("error", f"scenario {scenario_id} is wire {found.wire}, not mount {mount}")) + ), ) - tail = "/".join(segments[2:]) - return _render(scenario, stream=_request_wants_stream(tail, body), requested_model=_request_model(body)) + tail: Final = "/".join(segments[2:]) + return _render(found, stream=_request_wants_stream(tail, body), requested_model=_request_model(body)) class _ScriptedHandler(BaseHTTPRequestHandler): store: Final[_ScenarioStore] = _ScenarioStore() def _dispatch(self, method: str) -> None: - length = int(self.headers.get("content-length") or 0) - body = self.rfile.read(length) if length else b"" - rendered = handle_request(self.store, method, self.path, body) + length: Final = int(self.headers.get("content-length") or 0) + body: Final = self.rfile.read(length) if length else b"" + rendered: Final = handle_request(self.store, method, self.path, body) self.send_response(rendered.status_code) self.send_header("content-type", rendered.content_type) self.send_header("content-length", str(len(rendered.body))) @@ -621,11 +764,11 @@ DEFAULT_PORT: Final = 9100 def serve(port: int = DEFAULT_PORT, bind_host: str = "127.0.0.1") -> None: - server = ThreadingHTTPServer((bind_host, port), _ScriptedHandler) + server: Final = ThreadingHTTPServer((bind_host, port), _ScriptedHandler) sys.stderr.write(f"scripted-provider listening on http://{bind_host}:{port}\n") server.serve_forever() if __name__ == "__main__": - port_arg = int(sys.argv[1]) if len(sys.argv) > 1 else DEFAULT_PORT + port_arg: Final = int(sys.argv[1]) if len(sys.argv) > 1 else DEFAULT_PORT serve(port=port_arg) diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py index 8d7678cf9ca..e210dad94b1 100644 --- a/tests/e2e/cost_calculation/test_token_pricing_e2e.py +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -11,6 +11,7 @@ tests/e2e/cost_map.json. from __future__ import annotations import pytest +from typing import Final from conftest import CostCalcClient, cost_rows, register_scenario_deployment from cost_matrix import ( @@ -25,11 +26,11 @@ from e2e_config import unique_marker from lifecycle import ResourceManager from models import ChatBody, ChatMessage, ChatStreamOptions -pytestmark = [pytest.mark.e2e, pytest.mark.cost_map_stack] +pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark -_MATRIX: list[tuple[FrontierModel, Case]] = [ +_MATRIX: Final[tuple[tuple[FrontierModel, Case], ...]] = tuple( (model, case) for model in FRONTIER_MODELS for case in cases_for(model) -] +) def _case_id(param: tuple[FrontierModel, Case]) -> str: @@ -40,7 +41,7 @@ def _case_id(param: tuple[FrontierModel, Case]) -> str: def _chat_body(model_name: str, marker: str, case: Case) -> ChatBody: return ChatBody( model=model_name, - messages=[ChatMessage(role="user", content=f"{marker} scripted pricing call")], + messages=(ChatMessage(role="user", content=f"{marker} scripted pricing call"),), stream=case.stream, stream_options=ChatStreamOptions(include_usage=True) if case.stream else None, service_tier=case.service_tier, @@ -58,9 +59,9 @@ class TestTokenPricing: model_case: tuple[FrontierModel, Case], ) -> None: model, case = model_case - marker = unique_marker() + marker: Final = unique_marker() model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response = client.proxy.transport.send( + response: Final = client.proxy.transport.send( "/chat/completions", headers=client.proxy.transport.bearer(scoped_key), json=_chat_body(model_name, marker, case), @@ -71,7 +72,7 @@ class TestTokenPricing: ) assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" - expected = expected_cost(model, case) + expected: Final = expected_cost(model, case) if case.exact_spend and not case.stream: # Streamed responses commit headers before the bill is computed, so # the x-litellm-response-cost header is asserted only on non-stream @@ -82,7 +83,7 @@ class TestTokenPricing: f"x-litellm-response-cost {response.response_cost} != expected {expected}" ) - row = cost_rows.poll_cost_row_where( + row: Final = cost_rows.poll_cost_row_where( client.proxy, scoped_key, lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py index b1ef675d9ef..c0276cf370c 100644 --- a/tests/e2e/cost_calculation/test_wire_formats_e2e.py +++ b/tests/e2e/cost_calculation/test_wire_formats_e2e.py @@ -12,6 +12,9 @@ the proxy to POST /responses) and a streamed Anthropic-messages case. from __future__ import annotations import pytest +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final from conftest import CostCalcClient, cost_rows, register_scenario_deployment from cost_matrix import ( @@ -26,12 +29,14 @@ from lifecycle import ResourceManager from models import ChatBody, ChatMessage, ChatStreamOptions from scripted_provider import ScriptedUsage -pytestmark = [pytest.mark.e2e, pytest.mark.cost_map_stack] +pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark -_MODELS: dict[str, FrontierModel] = {model.map_key: model for model in FRONTIER_MODELS} +_MODELS: Final[Mapping[str, FrontierModel]] = MappingProxyType( + {model.map_key: model for model in FRONTIER_MODELS} +) # One scripted usage per wire, every reportable token kind nonzero. -_WIRE_USAGE: dict[str, tuple[str, ScriptedUsage]] = { +_WIRE_USAGE: Final[Mapping[str, tuple[str, ScriptedUsage]]] = MappingProxyType({ "openai_chat": ( "gpt-5.6", ScriptedUsage( @@ -89,7 +94,7 @@ _WIRE_USAGE: dict[str, tuple[str, ScriptedUsage]] = { "fireworks_ai/kimi-k3", ScriptedUsage(fresh_input_tokens=80, cache_read_tokens=40, output_tokens=25), ), -} +}) class TestWireFormats: @@ -103,22 +108,22 @@ class TestWireFormats: wire: str, ) -> None: map_key, usage = _WIRE_USAGE[wire] - model = _MODELS[map_key] - case = Case(name="basic", usage=usage) - marker = unique_marker() + model: Final = _MODELS[map_key] + case: Final = Case(name="basic", usage=usage) + marker: Final = unique_marker() model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response = client.proxy.transport.send( + response: Final = client.proxy.transport.send( "/chat/completions", headers=client.proxy.transport.bearer(scoped_key), json=ChatBody( model=model_name, - messages=[ChatMessage(role="user", content=f"{marker} scripted wire call")], + messages=(ChatMessage(role="user", content=f"{marker} scripted wire call"),), ), ) assert response.ok, f"{wire}: proxy returned {response.status_code}: {response.body[:400]}" - expected = expected_breakdown(model, case) - row = cost_rows.poll_cost_row_where( + expected: Final = expected_breakdown(model, case) + row: Final = cost_rows.poll_cost_row_where( client.proxy, scoped_key, lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, @@ -128,7 +133,7 @@ class TestWireFormats: f"{wire}: spend {row.spend} != expected {expected.total} " f"(breakdown {row.breakdown.model_dump()})" ) - breakdown = row.breakdown + breakdown: Final = row.breakdown assert breakdown.input_cost is not None and cost_rows.approx_equal( breakdown.input_cost, expected.input_cost ), ( @@ -153,16 +158,16 @@ class TestWireFormats: self, client: CostCalcClient, resources: ResourceManager, scoped_key: str ) -> None: map_key, usage = _WIRE_USAGE["anthropic_messages"] - model = _MODELS[map_key] - case = Case(name="stream", usage=usage, stream=True) - marker = unique_marker() + model: Final = _MODELS[map_key] + case: Final = Case(name="stream", usage=usage, stream=True) + marker: Final = unique_marker() model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response = client.proxy.transport.send( + response: Final = client.proxy.transport.send( "/chat/completions", headers=client.proxy.transport.bearer(scoped_key), json=ChatBody( model=model_name, - messages=[ChatMessage(role="user", content=f"{marker} scripted anthropic stream")], + messages=(ChatMessage(role="user", content=f"{marker} scripted anthropic stream"),), stream=True, stream_options=ChatStreamOptions(include_usage=True), ), @@ -172,8 +177,8 @@ class TestWireFormats: assert response.stream_done, "anthropic stream did not reach its terminal event" assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" - expected = expected_breakdown(model, case) - row = cost_rows.poll_cost_row_where( + expected: Final = expected_breakdown(model, case) + row: Final = cost_rows.poll_cost_row_where( client.proxy, scoped_key, lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, From 99bf8e9b2ffb6c647813029debba23c788e86b41 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 23:32:39 +0000 Subject: [PATCH 017/251] test(e2e): add cost calculation CI proxy config Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/CLAUDE.md | 2 +- tests/e2e/gateway/cost_calculation_ci_config.yml | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) create mode 100644 tests/e2e/gateway/cost_calculation_ci_config.yml diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index b6c3840f626..49cfc29aa17 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -21,7 +21,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see); needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in), and markerless harness unit tests for the locust, process-usage, and session-anomaly aggregation logic - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests -- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; asserts literal rate arithmetic on scripted usage across every provider wire and pricing component, deselected unless `E2E_COST_MAP_STACK` is set +- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; asserts literal rate arithmetic on scripted usage across every provider wire and pricing component, deselected unless `E2E_COST_MAP_STACK` is set, driven by the Buildkite `e2e-cost-calculation` step in project-releaser, which runs a proxy booted from `gateway/cost_calculation_ci_config.yml`, Postgres and the scripted provider co-located with pytest in one pod and sets the opt-in - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` diff --git a/tests/e2e/gateway/cost_calculation_ci_config.yml b/tests/e2e/gateway/cost_calculation_ci_config.yml new file mode 100644 index 00000000000..ac0603fa7c1 --- /dev/null +++ b/tests/e2e/gateway/cost_calculation_ci_config.yml @@ -0,0 +1,7 @@ +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY + database_url: os.environ/DATABASE_URL + store_model_in_db: true + proxy_batch_write_at: 5 + +model_list: [] From 415b06f5ff6d5959e2144ea826922694b2c8a60b Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 04:42:07 +0000 Subject: [PATCH 018/251] test(e2e): assert the real bill for the four fixed cost gaps Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cost_matrix.py | 22 +++++-------------- .../test_token_pricing_e2e.py | 5 ----- tests/e2e/cost_map.json | 15 +++++++++++++ 3 files changed, 21 insertions(+), 21 deletions(-) diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index e8b1d249559..8f39e89a358 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -186,15 +186,12 @@ _WIRE_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType({ } ), "openai_responses": frozenset({"cache_read", "reasoning", "web_search", "response_model", "absent_usage"}), - # Product gap: litellm hard-indexes message_delta["usage"] in - # anthropic/chat/handler.py, so a usage-absent anthropic stream raises - # KeyError; the real wire always carries it, so the case cannot be - # represented. - "anthropic_messages": frozenset({"cache_read", "cache_write_5m", "cache_write_1h", "web_search", "response_model"}), - # Product gap: the gemini transform sets ModelResponse.model from the - # request and drops the provider's modelVersion, so a response-model - # override can never be priced on this wire. - "gemini_generate": frozenset({"cache_read", "reasoning", "audio", "web_search", "absent_usage"}), + "anthropic_messages": frozenset( + {"cache_read", "cache_write_5m", "cache_write_1h", "web_search", "response_model", "absent_usage"} + ), + "gemini_generate": frozenset( + {"cache_read", "reasoning", "audio", "web_search", "response_model", "absent_usage"} + ), "together_chat": frozenset( { "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", @@ -240,9 +237,6 @@ class Case: billed_web_search_calls: int = 0 response_model_override: bool = False exact_spend: bool = True - # stream_usage=absent on a wire with no proxy-side token recount means the - # bill is exactly zero; asserted as such rather than skipped. - expect_zero_bill: bool = False def scenario(self, scenario_id: str, model: FrontierModel, text: str) -> Scenario: return Scenario( @@ -359,10 +353,6 @@ def cases_for(model: FrontierModel) -> tuple[Case, ...]: stream=True, stream_usage="absent", exact_spend=False, - # The responses surface bills only provider-reported usage; - # with no usage in the stream the spend row is zero. Other - # wires recount tokens proxy-side and bill a nonzero amount. - expect_zero_bill=model.wire == "openai_responses", ) if "absent_usage" in caps else None diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py index e210dad94b1..ead86931424 100644 --- a/tests/e2e/cost_calculation/test_token_pricing_e2e.py +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -90,11 +90,6 @@ class TestTokenPricing: ) assert row is not None, f"no spend row with a cost breakdown landed for {model.map_key}/{case.name}" - if not case.exact_spend and case.expect_zero_bill: - # The provider reported no usage and this wire has no proxy-side - # recount, so the bill is exactly zero. - assert row.spend is not None and row.spend == 0, f"no-usage stream billed {row.spend}: {row}" - return if not case.exact_spend: # stream_usage=absent: the provider reported no usage, so the row's # token counts are the proxy's own recount; only assert a bill landed. diff --git a/tests/e2e/cost_map.json b/tests/e2e/cost_map.json index b761710bae3..68d840870c9 100644 --- a/tests/e2e/cost_map.json +++ b/tests/e2e/cost_map.json @@ -63,13 +63,18 @@ "supports_web_search": true }, "fireworks_ai/deepseek-v4p1-flash": { + "cache_creation_input_token_cost": 0.00033, + "cache_creation_input_token_cost_above_1hr": 0.00044, "cache_read_input_token_cost": 1.4e-05, + "input_cost_per_audio_token": 0.00066, "input_cost_per_token": 0.00014000000000000001, "litellm_provider": "fireworks_ai", "max_input_tokens": 2000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_audio_token": 0.00077, + "output_cost_per_reasoning_token": 0.00055, "output_cost_per_token": 0.00028000000000000003, "search_context_cost_per_query": { "search_context_size_high": 0.03, @@ -82,13 +87,18 @@ "supports_web_search": true }, "fireworks_ai/kimi-k3": { + "cache_creation_input_token_cost": 0.00033, + "cache_creation_input_token_cost_above_1hr": 0.00044, "cache_read_input_token_cost": 1.2e-05, + "input_cost_per_audio_token": 0.00066, "input_cost_per_token": 0.00012000000000000002, "litellm_provider": "fireworks_ai", "max_input_tokens": 2000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_audio_token": 0.00077, + "output_cost_per_reasoning_token": 0.00055, "output_cost_per_token": 0.00024000000000000003, "search_context_cost_per_query": { "search_context_size_high": 0.03, @@ -101,13 +111,18 @@ "supports_web_search": true }, "fireworks_ai/qwen3p8-max": { + "cache_creation_input_token_cost": 0.00033, + "cache_creation_input_token_cost_above_1hr": 0.00044, "cache_read_input_token_cost": 1.3e-05, + "input_cost_per_audio_token": 0.00066, "input_cost_per_token": 0.00013000000000000002, "litellm_provider": "fireworks_ai", "max_input_tokens": 2000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", + "output_cost_per_audio_token": 0.00077, + "output_cost_per_reasoning_token": 0.00055, "output_cost_per_token": 0.00026000000000000003, "search_context_cost_per_query": { "search_context_size_high": 0.03, From 67778cfe2625eb97fd3d4733f1fae41aed7599fb Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:18:08 +0000 Subject: [PATCH 019/251] fix(cost): resolve dated openai/azure snapshots to their undated cost map entry Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/utils.py | 9 +++++++- tests/test_litellm/test_cost_calculator.py | 24 ++++++++++++++++++++++ tests/test_litellm/test_utils.py | 19 +++++++++++++++++ 3 files changed, 51 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 734522c0c6a..33fdfcc36ba 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5286,6 +5286,13 @@ def _strip_stable_vertex_version(model_name) -> str: return re.sub(r"-\d+$", "", model_name) +_DATED_SNAPSHOT_SUFFIX: Final = re.compile(r"-\d{4}-\d{2}-\d{2}$") + + +def _strip_dated_snapshot_suffix(model_name: str) -> str: + return _DATED_SNAPSHOT_SUFFIX.sub("", model_name) + + def _get_base_bedrock_model(model_name) -> str: """ Get the base model from the given model name. @@ -5333,7 +5340,7 @@ def _strip_model_name(model: str, custom_llm_provider: str | None) -> str: strip_finetune: Final = _strip_openai_finetune_model_name(model_name=model) return strip_finetune else: - return model + return _strip_dated_snapshot_suffix(model_name=model) # Global case-insensitive lookup map for model_cost (built eagerly at module import) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 7b53d3a58df..c5bc8d80fed 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -21,7 +21,9 @@ from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponseAPIUsage from litellm.types.rerank import RerankResponse from litellm.types.utils import ( CallTypes, + Choices, LiteLLMRealtimeStreamLoggingObject, + Message, ModelInfo, ModelResponse, PromptTokensDetailsWrapper, @@ -110,6 +112,28 @@ def test_completion_cost_uses_response_model_for_dynamic_routing(_local_model_co assert cost > 0, "Cost should be calculated using response model" +def test_completion_cost_strips_dated_azure_snapshot_model(_local_model_cost_map: None) -> None: + dated_response = ModelResponse( + model="gpt-5.6-luna-2026-07-09", + choices=[Choices(index=0, message=Message(role="assistant", content="hi"), finish_reason="stop")], + usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), + ) + dated_response._hidden_params = {"custom_llm_provider": "azure"} + + undated_response = ModelResponse( + model="gpt-5.6-luna", + choices=[Choices(index=0, message=Message(role="assistant", content="hi"), finish_reason="stop")], + usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), + ) + undated_response._hidden_params = {"custom_llm_provider": "azure"} + + dated_cost = litellm.completion_cost(completion_response=dated_response) + undated_cost = litellm.completion_cost(completion_response=undated_response) + + assert dated_cost == undated_cost + assert dated_cost > 0 + + def test_cost_calculator_with_response_cost_in_additional_headers(): class MockResponse(BaseModel): _hidden_params = {"additional_headers": {"llm_provider-x-litellm-response-cost": 1000}} diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 46149589371..2f9e27af797 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -186,6 +186,25 @@ def test_get_model_info_strips_openai_finetune_ids_without_a_custom_suffix(local assert info["key"] == "ft:gpt-4o-2024-08-06" +@pytest.mark.parametrize( + ("model", "custom_llm_provider", "expected_key"), + [ + ("gpt-5.6-luna-2026-07-09", "openai", "gpt-5.6-luna"), + ("gpt-5.6-luna-2026-07-09", "azure", "azure/gpt-5.6-luna"), + ], +) +def test_get_model_info_falls_back_from_dated_snapshot_to_undated_entry( + local_model_cost_map, model, custom_llm_provider, expected_key +): + info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + assert info["key"] == expected_key + + +def test_get_model_info_prefers_exact_dated_key_over_stripped(local_model_cost_map): + info = litellm.get_model_info(model="gpt-4o-2024-08-06", custom_llm_provider="openai") + assert info["key"] == "gpt-4o-2024-08-06" + + def test_check_provider_match_azure_ai_allows_openai_and_azure(): """ Test that azure_ai provider can match openai and azure models. From feb69c5f789d44a65dbbfa348ce39eaa3874b37f Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 16:38:15 +0000 Subject: [PATCH 020/251] test(e2e): add tool-call, terminal, and image-input shapes to the cost matrix Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cost_matrix.py | 184 +++++++- .../e2e/cost_calculation/scripted_provider.py | 408 +++++++++++++++--- .../test_token_pricing_e2e.py | 61 ++- .../cost_calculation/test_wire_formats_e2e.py | 111 ++++- 4 files changed, 688 insertions(+), 76 deletions(-) diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index 8f39e89a358..5f634712778 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -17,7 +17,11 @@ creation), the case is absent from the matrix rather than silently zero. from __future__ import annotations +import base64 import json +import random +import struct +import zlib from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path @@ -26,7 +30,7 @@ from typing import Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict, TypeAdapter -from scripted_provider import Scenario, ScriptedOutput, ScriptedUsage, Wire +from scripted_provider import Scenario, ScriptedOutput, ScriptedToolCall, ScriptedUsage, Wire COST_MAP_PATH: Final = Path(__file__).resolve().parent.parent / "cost_map.json" @@ -182,26 +186,37 @@ _WIRE_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType({ "openai_chat": frozenset( { "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", - "web_search", "response_model", "absent_usage", + "web_search", "response_model", "absent_usage", "tool_call", "image_input", + } + ), + "openai_responses": frozenset( + { + "cache_read", "reasoning", "web_search", "response_model", "absent_usage", + "tool_call", "image_input", "responses_terminal", } ), - "openai_responses": frozenset({"cache_read", "reasoning", "web_search", "response_model", "absent_usage"}), "anthropic_messages": frozenset( - {"cache_read", "cache_write_5m", "cache_write_1h", "web_search", "response_model", "absent_usage"} + { + "cache_read", "cache_write_5m", "cache_write_1h", "web_search", + "response_model", "absent_usage", "tool_call", "image_input", + } ), "gemini_generate": frozenset( - {"cache_read", "reasoning", "audio", "web_search", "response_model", "absent_usage"} + { + "cache_read", "reasoning", "audio", "web_search", "response_model", + "absent_usage", "tool_call", "image_input", "prompt_blocked", + } ), "together_chat": frozenset( { "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", - "web_search", "response_model", "absent_usage", + "web_search", "response_model", "absent_usage", "tool_call", "image_input", } ), "fireworks_chat": frozenset( { "cache_read", "cache_write_5m", "cache_write_1h", "reasoning", "audio", - "web_search", "response_model", "absent_usage", + "web_search", "response_model", "absent_usage", "tool_call", "image_input", } ), }) @@ -220,6 +235,15 @@ CaseName: TypeAlias = Literal[ "stream", "stream_no_usage", "response_model_override", + "stream_response_model_override", + "tool_call", + "stream_no_usage_tool_call", + "stream_no_usage_image_input", + "stream_no_usage_incomplete", + "stream_unvalidated", + "stream_no_usage_unvalidated", + "prompt_blocked", + "stream_prompt_blocked", ] @@ -237,6 +261,9 @@ class Case: billed_web_search_calls: int = 0 response_model_override: bool = False exact_spend: bool = True + tool_call: bool = False + image_input: bool = False + terminal: Literal["completed", "incomplete", "unvalidated", "prompt_blocked"] = "completed" def scenario(self, scenario_id: str, model: FrontierModel, text: str) -> Scenario: return Scenario( @@ -246,6 +273,10 @@ class Case: output=ScriptedOutput( text=text, response_model=model.override_model if self.response_model_override else None, + tool_call=ScriptedToolCall(name="get_weather", arguments=TOOL_CALL_ARGUMENTS) + if self.tool_call + else None, + terminal=self.terminal, ), stream_usage=self.stream_usage, service_tier=self.service_tier, @@ -254,6 +285,15 @@ class Case: _BASIC_USAGE: Final = ScriptedUsage(fresh_input_tokens=120, output_tokens=40) +TOOL_CALL_ARGUMENTS: Final = json.dumps({ + "city": "Berlin", + "days": 7, + "units": "metric", + "notes": "filler " * 30, +}) + +_PROMPT_BLOCKED_USAGE: Final = ScriptedUsage(fresh_input_tokens=1000, output_tokens=0) + def _web_search_case(model: FrontierModel) -> Case: counts_exactly: Final = model.wire in ("openai_responses", "anthropic_messages", "gemini_generate") @@ -362,6 +402,100 @@ def cases_for(model: FrontierModel) -> tuple[Case, ...]: if "response_model" in caps else None ), + ( + Case( + name="stream_response_model_override", + usage=_BASIC_USAGE, + stream=True, + response_model_override=True, + ) + if "response_model" in caps + else None + ), + ( + Case(name="tool_call", usage=_BASIC_USAGE, tool_call=True) + if "tool_call" in caps + else None + ), + ( + Case( + name="stream_no_usage_tool_call", + usage=_BASIC_USAGE, + stream=True, + stream_usage="absent", + tool_call=True, + exact_spend=False, + ) + if "absent_usage" in caps and "tool_call" in caps + else None + ), + ( + Case( + name="stream_no_usage_image_input", + usage=_BASIC_USAGE, + stream=True, + stream_usage="absent", + image_input=True, + exact_spend=False, + ) + if "absent_usage" in caps and "image_input" in caps + else None + ), + ( + Case( + name="stream_no_usage_incomplete", + usage=_BASIC_USAGE, + stream=True, + stream_usage="absent", + terminal="incomplete", + exact_spend=False, + ) + if "responses_terminal" in caps + else None + ), + ( + Case( + name="stream_unvalidated", + usage=_BASIC_USAGE, + stream=True, + terminal="unvalidated", + ) + if "responses_terminal" in caps + else None + ), + ( + Case( + name="stream_no_usage_unvalidated", + usage=_BASIC_USAGE, + stream=True, + stream_usage="absent", + terminal="unvalidated", + exact_spend=False, + ) + if "responses_terminal" in caps + else None + ), + ( + Case( + name="prompt_blocked", + usage=_PROMPT_BLOCKED_USAGE, + terminal="prompt_blocked", + response_model_override=True, + ) + if "prompt_blocked" in caps + else None + ), + ( + Case( + name="stream_prompt_blocked", + usage=_PROMPT_BLOCKED_USAGE, + stream=True, + terminal="prompt_blocked", + response_model_override=True, + ) + if "prompt_blocked" in caps + else None + ), ) return tuple(case for case in candidates if case is not None) @@ -436,6 +570,42 @@ def expected_cost(model: FrontierModel, case: Case) -> float: return expected_breakdown(model, case).total +def recount_cost( + model: FrontierModel, case: Case, prompt_tokens: int, completion_tokens: int +) -> float: + """What the proxy's own token recount should cost at the case's rates, + without pinning the tokenizer's exact counts.""" + rates: Final = model.override_rates if case.response_model_override else model.rates + return prompt_tokens * (rates.input_cost_per_token or 0.0) + completion_tokens * ( + rates.output_cost_per_token or 0.0 + ) + + +def _png_chunk(tag: bytes, payload: bytes) -> bytes: + return struct.pack(">I", len(payload)) + tag + payload + struct.pack(">I", zlib.crc32(tag + payload)) + + +def image_input_data_url() -> str: + """A deterministic 256x256 RGB noise PNG as a data URL; noise compresses + poorly on purpose so the base64 payload stays well above 100 KB and would + blow up the prompt recount if the URL were ever tokenized as text.""" + rng: Final = random.Random(0) + side: Final = 256 + raw: Final = b"".join( + b"\x00" + rng.randbytes(side * 3) for _ in range(side) + ) + png: Final = ( + b"\x89PNG\r\n\x1a\n" + + _png_chunk(b"IHDR", struct.pack(">IIBBBBB", side, side, 8, 2, 0, 0, 0)) + + _png_chunk(b"IDAT", zlib.compress(raw)) + + _png_chunk(b"IEND", b"") + ) + return "data:image/png;base64," + base64.b64encode(png).decode() + + +IMAGE_INPUT_DATA_URL: Final = image_input_data_url() + + def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: """(prompt_tokens, completion_tokens) the spend row should carry, per the wire's normalization: Anthropic folds cache read/write into prompt_tokens, diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py index f1deafd1bc5..e1a6c430307 100644 --- a/tests/e2e/cost_calculation/scripted_provider.py +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -38,7 +38,7 @@ from types import MappingProxyType from typing import Final, Literal, TypeAlias from urllib.parse import urlsplit -from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, model_validator Wire: TypeAlias = Literal[ "openai_chat", @@ -62,6 +62,26 @@ WIRE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType( StreamUsage: TypeAlias = Literal["final_chunk", "absent"] ServiceTier: TypeAlias = Literal["flex", "priority"] +TerminalKind: TypeAlias = Literal["completed", "incomplete", "unvalidated", "prompt_blocked"] + +# Which terminal variant each wire can represent. +_TERMINAL_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType( + { + "openai_responses": frozenset({"incomplete", "unvalidated"}), + "gemini_generate": frozenset({"prompt_blocked"}), + } +) + + +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 + for streams.""" + + model_config = ConfigDict(frozen=True) + + name: str + arguments: str class ScriptedUsage(BaseModel): @@ -96,6 +116,12 @@ class ScriptedOutput(BaseModel): # OpenAI-compatible providers can report a provider-computed cost; emitted as # the top-level "cost" field on the together/fireworks wire. provider_cost: float | None = None + # When set, the response is a tool call only: no text content on any wire. + 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; + # "prompt_blocked" is a Gemini promptFeedback-only body. + terminal: TerminalKind = "completed" class Scenario(BaseModel): @@ -108,6 +134,17 @@ class Scenario(BaseModel): stream_usage: StreamUsage = "final_chunk" service_tier: ServiceTier | None = None + @model_validator(mode="after") + def _check_terminal_supported(self) -> Scenario: + if ( + self.output.terminal != "completed" + and self.output.terminal not in _TERMINAL_CAPS.get(self.wire, frozenset()) + ): + raise ValueError( + f"wire {self.wire} cannot emit terminal={self.output.terminal}" + ) + return self + @property def mount(self) -> str: return WIRE_MOUNTS[self.wire] @@ -291,10 +328,38 @@ def _responses_usage(u: ScriptedUsage) -> Mapping[str, object]: # ---------- per-wire responses ---------- +def _split_arguments(arguments: str) -> tuple[str, ...]: + """Slice a tool-call arguments JSON string into 2-3 streamed deltas.""" + third: Final = max(1, len(arguments) // 3) + return tuple( + slice_ + for slice_ in (arguments[:third], arguments[third : 2 * third], arguments[2 * third :]) + if slice_ + ) + + def _openai_message(scenario: Scenario) -> Mapping[str, object]: + tool_call: Final = scenario.output.tool_call return _jobj_opt( ("role", "assistant"), - ("content", scenario.output.text), + ("content", None if tool_call is not None else scenario.output.text), + ( + ( + "tool_calls", + ( + _jobj( + ("id", f"call_{scenario.scenario_id}"), + ("type", "function"), + ( + "function", + _jobj(("name", tool_call.name), ("arguments", tool_call.arguments)), + ), + ), + ), + ) + if tool_call is not None + else None + ), ( ( "annotations", @@ -332,7 +397,12 @@ def _openai_chat_body(scenario: Scenario, requested_model: str) -> Mapping[str, _jobj( ("index", 0), ("message", _openai_message(scenario)), - ("finish_reason", scenario.output.finish_reason), + ( + "finish_reason", + "tool_calls" + if scenario.output.tool_call is not None + else scenario.output.finish_reason, + ), ), ), ), @@ -359,6 +429,7 @@ def _openai_chunk( def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: + tool_call: Final = scenario.output.tool_call delta: Final = _jobj_opt( ("role", "assistant"), ("content", scenario.output.text), @@ -368,6 +439,43 @@ def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: else None ), ) + body_deltas: Final[tuple[Mapping[str, object], ...]] = ( + ( + _jobj( + ("role", "assistant"), + ( + "tool_calls", + ( + _jobj( + ("index", 0), + ("id", f"call_{scenario.scenario_id}"), + ("type", "function"), + ( + "function", + _jobj(("name", tool_call.name), ("arguments", "")), + ), + ), + ), + ), + ), + *( + _jobj( + ( + "tool_calls", + ( + _jobj( + ("index", 0), + ("function", _jobj(("arguments", arguments_slice))), + ), + ), + ) + ) + for arguments_slice in _split_arguments(tool_call.arguments) + ), + ) + if tool_call is not None + else (delta,) + ) return _sse( ( ( @@ -378,13 +486,16 @@ def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: choices=(_jobj(("index", 0), ("delta", _jobj(("role", "assistant"))), ("finish_reason", None)),), ), ), - ( - None, - _openai_chunk( - scenario, - requested_model, - choices=(_jobj(("index", 0), ("delta", delta), ("finish_reason", None)),), - ), + *( + ( + None, + _openai_chunk( + scenario, + requested_model, + choices=(_jobj(("index", 0), ("delta", body_delta), ("finish_reason", None)),), + ), + ) + for body_delta in body_deltas ), ( None, @@ -395,7 +506,12 @@ def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: _jobj( ("index", 0), ("delta", _jobj()), - ("finish_reason", scenario.output.finish_reason), + ( + "finish_reason", + "tool_calls" + if tool_call is not None + else scenario.output.finish_reason, + ), ), ), ), @@ -410,17 +526,34 @@ def _openai_chat_sse(scenario: Scenario, requested_model: str) -> bytes: ) +def _anthropic_content(scenario: Scenario) -> tuple[Mapping[str, object], ...]: + tool_call: Final = scenario.output.tool_call + if tool_call is not None: + return ( + _jobj( + ("type", "tool_use"), + ("id", f"toolu_{scenario.scenario_id}"), + ("name", tool_call.name), + ("input", json.loads(tool_call.arguments)), + ), + ) + return (_jobj(("type", "text"), ("text", scenario.output.text)),) + + +def _anthropic_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 _anthropic_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: return _jobj( ("id", f"msg_{scenario.scenario_id}"), ("type", "message"), ("role", "assistant"), ("model", scenario.output.response_model or requested_model), - ("content", (_jobj(("type", "text"), ("text", scenario.output.text)),)), - ( - "stop_reason", - "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, - ), + ("content", _anthropic_content(scenario)), + ("stop_reason", _anthropic_stop_reason(scenario)), ("usage", _anthropic_usage(scenario.usage)), ) @@ -453,12 +586,7 @@ def _anthropic_sse(scenario: Scenario, requested_model: str) -> bytes: ("type", "message_delta"), ( "delta", - _jobj( - ( - "stop_reason", - "end_turn" if scenario.output.finish_reason == "stop" else scenario.output.finish_reason, - ) - ), + _jobj(("stop_reason", _anthropic_stop_reason(scenario))), ), ( ("usage", _jobj(("output_tokens", scenario.usage.output_tokens))) @@ -474,16 +602,45 @@ def _anthropic_sse(scenario: Scenario, requested_model: str) -> bytes: _jobj( ("type", "content_block_start"), ("index", 0), - ("content_block", _jobj(("type", "text"), ("text", ""))), + ( + "content_block", + _jobj( + ("type", "tool_use"), + ("id", f"toolu_{scenario.scenario_id}"), + ("name", scenario.output.tool_call.name), + ("input", _jobj()), + ) + if scenario.output.tool_call is not None + else _jobj(("type", "text"), ("text", "")), + ), ), ), - ( - "content_block_delta", - _jobj( - ("type", "content_block_delta"), - ("index", 0), - ("delta", _jobj(("type", "text_delta"), ("text", scenario.output.text))), - ), + *( + tuple( + ( + "content_block_delta", + _jobj( + ("type", "content_block_delta"), + ("index", 0), + ( + "delta", + _jobj(("type", "input_json_delta"), ("partial_json", arguments_slice)), + ), + ), + ) + for arguments_slice in _split_arguments(scenario.output.tool_call.arguments) + ) + if scenario.output.tool_call is not None + else ( + ( + "content_block_delta", + _jobj( + ("type", "content_block_delta"), + ("index", 0), + ("delta", _jobj(("type", "text_delta"), ("text", scenario.output.text))), + ), + ), + ) ), ("content_block_stop", _jobj(("type", "content_block_stop"), ("index", 0))), ("message_delta", message_delta), @@ -492,7 +649,49 @@ def _anthropic_sse(scenario: Scenario, requested_model: str) -> bytes: ) +def _gemini_prompt_blocked_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + return _jobj( + ( + "promptFeedback", + _jobj( + ("blockReason", "SAFETY"), + ( + "safetyRatings", + ( + _jobj( + ("category", "HARM_CATEGORY_HARASSMENT"), + ("probability", "HIGH"), + ("blocked", True), + ), + ), + ), + ), + ), + ("usageMetadata", _gemini_usage(scenario.usage)), + ("modelVersion", scenario.output.response_model or requested_model), + ) + + +def _gemini_parts(scenario: Scenario) -> tuple[Mapping[str, object], ...]: + tool_call: Final = scenario.output.tool_call + if tool_call is not None: + return ( + _jobj( + ( + "functionCall", + _jobj( + ("name", tool_call.name), + ("args", json.loads(tool_call.arguments)), + ), + ) + ), + ) + return (_jobj(("text", scenario.output.text)),) + + def _gemini_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + if scenario.output.terminal == "prompt_blocked": + return _gemini_prompt_blocked_body(scenario, requested_model) return _jobj( ( "candidates", @@ -501,7 +700,7 @@ def _gemini_body(scenario: Scenario, requested_model: str) -> Mapping[str, objec ( "content", _jobj( - ("parts", (_jobj(("text", scenario.output.text)),)), + ("parts", _gemini_parts(scenario)), ("role", "model"), ), ), @@ -559,67 +758,148 @@ def _gemini_sse(scenario: Scenario, requested_model: str) -> bytes: ) -def _responses_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: - return _jobj( - ("id", f"resp_{scenario.scenario_id}"), - ("object", "response"), - ("created_at", int(time.time())), - ("status", "completed"), - ("model", scenario.output.response_model or requested_model), - ( - "output", +def _responses_output(scenario: Scenario) -> tuple[Mapping[str, object], ...]: + tool_call: Final = scenario.output.tool_call + return ( + *( ( - *( - _jobj(("type", "web_search_call"), ("id", f"ws_{i}"), ("status", "completed")) - for i in range(scenario.usage.web_search_calls) - ), - _jobj( - ("type", "message"), - ("id", f"msg_{scenario.scenario_id}"), - ("status", "completed"), - ("role", "assistant"), - ( - "content", - ( - _jobj( - ("type", "output_text"), - ("text", scenario.output.text), - ("annotations", ()), - ), - ), + _jobj(("type", "scripted_future_item"), ("id", f"fut_{scenario.scenario_id}"), ("status", "completed")), + ) + if scenario.output.terminal == "unvalidated" + else () + ), + *( + _jobj(("type", "web_search_call"), ("id", f"ws_{i}"), ("status", "completed")) + for i in range(scenario.usage.web_search_calls) + ), + _jobj( + ("type", "function_call"), + ("id", f"fc_{scenario.scenario_id}"), + ("call_id", f"call_{scenario.scenario_id}"), + ("name", tool_call.name), + ("arguments", tool_call.arguments), + ("status", "completed"), + ) + if tool_call is not None + else _jobj( + ("type", "message"), + ("id", f"msg_{scenario.scenario_id}"), + ("status", "completed"), + ("role", "assistant"), + ( + "content", + ( + _jobj( + ("type", "output_text"), + ("text", scenario.output.text), + ("annotations", ()), ), ), ), ), + ) + + +def _responses_body(scenario: Scenario, requested_model: str) -> Mapping[str, object]: + incomplete: Final = scenario.output.terminal == "incomplete" + return _jobj_opt( + ("id", f"resp_{scenario.scenario_id}"), + ("object", "response"), + ( + "created_at", + "not-a-number" if scenario.output.terminal == "unvalidated" else int(time.time()), + ), + ("status", "incomplete" if incomplete else "completed"), + ( + ("incomplete_details", _jobj(("reason", "max_output_tokens"))) + if incomplete + else None + ), + ("model", scenario.output.response_model or requested_model), + ("output", _responses_output(scenario)), ("usage", _responses_usage(scenario.usage)), ) def _responses_sse(scenario: Scenario, requested_model: str) -> bytes: - completed: Final = ( + tool_call: Final = scenario.output.tool_call + terminal: Final = ( _jobj(*((key, value) for key, value in _responses_body(scenario, requested_model).items() if key != "usage")) if scenario.stream_usage == "absent" else _responses_body(scenario, requested_model) ) created: Final = _jobj( - *((key, value) for key, value in completed.items() if key not in ("status", "usage")), + *((key, value) for key, value in terminal.items() if key not in ("status", "usage")), ("status", "in_progress"), ("usage", None), ) - return _sse( + terminal_event: Final = ( + "response.incomplete" if scenario.output.terminal == "incomplete" else "response.completed" + ) + output_index: Final = ( + scenario.usage.web_search_calls + (1 if scenario.output.terminal == "unvalidated" else 0) + ) + middle_events: Final[tuple[tuple[str, Mapping[str, object]], ...]] = ( ( - ("response.created", _jobj(("type", "response.created"), ("response", created))), + ( + "response.output_item.added", + _jobj( + ("type", "response.output_item.added"), + ("output_index", output_index), + ( + "item", + _jobj( + ("type", "function_call"), + ("id", f"fc_{scenario.scenario_id}"), + ("call_id", f"call_{scenario.scenario_id}"), + ("name", tool_call.name), + ("arguments", ""), + ("status", "in_progress"), + ), + ), + ), + ), + *( + ( + "response.function_call_arguments.delta", + _jobj( + ("type", "response.function_call_arguments.delta"), + ("item_id", f"fc_{scenario.scenario_id}"), + ("output_index", output_index), + ("delta", arguments_slice), + ), + ) + for arguments_slice in _split_arguments(tool_call.arguments) + ), + ( + "response.function_call_arguments.done", + _jobj( + ("type", "response.function_call_arguments.done"), + ("item_id", f"fc_{scenario.scenario_id}"), + ("output_index", output_index), + ("arguments", tool_call.arguments), + ), + ), + ) + if tool_call is not None + else ( ( "response.output_text.delta", _jobj( ("type", "response.output_text.delta"), ("item_id", f"msg_{scenario.scenario_id}"), - ("output_index", scenario.usage.web_search_calls), + ("output_index", output_index), ("content_index", 0), ("delta", scenario.output.text), ), ), - ("response.completed", _jobj(("type", "response.completed"), ("response", completed))), + ) + ) + return _sse( + ( + ("response.created", _jobj(("type", "response.created"), ("response", created))), + *middle_events, + (terminal_event, _jobj(("type", terminal_event), ("response", terminal))), ) ) diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py index ead86931424..0b4f3e1fd37 100644 --- a/tests/e2e/cost_calculation/test_token_pricing_e2e.py +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -16,15 +16,26 @@ from typing import Final from conftest import CostCalcClient, cost_rows, register_scenario_deployment from cost_matrix import ( FRONTIER_MODELS, + IMAGE_INPUT_DATA_URL, Case, FrontierModel, cases_for, expected_cost, expected_token_columns, + recount_cost, ) from e2e_config import unique_marker from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatStreamOptions +from models import ( + ChatBody, + ChatMessage, + ChatStreamOptions, + ChatTool, + ChatToolFunction, + ImageContentPart, + ImageUrl, + TextContentPart, +) pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark @@ -41,10 +52,37 @@ def _case_id(param: tuple[FrontierModel, Case]) -> str: def _chat_body(model_name: str, marker: str, case: Case) -> ChatBody: return ChatBody( model=model_name, - messages=(ChatMessage(role="user", content=f"{marker} scripted pricing call"),), + messages=( + ChatMessage( + role="user", + content=( + [ + TextContentPart(text=f"{marker} scripted pricing call"), + ImageContentPart(image_url=ImageUrl(url=IMAGE_INPUT_DATA_URL)), + ] + if case.image_input + else f"{marker} scripted pricing call" + ), + ), + ), stream=case.stream, stream_options=ChatStreamOptions(include_usage=True) if case.stream else None, service_tier=case.service_tier, + tools=( + ( + ChatTool( + function=ChatToolFunction( + name="get_weather", + parameters={ + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + ) + ), + ) + if case.tool_call + else None + ), ) @@ -92,8 +130,23 @@ class TestTokenPricing: if not case.exact_spend: # stream_usage=absent: the provider reported no usage, so the row's - # token counts are the proxy's own recount; only assert a bill landed. - assert row.spend is not None and row.spend > 0, f"no-usage stream billed nothing: {row}" + # token counts are the proxy's own recount; assert the recount + # billed both directions at the case's rates. + assert row.prompt_tokens is not None and row.prompt_tokens > 0, ( + f"no-usage stream counted no input tokens: {row}" + ) + assert row.completion_tokens is not None and row.completion_tokens > 0, ( + f"no-usage stream counted no output tokens: {row}" + ) + if case.image_input: + assert row.prompt_tokens < 4000, ( + f"image data URL looks tokenized as text: prompt_tokens={row.prompt_tokens}" + ) + assert row.spend is not None and cost_rows.approx_equal( + row.spend, + recount_cost(model, case, row.prompt_tokens, row.completion_tokens), + ), f"no-usage stream spend {row.spend} != recount at map rates: {row}" + cost_rows.assert_total_is_sum_of_components(row) return assert row.spend is not None and cost_rows.approx_equal(row.spend, expected), ( diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py index c0276cf370c..3c6c34b24fb 100644 --- a/tests/e2e/cost_calculation/test_wire_formats_e2e.py +++ b/tests/e2e/cost_calculation/test_wire_formats_e2e.py @@ -26,7 +26,7 @@ from cost_matrix import ( ) from e2e_config import unique_marker from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatStreamOptions +from models import ChatBody, ChatMessage, ChatStreamOptions, ChatTool, ChatToolFunction from scripted_provider import ScriptedUsage pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark @@ -96,6 +96,53 @@ _WIRE_USAGE: Final[Mapping[str, tuple[str, ScriptedUsage]]] = MappingProxyType({ ), }) +_SHAPE_USAGE: Final = ScriptedUsage(fresh_input_tokens=80, output_tokens=25) + +# Renderer-level shapes the pricing matrix gates per cap, pinned here once per +# wire so the sidecar emits prove they survive the proxy end to end. +_SHAPES: Final[tuple[tuple[str, str, Case], ...]] = ( + *( + ( + f"tool_call_{'stream' if stream else 'sync'}", + wire, + Case(name="tool_call", usage=_SHAPE_USAGE, stream=stream, tool_call=True), + ) + for wire in _WIRE_USAGE + for stream in (False, True) + ), + ( + "responses_incomplete", + "openai_responses", + Case(name="stream_no_usage_incomplete", usage=_SHAPE_USAGE, stream=True, terminal="incomplete"), + ), + ( + "responses_unvalidated", + "openai_responses", + Case(name="stream_unvalidated", usage=_SHAPE_USAGE, stream=True, terminal="unvalidated"), + ), + ( + "gemini_prompt_blocked", + "gemini_generate", + Case( + name="prompt_blocked", + usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), + terminal="prompt_blocked", + response_model_override=True, + ), + ), + ( + "gemini_prompt_blocked_stream", + "gemini_generate", + Case( + name="stream_prompt_blocked", + usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), + stream=True, + terminal="prompt_blocked", + response_model_override=True, + ), + ), +) + class TestWireFormats: @pytest.mark.parametrize("wire", tuple(_WIRE_USAGE)) @@ -189,3 +236,65 @@ class TestWireFormats: f"(breakdown {row.breakdown.model_dump()})" ) cost_rows.assert_total_is_sum_of_components(row) + + @pytest.mark.parametrize("shape_wire_case", _SHAPES, ids=lambda entry: entry[0]) + @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") + def test_response_shape_bills_reported_usage( + self, + client: CostCalcClient, + resources: ResourceManager, + scoped_key: str, + shape_wire_case: tuple[str, str, Case], + ) -> None: + shape, wire, case = shape_wire_case + map_key, _usage = _WIRE_USAGE[wire] + model: Final = _MODELS[map_key] + marker: Final = unique_marker() + model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) + response: Final = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(scoped_key), + json=ChatBody( + model=model_name, + messages=(ChatMessage(role="user", content=f"{marker} scripted {shape}"),), + stream=case.stream, + stream_options=ChatStreamOptions(include_usage=True) if case.stream else None, + tools=( + ( + ChatTool( + function=ChatToolFunction( + name="get_weather", + parameters={"type": "object", "properties": {"city": {"type": "string"}}}, + ) + ), + ) + if case.tool_call + else None + ), + ), + stream=case.stream, + ) + assert response.ok, f"{shape}: proxy returned {response.status_code}: {response.body[:400]}" + if case.stream: + assert response.stream_done, f"{shape}: stream did not reach its terminal event" + assert response.stream_error is None, f"{shape}: stream error: {response.stream_error}" + + expected: Final = expected_breakdown(model, case) + row: Final = cost_rows.poll_cost_row_where( + client.proxy, + scoped_key, + lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, + ) + assert row is not None, f"{shape}: no spend row landed" + assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( + f"{shape}: spend {row.spend} != expected {expected.total} " + f"(breakdown {row.breakdown.model_dump()})" + ) + prompt_tokens, completion_tokens = expected_token_columns(model, case) + assert row.prompt_tokens == prompt_tokens, ( + f"{shape}: prompt_tokens {row.prompt_tokens} != {prompt_tokens}" + ) + assert row.completion_tokens == completion_tokens, ( + f"{shape}: completion_tokens {row.completion_tokens} != {completion_tokens}" + ) + cost_rows.assert_total_is_sum_of_components(row) From 5507de326e3e98f9069af5f9d1315c89bb3c3e25 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 16:45:09 +0000 Subject: [PATCH 021/251] test(e2e): type the wire-shape parametrize ids callback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/test_wire_formats_e2e.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py index 3c6c34b24fb..4da7b31a6ef 100644 --- a/tests/e2e/cost_calculation/test_wire_formats_e2e.py +++ b/tests/e2e/cost_calculation/test_wire_formats_e2e.py @@ -144,6 +144,10 @@ _SHAPES: Final[tuple[tuple[str, str, Case], ...]] = ( ) +def _shape_id(entry: tuple[str, str, Case]) -> str: + return entry[0] + + class TestWireFormats: @pytest.mark.parametrize("wire", tuple(_WIRE_USAGE)) @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") @@ -237,7 +241,7 @@ class TestWireFormats: ) cost_rows.assert_total_is_sum_of_components(row) - @pytest.mark.parametrize("shape_wire_case", _SHAPES, ids=lambda entry: entry[0]) + @pytest.mark.parametrize("shape_wire_case", _SHAPES, ids=_shape_id) @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") def test_response_shape_bills_reported_usage( self, From 9885dc89621697e235fc65e0e86e147da87398c1 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 22:52:42 +0000 Subject: [PATCH 022/251] 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): From 2466975d290576de9e89a1d5d69c9ca9a6aab1ab Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 22:57:23 +0000 Subject: [PATCH 023/251] test(e2e): clean cost map decimals and simplify scripted wire helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/conftest.py | 52 ++-- .../e2e/cost_calculation/scripted_provider.py | 39 ++- tests/e2e/cost_map.json | 270 +++++++++--------- 3 files changed, 177 insertions(+), 184 deletions(-) diff --git a/tests/e2e/cost_calculation/conftest.py b/tests/e2e/cost_calculation/conftest.py index 8c6db7c0010..3f3e9fd9243 100644 --- a/tests/e2e/cost_calculation/conftest.py +++ b/tests/e2e/cost_calculation/conftest.py @@ -11,6 +11,7 @@ Deselected unless E2E_COST_MAP_STACK is set (marker `cost_map_stack`). from __future__ import annotations +import functools import importlib.util import json import sys @@ -21,6 +22,8 @@ from types import ModuleType from typing import Final, Protocol, cast import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa from cost_matrix import Case, FrontierModel from e2e_config import COST_MAP_PROXY_URL, SCRIPTED_PROVIDER_PROXY_BASE @@ -112,33 +115,25 @@ def client() -> CostCalcClient: return CostCalcClient(proxy=proxy) -_vertex_key_pem: str | None = None +@functools.cache +def _vertex_private_key_pem() -> str: + return rsa.generate_private_key(public_exponent=65537, key_size=2048).private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ).decode() 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() - ) + access token without touching Google.""" return json.dumps( { "type": "service_account", "project_id": "cc-scripted-project", "private_key_id": "scripted", - "private_key": _vertex_key_pem, + "private_key": _vertex_private_key_pem(), "client_email": "scripted@cc-scripted-project.iam.gserviceaccount.com", "client_id": "0", "auth_uri": f"{SCRIPTED_PROVIDER_PROXY_BASE}/_oauth/authorize", @@ -162,20 +157,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() + params: Final = { + "model": model.litellm_model, + "api_key": model.api_key, + "api_base": handle.api_base(), + **model.litellm_params, + **( + {"vertex_credentials": _vertex_service_account_json()} + if model.wire == "vertex_generate" + else {} + ), + } model_id: Final = client.proxy.register_model( ModelNewBody( model_name=model_name, - litellm_params=LiteLLMParamsBody.model_validate( - { - "model": model.litellm_model, - "api_key": model.api_key, - "api_base": handle.api_base(), - **extra_params, - } - ), + litellm_params=LiteLLMParamsBody.model_validate(params), model_info=ModelInfoBody(base_model=model.base_model), ) ) diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py index 00230fabeba..982132ed8df 100644 --- a/tests/e2e/cost_calculation/scripted_provider.py +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -997,36 +997,33 @@ def _bedrock_body(scenario: Scenario) -> Mapping[str, object]: ) +def _aws_str_header(name: str, value: str) -> bytes: + """One eventstream header: 1-byte name len + name + type-7 marker + value.""" + 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 + ) + + 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") + _aws_str_header(":event-type", event_type) + + _aws_str_header(":content-type", "application/json") + + _aws_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) + prelude_crc: Final = struct.pack("!I", zlib.crc32(prelude) & 0xFFFFFFFF) message: Final = prelude + prelude_crc + headers_bytes + payload_bytes - return message + struct.pack("!I", _crc32(message, 0) & 0xFFFFFFFF) + return message + struct.pack("!I", zlib.crc32(message) & 0xFFFFFFFF) def _bedrock_eventstream(scenario: Scenario) -> bytes: diff --git a/tests/e2e/cost_map.json b/tests/e2e/cost_map.json index 4fba337b701..85cd5ade3d5 100644 --- a/tests/e2e/cost_map.json +++ b/tests/e2e/cost_map.json @@ -1,4 +1,79 @@ { + "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.00045, + "cache_creation_input_token_cost_above_1hr": 0.0006, + "cache_read_input_token_cost": 1.5e-05, + "input_cost_per_audio_token": 0.0009, + "input_cost_per_token": 0.00015, + "input_cost_per_token_above_200k_tokens": 0.0012, + "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.00105, + "output_cost_per_reasoning_token": 0.00075, + "output_cost_per_token": 0.0003, + "output_cost_per_token_above_200k_tokens": 0.00135, + "output_cost_per_token_flex": 0.000375, + "output_cost_per_token_priority": 0.000405, + "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 + }, "claude-haiku-4-5": { "cache_creation_input_token_cost": 0.00021, "cache_creation_input_token_cost_above_1hr": 0.00028000000000000003, @@ -134,6 +209,64 @@ "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.00105, + "output_cost_per_token": 0.00042, + "output_cost_per_token_above_200k_tokens": 0.00189, + "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.0014, + "output_cost_per_reasoning_token": 0.001, + "output_cost_per_token": 0.0004, + "output_cost_per_token_above_200k_tokens": 0.0018, + "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" + }, "gemini/gemini-3.1-pro-preview": { "cache_read_input_token_cost": 9e-06, "input_cost_per_audio_token": 0.00054, @@ -304,139 +437,6 @@ "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", @@ -508,7 +508,7 @@ "supports_web_search": true }, "us.anthropic.claude-opus-5-v1:0": { - "cache_creation_input_token_cost": 0.0005400000000000001, + "cache_creation_input_token_cost": 0.00054, "cache_creation_input_token_cost_above_1hr": 0.00072, "cache_read_input_token_cost": 1.8e-05, "input_cost_per_token": 0.00018, @@ -517,7 +517,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 0.00036000000000000004, + "output_cost_per_token": 0.00036, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true From 813d96f26ea6780bacb4b5ad562f1aecc3cb5069 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 23:18:26 +0000 Subject: [PATCH 024/251] fix(e2e): resolve remaining merge markers in e2e_config Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/e2e_config.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index b34eadd8744..e19cfaa684f 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -145,7 +145,6 @@ WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY" MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK" PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK" REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS" -<<<<<<< HEAD # The cost_calculation suite needs a proxy booted with LITELLM_MODEL_COST_MAP_URL # pointing at tests/e2e/cost_map.json (its whole map is test-owned rates) plus a # scripted-provider sidecar; deselected unless the opt-in env var is set. @@ -162,9 +161,6 @@ SCRIPTED_PROVIDER_CONTROL_URL = os.environ.get( SCRIPTED_PROVIDER_PROXY_BASE = os.environ.get( "E2E_SCRIPTED_PROVIDER_PROXY_BASE", SCRIPTED_PROVIDER_CONTROL_URL ).rstrip("/") -||||||| 930ec9643a -======= -CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM" CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM" ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) From bdfff602fb0325f88768f7c4411cce921ab28fcb Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 00:47:55 +0000 Subject: [PATCH 025/251] test(e2e): drive the cost matrix from cases.json and expected.json goldens Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/CLAUDE.md | 2 +- tests/e2e/cost_calculation/cases.json | 231 ++ tests/e2e/cost_calculation/conftest.py | 10 +- tests/e2e/cost_calculation/cost_matrix.py | 788 ++----- tests/e2e/cost_calculation/expected.json | 2004 +++++++++++++++++ .../e2e/cost_calculation/generate_expected.py | 189 ++ .../e2e/cost_calculation/test_matrix_data.py | 64 + .../test_token_pricing_e2e.py | 62 +- .../cost_calculation/test_wire_formats_e2e.py | 368 --- 9 files changed, 2761 insertions(+), 957 deletions(-) create mode 100644 tests/e2e/cost_calculation/cases.json create mode 100644 tests/e2e/cost_calculation/expected.json create mode 100644 tests/e2e/cost_calculation/generate_expected.py create mode 100644 tests/e2e/cost_calculation/test_matrix_data.py delete mode 100644 tests/e2e/cost_calculation/test_wire_formats_e2e.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 49cfc29aa17..707d35b4aa6 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -21,7 +21,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see); needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in), and markerless harness unit tests for the locust, process-usage, and session-anomaly aggregation logic - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests -- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; asserts literal rate arithmetic on scripted usage across every provider wire and pricing component, deselected unless `E2E_COST_MAP_STACK` is set, driven by the Buildkite `e2e-cost-calculation` step in project-releaser, which runs a proxy booted from `gateway/cost_calculation_ci_config.yml`, Postgres and the scripted provider co-located with pytest in one pod and sets the opt-in +- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; every cost-map entry is a deployment and the cases plus asserted goldens are data in `cases.json` and `expected.json` (regenerate with `generate_expected.py`), deselected unless `E2E_COST_MAP_STACK` is set, driven by the Buildkite `e2e-cost-calculation` step in project-releaser, which runs a proxy booted from `gateway/cost_calculation_ci_config.yml`, Postgres and the scripted provider co-located with pytest in one pod and sets the opt-in - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` diff --git a/tests/e2e/cost_calculation/cases.json b/tests/e2e/cost_calculation/cases.json new file mode 100644 index 00000000000..e898557ea35 --- /dev/null +++ b/tests/e2e/cost_calculation/cases.json @@ -0,0 +1,231 @@ +{ + "deployments": [ + { + "map_key": "azure/gpt-5.4-mini", + "litellm_model": "azure/cc-pinned-deployment", + "base_model": "azure/gpt-5.4-mini" + } + ], + "cases": [ + { + "name": "basic", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40} + }, + { + "name": "cache_read", + "usage": {"fresh_input_tokens": 100, "cache_read_tokens": 50, "output_tokens": 30}, + "requires_rates": ["cache_read_input_token_cost"], + "requires_caps": ["cache_read"] + }, + { + "name": "cache_write_5m", + "usage": {"fresh_input_tokens": 90, "cache_write_5m_tokens": 60, "output_tokens": 30}, + "requires_rates": ["cache_creation_input_token_cost"], + "requires_caps": ["cache_write_5m"] + }, + { + "name": "cache_write_1h", + "usage": {"fresh_input_tokens": 90, "cache_write_5m_tokens": 20, "cache_write_1h_tokens": 40, "output_tokens": 30}, + "requires_rates": ["cache_creation_input_token_cost_above_1hr", "cache_creation_input_token_cost"], + "requires_caps": ["cache_write_1h"] + }, + { + "name": "reasoning", + "usage": {"fresh_input_tokens": 100, "output_tokens": 30, "reasoning_tokens": 70}, + "requires_rates": ["output_cost_per_reasoning_token"], + "requires_caps": ["reasoning"] + }, + { + "name": "audio", + "usage": {"fresh_input_tokens": 100, "audio_input_tokens": 25, "output_tokens": 30, "audio_output_tokens": 15}, + "requires_rates": ["input_cost_per_audio_token", "output_cost_per_audio_token"], + "requires_caps": ["audio"] + }, + { + "name": "tiered", + "usage": {"fresh_input_tokens": 200001, "output_tokens": 30}, + "requires_rates": ["input_cost_per_token_above_200k_tokens", "output_cost_per_token_above_200k_tokens"] + }, + { + "name": "service_tier_flex", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "service_tier": "flex", + "requires_rates": ["input_cost_per_token_flex", "output_cost_per_token_flex"] + }, + { + "name": "service_tier_priority", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "service_tier": "priority", + "requires_rates": ["input_cost_per_token_priority", "output_cost_per_token_priority"] + }, + { + "name": "web_search", + "usage": {"fresh_input_tokens": 100, "output_tokens": 30, "web_search_calls": 3}, + "requires_rates": ["search_context_cost_per_query"], + "requires_caps": ["web_search"] + }, + { + "name": "stream", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true + }, + { + "name": "stream_no_usage", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "stream_usage": "absent", + "exact_spend": false, + "requires_caps": ["absent_usage"] + }, + { + "name": "response_model_override", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "response_model_override": true, + "requires_caps": ["response_model"] + }, + { + "name": "stream_response_model_override", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "response_model_override": true, + "requires_caps": ["response_model"] + }, + { + "name": "tool_call", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "tool_call": true, + "requires_caps": ["tool_call"] + }, + { + "name": "stream_tool_call", + "usage": {"fresh_input_tokens": 80, "output_tokens": 25}, + "stream": true, + "tool_call": true, + "requires_caps": ["tool_call"] + }, + { + "name": "stream_no_usage_tool_call", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "stream_usage": "absent", + "tool_call": true, + "exact_spend": false, + "requires_caps": ["absent_usage", "tool_call"] + }, + { + "name": "stream_no_usage_image_input", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "stream_usage": "absent", + "image_input": true, + "exact_spend": false, + "requires_caps": ["absent_usage", "image_input"] + }, + { + "name": "stream_incomplete", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "terminal": "incomplete", + "requires_caps": ["responses_terminal"] + }, + { + "name": "stream_no_usage_incomplete", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "stream_usage": "absent", + "terminal": "incomplete", + "exact_spend": false, + "requires_caps": ["responses_terminal"] + }, + { + "name": "stream_unvalidated", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "terminal": "unvalidated", + "requires_caps": ["responses_terminal"] + }, + { + "name": "stream_no_usage_unvalidated", + "usage": {"fresh_input_tokens": 120, "output_tokens": 40}, + "stream": true, + "stream_usage": "absent", + "terminal": "unvalidated", + "exact_spend": false, + "requires_caps": ["responses_terminal"] + }, + { + "name": "prompt_blocked", + "usage": {"fresh_input_tokens": 1000, "output_tokens": 0}, + "terminal": "prompt_blocked", + "response_model_override": true, + "requires_caps": ["prompt_blocked"] + }, + { + "name": "stream_prompt_blocked", + "usage": {"fresh_input_tokens": 1000, "output_tokens": 0}, + "stream": true, + "terminal": "prompt_blocked", + "response_model_override": true, + "requires_caps": ["prompt_blocked"] + }, + { + "name": "all_components_chat", + "usage": { + "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 + }, + "wires": ["openai_chat", "azure_chat", "together_chat"] + }, + { + "name": "all_components_fireworks", + "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25}, + "wires": ["fireworks_chat"] + }, + { + "name": "all_components_anthropic", + "usage": { + "fresh_input_tokens": 80, + "cache_read_tokens": 40, + "cache_write_5m_tokens": 20, + "cache_write_1h_tokens": 10, + "output_tokens": 25 + }, + "wires": ["anthropic_messages", "bedrock_converse"] + }, + { + "name": "all_components_anthropic_stream", + "usage": { + "fresh_input_tokens": 80, + "cache_read_tokens": 40, + "cache_write_5m_tokens": 20, + "cache_write_1h_tokens": 10, + "output_tokens": 25 + }, + "stream": true, + "wires": ["anthropic_messages"] + }, + { + "name": "all_components_gemini", + "usage": { + "fresh_input_tokens": 80, + "cache_read_tokens": 40, + "output_tokens": 25, + "reasoning_tokens": 15, + "audio_input_tokens": 5, + "audio_output_tokens": 3 + }, + "wires": ["gemini_generate", "vertex_generate"] + }, + { + "name": "all_components_responses", + "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25, "reasoning_tokens": 15}, + "wires": ["openai_responses"] + } + ] +} diff --git a/tests/e2e/cost_calculation/conftest.py b/tests/e2e/cost_calculation/conftest.py index 3f3e9fd9243..3de9786854e 100644 --- a/tests/e2e/cost_calculation/conftest.py +++ b/tests/e2e/cost_calculation/conftest.py @@ -1,10 +1,12 @@ """Cost-calculation suite fixtures. Runs against a dedicated proxy whose whole model cost map is the test-owned -``tests/e2e/cost_map.json`` (LITELLM_MODEL_COST_MAP_URL), so every deployment -bills at rates the test asserts literal arithmetic on. Provider calls are -answered by the scripted-provider sidecar (``scripted_provider.py``), registered -per scenario over its control API. +``tests/e2e/cost_map.json`` (LITELLM_MODEL_COST_MAP_URL); every map entry is a +deployment under test, the request shapes live in ``cases.json``, and the +asserted goldens live in ``expected.json`` (regenerate proposals with +``generate_expected.py``). Provider calls are answered by the +scripted-provider sidecar (``scripted_provider.py``), registered per scenario +over its control API. Deselected unless E2E_COST_MAP_STACK is set (marker `cost_map_stack`). """ diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index 37495f37b0b..b03c851d208 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -1,18 +1,16 @@ -"""The cost-calculation matrix: frontier model set, the pricing-component cases -each model runs, and the expected-cost arithmetic. +"""The cost-calculation matrix: the model set derived from the test cost map, +the request/response cases from ``cases.json``, and the loaders both use. -Rates come from ``tests/e2e/cost_map.json``, which the proxy under test loads as -its ENTIRE model cost map (LITELLM_MODEL_COST_MAP_URL), so an entry's rates are -exactly what the proxy bills and nothing in the suite depends on the bundled -map. Each model's rates are a distinct multiple of a shared base set, so a -component billed at the wrong model's rate (or the wrong case's rate) can never -coincidentally match. - -Case applicability is pricing-field-gated AND wire-gated: a case runs for a -model only when the entry carries the rate the case exercises and the wire can -report the token kind that rate prices. When the wire cannot report a kind -(e.g. Anthropic has no reasoning-token field, Responses reports no cache -creation), the case is absent from the matrix rather than silently zero. +Three data files drive the suite; nothing in Python lists models or cases: +- ``tests/e2e/cost_map.json`` is the proxy's ENTIRE model cost map + (LITELLM_MODEL_COST_MAP_URL); every entry becomes a deployment under test. +- ``tests/e2e/cost_calculation/cases.json`` is the case list; each case runs + for a model when the entry carries the rates it exercises (``requires_rates``) + and the wire can report the token kinds involved (``requires_caps`` / + ``wires``). +- ``tests/e2e/cost_calculation/expected.json`` holds the reviewed goldens; the + tests assert them verbatim and never compute a price themselves. The rate + arithmetic that proposes goldens lives in ``generate_expected.py``, not here. """ from __future__ import annotations @@ -26,13 +24,15 @@ from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path from types import MappingProxyType -from typing import Final, Literal, TypeAlias +from typing import Final, Literal from pydantic import BaseModel, ConfigDict, TypeAdapter from scripted_provider import Scenario, ScriptedOutput, ScriptedToolCall, ScriptedUsage, Wire COST_MAP_PATH: Final = Path(__file__).resolve().parent.parent / "cost_map.json" +CASES_PATH: Final = Path(__file__).resolve().parent / "cases.json" +EXPECTED_PATH: Final = Path(__file__).resolve().parent / "expected.json" class SearchContextCostPerQuery(BaseModel): @@ -77,119 +77,90 @@ _COST_MAP: Final[Mapping[str, CostMapEntry]] = MappingProxyType( TIER_THRESHOLD_TOKENS: Final = 200_000 -@dataclass(frozen=True, slots=True) -class FrontierModel: - """One deployment under test: the model_name the suite registers, the - provider-prefixed litellm model string, the wire the scripted upstream - speaks, its cost-map key, and the sibling map model the response_model - override case reports.""" +class DeploymentSpec(BaseModel): + """A deployment-level fact from cases.json: when a map key needs a + registered deployment name that is not its provider model (or a + model_info.base_model pin), the matrix uses these instead of the defaults.""" + + model_config = ConfigDict(frozen=True) - model_name: str - litellm_model: str - wire: Wire map_key: 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. + litellm_model: str | None = None 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: - return _COST_MAP[self.map_key] - - @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 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: - return self.rates.litellm_provider - - @property - def api_key(self) -> str: - # The scripted upstream ignores auth; a fixed bogus key proves the suite - # spends zero real provider calls. - return "sk-scripted-provider" -# Response-model override targets: emit a sibling's bare provider-facing name so -# the biller's provider-prefixed lookup lands on that sibling's map key. -_OVERRIDE_MODELS: Final[Mapping[str, str]] = MappingProxyType({ - "gpt-5.6": "gpt-5.4-mini", - "gpt-5.5-pro": "gpt-5.3-codex", - "gpt-5.3-codex": "gpt-5.5-pro", - "gpt-5.4-mini": "gpt-5.6", - "claude-opus-5": "claude-sonnet-5", - "claude-sonnet-5": "claude-opus-5", - "claude-haiku-4-5": "claude-sonnet-5", - "gemini/gemini-3.8-flash": "gemini-3.1-pro-preview", - "gemini/gemini-3.1-pro-preview": "gemini-3.8-flash", - "together_ai/moonshotai/Kimi-K3": "zai-org/GLM-5.3", - "together_ai/zai-org/GLM-5.3": "moonshotai/Kimi-K3", - "fireworks_ai/kimi-k3": "qwen3p8-max", - "fireworks_ai/qwen3p8-max": "kimi-k3", - "fireworks_ai/deepseek-v4p1-flash": "kimi-k3", -}) +class Case(BaseModel): + """One request/response shape from cases.json; gated onto a model by + ``requires_rates`` (entry must carry each rate field), ``requires_caps`` + (the wire must report the token kind) and ``wires`` (shape is wire-specific).""" -_OVERRIDE_MAP_KEYS: Final[Mapping[str, str]] = MappingProxyType({ - "gpt-5.4-mini": "gpt-5.4-mini", - "gpt-5.6": "gpt-5.6", - "gpt-5.3-codex": "gpt-5.3-codex", - "gpt-5.5-pro": "gpt-5.5-pro", - "claude-sonnet-5": "claude-sonnet-5", - "claude-opus-5": "claude-opus-5", - "gemini-3.1-pro-preview": "gemini/gemini-3.1-pro-preview", - "gemini-3.8-flash": "gemini/gemini-3.8-flash", - "zai-org/GLM-5.3": "together_ai/zai-org/GLM-5.3", - "moonshotai/Kimi-K3": "together_ai/moonshotai/Kimi-K3", - "qwen3p8-max": "fireworks_ai/qwen3p8-max", - "kimi-k3": "fireworks_ai/kimi-k3", -}) + model_config = ConfigDict(frozen=True) + + name: str + usage: ScriptedUsage + stream: bool = False + stream_usage: Literal["final_chunk", "absent"] = "final_chunk" + service_tier: Literal["flex", "priority"] | None = None + response_model_override: bool = False + exact_spend: bool = True + tool_call: bool = False + image_input: bool = False + terminal: Literal["completed", "incomplete", "unvalidated", "prompt_blocked"] = "completed" + requires_rates: tuple[str, ...] = () + requires_caps: tuple[str, ...] = () + wires: tuple[Wire, ...] | None = None + + def applies_to(self, model: FrontierModel) -> bool: + if self.wires is not None and model.wire not in self.wires: + return False + caps: Final = _WIRE_CAPS[model.wire] + if not frozenset(self.requires_caps) <= caps: + return False + return all( + getattr(model.rates, field, None) is not None for field in self.requires_rates + ) + + def scenario(self, scenario_id: str, model: FrontierModel, text: str) -> Scenario: + return Scenario( + 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, + tool_call=ScriptedToolCall(name="get_weather", arguments=TOOL_CALL_ARGUMENTS) + if self.tool_call + else None, + terminal=self.terminal, + ), + stream_usage=self.stream_usage, + service_tier=self.service_tier, + ) -_FRONTIER_SPECS: Final[tuple[tuple[str, str, Wire], ...]] = ( - ("gpt-5.6", "openai/gpt-5.6", "openai_chat"), - ("gpt-5.5-pro", "openai/gpt-5.5-pro", "openai_responses"), - ("gpt-5.3-codex", "openai/gpt-5.3-codex", "openai_responses"), - ("gpt-5.4-mini", "openai/gpt-5.4-mini", "openai_chat"), - ("claude-opus-5", "anthropic/claude-opus-5", "anthropic_messages"), - ("claude-sonnet-5", "anthropic/claude-sonnet-5", "anthropic_messages"), - ("claude-haiku-4-5", "anthropic/claude-haiku-4-5", "anthropic_messages"), - ("gemini/gemini-3.8-flash", "gemini/gemini-3.8-flash", "gemini_generate"), - ("gemini/gemini-3.1-pro-preview", "gemini/gemini-3.1-pro-preview", "gemini_generate"), - ("together_ai/moonshotai/Kimi-K3", "together_ai/moonshotai/Kimi-K3", "together_chat"), - ("together_ai/zai-org/GLM-5.3", "together_ai/zai-org/GLM-5.3", "together_chat"), - ("fireworks_ai/kimi-k3", "fireworks_ai/kimi-k3", "fireworks_chat"), - ("fireworks_ai/qwen3p8-max", "fireworks_ai/qwen3p8-max", "fireworks_chat"), - ("fireworks_ai/deepseek-v4p1-flash", "fireworks_ai/deepseek-v4p1-flash", "fireworks_chat"), +class _CasesFile(BaseModel): + model_config = ConfigDict(frozen=True) + + deployments: tuple[DeploymentSpec, ...] = () + cases: tuple[Case, ...] = () + + +_CASES_FILE: Final = _CasesFile.model_validate(json.loads(CASES_PATH.read_text())) +CASES: Final[tuple[Case, ...]] = _CASES_FILE.cases +_DEPLOYMENTS: Final[Mapping[str, DeploymentSpec]] = MappingProxyType( + {spec.map_key: spec for spec in _CASES_FILE.deployments} ) @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.""" +class _ProviderWiring: + """How a (litellm_provider, mode) pair maps to a sidecar wire, the provider + prefix on the registered litellm model string, and extra litellm_params.""" - 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({}) + model_prefix: str | None + litellm_params: Mapping[str, str] _AZURE_PARAMS: Final[Mapping[str, str]] = MappingProxyType({"api_version": "2025-04-01-preview"}) @@ -207,87 +178,134 @@ _VERTEX_PARAMS: Final[Mapping[str, str]] = MappingProxyType( } ) -_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, - ), +_PROVIDER_WIRING: Final[Mapping[tuple[str, str], _ProviderWiring]] = MappingProxyType( + { + ("openai", "chat"): _ProviderWiring("openai_chat", "openai", MappingProxyType({})), + ("openai", "responses"): _ProviderWiring( + "openai_responses", "openai", MappingProxyType({}) + ), + ("anthropic", "chat"): _ProviderWiring( + "anthropic_messages", "anthropic", MappingProxyType({}) + ), + ("gemini", "chat"): _ProviderWiring("gemini_generate", None, MappingProxyType({})), + ("together_ai", "chat"): _ProviderWiring("together_chat", None, MappingProxyType({})), + ("fireworks_ai", "chat"): _ProviderWiring("fireworks_chat", None, MappingProxyType({})), + ("azure", "chat"): _ProviderWiring("azure_chat", None, _AZURE_PARAMS), + ("bedrock_converse", "chat"): _ProviderWiring( + "bedrock_converse", "bedrock/converse", _BEDROCK_PARAMS + ), + ("vertex_ai-language-models", "chat"): _ProviderWiring( + "vertex_generate", "vertex_ai", _VERTEX_PARAMS + ), + } ) +@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.""" + + model_name: str + litellm_model: str + wire: Wire + map_key: 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 + litellm_params: Mapping[str, str] = MappingProxyType({}) + + @property + def rates(self) -> CostMapEntry: + return _COST_MAP[self.map_key] + + @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 provider_model(self) -> str: + """The bare provider-facing model name: litellm_model minus the provider + prefix and any routing segment (converse/, responses/).""" + return _provider_model(self.litellm_model) + + @property + def provider(self) -> str: + return self.rates.litellm_provider + + @property + def api_key(self) -> str: + # The scripted upstream ignores auth; a fixed bogus key proves the suite + # spends zero real provider calls. + return "sk-scripted-provider" + + +def _provider_model(litellm_model: str) -> str: + tail: Final = litellm_model.split("/")[1:] + 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: + return map_key + if map_key.startswith(f"{wiring.model_prefix}/"): + return map_key + return f"{wiring.model_prefix}/{map_key}" + + def _frontier() -> tuple[FrontierModel, ...]: - return tuple( - FrontierModel( - model_name=f"cc-{map_key.replace('/', '-').lower()}", - litellm_model=litellm_model, - 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 + groups: Final[Mapping[tuple[str, str], tuple[str, ...]]] = MappingProxyType( + { + pair: tuple(sorted(k for k, e in _COST_MAP.items() if (e.litellm_provider, e.mode) == pair)) + for pair in {(e.litellm_provider, e.mode) for e in _COST_MAP.values()} + } ) + models: list[FrontierModel] = [] # mutable-ok: accumulated once at import into a tuple + for map_key in sorted(_COST_MAP): + entry: Final = _COST_MAP[map_key] + pair: Final = (entry.litellm_provider, entry.mode) + wiring: Final = _PROVIDER_WIRING.get(pair) + if wiring is None: + raise ValueError( + f"cost_map entry {map_key} has no wiring for " + f"(litellm_provider={pair[0]}, mode={pair[1]}); add a " + f"_ProviderWiring row in cost_matrix.py" + ) + siblings: Final = groups[pair] + override_key: Final = ( + siblings[(siblings.index(map_key) + 1) % len(siblings)] if len(siblings) > 1 else None + ) + override_litellm: Final = ( + _litellm_model_for(override_key, wiring) if override_key is not None else None + ) + deployment: Final = _DEPLOYMENTS.get(map_key) + 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, + map_key=map_key, + override_model=( + _provider_model(override_litellm) + if override_litellm is not None + else None + ), + override_map_key=override_key, + base_model=deployment.base_model if deployment is not None else None, + litellm_params=wiring.litellm_params, + ) + ) + return tuple(models) FRONTIER_MODELS: Final[tuple[FrontierModel, ...]] = _frontier() @@ -350,71 +368,6 @@ _WIRE_CAPS: Final[Mapping[str, frozenset[str]]] = MappingProxyType({ ), }) -CaseName: TypeAlias = Literal[ - "basic", - "cache_read", - "cache_write_5m", - "cache_write_1h", - "reasoning", - "audio", - "tiered", - "service_tier_flex", - "service_tier_priority", - "web_search", - "stream", - "stream_no_usage", - "response_model_override", - "stream_response_model_override", - "tool_call", - "stream_no_usage_tool_call", - "stream_no_usage_image_input", - "stream_no_usage_incomplete", - "stream_unvalidated", - "stream_no_usage_unvalidated", - "prompt_blocked", - "stream_prompt_blocked", -] - - -@dataclass(frozen=True, slots=True) -class Case: - name: CaseName - usage: ScriptedUsage - stream: bool = False - stream_usage: Literal["final_chunk", "absent"] = "final_chunk" - service_tier: Literal["flex", "priority"] | None = None - # For web_search the wire's reported call count is not always what gets - # billed: chat-completions surfaces only expose url_citation annotations, so - # the biller floors to one call; responses/messages/gemini report a real - # count. - billed_web_search_calls: int = 0 - response_model_override: bool = False - exact_spend: bool = True - tool_call: bool = False - image_input: bool = False - terminal: Literal["completed", "incomplete", "unvalidated", "prompt_blocked"] = "completed" - - def scenario(self, scenario_id: str, model: FrontierModel, text: str) -> Scenario: - return Scenario( - 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, - tool_call=ScriptedToolCall(name="get_weather", arguments=TOOL_CALL_ARGUMENTS) - if self.tool_call - else None, - terminal=self.terminal, - ), - stream_usage=self.stream_usage, - service_tier=self.service_tier, - ) - - -_BASIC_USAGE: Final = ScriptedUsage(fresh_input_tokens=120, output_tokens=40) - TOOL_CALL_ARGUMENTS: Final = json.dumps({ "city": "Berlin", "days": 7, @@ -422,284 +375,9 @@ TOOL_CALL_ARGUMENTS: Final = json.dumps({ "notes": "filler " * 30, }) -_PROMPT_BLOCKED_USAGE: Final = ScriptedUsage(fresh_input_tokens=1000, output_tokens=0) - - -def _web_search_case(model: FrontierModel) -> Case: - 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), - billed_web_search_calls=3 if counts_exactly else 1, - ) - def cases_for(model: FrontierModel) -> tuple[Case, ...]: - rates: Final = model.rates - caps: Final = _WIRE_CAPS[model.wire] - candidates: Final[tuple[Case | None, ...]] = ( - Case(name="basic", usage=_BASIC_USAGE), - ( - Case(name="cache_read", usage=ScriptedUsage(fresh_input_tokens=100, cache_read_tokens=50, output_tokens=30)) - if rates.cache_read_input_token_cost is not None and "cache_read" in caps - else None - ), - ( - Case( - name="cache_write_5m", - usage=ScriptedUsage(fresh_input_tokens=90, cache_write_5m_tokens=60, output_tokens=30), - ) - if rates.cache_creation_input_token_cost is not None and "cache_write_5m" in caps - else None - ), - ( - Case( - name="cache_write_1h", - usage=ScriptedUsage( - fresh_input_tokens=90, - cache_write_5m_tokens=20, - cache_write_1h_tokens=40, - output_tokens=30, - ), - ) - if ( - rates.cache_creation_input_token_cost_above_1hr is not None - and rates.cache_creation_input_token_cost is not None - and "cache_write_1h" in caps - ) - else None - ), - ( - Case( - name="reasoning", - usage=ScriptedUsage(fresh_input_tokens=100, output_tokens=30, reasoning_tokens=70), - ) - if rates.output_cost_per_reasoning_token is not None and "reasoning" in caps - else None - ), - ( - Case( - name="audio", - usage=ScriptedUsage( - fresh_input_tokens=100, audio_input_tokens=25, output_tokens=30, audio_output_tokens=15 - ), - ) - if ( - rates.input_cost_per_audio_token is not None - and rates.output_cost_per_audio_token is not None - and "audio" in caps - ) - else None - ), - ( - Case( - name="tiered", - usage=ScriptedUsage( - fresh_input_tokens=TIER_THRESHOLD_TOKENS + 1, output_tokens=30 - ), - ) - if ( - rates.input_cost_per_token_above_200k_tokens is not None - and rates.output_cost_per_token_above_200k_tokens is not None - ) - else None - ), - ( - Case(name="service_tier_flex", usage=_BASIC_USAGE, service_tier="flex") - if rates.input_cost_per_token_flex is not None and rates.output_cost_per_token_flex is not None - else None - ), - ( - Case(name="service_tier_priority", usage=_BASIC_USAGE, service_tier="priority") - if rates.input_cost_per_token_priority is not None and rates.output_cost_per_token_priority is not None - else None - ), - _web_search_case(model) if rates.search_context_cost_per_query is not None and "web_search" in caps else None, - Case(name="stream", usage=_BASIC_USAGE, stream=True), - ( - Case( - name="stream_no_usage", - usage=_BASIC_USAGE, - stream=True, - stream_usage="absent", - exact_spend=False, - ) - if "absent_usage" in caps - else None - ), - ( - Case(name="response_model_override", usage=_BASIC_USAGE, response_model_override=True) - if "response_model" in caps - else None - ), - ( - Case( - name="stream_response_model_override", - usage=_BASIC_USAGE, - stream=True, - response_model_override=True, - ) - if "response_model" in caps - else None - ), - ( - Case(name="tool_call", usage=_BASIC_USAGE, tool_call=True) - if "tool_call" in caps - else None - ), - ( - Case( - name="stream_no_usage_tool_call", - usage=_BASIC_USAGE, - stream=True, - stream_usage="absent", - tool_call=True, - exact_spend=False, - ) - if "absent_usage" in caps and "tool_call" in caps - else None - ), - ( - Case( - name="stream_no_usage_image_input", - usage=_BASIC_USAGE, - stream=True, - stream_usage="absent", - image_input=True, - exact_spend=False, - ) - if "absent_usage" in caps and "image_input" in caps - else None - ), - ( - Case( - name="stream_no_usage_incomplete", - usage=_BASIC_USAGE, - stream=True, - stream_usage="absent", - terminal="incomplete", - exact_spend=False, - ) - if "responses_terminal" in caps - else None - ), - ( - Case( - name="stream_unvalidated", - usage=_BASIC_USAGE, - stream=True, - terminal="unvalidated", - ) - if "responses_terminal" in caps - else None - ), - ( - Case( - name="stream_no_usage_unvalidated", - usage=_BASIC_USAGE, - stream=True, - stream_usage="absent", - terminal="unvalidated", - exact_spend=False, - ) - if "responses_terminal" in caps - else None - ), - ( - Case( - name="prompt_blocked", - usage=_PROMPT_BLOCKED_USAGE, - terminal="prompt_blocked", - response_model_override=True, - ) - if "prompt_blocked" in caps - else None - ), - ( - Case( - name="stream_prompt_blocked", - usage=_PROMPT_BLOCKED_USAGE, - stream=True, - terminal="prompt_blocked", - response_model_override=True, - ) - if "prompt_blocked" in caps - else None - ), - ) - return tuple(case for case in candidates if case is not None) - - -@dataclass(frozen=True, slots=True) -class ExpectedCost: - """The expected bill split the way the spend row's cost_breakdown reports - it: the gross input component (cache reads/writes folded in), the output - component, and the tool-usage component.""" - - input_cost: float - output_cost: float - tool_cost: float - - @property - def total(self) -> float: - return self.input_cost + self.output_cost + self.tool_cost - - -def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: - """Literal arithmetic on the test-map rates over the scripted token counts. - - Input = fresh*in + read*read + 5m*create + 1h*create_1h + audio_in*audio_in; - output = text*out + reasoning*reasoning + audio_out*audio_out; plus the - billed web-search calls at the medium search-context rate. Above-threshold - swaps every input/output rate to its ``_above_200k_tokens`` variant when - total prompt tokens exceed the threshold; a service tier swaps input/output - to the tier's variants, falling back to the base rate when a variant is - unset -- mirroring _get_token_base_cost in litellm's cost calculator. - """ - rates: Final = model.override_rates if case.response_model_override else model.rates - u: Final = case.usage - prompt_tokens: Final = ( - u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens - + u.cache_write_1h_tokens + u.audio_input_tokens - ) - tiered: Final = prompt_tokens > TIER_THRESHOLD_TOKENS - in_rate: Final = ( - (rates.input_cost_per_token_above_200k_tokens if tiered else None) - or (rates.input_cost_per_token_priority if case.service_tier == "priority" else None) - or (rates.input_cost_per_token_flex if case.service_tier == "flex" else None) - or rates.input_cost_per_token - or 0.0 - ) - out_rate: Final = ( - (rates.output_cost_per_token_above_200k_tokens if tiered else None) - or (rates.output_cost_per_token_priority if case.service_tier == "priority" else None) - or (rates.output_cost_per_token_flex if case.service_tier == "flex" else None) - or rates.output_cost_per_token - or 0.0 - ) - input_cost: Final = ( - u.fresh_input_tokens * in_rate - + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) - + u.cache_write_5m_tokens * (rates.cache_creation_input_token_cost or 0.0) - + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or 0.0) - + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) - ) - output_cost: Final = ( - u.output_tokens * out_rate - + u.reasoning_tokens * (rates.output_cost_per_reasoning_token or out_rate) - + u.audio_output_tokens * (rates.output_cost_per_audio_token or out_rate) - ) - search: Final = rates.search_context_cost_per_query - tool_cost: Final = case.billed_web_search_calls * ( - search.search_context_size_medium if search and search.search_context_size_medium else 0.0 - ) - return ExpectedCost(input_cost=input_cost, output_cost=output_cost, tool_cost=tool_cost) - - -def expected_cost(model: FrontierModel, case: Case) -> float: - return expected_breakdown(model, case).total + return tuple(case for case in CASES if case.applies_to(model)) def recount_cost( @@ -738,31 +416,23 @@ def image_input_data_url() -> str: IMAGE_INPUT_DATA_URL: Final = image_input_data_url() -def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: - """(prompt_tokens, completion_tokens) the spend row should carry, per the - 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 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 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, - ) - if model.wire == "openai_responses": - return ( - u.fresh_input_tokens + u.cache_read_tokens, - u.output_tokens + u.reasoning_tokens, - ) - return ( - u.fresh_input_tokens - + u.cache_read_tokens - + u.cache_write_5m_tokens - + u.cache_write_1h_tokens - + u.audio_input_tokens, - u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, - ) +class _ExpectedCell(BaseModel): + model_config = ConfigDict(frozen=True) + + spend: float + input_cost: float + output_cost: float + prompt_tokens: int + completion_tokens: int + + +_EXPECTED_ADAPTER: Final = TypeAdapter(dict[str, _ExpectedCell]) +EXPECTED: Final[Mapping[str, _ExpectedCell]] = MappingProxyType( + _EXPECTED_ADAPTER.validate_python(json.loads(EXPECTED_PATH.read_text())) + if EXPECTED_PATH.exists() + else {} +) + + +def expected_key(model: FrontierModel, case: Case) -> str: + return f"{model.map_key}|{case.name}" diff --git a/tests/e2e/cost_calculation/expected.json b/tests/e2e/cost_calculation/expected.json new file mode 100644 index 00000000000..7a92fb2476f --- /dev/null +++ b/tests/e2e/cost_calculation/expected.json @@ -0,0 +1,2004 @@ +{ + "anthropic.claude-sonnet-5-v1:0|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.03128, + "output_cost": 0.0085, + "prompt_tokens": 150, + "spend": 0.03978 + }, + "anthropic.claude-sonnet-5-v1:0|basic": { + "completion_tokens": 40, + "input_cost": 0.0204, + "output_cost": 0.013600000000000001, + "prompt_tokens": 120, + "spend": 0.034 + }, + "anthropic.claude-sonnet-5-v1:0|cache_read": { + "completion_tokens": 30, + "input_cost": 0.01785, + "output_cost": 0.0102, + "prompt_tokens": 150, + "spend": 0.028050000000000002 + }, + "anthropic.claude-sonnet-5-v1:0|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.052700000000000004, + "output_cost": 0.0102, + "prompt_tokens": 150, + "spend": 0.06290000000000001 + }, + "anthropic.claude-sonnet-5-v1:0|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0459, + "output_cost": 0.0102, + "prompt_tokens": 150, + "spend": 0.056100000000000004 + }, + "anthropic.claude-sonnet-5-v1:0|stream": { + "completion_tokens": 40, + "input_cost": 0.0204, + "output_cost": 0.013600000000000001, + "prompt_tokens": 120, + "spend": 0.034 + }, + "anthropic.claude-sonnet-5-v1:0|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.013600000000000001, + "output_cost": 0.0085, + "prompt_tokens": 80, + "spend": 0.0221 + }, + "anthropic.claude-sonnet-5-v1:0|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0204, + "output_cost": 0.013600000000000001, + "prompt_tokens": 120, + "spend": 0.034 + }, + "azure/gpt-5.4-mini|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.03424, + "output_cost": 0.02336, + "prompt_tokens": 155, + "spend": 0.0576 + }, + "azure/gpt-5.4-mini|audio": { + "completion_tokens": 45, + "input_cost": 0.04, + "output_cost": 0.0264, + "prompt_tokens": 125, + "spend": 0.0664 + }, + "azure/gpt-5.4-mini|basic": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.4-mini|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0168, + "output_cost": 0.009600000000000001, + "prompt_tokens": 150, + "spend": 0.0264 + }, + "azure/gpt-5.4-mini|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.049600000000000005, + "output_cost": 0.009600000000000001, + "prompt_tokens": 150, + "spend": 0.0592 + }, + "azure/gpt-5.4-mini|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0432, + "output_cost": 0.009600000000000001, + "prompt_tokens": 150, + "spend": 0.0528 + }, + "azure/gpt-5.4-mini|reasoning": { + "completion_tokens": 100, + "input_cost": 0.016, + "output_cost": 0.0656, + "prompt_tokens": 100, + "spend": 0.0816 + }, + "azure/gpt-5.4-mini|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.4-mini|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0288, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.0448 + }, + "azure/gpt-5.4-mini|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.03264, + "output_cost": 0.01728, + "prompt_tokens": 120, + "spend": 0.049920000000000006 + }, + "azure/gpt-5.4-mini|stream": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.4-mini|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.4-mini|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0128, + "output_cost": 0.008, + "prompt_tokens": 80, + "spend": 0.0208 + }, + "azure/gpt-5.4-mini|tiered": { + "completion_tokens": 30, + "input_cost": 256.00128, + "output_cost": 0.0432, + "prompt_tokens": 200001, + "spend": 256.04448 + }, + "azure/gpt-5.4-mini|tool_call": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.4-mini|web_search": { + "completion_tokens": 30, + "input_cost": 0.016, + "output_cost": 0.009600000000000001, + "prompt_tokens": 100, + "spend": 0.0456 + }, + "azure/gpt-5.6|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.0321, + "output_cost": 0.0219, + "prompt_tokens": 155, + "spend": 0.05399999999999999 + }, + "azure/gpt-5.6|audio": { + "completion_tokens": 45, + "input_cost": 0.0375, + "output_cost": 0.02475, + "prompt_tokens": 125, + "spend": 0.06225 + }, + "azure/gpt-5.6|basic": { + "completion_tokens": 40, + "input_cost": 0.018, + "output_cost": 0.011999999999999999, + "prompt_tokens": 120, + "spend": 0.03 + }, + "azure/gpt-5.6|cache_read": { + "completion_tokens": 30, + "input_cost": 0.01575, + "output_cost": 0.009, + "prompt_tokens": 150, + "spend": 0.02475 + }, + "azure/gpt-5.6|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0465, + "output_cost": 0.009, + "prompt_tokens": 150, + "spend": 0.0555 + }, + "azure/gpt-5.6|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.040499999999999994, + "output_cost": 0.009, + "prompt_tokens": 150, + "spend": 0.049499999999999995 + }, + "azure/gpt-5.6|reasoning": { + "completion_tokens": 100, + "input_cost": 0.015, + "output_cost": 0.0615, + "prompt_tokens": 100, + "spend": 0.0765 + }, + "azure/gpt-5.6|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.6|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.027, + "output_cost": 0.015, + "prompt_tokens": 120, + "spend": 0.041999999999999996 + }, + "azure/gpt-5.6|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.030600000000000002, + "output_cost": 0.0162, + "prompt_tokens": 120, + "spend": 0.0468 + }, + "azure/gpt-5.6|stream": { + "completion_tokens": 40, + "input_cost": 0.018, + "output_cost": 0.011999999999999999, + "prompt_tokens": 120, + "spend": 0.03 + }, + "azure/gpt-5.6|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.019200000000000002, + "output_cost": 0.0128, + "prompt_tokens": 120, + "spend": 0.032 + }, + "azure/gpt-5.6|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.011999999999999999, + "output_cost": 0.0075, + "prompt_tokens": 80, + "spend": 0.019499999999999997 + }, + "azure/gpt-5.6|tiered": { + "completion_tokens": 30, + "input_cost": 240.00119999999998, + "output_cost": 0.0405, + "prompt_tokens": 200001, + "spend": 240.0417 + }, + "azure/gpt-5.6|tool_call": { + "completion_tokens": 40, + "input_cost": 0.018, + "output_cost": 0.011999999999999999, + "prompt_tokens": 120, + "spend": 0.03 + }, + "azure/gpt-5.6|web_search": { + "completion_tokens": 30, + "input_cost": 0.015, + "output_cost": 0.009, + "prompt_tokens": 100, + "spend": 0.044 + }, + "claude-haiku-4-5|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.012880000000000003, + "output_cost": 0.0035000000000000005, + "prompt_tokens": 150, + "spend": 0.016380000000000002 + }, + "claude-haiku-4-5|all_components_anthropic_stream": { + "completion_tokens": 25, + "input_cost": 0.012880000000000003, + "output_cost": 0.0035000000000000005, + "prompt_tokens": 150, + "spend": 0.016380000000000002 + }, + "claude-haiku-4-5|basic": { + "completion_tokens": 40, + "input_cost": 0.008400000000000001, + "output_cost": 0.005600000000000001, + "prompt_tokens": 120, + "spend": 0.014000000000000002 + }, + "claude-haiku-4-5|cache_read": { + "completion_tokens": 30, + "input_cost": 0.007350000000000001, + "output_cost": 0.004200000000000001, + "prompt_tokens": 150, + "spend": 0.011550000000000001 + }, + "claude-haiku-4-5|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.021700000000000004, + "output_cost": 0.004200000000000001, + "prompt_tokens": 150, + "spend": 0.025900000000000006 + }, + "claude-haiku-4-5|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0189, + "output_cost": 0.004200000000000001, + "prompt_tokens": 150, + "spend": 0.023100000000000002 + }, + "claude-haiku-4-5|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.006, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.01 + }, + "claude-haiku-4-5|stream": { + "completion_tokens": 40, + "input_cost": 0.008400000000000001, + "output_cost": 0.005600000000000001, + "prompt_tokens": 120, + "spend": 0.014000000000000002 + }, + "claude-haiku-4-5|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.006, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.01 + }, + "claude-haiku-4-5|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.005600000000000001, + "output_cost": 0.0035000000000000005, + "prompt_tokens": 80, + "spend": 0.0091 + }, + "claude-haiku-4-5|tool_call": { + "completion_tokens": 40, + "input_cost": 0.008400000000000001, + "output_cost": 0.005600000000000001, + "prompt_tokens": 120, + "spend": 0.014000000000000002 + }, + "claude-haiku-4-5|web_search": { + "completion_tokens": 30, + "input_cost": 0.007000000000000001, + "output_cost": 0.004200000000000001, + "prompt_tokens": 100, + "spend": 0.0712 + }, + "claude-opus-5|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.0092, + "output_cost": 0.0025, + "prompt_tokens": 150, + "spend": 0.0117 + }, + "claude-opus-5|all_components_anthropic_stream": { + "completion_tokens": 25, + "input_cost": 0.0092, + "output_cost": 0.0025, + "prompt_tokens": 150, + "spend": 0.0117 + }, + "claude-opus-5|basic": { + "completion_tokens": 40, + "input_cost": 0.006, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.01 + }, + "claude-opus-5|cache_read": { + "completion_tokens": 30, + "input_cost": 0.00525, + "output_cost": 0.003, + "prompt_tokens": 150, + "spend": 0.00825 + }, + "claude-opus-5|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0155, + "output_cost": 0.003, + "prompt_tokens": 150, + "spend": 0.0185 + }, + "claude-opus-5|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.013500000000000002, + "output_cost": 0.003, + "prompt_tokens": 150, + "spend": 0.0165 + }, + "claude-opus-5|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.007200000000000001, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 120, + "spend": 0.012 + }, + "claude-opus-5|stream": { + "completion_tokens": 40, + "input_cost": 0.006, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.01 + }, + "claude-opus-5|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.007200000000000001, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 120, + "spend": 0.012 + }, + "claude-opus-5|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.004, + "output_cost": 0.0025, + "prompt_tokens": 80, + "spend": 0.006500000000000001 + }, + "claude-opus-5|tool_call": { + "completion_tokens": 40, + "input_cost": 0.006, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.01 + }, + "claude-opus-5|web_search": { + "completion_tokens": 30, + "input_cost": 0.005, + "output_cost": 0.003, + "prompt_tokens": 100, + "spend": 0.068 + }, + "claude-sonnet-5|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.011040000000000001, + "output_cost": 0.0030000000000000005, + "prompt_tokens": 150, + "spend": 0.014040000000000002 + }, + "claude-sonnet-5|all_components_anthropic_stream": { + "completion_tokens": 25, + "input_cost": 0.011040000000000001, + "output_cost": 0.0030000000000000005, + "prompt_tokens": 150, + "spend": 0.014040000000000002 + }, + "claude-sonnet-5|basic": { + "completion_tokens": 40, + "input_cost": 0.007200000000000001, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 120, + "spend": 0.012 + }, + "claude-sonnet-5|cache_read": { + "completion_tokens": 30, + "input_cost": 0.006300000000000001, + "output_cost": 0.0036000000000000003, + "prompt_tokens": 150, + "spend": 0.0099 + }, + "claude-sonnet-5|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.018600000000000002, + "output_cost": 0.0036000000000000003, + "prompt_tokens": 150, + "spend": 0.0222 + }, + "claude-sonnet-5|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.016200000000000003, + "output_cost": 0.0036000000000000003, + "prompt_tokens": 150, + "spend": 0.0198 + }, + "claude-sonnet-5|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.008400000000000001, + "output_cost": 0.005600000000000001, + "prompt_tokens": 120, + "spend": 0.014000000000000002 + }, + "claude-sonnet-5|stream": { + "completion_tokens": 40, + "input_cost": 0.007200000000000001, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 120, + "spend": 0.012 + }, + "claude-sonnet-5|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.008400000000000001, + "output_cost": 0.005600000000000001, + "prompt_tokens": 120, + "spend": 0.014000000000000002 + }, + "claude-sonnet-5|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0030000000000000005, + "prompt_tokens": 80, + "spend": 0.007800000000000001 + }, + "claude-sonnet-5|tool_call": { + "completion_tokens": 40, + "input_cost": 0.007200000000000001, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 120, + "spend": 0.012 + }, + "claude-sonnet-5|web_search": { + "completion_tokens": 30, + "input_cost": 0.006000000000000001, + "output_cost": 0.0036000000000000003, + "prompt_tokens": 100, + "spend": 0.0696 + }, + "fireworks_ai/deepseek-v4p1-flash|all_components_fireworks": { + "completion_tokens": 25, + "input_cost": 0.011760000000000001, + "output_cost": 0.007000000000000001, + "prompt_tokens": 120, + "spend": 0.018760000000000002 + }, + "fireworks_ai/deepseek-v4p1-flash|audio": { + "completion_tokens": 45, + "input_cost": 0.030500000000000003, + "output_cost": 0.019950000000000002, + "prompt_tokens": 125, + "spend": 0.05045000000000001 + }, + "fireworks_ai/deepseek-v4p1-flash|basic": { + "completion_tokens": 40, + "input_cost": 0.016800000000000002, + "output_cost": 0.011200000000000002, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "fireworks_ai/deepseek-v4p1-flash|cache_read": { + "completion_tokens": 30, + "input_cost": 0.014700000000000001, + "output_cost": 0.008400000000000001, + "prompt_tokens": 150, + "spend": 0.023100000000000002 + }, + "fireworks_ai/deepseek-v4p1-flash|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0368, + "output_cost": 0.008400000000000001, + "prompt_tokens": 150, + "spend": 0.045200000000000004 + }, + "fireworks_ai/deepseek-v4p1-flash|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0324, + "output_cost": 0.008400000000000001, + "prompt_tokens": 150, + "spend": 0.0408 + }, + "fireworks_ai/deepseek-v4p1-flash|reasoning": { + "completion_tokens": 100, + "input_cost": 0.014000000000000002, + "output_cost": 0.0469, + "prompt_tokens": 100, + "spend": 0.060899999999999996 + }, + "fireworks_ai/deepseek-v4p1-flash|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.014400000000000001, + "output_cost": 0.009600000000000001, + "prompt_tokens": 120, + "spend": 0.024 + }, + "fireworks_ai/deepseek-v4p1-flash|stream": { + "completion_tokens": 40, + "input_cost": 0.016800000000000002, + "output_cost": 0.011200000000000002, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "fireworks_ai/deepseek-v4p1-flash|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.014400000000000001, + "output_cost": 0.009600000000000001, + "prompt_tokens": 120, + "spend": 0.024 + }, + "fireworks_ai/deepseek-v4p1-flash|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.011200000000000002, + "output_cost": 0.007000000000000001, + "prompt_tokens": 80, + "spend": 0.0182 + }, + "fireworks_ai/deepseek-v4p1-flash|tool_call": { + "completion_tokens": 40, + "input_cost": 0.016800000000000002, + "output_cost": 0.011200000000000002, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "fireworks_ai/deepseek-v4p1-flash|web_search": { + "completion_tokens": 30, + "input_cost": 0.014000000000000002, + "output_cost": 0.008400000000000001, + "prompt_tokens": 100, + "spend": 0.04240000000000001 + }, + "fireworks_ai/kimi-k3|all_components_fireworks": { + "completion_tokens": 25, + "input_cost": 0.01008, + "output_cost": 0.006000000000000001, + "prompt_tokens": 120, + "spend": 0.01608 + }, + "fireworks_ai/kimi-k3|audio": { + "completion_tokens": 45, + "input_cost": 0.028500000000000004, + "output_cost": 0.01875, + "prompt_tokens": 125, + "spend": 0.04725 + }, + "fireworks_ai/kimi-k3|basic": { + "completion_tokens": 40, + "input_cost": 0.014400000000000001, + "output_cost": 0.009600000000000001, + "prompt_tokens": 120, + "spend": 0.024 + }, + "fireworks_ai/kimi-k3|cache_read": { + "completion_tokens": 30, + "input_cost": 0.012600000000000002, + "output_cost": 0.007200000000000001, + "prompt_tokens": 150, + "spend": 0.0198 + }, + "fireworks_ai/kimi-k3|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.035, + "output_cost": 0.007200000000000001, + "prompt_tokens": 150, + "spend": 0.0422 + }, + "fireworks_ai/kimi-k3|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.030600000000000002, + "output_cost": 0.007200000000000001, + "prompt_tokens": 150, + "spend": 0.0378 + }, + "fireworks_ai/kimi-k3|reasoning": { + "completion_tokens": 100, + "input_cost": 0.012000000000000002, + "output_cost": 0.0457, + "prompt_tokens": 100, + "spend": 0.0577 + }, + "fireworks_ai/kimi-k3|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.015600000000000003, + "output_cost": 0.010400000000000001, + "prompt_tokens": 120, + "spend": 0.026000000000000002 + }, + "fireworks_ai/kimi-k3|stream": { + "completion_tokens": 40, + "input_cost": 0.014400000000000001, + "output_cost": 0.009600000000000001, + "prompt_tokens": 120, + "spend": 0.024 + }, + "fireworks_ai/kimi-k3|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.015600000000000003, + "output_cost": 0.010400000000000001, + "prompt_tokens": 120, + "spend": 0.026000000000000002 + }, + "fireworks_ai/kimi-k3|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.009600000000000001, + "output_cost": 0.006000000000000001, + "prompt_tokens": 80, + "spend": 0.015600000000000003 + }, + "fireworks_ai/kimi-k3|tool_call": { + "completion_tokens": 40, + "input_cost": 0.014400000000000001, + "output_cost": 0.009600000000000001, + "prompt_tokens": 120, + "spend": 0.024 + }, + "fireworks_ai/kimi-k3|web_search": { + "completion_tokens": 30, + "input_cost": 0.012000000000000002, + "output_cost": 0.007200000000000001, + "prompt_tokens": 100, + "spend": 0.0392 + }, + "fireworks_ai/qwen3p8-max|all_components_fireworks": { + "completion_tokens": 25, + "input_cost": 0.010920000000000001, + "output_cost": 0.006500000000000001, + "prompt_tokens": 120, + "spend": 0.01742 + }, + "fireworks_ai/qwen3p8-max|audio": { + "completion_tokens": 45, + "input_cost": 0.029500000000000002, + "output_cost": 0.01935, + "prompt_tokens": 125, + "spend": 0.048850000000000005 + }, + "fireworks_ai/qwen3p8-max|basic": { + "completion_tokens": 40, + "input_cost": 0.015600000000000003, + "output_cost": 0.010400000000000001, + "prompt_tokens": 120, + "spend": 0.026000000000000002 + }, + "fireworks_ai/qwen3p8-max|cache_read": { + "completion_tokens": 30, + "input_cost": 0.01365, + "output_cost": 0.007800000000000001, + "prompt_tokens": 150, + "spend": 0.021450000000000004 + }, + "fireworks_ai/qwen3p8-max|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0359, + "output_cost": 0.007800000000000001, + "prompt_tokens": 150, + "spend": 0.0437 + }, + "fireworks_ai/qwen3p8-max|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0315, + "output_cost": 0.007800000000000001, + "prompt_tokens": 150, + "spend": 0.0393 + }, + "fireworks_ai/qwen3p8-max|reasoning": { + "completion_tokens": 100, + "input_cost": 0.013000000000000001, + "output_cost": 0.0463, + "prompt_tokens": 100, + "spend": 0.059300000000000005 + }, + "fireworks_ai/qwen3p8-max|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.016800000000000002, + "output_cost": 0.011200000000000002, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "fireworks_ai/qwen3p8-max|stream": { + "completion_tokens": 40, + "input_cost": 0.015600000000000003, + "output_cost": 0.010400000000000001, + "prompt_tokens": 120, + "spend": 0.026000000000000002 + }, + "fireworks_ai/qwen3p8-max|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.016800000000000002, + "output_cost": 0.011200000000000002, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "fireworks_ai/qwen3p8-max|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.010400000000000001, + "output_cost": 0.006500000000000001, + "prompt_tokens": 80, + "spend": 0.016900000000000002 + }, + "fireworks_ai/qwen3p8-max|tool_call": { + "completion_tokens": 40, + "input_cost": 0.015600000000000003, + "output_cost": 0.010400000000000001, + "prompt_tokens": 120, + "spend": 0.026000000000000002 + }, + "fireworks_ai/qwen3p8-max|web_search": { + "completion_tokens": 30, + "input_cost": 0.013000000000000001, + "output_cost": 0.007800000000000001, + "prompt_tokens": 100, + "spend": 0.0408 + }, + "gemini-3.1-pro-preview|all_components_gemini": { + "completion_tokens": 43, + "input_cost": 0.023940000000000003, + "output_cost": 0.030660000000000003, + "prompt_tokens": 125, + "spend": 0.05460000000000001 + }, + "gemini-3.1-pro-preview|audio": { + "completion_tokens": 45, + "input_cost": 0.052500000000000005, + "output_cost": 0.03465, + "prompt_tokens": 125, + "spend": 0.08715 + }, + "gemini-3.1-pro-preview|basic": { + "completion_tokens": 40, + "input_cost": 0.0252, + "output_cost": 0.016800000000000002, + "prompt_tokens": 120, + "spend": 0.042 + }, + "gemini-3.1-pro-preview|cache_read": { + "completion_tokens": 30, + "input_cost": 0.02205, + "output_cost": 0.0126, + "prompt_tokens": 150, + "spend": 0.03465 + }, + "gemini-3.1-pro-preview|prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.2, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.2 + }, + "gemini-3.1-pro-preview|reasoning": { + "completion_tokens": 100, + "input_cost": 0.021, + "output_cost": 0.0861, + "prompt_tokens": 100, + "spend": 0.1071 + }, + "gemini-3.1-pro-preview|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.024, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.04 + }, + "gemini-3.1-pro-preview|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0378, + "output_cost": 0.020999999999999998, + "prompt_tokens": 120, + "spend": 0.0588 + }, + "gemini-3.1-pro-preview|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.04284, + "output_cost": 0.02268, + "prompt_tokens": 120, + "spend": 0.06552 + }, + "gemini-3.1-pro-preview|stream": { + "completion_tokens": 40, + "input_cost": 0.0252, + "output_cost": 0.016800000000000002, + "prompt_tokens": 120, + "spend": 0.042 + }, + "gemini-3.1-pro-preview|stream_prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.2, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.2 + }, + "gemini-3.1-pro-preview|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.024, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.04 + }, + "gemini-3.1-pro-preview|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.016800000000000002, + "output_cost": 0.0105, + "prompt_tokens": 80, + "spend": 0.027300000000000005 + }, + "gemini-3.1-pro-preview|tiered": { + "completion_tokens": 30, + "input_cost": 336.00168, + "output_cost": 0.0567, + "prompt_tokens": 200001, + "spend": 336.05838 + }, + "gemini-3.1-pro-preview|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0252, + "output_cost": 0.016800000000000002, + "prompt_tokens": 120, + "spend": 0.042 + }, + "gemini-3.1-pro-preview|web_search": { + "completion_tokens": 30, + "input_cost": 0.021, + "output_cost": 0.0126, + "prompt_tokens": 100, + "spend": 0.0936 + }, + "gemini-3.8-flash|all_components_gemini": { + "completion_tokens": 43, + "input_cost": 0.022799999999999997, + "output_cost": 0.0292, + "prompt_tokens": 125, + "spend": 0.052 + }, + "gemini-3.8-flash|audio": { + "completion_tokens": 45, + "input_cost": 0.05, + "output_cost": 0.033, + "prompt_tokens": 125, + "spend": 0.083 + }, + "gemini-3.8-flash|basic": { + "completion_tokens": 40, + "input_cost": 0.024, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.04 + }, + "gemini-3.8-flash|cache_read": { + "completion_tokens": 30, + "input_cost": 0.021, + "output_cost": 0.012, + "prompt_tokens": 150, + "spend": 0.033 + }, + "gemini-3.8-flash|prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.21000000000000002, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.21000000000000002 + }, + "gemini-3.8-flash|reasoning": { + "completion_tokens": 100, + "input_cost": 0.02, + "output_cost": 0.082, + "prompt_tokens": 100, + "spend": 0.10200000000000001 + }, + "gemini-3.8-flash|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0252, + "output_cost": 0.016800000000000002, + "prompt_tokens": 120, + "spend": 0.042 + }, + "gemini-3.8-flash|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.036, + "output_cost": 0.02, + "prompt_tokens": 120, + "spend": 0.055999999999999994 + }, + "gemini-3.8-flash|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.0408, + "output_cost": 0.0216, + "prompt_tokens": 120, + "spend": 0.062400000000000004 + }, + "gemini-3.8-flash|stream": { + "completion_tokens": 40, + "input_cost": 0.024, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.04 + }, + "gemini-3.8-flash|stream_prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.21000000000000002, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.21000000000000002 + }, + "gemini-3.8-flash|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0252, + "output_cost": 0.016800000000000002, + "prompt_tokens": 120, + "spend": 0.042 + }, + "gemini-3.8-flash|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.016, + "output_cost": 0.01, + "prompt_tokens": 80, + "spend": 0.026000000000000002 + }, + "gemini-3.8-flash|tiered": { + "completion_tokens": 30, + "input_cost": 320.0016, + "output_cost": 0.054, + "prompt_tokens": 200001, + "spend": 320.05559999999997 + }, + "gemini-3.8-flash|tool_call": { + "completion_tokens": 40, + "input_cost": 0.024, + "output_cost": 0.016, + "prompt_tokens": 120, + "spend": 0.04 + }, + "gemini-3.8-flash|web_search": { + "completion_tokens": 30, + "input_cost": 0.02, + "output_cost": 0.012, + "prompt_tokens": 100, + "spend": 0.092 + }, + "gemini/gemini-3.1-pro-preview|all_components_gemini": { + "completion_tokens": 43, + "input_cost": 0.010260000000000002, + "output_cost": 0.01314, + "prompt_tokens": 125, + "spend": 0.023400000000000004 + }, + "gemini/gemini-3.1-pro-preview|audio": { + "completion_tokens": 45, + "input_cost": 0.0225, + "output_cost": 0.014849999999999999, + "prompt_tokens": 125, + "spend": 0.037349999999999994 + }, + "gemini/gemini-3.1-pro-preview|basic": { + "completion_tokens": 40, + "input_cost": 0.0108, + "output_cost": 0.007200000000000001, + "prompt_tokens": 120, + "spend": 0.018000000000000002 + }, + "gemini/gemini-3.1-pro-preview|cache_read": { + "completion_tokens": 30, + "input_cost": 0.009450000000000002, + "output_cost": 0.0054, + "prompt_tokens": 150, + "spend": 0.014850000000000002 + }, + "gemini/gemini-3.1-pro-preview|prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.08, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.08 + }, + "gemini/gemini-3.1-pro-preview|reasoning": { + "completion_tokens": 100, + "input_cost": 0.009000000000000001, + "output_cost": 0.0369, + "prompt_tokens": 100, + "spend": 0.0459 + }, + "gemini/gemini-3.1-pro-preview|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.009600000000000001, + "output_cost": 0.0064, + "prompt_tokens": 120, + "spend": 0.016 + }, + "gemini/gemini-3.1-pro-preview|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0162, + "output_cost": 0.009000000000000001, + "prompt_tokens": 120, + "spend": 0.0252 + }, + "gemini/gemini-3.1-pro-preview|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.01836, + "output_cost": 0.00972, + "prompt_tokens": 120, + "spend": 0.02808 + }, + "gemini/gemini-3.1-pro-preview|stream": { + "completion_tokens": 40, + "input_cost": 0.0108, + "output_cost": 0.007200000000000001, + "prompt_tokens": 120, + "spend": 0.018000000000000002 + }, + "gemini/gemini-3.1-pro-preview|stream_prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.08, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.08 + }, + "gemini/gemini-3.1-pro-preview|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.009600000000000001, + "output_cost": 0.0064, + "prompt_tokens": 120, + "spend": 0.016 + }, + "gemini/gemini-3.1-pro-preview|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.007200000000000001, + "output_cost": 0.0045000000000000005, + "prompt_tokens": 80, + "spend": 0.011700000000000002 + }, + "gemini/gemini-3.1-pro-preview|tiered": { + "completion_tokens": 30, + "input_cost": 144.00072, + "output_cost": 0.024300000000000002, + "prompt_tokens": 200001, + "spend": 144.02502 + }, + "gemini/gemini-3.1-pro-preview|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0108, + "output_cost": 0.007200000000000001, + "prompt_tokens": 120, + "spend": 0.018000000000000002 + }, + "gemini/gemini-3.1-pro-preview|web_search": { + "completion_tokens": 30, + "input_cost": 0.009000000000000001, + "output_cost": 0.0054, + "prompt_tokens": 100, + "spend": 0.0744 + }, + "gemini/gemini-3.8-flash|all_components_gemini": { + "completion_tokens": 43, + "input_cost": 0.00912, + "output_cost": 0.01168, + "prompt_tokens": 125, + "spend": 0.0208 + }, + "gemini/gemini-3.8-flash|audio": { + "completion_tokens": 45, + "input_cost": 0.02, + "output_cost": 0.0132, + "prompt_tokens": 125, + "spend": 0.0332 + }, + "gemini/gemini-3.8-flash|basic": { + "completion_tokens": 40, + "input_cost": 0.009600000000000001, + "output_cost": 0.0064, + "prompt_tokens": 120, + "spend": 0.016 + }, + "gemini/gemini-3.8-flash|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0084, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 150, + "spend": 0.0132 + }, + "gemini/gemini-3.8-flash|prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.09000000000000001, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.09000000000000001 + }, + "gemini/gemini-3.8-flash|reasoning": { + "completion_tokens": 100, + "input_cost": 0.008, + "output_cost": 0.0328, + "prompt_tokens": 100, + "spend": 0.0408 + }, + "gemini/gemini-3.8-flash|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0108, + "output_cost": 0.007200000000000001, + "prompt_tokens": 120, + "spend": 0.018000000000000002 + }, + "gemini/gemini-3.8-flash|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0144, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.0224 + }, + "gemini/gemini-3.8-flash|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.01632, + "output_cost": 0.00864, + "prompt_tokens": 120, + "spend": 0.024960000000000003 + }, + "gemini/gemini-3.8-flash|stream": { + "completion_tokens": 40, + "input_cost": 0.009600000000000001, + "output_cost": 0.0064, + "prompt_tokens": 120, + "spend": 0.016 + }, + "gemini/gemini-3.8-flash|stream_prompt_blocked": { + "completion_tokens": 0, + "input_cost": 0.09000000000000001, + "output_cost": 0.0, + "prompt_tokens": 1000, + "spend": 0.09000000000000001 + }, + "gemini/gemini-3.8-flash|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0108, + "output_cost": 0.007200000000000001, + "prompt_tokens": 120, + "spend": 0.018000000000000002 + }, + "gemini/gemini-3.8-flash|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0064, + "output_cost": 0.004, + "prompt_tokens": 80, + "spend": 0.0104 + }, + "gemini/gemini-3.8-flash|tiered": { + "completion_tokens": 30, + "input_cost": 128.00064, + "output_cost": 0.0216, + "prompt_tokens": 200001, + "spend": 128.02224 + }, + "gemini/gemini-3.8-flash|tool_call": { + "completion_tokens": 40, + "input_cost": 0.009600000000000001, + "output_cost": 0.0064, + "prompt_tokens": 120, + "spend": 0.016 + }, + "gemini/gemini-3.8-flash|web_search": { + "completion_tokens": 30, + "input_cost": 0.008, + "output_cost": 0.0048000000000000004, + "prompt_tokens": 100, + "spend": 0.0728 + }, + "gpt-5.3-codex|all_components_responses": { + "completion_tokens": 40, + "input_cost": 0.00252, + "output_cost": 0.0037500000000000007, + "prompt_tokens": 120, + "spend": 0.006270000000000001 + }, + "gpt-5.3-codex|basic": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.3-codex|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0031500000000000005, + "output_cost": 0.0018000000000000002, + "prompt_tokens": 150, + "spend": 0.00495 + }, + "gpt-5.3-codex|reasoning": { + "completion_tokens": 100, + "input_cost": 0.0030000000000000005, + "output_cost": 0.0123, + "prompt_tokens": 100, + "spend": 0.015300000000000001 + }, + "gpt-5.3-codex|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.3-codex|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0054, + "output_cost": 0.003, + "prompt_tokens": 120, + "spend": 0.008400000000000001 + }, + "gpt-5.3-codex|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.00612, + "output_cost": 0.00324, + "prompt_tokens": 120, + "spend": 0.00936 + }, + "gpt-5.3-codex|stream": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.3-codex|stream_incomplete": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.3-codex|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.3-codex|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0015000000000000002, + "prompt_tokens": 80, + "spend": 0.0039000000000000007 + }, + "gpt-5.3-codex|stream_unvalidated": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.3-codex|tiered": { + "completion_tokens": 30, + "input_cost": 48.000240000000005, + "output_cost": 0.0081, + "prompt_tokens": 200001, + "spend": 48.008340000000004 + }, + "gpt-5.3-codex|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.3-codex|web_search": { + "completion_tokens": 30, + "input_cost": 0.0030000000000000005, + "output_cost": 0.0018000000000000002, + "prompt_tokens": 100, + "spend": 0.0648 + }, + "gpt-5.4-mini|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.00856, + "output_cost": 0.00584, + "prompt_tokens": 155, + "spend": 0.0144 + }, + "gpt-5.4-mini|audio": { + "completion_tokens": 45, + "input_cost": 0.01, + "output_cost": 0.0066, + "prompt_tokens": 125, + "spend": 0.0166 + }, + "gpt-5.4-mini|basic": { + "completion_tokens": 40, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0032, + "prompt_tokens": 120, + "spend": 0.008 + }, + "gpt-5.4-mini|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0042, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 150, + "spend": 0.0066 + }, + "gpt-5.4-mini|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.012400000000000001, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 150, + "spend": 0.0148 + }, + "gpt-5.4-mini|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0108, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 150, + "spend": 0.0132 + }, + "gpt-5.4-mini|reasoning": { + "completion_tokens": 100, + "input_cost": 0.004, + "output_cost": 0.0164, + "prompt_tokens": 100, + "spend": 0.0204 + }, + "gpt-5.4-mini|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0012000000000000001, + "output_cost": 0.0008, + "prompt_tokens": 120, + "spend": 0.002 + }, + "gpt-5.4-mini|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0072, + "output_cost": 0.004, + "prompt_tokens": 120, + "spend": 0.0112 + }, + "gpt-5.4-mini|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.00816, + "output_cost": 0.00432, + "prompt_tokens": 120, + "spend": 0.012480000000000002 + }, + "gpt-5.4-mini|stream": { + "completion_tokens": 40, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0032, + "prompt_tokens": 120, + "spend": 0.008 + }, + "gpt-5.4-mini|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0012000000000000001, + "output_cost": 0.0008, + "prompt_tokens": 120, + "spend": 0.002 + }, + "gpt-5.4-mini|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0032, + "output_cost": 0.002, + "prompt_tokens": 80, + "spend": 0.0052 + }, + "gpt-5.4-mini|tiered": { + "completion_tokens": 30, + "input_cost": 64.00032, + "output_cost": 0.0108, + "prompt_tokens": 200001, + "spend": 64.01112 + }, + "gpt-5.4-mini|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0032, + "prompt_tokens": 120, + "spend": 0.008 + }, + "gpt-5.4-mini|web_search": { + "completion_tokens": 30, + "input_cost": 0.004, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 100, + "spend": 0.0264 + }, + "gpt-5.5-pro|all_components_responses": { + "completion_tokens": 40, + "input_cost": 0.00168, + "output_cost": 0.0025, + "prompt_tokens": 120, + "spend": 0.00418 + }, + "gpt-5.5-pro|basic": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.5-pro|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0021, + "output_cost": 0.0012000000000000001, + "prompt_tokens": 150, + "spend": 0.0033 + }, + "gpt-5.5-pro|reasoning": { + "completion_tokens": 100, + "input_cost": 0.002, + "output_cost": 0.0082, + "prompt_tokens": 100, + "spend": 0.0102 + }, + "gpt-5.5-pro|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.5-pro|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0036, + "output_cost": 0.002, + "prompt_tokens": 120, + "spend": 0.0056 + }, + "gpt-5.5-pro|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.00408, + "output_cost": 0.00216, + "prompt_tokens": 120, + "spend": 0.006240000000000001 + }, + "gpt-5.5-pro|stream": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.5-pro|stream_incomplete": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.5-pro|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0036000000000000003, + "output_cost": 0.0024000000000000002, + "prompt_tokens": 120, + "spend": 0.006 + }, + "gpt-5.5-pro|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0016, + "output_cost": 0.001, + "prompt_tokens": 80, + "spend": 0.0026 + }, + "gpt-5.5-pro|stream_unvalidated": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.5-pro|tiered": { + "completion_tokens": 30, + "input_cost": 32.00016, + "output_cost": 0.0054, + "prompt_tokens": 200001, + "spend": 32.00556 + }, + "gpt-5.5-pro|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0024000000000000002, + "output_cost": 0.0016, + "prompt_tokens": 120, + "spend": 0.004 + }, + "gpt-5.5-pro|web_search": { + "completion_tokens": 30, + "input_cost": 0.002, + "output_cost": 0.0012000000000000001, + "prompt_tokens": 100, + "spend": 0.06319999999999999 + }, + "gpt-5.6|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.00214, + "output_cost": 0.00146, + "prompt_tokens": 155, + "spend": 0.0036 + }, + "gpt-5.6|audio": { + "completion_tokens": 45, + "input_cost": 0.0025, + "output_cost": 0.00165, + "prompt_tokens": 125, + "spend": 0.00415 + }, + "gpt-5.6|basic": { + "completion_tokens": 40, + "input_cost": 0.0012000000000000001, + "output_cost": 0.0008, + "prompt_tokens": 120, + "spend": 0.002 + }, + "gpt-5.6|cache_read": { + "completion_tokens": 30, + "input_cost": 0.00105, + "output_cost": 0.0006000000000000001, + "prompt_tokens": 150, + "spend": 0.00165 + }, + "gpt-5.6|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0031000000000000003, + "output_cost": 0.0006000000000000001, + "prompt_tokens": 150, + "spend": 0.0037 + }, + "gpt-5.6|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.0027, + "output_cost": 0.0006000000000000001, + "prompt_tokens": 150, + "spend": 0.0033 + }, + "gpt-5.6|reasoning": { + "completion_tokens": 100, + "input_cost": 0.001, + "output_cost": 0.0041, + "prompt_tokens": 100, + "spend": 0.0051 + }, + "gpt-5.6|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0032, + "prompt_tokens": 120, + "spend": 0.008 + }, + "gpt-5.6|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.0018, + "output_cost": 0.001, + "prompt_tokens": 120, + "spend": 0.0028 + }, + "gpt-5.6|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.00204, + "output_cost": 0.00108, + "prompt_tokens": 120, + "spend": 0.0031200000000000004 + }, + "gpt-5.6|stream": { + "completion_tokens": 40, + "input_cost": 0.0012000000000000001, + "output_cost": 0.0008, + "prompt_tokens": 120, + "spend": 0.002 + }, + "gpt-5.6|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0048000000000000004, + "output_cost": 0.0032, + "prompt_tokens": 120, + "spend": 0.008 + }, + "gpt-5.6|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0008, + "output_cost": 0.0005, + "prompt_tokens": 80, + "spend": 0.0013 + }, + "gpt-5.6|tiered": { + "completion_tokens": 30, + "input_cost": 16.00008, + "output_cost": 0.0027, + "prompt_tokens": 200001, + "spend": 16.00278 + }, + "gpt-5.6|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0012000000000000001, + "output_cost": 0.0008, + "prompt_tokens": 120, + "spend": 0.002 + }, + "gpt-5.6|web_search": { + "completion_tokens": 30, + "input_cost": 0.001, + "output_cost": 0.0006000000000000001, + "prompt_tokens": 100, + "spend": 0.0216 + }, + "meta.llama4-maverick-17b-instruct-v1:0|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.020900000000000002, + "output_cost": 0.0095, + "prompt_tokens": 150, + "spend": 0.030400000000000003 + }, + "meta.llama4-maverick-17b-instruct-v1:0|basic": { + "completion_tokens": 40, + "input_cost": 0.0228, + "output_cost": 0.015200000000000002, + "prompt_tokens": 120, + "spend": 0.038000000000000006 + }, + "meta.llama4-maverick-17b-instruct-v1:0|stream": { + "completion_tokens": 40, + "input_cost": 0.0228, + "output_cost": 0.015200000000000002, + "prompt_tokens": 120, + "spend": 0.038000000000000006 + }, + "meta.llama4-maverick-17b-instruct-v1:0|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.015200000000000002, + "output_cost": 0.0095, + "prompt_tokens": 80, + "spend": 0.0247 + }, + "meta.llama4-maverick-17b-instruct-v1:0|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0228, + "output_cost": 0.015200000000000002, + "prompt_tokens": 120, + "spend": 0.038000000000000006 + }, + "together_ai/moonshotai/Kimi-K3|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.0214, + "output_cost": 0.0146, + "prompt_tokens": 155, + "spend": 0.036 + }, + "together_ai/moonshotai/Kimi-K3|audio": { + "completion_tokens": 45, + "input_cost": 0.025, + "output_cost": 0.0165, + "prompt_tokens": 125, + "spend": 0.0415 + }, + "together_ai/moonshotai/Kimi-K3|basic": { + "completion_tokens": 40, + "input_cost": 0.012, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.02 + }, + "together_ai/moonshotai/Kimi-K3|cache_read": { + "completion_tokens": 30, + "input_cost": 0.0105, + "output_cost": 0.006, + "prompt_tokens": 150, + "spend": 0.0165 + }, + "together_ai/moonshotai/Kimi-K3|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.031, + "output_cost": 0.006, + "prompt_tokens": 150, + "spend": 0.037 + }, + "together_ai/moonshotai/Kimi-K3|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.027000000000000003, + "output_cost": 0.006, + "prompt_tokens": 150, + "spend": 0.033 + }, + "together_ai/moonshotai/Kimi-K3|reasoning": { + "completion_tokens": 100, + "input_cost": 0.01, + "output_cost": 0.041, + "prompt_tokens": 100, + "spend": 0.051000000000000004 + }, + "together_ai/moonshotai/Kimi-K3|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0132, + "output_cost": 0.0088, + "prompt_tokens": 120, + "spend": 0.022 + }, + "together_ai/moonshotai/Kimi-K3|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.018000000000000002, + "output_cost": 0.01, + "prompt_tokens": 120, + "spend": 0.028000000000000004 + }, + "together_ai/moonshotai/Kimi-K3|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.0204, + "output_cost": 0.0108, + "prompt_tokens": 120, + "spend": 0.031200000000000002 + }, + "together_ai/moonshotai/Kimi-K3|stream": { + "completion_tokens": 40, + "input_cost": 0.012, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.02 + }, + "together_ai/moonshotai/Kimi-K3|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.0132, + "output_cost": 0.0088, + "prompt_tokens": 120, + "spend": 0.022 + }, + "together_ai/moonshotai/Kimi-K3|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.008, + "output_cost": 0.005, + "prompt_tokens": 80, + "spend": 0.013000000000000001 + }, + "together_ai/moonshotai/Kimi-K3|tiered": { + "completion_tokens": 30, + "input_cost": 160.0008, + "output_cost": 0.027000000000000003, + "prompt_tokens": 200001, + "spend": 160.02779999999998 + }, + "together_ai/moonshotai/Kimi-K3|tool_call": { + "completion_tokens": 40, + "input_cost": 0.012, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.02 + }, + "together_ai/moonshotai/Kimi-K3|web_search": { + "completion_tokens": 30, + "input_cost": 0.01, + "output_cost": 0.006, + "prompt_tokens": 100, + "spend": 0.036000000000000004 + }, + "together_ai/zai-org/GLM-5.3|all_components_chat": { + "completion_tokens": 43, + "input_cost": 0.023540000000000002, + "output_cost": 0.01606, + "prompt_tokens": 155, + "spend": 0.0396 + }, + "together_ai/zai-org/GLM-5.3|audio": { + "completion_tokens": 45, + "input_cost": 0.027500000000000004, + "output_cost": 0.01815, + "prompt_tokens": 125, + "spend": 0.04565 + }, + "together_ai/zai-org/GLM-5.3|basic": { + "completion_tokens": 40, + "input_cost": 0.0132, + "output_cost": 0.0088, + "prompt_tokens": 120, + "spend": 0.022 + }, + "together_ai/zai-org/GLM-5.3|cache_read": { + "completion_tokens": 30, + "input_cost": 0.011550000000000001, + "output_cost": 0.0066, + "prompt_tokens": 150, + "spend": 0.01815 + }, + "together_ai/zai-org/GLM-5.3|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.034100000000000005, + "output_cost": 0.0066, + "prompt_tokens": 150, + "spend": 0.04070000000000001 + }, + "together_ai/zai-org/GLM-5.3|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.029699999999999997, + "output_cost": 0.0066, + "prompt_tokens": 150, + "spend": 0.0363 + }, + "together_ai/zai-org/GLM-5.3|reasoning": { + "completion_tokens": 100, + "input_cost": 0.011000000000000001, + "output_cost": 0.0451, + "prompt_tokens": 100, + "spend": 0.056100000000000004 + }, + "together_ai/zai-org/GLM-5.3|response_model_override": { + "completion_tokens": 40, + "input_cost": 0.012, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.02 + }, + "together_ai/zai-org/GLM-5.3|service_tier_flex": { + "completion_tokens": 40, + "input_cost": 0.019799999999999998, + "output_cost": 0.011000000000000001, + "prompt_tokens": 120, + "spend": 0.0308 + }, + "together_ai/zai-org/GLM-5.3|service_tier_priority": { + "completion_tokens": 40, + "input_cost": 0.022439999999999998, + "output_cost": 0.01188, + "prompt_tokens": 120, + "spend": 0.034319999999999996 + }, + "together_ai/zai-org/GLM-5.3|stream": { + "completion_tokens": 40, + "input_cost": 0.0132, + "output_cost": 0.0088, + "prompt_tokens": 120, + "spend": 0.022 + }, + "together_ai/zai-org/GLM-5.3|stream_response_model_override": { + "completion_tokens": 40, + "input_cost": 0.012, + "output_cost": 0.008, + "prompt_tokens": 120, + "spend": 0.02 + }, + "together_ai/zai-org/GLM-5.3|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.0088, + "output_cost": 0.0055000000000000005, + "prompt_tokens": 80, + "spend": 0.0143 + }, + "together_ai/zai-org/GLM-5.3|tiered": { + "completion_tokens": 30, + "input_cost": 176.00088, + "output_cost": 0.0297, + "prompt_tokens": 200001, + "spend": 176.03058 + }, + "together_ai/zai-org/GLM-5.3|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0132, + "output_cost": 0.0088, + "prompt_tokens": 120, + "spend": 0.022 + }, + "together_ai/zai-org/GLM-5.3|web_search": { + "completion_tokens": 30, + "input_cost": 0.011000000000000001, + "output_cost": 0.0066, + "prompt_tokens": 100, + "spend": 0.0376 + }, + "us.anthropic.claude-opus-5-v1:0|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.033120000000000004, + "output_cost": 0.009000000000000001, + "prompt_tokens": 150, + "spend": 0.042120000000000005 + }, + "us.anthropic.claude-opus-5-v1:0|basic": { + "completion_tokens": 40, + "input_cost": 0.0216, + "output_cost": 0.014400000000000001, + "prompt_tokens": 120, + "spend": 0.036000000000000004 + }, + "us.anthropic.claude-opus-5-v1:0|cache_read": { + "completion_tokens": 30, + "input_cost": 0.018900000000000004, + "output_cost": 0.0108, + "prompt_tokens": 150, + "spend": 0.029700000000000004 + }, + "us.anthropic.claude-opus-5-v1:0|cache_write_1h": { + "completion_tokens": 30, + "input_cost": 0.0558, + "output_cost": 0.0108, + "prompt_tokens": 150, + "spend": 0.0666 + }, + "us.anthropic.claude-opus-5-v1:0|cache_write_5m": { + "completion_tokens": 30, + "input_cost": 0.048600000000000004, + "output_cost": 0.0108, + "prompt_tokens": 150, + "spend": 0.05940000000000001 + }, + "us.anthropic.claude-opus-5-v1:0|stream": { + "completion_tokens": 40, + "input_cost": 0.0216, + "output_cost": 0.014400000000000001, + "prompt_tokens": 120, + "spend": 0.036000000000000004 + }, + "us.anthropic.claude-opus-5-v1:0|stream_tool_call": { + "completion_tokens": 25, + "input_cost": 0.014400000000000001, + "output_cost": 0.009000000000000001, + "prompt_tokens": 80, + "spend": 0.023400000000000004 + }, + "us.anthropic.claude-opus-5-v1:0|tool_call": { + "completion_tokens": 40, + "input_cost": 0.0216, + "output_cost": 0.014400000000000001, + "prompt_tokens": 120, + "spend": 0.036000000000000004 + } +} diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py new file mode 100644 index 00000000000..de979f272fe --- /dev/null +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -0,0 +1,189 @@ +"""Golden generator for the cost suite. Run: + + uv run python tests/e2e/cost_calculation/generate_expected.py + +Loads the derived matrix (models x applicable cases), computes the golden for +each exact-spend cell from the rate arithmetic, and writes ``expected.json`` +with sorted keys. Default behaviour adds missing cells and drops stale cells +but never overwrites an existing cell's values (a reviewed golden is +authoritative); ``--rewrite`` recomputes everything. Prints added/removed/kept +counts. +""" + +from __future__ import annotations + +import json +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from cost_matrix import ( # noqa: E402 # path bootstrap before package-local imports + EXPECTED_PATH, + FRONTIER_MODELS, + TIER_THRESHOLD_TOKENS, + Case, + CostMapEntry, + FrontierModel, + cases_for, + expected_key, +) + +# Wires whose response surface reports a real web-search call count; the +# chat-completions wires only expose url_citation annotations, so their billed +# count floors to one. +_EXACT_WEB_SEARCH_WIRES: Final = frozenset( + {"openai_responses", "anthropic_messages", "gemini_generate", "vertex_generate"} +) + + +def billed_web_search_calls(model: FrontierModel, case: Case) -> int: + if case.usage.web_search_calls == 0: + return 0 + return case.usage.web_search_calls if model.wire in _EXACT_WEB_SEARCH_WIRES else 1 + + +@dataclass(frozen=True, slots=True) +class ExpectedCost: + """The expected bill split the way the spend row's cost_breakdown reports + it: the gross input component (cache reads/writes folded in), the output + component, and the tool-usage component.""" + + input_cost: float + output_cost: float + tool_cost: float + + @property + def total(self) -> float: + return self.input_cost + self.output_cost + self.tool_cost + + +def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: + """Literal arithmetic on the test-map rates over the scripted token counts. + + Input = fresh*in + read*read + 5m*create + 1h*create_1h + audio_in*audio_in; + output = text*out + reasoning*reasoning + audio_out*audio_out; plus the + billed web-search calls at the medium search-context rate. Above-threshold + swaps every input/output rate to its ``_above_200k_tokens`` variant when + total prompt tokens exceed the threshold; a service tier swaps input/output + to the tier's variants, falling back to the base rate when a variant is + unset -- mirroring _get_token_base_cost in litellm's cost calculator. + """ + rates: Final[CostMapEntry] = model.override_rates if case.response_model_override else model.rates + u: Final = case.usage + prompt_tokens: Final = ( + u.fresh_input_tokens + u.cache_read_tokens + u.cache_write_5m_tokens + + u.cache_write_1h_tokens + u.audio_input_tokens + ) + tiered: Final = prompt_tokens > TIER_THRESHOLD_TOKENS + in_rate: Final = ( + (rates.input_cost_per_token_above_200k_tokens if tiered else None) + or (rates.input_cost_per_token_priority if case.service_tier == "priority" else None) + or (rates.input_cost_per_token_flex if case.service_tier == "flex" else None) + or rates.input_cost_per_token + or 0.0 + ) + out_rate: Final = ( + (rates.output_cost_per_token_above_200k_tokens if tiered else None) + or (rates.output_cost_per_token_priority if case.service_tier == "priority" else None) + or (rates.output_cost_per_token_flex if case.service_tier == "flex" else None) + or rates.output_cost_per_token + or 0.0 + ) + # The biller charges cache writes at the input rate when the entry carries + # no cache_creation rate (cost_calculator.py:2452), and at the 5m write + # rate when the 1h variant is unset; cache reads bill only at their own + # rate (zero when the entry lacks one). + write_5m_rate: Final = rates.cache_creation_input_token_cost or in_rate + input_cost: Final = ( + u.fresh_input_tokens * in_rate + + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) + + u.cache_write_5m_tokens * write_5m_rate + + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or write_5m_rate) + + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) + ) + output_cost: Final = ( + u.output_tokens * out_rate + + u.reasoning_tokens * (rates.output_cost_per_reasoning_token or out_rate) + + u.audio_output_tokens * (rates.output_cost_per_audio_token or out_rate) + ) + search: Final = rates.search_context_cost_per_query + tool_cost: Final = billed_web_search_calls(model, case) * ( + search.search_context_size_medium if search and search.search_context_size_medium else 0.0 + ) + return ExpectedCost(input_cost=input_cost, output_cost=output_cost, tool_cost=tool_cost) + + +def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: + """(prompt_tokens, completion_tokens) the spend row should carry, per the + 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 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 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, + ) + if model.wire == "openai_responses": + return ( + u.fresh_input_tokens + u.cache_read_tokens, + u.output_tokens + u.reasoning_tokens, + ) + return ( + u.fresh_input_tokens + + u.cache_read_tokens + + u.cache_write_5m_tokens + + u.cache_write_1h_tokens + + u.audio_input_tokens, + u.output_tokens + u.reasoning_tokens + u.audio_output_tokens, + ) + + +def _proposed() -> dict[str, dict[str, object]]: + return { + expected_key(model, case): ( + lambda breakdown, tokens: { + "spend": breakdown.total, + "input_cost": breakdown.input_cost, + "output_cost": breakdown.output_cost, + "prompt_tokens": tokens[0], + "completion_tokens": tokens[1], + } + )(expected_breakdown(model, case), expected_token_columns(model, case)) + for model in FRONTIER_MODELS + for case in cases_for(model) + if case.exact_spend + } + + +def main() -> None: + rewrite: Final = "--rewrite" in sys.argv[1:] + proposed: Final = _proposed() + existing: Final = ( + json.loads(EXPECTED_PATH.read_text()) if EXPECTED_PATH.exists() else {} + ) + merged: Final = { + key: (proposed[key] if rewrite or key not in existing else existing[key]) + for key in sorted(proposed) + } + added: Final = sum(1 for key in proposed if key not in existing) + removed: Final = sum(1 for key in existing if key not in proposed) + kept: Final = sum(1 for key in proposed if key in existing and not rewrite) + rewritten: Final = sum(1 for key in proposed if key in existing and rewrite) + EXPECTED_PATH.write_text(json.dumps(merged, indent=2, sort_keys=True) + "\n") + print( + f"expected.json: {added} added, {removed} removed, {kept} kept, " + f"{rewritten} rewritten ({len(merged)} cells)" + ) + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/cost_calculation/test_matrix_data.py b/tests/e2e/cost_calculation/test_matrix_data.py new file mode 100644 index 00000000000..fdbb6ddd293 --- /dev/null +++ b/tests/e2e/cost_calculation/test_matrix_data.py @@ -0,0 +1,64 @@ +"""Freshness checks for the cost suite's data files; markerless, so it runs on +any pytest invocation of the folder without the stack. expected.json is the +oracle: these tests check its key set against the derived matrix, never its +values (the generator proposes, the file decides).""" + +from __future__ import annotations + +from typing import Final + +import pytest + +from cost_matrix import ( + _CASES_FILE, + _COST_MAP, + CASES, + EXPECTED, + FRONTIER_MODELS, + CostMapEntry, + cases_for, + expected_key, +) + + +def test_expected_keys_match_derived_exact_cells() -> None: + derived: Final = { + expected_key(model, case) + for model in FRONTIER_MODELS + for case in cases_for(model) + if case.exact_spend + } + golden: Final = set(EXPECTED) + if derived != golden: + missing: Final = sorted(derived - golden) + stale: Final = sorted(golden - derived) + pytest.fail( + "expected.json is out of sync with the derived matrix; run " + "uv run python tests/e2e/cost_calculation/generate_expected.py " + f"(missing: {missing}; stale: {stale})" + ) + + +def test_deployments_reference_existing_map_keys() -> None: + unknown: Final = sorted( + spec.map_key for spec in _CASES_FILE.deployments if spec.map_key not in _COST_MAP + ) + assert not unknown, f"deployments entries name map keys absent from cost_map.json: {unknown}" + + +def test_requires_rates_are_cost_map_fields() -> None: + fields: Final = set(CostMapEntry.model_fields) + unknown: Final = sorted( + {field for case in CASES for field in case.requires_rates} - fields + ) + assert not unknown, f"requires_rates names that are not CostMapEntry fields: {unknown}" + + +def test_no_two_entries_share_input_rate() -> None: + rates: Final = [ + entry.input_cost_per_token for entry in _COST_MAP.values() + ] + assert len(rates) == len(set(rates)), ( + "two cost_map entries share input_cost_per_token; the suite relies on " + "distinct rates so a wrong-model bill can never coincidentally match" + ) diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py index 0b4f3e1fd37..7cd128ad6fb 100644 --- a/tests/e2e/cost_calculation/test_token_pricing_e2e.py +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -1,7 +1,7 @@ -"""Token-pricing e2e: every (frontier model, pricing-component case) cell runs a -scripted-usage call through a deployment registered on the cost-map proxy, and -the spend row plus response-cost header must equal literal arithmetic on the -test map's rates. +"""Token-pricing e2e: every (map entry, case) cell derived from cost_map.json x +cases.json runs a scripted-usage call through a deployment registered on the +cost-map proxy, and the spend row plus response-cost header must equal the +reviewed golden in expected.json verbatim -- no rate arithmetic lives here. Nothing here touches a real provider or the bundled cost map: the proxy's upstream is the scripted-provider sidecar and its entire cost map is @@ -15,13 +15,13 @@ from typing import Final from conftest import CostCalcClient, cost_rows, register_scenario_deployment from cost_matrix import ( + EXPECTED, FRONTIER_MODELS, IMAGE_INPUT_DATA_URL, Case, FrontierModel, cases_for, - expected_cost, - expected_token_columns, + expected_key, recount_cost, ) from e2e_config import unique_marker @@ -110,17 +110,6 @@ class TestTokenPricing: ) assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" - expected: Final = expected_cost(model, case) - if case.exact_spend and not case.stream: - # Streamed responses commit headers before the bill is computed, so - # the x-litellm-response-cost header is asserted only on non-stream - # calls. - assert response.response_cost is not None and cost_rows.approx_equal( - response.response_cost, expected - ), ( - f"x-litellm-response-cost {response.response_cost} != expected {expected}" - ) - row: Final = cost_rows.poll_cost_row_where( client.proxy, scoped_key, @@ -149,16 +138,39 @@ class TestTokenPricing: cost_rows.assert_total_is_sum_of_components(row) return - assert row.spend is not None and cost_rows.approx_equal(row.spend, expected), ( - f"{model.map_key}/{case.name}: spend {row.spend} != expected {expected} " + golden: Final = EXPECTED[expected_key(model, case)] + + if not case.stream: + # Streamed responses commit headers before the bill is computed, so + # the x-litellm-response-cost header is asserted only on non-stream + # calls. + assert response.response_cost is not None and cost_rows.approx_equal( + response.response_cost, golden.spend + ), ( + f"x-litellm-response-cost {response.response_cost} != golden {golden.spend}" + ) + + assert row.spend is not None and cost_rows.approx_equal(row.spend, golden.spend), ( + f"{model.map_key}/{case.name}: spend {row.spend} != golden {golden.spend} " f"(breakdown {row.breakdown.model_dump()})" ) - - prompt_tokens, completion_tokens = expected_token_columns(model, case) - assert row.prompt_tokens == prompt_tokens, ( - f"prompt_tokens {row.prompt_tokens} != {prompt_tokens}" + breakdown: Final = row.breakdown + assert breakdown.input_cost is not None and cost_rows.approx_equal( + breakdown.input_cost, golden.input_cost + ), ( + f"{model.map_key}/{case.name}: gross input_cost {breakdown.input_cost} " + f"!= golden {golden.input_cost}; cached/written tokens billed at the input rate" ) - assert row.completion_tokens == completion_tokens, ( - f"completion_tokens {row.completion_tokens} != {completion_tokens}" + assert breakdown.output_cost is not None and cost_rows.approx_equal( + breakdown.output_cost, golden.output_cost + ), ( + f"{model.map_key}/{case.name}: output_cost {breakdown.output_cost} " + f"!= golden {golden.output_cost}" + ) + assert row.prompt_tokens == golden.prompt_tokens, ( + f"prompt_tokens {row.prompt_tokens} != {golden.prompt_tokens}" + ) + assert row.completion_tokens == golden.completion_tokens, ( + f"completion_tokens {row.completion_tokens} != {golden.completion_tokens}" ) cost_rows.assert_total_is_sum_of_components(row) diff --git a/tests/e2e/cost_calculation/test_wire_formats_e2e.py b/tests/e2e/cost_calculation/test_wire_formats_e2e.py deleted file mode 100644 index a36bb1a8662..00000000000 --- a/tests/e2e/cost_calculation/test_wire_formats_e2e.py +++ /dev/null @@ -1,368 +0,0 @@ -"""Wire-format e2e: one scripted upstream per provider wire, answering with a -usage payload where every token kind the wire can report is nonzero. The spend -row's gross input cost must equal fresh tokens at the input rate plus each cache -and audio component at its own rate -- proving the wire's usage shape landed the -cached tokens inside the total (OpenAI/Gemini) or as separate fields -(Anthropic), and that the biller subtracted them before billing fresh tokens. - -Also covers the Responses API wire (an openai/gpt-5.5-pro deployment bridged by -the proxy to POST /responses) and a streamed Anthropic-messages case. -""" - -from __future__ import annotations - -import pytest -from collections.abc import Mapping -from types import MappingProxyType -from typing import Final - -from conftest import CostCalcClient, cost_rows, register_scenario_deployment -from cost_matrix import ( - FRONTIER_MODELS, - Case, - FrontierModel, - expected_breakdown, - expected_token_columns, -) -from e2e_config import unique_marker -from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatStreamOptions, ChatTool, ChatToolFunction -from scripted_provider import ScriptedUsage - -pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark - -_MODELS: Final[Mapping[str, FrontierModel]] = MappingProxyType( - {model.map_key: model for model in FRONTIER_MODELS} -) - -# One scripted usage per wire, every reportable token kind nonzero. -_WIRE_USAGE: Final[Mapping[str, tuple[str, ScriptedUsage]]] = MappingProxyType({ - "openai_chat": ( - "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, - ), - ), - "openai_responses": ( - "gpt-5.5-pro", - ScriptedUsage( - fresh_input_tokens=80, cache_read_tokens=40, output_tokens=25, reasoning_tokens=15 - ), - ), - "anthropic_messages": ( - "claude-sonnet-5", - ScriptedUsage( - fresh_input_tokens=80, - cache_read_tokens=40, - cache_write_5m_tokens=20, - cache_write_1h_tokens=10, - output_tokens=25, - ), - ), - "gemini_generate": ( - "gemini/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, - ), - ), - "together_chat": ( - "together_ai/moonshotai/Kimi-K3", - 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, - ), - ), - "fireworks_chat": ( - "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) - -# Renderer-level shapes the pricing matrix gates per cap, pinned here once per -# wire so the sidecar emits prove they survive the proxy end to end. -_SHAPES: Final[tuple[tuple[str, str, Case], ...]] = ( - *( - ( - f"tool_call_{'stream' if stream else 'sync'}", - wire, - Case(name="tool_call", usage=_SHAPE_USAGE, stream=stream, tool_call=True), - ) - for wire in _WIRE_USAGE - for stream in (False, True) - ), - ( - "responses_incomplete", - "openai_responses", - Case(name="stream_no_usage_incomplete", usage=_SHAPE_USAGE, stream=True, terminal="incomplete"), - ), - ( - "responses_unvalidated", - "openai_responses", - Case(name="stream_unvalidated", usage=_SHAPE_USAGE, stream=True, terminal="unvalidated"), - ), - ( - "gemini_prompt_blocked", - "gemini_generate", - Case( - name="prompt_blocked", - usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), - terminal="prompt_blocked", - response_model_override=True, - ), - ), - ( - "gemini_prompt_blocked_stream", - "gemini_generate", - Case( - name="stream_prompt_blocked", - usage=ScriptedUsage(fresh_input_tokens=1000, output_tokens=0), - stream=True, - terminal="prompt_blocked", - 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, - ), - ), -) - - -def _shape_id(entry: tuple[str, str, Case]) -> str: - return entry[0] - - -class TestWireFormats: - @pytest.mark.parametrize("wire", tuple(_WIRE_USAGE)) - @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") - def test_wire_usage_shape_bills_each_component( - self, - client: CostCalcClient, - resources: ResourceManager, - scoped_key: str, - wire: str, - ) -> None: - map_key, usage = _WIRE_USAGE[wire] - model: Final = _MODELS[map_key] - case: Final = Case(name="basic", usage=usage) - marker: Final = unique_marker() - model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response: Final = client.proxy.transport.send( - "/chat/completions", - headers=client.proxy.transport.bearer(scoped_key), - json=ChatBody( - model=model_name, - messages=(ChatMessage(role="user", content=f"{marker} scripted wire call"),), - ), - ) - assert response.ok, f"{wire}: proxy returned {response.status_code}: {response.body[:400]}" - - expected: Final = expected_breakdown(model, case) - row: Final = cost_rows.poll_cost_row_where( - client.proxy, - scoped_key, - lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, - ) - assert row is not None, f"{wire}: no spend row landed" - assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( - f"{wire}: spend {row.spend} != expected {expected.total} " - f"(breakdown {row.breakdown.model_dump()})" - ) - breakdown: Final = row.breakdown - assert breakdown.input_cost is not None and cost_rows.approx_equal( - breakdown.input_cost, expected.input_cost - ), ( - f"{wire}: gross input_cost {breakdown.input_cost} != expected {expected.input_cost}; " - "cached/written tokens billed at the input rate" - ) - assert breakdown.output_cost is not None and cost_rows.approx_equal( - breakdown.output_cost, expected.output_cost - ), f"{wire}: output_cost {breakdown.output_cost} != expected {expected.output_cost}" - - prompt_tokens, completion_tokens = expected_token_columns(model, case) - assert row.prompt_tokens == prompt_tokens, ( - f"{wire}: prompt_tokens {row.prompt_tokens} != {prompt_tokens}" - ) - assert row.completion_tokens == completion_tokens, ( - f"{wire}: completion_tokens {row.completion_tokens} != {completion_tokens}" - ) - cost_rows.assert_total_is_sum_of_components(row) - - @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") - def test_anthropic_streamed_usage_bills_each_component( - self, client: CostCalcClient, resources: ResourceManager, scoped_key: str - ) -> None: - map_key, usage = _WIRE_USAGE["anthropic_messages"] - model: Final = _MODELS[map_key] - case: Final = Case(name="stream", usage=usage, stream=True) - marker: Final = unique_marker() - model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response: Final = client.proxy.transport.send( - "/chat/completions", - headers=client.proxy.transport.bearer(scoped_key), - json=ChatBody( - model=model_name, - messages=(ChatMessage(role="user", content=f"{marker} scripted anthropic stream"),), - stream=True, - stream_options=ChatStreamOptions(include_usage=True), - ), - stream=True, - ) - assert response.ok, f"anthropic stream: proxy returned {response.status_code}: {response.body[:400]}" - assert response.stream_done, "anthropic stream did not reach its terminal event" - assert response.stream_error is None, f"stream carried an error event: {response.stream_error}" - - expected: Final = expected_breakdown(model, case) - row: Final = cost_rows.poll_cost_row_where( - client.proxy, - scoped_key, - lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, - ) - assert row is not None, "anthropic stream: no spend row landed" - assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( - f"anthropic stream: spend {row.spend} != expected {expected.total} " - f"(breakdown {row.breakdown.model_dump()})" - ) - cost_rows.assert_total_is_sum_of_components(row) - - @pytest.mark.parametrize("shape_wire_case", _SHAPES, ids=_shape_id) - @pytest.mark.covers("quota_management.spend_tracking.scripted_wire.logs_cost") - def test_response_shape_bills_reported_usage( - self, - client: CostCalcClient, - resources: ResourceManager, - scoped_key: str, - shape_wire_case: tuple[str, str, Case], - ) -> None: - shape, wire, case = shape_wire_case - map_key, _usage = _WIRE_USAGE[wire] - model: Final = _MODELS[map_key] - marker: Final = unique_marker() - model_name, _handle = register_scenario_deployment(client, resources, model, case, marker) - response: Final = client.proxy.transport.send( - "/chat/completions", - headers=client.proxy.transport.bearer(scoped_key), - json=ChatBody( - model=model_name, - messages=(ChatMessage(role="user", content=f"{marker} scripted {shape}"),), - stream=case.stream, - stream_options=ChatStreamOptions(include_usage=True) if case.stream else None, - tools=( - ( - ChatTool( - function=ChatToolFunction( - name="get_weather", - parameters={"type": "object", "properties": {"city": {"type": "string"}}}, - ) - ), - ) - if case.tool_call - else None - ), - ), - stream=case.stream, - ) - assert response.ok, f"{shape}: proxy returned {response.status_code}: {response.body[:400]}" - if case.stream: - assert response.stream_done, f"{shape}: stream did not reach its terminal event" - assert response.stream_error is None, f"{shape}: stream error: {response.stream_error}" - - expected: Final = expected_breakdown(model, case) - row: Final = cost_rows.poll_cost_row_where( - client.proxy, - scoped_key, - lambda r: r.metadata is not None and r.metadata.cost_breakdown is not None, - ) - assert row is not None, f"{shape}: no spend row landed" - assert row.spend is not None and cost_rows.approx_equal(row.spend, expected.total), ( - f"{shape}: spend {row.spend} != expected {expected.total} " - f"(breakdown {row.breakdown.model_dump()})" - ) - prompt_tokens, completion_tokens = expected_token_columns(model, case) - assert row.prompt_tokens == prompt_tokens, ( - f"{shape}: prompt_tokens {row.prompt_tokens} != {prompt_tokens}" - ) - assert row.completion_tokens == completion_tokens, ( - f"{shape}: completion_tokens {row.completion_tokens} != {completion_tokens}" - ) - cost_rows.assert_total_is_sum_of_components(row) From 522a7f569283b9a3bc0fed2a66f1c7545f30b12b Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 00:51:43 +0000 Subject: [PATCH 026/251] test(e2e): gate all_components cases by rates and tidy cost matrix names Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cases.json | 26 ++++++++ tests/e2e/cost_calculation/cost_matrix.py | 31 +++++---- tests/e2e/cost_calculation/expected.json | 7 --- .../e2e/cost_calculation/generate_expected.py | 63 ++++++++++--------- .../e2e/cost_calculation/test_matrix_data.py | 11 ++-- 5 files changed, 79 insertions(+), 59 deletions(-) diff --git a/tests/e2e/cost_calculation/cases.json b/tests/e2e/cost_calculation/cases.json index e898557ea35..3dc4fc4d99c 100644 --- a/tests/e2e/cost_calculation/cases.json +++ b/tests/e2e/cost_calculation/cases.json @@ -180,11 +180,20 @@ "audio_input_tokens": 5, "audio_output_tokens": 3 }, + "requires_rates": [ + "cache_read_input_token_cost", + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr", + "output_cost_per_reasoning_token", + "input_cost_per_audio_token", + "output_cost_per_audio_token" + ], "wires": ["openai_chat", "azure_chat", "together_chat"] }, { "name": "all_components_fireworks", "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25}, + "requires_rates": ["cache_read_input_token_cost"], "wires": ["fireworks_chat"] }, { @@ -196,6 +205,11 @@ "cache_write_1h_tokens": 10, "output_tokens": 25 }, + "requires_rates": [ + "cache_read_input_token_cost", + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr" + ], "wires": ["anthropic_messages", "bedrock_converse"] }, { @@ -208,6 +222,11 @@ "output_tokens": 25 }, "stream": true, + "requires_rates": [ + "cache_read_input_token_cost", + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr" + ], "wires": ["anthropic_messages"] }, { @@ -220,11 +239,18 @@ "audio_input_tokens": 5, "audio_output_tokens": 3 }, + "requires_rates": [ + "cache_read_input_token_cost", + "output_cost_per_reasoning_token", + "input_cost_per_audio_token", + "output_cost_per_audio_token" + ], "wires": ["gemini_generate", "vertex_generate"] }, { "name": "all_components_responses", "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25, "reasoning_tokens": 15}, + "requires_rates": ["cache_read_input_token_cost", "output_cost_per_reasoning_token"], "wires": ["openai_responses"] } ] diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index b03c851d208..a8f60b79ae7 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -27,7 +27,6 @@ from types import MappingProxyType from typing import Final, Literal from pydantic import BaseModel, ConfigDict, TypeAdapter - from scripted_provider import Scenario, ScriptedOutput, ScriptedToolCall, ScriptedUsage, Wire COST_MAP_PATH: Final = Path(__file__).resolve().parent.parent / "cost_map.json" @@ -69,9 +68,9 @@ class CostMapEntry(BaseModel): web_search_billing_unit: str | None = None -_COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, CostMapEntry]) -_COST_MAP: Final[Mapping[str, CostMapEntry]] = MappingProxyType( - _COST_MAP_ADAPTER.validate_python(json.loads(COST_MAP_PATH.read_text())) +COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, CostMapEntry]) +COST_MAP: Final[Mapping[str, CostMapEntry]] = MappingProxyType( + COST_MAP_ADAPTER.validate_python(json.loads(COST_MAP_PATH.read_text())) ) TIER_THRESHOLD_TOKENS: Final = 200_000 @@ -146,10 +145,10 @@ class _CasesFile(BaseModel): cases: tuple[Case, ...] = () -_CASES_FILE: Final = _CasesFile.model_validate(json.loads(CASES_PATH.read_text())) -CASES: Final[tuple[Case, ...]] = _CASES_FILE.cases +CASES_FILE: Final = _CasesFile.model_validate(json.loads(CASES_PATH.read_text())) +CASES: Final[tuple[Case, ...]] = CASES_FILE.cases _DEPLOYMENTS: Final[Mapping[str, DeploymentSpec]] = MappingProxyType( - {spec.map_key: spec for spec in _CASES_FILE.deployments} + {spec.map_key: spec for spec in CASES_FILE.deployments} ) @@ -221,13 +220,13 @@ class FrontierModel: @property def rates(self) -> CostMapEntry: - return _COST_MAP[self.map_key] + return COST_MAP[self.map_key] @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] + return COST_MAP[self.override_map_key] @property def provider_model(self) -> str: @@ -262,13 +261,13 @@ def _litellm_model_for(map_key: str, wiring: _ProviderWiring) -> str: def _frontier() -> tuple[FrontierModel, ...]: groups: Final[Mapping[tuple[str, str], tuple[str, ...]]] = MappingProxyType( { - pair: tuple(sorted(k for k, e in _COST_MAP.items() if (e.litellm_provider, e.mode) == pair)) - for pair in {(e.litellm_provider, e.mode) for e in _COST_MAP.values()} + pair: tuple(sorted(k for k, e in COST_MAP.items() if (e.litellm_provider, e.mode) == pair)) + for pair in {(e.litellm_provider, e.mode) for e in COST_MAP.values()} } ) models: list[FrontierModel] = [] # mutable-ok: accumulated once at import into a tuple - for map_key in sorted(_COST_MAP): - entry: Final = _COST_MAP[map_key] + for map_key in sorted(COST_MAP): + entry: Final = COST_MAP[map_key] pair: Final = (entry.litellm_provider, entry.mode) wiring: Final = _PROVIDER_WIRING.get(pair) if wiring is None: @@ -416,7 +415,7 @@ def image_input_data_url() -> str: IMAGE_INPUT_DATA_URL: Final = image_input_data_url() -class _ExpectedCell(BaseModel): +class ExpectedCell(BaseModel): model_config = ConfigDict(frozen=True) spend: float @@ -426,8 +425,8 @@ class _ExpectedCell(BaseModel): completion_tokens: int -_EXPECTED_ADAPTER: Final = TypeAdapter(dict[str, _ExpectedCell]) -EXPECTED: Final[Mapping[str, _ExpectedCell]] = MappingProxyType( +_EXPECTED_ADAPTER: Final = TypeAdapter(dict[str, ExpectedCell]) +EXPECTED: Final[Mapping[str, ExpectedCell]] = MappingProxyType( _EXPECTED_ADAPTER.validate_python(json.loads(EXPECTED_PATH.read_text())) if EXPECTED_PATH.exists() else {} diff --git a/tests/e2e/cost_calculation/expected.json b/tests/e2e/cost_calculation/expected.json index 7a92fb2476f..3b18e9ed9f4 100644 --- a/tests/e2e/cost_calculation/expected.json +++ b/tests/e2e/cost_calculation/expected.json @@ -1686,13 +1686,6 @@ "prompt_tokens": 100, "spend": 0.0216 }, - "meta.llama4-maverick-17b-instruct-v1:0|all_components_anthropic": { - "completion_tokens": 25, - "input_cost": 0.020900000000000002, - "output_cost": 0.0095, - "prompt_tokens": 150, - "spend": 0.030400000000000003 - }, "meta.llama4-maverick-17b-instruct-v1:0|basic": { "completion_tokens": 40, "input_cost": 0.0228, diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py index de979f272fe..c093ecbe0ea 100644 --- a/tests/e2e/cost_calculation/generate_expected.py +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -14,8 +14,10 @@ from __future__ import annotations import json import sys +from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path +from types import MappingProxyType from typing import Final sys.path.insert(0, str(Path(__file__).resolve().parent)) @@ -27,6 +29,7 @@ from cost_matrix import ( # noqa: E402 # path bootstrap before package-local i TIER_THRESHOLD_TOKENS, Case, CostMapEntry, + ExpectedCell, FrontierModel, cases_for, expected_key, @@ -93,16 +96,11 @@ def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: or rates.output_cost_per_token or 0.0 ) - # The biller charges cache writes at the input rate when the entry carries - # no cache_creation rate (cost_calculator.py:2452), and at the 5m write - # rate when the 1h variant is unset; cache reads bill only at their own - # rate (zero when the entry lacks one). - write_5m_rate: Final = rates.cache_creation_input_token_cost or in_rate input_cost: Final = ( u.fresh_input_tokens * in_rate + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) - + u.cache_write_5m_tokens * write_5m_rate - + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or write_5m_rate) + + u.cache_write_5m_tokens * (rates.cache_creation_input_token_cost or 0.0) + + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or 0.0) + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) ) output_cost: Final = ( @@ -147,39 +145,46 @@ def expected_token_columns(model: FrontierModel, case: Case) -> tuple[int, int]: ) -def _proposed() -> dict[str, dict[str, object]]: - return { - expected_key(model, case): ( - lambda breakdown, tokens: { - "spend": breakdown.total, - "input_cost": breakdown.input_cost, - "output_cost": breakdown.output_cost, - "prompt_tokens": tokens[0], - "completion_tokens": tokens[1], - } - )(expected_breakdown(model, case), expected_token_columns(model, case)) - for model in FRONTIER_MODELS - for case in cases_for(model) - if case.exact_spend - } +def _cell(model: FrontierModel, case: Case) -> ExpectedCell: + breakdown: Final = expected_breakdown(model, case) + prompt_tokens, completion_tokens = expected_token_columns(model, case) + return ExpectedCell( + spend=breakdown.total, + input_cost=breakdown.input_cost, + output_cost=breakdown.output_cost, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ) + + +def _proposed() -> Mapping[str, ExpectedCell]: + return MappingProxyType( + { + expected_key(model, case): _cell(model, case) + for model in FRONTIER_MODELS + for case in cases_for(model) + if case.exact_spend + } + ) def main() -> None: rewrite: Final = "--rewrite" in sys.argv[1:] proposed: Final = _proposed() + proposed_values: Final = {key: cell.model_dump() for key, cell in proposed.items()} existing: Final = ( json.loads(EXPECTED_PATH.read_text()) if EXPECTED_PATH.exists() else {} ) merged: Final = { - key: (proposed[key] if rewrite or key not in existing else existing[key]) - for key in sorted(proposed) + key: (proposed_values[key] if rewrite or key not in existing else existing[key]) + for key in sorted(proposed_values) } - added: Final = sum(1 for key in proposed if key not in existing) - removed: Final = sum(1 for key in existing if key not in proposed) - kept: Final = sum(1 for key in proposed if key in existing and not rewrite) - rewritten: Final = sum(1 for key in proposed if key in existing and rewrite) + added: Final = sum(1 for key in proposed_values if key not in existing) + removed: Final = sum(1 for key in existing if key not in proposed_values) + kept: Final = sum(1 for key in proposed_values if key in existing and not rewrite) + rewritten: Final = sum(1 for key in proposed_values if key in existing and rewrite) EXPECTED_PATH.write_text(json.dumps(merged, indent=2, sort_keys=True) + "\n") - print( + print( # noqa: T201 # CLI summary is the tool output f"expected.json: {added} added, {removed} removed, {kept} kept, " f"{rewritten} rewritten ({len(merged)} cells)" ) diff --git a/tests/e2e/cost_calculation/test_matrix_data.py b/tests/e2e/cost_calculation/test_matrix_data.py index fdbb6ddd293..8340257939c 100644 --- a/tests/e2e/cost_calculation/test_matrix_data.py +++ b/tests/e2e/cost_calculation/test_matrix_data.py @@ -8,11 +8,10 @@ from __future__ import annotations from typing import Final import pytest - from cost_matrix import ( - _CASES_FILE, - _COST_MAP, CASES, + CASES_FILE, + COST_MAP, EXPECTED, FRONTIER_MODELS, CostMapEntry, @@ -41,7 +40,7 @@ def test_expected_keys_match_derived_exact_cells() -> None: def test_deployments_reference_existing_map_keys() -> None: unknown: Final = sorted( - spec.map_key for spec in _CASES_FILE.deployments if spec.map_key not in _COST_MAP + spec.map_key for spec in CASES_FILE.deployments if spec.map_key not in COST_MAP ) assert not unknown, f"deployments entries name map keys absent from cost_map.json: {unknown}" @@ -55,9 +54,7 @@ def test_requires_rates_are_cost_map_fields() -> None: def test_no_two_entries_share_input_rate() -> None: - rates: Final = [ - entry.input_cost_per_token for entry in _COST_MAP.values() - ] + rates: Final = tuple(entry.input_cost_per_token for entry in COST_MAP.values()) assert len(rates) == len(set(rates)), ( "two cost_map entries share input_cost_per_token; the suite relies on " "distinct rates so a wrong-model bill can never coincidentally match" From e1c9ae5ae45e3b66041a25a6ae6ca9c7633944b7 Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 00:56:39 +0000 Subject: [PATCH 027/251] test(e2e): drop needless sys.path bootstrap from golden generator Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/generate_expected.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py index c093ecbe0ea..e243e477839 100644 --- a/tests/e2e/cost_calculation/generate_expected.py +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -16,14 +16,10 @@ import json import sys from collections.abc import Mapping from dataclasses import dataclass -from pathlib import Path from types import MappingProxyType from typing import Final -sys.path.insert(0, str(Path(__file__).resolve().parent)) -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) - -from cost_matrix import ( # noqa: E402 # path bootstrap before package-local imports +from cost_matrix import ( EXPECTED_PATH, FRONTIER_MODELS, TIER_THRESHOLD_TOKENS, From 3e11c986766ed7a32ead704e5284fbfeaf889c6b Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 01:02:41 +0000 Subject: [PATCH 028/251] test(e2e): satisfy pyright in cost matrix derivation and golden generator Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cost_matrix.py | 14 +++++++------- tests/e2e/cost_calculation/generate_expected.py | 16 +++++++++++++--- 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index a8f60b79ae7..35344a099be 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -267,23 +267,23 @@ def _frontier() -> tuple[FrontierModel, ...]: ) models: list[FrontierModel] = [] # mutable-ok: accumulated once at import into a tuple for map_key in sorted(COST_MAP): - entry: Final = COST_MAP[map_key] - pair: Final = (entry.litellm_provider, entry.mode) - wiring: Final = _PROVIDER_WIRING.get(pair) + entry = COST_MAP[map_key] + pair = (entry.litellm_provider, entry.mode) + wiring = _PROVIDER_WIRING.get(pair) if wiring is None: raise ValueError( f"cost_map entry {map_key} has no wiring for " f"(litellm_provider={pair[0]}, mode={pair[1]}); add a " f"_ProviderWiring row in cost_matrix.py" ) - siblings: Final = groups[pair] - override_key: Final = ( + siblings = groups[pair] + override_key = ( siblings[(siblings.index(map_key) + 1) % len(siblings)] if len(siblings) > 1 else None ) - override_litellm: Final = ( + override_litellm = ( _litellm_model_for(override_key, wiring) if override_key is not None else None ) - deployment: Final = _DEPLOYMENTS.get(map_key) + deployment = _DEPLOYMENTS.get(map_key) models.append( FrontierModel( model_name=f"cc-{map_key.replace('/', '-').replace(':', '-').replace('.', '-').lower()}", diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py index e243e477839..f5514092c88 100644 --- a/tests/e2e/cost_calculation/generate_expected.py +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -19,6 +19,8 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final +from pydantic import TypeAdapter + from cost_matrix import ( EXPECTED_PATH, FRONTIER_MODELS, @@ -168,11 +170,19 @@ def main() -> None: rewrite: Final = "--rewrite" in sys.argv[1:] proposed: Final = _proposed() proposed_values: Final = {key: cell.model_dump() for key, cell in proposed.items()} - existing: Final = ( - json.loads(EXPECTED_PATH.read_text()) if EXPECTED_PATH.exists() else {} + existing: Final[Mapping[str, ExpectedCell]] = ( + TypeAdapter(dict[str, ExpectedCell]).validate_python( + json.loads(EXPECTED_PATH.read_text()) + ) + if EXPECTED_PATH.exists() + else {} ) merged: Final = { - key: (proposed_values[key] if rewrite or key not in existing else existing[key]) + key: ( + proposed_values[key] + if rewrite or key not in existing + else existing[key].model_dump() + ) for key in sorted(proposed_values) } added: Final = sum(1 for key in proposed_values if key not in existing) From fc0cce553a631e912e7892893bde188e9716b415 Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 01:13:36 +0000 Subject: [PATCH 029/251] test(e2e): derive cache rates from first principles and ungate all_components cases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cases.json | 17 +---------------- tests/e2e/cost_calculation/expected.json | 7 +++++++ tests/e2e/cost_calculation/generate_expected.py | 17 ++++++++++++++--- 3 files changed, 22 insertions(+), 19 deletions(-) diff --git a/tests/e2e/cost_calculation/cases.json b/tests/e2e/cost_calculation/cases.json index 3dc4fc4d99c..e01ac97e9ff 100644 --- a/tests/e2e/cost_calculation/cases.json +++ b/tests/e2e/cost_calculation/cases.json @@ -181,9 +181,6 @@ "audio_output_tokens": 3 }, "requires_rates": [ - "cache_read_input_token_cost", - "cache_creation_input_token_cost", - "cache_creation_input_token_cost_above_1hr", "output_cost_per_reasoning_token", "input_cost_per_audio_token", "output_cost_per_audio_token" @@ -193,7 +190,6 @@ { "name": "all_components_fireworks", "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25}, - "requires_rates": ["cache_read_input_token_cost"], "wires": ["fireworks_chat"] }, { @@ -205,11 +201,6 @@ "cache_write_1h_tokens": 10, "output_tokens": 25 }, - "requires_rates": [ - "cache_read_input_token_cost", - "cache_creation_input_token_cost", - "cache_creation_input_token_cost_above_1hr" - ], "wires": ["anthropic_messages", "bedrock_converse"] }, { @@ -222,11 +213,6 @@ "output_tokens": 25 }, "stream": true, - "requires_rates": [ - "cache_read_input_token_cost", - "cache_creation_input_token_cost", - "cache_creation_input_token_cost_above_1hr" - ], "wires": ["anthropic_messages"] }, { @@ -240,7 +226,6 @@ "audio_output_tokens": 3 }, "requires_rates": [ - "cache_read_input_token_cost", "output_cost_per_reasoning_token", "input_cost_per_audio_token", "output_cost_per_audio_token" @@ -250,7 +235,7 @@ { "name": "all_components_responses", "usage": {"fresh_input_tokens": 80, "cache_read_tokens": 40, "output_tokens": 25, "reasoning_tokens": 15}, - "requires_rates": ["cache_read_input_token_cost", "output_cost_per_reasoning_token"], + "requires_rates": ["output_cost_per_reasoning_token"], "wires": ["openai_responses"] } ] diff --git a/tests/e2e/cost_calculation/expected.json b/tests/e2e/cost_calculation/expected.json index 3b18e9ed9f4..caea2c3c764 100644 --- a/tests/e2e/cost_calculation/expected.json +++ b/tests/e2e/cost_calculation/expected.json @@ -1686,6 +1686,13 @@ "prompt_tokens": 100, "spend": 0.0216 }, + "meta.llama4-maverick-17b-instruct-v1:0|all_components_anthropic": { + "completion_tokens": 25, + "input_cost": 0.0285, + "output_cost": 0.0095, + "prompt_tokens": 150, + "spend": 0.038 + }, "meta.llama4-maverick-17b-instruct-v1:0|basic": { "completion_tokens": 40, "input_cost": 0.0228, diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py index f5514092c88..a6eabcc7286 100644 --- a/tests/e2e/cost_calculation/generate_expected.py +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -94,11 +94,22 @@ def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: or rates.output_cost_per_token or 0.0 ) + write_rate: Final = ( + rates.cache_creation_input_token_cost + if rates.cache_creation_input_token_cost is not None + else in_rate + ) input_cost: Final = ( u.fresh_input_tokens * in_rate - + u.cache_read_tokens * (rates.cache_read_input_token_cost or 0.0) - + u.cache_write_5m_tokens * (rates.cache_creation_input_token_cost or 0.0) - + u.cache_write_1h_tokens * (rates.cache_creation_input_token_cost_above_1hr or 0.0) + + u.cache_read_tokens + * (rates.cache_read_input_token_cost if rates.cache_read_input_token_cost is not None else in_rate) + + u.cache_write_5m_tokens * write_rate + + u.cache_write_1h_tokens + * ( + rates.cache_creation_input_token_cost_above_1hr + if rates.cache_creation_input_token_cost_above_1hr is not None + else write_rate + ) + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) ) output_cost: Final = ( From 072b32baf2097e5421672956ce34809106f592aa Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 01:18:02 +0000 Subject: [PATCH 030/251] test(e2e): derive goldens from first-principles rate selection Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/cases.json | 10 ++- tests/e2e/cost_calculation/expected.json | 18 ++-- .../e2e/cost_calculation/generate_expected.py | 86 +++++++++---------- 3 files changed, 61 insertions(+), 53 deletions(-) diff --git a/tests/e2e/cost_calculation/cases.json b/tests/e2e/cost_calculation/cases.json index e01ac97e9ff..49eebc85231 100644 --- a/tests/e2e/cost_calculation/cases.json +++ b/tests/e2e/cost_calculation/cases.json @@ -62,7 +62,15 @@ "name": "web_search", "usage": {"fresh_input_tokens": 100, "output_tokens": 30, "web_search_calls": 3}, "requires_rates": ["search_context_cost_per_query"], - "requires_caps": ["web_search"] + "requires_caps": ["web_search"], + "wires": ["openai_responses", "anthropic_messages", "gemini_generate", "vertex_generate"] + }, + { + "name": "web_search_single", + "usage": {"fresh_input_tokens": 100, "output_tokens": 30, "web_search_calls": 1}, + "requires_rates": ["search_context_cost_per_query"], + "requires_caps": ["web_search"], + "wires": ["openai_chat", "together_chat", "fireworks_chat", "azure_chat"] }, { "name": "stream", diff --git a/tests/e2e/cost_calculation/expected.json b/tests/e2e/cost_calculation/expected.json index caea2c3c764..984b670a82c 100644 --- a/tests/e2e/cost_calculation/expected.json +++ b/tests/e2e/cost_calculation/expected.json @@ -160,7 +160,7 @@ "prompt_tokens": 120, "spend": 0.032 }, - "azure/gpt-5.4-mini|web_search": { + "azure/gpt-5.4-mini|web_search_single": { "completion_tokens": 30, "input_cost": 0.016, "output_cost": 0.009600000000000001, @@ -272,7 +272,7 @@ "prompt_tokens": 120, "spend": 0.03 }, - "azure/gpt-5.6|web_search": { + "azure/gpt-5.6|web_search_single": { "completion_tokens": 30, "input_cost": 0.015, "output_cost": 0.009, @@ -615,7 +615,7 @@ "prompt_tokens": 120, "spend": 0.028000000000000004 }, - "fireworks_ai/deepseek-v4p1-flash|web_search": { + "fireworks_ai/deepseek-v4p1-flash|web_search_single": { "completion_tokens": 30, "input_cost": 0.014000000000000002, "output_cost": 0.008400000000000001, @@ -706,7 +706,7 @@ "prompt_tokens": 120, "spend": 0.024 }, - "fireworks_ai/kimi-k3|web_search": { + "fireworks_ai/kimi-k3|web_search_single": { "completion_tokens": 30, "input_cost": 0.012000000000000002, "output_cost": 0.007200000000000001, @@ -797,7 +797,7 @@ "prompt_tokens": 120, "spend": 0.026000000000000002 }, - "fireworks_ai/qwen3p8-max|web_search": { + "fireworks_ai/qwen3p8-max|web_search_single": { "completion_tokens": 30, "input_cost": 0.013000000000000001, "output_cost": 0.007800000000000001, @@ -1462,7 +1462,7 @@ "prompt_tokens": 120, "spend": 0.008 }, - "gpt-5.4-mini|web_search": { + "gpt-5.4-mini|web_search_single": { "completion_tokens": 30, "input_cost": 0.004, "output_cost": 0.0024000000000000002, @@ -1679,7 +1679,7 @@ "prompt_tokens": 120, "spend": 0.002 }, - "gpt-5.6|web_search": { + "gpt-5.6|web_search_single": { "completion_tokens": 30, "input_cost": 0.001, "output_cost": 0.0006000000000000001, @@ -1826,7 +1826,7 @@ "prompt_tokens": 120, "spend": 0.02 }, - "together_ai/moonshotai/Kimi-K3|web_search": { + "together_ai/moonshotai/Kimi-K3|web_search_single": { "completion_tokens": 30, "input_cost": 0.01, "output_cost": 0.006, @@ -1938,7 +1938,7 @@ "prompt_tokens": 120, "spend": 0.022 }, - "together_ai/zai-org/GLM-5.3|web_search": { + "together_ai/zai-org/GLM-5.3|web_search_single": { "completion_tokens": 30, "input_cost": 0.011000000000000001, "output_cost": 0.0066, diff --git a/tests/e2e/cost_calculation/generate_expected.py b/tests/e2e/cost_calculation/generate_expected.py index a6eabcc7286..64abdb14c99 100644 --- a/tests/e2e/cost_calculation/generate_expected.py +++ b/tests/e2e/cost_calculation/generate_expected.py @@ -19,8 +19,6 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final -from pydantic import TypeAdapter - from cost_matrix import ( EXPECTED_PATH, FRONTIER_MODELS, @@ -32,19 +30,11 @@ from cost_matrix import ( cases_for, expected_key, ) - -# Wires whose response surface reports a real web-search call count; the -# chat-completions wires only expose url_citation annotations, so their billed -# count floors to one. -_EXACT_WEB_SEARCH_WIRES: Final = frozenset( - {"openai_responses", "anthropic_messages", "gemini_generate", "vertex_generate"} -) +from pydantic import TypeAdapter -def billed_web_search_calls(model: FrontierModel, case: Case) -> int: - if case.usage.web_search_calls == 0: - return 0 - return case.usage.web_search_calls if model.wire in _EXACT_WEB_SEARCH_WIRES else 1 +def _first_present(*rates: float | None) -> float | None: + return next((rate for rate in rates if rate is not None), None) @dataclass(frozen=True, slots=True) @@ -67,11 +57,14 @@ def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: Input = fresh*in + read*read + 5m*create + 1h*create_1h + audio_in*audio_in; output = text*out + reasoning*reasoning + audio_out*audio_out; plus the - billed web-search calls at the medium search-context rate. Above-threshold - swaps every input/output rate to its ``_above_200k_tokens`` variant when - total prompt tokens exceed the threshold; a service tier swaps input/output - to the tier's variants, falling back to the base rate when a variant is - unset -- mirroring _get_token_base_cost in litellm's cost calculator. + billed web-search calls at the medium search-context rate. Every billed + token is a token the provider charged for: a component whose entry has no + dedicated rate bills at the ordinary input or output rate, and a present + rate (including an explicit 0.0) is authoritative. When the total prompt + tokens exceed the threshold, input/output rates come from the + ``_above_200k_tokens`` variants; a service tier takes its ``_priority`` or + ``_flex`` variant when the entry carries one, and otherwise bills at the + base rate. """ rates: Final[CostMapEntry] = model.override_rates if case.response_model_override else model.rates u: Final = case.usage @@ -81,46 +74,53 @@ def expected_breakdown(model: FrontierModel, case: Case) -> ExpectedCost: ) tiered: Final = prompt_tokens > TIER_THRESHOLD_TOKENS in_rate: Final = ( - (rates.input_cost_per_token_above_200k_tokens if tiered else None) - or (rates.input_cost_per_token_priority if case.service_tier == "priority" else None) - or (rates.input_cost_per_token_flex if case.service_tier == "flex" else None) - or rates.input_cost_per_token + _first_present( + rates.input_cost_per_token_above_200k_tokens if tiered else None, + rates.input_cost_per_token_priority if case.service_tier == "priority" else None, + rates.input_cost_per_token_flex if case.service_tier == "flex" else None, + rates.input_cost_per_token, + ) or 0.0 ) out_rate: Final = ( - (rates.output_cost_per_token_above_200k_tokens if tiered else None) - or (rates.output_cost_per_token_priority if case.service_tier == "priority" else None) - or (rates.output_cost_per_token_flex if case.service_tier == "flex" else None) - or rates.output_cost_per_token + _first_present( + rates.output_cost_per_token_above_200k_tokens if tiered else None, + rates.output_cost_per_token_priority if case.service_tier == "priority" else None, + rates.output_cost_per_token_flex if case.service_tier == "flex" else None, + rates.output_cost_per_token, + ) or 0.0 ) - write_rate: Final = ( - rates.cache_creation_input_token_cost - if rates.cache_creation_input_token_cost is not None - else in_rate + read_rate: Final = _first_present(rates.cache_read_input_token_cost, in_rate) or 0.0 + write_rate: Final = _first_present(rates.cache_creation_input_token_cost, in_rate) or 0.0 + write_1h_rate: Final = ( + _first_present(rates.cache_creation_input_token_cost_above_1hr, write_rate) or 0.0 ) + audio_in_rate: Final = _first_present(rates.input_cost_per_audio_token, in_rate) or 0.0 + reasoning_rate: Final = _first_present(rates.output_cost_per_reasoning_token, out_rate) or 0.0 + audio_out_rate: Final = _first_present(rates.output_cost_per_audio_token, out_rate) or 0.0 input_cost: Final = ( u.fresh_input_tokens * in_rate - + u.cache_read_tokens - * (rates.cache_read_input_token_cost if rates.cache_read_input_token_cost is not None else in_rate) + + u.cache_read_tokens * read_rate + u.cache_write_5m_tokens * write_rate - + u.cache_write_1h_tokens - * ( - rates.cache_creation_input_token_cost_above_1hr - if rates.cache_creation_input_token_cost_above_1hr is not None - else write_rate - ) - + u.audio_input_tokens * (rates.input_cost_per_audio_token or 0.0) + + u.cache_write_1h_tokens * write_1h_rate + + u.audio_input_tokens * audio_in_rate ) output_cost: Final = ( u.output_tokens * out_rate - + u.reasoning_tokens * (rates.output_cost_per_reasoning_token or out_rate) - + u.audio_output_tokens * (rates.output_cost_per_audio_token or out_rate) + + u.reasoning_tokens * reasoning_rate + + u.audio_output_tokens * audio_out_rate ) search: Final = rates.search_context_cost_per_query - tool_cost: Final = billed_web_search_calls(model, case) * ( - search.search_context_size_medium if search and search.search_context_size_medium else 0.0 + medium_rate: Final = ( + search.search_context_size_medium if search is not None else None ) + if u.web_search_calls and medium_rate is None: + raise ValueError( + f"{model.map_key}: case {case.name} bills {u.web_search_calls} web-search " + "calls but the entry has no search_context_cost_per_query medium rate" + ) + tool_cost: Final = u.web_search_calls * (medium_rate if medium_rate is not None else 0.0) return ExpectedCost(input_cost=input_cost, output_cost=output_cost, tool_cost=tool_cost) From 1de633ac36644c5774cf629a793b6716a98b7580 Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 01:22:01 +0000 Subject: [PATCH 031/251] test(e2e): move matrix data freshness checks to collection time Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/CLAUDE.md | 2 +- tests/e2e/cost_calculation/cost_matrix.py | 48 +++++++++++++++ .../e2e/cost_calculation/test_matrix_data.py | 61 ------------------- .../test_token_pricing_e2e.py | 4 ++ 4 files changed, 53 insertions(+), 62 deletions(-) delete mode 100644 tests/e2e/cost_calculation/test_matrix_data.py diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 707d35b4aa6..f89b3203622 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -21,7 +21,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see); needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in), and markerless harness unit tests for the locust, process-usage, and session-anomaly aggregation logic - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `gateway/` - proxy configuration only (`litellm-config.yml`); no tests -- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; every cost-map entry is a deployment and the cases plus asserted goldens are data in `cases.json` and `expected.json` (regenerate with `generate_expected.py`), deselected unless `E2E_COST_MAP_STACK` is set, driven by the Buildkite `e2e-cost-calculation` step in project-releaser, which runs a proxy booted from `gateway/cost_calculation_ci_config.yml`, Postgres and the scripted provider co-located with pytest in one pod and sets the opt-in +- `cost_calculation/` - cost accounting against a dedicated proxy whose whole model cost map is the test-owned `tests/e2e/cost_map.json` (loaded via `LITELLM_MODEL_COST_MAP_URL`), with provider calls answered by the scripted-provider sidecar in `scripted_provider.py`; every cost-map entry is a deployment and the cases plus asserted goldens are data in `cases.json` and `expected.json` (regenerate with `generate_expected.py`; `cost_matrix.matrix_data_errors()` runs at collection time so a stale key set fails the suite's collection loudly), deselected unless `E2E_COST_MAP_STACK` is set, driven by the Buildkite `e2e-cost-calculation` step in project-releaser, which runs a proxy booted from `gateway/cost_calculation_ci_config.yml`, Postgres and the scripted provider co-located with pytest in one pod and sets the opt-in - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream; see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport; the Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json` diff --git a/tests/e2e/cost_calculation/cost_matrix.py b/tests/e2e/cost_calculation/cost_matrix.py index 35344a099be..68f3186809d 100644 --- a/tests/e2e/cost_calculation/cost_matrix.py +++ b/tests/e2e/cost_calculation/cost_matrix.py @@ -435,3 +435,51 @@ EXPECTED: Final[Mapping[str, ExpectedCell]] = MappingProxyType( def expected_key(model: FrontierModel, case: Case) -> str: return f"{model.map_key}|{case.name}" + + +def matrix_data_errors() -> tuple[str, ...]: + """Freshness findings for the data files, as human-readable strings. + + Called at collection time by the e2e suite; also usable from + generate_expected.py's context without importing pytest. + """ + derived: Final = { + expected_key(model, case) + for model in FRONTIER_MODELS + for case in cases_for(model) + if case.exact_spend + } + golden: Final = set(EXPECTED) + unknown_deployments: Final = sorted( + spec.map_key for spec in CASES_FILE.deployments if spec.map_key not in COST_MAP + ) + unknown_rates: Final = sorted( + {field for case in CASES for field in case.requires_rates} - set(CostMapEntry.model_fields) + ) + input_rates: Final = tuple(entry.input_cost_per_token for entry in COST_MAP.values()) + findings: Final = ( + ( + "expected.json is out of sync with the derived matrix; run " + "uv run python tests/e2e/cost_calculation/generate_expected.py " + f"(missing: {sorted(derived - golden)}; stale: {sorted(golden - derived)})" + ) + if derived != golden + else None, + ( + f"deployments entries name map keys absent from cost_map.json: {unknown_deployments}" + if unknown_deployments + else None + ), + ( + f"requires_rates names that are not CostMapEntry fields: {unknown_rates}" + if unknown_rates + else None + ), + ( + "two cost_map entries share input_cost_per_token; the suite relies on " + "distinct rates so a wrong-model bill can never coincidentally match" + if len(input_rates) != len(set(input_rates)) + else None + ), + ) + return tuple(finding for finding in findings if finding is not None) diff --git a/tests/e2e/cost_calculation/test_matrix_data.py b/tests/e2e/cost_calculation/test_matrix_data.py deleted file mode 100644 index 8340257939c..00000000000 --- a/tests/e2e/cost_calculation/test_matrix_data.py +++ /dev/null @@ -1,61 +0,0 @@ -"""Freshness checks for the cost suite's data files; markerless, so it runs on -any pytest invocation of the folder without the stack. expected.json is the -oracle: these tests check its key set against the derived matrix, never its -values (the generator proposes, the file decides).""" - -from __future__ import annotations - -from typing import Final - -import pytest -from cost_matrix import ( - CASES, - CASES_FILE, - COST_MAP, - EXPECTED, - FRONTIER_MODELS, - CostMapEntry, - cases_for, - expected_key, -) - - -def test_expected_keys_match_derived_exact_cells() -> None: - derived: Final = { - expected_key(model, case) - for model in FRONTIER_MODELS - for case in cases_for(model) - if case.exact_spend - } - golden: Final = set(EXPECTED) - if derived != golden: - missing: Final = sorted(derived - golden) - stale: Final = sorted(golden - derived) - pytest.fail( - "expected.json is out of sync with the derived matrix; run " - "uv run python tests/e2e/cost_calculation/generate_expected.py " - f"(missing: {missing}; stale: {stale})" - ) - - -def test_deployments_reference_existing_map_keys() -> None: - unknown: Final = sorted( - spec.map_key for spec in CASES_FILE.deployments if spec.map_key not in COST_MAP - ) - assert not unknown, f"deployments entries name map keys absent from cost_map.json: {unknown}" - - -def test_requires_rates_are_cost_map_fields() -> None: - fields: Final = set(CostMapEntry.model_fields) - unknown: Final = sorted( - {field for case in CASES for field in case.requires_rates} - fields - ) - assert not unknown, f"requires_rates names that are not CostMapEntry fields: {unknown}" - - -def test_no_two_entries_share_input_rate() -> None: - rates: Final = tuple(entry.input_cost_per_token for entry in COST_MAP.values()) - assert len(rates) == len(set(rates)), ( - "two cost_map entries share input_cost_per_token; the suite relies on " - "distinct rates so a wrong-model bill can never coincidentally match" - ) diff --git a/tests/e2e/cost_calculation/test_token_pricing_e2e.py b/tests/e2e/cost_calculation/test_token_pricing_e2e.py index 7cd128ad6fb..346a55aa22d 100644 --- a/tests/e2e/cost_calculation/test_token_pricing_e2e.py +++ b/tests/e2e/cost_calculation/test_token_pricing_e2e.py @@ -22,6 +22,7 @@ from cost_matrix import ( FrontierModel, cases_for, expected_key, + matrix_data_errors, recount_cost, ) from e2e_config import unique_marker @@ -39,6 +40,9 @@ from models import ( pytestmark: Final = [pytest.mark.e2e, pytest.mark.cost_map_stack] # mutable-ok: pytest only accepts a list for pytestmark +if _data_errors := matrix_data_errors(): + raise ValueError("\n".join(_data_errors)) + _MATRIX: Final[tuple[tuple[FrontierModel, Case], ...]] = tuple( (model, case) for model in FRONTIER_MODELS for case in cases_for(model) ) From ef1f306a7dc0777276166859b3fa4d2e6272cefb Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 01:53:52 +0000 Subject: [PATCH 032/251] test(e2e): emit gemini stream usage only on the final chunk Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/cost_calculation/scripted_provider.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/e2e/cost_calculation/scripted_provider.py b/tests/e2e/cost_calculation/scripted_provider.py index 982132ed8df..90d95441e5c 100644 --- a/tests/e2e/cost_calculation/scripted_provider.py +++ b/tests/e2e/cost_calculation/scripted_provider.py @@ -750,10 +750,8 @@ def _gemini_body(scenario: Scenario, requested_model: str) -> Mapping[str, objec def _gemini_sse(scenario: Scenario, requested_model: str) -> bytes: emit_usage: Final = scenario.stream_usage == "final_chunk" - first: Final = ( - _jobj(*((key, value) for key, value in _gemini_body(scenario, requested_model).items() if key != "usageMetadata")) - if scenario.stream_usage == "absent" - else _gemini_body(scenario, requested_model) + first: Final = _jobj( + *((key, value) for key, value in _gemini_body(scenario, requested_model).items() if key != "usageMetadata") ) return _sse( ( From e254377049ca6f7087693cff4b99b291c9ba6010 Mon Sep 17 00:00:00 2001 From: David Steele Date: Thu, 17 Sep 2026 07:07:57 +0100 Subject: [PATCH 033/251] fix(azure): drop tool_choice without tools DEVX-829 --- litellm/llms/azure/chat/gpt_transformation.py | 9 +- .../test_azure_chat_gpt_transformation.py | 123 ++++++++++++++++++ ...test_azure_chat_o_series_transformation.py | 3 +- 3 files changed, 133 insertions(+), 2 deletions(-) diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 6d17a1359bc..0debbe4f74d 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -280,10 +280,17 @@ class AzureOpenAIConfig(BaseConfig): ordered_messages: Final = system_messages_first(messages) if litellm.openai_system_messages_first else messages stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(ordered_messages) azure_messages: Final = convert_to_azure_openai_messages(hoist_images_from_tool_messages(stripped_messages)) + request_params: Final = { + key: value + for key, value in optional_params.items() + if key != "tool_choice" + or optional_params.get("tools") + or optional_params.get("functions") + } return { "model": model, "messages": azure_messages, - **optional_params, + **request_params, **sanitized_tools_update(optional_params), } diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py index e8b98c696e1..92a8124d8a9 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -333,3 +333,126 @@ class TestAzureToolSchemaCombinatorFlattening: ) assert "tools" not in request assert request["temperature"] == 0.2 + + +@pytest.mark.parametrize("tool_choice", ["none", "auto"]) +def test_azure_drops_tool_choice_without_tools_or_functions(tool_choice: str) -> None: + optional_params = {"tool_choice": tool_choice, "temperature": 0.2} + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params=optional_params, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert "tool_choice" not in request + assert request["temperature"] == 0.2 + assert optional_params["tool_choice"] == tool_choice + + +def test_azure_tools_empty_drops_tool_choice() -> None: + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"tools": [], "tool_choice": "auto"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["tools"] == [] + assert "tool_choice" not in request + + +def test_azure_functions_empty_drops_tool_choice() -> None: + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"functions": [], "tool_choice": "none"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["functions"] == [] + assert "tool_choice" not in request + + +def test_azure_preserves_tool_choice_with_tools() -> None: + tools = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}] + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"tools": tools, "tool_choice": "auto"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["tools"] == tools + assert request["tool_choice"] == "auto" + + +def test_azure_preserves_tool_choice_with_legacy_functions() -> None: + functions = [{"name": "get_weather", "parameters": {}}] + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"functions": functions, "tool_choice": "auto"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["functions"] == functions + assert request["tool_choice"] == "auto" + + +def test_azure_preserves_function_call_without_tools() -> None: + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"function_call": "none", "tool_choice": "auto"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["function_call"] == "none" + assert "tool_choice" not in request + + +def test_azure_gpt5_drops_tool_choice_without_tools() -> None: + request = AzureOpenAIGPT5Config().transform_request( + model="gpt5_series/gpt-5.6-sol", + messages=[{"role": "user", "content": "hi"}], + optional_params={"tool_choice": "none"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["model"] == "gpt-5.6-sol" + assert "tool_choice" not in request + + +@pytest.mark.asyncio +async def test_azure_async_transform_drops_tool_choice_without_tools() -> None: + request = await AzureOpenAIConfig().async_transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"tool_choice": "none"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert "tool_choice" not in request + + +@pytest.mark.asyncio +async def test_azure_gpt5_async_transform_drops_tool_choice_without_tools() -> None: + request = await AzureOpenAIGPT5Config().async_transform_request( + model="gpt5_series/gpt-5.6-sol", + messages=[{"role": "user", "content": "hi"}], + optional_params={"tool_choice": "auto"}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + assert request["model"] == "gpt-5.6-sol" + assert "tool_choice" not in request diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py index 9db9ab971a0..57d60df3a11 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py @@ -14,7 +14,7 @@ async def test_azure_chat_o_series_transformation(): provider_config = AzureOpenAIO1Config() model = "o_series/web-interface-o1-mini" messages = [{"role": "user", "content": "Hello, how are you?"}] - optional_params = {} + optional_params = {"tool_choice": "none"} litellm_params = {} headers = {} @@ -23,6 +23,7 @@ async def test_azure_chat_o_series_transformation(): ) print(response) assert response["model"] == "web-interface-o1-mini" + assert "tool_choice" not in response def test_azure_o_series_transform_request_flattens_top_level_anyof(): From 48712f733a641f16b0b4fe60c221fd9f1c7076fa Mon Sep 17 00:00:00 2001 From: David Steele Date: Thu, 17 Sep 2026 07:30:06 +0100 Subject: [PATCH 034/251] style(azure): format request parameter filter DEVX-829 Co-Authored-By: Claude Code --- litellm/llms/azure/chat/gpt_transformation.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 0debbe4f74d..7cb50ee5348 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -283,9 +283,7 @@ class AzureOpenAIConfig(BaseConfig): request_params: Final = { key: value for key, value in optional_params.items() - if key != "tool_choice" - or optional_params.get("tools") - or optional_params.get("functions") + if key != "tool_choice" or optional_params.get("tools") or optional_params.get("functions") } return { "model": model, From bd222bd8d9f6d083b8058c5fef3e998b4f92b3af Mon Sep 17 00:00:00 2001 From: David Steele Date: Thu, 17 Sep 2026 08:18:36 +0100 Subject: [PATCH 035/251] fix(azure): avoid mutable request mapping DEVX-829 Co-Authored-By: Claude Code --- litellm/llms/azure/chat/gpt_transformation.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 7cb50ee5348..424422612db 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -280,11 +280,13 @@ class AzureOpenAIConfig(BaseConfig): ordered_messages: Final = system_messages_first(messages) if litellm.openai_system_messages_first else messages stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(ordered_messages) azure_messages: Final = convert_to_azure_openai_messages(hoist_images_from_tool_messages(stripped_messages)) - request_params: Final = { - key: value - for key, value in optional_params.items() - if key != "tool_choice" or optional_params.get("tools") or optional_params.get("functions") - } + request_params: Final = MappingProxyType( + { + key: value + for key, value in optional_params.items() + if key != "tool_choice" or optional_params.get("tools") or optional_params.get("functions") + } + ) return { "model": model, "messages": azure_messages, From 2fa115db2bac68ac90b0b900dc9eecb728c3bd4d Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Thu, 17 Sep 2026 12:53:15 -0400 Subject: [PATCH 036/251] feat(router): add maintained Fuse model and harness presets --- .../public_endpoints/public_endpoints.py | 9 + .../complexity_router/README.md | 45 +++++ .../complexity_router/fuse_presets.json | 100 +++++++++++ .../complexity_router/fuse_presets.py | 52 ++++++ .../complexity_router/llm_v2.py | 37 +++- pyproject.toml | 1 + .../public_endpoints/test_public_endpoints.py | 11 ++ .../router_strategy/test_fuse_presets.py | 43 +++++ .../router_strategy/test_llm_v2.py | 110 ++++++++++++ .../test_auto_router_model_naming.py | 49 ++++++ ...ecastClassifierConfig.integration.test.tsx | 163 +++++++++++++++++- .../add_model/ForecastClassifierConfig.tsx | 26 +-- .../add_model/FuseProfilePresets.tsx | 117 +++++++++++++ .../build_complexity_router_config.test.ts | 17 ++ .../forecast_classifier_config.test.ts | 65 ++++++- .../add_model/forecast_classifier_config.ts | 33 +++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 82 ++++++++- 17 files changed, 917 insertions(+), 43 deletions(-) create mode 100644 litellm/router_strategy/complexity_router/fuse_presets.json create mode 100644 litellm/router_strategy/complexity_router/fuse_presets.py create mode 100644 tests/test_litellm/router_strategy/test_fuse_presets.py create mode 100644 ui/litellm-dashboard/src/components/add_model/FuseProfilePresets.tsx diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 94a59828451..e395f56194f 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -23,6 +23,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.utils import get_custom_url from litellm.repositories.table_repositories import ClaudeCodePluginRepository +from litellm.router_strategy.complexity_router.fuse_presets import FusePresetCatalog, get_fuse_presets from litellm.types.agents import AgentCard from litellm.types.mcp import MCPPublicServer from litellm.types.proxy.management_endpoints.model_management_endpoints import ( @@ -424,6 +425,14 @@ async def get_complexity_scorer_defaults() -> ComplexityScorerDefaults: ) +@router.get( + "/public/complexity_router/fuse_presets", + response_model=FusePresetCatalog, +) +async def get_public_fuse_presets() -> FusePresetCatalog: + return get_fuse_presets() + + @router.get( "/public/litellm_model_cost_map", tags=["public", "model management"], diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 6505746bca1..d9159cea426 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -179,6 +179,51 @@ Configure capability forecasting through YAML or the model-management API. The dashboard preserves its classifier and calibration on an untouched save; it does not provide a capability-card editor +### Fuse v2 profile presets + +Fuse v2 accepts maintained model and runtime descriptions instead of requiring +custom prose for both solvers and the harness. Select profiles explicitly for +all deployments behind your configured model groups and their actual settings. +Group names do not select profiles automatically + +```yaml +complexity_router_config: + classifier_type: llm_v2 + classifier_llm_config: + model: your-judge-group + tiers: + SIMPLE: your-efficient-group + REASONING: your-capable-group + llm_v2_config: + efficient_profile_preset: claude-sonnet-5-v1 + capable_profile_preset: claude-fable-5-1-v1 + harness_preset: claude-code-v1 + max_quality_gap: 0.05 +``` + +`GET /public/complexity_router/fuse_presets` returns the catalog version, model +profiles, and runtime descriptions, including source URLs. The bundled catalog +is loaded once per process without network requests. Sources are citations only + +Each of `efficient_profile`, `capable_profile`, and `harness` requires either +nonblank custom text or its corresponding preset reference. Custom text wins +when both are supplied, but an unknown or wrong-kind preset is still rejected. +Explicit blank text is invalid even with a valid preset. Custom text remains +limited to 4000 characters + +Saved configurations retain preset references and explicit text separately. +Preset text is resolved when building the classifier prompt, not copied into +stored custom fields. Existing all-custom configurations keep the same prompt. +Versioned preset IDs identify immutable content: revised wording receives a new +ID, and older referenced entries must remain available + +The runtime presets do not imply a repository, runnable tests, network access, +additional tools, or a step, time, or spending budget. mini-SWE-agent describes +an agent interface, not a SWE-bench task. Model descriptions summarize provider +positioning without solve rates or guaranteed rankings. Wording is an evaluation +input, not a calibrated quality claim. Existing Fuse licensing, policy, +calibration, and prompt version are unchanged + ### Heuristic v2 Set `classifier_type: heuristic_v2` to classify with the bundled calibrated diff --git a/litellm/router_strategy/complexity_router/fuse_presets.json b/litellm/router_strategy/complexity_router/fuse_presets.json new file mode 100644 index 00000000000..4006366dc25 --- /dev/null +++ b/litellm/router_strategy/complexity_router/fuse_presets.json @@ -0,0 +1,100 @@ +{ + "version": "2026-09-17-v1", + "models": [ + { + "id": "gpt-6-astra-v1", + "label": "GPT-6 Astra", + "model": "gpt-6-astra", + "text": "OpenAI model for demanding end-to-end work, including reasoning, coding, research, and document tasks", + "sources": ["https://developers.openai.com/api/docs/models/gpt-6-astra"] + }, + { + "id": "gpt-5.6-sol-v1", + "label": "GPT-5.6 Sol", + "model": "gpt-5.6-sol", + "text": "OpenAI model for complex professional work, supporting reasoning and tool calling", + "sources": ["https://developers.openai.com/api/docs/models/gpt-5.6-sol"] + }, + { + "id": "gpt-5.6-luna-v1", + "label": "GPT-5.6 Luna", + "model": "gpt-5.6-luna", + "text": "OpenAI model for high-volume workloads, supporting reasoning and tool calling", + "sources": ["https://developers.openai.com/api/docs/models/gpt-5.6-luna"] + }, + { + "id": "gpt-5.6-terra-v1", + "label": "GPT-5.6 Terra", + "model": "gpt-5.6-terra", + "text": "OpenAI general-purpose model supporting reasoning, text and image input, and tool calling", + "sources": ["https://developers.openai.com/api/docs/models/gpt-5.6-terra"] + }, + { + "id": "claude-haiku-4-5-v1", + "label": "Claude Haiku 4.5", + "model": "claude-haiku-4-5", + "text": "Anthropic latency-focused model supporting text and image input, tool use, and extended thinking", + "sources": ["https://platform.claude.com/docs/en/models/haiku-4-5/overview"] + }, + { + "id": "claude-sonnet-5-v1", + "label": "Claude Sonnet 5", + "model": "claude-sonnet-5", + "text": "Anthropic model balancing speed and capability, with adaptive thinking and tool use", + "sources": ["https://platform.claude.com/docs/en/models/sonnet-5/overview"] + }, + { + "id": "claude-opus-5-v1", + "label": "Claude Opus 5", + "model": "claude-opus-5", + "text": "Anthropic model for complex agentic coding and enterprise work, with adaptive thinking", + "sources": ["https://platform.claude.com/docs/en/models/opus-5/overview"] + }, + { + "id": "claude-fable-5-v1", + "label": "Claude Fable 5", + "model": "claude-fable-5", + "text": "Anthropic model for demanding reasoning and long-running agent tasks, with always-on adaptive thinking", + "sources": ["https://platform.claude.com/docs/en/models/fable-5/introducing-claude-fable-5-and-claude-mythos-5"] + }, + { + "id": "claude-fable-5-1-v1", + "label": "Claude Fable 5.1", + "model": "claude-fable-5-1", + "text": "Anthropic model for demanding reasoning, long-running agentic coding, and multistep research, with always-on adaptive thinking", + "sources": ["https://platform.claude.com/docs/en/models/fable-5-1/overview"] + } + ], + "harnesses": [ + { + "id": "unspecified-v1", + "label": "Unspecified runtime", + "text": "Agent runtime is unspecified. Assess the task using the supplied context without assuming repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://code.claude.com/docs/en/how-claude-code-works", "https://mini-swe-agent.com/latest/faq/"] + }, + { + "id": "claude-code-v1", + "label": "Claude Code", + "text": "Claude Code supplies an agent loop with context management and configured tools. Available actions depend on the session's tools, permissions, and execution environment. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://code.claude.com/docs/en/how-claude-code-works"] + }, + { + "id": "codex-cli-v1", + "label": "Codex CLI", + "text": "Codex CLI supplies a terminal-based coding agent. File operations, command execution, and integrations depend on the session's tools, permissions, and sandbox. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://learn.chatgpt.com/docs/codex/cli", "https://learn.chatgpt.com/codex/permissions"] + }, + { + "id": "opencode-v1", + "label": "OpenCode", + "text": "OpenCode supplies a configurable agent runtime. Available actions depend on the selected agent, tools, permissions, and execution environment. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://opencode.ai/docs/agents/"] + }, + { + "id": "mini-swe-agent-v1", + "label": "mini-SWE-agent", + "text": "The standard mini-SWE-agent setup uses a bash-only action interface and separate command executions. Available commands and resources depend on its configured environment. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget", + "sources": ["https://mini-swe-agent.com/latest/faq/"] + } + ] +} diff --git a/litellm/router_strategy/complexity_router/fuse_presets.py b/litellm/router_strategy/complexity_router/fuse_presets.py new file mode 100644 index 00000000000..66a96ec5ad5 --- /dev/null +++ b/litellm/router_strategy/complexity_router/fuse_presets.py @@ -0,0 +1,52 @@ +from functools import lru_cache +from importlib.resources import files +from typing import Annotated, Final, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, StringConstraints + +ProfileText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=4000)] + + +class FuseModelPreset(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + id: str + label: str + text: ProfileText + sources: tuple[str, ...] = Field(min_length=1) + model: str + + +class FuseHarnessPreset(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + id: str + label: str + text: ProfileText + sources: tuple[str, ...] = Field(min_length=1) + + +class FusePresetCatalog(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + version: str + models: tuple[FuseModelPreset, ...] + harnesses: tuple[FuseHarnessPreset, ...] + + +@lru_cache(maxsize=1) +def get_fuse_presets() -> FusePresetCatalog: + return FusePresetCatalog.model_validate_json( + files(__package__).joinpath("fuse_presets.json").read_text(encoding="utf-8") + ) + + +def resolve_fuse_profile(text: str | None, preset_id: str | None, kind: Literal["model", "harness"]) -> str | None: + if preset_id is None: + return text + catalog: Final = get_fuse_presets() + presets: Final = catalog.models if kind == "model" else catalog.harnesses + preset: Final = next((entry for entry in presets if entry.id == preset_id), None) + if preset is None: + return None + return text if text is not None else preset.text diff --git a/litellm/router_strategy/complexity_router/llm_v2.py b/litellm/router_strategy/complexity_router/llm_v2.py index 2f545a65aaa..18351237e65 100644 --- a/litellm/router_strategy/complexity_router/llm_v2.py +++ b/litellm/router_strategy/complexity_router/llm_v2.py @@ -7,15 +7,15 @@ from dataclasses import dataclass from sys import float_info from typing import Annotated, Final, Literal, TypeAlias -from pydantic import BaseModel, ConfigDict, Field, StrictFloat, StringConstraints, TypeAdapter +from pydantic import BaseModel, ConfigDict, Field, StrictFloat, StringConstraints, TypeAdapter, model_validator from typing_extensions import ReadOnly, TypedDict from litellm.llms.base_llm.base_utils import ( type_to_response_format_param, # pyright: ignore[reportUnknownVariableType] # legacy output validated below ) +from litellm.router_strategy.complexity_router.fuse_presets import ProfileText, resolve_fuse_profile ShortText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=512)] -ProfileText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=4000)] class _SolverProfile(TypedDict): @@ -139,20 +139,41 @@ class LLMV2Config(BaseModel): efficient_tier: str = "SIMPLE" capable_tier: str = "REASONING" - efficient_profile: ProfileText - capable_profile: ProfileText - harness: ProfileText + efficient_profile: ProfileText | None = None + capable_profile: ProfileText | None = None + harness: ProfileText | None = None + efficient_profile_preset: str | None = None + capable_profile_preset: str | None = None + harness_preset: str | None = None max_quality_gap: float = Field(ge=0.0, le=1.0, description="Maximum estimated success loss allowed for efficient.") max_output_tokens: int = Field(default=1024, ge=1) response_format: Literal["json_schema", "json_object"] = "json_schema" calibration: LLMV2Calibration | None = None + @model_validator(mode="after") + def validate_profiles(self) -> LLMV2Config: + self._profile_texts() + return self + + def _profile_texts(self) -> tuple[str, str, str]: + efficient: Final = resolve_fuse_profile(self.efficient_profile, self.efficient_profile_preset, "model") + capable: Final = resolve_fuse_profile(self.capable_profile, self.capable_profile_preset, "model") + harness: Final = resolve_fuse_profile(self.harness, self.harness_preset, "harness") + if efficient is None: + raise ValueError("efficient_profile requires text or a known efficient_profile_preset") + if capable is None: + raise ValueError("capable_profile requires text or a known capable_profile_preset") + if harness is None: + raise ValueError("harness requires text or a known harness_preset") + return efficient, capable, harness + def system_prompt(self, efficient_model: str, capable_model: str) -> str: + efficient, capable, harness = self._profile_texts() profiles: Final[_SolverProfiles] = { "prompt_version": LLM_V2_PROMPT_VERSION, - "harness": self.harness, - "efficient": {"model": efficient_model, "profile": self.efficient_profile}, - "capable": {"model": capable_model, "profile": self.capable_profile}, + "harness": harness, + "efficient": {"model": efficient_model, "profile": efficient}, + "capable": {"model": capable_model, "profile": capable}, } schema: Final = ( "\n\nResponse JSON schema:\n" + json.dumps(LLMV2Verdict.model_json_schema()) diff --git a/pyproject.toml b/pyproject.toml index 93ff55c4069..2d6d133cd0e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -290,6 +290,7 @@ editable-profile = "dev" include = [ "litellm/proxy/_experimental/out/**", "litellm/router_strategy/complexity_router/artifacts/*.json", + "litellm/router_strategy/complexity_router/fuse_presets.json", "litellm/proxy/client/cli/commands/codex_base_instructions.md", ] exclude = [ diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 0d82ed778f5..b830eb588a5 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -11,12 +11,23 @@ from fastapi.testclient import TestClient from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.public_endpoints import router +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets from litellm.types.proxy.management_endpoints.model_management_endpoints import ( ModelGroupInfoProxy, ) from litellm.types.utils import LlmProviders +def test_fuse_presets_route_serves_the_shared_catalog_without_authentication() -> None: + app: Final = FastAPI() + app.include_router(router) + client: Final = TestClient(app) + response: Final = client.get("/public/complexity_router/fuse_presets") + assert response.status_code == 200 + assert response.json() == get_fuse_presets().model_dump(mode="json") + assert client.get("/public/complexity_router/fuse_presets").json() == response.json() + + def test_get_supported_providers_returns_enum_values(): app_instance = FastAPI() app_instance.include_router(router) diff --git a/tests/test_litellm/router_strategy/test_fuse_presets.py b/tests/test_litellm/router_strategy/test_fuse_presets.py new file mode 100644 index 00000000000..0b8d936383b --- /dev/null +++ b/tests/test_litellm/router_strategy/test_fuse_presets.py @@ -0,0 +1,43 @@ +import json +from importlib.resources import files +from typing import Final + +import pytest +from pydantic import ValidationError + +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets, resolve_fuse_profile + + +def test_catalog_is_loaded_once_and_preserves_bundled_content() -> None: + get_fuse_presets.cache_clear() + first: Final = get_fuse_presets() + second: Final = get_fuse_presets() + assert first is second + bundled: Final = json.loads( + files("litellm.router_strategy.complexity_router").joinpath("fuse_presets.json").read_text(encoding="utf-8") + ) + assert first.model_dump(mode="json") == bundled + entries: Final = (*first.models, *first.harnesses) + assert len({entry.id for entry in entries}) == len(entries) + assert len(first.models) == 9 + assert len(first.harnesses) == 5 + assert all(entry.sources and all(source.startswith("https://") for source in entry.sources) for entry in entries) + + +def test_every_catalog_entry_resolves_without_changing_custom_ownership() -> None: + catalog: Final = get_fuse_presets() + for entry in catalog.models: + assert resolve_fuse_profile(None, entry.id, "model") == entry.text + assert resolve_fuse_profile("Custom text", entry.id, "model") == "Custom text" + for entry in catalog.harnesses: + assert resolve_fuse_profile(None, entry.id, "harness") == entry.text + assert resolve_fuse_profile("Custom text", entry.id, "harness") == "Custom text" + assert resolve_fuse_profile("Custom text", None, "model") == "Custom text" + assert resolve_fuse_profile("Custom text", None, "harness") == "Custom text" + + +def test_cached_catalog_and_records_cannot_be_modified() -> None: + catalog: Final = get_fuse_presets() + for record, field in ((catalog, "version"), (catalog.models[0], "text"), (catalog.harnesses[0], "text")): + with pytest.raises(ValidationError, match="frozen"): + setattr(record, field, "Changed") diff --git a/tests/test_litellm/router_strategy/test_llm_v2.py b/tests/test_litellm/router_strategy/test_llm_v2.py index 5447c8b43ce..27d31cbe640 100644 --- a/tests/test_litellm/router_strategy/test_llm_v2.py +++ b/tests/test_litellm/router_strategy/test_llm_v2.py @@ -11,8 +11,10 @@ from litellm import ModelResponse, Router from litellm.caching.dual_cache import DualCache from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig, ComplexityTier +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets from litellm.router_strategy.complexity_router.llm_v2 import ( LLM_V2_PROMPT_VERSION, + LLM_V2_SYSTEM_PROMPT, LLMV2Calibration, LLMV2Config, LLMV2ProbabilityCalibration, @@ -174,6 +176,114 @@ def test_invalid_forecast_settings_are_rejected(overrides: dict[str, object]) -> LLMV2Config.model_validate({**base.model_dump(), **overrides}) +def _preset_config(**overrides: object) -> LLMV2Config: + catalog: Final = get_fuse_presets() + return LLMV2Config.model_validate( + { + "efficient_profile_preset": catalog.models[0].id, + "capable_profile_preset": catalog.models[-1].id, + "harness_preset": catalog.harnesses[-1].id, + "max_quality_gap": 0.05, + **overrides, + } + ) + + +def test_preset_roundtrip_keeps_references_without_materializing_text() -> None: + config: Final = _preset_config() + serialized: Final = config.model_dump(exclude_none=True) + assert serialized["efficient_profile_preset"] == config.efficient_profile_preset + assert serialized["capable_profile_preset"] == config.capable_profile_preset + assert serialized["harness_preset"] == config.harness_preset + assert not {"efficient_profile", "capable_profile", "harness"}.intersection(serialized) + assert LLMV2Config.model_validate(config.model_dump()) == config + assert LLMV2Config.model_validate_json(config.model_dump_json()) == config + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +def test_preset_explicit_override_wins_and_survives_roundtrip(field: str) -> None: + config: Final = _preset_config(**{field: " Operator description "}) + roundtrip: Final = LLMV2Config.model_validate_json(config.model_dump_json()) + assert roundtrip.model_dump()[field] == "Operator description" + assert roundtrip.efficient_profile_preset == config.efficient_profile_preset + assert roundtrip.capable_profile_preset == config.capable_profile_preset + assert roundtrip.harness_preset == config.harness_preset + payload: Final = json.loads( + roundtrip.system_prompt("opaque-efficient", "opaque-capable").split("Configured solver profiles:\n")[1] + ) + if field == "harness": + assert payload["harness"] == "Operator description" + else: + assert payload[field.removesuffix("_profile")]["profile"] == "Operator description" + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +@pytest.mark.parametrize("invalid", ("", " \n\t", "x" * 4001)) +def test_preset_does_not_bypass_supplied_text_bounds(field: str, invalid: str) -> None: + with pytest.raises(ValidationError, match=field): + _preset_config(**{field: invalid}) + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +@pytest.mark.parametrize("override", (None, "Custom override")) +@pytest.mark.parametrize("invalid_id", ("missing-v1", "")) +def test_preset_unknown_reference_rejects_even_when_overridden( + field: str, override: str | None, invalid_id: str +) -> None: + with pytest.raises(ValidationError, match=f"{field}.*preset"): + _preset_config(**{field: override, f"{field}_preset": invalid_id}) + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +def test_preset_missing_text_and_reference_rejects(field: str) -> None: + with pytest.raises(ValidationError, match=field): + _preset_config(**{f"{field}_preset": None}) + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +def test_preset_reference_rejects_the_wrong_catalog_kind(field: str) -> None: + catalog: Final = get_fuse_presets() + wrong_id: Final = catalog.models[0].id if field == "harness" else catalog.harnesses[0].id + with pytest.raises(ValidationError, match=field): + _preset_config(**{f"{field}_preset": wrong_id}) + + +@pytest.mark.parametrize("mode", ("json_schema", "json_object")) +def test_custom_profile_prompt_bytes_are_unchanged(mode: str) -> None: + base: Final = _config().llm_v2_config + assert base is not None + config: Final = LLMV2Config.model_validate({**base.model_dump(), "response_format": mode}) + old_payload: Final = { + "prompt_version": LLM_V2_PROMPT_VERSION, + "harness": config.harness, + "efficient": {"model": "opaque-efficient", "profile": config.efficient_profile}, + "capable": {"model": "opaque-capable", "profile": config.capable_profile}, + } + schema: Final = ( + "\n\nResponse JSON schema:\n" + json.dumps(LLMV2Verdict.model_json_schema()) if mode == "json_object" else "" + ) + assert config.system_prompt("opaque-efficient", "opaque-capable") == ( + LLM_V2_SYSTEM_PROMPT + "\n\nConfigured solver profiles:\n" + json.dumps(old_payload) + schema + ) + + +@pytest.mark.asyncio +async def test_preset_router_passes_catalog_text_and_opaque_group_names_to_judge() -> None: + catalog: Final = get_fuse_presets() + config: Final = _config(llm_v2_config=_preset_config().model_dump()) + router, client = _router(_verdict().model_dump_json(), config) + outcome: Final = await router.aclassify("Complete the supplied task") + assert outcome.tier == ComplexityTier.SIMPLE + prompt: Final = client.acompletion.call_args.kwargs["messages"][0]["content"] + payload: Final = json.loads(prompt.split("Configured solver profiles:\n")[1]) + assert payload == { + "prompt_version": LLM_V2_PROMPT_VERSION, + "harness": catalog.harnesses[-1].text, + "efficient": {"model": "efficient", "profile": catalog.models[0].text}, + "capable": {"model": "capable", "profile": catalog.models[-1].text}, + } + + @pytest.mark.asyncio async def test_one_judge_fuses_whole_task_and_keeps_caller_text_out_of_system_prompt() -> None: router, client = _router(_verdict().model_dump_json()) diff --git a/tests/test_litellm/router_utils/test_auto_router_model_naming.py b/tests/test_litellm/router_utils/test_auto_router_model_naming.py index 3dcb8d5af94..7d59a0590f2 100644 --- a/tests/test_litellm/router_utils/test_auto_router_model_naming.py +++ b/tests/test_litellm/router_utils/test_auto_router_model_naming.py @@ -1,7 +1,10 @@ from collections.abc import Mapping +from typing import Final import pytest +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets + from litellm.router_utils.auto_router_model_naming import ( carries_complexity_router_settings, classify_strategy_router_model, @@ -171,6 +174,52 @@ def test_validate_accepts_loadable_complexity_config(complexity_router_config): assert validate_complexity_router_config_write(complexity_router_config=complexity_router_config) is None +def _fuse_write_config(profiles: Mapping[str, object]) -> Mapping[str, object]: + return { + "classifier_type": "llm_v2", + "classifier_llm_config": {"model": "judge"}, + "tiers": {"SIMPLE": ["opaque-efficient"], "REASONING": ["opaque-capable"]}, + "llm_v2_config": {"max_quality_gap": 0.05, **profiles}, + } + + +def test_fuse_write_accepts_presets_and_custom_text_with_the_same_entitlement() -> None: + catalog: Final = get_fuse_presets() + presets: Final = _fuse_write_config( + { + "efficient_profile_preset": catalog.models[0].id, + "capable_profile_preset": catalog.models[-1].id, + "harness_preset": catalog.harnesses[0].id, + } + ) + custom: Final = _fuse_write_config( + { + "efficient_profile": catalog.models[0].text, + "capable_profile": catalog.models[-1].text, + "harness": catalog.harnesses[0].text, + } + ) + assert validate_complexity_router_config_write(presets) is None + assert validate_complexity_router_config_write(custom) is None + assert claimed_capability(presets) is claimed_capability(custom) + assert claimed_capability(presets) is not None + + +@pytest.mark.parametrize("field", ("efficient_profile", "capable_profile", "harness")) +def test_fuse_write_rejects_unknown_preset_even_with_custom_text(field: str) -> None: + config: Final = _fuse_write_config( + { + "efficient_profile": "Custom efficient solver", + "capable_profile": "Custom capable solver", + "harness": "Custom runtime", + f"{field}_preset": "unknown-v1", + } + ) + violation: Final = validate_complexity_router_config_write(config) + assert violation is not None + assert f"{field}_preset" in violation + + def test_naming_check_ignores_the_config_entirely(): """The naming contract and the config's contents are separate questions with separate owners; a write may carry a config without naming a model, so neither can stand in for the other.""" diff --git a/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx b/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx index 4a574ac736d..f9b0edf9508 100644 --- a/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ForecastClassifierConfig.integration.test.tsx @@ -1,7 +1,7 @@ import React, { useState } from "react"; -import { describe, expect, it, vi } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import userEvent from "@testing-library/user-event"; -import { fireEvent, renderWithProviders, screen } from "../../../tests/test-utils"; +import { act, fireEvent, renderWithProviders, screen, testQueryClient, waitFor } from "../../../tests/test-utils"; import ClassificationMethodConfig from "./ClassificationMethodConfig"; import AutoRouterClassifierTabs from "./AutoRouterClassifierTabs"; import ForecastClassifierConfig from "./ForecastClassifierConfig"; @@ -37,6 +37,53 @@ const fuseInitial: ComplexityRouterConfigValue = { }, }; const options = ["judge", "efficient", "capable"].map((model) => ({ value: model, label: model })); +const catalog = { + version: "catalog-v1", + models: [ + { + id: "efficient-v1", + label: "Efficient preset", + text: "Maintained efficient profile", + sources: ["https://example.com/efficient"], + model: "efficient-model", + }, + { + id: "capable-v1", + label: "Capable preset", + text: "Maintained capable profile", + sources: ["https://example.com/capable"], + model: "capable-model", + }, + ], + harnesses: [ + { + id: "runtime-v1", + label: "Runtime preset", + text: "Maintained runtime profile", + sources: ["https://example.com/runtime"], + }, + ], +}; +const presetConfig = { + efficient_profile_preset: catalog.models[0].id, + capable_profile_preset: catalog.models[1].id, + harness_preset: catalog.harnesses[0].id, + max_quality_gap: 0.05, +}; +const presetInitial = { ...fuseInitial, llm_v2_config: presetConfig }; + +beforeEach(() => { + testQueryClient.clear(); + vi.stubGlobal( + "fetch", + vi.fn().mockImplementation(async () => Response.json(catalog)), + ); +}); + +afterEach(() => { + testQueryClient.clear(); + vi.unstubAllGlobals(); +}); function Form({ initialValue = initial }: { initialValue?: ComplexityRouterConfigValue }) { const [value, setValue] = useState(initialValue); @@ -72,6 +119,118 @@ function Form({ initialValue = initial }: { initialValue?: ComplexityRouterConfi } describe("forecast classifier form", () => { + it("selects all three maintained presets, previews provenance, and saves only references", async () => { + const user = userEvent.setup(); + renderWithProviders(
); + await user.click(screen.getByRole("combobox", { name: "Efficient solver profile preset" })); + await user.click(await screen.findByRole("option", { name: /^Efficient preset/ })); + await user.click(screen.getByRole("combobox", { name: "Capable solver profile preset" })); + await user.click(screen.getByRole("option", { name: /^Capable preset/ })); + await user.click(screen.getByRole("combobox", { name: "Harness and budget preset" })); + await user.click(screen.getByRole("option", { name: /^Runtime preset/ })); + expect(screen.getByLabelText("Efficient solver profile")).toHaveValue(catalog.models[0].text); + expect(screen.getByLabelText("Efficient solver profile")).toHaveAttribute("readonly"); + expect(screen.getByLabelText("Capable solver profile")).toHaveValue(catalog.models[1].text); + expect(screen.getByLabelText("Harness and budget")).toHaveValue(catalog.harnesses[0].text); + expect(screen.getAllByText(`Catalog version: ${catalog.version}`)).toHaveLength(3); + expect(screen.getByText(`Model: ${catalog.models[0].model}`)).toBeInTheDocument(); + expect(screen.getAllByRole("link", { name: "Source 1" }).map((link) => link.getAttribute("href"))).toEqual([ + catalog.models[0].sources[0], + catalog.models[1].sources[0], + catalog.harnesses[0].sources[0], + ]); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + expect(JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config).toEqual( + presetConfig, + ); + expect(fetch).toHaveBeenCalledTimes(1); + expect(fetch).toHaveBeenCalledWith( + expect.objectContaining({ url: expect.stringMatching(/\/public\/complexity_router\/fuse_presets$/) }), + ); + }); + + it.each([undefined, null, "Explicit override"])( + "copies effective text to Custom and clears only that reference, override=%s", + async (override) => { + const user = userEvent.setup(); + renderWithProviders( + , + ); + const effectiveText = override ?? catalog.models[0].text; + await waitFor(() => expect(screen.getByLabelText("Efficient solver profile")).toHaveValue(effectiveText)); + await user.click(screen.getByRole("combobox", { name: "Efficient solver profile preset" })); + await user.click(screen.getByRole("option", { name: "Custom", exact: true })); + expect(screen.getByLabelText("Efficient solver profile")).not.toHaveAttribute("readonly"); + expect(screen.getByLabelText("Efficient solver profile")).toHaveValue(effectiveText); + fireEvent.change(screen.getByLabelText("Efficient solver profile"), { target: { value: "Custom budget" } }); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + const { efficient_profile_preset: _preset, ...rest } = presetConfig; + expect( + JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config, + ).toEqual({ + ...rest, + efficient_profile: "Custom budget", + }); + }, + ); + + it.each([ + { ...fuseInitial.llm_v2_config!, efficient_profile: catalog.models[0].text }, + { + ...presetConfig, + efficient_profile: "Explicit override", + capable_profile: "Capable override", + harness: "Harness override", + }, + ])("keeps existing custom ownership and references on an unchanged save: %j", async (settings) => { + renderWithProviders(); + await waitFor(() => expect(screen.queryByText(/Loading profile presets/)).not.toBeInTheDocument()); + expect(screen.getByRole("combobox", { name: "Efficient solver profile preset" })).toHaveValue("Custom"); + expect(screen.getByLabelText("Efficient solver profile")).toHaveValue(settings.efficient_profile); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + expect(JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config).toEqual( + settings, + ); + }); + + it.each([true, false])( + "keeps edits and stored IDs while the pending catalog settles, success=%s", + async (success) => { + let resolveCatalog: (response: Response) => void = () => {}; + vi.mocked(fetch).mockReturnValue( + new Promise((resolve) => { + resolveCatalog = resolve; + }), + ); + const settings = { ...presetConfig, efficient_profile: "Original override" }; + renderWithProviders(); + expect(screen.getByText(/Loading profile presets/)).toBeInTheDocument(); + fireEvent.change(screen.getByLabelText("Efficient solver profile"), { target: { value: "Typed while loading" } }); + await act(async () => + resolveCatalog(success ? Response.json(catalog) : Response.json({ error: "unavailable" }, { status: 503 })), + ); + if (success) await screen.findAllByText(`Catalog version: ${catalog.version}`); + else expect(await screen.findByText(/Profile presets could not be loaded/)).toBeInTheDocument(); + expect(screen.getByLabelText("Efficient solver profile")).toHaveValue("Typed while loading"); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + expect( + JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config, + ).toEqual({ ...settings, efficient_profile: "Typed while loading" }); + }, + ); + + it("keeps unknown saved IDs visible with unavailable previews rather than replacing them", async () => { + const settings = { ...presetConfig, efficient_profile_preset: "unavailable-v8" }; + renderWithProviders(); + await screen.findAllByText(`Catalog version: ${catalog.version}`); + expect(screen.getByRole("combobox", { name: "Efficient solver profile preset" })).toHaveValue("unavailable-v8"); + expect(screen.getByText("Preset preview unavailable. The saved reference is preserved")).toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "Save configuration" })); + expect(JSON.parse(screen.getByRole("status", { name: "Saved configuration" }).textContent!).llm_v2_config).toEqual( + settings, + ); + }); + it("switches a populated standard router to Capability without saving hidden pools or their overrides", () => { renderWithProviders( ) : ( <> - {(["efficient_profile", "capable_profile", "harness"] as const).map((field) => { - const label = { - efficient_profile: "Efficient solver profile", - capable_profile: "Capable solver profile", - harness: "Harness and budget", - }[field]; - return ( -
- -