From ff91b89b8d57771edd4d8ddb8b020093d45fcf5d Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Mon, 2 Nov 2026 05:47:04 -0800 Subject: [PATCH 01/13] fix(router): keep encrypted reasoning when a fallback hops between same-key deployments with no api_base (#49396) * fix(router): keep encrypted reasoning on a hop between same-key deployments on the provider's default endpoint * test(router): type the default-endpoint boundary test helpers * fix(guardrails): restore the google_text_moderation embeddings fix an earlier main merge dropped * fix(router): match encryption boundaries on the key, account, and region the request really uses A deployment's named credential now fills only the values the deployment leaves out, through the same helper the request path uses, so the boundary names the key and base the request sends. A deployment with no api_key takes the provider's default key, the boundary carries the OpenAI organization and project a deployment sends, and a Bedrock Mantle region named in the model routes through the provider resolver before the Responses URL builder * fix(router): keep header-signed accounts apart on a shared boundary and skip boundary work no reasoning needs Headers the request logger masks as secrets (Authorization, api-key, x-api-key) now join the boundary's account scope, with extra_headers overriding headers the way the request merges them, so deployments that sign in through a header on one base never share encrypted reasoning. Target boundaries are built only when a reasoning item from another deployment needs the comparison, and the Responses input strip asks only about encrypted reasoning items, so a pinned turn or fallback hop with nothing to compare makes no secret-manager reads. * fix(router): resolve default boundaries only for openai and drop the unsent organization param Aggregators and Bedrock Mantle serve several vendors on one key, so a missing key or base there no longer yields a shared boundary. The Responses request never sends the organization param, so it no longer splits the boundary, while a Cookie header now does. * fix(router): match explicit pairs as on main and key default boundaries by model An explicit api_base plus api_key pair is the whole boundary again, so tag headers no longer split deployments that main kept together. The default OpenAI boundary now carries the model, since one OPENAI_BASE_URL can front several vendors, and a deployment that declares another provider gets no boundary before any secret lookup. * fix(router): match an explicit key and the environment key on one api_base A deployment with an explicit api_base and api_key now shares a boundary with one on the same api_base that takes the same key from OPENAI_API_KEY, when the model and account headers also match. Two explicit pairs still match on base and key alone, as on main. * fix(router): skip a deployment whose params cannot be read in boundary checks A named credential with a value of the wrong type, or an extra header whose name is not a string, made the boundary check raise and fail the whole request. That deployment now has no boundary, so the others still keep the reasoning they can decrypt. Two explicit api_base plus api_key pairs also compare without resolving the provider, so they read no secrets, as on main. * fix(router): ask a secret manager for default boundaries at most once a minute per deployment * fix(router): resolve a changed named credential at once in the cached boundary and take an injected clock --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../get_llm_provider_logic.py | 6 +- litellm/responses/utils.py | 15 +- .../encrypted_content_affinity_check.py | 285 ++++++-- ...allback_hop_foreign_encrypted_reasoning.py | 36 +- .../test_encrypted_content_affinity_check.py | 681 ++++++++++++++++-- 5 files changed, 917 insertions(+), 106 deletions(-) diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 38242291b71..94ba02bfbd7 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -146,6 +146,10 @@ def handle_anthropic_text_model_custom_llm_provider( return model, custom_llm_provider +def declared_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None: + return custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None) + + def declared_authenticating_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None: """The authenticating provider this pair already names, or None. @@ -154,7 +158,7 @@ def declared_authenticating_provider(model: str | None, custom_llm_provider: str and for a declared pair the resolver's answer is the declaration itself, so metadata callers adopt the declaration instead of resolving. """ - declared: Final = custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None) + declared: Final = declared_provider(model, custom_llm_provider) return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 1a97c1cb48c..43ed9a3b96f 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -678,20 +678,25 @@ class ResponsesAPIRequestUtils: return items: Final = cast(list[object], request_input) # cast-ok: untyped client json stripped: Final = tuple( - ResponsesAPIRequestUtils._without_encrypted_reasoning(item) - if should_strip is None or (isinstance(item, Mapping) and should_strip(cast(Mapping[str, object], item))) + ResponsesAPIRequestUtils._without_encrypted_reasoning(reasoning) + if (reasoning := ResponsesAPIRequestUtils._encrypted_reasoning_item(item)) is not None + and (should_strip is None or should_strip(reasoning)) else item for item in items ) items[:] = (item for item in stripped if item is not None) @staticmethod - def _without_encrypted_reasoning(item: object) -> object | None: + def _encrypted_reasoning_item(item: object) -> Mapping[str, object] | None: if not isinstance(item, dict): - return item + return None reasoning: Final = cast(Mapping[str, object], item) # cast-ok: untyped client json if reasoning.get("type") != "reasoning" or not reasoning.get("encrypted_content"): - return reasoning + return None + return reasoning + + @staticmethod + def _without_encrypted_reasoning(reasoning: Mapping[str, object]) -> object | None: readable: Final = any( ResponsesAPIRequestUtils._has_readable_text(reasoning.get(key)) for key in ("summary", "content") ) diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index 6f3f0794371..f3425596afa 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -36,31 +36,75 @@ Safe to enable globally: - No cache required. """ -from collections.abc import Iterator, Mapping, Sequence -from functools import cache +import hashlib +from collections.abc import Callable, Iterator, Mapping, Sequence +from dataclasses import dataclass, field +from functools import cache, lru_cache, partial +from time import monotonic from typing import TYPE_CHECKING, Final, Optional, cast -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter, ValidationError +import litellm from litellm._logging import verbose_router_logger +from litellm.exceptions import AuthenticationError, BadRequestError from litellm.integrations.custom_logger import CustomLogger, Span -from litellm.litellm_core_utils.credential_accessor import CredentialAccessor +from litellm.litellm_core_utils.get_llm_provider_logic import declared_provider, get_llm_provider +from litellm.litellm_core_utils.litellm_logging import has_sensitive_keyword from litellm.litellm_core_utils.prompt_templates.common_utils import ( anthropic_content_lists, encrypted_content_of_block, strip_encrypted_reasoning_from_messages, ) from litellm.litellm_core_utils.routing_common_utils import as_excluded_deployment_ids, as_retry_skipped_deployment_ids +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_deprecated import RenamedRouterAttribute, renamed_router_param +from litellm.secret_managers.main import secret_manager_would_be_consulted from litellm.types.llms.openai import AllMessageValues -from litellm.types.router import Deployment -from litellm.utils import get_order_filtered_deployments +from litellm.types.router import Deployment, GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager, get_order_filtered_deployments, with_named_credential if TYPE_CHECKING: from litellm.router_api import RoutingContext _REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_HEADERS_ADAPTER: Final = TypeAdapter(Mapping[object, object]) +_ACCOUNT_SCOPE_HEADERS: Final = frozenset({"openai-organization", "openai-project", "cookie"}) +_SINGLE_VENDOR_PROVIDERS: Final = frozenset({LlmProviders.OPENAI.value}) +_DEFAULT_ROUTE_SECRETS: Final = ("USE_LITELLM_PROXY", "OPENAI_API_KEY", "OPENAI_BASE_URL", "OPENAI_API_BASE") +_SECRET_MANAGER_REFRESH_SECONDS: Final = 60 + + +@dataclass(frozen=True, slots=True) +class _ProviderRoute: + model: str + provider: str + api_base: str | None + + +@dataclass(frozen=True, slots=True) +class _EncryptionBoundary: + api_base: object + api_key: object + model: Callable[[], str | None] = field(compare=False) + account_scope: frozenset[tuple[str, object]] + explicit_pair: bool + + def is_shared_with(self, other: "_EncryptionBoundary") -> bool: + if (self.api_base, self.api_key) != (other.api_base, other.api_key): + return False + if self.explicit_pair and other.explicit_pair: + return True + return self.account_scope == other.account_scope and self.model() == other.model() + + +@dataclass(frozen=True, slots=True) +class _ParamsInRefreshWindow: + digest: str + window: int + params: Mapping[str, object] = field(compare=False) class EncryptedContentAffinityCheck(CustomLogger): @@ -83,9 +127,11 @@ class EncryptedContentAffinityCheck(CustomLogger): ctx: Optional["RoutingContext"] = None, enable_global_affinity: bool = True, model_group_affinity_config: dict[str, list[str]] | None = None, + clock: Callable[[], float] = monotonic, ) -> None: super().__init__() self.ctx = ctx + self._clock = clock self.enable_global_affinity = enable_global_affinity self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {} @@ -209,44 +255,156 @@ class EncryptedContentAffinityCheck(CustomLogger): @staticmethod def _encryption_boundary_key( litellm_params: object, - ) -> tuple[object, object] | None: + clock: Callable[[], float] = monotonic, + ) -> _EncryptionBoundary | None: """ - ``(api_base, api_key)`` identifies an upstream encryption boundary. - The values are resolved from the deployment and its named credential - without modifying the deployment. + The upstream encryption boundary a deployment's encrypted reasoning belongs to. The values + are resolved the way the request resolves them, the deployment's own values winning over its + named credential's, without modifying the deployment. Two explicit ``(api_base, api_key)`` + pairs match on those values alone. On a provider that serves one vendor's models, a deployment + with no ``api_key`` takes the provider's default key, one with no ``api_base`` takes the endpoint + its Responses URL builder resolves, and a match involving either also needs the same model and + the same headers that pick an account. - Accepts any object exposing dict-style ``.get(key, default)``: plain - dicts (the common case in ``healthy_deployments``) as well as - ``LiteLLM_Params``-style Pydantic instances, which define a custom - ``.get()``. A stricter ``isinstance(dict)`` guard would silently drop - the latter from boundary matching and fall back to the full pool — - i.e. trigger the exact ``invalid_encrypted_content`` failure this - check exists to prevent. + Accepts plain dicts (the common case in ``healthy_deployments``) as + well as ``LiteLLM_Params`` Pydantic instances. A stricter + ``isinstance(dict)`` guard would silently drop the latter from + boundary matching and fall back to the full pool — i.e. trigger the + exact ``invalid_encrypted_content`` failure this check exists to prevent. + + When a secret manager would be asked for the provider defaults, the boundary of one set of + deployment and named credential values is reused for up to a minute, so a run of pin misses + does not ask it again for every candidate, while a changed credential resolves at once. """ - getter: Final = getattr(litellm_params, "get", None) - if not callable(getter): + params: Final = EncryptedContentAffinityCheck._params_mapping(litellm_params) + if params is None: return None - api_base: Final = getter("api_base") - api_key: Final = getter("api_key") - credential_name: Final = getter("litellm_credential_name") - credential_values: Final[Mapping[str, object] | None] = ( - CredentialAccessor.get_credential_values(credential_name) - if isinstance(credential_name, str) and credential_name - else None + effective_params: Final = with_named_credential(params) + if not any(secret_manager_would_be_consulted(name) for name in _DEFAULT_ROUTE_SECRETS): + return EncryptedContentAffinityCheck._resolve_encryption_boundary(effective_params) + return EncryptedContentAffinityCheck._encryption_boundary_in_refresh_window( + _ParamsInRefreshWindow( + hashlib.sha256(repr(effective_params).encode()).hexdigest(), + int(clock() // _SECRET_MANAGER_REFRESH_SECONDS), + effective_params, + ) ) - effective_api_base: Final = ( - credential_values.get("api_base") - if credential_values is not None and "api_base" in credential_values - else api_base - ) - effective_api_key: Final = ( - credential_values.get("api_key") - if credential_values is not None and "api_key" in credential_values - else api_key - ) - if not effective_api_base or not effective_api_key: + + @staticmethod + @lru_cache(maxsize=1024) + def _encryption_boundary_in_refresh_window(key: _ParamsInRefreshWindow) -> _EncryptionBoundary | None: + return EncryptedContentAffinityCheck._resolve_encryption_boundary(key.params) + + @staticmethod + def _resolve_encryption_boundary(effective_params: Mapping[str, object]) -> _EncryptionBoundary | None: + explicit_api_key: Final = effective_params.get("api_key") + explicit_api_base: Final = effective_params.get("api_base") + account_scope: Final = EncryptedContentAffinityCheck._account_scope(effective_params) + if explicit_api_key and explicit_api_base: + return _EncryptionBoundary( + explicit_api_base, + explicit_api_key, + cache(partial(EncryptedContentAffinityCheck._single_vendor_model, effective_params)), + account_scope, + explicit_pair=True, + ) + route: Final = EncryptedContentAffinityCheck._single_vendor_route(effective_params) + if route is None: + return None + api_key: Final = explicit_api_key or EncryptedContentAffinityCheck._default_api_key(route) + if not api_key: + return None + api_base: Final = explicit_api_base or EncryptedContentAffinityCheck._default_responses_endpoint( + route, effective_params + ) + if not api_base: + return None + return _EncryptionBoundary(api_base, api_key, lambda: route.model, account_scope, explicit_pair=False) + + @staticmethod + def _params_mapping(litellm_params: object) -> Mapping[str, object] | None: + if isinstance(litellm_params, BaseModel): + return litellm_params.model_dump(exclude_none=True) + if isinstance(litellm_params, Mapping): + return _REQUEST_KWARGS_ADAPTER.validate_python(litellm_params) + return None + + @staticmethod + def _account_scope(litellm_params: Mapping[str, object]) -> frozenset[tuple[str, object]]: + return frozenset( + (name.lower(), str(value)) + for name, value in EncryptedContentAffinityCheck._header_items(litellm_params) + if name.lower() in _ACCOUNT_SCOPE_HEADERS or has_sensitive_keyword(name) + ) + + @staticmethod + def _header_items(litellm_params: Mapping[str, object]) -> Iterator[tuple[str, object]]: + headers: Final = EncryptedContentAffinityCheck._header_mapping(litellm_params.get("headers")) + extra_headers: Final = EncryptedContentAffinityCheck._header_mapping(litellm_params.get("extra_headers")) + overridden: Final = frozenset(name.lower() for name in extra_headers) + yield from ((name, value) for name, value in headers.items() if name.lower() not in overridden) + yield from extra_headers.items() + + @staticmethod + def _header_mapping(headers: object) -> Mapping[str, object]: + if isinstance(headers, Mapping): + return {str(name): value for name, value in _HEADERS_ADAPTER.validate_python(headers).items()} + return {} + + @staticmethod + def _single_vendor_route(litellm_params: Mapping[str, object]) -> _ProviderRoute | None: + route: Final = EncryptedContentAffinityCheck._provider_route(litellm_params) + return route if route is not None and route.provider in _SINGLE_VENDOR_PROVIDERS else None + + @staticmethod + def _single_vendor_model(litellm_params: Mapping[str, object]) -> str | None: + route: Final = EncryptedContentAffinityCheck._single_vendor_route(litellm_params) + return route.model if route is not None else None + + @staticmethod + def _provider_route(litellm_params: Mapping[str, object]) -> _ProviderRoute | None: + model: Final = litellm_params.get("model") + if not isinstance(model, str): + return None + custom_llm_provider: Final = litellm_params.get("custom_llm_provider") + declared: Final = declared_provider( + model, custom_llm_provider if isinstance(custom_llm_provider, str) else None + ) + if declared in litellm.provider_list and declared not in _SINGLE_VENDOR_PROVIDERS: + return None + try: + provider_model, provider, _, provider_api_base = get_llm_provider( + model=model, litellm_params=GenericLiteLLMParams.model_validate(dict(litellm_params)) + ) + except (BadRequestError, ValidationError): + return None + return _ProviderRoute(provider_model, provider, provider_api_base) + + @staticmethod + def _default_api_key(route: _ProviderRoute) -> str | None: + model_info: Final = ProviderConfigManager.get_provider_model_info( + model=route.model, provider=LlmProviders(route.provider) + ) + if model_info is None: + return None + try: + return model_info.get_api_key() + except NotImplementedError: + return None + + @staticmethod + def _default_responses_endpoint(route: _ProviderRoute, litellm_params: Mapping[str, object]) -> str | None: + responses_config: Final = ProviderConfigManager.get_provider_responses_api_config( + provider=route.provider, model=route.model + ) + if responses_config is None: + return None + try: + return responses_config.get_complete_url( # pyright: ignore[reportUnknownMemberType] # the base types litellm_params as a bare dict + api_base=route.api_base, litellm_params=dict(litellm_params) + ) + except (ValueError, BaseLLMException, AuthenticationError): return None - return (effective_api_base, effective_api_key) def _find_deployments_on_same_encryption_boundary( self, @@ -255,7 +413,7 @@ class EncryptedContentAffinityCheck(CustomLogger): ) -> tuple[list[dict], Deployment | None]: """ Deployments in ``healthy_deployments`` sharing the originating - deployment's ``(api_base, api_key)``, alongside the originating + deployment's encryption boundary, alongside the originating deployment object (or ``None`` if it was removed / router unavailable). """ if self.ctx is None: @@ -263,11 +421,16 @@ class EncryptedContentAffinityCheck(CustomLogger): originating: Final = self.ctx.registry().deployment(model_id) if originating is None: return [], None - boundary: Final = self._encryption_boundary_key(originating.litellm_params.model_dump(exclude_none=True)) + boundary: Final = self._encryption_boundary_key( + originating.litellm_params.model_dump(exclude_none=True), self._clock + ) if boundary is None: return [], originating matches: Final = [ - d for d in healthy_deployments if self._encryption_boundary_key(d.get("litellm_params", {})) == boundary + d + for d in healthy_deployments + if (candidate := self._encryption_boundary_key(d.get("litellm_params", {}), self._clock)) is not None + and boundary.is_shared_with(candidate) ] return matches, originating @@ -280,6 +443,7 @@ class EncryptedContentAffinityCheck(CustomLogger): target_deployments: Sequence[Mapping[str, object]], *, unmarked_origin: str | None, + clock: Callable[[], float] = monotonic, ) -> None: """ Drop the encrypted reasoning that none of ``target_deployments`` minted or shares an @@ -293,12 +457,19 @@ class EncryptedContentAffinityCheck(CustomLogger): for target in target_deployments if isinstance((model_info := target.get("model_info")), Mapping) and model_info.get("id") is not None ) - target_boundaries: Final = frozenset( - boundary - for target in target_deployments - if (boundary := EncryptedContentAffinityCheck._encryption_boundary_key(target.get("litellm_params"))) - is not None - ) + + @cache + def target_boundaries() -> tuple[_EncryptionBoundary, ...]: + return tuple( + boundary + for target in target_deployments + if ( + boundary := EncryptedContentAffinityCheck._encryption_boundary_key( + target.get("litellm_params"), clock + ) + ) + is not None + ) @cache def target_can_decrypt(marked_origin: str | None) -> bool: @@ -312,12 +483,14 @@ class EncryptedContentAffinityCheck(CustomLogger): origin: Final = ctx.registry().deployment(origin_model_id) origin_boundary: Final = ( EncryptedContentAffinityCheck._encryption_boundary_key( - origin.litellm_params.model_dump(exclude_none=True) + origin.litellm_params.model_dump(exclude_none=True), clock ) if origin is not None else None ) - return origin_boundary is not None and origin_boundary in target_boundaries + return origin_boundary is not None and any( + origin_boundary.is_shared_with(target_boundary) for target_boundary in target_boundaries() + ) def should_strip_input_item(item: Mapping[str, object]) -> bool: return not target_can_decrypt(EncryptedContentAffinityCheck._model_id_of_input_item(item)) @@ -424,7 +597,12 @@ class EncryptedContentAffinityCheck(CustomLogger): ) request_kwargs["_encrypted_content_affinity_pinned"] = True self.strip_reasoning_the_targets_cannot_decrypt( - self.ctx, request_input, anthropic_messages, (deployment,), unmarked_origin=model_id + self.ctx, + request_input, + anthropic_messages, + (deployment,), + unmarked_origin=model_id, + clock=self._clock, ) return [deployment] @@ -442,7 +620,12 @@ class EncryptedContentAffinityCheck(CustomLogger): ) request_kwargs["_encrypted_content_affinity_pinned"] = True self.strip_reasoning_the_targets_cannot_decrypt( - self.ctx, request_input, anthropic_messages, boundary_matches, unmarked_origin=model_id + self.ctx, + request_input, + anthropic_messages, + boundary_matches, + unmarked_origin=model_id, + clock=self._clock, ) return boundary_matches diff --git a/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py b/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py index e6ec3444369..ef2a9e5b72f 100644 --- a/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py +++ b/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py @@ -303,7 +303,7 @@ def scenario_api_base(scenario: Scenario, scenario_id: str, response: StoredResp def deployment( scenario: Scenario, group: str, - api_base: str, + api_base: str | None, *, order: int | None, model: str = RESPONSES_MODEL, @@ -316,7 +316,7 @@ def deployment( "litellm_params": { "model": model, "api_key": PROVIDER_KEY, - "api_base": api_base, + **({} if api_base is None else {"api_base": api_base}), **({} if order is None else {"order": order}), **extra_params, }, @@ -863,6 +863,38 @@ def test_hop_inside_one_encryption_boundary_keeps_the_reasoning(gateway: Gateway assert [body.get("input") for body in bodies] == [history(mark), history(mark)], bodies +def test_hop_between_same_key_deployments_on_the_providers_default_endpoint_keeps_the_reasoning( + gateway: Gateway, tmp_path: Path +) -> None: + shared: Final = f"default-{uuid.uuid4().hex[:12]}" + with ( + scratch_database() as database, + owned_proxy_process( + gateway, + tmp_path, + {"DATABASE_URL": database, "OPENAI_BASE_URL": f"{gateway.upstream_url}/{shared}"}, + remove_environment=("DIRECT_URL", "DATABASE_URL_READ_REPLICA", "OPENAI_API_BASE"), + ) as owned, + owned.gateway.scenario() as scenario, + ): + owned.gateway.post("/config/update", {"router_settings": {"num_retries": 0}}) + one: Final = scenario_api_base(scenario, f"{shared}-o1", failing()).rsplit("/", 1)[-1] + two: Final = scenario_api_base(scenario, f"{shared}-o2", healthy_json()).rsplit("/", 1)[-1] + order_one: Final = deployment( + scenario, shared, None, order=1, extra_params={"extra_headers": {"x-scripted-scenario": one}} + ) + order_two: Final = deployment( + scenario, shared, None, order=2, extra_params={"extra_headers": {"x-scripted-scenario": two}} + ) + mark: Final = marker() + answer: Final = responses_answer( + owned.gateway, HTTPX, Target(shared, frozenset({order_one, order_two})), history(mark) + ) + assert_served_by(answer, order_two) + bodies: Final = bodies_for(observed(owned.gateway), shared) + assert [body.get("input") for body in bodies] == [history(mark), history(mark)], bodies + + def assert_cross_group_hop_stripped(gateway: Gateway, source_scenario: str, target_scenario: str, mark: str) -> None: records: Final = observed(gateway) assert [body.get("input") for body in bodies_for(records, source_scenario)] == [history(mark)], records diff --git a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index d184bf84035..a4d279e32e6 100644 --- a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -15,17 +15,21 @@ The mechanism works without any cache and supports two encoding strategies: encrypted_content back to their original forms before sending to the upstream provider. """ +import datetime import time +from collections.abc import Mapping from typing import Final from unittest.mock import AsyncMock, patch import pytest import litellm +from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.models.credentials import CredentialItem from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_api import RoutingContext from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem @pytest.fixture(autouse=True) @@ -1100,14 +1104,7 @@ def test_boundary_key_accepts_pydantic_litellm_params_instance(): plain_key = EncryptedContentAffinityCheck._encryption_boundary_key(plain_params) assert pydantic_key is not None - assert ( - pydantic_key - == plain_key - == ( - "https://mateo-resource.openai.azure.com", - "fake-azure-resource-key-a", - ) - ) + assert pydantic_key == plain_key def test_boundary_key_resolves_missing_values_from_named_credential(): @@ -1131,14 +1128,35 @@ def test_boundary_key_resolves_missing_values_from_named_credential(): ): boundary = EncryptedContentAffinityCheck._encryption_boundary_key({"litellm_credential_name": "account-a"}) - assert boundary == ("https://account-a.example.com", "credential-key-a") + assert boundary is not None + assert boundary == EncryptedContentAffinityCheck._encryption_boundary_key( + {"api_base": "https://account-a.example.com", "api_key": "credential-key-a"} + ) + assert boundary != EncryptedContentAffinityCheck._encryption_boundary_key( + {"api_base": "https://account-a.example.com", "api_key": "credential-key-b"} + ) -def test_boundary_key_matches_named_credential_precedence(): +@pytest.mark.parametrize( + "deployment_values", + [ + pytest.param({"api_base": "https://deployment.example.com", "api_key": "deployment-key"}, id="set values"), + pytest.param({"api_base": "", "api_key": ""}, id="empty values"), + ], +) +@pytest.mark.usefixtures("no_process_wide_api_base") +def test_boundary_key_resolves_named_credential_values_the_way_the_request_does( + monkeypatch: pytest.MonkeyPatch, deployment_values: dict[str, str] +): from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, ) + from litellm.utils import with_named_credential + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.setenv("OPENAI_API_KEY", "env-key") + deployment: Final = {"model": "openai/gpt-5.5", **deployment_values, "litellm_credential_name": "account-a"} with patch.object( litellm, "credential_list", @@ -1153,46 +1171,14 @@ def test_boundary_key_matches_named_credential_precedence(): ) ], ): - boundary = EncryptedContentAffinityCheck._encryption_boundary_key( - { - "api_base": "https://deployment.example.com", - "api_key": "deployment-key", - "litellm_credential_name": "account-a", - } - ) + boundary: Final = EncryptedContentAffinityCheck._encryption_boundary_key(deployment) + sent_params: Final = with_named_credential(deployment) - assert boundary == ("https://credential.example.com", "credential-key-a") - - -def test_boundary_key_resolves_credential_when_explicit_values_are_empty(): - from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( - EncryptedContentAffinityCheck, + assert boundary is not None + assert boundary == EncryptedContentAffinityCheck._encryption_boundary_key( + {key: value for key, value in sent_params.items() if key != "litellm_credential_name"} ) - with patch.object( - litellm, - "credential_list", - [ - CredentialItem( - credential_name="account-a", - credential_values={ - "api_base": "https://credential.example.com", - "api_key": "credential-key-a", - }, - credential_info={}, - ) - ], - ): - boundary = EncryptedContentAffinityCheck._encryption_boundary_key( - { - "api_base": "", - "api_key": "", - "litellm_credential_name": "account-a", - } - ) - - assert boundary == ("https://credential.example.com", "credential-key-a") - def test_boundary_fallback_matches_deployments_with_same_named_credential_values(): from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( @@ -2728,3 +2714,604 @@ async def test_affinity_keeps_reasoning_when_every_deployment_already_refused_th assert [deployment["model_info"]["id"] for deployment in result] == ["region-a"] assert reasoning_input[0] == original_reasoning + + +def _default_endpoint_deployment(model_id: str, litellm_params: Mapping[str, object]) -> dict[str, object]: + return {"model_name": "gpt-5.5", "litellm_params": dict(litellm_params), "model_info": {"id": model_id}} + + +def _deployment_of(ctx: RoutingContext, model_id: str) -> dict[str, object]: + deployment: Final = ctx.registry().deployment(model_id) + assert deployment is not None + return deployment.model_dump(exclude_none=True) + + +SUMMARY_ONLY_REASONING: Final = {"type": "reasoning", "summary": [{"type": "summary_text", "text": "scattering"}]} +UNMARKED_REASONING: Final = {**SUMMARY_ONLY_REASONING, "id": "rs_test", "encrypted_content": "gAAAAA-blob"} +OPENAI_SHARED_KEY: Final = {"model": "openai/gpt-5.5", "api_key": "shared-key"} + + +@pytest.fixture +def no_process_wide_api_base(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.delenv("AZURE_API_BASE", raising=False) + + +@pytest.mark.usefixtures("no_process_wide_api_base") +@pytest.mark.parametrize("marked", [True, False], ids=["marked origin", "unmarked origin"]) +@pytest.mark.parametrize( + "origin_params, target_params, target_can_decrypt", + [ + pytest.param(OPENAI_SHARED_KEY, OPENAI_SHARED_KEY, True, id="same key on the provider default endpoint"), + pytest.param( + {"model": "gpt-5.5", "custom_llm_provider": "openai", "api_key": "shared-key"}, + OPENAI_SHARED_KEY, + True, + id="provider named apart from the model", + ), + pytest.param(OPENAI_SHARED_KEY, {**OPENAI_SHARED_KEY, "api_key": "other-key"}, False, id="different key"), + pytest.param( + OPENAI_SHARED_KEY, + {**OPENAI_SHARED_KEY, "model": "openai/gpt-5.5-mini"}, + False, + id="same key on the default endpoint for another model", + ), + pytest.param( + {"model": "bedrock_mantle/us-west-2/openai.gpt-6-astra", "api_key": "shared-key"}, + {"model": "bedrock_mantle/us-west-2/openai.gpt-6-astra", "api_key": "shared-key"}, + False, + id="same key on Bedrock Mantle, which serves several vendors", + ), + pytest.param( + {"model": "openrouter/openai/gpt-6.1-sol", "api_key": "shared-key"}, + {"model": "openrouter/anthropic/claude-haiku-5.5", "api_key": "shared-key"}, + False, + id="same key on an aggregator that serves several vendors", + ), + pytest.param( + {"model": "anthropic/claude-opus-5-5", "api_key": "shared-key"}, + {"model": "anthropic/claude-opus-5-5", "api_key": "shared-key"}, + False, + id="provider without a Responses URL builder", + ), + pytest.param( + {"model": "azure/gpt-5.5", "api_key": "shared-key"}, + {"model": "azure/gpt-5.5", "api_key": "shared-key"}, + False, + id="provider whose builder needs an api_base", + ), + pytest.param( + {"model": "github_copilot/gpt-5.5", "api_key": "shared-key"}, + {"model": "github_copilot/gpt-5.5", "api_key": "shared-key"}, + False, + id="same key on a provider that authenticates with its own token", + ), + pytest.param( + {"model": "groq/llama-4-scout"}, + {"model": "groq/llama-4-scout"}, + False, + id="provider with no default key", + ), + pytest.param( + {"model": "vertex_ai/gemini-3-pro"}, + {"model": "vertex_ai/gemini-3-pro"}, + False, + id="provider whose default key is a credentials file", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "api_base": "https://openai-proxy.example.com/v1"}, + OPENAI_SHARED_KEY, + False, + id="explicit api_base beside the default endpoint", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "extra_headers": {"OpenAI-Organization": "org-a"}}, + {**OPENAI_SHARED_KEY, "extra_headers": {"openai-organization": "org-b"}}, + False, + id="same key sent to different OpenAI organizations", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "extra_headers": {"OpenAI-Project": "proj-a"}}, + OPENAI_SHARED_KEY, + False, + id="same key with an OpenAI project beside one without", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "extra_headers": {"Cookie": "session=account-a"}}, + {**OPENAI_SHARED_KEY, "extra_headers": {"cookie": "session=account-b"}}, + False, + id="same key signed in to different accounts by cookie", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "headers": {"api-key": "account-a"}}, + {**OPENAI_SHARED_KEY, "headers": {"api-key": "account-b"}}, + False, + id="explicit key beside different api-key headers", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "organization": "org-a"}, + {**OPENAI_SHARED_KEY, "organization": "org-b"}, + True, + id="organization param the Responses request never sends", + ), + pytest.param( + {**OPENAI_SHARED_KEY, "extra_headers": {"x-team": "a", "OpenAI-Organization": "org-a"}}, + {**OPENAI_SHARED_KEY, "extra_headers": {"x-team": "b", "OpenAI-Organization": "org-a"}}, + True, + id="same OpenAI organization beside unrelated headers", + ), + pytest.param( + { + **OPENAI_SHARED_KEY, + "headers": {"OpenAI-Organization": "org-a"}, + "extra_headers": {"openai-organization": "org-b"}, + }, + {**OPENAI_SHARED_KEY, "extra_headers": {"OpenAI-Organization": "org-b"}}, + True, + id="extra headers override the organization the headers set", + ), + ], +) +def test_hop_strip_keeps_reasoning_between_same_key_deployments_on_the_default_endpoint( + origin_params: dict[str, object], target_params: dict[str, object], target_can_decrypt: bool, marked: bool +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", origin_params), + _default_endpoint_deployment("target", target_params), + ] + ) + ctx: Final = RoutingContext(router) + reasoning: Final = _reasoning_input_minted_by("origin")[0] if marked else dict(UNMARKED_REASONING) + original_reasoning: Final = dict(reasoning) + request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, + request_input, + None, + (_deployment_of(ctx, "target"),), + unmarked_origin=None if marked else "origin", + ) + + assert request_input[0] == (original_reasoning if target_can_decrypt else SUMMARY_ONLY_REASONING) + + +@pytest.mark.usefixtures("no_process_wide_api_base") +@pytest.mark.parametrize( + "target_api_key, target_can_decrypt", + [ + pytest.param(None, True, id="both on the environment key"), + pytest.param("env-key", True, id="explicit key equal to the environment key"), + pytest.param("other-key", False, id="explicit key other than the environment key"), + ], +) +def test_hop_strip_resolves_a_missing_key_from_the_environment( + monkeypatch: pytest.MonkeyPatch, target_api_key: str | None, target_can_decrypt: bool +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.setenv("OPENAI_API_KEY", "env-key") + target_params: Final = {"model": "openai/gpt-5.5", **({"api_key": target_api_key} if target_api_key else {})} + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", {"model": "openai/gpt-5.5"}), + _default_endpoint_deployment("target", target_params), + ] + ) + ctx: Final = RoutingContext(router) + reasoning: Final = _reasoning_input_minted_by("origin")[0] + original_reasoning: Final = dict(reasoning) + request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, request_input, None, (_deployment_of(ctx, "target"),), unmarked_origin=None + ) + + assert request_input[0] == (original_reasoning if target_can_decrypt else SUMMARY_ONLY_REASONING) + + +def test_hop_strip_leaves_input_items_that_are_not_objects_as_sent(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + request_input: Final = ["and sunsets?", *_reasoning_input_minted_by("origin")] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + None, request_input, None, (), unmarked_origin=None + ) + + assert request_input == ["and sunsets?", SUMMARY_ONLY_REASONING, {"role": "user", "content": "and sunsets?"}] + + +GATEWAY_BASE: Final = {"model": "openai/gpt-5.5", "api_base": "https://gateway.example.com/v1"} + + +@pytest.mark.parametrize( + "origin_params, target_params, target_can_decrypt", + [ + pytest.param( + {**GATEWAY_BASE, "extra_headers": {"Authorization": "Bearer sk-account-a"}}, + {**GATEWAY_BASE, "extra_headers": {"Authorization": "Bearer sk-account-b"}}, + False, + id="environment key replaced by different Authorization headers", + ), + pytest.param( + {**GATEWAY_BASE, "api_key": "shared-key", "extra_headers": {"x-idempotency-key": "a"}}, + {**GATEWAY_BASE, "api_key": "shared-key", "extra_headers": {"x-idempotency-key": "b"}}, + True, + id="explicit pair beside different tag headers", + ), + pytest.param( + {**GATEWAY_BASE, "extra_headers": {"Authorization": "Bearer sk-account-a"}}, + {**GATEWAY_BASE, "extra_headers": {"authorization": "Bearer sk-account-a"}}, + True, + id="same Authorization header on both", + ), + ], +) +def test_hop_strip_keeps_accounts_that_sign_in_through_headers_apart( + monkeypatch: pytest.MonkeyPatch, + origin_params: dict[str, object], + target_params: dict[str, object], + target_can_decrypt: bool, +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + monkeypatch.setenv("OPENAI_API_KEY", "env-key") + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", origin_params), + _default_endpoint_deployment("target", target_params), + ] + ) + ctx: Final = RoutingContext(router) + reasoning: Final = _reasoning_input_minted_by("origin")[0] + original_reasoning: Final = dict(reasoning) + request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, request_input, None, (_deployment_of(ctx, "target"),), unmarked_origin=None + ) + + assert request_input[0] == (original_reasoning if target_can_decrypt else SUMMARY_ONLY_REASONING) + + +@pytest.mark.parametrize( + "environment_key, target_model, target_can_decrypt", + [ + pytest.param("shared-key", GATEWAY_BASE["model"], True, id="environment key equal to the explicit key"), + pytest.param("other-key", GATEWAY_BASE["model"], False, id="environment key other than the explicit key"), + pytest.param("shared-key", "openai/gpt-5.5-mini", False, id="equal keys serving another model"), + ], +) +def test_hop_strip_matches_an_explicit_key_and_the_environment_key_on_one_api_base( + monkeypatch: pytest.MonkeyPatch, environment_key: str, target_model: str, target_can_decrypt: bool +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.setenv("OPENAI_API_KEY", environment_key) + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", {**GATEWAY_BASE, "api_key": "shared-key"}), + _default_endpoint_deployment("target", {**GATEWAY_BASE, "model": target_model}), + ] + ) + ctx: Final = RoutingContext(router) + reasoning: Final = _reasoning_input_minted_by("origin")[0] + original_reasoning: Final = dict(reasoning) + request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, request_input, None, (_deployment_of(ctx, "target"),), unmarked_origin=None + ) + + assert request_input[0] == (original_reasoning if target_can_decrypt else SUMMARY_ONLY_REASONING) + + +class _SecretReadCounter(CustomSecretManager): + def __init__(self) -> None: + super().__init__(secret_manager_name="read-counter") + self.reads: int = 0 + + async def async_read_secret( + self, secret_name: str, optional_params: Mapping[str, object] | None = None, timeout: object = None + ) -> str | None: + self.reads += 1 + return None + + def sync_read_secret( + self, secret_name: str, optional_params: Mapping[str, object] | None = None, timeout: object = None + ) -> str | None: + self.reads += 1 + return None + + +@pytest.mark.parametrize( + "deployment_params, reasoning_minted_by", + [ + pytest.param(OPENAI_SHARED_KEY, None, id="no reasoning in the input"), + pytest.param(OPENAI_SHARED_KEY, "target", id="reasoning the hop's target minted"), + pytest.param( + {"model": "groq/llama-4-scout"}, "origin", id="reasoning from a provider with no default boundary" + ), + pytest.param({**GATEWAY_BASE, "api_key": "shared-key"}, "origin", id="reasoning between explicit pairs"), + ], +) +def test_hop_strip_reads_no_secrets_when_no_default_boundary_needs_resolving( + monkeypatch: pytest.MonkeyPatch, deployment_params: dict[str, object], reasoning_minted_by: str | None +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", deployment_params), + _default_endpoint_deployment("target", deployment_params), + ] + ) + ctx: Final = RoutingContext(router) + secret_reads: Final = _SecretReadCounter() + monkeypatch.setattr(litellm, "secret_manager_client", secret_reads) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + request_input: Final = ( + _reasoning_input_minted_by(reasoning_minted_by) + if reasoning_minted_by + else [{"role": "user", "content": "and sunsets?"}] + ) + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, request_input, None, (_deployment_of(ctx, "target"),), unmarked_origin="origin" + ) + + assert secret_reads.reads == 0 + + +@pytest.mark.parametrize( + "unreadable_params", + [ + pytest.param( + {"model": "openai/gpt-5.5", "litellm_credential_name": "mistyped"}, + id="named credential with a value of the wrong type", + ), + pytest.param( + {**GATEWAY_BASE, "api_key": "shared-key", "extra_headers": {1: "one"}}, + id="extra header whose name is not a string", + ), + ], +) +def test_hop_strip_keeps_reasoning_a_target_can_decrypt_beside_a_target_with_unreadable_params( + monkeypatch: pytest.MonkeyPatch, unreadable_params: dict[str, object] +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="mistyped", + credential_values={"api_key": "shared-key", "api_version": datetime.date(2024, 10, 21)}, + credential_info={}, + ) + ], + ) + shared_pair: Final = {**GATEWAY_BASE, "api_key": "shared-key"} + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", shared_pair), + _default_endpoint_deployment("unreadable", unreadable_params), + _default_endpoint_deployment("target", shared_pair), + ] + ) + ctx: Final = RoutingContext(router) + reasoning: Final = _reasoning_input_minted_by("origin")[0] + original_reasoning: Final = dict(reasoning) + request_input: Final = [reasoning, {"role": "user", "content": "and sunsets?"}] + + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + ctx, + request_input, + None, + (_deployment_of(ctx, "unreadable"), _deployment_of(ctx, "target")), + unmarked_origin=None, + ) + + assert request_input[0] == original_reasoning + + +def test_boundary_key_puts_pydantic_and_plain_params_on_one_default_endpoint_boundary(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + from litellm.types.router import LiteLLM_Params + + pydantic_key: Final = EncryptedContentAffinityCheck._encryption_boundary_key(LiteLLM_Params(**OPENAI_SHARED_KEY)) + plain_key: Final = EncryptedContentAffinityCheck._encryption_boundary_key(dict(OPENAI_SHARED_KEY)) + + assert pydantic_key is not None + assert plain_key is not None + assert pydantic_key.is_shared_with(plain_key) + assert plain_key.is_shared_with(pydantic_key) + + +def test_boundary_key_has_no_default_endpoint_for_a_model_no_provider_claims(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + unmapped: Final = {"model": "unmapped-model-name", "api_key": "shared-key"} + + assert EncryptedContentAffinityCheck._encryption_boundary_key(unmapped) is None + + +def test_boundary_key_has_no_default_key_for_a_custom_provider(monkeypatch: pytest.MonkeyPatch): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list, "qa-custom"]) + monkeypatch.setattr(litellm, "_custom_providers", ["qa-custom"]) + + assert EncryptedContentAffinityCheck._encryption_boundary_key({"model": "qa-custom/some-model"}) is None + + +@pytest.mark.asyncio +async def test_affinity_pins_to_the_same_key_peer_on_the_default_endpoint_when_the_origin_cannot_serve(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", OPENAI_SHARED_KEY), + _default_endpoint_deployment("peer", OPENAI_SHARED_KEY), + _default_endpoint_deployment("other-org", {**OPENAI_SHARED_KEY, "api_key": "other-key"}), + ] + ) + ctx: Final = RoutingContext(router) + peer: Final = _deployment_of(ctx, "peer") + reasoning_input: Final = _reasoning_input_minted_by("origin") + original_reasoning: Final = dict(reasoning_input[0]) + + result: Final[list[dict[str, object]]] = await EncryptedContentAffinityCheck(ctx=ctx).async_filter_deployments( + model="gpt-5.5", + healthy_deployments=[peer, _deployment_of(ctx, "other-org")], + messages=None, + request_kwargs={"litellm_metadata": {}, "input": reasoning_input}, + ) + + assert result == [peer] + assert reasoning_input[0] == original_reasoning + + +class _HandAdvancedClock: + def __init__(self) -> None: + self.now: float = 0.0 + + def __call__(self) -> float: + return self.now + + +def _count_secret_reads(monkeypatch: pytest.MonkeyPatch) -> _SecretReadCounter: + secret_reads: Final = _SecretReadCounter() + monkeypatch.setattr(litellm, "secret_manager_client", secret_reads) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + return secret_reads + + +async def _pin_a_turn_minted_by_the_origin( + affinity: object, healthy_deployments: list[dict[str, object]] +) -> tuple[list[dict[str, object]], Mapping[str, object]]: + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + assert isinstance(affinity, EncryptedContentAffinityCheck) + reasoning_input: Final = _reasoning_input_minted_by("origin") + result: Final[list[dict[str, object]]] = await affinity.async_filter_deployments( + model="gpt-5.5", + healthy_deployments=list(healthy_deployments), + messages=None, + request_kwargs={"litellm_metadata": {}, "input": reasoning_input}, + ) + return result, reasoning_input[0] + + +@pytest.mark.asyncio +async def test_repeated_pins_to_a_default_endpoint_peer_ask_the_secret_manager_once_a_minute( + monkeypatch: pytest.MonkeyPatch, +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment("origin", OPENAI_SHARED_KEY), + _default_endpoint_deployment("peer", OPENAI_SHARED_KEY), + _default_endpoint_deployment("other-org", {**OPENAI_SHARED_KEY, "api_key": "other-key"}), + ] + ) + ctx: Final = RoutingContext(router) + peer: Final = _deployment_of(ctx, "peer") + healthy_deployments: Final = [peer, _deployment_of(ctx, "other-org")] + secret_reads: Final = _count_secret_reads(monkeypatch) + EncryptedContentAffinityCheck._encryption_boundary_in_refresh_window.cache_clear() + clock: Final = _HandAdvancedClock() + affinity: Final = EncryptedContentAffinityCheck(ctx=ctx, clock=clock) + + first_pin, first_reasoning = await _pin_a_turn_minted_by_the_origin(affinity, healthy_deployments) + reads_for_the_first_pin: Final = secret_reads.reads + clock.now = 59.0 + second_pin, second_reasoning = await _pin_a_turn_minted_by_the_origin(affinity, healthy_deployments) + reads_within_the_minute: Final = secret_reads.reads + clock.now = 60.0 + third_pin, third_reasoning = await _pin_a_turn_minted_by_the_origin(affinity, healthy_deployments) + + assert first_pin == second_pin == third_pin == [peer] + assert first_reasoning == second_reasoning == third_reasoning == _reasoning_input_minted_by("origin")[0] + assert reads_for_the_first_pin > 0 + assert reads_within_the_minute == reads_for_the_first_pin + assert secret_reads.reads > reads_within_the_minute + + +@pytest.mark.asyncio +async def test_a_changed_named_credential_moves_the_pin_within_the_secret_manager_refresh_minute( + monkeypatch: pytest.MonkeyPatch, +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + def origin_credential(api_key: str) -> list[CredentialItem]: + return [ + CredentialItem(credential_name="origin-account", credential_values={"api_key": api_key}, credential_info={}) + ] + + monkeypatch.setattr(litellm, "credential_list", origin_credential("shared-key")) + router: Final = litellm.Router( + model_list=[ + _default_endpoint_deployment( + "origin", {"model": "openai/gpt-5.5", "litellm_credential_name": "origin-account"} + ), + _default_endpoint_deployment("peer", OPENAI_SHARED_KEY), + _default_endpoint_deployment("other-org", {**OPENAI_SHARED_KEY, "api_key": "other-key"}), + ] + ) + ctx: Final = RoutingContext(router) + peer: Final = _deployment_of(ctx, "peer") + other_org: Final = _deployment_of(ctx, "other-org") + _count_secret_reads(monkeypatch) + EncryptedContentAffinityCheck._encryption_boundary_in_refresh_window.cache_clear() + affinity: Final = EncryptedContentAffinityCheck(ctx=ctx, clock=_HandAdvancedClock()) + + pin_before_the_change, _ = await _pin_a_turn_minted_by_the_origin(affinity, [peer, other_org]) + monkeypatch.setattr(litellm, "credential_list", origin_credential("other-key")) + pin_after_the_change, reasoning_after_the_change = await _pin_a_turn_minted_by_the_origin( + affinity, [peer, other_org] + ) + + assert pin_before_the_change == [peer] + assert pin_after_the_change == [other_org] + assert reasoning_after_the_change == _reasoning_input_minted_by("origin")[0] From 1046716b6d3d45f30d5efad42becbd440d6187ee Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Mon, 2 Nov 2026 05:57:51 -0800 Subject: [PATCH 02/13] fix(anthropic)!: sync the beta header allowlist with what anthropic, vertex_ai, and azure_ai accept (#49409) * fix(anthropic): sync the beta header allowlist with what anthropic, vertex_ai, and azure_ai accept * test(e2e): send the azure_ai beta check without the code execution tool * test(anthropic): require a vertex_ai and azure_ai decision for every anthropic beta * fix(anthropic): forward the betas Claude Code sends that anthropic, vertex_ai, and azure_ai accept * test(e2e): send Vertex's header-only betas in the beta acceptance case * fix(anthropic): keep forwarding Vertex computer-use betas and point beta tests at platforms that still lack them A live probe found each Claude 4.5 and 4.6 model on Vertex needs one of the two computer-use betas for the computer tool, so dropping both broke it. The e2e acceptance test leaves those two out of its all-at-once Vertex call, and the integration and unit tests that relied on a beta being dropped now use one the platform still lacks. * test(anthropic): point unsupported-beta tests at betas the platforms still drop The Azure AI handler test used claude-code-20250219 as the beta Azure drops, and the Rust parse test used it as an unknown value. This PR forwards it on Azure AI and maps it, so both now use values that stay dropped or unknown. * fix(anthropic): add computer-use-2025-11-24 for the computer_20251124 tool The computer_20251124 tool fell through to the default computer-use-2024-10-22 beta, which the provider rejects for that tool version. Map it to computer-use-2025-11-24, which the beta map now forwards on anthropic and azure_ai * fix(anthropic): send no computer use beta for an unknown computer tool version An unknown computer tool version fell back to computer-use-2024-10-22, which the anthropic and azure_ai columns now forward. It now gets no beta, as every platform effectively sent before. Every computer tool version also skips defer_loading and allowed_callers, computer_20251124 included * test(azure_ai): read the branch's beta map in the auto-injected beta wire test * fix(anthropic): keep defer_loading and allowed_callers on the computer_20251124 chat tool * test(anthropic): read the branch's beta map in the azure_ai safeguards beta wire test --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../src/providers/anthropic/beta.rs | 293 +++++++++++++----- .../crates/llms/src/anthropic/beta_headers.rs | 5 +- litellm/anthropic_beta_headers_config.json | 175 ++++++++++- litellm/llms/anthropic/common_utils.py | 20 +- .../test_messages_beta_headers_e2e.py | 74 +++-- .../test_anthropic_safeguards_beta_wire.py | 27 +- .../test_beta_gated_fields_wire.py | 19 +- .../test_azure_ai_claude_messages_wire.py | 29 +- ...vertex_anthropic_inline_tools_beta_wire.py | 2 +- .../test_anthropic_chat_transformation.py | 102 +++++- .../custom_httpx/test_llm_http_handler.py | 2 +- ...partner_models_anthropic_transformation.py | 38 +-- .../test_anthropic_beta_headers_filtering.py | 67 +++- 13 files changed, 674 insertions(+), 179 deletions(-) diff --git a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs index 01ed023999d..46c14a21c59 100644 --- a/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs +++ b/litellm-rust/crates/llms-types/src/providers/anthropic/beta.rs @@ -124,12 +124,86 @@ pub enum AnthropicBeta { ToolSearchTool20251019, #[strum(serialize = "inline-tools-2026-09-15")] InlineTools20260915, + #[strum(serialize = "afk-mode-2026-01-31")] + AfkMode20260131, + #[strum(serialize = "agent-memory-2026-07-22")] + AgentMemory20260722, + #[strum(serialize = "analytics-usage-cost-2026-08-24")] + AnalyticsUsageCost20260824, + #[strum(serialize = "cache-diagnosis-2026-04-07")] + CacheDiagnosis20260407, + #[strum(serialize = "ce-plugins-2026-09-01")] + CePlugins20260901, + #[strum(serialize = "ce-user-management-2026-07-13")] + CeUserManagement20260713, + #[strum(serialize = "code-execution-2025-05-22")] + CodeExecution20250522, + #[strum(serialize = "computer-use-2024-10-22")] + ComputerUse20241022, + #[strum(serialize = "dev-full-thinking-2025-05-14")] + DevFullThinking20250514, + #[strum(serialize = "dreaming-2026-04-21")] + Dreaming20260421, + #[strum(serialize = "experimental-claude-tag-2026-10-21")] + ExperimentalClaudeTag20261021, + #[strum(serialize = "extended-cache-ttl-2025-04-11")] + ExtendedCacheTtl20250411, + #[strum(serialize = "fallback-credit-2026-06-01")] + FallbackCredit20260601, + #[strum(serialize = "fallback-credit-2026-07-01")] + FallbackCredit20260701, + #[strum(serialize = "managed-agents-2026-04-01")] + ManagedAgents20260401, + #[strum(serialize = "mcp-client-2026-09-15")] + McpClient20260915, + #[strum(serialize = "mcp-tunnels-2026-06-22")] + McpTunnels20260622, + #[strum(serialize = "message-batches-2024-09-24")] + MessageBatches20240924, + #[strum(serialize = "message-threads-2026-08-12")] + MessageThreads20260812, + #[strum(serialize = "mid-conversation-system-clear-at-2026-08-21")] + MidConversationSystemClearAt20260821, + #[strum(serialize = "model-context-window-exceeded-2025-08-26")] + ModelContextWindowExceeded20250826, + #[strum(serialize = "output-300k-2026-03-24")] + Output300k20260324, + #[strum(serialize = "pdfs-2024-09-25")] + Pdfs20240925, + #[strum(serialize = "prompt-caching-2024-07-31")] + PromptCaching20240731, + #[strum(serialize = "redact-thinking-2026-02-12")] + RedactThinking20260212, + #[strum(serialize = "server-side-fallback-2026-06-01")] + ServerSideFallback20260601, + #[strum(serialize = "server-side-fallback-2026-07-01")] + ServerSideFallback20260701, + #[strum(serialize = "spend-limit-reads-2026-09-26")] + SpendLimitReads20260926, + #[strum(serialize = "structured-outputs-2025-12-15")] + StructuredOutputs20251215, + #[strum(serialize = "task-budgets-2026-03-13")] + TaskBudgets20260313, + #[strum(serialize = "telemetry-destinations-2026-08-11")] + TelemetryDestinations20260811, + #[strum(serialize = "thinking-token-count-2026-05-13")] + ThinkingTokenCount20260513, + #[strum(serialize = "token-counting-2024-11-01")] + TokenCounting20241101, + #[strum(serialize = "user-profiles-2026-03-24")] + UserProfiles20260324, + #[strum(serialize = "user-profiles-2026-08-18")] + UserProfiles20260818, + #[strum(serialize = "user-profiles-2026-09-04")] + UserProfiles20260904, + #[strum(serialize = "webhooks-2026-10-07")] + Webhooks20261007, #[strum(default, transparent)] Other(String), } impl AnthropicBeta { - pub const KNOWN: [Self; 40] = [ + pub const KNOWN: [Self; 77] = [ Self::Oauth20250420, Self::WebFetch20250910, Self::WebSearch20250305, @@ -170,6 +244,43 @@ impl AnthropicBeta { Self::ToolExamples20251029, Self::ToolSearchTool20251019, Self::InlineTools20260915, + Self::AfkMode20260131, + Self::AgentMemory20260722, + Self::AnalyticsUsageCost20260824, + Self::CacheDiagnosis20260407, + Self::CePlugins20260901, + Self::CeUserManagement20260713, + Self::CodeExecution20250522, + Self::ComputerUse20241022, + Self::DevFullThinking20250514, + Self::Dreaming20260421, + Self::ExperimentalClaudeTag20261021, + Self::ExtendedCacheTtl20250411, + Self::FallbackCredit20260601, + Self::FallbackCredit20260701, + Self::ManagedAgents20260401, + Self::McpClient20260915, + Self::McpTunnels20260622, + Self::MessageBatches20240924, + Self::MessageThreads20260812, + Self::MidConversationSystemClearAt20260821, + Self::ModelContextWindowExceeded20250826, + Self::Output300k20260324, + Self::Pdfs20240925, + Self::PromptCaching20240731, + Self::RedactThinking20260212, + Self::ServerSideFallback20260601, + Self::ServerSideFallback20260701, + Self::SpendLimitReads20260926, + Self::StructuredOutputs20251215, + Self::TaskBudgets20260313, + Self::TelemetryDestinations20260811, + Self::ThinkingTokenCount20260513, + Self::TokenCounting20241101, + Self::UserProfiles20260324, + Self::UserProfiles20260818, + Self::UserProfiles20260904, + Self::Webhooks20261007, ]; pub fn as_str(&self) -> &str { @@ -195,10 +306,9 @@ impl AnthropicBeta { fn rejected_by(&self) -> &'static [BetaProvider] { match self { - Self::AdvancedToolUse20251120 => { + Self::AdvancedToolUse20251120 | Self::ToolSearchTool20251019 => { &[BetaProvider::BedrockConverse, BetaProvider::Databricks] } - Self::ContextManagement20250627 => &[BetaProvider::BedrockConverse], Self::AdvisorTool20260301 => &[ BetaProvider::AzureAi, BetaProvider::BedrockConverse, @@ -207,98 +317,99 @@ impl AnthropicBeta { BetaProvider::VertexAi, BetaProvider::Databricks, ], - Self::Compact20260904 => &[ + Self::AfkMode20260131 + | Self::CePlugins20260901 + | Self::CeUserManagement20260713 + | Self::Compact20260904 + | Self::Dreaming20260421 + | Self::ExperimentalClaudeTag20261021 + | Self::ExtendedCacheTtl20250411 + | Self::FallbackCredit20260601 + | Self::FallbackCredit20260701 + | Self::InlineTools20260915 + | Self::ManagedAgents20260401 + | Self::McpTunnels20260622 + | Self::MidConversationSystemClearAt20260821 + | Self::ModelContextWindowExceeded20250826 + | Self::Oauth20250420 + | Self::Pdfs20240925 + | Self::PromptCaching20240731 + | Self::PromptCachingScope20260105 + | Self::RedactThinking20260212 + | Self::ServerSideFallback20260601 + | Self::SpendLimitReads20260926 + | Self::StructuredOutputs20251215 + | Self::TaskBudgets20260313 + | Self::TelemetryDestinations20260811 + | Self::ThinkingTokenCount20260513 + | Self::TokenCounting20241101 + | Self::UserProfiles20260324 + | Self::UserProfiles20260818 + | Self::UserProfiles20260904 => &[ BetaProvider::BedrockConverse, BetaProvider::Bedrock, BetaProvider::BedrockMantle, BetaProvider::Databricks, ], - Self::Bash20241022 + Self::AgentMemory20260722 + | Self::Bash20241022 | Self::Bash20250124 + | Self::DevFullThinking20250514 | Self::McpServers20251204 | Self::StructuredOutput20240301 | Self::TextEditor20241022 | Self::TextEditor20250124 => BetaProvider::VARIANTS, - Self::ClaudeCode20250219 | Self::ToolExamples20251029 => &[ - BetaProvider::Anthropic, - BetaProvider::AzureAi, - BetaProvider::BedrockConverse, - BetaProvider::Bedrock, - BetaProvider::VertexAi, - BetaProvider::Databricks, - ], - Self::CodeExecution20250825 + Self::AnalyticsUsageCost20260824 + | Self::CacheDiagnosis20260407 + | Self::CodeExecution20250522 + | Self::CodeExecution20250825 + | Self::ComputerUse20241022 + | Self::FastMode20260201 | Self::FilesApi20250414 | Self::McpClient20250404 | Self::McpClient20251120 - | Self::PromptCachingScope20260105 + | Self::McpClient20260915 + | Self::MessageBatches20240924 + | Self::MessageThreads20260812 + | Self::Output300k20260324 + | Self::ServerSideFallback20260701 | Self::Skills20251002 - | Self::WebFetch20250910 => &[ + | Self::WebFetch20250910 + | Self::Webhooks20261007 => &[ BetaProvider::BedrockConverse, BetaProvider::Bedrock, BetaProvider::BedrockMantle, BetaProvider::VertexAi, BetaProvider::Databricks, ], - Self::Compact20260112 => &[BetaProvider::AzureAi, BetaProvider::BedrockConverse], + Self::ClaudeCode20250219 + | Self::PerTurnControl20260701 + | Self::ToolExamples20251029 => &[ + BetaProvider::BedrockConverse, + BetaProvider::Bedrock, + BetaProvider::Databricks, + ], + Self::Compact20260112 + | Self::ContextManagement20250627 + | Self::DangerousToolUse20260903 => &[BetaProvider::BedrockConverse], Self::ComputerUse20250124 | Self::ComputerUse20251124 | Self::Context1m20250807 + | Self::Effort20251124 + | Self::FineGrainedToolStreaming20250514 | Self::ThinkingBindingControls20260801 | Self::ThinkingDisplayUpdates20260818 => &[], - Self::DangerousToolUse20260903 => &[BetaProvider::BedrockConverse], - Self::Effort20251124 => &[BetaProvider::VertexAi], - Self::FastMode20260201 | Self::Oauth20250420 => &[ - BetaProvider::AzureAi, - BetaProvider::BedrockConverse, - BetaProvider::Bedrock, - BetaProvider::BedrockMantle, - BetaProvider::VertexAi, - BetaProvider::Databricks, - ], - Self::FineGrainedToolStreaming20250514 => { - &[BetaProvider::AzureAi, BetaProvider::VertexAi] - } - Self::InterleavedThinking20250514 | Self::WebSearch20250305 => { - &[BetaProvider::BedrockConverse, BetaProvider::Bedrock] - } - Self::MidConversationOutputConfig20260701 => &[ - BetaProvider::AzureAi, - BetaProvider::BedrockConverse, - BetaProvider::BedrockMantle, - BetaProvider::VertexAi, - BetaProvider::Databricks, - ], - Self::MidConversationToolChanges20260701 => &[ + Self::InterleavedThinking20250514 + | Self::Output128k20250219 + | Self::TokenEfficientTools20250219 + | Self::WebSearch20250305 => &[BetaProvider::BedrockConverse, BetaProvider::Bedrock], + Self::MidConversationOutputConfig20260701 + | Self::MidConversationToolChanges20260701 => &[ BetaProvider::BedrockConverse, BetaProvider::BedrockMantle, BetaProvider::Databricks, ], - Self::Output128k20250219 | Self::TokenEfficientTools20250219 => &[ - BetaProvider::AzureAi, - BetaProvider::BedrockConverse, - BetaProvider::Bedrock, - BetaProvider::VertexAi, - ], - Self::PerTurnControl20260701 => &[ - BetaProvider::BedrockConverse, - BetaProvider::Bedrock, - BetaProvider::Databricks, - ], - Self::StructuredOutputs20251113 => &[BetaProvider::Bedrock, BetaProvider::VertexAi], - Self::ToolSearchTool20251019 => &[ - BetaProvider::Anthropic, - BetaProvider::AzureAi, - BetaProvider::BedrockConverse, - BetaProvider::Databricks, - ], - Self::InlineTools20260915 => &[ - BetaProvider::AzureAi, - BetaProvider::BedrockConverse, - BetaProvider::Bedrock, - BetaProvider::BedrockMantle, - BetaProvider::Databricks, - ], + Self::StructuredOutputs20251113 => &[BetaProvider::Bedrock], Self::Other(_) => BetaProvider::VARIANTS, } } @@ -451,7 +562,44 @@ mod tests { AnthropicBeta::TokenEfficientTools20250219, AnthropicBeta::ToolExamples20251029, AnthropicBeta::ToolSearchTool20251019, - AnthropicBeta::InlineTools20260915 + AnthropicBeta::InlineTools20260915, + AnthropicBeta::AfkMode20260131, + AnthropicBeta::AgentMemory20260722, + AnthropicBeta::AnalyticsUsageCost20260824, + AnthropicBeta::CacheDiagnosis20260407, + AnthropicBeta::CePlugins20260901, + AnthropicBeta::CeUserManagement20260713, + AnthropicBeta::CodeExecution20250522, + AnthropicBeta::ComputerUse20241022, + AnthropicBeta::DevFullThinking20250514, + AnthropicBeta::Dreaming20260421, + AnthropicBeta::ExperimentalClaudeTag20261021, + AnthropicBeta::ExtendedCacheTtl20250411, + AnthropicBeta::FallbackCredit20260601, + AnthropicBeta::FallbackCredit20260701, + AnthropicBeta::ManagedAgents20260401, + AnthropicBeta::McpClient20260915, + AnthropicBeta::McpTunnels20260622, + AnthropicBeta::MessageBatches20240924, + AnthropicBeta::MessageThreads20260812, + AnthropicBeta::MidConversationSystemClearAt20260821, + AnthropicBeta::ModelContextWindowExceeded20250826, + AnthropicBeta::Output300k20260324, + AnthropicBeta::Pdfs20240925, + AnthropicBeta::PromptCaching20240731, + AnthropicBeta::RedactThinking20260212, + AnthropicBeta::ServerSideFallback20260601, + AnthropicBeta::ServerSideFallback20260701, + AnthropicBeta::SpendLimitReads20260926, + AnthropicBeta::StructuredOutputs20251215, + AnthropicBeta::TaskBudgets20260313, + AnthropicBeta::TelemetryDestinations20260811, + AnthropicBeta::ThinkingTokenCount20260513, + AnthropicBeta::TokenCounting20241101, + AnthropicBeta::UserProfiles20260324, + AnthropicBeta::UserProfiles20260818, + AnthropicBeta::UserProfiles20260904, + AnthropicBeta::Webhooks20261007 )] beta: AnthropicBeta, ) { @@ -551,7 +699,7 @@ mod tests { BetaProvider::Anthropic, Some(AnthropicBeta::Oauth20250420) )] - #[case::rejected("effort-2025-11-24", BetaProvider::VertexAi, None)] + #[case::rejected("skills-2025-10-02", BetaProvider::VertexAi, None)] #[case::unknown("example-beta-2099-01-01", BetaProvider::Anthropic, None)] fn on_resolves_known_spellings_held_as_other( #[case] raw: &str, @@ -580,12 +728,11 @@ mod tests { #[test] fn unknown_values_are_kept_verbatim() { - let parsed: AnthropicBeta = "claude-code-20250219".parse().unwrap(); - assert_eq!( - parsed, - AnthropicBeta::Other("claude-code-20250219".to_string()) + let parsed: AnthropicBeta = "example-beta-2099-01-01".parse().unwrap(); + assert!( + matches!(&parsed, AnthropicBeta::Other(value) if value == "example-beta-2099-01-01") ); - assert_eq!(parsed.to_string(), "claude-code-20250219"); + assert_eq!(parsed.to_string(), "example-beta-2099-01-01"); } #[test] diff --git a/litellm-rust/crates/llms/src/anthropic/beta_headers.rs b/litellm-rust/crates/llms/src/anthropic/beta_headers.rs index a0843621043..7842b478dce 100644 --- a/litellm-rust/crates/llms/src/anthropic/beta_headers.rs +++ b/litellm-rust/crates/llms/src/anthropic/beta_headers.rs @@ -96,7 +96,10 @@ mod tests { } #[rstest] - #[case::azure_ai(BetaProvider::AzureAi, "fast-mode-2026-02-01,example-beta-2099-01-01")] + #[case::azure_ai( + BetaProvider::AzureAi, + "advisor-tool-2026-03-01,example-beta-2099-01-01" + )] #[case::anthropic(BetaProvider::Anthropic, "bash_20241022")] #[case::vertex_ai(BetaProvider::VertexAi, " , ")] fn filter_removes_header_when_nothing_survives( diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index 253074b0e1f..bdfddb48522 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -37,7 +37,47 @@ "mid-conversation-output-config-2026-07-01": "mid-conversation-output-config-2026-07-01", "thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18", "mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01", - "inline-tools-2026-09-15": "inline-tools-2026-09-15" + "inline-tools-2026-09-15": "inline-tools-2026-09-15", + "agent-memory-2026-07-22": null, + "analytics-usage-cost-2026-08-24": "analytics-usage-cost-2026-08-24", + "cache-diagnosis-2026-04-07": "cache-diagnosis-2026-04-07", + "ce-plugins-2026-09-01": "ce-plugins-2026-09-01", + "ce-user-management-2026-07-13": "ce-user-management-2026-07-13", + "claude-code-20250219": "claude-code-20250219", + "code-execution-2025-05-22": "code-execution-2025-05-22", + "computer-use-2024-10-22": "computer-use-2024-10-22", + "dev-full-thinking-2025-05-14": null, + "dreaming-2026-04-21": "dreaming-2026-04-21", + "experimental-claude-tag-2026-10-21": "experimental-claude-tag-2026-10-21", + "extended-cache-ttl-2025-04-11": "extended-cache-ttl-2025-04-11", + "fallback-credit-2026-06-01": "fallback-credit-2026-06-01", + "fallback-credit-2026-07-01": "fallback-credit-2026-07-01", + "managed-agents-2026-04-01": "managed-agents-2026-04-01", + "mcp-client-2026-09-15": "mcp-client-2026-09-15", + "mcp-tunnels-2026-06-22": "mcp-tunnels-2026-06-22", + "message-batches-2024-09-24": "message-batches-2024-09-24", + "message-threads-2026-08-12": "message-threads-2026-08-12", + "mid-conversation-system-clear-at-2026-08-21": "mid-conversation-system-clear-at-2026-08-21", + "model-context-window-exceeded-2025-08-26": "model-context-window-exceeded-2025-08-26", + "output-300k-2026-03-24": "output-300k-2026-03-24", + "pdfs-2024-09-25": "pdfs-2024-09-25", + "prompt-caching-2024-07-31": "prompt-caching-2024-07-31", + "server-side-fallback-2026-06-01": "server-side-fallback-2026-06-01", + "server-side-fallback-2026-07-01": "server-side-fallback-2026-07-01", + "spend-limit-reads-2026-09-26": "spend-limit-reads-2026-09-26", + "task-budgets-2026-03-13": "task-budgets-2026-03-13", + "telemetry-destinations-2026-08-11": "telemetry-destinations-2026-08-11", + "thinking-token-count-2026-05-13": "thinking-token-count-2026-05-13", + "token-counting-2024-11-01": "token-counting-2024-11-01", + "tool-examples-2025-10-29": "tool-examples-2025-10-29", + "tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19", + "user-profiles-2026-03-24": "user-profiles-2026-03-24", + "user-profiles-2026-08-18": "user-profiles-2026-08-18", + "user-profiles-2026-09-04": "user-profiles-2026-09-04", + "webhooks-2026-10-07": "webhooks-2026-10-07", + "afk-mode-2026-01-31": "afk-mode-2026-01-31", + "redact-thinking-2026-02-12": "redact-thinking-2026-02-12", + "structured-outputs-2025-12-15": "structured-outputs-2025-12-15" }, "azure_ai": { "advisor-tool-2026-03-01": null, @@ -45,23 +85,21 @@ "bash_20241022": null, "bash_20250124": null, "code-execution-2025-08-25": "code-execution-2025-08-25", - "compact-2026-01-12": null, - "compact-2026-09-04": "compact-2026-09-04", + "compact-2026-01-12": "compact-2026-01-12", "computer-use-2025-01-24": "computer-use-2025-01-24", "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": "context-management-2025-06-27", "dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03", "effort-2025-11-24": "effort-2025-11-24", - "fast-mode-2026-02-01": null, + "fast-mode-2026-02-01": "fast-mode-2026-02-01", "files-api-2025-04-14": "files-api-2025-04-14", - "fine-grained-tool-streaming-2025-05-14": null, + "fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14", "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14", "mcp-client-2025-11-20": "mcp-client-2025-11-20", "mcp-client-2025-04-04": "mcp-client-2025-04-04", "mcp-servers-2025-12-04": null, - "mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01", - "output-128k-2025-02-19": null, + "output-128k-2025-02-19": "output-128k-2025-02-19", "structured-output-2024-03-01": null, "per-turn-control-2026-07-01": "per-turn-control-2026-07-01", "prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05", @@ -71,9 +109,54 @@ "text_editor_20250124": null, "thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01", "thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18", - "token-efficient-tools-2025-02-19": null, + "token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19", "web-fetch-2025-09-10": "web-fetch-2025-09-10", - "web-search-2025-03-05": "web-search-2025-03-05" + "web-search-2025-03-05": "web-search-2025-03-05", + "agent-memory-2026-07-22": null, + "analytics-usage-cost-2026-08-24": "analytics-usage-cost-2026-08-24", + "cache-diagnosis-2026-04-07": "cache-diagnosis-2026-04-07", + "ce-plugins-2026-09-01": "ce-plugins-2026-09-01", + "ce-user-management-2026-07-13": "ce-user-management-2026-07-13", + "claude-code-20250219": "claude-code-20250219", + "code-execution-2025-05-22": "code-execution-2025-05-22", + "compact-2026-09-04": "compact-2026-09-04", + "computer-use-2024-10-22": "computer-use-2024-10-22", + "dev-full-thinking-2025-05-14": null, + "dreaming-2026-04-21": "dreaming-2026-04-21", + "experimental-claude-tag-2026-10-21": "experimental-claude-tag-2026-10-21", + "extended-cache-ttl-2025-04-11": "extended-cache-ttl-2025-04-11", + "fallback-credit-2026-06-01": "fallback-credit-2026-06-01", + "fallback-credit-2026-07-01": "fallback-credit-2026-07-01", + "inline-tools-2026-09-15": "inline-tools-2026-09-15", + "managed-agents-2026-04-01": "managed-agents-2026-04-01", + "mcp-client-2026-09-15": "mcp-client-2026-09-15", + "mcp-tunnels-2026-06-22": "mcp-tunnels-2026-06-22", + "message-batches-2024-09-24": "message-batches-2024-09-24", + "message-threads-2026-08-12": "message-threads-2026-08-12", + "mid-conversation-output-config-2026-07-01": "mid-conversation-output-config-2026-07-01", + "mid-conversation-system-clear-at-2026-08-21": "mid-conversation-system-clear-at-2026-08-21", + "mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01", + "model-context-window-exceeded-2025-08-26": "model-context-window-exceeded-2025-08-26", + "oauth-2025-04-20": "oauth-2025-04-20", + "output-300k-2026-03-24": "output-300k-2026-03-24", + "pdfs-2024-09-25": "pdfs-2024-09-25", + "prompt-caching-2024-07-31": "prompt-caching-2024-07-31", + "server-side-fallback-2026-06-01": "server-side-fallback-2026-06-01", + "server-side-fallback-2026-07-01": "server-side-fallback-2026-07-01", + "spend-limit-reads-2026-09-26": "spend-limit-reads-2026-09-26", + "task-budgets-2026-03-13": "task-budgets-2026-03-13", + "telemetry-destinations-2026-08-11": "telemetry-destinations-2026-08-11", + "thinking-token-count-2026-05-13": "thinking-token-count-2026-05-13", + "token-counting-2024-11-01": "token-counting-2024-11-01", + "tool-examples-2025-10-29": "tool-examples-2025-10-29", + "tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19", + "user-profiles-2026-03-24": "user-profiles-2026-03-24", + "user-profiles-2026-08-18": "user-profiles-2026-08-18", + "user-profiles-2026-09-04": "user-profiles-2026-09-04", + "webhooks-2026-10-07": "webhooks-2026-10-07", + "afk-mode-2026-01-31": "afk-mode-2026-01-31", + "redact-thinking-2026-02-12": "redact-thinking-2026-02-12", + "structured-outputs-2025-12-15": "structured-outputs-2025-12-15" }, "bedrock_converse": { "advisor-tool-2026-03-01": null, @@ -196,34 +279,92 @@ "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": "context-management-2025-06-27", "dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03", - "effort-2025-11-24": null, + "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": null, "files-api-2025-04-14": null, - "fine-grained-tool-streaming-2025-05-14": null, + "fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14", "interleaved-thinking-2025-05-14": "interleaved-thinking-2025-05-14", "mcp-client-2025-11-20": null, "mcp-client-2025-04-04": null, "mcp-servers-2025-12-04": null, - "mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01", - "output-128k-2025-02-19": null, + "output-128k-2025-02-19": "output-128k-2025-02-19", "structured-output-2024-03-01": null, "per-turn-control-2026-07-01": "per-turn-control-2026-07-01", "prompt-caching-scope-2026-01-05": null, "skills-2025-10-02": null, - "structured-outputs-2025-11-13": null, + "structured-outputs-2025-11-13": "structured-outputs-2025-11-13", "text_editor_20241022": null, "text_editor_20250124": null, "thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01", "thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18", - "token-efficient-tools-2025-02-19": null, + "token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19", "tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19", "web-fetch-2025-09-10": null, "web-search-2025-03-05": "web-search-2025-03-05", - "inline-tools-2026-09-15": "inline-tools-2026-09-15" + "inline-tools-2026-09-15": "inline-tools-2026-09-15", + "agent-memory-2026-07-22": null, + "analytics-usage-cost-2026-08-24": null, + "cache-diagnosis-2026-04-07": null, + "ce-plugins-2026-09-01": null, + "ce-user-management-2026-07-13": null, + "claude-code-20250219": "claude-code-20250219", + "code-execution-2025-05-22": null, + "compact-2026-09-04": null, + "computer-use-2024-10-22": null, + "dev-full-thinking-2025-05-14": null, + "dreaming-2026-04-21": null, + "experimental-claude-tag-2026-10-21": null, + "extended-cache-ttl-2025-04-11": null, + "fallback-credit-2026-06-01": "fallback-credit-2026-06-01", + "fallback-credit-2026-07-01": "fallback-credit-2026-07-01", + "managed-agents-2026-04-01": null, + "mcp-client-2026-09-15": null, + "mcp-tunnels-2026-06-22": null, + "message-batches-2024-09-24": null, + "message-threads-2026-08-12": null, + "mid-conversation-output-config-2026-07-01": "mid-conversation-output-config-2026-07-01", + "mid-conversation-system-clear-at-2026-08-21": "mid-conversation-system-clear-at-2026-08-21", + "mid-conversation-tool-changes-2026-07-01": "mid-conversation-tool-changes-2026-07-01", + "model-context-window-exceeded-2025-08-26": "model-context-window-exceeded-2025-08-26", + "oauth-2025-04-20": null, + "output-300k-2026-03-24": null, + "pdfs-2024-09-25": "pdfs-2024-09-25", + "prompt-caching-2024-07-31": null, + "server-side-fallback-2026-06-01": "server-side-fallback-2026-06-01", + "server-side-fallback-2026-07-01": null, + "spend-limit-reads-2026-09-26": null, + "task-budgets-2026-03-13": "task-budgets-2026-03-13", + "telemetry-destinations-2026-08-11": null, + "thinking-token-count-2026-05-13": "thinking-token-count-2026-05-13", + "token-counting-2024-11-01": "token-counting-2024-11-01", + "tool-examples-2025-10-29": "tool-examples-2025-10-29", + "user-profiles-2026-03-24": null, + "user-profiles-2026-08-18": null, + "user-profiles-2026-09-04": null, + "webhooks-2026-10-07": null, + "afk-mode-2026-01-31": "afk-mode-2026-01-31", + "redact-thinking-2026-02-12": null, + "structured-outputs-2025-12-15": "structured-outputs-2025-12-15" }, "vertex_ai_header_only": { + "ce-plugins-2026-09-01": "ce-plugins-2026-09-01", + "ce-user-management-2026-07-13": "ce-user-management-2026-07-13", "compact-2026-09-04": "compact-2026-09-04", - "inline-tools-2026-09-15": "inline-tools-2026-09-15" + "dreaming-2026-04-21": "dreaming-2026-04-21", + "experimental-claude-tag-2026-10-21": "experimental-claude-tag-2026-10-21", + "extended-cache-ttl-2025-04-11": "extended-cache-ttl-2025-04-11", + "inline-tools-2026-09-15": "inline-tools-2026-09-15", + "managed-agents-2026-04-01": "managed-agents-2026-04-01", + "mcp-tunnels-2026-06-22": "mcp-tunnels-2026-06-22", + "oauth-2025-04-20": "oauth-2025-04-20", + "prompt-caching-2024-07-31": "prompt-caching-2024-07-31", + "prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05", + "redact-thinking-2026-02-12": "redact-thinking-2026-02-12", + "spend-limit-reads-2026-09-26": "spend-limit-reads-2026-09-26", + "telemetry-destinations-2026-08-11": "telemetry-destinations-2026-08-11", + "user-profiles-2026-03-24": "user-profiles-2026-03-24", + "user-profiles-2026-08-18": "user-profiles-2026-08-18", + "user-profiles-2026-09-04": "user-profiles-2026-09-04" }, "databricks": { "advisor-tool-2026-03-01": null, diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 4f5dc678bae..c9df903cbd6 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1297,7 +1297,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): return None return anthropic_beta_header.split(",") - def get_computer_tool_beta_header(self, computer_tool_version: str) -> str: + def get_computer_tool_beta_header(self, computer_tool_version: str) -> str | None: """ Get the appropriate beta header for a given computer tool version. @@ -1305,16 +1305,14 @@ class AnthropicModelInfo(BaseLLMModelInfo): computer_tool_version: The computer tool version (e.g., 'computer_20250124', 'computer_20241022') Returns: - The corresponding beta header string + The corresponding beta header string, or None for a version with no known beta """ computer_tool_beta_mapping: Final = { + "computer_20251124": "computer-use-2025-11-24", "computer_20250124": "computer-use-2025-01-24", "computer_20241022": "computer-use-2024-10-22", } - return computer_tool_beta_mapping.get( - computer_tool_version, - "computer-use-2024-10-22", # Default fallback - ) + return computer_tool_beta_mapping.get(computer_tool_version) def get_anthropic_beta_list( self, @@ -1346,9 +1344,8 @@ class AnthropicModelInfo(BaseLLMModelInfo): if effort_used: betas.append(ANTHROPIC_EFFORT_BETA_HEADER) # effort-2025-11-24 - if computer_tool_used: - beta_header: Final = self.get_computer_tool_beta_header(computer_tool_used) - betas.append(beta_header) + if computer_tool_used and (computer_tool_beta := self.get_computer_tool_beta_header(computer_tool_used)): + betas.append(computer_tool_beta) # Anthropic no longer requires the prompt-caching beta header # Prompt caching now works automatically when cache_control is used in messages @@ -1437,9 +1434,8 @@ class AnthropicModelInfo(BaseLLMModelInfo): # Anthropic no longer requires the prompt-caching beta header # Prompt caching now works automatically when cache_control is used in messages # Reference: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching - if computer_tool_used: - beta_header: Final = self.get_computer_tool_beta_header(computer_tool_used) - betas.add(beta_header) + if computer_tool_used and (computer_tool_beta := self.get_computer_tool_beta_header(computer_tool_used)): + betas.add(computer_tool_beta) # if pdf_used: # betas.add("pdfs-2024-09-25") if file_id_used: diff --git a/tests/e2e/llm_translation/test_messages_beta_headers_e2e.py b/tests/e2e/llm_translation/test_messages_beta_headers_e2e.py index 91a31201b0f..59035f8f6b6 100644 --- a/tests/e2e/llm_translation/test_messages_beta_headers_e2e.py +++ b/tests/e2e/llm_translation/test_messages_beta_headers_e2e.py @@ -15,48 +15,85 @@ pytestmark = pytest.mark.e2e _BETA_CONFIG: Final = files("litellm").joinpath("anthropic_beta_headers_config.json") _ANTHROPIC_MODEL: Final = "anthropic/claude-sonnet-5-5" +_VERTEX_MODEL: Final = "vertex_ai/claude-sonnet-5-5" +_VERTEX_LOCATION: Final = "global" +_AZURE_AI_MODEL: Final = "azure_ai/claude-fable-5-1" _BEDROCK_MODEL_ID: Final = "us.anthropic.claude-fable-5-1" _BETA_CONFIG_FILE: Final = TypeAdapter(dict[str, JsonValue]) _BETA_MAPPING: Final = TypeAdapter(dict[str, str | None]) _CODE_EXECUTION: Final[ToolUnionParam] = {"type": "code_execution_20250825", "name": "code_execution"} +# A live Vertex probe on 2026-10-31: each Claude 4.5 and 4.6 model takes one of these and Claude 5.5 takes neither +_VERTEX_PER_MODEL_BETAS: Final = frozenset({"computer-use-2025-01-24", "computer-use-2025-11-24"}) +_NO_PER_MODEL_BETAS: Final[frozenset[str]] = frozenset() -def _mapped_beta_headers(provider: str) -> str: - mapping: Final = _BETA_MAPPING.validate_python(_BETA_CONFIG_FILE.validate_json(_BETA_CONFIG.read_text())[provider]) - return ",".join(name for name, value in mapping.items() if value is not None) +def _mapped_beta_headers(provider: str, per_model: frozenset[str]) -> str: + config: Final = _BETA_CONFIG_FILE.validate_json(_BETA_CONFIG.read_text()) + mapping: Final = _BETA_MAPPING.validate_python(config[provider]) + header_only: Final = _BETA_MAPPING.validate_python(config.get(f"{provider}_header_only", {})) + forwarded: Final = {name for name, value in {**mapping, **header_only}.items() if value is not None} + assert per_model <= forwarded, per_model - forwarded + return ",".join(sorted(forwarded - per_model)) def _deployment(backend: str) -> LiteLLMParamsBody: - if backend.startswith("anthropic/"): - return LiteLLMParamsBody(model=backend, api_key="os.environ/ANTHROPIC_API_KEY") - return LiteLLMParamsBody( - model=backend, - aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", - aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", - aws_region_name="us-east-1", - ) + match backend.split("/", 1)[0]: + case "anthropic": + return LiteLLMParamsBody(model=backend, api_key="os.environ/ANTHROPIC_API_KEY") + case "vertex_ai": + return LiteLLMParamsBody( + model=backend, + vertex_project="os.environ/VERTEXAI_PROJECT", + vertex_location=_VERTEX_LOCATION, + vertex_credentials="os.environ/VERTEXAI_CREDENTIALS", + ) + case "azure_ai": + return LiteLLMParamsBody( + model=backend, + api_base="os.environ/AZURE_AI_API_BASE", + api_key="os.environ/AZURE_AI_API_KEY", + ) + case _: + return LiteLLMParamsBody( + model=backend, + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="us-east-1", + ) class TestAnthropicBetaHeaders: @pytest.mark.parametrize( - ("backend", "beta_provider", "tools"), + ("backend", "beta_provider", "tools", "per_model"), [ - pytest.param(_ANTHROPIC_MODEL, "anthropic", [_CODE_EXECUTION], id="anthropic"), - pytest.param(f"bedrock/invoke/{_BEDROCK_MODEL_ID}", "bedrock", [], id="bedrock-invoke"), - pytest.param(f"bedrock/converse/{_BEDROCK_MODEL_ID}", "bedrock_converse", [], id="bedrock-converse"), + pytest.param(_ANTHROPIC_MODEL, "anthropic", [_CODE_EXECUTION], _NO_PER_MODEL_BETAS, id="anthropic"), + pytest.param(_VERTEX_MODEL, "vertex_ai", [], _VERTEX_PER_MODEL_BETAS, id="vertex-ai"), + pytest.param(_AZURE_AI_MODEL, "azure_ai", [], _NO_PER_MODEL_BETAS, id="azure-ai"), + pytest.param( + f"bedrock/invoke/{_BEDROCK_MODEL_ID}", "bedrock", [], _NO_PER_MODEL_BETAS, id="bedrock-invoke" + ), + pytest.param( + f"bedrock/converse/{_BEDROCK_MODEL_ID}", + "bedrock_converse", + [], + _NO_PER_MODEL_BETAS, + id="bedrock-converse", + ), ], ) @pytest.mark.covers( "llm.messages.anthropic.basic.nonstream.works", "llm.messages.bedrock_invoke.basic.nonstream.works", "llm.messages.bedrock_converse.basic.nonstream.works", + "llm.messages.vertex.basic.nonstream.works", + "llm.messages.azure_foundry.basic.nonstream.works", ) @meta( Subject( domain=Domain.LLM_TRANSLATION, route=Route.MESSAGES, - providers=(Provider.ANTHROPIC, Provider.BEDROCK), - models=(_ANTHROPIC_MODEL, _BEDROCK_MODEL_ID), + providers=(Provider.ANTHROPIC, Provider.BEDROCK, Provider.VERTEX_AI, Provider.AZURE_AI), + models=(_ANTHROPIC_MODEL, _BEDROCK_MODEL_ID, _VERTEX_MODEL, _AZURE_AI_MODEL), mode=Mode.NONSTREAM, ) ) @@ -68,6 +105,7 @@ class TestAnthropicBetaHeaders: backend: str, beta_provider: str, tools: list[ToolUnionParam], + per_model: frozenset[str], ) -> None: model: Final = f"e2e-messages-beta-{unique_marker()}" model_id: Final = proxy.create_model(model, _deployment(backend)) @@ -78,7 +116,7 @@ class TestAnthropicBetaHeaders: max_tokens=64, messages=[{"role": "user", "content": "Say 'hello' and nothing else"}], tools=tools, - extra_headers={"anthropic-beta": _mapped_beta_headers(beta_provider)}, + extra_headers={"anthropic-beta": _mapped_beta_headers(beta_provider, per_model)}, extra_body=NO_PROXY_CACHE, ) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py index 3ac1bdb0334..a1c5f6964bf 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py @@ -3,7 +3,7 @@ import json import re import time import uuid -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass from pathlib import Path from types import MappingProxyType @@ -13,7 +13,7 @@ import anthropic import httpx import openai import pytest -from integration._support.client import Gateway, eventually, object_value +from integration._support.client import Gateway, eventually, gateway_from_environment, object_value from integration._support.database import read_rows from integration._support.process import owned_proxy from integration._support.vertex import service_account_json @@ -417,14 +417,31 @@ def test_a_deployment_dropping_context_management_sends_neither_the_field_nor_it _spend_row(f"msg_{marker}") +@pytest.fixture(scope="module") +def local_beta_map_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + with ( + gateway_from_environment() as shared, + owned_proxy( + shared, + tmp_path_factory.mktemp("safeguards-local-beta-map"), + {"LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS": "true"}, + ) as owned, + ): + yield owned + + @pytest.mark.parametrize("row", tuple(_FOUNDRY_SHAPES)) -def test_azure_ai_messages_forward_the_mapped_client_betas_beside_the_derived_one(gateway: Gateway, row: str) -> None: +def test_azure_ai_messages_forward_the_mapped_client_betas_beside_the_derived_one( + local_beta_map_gateway: Gateway, row: str +) -> None: case: Final = _FOUNDRY_SHAPES[row] marker: Final = uuid.uuid4().hex - with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + with wire_server(anthropic_peer) as wire, local_beta_map_gateway.scenario() as scenario: model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) response: Final = _post( - gateway, messages_body(model, marker, safeguards=case.sent, stream=case.stream), beta=case.client_beta + local_beta_map_gateway, + messages_body(model, marker, safeguards=case.sent, stream=case.stream), + beta=case.client_beta, ) if case.stream: _assert_streamed(response, marker) diff --git a/tests/integration/messages_endpoint/test_beta_gated_fields_wire.py b/tests/integration/messages_endpoint/test_beta_gated_fields_wire.py index 26ad9f20897..6d36003dbfa 100644 --- a/tests/integration/messages_endpoint/test_beta_gated_fields_wire.py +++ b/tests/integration/messages_endpoint/test_beta_gated_fields_wire.py @@ -15,8 +15,8 @@ from pydantic import JsonValue, TypeAdapter pytestmark: Final = pytest.mark.shares_worker_state ANTHROPIC_MODEL: Final = "beta-gated-anthropic" -FOUNDRY_MODEL: Final = "beta-gated-foundry" -FOUNDRY_DROP_PARAMS_MODEL: Final = "beta-gated-foundry-drop-params" +DATABRICKS_MODEL: Final = "beta-gated-databricks" +DATABRICKS_DROP_PARAMS_MODEL: Final = "beta-gated-databricks-drop-params" CLEAR_AT_BETA: Final = "mid-conversation-system-clear-at-2026-08-21" _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _NONCE: Final = re.compile(rb"nonce-([0-9a-f]{32})") @@ -71,9 +71,16 @@ def write_config(directory: Path, wire_url: str) -> Path: { "model_list": [ deployment(ANTHROPIC_MODEL, "anthropic/claude-opus-4-8", f"{wire_url}/anthropic"), - deployment(FOUNDRY_MODEL, "azure_ai/claude-opus-4-8", f"{wire_url}/foundry"), deployment( - FOUNDRY_DROP_PARAMS_MODEL, "azure_ai/claude-opus-4-8", f"{wire_url}/foundry", drop_params=True + DATABRICKS_MODEL, + "databricks/databricks-claude-opus-4-8", + f"{wire_url}/databricks/serving-endpoints", + ), + deployment( + DATABRICKS_DROP_PARAMS_MODEL, + "databricks/databricks-claude-opus-4-8", + f"{wire_url}/databricks/serving-endpoints", + drop_params=True, ), ], "general_settings": { @@ -150,7 +157,7 @@ def system_message(body: bytes) -> dict[str, JsonValue]: def test_a_drop_params_leg_whose_row_lacks_the_beta_sends_neither_the_beta_nor_clear_at(rig: Rig) -> None: request_nonce: Final = nonce() - status, raw = rig.send(FOUNDRY_DROP_PARAMS_MODEL, request_nonce) + status, raw = rig.send(DATABRICKS_DROP_PARAMS_MODEL, request_nonce) assert reply_text(status, raw) == f"answered {request_nonce.removeprefix('nonce-')}" upstream: Final = rig.received(request_nonce) assert CLEAR_AT_BETA not in beta_header(upstream.headers), upstream.headers @@ -159,7 +166,7 @@ def test_a_drop_params_leg_whose_row_lacks_the_beta_sends_neither_the_beta_nor_c def test_a_leg_whose_row_lacks_the_beta_forwards_clear_at_as_sent_without_drop_params(rig: Rig) -> None: request_nonce: Final = nonce() - status, raw = rig.send(FOUNDRY_MODEL, request_nonce) + status, raw = rig.send(DATABRICKS_MODEL, request_nonce) assert reply_text(status, raw) == f"answered {request_nonce.removeprefix('nonce-')}" upstream: Final = rig.received(request_nonce) assert CLEAR_AT_BETA not in beta_header(upstream.headers), upstream.headers diff --git a/tests/integration/providers/test_azure_ai_claude_messages_wire.py b/tests/integration/providers/test_azure_ai_claude_messages_wire.py index c150de7520b..d39c6d2d2c1 100644 --- a/tests/integration/providers/test_azure_ai_claude_messages_wire.py +++ b/tests/integration/providers/test_azure_ai_claude_messages_wire.py @@ -1,10 +1,12 @@ +from collections.abc import Iterator from typing import Final import anthropic import pytest from anthropic.types import Message from integration._support import claude_code as cc -from integration._support.client import Gateway +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter @@ -144,8 +146,23 @@ def test_azure_ai_claude_messages_preserves_target_auth_beta_and_typed_response( assert _JSON_OBJECT.validate_json(provider_request.body) == expected_provider_body, provider_request.body +@pytest.fixture(scope="module") +def local_beta_map_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + with ( + gateway_from_environment() as shared, + owned_proxy( + shared, + tmp_path_factory.mktemp("azure-ai-local-beta-map"), + {"LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS": "true"}, + ) as owned, + ): + yield owned + + @pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) -def test_azure_ai_claude_messages_filters_auto_injected_betas_by_azure_support(gateway: Gateway, stream: bool) -> None: +def test_azure_ai_claude_messages_forwards_the_auto_injected_betas_azure_accepts( + local_beta_map_gateway: Gateway, stream: bool +) -> None: identity: Final = f"msg_azure_claude_beta_filter_{stream}" api_path: Final = "/anthropic/v1/messages" reply_body: Final = cc.message_reply(identity, _BACKEND, _CONTENT, _USAGE) @@ -172,7 +189,7 @@ def test_azure_ai_claude_messages_filters_auto_injected_betas_by_azure_support(g "context_management": context_management, } expected_provider_body: Final = {**request_body_template, "model": _BACKEND} - expected_beta: Final = "structured-outputs-2025-11-13" + expected_beta: Final = "compact-2026-01-12,structured-outputs-2025-11-13" def respond(request: Request) -> Reply: assert request.method == "POST" and request.target == api_path, request.target @@ -185,7 +202,7 @@ def test_azure_ai_claude_messages_filters_auto_injected_betas_by_azure_support(g return Reply(content_type="text/event-stream", chunks=stream_reply) return Reply(body=reply_body) - with wire_server(respond) as wire, gateway.scenario() as scenario: + with wire_server(respond) as wire, local_beta_map_gateway.scenario() as scenario: model: Final = scenario.model( model=f"azure_ai/{_BACKEND}", api_base=wire.url, @@ -193,8 +210,8 @@ def test_azure_ai_claude_messages_filters_auto_injected_betas_by_azure_support(g num_retries=0, ) with anthropic.Anthropic( - api_key=gateway.key, - base_url=str(gateway.client.base_url), + api_key=local_beta_map_gateway.key, + base_url=str(local_beta_map_gateway.client.base_url), max_retries=0, ) as client: if stream: diff --git a/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py b/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py index 3175f2ec075..00f861678cb 100644 --- a/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py +++ b/tests/integration/providers/test_vertex_anthropic_inline_tools_beta_wire.py @@ -25,7 +25,7 @@ from pydantic import JsonValue, TypeAdapter _INLINE_TOOLS: Final = "inline-tools-2026-09-15" _THINKING: Final = "interleaved-thinking-2025-05-14" -_VERTEX_UNSUPPORTED: Final = "effort-2025-11-24" +_VERTEX_UNSUPPORTED: Final = "skills-2025-10-02" _UNKNOWN: Final = "not-a-beta-2099-01-01" _BACKEND: Final = "claude-opus-5-5" _PROJECT: Final = "scripted-project" diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 636e23168cf..63a304ab77d 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -6788,6 +6788,34 @@ def test_eager_input_streaming_not_set_on_computer_use_tool(): assert "eager_input_streaming" not in mapped_tool +def _computer_tool_with_defer_loading_and_allowed_callers(tool_type: str) -> dict[str, object]: + return { + "type": tool_type, + "function": {"name": "computer", "parameters": {"display_width_px": 1024, "display_height_px": 768}}, + "defer_loading": True, + "allowed_callers": ["code_execution_20250825"], + } + + +@pytest.mark.parametrize("tool_type", ["computer_20241022", "computer_20250124"]) +def test_legacy_computer_tools_drop_defer_loading_and_allowed_callers(tool_type: str): + mapped_tool, _ = AnthropicConfig().map_tool_helper(_computer_tool_with_defer_loading_and_allowed_callers(tool_type)) + + assert mapped_tool["type"] == tool_type + assert "defer_loading" not in mapped_tool + assert "allowed_callers" not in mapped_tool + + +def test_computer_20251124_keeps_defer_loading_and_allowed_callers(): + mapped_tool, _ = AnthropicConfig().map_tool_helper( + _computer_tool_with_defer_loading_and_allowed_callers("computer_20251124") + ) + + assert mapped_tool["type"] == "computer_20251124" + assert mapped_tool["defer_loading"] is True + assert mapped_tool["allowed_callers"] == ["code_execution_20250825"] + + def test_eager_input_streaming_reaches_anthropic_request_tools(): result = AnthropicConfig().map_openai_params( non_default_params={"tools": [_eager_chat_tool(eager_input_streaming=True)], "stream": True}, @@ -7520,10 +7548,10 @@ def test_process_anthropic_headers_with_no_matching_headers(): @pytest.mark.parametrize( "computer_tool_used, prompt_caching_set, expected_beta_header", [ - (True, False, True), - (False, True, False), - (True, True, True), - (False, False, False), + ("computer_20250124", False, True), + (None, True, False), + ("computer_20250124", True, True), + (None, False, False), ], ) def test_anthropic_beta_header(computer_tool_used, prompt_caching_set, expected_beta_header): @@ -9233,6 +9261,7 @@ async def test_litellm_anthropic_prompt_caching_system( } +@pytest.mark.usefixtures("local_beta_headers_config") def test_completion_forwards_caller_anthropic_beta_header_with_computer_tool(respx_mock: respx.MockRouter): route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( return_value=httpx.Response( @@ -9271,7 +9300,7 @@ def test_completion_forwards_caller_anthropic_beta_header_with_computer_tool(res sent_betas: Final = frozenset( beta.strip() for beta in route.calls.last.request.headers["anthropic-beta"].split(",") ) - assert sent_betas == frozenset({"computer-use-2025-01-24"}) + assert sent_betas == frozenset({"computer-use-2024-10-22", "computer-use-2025-01-24"}) assert json.loads(route.calls.last.request.content)["tools"] == [ { "type": "computer_20241022", @@ -9283,6 +9312,69 @@ def test_completion_forwards_caller_anthropic_beta_header_with_computer_tool(res ] +# Tool-to-beta pairs from Anthropic's computer use docs; api.anthropic.com answered 400 to +# computer_20251124 without computer-use-2025-11-24 and 200 with it on 2026-10-31 +@pytest.mark.usefixtures("local_beta_headers_config") +@pytest.mark.parametrize( + ("tool_type", "expected_beta"), + [ + ("computer_20241022", "computer-use-2024-10-22"), + ("computer_20250124", "computer-use-2025-01-24"), + ("computer_20251124", "computer-use-2025-11-24"), + ], +) +def test_completion_adds_the_beta_matching_the_computer_tool_version( + respx_mock: respx.MockRouter, tool_type: str, expected_beta: str +): + route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=httpx.Response(200, json=_anthropic_message_body([{"type": "text", "text": "hi"}], "end_turn")) + ) + + litellm.completion( + model="anthropic/claude-opus-5", + api_key="mock_api_key", + messages=[{"role": "user", "content": "Say hi"}], + tools=[ + { + "type": tool_type, + "function": { + "name": "computer", + "parameters": {"display_height_px": 768, "display_width_px": 1024}, + }, + } + ], + ) + + sent_betas: Final = frozenset( + beta.strip() for beta in route.calls.last.request.headers["anthropic-beta"].split(",") + ) + assert sent_betas == frozenset({expected_beta}) + + +@pytest.mark.usefixtures("local_beta_headers_config") +def test_completion_sends_no_computer_use_beta_for_an_unknown_computer_tool_version(respx_mock: respx.MockRouter): + route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=httpx.Response(200, json=_anthropic_message_body([{"type": "text", "text": "hi"}], "end_turn")) + ) + + litellm.completion( + model="anthropic/claude-opus-5", + api_key="mock_api_key", + messages=[{"role": "user", "content": "Say hi"}], + tools=[ + { + "type": "computer_20991231", + "function": { + "name": "computer", + "parameters": {"display_height_px": 768, "display_width_px": 1024}, + }, + } + ], + ) + + assert "anthropic-beta" not in route.calls.last.request.headers + + def _anthropic_message_body(content: list[dict[str, object]], stop_reason: str) -> dict[str, object]: return { "id": "msg_01XrAv7gc5tQNDuoADra7vB4", diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index a8f6772a76e..464ff01d16b 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -1668,7 +1668,7 @@ _BOTH_FOUNDRY_BETAS: Final = ("context-1m-2025-08-07,context-management-2025-06- pytest.param( None, {"anthropic-beta": "context-1m-2025-08-07"}, - "claude-code-20250219", + "mcp-servers-2025-12-04", ("context-1m-2025-08-07",), id="unsupported-client-beta-is-dropped-and-deployment-beta-kept", ), diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index 0a582c95749..2be6794186e 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -640,36 +640,16 @@ async def test_vertex_ai_anthropic_prompt_caching_does_not_add_beta_header( assert "anthropic_beta" not in json.loads(request.content) -def test_vertex_ai_partner_models_anthropic_remove_prompt_caching_scope_beta_header(): - """ - Test that remove_unsupported_beta correctly filters out prompt-caching-scope-2026-01-05 - from the anthropic-beta headers. - """ +def test_vertex_ai_partner_models_anthropic_forwards_prompt_caching_scope_beta_header( + local_beta_headers_config: None, +) -> None: + """A live Vertex rawPredict probe on 2026-10-28 accepted prompt-caching-scope-2026-01-05, so it goes upstream + while a beta missing from the vertex_ai column is still dropped.""" + headers: Final = update_headers_with_filtered_beta( + {"anthropic-beta": "other-feature,prompt-caching-scope-2026-01-05,web-search-2025-03-05"}, "vertex_ai" + ) - # This beta header should be removed - PROMPT_CACHING_BETA_HEADER = "prompt-caching-scope-2026-01-05" - headers = { - "anthropic-beta": f"other-feature,{PROMPT_CACHING_BETA_HEADER},web-search-2025-03-05" - } - - headers = update_headers_with_filtered_beta(headers, "vertex_ai") - - beta_header = headers.get("anthropic-beta") - assert PROMPT_CACHING_BETA_HEADER not in ( - beta_header or "" - ), f"{PROMPT_CACHING_BETA_HEADER} should be filtered out" - assert "other-feature" not in ( - beta_header or "" - ), "Other non-excluded beta headers should remain" - assert "web-search-2025-03-05" in ( - beta_header or "" - ), "Other non-excluded beta headers should remain" - # If prompt-caching was the only value, header should be removed completely - headers2 = {"anthropic-beta": PROMPT_CACHING_BETA_HEADER} - headers2 = update_headers_with_filtered_beta(headers2, "vertex_ai") - assert ( - "anthropic-beta" not in headers2 - ), "Header should be removed if no supported values remain" + assert headers == {"anthropic-beta": "prompt-caching-scope-2026-01-05,web-search-2025-03-05"} def test_vertex_ai_anthropic_output_config_effort_only_forwarded(): diff --git a/tests/unit/test_anthropic_beta_headers_filtering.py b/tests/unit/test_anthropic_beta_headers_filtering.py index c8fef95cf98..c33f207da74 100644 --- a/tests/unit/test_anthropic_beta_headers_filtering.py +++ b/tests/unit/test_anthropic_beta_headers_filtering.py @@ -10,10 +10,11 @@ This test validates: import json import os -from typing import Dict, Final, List +from typing import Dict, Final, List, get_args from unittest.mock import AsyncMock, MagicMock, patch import pytest +from anthropic.types.anthropic_beta_param import AnthropicBetaParam import litellm from litellm.anthropic_beta_headers_manager import ( @@ -25,6 +26,13 @@ from litellm.anthropic_beta_headers_manager import ( ) +def _literal_strings(annotation: object) -> frozenset[str]: + args: Final = get_args(annotation) + if not args: + return frozenset({annotation}) if isinstance(annotation, str) else frozenset() + return frozenset[str]().union(*(_literal_strings(arg) for arg in args)) + + class TestAnthropicBetaHeadersFiltering: """Test beta header filtering and mapping for all providers.""" @@ -50,21 +58,24 @@ class TestAnthropicBetaHeadersFiltering: """Get all beta headers from the anthropic provider config.""" return list(self.config.get("anthropic", {}).keys()) + def get_header_column(self, provider: str) -> Dict[str, str | None]: + return {**self.config.get(provider, {}), **self.config.get(f"{provider}_header_only", {})} + def get_supported_headers(self, provider: str) -> List[str]: """Get headers with non-null values for a provider.""" - provider_config = self.config.get(provider, {}) + provider_config = self.get_header_column(provider) return [ header for header, value in provider_config.items() if value is not None ] def get_unsupported_headers(self, provider: str) -> List[str]: """Get headers with null values for a provider.""" - provider_config = self.config.get(provider, {}) + provider_config = self.get_header_column(provider) return [header for header, value in provider_config.items() if value is None] def get_mapped_headers(self, provider: str) -> Dict[str, str]: """Get mapping of input headers to provider-specific headers.""" - provider_config = self.config.get(provider, {}) + provider_config = self.get_header_column(provider) return { header: value for header, value in provider_config.items() @@ -407,7 +418,15 @@ class TestAnthropicBetaHeadersFiltering: test_case["expected"] in filtered ), f"Header '{test_case['input']}' should be mapped to '{test_case['expected']}' for {test_case['provider']}, but got: {filtered}" - @pytest.mark.parametrize("beta", ["compact-2026-09-04", "inline-tools-2026-09-15"]) + @pytest.mark.parametrize( + "beta", + [ + "compact-2026-09-04", + "inline-tools-2026-09-15", + "prompt-caching-scope-2026-01-05", + "redact-thinking-2026-02-12", + ], + ) def test_vertex_ai_header_only_beta_is_forwarded_as_the_header_and_dropped_from_the_body( self, beta: str ) -> None: @@ -509,6 +528,44 @@ class TestAnthropicBetaHeadersFiltering: assert filtered == ["inline-tools-2026-09-15"] + @pytest.mark.parametrize("provider", ["anthropic", "azure_ai", "vertex_ai"]) + def test_task_budgets_and_computer_use_betas_reach_every_platform(self, provider: str) -> None: + """Live probes on 2026-10-28 found all three platforms accept task-budgets-2026-03-13. On 2026-10-31 Vertex + ran the computer tool only with the beta each model takes (computer-use-2025-01-24 on Claude Sonnet 4.5 and + Haiku 4.5, computer-use-2025-11-24 on Claude Opus 4.5 and the 4.6 models) and answered 400 to the other, so + both reach Vertex and the model decides, the same as calling Vertex directly.""" + sent: Final = frozenset(("computer-use-2025-01-24", "computer-use-2025-11-24", "task-budgets-2026-03-13")) + + headers: Final = update_headers_with_filtered_beta({"anthropic-beta": ",".join(sorted(sent))}, provider) + + assert frozenset(headers["anthropic-beta"].split(",")) == sent + + @pytest.mark.parametrize("provider", ["anthropic", "azure_ai", "vertex_ai"]) + def test_claude_code_betas_reach_the_probed_platforms(self, provider: str) -> None: + """Claude Code 2.1.299 sends these betas outside the anthropic SDK's list, and live probes on + 2026-10-28 found Anthropic, Vertex, and Azure AI Foundry accept each of them as the anthropic-beta header.""" + claude_code_betas: Final = frozenset( + ("afk-mode-2026-01-31", "redact-thinking-2026-02-12", "structured-outputs-2025-12-15") + ) + + headers: Final = update_headers_with_filtered_beta({"anthropic-beta": ",".join(claude_code_betas)}, provider) + + assert frozenset(headers["anthropic-beta"].split(",")) == claude_code_betas + + @pytest.mark.parametrize("provider", ["anthropic", "azure_ai", "vertex_ai"]) + def test_every_anthropic_sdk_beta_has_a_decision_for_the_probed_platforms(self, provider: str) -> None: + sdk_betas: Final = _literal_strings(AnthropicBetaParam) + provider_column: Final = GetAnthropicBetaHeadersConfig.load_local_beta_headers_config()[provider] + + assert sdk_betas + assert sdk_betas - provider_column.keys() == frozenset() + + @pytest.mark.parametrize("provider", ["azure_ai", "vertex_ai"]) + def test_every_anthropic_beta_has_a_decision_for_vertex_ai_and_azure_ai(self, provider: str) -> None: + config: Final = GetAnthropicBetaHeadersConfig.load_local_beta_headers_config() + + assert config["anthropic"].keys() - config[provider].keys() == set() + def test_null_value_headers_filtered(self): """Test that headers with null values are always filtered out.""" for provider in [ From 6f3251c7e31fb6fed73a786becc6c62f2df56ce7 Mon Sep 17 00:00:00 2001 From: Kaustubh Anil Jain <78426775+kaust1495@users.noreply.github.com> Date: Mon, 2 Nov 2026 22:31:58 +0530 Subject: [PATCH 03/13] fix(responses): serialize sequence_number on every bridged streaming event (#44944) * fix(responses): serialize sequence_number on every bridged streaming event * test(responses): cover sequence_number on bridged streaming events * fix(responses): read undeclared fields through dict accessors and number guardrail block continuations * fix(responses): scope dict reads that see undeclared fields to Responses stream events Moves the extras-aware __getitem__, get, __contains__ and items from the shared response base onto a new base for the 37 stream event classes, so usage, Gemini, Interactions and moderation objects keep reading __dict__ and Gemini responses keep _hidden_params out of .get * fix(router): renumber Responses stream events only after a fallback takes over A stream that never falls back now keeps its provider's own sequence numbers, repeats included, so a starting_after resume still matches them. After a fallback the wrapper continues above the highest number the client saw * fix(guardrails): keep Responses stream numbers unique when LLM Shield flushes held text The restorer's flush copied the last delta and kept its sequence_number, so a native stream repeated a number and a bridged stream sent the flush unnumbered. A flush now takes the done event's slot and later events move up by one, and a stream the provider never closed numbers its flush after the last event. Also pins the created_at in the streaming iterator test fixture instead of reading the clock. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../guardrail_translation/handler.py | 6 +- .../llm_shield_proxy/stream_restorers.py | 30 +++- .../streaming_iterator.py | 125 ++++++--------- litellm/responses/streaming_iterator.py | 12 +- litellm/router_python_core.py | 19 ++- litellm/types/llms/openai.py | 100 +++++++----- ...test_openai_responses_guardrail_handler.py | 28 ++++ .../guardrail_hooks/test_llm_shield_proxy.py | 42 +++++ .../test_streaming_iterator_transformation.py | 66 +++++++- .../unit/responses/test_streaming_iterator.py | 51 ++++++ tests/unit/test_router/test_router.py | 145 ++++++++++++++++++ .../unit/types/llms/test_types_llms_openai.py | 31 ++++ 12 files changed, 511 insertions(+), 144 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index a905bff3521..a18cbe13925 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -1988,7 +1988,7 @@ class OpenAIResponsesHandler(BaseTranslation): "annotations": (), "logprobs": None, } - return ( + events: Final = ( *_open_item_closing_events(responses_so_far), OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, @@ -2033,6 +2033,10 @@ class OpenAIResponsesHandler(BaseTranslation): response=_blocked_response(exc, response_id=response_id, model=model, output_item=item), ), ) + return tuple( + event.model_copy(update={"sequence_number": number}) + for number, event in enumerate(events, start=_next_stream_sequence_number(responses_so_far)) + ) class _BlockedContentPart(TypedDict): diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py index ef007155c46..66a2d427944 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/stream_restorers.py @@ -261,7 +261,8 @@ class ResponsesStreamRestorer: own window, keyed by the family, the item id and the part index. - Any `*.done` event closes the stream of the same family. Whatever its window still holds goes out first, as a copy of that stream's last delta event -- so it carries - the stream's own ids, and repeats that event's `sequence_number`. Then every text + the stream's own ids. It takes the done event's `sequence_number`, and every later + event's number moves up by one, so the numbers stay unique and in order. Then every text field on the done event is restored in full: its string fields other than identifiers, plus any `part` or `item` it repeats. - `response.completed` / `response.incomplete` repeat the whole reply, and are @@ -273,6 +274,8 @@ class ResponsesStreamRestorer: self._rehydrate: Final = rehydrate self._carries: Final[dict[ResponsesStreamKey, str]] = {} # mutable-ok: per-stream windows advanced in place. self._last_deltas: Final[dict[ResponsesStreamKey, object]] = {} # mutable-ok: newest delta per stream. + self._inserted: int = 0 + self._last_sequence_number: int | None = None async def restore(self, event: object) -> tuple[object, ...]: """The events to emit in place of `event`: any flush, then the event itself.""" @@ -281,7 +284,7 @@ class ResponsesStreamRestorer: return (event,) if kind.endswith(".delta") and kind not in RESPONSES_BINARY_DELTAS: await self._restore_delta(event, kind) - return (event,) + return self._numbered((), event) slots: Final[SlotSink] = [] flushed: Final = await self._flush(responses_stream_key(event, kind)) if kind.endswith(".done") else () if kind.endswith(".done"): @@ -294,12 +297,29 @@ class ResponsesStreamRestorer: for item in read_list(read_field(event, "response"), "output"): collect_response_item(item, slots) await rehydrate_slots(slots, self._rehydrate) - return (*flushed, event) + return self._numbered(flushed, event) async def finish(self) -> tuple[object, ...]: """Flushes every stream the provider never closed, e.g. a truncated reply.""" - flushed: Final = tuple([await self._flush(key) for key in tuple(self._carries)]) - return tuple(itertools.chain.from_iterable(flushed)) + flushed: Final = tuple(itertools.chain.from_iterable([await self._flush(key) for key in tuple(self._carries)])) + last: Final = self._last_sequence_number + if last is not None: + for offset, flush in enumerate(flushed, start=1): + write_field(flush, "sequence_number", last + offset) + return flushed + + def _numbered(self, flushed: tuple[object, ...], event: object) -> tuple[object, ...]: + """Gives each flush the slot `event` arrived in and moves `event` and later events up.""" + number: Final = read_field(event, "sequence_number") + if not isinstance(number, int): + return (*flushed, event) + for offset, flush in enumerate(flushed): + write_field(flush, "sequence_number", number + self._inserted + offset) + self._inserted += len(flushed) + self._last_sequence_number = number + self._inserted + if self._inserted: + write_field(event, "sequence_number", self._last_sequence_number) + return (*flushed, event) async def _restore_delta(self, event: object, kind: str) -> None: text: Final = read_field(event, "delta") diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index ddcdc66ff20..521a698b276 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -27,7 +27,11 @@ from litellm.responses.litellm_completion_transformation.tool_search import ( from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) -from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator, stamp_responses_usage_cost +from litellm.responses.streaming_iterator import ( + SEQUENCE_NUMBER_STAMPED_ON_EMISSION, + ResponsesAPIStreamingIterator, + stamp_responses_usage_cost, +) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( PART_UNION_TYPES, @@ -239,14 +243,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return execution def _next_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> OutputItemAddedEvent: - self._sequence_number += 1 - added: Final = OutputItemAddedEvent( + return OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=output_index, item=item, ) - added.__dict__["sequence_number"] = self._sequence_number - return added def _queue_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> None: self._pending_tool_events.append(self._next_output_item_added(output_index, item)) @@ -268,12 +269,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): item: Final = build_tool_search_call( call_id, final_args or self._tool_args_by_call_id.get(call_id, ""), "completed", execution ) - self._sequence_number += 1 self._pending_tool_events.append( OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=output_index, - sequence_number=self._sequence_number, item=BaseLiteLLMOpenAIResponseObject(**item.model_dump()), ) ) @@ -346,12 +345,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if streamed_item is None: return self._sent_compaction_item_done = True - self._sequence_number += 1 self._pending_response_events.append( OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=self._compaction_output_index, - sequence_number=self._sequence_number, item=streamed_item, ) ) @@ -494,15 +491,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): chunk_size = 10 # Match typical OpenAI delta size for i in range(0, len(fn_args_delta), chunk_size): delta_chunk = fn_args_delta[i : i + chunk_size] - self._sequence_number += 1 delta_event: BaseLiteLLMOpenAIResponseObject = FunctionCallArgumentsDeltaEvent( type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, item_id=self._tool_item_id_by_call_id.get(call_id, call_id), output_index=output_index, delta=delta_chunk, ) - # Add sequence_number as extra field (BaseLiteLLMOpenAIResponseObject allows extra fields) - delta_event.__dict__["sequence_number"] = self._sequence_number self._pending_tool_events.append(delta_event) def _queue_final_tool_call_done_events(self, litellm_complete_object: ModelResponse) -> None: @@ -567,33 +561,27 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): chunk_size = 10 # Match typical OpenAI delta size for i in range(0, len(remaining_args), chunk_size): delta_chunk = remaining_args[i : i + chunk_size] - self._sequence_number += 1 delta_event = FunctionCallArgumentsDeltaEvent( type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, item_id=self._tool_item_id_by_call_id.get(call_id, call_id), output_index=output_index, delta=delta_chunk, ) - delta_event.__dict__["sequence_number"] = self._sequence_number self._pending_tool_events.append(delta_event) - self._sequence_number += 1 done_event = FunctionCallArgumentsDoneEvent( type=ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE, item_id=self._tool_item_id_by_call_id.get(call_id, call_id), output_index=output_index, arguments=final_args, ) - done_event.__dict__["sequence_number"] = self._sequence_number self._pending_tool_events.append(done_event) - self._sequence_number += 1 item_kwargs = self._tool_call_item_kwargs(call_id, fn_name, final_args, "completed") item_kwargs["id"] = self._tool_item_id_by_call_id.setdefault(call_id, item_kwargs["id"]) item_done_event = OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=output_index, - sequence_number=self._sequence_number, item=BaseLiteLLMOpenAIResponseObject( **item_kwargs, **(_tool_call_provider_specific_fields(tc) if self.tool_call_provider_fields else {}), @@ -625,16 +613,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): _WEB_SEARCH_PROGRESS_EVENTS[:1] if item.status == "in_progress" else _WEB_SEARCH_PROGRESS_EVENTS ) for event_type, event_class in progress_events: - self._sequence_number += 1 - event = event_class(type=event_type, output_index=output_index, item_id=item.id) - event.__dict__["sequence_number"] = self._sequence_number - self._pending_tool_events.append(event) - self._sequence_number += 1 + self._pending_tool_events.append(event_class(type=event_type, output_index=output_index, item_id=item.id)) self._pending_tool_events.append( OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=output_index, - sequence_number=self._sequence_number, item=BaseLiteLLMOpenAIResponseObject(**item.model_dump()), ) ) @@ -770,31 +753,24 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): """ response_created_event_data: Final = self._default_response_created_event_data() - self._sequence_number += 1 - event: Final = ResponseCreatedEvent( + return ResponseCreatedEvent( type=ResponsesAPIStreamEvents.RESPONSE_CREATED, response=self._with_encoded_response_id(ResponsesAPIResponse(**response_created_event_data)), ) - event.__dict__["sequence_number"] = self._sequence_number - return event def create_response_in_progress_event(self) -> ResponseInProgressEvent: response_in_progress_event_data: Final = self._default_response_created_event_data() response_in_progress_event_data["status"] = "in_progress" - self._sequence_number += 1 - event: Final = ResponseInProgressEvent( + return ResponseInProgressEvent( type=ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, response=self._with_encoded_response_id(ResponsesAPIResponse(**response_in_progress_event_data)), ) - event.__dict__["sequence_number"] = self._sequence_number - return event def create_output_item_added_event(self) -> OutputItemAddedEvent: if self._cached_item_id is None: self._cached_item_id = f"msg_{uuid.uuid4()}" - self._sequence_number += 1 - event: Final = OutputItemAddedEvent( + return OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=self._message_index(), item=BaseLiteLLMOpenAIResponseObject( @@ -807,23 +783,18 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): } ), ) - event.__dict__["sequence_number"] = self._sequence_number - return event def create_content_part_added_event(self) -> ContentPartAddedEvent: if self._cached_item_id is None: self._cached_item_id = f"msg_{uuid.uuid4()}" - self._sequence_number += 1 - event: Final = ContentPartAddedEvent( + return ContentPartAddedEvent( type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED, item_id=self._cached_item_id, output_index=self._message_index(), content_index=0, part=BaseLiteLLMOpenAIResponseObject(**{"type": "output_text", "text": "", "annotations": []}), ) - event.__dict__["sequence_number"] = self._sequence_number - return event def _queue_message_item_added_events(self) -> None: if not self.sent_message_item_added_event: @@ -834,7 +805,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def _queue_message_item_added_event(self) -> None: item_id: Final = self._message_item_id() self.sent_message_item_added_event = True - self._sequence_number += 1 event: Final = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=self._message_index(), @@ -848,7 +818,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): } ), ) - event.__dict__["sequence_number"] = self._sequence_number self._pending_response_events.append(event) def _queue_refusal_events(self, refusal: str) -> None: @@ -885,12 +854,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): part=ContentPartDonePartRefusal(type="refusal", refusal=refusal), ), ) - first_sequence_number: Final = self._sequence_number + 1 - self._sequence_number += len(events) - self._pending_response_events.extend( - event.model_copy(update={"sequence_number": sequence_number}) - for sequence_number, event in enumerate(events, start=first_sequence_number) - ) + self._pending_response_events.extend(events) def _merge_provider_specific_fields(self, src: dict) -> None: """Merge provider_specific_fields using last-value-wins for lists. @@ -944,7 +908,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self, reasoning_item_id: str, reasoning_content: str, - sequence_number: int, ) -> ReasoningSummaryTextDoneEvent: """ Create response.reasoning_summary_text.done event. @@ -954,7 +917,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): "type": "response.reasoning_summary_text.done", "item_id": "rs_0c5dae30e53172980069708ba2f59c8197b71ca9820edad07c", "output_index": 0, - "sequence_number": 97, "summary_index": 0, "text": "**Clarifying the first humans**\n\nThe I'm addressing the user's specific interest." } @@ -963,7 +925,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): type=ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE, item_id=reasoning_item_id, output_index=self._reasoning_index(), - sequence_number=sequence_number, + sequence_number=SEQUENCE_NUMBER_STAMPED_ON_EMISSION, summary_index=0, text=reasoning_content, ) @@ -972,7 +934,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self, reasoning_item_id: str, reasoning_content: str, - sequence_number: int, ) -> ReasoningSummaryPartDoneEvent: """ Create response.reasoning_summary_part.done event. @@ -986,7 +947,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): "type": "summary_text", "text": "**Clarifying the first humans**\n\nThe earlier hominins. It feels important to ensure I'm addressing the user's specific interest." }, - "sequence_number": 98, "summary_index": 0 } """ @@ -994,7 +954,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): type=ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE, item_id=reasoning_item_id, output_index=self._reasoning_index(), - sequence_number=sequence_number, + sequence_number=SEQUENCE_NUMBER_STAMPED_ON_EMISSION, summary_index=0, part=BaseLiteLLMOpenAIResponseObject( **{ @@ -1066,7 +1026,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=self._message_index(), - sequence_number=1, item=BaseLiteLLMOpenAIResponseObject( **{ "id": self._cached_item_id, @@ -1082,7 +1041,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self, reasoning_item_id: str, reasoning_content: str, - sequence_number: int, ) -> OutputItemDoneEvent: """ Create response.output_item.done event for reasoning items. @@ -1091,7 +1049,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): { "type": "response.output_item.done", "output_index": 0, - "sequence_number": 99, "item": { "id": "rs_0c5dae30e53172980069708ba2f59c8197b71ca9820edad07c", "type": "reasoning", @@ -1107,7 +1064,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=self._reasoning_index(), - sequence_number=sequence_number, item=BaseLiteLLMOpenAIResponseObject( **{ "id": reasoning_item_id, @@ -1223,7 +1179,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if chunk.choices[0].finish_reason is None and not _delta_carries_output(delta): return - self._sequence_number += 1 self.sent_output_item_added_event = True # Reasoning-first @@ -1233,20 +1188,20 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._cached_reasoning_item_id = mint_reasoning_item_id() self._reasoning_item_id = self._cached_reasoning_item_id - event = OutputItemAddedEvent( - type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, - output_index=self._reasoning_index(), - item=BaseLiteLLMOpenAIResponseObject( - **{ - "id": self._cached_reasoning_item_id, - "type": "reasoning", - "status": "in_progress", - "summary": None, - } - ), + self._pending_response_events.append( + OutputItemAddedEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + output_index=self._reasoning_index(), + item=BaseLiteLLMOpenAIResponseObject( + **{ + "id": self._cached_reasoning_item_id, + "type": "reasoning", + "status": "in_progress", + "summary": None, + } + ), + ) ) - event.__dict__["sequence_number"] = self._sequence_number - self._pending_response_events.append(event) return # Tool-first @@ -1279,30 +1234,34 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) reasoning_item_id: Final = self._cached_reasoning_item_id self._queue_compaction_item_done() - self._sequence_number += 1 text_done_event: Final = self.create_reasoning_summary_text_done_event( reasoning_item_id=reasoning_item_id, reasoning_content=reasoning_content, - sequence_number=self._sequence_number, ) - self._sequence_number += 1 part_done_event: Final = self.create_reasoning_summary_part_done_event( reasoning_item_id=reasoning_item_id, reasoning_content=reasoning_content, - sequence_number=self._sequence_number, ) - self._sequence_number += 1 reasoning_output_item_done_event: Final = self.create_reasoning_output_item_done_event( reasoning_item_id=reasoning_item_id, reasoning_content=reasoning_content, - sequence_number=self._sequence_number, ) self._pending_response_events.extend([text_done_event, part_done_event, reasoning_output_item_done_event]) self._reasoning_done_emitted = True self._reasoning_active = False + def _stamp_sequence_number(self, event: BaseLiteLLMOpenAIResponseObject) -> BaseLiteLLMOpenAIResponseObject: + stamped: Final = event.model_copy(update={"sequence_number": self._sequence_number}) + self._sequence_number += 1 + return stamped + async def __anext__( self, + ) -> ResponsesAPIStreamingResponse | ResponseCompletedEvent | BaseLiteLLMOpenAIResponseObject: + return self._stamp_sequence_number(await self._anext_unstamped_event()) + + async def _anext_unstamped_event( + self, ) -> ResponsesAPIStreamingResponse | ResponseCompletedEvent | BaseLiteLLMOpenAIResponseObject: try: while True: @@ -1363,6 +1322,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def __next__( self, + ) -> ResponsesAPIStreamingResponse | ResponseCompletedEvent | BaseLiteLLMOpenAIResponseObject: + return self._stamp_sequence_number(self._next_unstamped_event()) + + def _next_unstamped_event( + self, ) -> ResponsesAPIStreamingResponse | ResponseCompletedEvent | BaseLiteLLMOpenAIResponseObject: try: while True: @@ -1504,16 +1468,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if not self.sent_content_part_added_event: self._queue_message_item_added_events() self._hold_tool_calls_behind_text(getattr(chunk.choices[0].delta, "tool_calls", None)) - self._sequence_number += 1 - text_delta_event: Final = OutputTextDeltaEvent( + return OutputTextDeltaEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id=item_id, output_index=self._message_index(), content_index=0, delta=delta_content, ) - text_delta_event.__dict__["sequence_number"] = self._sequence_number - return text_delta_event # Priority 3: Handle tool call deltas (if any) -> queue events and emit them # For each tool call delta, we emit events one at a time to match OpenAI's streaming behavior diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 09ee6828157..da8ff2cd207 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -379,6 +379,7 @@ def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool: PRE_OUTPUT_LIFECYCLE_EVENT_TYPES: Final = frozenset({"response.created", "response.in_progress", "response.queued"}) +SEQUENCE_NUMBER_STAMPED_ON_EMISSION: Final = -1 class BaseResponsesAPIStreamingIterator: @@ -1933,7 +1934,6 @@ def build_synthetic_response_events( _build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed), ] - sequence_number = 0 output_items: Final[Sequence[object]] = getattr(transformed, "output", []) or [] for output_index, output_item in enumerate(output_items): output_item_payload = _dump_response_object(output_item) @@ -2043,35 +2043,31 @@ def build_synthetic_response_events( delta=summary_text[i : i + chunk_size], ) ) - sequence_number += 1 events.append( openai_types.ReasoningSummaryTextDoneEvent( type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE, item_id=item_id, output_index=output_index, - sequence_number=sequence_number, + sequence_number=SEQUENCE_NUMBER_STAMPED_ON_EMISSION, summary_index=summary_index, text=summary_text, ) ) - sequence_number += 1 events.append( openai_types.ReasoningSummaryPartDoneEvent( type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE, item_id=item_id, output_index=output_index, - sequence_number=sequence_number, + sequence_number=SEQUENCE_NUMBER_STAMPED_ON_EMISSION, summary_index=summary_index, part=openai_types.BaseLiteLLMOpenAIResponseObject(**summary_payload), ) ) - sequence_number += 1 events.append( openai_types.OutputItemDoneEvent( type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=output_index, - sequence_number=sequence_number, item=openai_types.BaseLiteLLMOpenAIResponseObject(**output_item_payload), ) ) @@ -2082,7 +2078,7 @@ def build_synthetic_response_events( response=transformed, ) ) - return events + return [event.model_copy(update={"sequence_number": number}) for number, event in enumerate(events)] # --------------------------------------------------------------------------- diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 30b05456987..c63e538c998 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -544,6 +544,7 @@ if TYPE_CHECKING: ResponseAPIUsage, ResponseIncompleteEvent, ResponsesAPIResponse, + ResponsesAPIStreamingResponse, ) Span = _Span @@ -4385,12 +4386,13 @@ class PythonCore(DeprecatedRouterMembers): fallback_headers_adopted: bool = False - def __init__(self, async_generator: AsyncGenerator): + def __init__(self, async_generator: "AsyncGenerator[ResponsesAPIStreamingResponse, None]"): import time from datetime import datetime self._async_generator = async_generator self._billing_source: object = source_iterator + self._last_sequence_number = -1 # Mirror every attribute BaseResponsesAPIStreamingIterator.__init__ # would have set. The wrapper bypasses super().__init__ (it has no # httpx.Response of its own and no provider config to drive), so @@ -4497,7 +4499,16 @@ class PythonCore(DeprecatedRouterMembers): if self.completed_response is None and getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES: self.completed_response = chunk self._yielded_first_chunk = True - return chunk + sequence_number: Final = getattr(chunk, "sequence_number", None) + if not isinstance(sequence_number, int): + return chunk + if not self.fallback_headers_adopted: + self._last_sequence_number = max(self._last_sequence_number, sequence_number) + return chunk + self._last_sequence_number = max(self._last_sequence_number + 1, sequence_number) + if self._last_sequence_number == sequence_number: + return chunk + return chunk.model_copy(update={"sequence_number": self._last_sequence_number}) async def aclose(self): # async generators always expose aclose — no defensive check needed. @@ -4559,7 +4570,9 @@ class PythonCore(DeprecatedRouterMembers): exc, ) - wrapper: Final = FallbackResponsesStreamWrapper(stream_with_fallbacks()) + wrapper: Final = FallbackResponsesStreamWrapper( + stream_with_fallbacks() # pyright: ignore[reportArgumentType] # relays response events typed as object + ) return wrapper def _aresponses_background_stream_can_release( diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index d7e55b57a07..43d71a6c0b4 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1,8 +1,9 @@ import builtins import math -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import ItemsView, Iterable, Mapping, Sequence from enum import Enum from os import PathLike +from types import MappingProxyType from typing import IO, TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeAlias, Union import httpx @@ -1754,17 +1755,40 @@ class ResponsesAPIStreamEvents(str, Enum): ERROR = "error" -class ResponseCreatedEvent(BaseLiteLLMOpenAIResponseObject): +_NO_EXTRAS: Final[Mapping[str, object]] = MappingProxyType({}) + + +class BaseResponsesStreamEvent(BaseLiteLLMOpenAIResponseObject): # frozen-ok: guardrails rewrite events in place + def __getitem__(self, key: str) -> object: + return self._mapping_holding(key)[key] + + def get(self, key: str, default: object = None) -> object: + return self._mapping_holding(key).get(key, default) + + def __contains__(self, key: object) -> bool: + return key in self.__dict__ or key in self._extras() + + def items(self) -> ItemsView[str, object]: + return MappingProxyType({**self.__dict__, **self._extras()}).items() + + def _extras(self) -> Mapping[str, object]: + return self.__pydantic_extra__ or _NO_EXTRAS + + def _mapping_holding(self, key: str) -> Mapping[str, object]: + return self.__dict__ if key in self.__dict__ else self._extras() + + +class ResponseCreatedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.RESPONSE_CREATED] response: ResponsesAPIResponse -class ResponseInProgressEvent(BaseLiteLLMOpenAIResponseObject): +class ResponseInProgressEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS] response: ResponsesAPIResponse -class ResponseCompletedEvent(BaseLiteLLMOpenAIResponseObject): +class ResponseCompletedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.RESPONSE_COMPLETED] response: ResponsesAPIResponse _hidden_params: dict = PrivateAttr(default_factory=dict) @@ -1778,12 +1802,12 @@ class ResponseCompletedEvent(BaseLiteLLMOpenAIResponseObject): self._hidden_params = hidden_params -class ResponseFailedEvent(BaseLiteLLMOpenAIResponseObject): +class ResponseFailedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.RESPONSE_FAILED] response: ResponsesAPIResponse -class ResponseIncompleteEvent(BaseLiteLLMOpenAIResponseObject): +class ResponseIncompleteEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE] response: ResponsesAPIResponse @@ -1791,7 +1815,7 @@ class ResponseIncompleteEvent(BaseLiteLLMOpenAIResponseObject): ResponsesTerminalEvent: TypeAlias = ResponseCompletedEvent | ResponseIncompleteEvent | ResponseFailedEvent -class ResponsePartAddedEvent(BaseLiteLLMOpenAIResponseObject): +class ResponsePartAddedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.RESPONSE_PART_ADDED] item_id: str output_index: int @@ -1799,7 +1823,7 @@ class ResponsePartAddedEvent(BaseLiteLLMOpenAIResponseObject): part: dict -class ReasoningSummaryTextDeltaEvent(BaseLiteLLMOpenAIResponseObject): +class ReasoningSummaryTextDeltaEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA] item_id: str output_index: int @@ -1807,7 +1831,7 @@ class ReasoningSummaryTextDeltaEvent(BaseLiteLLMOpenAIResponseObject): delta: str -class ReasoningSummaryTextDoneEvent(BaseLiteLLMOpenAIResponseObject): +class ReasoningSummaryTextDoneEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE] item_id: str output_index: int @@ -1816,7 +1840,7 @@ class ReasoningSummaryTextDoneEvent(BaseLiteLLMOpenAIResponseObject): text: str -class ReasoningSummaryPartDoneEvent(BaseLiteLLMOpenAIResponseObject): +class ReasoningSummaryPartDoneEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE] item_id: str output_index: int @@ -1825,13 +1849,13 @@ class ReasoningSummaryPartDoneEvent(BaseLiteLLMOpenAIResponseObject): part: BaseLiteLLMOpenAIResponseObject -class OutputItemAddedEvent(BaseLiteLLMOpenAIResponseObject): +class OutputItemAddedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED] output_index: int item: BaseLiteLLMOpenAIResponseObject | None -class OutputItemDoneEvent(BaseLiteLLMOpenAIResponseObject): +class OutputItemDoneEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE] output_index: int sequence_number: int = 1 @@ -1851,7 +1875,7 @@ class OpenAIChatCompletionLogprobsContent(TypedDict, total=False): top_logprobs: list[OpenAIChatCompletionLogprobsContentTopLogprobs] -class ContentPartAddedEvent(BaseLiteLLMOpenAIResponseObject): +class ContentPartAddedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.CONTENT_PART_ADDED] item_id: str output_index: int @@ -1879,7 +1903,7 @@ class ContentPartDonePartReasoningText(BaseLiteLLMOpenAIResponseObject): PART_UNION_TYPES = ContentPartDonePartOutputText | ContentPartDonePartRefusal | ContentPartDonePartReasoningText -class ContentPartDoneEvent(BaseLiteLLMOpenAIResponseObject): +class ContentPartDoneEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.CONTENT_PART_DONE] item_id: str output_index: int @@ -1887,7 +1911,7 @@ class ContentPartDoneEvent(BaseLiteLLMOpenAIResponseObject): part: PART_UNION_TYPES -class OutputTextDeltaEvent(BaseLiteLLMOpenAIResponseObject): +class OutputTextDeltaEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA] item_id: str output_index: int @@ -1895,7 +1919,7 @@ class OutputTextDeltaEvent(BaseLiteLLMOpenAIResponseObject): delta: str -class OutputTextAnnotationAddedEvent(BaseLiteLLMOpenAIResponseObject): +class OutputTextAnnotationAddedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED] item_id: str output_index: int @@ -1904,7 +1928,7 @@ class OutputTextAnnotationAddedEvent(BaseLiteLLMOpenAIResponseObject): annotation: dict -class OutputTextDoneEvent(BaseLiteLLMOpenAIResponseObject): +class OutputTextDoneEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE] item_id: str output_index: int @@ -1912,7 +1936,7 @@ class OutputTextDoneEvent(BaseLiteLLMOpenAIResponseObject): text: str -class RefusalDeltaEvent(BaseLiteLLMOpenAIResponseObject): +class RefusalDeltaEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.REFUSAL_DELTA] item_id: str output_index: int @@ -1920,7 +1944,7 @@ class RefusalDeltaEvent(BaseLiteLLMOpenAIResponseObject): delta: str -class RefusalDoneEvent(BaseLiteLLMOpenAIResponseObject): +class RefusalDoneEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.REFUSAL_DONE] item_id: str output_index: int @@ -1928,72 +1952,72 @@ class RefusalDoneEvent(BaseLiteLLMOpenAIResponseObject): refusal: str -class FunctionCallArgumentsDeltaEvent(BaseLiteLLMOpenAIResponseObject): +class FunctionCallArgumentsDeltaEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA] item_id: str output_index: int delta: str -class FunctionCallArgumentsDoneEvent(BaseLiteLLMOpenAIResponseObject): +class FunctionCallArgumentsDoneEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE] item_id: str output_index: int arguments: str -class FileSearchCallInProgressEvent(BaseLiteLLMOpenAIResponseObject): +class FileSearchCallInProgressEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.FILE_SEARCH_CALL_IN_PROGRESS] output_index: int item_id: str -class FileSearchCallSearchingEvent(BaseLiteLLMOpenAIResponseObject): +class FileSearchCallSearchingEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.FILE_SEARCH_CALL_SEARCHING] output_index: int item_id: str -class FileSearchCallCompletedEvent(BaseLiteLLMOpenAIResponseObject): +class FileSearchCallCompletedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.FILE_SEARCH_CALL_COMPLETED] output_index: int item_id: str -class WebSearchCallInProgressEvent(BaseLiteLLMOpenAIResponseObject): +class WebSearchCallInProgressEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.WEB_SEARCH_CALL_IN_PROGRESS] output_index: int item_id: str -class WebSearchCallSearchingEvent(BaseLiteLLMOpenAIResponseObject): +class WebSearchCallSearchingEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.WEB_SEARCH_CALL_SEARCHING] output_index: int item_id: str -class WebSearchCallCompletedEvent(BaseLiteLLMOpenAIResponseObject): +class WebSearchCallCompletedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.WEB_SEARCH_CALL_COMPLETED] output_index: int item_id: str # MCP List Tools Events -class MCPListToolsInProgressEvent(BaseLiteLLMOpenAIResponseObject): +class MCPListToolsInProgressEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.MCP_LIST_TOOLS_IN_PROGRESS] sequence_number: int output_index: int item_id: str -class MCPListToolsCompletedEvent(BaseLiteLLMOpenAIResponseObject): +class MCPListToolsCompletedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.MCP_LIST_TOOLS_COMPLETED] sequence_number: int output_index: int item_id: str -class MCPListToolsFailedEvent(BaseLiteLLMOpenAIResponseObject): +class MCPListToolsFailedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.MCP_LIST_TOOLS_FAILED] sequence_number: int output_index: int @@ -2001,14 +2025,14 @@ class MCPListToolsFailedEvent(BaseLiteLLMOpenAIResponseObject): # MCP Call Events -class MCPCallInProgressEvent(BaseLiteLLMOpenAIResponseObject): +class MCPCallInProgressEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.MCP_CALL_IN_PROGRESS] sequence_number: int output_index: int item_id: str -class MCPCallArgumentsDeltaEvent(BaseLiteLLMOpenAIResponseObject): +class MCPCallArgumentsDeltaEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA] output_index: int item_id: str @@ -2016,7 +2040,7 @@ class MCPCallArgumentsDeltaEvent(BaseLiteLLMOpenAIResponseObject): sequence_number: int -class MCPCallArgumentsDoneEvent(BaseLiteLLMOpenAIResponseObject): +class MCPCallArgumentsDoneEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DONE] output_index: int item_id: str @@ -2024,21 +2048,21 @@ class MCPCallArgumentsDoneEvent(BaseLiteLLMOpenAIResponseObject): sequence_number: int -class MCPCallCompletedEvent(BaseLiteLLMOpenAIResponseObject): +class MCPCallCompletedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.MCP_CALL_COMPLETED] sequence_number: int item_id: str output_index: int -class MCPCallFailedEvent(BaseLiteLLMOpenAIResponseObject): +class MCPCallFailedEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.MCP_CALL_FAILED] sequence_number: int item_id: str output_index: int -class ImageGenerationPartialImageEvent(BaseLiteLLMOpenAIResponseObject): +class ImageGenerationPartialImageEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.IMAGE_GENERATION_PARTIAL_IMAGE] partial_image_index: int b64_json: str @@ -2053,7 +2077,7 @@ class ErrorEventError(BaseLiteLLMOpenAIResponseObject): param: str | dict[str, object] | None = None -class ErrorEvent(BaseLiteLLMOpenAIResponseObject): +class ErrorEvent(BaseResponsesStreamEvent): type: Literal[ResponsesAPIStreamEvents.ERROR] sequence_number: int error: ErrorEventError | None = None @@ -2062,7 +2086,7 @@ class ErrorEvent(BaseLiteLLMOpenAIResponseObject): param: str | Mapping[str, object] | None = None -class GenericEvent(BaseLiteLLMOpenAIResponseObject): +class GenericEvent(BaseResponsesStreamEvent): type: str model_config = ConfigDict(extra="allow", protected_namespaces=()) diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index d80a5e835c9..d8954c4072a 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -4547,6 +4547,34 @@ class TestBuildBlockSseChunks: assert completed["output"][0]["content"][0]["text"] == "Blocked by policy." assert completed["usage"] == {"input_tokens": 7, "output_tokens": 21, "total_tokens": 28} + def test_continuation_numbers_every_event_after_the_last_one_relayed(self): + yielded: Final = [ + {"type": "response.created", "sequence_number": 0, "response": {"id": "resp_live", "output": []}}, + { + "type": "response.output_item.added", + "sequence_number": 1, + "output_index": 0, + "item": {"id": "msg_orig", "type": "message", "role": "assistant"}, + }, + { + "type": "response.output_text.delta", + "sequence_number": 2, + "item_id": "msg_orig", + "output_index": 0, + "content_index": 0, + "delta": "partial", + }, + ] + + payloads: Final = self._payloads( + OpenAIResponsesHandler().build_block_sse_chunks( + self._exc(original_response=yielded), stream_started=True, responses_so_far=yielded + ) + ) + + assert payloads[-1]["type"] == "response.completed" + assert [payload.get("sequence_number") for payload in payloads] == list(range(3, 3 + len(payloads))) + def test_continuation_reads_usage_from_typed_completed_event(self): from litellm.types.llms.openai import ( ResponseCompletedEvent, diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index 3da77a8707e..2603cf92f30 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -1694,6 +1694,48 @@ class TestResponsesStreamRestoration: assert "".join(e.delta for e in out) == "see a@example.com a[" + @pytest.mark.asyncio + async def test_a_flushed_delta_takes_the_done_events_number_and_later_events_move_up(self): + guardrail, _ = _shielded(self.VAULT) + done = OutputTextDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, + item_id="msg_1", + output_index=0, + content_index=0, + text="Mail [EMAIL_1] x[", + sequence_number=3, + ) + + out = await _restore_stream( + guardrail, + [ + self._text_delta("Mail [EMA", 1), + self._text_delta("IL_1] x[", 2), + done, + self._text_delta("next", 4, content_index=1), + ], + ) + + assert [(e.type, e.sequence_number) for e in out] == [ + (ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, 1), + (ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, 2), + (ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, 3), + (ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, 4), + (ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, 5), + ] + assert out[2].delta == "[" + + @pytest.mark.asyncio + async def test_a_truncated_stream_flush_continues_after_the_last_number(self): + guardrail, _ = _shielded(self.VAULT) + + out = await _restore_stream( + guardrail, [self._text_delta("see [EMAIL_1] a[", 7), self._text_delta(" b", 8, content_index=1)] + ) + + assert [e.sequence_number for e in out] == [7, 8, 9] + assert out[-1].delta == "[" + @pytest.mark.asyncio async def test_completed_response_is_restored(self): """The terminal event repeats the whole reply, and clients read it as the answer.""" diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 684f5a89e80..87858741908 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -477,7 +477,6 @@ def test_tool_call_arguments_are_chunked_to_match_openai_behavior(): assert evt is not None assert evt.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED assert evt.output_index == 0 - assert hasattr(evt, "__dict__") and "sequence_number" in evt.__dict__ # Collect all remaining delta events from the pending queue by creating empty chunks delta_events = [] @@ -511,17 +510,11 @@ def test_tool_call_arguments_are_chunked_to_match_openai_behavior(): assert len(evt.delta) <= 10 assert evt.item_id == "fc_call_test" assert evt.output_index == 0 - assert hasattr(evt, "__dict__") and "sequence_number" in evt.__dict__ # Verify all deltas concatenated equal the original arguments concatenated = "".join(evt.delta for evt in delta_events) assert concatenated == large_arguments - # Verify sequence numbers are increasing - sequence_numbers = [evt.__dict__["sequence_number"] for evt in delta_events] - assert sequence_numbers == sorted(sequence_numbers) - assert len(set(sequence_numbers)) == len(sequence_numbers) # All unique - def test_tool_call_delta_without_id_uses_index_mapping(): iterator = LiteLLMCompletionStreamingIterator( @@ -1990,6 +1983,65 @@ def test_the_partial_response_billed_on_disconnect_carries_the_logging_objects_b ), "the bill is priced on the partial's own usage" +def _serialized_sequence_numbers(events: Iterable[BaseLiteLLMOpenAIResponseObject]) -> list[int]: + payloads: Final = tuple(json.loads(event.model_dump_json()) for event in events) + assert all("sequence_number" in payload for payload in payloads), [ + payload["type"] for payload in payloads if "sequence_number" not in payload + ] + return [payload["sequence_number"] for payload in payloads] + + +def _long_text_chunks() -> tuple[ModelResponseStream, ...]: + return (*(_chunk(f"word{i} ") for i in range(300)), _chunk("end", finish_reason="stop")) + + +def _text_then_tool_call_chunks() -> tuple[ModelResponseStream, ...]: + return (_text_opening_a_tool_call_chunk("Checking. "), _tool_arguments_chunk('{"x": 1}', finish_reason="tool_calls")) + + +@pytest.mark.parametrize( + "stream", + [ + lambda: (_chunk("Hello"), _chunk("!", finish_reason="stop")), + _long_text_chunks, + _text_then_tool_call_chunks, + _reasoning_then_text_chunks, + _reasoning_and_text_in_one_chunk, + _reasoning_citation_then_tool_call_finishing_chunks, + lambda: (_tool_call_chunk(), _chunk("", finish_reason="tool_calls")), + lambda: (_refusal_chunk("I can't "), _refusal_chunk("help."), _chunk("", finish_reason="stop")), + lambda: (_chunk("Sure. "), _refusal_chunk("I can't help."), _chunk("", finish_reason="stop")), + ], + ids=[ + "text", + "long_text", + "text_then_tool", + "reasoning_then_text", + "reasoning_and_text", + "citation_tool", + "tool", + "refusal", + "text_then_refusal", + ], +) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_every_event_serializes_its_position_in_the_stream_from_zero(stream: StreamChunks, sync_mode: bool): + iterator: Final = _build_iterator(stream()) + events: Final = list(iterator) if sync_mode else [event async for event in iterator] + + assert _serialized_sequence_numbers(events) == list(range(len(events))) + + +def test_bridged_events_answer_dict_reads_of_their_number(): + events: Final = list(_build_iterator((_chunk("Hello"), _chunk("!", finish_reason="stop")))) + positions: Final = list(range(len(events))) + + assert [event["sequence_number"] for event in events] == positions + assert [event.get("sequence_number") for event in events] == positions + assert all("sequence_number" in event for event in events) + + def _refusal_chunk(refusal: str, finish_reason: str | None = None) -> ModelResponseStream: return ModelResponseStream( id=CHAT_COMPLETION_ID, diff --git a/tests/unit/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py index f98d7ca188e..5072f3d3108 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/tests/unit/responses/test_streaming_iterator.py @@ -2867,3 +2867,54 @@ async def test_websocket_turns_logging_survives_garbage_collection(terminal_even await _run_websocket_turn(logger, terminal_event) assert await _logging_tasks_left_after_garbage_collection(logger) == 1 + + +@pytest.mark.parametrize( + "output", + [ + [ + { + "type": "message", + "id": "msg_seq", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "hello world", "annotations": []}], + } + ], + [ + { + "type": "reasoning", + "id": "rs_seq", + "summary": [{"type": "summary_text", "text": "thinking about it"}], + }, + { + "type": "function_call", + "id": "fc_seq", + "call_id": "call_seq", + "name": "lookup", + "arguments": '{"id":1}', + }, + ], + ], + ids=["message", "reasoning_and_tool"], +) +def test_synthetic_stream_serializes_every_event_position_from_zero(output: list[dict[str, object]]): + transformed: Final = ResponsesAPIResponse.model_validate( + { + "id": "resp_seq", + "created_at": 1741476542, + "status": "completed", + "model": "claude-sonnet-5-5", + "object": "response", + "output": output, + } + ) + + events: Final = streaming_module.build_synthetic_response_events( + transformed=transformed, + logging_obj=None, + chunk_size=5, + ) + + payloads: Final = tuple(json.loads(event.model_dump_json()) for event in events) + assert [payload.get("sequence_number") for payload in payloads] == list(range(len(payloads))) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 9996ffcb098..4f35efdfd6e 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -6776,6 +6776,151 @@ async def test_aresponses_streaming_content_policy_error_event_routes_to_content assert fallback_call.await_args.kwargs["model"] == "gemini/gemini-2.5-flash" +@pytest.mark.asyncio +async def test_aresponses_stream_keeps_numbering_after_a_mid_stream_fallback(): + from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents + + def delta(sequence_number: int, text: str) -> OutputTextDeltaEvent: + return OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta=text, + sequence_number=sequence_number, + ) + + from litellm.exceptions import MidStreamFallbackError + + router: Final = _InjectedFallbackRouter(_AsyncList([delta(0, "c"), delta(1, "d")])) + primary: Final = _make_responses_iterator( + chunks=[delta(0, "a"), delta(1, "b")], + error=MidStreamFallbackError( + message="primary died", + model="primary", + llm_provider="anthropic", + is_pre_first_chunk=False, + generated_content="ab", + ), + model="primary", + ) + + wrapped: Final = await router._aresponses_streaming_iterator( + response=primary, + initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": litellm.aresponses}, + ) + events: Final = [event async for event in wrapped] + + assert [(event.delta, event.sequence_number) for event in events] == [("a", 0), ("b", 1), ("c", 2), ("d", 3)] + + + +@pytest.mark.asyncio +async def test_aresponses_stream_fallback_keeps_higher_numbers_and_passes_unnumbered_events(): + from litellm.exceptions import MidStreamFallbackError + from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents + + def delta(text: str) -> OutputTextDeltaEvent: + return OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id="msg_1", output_index=0, content_index=0, delta=text + ) + + def numbered(sequence_number: int, text: str) -> OutputTextDeltaEvent: + return delta(text).model_copy(update={"sequence_number": sequence_number}) + + router: Final = _InjectedFallbackRouter(_AsyncList([numbered(7, "c"), delta("d"), numbered(8, "e")])) + primary: Final = _make_responses_iterator( + chunks=[numbered(0, "a"), numbered(1, "b")], + error=MidStreamFallbackError( + message="primary died", + model="primary", + llm_provider="anthropic", + is_pre_first_chunk=False, + generated_content="ab", + ), + model="primary", + ) + + wrapped: Final = await router._aresponses_streaming_iterator( + response=primary, + initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": litellm.aresponses}, + ) + events: Final = [event async for event in wrapped] + + assert [(event.delta, event.get("sequence_number")) for event in events] == [ + ("a", 0), + ("b", 1), + ("c", 7), + ("d", None), + ("e", 8), + ] + + +@pytest.mark.asyncio +async def test_aresponses_stream_without_fallback_passes_upstream_numbers_unchanged(): + from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents + + def delta(sequence_number: int, text: str) -> OutputTextDeltaEvent: + return OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta=text, + sequence_number=sequence_number, + ) + + router: Final = _InjectedFallbackRouter(_AsyncList([])) + primary: Final = _make_responses_iterator( + chunks=[delta(0, "a"), delta(1, "b"), delta(1, "c"), delta(0, "d")], model="primary" + ) + + wrapped: Final = await router._aresponses_streaming_iterator( + response=primary, + initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": litellm.aresponses}, + ) + events: Final = [event async for event in wrapped] + + assert [(event.delta, event.sequence_number) for event in events] == [("a", 0), ("b", 1), ("c", 1), ("d", 0)] + + +@pytest.mark.asyncio +async def test_aresponses_stream_fallback_continues_above_the_highest_number_the_primary_sent(): + from litellm.exceptions import MidStreamFallbackError + from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents + + def delta(sequence_number: int, text: str) -> OutputTextDeltaEvent: + return OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta=text, + sequence_number=sequence_number, + ) + + router: Final = _InjectedFallbackRouter(_AsyncList([delta(0, "d")])) + primary: Final = _make_responses_iterator( + chunks=[delta(0, "a"), delta(3, "b"), delta(1, "c")], + error=MidStreamFallbackError( + message="primary died", + model="primary", + llm_provider="anthropic", + is_pre_first_chunk=False, + generated_content="abc", + ), + model="primary", + ) + + wrapped: Final = await router._aresponses_streaming_iterator( + response=primary, + initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": litellm.aresponses}, + ) + events: Final = [event async for event in wrapped] + + assert [(event.delta, event.sequence_number) for event in events] == [("a", 0), ("b", 3), ("c", 1), ("d", 4)] + + @pytest.mark.asyncio async def test_aresponses_streaming_iterator_unwraps_content_policy_trigger_for_fallback_dispatch(): """The fallback dispatch matches on the trigger's own type, so the wrapper must hand it the diff --git a/tests/unit/types/llms/test_types_llms_openai.py b/tests/unit/types/llms/test_types_llms_openai.py index 05801c9feec..8c37e86500b 100644 --- a/tests/unit/types/llms/test_types_llms_openai.py +++ b/tests/unit/types/llms/test_types_llms_openai.py @@ -16,8 +16,10 @@ from litellm.types.llms.openai import ( OpenAIChatCompletionLogprobsContent, OpenAIModerationResponse, OpenAIVideoObject, + OutputTextDeltaEvent, ResponseCompletedEvent, ResponsesAPIResponse, + ResponsesAPIStreamEvents, ) @@ -816,3 +818,32 @@ def test_rebuilding_a_dumped_responses_answer_keeps_the_tool_search_call_keys(ex rebuilt: Final = ResponsesAPIResponse(**answer.model_dump()) assert rebuilt.model_dump()["output"] == [search_call] + + +def _text_delta() -> OutputTextDeltaEvent: + return OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id="msg_1", output_index=0, content_index=0, delta="hi" + ) + + +def test_stream_event_dict_reads_see_a_field_set_after_construction(): + event: Final = _text_delta().model_copy(update={"sequence_number": 3}) + + assert event["sequence_number"] == 3 + assert event.get("sequence_number") == 3 + assert "sequence_number" in event + assert dict(event.items())["sequence_number"] == 3 + assert event["delta"] == "hi" + + +def test_stream_event_dict_reads_see_an_undeclared_field_named_keys(): + assert _text_delta().model_copy(update={"keys": "value"})["keys"] == "value" + + +def test_stream_event_dict_reads_still_miss_absent_keys(): + event: Final = _text_delta() + + assert event.get("sequence_number", "fallback") == "fallback" + assert "sequence_number" not in event + with pytest.raises(KeyError): + event["sequence_number"] From 3b62b6311cb5d67cd798e4cc55525440e8b60742 Mon Sep 17 00:00:00 2001 From: jesus-berri Date: Mon, 2 Nov 2026 09:14:10 -0800 Subject: [PATCH 04/13] fix(proxy): make /update/sso_settings a true PATCH and send only changed SSO fields from the UI (#49527) * fix(proxy): make /update/sso_settings a true PATCH and send only changed SSO fields from the UI * fix(proxy): mask SSO client secrets in the /update/sso_settings response * fix(proxy): keep SSO values that only live in the environment_variables row on a partial SSO PATCH The SSO PATCH strips every SSO env var from the environment_variables config row. Now that the UI only sends changed fields, an omitted value that lived only in that row was dropped from both tables and lost on restart. Carry those values into the SSO row before stripping them * test(ui): pin CLEARED_SSO_SETTINGS to exactly the backend SSOConfig fields With PATCH semantics an omitted key keeps its stored value, so a field missing from the Clear list would silently survive Clear SSO settings --- .../proxy_setting_endpoints.py | 88 ++++++--- .../test_proxy_setting_endpoints.py | 168 +++++++++++++++++- .../src/components/SSOModals.test.tsx | 23 +-- .../src/components/SSOModals.tsx | 134 +++++++------- .../Modals/AddSSOSettingsModal.tsx | 7 +- .../Modals/DeleteSSOSettingsModal.tsx | 26 +-- .../EditSSOSettingsModal.integration.test.tsx | 52 +++++- .../Modals/EditSSOSettingsModal.test.tsx | 2 + .../Modals/EditSSOSettingsModal.tsx | 7 +- .../AdminSettings/SSOSettings/utils.test-d.ts | 12 ++ .../AdminSettings/SSOSettings/utils.test.ts | 57 +++++- .../AdminSettings/SSOSettings/utils.ts | 43 +++++ 12 files changed, 458 insertions(+), 161 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test-d.ts diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index e5348684d78..6d31fa3f226 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -16,6 +16,7 @@ from typing import ( from urllib.parse import urlparse from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile +from prisma.errors import PrismaError from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, create_model, field_validator from pydantic.fields import FieldInfo, PydanticUndefined from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -26,6 +27,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.common_utils.validation_error_body import public_validation_errors from litellm.proxy.config_resolvers import FieldSource, SettingsStore, source_for from litellm.proxy.config_resolvers.settings_store import ConfigOwnedKeyError @@ -1247,6 +1249,37 @@ async def get_sso_settings(): } +_ENVIRONMENT_VARIABLES_PARAM: Final = TypeAdapter(dict[str, object]) + + +def _parse_environment_variables_param(param_value: str | Mapping[str, object] | None) -> Mapping[str, object]: + if isinstance(param_value, str): + return _ENVIRONMENT_VARIABLES_PARAM.validate_python(json.loads(param_value)) + return param_value or {} + + +async def _load_environment_variables_param( + repo: ConfigRepository, +) -> Mapping[str, object] | None | PrismaError | ValueError: + try: + entry: Final = await _config_param_db(repo).find_unique(where={"param_name": "environment_variables"}) + param_value: Final = None if entry is None else _parse_environment_variables_param(entry.param_value) + except (PrismaError, ValueError) as e: + return e + else: + return param_value + + +def _decrypted_sso_fields(environment_variables: Mapping[str, object]) -> Mapping[str, str | None]: + return MappingProxyType( + { + field_name: decrypt_value_helper(value=value, key=env_var, return_original_value=True) + for field_name, env_var in SSO_FIELD_ENV_VARS.items() + if isinstance(value := environment_variables.get(env_var), str) + } + ) + + @router.patch( "/update/sso_settings", tags=["SSO Settings"], @@ -1310,9 +1343,24 @@ async def update_sso_settings( if "general_settings" not in config: config["general_settings"] = {} - # Update environment variables in config and in memory - sso_data: Final = sso_config.model_dump() - for field_name, value in sso_data.items(): + # PATCH semantics: omitted fields keep their stored value, an explicit null clears it. SSO values that + # still live only in the environment_variables row are carried over, since that row is stripped below + environment_variables_param: Final = await _load_environment_variables_param( + ConfigRepository(prisma_client, use_writer=True) + ) + legacy_env_values: Final[Mapping[str, str | None]] = ( + _decrypted_sso_fields(environment_variables_param) + if isinstance(environment_variables_param, Mapping) + else MappingProxyType({}) + ) + stored_values: Final = {k: v for k, v in (before_sso_data or {}).items() if v is not None} + fallback_values: Final = {**legacy_env_values, **stored_values} + submitted: Final = sso_config.model_dump(include=sso_config.model_fields_set) + sso_data: Final = { + field_name: submitted[field_name] if field_name in submitted else fallback_values.get(field_name) + for field_name in SSOConfig.model_fields + } + for field_name, value in submitted.items(): if field_name in SSO_FIELD_ENV_VARS: env_var_name = SSO_FIELD_ENV_VARS[field_name] if value: @@ -1348,25 +1396,17 @@ async def update_sso_settings( ) ) - # Remove SSO-related env vars from config.environment_variables - try: - env_var_entry: Final = await _config_param_db(ConfigRepository(prisma_client, use_writer=True)).find_unique( - where={"param_name": "environment_variables"} + if isinstance(environment_variables_param, (PrismaError, ValueError)): + raise HTTPException( + status_code=500, + detail={"error": f"Error updating environment_variables: {environment_variables_param}"}, ) - # If no environment_variables entry exists, nothing to clean up - if env_var_entry is not None: - if env_var_entry.param_value is not None: - if isinstance(env_var_entry.param_value, str): - environment_variables: Mapping[str, object] = json.loads(env_var_entry.param_value) - else: - environment_variables = dict(env_var_entry.param_value) - else: - environment_variables = {} - + if environment_variables_param is not None: + try: env_vars_to_remove: Final = set(SSO_FIELD_ENV_VARS.values()) filtered_env_vars: Final = { - key: value for key, value in environment_variables.items() if key not in env_vars_to_remove + key: value for key, value in environment_variables_param.items() if key not in env_vars_to_remove } await _config_param_db(ConfigRepository(prisma_client)).update( @@ -1376,16 +1416,16 @@ async def update_sso_settings( }, ) await invalidate_config_param("environment_variables") - except Exception as e: - raise HTTPException( - status_code=500, - detail={"error": f"Error updating environment_variables: {e}"}, - ) + except Exception as e: + raise HTTPException( + status_code=500, + detail={"error": f"Error updating environment_variables: {e}"}, + ) return { "message": "SSO settings updated successfully", "status": "success", - "settings": sso_data, + "settings": mask_sensitive_keys(sso_data, set(SSO_SECRET_FIELDS)), } diff --git a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index ea39d5d59e2..e66050877df 100644 --- a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -689,16 +689,11 @@ class TestProxySettingEndpoints: # Verify settings were updated settings = data["settings"] assert settings["google_client_id"] == new_sso_settings["google_client_id"] - assert ( - settings["google_client_secret"] == new_sso_settings["google_client_secret"] - ) + assert settings["google_client_secret"] == "new_****************cret" assert ( settings["microsoft_client_id"] == new_sso_settings["microsoft_client_id"] ) - assert ( - settings["microsoft_client_secret"] - == new_sso_settings["microsoft_client_secret"] - ) + assert settings["microsoft_client_secret"] == "new_*******************cret" assert settings["proxy_base_url"] == new_sso_settings["proxy_base_url"] assert settings["user_email"] == new_sso_settings["user_email"] @@ -722,6 +717,163 @@ class TestProxySettingEndpoints: create_sso_settings = json.loads(create_data["sso_settings"]) assert create_sso_settings["google_client_id"] == "new_google_client_id" + def _patch_sso_settings(self, monkeypatch, stored, body, environment_variables=None): + """PATCH /update/sso_settings against a stored row; returns (response, saved row).""" + import json + from unittest.mock import AsyncMock, MagicMock + + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + mock_prisma = MagicMock() + stored_record = None + if stored is not None: + stored_record = MagicMock() + stored_record.sso_settings = stored + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock( + return_value=stored_record + ) + mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() + env_var_entry = None + if environment_variables is not None: + env_var_entry = MagicMock() + env_var_entry.param_value = environment_variables + mock_prisma.writer_db.litellm_config.find_unique = AsyncMock(return_value=env_var_entry) + mock_prisma.db.litellm_config.update = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + from litellm.proxy.proxy_server import proxy_config + + monkeypatch.setattr( + proxy_config, "encrypt_env_variables", lambda environment_variables: environment_variables + ) + monkeypatch.setattr(proxy_config, "decrypt_db_variables", lambda variables: variables) + + response = client.patch("/update/sso_settings", json=body) + upsert_kwargs = mock_prisma.db.litellm_ssoconfig.upsert.call_args.kwargs + return response, json.loads(upsert_kwargs["data"]["update"]["sso_settings"]) + + _STORED_GENERIC_SSO = { + "generic_client_id": "stored-client-id", + "generic_client_secret": "stored-client-secret", + "generic_authorization_endpoint": "https://idp.example.com/authorize", + "user_email": "admin@example.com", + "role_mappings": { + "provider": "generic", + "group_claim": "groups", + "default_role": "internal_user", + "roles": {"proxy_admin": ["admins"]}, + }, + } + + def test_update_sso_settings_omitted_fields_keep_stored_values( + self, mock_proxy_config, mock_auth, monkeypatch + ): + monkeypatch.setenv("GENERIC_CLIENT_SECRET", "stored-client-secret") + new_role_mappings = { + "provider": "generic", + "group_claim": "roles", + "default_role": "internal_user", + "roles": {"proxy_admin": ["platform-admins"]}, + } + + response, saved = self._patch_sso_settings( + monkeypatch, self._STORED_GENERIC_SSO, {"role_mappings": new_role_mappings} + ) + + assert response.status_code == 200 + assert saved["generic_client_secret"] == "stored-client-secret" + assert saved["generic_client_id"] == "stored-client-id" + assert saved["user_email"] == "admin@example.com" + assert saved["role_mappings"]["group_claim"] == "roles" + assert os.environ["GENERIC_CLIENT_SECRET"] == "stored-client-secret" + + def test_update_sso_settings_null_clears_only_that_field( + self, mock_proxy_config, mock_auth, monkeypatch + ): + monkeypatch.setenv("GENERIC_CLIENT_SECRET", "stored-client-secret") + monkeypatch.setenv("GENERIC_CLIENT_ID", "stored-client-id") + + response, saved = self._patch_sso_settings( + monkeypatch, + self._STORED_GENERIC_SSO, + {"generic_client_secret": None, "role_mappings": None}, + ) + + assert response.status_code == 200 + assert saved["generic_client_secret"] is None + assert saved["role_mappings"] is None + assert saved["generic_client_id"] == "stored-client-id" + assert "GENERIC_CLIENT_SECRET" not in os.environ + assert os.environ["GENERIC_CLIENT_ID"] == "stored-client-id" + + def test_update_sso_settings_new_value_replaces_stored_value( + self, mock_proxy_config, mock_auth, monkeypatch + ): + response, saved = self._patch_sso_settings( + monkeypatch, self._STORED_GENERIC_SSO, {"generic_client_secret": "rotated-secret"} + ) + + assert response.status_code == 200 + assert saved["generic_client_secret"] == "rotated-secret" + assert saved["generic_client_id"] == "stored-client-id" + assert os.environ["GENERIC_CLIENT_SECRET"] == "rotated-secret" + + def test_update_sso_settings_omitted_field_leaves_env_only_value_alone( + self, mock_proxy_config, mock_auth, monkeypatch + ): + monkeypatch.setenv("GENERIC_CLIENT_SECRET", "env-only-secret") + + response, saved = self._patch_sso_settings( + monkeypatch, None, {"user_email": "admin@example.com"} + ) + + assert response.status_code == 200 + assert saved["user_email"] == "admin@example.com" + assert saved["generic_client_secret"] is None + assert os.environ["GENERIC_CLIENT_SECRET"] == "env-only-secret" + + def test_update_sso_settings_response_masks_stored_secrets( + self, mock_proxy_config, mock_auth, monkeypatch + ): + response, saved = self._patch_sso_settings( + monkeypatch, self._STORED_GENERIC_SSO, {} + ) + + assert response.status_code == 200 + assert saved["generic_client_secret"] == "stored-client-secret" + returned = response.json()["settings"] + assert returned["generic_client_id"] == "stored-client-id" + assert returned["generic_client_secret"] != "stored-client-secret" + assert "stored-client-secret" not in response.text + + def test_update_sso_settings_moves_omitted_env_row_values_into_sso_row( + self, mock_proxy_config, mock_auth, monkeypatch + ): + import json + + environment_variables = { + "MICROSOFT_TENANT": "env-row-tenant", + "MICROSOFT_CLIENT_ID": "env-row-client-id", + "GENERIC_CLIENT_ID": "env-row-generic-id", + "LANGFUSE_PUBLIC_KEY": "unrelated", + } + + response, saved = self._patch_sso_settings( + monkeypatch, + self._STORED_GENERIC_SSO, + {"user_email": "new@example.com", "microsoft_client_id": None}, + environment_variables=environment_variables, + ) + + from litellm.proxy.proxy_server import prisma_client + + assert response.status_code == 200 + assert saved["microsoft_tenant"] == "env-row-tenant" + assert saved["microsoft_client_id"] is None + assert saved["generic_client_id"] == "stored-client-id" + assert saved["user_email"] == "new@example.com" + env_row_update = prisma_client.db.litellm_config.update.call_args.kwargs + assert json.loads(env_row_update["data"]["param_value"]) == {"LANGFUSE_PUBLIC_KEY": "unrelated"} + def test_update_sso_settings_maps_saml_fields_to_env_vars( self, mock_proxy_config, mock_auth, monkeypatch ): @@ -804,7 +956,7 @@ class TestProxySettingEndpoints: mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() mock_prisma.db.litellm_config = MagicMock() - mock_prisma.db.litellm_config.find_unique = AsyncMock( + mock_prisma.writer_db.litellm_config.find_unique = AsyncMock( side_effect=ValueError("cleanup failed") ) mock_prisma.db.litellm_config.update = AsyncMock() diff --git a/ui/litellm-dashboard/src/components/SSOModals.test.tsx b/ui/litellm-dashboard/src/components/SSOModals.test.tsx index bee5f13c626..d78de8adaef 100644 --- a/ui/litellm-dashboard/src/components/SSOModals.test.tsx +++ b/ui/litellm-dashboard/src/components/SSOModals.test.tsx @@ -19,6 +19,7 @@ vi.mock("./shared/errorUtils", () => ({ import { toast } from "@/lib/toast"; import { getSSOSettings, updateSSOSettings } from "./networking"; +import { CLEARED_SSO_SETTINGS } from "./Settings/AdminSettings/SSOSettings/utils"; describe("SSOModals", () => { it("should render the SSOModals component", () => { @@ -515,27 +516,7 @@ describe("SSOModals", () => { // Verify the clear API was called with null values await waitFor(() => { - expect(updateSSOSettings).toHaveBeenCalledWith("test-token", { - google_client_id: null, - google_client_secret: null, - microsoft_client_id: null, - microsoft_client_secret: null, - microsoft_tenant: null, - generic_client_id: null, - generic_client_secret: null, - generic_authorization_endpoint: null, - generic_token_endpoint: null, - generic_userinfo_endpoint: null, - saml_idp_metadata_url: null, - saml_idp_metadata_xml: null, - saml_sp_entity_id: null, - saml_allow_unsolicited: null, - generic_scope: null, - proxy_base_url: null, - user_email: null, - sso_provider: null, - role_mappings: null, - }); + expect(updateSSOSettings).toHaveBeenCalledWith("test-token", CLEARED_SSO_SETTINGS); }); expect(toast.success).toHaveBeenCalledWith("SSO settings cleared successfully"); diff --git a/ui/litellm-dashboard/src/components/SSOModals.tsx b/ui/litellm-dashboard/src/components/SSOModals.tsx index 609c0efcd99..5c4825698a0 100644 --- a/ui/litellm-dashboard/src/components/SSOModals.tsx +++ b/ui/litellm-dashboard/src/components/SSOModals.tsx @@ -1,6 +1,7 @@ import React, { useEffect, useState } from "react"; import { FormProvider, useWatch, type UseFormReturn } from "react-hook-form"; import { getSSOSettings, updateSSOSettings } from "./networking"; +import { CLEARED_SSO_SETTINGS, diffSSOSettingsPayload } from "./Settings/AdminSettings/SSOSettings/utils"; import { toast } from "@/lib/toast"; import { parseErrorMessage } from "./shared/errorUtils"; import { Button } from "@/components/ui/button"; @@ -32,6 +33,60 @@ interface SSOModalsProps { ssoConfigured?: boolean; // Add optional prop to indicate if SSO is configured } +const buildLegacySSOPayload = (formValues: SSOSettingsFormValues): Record => { + const { + proxy_admin_teams, + admin_viewer_teams, + internal_user_teams, + internal_viewer_teams, + default_role, + group_claim, + use_role_mappings, + ...rest + } = formValues; + + const payload: Record = { + ...rest, + }; + + if (typeof payload.saml_allow_unsolicited === "boolean") { + payload.saml_allow_unsolicited = payload.saml_allow_unsolicited ? "true" : "false"; + } + + // Add role mappings if use_role_mappings is checked + if (use_role_mappings) { + // Helper function to split comma-separated string into array + const splitTeams = (teams: string | undefined): string[] => { + if (!teams || teams.trim() === "") return []; + return teams + .split(",") + .map((team) => team.trim()) + .filter((team) => team.length > 0); + }; + + // Map default role display values to backend values + const defaultRoleMapping: Record = { + internal_user_viewer: "internal_user_viewer", + internal_user: "internal_user", + proxy_admin_viewer: "proxy_admin_viewer", + proxy_admin: "proxy_admin", + }; + + payload.role_mappings = { + provider: "generic", + group_claim, + default_role: (default_role ? defaultRoleMapping[default_role] : undefined) || "internal_user", + roles: { + proxy_admin: splitTeams(proxy_admin_teams), + proxy_admin_viewer: splitTeams(admin_viewer_teams), + internal_user: splitTeams(internal_user_teams), + internal_user_viewer: splitTeams(internal_viewer_teams), + }, + }; + } + return payload; +}; + const detectSSOProvider = (values: Record): string | null => { if (values.google_client_id) return "google"; if (values.microsoft_client_id) return "microsoft"; @@ -134,56 +189,10 @@ const SSOModals: React.FC = ({ } try { - const { - proxy_admin_teams, - admin_viewer_teams, - internal_user_teams, - internal_viewer_teams, - default_role, - group_claim, - use_role_mappings, - ...rest - } = formValues; - - const payload: Record = { - ...rest, - }; - - if (typeof payload.saml_allow_unsolicited === "boolean") { - payload.saml_allow_unsolicited = payload.saml_allow_unsolicited ? "true" : "false"; - } - - // Add role mappings if use_role_mappings is checked - if (use_role_mappings) { - // Helper function to split comma-separated string into array - const splitTeams = (teams: string | undefined): string[] => { - if (!teams || teams.trim() === "") return []; - return teams - .split(",") - .map((team) => team.trim()) - .filter((team) => team.length > 0); - }; - - // Map default role display values to backend values - const defaultRoleMapping: Record = { - internal_user_viewer: "internal_user_viewer", - internal_user: "internal_user", - proxy_admin_viewer: "proxy_admin_viewer", - proxy_admin: "proxy_admin", - }; - - payload.role_mappings = { - provider: "generic", - group_claim, - default_role: (default_role ? defaultRoleMapping[default_role] : undefined) || "internal_user", - roles: { - proxy_admin: splitTeams(proxy_admin_teams), - proxy_admin_viewer: splitTeams(admin_viewer_teams), - internal_user: splitTeams(internal_user_teams), - internal_user_viewer: splitTeams(internal_viewer_teams), - }, - }; - } + const payload = diffSSOSettingsPayload( + buildLegacySSOPayload({ ...emptySSOSettingsFormValues, ...form.formState.defaultValues }), + buildLegacySSOPayload(formValues), + ); // Save SSO settings using the new API await updateSSOSettings(accessToken, payload); @@ -203,30 +212,7 @@ const SSOModals: React.FC = ({ } try { - // Clear all SSO settings - const clearSettings = { - google_client_id: null, - google_client_secret: null, - microsoft_client_id: null, - microsoft_client_secret: null, - microsoft_tenant: null, - generic_client_id: null, - generic_client_secret: null, - generic_authorization_endpoint: null, - generic_token_endpoint: null, - generic_userinfo_endpoint: null, - saml_idp_metadata_url: null, - saml_idp_metadata_xml: null, - saml_sp_entity_id: null, - saml_allow_unsolicited: null, - generic_scope: null, - proxy_base_url: null, - user_email: null, - sso_provider: null, - role_mappings: null, - }; - - await updateSSOSettings(accessToken, clearSettings); + await updateSSOSettings(accessToken, CLEARED_SSO_SETTINGS); // Clear the form form.reset(emptySSOSettingsFormValues); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.tsx index 7fa16515bf5..a16598653cf 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.tsx @@ -12,7 +12,7 @@ import BaseSSOSettingsForm, { import { Button } from "@/components/ui/button"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { useEditSSOSettings } from "@/app/(dashboard)/hooks/sso/useEditSSOSettings"; -import { processSSOSettingsPayload } from "../utils"; +import { diffSSOSettingsPayload, processSSOSettingsPayload } from "../utils"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; interface AddSSOSettingsModalProps { @@ -26,7 +26,10 @@ const AddSSOSettingsModal: React.FC = ({ isVisible, on const { mutateAsync, isPending } = useEditSSOSettings(); const handleFormSubmit = async (formValues: SSOSettingsFormValues) => { - const payload = processSSOSettingsPayload(formValues); + const payload = diffSSOSettingsPayload( + processSSOSettingsPayload(emptySSOSettingsFormValues), + processSSOSettingsPayload(formValues), + ); await mutateAsync(payload, { onSuccess: () => { diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.tsx index 053dac499bd..416ec4d2432 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.tsx @@ -4,7 +4,7 @@ import React from "react"; import DeleteResourceModal from "../../../../common_components/DeleteResourceModal"; import { toast } from "@/lib/toast"; import { parseErrorMessage } from "../../../../shared/errorUtils"; -import { detectSSOProvider } from "../utils"; +import { CLEARED_SSO_SETTINGS, detectSSOProvider } from "../utils"; interface DeleteSSOSettingsModalProps { isVisible: boolean; @@ -18,29 +18,7 @@ const DeleteSSOSettingsModal: React.FC = ({ isVisib // Handle clearing SSO settings const handleClearSSO = async () => { - const clearSettings = { - google_client_id: null, - google_client_secret: null, - microsoft_client_id: null, - microsoft_client_secret: null, - microsoft_tenant: null, - generic_client_id: null, - generic_client_secret: null, - generic_authorization_endpoint: null, - generic_token_endpoint: null, - generic_userinfo_endpoint: null, - saml_idp_metadata_url: null, - saml_idp_metadata_xml: null, - saml_sp_entity_id: null, - saml_allow_unsolicited: null, - proxy_base_url: null, - user_email: null, - sso_provider: null, - role_mappings: null, - team_mappings: null, - }; - - await editSSOSettings(clearSettings, { + await editSSOSettings(CLEARED_SSO_SETTINGS, { onSuccess: () => { toast.success("SSO settings cleared successfully"); onCancel(); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.integration.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.integration.test.tsx index 14bfe5d6ac9..cddcffc6926 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.integration.test.tsx @@ -49,7 +49,7 @@ describe("EditSSOSettingsModal (real form tree)", () => { expect(mutateAsync).not.toHaveBeenCalled(); }); - it("sends only the mounted fields when the form is valid", async () => { + it("sends only the changed field and leaves the masked secret out", async () => { const user = userEvent.setup({ pointerEventsCheck: 0 }); renderWithProviders(); @@ -58,12 +58,54 @@ describe("EditSSOSettingsModal (real form tree)", () => { await user.click(saveButton()); await waitFor(() => expect(mutateAsync).toHaveBeenCalledTimes(1)); - expect(mutateAsync.mock.calls[0][0]).toEqual({ - sso_provider: "google", - google_client_id: "rotated-client-id", - google_client_secret: "stored-client-secret", + expect(mutateAsync.mock.calls[0][0]).toEqual({ google_client_id: "rotated-client-id" }); + }); + + it("sends an empty patch when nothing changed", async () => { + const user = userEvent.setup({ pointerEventsCheck: 0 }); + renderWithProviders(); + + await screen.findByLabelText("Google Client ID"); + await user.click(saveButton()); + + await waitFor(() => expect(mutateAsync).toHaveBeenCalledTimes(1)); + expect(mutateAsync.mock.calls[0][0]).toEqual({}); + }); + + it("sends a new secret when the admin types one", async () => { + const user = userEvent.setup({ pointerEventsCheck: 0 }); + renderWithProviders(); + + fireEvent.change(await screen.findByLabelText("Google Client Secret"), { target: { value: "new-secret" } }); + await user.click(saveButton()); + + await waitFor(() => expect(mutateAsync).toHaveBeenCalledTimes(1)); + expect(mutateAsync.mock.calls[0][0]).toEqual({ google_client_secret: "new-secret" }); + }); + + it("sends role_mappings null when the admin unchecks role mappings", async () => { + seed({ + generic_client_id: "stored-client-id", + generic_client_secret: "kc-s****6e0d", + generic_authorization_endpoint: "https://idp.example.com/authorize", + generic_token_endpoint: "https://idp.example.com/token", + generic_userinfo_endpoint: "https://idp.example.com/userinfo", user_email: "admin@example.com", proxy_base_url: "https://gateway.example.com", + role_mappings: { + provider: "generic", + group_claim: "groups", + default_role: "internal_user", + roles: { proxy_admin: ["admins"], proxy_admin_viewer: [], internal_user: [], internal_user_viewer: [] }, + }, }); + const user = userEvent.setup({ pointerEventsCheck: 0 }); + renderWithProviders(); + + await user.click((await screen.findAllByLabelText("Use Role Mappings"))[0]); + await user.click(saveButton()); + + await waitFor(() => expect(mutateAsync).toHaveBeenCalledTimes(1)); + expect(mutateAsync.mock.calls[0][0]).toEqual({ role_mappings: null }); }); }); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx index ef78e3827d8..132800faf9f 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx @@ -293,6 +293,8 @@ describe("EditSSOSettingsModal", () => { useEditSSOSettings: { mutateAsync: mockMutateAsync, isPending: false }, }); + (processSSOSettingsPayload as any).mockReturnValueOnce({}).mockReturnValueOnce(processedPayload); + renderComponent(); fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx index 58c03eda979..706597e305b 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx @@ -11,7 +11,7 @@ import { Button } from "@/components/ui/button"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { toast } from "@/lib/toast"; import { parseErrorMessage } from "@/components/shared/errorUtils"; -import { detectSSOProvider, processSSOSettingsPayload } from "../utils"; +import { detectSSOProvider, diffSSOSettingsPayload, processSSOSettingsPayload } from "../utils"; import { useSSOSettings, type SSOSettingsValues } from "@/app/(dashboard)/hooks/sso/useSSOSettings"; import { useEditSSOSettings } from "@/app/(dashboard)/hooks/sso/useEditSSOSettings"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; @@ -86,7 +86,10 @@ const EditSSOSettingsModal: React.FC = ({ isVisible, const handleFormSubmit = async (formValues: SSOSettingsFormValues) => { try { - const payload = processSSOSettingsPayload(formValues); + const payload = diffSSOSettingsPayload( + processSSOSettingsPayload(seededValues), + processSSOSettingsPayload(formValues), + ); await mutateAsync(payload, { onSuccess: () => { diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test-d.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test-d.ts new file mode 100644 index 00000000000..c5f50e391ac --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test-d.ts @@ -0,0 +1,12 @@ +import { describe, expectTypeOf, test } from "vitest"; + +import type { components } from "@/lib/http/schema"; +import { CLEARED_SSO_SETTINGS } from "./utils"; + +describe("SSO settings utils types", () => { + test("clearing SSO settings sends null for exactly the fields the backend declares", () => { + expectTypeOf(CLEARED_SSO_SETTINGS).toEqualTypeOf<{ + readonly [K in keyof components["schemas"]["SSOConfig"]]-?: null; + }>(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test.ts index b280f5e92ec..b0d519afcbc 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test.ts +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test.ts @@ -1,4 +1,4 @@ -import { detectSSOProvider, processSSOSettingsPayload } from "./utils"; +import { detectSSOProvider, diffSSOSettingsPayload, processSSOSettingsPayload } from "./utils"; import { describe, it, expect } from "vitest"; import type { SSOSettingsValues } from "@/app/(dashboard)/hooks/sso/useSSOSettings"; @@ -452,3 +452,58 @@ describe("processSSOSettingsPayload with SAML", () => { ).toBe("false"); }); }); + +describe("diffSSOSettingsPayload", () => { + const stored = { + sso_provider: "generic", + generic_client_id: "client-id", + generic_client_secret: "kc-s****6e0d", + user_email: "admin@example.com", + role_mappings: { group_claim: "groups", roles: { proxy_admin: ["admins"] } }, + }; + + it("returns an empty patch when nothing changed", () => { + expect(diffSSOSettingsPayload(stored, { ...stored })).toEqual({}); + }); + + it("leaves an untouched masked secret out of the patch", () => { + expect(diffSSOSettingsPayload(stored, { ...stored, user_email: "new@example.com" })).toEqual({ + user_email: "new@example.com", + }); + }); + + it("sends null for a field that was emptied", () => { + expect(diffSSOSettingsPayload(stored, { ...stored, generic_client_secret: "" })).toEqual({ + generic_client_secret: null, + }); + }); + + it("sends null for fields dropped from the payload, like another provider's or disabled mappings", () => { + const { generic_client_id, generic_client_secret, role_mappings, ...rest } = stored; + expect(diffSSOSettingsPayload(stored, { ...rest, sso_provider: "google", google_client_id: "g-id" })).toEqual({ + sso_provider: "google", + google_client_id: "g-id", + generic_client_id: null, + generic_client_secret: null, + role_mappings: null, + }); + }); + + it("sends the whole mapping object when one part of it changed", () => { + const role_mappings = { group_claim: "roles", roles: { proxy_admin: ["admins"] } }; + expect(diffSSOSettingsPayload(stored, { ...stored, role_mappings })).toEqual({ role_mappings }); + }); + + it("treats empty values on both sides as unchanged", () => { + expect(diffSSOSettingsPayload({ generic_scope: undefined, proxy_base_url: "" }, { proxy_base_url: null })).toEqual( + {}, + ); + }); + + it("treats an unchecked saml_allow_unsolicited as unset, but sends false when it was true", () => { + expect(diffSSOSettingsPayload({ saml_allow_unsolicited: "false" }, {})).toEqual({}); + expect(diffSSOSettingsPayload({ saml_allow_unsolicited: "true" }, { saml_allow_unsolicited: "false" })).toEqual({ + saml_allow_unsolicited: "false", + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.ts index 768fc7e8e0e..93af987adde 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.ts +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.ts @@ -71,6 +71,49 @@ export const processSSOSettingsPayload = (formValues: Record): Reco return payload; }; +const isEmptySSOValue = (value: unknown): boolean => value === undefined || value === null || value === ""; + +// An unchecked saml_allow_unsolicited is the same as never having set it +const comparableSSOValue = (key: string, value: unknown): unknown => + isEmptySSOValue(value) || (key === "saml_allow_unsolicited" && value === "false") ? undefined : value; + +// /update/sso_settings is a PATCH: send only the fields that changed, with null for the ones that were emptied +export const diffSSOSettingsPayload = ( + before: Record, + after: Record, +): Record => + Object.fromEntries( + [...new Set([...Object.keys(before), ...Object.keys(after)])] + .filter( + (key) => + JSON.stringify(comparableSSOValue(key, before[key])) !== JSON.stringify(comparableSSOValue(key, after[key])), + ) + .map((key) => [key, isEmptySSOValue(after[key]) ? null : after[key]]), + ); + +export const CLEARED_SSO_SETTINGS = { + google_client_id: null, + google_client_secret: null, + microsoft_client_id: null, + microsoft_client_secret: null, + microsoft_tenant: null, + generic_client_id: null, + generic_client_secret: null, + generic_authorization_endpoint: null, + generic_token_endpoint: null, + generic_userinfo_endpoint: null, + generic_scope: null, + saml_idp_metadata_url: null, + saml_idp_metadata_xml: null, + saml_sp_entity_id: null, + saml_allow_unsolicited: null, + proxy_base_url: null, + user_email: null, + ui_access_mode: null, + role_mappings: null, + team_mappings: null, +} as const; + // Determine the SSO provider based on the configuration export const detectSSOProvider = (values: SSOSettingsValues): string | null => { if (values.google_client_id) return "google"; From 229f5a9745bc3fb7eaf275cd6f040379380f22f3 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 2 Nov 2026 09:23:17 -0800 Subject: [PATCH 05/13] test(integration): fix four flaky integration tests (#49554) Scope the Prisma migration lock test's advisory lock check to its own database, so a lock another xdist worker holds in its database no longer trips it Keep in-process MCP peers' idle connections open for 30s instead of 1s, so the gateway's pooled HTTP client doesn't reuse a connection the peer is closing and hang the tool call Delete database guardrails in the logging-only scope tests through DELETE /guardrails/{id} on the shared gateway instead of raw SQL, so a default-on guardrail stops applying to the next test on the same worker right away Bound the wedged coordination Redis /customer/update check by the shortest Redis wait the proxy uses, like the burst check, so a slow runner can't fail it while a handler that waits on Redis still does --- tests/integration/_support/asgi.py | 2 +- .../database/test_prisma_migration_lock.py | 3 +- ..._user_updates_wedged_coordination_redis.py | 6 ++-- .../_logging_only_scope_support.py | 13 +++++--- .../test_logging_only_scope_config.py | 30 +++++++++---------- .../test_logging_only_scope_runtime.py | 2 +- 6 files changed, 32 insertions(+), 24 deletions(-) diff --git a/tests/integration/_support/asgi.py b/tests/integration/_support/asgi.py index e84dc9439d3..30eb9f21760 100644 --- a/tests/integration/_support/asgi.py +++ b/tests/integration/_support/asgi.py @@ -15,7 +15,7 @@ from starlette.types import ASGIApp @contextmanager def asgi_server( - app: ASGIApp, *, before_stop: Callable[[], None] | None = None, port: int = 0, keepalive_seconds: int = 1 + app: ASGIApp, *, before_stop: Callable[[], None] | None = None, port: int = 0, keepalive_seconds: int = 30 ) -> Iterator[str]: with socket.socket() as listener: listener.bind(("127.0.0.1", port)) diff --git a/tests/integration/database/test_prisma_migration_lock.py b/tests/integration/database/test_prisma_migration_lock.py index 7482e96d3d3..f4a2a022668 100644 --- a/tests/integration/database/test_prisma_migration_lock.py +++ b/tests/integration/database/test_prisma_migration_lock.py @@ -19,7 +19,8 @@ def _can_take_prismas_session_lock(connection: psycopg.Connection[tuple[object, def _granted_advisory_keys(connection: psycopg.Connection[tuple[object, ...]]) -> frozenset[int]: rows: Final = connection.execute( - "SELECT (classid::bigint << 32) | objid::bigint FROM pg_locks WHERE locktype = 'advisory' AND granted" + "SELECT (classid::bigint << 32) | objid::bigint FROM pg_locks WHERE locktype = 'advisory' AND granted " + "AND database = (SELECT oid FROM pg_database WHERE datname = current_database())" ).fetchall() return frozenset(int(row[0]) for row in rows) diff --git a/tests/integration/management/test_user_updates_wedged_coordination_redis.py b/tests/integration/management/test_user_updates_wedged_coordination_redis.py index 2b05d35d553..bb5766e5816 100644 --- a/tests/integration/management/test_user_updates_wedged_coordination_redis.py +++ b/tests/integration/management/test_user_updates_wedged_coordination_redis.py @@ -193,8 +193,10 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( timings["s6_customer_update"] = _timed_post( candidate, "/customer/update", {"user_id": customer_id, "max_budget": 6.0} ) - assert timings["s6_customer_update"] < _HANDLER_BUDGET_SECONDS, ( - f"/customer/update took {timings['s6_customer_update']:.3f}s with a wedged coordination Redis" + assert timings["s6_customer_update"] < _SHORTEST_REDIS_WAIT_SECONDS, ( + f"/customer/update took {timings['s6_customer_update']:.3f}s, reaching the " + f"{_SHORTEST_REDIS_WAIT_SECONDS}s shortest wait the proxy puts on Redis, so it waited on the " + "wedged coordination Redis" ) timings["s7_reset_spend"] = _timed_post(candidate, f"/key/{key}/reset_spend", {"reset_to": 0}) assert timings["s7_reset_spend"] < _HANDLER_BUDGET_SECONDS, ( diff --git a/tests/integration/observability/_logging_only_scope_support.py b/tests/integration/observability/_logging_only_scope_support.py index 3d29f7ab8c5..3b14c502623 100644 --- a/tests/integration/observability/_logging_only_scope_support.py +++ b/tests/integration/observability/_logging_only_scope_support.py @@ -22,7 +22,7 @@ import yaml from anthropic import Anthropic, AsyncAnthropic from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa -from integration._support.client import Gateway, Scenario, eventually, object_value +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value from integration._support.database import read_rows, write_rows from integration._support.upstream import delete_scenario, register_scenario from integration._support.wire import Reply, Request, Wire @@ -798,6 +798,7 @@ def _insert_database_guardrail( @contextmanager def _database_guardrail( + gateway: Gateway, identity: str, policy_url: str, scope: str | None, @@ -809,11 +810,15 @@ def _database_guardrail( try: yield finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) -def _delete_database_guardrail(identity: str) -> None: - write_rows('DELETE FROM "LiteLLM_GuardrailsTable" WHERE guardrail_name=%s', (identity,)) +def _delete_database_guardrail(gateway: Gateway, identity: str) -> None: + rows: Final = read_rows('SELECT guardrail_id FROM "LiteLLM_GuardrailsTable" WHERE guardrail_name=%s', (identity,)) + responses: Final = tuple( + gateway.request("DELETE", f"/guardrails/{string_value(row['guardrail_id'])}") for row in rows + ) + assert all(response.status_code == 200 for response in responses), [response.text for response in responses] def _post_guardrail_body( diff --git a/tests/integration/observability/test_logging_only_scope_config.py b/tests/integration/observability/test_logging_only_scope_config.py index 0841c91f8ca..ea524c15c92 100644 --- a/tests/integration/observability/test_logging_only_scope_config.py +++ b/tests/integration/observability/test_logging_only_scope_config.py @@ -115,7 +115,7 @@ def test_G7_database_invalid_scope_reads_preserve_pre_call_blocking(gateway: Gat assert blocked_response.status_code == 400, blocked_response.text assert _upstream_requests(gateway.upstream_url) == () finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) finally: delete_scenario(upstream_handle) @@ -678,7 +678,7 @@ def test_database_guardrail_load_keeps_pre_call_blocking( for entry in entries ) == ((identity, "pre_call", "guardrail_intervened"),), entries finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) finally: delete_scenario(upstream_handle) @@ -799,7 +799,7 @@ def test_G6_database_guardrail_polling_normalizes_invalid_scope_without_reinitia ) == ((identity, "pre_call", "guardrail_intervened"),), (call_id, entries) assert _upstream_requests(gateway.upstream_url) == () finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) finally: delete_scenario(upstream_handle) @@ -952,7 +952,7 @@ def test_management_post_validates_logging_only_scope( assert isinstance(after_list, list), after_list assert identity not in {object_value(item)["guardrail_name"] for item in after_list}, after_list finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) def test_H5_management_put_rejection_preserves_database_and_runtime(gateway: Gateway, tmp_path: Path) -> None: @@ -1074,7 +1074,7 @@ def test_H5_management_put_rejection_preserves_database_and_runtime(gateway: Gat for entry in entries ) == ((identity, "pre_call", "guardrail_intervened"),), (call_id, entries) finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) delete_scenario(upstream_handle) @@ -1153,7 +1153,7 @@ def test_H6_management_patch_clears_scope_when_switching_to_blocking_only_mode( for entry in entries ) == ((identity, ("pre_call",), "guardrail_intervened"),), entries finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) delete_scenario(upstream_handle) @@ -1240,7 +1240,7 @@ def test_H7_management_patch_rejection_preserves_pre_call_guardrail(gateway: Gat for entry in entries ) == ((identity, "pre_call", "guardrail_intervened"),), (call_id, entries) finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) delete_scenario(upstream_handle) @@ -1331,7 +1331,7 @@ def test_management_patch_default_on_heals_stored_scope( for entry in entries ) == ((identity, "pre_call", "guardrail_intervened"),), entries finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) finally: delete_scenario(upstream_handle) @@ -1429,7 +1429,7 @@ def test_H10_management_patch_null_scope_restores_both_directions(gateway: Gatew (identity, "logging_only", "guardrail_intervened"), ), entries finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) delete_scenario(upstream_handle) @@ -1517,7 +1517,7 @@ def test_H11_management_patch_same_scope_is_idempotent_and_output_only(gateway: (entry["guardrail_name"], entry["guardrail_mode"], entry["guardrail_status"]) for entry in entries ) == tuple((identity, "logging_only", "guardrail_intervened") for _ in expected_directions), entries finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) delete_scenario(upstream_handle) @@ -1550,7 +1550,7 @@ def test_H12_management_reads_expose_typed_logging_only_scope(gateway: Gateway, assert object_value(info["litellm_params"])["logging_only_scope"] == "input", info assert object_value(matches[0]["litellm_params"])["logging_only_scope"] == "input", matches finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) def test_H13_unauthenticated_management_and_chat_requests_do_not_scan(gateway: Gateway, tmp_path: Path) -> None: @@ -1629,7 +1629,7 @@ def test_K1_post_rejects_continue_flag_outside_logging_only(gateway: Gateway, tm assert after_list == before_list, (before_list, after_list) assert identity not in {object_value(item)["guardrail_name"] for item in after_list}, after_list finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) def test_K3_patch_flag_rejection_preserves_raw_stored_row(gateway: Gateway, tmp_path: Path) -> None: @@ -1675,7 +1675,7 @@ def test_K3_patch_flag_rejection_preserves_raw_stored_row(gateway: Gateway, tmp_ } == {key: after_params[key] for key in after_params if key != "api_key"}, (before_params, after_params) assert str(after_params["api_key"]).startswith("litellm_enc::"), after_params finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) def test_K4_patch_dropping_logging_only_mode_clears_stored_flag(gateway: Gateway, tmp_path: Path) -> None: @@ -1713,7 +1713,7 @@ def test_K4_patch_dropping_logging_only_mode_clears_stored_flag(gateway: Gateway assert params.get("logging_only_continue_on_input_failure") is None, persisted assert params["mode"] == ["pre_call"], persisted finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) def test_K5_yaml_flag_on_blocking_guardrail_boots_and_keeps_blocking(gateway: Gateway, tmp_path: Path) -> None: @@ -1759,5 +1759,5 @@ def test_K5_yaml_flag_on_blocking_guardrail_boots_and_keeps_blocking(gateway: Ga log_text: Final = owned.log.read_text() assert "Ignoring logging_only_continue_on_input_failure" in log_text, log_text finally: - _delete_database_guardrail(identity) + _delete_database_guardrail(gateway, identity) delete_scenario(upstream_handle) diff --git a/tests/integration/observability/test_logging_only_scope_runtime.py b/tests/integration/observability/test_logging_only_scope_runtime.py index 8da26f20eae..271d8d0e24d 100644 --- a/tests/integration/observability/test_logging_only_scope_runtime.py +++ b/tests/integration/observability/test_logging_only_scope_runtime.py @@ -1098,7 +1098,7 @@ def test_X2_scope_patch_toggles_during_concurrent_requests_keep_one_scan_per_res with wire_server(policy) as edge: config: Final = _empty_proxy_configuration(tmp_path, identity) with ( - _database_guardrail(identity, edge.url, "output", mode="logging_only"), + _database_guardrail(gateway, identity, edge.url, "output", mode="logging_only"), owned_proxy(gateway, tmp_path, {}, config=config, workers=1) as candidate, ): guardrails: Final = candidate.get("/v2/guardrails/list")["guardrails"] From 40e4e512d05abe72a9810ba4c2f2cfeb85a19db6 Mon Sep 17 00:00:00 2001 From: mika-lt <349045450+mika-lt@users.noreply.github.com> Date: Mon, 2 Nov 2026 09:31:05 -0800 Subject: [PATCH 06/13] fix(anthropic): stream azure_ai natively on an Azure OpenAI host on /v1/messages (#49216) * fix(anthropic): stream azure_ai natively on an Azure OpenAI host on /v1/messages The /v1/messages bridges got the requested model (azure_ai/gpt-6-luna) with the resolved provider (azure), so the Responses bridge faked the stream and the chat bridge called a deployment named azure_ai/gpt-6-luna (404). Pass the resolved provider prefix when the requested prefix is a different known provider. * test(anthropic): cover the bridged model with respx and a regional bedrock_mantle id * fix(anthropic): carry azure_ai env credentials into bridged /v1/messages calls When an azure_ai deployment's api_base and key come only from AZURE_AI_API_BASE and AZURE_AI_API_KEY, the provider the async wrapper derived from the model alone replaced the model's own resolution, so the bridged call lost both and failed. The model's own resolution now wins when it names the same provider, an explicit provider that disagrees still wins, and the bridge carries the resolved key and base. * fix(anthropic): keep the caller's key, proxy routing and MCP credentials on bridged /v1/messages calls * fix(anthropic): keep an azure deployment's Entra credential over the env key on bridged /v1/messages calls * fix(anthropic): skip the sign-in lookup for chatgpt and github_copilot model prefixes on bridged /v1/messages calls * fix(anthropic): keep the Azure credential of an azure_ai deployment that sets its own api_base on bridged /v1/messages calls * test(anthropic): annotate the parameters of the bridged /v1/messages tests --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../llms/anthropic/messages/native/handler.py | 100 +++++-- .../native/test_anthropic_messages_handler.py | 257 ++++++++++++++++++ .../messages/native/test_mcp_handler.py | 28 ++ 3 files changed, 369 insertions(+), 16 deletions(-) diff --git a/litellm/llms/anthropic/messages/native/handler.py b/litellm/llms/anthropic/messages/native/handler.py index 1bf00cbdb2f..922de386fe4 100644 --- a/litellm/llms/anthropic/messages/native/handler.py +++ b/litellm/llms/anthropic/messages/native/handler.py @@ -8,6 +8,7 @@ import asyncio import contextvars from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence +from dataclasses import dataclass from functools import partial from typing import Final, cast @@ -15,6 +16,7 @@ from pydantic import TypeAdapter import litellm from litellm.anthropic_beta_headers_manager import aload_beta_headers_config +from litellm.constants import PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.exception_mapping_utils import exception_type @@ -44,7 +46,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import CallTypes, LlmProviders +from litellm.types.utils import CallTypes, LlmProviders, LlmProvidersSet from litellm.utils import ProviderConfigManager, client from ..adapters.handler import LiteLLMMessagesToCompletionTransformationHandler, request_safeguards_evaluator @@ -94,6 +96,68 @@ def _responses_mode_is_lost_by_prefix_strip( ) and not _bridges_to_responses_api(resolved_model, custom_llm_provider, deployment_model_info) +@dataclass(frozen=True, slots=True) +class _BridgeTarget: + model: str + api_key: str | None + api_base: str | None + + +def _resolve_requested_provider( + model: str, custom_llm_provider: str | None, litellm_params: GenericLiteLLMParams +) -> tuple[str, str, str | None, str | None]: + resolved: Final = litellm.get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider, + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + litellm_params=litellm_params, + ) + requested_provider: Final = model.partition("/")[0] + if ( + custom_llm_provider in (None, requested_provider) + or requested_provider not in LlmProvidersSet + or requested_provider in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO + ): + return resolved + own: Final = _resolve_by_model_prefix(model, litellm_params) + return own if own is not None and own[1] == custom_llm_provider else resolved + + +def _resolve_by_model_prefix( + model: str, litellm_params: GenericLiteLLMParams +) -> tuple[str, str, str | None, str | None] | None: + try: + return litellm.get_llm_provider( + model=model, + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + litellm_params=litellm_params.model_copy(update={"custom_llm_provider": None}), + ) + except litellm.BadRequestError: + return None + + +def _sets_own_entra_credential(litellm_params: GenericLiteLLMParams) -> bool: + return bool( + litellm_params.azure_ad_token + or (litellm_params.model_extra or {}).get("azure_ad_token_provider") is not None + or (litellm_params.tenant_id and litellm_params.client_id and litellm_params.client_secret) + or (litellm_params.azure_username and litellm_params.azure_password and litellm_params.client_id) + ) + + +def _bridge_target(requested: _BridgeTarget, resolved: _BridgeTarget, custom_llm_provider: str) -> _BridgeTarget: + requested_provider: Final = requested.model.partition("/")[0] + if requested_provider == custom_llm_provider or requested_provider not in LlmProvidersSet: + return requested + return _BridgeTarget( + model=f"{custom_llm_provider}/{resolved.model}", + api_key=requested.api_key if requested.api_base else requested.api_key or resolved.api_key, + api_base=resolved.api_base, + ) + + def _should_route_to_responses_api( custom_llm_provider: str | None, requested_model: str | None = None, @@ -624,13 +688,7 @@ def anthropic_messages_handler( custom_llm_provider, dynamic_api_key, dynamic_api_base, - ) = litellm.get_llm_provider( - model=model, - custom_llm_provider=custom_llm_provider, - api_base=litellm_params.api_base, - api_key=litellm_params.api_key, - litellm_params=litellm_params, - ) + ) = _resolve_requested_provider(model, custom_llm_provider, litellm_params) # Store agentic loop params in logging object for agentic hooks # This provides original request context needed for follow-up calls @@ -656,6 +714,16 @@ def anthropic_messages_handler( mock_response=litellm_params.mock_response, ) + bridge: Final = _bridge_target( + _BridgeTarget(model=original_model, api_key=api_key, api_base=api_base), + _BridgeTarget( + model=model, + api_key=None if _sets_own_entra_credential(litellm_params) else dynamic_api_key, + api_base=dynamic_api_base, + ), + custom_llm_provider, + ) + # Expand litellm_proxy MCP references through the MCP gateway before dispatch, so every # downstream path (native passthrough and both bridges) gets real tools rather than a # reference the provider cannot resolve. Popped from kwargs so it never reaches the provider. @@ -682,8 +750,8 @@ def anthropic_messages_handler( top_k=top_k, top_p=top_p, container=container, - api_key=api_key, - api_base=api_base, + api_key=bridge.api_key, + api_base=bridge.api_base, client=client, custom_llm_provider=custom_llm_provider, **kwargs, @@ -700,7 +768,7 @@ def anthropic_messages_handler( return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler( max_tokens=max_tokens, messages=messages, - model=original_model, + model=bridge.model, metadata=metadata, stop_sequences=stop_sequences, stream=stream, @@ -716,8 +784,8 @@ def anthropic_messages_handler( top_k=top_k, top_p=top_p, _is_async=is_async, - api_key=api_key, - api_base=api_base, + api_key=bridge.api_key, + api_base=bridge.api_base, client=client, custom_llm_provider=custom_llm_provider, **kwargs, @@ -730,7 +798,7 @@ def anthropic_messages_handler( return LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( max_tokens=max_tokens, messages=messages, - model=original_model, + model=bridge.model, metadata=metadata, stop_sequences=stop_sequences, stream=stream, @@ -744,8 +812,8 @@ def anthropic_messages_handler( top_k=top_k, top_p=top_p, _is_async=is_async, - api_key=api_key, - api_base=api_base, + api_key=bridge.api_key, + api_base=bridge.api_base, client=client, custom_llm_provider=custom_llm_provider, **kwargs, diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py index af3108e5072..8a924dd32d6 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py @@ -3,10 +3,13 @@ import copy import json import os import uuid +from collections.abc import Mapping from dataclasses import dataclass +from pathlib import Path from typing import Any, Dict, Final, List import httpx +import openai import responses import respx import pytest @@ -1136,6 +1139,260 @@ async def test_messages_strips_provider_prefix_exactly_once(requested_model, exp assert captured["url"] == expected_url +_GET_WEATHER_TOOL: Final = { + "name": "get_weather", + "description": "Get the current weather for a city", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tools, expected_path", + [ + ([_GET_WEATHER_TOOL], "/openai/v1/responses"), + (None, "/openai/deployments/gpt-6-luna/chat/completions"), + ], + ids=["function_tool_bridged_to_responses", "no_tools_on_chat"], +) +@pytest.mark.parametrize("credentials_from_env", [False, True], ids=["credentials_in_params", "credentials_from_env"]) +async def test_azure_ai_on_azure_openai_host_streams_natively_with_the_bare_deployment( + tools: list[dict[str, object]] | None, + expected_path: str, + credentials_from_env: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.delenv("AZURE_API_BASE", raising=False) + monkeypatch.delenv("AZURE_API_KEY", raising=False) + monkeypatch.setenv("AZURE_AI_API_BASE", "https://qa-resource.cognitiveservices.azure.com") + monkeypatch.setenv("AZURE_AI_API_KEY", "env-api-key") + credentials: Final = ( + {} + if credentials_from_env + else {"api_base": "https://qa-resource.cognitiveservices.azure.com", "api_key": "params-api-key"} + ) + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.route(host="qa-resource.cognitiveservices.azure.com").mock( + side_effect=httpx.ConnectError("cut at the wire") + ) + with pytest.raises(openai.APIError): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "What is the weather in Berlin?"}], + model="azure_ai/gpt-6-luna", + stream=True, + tools=tools, + **credentials, + ) + + bodies: Final = [json.loads(call.request.content) for call in upstream.calls] + assert {call.request.url.path for call in upstream.calls} == {expected_path} + assert {call.request.headers["api-key"] for call in upstream.calls} == {credentials.get("api_key", "env-api-key")} + assert {body["model"] for body in bodies} == {"gpt-6-luna"} + assert all(body["stream"] is True for body in bodies) + + +@pytest.mark.asyncio +async def test_bridged_messages_keep_the_region_of_a_regional_bedrock_mantle_model( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire")) + with pytest.raises(litellm.exceptions.ServiceUnavailableError): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model="bedrock_mantle/eu-west-1/openai.gpt-oss-120b", + api_key="test-api-key", + ) + + assert {call.request.url.host for call in upstream.calls} == {"bedrock-mantle.eu-west-1.api.aws"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["anthropic/claude-sonnet-5-5", "ragflow/gpt-4o"]) +async def test_bridged_messages_keep_the_callers_explicit_provider_over_the_model_prefix( + model: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire")) + with pytest.raises(openai.APIError): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model=model, + custom_llm_provider="openai", + api_base="https://gateway.example.com/v1", + api_key="test-api-key", + ) + + assert {call.request.url.host for call in upstream.calls} == {"gateway.example.com"} + assert {json.loads(call.request.content)["model"] for call in upstream.calls} == {model} + + +@pytest.mark.parametrize("model", ["chatgpt/gpt-5", "github_copilot/gpt-5"]) +def test_bridged_messages_reach_the_callers_provider_without_signing_in_to_the_prefix_provider( + model: str, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("OPENAI_API_KEY", "server-openai-key") + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire")) + with pytest.raises(openai.APIError): + litellm.anthropic.messages.create( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model=model, + custom_llm_provider="openai", + ) + + assert {call.request.url.host for call in upstream.calls} == {"api.openai.com"} + assert {call.request.headers["authorization"] for call in upstream.calls} == {"Bearer server-openai-key"} + + +@pytest.mark.asyncio +async def test_bridged_messages_send_the_deployments_key_for_azure_models_served_as_openai( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("OPENAI_API_KEY", "server-openai-key") + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire")) + with pytest.raises(openai.APIError): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model="azure/mistral-large-latest", + api_base="https://example-resource.services.ai.azure.com/models", + api_key="deployment-key", + ) + + assert {call.request.headers["authorization"] for call in upstream.calls} == {"Bearer deployment-key"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "entra_credential", + [ + {"azure_ad_token": "entra-token"}, + {"azure_ad_token_provider": lambda: "entra-token"}, + {"tenant_id": "tenant", "client_id": "client", "client_secret": "secret"}, + {"azure_username": "user", "azure_password": "password", "client_id": "client"}, + ], +) +@pytest.mark.parametrize("api_base_from_env", [False, True], ids=["api_base_in_params", "api_base_from_env"]) +async def test_bridged_messages_keep_the_deployments_entra_credential_over_the_env_key( + monkeypatch: pytest.MonkeyPatch, entra_credential: Mapping[str, object], api_base_from_env: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("AZURE_AI_API_BASE", "https://example-resource.cognitiveservices.azure.com") + monkeypatch.setenv("AZURE_AI_API_KEY", "azure-env-key") + litellm.in_memory_llm_clients_cache.flush_cache() + api_base: Final = {} if api_base_from_env else {"api_base": "https://example-resource.cognitiveservices.azure.com"} + + with ( + patch("litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id", return_value=lambda: "entra-token"), # test-quality-ok: stands in for the Entra token fetch so the test can see which credential reaches Azure + patch( + "litellm.llms.azure.common_utils.get_azure_ad_token_from_username_password", + return_value=lambda: "entra-token", + ), # test-quality-ok: stands in for the Entra token fetch so the test can see which credential reaches Azure + respx.mock(assert_all_called=True) as respx_mock, + ): + upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire")) + with pytest.raises(openai.APIError): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model="azure_ai/gpt-6-luna", + tools=[{"name": "get_weather", "input_schema": {"type": "object", "properties": {}}}], + **api_base, + **entra_credential, + ) + + assert {call.request.headers.get("authorization") for call in upstream.calls} == {"Bearer entra-token"} + assert {call.request.headers.get("api-key") for call in upstream.calls} == {None} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "azure_env_credential, expected_header, expected_value", + [ + ({"AZURE_API_KEY": "azure-env-key"}, "api-key", "azure-env-key"), + ( + {"AZURE_TENANT_ID": "tenant", "AZURE_CLIENT_ID": "client", "AZURE_CLIENT_SECRET": "secret"}, + "authorization", + "Bearer entra-token", + ), + ], + ids=["azure_env_key", "azure_env_entra"], +) +async def test_bridged_messages_keep_the_azure_env_credential_when_the_deployment_sets_its_own_api_base( + azure_env_credential: Mapping[str, str], expected_header: str, expected_value: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + for name in ("AZURE_API_KEY", "AZURE_OPENAI_API_KEY", "AZURE_AD_TOKEN", "AZURE_TENANT_ID", "AZURE_CLIENT_ID"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("AZURE_AI_API_KEY", "key-for-another-resource") + for name, value in azure_env_credential.items(): + monkeypatch.setenv(name, value) + litellm.in_memory_llm_clients_cache.flush_cache() + + with ( + patch("litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id", return_value=lambda: "entra-token"), # test-quality-ok: stands in for the Entra token fetch so the test can see which credential reaches Azure + respx.mock(assert_all_called=True) as respx_mock, + ): + upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire")) + with pytest.raises(openai.APIError): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model="azure_ai/gpt-6-luna", + api_base="https://example-resource.cognitiveservices.azure.com", + tools=[{"name": "get_weather", "input_schema": {"type": "object", "properties": {}}}], + ) + + assert {call.request.headers.get(expected_header) for call in upstream.calls} == {expected_value} + + +@pytest.mark.asyncio +async def test_bridged_messages_still_route_through_the_litellm_proxy_when_asked( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("LITELLM_PROXY_API_BASE", "https://proxy.example.com") + monkeypatch.setenv("LITELLM_PROXY_API_KEY", "proxy-key") + monkeypatch.setenv("AZURE_AI_API_BASE", "https://example-resource.cognitiveservices.azure.com") + monkeypatch.setenv("AZURE_AI_API_KEY", "azure-env-key") + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire")) + with pytest.raises(openai.APIError): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model="azure_ai/gpt-6-luna", + use_litellm_proxy=True, + ) + + assert {call.request.url.host for call in upstream.calls} == {"proxy.example.com"} + + @pytest.mark.asyncio @pytest.mark.parametrize( "model_info, expected_url", diff --git a/tests/unit/llms/anthropic/messages/native/test_mcp_handler.py b/tests/unit/llms/anthropic/messages/native/test_mcp_handler.py index 73c2edbce19..c0e0ea9f803 100644 --- a/tests/unit/llms/anthropic/messages/native/test_mcp_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_mcp_handler.py @@ -2,8 +2,12 @@ from importlib import import_module from typing import Final from unittest.mock import AsyncMock, patch +import httpx +import openai import pytest +import respx +import litellm from litellm.llms.anthropic.messages.adapters.transformation import ( LiteLLMAnthropicMessagesAdapter, ) @@ -56,6 +60,30 @@ def test_anthropic_messages_handler_routes_litellm_proxy_mcp_to_the_gateway(): assert result is not None +@pytest.mark.asyncio +async def test_anthropic_messages_with_mcp_reaches_an_env_only_deployment_with_its_env_credentials( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("AZURE_AI_API_BASE", "https://example-resource.cognitiveservices.azure.com") + monkeypatch.setenv("AZURE_AI_API_KEY", "azure-env-key") + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.route().mock(side_effect=httpx.ConnectError("cut at the wire")) + with pytest.raises(openai.APIError): + await litellm.anthropic.messages.acreate( + max_tokens=100, + messages=[{"role": "user", "content": "ping"}], + model="azure_ai/gpt-6-luna", + tools=[{**MCP_REFERENCE, "server_url": "litellm_proxy"}], + ) + + assert {call.request.url.host for call in upstream.calls} == {"example-resource.cognitiveservices.azure.com"} + assert {call.request.headers.get("api-key") for call in upstream.calls} == {"azure-env-key"} + assert not any(b"litellm_proxy" in call.request.content for call in upstream.calls) + + def test_anthropic_messages_handler_skips_the_gateway_on_recursion(): """The gateway's own follow-up call must not re-enter the gateway.""" with patch( From 33457d3ac925c5e97c7a0c4765c2ca79940a899f Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Mon, 2 Nov 2026 09:48:59 -0800 Subject: [PATCH 07/13] fix(proxy): check the model in every OpenAI websocket passthrough frame (#48979) * test(proxy): cover per-frame model checks and session pricing on the OpenAI websocket passthrough * fix(proxy): check the model in every OpenAI websocket passthrough frame and price the session * fix(proxy): bill each OpenAI websocket turn on its own row and run full auth per frame Each finished turn (completed, incomplete, or failed) now logs its own spend row under the model and access groups that allowed it, so group budgets see spend as turns end instead of at close. Every model a frame names goes through the key's full auth, and a stored prompt reference with no model is refused for a key held to some models. A frame still in flight when the client leaves is recorded before the relay gives up on the send * fix(proxy): keep the stored prompt refusal's first sentence inside the websocket close reason The close reason carries the refusal's first sentence, and the stored prompt sentence ran past the 123 bytes a close frame allows, so it was cut mid-word. The sentence is shorter now and the rest of the explanation stays in the error event. * fix(proxy): bill pipelined websocket turns to their own frame and close the review gaps * fix(proxy): log no extra websocket spend row for frames after the last turn * fix(proxy): price websocket Responses turns by served tier and hosted tools, and bill turns a client walked out of A websocket Responses turn is now priced through the same Responses cost path the HTTP passthrough uses, so a priority-tier answer and a hosted web search cost what they cost over HTTP. When the client leaves mid-answer, the gateway keeps reading the provider's socket until that turn ends, up to 10 minutes, so the turn the provider still bills reaches spend * fix(proxy): bill websocket turns by lane, skip turns OpenAI refuses, and stop frame checks counting the session's own reservation An error event refuses a queued response.create, so the next turn is no longer charged to the refused frame's model. Turns on separate stream_id lanes are charged to the frame that started each one. The frame model gate settles the connect reservation before it re-runs auth, so a key with less budget left than one turn's max cost is no longer refused on its first frame. * fix(proxy): bill each websocket lane cut off at session end on its own row When a client leaves while turns on several stream_id lanes are still running, each unfinished turn now gets its own spend row under the groups and model its own frame asked for, instead of one merged row. Also pins the PR's new single-assigned test variables as Final and swaps datetime.now() in the logging-handler and success-handler fixtures for fixed timestamps. * test(proxy): mark the remaining new websocket test variables Final Builds the logging mocks with model_call_details in the constructor instead of setting it afterwards. * fix(proxy): check every model the realtime URL names before the OpenAI websocket relay connects * fix(proxy): keep reading running answers after a refused websocket frame so they are billed * fix(proxy): bill realtime transcription sessions on the OpenAI websocket passthrough A transcription session names its model only in session.updated under audio.input.transcription, and reports usage on each completed transcription instead of response.done, so its row logged unknown at $0 and a budgeted key never spent. * test(proxy): mark the turn-taking client's delivered list Final * fix(proxy): hold a websocket session's parallel slot until the socket closes Each finished turn logged its spend row under the session's own call id, so the first turn's success callback freed the key's max_parallel_requests slot while the socket was still open, and a key limited to one session could open several at once. The session now claims its slot under its own call id, every turn row logs under a call id of its own, and the slot is released when the session ends, as before Priced turn rows also take their row's call id as their id, so two turns that start in the same instant no longer share a spend log id and drop one row * fix(proxy): bill realtime transcription clips that finish outside an answer A conversation.item.input_audio_transcription.completed frame that arrived after a session's last response.done was never billed, because the session end only owes a row for a turn still running. A clip that completes while no answer is running now logs its own row, and a clip inside a running answer stays in that answer's row Each row also carries the session's latest session.updated frame, since later rows lost the transcription model it names and priced clips against the realtime model. A newer session.updated replaces the older one, so the stale model is never read first It also drops the docstrings this PR had added, since the names already say what each helper does * fix(cost): price a realtime transcription clip at the transcription model the session set A realtime conversation session opens with session.created naming the realtime model and no transcription model, and the client sets the transcription model later in session.update. The cost code took the first session frame and fell back to its session model, so a clip was priced at the realtime model. That model has no per-second rate, so a gpt-transcribe clip cost $0, and a token-billed clip used the realtime model's audio rate instead of the transcription model's. Now the newest session frame that names a transcription model sets the price, and the first session model is only the fallback when no frame names one, as before * fix(cost): price each realtime transcription clip at the model set when it ran A session that changes its transcription model mid-way kept every logged clip on the newest model. Each clip now takes the transcription model from the last session event before it, and falls back to the session's own model * fix(passthrough): keep the session update a buffered clip ran under until its answer is billed A transcription that completes while an answer streams waits in the buffer for that answer's response.done. A session.updated in between used to drop the earlier update, so the clip was priced at the opening configuration. An earlier update now stays while an unbilled clip follows it, and drops at the next update once that clip is billed. * fix(passthrough): keep reading a websocket turn the client left before OpenAI started it A client that sends response.create and hangs up before OpenAI answers with response.created made the walkout checks see no turn in flight, so the gateway closed the provider socket and the turn OpenAI still ran and billed never reached spend. The relay now counts every forwarded create until the provider starts or refuses it, and keeps draining while one is owed * fix(passthrough): settle a websocket turn that ends without its start frame A turn the provider ended without sending response.created first left its create counted as not started, so a client leaving later held the provider socket for the whole walkout drain with nothing left to bill. A turn's end now settles one waiting create when its lane saw no start frame * fix(proxy): refuse unparseable and unroutable websocket frames and hold legacy slots per socket A client frame too deep or too long for the JSON parser closed the provider socket mid-turn, so the turn in flight was never billed. The gate now refuses that frame with an error event and keeps draining the open turn. A stream_id OpenAI rejects comes back as an error that names no lane, which dropped another turn's authorization. The gate now refuses those frames before they reach OpenAI. With the legacy parallel request limiter, each per-turn spend row freed the key's slot while the socket stayed open. Websocket passthrough rows now free nothing, and the socket frees its slot once when it closes. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/cost_calculator.py | 76 +- litellm/litellm_core_utils/core_helpers.py | 6 +- litellm/proxy/auth/auth_checks.py | 73 - litellm/proxy/common_utils/realtime_utils.py | 3 +- .../proxy/hooks/parallel_request_limiter.py | 11 +- .../llm_passthrough_endpoints.py | 197 ++- ...i_websocket_passthrough_logging_handler.py | 262 +++ .../pass_through_endpoints.py | 734 ++++++-- .../pass_through_endpoints/success_handler.py | 70 +- .../spend_tracking/budget_reservation.py | 2 +- .../pass_through_endpoints.py | 14 +- ...st_auth_checks_object_access_and_lookup.py | 133 -- .../proxy/common_utils/test_realtime_cache.py | 9 +- ...i_websocket_passthrough_logging_handler.py | 520 ++++++ .../test_pass_through_endpoints.py | 1535 ++++++++++++++++- .../test_success_handler.py | 53 + .../test_openai_ws_passthrough_routes.py | 439 ++++- tests/unit/test_cost_calculator.py | 32 + 18 files changed, 3642 insertions(+), 527 deletions(-) create mode 100644 litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py create mode 100644 tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index e66be233be4..3c028278b2d 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -6,6 +6,7 @@ import time from collections.abc import Mapping, Sequence from datetime import datetime from functools import lru_cache +from itertools import accumulate from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast @@ -3850,44 +3851,65 @@ def handle_realtime_transcription_cost_calculation( - {"type": "duration", "seconds": } → priced via input_cost_per_second - {"type": "tokens", "input_tokens": ...} → priced via input/audio token cost """ - completed_events: Final = [ - cast(dict, result) for result in results if result.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE - ] - if not completed_events: + configured_models: Final = tuple(accumulate(map(_session_transcription_model, results), _newer_model)) + clips: Final = tuple( + (as_str_mapping(result.get("usage")) or {}, configured) + for result, configured in zip(results, configured_models, strict=True) + if result.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE + ) + if not clips: return 0.0 - model_name: Final = _get_transcription_model_name_from_results(results) or litellm_model_name - model_info: Final = _cost_map_model_info(model_name, custom_llm_provider) + fallback_model: Final = _get_transcription_model_name_from_results(results) or litellm_model_name + clip_models: Final = tuple((usage, configured or fallback_model) for usage, configured in clips) + model_infos: Final = MappingProxyType( + {model: _cost_map_model_info(model, custom_llm_provider) for model in frozenset(m for _, m in clip_models)} + ) override_info: Final = ( _cost_map_model_info(custom_pricing_model, custom_llm_provider) if custom_pricing_model is not None else None ) + return sum( + (_transcription_usage_cost(dict(usage), model_infos[model], override_info) for usage, model in clip_models), + 0.0, + ) - total_cost = 0.0 - for event in completed_events: - usage = event.get("usage") or {} - total_cost += _transcription_usage_cost(usage, model_info, override_info) - return total_cost + +_TRANSCRIPTION_SESSION_EVENT_TYPES: Final = frozenset( + {"transcription_session.created", "transcription_session.updated", "session.created", "session.updated"} +) + + +def _str_field(mapping: Mapping[str, object] | None, key: str) -> str | None: + value: Final = None if mapping is None else mapping.get(key) + return value if isinstance(value, str) and value else None + + +def _configured_transcription_model(session: Mapping[str, object]) -> str | None: + audio_input: Final = as_str_mapping((as_str_mapping(session.get("audio")) or {}).get("input")) + nested: Final = _str_field(as_str_mapping((audio_input or {}).get("transcription")), "model") + return nested or _str_field(as_str_mapping(session.get("input_audio_transcription")), "model") + + +def _session_of(result: Mapping[str, object]) -> Mapping[str, object] | None: + return as_str_mapping(result.get("session")) if result.get("type") in _TRANSCRIPTION_SESSION_EVENT_TYPES else None + + +def _session_transcription_model(result: Mapping[str, object]) -> str | None: + session: Final = _session_of(result) + return None if session is None else _configured_transcription_model(session) + + +def _newer_model(current: str | None, newer: str | None) -> str | None: + return newer or current def _get_transcription_model_name_from_results( results: OpenAIRealtimeStreamList, ) -> str | None: - """Resolve the ASR model from a transcription_session.* / session.* event.""" - for result in results: - if result.get("type") in ( - "transcription_session.created", - "transcription_session.updated", - "session.created", - "session.updated", - ): - session = cast(dict, result).get("session", {}) or {} - transcription = ((session.get("audio", {}) or {}).get("input", {}) or {}).get( - "transcription", {} - ) or session.get("input_audio_transcription", {}) - model = (transcription or {}).get("model") or session.get("model") - if model: - return model - return None + return next( + (model for result in results if (model := _str_field(_session_of(result), "model"))), + None, + ) _TranscriptionRateKey: TypeAlias = Literal[ diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 58189140a98..6fd6e91a627 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -4,7 +4,7 @@ import copy import logging import math import re -from collections.abc import Collection, Iterable, Iterator, Mapping +from collections.abc import Collection, Iterable, Iterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable @@ -104,6 +104,10 @@ def as_str_mapping(value: object) -> Mapping[str, object] | None: return value if isinstance(value, Mapping) else None +def as_object_list(value: object) -> Sequence[object] | None: + return value if isinstance(value, list) else None + + def decoded_json_object(value: object) -> Mapping[str, object]: mapping: Final = as_str_mapping(value) if mapping is not None: diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e2ae25686da..f5d3a80ee91 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5378,79 +5378,6 @@ async def _granted_model_lists( ) -async def _user_object_or_none( - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> LiteLLM_UserTable | None: - try: - return await get_user_object( - user_id=valid_token.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - proxy_logging_obj=proxy_logging_obj, - ) - except UserNotFoundError: - return None - - -async def enforced_model_allowlists( - valid_token: UserAPIKeyAuth, - prisma_client: PrismaClient | None, - user_api_key_cache: UserApiKeyCache, - proxy_logging_obj: ProxyLogging, -) -> tuple[Sequence[str], ...]: - """One model allowlist per level that ``common_checks`` enforces on a request from this identity.""" - key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token) - if prisma_client is None: - return (key_models, tuple(valid_token.team_models or ())) - team_object: Final = ( - None - if valid_token.team_id is None - else await get_team_object( - team_id=valid_token.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - user_object: Final = ( - None - if team_object is not None - else await _user_object_or_none( - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - project_object: Final = ( - None - if valid_token.project_id is None - else await get_project_object( - project_id=valid_token.project_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - ) - return ( - key_models, - team_object.models if team_object is not None else (), - await _team_member_granted_models( - valid_token=valid_token, - team_object=team_object, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ), - user_object.models if user_object is not None else (), - project_object.models if project_object is not None else (), - ) - - async def collect_matched_model_access_groups( model: str | Sequence[str] | None, valid_token: UserAPIKeyAuth | None, diff --git a/litellm/proxy/common_utils/realtime_utils.py b/litellm/proxy/common_utils/realtime_utils.py index c641748843e..e96685bdeae 100644 --- a/litellm/proxy/common_utils/realtime_utils.py +++ b/litellm/proxy/common_utils/realtime_utils.py @@ -1,3 +1,4 @@ +import json from functools import lru_cache from typing import Final @@ -10,7 +11,7 @@ def realtime_request_body(model: str | None) -> bytes: Generate the realtime websocket request body. Cached with LRU semantics to avoid repeated string formatting work while keeping memory usage bounded. """ - return f'{{"model": "{model or ""}"}}'.encode() + return json.dumps({"model": model or ""}).encode() _realtime_request_body: Final = realtime_request_body diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index cd01d15d834..f9e85c31d46 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -13,7 +13,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.exceptions import RateLimitType from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs +from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_parent_otel_span_from_kwargs from litellm.proxy._types import CommonProxyErrors, CurrentItemRateLimit, UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( get_key_model_rpm_limit, @@ -53,9 +53,16 @@ def _response_total_tokens(kwargs: Mapping[str, object], response_obj: object) - _CONNECTION_SCOPED_CALL_TYPES: Final = frozenset({CallTypes.aresponses_websocket.value}) +def _is_websocket_passthrough_row(kwargs: Mapping[str, object]) -> bool: + passthrough_payload: Final = as_str_mapping(kwargs.get("passthrough_logging_payload")) + return passthrough_payload is not None and passthrough_payload.get("request_method") == "WEBSOCKET" + + def _slots_released(kwargs: Mapping[str, object]) -> int: """A websocket turn callback releases nothing; the socket's slot is released once at disconnect.""" - return 0 if kwargs.get("call_type") in _CONNECTION_SCOPED_CALL_TYPES else 1 + if kwargs.get("call_type") in _CONNECTION_SCOPED_CALL_TYPES or _is_websocket_passthrough_row(kwargs): + return 0 + return 1 class CacheObject(TypedDict): diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index c7983f40aec..be9f80694b8 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -15,11 +15,11 @@ import os import posixpath import re import sys -from collections.abc import AsyncGenerator, Callable, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from functools import partial from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, cast import httpx import openai @@ -51,6 +51,7 @@ from litellm.constants import ( from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming, is_guardrail_intervention from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping, attribute_or_none from litellm.litellm_core_utils.fallback_generalizations import match_routing_generalization from litellm.llms.anthropic.common_utils import ( AnthropicModelInfo, @@ -85,7 +86,12 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.llms.vllm.common_utils import endpoint_under_api_base from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * -from litellm.proxy.auth.auth_checks import enforced_model_allowlists +from litellm.proxy.auth.auth_checks import ( + can_agent_call_model, + can_key_call_resolved_model, + key_model_aliases_for_auth_check, + team_model_aliases_for_auth_check, +) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports @@ -93,6 +99,7 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy modul get_bearer_token, is_no_auth_dev_mode, user_api_key_auth, + user_api_key_auth_accepted_websocket_for_model, user_api_key_auth_websocket, user_api_key_auth_websocket_for_model, ) @@ -114,6 +121,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg safe_get_request_headers, safe_set_request_parsed_body, ) +from litellm.proxy.common_utils.model_listing_utils import alias_map from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.proxy.common_utils.sse_keepalive import ( @@ -129,16 +137,26 @@ from litellm.proxy.pass_through_endpoints.injected_stream_usage import ( stream_usage_relay, without_injected_usage_frames, ) +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + openai_websocket_frame_kind, +) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + STORED_PROMPT_PROBE_MODEL, HttpPassThroughEndpointHelpers, + WebsocketFrameModelAllowed, + WebsocketFrameModelGate, + WebsocketFrameModelRefused, + WebsocketFrameModelVerdict, + WebsocketUpstreamFrameClassifier, YieldsToConfiguredPassThroughRoute, create_pass_through_route, create_websocket_passthrough_route, + refuse_websocket, truncated_close_reason, websocket_passthrough_request, ) from litellm.proxy.utils import ProxyLogging as ProxyLoggingType -from litellm.proxy.utils import is_known_model +from litellm.proxy.utils import is_known_model, normalize_route_for_root_path from litellm.proxy.vector_store_endpoints.utils import ( assert_proxy_admin_for_vector_store_index_management, assert_user_can_access_vector_store, @@ -198,7 +216,7 @@ def is_passthrough_request_using_router_model(request_body: dict, llm_router: li """ try: model: Final = request_body.get("model") - return is_known_model(model, llm_router) + return isinstance(model, str) and is_known_model(model, llm_router) except Exception: return False @@ -228,9 +246,8 @@ def is_passthrough_request_streaming(request_body: object) -> bool: A JSON body need not be an object, so a list or scalar can reach here; it carries no streaming flag. """ - if not isinstance(request_body, dict): - return False - return bool(request_body.get("stream", False)) + body: Final = as_str_mapping(request_body) + return body is not None and bool(body.get("stream", False)) def _optional_str(value: object) -> str | None: @@ -3387,19 +3404,6 @@ def _join_url_paths(base_url: httpx.URL, path: str, custom_llm_provider: litellm return joined_path_str -_OPENAI_WS_ALL_MODEL_ACCESS: Final = frozenset( - { - SpecialModelNames.all_proxy_models.value, - SpecialModelNames.all_team_models.value, - "*", - } -) - - -def _has_model_restrictions(model_allowlists: tuple[Sequence[str], ...]) -> bool: - return any(str(model) not in _OPENAI_WS_ALL_MODEL_ACCESS for allowlist in model_allowlists for model in allowlist) - - @dataclass(frozen=True, slots=True) class _OpenAIWebsocketRefusal: close_reason: str @@ -3416,6 +3420,16 @@ class _OpenAIWebsocketErrorFrame(TypedDict): error: ReadOnly[_OpenAIWebsocketErrorDetail] +_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL: Final = ( + "This gateway could not authorize the model named in this frame, so the session was closed. Try again, " + "and ask a proxy admin to check the gateway logs if it keeps happening." +) + +_OPENAI_WS_ALIASED_MODEL_REFUSAL: Final = ( + "This key cannot use a model alias over OpenAI websocket passthrough. The passthrough sends each frame to " + "OpenAI unchanged, so name the model this key may use directly." +) + _OPENAI_WS_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal( close_reason="OpenAI websocket passthrough is disabled", message=( @@ -3424,14 +3438,6 @@ _OPENAI_WS_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal( ), ) -_OPENAI_WS_MODEL_RESTRICTED_REFUSAL: Final = _OpenAIWebsocketRefusal( - close_reason="Keys with model restrictions cannot use OpenAI websocket passthrough", - message=( - "Keys with model restrictions cannot use OpenAI websocket passthrough, because this route " - "relays frames to the provider without reading which model they ask for." - ), -) - def _is_openai_websocket_passthrough_enabled(general_settings: Mapping[str, object]) -> bool: setting: Final = general_settings.get("enable_openai_websocket_passthrough") @@ -3440,19 +3446,9 @@ def _is_openai_websocket_passthrough_enabled(general_settings: Mapping[str, obje return setting is True -class _OpenAIWebsocketModelAllowlists(Protocol): - async def __call__(self, valid_token: UserAPIKeyAuth, /) -> tuple[Sequence[str], ...]: ... - - -async def _openai_websocket_refusal( - user_api_key_dict: UserAPIKeyAuth, - general_settings: Mapping[str, object], - model_allowlists: _OpenAIWebsocketModelAllowlists, -) -> _OpenAIWebsocketRefusal | None: +def _openai_websocket_refusal(general_settings: Mapping[str, object]) -> _OpenAIWebsocketRefusal | None: if not _is_openai_websocket_passthrough_enabled(general_settings): return _OPENAI_WS_DISABLED_REFUSAL - if _has_model_restrictions(await model_allowlists(user_api_key_dict)): - return _OPENAI_WS_MODEL_RESTRICTED_REFUSAL return None @@ -3467,6 +3463,8 @@ class _WebsocketRelay(Protocol): forward_headers: bool, endpoint: str, accept_websocket: bool, + client_frame_model_gate: WebsocketFrameModelGate | None = None, + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None, meter: WebSocketSessionMeter | None = None, ) -> None: ... @@ -3475,18 +3473,95 @@ def _websocket_relay() -> _WebsocketRelay: return websocket_passthrough_request -def _proxy_model_allowlists() -> _OpenAIWebsocketModelAllowlists: - from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache +def _proxy_error_type(denial: ProxyException) -> str: + return denial.type.value if isinstance(denial.type, ProxyErrorTypes) else denial.type - async def resolve(valid_token: UserAPIKeyAuth, /) -> tuple[Sequence[str], ...]: - return await enforced_model_allowlists( - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, + +FrameModelAuthorizer: TypeAlias = Callable[[WebSocket, str], Awaitable[UserAPIKeyAuth]] + + +def _frame_model_authorizer() -> FrameModelAuthorizer: + return user_api_key_auth_accepted_websocket_for_model + + +def _proxy_model_list() -> Sequence[object] | None: + from litellm.proxy import proxy_server + + return as_object_list(attribute_or_none(proxy_server, "llm_model_list")) + + +def _gateway_remaps_model(model: str, frame_auth: UserAPIKeyAuth, llm_router: Router | None) -> bool: + alias_maps: Final = ( + key_model_aliases_for_auth_check(frame_auth) or {}, + team_model_aliases_for_auth_check(frame_auth) or {}, + alias_map(litellm.model_alias_map), + ) + targets: Final = ( + *(aliases.get(model) for aliases in alias_maps), + llm_router.registry().resolve(model, kind="alias") if llm_router is not None else None, + ) + return any(target not in (None, model) for target in targets) + + +def _aliased_model_verdict( + allowed: WebsocketFrameModelAllowed, probe: WebsocketFrameModelVerdict +) -> WebsocketFrameModelVerdict: + match probe: + case WebsocketFrameModelRefused(error_type=error_type) if error_type.endswith("_model_access_denied"): + return WebsocketFrameModelRefused(message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=error_type) + case WebsocketFrameModelRefused(): + return probe + case WebsocketFrameModelAllowed(): + return allowed + + +def _proxy_frame_model_gate( + authorize_frame_model: Annotated[FrameModelAuthorizer, Depends(_frame_model_authorizer)], +) -> WebsocketFrameModelGate: + from litellm.proxy import proxy_server + from litellm.proxy.spend_tracking.budget_reservation import release_or_invalidate_budget_reservation + + async def authorize(websocket: WebSocket, model: str) -> WebsocketFrameModelVerdict: + connect_reservation: Final = attribute_or_none(websocket.state, "budget_reservation") + await release_or_invalidate_budget_reservation( + budget_reservation=connect_reservation if isinstance(connect_reservation, dict) else None ) + frame_auth: Final = await authorize_frame_model(websocket, model) + access_groups: Final = tuple(frame_auth.matched_model_access_groups or ()) + llm_router: Final = proxy_server.llm_router + try: + await can_key_call_resolved_model( + model=model, llm_model_list=_proxy_model_list(), valid_token=frame_auth, llm_router=llm_router + ) + await can_agent_call_model( + model=model, + valid_token=frame_auth, + llm_router=llm_router, + prisma_client=proxy_server.prisma_client, + user_api_key_cache=proxy_server.user_api_key_cache, + proxy_logging_obj=proxy_server.proxy_logging_obj, + ) + finally: + await release_or_invalidate_budget_reservation(budget_reservation=frame_auth.budget_reservation) + allowed: Final = WebsocketFrameModelAllowed(access_groups=access_groups) + if not _gateway_remaps_model(model, frame_auth, llm_router): + return allowed + return _aliased_model_verdict(allowed, await gate(websocket, STORED_PROMPT_PROBE_MODEL)) - return resolve + async def gate(websocket: WebSocket, model: str, /) -> WebsocketFrameModelVerdict: + try: + return await authorize(websocket, model) + except ProxyException as denial: + return WebsocketFrameModelRefused(message=denial.message, error_type=_proxy_error_type(denial)) + except HTTPException as denial: + return WebsocketFrameModelRefused(message=str(denial.detail), error_type="invalid_request_error") + except litellm.BudgetExceededError as exhausted: + return WebsocketFrameModelRefused(message=exhausted.message, error_type="budget_exceeded") + except Exception: # noqa: BLE001 # fail closed: a gate that cannot answer refuses, and says so in band + verbose_proxy_logger.exception("OpenAI websocket passthrough: the frame model gate failed on %s", model) + return WebsocketFrameModelRefused(message=_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, error_type="server_error") + + return gate def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None: @@ -3498,6 +3573,14 @@ def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None: return requested_subprotocols[0] if requested_subprotocols else None +async def _opening_model_refusal( + websocket: WebSocket, query_string: str, frame_model_gate: WebsocketFrameModelGate +) -> WebsocketFrameModelRefused | None: + opening_models: Final = dict.fromkeys(httpx.QueryParams(query_string).get_list("model")) + verdicts: Final = tuple([await frame_model_gate(websocket, model) for model in opening_models]) + return next((verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelRefused)), None) + + @router.websocket("/openai_passthrough/{endpoint:path}") @router.websocket("/openai/{endpoint:path}") async def openai_websocket_proxy_route( @@ -3506,12 +3589,12 @@ async def openai_websocket_proxy_route( user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)], general_settings: Annotated[Mapping[str, object], Depends(_proxy_general_settings)], relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)], - model_allowlists: Annotated[_OpenAIWebsocketModelAllowlists, Depends(_proxy_model_allowlists)], + frame_model_gate: Annotated[WebsocketFrameModelGate, Depends(_proxy_frame_model_gate)], ) -> None: """WebSocket passthrough for OpenAI prefixes (realtime / responses.connect).""" negotiated_subprotocol: Final = _negotiated_websocket_subprotocol(websocket) - refusal: Final = await _openai_websocket_refusal(user_api_key_dict, general_settings, model_allowlists) + refusal: Final = _openai_websocket_refusal(general_settings) if refusal is not None: await websocket.accept(subprotocol=negotiated_subprotocol) error_frame: Final[_OpenAIWebsocketErrorFrame] = { @@ -3557,14 +3640,21 @@ async def openai_websocket_proxy_route( await websocket.accept(subprotocol=negotiated_subprotocol) + opening_refusal: Final = await _opening_model_refusal(websocket, query_string, frame_model_gate) + if opening_refusal is not None: + await refuse_websocket(websocket, opening_refusal) + return + await relay( websocket=websocket, target=wss_target, custom_headers=custom_headers, user_api_key_dict=user_api_key_dict, forward_headers=False, - endpoint=websocket.url.path, + endpoint=normalize_route_for_root_path(websocket.url.path) or websocket.url.path, accept_websocket=False, + client_frame_model_gate=frame_model_gate, + upstream_frame_kind=openai_websocket_frame_kind, ) @@ -3882,7 +3972,8 @@ async def _tinyfish_json_object_field_names(request: Request) -> frozenset[str] parsed: Final = _DECODED_JSON.validate_python(json.loads(raw_body)) except (json.JSONDecodeError, UnicodeDecodeError): return None - return frozenset(parsed) if isinstance(parsed, dict) else None + fields: Final = as_str_mapping(parsed) + return None if fields is None else frozenset(fields) def _tinyfish_route_timeout() -> float | None: diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py new file mode 100644 index 00000000000..f3bad54ffb2 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py @@ -0,0 +1,262 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime +from types import MappingProxyType +from typing import Final, Literal, TypeAlias, cast # noqa: TID251 # cost helpers type raw frames as a TypedDict list + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.cost_calculator import RealtimeAPITokenUsageProcessor, handle_realtime_stream_cost_calculation +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._types import PassThroughEndpointLoggingTypedDict +from litellm.responses.utils import ResponseAPILoggingUtils +from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponsesAPIResponse +from litellm.types.passthrough_endpoints.pass_through_endpoints import WebsocketUpstreamFrameKind +from litellm.types.utils import LlmProviders, ModelResponse, Usage + +WebsocketMessages: TypeAlias = Sequence[Mapping[str, object]] + +OPENAI_WEBSOCKET_PROVIDER: Final = LlmProviders.OPENAI.value + +_SESSION_CREATED: Final = "session.created" +_SESSION_UPDATED: Final = "session.updated" +_RESPONSE_CREATED: Final = "response.created" +_ERROR: Final = "error" +_REALTIME_RESPONSE_DONE: Final = "response.done" +_TRANSCRIPTION_COMPLETED: Final = "conversation.item.input_audio_transcription.completed" +_RESPONSES_TURN_ENDS: Final = frozenset({"response.completed", "response.incomplete", "response.failed"}) +_TURN_ENDS: Final = _RESPONSES_TURN_ENDS | {_REALTIME_RESPONSE_DONE} +_MODEL_NAMING_EVENTS: Final = _TURN_ENDS | {_SESSION_CREATED, _SESSION_UPDATED, _RESPONSE_CREATED} +_REALTIME_BILLED_EVENTS: Final = frozenset({_REALTIME_RESPONSE_DONE, _TRANSCRIPTION_COMPLETED}) +_RESPONSE_COST_HEADER: Final = "llm_provider-x-litellm-response-cost" +_FRAME_KINDS: Final[Mapping[str, WebsocketUpstreamFrameKind]] = MappingProxyType( + { + _SESSION_CREATED: "opens_session", + _SESSION_UPDATED: "updates_session", + _RESPONSE_CREATED: "starts_turn", + _TRANSCRIPTION_COMPLETED: "bills_usage", + _ERROR: "refuses_turn", + } +) + + +class _NamedModel(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + model: str | None = None + + +class _AudioInput(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + transcription: _NamedModel | None = None + + +class _SessionAudio(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + input: _AudioInput | None = None + + +class _NamedSession(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + model: str | None = None + audio: _SessionAudio | None = None + + +class _ModelNamingEvent(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + session: _NamedSession | None = None + response: _NamedModel | None = None + + +class _ResponsesTurnEndEvent(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + type: Literal["response.completed", "response.incomplete", "response.failed"] + response: ResponsesAPIResponse + + +_MODEL_NAMING_EVENT: Final = TypeAdapter(_ModelNamingEvent) +_RESPONSES_TURN_END_EVENT: Final = TypeAdapter(_ResponsesTurnEndEvent) + + +@dataclass(frozen=True, slots=True) +class _Billing: + usage: Usage + cost: float + + +@dataclass(frozen=True, slots=True) +class _CompletedTurn: + model: str + response: ResponsesAPIResponse + usage: Usage + + +def _event_model(event: Mapping[str, object]) -> str | None: + if event.get("type") not in _MODEL_NAMING_EVENTS: + return None + try: + naming: Final = _MODEL_NAMING_EVENT.validate_python(event) + except ValidationError: + return None + session: Final = naming.session or _NamedSession() + response: Final = naming.response or _NamedModel() + audio_input: Final = (session.audio or _SessionAudio()).input or _AudioInput() + transcription: Final = audio_input.transcription or _NamedModel() + return session.model or response.model or transcription.model or None + + +def observed_websocket_model(messages: WebsocketMessages) -> str | None: + return next((model for model in map(_event_model, messages) if model is not None), None) + + +def openai_websocket_frame_kind(frame: Mapping[str, object]) -> WebsocketUpstreamFrameKind: + frame_type: Final = frame.get("type") + if frame_type in _TURN_ENDS: + return "ends_turn" + if not isinstance(frame_type, str): + return "relayed" + return _FRAME_KINDS.get(frame_type, "relayed") + + +def _realtime_billing(messages: WebsocketMessages, model: str, logging_obj: LiteLLMLoggingObj) -> _Billing | None: + if not any(event.get("type") in _REALTIME_BILLED_EVENTS for event in messages): + return None + results: Final = cast( # cast-ok: the realtime cost helpers read raw provider frames through their TypedDict view + OpenAIRealtimeStreamList, tuple(event for event in messages if "type" in event) + ) + usage: Final = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results) + cost: Final = handle_realtime_stream_cost_calculation( + results=results, + combined_usage_object=usage, + custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER, + litellm_model_name=model, + litellm_logging_obj=logging_obj, + ) + return _Billing(usage=usage, cost=cost) + + +def _responses_turn(event: Mapping[str, object], session_model: str) -> _CompletedTurn | None: + if event.get("type") not in _RESPONSES_TURN_ENDS: + return None + try: + ended: Final = _RESPONSES_TURN_END_EVENT.validate_python(event) + except ValidationError: + verbose_proxy_logger.warning( + "OpenAI websocket passthrough: a %s frame carried no readable usage, so it is not billed", event.get("type") + ) + return None + if ended.response.usage is None: + return None + return _CompletedTurn( + model=ended.response.model or session_model, + response=ended.response, + usage=ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # the one shared Responses usage converter + ended.response.usage + ), + ) + + +def _responses_billing(messages: WebsocketMessages, model: str, requested_model: str | None) -> _Billing | None: + turns: Final = tuple(turn for turn in (_responses_turn(event, model) for event in messages) if turn is not None) + if not turns: + return None + usage: Final = RealtimeAPITokenUsageProcessor.combine_usage_objects([turn.usage for turn in turns]) + return _Billing( + usage=usage, + cost=sum(_responses_cost(turn.response, (turn.model, requested_model)) for turn in turns), + ) + + +def _model_cost(response: ResponsesAPIResponse, model: str) -> float | None: + try: + return litellm.completion_cost( + completion_response=response, + model=model, + custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER, + call_type="responses", + ) + except Exception: # noqa: BLE001 # completion_cost raises a bare Exception for an unpriced model, and the turn must still log + return None + + +def _responses_cost(response: ResponsesAPIResponse, price_as: Sequence[str | None]) -> float: + costs: Final = (_model_cost(response, model) for model in dict.fromkeys(price_as) if model is not None) + cost: Final = next((cost for cost in costs if cost), None) + if cost is None: + verbose_proxy_logger.warning( + "OpenAI websocket passthrough: could not price a Responses turn on %s, logging it at zero", price_as + ) + return 0.0 + return cost + + +def _requested_model(kwargs: Mapping[str, object]) -> str | None: + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) + metadata: Final = None if litellm_params is None else as_str_mapping(litellm_params.get("metadata")) + model_group: Final = None if metadata is None else metadata.get("model_group") + return model_group if isinstance(model_group, str) else None + + +def _priced_result(model: str, billing: _Billing, start_time: datetime, call_id: str) -> ModelResponse: + result: Final = ModelResponse( + id=call_id, + object="chat.completion", + created=int(start_time.timestamp()), + model=model, + usage=billing.usage, + ) + hidden: Final = {"additional_headers": MappingProxyType({_RESPONSE_COST_HEADER: billing.cost})} + result._hidden_params = hidden # pyright: ignore[reportPrivateUsage] # the one settled-cost channel + return result + + +class OpenAIWebsocketPassthroughLoggingHandler: + def openai_websocket_passthrough_handler( + self, + websocket_messages: WebsocketMessages, + logging_obj: LiteLLMLoggingObj, + url_route: str, + start_time: datetime, + kwargs: Mapping[str, object], + ) -> PassThroughEndpointLoggingTypedDict: + model: Final = observed_websocket_model(websocket_messages) + if model is None: + verbose_proxy_logger.debug("OpenAI websocket passthrough (%s): no model named, skipping cost", url_route) + unpriced: Final[PassThroughEndpointLoggingTypedDict] = {"result": None, "kwargs": {**kwargs}} + return unpriced + + logging_obj.model = model # rebind-ok: the spend logger reads model and cost off the shared logging object + logging_obj.model_call_details.update( + MappingProxyType({"model": model, "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER}) + ) + + billing: Final = _realtime_billing(websocket_messages, model, logging_obj) or _responses_billing( + websocket_messages, model, _requested_model(kwargs) + ) + if billing is None: + verbose_proxy_logger.debug("OpenAI websocket passthrough (%s): no usage reported, skipping cost", url_route) + unbilled: Final[PassThroughEndpointLoggingTypedDict] = { + "result": None, + "kwargs": {**kwargs, "model": model, "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER}, + } + return unbilled + + logging_obj.model_call_details.update(MappingProxyType({"response_cost": billing.cost})) + priced: Final[PassThroughEndpointLoggingTypedDict] = { + "result": _priced_result(model, billing, start_time, logging_obj.litellm_call_id), + "kwargs": { + **kwargs, + "model": model, + "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER, + "response_cost": billing.cost, + }, + } + return priced diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 970cd466337..7f5983d5799 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -6,15 +6,17 @@ import posixpath import re import traceback from base64 import b64encode -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, MutableMapping, Sequence from dataclasses import dataclass from datetime import datetime -from itertools import chain, count, groupby +from functools import reduce +from itertools import chain, count, groupby, takewhile from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypedDict, cast -from urllib.parse import urlencode, urlparse +from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict, TypeVar, cast +from urllib.parse import parse_qs, urlencode, urlparse import httpx +import orjson from fastapi import ( APIRouter, Depends, @@ -24,6 +26,7 @@ from fastapi import ( Response, UploadFile, WebSocket, + WebSocketDisconnect, status, ) from fastapi.responses import StreamingResponse @@ -34,7 +37,7 @@ from starlette.routing import BaseRoute, Match, Route from starlette.types import Scope from starlette.websockets import WebSocketState from typing_extensions import ReadOnly -from websockets.asyncio.client import connect +from websockets.asyncio.client import ClientConnection, connect from websockets.exceptions import ( ConnectionClosedError, ConnectionClosedOK, @@ -60,6 +63,8 @@ from litellm.integrations.custom_guardrail import ( from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation from litellm.litellm_core_utils.core_helpers import ( + as_object_list, + as_str_mapping, bind_budget_reservation_to_callbacks, get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, @@ -69,6 +74,7 @@ from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import get_masked_values from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.realtime_errors import realtime_error_event from litellm.litellm_core_utils.redact_messages import should_redact_message_logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.anthropic.common_utils import ( @@ -154,6 +160,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( PassthroughStandardLoggingPayload, WebSocketSessionMeter, WebSocketSessionUsage, + WebsocketUpstreamFrameKind, ) from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, CallTypes, Usage @@ -2574,11 +2581,16 @@ def _resolved_vertex_live_setup( return {**setup_data, "model": setup_model_rewriter(setup_model)} +def _decoded_frame(frame: str | bytes) -> object: + try: + return _DECODED_JSON.validate_python(orjson.loads(frame)) + except orjson.JSONDecodeError: + return _DECODED_JSON.validate_python(json.loads(frame if isinstance(frame, str) else frame.decode("utf-8"))) + + def _json_object_frame(frame: str | bytes) -> dict[str, object] | None: try: - decoded: Final[object] = _DECODED_JSON.validate_python( - json.loads(frame if isinstance(frame, str) else frame.decode("utf-8")) - ) + decoded: Final = _decoded_frame(frame) except (json.JSONDecodeError, UnicodeDecodeError): return None return decoded if isinstance(decoded, dict) else None @@ -2637,6 +2649,359 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict return inject_trace_context(headers, parent_span=parent_span) +@dataclass(frozen=True, slots=True) +class WebsocketFrameModelAllowed: + access_groups: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class WebsocketFrameModelRefused: + message: str + error_type: str + + +WebsocketFrameModelVerdict: TypeAlias = WebsocketFrameModelAllowed | WebsocketFrameModelRefused + + +class WebsocketFrameModelGate(Protocol): + async def __call__(self, websocket: WebSocket, model: str, /) -> WebsocketFrameModelVerdict: ... + + +WebsocketUpstreamFrameClassifier: TypeAlias = Callable[[Mapping[str, object]], WebsocketUpstreamFrameKind] +_SESSION_FRAMES_EVERY_ROW_CARRIES: Final[frozenset[WebsocketUpstreamFrameKind]] = frozenset( + {"opens_session", "updates_session"} +) +_FrameT: Final = TypeVar("_FrameT", bound=Mapping[str, object]) + + +@dataclass(frozen=True, slots=True) +class _TurnAuthorization: + requested_model: str | None + access_groups: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class _ForwardClientFrame: + starts_turn: bool + turn: _TurnAuthorization | None + lane: str | None + + +@dataclass(frozen=True, slots=True) +class _RefuseClientFrame: + refusal: WebsocketFrameModelRefused + + +@dataclass(frozen=True, slots=True) +class _ClientFrameRequests: + starts_turn: bool + lane: str | None + turn_model: str | None + models: tuple[str, ...] + unverifiable_model_refusal: str | None + + +_NO_CLIENT_FRAME_REQUESTS: Final = _ClientFrameRequests( + starts_turn=False, lane=None, turn_model=None, models=(), unverifiable_model_refusal=None +) +_FORWARD_WITHOUT_TURN: Final = _ForwardClientFrame(starts_turn=False, turn=None, lane=None) +_UNREADABLE_CLIENT_FRAME: Final = _RefuseClientFrame( + refusal=WebsocketFrameModelRefused( + message="The gateway cannot parse this frame, so it cannot check which model it asks for.", + error_type="invalid_request_error", + ) +) +_TURN_STARTING_FRAME: Final = "response.create" +_FRAME_TEXT_WORTH_PARSING: Final = (f'"{_TURN_STARTING_FRAME}"', '"model"', '"prompt"', '"tools"', "\\") +_TURN_MODEL_PLACEMENTS: Final = (("model",), ("response", "model"), ("session", "model")) +_TRANSCRIPTION_MODEL_PLACEMENTS: Final = ( + ("session", "input_audio_transcription", "model"), + ("session", "audio", "input", "transcription", "model"), +) +_STORED_PROMPT_PLACEMENTS: Final = (("prompt",), ("response", "prompt"), ("session", "prompt")) +_TOOL_LIST_PLACEMENTS: Final = (("tools",), ("response", "tools"), ("session", "tools")) +_TOOL_THAT_RUNS_A_MODEL: Final = "image_generation" +STORED_PROMPT_PROBE_MODEL: Final = "litellm-unverifiable-stored-prompt-model" +_STORED_PROMPT_REFUSAL: Final = ( + "This key is limited to specific models, so it cannot run a stored prompt. " + "The gateway cannot check which model a stored prompt runs." +) +_UNNAMED_TOOL_MODEL_REFUSAL: Final = ( + "This key is limited to specific models, so an image_generation tool has to name its model. " + "The gateway cannot check which model an unnamed one runs." +) +_FIRST_SENTENCE: Final = re.compile(r".*?[.!?](?=\s|$)") +_LANE_NAME: Final = re.compile(r"[A-Za-z0-9_.-]{1,256}") +_UNROUTABLE_LANE_REFUSAL: Final = _RefuseClientFrame( + refusal=WebsocketFrameModelRefused( + message=( + "This stream_id is not a lane name the provider accepts. " + "The provider's error for it names no lane, so the gateway cannot tell which turn it refuses." + ), + error_type="invalid_request_error", + ) +) + + +def _may_name_a_model(text_data: str) -> bool: + return any(marker in text_data for marker in _FRAME_TEXT_WORTH_PARSING) + + +def _frame_lane(frame: Mapping[str, object]) -> str | None: + stream_id: Final = frame.get("stream_id") + return stream_id if isinstance(stream_id, str) else None + + +def _names_an_unroutable_lane(frame: Mapping[str, object]) -> bool: + stream_id: Final = frame.get("stream_id") + return stream_id is not None and not (isinstance(stream_id, str) and _LANE_NAME.fullmatch(stream_id) is not None) + + +def _frames_of_lane(lane: str | None) -> Callable[[Mapping[str, object]], bool]: + def in_lane(frame: Mapping[str, object]) -> bool: + return _frame_lane(frame) == lane + + return in_lane + + +def _every_frame(_frame: Mapping[str, object]) -> bool: + return True + + +def _without_earlier_session_updates( + upstream_frame_kind: WebsocketUpstreamFrameClassifier, + frames: Sequence[_FrameT], + latest_update: _FrameT, +) -> tuple[_FrameT, ...]: + return tuple( + frame + for index, frame in enumerate(frames) + if frame is latest_update + or upstream_frame_kind(frame) != "updates_session" + or _unbilled_usage_ran_under(map(upstream_frame_kind, frames[index + 1 :])) + ) + + +def _unbilled_usage_ran_under(later_kinds: Iterable[WebsocketUpstreamFrameKind]) -> bool: + return "bills_usage" in takewhile(lambda kind: kind != "updates_session", later_kinds) + + +def _nested_frame(frame: Mapping[str, object] | None, key: str) -> Mapping[str, object] | None: + return None if frame is None else as_str_mapping(frame.get(key)) + + +def _placed_value(frame: Mapping[str, object], placement: Sequence[str]) -> object: + *parents, leaf = placement + container: Final = reduce(_nested_frame, parents, frame) + return None if container is None else container.get(leaf) + + +def _placed_models(frame: Mapping[str, object], placements: Iterable[Sequence[str]]) -> tuple[str, ...]: + placed: Final = (_placed_value(frame, placement) for placement in placements) + return tuple(value for value in placed if isinstance(value, str)) + + +def _placed_tools(frame: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + tool_lists: Final = (_placed_value(frame, placement) for placement in _TOOL_LIST_PLACEMENTS) + tools: Final = chain.from_iterable(as_object_list(tool_list) or () for tool_list in tool_lists) + return tuple(tool for tool in map(as_str_mapping, tools) if tool is not None) + + +def _unverifiable_model_refusal(frame: Mapping[str, object], tools: Sequence[Mapping[str, object]]) -> str | None: + prompts: Final = (_placed_value(frame, placement) for placement in _STORED_PROMPT_PLACEMENTS) + if any(isinstance(prompt, Mapping) for prompt in prompts): + return _STORED_PROMPT_REFUSAL + if any(tool.get("type") == _TOOL_THAT_RUNS_A_MODEL and not isinstance(tool.get("model"), str) for tool in tools): + return _UNNAMED_TOOL_MODEL_REFUSAL + return None + + +def _client_frame_requests(frame_data: str | bytes) -> _ClientFrameRequests | _RefuseClientFrame: + if isinstance(frame_data, str) and not _may_name_a_model(frame_data): + return _NO_CLIENT_FRAME_REQUESTS + try: + frame: Final = _json_object_frame(frame_data) + except (RecursionError, ValueError): + return _UNREADABLE_CLIENT_FRAME + if frame is None: + return _NO_CLIENT_FRAME_REQUESTS + if _names_an_unroutable_lane(frame): + return _UNROUTABLE_LANE_REFUSAL + turn_models: Final = _placed_models(frame, _TURN_MODEL_PLACEMENTS) + tools: Final = _placed_tools(frame) + tool_models: Final = (model for model in (tool.get("model") for tool in tools) if isinstance(model, str)) + return _ClientFrameRequests( + starts_turn=frame.get("type") == _TURN_STARTING_FRAME, + lane=_frame_lane(frame), + turn_model=next(iter(turn_models), None), + models=tuple( + dict.fromkeys((*turn_models, *_placed_models(frame, _TRANSCRIPTION_MODEL_PLACEMENTS), *tool_models)) + ), + unverifiable_model_refusal=_unverifiable_model_refusal(frame, tools), + ) + + +def _unverifiable_model_probe_refusal(probe: WebsocketFrameModelRefused, message: str) -> WebsocketFrameModelRefused: + if probe.error_type.endswith("_model_access_denied"): + return WebsocketFrameModelRefused(message=message, error_type=probe.error_type) + return probe + + +async def _client_frame_verdict( + websocket: WebSocket, frame_data: str | bytes, gate: WebsocketFrameModelGate | None +) -> _ForwardClientFrame | _RefuseClientFrame: + if gate is None: + return _FORWARD_WITHOUT_TURN + requests: Final = _client_frame_requests(frame_data) + if isinstance(requests, _RefuseClientFrame): + return requests + verdicts: Final = tuple([await gate(websocket, model) for model in requests.models]) + refusal: Final = next((verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelRefused)), None) + if refusal is not None: + return _RefuseClientFrame(refusal=refusal) + if requests.unverifiable_model_refusal is not None and isinstance( + probe := await gate(websocket, STORED_PROMPT_PROBE_MODEL), WebsocketFrameModelRefused + ): + return _RefuseClientFrame(refusal=_unverifiable_model_probe_refusal(probe, requests.unverifiable_model_refusal)) + allowed: Final = tuple(verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelAllowed)) + return _ForwardClientFrame( + starts_turn=requests.starts_turn, + lane=requests.lane, + turn=( + _TurnAuthorization(requested_model=requests.turn_model, access_groups=allowed[0].access_groups) + if requests.turn_model is not None and allowed + else None + ), + ) + + +def _refusal_close_reason(message: str) -> str: + first_sentence: Final = _FIRST_SENTENCE.match(message) + return truncated_close_reason(message if first_sentence is None else first_sentence.group(0)) + + +async def refuse_websocket(websocket: WebSocket, refusal: WebsocketFrameModelRefused) -> None: + await websocket.send_text(realtime_error_event(refusal.message, error_type=refusal.error_type)) + await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(refusal.message)) + + +def _stamp_billed_turn(kwargs: Mapping[str, object], turn: _TurnAuthorization) -> None: + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) + metadata: Final = None if litellm_params is None else litellm_params.get("metadata") + if not isinstance(metadata, MutableMapping): + return + metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list(turn.access_groups) + if turn.requested_model is not None: + metadata["model_group"] = turn.requested_model + + +def _opening_model(target: str) -> str | None: + return next(iter(parse_qs(urlparse(target).query).get("model", ())), None) + + +class _WebsocketTurnLedger: + def __init__(self, opening: _TurnAuthorization, started: datetime) -> None: + self.opening = opening + self.queued_turns: Mapping[str | None, tuple[_TurnAuthorization | None, ...]] = MappingProxyType({}) + self.session_opened = False + self.started = started + self.logged_turns = 0 + self.unstarted_turns = 0 + + def current(self, lane: str | None) -> _TurnAuthorization: + oldest: Final = next(iter(self.queued_turns.get(lane, ())), None) + return self.opening if oldest is None else oldest + + def queue_turn(self, lane: str | None, turn: _TurnAuthorization | None) -> None: + self.unstarted_turns += 1 + if not self.session_opened: + self.queued_turns = MappingProxyType({**self.queued_turns, lane: (*self.queued_turns.get(lane, ()), turn)}) + + def open_session(self) -> None: + self.session_opened = True + self.queued_turns = MappingProxyType({}) + + def close_turn(self, lane: str | None, end_time: datetime) -> None: + self.logged_turns += 1 + self.started = end_time + self._forget_turn(lane, 0) + + def start_turn(self) -> None: + self.unstarted_turns = max(0, self.unstarted_turns - 1) + + def drop_refused_turn(self, lane: str | None, lane_is_running_a_turn: bool) -> None: + self.start_turn() + self._forget_turn(lane, 1 if lane_is_running_a_turn else 0) + + def _forget_turn(self, lane: str | None, position: int) -> None: + turns: Final = self.queued_turns.get(lane, ()) + self.queued_turns = MappingProxyType({**self.queued_turns, lane: (*turns[:position], *turns[position + 1 :])}) + + +def _session_end_needs_a_row( + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None, + unlogged_frames: Sequence[Mapping[str, object]], + ledger: _WebsocketTurnLedger, +) -> bool: + if upstream_frame_kind is None: + return True + kinds: Final = frozenset(map(upstream_frame_kind, unlogged_frames)) + return "starts_turn" in kinds or (ledger.logged_turns == 0 and bool(kinds - {"opens_session"})) + + +def _turn_in_flight( + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None, unlogged_frames: Sequence[Mapping[str, object]] +) -> bool: + return upstream_frame_kind is not None and any( + upstream_frame_kind(frame) == "starts_turn" for frame in unlogged_frames + ) + + +def _turn_owed( + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None, + unlogged_frames: Sequence[Mapping[str, object]], + ledger: _WebsocketTurnLedger, +) -> bool: + return _turn_in_flight(upstream_frame_kind, unlogged_frames) or ( + upstream_frame_kind is not None and ledger.unstarted_turns > 0 + ) + + +def _lanes_in_flight( + upstream_frame_kind: WebsocketUpstreamFrameClassifier, unlogged_frames: Sequence[Mapping[str, object]] +) -> tuple[str | None, ...]: + started: Final = (frame for frame in unlogged_frames if upstream_frame_kind(frame) == "starts_turn") + return tuple(dict.fromkeys(map(_frame_lane, started))) + + +async def _relay_client_frame( + websocket: WebSocket, + upstream_ws: ClientConnection, + checked_frame: str | bytes, + forwarded_frame: str | bytes, + gate: WebsocketFrameModelGate | None, + ledger: _WebsocketTurnLedger, +) -> _ForwardClientFrame | _RefuseClientFrame: + verdict: Final = await _client_frame_verdict(websocket, checked_frame, gate) + match verdict: + case _RefuseClientFrame(refusal=refusal): + await refuse_websocket(websocket, refusal) + case _ForwardClientFrame(starts_turn=True, turn=turn, lane=lane): + ledger.queue_turn(lane, turn) + await upstream_ws.send(forwarded_frame) + case _ForwardClientFrame(): + await upstream_ws.send(forwarded_frame) + return verdict + + +_EXPLICIT_SUCCESS_KWARGS: Final = frozenset({"logging_obj", "passthrough_logging_payload", "custom_llm_provider"}) +WALKOUT_TURN_DRAIN_SECONDS: Final = 600.0 + + +def _metered_client_text(meter: WebSocketSessionMeter | None, client_text: str) -> int: + return 0 if meter is None else meter.measure(client_text) + + def _log_session_usage(logging_obj: LiteLLMLoggingObj, usage: WebSocketSessionUsage) -> None: logging_obj.model = usage.model # rebind-ok: the spend logger reads model and cost off the shared logging object logging_obj.model_call_details["model"] = usage.model # rebind-ok: same shared logging object @@ -2654,6 +3019,9 @@ async def websocket_passthrough_request( cost_per_request: float | None = None, accept_websocket: bool = True, setup_model_rewriter: Callable[[str], str] | None = None, + client_frame_model_gate: WebsocketFrameModelGate | None = None, + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None, + walkout_turn_drain_seconds: float = WALKOUT_TURN_DRAIN_SECONDS, builtin_provider_route: bool = False, model: str | None = None, meter: WebSocketSessionMeter | None = None, @@ -2670,6 +3038,11 @@ async def websocket_passthrough_request( endpoint: The endpoint path (for logging purposes) cost_per_request: Optional field - cost per request to the target endpoint setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream + client_frame_model_gate: Optional per-frame check of every model a client frame names + upstream_frame_kind: Optional sort of upstream frames into the one that ends a billable turn and the ones + that open the session, so each turn logs its own spend and keeps the session's provider-named model + walkout_turn_drain_seconds: How long to keep reading a turn the provider is still running after the client + left or was refused, so the provider's usage for it still reaches spend builtin_provider_route: Whether a built-in provider route serves the connection, so its URL names the provider model: The model the caller named before connecting, logged on a failure before the setup frame arrives meter: Prices the session from the client text frames delivered upstream, for an upstream that reports no usage @@ -2709,26 +3082,31 @@ async def websocket_passthrough_request( proxy_config=proxy_config, route_description="websocket_passthrough", ) - logging_obj: Final = Logging( - model="unknown", - messages=[{"role": "user", "content": "WebSocket connection"}], - stream=True, # WebSockets are inherently streaming - call_type="pass_through_endpoint", - start_time=start_time, - litellm_call_id=litellm_call_id, - function_id="websocket_passthrough", - dynamic_success_callbacks=team_callbacks.success_callbacks, - dynamic_failure_callbacks=team_callbacks.failure_callbacks, - kwargs=team_callbacks.logging_kwargs, - ) - # Create passthrough logging payload - passthrough_logging_payload: Final = PassthroughStandardLoggingPayload( - url=target, - request_body={}, # WebSocket doesn't have a traditional request body - request_method="WEBSOCKET", - cost_per_request=cost_per_request, - ) + def new_call_logging(call_start: datetime, call_id: str) -> LiteLLMLoggingObj: + return Logging( + model="unknown", + messages=[{"role": "user", "content": "WebSocket connection"}], + stream=True, + call_type="pass_through_endpoint", + start_time=call_start, + litellm_call_id=call_id, + function_id="websocket_passthrough", + dynamic_success_callbacks=team_callbacks.success_callbacks, + dynamic_failure_callbacks=team_callbacks.failure_callbacks, + kwargs=team_callbacks.logging_kwargs, + ) + + def new_call_payload() -> PassthroughStandardLoggingPayload: + return PassthroughStandardLoggingPayload( + url=target, + request_body={}, + request_method="WEBSOCKET", + cost_per_request=cost_per_request, + ) + + logging_obj: Final = new_call_logging(start_time, litellm_call_id) + passthrough_logging_payload: Final = new_call_payload() # Create a dummy request object for WebSocket connections to maintain compatibility # with the existing _init_kwargs_for_pass_through_endpoint function @@ -2750,38 +3128,136 @@ async def websocket_passthrough_request( ) # Initialize kwargs for logging using the same pattern as HTTP passthrough - kwargs: Final = HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint( - user_api_key_dict=user_api_key_dict, - _parsed_body={}, # WebSocket doesn't have a traditional request body - passthrough_logging_payload=passthrough_logging_payload, - litellm_call_id=litellm_call_id, - request=dummy_request, - logging_obj=logging_obj, + def init_call_kwargs( + token: UserAPIKeyAuth, + call_payload: PassthroughStandardLoggingPayload, + call_id: str, + call_logging: LiteLLMLoggingObj, + ) -> dict[str, object]: # mutable-ok: the Vertex relay stamps the model from setup frames into these kwargs + return HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint( + user_api_key_dict=token, + _parsed_body={}, # WebSocket doesn't have a traditional request body + passthrough_logging_payload=call_payload, + litellm_call_id=call_id, + request=dummy_request, + logging_obj=call_logging, + ) + + kwargs: Final = init_call_kwargs(user_api_key_dict, passthrough_logging_payload, litellm_call_id, logging_obj) + + def start_call_logging(call_logging: LiteLLMLoggingObj, call_kwargs: Mapping[str, object]) -> None: + call_logging.update_environment_variables( + model="unknown", + user="unknown", + optional_params={}, + litellm_params=dict(as_str_mapping(call_kwargs.get("litellm_params")) or {}), + call_type="pass_through_endpoint", + ) + call_logging.pre_call( + input=[{"role": "user", "content": "WebSocket connection"}], + api_key="", + additional_args={ + "complete_input_dict": {}, + "api_base": target, + "headers": upstream_headers, + }, + ) + + start_call_logging(logging_obj, kwargs) + + class MockWebSocketRequest: + def __init__(self, target_url: str): + self.method = "WEBSOCKET" + self.url = target_url + + class MockWebSocketResponse: + def __init__(self, target_url: str): + self.status_code = 200 + self.text = "WebSocket connection successful" + self.headers: dict[str, str] = {} + self.request = MockWebSocketRequest(target_url) + + ledger: Final = _WebsocketTurnLedger( + opening=_TurnAuthorization( + requested_model=_opening_model(target), + access_groups=tuple(user_api_key_dict.matched_model_access_groups or ()), + ), + started=start_time, ) - # Update logging environment variables - logging_obj.update_environment_variables( - model="unknown", - user="unknown", - optional_params={}, - litellm_params=dict(kwargs.get("litellm_params", {})), - call_type="pass_through_endpoint", - ) - logging_obj.model_call_details["litellm_call_id"] = litellm_call_id + def log_frames( + call_logging: LiteLLMLoggingObj, + call_payload: PassthroughStandardLoggingPayload, + call_kwargs: Mapping[str, object], + end_time: datetime, + lane: str | None, + in_row: Callable[[Mapping[str, object]], bool], + ) -> None: + all_frames: Final = list(websocket_messages) + frames: Final = [frame for frame in all_frames if in_row(frame)] + websocket_messages[:] = ( + () + if upstream_frame_kind is None + else ( + frame + for frame in all_frames + if not in_row(frame) or upstream_frame_kind(frame) in _SESSION_FRAMES_EVERY_ROW_CARRIES + ) + ) + logged_payload: Final[PassthroughStandardLoggingPayload] = { + **call_payload, + "response_body": frames, + "end_time": end_time, + } + if client_frame_model_gate is not None: + _stamp_billed_turn(call_kwargs, ledger.current(lane)) + success_kwargs: Final = { + key: value for key, value in call_kwargs.items() if key not in _EXPLICIT_SUCCESS_KWARGS + } + custom_llm_provider: Final = call_kwargs.get("custom_llm_provider") + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( + async_coroutine=pass_through_endpoint_logging.pass_through_async_success_handler( + httpx_response=MockWebSocketResponse(target), + response_body=frames, + url_route=endpoint or "", + result="websocket_connection_successful", + start_time=ledger.started, + end_time=end_time, + logging_obj=call_logging, + cache_hit=False, + request_body={}, + passthrough_logging_payload=logged_payload, + custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None, + finalize_before_dispatch=False, + **success_kwargs, + ) + ) + bind_budget_reservation_to_callbacks(call_logging.litellm_params) - # Pre-call logging - logging_obj.pre_call( - input=[{"role": "user", "content": "WebSocket connection"}], - api_key="", - additional_args={ - "complete_input_dict": {}, - "api_base": target, - "headers": upstream_headers, - }, - ) + def log_turn(lane: str | None, in_row: Callable[[Mapping[str, object]], bool]) -> None: + end_time: Final = datetime.now() + if upstream_frame_kind is None: + log_frames(logging_obj, passthrough_logging_payload, kwargs, end_time, lane, in_row) + else: + turn_call_id: Final = str(uuid.uuid4()) + turn_logging: Final = new_call_logging(ledger.started, turn_call_id) + turn_payload: Final = new_call_payload() + turn_kwargs: Final = init_call_kwargs( + user_api_key_dict + if ledger.logged_turns == 0 + else user_api_key_dict.model_copy(update={"budget_reservation": None}), + turn_payload, + turn_call_id, + turn_logging, + ) + start_call_logging(turn_logging, turn_kwargs) + log_frames(turn_logging, turn_payload, turn_kwargs, end_time, lane, in_row) + ledger.close_turn(lane, end_time) ### CALL HOOKS ### - modify incoming data / reject request before calling the model - websocket_hook_data: Final = guardrail_request_data_with_streaming(MappingProxyType({}), is_streaming=True) + websocket_hook_data: Final = guardrail_request_data_with_streaming( + MappingProxyType({"litellm_call_id": litellm_call_id}), is_streaming=True + ) await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=websocket_hook_data, @@ -2808,7 +3284,6 @@ async def websocket_passthrough_request( message = await websocket.receive() message_type = message.get("type") if message_type == "websocket.disconnect": - await upstream_ws.close() break text_data: str | None = message.get("text") @@ -2871,11 +3346,29 @@ async def websocket_passthrough_request( ) # Not a JSON message or doesn't contain setup data - await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter)) - if meter is not None: - metered_total += meter.measure(text_data) + if isinstance( + await _relay_client_frame( + websocket, + upstream_ws, + text_data, + _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), + client_frame_model_gate, + ledger, + ), + _RefuseClientFrame, + ): + break + metered_total += _metered_client_text(meter, text_data) elif bytes_data is not None: - await upstream_ws.send(bytes_data) + if isinstance( + await _relay_client_frame( + websocket, upstream_ws, bytes_data, bytes_data, client_frame_model_gate, ledger + ), + _RefuseClientFrame, + ): + break + if not _turn_owed(upstream_frame_kind, websocket_messages, ledger): + await upstream_ws.close() except asyncio.CancelledError: raise except Exception: @@ -2902,11 +3395,7 @@ async def websocket_passthrough_request( is_vertex_live: Final = bool(endpoint and "/vertex_ai/live" in endpoint) json_frame_ordinal: Final = count() - async def relay_upstream_frame(upstream_message: str | bytes) -> None: - if isinstance(upstream_message, bytes): - await websocket.send_bytes(upstream_message) - else: - await websocket.send_text(upstream_message) + def record_upstream_frame(upstream_message: str | bytes) -> None: message_data: Final = _json_object_frame(upstream_message) if message_data is None: return @@ -2914,6 +3403,43 @@ async def websocket_passthrough_request( _extract_vertex_live_model_from_setup_response(message_data) return websocket_messages.append(message_data) + match None if upstream_frame_kind is None else upstream_frame_kind(message_data): + case "opens_session": + ledger.open_session() + case "starts_turn": + ledger.start_turn() + case "updates_session" if upstream_frame_kind is not None: + websocket_messages[:] = _without_earlier_session_updates( + upstream_frame_kind, websocket_messages, message_data + ) + case "ends_turn" if upstream_frame_kind is not None: + if _frame_lane(message_data) not in _lanes_in_flight(upstream_frame_kind, websocket_messages): + ledger.start_turn() + log_turn(_frame_lane(message_data), _frames_of_lane(_frame_lane(message_data))) + case "bills_usage" if not _turn_in_flight(upstream_frame_kind, websocket_messages): + log_turn(_frame_lane(message_data), _frames_of_lane(_frame_lane(message_data))) + case "refuses_turn" if upstream_frame_kind is not None: + lane: Final = _frame_lane(message_data) + ledger.drop_refused_turn( + lane, lane in _lanes_in_flight(upstream_frame_kind, websocket_messages) + ) + case _: + pass + + async def relay_upstream_frame(upstream_message: str | bytes) -> None: + record_upstream_frame(upstream_message) + if not _client_socket_is_open(websocket): + if not _turn_owed(upstream_frame_kind, websocket_messages, ledger): + await upstream_ws.close() + return + try: + if isinstance(upstream_message, bytes): + await websocket.send_bytes(upstream_message) + else: + await websocket.send_text(upstream_message) + except WebSocketDisconnect: + if not _turn_owed(upstream_frame_kind, websocket_messages, ledger): + raise async def forward_upstream_to_client() -> Close | None: try: @@ -2922,6 +3448,11 @@ async def websocket_passthrough_request( except (ConnectionClosedOK, ConnectionClosedError) as e: verbose_proxy_logger.debug("Upstream WebSocket connection closed: %s", e) return e.rcvd + except WebSocketDisconnect: + verbose_proxy_logger.debug( + "WebSocket passthrough (%s): client left while upstream was still sending", endpoint + ) + return None except asyncio.CancelledError: verbose_proxy_logger.debug("asyncio.CancelledError in forward_upstream_to_client") raise @@ -2933,12 +3464,15 @@ async def websocket_passthrough_request( raise # Create tasks for bidirectional message forwarding - tasks: Final = [ - asyncio.create_task(forward_client_to_upstream()), - asyncio.create_task(forward_upstream_to_client()), - ] + client_task: Final = asyncio.create_task(forward_client_to_upstream()) + upstream_task: Final = asyncio.create_task(forward_upstream_to_client()) + tasks: Final = (client_task, upstream_task) - done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + if client_task.done() and _turn_owed(upstream_frame_kind, websocket_messages, ledger): + await asyncio.wait((upstream_task,), timeout=walkout_turn_drain_seconds) + done: Final = tuple(task for task in tasks if task.done()) + pending: Final = tuple(task for task in tasks if not task.done()) # Cancel remaining tasks for task in pending: @@ -2961,64 +3495,17 @@ async def websocket_passthrough_request( reason=truncated_close_reason(upstream_close.reason), ) - end_time: Final = datetime.now() - - # Update passthrough logging payload with response data - passthrough_logging_payload["response_body"] = websocket_messages - passthrough_logging_payload["end_time"] = end_time - - # Remove logging_obj from kwargs to avoid duplicate keyword argument - success_kwargs: Final = kwargs.copy() - success_kwargs.pop("logging_obj", None) - - # # Add user authentication context for database logging - # if user_api_key_dict: - # success_kwargs.setdefault('litellm_params', {}) - # success_kwargs['litellm_params'].update({ - # 'proxy_server_request': { - # 'body': { - # 'user': user_api_key_dict.user_id, - # 'team_id': user_api_key_dict.team_id, - # 'end_user_id': user_api_key_dict.end_user_id, - # } - # } - # }) - # # Also add the user_api_key for direct access - # success_kwargs['user_api_key'] = user_api_key_dict.api_key - - # Create a dummy httpx.Response for WebSocket connections - class MockWebSocketResponse: - def __init__(self, target_url: str): - self.status_code = 200 - self.text = "WebSocket connection successful" - self.headers: dict[str, str] = {} - self.request = MockWebSocketRequest(target_url) - - class MockWebSocketRequest: - def __init__(self, target_url: str): - self.method = "WEBSOCKET" - self.url = target_url - - mock_response: Final = MockWebSocketResponse(target) if meter is not None: _log_session_usage(logging_obj, meter.price(metered_total)) - # Use the same success handler as HTTP passthrough endpoints - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - async_coroutine=pass_through_endpoint_logging.pass_through_async_success_handler( - httpx_response=mock_response, - response_body=websocket_messages, - url_route=endpoint or "", - result="websocket_connection_successful", - start_time=start_time, - end_time=end_time, - logging_obj=logging_obj, - cache_hit=False, - request_body={}, - **success_kwargs, + if _session_end_needs_a_row(upstream_frame_kind, websocket_messages, ledger): + lanes_in_flight: Final = ( + () if upstream_frame_kind is None else _lanes_in_flight(upstream_frame_kind, websocket_messages) ) - ) - bind_budget_reservation_to_callbacks(logging_obj.litellm_params) + *earlier_lanes, last_lane = lanes_in_flight or (None,) + for lane in earlier_lanes: + log_turn(lane, _frames_of_lane(lane)) + log_turn(last_lane, _every_frame) # Call the proxy logging success hook if proxy_logging_obj: @@ -3076,6 +3563,7 @@ async def websocket_passthrough_request( if _client_socket_is_open(websocket): await websocket.close(code=1011, reason="WebSocket passthrough error") finally: + await proxy_logging_obj.arelease_responses_websocket_slot(user_api_key_dict) if _client_socket_is_open(websocket): await websocket.close() diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 8b09b1cf7cd..8f55906498b 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -1,5 +1,6 @@ import asyncio import json +from collections.abc import Mapping, Sequence from contextlib import suppress from datetime import datetime from types import MappingProxyType @@ -13,7 +14,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import AZURE_SPEECH_CUSTOM_LLM_PROVIDER from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.proxy._types import PassThroughEndpointLoggingResultValues +from litellm.proxy._types import PassThroughEndpointLoggingResultValues, PassThroughEndpointLoggingTypedDict from litellm.types.passthrough_endpoints.pass_through_endpoints import ( PassthroughStandardLoggingPayload, ) @@ -70,6 +71,14 @@ def _safe_response_text(httpx_response: httpx.Response) -> str: return "" +def _websocket_frames( + response_body: Mapping[str, object] | Sequence[Mapping[str, object]] | None, +) -> tuple[Mapping[str, object], ...]: + return tuple( + message for message in (response_body if isinstance(response_body, list) else ()) if isinstance(message, dict) + ) + + class PassThroughEndpointLogging: def __init__( self, @@ -120,6 +129,7 @@ class PassThroughEndpointLogging: # Vertex AI Live API WebSocket self.TRACKED_VERTEX_AI_LIVE_ROUTES = ["/vertex_ai/live"] + self.TRACKED_OPENAI_WEBSOCKET_ROUTES: Final = ("/openai/", "/openai_passthrough/") @property def _log_dispatch(self) -> PassThroughLogDispatch: @@ -422,21 +432,18 @@ class PassThroughEndpointLogging: standard_logging_response_object = vertex_ai_live_handler_result["result"] kwargs = vertex_ai_live_handler_result["kwargs"] - elif DeepgramListenPassthroughLoggingHandler.is_deepgram_listen_route(url_route): - deepgram_handler_result: Final = ( - DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( - websocket_messages=tuple( - message - for message in (response_body if isinstance(response_body, list) else ()) - if isinstance(message, dict) - ), - logging_obj=logging_obj, - upstream_url=str(httpx_response.request.url), - kwargs=kwargs, - ) + elif ( + websocket_relay_result := self._websocket_relay_logging_result( + url_route=url_route, + websocket_frames=_websocket_frames(response_body), + logging_obj=logging_obj, + httpx_response=httpx_response, + start_time=start_time, + kwargs=kwargs, ) - standard_logging_response_object = deepgram_handler_result["result"] # rebind-ok: elif-chain - kwargs = deepgram_handler_result["kwargs"] # rebind-ok: elif-chain contract + ) is not None: + standard_logging_response_object = websocket_relay_result["result"] # rebind-ok: elif-chain + kwargs = websocket_relay_result["kwargs"] # rebind-ok: elif-chain contract elif FalAIPassthroughLoggingHandler.is_fal_ai_route(url_route, custom_llm_provider): fal_ai_handler_result: Final = FalAIPassthroughLoggingHandler().fal_ai_passthrough_handler( response_body=response_body if isinstance(response_body, dict) else MappingProxyType({}), @@ -452,6 +459,36 @@ class PassThroughEndpointLogging: return_dict["kwargs"] = kwargs return return_dict + def _websocket_relay_logging_result( + self, + url_route: str, + websocket_frames: tuple[Mapping[str, object], ...], + logging_obj: LiteLLMLoggingObj, + httpx_response: httpx.Response, + start_time: datetime, + kwargs: Mapping[str, object], + ) -> PassThroughEndpointLoggingTypedDict | None: + if self.is_openai_websocket_route(url_route): + from .llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + OpenAIWebsocketPassthroughLoggingHandler, + ) + + return OpenAIWebsocketPassthroughLoggingHandler().openai_websocket_passthrough_handler( + websocket_messages=websocket_frames, + logging_obj=logging_obj, + url_route=url_route, + start_time=start_time, + kwargs=kwargs, + ) + if DeepgramListenPassthroughLoggingHandler.is_deepgram_listen_route(url_route): + return DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler( + websocket_messages=websocket_frames, + logging_obj=logging_obj, + upstream_url=str(httpx_response.request.url), + kwargs=kwargs, + ) + return None + async def pass_through_async_success_handler( self, httpx_response: httpx.Response, @@ -625,6 +662,9 @@ class PassThroughEndpointLogging: return True return False + def is_openai_websocket_route(self, url_route: str) -> bool: + return url_route.startswith(self.TRACKED_OPENAI_WEBSOCKET_ROUTES) + def is_cursor_route(self, url_route: str, custom_llm_provider: str | None = None): """Check if the URL route is a Cursor Cloud Agents API route.""" if custom_llm_provider == "cursor": diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index fc6fb266d04..dcfd6148f54 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -432,7 +432,7 @@ async def invalidate_budget_reservation_counters( async def release_or_invalidate_budget_reservation( - budget_reservation: dict | None, # mutable-ok: stamps finalized on the caller's shared reservation dict + budget_reservation: dict[str, object] | None, # mutable-ok: stamps finalized on the caller's reservation ) -> None: """Reconcile a still-open reservation on a terminal path that settles no cost. diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index 2f269dc87d9..b6bab77ab33 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -1,9 +1,11 @@ +from collections.abc import Mapping, Sequence from dataclasses import dataclass +from datetime import datetime from enum import Enum -from typing import Final, Protocol +from typing import Final, Literal, Protocol, TypeAlias from pydantic import Field -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -28,6 +30,10 @@ LITELLM_PASS_THROUGH_ENDPOINT_MARKER: Final = "__litellm_pass_through_endpoint__ SETTLES_BACKGROUND_RESPONSE_KWARG: Final = "settles_background_response" +WebsocketUpstreamFrameKind: TypeAlias = Literal[ + "starts_turn", "ends_turn", "opens_session", "updates_session", "bills_usage", "refuses_turn", "relayed" +] + class BackgroundResponseSettlement(str, Enum): SETTLES = "settles" @@ -63,7 +69,7 @@ class PassthroughStandardLoggingPayload(TypedDict, total=False): """ The body of the request """ - response_body: dict | None # only tracked for non-streaming responses + response_body: dict | Sequence[Mapping[str, object]] | None # only tracked for non-streaming responses """ The body of the response """ @@ -75,6 +81,8 @@ class PassthroughStandardLoggingPayload(TypedDict, total=False): Optional field, we use this for cost tracking only if it's set. """ + end_time: ReadOnly[datetime] + @dataclass(frozen=True, slots=True) class WebSocketSessionUsage: diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 8e9e87f370a..7a060a214f7 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -10998,139 +10998,6 @@ async def test_key_budget_error_keeps_the_masked_key_name(key_name): assert f"Key=prod-key ({key_name}) Current cost" in message -class _UntouchedPrisma: - def __getattr__(self, name: str) -> object: - raise AssertionError(f"database reached through {name}") - - -class _MissingUserPrisma: - class db: - class litellm_usertable: - @staticmethod - async def find_unique(where: dict[str, str], include: dict[str, bool]) -> None: - return None - - -@pytest.mark.asyncio -async def test_enforced_model_allowlists_treats_a_missing_user_row_as_unrestricted(): - from litellm.proxy.auth.auth_checks import enforced_model_allowlists - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache - from litellm.proxy.utils import ProxyLogging - - cache = UserApiKeyCache() - scopes = await enforced_model_allowlists( - valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="default_user_id"), - prisma_client=_MissingUserPrisma(), - user_api_key_cache=cache, - proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), - ) - - assert [list(scope) for scope in scopes] == [[], [], [], [], []] - - -class _UnreachableUserPrisma: - class db: - class litellm_usertable: - @staticmethod - async def find_unique(where: dict[str, str], include: dict[str, bool]) -> None: - raise RuntimeError("database gone") - - -@pytest.mark.asyncio -async def test_enforced_model_allowlists_surfaces_a_failed_user_lookup(): - from litellm.proxy.auth.auth_checks import enforced_model_allowlists - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache - from litellm.proxy.utils import ProxyLogging - - cache = UserApiKeyCache() - with pytest.raises(ValueError, match="database gone"): - await enforced_model_allowlists( - valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="user-fake"), - prisma_client=_UnreachableUserPrisma(), - user_api_key_cache=cache, - proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), - ) - - -@pytest.mark.asyncio -async def test_enforced_model_allowlists_reads_every_level_from_cache(): - from litellm.proxy._types import ( - LiteLLM_BudgetTable, - LiteLLM_ProjectTableCachedObj, - LiteLLM_TeamMembership, - LiteLLM_TeamTableCachedObj, - LiteLLM_UserTable, - ) - from litellm.proxy.auth.auth_checks import enforced_model_allowlists - from litellm.proxy.common_utils.user_api_key_cache import ( - UserApiKeyCache, - team_membership_reservation_cache_key, - ) - from litellm.proxy.utils import ProxyLogging - - cache = UserApiKeyCache() - proxy_logging_obj = ProxyLogging(user_api_key_cache=cache) - await cache.async_set_cache( - key="team_id:team-fake", value=LiteLLM_TeamTableCachedObj(team_id="team-fake", models=["gpt-4o"]) - ) - await cache.async_set_cache( - key=team_membership_reservation_cache_key(user_id="user-fake", team_id="team-fake"), - value=LiteLLM_TeamMembership( - user_id="user-fake", - team_id="team-fake", - litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["gpt-4o-mini"]), - ), - ) - await cache.async_set_cache( - key="project_id:project-fake", - value=LiteLLM_ProjectTableCachedObj(project_id="project-fake", models=["gpt-4.1"]), - ) - await cache.async_set_cache(key="user-fake", value=LiteLLM_UserTable(user_id="user-fake", models=["o3"])) - prisma_client = _UntouchedPrisma() - - team_scoped = await enforced_model_allowlists( - valid_token=UserAPIKeyAuth( - token="hashed-fake", - models=["all-team-models"], - team_models=["gpt-4o", "gpt-4o-mini"], - user_id="user-fake", - team_id="team-fake", - project_id="project-fake", - ), - prisma_client=prisma_client, - user_api_key_cache=cache, - proxy_logging_obj=proxy_logging_obj, - ) - personal = await enforced_model_allowlists( - valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="user-fake"), - prisma_client=prisma_client, - user_api_key_cache=cache, - proxy_logging_obj=proxy_logging_obj, - ) - without_database = await enforced_model_allowlists( - valid_token=UserAPIKeyAuth( - token="hashed-fake", - models=["gpt-4o"], - team_models=["gpt-4o-mini"], - user_id="user-fake", - team_id="team-fake", - ), - prisma_client=None, - user_api_key_cache=cache, - proxy_logging_obj=proxy_logging_obj, - ) - - assert [list(scope) for scope in team_scoped] == [ - ["gpt-4o", "gpt-4o-mini"], - ["gpt-4o"], - ["gpt-4o-mini"], - [], - ["gpt-4.1"], - ] - assert [list(scope) for scope in personal] == [[], [], [], ["o3"], []] - assert [list(scope) for scope in without_database] == [["gpt-4o"], ["gpt-4o-mini"]] - - @pytest.mark.asyncio @pytest.mark.parametrize("channel", ["team", "key"]) async def test_access_group_model_fallback_uses_the_injected_database(channel: str) -> None: diff --git a/tests/unit/proxy/common_utils/test_realtime_cache.py b/tests/unit/proxy/common_utils/test_realtime_cache.py index 6f5751fa923..9e76caff054 100644 --- a/tests/unit/proxy/common_utils/test_realtime_cache.py +++ b/tests/unit/proxy/common_utils/test_realtime_cache.py @@ -1,4 +1,5 @@ -from typing import Any, cast +import json +from typing import Any, Final, cast import pytest @@ -40,6 +41,12 @@ def test_realtime_request_body_caches_each_model_separately(): assert gpt4o_body_first is not gpt4o_mini_body +def test_realtime_request_body_keeps_a_model_name_with_quotes_inside_its_own_field(): + body: Final = realtime_request_body('gpt-4o", "user": "someone-else') + + assert json.loads(body) == {"model": 'gpt-4o", "user": "someone-else'} + + def test_realtime_query_params_template_caches_each_pair_separately(): params_with_intent_first = _realtime_query_params_template("gpt-4o", "intent-a") params_with_intent_second = _realtime_query_params_template("gpt-4o", "intent-a") diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py new file mode 100644 index 00000000000..e7b4d087ed3 --- /dev/null +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py @@ -0,0 +1,520 @@ +import uuid +from collections.abc import Mapping, Sequence +from datetime import datetime +from types import MappingProxyType +from typing import Final +from unittest.mock import MagicMock + +import httpx +import pytest + +import litellm +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import ( + OpenAIPassthroughLoggingHandler, +) +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + OpenAIWebsocketPassthroughLoggingHandler, + openai_websocket_frame_kind, +) +from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id +from litellm.types.utils import CallTypes + +REALTIME_MODEL: Final = "gpt-realtime-2.1" +RESPONSES_MODEL: Final = "gpt-5.4-mini" +REASONING_MODEL: Final = "gpt-5.4" +TRANSCRIPTION_MODEL: Final = "gpt-4o-mini-transcribe" +SESSION_START: Final = datetime(2026, 1, 1, 12, 0, 0) + + +def _logging_obj() -> MagicMock: + logging_obj: Final = MagicMock(model_call_details={}, litellm_call_id=str(uuid.uuid4())) + return logging_obj + + +def _session_created(model: str) -> dict[str, object]: + return {"type": "session.created", "session": {"id": "sess_fake", "model": model}} + + +def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_out: int) -> dict[str, object]: + return { + "type": "response.done", + "response": { + "id": "resp_fake", + "status": "completed", + "usage": { + "input_tokens": text_in + audio_in, + "output_tokens": text_out + audio_out, + "total_tokens": text_in + audio_in + text_out + audio_out, + "input_token_details": {"text_tokens": text_in, "audio_tokens": audio_in, "cached_tokens": 0}, + "output_token_details": {"text_tokens": text_out, "audio_tokens": audio_out}, + }, + }, + } + + +def _responses_completed( + model: str, + input_tokens: int, + output_tokens: int, + event: str = "response.completed", + served: Mapping[str, object] = MappingProxyType({}), +) -> dict[str, object]: + return { + "type": event, + "response": { + "id": "resp_fake", + "object": "response", + "created_at": 1761700000, + "model": model, + "status": "completed", + "output": [], + **served, + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + }, + } + + +def _handle( + messages: Sequence[Mapping[str, object]], + kwargs: Mapping[str, object] = MappingProxyType({}), + logging_obj: MagicMock | None = None, + url_route: str = "/openai/v1/realtime", +): + return OpenAIWebsocketPassthroughLoggingHandler().openai_websocket_passthrough_handler( + websocket_messages=messages, + logging_obj=logging_obj if logging_obj is not None else _logging_obj(), + url_route=url_route, + start_time=SESSION_START, + kwargs=kwargs, + ) + + +def _realtime_cost(model: str, text_in: int, audio_in: int, text_out: int, audio_out: int) -> float: + prices: Final = litellm.model_cost[model] + return ( + text_in * prices["input_cost_per_token"] + + audio_in * prices["input_cost_per_audio_token"] + + text_out * prices["output_cost_per_token"] + + audio_out * prices["output_cost_per_audio_token"] + ) + + +def test_realtime_session_is_priced_on_the_model_the_session_reported(): + logging_obj: Final = _logging_obj() + messages: Final = [ + _session_created(REALTIME_MODEL), + {"type": "response.created", "response": {"id": "resp_fake"}}, + _realtime_response_done(text_in=33, audio_in=120, text_out=16, audio_out=240), + ] + + handled: Final = _handle(messages, logging_obj=logging_obj) + + expected_cost: Final = _realtime_cost(REALTIME_MODEL, 33, 120, 16, 240) + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["kwargs"]["model"] == REALTIME_MODEL + assert handled["kwargs"]["custom_llm_provider"] == "openai" + assert handled["result"].model == REALTIME_MODEL + assert handled["result"].usage.total_tokens == 409 + assert logging_obj.model_call_details["response_cost"] == pytest.approx(expected_cost) + assert logging_obj.model_call_details["model"] == REALTIME_MODEL + + +def test_realtime_cost_covers_every_turn_of_the_session(): + one_turn: Final = _handle([_session_created(REALTIME_MODEL), _realtime_response_done(10, 0, 5, 0)]) + two_turns: Final = _handle( + [ + _session_created(REALTIME_MODEL), + _realtime_response_done(10, 0, 5, 0), + _realtime_response_done(10, 0, 5, 0), + ] + ) + + assert two_turns["kwargs"]["response_cost"] == pytest.approx(2 * one_turn["kwargs"]["response_cost"]) + assert two_turns["result"].usage.total_tokens == 30 + + +def test_responses_turn_is_priced_on_the_model_the_response_reported(): + handled: Final = _handle( + [_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)], + url_route="/openai_passthrough/v1/responses", + ) + + prices: Final = litellm.model_cost[RESPONSES_MODEL] + expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["kwargs"]["model"] == RESPONSES_MODEL + assert handled["result"].usage.prompt_tokens == 13 + assert handled["result"].usage.completion_tokens == 5 + + +def test_responses_cost_covers_every_turn_of_the_connection(): + handled: Final = _handle( + [ + _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5), + _responses_completed(RESPONSES_MODEL, input_tokens=20, output_tokens=7), + ], + url_route="/openai_passthrough/v1/responses", + ) + + prices: Final = litellm.model_cost[RESPONSES_MODEL] + expected_cost: Final = 33 * prices["input_cost_per_token"] + 12 * prices["output_cost_per_token"] + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["result"].usage.total_tokens == 45 + + +def test_a_model_the_caller_asked_for_never_prices_a_session_the_provider_never_named(): + """ + The caller picks what to ask for, and the provider picks what to run. Only + the provider's own frames may set the price of a session, so a model the + caller put in the request cannot stand in for one the provider never named. + """ + logging_obj: Final = _logging_obj() + + handled: Final = _handle( + [_realtime_response_done(text_in=10, audio_in=0, text_out=5, audio_out=0)], + kwargs=MappingProxyType({"model": REALTIME_MODEL}), + logging_obj=logging_obj, + ) + + assert handled["result"] is None + assert handled["kwargs"] == {"model": REALTIME_MODEL} + assert logging_obj.model_call_details == {} + + +def test_each_responses_turn_is_priced_on_the_model_that_turn_ran(): + """ + One socket can run several models one after another, so a connection that + starts on a big model and finishes on a small one has to pay the big + model's price for the turn that used it. + """ + handled: Final = _handle( + [ + _responses_completed(REASONING_MODEL, input_tokens=13, output_tokens=5), + _responses_completed(RESPONSES_MODEL, input_tokens=20, output_tokens=7), + ], + url_route="/openai_passthrough/v1/responses", + ) + + first: Final = litellm.model_cost[REASONING_MODEL] + second: Final = litellm.model_cost[RESPONSES_MODEL] + expected_cost: Final = ( + 13 * first["input_cost_per_token"] + + 5 * first["output_cost_per_token"] + + 20 * second["input_cost_per_token"] + + 7 * second["output_cost_per_token"] + ) + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["result"].usage.total_tokens == 45 + + +def test_a_responses_turn_that_names_no_model_is_priced_on_the_session_model(): + handled: Final = _handle( + [ + _session_created(RESPONSES_MODEL), + _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served={"model": None}), + ], + url_route="/openai_passthrough/v1/responses", + ) + + prices: Final = litellm.model_cost[RESPONSES_MODEL] + expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + + +def test_a_session_nobody_named_a_model_for_is_left_unpriced(): + logging_obj: Final = _logging_obj() + handled: Final = _handle( + [{"type": "response.done", "response": {"usage": {"input_tokens": 10, "output_tokens": 5}}}], + kwargs=MappingProxyType({"litellm_call_id": "call-fake"}), + logging_obj=logging_obj, + ) + + assert handled["result"] is None + assert handled["kwargs"] == {"litellm_call_id": "call-fake"} + assert logging_obj.model_call_details == {} + + +def test_a_session_that_reported_no_usage_still_records_its_model(): + logging_obj: Final = _logging_obj() + + handled: Final = _handle([_session_created(REALTIME_MODEL)], logging_obj=logging_obj) + + assert handled["result"] is None + assert handled["kwargs"] == {"model": REALTIME_MODEL, "custom_llm_provider": "openai"} + assert logging_obj.model_call_details["model"] == REALTIME_MODEL + assert logging_obj.model_call_details["custom_llm_provider"] == "openai" + + +def test_a_responses_turn_without_readable_usage_is_not_billed(): + handled: Final = _handle( + [{"type": "response.completed", "response": {"model": RESPONSES_MODEL, "status": "completed"}}], + url_route="/openai_passthrough/v1/responses", + ) + + assert handled["result"] is None + assert handled["kwargs"] == {"model": RESPONSES_MODEL, "custom_llm_provider": "openai"} + + +def test_the_price_of_a_multi_model_connection_reaches_the_spend_row(): + """ + The logging layer prices whatever result it is handed, and one result can + only name one model, so a connection that ran two models has to carry the + price it already worked out or the cheaper model pays for both turns. The + row it lands in names the model the connection opened on. + """ + handled: Final = _handle( + [ + _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5), + _responses_completed(REASONING_MODEL, input_tokens=20, output_tokens=7), + ], + url_route="/openai_passthrough/v1/responses", + ) + + logged_cost: Final = litellm.response_cost_calculator( + response_object=handled["result"], + model=handled["kwargs"]["model"], + custom_llm_provider="openai", + call_type="completion", + optional_params={}, + ) + + session_prices: Final = litellm.model_cost[RESPONSES_MODEL] + priced_on_the_session_model: Final = ( + 33 * session_prices["input_cost_per_token"] + 12 * session_prices["output_cost_per_token"] + ) + assert logged_cost == pytest.approx(handled["kwargs"]["response_cost"]) + assert logged_cost > priced_on_the_session_model + assert handled["kwargs"]["model"] == RESPONSES_MODEL + + +def test_turns_that_start_in_the_same_instant_each_keep_their_own_spend_row(): + """ + Spend logs key each row by the id its result carries and skip a repeated id + without a word, so two turns that start in the same instant, on one socket + or on two, each need an id of their own or one of them is never billed. + """ + turns: Final = [ + _handle( + [_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)], + url_route="/openai_passthrough/v1/responses", + ) + for _ in range(2) + ] + + spend_log_ids: Final = { + get_spend_logs_id(CallTypes.pass_through.value, turn["result"].model_dump(), turn["kwargs"]) for turn in turns + } + assert len(spend_log_ids) == 2 + assert None not in spend_log_ids + + +def test_a_turn_the_caller_walked_out_of_still_names_the_model_it_ran(): + """ + A client that closes before response.completed leaves no usage behind to + bill, but response.created already named the model, so the row says which + model answered instead of falling back to unknown. + """ + logging_obj: Final = _logging_obj() + + handled: Final = _handle( + [{"type": "response.created", "response": {"id": "resp_fake", "model": REASONING_MODEL}}], + logging_obj=logging_obj, + url_route="/openai_passthrough/v1/responses", + ) + + assert handled["result"] is None + assert handled["kwargs"] == {"model": REASONING_MODEL, "custom_llm_provider": "openai"} + assert logging_obj.model_call_details["model"] == REASONING_MODEL + + +def _transcription_session(model: str) -> tuple[dict[str, object], ...]: + return ( + {"type": "session.created", "session": {"type": "transcription", "audio": {"input": {"transcription": None}}}}, + { + "type": "session.updated", + "session": {"type": "transcription", "audio": {"input": {"transcription": {"model": model}}}}, + }, + ) + + +def _transcription_completed(usage: Mapping[str, object]) -> dict[str, object]: + return { + "type": "conversation.item.input_audio_transcription.completed", + "item_id": "item_fake", + "transcript": "hello", + "usage": usage, + } + + +@pytest.mark.parametrize( + ("usage", "billed_units"), + [ + ({"type": "duration", "seconds": 3}, ((3, "input_cost_per_second"),)), + ( + { + "type": "tokens", + "input_tokens": 30, + "output_tokens": 11, + "input_token_details": {"text_tokens": 0, "audio_tokens": 30}, + }, + ((30, "input_cost_per_audio_token"), (11, "output_cost_per_token")), + ), + ], +) +def test_a_transcription_session_is_billed_for_every_clip_it_transcribed(usage, billed_units): + handled: Final = _handle( + [ + *_transcription_session(TRANSCRIPTION_MODEL), + _transcription_completed(usage), + _transcription_completed(usage), + ], + url_route="/openai_passthrough/v1/realtime", + ) + + logged_cost: Final = litellm.response_cost_calculator( + response_object=handled["result"], + model=handled["kwargs"]["model"], + custom_llm_provider="openai", + call_type="completion", + optional_params={}, + ) + + prices: Final = litellm.model_cost[TRANSCRIPTION_MODEL] + expected_cost: Final = 2 * sum(units * prices[rate] for units, rate in billed_units) + assert expected_cost > 0 + assert handled["kwargs"]["model"] == TRANSCRIPTION_MODEL + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert logged_cost == pytest.approx(expected_cost) + + +def test_a_conversation_that_also_transcribes_the_caller_is_logged_on_its_own_model(): + handled: Final = _handle( + [ + { + "type": "session.created", + "session": { + "model": REALTIME_MODEL, + "audio": {"input": {"transcription": {"model": TRANSCRIPTION_MODEL}}}, + }, + }, + _realtime_response_done(text_in=10, audio_in=0, text_out=5, audio_out=0), + ] + ) + + assert handled["kwargs"]["model"] == REALTIME_MODEL + + +@pytest.mark.parametrize("event", ["response.incomplete", "response.failed"]) +def test_a_responses_turn_that_ended_early_is_billed_for_the_tokens_it_used(event): + """ + A turn cut short by max_output_tokens or a provider failure still reports + the tokens it burned, and OpenAI bills them, so the row has to as well. + """ + handled: Final = _handle( + [_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, event=event)], + url_route="/openai_passthrough/v1/responses", + ) + + prices: Final = litellm.model_cost[RESPONSES_MODEL] + expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["result"].usage.total_tokens == 18 + + +def test_a_failed_responses_turn_that_reported_no_usage_is_not_billed(): + handled: Final = _handle( + [{"type": "response.failed", "response": {"model": RESPONSES_MODEL, "status": "failed", "usage": None}}], + url_route="/openai_passthrough/v1/responses", + ) + + assert handled["result"] is None + assert handled["kwargs"] == {"model": RESPONSES_MODEL, "custom_llm_provider": "openai"} + + +@pytest.mark.parametrize( + ("turn_model", "requested_model", "priced_as"), + [ + ("gpt-unreleased-qa-snapshot", RESPONSES_MODEL, RESPONSES_MODEL), + ("gpt-5.4-mini-qa-preview", RESPONSES_MODEL, RESPONSES_MODEL), + (RESPONSES_MODEL, REASONING_MODEL, RESPONSES_MODEL), + ], + ids=["unmapped_model", "model_the_cost_map_prices_at_zero", "priced_turn_model"], +) +def test_a_responses_turn_falls_back_to_the_requested_model_only_when_its_own_has_no_price( + turn_model: str, requested_model: str, priced_as: str +): + """ + The provider can answer under a model name the cost map does not price + yet. That turn still ran the model the key was allowed to ask for, so it is + billed at that model's price instead of $0, while a turn whose own model + has a price keeps it. + """ + handled: Final = _handle( + [_responses_completed(turn_model, input_tokens=13, output_tokens=5)], + kwargs=MappingProxyType({"litellm_params": {"metadata": {"model_group": requested_model}}}), + url_route="/openai_passthrough/v1/responses", + ) + + prices: Final = litellm.model_cost[priced_as] + expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] + assert expected_cost > 0 + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["kwargs"]["model"] == turn_model + + +def _http_passthrough_cost(response: Mapping[str, object]) -> float: + http_response: Final = httpx.Response( + 200, json=response, request=httpx.Request("POST", "https://api.openai.com/v1/responses") + ) + _, cost = OpenAIPassthroughLoggingHandler._build_responses_api_response_and_cost( + RESPONSES_MODEL, http_response, _logging_obj(), "openai" + ) + return cost + + +@pytest.mark.parametrize( + "served", + [ + {"service_tier": "priority"}, + {"output": [{"type": "web_search_call", "id": "ws_fake", "status": "completed", "action": {"type": "search"}}]}, + ], + ids=["priority_tier", "hosted_web_search"], +) +def test_a_responses_turn_costs_what_the_http_passthrough_charges_for_the_same_response(served): + """ + The provider bills a turn by the tier that served it and the hosted tools it + called, not only its tokens, so a websocket turn has to cost exactly what the + same response costs over HTTP, never the cheaper token-only price. + """ + turn: Final = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served=served) + plain_turn: Final = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5) + + handled: Final = _handle([turn], url_route="/openai_passthrough/v1/responses") + + assert handled["kwargs"]["response_cost"] == pytest.approx(_http_passthrough_cost(turn["response"])) + assert handled["kwargs"]["response_cost"] > _http_passthrough_cost(plain_turn["response"]) + + +@pytest.mark.parametrize( + ("frame_type", "kind"), + [ + ("response.completed", "ends_turn"), + ("response.incomplete", "ends_turn"), + ("response.failed", "ends_turn"), + ("response.done", "ends_turn"), + ("session.created", "opens_session"), + ("session.updated", "updates_session"), + ("conversation.item.input_audio_transcription.completed", "bills_usage"), + ("response.created", "starts_turn"), + ("response.output_text.delta", "relayed"), + ("error", "refuses_turn"), + ], +) +def test_each_frame_is_classified_by_what_it_does_to_a_turn(frame_type, kind): + assert openai_websocket_frame_kind({"type": frame_type}) == kind diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 4fe358c38a3..824f04aee75 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -1,5 +1,6 @@ import asyncio import base64 +import contextvars import gzip import json import logging @@ -7,20 +8,20 @@ import os import sys import uuid import zlib -from collections.abc import AsyncIterator, Callable, Mapping +from collections.abc import AsyncIterator, Callable, Mapping, Sequence from contextlib import ExitStack, contextmanager from dataclasses import dataclass from datetime import datetime from io import BytesIO from types import MappingProxyType, ModuleType, SimpleNamespace from typing import Final, Optional, cast -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch from uuid import UUID import httpx import pytest import respx -from fastapi import APIRouter, FastAPI, HTTPException, Request, Response, UploadFile +from fastapi import APIRouter, FastAPI, HTTPException, Request, Response, UploadFile, WebSocketDisconnect from fastapi.responses import StreamingResponse from fastapi.routing import APIRoute from fastapi.testclient import TestClient @@ -35,9 +36,13 @@ from litellm.constants import ANTHROPIC_COMPATIBLE_MAX_USER_ID_LENGTH, ANTHROPIC from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.parallel_request_limiter import PROXY_MaxParallelRequestsHandler +from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( @@ -46,9 +51,15 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrou from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, + STORED_PROMPT_PROBE_MODEL, + WALKOUT_TURN_DRAIN_SECONDS, HttpPassThroughEndpointHelpers, InitPassThroughEndpointHelpers, SafeRouteAdder, + WebsocketFrameModelAllowed, + WebsocketFrameModelGate, + WebsocketFrameModelRefused, + WebsocketFrameModelVerdict, _build_passthrough_failure_request_payload, _registered_pass_through_routes, _truncate_upstream_error_body, @@ -63,12 +74,15 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( set_env_variables_in_header, websocket_passthrough_request, ) +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + openai_websocket_frame_kind, +) from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token from litellm.types import utils as types_utils from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypes @@ -7382,6 +7396,7 @@ async def test_websocket_passthrough_forwards_non_ascii_first_frame(): mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.arelease_responses_websocket_slot = AsyncMock() mock_worker.ensure_initialized_and_enqueue = MagicMock( side_effect=lambda async_coroutine: async_coroutine.close() ) @@ -7434,6 +7449,7 @@ async def test_websocket_passthrough_propagates_active_trace_context( mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.arelease_responses_websocket_slot = AsyncMock() mock_worker = MagicMock() mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=lambda async_coroutine: async_coroutine.close()) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging) @@ -7500,6 +7516,7 @@ async def test_websocket_passthrough_never_forwards_caller_credentials_upstream( mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.arelease_responses_websocket_slot = AsyncMock() mock_worker = MagicMock() mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=lambda async_coroutine: async_coroutine.close()) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging) @@ -7593,10 +7610,11 @@ def _patched_websocket_passthrough_environment(upstream_ws): mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.arelease_responses_websocket_slot = AsyncMock() mock_worker.ensure_initialized_and_enqueue = MagicMock( side_effect=lambda async_coroutine: async_coroutine.close() ) - yield + yield mock_proxy_logging async def _pending_receive(): @@ -7901,6 +7919,1487 @@ async def test_websocket_passthrough_does_not_close_twice_when_success_logging_f websocket.close.assert_awaited_once_with(code=1008, reason=upstream_reason) +SUCCESS_HANDLER: Final = ( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints." + "pass_through_endpoint_logging.pass_through_async_success_handler" +) + + +KEY_MODEL_REFUSAL: Final = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']." +KEY_MODEL_CLOSE_REASON: Final = "key not allowed to access model." + + +class RecordingFrameModelGate: + def __init__( + self, + refusals: Mapping[str, str] = MappingProxyType({}), + groups: Mapping[str, tuple[str, ...]] = MappingProxyType({}), + error_type: str = "key_model_access_denied", + ): + self.refusals = refusals + self.groups = groups + self.error_type = error_type + self.checked: list[str] = [] + + async def __call__(self, websocket: object, model: str, /) -> WebsocketFrameModelVerdict: + self.checked.append(model) + refusal: Final = self.refusals.get(model) + if refusal is not None: + return WebsocketFrameModelRefused(message=refusal, error_type=self.error_type) + return WebsocketFrameModelAllowed(access_groups=self.groups.get(model, ())) + + +def _received_frame(frame: str | bytes) -> Mapping[str, str | bytes]: + if isinstance(frame, bytes): + return {"type": "websocket.receive", "bytes": frame} + return {"type": "websocket.receive", "text": frame} + + +async def _relay_client_frames( + frames: Sequence[str | bytes], + gate: WebsocketFrameModelGate | None, + user_api_key_dict: UserAPIKeyAuth | None = None, +) -> SimpleNamespace: + upstream_ws: Final = RecordingUpstreamWebSocket() + websocket: Final = _client_websocket( + AsyncMock(side_effect=[_received_frame(frame) for frame in frames] + [{"type": "websocket.disconnect"}]) + ) + + with ( + _patched_websocket_passthrough_environment(upstream_ws), + patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler, + ): + await websocket_passthrough_request( + websocket=websocket, + target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1", + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=user_api_key_dict or UserAPIKeyAuth(), + forward_headers=False, + endpoint="/openai/v1/realtime", + accept_websocket=False, + client_frame_model_gate=gate, + ) + + return SimpleNamespace(upstream=upstream_ws, client=websocket, success_handler=success_handler) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "frame, expected_model", + [ + ({"type": "response.create", "model": "gpt-5.4-mini"}, "gpt-5.4-mini"), + ({"type": "response.create", "response": {"model": "gpt-5.4"}}, "gpt-5.4"), + ({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}, "gpt-realtime-2.1"), + ( + {"type": "session.update", "session": {"input_audio_transcription": {"model": "gpt-4o-transcribe"}}}, + "gpt-4o-transcribe", + ), + ( + { + "type": "session.update", + "session": {"audio": {"input": {"transcription": {"model": "gpt-transcribe"}}}}, + }, + "gpt-transcribe", + ), + ], +) +async def test_websocket_passthrough_checks_the_model_a_client_frame_asks_for(frame, expected_model): + gate: Final = RecordingFrameModelGate() + + relayed: Final = await _relay_client_frames([json.dumps(frame)], gate) + + assert gate.checked == [expected_model] + assert [call.args[0] for call in relayed.upstream.send.await_args_list] == [json.dumps(frame)] + assert relayed.client.close.await_args_list == [call()] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_refuses_a_frame_naming_a_model_the_key_cannot_call(): + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) + frame: Final = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}) + + relayed: Final = await _relay_client_frames([frame], gate) + + assert relayed.upstream.send.await_args_list == [] + assert json.loads(relayed.client.send_text.await_args.args[0]) == { + "type": "error", + "error": {"type": "key_model_access_denied", "message": KEY_MODEL_REFUSAL}, + } + assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)] + relayed.upstream.close.assert_awaited() + + +@pytest.mark.asyncio +async def test_websocket_passthrough_closes_with_the_whole_first_sentence_of_the_refusal(): + """ + The close frame only has room for 123 bytes, so a refusal cut at that limit + reads as a truncated number. The first sentence names the reason whole, and + the error event before the close still carries every word with its type. + """ + refusal: Final = "Budget has been exceeded! Key=" + "f" * 64 + " Spend=1.23456789, Budget=1.0" + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": refusal}), error_type="budget_exceeded") + + relayed: Final = await _relay_client_frames([json.dumps({"type": "response.create", "model": "gpt-5.4"})], gate) + + assert relayed.client.close.await_args_list == [call(code=1008, reason="Budget has been exceeded!")] + assert json.loads(relayed.client.send_text.await_args.args[0])["error"] == { + "type": "budget_exceeded", + "message": refusal, + } + + +@pytest.mark.asyncio +async def test_websocket_passthrough_refuses_the_whole_frame_when_any_named_model_is_denied(): + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-4o-transcribe": KEY_MODEL_REFUSAL})) + frame: Final = json.dumps( + { + "type": "session.update", + "session": {"model": "gpt-realtime-2.1", "input_audio_transcription": {"model": "gpt-4o-transcribe"}}, + } + ) + + relayed: Final = await _relay_client_frames([frame], gate) + + assert gate.checked == ["gpt-realtime-2.1", "gpt-4o-transcribe"] + assert relayed.upstream.send.await_args_list == [] + assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_stops_relaying_after_a_refused_frame(): + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) + frames: Final = [ + json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}), + json.dumps({"type": "response.create"}), + ] + + relayed: Final = await _relay_client_frames(frames, gate) + + assert relayed.upstream.send.await_args_list == [] + assert gate.checked == ["gpt-realtime-2.1"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "frame", + [ + '{"type": "input_audio_buffer.append", "audio": "AAAA"}', + '{"type": "response.create", "response": {"modalities": ["audio"]}}', + '{"type": "session.update", "session": {"model": 42}}', + 'not json at all, but it does say "model"', + '{"type": "session.update", "session": "model"}', + ], +) +async def test_websocket_passthrough_forwards_frames_that_name_no_model(frame): + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"})) + + relayed: Final = await _relay_client_frames([frame], gate) + + assert gate.checked == [] + assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [frame] + assert relayed.client.close.await_args_list == [call()] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_never_bills_the_model_a_client_frame_named(): + """ + A frame is what the caller asked for, not what the provider ran. A caller + who names a cheap model in a frame the provider ignores must not move the + session onto that model's price, so only the provider's own frames decide + what a session is billed as. + """ + gate: Final = RecordingFrameModelGate() + frames: Final = [ + json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}), + json.dumps({"type": "response.create", "response": {"model": "gpt-realtime-2.1-mini"}}), + ] + + relayed: Final = await _relay_client_frames(frames, gate) + + assert gate.checked == ["gpt-realtime-2.1", "gpt-realtime-2.1-mini"] + assert "model" not in relayed.success_handler.call_args.kwargs + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "frame, expected_model", + [ + ('{"type": "response.create", "\\u006dodel": "gpt-5.4"}', "gpt-5.4"), + ('{"type": "session.update", "session": {"\\u006d\\u006fdel": "gpt-realtime-2.1"}}', "gpt-realtime-2.1"), + ], +) +async def test_websocket_passthrough_checks_a_model_named_with_escaped_json(frame, expected_model): + """ + JSON lets a caller spell the same key several ways, so a frame carrying + "\\u006dodel" asks for a model just as plainly as one carrying "model" and + has to face the same check. + """ + gate: Final = RecordingFrameModelGate(MappingProxyType({expected_model: KEY_MODEL_REFUSAL})) + + relayed: Final = await _relay_client_frames([frame], gate) + + assert gate.checked == [expected_model] + assert relayed.upstream.send.await_args_list == [] + assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_checks_the_model_a_binary_frame_asks_for(): + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL})) + frame: Final = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}).encode() + + relayed: Final = await _relay_client_frames([frame], gate) + + assert gate.checked == ["gpt-realtime-2.1"] + assert relayed.upstream.send.await_args_list == [] + assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)] + + +STORED_PROMPT_FRAMES: Final = ( + {"type": "response.create", "prompt": {"id": "pmpt_fake", "version": "2"}}, + {"type": "response.create", "response": {"prompt": {"id": "pmpt_fake"}}}, + {"type": "session.update", "session": {"prompt": {"id": "pmpt_fake"}}}, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("frame", STORED_PROMPT_FRAMES) +async def test_websocket_passthrough_refuses_a_stored_prompt_for_a_key_limited_to_some_models(frame): + """ + A stored prompt carries its own model on the provider side, so a key held to + a model list could reach any model through one. The gateway cannot see that + model, so a key that is not allowed every model cannot run a stored prompt. + """ + gate: Final = RecordingFrameModelGate(MappingProxyType({STORED_PROMPT_PROBE_MODEL: KEY_MODEL_REFUSAL})) + + relayed: Final = await _relay_client_frames([json.dumps(frame)], gate) + + error_message: Final = json.loads(relayed.client.send_text.await_args.args[0])["error"]["message"] + close_reason: Final = relayed.client.close.await_args.kwargs["reason"] + assert gate.checked == [STORED_PROMPT_PROBE_MODEL] + assert relayed.upstream.send.await_args_list == [] + assert relayed.client.close.await_args.kwargs["code"] == 1008 + assert "stored prompt" in close_reason + assert close_reason.endswith(".") and error_message.startswith(close_reason) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("frame", STORED_PROMPT_FRAMES) +async def test_websocket_passthrough_forwards_a_stored_prompt_for_a_key_allowed_every_model(frame): + gate: Final = RecordingFrameModelGate() + + relayed: Final = await _relay_client_frames([json.dumps(frame)], gate) + + assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [json.dumps(frame)] + assert relayed.client.close.await_args_list == [call()] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_keeps_the_budget_reason_when_a_stored_prompt_is_refused_for_spend(): + refusal: Final = "Budget has been exceeded! Key=hashed Spend=2.0, Budget=1.0" + gate: Final = RecordingFrameModelGate( + MappingProxyType({STORED_PROMPT_PROBE_MODEL: refusal}), error_type="budget_exceeded" + ) + + relayed: Final = await _relay_client_frames([json.dumps(STORED_PROMPT_FRAMES[0])], gate) + + assert json.loads(relayed.client.send_text.await_args.args[0])["error"] == { + "type": "budget_exceeded", + "message": refusal, + } + + +@pytest.mark.asyncio +async def test_websocket_passthrough_checks_the_model_a_hosted_tool_names(): + """A hosted tool such as image generation runs its own model, so that model gets the same check as the turn's""" + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-image-2": KEY_MODEL_REFUSAL})) + frame: Final = { + "type": "response.create", + "model": "gpt-5.4-mini", + "tools": [{"type": "web_search"}, {"type": "image_generation", "model": "gpt-image-2"}], + } + + relayed: Final = await _relay_client_frames([json.dumps(frame)], gate) + + assert gate.checked == ["gpt-5.4-mini", "gpt-image-2"] + assert relayed.upstream.send.await_args_list == [] + assert relayed.client.close.await_args.kwargs["code"] == 1008 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "frame", + [ + {"type": "response.create", "model": "gpt-5.4-mini", "tools": [{"type": "image_generation"}]}, + {"type": "response.create", "response": {"tools": [{"type": "image_generation", "model": None}]}}, + ], +) +async def test_websocket_passthrough_refuses_an_image_tool_without_a_model_for_a_key_limited_to_some_models(frame): + gate: Final = RecordingFrameModelGate(MappingProxyType({STORED_PROMPT_PROBE_MODEL: KEY_MODEL_REFUSAL})) + + relayed: Final = await _relay_client_frames([json.dumps(frame)], gate) + + assert gate.checked[-1] == STORED_PROMPT_PROBE_MODEL + assert relayed.upstream.send.await_args_list == [] + assert "image_generation tool has to name its model" in relayed.client.close.await_args.kwargs["reason"] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_checks_a_frame_only_the_standard_json_parser_reads(): + """A lone surrogate escape is valid JSON syntax the fast parser refuses, so the frame is still checked""" + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL})) + frame: Final = '{"type": "response.create", "model": "gpt-5.4", "metadata": {"note": "\\ud800"}}' + + relayed: Final = await _relay_client_frames([frame], gate) + + assert gate.checked == ["gpt-5.4"] + assert relayed.upstream.send.await_args_list == [] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_forwards_binary_audio_untouched(): + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"})) + audio: Final = bytes(range(256)) * 4 + + relayed: Final = await _relay_client_frames([audio], gate) + + assert gate.checked == [] + assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [audio] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_without_a_frame_gate_forwards_every_frame(): + frame: Final = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}) + + relayed: Final = await _relay_client_frames([frame], None) + + assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [frame] + assert "model" not in relayed.success_handler.call_args.kwargs + + +REALTIME_SESSION_CREATED: Final = json.dumps( + {"type": "session.created", "session": {"id": "sess_fake", "model": "gpt-realtime-2.1"}} +) +REALTIME_RESPONSE_DONE: Final = json.dumps( + { + "type": "response.done", + "response": {"status": "completed", "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}}, + } +) + + +@pytest.mark.asyncio +async def test_websocket_passthrough_still_bills_a_session_the_client_walked_out_of(): + """ + A realtime client closes the moment it has the answer it asked for, so the + next frame the provider sends lands on a socket that is already gone. That + is how a session ends, not a gateway failure, and the turns the caller + already used still have to reach spend, including the last frame the client + never collected. + """ + upstream_ws: Final = FakeUpstreamWebSocket( + REALTIME_SESSION_CREATED, + REALTIME_RESPONSE_DONE, + json.dumps({"type": "response.output_text.delta", "delta": "too late"}), + ) + websocket: Final = _client_websocket(_pending_receive) + websocket.send_text = AsyncMock(side_effect=[None, None, WebSocketDisconnect(code=1006)]) + + with ( + _patched_websocket_passthrough_environment(upstream_ws), + patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler, + ): + await websocket_passthrough_request( + websocket=websocket, + target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1", + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/openai_passthrough/v1/realtime", + accept_websocket=False, + ) + + relayed_frames: Final = success_handler.call_args.kwargs["response_body"] + assert [frame["type"] for frame in relayed_frames] == [ + "session.created", + "response.done", + "response.output_text.delta", + ] + + +class ScriptedUpstreamWebSocket: + """Answers each client frame with its scripted reply""" + + def __init__(self, *opening: str, replies: Sequence[Sequence[str]] = ()): + self._frames: asyncio.Queue[str | None] = asyncio.Queue() + for frame in opening: + self._frames.put_nowait(frame) + self._replies = iter(replies) + self._closed = False + self.send = AsyncMock(side_effect=self._reply) + self.close = AsyncMock(side_effect=self._flush_and_close) + + async def _reply(self, frame: str | bytes) -> None: + for reply in next(self._replies, ()): + self._frames.put_nowait(reply) + + async def _flush_and_close(self, *args: object, **kwargs: object) -> None: + if self._closed: + return + self._closed = True + self._frames.put_nowait(None) + while not self._frames.empty(): + await asyncio.sleep(0) + + async def recv(self, decode: bool | None = None) -> str: + from websockets.exceptions import ConnectionClosedOK + from websockets.frames import Close + + frame: Final = await self._frames.get() + if frame is None: + raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True) + return frame + + +def _turn_taking_client( + frames: Sequence[str], answered_by: str, answers_due: Sequence[int] | None = None +) -> MagicMock: + """A client that sends each frame, and then leaves, once the answers due by then reached it (by default the + answer to every earlier frame), and that Starlette refuses to send to once it is closed""" + from starlette.websockets import WebSocketState + + delivered: Final[list[str]] = [] + script: Final = iter(frames) + receives: Final = iter(range(len(frames) + 1) if answers_due is None else answers_due) + + async def receive(): + answers_due: Final = next(receives) + while sum(answered_by in text for text in delivered) < answers_due: + await asyncio.sleep(0) + frame: Final = next(script, None) + return {"type": "websocket.disconnect"} if frame is None else _received_frame(frame) + + websocket: Final = _client_websocket(receive) + + async def send_text(text: str) -> None: + if websocket.application_state == WebSocketState.DISCONNECTED: + raise RuntimeError('Cannot call "send" once a close message has been sent.') + delivered.append(text) + + websocket.send_text = AsyncMock(side_effect=send_text) + return websocket + + +async def _relay_turns( + client: MagicMock, + upstream_ws: ScriptedUpstreamWebSocket, + gate: WebsocketFrameModelGate, + endpoint: str = "/openai_passthrough/v1/responses", + target: str = "wss://api.openai.com/v1/responses", + opening_access_groups: Sequence[str] = (), + walkout_turn_drain_seconds: float = 0.05, +) -> SimpleNamespace: + with ( + _patched_websocket_passthrough_environment(upstream_ws) as proxy_logging, + patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler, + ): + await asyncio.wait_for( + websocket_passthrough_request( + websocket=client, + target=target, + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=UserAPIKeyAuth( + api_key="hashed-key", matched_model_access_groups=list(opening_access_groups) + ), + forward_headers=False, + endpoint=endpoint, + accept_websocket=False, + client_frame_model_gate=gate, + upstream_frame_kind=openai_websocket_frame_kind, + walkout_turn_drain_seconds=walkout_turn_drain_seconds, + ), + timeout=10, + ) + return SimpleNamespace(success_handler=success_handler, proxy_logging=proxy_logging) + + +def _responses_turn(model: str) -> list[str]: + return [_responses_turn_started(model), _responses_turn_ended(model)] + + +def _responses_turn_started(model: str) -> str: + return json.dumps({"type": "response.created", "response": {"model": model}}) + + +def _responses_turn_ended(model: str) -> str: + return ( + + json.dumps( + { + "type": "response.completed", + "response": {"model": model, "usage": {"input_tokens": 9, "output_tokens": 3, "total_tokens": 12}}, + } + ) + ) + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_each_turn_on_its_own_row_charged_to_the_groups_that_allowed_it(): + """ + One socket can run a model from one access group and then a model from + another. Each turn has to reach spend as it finishes, under its own call, + and be charged to the group that let its model through, so a group budget + sees its share during the session and never pays for the other group's turn. + """ + gate: Final = RecordingFrameModelGate( + groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), + ) + client: Final = _turn_taking_client( + [ + json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), + json.dumps({"type": "response.create", "model": "gpt-5.4"}), + ], + answered_by='"response.completed"', + ) + upstream_ws: Final = ScriptedUpstreamWebSocket(replies=[_responses_turn("gpt-5.4-mini"), _responses_turn("gpt-5.4")]) + + relayed: Final = await _relay_turns(client, upstream_ws, gate) + + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] + assert [[frame["response"]["model"] for frame in row["response_body"]] for row in rows] == [ + ["gpt-5.4-mini", "gpt-5.4-mini"], + ["gpt-5.4", "gpt-5.4"], + ] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [ + ["tier-a"], + ["tier-b"], + ] + assert [row["litellm_params"]["metadata"]["user_api_key"] for row in rows] == ["hashed-key", "hashed-key"] + assert rows[0]["litellm_call_id"] != rows[1]["litellm_call_id"] + assert rows[1]["start_time"] == rows[0]["end_time"] + relayed.proxy_logging.post_call_failure_hook.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_websocket_passthrough_holds_the_keys_parallel_slot_until_the_session_ends(): + """ + A key allowed one parallel request holds that slot for as long as its + socket stays open. A turn that finishes and reaches spend mid-session must + not free it, or a second session gets in while the first is still running. + """ + limiter_cache: Final = DualCache() + limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(limiter_cache)) + key: Final = UserAPIKeyAuth(api_key=hash_token("sk-one-socket"), max_parallel_requests=1) + logged_rows: Final[list[asyncio.Task[None]]] = [] + + async def admit(user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str) -> None: + await limiter.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=limiter_cache, data=data, call_type=call_type + ) + + async def release(data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: object) -> None: + await limiter.async_post_call_success_hook(data=data, user_api_key_dict=user_api_key_dict, response=response) + + async def log_row(**row: object) -> None: + await limiter.async_log_success_event( + kwargs=row, response_obj=None, start_time=row["start_time"], end_time=row["end_time"] + ) + + async def opens_another_session() -> bool: + await asyncio.gather(*logged_rows) + try: + await asyncio.create_task( + admit(key, {"litellm_call_id": str(uuid.uuid4())}, "pass_through_endpoint"), + context=contextvars.Context(), + ) + except ProxyRateLimitError: + return False + return True + + client: Final = _turn_taking_client( + [json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})], answered_by='"response.completed"' + ) + take_frame: Final = client.receive + while_open: Final[list[bool]] = [] + + async def receive() -> Mapping[str, str | bytes]: + message: Final = await take_frame() + if message["type"] == "websocket.disconnect": + while_open.append(await opens_another_session()) + return message + + client.receive = receive + + with ( + _patched_websocket_passthrough_environment( + ScriptedUpstreamWebSocket(replies=[_responses_turn("gpt-5.4-mini")]) + ) as proxy_logging, + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as worker, + patch(SUCCESS_HANDLER, new_callable=AsyncMock, side_effect=log_row), + ): + proxy_logging.pre_call_hook = AsyncMock(side_effect=admit) + proxy_logging.post_call_success_hook = AsyncMock(side_effect=release) + worker.ensure_initialized_and_enqueue = MagicMock( + side_effect=lambda async_coroutine: logged_rows.append(asyncio.create_task(async_coroutine)) + ) + await asyncio.wait_for( + websocket_passthrough_request( + websocket=client, + target="wss://api.openai.com/v1/responses", + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=key, + forward_headers=False, + endpoint="/openai_passthrough/v1/responses", + accept_websocket=False, + client_frame_model_gate=RecordingFrameModelGate(), + upstream_frame_kind=openai_websocket_frame_kind, + walkout_turn_drain_seconds=0.05, + ), + timeout=10, + ) + after_close: Final = await opens_another_session() + + assert len(logged_rows) == 1 + assert (while_open, after_close) == ([False], True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "models, upstream_refuses, while_open", + [ + ((), False, [False]), + (("gpt-5.4-mini",), False, [False]), + (("gpt-5.4-mini", "gpt-5.4"), False, [False]), + ((), True, []), + ], + ids=["no-turn", "one-turn", "two-turns", "upstream-refused"], +) +async def test_websocket_passthrough_frees_one_legacy_parallel_slot_only_when_the_session_ends( + models: tuple[str, ...], upstream_refuses: bool, while_open: list[bool] +): + limiter: Final = PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), clock=lambda: datetime(2026, 1, 31, 23, 59, 30) + ) + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = limiter + key: Final = UserAPIKeyAuth(api_key=hash_token("sk-legacy-socket"), max_parallel_requests=2) + logged_rows: Final[list[asyncio.Task[None]]] = [] + seen_while_open: Final[list[bool]] = [] + + async def admit(user_api_key_dict: UserAPIKeyAuth, data: Mapping[str, object], call_type: str) -> None: + await limiter.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=DualCache(), data=data, call_type=call_type + ) + + async def log_row(**row: object) -> None: + await limiter.async_log_success_event( + kwargs=row, response_obj=None, start_time=row["start_time"], end_time=row["end_time"] + ) + + async def opens_another_session() -> bool: + await asyncio.gather(*logged_rows) + try: + await admit(key, {}, "pass_through_endpoint") + except ProxyRateLimitError: + return False + return True + + client: Final = _turn_taking_client( + [json.dumps({"type": "response.create", "model": model}) for model in models], + answered_by='"response.completed"', + ) + take_frame: Final = client.receive + + async def receive() -> Mapping[str, str | bytes]: + message: Final = await take_frame() + if message["type"] == "websocket.disconnect": + seen_while_open.append(await opens_another_session()) + return message + + client.receive = receive + await admit(key, {}, "pass_through_endpoint") + + with ExitStack() as stack: + proxy_logging: Final = stack.enter_context( + _patched_websocket_passthrough_environment( + ScriptedUpstreamWebSocket(replies=[_responses_turn(model) for model in models]) + ) + ) + worker: Final = stack.enter_context( + patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") + ) + stack.enter_context(patch(SUCCESS_HANDLER, new_callable=AsyncMock, side_effect=log_row)) + if upstream_refuses: + stack.enter_context( + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", + side_effect=OSError("upstream refused the socket"), + ) + ) + proxy_logging.pre_call_hook = AsyncMock(side_effect=admit) + proxy_logging.arelease_responses_websocket_slot = proxy_logging_obj.arelease_responses_websocket_slot + worker.ensure_initialized_and_enqueue = MagicMock( + side_effect=lambda async_coroutine: logged_rows.append(asyncio.create_task(async_coroutine)) + ) + await asyncio.wait_for( + websocket_passthrough_request( + websocket=client, + target="wss://api.openai.com/v1/responses", + custom_headers={"Authorization": "Bearer sk-provider"}, + user_api_key_dict=key, + forward_headers=False, + endpoint="/openai_passthrough/v1/responses", + accept_websocket=False, + client_frame_model_gate=RecordingFrameModelGate(), + upstream_frame_kind=openai_websocket_frame_kind, + walkout_turn_drain_seconds=0.05, + ), + timeout=10, + ) + after_close: Final = (await opens_another_session(), await opens_another_session()) + + assert len(logged_rows) == len(models) + assert (seen_while_open, after_close) == (while_open, (True, False)) + + +@pytest.mark.asyncio +async def test_websocket_passthrough_charges_pipelined_turns_to_the_frame_that_asked_for_each(): + """ + A client can send its next response.create before the provider finished the + last one, and the provider then runs them one after the other. Each turn + still has to reach spend under the groups that let its own frame through, + and under the model that frame asked for, even when the provider answers + under a dated snapshot name a per-model budget was never set on. + """ + gate: Final = RecordingFrameModelGate( + groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), + ) + client: Final = _turn_taking_client( + [ + json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), + json.dumps({"type": "response.create", "model": "gpt-5.4"}), + ], + answered_by='"response.completed"', + answers_due=(0, 0, 2), + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + replies=[_responses_turn("gpt-5.4-mini-2099-01-01"), _responses_turn("gpt-5.4-2099-01-01")] + ) + + relayed: Final = await _relay_turns(client, upstream_ws, gate) + + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] + assert [[frame["response"]["model"] for frame in row["response_body"]] for row in rows] == [ + ["gpt-5.4-mini-2099-01-01", "gpt-5.4-mini-2099-01-01"], + ["gpt-5.4-2099-01-01", "gpt-5.4-2099-01-01"], + ] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [ + ["tier-a"], + ["tier-b"], + ] + assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-5.4-mini", "gpt-5.4"] + + +def _responses_error(code: str) -> str: + return json.dumps({"type": "error", "status": 400, "error": {"type": "invalid_request_error", "code": code}}) + + +def _on_lane(frame: str, stream_id: str) -> str: + return json.dumps({**json.loads(frame), "stream_id": stream_id}) + + +@pytest.mark.asyncio +async def test_websocket_passthrough_charges_turns_after_a_refused_frame_to_the_frames_that_asked_for_them(): + """ + The provider answers a response.create it cannot run, one naming no model + or a stale previous_response_id, with an error event and keeps the socket + open. Those frames never run, so each later turn still has to reach spend + under the groups and model its own frame asked for. + """ + gate: Final = RecordingFrameModelGate( + groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), + ) + client: Final = _turn_taking_client( + [ + json.dumps({"type": "response.create", "input": "hi"}), + json.dumps({"type": "response.create", "model": "gpt-5.4", "previous_response_id": "resp_gone"}), + json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), + json.dumps({"type": "response.create", "model": "gpt-5.4"}), + ], + answered_by='"response.completed"', + answers_due=(0, 0, 0, 0, 2), + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + replies=[ + [_responses_error("missing_required_parameter")], + [_responses_error("previous_response_not_found")], + _responses_turn("gpt-5.4-mini-2099-01-01"), + _responses_turn("gpt-5.4-2099-01-01"), + ] + ) + + relayed: Final = await _relay_turns(client, upstream_ws, gate) + + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] + assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-5.4-mini", "gpt-5.4"] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [ + ["tier-a"], + ["tier-b"], + ] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_charges_turns_on_parallel_lanes_to_the_frame_that_asked_for_each(): + """ + A client can put each response.create on its own stream_id lane, and the + provider runs lanes side by side, so a later frame's turn can end first. + Each turn still has to reach spend on its own row, under the groups and + model its own frame asked for. + """ + gate: Final = RecordingFrameModelGate( + groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), + ) + client: Final = _turn_taking_client( + [ + json.dumps({"type": "response.create", "stream_id": "slow", "model": "gpt-5.4"}), + json.dumps({"type": "response.create", "stream_id": "fast", "model": "gpt-5.4-mini"}), + ], + answered_by='"response.completed"', + answers_due=(0, 0, 2), + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + replies=[ + [_on_lane(_responses_turn_started("gpt-5.4"), "slow")], + [ + *(_on_lane(frame, "fast") for frame in _responses_turn("gpt-5.4-mini")), + _on_lane(_responses_turn_ended("gpt-5.4"), "slow"), + ], + ] + ) + + relayed: Final = await _relay_turns(client, upstream_ws, gate) + + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] + assert [[(frame["stream_id"], frame["response"]["model"]) for frame in row["response_body"]] for row in rows] == [ + [("fast", "gpt-5.4-mini"), ("fast", "gpt-5.4-mini")], + [("slow", "gpt-5.4"), ("slow", "gpt-5.4")], + ] + assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-5.4-mini", "gpt-5.4"] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [ + ["tier-a"], + ["tier-b"], + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream_id", ["", 7]) +async def test_websocket_passthrough_refuses_a_frame_on_a_lane_the_provider_error_cannot_name(stream_id): + """ + The provider refuses a frame with an invalid stream_id through an error that + names no lane (OpenAI WebSocket mode guide, "Errors to handle", read + 2026-11-02, which also says an empty stream_id is invalid). Relayed, that + error would refuse a turn waiting on the default lane instead, and that turn + would reach spend without the groups its own frame asked for. + """ + gate: Final = RecordingFrameModelGate() + allowed: Final = json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}) + + relayed: Final = await _relay_client_frames( + [allowed, json.dumps({"type": "response.create", "stream_id": stream_id, "model": "gpt-5.4-mini"})], gate + ) + + assert [sent.args[0] for sent in relayed.upstream.send.await_args_list] == [allowed] + assert json.loads(relayed.client.send_text.await_args.args[0])["type"] == "error" + assert relayed.client.close.await_args.kwargs["code"] == 1008 + + +@pytest.mark.asyncio +async def test_websocket_passthrough_charges_realtime_turns_to_the_session_the_key_opened(): + """ + A realtime session runs the model it was opened on, so its turns stay + charged to the groups and the model the key opened it with, whatever model + a later session.update or response.create names. + """ + gate: Final = RecordingFrameModelGate(groups=MappingProxyType({"gpt-realtime-2.1-mini": ("tier-b",)})) + client: Final = _turn_taking_client( + [ + json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1-mini"}}), + json.dumps({"type": "response.create", "response": {"model": "gpt-realtime-2.1-mini"}}), + ], + answered_by='"type"', + answers_due=(1, 2, 4), + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, + replies=[ + [json.dumps({"type": "session.updated"})], + [json.dumps({"type": "response.created", "response": {"status": "in_progress"}}), REALTIME_RESPONSE_DONE], + ], + ) + + relayed: Final = await _relay_turns( + client, + upstream_ws, + gate, + "/openai/v1/realtime", + target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1", + opening_access_groups=("tier-a",), + ) + + assert gate.checked == ["gpt-realtime-2.1-mini", "gpt-realtime-2.1-mini"] + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [["tier-a"]] + assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-realtime-2.1"] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_no_spend_row_for_a_session_refused_before_any_answer(): + """A session refused at its first frame ran nothing, so its only record is the refusal, not an extra $0 row""" + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL})) + client: Final = _turn_taking_client([json.dumps({"type": "response.create", "model": "gpt-5.4"})], answered_by='"error"') + + relayed: Final = await _relay_turns(client, ScriptedUpstreamWebSocket(), gate) + + assert relayed.success_handler.call_args_list == [] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_the_unfinished_turn_a_client_walked_out_of(): + """ + A client can leave while a later turn is still running. That turn never + ends, but it was started and may be billed, so it still gets its own row. + """ + client: Final = _turn_taking_client( + [ + json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), + json.dumps({"type": "response.create", "model": "gpt-5.4"}), + ], + answered_by='"response.created"', + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + replies=[_responses_turn("gpt-5.4-mini"), [_responses_turn_started("gpt-5.4")]] + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate()) + + assert [ + [json.loads(frame)["type"] if isinstance(frame, str) else frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] == [["response.created", "response.completed"], ["response.created"]] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_charges_each_lane_a_client_walked_out_of_to_the_frame_that_started_it(): + """ + A client can leave while turns on several stream_id lanes are still + running. Each unfinished turn still gets its own row, under the groups and + model its own frame asked for. + """ + gate: Final = RecordingFrameModelGate( + groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), + ) + client: Final = _turn_taking_client( + [ + json.dumps({"type": "response.create", "stream_id": "slow", "model": "gpt-5.4"}), + json.dumps({"type": "response.create", "stream_id": "fast", "model": "gpt-5.4-mini"}), + ], + answered_by='"response.created"', + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + replies=[ + [_on_lane(_responses_turn_started("gpt-5.4"), "slow")], + [_on_lane(_responses_turn_started("gpt-5.4-mini"), "fast")], + ] + ) + + relayed: Final = await _relay_turns(client, upstream_ws, gate) + + rows: Final = [logged.kwargs for logged in relayed.success_handler.call_args_list] + assert [[frame["stream_id"] for frame in row["response_body"]] for row in rows] == [["slow"], ["fast"]] + assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-5.4", "gpt-5.4-mini"] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [ + ["tier-b"], + ["tier-a"], + ] + + +class UpstreamFinishingTheTurnAfterTheClientLeft: + """Starts a turn for each client frame, unless it only answers once the client left, sends the rest of the session + once the client left, and stops once closed""" + + def __init__( + self, + client_left: asyncio.Event, + started: str | None, + rest_of_the_turn: Sequence[str], + opening: Sequence[str] = (), + ): + self._frames: asyncio.Queue[str | None] = asyncio.Queue() + for frame in opening: + self._frames.put_nowait(frame) + self._client_left = client_left + self._started = started + self._rest_of_the_turn = rest_of_the_turn + self._closed = False + self._finishing: list[asyncio.Task[None]] = [] + self.send = AsyncMock(side_effect=self._start_turn) + self.close = AsyncMock(side_effect=self._close) + + def _close(self, *args: object, **kwargs: object) -> None: + self._closed = True + self._frames.put_nowait(None) + + async def _start_turn(self, frame: str | bytes) -> None: + if self._started is not None: + self._frames.put_nowait(self._started) + if not self._finishing: + self._finishing.append(asyncio.create_task(self._finish_once_the_client_left())) + + async def _finish_once_the_client_left(self) -> None: + await self._client_left.wait() + for frame in self._rest_of_the_turn: + await asyncio.sleep(0) + if self._closed: + return + self._frames.put_nowait(frame) + + async def recv(self, decode: bool | None = None) -> str: + from websockets.exceptions import ConnectionClosedOK + from websockets.frames import Close + + frame: Final = await self._frames.get() + if frame is None: + raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True) + return frame + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hang_up_read_before_the_next_send", [True, False], ids=["read_first", "found_by_a_send"]) +async def test_websocket_passthrough_bills_the_turn_a_client_walked_out_of_once_the_provider_ends_it( + hang_up_read_before_the_next_send: bool, +): + """ + Leaving mid-answer must not make a turn free. The provider keeps running and + bills it, so the gateway keeps reading until the turn ends and logs its usage, + whether it reads the hang-up first or finds it by failing to relay a frame. + """ + client_left: Final = asyncio.Event() + client: Final = _walking_out_client( + [json.dumps({"type": "response.create", "model": "gpt-5.4"})], + '"response.created"', + None, + client_left, + hang_up_read_before_the_next_send, + ) + upstream_ws: Final = UpstreamFinishingTheTurnAfterTheClientLeft( + client_left, + _responses_turn_started("gpt-5.4"), + (json.dumps({"type": "response.output_text.delta", "delta": "Hi"}), _responses_turn_ended("gpt-5.4")), + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), walkout_turn_drain_seconds=5) + + assert [ + [json.loads(frame)["type"] if isinstance(frame, str) else frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] == [["response.created", "response.output_text.delta", "response.completed"]] + upstream_ws.close.assert_awaited() + + +def _walking_out_client( + frames: Sequence[str], + answered_by: str, + answers_due: Sequence[int] | None, + client_left: asyncio.Event, + hang_up_read_before_the_next_send: bool = True, +) -> MagicMock: + """A turn-taking client that sets client_left as it hangs up, which the gateway either reads at once or only finds + when its next send to the client fails""" + from starlette.websockets import WebSocketDisconnect, WebSocketState + + send_failed: Final = asyncio.Event() + client: Final = _turn_taking_client(frames, answered_by, answers_due) + receive: Final = client.receive + send_text: Final = client.send_text + + async def receive_until_leaving() -> Mapping[str, object]: + message: Final = await receive() + if message["type"] == "websocket.disconnect": + client_left.set() + if not hang_up_read_before_the_next_send: + await send_failed.wait() + client.client_state = WebSocketState.DISCONNECTED + return message + + async def send_text_until_leaving(text: str) -> None: + if client_left.is_set(): + send_failed.set() + raise WebSocketDisconnect(code=1006) + await send_text(text) + + client.receive = receive_until_leaving + client.send_text = AsyncMock(side_effect=send_text_until_leaving) + return client + + +@pytest.mark.asyncio +@pytest.mark.parametrize("turns", [1, 2], ids=["one_turn", "two_turns_asked_at_once"]) +async def test_websocket_passthrough_bills_the_turn_a_client_walked_out_of_before_the_provider_started_it(turns: int): + """ + A client can ask for turns and leave before the provider says they started. + The provider still runs and bills each one, so the gateway keeps reading + until the last one ends and logs each one's usage. + """ + client_left: Final = asyncio.Event() + client: Final = _walking_out_client( + [json.dumps({"type": "response.create", "model": "gpt-5.4"})] * turns, + '"response.created"', + (0,) * (turns + 1), + client_left, + ) + upstream_ws: Final = UpstreamFinishingTheTurnAfterTheClientLeft( + client_left, None, (_responses_turn_started("gpt-5.4"), _responses_turn_ended("gpt-5.4")) * turns + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), walkout_turn_drain_seconds=5) + + assert _logged_frame_types(relayed) == [["response.created", "response.completed"]] * turns + upstream_ws.close.assert_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hang_up_read_before_the_next_send", [True, False], ids=["read_first", "found_by_a_send"]) +async def test_websocket_passthrough_bills_a_realtime_turn_a_client_walked_out_of_while_the_provider_stored_its_input( + hang_up_read_before_the_next_send: bool, +): + """ + A realtime client can add an item, ask for a turn, and leave at once. The + provider confirms the item before it starts the turn, and that confirmation + must not end the session while the turn it was asked for is still owed. + """ + client_left: Final = asyncio.Event() + client: Final = _walking_out_client( + [ + json.dumps({"type": "conversation.item.create", "item": {"type": "message", "role": "user"}}), + json.dumps({"type": "response.create"}), + ], + '"session.created"', + (1, 1, 1), + client_left, + hang_up_read_before_the_next_send, + ) + upstream_ws: Final = UpstreamFinishingTheTurnAfterTheClientLeft( + client_left, + None, + ( + json.dumps({"type": "conversation.item.created", "item": {"type": "message", "role": "user"}}), + json.dumps({"type": "response.created", "response": {"status": "in_progress"}}), + REALTIME_RESPONSE_DONE, + ), + opening=(REALTIME_SESSION_CREATED,), + ) + + relayed: Final = await _relay_turns( + client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime", walkout_turn_drain_seconds=5 + ) + + assert _logged_frame_types(relayed) == [ + ["session.created", "conversation.item.created", "response.created", "response.done"] + ] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_bills_a_realtime_turn_a_client_walked_out_of_after_the_provider_started_one_itself(): + """ + A realtime provider can start a turn on its own when voice detection hears + the user stop. That turn answers no client frame, so a turn the client asks + for next and leaves before it starts is still owed and still billed. + """ + realtime_response_created: Final = json.dumps({"type": "response.created", "response": {"status": "in_progress"}}) + client_left: Final = asyncio.Event() + client: Final = _walking_out_client( + [json.dumps({"type": "response.create"})], '"response.done"', (1, 1), client_left + ) + upstream_ws: Final = UpstreamFinishingTheTurnAfterTheClientLeft( + client_left, + None, + (realtime_response_created, REALTIME_RESPONSE_DONE), + opening=(REALTIME_SESSION_CREATED, realtime_response_created, REALTIME_RESPONSE_DONE), + ) + + relayed: Final = await _relay_turns( + client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime", walkout_turn_drain_seconds=5 + ) + + assert _logged_frame_types(relayed) == [ + ["session.created", "response.created", "response.done"], + ["session.created", "response.created", "response.done"], + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "answer", + [ + _responses_turn("gpt-5.4"), + [_responses_turn_ended("gpt-5.4")], + [_responses_error("previous_response_not_found")], + ], + ids=["finished", "finished_with_no_start_frame", "refused"], +) +async def test_websocket_passthrough_lets_go_of_the_provider_at_once_when_a_client_leaves_owing_no_turn( + answer: Sequence[str], +): + """ + Once every turn a client asked for has ended or been refused, nothing more + can be billed, so a client leaving closes the provider socket right away + instead of holding it for the walkout drain. + """ + client: Final = _turn_taking_client( + [json.dumps({"type": "response.create", "model": "gpt-5.4"})], + answered_by=json.dumps(json.loads(answer[-1])["type"]), + ) + upstream_ws: Final = ScriptedUpstreamWebSocket(replies=[answer]) + + relayed: Final = await _relay_turns( + client, upstream_ws, RecordingFrameModelGate(), walkout_turn_drain_seconds=WALKOUT_TURN_DRAIN_SECONDS + ) + + upstream_ws.close.assert_awaited() + assert len(relayed.success_handler.call_args_list) == 1 + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_no_extra_row_for_what_the_provider_sends_after_the_last_turn(): + """ + A realtime provider follows each finished turn with frames that start no + turn, such as rate_limits.updated. The session ends with nothing running, + so its turn is its only row and the request count stays one per turn. + """ + client: Final = _turn_taking_client([json.dumps({"type": "response.create"})], answered_by='"rate_limits.updated"') + upstream_ws: Final = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, + replies=[[REALTIME_RESPONSE_DONE, json.dumps({"type": "rate_limits.updated", "rate_limits": []})]], + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime") + + assert [ + [frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] == [["session.created", "response.done"]] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_carries_the_session_the_provider_opened_into_every_later_turn(): + """ + A realtime provider names the session's model once, when it opens the + session, and never again. Each later turn logs on its own, so it has to + carry that opening frame or it would be priced as a model nobody named. + """ + client: Final = _turn_taking_client( + [json.dumps({"type": "response.create"}), json.dumps({"type": "response.create"})], + answered_by='"response.done"', + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, replies=[[REALTIME_RESPONSE_DONE], [REALTIME_RESPONSE_DONE]] + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime") + + assert [ + [frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] == [["session.created", "response.done"], ["session.created", "response.done"]] + + +REALTIME_TRANSCRIPTION_COMPLETED: Final = json.dumps( + { + "type": "conversation.item.input_audio_transcription.completed", + "item_id": "item_fake", + "transcript": "hello", + "usage": {"type": "duration", "seconds": 2}, + } +) +TRANSCRIPTION_COMPLETED: Final = "conversation.item.input_audio_transcription.completed" + + +def _transcription_session_update(frame_type: str, model: str) -> str: + return json.dumps({"type": frame_type, "session": {"audio": {"input": {"transcription": {"model": model}}}}}) + + +def _logged_frame_types(relayed: SimpleNamespace) -> list[list[str]]: + return [ + [frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_bills_a_clip_transcribed_after_the_last_answer(): + client: Final = _turn_taking_client( + [_transcription_session_update("session.update", "gpt-transcribe"), json.dumps({"type": "response.create"})], + answered_by='"type"', + answers_due=(1, 2, 4), + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, + replies=[ + [_transcription_session_update("session.updated", "gpt-transcribe")], + [REALTIME_RESPONSE_DONE, REALTIME_TRANSCRIPTION_COMPLETED], + ], + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime") + + assert _logged_frame_types(relayed) == [ + ["session.created", "session.updated", "response.done"], + ["session.created", "session.updated", TRANSCRIPTION_COMPLETED], + ] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_bills_a_clip_transcribed_during_an_answer_in_that_answer_row(): + client: Final = _turn_taking_client( + [json.dumps({"type": "response.create"})], answered_by='"type"', answers_due=(1, 4) + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, + replies=[ + [ + json.dumps({"type": "response.created", "response": {"status": "in_progress"}}), + REALTIME_TRANSCRIPTION_COMPLETED, + REALTIME_RESPONSE_DONE, + ] + ], + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime") + + assert _logged_frame_types(relayed) == [ + ["session.created", "response.created", TRANSCRIPTION_COMPLETED, "response.done"] + ] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_carries_only_the_latest_session_update_into_later_rows(): + client: Final = _turn_taking_client( + [ + _transcription_session_update("session.update", "gpt-transcribe-a"), + json.dumps({"type": "response.create"}), + _transcription_session_update("session.update", "gpt-transcribe-b"), + json.dumps({"type": "response.create"}), + ], + answered_by='"type"', + answers_due=(1, 2, 3, 4, 5), + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, + replies=[ + [_transcription_session_update("session.updated", "gpt-transcribe-a")], + [REALTIME_RESPONSE_DONE], + [_transcription_session_update("session.updated", "gpt-transcribe-b")], + [REALTIME_RESPONSE_DONE], + ], + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime") + + last_row: Final = relayed.success_handler.call_args_list[-1].kwargs["response_body"] + assert [frame for frame in last_row if frame["type"] == "session.updated"] == [ + json.loads(_transcription_session_update("session.updated", "gpt-transcribe-b")) + ] + assert len(relayed.success_handler.call_args_list) == 2 + + +@pytest.mark.asyncio +async def test_websocket_passthrough_keeps_the_session_update_a_clip_ran_under_until_its_answer_is_billed(): + client: Final = _turn_taking_client( + [ + _transcription_session_update("session.update", "gpt-transcribe-a"), + json.dumps({"type": "response.create"}), + _transcription_session_update("session.update", "gpt-transcribe-b"), + ], + answered_by='"type"', + answers_due=(1, 2, 4, 6), + ) + upstream_ws: Final = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, + replies=[ + [_transcription_session_update("session.updated", "gpt-transcribe-a")], + [ + json.dumps({"type": "response.created", "response": {"status": "in_progress"}}), + REALTIME_TRANSCRIPTION_COMPLETED, + ], + [_transcription_session_update("session.updated", "gpt-transcribe-b"), REALTIME_RESPONSE_DONE], + ], + ) + + relayed: Final = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime") + + assert _logged_frame_types(relayed) == [ + [ + "session.created", + "session.updated", + "response.created", + TRANSCRIPTION_COMPLETED, + "session.updated", + "response.done", + ] + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "later_frame", + [ + json.dumps({"type": "response.create", "model": "gpt-5.4"}), + '{"type": "response.create", "model": "gpt-5.4-mini", "metadata": ' + "[" * 100_000 + "]" * 100_000 + "}", + ( + '{"type": "response.create", "model": "gpt-5.4-mini", "seed": ' + + "9" * (sys.get_int_max_str_digits() + 1) + + "}" + ), + ], + ids=["names_a_denied_model", "nests_too_deep_to_parse", "holds_a_number_too_long_to_parse"], +) +async def test_websocket_passthrough_still_bills_the_answer_in_flight_when_a_later_frame_is_refused(later_frame: str): + """ + A refused frame closes the client while the provider may still be sending the + answer to an earlier, allowed frame. That answer cannot go to a closed client, + but the provider bills it once it ends, so the gateway keeps reading until it + ends and its usage reaches spend. + """ + gate: Final = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL})) + client: Final = _turn_taking_client( + [json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), later_frame], + answered_by='"response.created"', + ) + refused: Final = asyncio.Event() + close_client: Final = client.close.side_effect + + def close_client_and_note_the_refusal(*args: object, **kwargs: object) -> None: + close_client(*args, **kwargs) + refused.set() + + client.close = AsyncMock(side_effect=close_client_and_note_the_refusal) + upstream_ws: Final = UpstreamFinishingTheTurnAfterTheClientLeft( + refused, _responses_turn_started("gpt-5.4-mini"), (_responses_turn_ended("gpt-5.4-mini"),) + ) + + relayed: Final = await _relay_turns(client, upstream_ws, gate, walkout_turn_drain_seconds=5) + + assert [json.loads(text.args[0])["type"] for text in client.send_text.await_args_list] == [ + "response.created", + "error", + ] + assert [ + [frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] == [["response.created", "response.completed"]] + relayed.proxy_logging.post_call_failure_hook.assert_not_awaited() + + DEEPGRAM_LISTEN_TARGET = "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16&sample_rate=16000" DEEPGRAM_INTERIM_FRAME = json.dumps( { @@ -8962,6 +10461,32 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key ) +async def _billed_metadata(gate: WebsocketFrameModelGate) -> Mapping[str, object]: + served: Final = await _relay_client_frames( + [json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})], + gate, + UserAPIKeyAuth(api_key="hashed-key"), + ) + return served.success_handler.call_args.kwargs["litellm_params"]["metadata"] + + +@pytest.mark.asyncio +async def test_websocket_session_is_charged_to_the_group_its_frames_were_authorized_through(): + """A websocket names its model in a frame, long after the relay built its logging metadata, so the + groups the gate matched have to be picked up at the end or the group pays for nothing it allowed.""" + metadata: Final = await _billed_metadata(RecordingFrameModelGate(groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",)}))) + + assert metadata[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"] + assert metadata["user_api_key"] == "hashed-key" + + +@pytest.mark.asyncio +async def test_websocket_session_that_matched_no_group_is_charged_to_none(): + metadata: Final = await _billed_metadata(RecordingFrameModelGate()) + + assert not metadata[MODEL_ACCESS_GROUP_METADATA_KEY] + + _GEMINI_FLASH_TARGET: Final = "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.8-flash:generateContent" _VERTEX_FLASH_TARGET: Final = ( "https://us-east5-aiplatform.googleapis.com/v1/projects/p/locations/us-east5" diff --git a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py index cf00545f331..4934573a777 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py @@ -652,6 +652,59 @@ def test_vertex_ai_live_route_without_usage_frames_yields_no_logging_response( assert normalized["kwargs"] == {"model": _LIVE_MODEL} +REALTIME_FRAMES: Final = [ + {"type": "session.created", "session": {"id": "sess_fake", "model": "gpt-realtime-2.1"}}, + { + "type": "response.done", + "response": { + "usage": { + "input_tokens": 33, + "output_tokens": 16, + "total_tokens": 49, + "input_token_details": {"text_tokens": 33, "audio_tokens": 0, "cached_tokens": 0}, + "output_token_details": {"text_tokens": 16, "audio_tokens": 0}, + } + }, + }, +] + + +def _normalize(url_route: str, response_body: dict[str, object] | list[dict[str, object]]) -> dict: + logging_obj: Final = MagicMock(model_call_details={}, litellm_call_id="call-fake") + return PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( + httpx_response=MagicMock(spec=httpx.Response), + response_body=response_body, + request_body={}, + logging_obj=logging_obj, + url_route=url_route, + result="websocket_connection_successful", + start_time=_START, + end_time=_END, + cache_hit=False, + litellm_call_id="call-fake", + ) + + +@pytest.mark.parametrize("url_route", ["/openai/v1/realtime", "/openai_passthrough/v1/realtime"]) +def test_openai_websocket_frames_reach_the_websocket_cost_handler(url_route): + normalized: Final = _normalize(url_route, REALTIME_FRAMES) + + assert normalized["kwargs"]["model"] == "gpt-realtime-2.1" + assert normalized["kwargs"]["custom_llm_provider"] == "openai" + prices: Final = litellm.model_cost["gpt-realtime-2.1"] + assert normalized["kwargs"]["response_cost"] == pytest.approx( + 33 * prices["input_cost_per_token"] + 16 * prices["output_cost_per_token"] + ) + assert normalized["standard_logging_response_object"].usage.total_tokens == 49 + + +def test_openai_http_passthrough_bodies_are_left_to_the_http_handlers(): + normalized: Final = _normalize("/openai/v1/responses", {"model": "gpt-5.4-mini", "usage": {"input_tokens": 13}}) + + assert normalized["standard_logging_response_object"] is None + assert normalized["kwargs"] == {"litellm_call_id": "call-fake"} + + @pytest.mark.asyncio @pytest.mark.parametrize("cancel_mode", ["once", "twice", "scope", "none"]) @pytest.mark.parametrize("cancel_phase", ["allocation", "lock"]) diff --git a/tests/unit/proxy/test_openai_ws_passthrough_routes.py b/tests/unit/proxy/test_openai_ws_passthrough_routes.py index 6a9cd972dc2..839ab49408f 100644 --- a/tests/unit/proxy/test_openai_ws_passthrough_routes.py +++ b/tests/unit/proxy/test_openai_ws_passthrough_routes.py @@ -1,31 +1,41 @@ -"""OpenAI passthrough WebSocket route: registration, opt-in gating, and refusals.""" +"""OpenAI passthrough WebSocket route: registration, opt-in gating, refusals, and per-frame model checks.""" import json -from collections.abc import Mapping, Sequence +from collections.abc import Mapping from dataclasses import dataclass from types import MappingProxyType, SimpleNamespace from typing import Final -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import httpx import pytest import respx +from fastapi import HTTPException from starlette.routing import WebSocketRoute +from starlette.websockets import WebSocket import litellm from litellm.llms.openai.workload_identity import _workload_identity_auth -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _OPENAI_WS_ALIASED_MODEL_REFUSAL, _OPENAI_WS_DISABLED_REFUSAL, - _OPENAI_WS_MODEL_RESTRICTED_REFUSAL, - _has_model_restrictions, + _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, _openai_websocket_refusal, - _proxy_model_allowlists, + _proxy_frame_model_gate, openai_websocket_proxy_route, router, ) - -Scopes = tuple[Sequence[str], ...] +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import ( + openai_websocket_frame_kind, +) +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + WebsocketFrameModelAllowed, + WebsocketFrameModelGate, + WebsocketFrameModelRefused, + WebsocketFrameModelVerdict, + WebsocketUpstreamFrameClassifier, +) ENABLED: Final = MappingProxyType({"enable_openai_websocket_passthrough": True}) DISABLED_SETTINGS: Final = ( @@ -76,6 +86,8 @@ class _RelayCall: forward_headers: bool endpoint: str accept_websocket: bool + client_frame_model_gate: WebsocketFrameModelGate | None + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None class _FakeRelay: @@ -92,6 +104,8 @@ class _FakeRelay: forward_headers: bool, endpoint: str, accept_websocket: bool, + client_frame_model_gate: WebsocketFrameModelGate | None = None, + upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None, ) -> None: self.calls.append( _RelayCall( @@ -100,24 +114,26 @@ class _FakeRelay: forward_headers=forward_headers, endpoint=endpoint, accept_websocket=accept_websocket, + client_frame_model_gate=client_frame_model_gate, + upstream_frame_kind=upstream_frame_kind, ) ) -class _FakeModelAllowlists: - def __init__(self, scopes: Scopes) -> None: - self.scopes = scopes - self.calls: list[UserAPIKeyAuth] = [] +class _FakeFrameModelGate: + def __init__(self, refusals: Mapping[str, WebsocketFrameModelRefused] = MappingProxyType({})) -> None: + self.refusals = refusals + self.checked: list[str] = [] - async def __call__(self, valid_token: UserAPIKeyAuth, /) -> Scopes: - self.calls.append(valid_token) - return self.scopes + async def __call__(self, websocket: object, model: str, /) -> WebsocketFrameModelVerdict: + self.checked.append(model) + return self.refusals.get(model) or WebsocketFrameModelAllowed(access_groups=()) @dataclass(frozen=True, slots=True) class _Served: relay: _FakeRelay - allowlists: _FakeModelAllowlists + frame_model_gate: _FakeFrameModelGate async def _serve( @@ -125,16 +141,16 @@ async def _serve( endpoint: str, user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object], - scopes: Scopes = (), + frame_model_gate: _FakeFrameModelGate | None = None, ) -> _Served: - served = _Served(relay=_FakeRelay(), allowlists=_FakeModelAllowlists(scopes)) + served: Final = _Served(relay=_FakeRelay(), frame_model_gate=frame_model_gate or _FakeFrameModelGate()) await openai_websocket_proxy_route( websocket=websocket, endpoint=endpoint, user_api_key_dict=user_api_key_dict, general_settings=general_settings, relay=served.relay, - model_allowlists=served.allowlists, + frame_model_gate=served.frame_model_gate, ) return served @@ -143,18 +159,20 @@ async def _serve( @pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"]) async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, monkeypatch): monkeypatch.delenv("OPENAI_API_BASE", raising=False) - websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview") + websocket: Final = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1") with patch(GET_CREDENTIALS, return_value="sk-provider"): served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) assert served.relay.calls == [ _RelayCall( - target="wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview", + target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1", custom_headers=MappingProxyType({"Authorization": "Bearer sk-provider"}), forward_headers=False, endpoint=f"/{prefix}/v1/realtime", accept_websocket=False, + client_frame_model_gate=served.frame_model_gate, + upstream_frame_kind=openai_websocket_frame_kind, ) ] assert websocket.accepts == [None] @@ -162,11 +180,23 @@ async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, m assert websocket.closed is None +@pytest.mark.asyncio +async def test_openai_websocket_reports_its_route_without_the_proxy_root_path(monkeypatch): + """Route allowlists and spend rows key on the route the proxy serves, which a root path prefix is not part of""" + monkeypatch.setenv("SERVER_ROOT_PATH", "/litellm") + websocket: Final = _FakeWebSocket("/litellm/openai/v1/responses", "") + + with patch(GET_CREDENTIALS, return_value="sk-provider"): + served: Final = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED) + + assert [call.endpoint for call in served.relay.calls] == ["/openai/v1/responses"] + + @pytest.mark.asyncio async def test_openai_websocket_accepts_first_client_subprotocol(): websocket = _FakeWebSocket( "/openai/v1/realtime", - "model=gpt-4o-realtime-preview", + "model=gpt-realtime-2.1", subprotocols="realtime, openai-insecure-api-key.sk-abc, openai-beta.realtime-v1", ) @@ -239,7 +269,7 @@ async def test_openai_websocket_closes_cleanly_when_workload_identity_exchange_f @pytest.mark.asyncio async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing(): - websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview") + websocket: Final = _FakeWebSocket("/openai/v1/realtime", "model=gpt-realtime-2.1") with patch(GET_CREDENTIALS, return_value=None): served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED) @@ -255,7 +285,7 @@ async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing @pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"]) @pytest.mark.parametrize("general_settings", DISABLED_SETTINGS) async def test_openai_websocket_refused_unless_explicitly_enabled(prefix, general_settings): - websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview") + websocket: Final = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1") served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), general_settings) @@ -265,25 +295,21 @@ async def test_openai_websocket_refused_unless_explicitly_enabled(prefix, genera assert served.relay.calls == [] -@pytest.mark.asyncio @pytest.mark.parametrize("general_settings", DISABLED_SETTINGS) -async def test_openai_websocket_refusal_is_disabled_for_falsy_settings(general_settings): - refusal = await _openai_websocket_refusal(UserAPIKeyAuth(), general_settings, _FakeModelAllowlists(())) - assert refusal is _OPENAI_WS_DISABLED_REFUSAL +def test_openai_websocket_refusal_is_disabled_for_falsy_settings(general_settings): + assert _openai_websocket_refusal(general_settings) is _OPENAI_WS_DISABLED_REFUSAL -@pytest.mark.asyncio @pytest.mark.parametrize("value", [True, "true", "True"]) -async def test_openai_websocket_refusal_is_none_for_truthy_settings(value): - settings = MappingProxyType({"enable_openai_websocket_passthrough": value}) - assert await _openai_websocket_refusal(UserAPIKeyAuth(), settings, _FakeModelAllowlists(())) is None +def test_openai_websocket_refusal_is_none_for_truthy_settings(value): + assert _openai_websocket_refusal(MappingProxyType({"enable_openai_websocket_passthrough": value})) is None @pytest.mark.asyncio async def test_openai_websocket_refusal_echoes_requested_subprotocol(): websocket = _FakeWebSocket( "/openai_passthrough/v1/realtime", - "model=gpt-4o-realtime-preview", + "model=gpt-realtime-2.1", subprotocols="realtime, openai-beta.realtime-v1", ) @@ -294,68 +320,303 @@ async def test_openai_websocket_refusal_echoes_requested_subprotocol(): assert served.relay.calls == [] -RESTRICTED_SCOPES: Final[tuple[Scopes, ...]] = ( - (("gpt-4o",),), - ((), ("gpt-4o-realtime-preview",)), - (("all-team-models",), ("gpt-4o",)), - ((), ("all-proxy-models",), ("gpt-4o",)), - ((), (), (), ("gpt-4o",)), - (("*",), (), (), (), ("gpt-4o",)), +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_models", + [["gpt-5.4-mini"], ["gpt-realtime-2.1", "gpt-5.4"], []], ) -UNRESTRICTED_SCOPES: Final[tuple[Scopes, ...]] = ( - (), - ((),), - (("all-proxy-models",),), - (("*",),), - (("all-team-models",), ("all-proxy-models",)), - ((), (), (), (), ()), - (("*",), ("all-proxy-models",), ("all-team-models",), (), ()), -) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("scopes", RESTRICTED_SCOPES) -async def test_openai_websocket_rejects_model_restricted_identities(scopes): - websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview") - user_api_key_dict = UserAPIKeyAuth(token="hashed-fake", user_id="user-fake", team_id="team-fake") - - served = await _serve(websocket, "v1/realtime", user_api_key_dict, ENABLED, scopes) - - assert "model restrictions" in websocket.error_message() - assert websocket.closed == (1008, _OPENAI_WS_MODEL_RESTRICTED_REFUSAL.close_reason) - assert served.relay.calls == [] - assert served.allowlists.calls == [user_api_key_dict] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("scopes", RESTRICTED_SCOPES) -async def test_openai_websocket_disabled_refusal_skips_allowlist_lookups(scopes): - allowlists = _FakeModelAllowlists(scopes) - - refusal = await _openai_websocket_refusal(UserAPIKeyAuth(), MappingProxyType({}), allowlists) - - assert refusal is _OPENAI_WS_DISABLED_REFUSAL - assert allowlists.calls == [] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("scopes", UNRESTRICTED_SCOPES) -async def test_openai_websocket_allows_unrestricted_identities(scopes): - websocket = _FakeWebSocket("/openai/v1/responses", "") +async def test_openai_websocket_connects_model_restricted_keys(key_models): + websocket: Final = _FakeWebSocket("/openai/v1/responses", "") + user_api_key_dict: Final = UserAPIKeyAuth(models=key_models, token="hashed-fake", user_id="user-fake") with patch(GET_CREDENTIALS, return_value="sk-provider"): - served = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED, scopes) + served: Final = await _serve(websocket, "v1/responses", user_api_key_dict, ENABLED) assert len(served.relay.calls) == 1 + assert served.relay.calls[0].client_frame_model_gate is served.frame_model_gate assert websocket.sent == [] assert websocket.closed is None -@pytest.mark.asyncio -async def test_proxy_model_allowlists_reads_the_token_scopes_without_a_database(): - token: Final = UserAPIKeyAuth(models=[], team_id="team-fake", team_models=["gpt-4o"]) - with patch("litellm.proxy.proxy_server.prisma_client", None): - scopes = await _proxy_model_allowlists()(token) +ALIAS_REFUSAL: Final = WebsocketFrameModelRefused( + message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=ProxyErrorTypes.key_model_access_denied.value +) - assert tuple(tuple(scope) for scope in scopes) == ((), ("gpt-4o",)) - assert _has_model_restrictions(scopes) + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "query", + [ + "model=gpt-realtime-2.1", + "model=gpt-realtime-2.1&model=gpt-realtime-2.1-mini", + "model=gpt-realtime-2.1-mini&model=gpt-realtime-2.1", + ], + ids=["one_model", "refused_model_first", "refused_model_last"], +) +async def test_openai_websocket_refuses_an_opening_model_the_frame_gate_refuses(query): + """ + The relay sends the connection URL to OpenAI unchanged, so every model it names has to pass the + same gate as a frame, or an alias the upgrade resolves to an allowed model opens the session on the + alias name itself + """ + websocket: Final = _FakeWebSocket("/openai/v1/realtime", query) + gate: Final = _FakeFrameModelGate(MappingProxyType({"gpt-realtime-2.1": ALIAS_REFUSAL})) + + with patch(GET_CREDENTIALS, return_value="sk-provider"): + served: Final = await _serve( + websocket, "v1/realtime", UserAPIKeyAuth(models=["gpt-realtime-2.1-mini"]), ENABLED, gate + ) + + assert websocket.error_message() == _OPENAI_WS_ALIASED_MODEL_REFUSAL + assert websocket.closed == (1008, "This key cannot use a model alias over OpenAI websocket passthrough.") + assert served.relay.calls == [] + assert "gpt-realtime-2.1" in gate.checked + + +@pytest.mark.asyncio +async def test_openai_websocket_relays_a_session_once_every_opening_model_passes_the_frame_gate(): + websocket: Final = _FakeWebSocket("/openai/v1/realtime", "model=gpt-realtime-2.1-mini") + + with patch(GET_CREDENTIALS, return_value="sk-provider"): + served: Final = await _serve( + websocket, "v1/realtime", UserAPIKeyAuth(models=["gpt-realtime-2.1-mini"]), ENABLED + ) + + assert served.frame_model_gate.checked == ["gpt-realtime-2.1-mini"] + assert len(served.relay.calls) == 1 + assert websocket.closed is None + + +class _FrameAuthorizer: + """Stands in for the full key, team, user and budget auth the gate re-runs for every model a frame names.""" + + def __init__(self, token: UserAPIKeyAuth | None = None, denial: Exception | None = None) -> None: + self.token = token or UserAPIKeyAuth() + self.denial = denial + self.authorized: list[str] = [] + + async def __call__(self, websocket: object, model: str) -> UserAPIKeyAuth: + self.authorized.append(model) + if self.denial is not None: + raise self.denial + return self.token + + +SOCKET: Final = WebSocket({"type": "websocket"}, receive=AsyncMock(), send=AsyncMock()) + + +async def _gate_verdict(authorizer: _FrameAuthorizer, model: str) -> WebsocketFrameModelVerdict: + with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the model checks need no database + return await _proxy_frame_model_gate(authorizer)(SOCKET, model) + + +def _refusal(verdict: WebsocketFrameModelVerdict) -> WebsocketFrameModelRefused: + assert isinstance(verdict, WebsocketFrameModelRefused) + return verdict + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_runs_full_auth_for_the_model_the_frame_names(): + authorizer: Final = _FrameAuthorizer(UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"])) + + verdict: Final = await _gate_verdict(authorizer, "gpt-realtime-2.1") + + assert verdict == WebsocketFrameModelAllowed(access_groups=()) + assert authorizer.authorized == ["gpt-realtime-2.1"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("requested_model", ["gpt-5.4", "gpt-realtime-2.1-mini"]) +async def test_proxy_frame_model_gate_refuses_a_model_the_key_lacks(requested_model): + authorizer: Final = _FrameAuthorizer(UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"])) + + refusal: Final = _refusal(await _gate_verdict(authorizer, requested_model)) + + assert requested_model in refusal.message + assert refusal.error_type == ProxyErrorTypes.key_model_access_denied.value + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "granted_models, refusal_names", + [(["gpt-5.4-mini"], "gpt-5.4"), ([], "no model grants")], + ids=["model_outside_the_grant", "no_grants"], +) +async def test_proxy_frame_model_gate_refuses_a_model_the_agent_was_not_granted(granted_models, refusal_names): + from litellm.types.agents import AgentResponse + + token: Final = UserAPIKeyAuth(models=[], agent_id="agent-fake") + token.managed_agent_policy = AgentResponse( + agent_id="agent-fake", + agent_name="Agent", + agent_card_params={}, + object_permission={"object_permission_id": "grant-fake", "models": granted_models}, + ) + + refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4")) + + assert refusal_names in refusal.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "token", + [ + UserAPIKeyAuth(models=["gpt-5.4-mini"], aliases={"gpt-5.4": "gpt-5.4-mini"}), + UserAPIKeyAuth(models=["gpt-5.4-mini"], team_model_aliases={"gpt-5.4": "gpt-5.4-mini"}), + ], + ids=["key_alias", "team_alias"], +) +async def test_proxy_frame_model_gate_refuses_an_alias_that_would_reach_openai_unchanged(token): + """ + The passthrough sends each frame to OpenAI as written, so OpenAI runs the + alias name itself. An alias pointing at an allowed model must not let a key + limited to that model run the model the alias is named after. + """ + refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4")) + + assert refusal == WebsocketFrameModelRefused( + message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=ProxyErrorTypes.key_model_access_denied.value + ) + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_refuses_a_proxy_wide_alias_for_a_key_limited_to_its_target(monkeypatch): + monkeypatch.setattr(litellm, "model_alias_map", {"gpt-5.4": "gpt-5.4-mini"}) + + refusal: Final = _refusal( + await _gate_verdict(_FrameAuthorizer(UserAPIKeyAuth(models=["gpt-5.4-mini"])), "gpt-5.4") + ) + + assert refusal.message == _OPENAI_WS_ALIASED_MODEL_REFUSAL + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_lets_a_key_that_may_use_any_model_send_an_alias(): + token: Final = UserAPIKeyAuth(models=[], aliases={"fast": "gpt-5.4-mini"}) + + verdict: Final = await _gate_verdict(_FrameAuthorizer(token), "fast") + + assert verdict == WebsocketFrameModelAllowed(access_groups=()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("unrestricted_models", [[], ["all-proxy-models"], ["*"]]) +async def test_proxy_frame_model_gate_allows_every_model_for_unrestricted_keys(unrestricted_models): + verdict: Final = await _gate_verdict(_FrameAuthorizer(UserAPIKeyAuth(models=unrestricted_models)), "gpt-5.4") + + assert verdict == WebsocketFrameModelAllowed(access_groups=()) + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_hands_back_the_access_groups_that_authorized_the_model(): + """The turn a frame starts is billed to the groups that let its model through, so the gate has to report them.""" + token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"], matched_model_access_groups=["tier-a"]) + + verdict: Final = await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4-mini") + + assert verdict == WebsocketFrameModelAllowed(access_groups=("tier-a",)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "requested_model", + ["gpt-5.4-mini", "gpt-5.4"], + ids=["allowed", "refused_by_the_model_check"], +) +async def test_proxy_frame_model_gate_releases_the_budget_it_reserved_for_the_check(requested_model): + """ + Full auth reserves budget for the request it authorizes, but a frame check is + not a request, so a reservation the gate leaves open would hold the key's + spend counter above its real spend and refuse later calls on every route. + """ + reservation: Final[dict[str, object]] = {"reserved_cost": 0.5, "entries": []} + token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"], budget_reservation=reservation) + + await _gate_verdict(_FrameAuthorizer(token), requested_model) + + assert token.budget_reservation == {**reservation, "finalized": True} + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_does_not_count_the_sessions_own_reservation_as_key_spend(): + """ + The upgrade reserves up to the key's remaining budget for the session, and the key budget check + counts open reservations as spend, so a frame checked while that reservation is open is refused + on a key with budget left, where an HTTP call with the same key goes through. + """ + connect_reservation: Final[dict[str, object]] = {"reserved_cost": 0.0001, "entries": []} + socket: Final = WebSocket( + {"type": "websocket", "state": {"budget_reservation": connect_reservation}}, receive=AsyncMock(), send=AsyncMock() + ) + + async def key_budget_check(websocket: object, model: str) -> UserAPIKeyAuth: + if connect_reservation.get("finalized") is not True: + raise litellm.BudgetExceededError(current_cost=0.0001, max_budget=0.0001) + return UserAPIKeyAuth(models=["gpt-5.4-mini"]) + + with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the model checks need no database + verdict: Final = await _proxy_frame_model_gate(key_budget_check)(socket, "gpt-5.4-mini") + + assert verdict == WebsocketFrameModelAllowed(access_groups=()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "denial, error_type, message", + [ + ( + ProxyException( + message="Budget has been exceeded! Key=hashed Spend=2.0, Budget=1.0", + type=ProxyErrorTypes.budget_exceeded, + param=None, + code=400, + ), + "budget_exceeded", + "Budget has been exceeded! Key=hashed", + ), + ( + litellm.BudgetExceededError( + current_cost=2.0, + max_budget=1.0, + message="Budget has been exceeded! Model access group=tier-a Current cost: 2.0, Max budget: 1.0", + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP.value, + entity_id="tier-a", + ), + "budget_exceeded", + "Model access group=tier-a", + ), + ( + ProxyException( + message="team not allowed to access model. This team can only access models=['gpt-5.4-mini']", + type=ProxyErrorTypes.team_model_access_denied, + param="model", + code=401, + ), + "team_model_access_denied", + "team not allowed to access model", + ), + (HTTPException(status_code=403, detail="Invalid Authorization header format"), "invalid_request_error", ""), + ], + ids=["key_budget", "group_budget", "team_allowlist", "malformed_auth"], +) +async def test_proxy_frame_model_gate_refuses_what_full_auth_denies_with_its_reason(denial, error_type, message): + """A frame is denied for the same reasons a request is, and the caller reads which reason in the error event.""" + refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(denial=denial), "gpt-5.4-mini")) + + assert refusal.error_type == error_type + assert message in refusal.message + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_refuses_out_loud_when_auth_cannot_answer(): + """ + A gate that cannot answer has to close the session with a reason the caller can read. + Letting the failure escape drops the socket with no frame at all, which the caller + cannot tell apart from the network going away. + """ + refusal: Final = _refusal( + await _gate_verdict(_FrameAuthorizer(denial=ConnectionError("connection pool is exhausted")), "gpt-5.4-mini") + ) + + assert refusal == WebsocketFrameModelRefused(message=_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, error_type="server_error") diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index cca138e0bb5..aa759f2c4d9 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -1451,6 +1451,38 @@ def test_transcription_usage_cost_returns_zero_for_unknown_type(): assert _transcription_usage_cost({}, {}) == 0.0 +def test_realtime_transcription_prices_each_clip_at_the_transcription_model_set_when_it_ran(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + from litellm.cost_calculator import handle_realtime_transcription_cost_calculation + + realtime_model = "gpt-realtime-2.1-mini" + first_asr, second_asr = "gpt-4o-transcribe", "gpt-transcribe" + first_rate = litellm.model_cost[first_asr]["input_cost_per_second"] + second_rate = litellm.model_cost[second_asr]["input_cost_per_second"] + assert first_rate != second_rate and first_rate > 0 and second_rate > 0 + + def session_frame(frame_type: str, asr_model: str | None) -> dict[str, object]: + transcription = None if asr_model is None else {"model": asr_model} + return {"type": frame_type, "session": {"model": realtime_model, "audio": {"input": {"transcription": transcription}}}} + + clip = {"type": "conversation.item.input_audio_transcription.completed", "usage": {"type": "duration", "seconds": 5}} + results: OpenAIRealtimeStreamList = [ + session_frame("session.created", None), + session_frame("session.updated", first_asr), + clip, + session_frame("session.updated", second_asr), + clip, + ] + + cost = handle_realtime_transcription_cost_calculation( + results=results, custom_llm_provider="openai", litellm_model_name=realtime_model + ) + + assert cost == pytest.approx(5 * first_rate + 5 * second_rate, rel=1e-9) + + def test_get_transcription_model_falls_back_to_session_model(monkeypatch): """session.model is used when transcription-specific model fields are absent.""" monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") From 16ee2c71a78738cd24861ff5a89225040bd6cb2f Mon Sep 17 00:00:00 2001 From: quietjules <349048734+quietjules@users.noreply.github.com> Date: Mon, 2 Nov 2026 09:55:38 -0800 Subject: [PATCH 08/13] fix(proxy): keep provider fields on streamed chat completion chunks (#49398) * fix(proxy): keep provider fields on streamed chat completion chunks The hand-written serializer for plain text chunks dropped every field it did not list, so Azure's per-choice content_filter_results, the chunk's obfuscation, and delta reasoning_details and reasoning_items never reached clients since v1.86.0. Every chunk now serializes with model_dump_json, which keeps them and is faster per chunk. * fix(proxy): send streamed chunk frames as bytes so plain text frames keep their type * fix(proxy): send null for chat chunk values JSON cannot encode instead of failing the stream * fix(proxy): drop chat chunk fields JSON cannot encode instead of failing the stream A callback or provider value pydantic cannot encode, like invalid UTF-8 bytes, a lone surrogate, a circular dict, or an unknown type, now leaves only that field out of the frame with a warning naming it. Responses API events keep raising. * fix(proxy): drop only the nested chunk value JSON cannot encode and keep lone surrogates Main removed the StreamingChoices import, so the field drop helpers now take any BaseModel. A value JSON cannot encode is left out down to the nested key or list item that holds it, and a lone surrogate goes out as an ASCII JSON escape. * fix(proxy): keep streaming when an unencodable value sits under non-string dict keys --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 141 ++++----- .../proxy_server/test_streaming_helpers.py | 286 ++++++++++++++---- .../proxy/test_response_model_sanitization.py | 33 +- 3 files changed, 286 insertions(+), 174 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c5fcbfb2564..7eace4ade9b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -61,6 +61,8 @@ import websockets import websockets.exceptions from pydantic import BaseModel, ConfigDict, Json, JsonValue, TypeAdapter, ValidationError, with_config from pydantic.fields import FieldInfo, PydanticUndefined +from pydantic.main import IncEx +from pydantic_core import PydanticSerializationError, to_jsonable_python from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never from litellm._uuid import uuid @@ -10186,6 +10188,8 @@ def _event_with_restamped_nested_model(event: Mapping[str, object], target_model _SSE_DATA_FIELD: Final = b"data:" _SSE_NESTED_RESPONSE_MARKER: Final = b'"response"' _SSE_EVENT_OBJECT: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_JSON_OBJECT: Final = TypeAdapter(dict[object, object]) +_JSON_ARRAY: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...]) def _restamp_nested_response_model_in_sse_frame(frame: bytes, target_model: str) -> bytes: @@ -10211,91 +10215,64 @@ def _restamp_nested_response_model_in_sse_line(line: bytes, target_model: str) - return _SSE_DATA_FIELD + b" " + serialized_event.encode() + line_ending -def _fast_serialize_simple_model_response_stream( - chunk: ModelResponseStream, -) -> bytes | None: - """ - Serialize the common OpenAI text streaming chunk without the full Pydantic - serializer. Fall back for richer chunks so tool calls, logprobs, usage, and - provider-specific fields keep the canonical model_dump_json behavior. - """ - if ( - getattr(chunk, "provider_specific_fields", None) is not None - or getattr(chunk, "system_fingerprint", None) is not None - or getattr(chunk, "usage", None) is not None - ): - return None - - choices: Final = getattr(chunk, "choices", None) - if not isinstance(choices, list) or len(choices) != 1: - return None - - choice: Final = choices[0] - if getattr(choice, "logprobs", None) is not None or getattr(choice, "enhancements", None) is not None: - return None - - delta: Final[object] = getattr(choice, "delta", None) - if delta is None: - return None - - unsupported_delta_fields: Final = ( - "function_call", - "tool_calls", - "audio", - "images", - "annotations", - "reasoning_content", - "thinking_blocks", - "provider_specific_fields", - "refusal", - ) - if any(getattr(delta, field, None) is not None for field in unsupported_delta_fields): - return None - - delta_dict: Final[dict] = {} - role: Final = getattr(delta, "role", None) - content: Final = getattr(delta, "content", None) - if role is not None: - delta_dict["role"] = role - if content is not None: - delta_dict["content"] = content - - choice_dict: Final = {"index": getattr(choice, "index", 0), "delta": delta_dict} - finish_reason: Final = getattr(choice, "finish_reason", None) - if finish_reason is not None: - choice_dict["finish_reason"] = finish_reason - - # Match the canonical ``model_dump_json(exclude_none=True)`` shape — if a - # field is None, omit it entirely rather than emitting ``"key": null``. - # Strict OpenAI-compatible clients reject ``null`` for optional fields like - # ``model``, so diverging here would surface as a client-side regression - # only on the fast path. Fall back to the slow path if a required-looking - # top-level identifier is missing. - model: Final = getattr(chunk, "model", None) - if model is None: - return None - - payload: Final[dict] = { - "id": getattr(chunk, "id", None), - "object": getattr(chunk, "object", None), - "created": getattr(chunk, "created", None), - "model": model, - "service_tier": getattr(chunk, "service_tier", None), - "choices": [choice_dict], - } - for top_level_key in ("id", "object", "created", "service_tier"): - if payload[top_level_key] is None: - payload.pop(top_level_key) - return orjson.dumps(payload) +def _json_encodes(value: object) -> bool: + try: + to_jsonable_python(value) + except ValueError: + return False + return True -def _serialize_streaming_chunk(chunk: BaseModel) -> str | bytes: +def _unencodable_part(value: object, ancestors: frozenset[int]) -> IncEx | bool: + if id(value) in ancestors: + return True + lineage: Final = ancestors | {id(value)} + if isinstance(value, BaseModel): + return _unencodable_keys(iter(value), lineage) or True + if isinstance(value, dict): + return _unencodable_entries(_JSON_OBJECT.validate_python(value), lineage) + if isinstance(value, (list, tuple)): + return _unencodable_indexes(enumerate(_JSON_ARRAY.validate_python(value)), lineage) or True + return True + + +def _unencodable_entries(entries: Mapping[object, object], lineage: frozenset[int]) -> IncEx | bool: + by_name: Final = {key: item for key, item in entries.items() if isinstance(key, str)} + if len(by_name) == len(entries): + return _unencodable_keys(by_name.items(), lineage) or True + by_index: Final = {key: item for key, item in entries.items() if isinstance(key, int)} + if len(by_index) == len(entries): + return _unencodable_indexes(by_index.items(), lineage) or True + return True + + +def _unencodable_keys(items: Iterable[tuple[str, object]], lineage: frozenset[int]) -> Mapping[str, IncEx | bool]: + return {key: _unencodable_part(item, lineage) for key, item in items if not _json_encodes(item)} + + +def _unencodable_indexes(items: Iterable[tuple[int, object]], lineage: frozenset[int]) -> Mapping[int, IncEx | bool]: + return {index: _unencodable_part(item, lineage) for index, item in items if not _json_encodes(item)} + + +def _serialize_chat_chunk_leniently(chunk: ModelResponseStream) -> bytes: + exclude: Final = _unencodable_keys(iter(chunk), frozenset({id(chunk)})) + if exclude: + verbose_proxy_logger.warning("Dropping streamed chunk fields that JSON cannot encode: %s", exclude) + jsonable: Final = chunk.model_dump(mode="json", exclude_none=True, exclude_unset=True, exclude=exclude) + return json.dumps(jsonable, ensure_ascii=True, separators=(",", ":")).encode() + + +def _serialize_chat_chunk(chunk: ModelResponseStream) -> bytes: + try: + return chunk.model_dump_json(exclude_none=True, exclude_unset=True).encode() + except PydanticSerializationError: + return _serialize_chat_chunk_leniently(chunk) + + +def _serialize_streaming_chunk(chunk: BaseModel) -> bytes: if isinstance(chunk, ModelResponseStream): - serialized_chunk: Final = _fast_serialize_simple_model_response_stream(chunk) - if serialized_chunk is not None: - return serialized_chunk - - return chunk.model_dump_json(exclude_none=True, exclude_unset=True) + return _serialize_chat_chunk(chunk) + return chunk.model_dump_json(exclude_none=True, exclude_unset=True).encode() async def _apply_streaming_chunk_hooks( diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index c37475bd539..c7870dec17a 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -5,7 +5,6 @@ Pins covered: - ``async_assistants_data_generator`` - ``_get_client_requested_model_for_streaming`` - ``_restamp_streaming_chunk_model`` -- ``_fast_serialize_simple_model_response_stream`` - ``_serialize_streaming_chunk`` - ``_apply_streaming_chunk_hooks`` - ``_format_streaming_sse_chunk`` @@ -16,8 +15,11 @@ Pins covered: from __future__ import annotations import asyncio +import hashlib import json +import logging from collections.abc import AsyncIterator, Callable +from types import MappingProxyType from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock @@ -36,7 +38,6 @@ from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.common_request_processing import create_response from litellm.proxy.proxy_server import ( _apply_streaming_chunk_hooks, - _fast_serialize_simple_model_response_stream, _format_fallback_metadata_sse_event, _format_streaming_sse_chunk, _get_client_requested_model_for_streaming, @@ -55,7 +56,14 @@ from litellm.types.llms.openai import ( ResponseFailedEvent, ResponsesAPIResponse, ) -from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage +from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + Usage, +) from tests._support.clock_loops import AutojumpClockLoop from .conftest import normalize, route_entry_points_to_stubs @@ -81,6 +89,10 @@ def _simple_chunk(model: str = "gpt-4", content: str = "hi") -> ModelResponseStr ) +def _sse_text(frame: str | bytes) -> str: + return frame.decode() if isinstance(frame, bytes) else frame + + async def _async_iter(items): for it in items: yield it @@ -520,48 +532,13 @@ def test_format_fallback_metadata_sse_event(): assert isinstance(payload["created"], int) -# --------------------------------------------------------------------------- -# _fast_serialize_simple_model_response_stream -# --------------------------------------------------------------------------- - - -def test_fast_serialize_simple_model_response_stream_returns_bytes_payload(): - chunk = _simple_chunk() - result = _fast_serialize_simple_model_response_stream(chunk) - assert isinstance(result, bytes) - payload = json.loads(result) - assert normalize(payload) == { - "id": "", - "object": "chat.completion.chunk", - "created": "", - "model": "gpt-4", - "choices": [ - { - "index": 0, - "delta": {"role": "assistant", "content": "hi"}, - } - ], - } - - -def test_fast_serialize_simple_model_response_stream_with_usage_returns_none_invalid(): - """Fast path bails (returns None) when ``usage`` is populated — the slow - path is required to preserve usage fields. Returning None here is the - "I cannot handle this" sentinel, not a hard error.""" - chunk = _simple_chunk() - chunk.usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) - assert _fast_serialize_simple_model_response_stream(chunk) is None - - # --------------------------------------------------------------------------- # _serialize_streaming_chunk # --------------------------------------------------------------------------- -def test_serialize_streaming_chunk_simple_uses_fast_path_bytes(): - result = _serialize_streaming_chunk(_simple_chunk()) - assert isinstance(result, bytes) - payload = json.loads(result) +def test_serialize_streaming_chunk_simple_payload(): + payload = json.loads(_serialize_streaming_chunk(_simple_chunk())) assert normalize(payload) == { "id": "", "object": "chat.completion.chunk", @@ -700,10 +677,7 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch): out.append(line) assert out[-1] == "data: [DONE]\n\n" - # First chunk is bytes (fast path) wrapped via _format_streaming_sse_chunk. - first = out[0] - assert isinstance(first, bytes) - payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n")) + payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n")) assert normalize(payload) == { "id": "", "object": "chat.completion.chunk", @@ -718,6 +692,205 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch): } +_AZURE_CONTENT_FILTER_RESULTS: Final = { + "hate": {"filtered": False, "severity": "safe"}, + "protected_material_text": {"detected": False, "filtered": False}, +} + + +def _azure_chunk(choice_fields: dict[str, object], chunk_fields: dict[str, object]) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-azure", + choices=[StreamingChoices(index=0, delta=Delta(content="Red"), **choice_fields)], + created=0, + model="gpt-4", + object="chat.completion.chunk", + **chunk_fields, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("chunk", "read_field", "sent"), + [ + pytest.param( + _azure_chunk({"content_filter_results": _AZURE_CONTENT_FILTER_RESULTS}, {}), + lambda payload: payload["choices"][0]["content_filter_results"], + _AZURE_CONTENT_FILTER_RESULTS, + id="choice-content-filter-results", + ), + pytest.param( + _azure_chunk({}, {"obfuscation": "Ab3"}), + lambda payload: payload["obfuscation"], + "Ab3", + id="chunk-obfuscation", + ), + pytest.param( + ModelResponseStream( + id="chatcmpl-or", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="Red", reasoning_details=[{"type": "reasoning.text", "text": "colors"}]), + ) + ], + created=0, + model="gpt-4", + object="chat.completion.chunk", + ), + lambda payload: payload["choices"][0]["delta"]["reasoning_details"], + [{"type": "reasoning.text", "text": "colors"}], + id="delta-reasoning-details", + ), + pytest.param( + ModelResponseStream( + id="chatcmpl-resp", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="Red", reasoning_items=[{"type": "reasoning", "id": "rs_1", "summary": []}]), + ) + ], + created=0, + model="gpt-4", + object="chat.completion.chunk", + ), + lambda payload: payload["choices"][0]["delta"]["reasoning_items"], + [{"type": "reasoning", "id": "rs_1", "summary": []}], + id="delta-reasoning-items", + ), + ], +) +async def test_async_data_generator_sse_frame_keeps_provider_fields_on_a_text_chunk( + monkeypatch, chunk: ModelResponseStream, read_field: Callable[[dict[str, object]], object], sent: object +): + _patch_logging_flags(monkeypatch) + + out = [ + line + async for line in async_data_generator( + response=_async_iter([chunk]), + user_api_key_dict=_user_auth(), + request_data={"model": "gpt-4"}, + ) + ] + + assert read_field(json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n"))) == sent + + +def _circular_dict() -> dict[str, object]: + looped: Final[dict[str, object]] = {"keep": 1} # mutable-ok: a circular dict has to be inserted into itself + looped["self"] = looped + return looped + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "unencodable", + [ + pytest.param(object(), id="unknown-type"), + pytest.param(hashlib.sha256(b"trace").digest(), id="invalid-utf8-bytes"), + ], +) +async def test_async_data_generator_streams_a_text_chunk_a_hook_tagged_with_a_value_json_cannot_encode( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, unencodable: object +) -> None: + _patch_logging_flags(monkeypatch) + chunk: Final = _azure_chunk({"content_filter_results": _AZURE_CONTENT_FILTER_RESULTS}, {"obfuscation": "Ab3"}) + setattr(chunk, "trace_chunk", unencodable) + setattr(chunk.choices[0], "trace_choice", unencodable) + setattr(chunk.choices[0].delta, "trace_delta", unencodable) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out: Final = [ + line + async for line in async_data_generator( + response=_async_iter([chunk]), + user_api_key_dict=_user_auth(), + request_data={"model": "gpt-4"}, + ) + ] + + payload: Final = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n")) + assert payload["obfuscation"] == "Ab3" + assert payload["choices"][0]["content_filter_results"] == _AZURE_CONTENT_FILTER_RESULTS + assert payload["choices"][0]["delta"] == {"content": "Red"} + assert "trace_chunk" not in payload + assert "trace_choice" not in payload["choices"][0] + assert all(name in caplog.text for name in ("trace_chunk", "trace_choice", "trace_delta")) + + +def test_serialize_streaming_chunk_keeps_lone_surrogates_a_provider_split_across_frames( + caplog: pytest.LogCaptureFixture, +) -> None: + tool_call: Final = ChatCompletionDeltaToolCall( + id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Par\udc00is"}') + ) + chunk: Final = ModelResponseStream( + id="chatcmpl-split", + choices=[StreamingChoices(index=0, delta=Delta(content="Hi \ud83d", tool_calls=[tool_call]))], + created=0, + model="gpt-4", + object="chat.completion.chunk", + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delta: Final = json.loads(_serialize_streaming_chunk(chunk))["choices"][0]["delta"] + + assert caplog.text == "" + assert delta["content"] == "Hi \ud83d" + assert delta["tool_calls"][0]["id"] == "call_1" + assert delta["tool_calls"][0]["function"] == {"name": "get_weather", "arguments": '{"city": "Par\udc00is"}'} + + +def test_serialize_streaming_chunk_drops_only_the_nested_value_json_cannot_encode() -> None: + digest: Final = hashlib.sha256(b"trace").digest() + chunk: Final = ModelResponseStream( + id="chatcmpl-nested", + choices=[ + StreamingChoices(index=0, delta=Delta(content="a")), + StreamingChoices( + index=1, delta=Delta(content="b", provider_specific_fields={"citations": ["ok"], "raw": digest}) + ), + ], + created=0, + model="gpt-4", + object="chat.completion.chunk", + usage=Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5), + ) + setattr(chunk.usage, "raw", digest) + setattr(chunk, "trace_chunk", _circular_dict()) + + payload: Final = json.loads(_serialize_streaming_chunk(chunk)) + + assert payload["choices"][0]["delta"] == {"content": "a"} + assert payload["choices"][1]["delta"] == {"content": "b", "provider_specific_fields": {"citations": ["ok"]}} + assert payload["usage"] == {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5} + assert payload["trace_chunk"] == {"keep": 1} + + +def test_serialize_streaming_chunk_keeps_streaming_when_unencodable_values_sit_under_non_string_keys() -> None: + chunk: Final = ModelResponseStream( + id="chatcmpl-keys", + choices=[StreamingChoices(index=0, delta=Delta(content="still streaming"))], + created=0, + model="gpt-4", + object="chat.completion.chunk", + ) + setattr(chunk, "by_index", {0: "kept", 1: object()}) + setattr(chunk, "mixed_keys", {"name": "x", 2: object()}) + setattr(chunk, "bytes_keys", {b"raw": object()}) + setattr(chunk, "read_only_view", MappingProxyType({"name": object()})) + + payload: Final = json.loads(_serialize_streaming_chunk(chunk)) + + assert payload["choices"][0]["delta"] == {"content": "still streaming"} + assert payload["by_index"] == {"0": "kept"} + assert "mixed_keys" not in payload + assert "bytes_keys" not in payload + assert "read_only_view" not in payload + + @pytest.mark.asyncio async def test_async_data_generator_uses_response_fallback_metadata(monkeypatch): _patch_logging_flags(monkeypatch) @@ -739,9 +912,7 @@ async def test_async_data_generator_uses_response_fallback_metadata(monkeypatch) ): out.append(line) - first = out[0] - assert isinstance(first, bytes) - payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n")) + payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n")) assert payload["model"] == "fallback-model" @@ -764,9 +935,7 @@ async def test_async_data_generator_uses_chunk_fallback_metadata(monkeypatch): ): out.append(line) - first = out[0] - assert isinstance(first, bytes) - payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n")) + payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n")) assert payload["model"] == "fallback-model" @@ -793,8 +962,8 @@ async def test_async_data_generator_switches_model_mid_stream_on_fallback(monkey ): out.append(line) - first_payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n")) - second_payload = json.loads(out[1].removeprefix(b"data: ").removesuffix(b"\n\n")) + first_payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n")) + second_payload = json.loads(_sse_text(out[1]).removeprefix("data: ").removesuffix("\n\n")) assert first_payload["model"] == "primary-model" assert second_payload["model"] == "fallback-model" @@ -830,8 +999,7 @@ async def test_async_data_generator_emits_fallback_error_metadata_event(monkeypa ): out.append(line) - assert isinstance(out[0], bytes) - chunk_payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n")) + chunk_payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n")) assert chunk_payload["model"] == "fallback-model" assert isinstance(out[1], str) assert out[1].startswith("data: ") @@ -879,8 +1047,7 @@ async def test_async_data_generator_skips_fallback_error_event_without_opt_in( ): out.append(line) - assert isinstance(out[0], bytes) - payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n")) + payload = json.loads(_sse_text(out[0]).removeprefix("data: ").removesuffix("\n\n")) assert payload["model"] == "fallback-model" @@ -2434,14 +2601,11 @@ async def test_queue_request_stream_is_untouched_while_keepalives_are_unconfigur assert chunks[-1] == b"data: [DONE]\n\n" -def test_fast_serialize_simple_model_response_stream_keeps_served_service_tier(): +def test_serialize_streaming_chunk_keeps_served_service_tier(): chunk = _simple_chunk() chunk.service_tier = "priority" - result = _fast_serialize_simple_model_response_stream(chunk) - - assert result is not None - assert json.loads(result)["service_tier"] == "priority" + assert json.loads(_serialize_streaming_chunk(chunk))["service_tier"] == "priority" @pytest.mark.asyncio diff --git a/tests/unit/proxy/test_response_model_sanitization.py b/tests/unit/proxy/test_response_model_sanitization.py index f52736453c1..b3ad53221eb 100644 --- a/tests/unit/proxy/test_response_model_sanitization.py +++ b/tests/unit/proxy/test_response_model_sanitization.py @@ -98,42 +98,13 @@ def test_restamp_streaming_chunk_skips_matching_model(): assert model_mismatch_logged is False -def test_fast_serialize_simple_streaming_chunk_matches_model_dump_json(): +def test_serialize_streaming_chunk_omits_a_missing_model_instead_of_null(): from litellm.proxy.proxy_server import _serialize_streaming_chunk - chunk = _make_model_response_stream_chunk("client-model") - - assert json.loads(_serialize_streaming_chunk(chunk)) == json.loads( - chunk.model_dump_json(exclude_none=True, exclude_unset=True) - ) - - -def test_fast_serialize_returns_none_when_model_field_is_missing(): - """ - The fast path must mirror ``model_dump_json(exclude_none=True)``: when - ``chunk.model`` is ``None`` the slow path omits the field entirely. - Emitting ``"model": null`` would diverge and trip strict OpenAI- - compatible clients that reject ``null`` for optional string fields. - Falling back to ``None`` lets the canonical serializer handle the edge. - """ - from litellm.proxy.proxy_server import ( - _fast_serialize_simple_model_response_stream, - _serialize_streaming_chunk, - ) - chunk = _make_model_response_stream_chunk("client-model") chunk.model = None # type: ignore[assignment] - assert _fast_serialize_simple_model_response_stream(chunk) is None - - # Going through the public ``_serialize_streaming_chunk`` should still - # produce a serialized result via the slow-path fallback, and it must - # not contain ``"model": null``. - serialized = _serialize_streaming_chunk(chunk) - payload_str = serialized.decode("utf-8") if isinstance(serialized, bytes) else serialized - assert '"model": null' not in payload_str - assert '"model":null' not in payload_str - assert json.loads(payload_str) == json.loads(chunk.model_dump_json(exclude_none=True, exclude_unset=True)) + assert "model" not in json.loads(_serialize_streaming_chunk(chunk)) def test_proxy_chat_completion_does_not_return_provider_prefixed_model(tmp_path, monkeypatch): From a0b1edc0a0917fc41b44821b9a7bea361beb553f Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Mon, 2 Nov 2026 10:00:00 -0800 Subject: [PATCH 09/13] fix(providers): trim prompt_cache_breakpoint to what each API takes under drop_params (#49064) * fix(openai): trim a malformed prompt_cache_breakpoint under drop_params on chat completions * fix(openai): leave a non-object content part for OpenAI to reject under drop_params * fix(openai): trim an unsupported prompt_cache_breakpoint ttl on every path and on azure_ai chat under drop_params * fix: trim prompt_cache_breakpoint to what Azure and Mantle accept under drop_params Azure OpenAI chat now gets the same content part trim as OpenAI chat, and the Azure and Foundry gates also read a deployment's declared base_model. Each Responses config keeps only the ttl its API takes: Azure keeps 30m, and Mantle keeps no ttl. A trimmed ttl is logged at debug. * fix(responses): trim prompt_cache_breakpoint only on APIs known to read it Other Responses providers get the marker whole or not at all, as on main. Azure chat now trims for a deployment with no base_model unless the cost map marks the model as not reading breakpoints, and Azure AI Foundry Responses keeps Azure's 30m-only shape. * fix(azure): trim prompt cache breakpoints only for deployments known to read them Azure and Azure AI Responses keep main's whole-or-nothing rule for unmapped deployment names, an opted-in custom OpenAI host trims the ttl and keeps the mode, and Azure chat strips every marker for a known non-reader under drop_params * fix(providers): keep main's marker handling where the target won't trim it A /v1/responses model with unknown breakpoint support on OpenAI's host gets a marker whole or not at all, as on main. The chat to Responses bridge keeps main's whole-marker check when a non-bool global drop_params leaves the target's trim off. A hyphenated name like gpt-5-6-sol now reads as gpt-5.6. * fix(providers): read a non-version hyphen suffix like -32k as part of the model name * fix(providers): keep main's reading of a dotted minor version like gpt-5.6o * fix(providers): keep main's breakpoint shape for unmapped Mantle models and skip the Azure lookup without drop_params --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../transformation.py | 47 +++++-- .../anthropic_cache_control_hook.py | 30 ++++- .../prompt_templates/common_utils.py | 85 ++++++++++++- .../messages/responses_adapters/handler.py | 4 +- .../chat/gpt_reasoning_transformation.py | 4 +- litellm/llms/azure/chat/gpt_transformation.py | 21 +++- .../llms/azure/responses/transformation.py | 12 +- .../responses/transformation.py | 15 ++- .../llms/openai/chat/gpt_transformation.py | 18 ++- .../llms/openai/responses/transformation.py | 64 +++++----- litellm/types/llms/openai.py | 9 ++ ...penai_chat_prompt_cache_breakpoint_wire.py | 57 +++++++++ ...ses_bridge_prompt_cache_breakpoint_wire.py | 24 +++- ...responses_transformation_transformation.py | 47 ++++++- .../test_anthropic_cache_control_hook.py | 11 ++ ...ore_utils_prompt_templates_common_utils.py | 18 ++- .../test_azure_chat_gpt_transformation.py | 86 +++++++++++++ .../chat/test_azure_ai_transformation.py | 6 + ...bedrock_mantle_responses_transformation.py | 65 ++++++++++ .../chat/test_openai_gpt_transformation.py | 119 ++++++++++++++++++ .../test_openai_responses_transformation.py | 112 ++++++++++++++++- 21 files changed, 777 insertions(+), 77 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index bb59b4617af..c9d9cc25049 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -34,10 +34,12 @@ from litellm.completion_extras.litellm_responses_transformation.history import ( requested_phase, ) from litellm.integrations.anthropic_cache_control_hook import supports_openai_prompt_cache_breakpoint +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.hidden_params import get_hidden_params, get_or_create_hidden_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( - prompt_cache_breakpoint_for_wire, + prompt_cache_breakpoint_with_readable_mode, responses_reasoning_items_from_thinking_blocks, + whole_prompt_cache_breakpoint_or_none, with_prompt_cache_breakpoint, ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator @@ -97,6 +99,12 @@ _RESPONSES_API_ONLY_FIELDS: Final = frozenset((*Response.model_fields, *Response ) +def _bridged_prompt_cache_breakpoint(marker: object, drop_params: bool, target_trims_breakpoints: bool) -> object: + if target_trims_breakpoints: + return prompt_cache_breakpoint_with_readable_mode(marker, drop_params) + return whole_prompt_cache_breakpoint_or_none(marker) if drop_params else marker + + def _strip_prompt_cache_breakpoint_from_content_block(value: object) -> object: if not isinstance(value, dict): return value @@ -844,6 +852,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): messages: list["AllMessageValues"], *, drop_params: bool = False, + target_trims_breakpoints: bool = True, keep_prompt_cache_breakpoints: bool = False, replay_tool_search: bool = False, include_namespace: bool = False, @@ -852,6 +861,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): converted_input_items, instructions = self._convert_chat_completion_messages_to_responses_input( messages, drop_params=drop_params, + target_trims_breakpoints=target_trims_breakpoints, replay_tool_search=replay_tool_search, include_namespace=include_namespace, replay_openai_history=replay_openai_history, @@ -868,6 +878,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): messages: list["AllMessageValues"], *, drop_params: bool = False, + target_trims_breakpoints: bool = True, replay_tool_search: bool = False, include_namespace: bool = False, replay_openai_history: bool = False, @@ -921,6 +932,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): content, role, drop_params=drop_params, + target_trims_breakpoints=target_trims_breakpoints, ), } ) @@ -940,6 +952,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): content, "user", # Use "user" role to get input_* types drop_params=drop_params, + target_trims_breakpoints=target_trims_breakpoints, ) else: # Fallback: convert unexpected types to input_text @@ -968,7 +981,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): "type": "message", "role": "assistant", "content": self._convert_content_to_responses_format( - content, "assistant", drop_params=drop_params + content, + "assistant", + drop_params=drop_params, + target_trims_breakpoints=target_trims_breakpoints, ), **_assistant_phase_field(msg, replay_openai_history, drop_params), } @@ -1009,7 +1025,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): "type": "message", "role": role, "content": self._convert_content_to_responses_format( - content, cast(str, role), drop_params=drop_params + content, + cast(str, role), + drop_params=drop_params, + target_trims_breakpoints=target_trims_breakpoints, ), **_assistant_phase_field(msg, replay_openai_history and role == "assistant", drop_params), } @@ -1135,6 +1154,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): converted_input_items, converted_instructions = self.convert_chat_completion_messages_to_responses_api( messages, drop_params=bool(litellm_params.get("drop_params") or litellm.drop_params), + target_trims_breakpoints=drops_unsupported_params(litellm_params.get("drop_params")), keep_prompt_cache_breakpoints=supports_prompt_cache_breakpoint, replay_tool_search=_requests_tool_search(cast(object, optional_params)), # cast-ok: upcast include_namespace=True, @@ -1640,6 +1660,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): | None, role: str, drop_params: bool = False, + target_trims_breakpoints: bool = True, ) -> list[dict[str, object]]: """Convert chat completion content to responses API format""" from litellm.types.llms.openai import ChatCompletionImageObject @@ -1666,7 +1687,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if original_type == "text": converted = with_prompt_cache_breakpoint( self._convert_content_str_to_input_text(item.get("text", ""), role), - prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), + _bridged_prompt_cache_breakpoint( + item.get("prompt_cache_breakpoint"), drop_params, target_trims_breakpoints + ), ) result.append(converted) verbose_logger.debug("Chat provider: text -> %s", converted) @@ -1679,14 +1702,18 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): cast(ChatCompletionImageObject, item), role ), ), - prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), + _bridged_prompt_cache_breakpoint( + item.get("prompt_cache_breakpoint"), drop_params, target_trims_breakpoints + ), ) result.append(converted) verbose_logger.debug("Chat provider: image_url -> %s", converted) elif original_type == "input_audio": converted = with_prompt_cache_breakpoint( {"type": "input_audio", "input_audio": item.get("input_audio")}, - prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), + _bridged_prompt_cache_breakpoint( + item.get("prompt_cache_breakpoint"), drop_params, target_trims_breakpoints + ), ) result.append(converted) verbose_logger.debug("Chat provider: input_audio -> %s", converted) @@ -1702,7 +1729,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): _input_file_from_file_value( cast("ChatCompletionFileObject", item).get("file"), # cast-ok: type tag checked ), - prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), + _bridged_prompt_cache_breakpoint( + item.get("prompt_cache_breakpoint"), drop_params, target_trims_breakpoints + ), ) result.append(converted) verbose_logger.debug("Chat provider: file -> %s", converted) @@ -1726,7 +1755,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Default to input_text for unknown types converted = with_prompt_cache_breakpoint( self._convert_content_str_to_input_text(str(item.get("text", item)), role), - prompt_cache_breakpoint_for_wire(item.get("prompt_cache_breakpoint"), drop_params), + _bridged_prompt_cache_breakpoint( + item.get("prompt_cache_breakpoint"), drop_params, target_trims_breakpoints + ), ) result.append(converted) verbose_logger.debug("Chat provider: unknown(%s) -> %s", original_type, converted) diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index c05407b7972..01be455e229 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -17,7 +17,7 @@ from collections.abc import Callable, Container, Iterable, Iterator, Mapping, Se from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypeVar, cast from urllib.parse import urlparse -from pydantic import TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger @@ -73,7 +73,7 @@ MAX_CACHE_CONTROL_BLOCKS: Final = 4 CACHE_BREAKPOINT_KEYS: Final = ("cache_control", "prompt_cache_breakpoint") OPENAI_PROMPT_CACHE_BREAKPOINT_MIN_GPT_VERSION: Final = (5, 6) -_GPT_VERSION_PATTERN: Final = re.compile(r"^gpt-(\d+)(?:\.(\d+))?") +_GPT_VERSION_PATTERN: Final = re.compile(r"^gpt-(\d+)(?:\.(\d+)|-(0|[1-9]\d?)(?![\da-z]))?") _FINE_TUNED_BASE_PATTERN: Final = re.compile(r"^ft:([^:]+)") _O_SERIES_PATTERN: Final = re.compile(r"^o\d+(?:-|$)") OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES: Final = frozenset( @@ -335,6 +335,30 @@ def openai_prompt_cache_breakpoint_support(model: str, custom_llm_provider: str return _named_model_breakpoint_support(model) +class _DeclaredBaseModels(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True) + + base_model: object = None + model_info: object = None + + def names(self) -> tuple[str, ...]: + model_info: Final = _validated_object_mapping(self.model_info) or {} + candidates: Final = (self.base_model, model_info.get("base_model")) + return tuple(name for name in candidates if isinstance(name, str) and name) + + +def deployment_prompt_cache_breakpoint_support( + model: str, custom_llm_provider: str, litellm_params: object +) -> bool | None: + """``openai_prompt_cache_breakpoint_support`` for a deployment, read from its declared base model first.""" + declared: Final = _DeclaredBaseModels.model_validate(litellm_params) + supports: Final = ( + openai_prompt_cache_breakpoint_support(name, custom_llm_provider) + for name in dict.fromkeys((*declared.names(), model)) + ) + return next((support for support in supports if support is not None), None) + + def _named_model_breakpoint_support(model: str) -> bool | None: bare_model: Final = model.rsplit("/", 1)[-1].lower() if _O_SERIES_PATTERN.match(bare_model) is not None: @@ -342,7 +366,7 @@ def _named_model_breakpoint_support(model: str) -> bool | None: version_match: Final = _GPT_VERSION_PATTERN.match(bare_model) if version_match is None: return None - version: Final = (int(version_match.group(1)), int(version_match.group(2) or 0)) + version: Final = (int(version_match.group(1)), int(version_match.group(2) or version_match.group(3) or 0)) return version >= OPENAI_PROMPT_CACHE_BREAKPOINT_MIN_GPT_VERSION diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 938ab585b72..6287dc17e8f 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -29,6 +29,7 @@ from litellm.types.llms.openai import ( ChatCompletionAudioObject, ChatCompletionFileObject, ChatCompletionImageObject, + ChatCompletionPromptCacheBreakpoint, ChatCompletionReasoningItem, ChatCompletionReasoningSummaryTextBlock, ChatCompletionRedactedThinkingBlock, @@ -2193,16 +2194,62 @@ def with_prompt_cache_breakpoint(target: _MarkedT, marker: object) -> _MarkedT: _PROMPT_CACHE_BREAKPOINT: Final = TypeAdapter(PromptCacheBreakpoint) +_CHAT_COMPLETION_PROMPT_CACHE_BREAKPOINT: Final = TypeAdapter(ChatCompletionPromptCacheBreakpoint) +_AcceptedBreakpointT: Final = TypeVar("_AcceptedBreakpointT", bound=Mapping[str, object]) + + +def _accepted_prompt_cache_breakpoint( + accepted: TypeAdapter[_AcceptedBreakpointT], marker: object +) -> _AcceptedBreakpointT | None: + try: + return accepted.validate_python(marker) + except ValidationError: + return None + + +def prompt_cache_breakpoint_within( + accepted: TypeAdapter[_AcceptedBreakpointT], marker: object +) -> _AcceptedBreakpointT | None: + """The client's marker cut down to the ``accepted`` shape: an unsupported ttl is dropped and the mode kept, + and a marker whose mode the target cannot read is dropped whole.""" + whole: Final = _accepted_prompt_cache_breakpoint(accepted, marker) + if whole is not None: + return whole + fields: Final = _breakpoint_fields_or_none(marker) + without_ttl: Final = ( + None + if fields is None + else _accepted_prompt_cache_breakpoint(accepted, {key: item for key, item in fields.items() if key != "ttl"}) + ) + verbose_logger.debug( + "Dropping malformed prompt_cache_breakpoint %r under drop_params" + if without_ttl is None + else "Dropping the unsupported ttl of prompt_cache_breakpoint %r under drop_params", + marker, + ) + return without_ttl def prompt_cache_breakpoint_for_wire(marker: object, drop_params: bool) -> object: if marker is None or not drop_params: return marker - try: - return _PROMPT_CACHE_BREAKPOINT.validate_python(marker) - except ValidationError: - verbose_logger.debug("Dropping malformed prompt_cache_breakpoint %r under drop_params", marker) - return None + return prompt_cache_breakpoint_within(_PROMPT_CACHE_BREAKPOINT, marker) + + +def prompt_cache_breakpoint_with_readable_mode(marker: object, drop_params: bool) -> object: + """The client's marker as written, or None under drop_params when no target could read its mode, so the + target's own config decides which ttl and keys it keeps.""" + return None if prompt_cache_breakpoint_for_wire(marker, drop_params) is None else marker + + +def whole_prompt_cache_breakpoint_or_none(marker: object) -> PromptCacheBreakpoint | None: + return _accepted_prompt_cache_breakpoint(_PROMPT_CACHE_BREAKPOINT, marker) + + +def chat_completion_prompt_cache_breakpoint_for_wire(marker: object, drop_params: bool) -> object: + if marker is None or not drop_params: + return marker + return prompt_cache_breakpoint_within(_CHAT_COMPLETION_PROMPT_CACHE_BREAKPOINT, marker) def cache_control_as_prompt_cache_breakpoint(block: _MarkedT) -> _MarkedT: @@ -2220,6 +2267,13 @@ _BREAKPOINT_FIELDS: Final = TypeAdapter(Mapping[str, object]) _BREAKPOINT_BLOCKS: Final = TypeAdapter(list[object]) +def _breakpoint_fields_or_none(value: object) -> Mapping[str, object] | None: + try: + return _BREAKPOINT_FIELDS.validate_python(value) + except ValidationError: + return None + + def _without_prompt_cache_breakpoint(value: object) -> object: if not isinstance(value, Mapping): return value @@ -2242,6 +2296,27 @@ def without_prompt_cache_breakpoints(block: _MarkedT) -> _MarkedT: return cast(_MarkedT, stripped) # cast-ok: same block shape as the input minus the marker keys +def _part_with_chat_completion_prompt_cache_breakpoint(part: object) -> object: + if not isinstance(part, Mapping): + return part + fields: Final = _BREAKPOINT_FIELDS.validate_python(part) + marker: Final = fields.get("prompt_cache_breakpoint") + if marker is None: + return fields + unmarked: Final = {key: item for key, item in fields.items() if key != "prompt_cache_breakpoint"} + return with_prompt_cache_breakpoint(unmarked, chat_completion_prompt_cache_breakpoint_for_wire(marker, True)) + + +def chat_completion_prompt_cache_breakpoints_for_wire(message: _MarkedT, drop_params: bool) -> _MarkedT: + content: Final = message.get("content") + if not drop_params or not isinstance(content, list): + return message + parts: Final = [ + _part_with_chat_completion_prompt_cache_breakpoint(part) for part in _BREAKPOINT_BLOCKS.validate_python(content) + ] + return cast(_MarkedT, {**message, "content": parts}) # cast-ok: same message shape with its parts' markers trimmed + + LITELLM_RESPONSE_ONLY_MESSAGE_FIELDS: Final = frozenset( {"thinking_blocks", "reasoning_items", "provider_specific_fields"} ) diff --git a/litellm/llms/anthropic/messages/responses_adapters/handler.py b/litellm/llms/anthropic/messages/responses_adapters/handler.py index 0eef3fd87c0..08b00dd5fcf 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/handler.py +++ b/litellm/llms/anthropic/messages/responses_adapters/handler.py @@ -21,7 +21,7 @@ from litellm.litellm_core_utils.asyncify import ( from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( cache_control_as_prompt_cache_breakpoint, - prompt_cache_breakpoint_for_wire, + prompt_cache_breakpoint_with_readable_mode, with_prompt_cache_breakpoint, ) from litellm.responses.utils import auto_reasoning_summary @@ -207,7 +207,7 @@ def _without_prompt_cache_breakpoint(block: Mapping[str, object]) -> Mapping[str def _cache_control_as_wire_safe_prompt_cache_breakpoint(block: Mapping[str, object]) -> Mapping[str, object]: if "prompt_cache_breakpoint" not in block: return cache_control_as_prompt_cache_breakpoint(block) - wire_marker: Final = prompt_cache_breakpoint_for_wire(block["prompt_cache_breakpoint"], drop_params=True) + wire_marker: Final = prompt_cache_breakpoint_with_readable_mode(block["prompt_cache_breakpoint"], drop_params=True) return cache_control_as_prompt_cache_breakpoint( with_prompt_cache_breakpoint(_without_prompt_cache_breakpoint(block), wire_marker) ) diff --git a/litellm/llms/azure/chat/gpt_reasoning_transformation.py b/litellm/llms/azure/chat/gpt_reasoning_transformation.py index 72a0500e0af..29305866e23 100644 --- a/litellm/llms/azure/chat/gpt_reasoning_transformation.py +++ b/litellm/llms/azure/chat/gpt_reasoning_transformation.py @@ -1,6 +1,6 @@ """Support for Azure OpenAI gpt-5 model family.""" -from collections.abc import Collection +from collections.abc import Collection, Mapping from typing import ClassVar, Final import litellm @@ -105,7 +105,7 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): model: str, messages: list[AllMessageValues], optional_params: dict, - litellm_params: dict, + litellm_params: Mapping[str, object], headers: dict, ) -> dict: model = model.replace(self.GPT5_SERIES_ROUTE, "") diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 990c8869de3..fdb012756fc 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -5,12 +5,16 @@ from typing import TYPE_CHECKING, Any, Final from httpx._models import Headers, Response import litellm +from litellm.integrations.anthropic_cache_control_hook import deployment_prompt_cache_breakpoint_support +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( + chat_completion_prompt_cache_breakpoints_for_wire, drop_tool_reference_parts_from_tool_messages, flatten_combinators_and_drop_non_python_regex_patterns, hoist_images_from_tool_messages, system_messages_first, tool_with_sanitized_parameters, + without_prompt_cache_breakpoints, ) from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_azure_openai_messages, @@ -290,11 +294,16 @@ class AzureOpenAIConfig(BaseConfig): model: str, messages: list[AllMessageValues], optional_params: dict, - litellm_params: dict, + litellm_params: Mapping[str, object], headers: dict, ) -> dict: 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) + wire_messages: Final = ( + self._with_wire_safe_prompt_cache_breakpoints(model, ordered_messages, litellm_params) + if drops_unsupported_params(litellm_params.get("drop_params")) + else ordered_messages + ) + stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(wire_messages) azure_messages: Final = convert_to_azure_openai_messages(hoist_images_from_tool_messages(stripped_messages)) request_params: Final = MappingProxyType( { @@ -310,6 +319,14 @@ class AzureOpenAIConfig(BaseConfig): **sanitized_tools_update(optional_params), } + @staticmethod + def _with_wire_safe_prompt_cache_breakpoints( + model: str, messages: list[AllMessageValues], litellm_params: Mapping[str, object] + ) -> list[AllMessageValues]: + if deployment_prompt_cache_breakpoint_support(model, "azure", litellm_params) is False: + return [without_prompt_cache_breakpoints(message) for message in messages] + return [chat_completion_prompt_cache_breakpoints_for_wire(message, drop_params=True) for message in messages] + def transform_response( self, model: str, diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 2b757bb10d9..6c954d00e64 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -1,6 +1,7 @@ import re -from collections.abc import Mapping +from collections.abc import Callable, Mapping from copy import deepcopy +from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal from urllib.parse import urlsplit @@ -11,6 +12,7 @@ from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.constants import AZURE_DEFAULT_RESPONSES_API_VERSION +from litellm.litellm_core_utils.prompt_templates.common_utils import prompt_cache_breakpoint_within from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.azure.chat.gpt_reasoning_transformation import AzureOpenAIGPT5Config from litellm.llms.azure.common_utils import BaseAzureLLM @@ -29,6 +31,7 @@ else: LiteLLMLoggingObj = Any _DECODED_JSON: Final = TypeAdapter(object) +_AZURE_RESPONSES_PROMPT_CACHE_BREAKPOINT: Final = TypeAdapter(AzureResponsesPromptCacheBreakpoint) _DATED_GA_API_VERSION: Final = re.compile(r"\d{4}-\d{2}-\d{2}") @@ -81,6 +84,13 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): def _is_gpt_5_model(model: str) -> bool: return AzureOpenAIGPT5Config.is_model_gpt_5_model(model) + def _prompt_cache_breakpoint_to_wire( + self, litellm_params: GenericLiteLLMParams, reads_breakpoints: bool, known_support: bool | None + ) -> Callable[[object], object]: + if reads_breakpoints or known_support is True: + return partial(prompt_cache_breakpoint_within, _AZURE_RESPONSES_PROMPT_CACHE_BREAKPOINT) + return super()._prompt_cache_breakpoint_to_wire(litellm_params, reads_breakpoints, known_support) + @staticmethod def _supports_reasoning_effort_none(model: str) -> bool: return AzureOpenAIGPT5Config.supports_reasoning_effort_level(model, "none") diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 05855c307e9..1f7b1712d68 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -15,7 +15,8 @@ role / access key / profile / web identity), signed via the shared BaseAWSLLM._sign_request after the request body is finalized. """ -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence +from functools import partial from typing import Final, cast # noqa: TID251 # map_openai_params returns the filtered params as a bare dict import httpx @@ -23,6 +24,7 @@ import httpx import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import drops_unsupported_params +from litellm.litellm_core_utils.prompt_templates.common_utils import chat_completion_prompt_cache_breakpoint_for_wire from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.responses.codex_compat import drop_unsupported_tools, normalize_codex_input_items @@ -62,6 +64,10 @@ _BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"}) _BEDROCK_MANTLE_OPENAI_PATH_SUPPORTED_REASONING_SUMMARIES: Final = frozenset({"auto"}) +def _prompt_cache_breakpoint_as_sent_or_none(marker: object) -> object: + return marker if chat_completion_prompt_cache_breakpoint_for_wire(marker, drop_params=True) == marker else None + + class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPIConfig): def __init__( self, @@ -113,6 +119,13 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI headers["OpenAI-Project"] = litellm_params.aws_bedrock_project_id return headers + def _prompt_cache_breakpoint_to_wire( + self, litellm_params: GenericLiteLLMParams, reads_breakpoints: bool, known_support: bool | None + ) -> Callable[[object], object]: + if reads_breakpoints or known_support is True: + return partial(chat_completion_prompt_cache_breakpoint_for_wire, drop_params=True) + return _prompt_cache_breakpoint_as_sent_or_none + def supports_native_file_search(self) -> bool: return False diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 9d551f6f859..4bfc9f28939 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -21,6 +21,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo should_convert_tool_call_to_json_mode, ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( + chat_completion_prompt_cache_breakpoints_for_wire, drop_non_python_regex_patterns, drop_tool_reference_parts_from_tool_messages, flatten_combinators_and_drop_non_python_regex_patterns, @@ -556,9 +557,10 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): """ A deployment that reads OpenAI's ``prompt_cache_breakpoint`` (gpt-5.6+ on api.openai.com, or one opted in through ``prompt_cache_options`` or the cost - map) gets each Anthropic ``cache_control`` mark rewritten in place. Any other - endpoint keeps or strips the marks as before, and under drop_params also loses - every ``prompt_cache_breakpoint`` it cannot read. + map) gets each Anthropic ``cache_control`` mark rewritten in place, and under + drop_params each content part's ``prompt_cache_breakpoint`` cut down to what chat + completions accepts. Any other endpoint keeps or strips the marks as before, and + under drop_params also loses every ``prompt_cache_breakpoint`` it cannot read. """ from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook @@ -572,9 +574,15 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): _claude_messages_within_cache_control_cap(model, messages, tools, litellm_params, request_params), tools, ) - drops_breakpoints: Final = not reads_breakpoints and _drops_params(litellm_params) + drop_params: Final = _drops_params(litellm_params) + drops_breakpoints: Final = not reads_breakpoints and drop_params renamed: Final = ( - [AnthropicCacheControlHook.cache_control_as_prompt_cache_breakpoints(message) for message in messages] + [ + AnthropicCacheControlHook.cache_control_as_prompt_cache_breakpoints( + chat_completion_prompt_cache_breakpoints_for_wire(message, drop_params) + ) + for message in messages + ] if reads_breakpoints else [without_prompt_cache_breakpoints(message) for message in messages] if drops_breakpoints diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 7cca31c5d3f..ef3b468dc1a 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -1,5 +1,5 @@ from collections.abc import Callable, Collection, Iterable, Mapping, Sequence -from functools import lru_cache +from functools import lru_cache, partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_type_hints @@ -12,7 +12,7 @@ import litellm from litellm._logging import verbose_logger from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, - openai_prompt_cache_breakpoint_support, + deployment_prompt_cache_breakpoint_support, targets_openai_api, ) from litellm.litellm_core_utils.core_helpers import normalize_drop_params, process_response_headers @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( drop_non_python_regex_patterns, flatten_combinators_and_drop_non_python_regex_patterns, prompt_cache_breakpoint_for_wire, + whole_prompt_cache_breakpoint_or_none, with_prompt_cache_breakpoint, ) from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -103,13 +104,6 @@ _INPUT_ITEM: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, _CONTENT_PARTS: Final[TypeAdapter[Sequence[Mapping[str, object]]]] = TypeAdapter(Sequence[Mapping[str, object]]) -class _DeclaredModel(BaseModel): - model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True) - - base_model: object = None - model_info: object = None - - class _DeploymentHost(BaseModel): """`responses()` builds the params from raw kwargs, where callers may spell the host `base_url`.""" @@ -241,12 +235,12 @@ def _json_array_or_empty(value: object) -> Sequence[object]: return value if isinstance(value, (list, tuple)) else () -def _item_with_wire_safe_prompt_cache_breakpoints(item: object, keeps_markers: bool) -> object: +def _item_with_wire_safe_prompt_cache_breakpoints(item: object, to_wire: Callable[[object], object] | None) -> object: fields: Final = _json_object_or_none(item) if fields is None or not _carries_prompt_cache_breakpoint(fields): return item wire_safe_fields: Final = { - key: _parts_with_wire_safe_prompt_cache_breakpoints(key, value, keeps_markers) + key: _parts_with_wire_safe_prompt_cache_breakpoints(key, value, to_wire) for key, value in fields.items() if key != "prompt_cache_breakpoint" } @@ -270,21 +264,23 @@ def _parts_carry_prompt_cache_breakpoint(parts: object) -> bool: return any(isinstance(part, Mapping) and "prompt_cache_breakpoint" in part for part in _json_array_or_empty(parts)) -def _parts_with_wire_safe_prompt_cache_breakpoints(key: str, value: object, keeps_markers: bool) -> object: +def _parts_with_wire_safe_prompt_cache_breakpoints( + key: str, value: object, to_wire: Callable[[object], object] | None +) -> object: if key not in _INPUT_PART_LIST_KEYS or not _parts_carry_prompt_cache_breakpoint(value): return value parts: Final = _json_array_or_empty(value) - wire_safe_parts: Final = tuple(_part_with_wire_safe_prompt_cache_breakpoint(part, keeps_markers) for part in parts) + wire_safe_parts: Final = tuple(_part_with_wire_safe_prompt_cache_breakpoint(part, to_wire) for part in parts) return value if _same_objects(wire_safe_parts, parts) else list(wire_safe_parts) -def _part_with_wire_safe_prompt_cache_breakpoint(part: object, keeps_markers: bool) -> object: +def _part_with_wire_safe_prompt_cache_breakpoint(part: object, to_wire: Callable[[object], object] | None) -> object: fields: Final = _json_object_or_none(part) if fields is None or "prompt_cache_breakpoint" not in fields: return part marker: Final = fields["prompt_cache_breakpoint"] - wire_marker: Final = prompt_cache_breakpoint_for_wire(marker, drop_params=True) if keeps_markers else None - if keeps_markers and wire_marker == marker: + wire_marker: Final = None if to_wire is None else to_wire(marker) + if to_wire is not None and wire_marker == marker: return part unmarked: Final = {key: value for key, value in fields.items() if key != "prompt_cache_breakpoint"} return with_prompt_cache_breakpoint(unmarked, wire_marker) @@ -681,24 +677,31 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): ) -> str | ResponseInputParam: if not isinstance(input, list) or not drops_params(litellm_params): return input - keeps_markers: Final = reads_breakpoints or self._keeps_prompt_cache_breakpoints( - model=model, litellm_params=litellm_params - ) - wire_safe_items: Final = tuple( - _item_with_wire_safe_prompt_cache_breakpoints(item, keeps_markers) for item in input + known_support: Final = self._prompt_cache_breakpoint_support(model=model, litellm_params=litellm_params) + to_wire: Final = ( + self._prompt_cache_breakpoint_to_wire(litellm_params, reads_breakpoints, known_support) + if reads_breakpoints or known_support is not False + else None ) + wire_safe_items: Final = tuple(_item_with_wire_safe_prompt_cache_breakpoints(item, to_wire) for item in input) if _same_objects(wire_safe_items, input): return input wire_safe_input: Final = list(wire_safe_items) return cast("ResponseInputParam", wire_safe_input) # cast-ok: items keep their shape, minus a rejected marker - def _keeps_prompt_cache_breakpoints(self, model: str, litellm_params: GenericLiteLLMParams) -> bool: - provider: Final = _provider_name(self.custom_llm_provider) - base_model: Final = self._declared_base_model(model=model, litellm_params=litellm_params) - known_support: Final = ( - openai_prompt_cache_breakpoint_support(name, provider) for name in dict.fromkeys((base_model, model)) + def _prompt_cache_breakpoint_to_wire( + self, litellm_params: GenericLiteLLMParams, reads_breakpoints: bool, known_support: bool | None + ) -> Callable[[object], object]: + opted_in_openai_host: Final = reads_breakpoints and self.custom_llm_provider == LlmProviders.OPENAI + known_openai_reader: Final = known_support is True and self._targets_openai_api(litellm_params) + if opted_in_openai_host or known_openai_reader: + return partial(prompt_cache_breakpoint_for_wire, drop_params=True) + return whole_prompt_cache_breakpoint_or_none + + def _prompt_cache_breakpoint_support(self, model: str, litellm_params: GenericLiteLLMParams) -> bool | None: + return deployment_prompt_cache_breakpoint_support( + model, _provider_name(self.custom_llm_provider), litellm_params ) - return next((support for support in known_support if support is not None), None) is not False def _targets_openai_api(self, litellm_params: GenericLiteLLMParams) -> bool: host: Final = _DeploymentHost.model_validate(litellm_params) @@ -872,13 +875,6 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): base_model: Final[object] = model_info.get("base_model") if isinstance(model_info, dict) else None return base_model if isinstance(base_model, str) and base_model else model - @staticmethod - def _declared_base_model(model: str, litellm_params: GenericLiteLLMParams) -> str: - declared: Final = _DeclaredModel.model_validate(litellm_params) - model_info: Final = _json_object_or_none(declared.model_info) - candidates: Final = (declared.base_model, None if model_info is None else model_info.get("base_model")) - return next((name for name in candidates if isinstance(name, str) and name), model) - @staticmethod def _sanitized_tool_entry( entry: Mapping[str, object], diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 43d71a6c0b4..1d1628035ec 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -693,6 +693,15 @@ class PromptCacheBreakpoint(TypedDict): ttl: NotRequired[ReadOnly[Literal["30m", "12h"]]] +class ChatCompletionPromptCacheBreakpoint(TypedDict): + mode: ReadOnly[Literal["explicit"]] + + +class AzureResponsesPromptCacheBreakpoint(TypedDict): + mode: ReadOnly[Literal["explicit"]] + ttl: NotRequired[ReadOnly[Literal["30m"]]] + + class PromptCacheOptions(TypedDict, total=False): mode: ReadOnly[Literal["implicit", "explicit"]] ttl: ReadOnly[Literal["30m", "12h"]] diff --git a/tests/integration/providers/test_openai_chat_prompt_cache_breakpoint_wire.py b/tests/integration/providers/test_openai_chat_prompt_cache_breakpoint_wire.py index 807a7461675..858599829ad 100644 --- a/tests/integration/providers/test_openai_chat_prompt_cache_breakpoint_wire.py +++ b/tests/integration/providers/test_openai_chat_prompt_cache_breakpoint_wire.py @@ -41,6 +41,7 @@ class _Proxy: openai: Wire gpt_5_6_with_prompt_cache_options: str gpt_5_6_without_prompt_cache_options: str + gpt_5_6_dropping_params: str def openai_receives( self, model: str, messages: Sequence[JsonValue], *, stream: bool = False, tools: Sequence[JsonValue] = () @@ -79,6 +80,9 @@ def proxy() -> Iterator[_Proxy]: model="openai/gpt-5.6", api_base=api_base, prompt_cache_options={"mode": "explicit"} ), gpt_5_6_without_prompt_cache_options=scenario.model(model="openai/gpt-5.6", api_base=api_base), + gpt_5_6_dropping_params=scenario.model( + model="openai/gpt-5.6", api_base=api_base, prompt_cache_options={"mode": "explicit"}, drop_params=True + ), ) @@ -202,3 +206,56 @@ def test_custom_api_base_without_prompt_cache_options_forwards_cache_control_unc assert received["messages"] == _SYSTEM_AND_USER_WITH_CACHE_CONTROL assert "prompt_cache_options" not in received assert "prompt_cache_breakpoint" not in str(received) + + +_CLIENT_BREAKPOINTS_ON_PARTS: Final[list[JsonValue]] = [ + {"type": "text", "text": "Stable rules", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "1h"}}, + {"type": "text", "text": "More rules", "prompt_cache_breakpoint": {"mode": "implicit"}}, + {"type": "text", "text": "Describe this", "prompt_cache_breakpoint": {"mode": "explicit"}}, +] +_CLIENT_BREAKPOINTS_AS_CHAT_COMPLETIONS_TAKES_THEM: Final[list[JsonValue]] = [ + {"type": "text", "text": "Stable rules", "prompt_cache_breakpoint": {"mode": "explicit"}}, + {"type": "text", "text": "More rules"}, + {"type": "text", "text": "Describe this", "prompt_cache_breakpoint": {"mode": "explicit"}}, +] +_CHAT_COMPLETIONS_REJECTS_TTL_AND_IMPLICIT: Final = ( + "live gpt-6.1-sol Chat Completions answers 400 on a part breakpoint with any ttl or mode implicit (2026-10-28)" +) + + +@pytest.mark.parametrize("stream", [False, True], ids=["unary", "stream"]) +def test_drop_params_cuts_client_breakpoints_down_to_what_chat_completions_takes(proxy: _Proxy, stream: bool) -> None: + received: Final = proxy.openai_receives( + proxy.gpt_5_6_dropping_params, [{"role": "user", "content": _CLIENT_BREAKPOINTS_ON_PARTS}], stream=stream + ) + + assert received["messages"] == [ + {"role": "user", "content": _CLIENT_BREAKPOINTS_AS_CHAT_COMPLETIONS_TAKES_THEM} + ], _CHAT_COMPLETIONS_REJECTS_TTL_AND_IMPLICIT + + +def test_without_drop_params_client_breakpoints_reach_openai_unchanged(proxy: _Proxy) -> None: + received: Final = proxy.openai_receives( + proxy.gpt_5_6_with_prompt_cache_options, [{"role": "user", "content": _CLIENT_BREAKPOINTS_ON_PARTS}] + ) + + assert received["messages"] == [{"role": "user", "content": _CLIENT_BREAKPOINTS_ON_PARTS}] + + +def test_drop_params_cuts_client_breakpoints_on_v1_messages_routed_to_chat_completions(proxy: _Proxy) -> None: + _ = proxy.openai.drain() + response: Final = proxy.gateway.request( + "POST", + "/v1/messages", + { + "model": proxy.gpt_5_6_dropping_params, + "max_tokens": 16, + "stop_sequences": ["END"], + "messages": [{"role": "user", "content": _CLIENT_BREAKPOINTS_ON_PARTS}], + }, + ) + assert response.status_code == 200, response.text + + assert proxy.received_by_openai()["messages"] == [ + {"role": "user", "content": _CLIENT_BREAKPOINTS_AS_CHAT_COMPLETIONS_TAKES_THEM} + ], _CHAT_COMPLETIONS_REJECTS_TTL_AND_IMPLICIT diff --git a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire.py b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire.py index 33338cae95d..7ea02435369 100644 --- a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire.py +++ b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire.py @@ -27,8 +27,9 @@ _MALFORMED: Final[tuple[tuple[str, JsonValue], ...]] = ( ("5kb-string", "x" * 5000), ("empty-object", {}), ("bogus-mode", {"mode": "bogus"}), - ("bad-ttl", {"mode": "explicit", "ttl": "1h"}), + ("bogus-mode-with-ttl", {"mode": "bogus", "ttl": "30m"}), ) +_UNSUPPORTED_TTL: Final[Mapping[str, JsonValue]] = {"mode": "explicit", "ttl": "1h"} _MALFORMED_IDS: Final = tuple(name for name, _ in _MALFORMED) _MALFORMED_VALUES: Final = tuple(value for _, value in _MALFORMED) _CASES: Final[tuple[tuple[str, Mode, JsonValue, JsonValue], ...]] = ( @@ -192,6 +193,27 @@ def test_marker_with_an_unknown_key(bridge: _Bridge, mode: Mode, expected: JsonV bridge.spend.landed(bridge.model(mode), call_id, marker) +@pytest.mark.parametrize( + ("mode", "expected"), (("on", None), ("off", _UNSUPPORTED_TTL)), ids=("dropped-on", "verbatim-off") +) +def test_marker_with_an_unsupported_ttl(bridge: _Bridge, mode: Mode, expected: JsonValue) -> None: + marker: Final = uuid.uuid4().hex + call_id: Final = _completion(_chat(bridge, bridge.model(mode), [_user(marker, _UNSUPPORTED_TTL)]), marker) + pcb.assert_marker(_user_block_on_wire(bridge, marker), expected) + bridge.spend.landed(bridge.model(mode), call_id, marker) + + +def test_marker_with_an_unsupported_ttl_keeps_its_mode_on_a_host_opted_in_with_prompt_cache_options( + bridge: _Bridge, +) -> None: + marker: Final = uuid.uuid4().hex + with bridge.gateway.scenario() as scenario: + model: Final = scenario.model(model=pcb.MODEL, api_base=bridge.api_base, drop_params=True, **_OPTED_IN) + call_id: Final = _completion(_chat(bridge, model, [_user(marker, _UNSUPPORTED_TTL)]), marker) + pcb.assert_marker(_user_block_on_wire(bridge, marker), pcb.EXPLICIT) + bridge.spend.landed(model, call_id, marker) + + def _second_block_on_wire(bridge: _Bridge, marker: str) -> dict[str, JsonValue]: request: Final = pcb.posted(bridge.wire, marker) _wire_body(request) diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 4e6ed16e360..558350a790a 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -5884,9 +5884,9 @@ _HAND_WRITTEN_PROMPT_CACHE_BREAKPOINT_BLOCKS: Final = ( ) -def test_convert_chat_completion_messages_to_responses_api_drops_malformed_prompt_cache_breakpoint_under_drop_params(): - """OpenAI's Responses API answered "Supported values are: '30m'" for a 1h breakpoint ttl on 2026-10-07, - so an unsupported ttl drops the marker as a unit while an unknown key is dropped from a valid one.""" +def test_convert_chat_completion_messages_to_responses_api_drops_only_an_unreadable_prompt_cache_breakpoint_under_drop_params(): + """The target's Responses config decides which ttl and keys its API takes, so the bridge keeps every marker + whose mode some target reads.""" handler = LiteLLMResponsesTransformationHandler() messages = [{"role": "user", "content": list(_HAND_WRITTEN_PROMPT_CACHE_BREAKPOINT_BLOCKS)}] @@ -5899,12 +5899,12 @@ def test_convert_chat_completion_messages_to_responses_api_drops_malformed_promp None, None, None, - None, - {"mode": "explicit"}, + {"mode": "explicit", "ttl": "1h"}, + {"mode": "explicit", "scope": "all"}, {"mode": "explicit", "ttl": "30m"}, {"mode": "explicit"}, ] - assert all("prompt_cache_breakpoint" not in block for block in content[:4]) + assert all("prompt_cache_breakpoint" not in block for block in content[:3]) def test_convert_chat_completion_messages_to_responses_api_keeps_malformed_prompt_cache_breakpoint_by_default(): @@ -5937,6 +5937,41 @@ def test_transform_request_drop_params_in_litellm_params_gates_the_prompt_cache_ assert "prompt_cache_breakpoint" not in result["input"][0]["content"][0] +_TTL_MARKED_CONTENT: Final = ( + {"type": "text", "text": "unsupported ttl", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "1h"}}, + {"type": "text", "text": "well formed", "prompt_cache_breakpoint": {"mode": "explicit"}}, +) + + +@pytest.mark.parametrize( + "global_drop_params,deployment_drop_params,expected", + [ + ("true", None, [None, {"mode": "explicit"}]), + (False, True, [{"mode": "explicit", "ttl": "1h"}, {"mode": "explicit"}]), + (True, None, [{"mode": "explicit", "ttl": "1h"}, {"mode": "explicit"}]), + ], + ids=("string-global-target-keeps-raw-markers", "deployment-drop-params", "global-drop-params"), +) +def test_transform_request_hands_a_marker_to_the_target_whole_only_when_the_target_trims_it( + monkeypatch: pytest.MonkeyPatch, + global_drop_params: object, + deployment_drop_params: bool | None, + expected: list[object], +) -> None: + monkeypatch.setattr(litellm, "drop_params", global_drop_params) + request: Final = LiteLLMResponsesTransformationHandler().transform_request( + model="gpt-6.1-sol", + messages=[{"role": "user", "content": list(_TTL_MARKED_CONTENT)}], + optional_params={}, + litellm_params={"custom_llm_provider": "openai", "drop_params": deployment_drop_params}, + headers={}, + litellm_logging_obj=Mock(), + ) + (item,) = cast(list[dict[str, object]], request["input"]) # cast-ok: the bridge emits message item mappings + parts: Final = cast(list[dict[str, object]], item["content"]) # cast-ok: message items carry content parts + assert [part.get("prompt_cache_breakpoint") for part in parts] == expected + + def test_convert_response_output_folds_two_messages_and_function_call_into_one_choice() -> None: first: Final = ResponseOutputMessage.model_validate( { diff --git a/tests/unit/integrations/test_anthropic_cache_control_hook.py b/tests/unit/integrations/test_anthropic_cache_control_hook.py index 778c8dfe2f5..0a6f37e4c15 100644 --- a/tests/unit/integrations/test_anthropic_cache_control_hook.py +++ b/tests/unit/integrations/test_anthropic_cache_control_hook.py @@ -3485,12 +3485,23 @@ class TestOpenAIPromptCacheBreakpoint: ("gpt-5.6-sol", True), ("gpt-5.6-luna", True), ("gpt-5.7", True), + ("gpt-5.6o", True), + ("gpt-5.10", True), + ("gpt-5.06", True), ("gpt-6", True), ("GPT-5.6", True), ("gpt-5.5", False), ("gpt-5", False), ("gpt-5-chat-latest", False), ("gpt-4.1", False), + ("gpt-5-6-sol", True), + ("gpt-6-1", True), + ("gpt-5-4-mini", False), + ("gpt-5-32k", False), + ("gpt-5-6b", False), + ("gpt-5-06", False), + ("gpt-5-2026-03-01", False), + ("gpt-4-1106-preview", False), ("o3", False), ("claude-sonnet-4-5", False), ], diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 1ef97370b2b..146bb08a4fd 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -2343,7 +2343,8 @@ _MALFORMED_PROMPT_CACHE_BREAKPOINTS: Final = ( "", {}, {"mode": "bogus"}, - {"mode": "explicit", "ttl": "1h"}, + {"mode": "bogus", "ttl": "30m"}, + {1: "explicit"}, ) _MALFORMED_PROMPT_CACHE_BREAKPOINT_IDS: Final = ( "string", @@ -2352,7 +2353,8 @@ _MALFORMED_PROMPT_CACHE_BREAKPOINT_IDS: Final = ( "empty-string", "empty-object", "bogus-mode", - "bad-ttl", + "bogus-mode-with-ttl", + "non-string-key", ) _VALID_PROMPT_CACHE_BREAKPOINTS: Final = ({"mode": "explicit"}, {"mode": "explicit", "ttl": "30m"}) @@ -2382,6 +2384,18 @@ def test_prompt_cache_breakpoint_for_wire_strips_an_unknown_key_under_drop_param assert prompt_cache_breakpoint_for_wire({"mode": "explicit", "note": "kept"}, True) == {"mode": "explicit"} +@pytest.mark.parametrize( + "marker", + ({"mode": "explicit", "ttl": "1h"}, {"mode": "explicit", "ttl": "1h", "note": "kept"}), + ids=("bad-ttl", "bad-ttl-and-unknown-key"), +) +def test_prompt_cache_breakpoint_for_wire_drops_an_unsupported_ttl_and_keeps_the_mode_under_drop_params( + marker: dict[str, str], +) -> None: + assert prompt_cache_breakpoint_for_wire(marker, True) == {"mode": "explicit"} + assert prompt_cache_breakpoint_for_wire(marker, False) is marker + + def test_cache_control_is_replaced_by_an_explicit_prompt_cache_breakpoint_on_the_same_block() -> None: marked: Final = {"type": "text", "text": "rules", "cache_control": {"type": "ephemeral", "ttl": "1h"}} diff --git a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py index f6dae2c304f..1f1ff6a5748 100644 --- a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -806,3 +806,89 @@ def test_azure_keeps_a_later_developer_message_in_place_as_system(config: AzureO ) assert [message["role"] for message in messages] == ["system", "user", "assistant", "system", "user"] assert messages[3]["content"] == "Answer with exactly one word." + + +_MARKED_AZURE_PARTS: Final = ( + {"type": "text", "text": "Stable prefix", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"}}, + {"type": "text", "text": "Implicit", "prompt_cache_breakpoint": {"mode": "implicit"}}, +) +_TRIMMED_AZURE_PARTS: Final = [ + {"type": "text", "text": "Stable prefix", "prompt_cache_breakpoint": {"mode": "explicit"}}, + {"type": "text", "text": "Implicit"}, +] + + +@pytest.fixture +def registered_azure_breakpoint_reader(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setitem( + litellm.model_cost, + "azure/unit-breakpoint-reader", + {"litellm_provider": "azure", "mode": "chat", "supports_prompt_cache_breakpoint": True}, + ) + monkeypatch.setitem( + litellm.model_cost, + "azure/unit-breakpoint-ignorer", + {"litellm_provider": "azure", "mode": "chat", "supports_prompt_cache_breakpoint": False}, + ) + + +def _azure_request_parts( + config: AzureOpenAIConfig, model: str, litellm_params: dict[str, object] +) -> object: + request: Final = config.transform_request( + model=model, + messages=[{"role": "user", "content": list(_MARKED_AZURE_PARTS)}], + optional_params={}, + litellm_params=litellm_params, + headers={}, + ) + return request["messages"][0]["content"] + + +@pytest.mark.parametrize("config", [AzureOpenAIConfig(), AzureOpenAIGPT5Config()]) +@pytest.mark.parametrize( + "model, litellm_params", + [ + ("unit-breakpoint-reader", {"drop_params": True}), + ("my-azure-alias", {"drop_params": True, "base_model": "azure/unit-breakpoint-reader"}), + ("my-azure-alias", {"drop_params": True, "model_info": {"base_model": "azure/unit-breakpoint-reader"}}), + ("my-azure-alias", {"drop_params": True}), + ], + ids=("deployment-name", "litellm-params-base-model", "model-info-base-model", "unknown-deployment"), +) +def test_azure_drop_params_cuts_part_breakpoints_down_to_what_azure_chat_accepts( + registered_azure_breakpoint_reader: None, + config: AzureOpenAIConfig, + model: str, + litellm_params: dict[str, object], +) -> None: + """Azure OpenAI chat answered 400 for a content part's breakpoint ttl and for implicit mode on 2026-10-29, + the same as OpenAI chat completions.""" + assert _azure_request_parts(config, model, litellm_params) == _TRIMMED_AZURE_PARTS + + +@pytest.mark.parametrize("litellm_params", [{}, {"drop_params": False}], ids=("no-drop-params", "drop-params-false")) +def test_azure_sends_part_breakpoints_as_written_without_drop_params( + registered_azure_breakpoint_reader: None, litellm_params: dict[str, object] +) -> None: + assert _azure_request_parts(AzureOpenAIConfig(), "unit-breakpoint-reader", litellm_params) == list( + _MARKED_AZURE_PARTS + ) + + +@pytest.mark.parametrize( + "model, litellm_params", + [ + ("unit-breakpoint-ignorer", {"drop_params": True}), + ("unit-breakpoint-reader", {"drop_params": True, "base_model": "azure/unit-breakpoint-ignorer"}), + ("my-azure-alias", {"drop_params": True, "model_info": {"base_model": "azure/unit-breakpoint-ignorer"}}), + ], + ids=("known-non-reader", "base-model-known-non-reader", "model-info-base-model-known-non-reader"), +) +def test_azure_drop_params_strips_part_breakpoints_from_a_known_non_reader( + registered_azure_breakpoint_reader: None, model: str, litellm_params: dict[str, object] +) -> None: + assert _azure_request_parts(AzureOpenAIConfig(), model, litellm_params) == [ + {"type": "text", "text": "Stable prefix"}, + {"type": "text", "text": "Implicit"}, + ] diff --git a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py index 9f68a4fb73f..b1d52c9dd51 100644 --- a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -1271,6 +1271,12 @@ def test_azure_ai_falls_back_to_the_model_name_when_the_base_model_row_does_not_ {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}}, ), ({"type": "ephemeral"}, {"drop_params": True}, {"type": "text", "text": "stable"}), + ({"mode": "implicit"}, {"drop_params": True}, {"type": "text", "text": "stable"}), + ( + {"mode": "explicit", "ttl": "1h"}, + {"drop_params": True}, + {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}}, + ), ( {"mode": "explicit", "ttl": "30m"}, {}, diff --git a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index 0ae6a73968d..f8f18852543 100644 --- a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -1922,3 +1922,68 @@ class TestBedrockMantleResponsesPricing: def test_models_registered(self, local_cost_map): assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models + + +def _mantle_marked_part( + litellm_params: GenericLiteLLMParams, marker: object, model: str = "openai.gpt-6.1-sol" +) -> object: + body: Final = BedrockMantleResponsesAPIConfig().transform_responses_api_request( + model=model, + input=[ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Stable prefix", "prompt_cache_breakpoint": marker}], + } + ], + response_api_optional_request_params={}, + litellm_params=litellm_params, + headers={}, + ) + return body["input"][0]["content"][0] + + +@pytest.mark.parametrize( + "marker, wire_part", + [ + ( + {"mode": "explicit", "ttl": "30m"}, + {"type": "input_text", "text": "Stable prefix", "prompt_cache_breakpoint": {"mode": "explicit"}}, + ), + ({"mode": "implicit"}, {"type": "input_text", "text": "Stable prefix"}), + ], + ids=("ttl-dropped", "implicit-dropped"), +) +def test_mantle_drop_params_cuts_a_responses_breakpoint_down_to_its_mode( + monkeypatch: pytest.MonkeyPatch, marker: dict[str, str], wire_part: dict[str, object] +) -> None: + """Bedrock Mantle Responses answered 400 for a 30m breakpoint ttl and 200 for a bare explicit mode on 2026-10-29.""" + monkeypatch.setattr(litellm, "drop_params", False) + assert _mantle_marked_part(GenericLiteLLMParams(drop_params=True), marker) == wire_part + assert _mantle_marked_part(GenericLiteLLMParams(), marker) == { + "type": "input_text", + "text": "Stable prefix", + "prompt_cache_breakpoint": marker, + } + + +@pytest.mark.parametrize( + "marker, wire_part", + [ + ({"mode": "explicit", "ttl": "1h"}, {"type": "input_text", "text": "Stable prefix"}), + ({"mode": "explicit", "ttl": "30m"}, {"type": "input_text", "text": "Stable prefix"}), + ( + {"mode": "explicit"}, + {"type": "input_text", "text": "Stable prefix", "prompt_cache_breakpoint": {"mode": "explicit"}}, + ), + ], + ids=("1h-dropped", "30m-dropped", "explicit-kept"), +) +def test_mantle_drop_params_sends_an_unlisted_model_only_a_bare_explicit_breakpoint( + monkeypatch: pytest.MonkeyPatch, marker: dict[str, str], wire_part: dict[str, object] +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + assert ( + _mantle_marked_part(GenericLiteLLMParams(drop_params=True), marker, model="openai.not-in-the-cost-map") + == wire_part + ) diff --git a/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py b/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py index 79fe7061297..4aa93848c79 100644 --- a/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py @@ -1659,6 +1659,125 @@ class TestOpenAIPromptCacheBreakpointChatPath: ] +_CLIENT_BREAKPOINTS_AND_THEIR_CHAT_WIRE_FORM: Final = ( + ({"mode": "explicit", "ttl": "1h"}, {"mode": "explicit"}), + ({"mode": "explicit", "ttl": "30m"}, {"mode": "explicit"}), + ({"mode": "explicit"}, {"mode": "explicit"}), + ({"mode": "implicit"}, None), + ({}, None), + ("explicit", None), +) +_CLIENT_BREAKPOINT_IDS: Final = ("ttl-1h", "ttl-30m", "explicit", "implicit", "empty", "string") + + +class TestChatCompletionsClientBreakpointUnderDropParams: + """ + Live OpenAI Chat Completions (gpt-6.1-sol, gpt-6.1-terra, gpt-6-sol, gpt-5.6, 2026-10-28) answers 400 on a + content part's prompt_cache_breakpoint carrying any ttl or a mode other than "explicit". + """ + + @staticmethod + def _request(marker: object, drop_params: bool) -> dict[str, object]: + return OpenAIGPTConfig().transform_request( + model="gpt-5.6", + messages=[ + {"role": "system", "content": "Answer briefly"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this", "prompt_cache_breakpoint": marker}, + {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}, + {"type": "text", "text": "and nothing else"}, + ], + }, + ], + optional_params={}, + litellm_params={"custom_llm_provider": "openai", "drop_params": drop_params}, + headers={}, + ) + + @pytest.mark.parametrize( + ("marker", "wire_marker"), _CLIENT_BREAKPOINTS_AND_THEIR_CHAT_WIRE_FORM, ids=_CLIENT_BREAKPOINT_IDS + ) + def test_drop_params_cuts_the_breakpoint_down_to_what_chat_completions_accepts( + self, marker: object, wire_marker: dict[str, str] | None + ): + text_part: Final = {"type": "text", "text": "Describe this"} + + assert self._request(marker, drop_params=True)["messages"] == [ + {"role": "system", "content": "Answer briefly"}, + { + "role": "user", + "content": [ + text_part if wire_marker is None else {**text_part, "prompt_cache_breakpoint": wire_marker}, + {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}, + {"type": "text", "text": "and nothing else"}, + ], + }, + ] + + @pytest.mark.parametrize( + "marker", [marker for marker, _ in _CLIENT_BREAKPOINTS_AND_THEIR_CHAT_WIRE_FORM], ids=_CLIENT_BREAKPOINT_IDS + ) + def test_without_drop_params_the_client_breakpoint_reaches_openai_unchanged(self, marker: object): + assert self._request(marker, drop_params=False)["messages"][1]["content"][0] == { + "type": "text", + "text": "Describe this", + "prompt_cache_breakpoint": marker, + } + + def test_cache_control_still_marks_a_part_whose_client_breakpoint_was_dropped(self): + request: Final = OpenAIGPTConfig().transform_request( + model="gpt-5.6", + messages=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Describe this", + "cache_control": {"type": "ephemeral"}, + "prompt_cache_breakpoint": {"mode": "implicit"}, + } + ], + } + ], + optional_params={}, + litellm_params={"custom_llm_provider": "openai", "drop_params": True}, + headers={}, + ) + + assert request["messages"] == [ + { + "role": "user", + "content": [{"type": "text", "text": "Describe this", "prompt_cache_breakpoint": {"mode": "explicit"}}], + } + ] + + + @pytest.mark.parametrize("stray_part", ["Answer briefly", None], ids=["string", "null"]) + def test_drop_params_leaves_a_part_that_is_not_an_object_for_openai_to_reject(self, stray_part: object): + assert OpenAIGPTConfig().transform_request( + model="gpt-5.6", + messages=[ + { + "role": "system", + "content": [ + stray_part, + {"type": "text", "text": "and nothing else", "prompt_cache_breakpoint": {"mode": "implicit"}}, + ], + }, + {"role": "user", "content": "Describe this"}, + ], + optional_params={}, + litellm_params={"custom_llm_provider": "openai", "drop_params": True}, + headers={}, + )["messages"] == [ + {"role": "system", "content": [stray_part, {"type": "text", "text": "and nothing else"}]}, + {"role": "user", "content": "Describe this"}, + ] + + _ARTIFACT_FIELD_PATTERN: Final = r'^(?!__.*__$)[^\p{Cc}\p{Cf}\p{Zl}\p{Zp}"\\./[\]]{1,200}$' diff --git a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py index 41c54bf6cc9..6b36a683df7 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py @@ -4393,8 +4393,9 @@ async def test_extra_body_merges_with_request_data(extra_body_mock_response_data assert request_body["custom_field"] == "custom_value" -_MALFORMED_MARKERS: Final = ("yes", 1, {"mode": "bogus"}, {"mode": "explicit", "ttl": "1h"}) -_MALFORMED_MARKER_IDS: Final = ("string", "int", "bogus-mode", "bad-ttl") +_MALFORMED_MARKERS: Final = ("yes", 1, {"mode": "bogus"}, {"mode": "bogus", "ttl": "30m"}) +_MALFORMED_MARKER_IDS: Final = ("string", "int", "bogus-mode", "bogus-mode-with-ttl") +_UNSUPPORTED_TTL_MARKER: Final = {"mode": "explicit", "ttl": "1h"} _VALID_MARKERS: Final = ({"mode": "explicit"}, {"mode": "explicit", "ttl": "30m"}) _VALID_MARKER_IDS: Final = ("explicit", "explicit-30m") _PARTS: Final[Mapping[str, Mapping[str, object]]] = { @@ -4480,6 +4481,76 @@ def test_valid_prompt_cache_breakpoint_reaches_the_wire_under_both_settings( assert _transformed_input(_CONFIGS[0], input, GenericLiteLLMParams(drop_params=drop_params)) == input +@pytest.mark.parametrize("config", _CONFIGS, ids=_CONFIG_IDS) +@pytest.mark.parametrize("part_type", tuple(_PARTS)) +def test_an_unsupported_ttl_is_dropped_and_the_explicit_mode_kept_under_drop_params( + monkeypatch: pytest.MonkeyPatch, config: BaseResponsesAPIConfig, part_type: str +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + input: Final = _marked_message(part_type, _UNSUPPORTED_TTL_MARKER) + trimmed: Final = _transformed_input(config, input, GenericLiteLLMParams(drop_params=True)) + assert _single_part(trimmed) == {**_PARTS[part_type], "prompt_cache_breakpoint": {"mode": "explicit"}} + assert _transformed_input(config, input, GenericLiteLLMParams(drop_params=False)) == input + + +@pytest.mark.parametrize( + "config, marker, wire_marker", + [ + (OpenAIResponsesAPIConfig(), {"mode": "explicit", "ttl": "12h"}, {"mode": "explicit", "ttl": "12h"}), + (AzureOpenAIResponsesAPIConfig(), {"mode": "explicit", "ttl": "30m"}, {"mode": "explicit", "ttl": "30m"}), + (AzureOpenAIResponsesAPIConfig(), {"mode": "explicit", "ttl": "12h"}, {"mode": "explicit"}), + (AzureAIResponsesAPIConfig(), {"mode": "explicit", "ttl": "12h"}, {"mode": "explicit"}), + ], + ids=("openai-keeps-12h", "azure-keeps-30m", "azure-drops-12h", "azure-ai-drops-12h"), +) +def test_drop_params_keeps_only_a_ttl_the_target_responses_api_accepts( + monkeypatch: pytest.MonkeyPatch, + config: BaseResponsesAPIConfig, + marker: Mapping[str, str], + wire_marker: Mapping[str, str], +) -> None: + """Azure OpenAI Responses answered 400 for a 12h breakpoint ttl and 200 for 30m on 2026-10-29, and Azure AI + Foundry Responses did the same on 2026-10-30.""" + monkeypatch.setattr(litellm, "drop_params", False) + input: Final = _marked_message("input_text", marker) + trimmed: Final = _transformed_input(config, input, GenericLiteLLMParams(drop_params=True)) + assert _single_part(trimmed) == {**_PARTS["input_text"], "prompt_cache_breakpoint": wire_marker} + assert _transformed_input(config, input, GenericLiteLLMParams(drop_params=False)) == input + + +@pytest.mark.parametrize( + "config, litellm_params", + [ + (XAIResponsesAPIConfig(), GenericLiteLLMParams(drop_params=True)), + (HostedVLLMResponsesAPIConfig(), GenericLiteLLMParams(drop_params=True, api_base="http://localhost:8000/v1")), + (OpenAIResponsesAPIConfig(), GenericLiteLLMParams(drop_params=True, api_base="http://localhost:8000/v1")), + ], + ids=("xai", "hosted-vllm", "openai-compatible-host"), +) +def test_a_responses_api_not_known_to_take_breakpoints_gets_a_marker_whole_or_not_at_all( + monkeypatch: pytest.MonkeyPatch, config: BaseResponsesAPIConfig, litellm_params: GenericLiteLLMParams +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + unsupported_ttl: Final = _marked_message("input_text", _UNSUPPORTED_TTL_MARKER) + assert _single_part(_transformed_input(config, unsupported_ttl, litellm_params)) == _PARTS["input_text"] + openai_ttl: Final = _marked_message("input_text", {"mode": "explicit", "ttl": "12h"}) + assert _transformed_input(config, openai_ttl, litellm_params) == openai_ttl + + +def test_a_custom_host_opted_in_with_prompt_cache_options_keeps_the_mode_of_a_marker_with_an_unsupported_ttl( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + request: Final = OpenAIResponsesAPIConfig().transform_responses_api_request( + model="gpt-6.1-sol", + input=cast("ResponseInputParam", _marked_message("input_text", _UNSUPPORTED_TTL_MARKER)), + response_api_optional_request_params={"prompt_cache_options": {"mode": "explicit"}}, + litellm_params=GenericLiteLLMParams(drop_params=True, api_base="https://gateway.example/v1"), + headers={}, + ) + assert _single_part(request["input"]) == {**_PARTS["input_text"], "prompt_cache_breakpoint": {"mode": "explicit"}} + + def test_prompt_cache_breakpoint_unknown_key_is_trimmed_under_drop_params(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "drop_params", False) input: Final = _marked_message("input_text", {"mode": "explicit", "note": "kept"}) @@ -4569,12 +4640,43 @@ def test_a_valid_marker_survives_drop_params_on_a_deployment_name_with_unknown_b monkeypatch.setattr(litellm, "drop_params", False) litellm_params: Final = GenericLiteLLMParams(drop_params=True) valid: Final = _marked_message("input_text", {"mode": "explicit"}) - malformed: Final = _marked_message("input_text", {"mode": "explicit", "ttl": "1h"}) + malformed: Final = _marked_message("input_text", {"mode": "bogus"}) assert _transformed_input_for("litellm-qa-sol-alias", valid, litellm_params, config) == valid dropped: Final = _transformed_input_for("litellm-qa-sol-alias", malformed, litellm_params, config) assert _single_part(dropped) == _PARTS["input_text"] +@pytest.mark.parametrize("ttl", ("30m", "12h")) +@pytest.mark.parametrize( + "config", + (OpenAIResponsesAPIConfig(), AzureOpenAIResponsesAPIConfig(), AzureAIResponsesAPIConfig()), + ids=("openai", "azure", "azure_ai"), +) +def test_a_deployment_with_unknown_breakpoint_support_gets_a_marker_whole_or_not_at_all( + monkeypatch: pytest.MonkeyPatch, config: BaseResponsesAPIConfig, ttl: str +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + litellm_params: Final = GenericLiteLLMParams(drop_params=True) + supported_ttl: Final = _marked_message("input_text", {"mode": "explicit", "ttl": ttl}) + unsupported_ttl: Final = _marked_message("input_text", _UNSUPPORTED_TTL_MARKER) + assert _transformed_input_for("litellm-qa-sol-alias", supported_ttl, litellm_params, config) == supported_ttl + dropped: Final = _transformed_input_for("litellm-qa-sol-alias", unsupported_ttl, litellm_params, config) + assert _single_part(dropped) == _PARTS["input_text"] + + +@pytest.mark.parametrize( + "config", (AzureOpenAIResponsesAPIConfig(), AzureAIResponsesAPIConfig()), ids=("azure", "azure_ai") +) +def test_an_azure_deployment_declaring_a_breakpoint_reader_keeps_the_mode_of_a_marker_with_an_unsupported_ttl( + monkeypatch: pytest.MonkeyPatch, config: BaseResponsesAPIConfig +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + litellm_params: Final = GenericLiteLLMParams(drop_params=True, base_model="gpt-6.1-sol") + unsupported_ttl: Final = _marked_message("input_text", _UNSUPPORTED_TTL_MARKER) + trimmed: Final = _transformed_input_for("litellm-qa-sol-alias", unsupported_ttl, litellm_params, config) + assert _single_part(trimmed) == {**_PARTS["input_text"], "prompt_cache_breakpoint": {"mode": "explicit"}} + + def _declaring_base_model(where: str, base_model: str) -> GenericLiteLLMParams: if where == "litellm_params": return GenericLiteLLMParams(drop_params=True, base_model=base_model) @@ -4615,10 +4717,10 @@ def test_a_model_that_takes_markers_sends_kept_parts_byte_for_byte_and_drops_onl {"prompt_cache_breakpoint": {"mode": "explicit"}, "type": "input_text", "text": "b"}, ] rejected: Final = {"type": "input_text", "text": "c"} - bad_ttl: Final = {"prompt_cache_breakpoint": {"mode": "explicit", "ttl": "1h"}} + bogus_mode: Final = {"prompt_cache_breakpoint": {"mode": "bogus"}} input: Final = [ {"type": "message", "role": "user", "content": kept}, - {"type": "message", "role": "user", "content": [{**bad_ttl, **rejected}]}, + {"type": "message", "role": "user", "content": [{**bogus_mode, **rejected}]}, ] transformed: Final = _transformed_input_for("gpt-6.1-sol", input, GenericLiteLLMParams(drop_params=True)) assert json.dumps(transformed) == json.dumps( From aa95ac4cee62c14b798a7d43022599d61e467530 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 2 Nov 2026 18:17:20 +0000 Subject: [PATCH 10/13] fix(bedrock): price Web Search on the GPT-5.4 and later Mantle rows and the GovCloud West Mantle rows (#49687) * fix(bedrock): price Web Search on the GPT-6 Mantle rows and the GovCloud West Mantle rows * fix(bedrock): price Web Search on the GPT-5.6 cyber and daybreak-blue Mantle rows --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 36 ++++++++++++++----- model_prices_and_context_window.json | 36 ++++++++++++++----- 2 files changed, 56 insertions(+), 16 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d68eb3ce775..f2a3bb1a47e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -59875,6 +59875,7 @@ "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-cyber.html", "use_openai_responses_path": true, "supported_endpoints": [ @@ -59896,7 +59897,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-daybreak-blue-5.6-sol": { "input_cost_per_token": 4.4e-06, @@ -59911,6 +59913,7 @@ "max_input_tokens": 1050000, "max_output_tokens": 128000, "mode": "responses", + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "use_openai_responses_path": true, "supported_endpoints": [ "/v1/responses" @@ -59930,7 +59933,8 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html", + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.6-luna": { "input_cost_per_token": 2.2e-07, @@ -60336,6 +60340,7 @@ "mode": "responses", "output_cost_per_token_above_272k_tokens_ultrafast": 0.000495, "output_cost_per_token_ultrafast": 0.00033, + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60359,7 +60364,8 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html", + "supports_web_search": true }, "bedrock_mantle/openai.gpt-6-sol": { "input_cost_per_token": 2.2e-06, @@ -60374,6 +60380,7 @@ "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60397,7 +60404,8 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html", + "supports_web_search": true }, "bedrock_mantle/openai.gpt-6-luna": { "input_cost_per_token": 1.1e-07, @@ -60412,6 +60420,7 @@ "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60435,7 +60444,8 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html", + "supports_web_search": true }, "us.openai.gpt-6-astra": { "supports_bedrock_runtime_chat_completions_response_format": true, @@ -68633,7 +68643,9 @@ "cache_read_input_token_cost_above_272k_tokens": 5.4e-07, "output_cost_per_token": 1.62e-05, "output_cost_per_token_above_272k_tokens": 2.43e-05, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html" + "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html", + "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-luna": { "litellm_provider": "bedrock_mantle", @@ -68668,7 +68680,9 @@ "cache_read_input_token_cost_above_272k_tokens": 5.4e-08, "output_cost_per_token": 1.62e-06, "output_cost_per_token_above_272k_tokens": 2.43e-06, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html" + "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html", + "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.4": { "litellm_provider": "bedrock_mantle", @@ -68697,7 +68711,9 @@ "input_cost_per_token": 3.375e-06, "cache_read_input_token_cost": 3.375e-07, "output_cost_per_token": 2.025e-05, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html" + "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", + "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/xai.grok-4.3": { "use_openai_responses_path": true, @@ -75312,6 +75328,7 @@ "output_cost_per_token_above_272k_tokens": 1.65e-05, "output_cost_per_token_above_272k_tokens_ultrafast": 9.9e-05, "output_cost_per_token_ultrafast": 6.6e-05, + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-sol.html", "supported_endpoints": [ "/v1/chat/completions", @@ -75334,6 +75351,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true }, @@ -79022,6 +79040,7 @@ "mode": "responses", "output_cost_per_token": 5.5e-06, "output_cost_per_token_above_272k_tokens": 8.25e-06, + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-terra.html", "supported_endpoints": [ "/v1/chat/completions", @@ -79044,6 +79063,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true }, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d68eb3ce775..f2a3bb1a47e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -59875,6 +59875,7 @@ "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-cyber.html", "use_openai_responses_path": true, "supported_endpoints": [ @@ -59896,7 +59897,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-daybreak-blue-5.6-sol": { "input_cost_per_token": 4.4e-06, @@ -59911,6 +59913,7 @@ "max_input_tokens": 1050000, "max_output_tokens": 128000, "mode": "responses", + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "use_openai_responses_path": true, "supported_endpoints": [ "/v1/responses" @@ -59930,7 +59933,8 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-daybreak-blue-56-sol.html", + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.6-luna": { "input_cost_per_token": 2.2e-07, @@ -60336,6 +60340,7 @@ "mode": "responses", "output_cost_per_token_above_272k_tokens_ultrafast": 0.000495, "output_cost_per_token_ultrafast": 0.00033, + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60359,7 +60364,8 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html", + "supports_web_search": true }, "bedrock_mantle/openai.gpt-6-sol": { "input_cost_per_token": 2.2e-06, @@ -60374,6 +60380,7 @@ "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60397,7 +60404,8 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html", + "supports_web_search": true }, "bedrock_mantle/openai.gpt-6-luna": { "input_cost_per_token": 1.1e-07, @@ -60412,6 +60420,7 @@ "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "responses", + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "use_openai_responses_path": true, "supported_endpoints": [ "/v1/chat/completions", @@ -60435,7 +60444,8 @@ "supports_tool_choice": true, "supports_xhigh_reasoning_effort": true, "supports_vision": true, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html", + "supports_web_search": true }, "us.openai.gpt-6-astra": { "supports_bedrock_runtime_chat_completions_response_format": true, @@ -68633,7 +68643,9 @@ "cache_read_input_token_cost_above_272k_tokens": 5.4e-07, "output_cost_per_token": 1.62e-05, "output_cost_per_token_above_272k_tokens": 2.43e-05, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html" + "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html", + "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-luna": { "litellm_provider": "bedrock_mantle", @@ -68668,7 +68680,9 @@ "cache_read_input_token_cost_above_272k_tokens": 5.4e-08, "output_cost_per_token": 1.62e-06, "output_cost_per_token_above_272k_tokens": 2.43e-06, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html" + "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html", + "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.4": { "litellm_provider": "bedrock_mantle", @@ -68697,7 +68711,9 @@ "input_cost_per_token": 3.375e-06, "cache_read_input_token_cost": 3.375e-07, "output_cost_per_token": 2.025e-05, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html" + "search_context_cost_per_query": {"search_context_size_high":0.0144,"search_context_size_low":0.0144,"search_context_size_medium":0.0144}, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", + "supports_web_search": true }, "bedrock_mantle/us-gov-west-1/xai.grok-4.3": { "use_openai_responses_path": true, @@ -75312,6 +75328,7 @@ "output_cost_per_token_above_272k_tokens": 1.65e-05, "output_cost_per_token_above_272k_tokens_ultrafast": 9.9e-05, "output_cost_per_token_ultrafast": 6.6e-05, + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-sol.html", "supported_endpoints": [ "/v1/chat/completions", @@ -75334,6 +75351,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true }, @@ -79022,6 +79040,7 @@ "mode": "responses", "output_cost_per_token": 5.5e-06, "output_cost_per_token_above_272k_tokens": 8.25e-06, + "search_context_cost_per_query": {"search_context_size_high":0.012,"search_context_size_low":0.012,"search_context_size_medium":0.012}, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-1-terra.html", "supported_endpoints": [ "/v1/chat/completions", @@ -79044,6 +79063,7 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, + "supports_web_search": true, "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true }, From 97cd5d0c6a870ff9722cd50222423161b042d341 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Mon, 2 Nov 2026 10:20:23 -0800 Subject: [PATCH 11/13] fix(model-catalog): declare the tool and Responses token count framing keys (#49683) * fix(model-catalog): declare the tool and Responses token count framing keys * fix(model-catalog): read token count framings as strict integers The cost map guard read the framing keys as lax counts, so a text "214" or a true passed while litellm's reader drops both and falls back to the plain framing. Read all four framing keys as plain u64, matching the reader * fix(model-catalog): accept whole-float token count framings the way the Python reader does --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../model-catalog/src/wire/attributes.rs | 19 ++++++- .../crates/model-catalog/src/wire/lax.rs | 55 +++++++++++++++++++ 2 files changed, 72 insertions(+), 2 deletions(-) diff --git a/litellm-rust/crates/model-catalog/src/wire/attributes.rs b/litellm-rust/crates/model-catalog/src/wire/attributes.rs index 93c10a81524..6ac7b4bb304 100644 --- a/litellm-rust/crates/model-catalog/src/wire/attributes.rs +++ b/litellm-rust/crates/model-catalog/src/wire/attributes.rs @@ -11,7 +11,7 @@ use crate::error::{DateCause, Error, Repr}; use crate::vocabulary::{ AudioFormat, EffortLevel, Endpoint, InputModality, Mode, OutputModality, VertexAiAudioApi, }; -use crate::wire::lax::Count; +use crate::wire::lax::{Count, NumericCount}; macro_rules! attributes { ($($(#[$doc:meta])* $name:ident: $ty:ty $(as $via:ty)?),* $(,)?) => { @@ -104,9 +104,12 @@ attributes! { /// Cloud regions the model is available in ('global' or region ids). supported_regions: Vec, thinking_always_on: bool, - token_count_framing: u64 as Count, + token_count_framing: u64 as NumericCount, token_count_model: TokenCountModel, + token_count_responses_framing: u64 as NumericCount, + token_count_responses_tool_framing: u64 as NumericCount, token_count_tokenizer: TokenCountTokenizer, + token_count_tool_framing: u64 as NumericCount, /// Provider default tokens-per-minute limit. tpm: u64 as Count, use_openai_responses_path: bool, @@ -302,6 +305,18 @@ mod tests { #[case::reasoning_modes("reasoning_modes", json!(["none", "adaptive"]), true)] #[case::reasoning_context_not_a_list("reasoning_contexts", json!("reasoning"), false)] #[case::reasoning_mode_not_a_string("reasoning_modes", json!([1]), false)] + #[case::framing("token_count_framing", json!(12), true)] + #[case::framing_written_as_text("token_count_framing", json!("12"), false)] + #[case::framing_written_as_whole_float("token_count_framing", json!(12.0), true)] + #[case::fractional_tool_framing("token_count_tool_framing", json!(1.5), false)] + #[case::tool_framing("token_count_tool_framing", json!(214), true)] + #[case::tool_framing_written_as_text("token_count_tool_framing", json!("214"), false)] + #[case::tool_framing_as_a_flag("token_count_tool_framing", json!(true), false)] + #[case::negative_tool_framing("token_count_tool_framing", json!(-1), false)] + #[case::responses_framing("token_count_responses_framing", json!(38), true)] + #[case::responses_framing_as_a_flag("token_count_responses_framing", json!(false), false)] + #[case::responses_tool_framing("token_count_responses_tool_framing", json!(267), true)] + #[case::responses_tool_framing_written_as_text("token_count_responses_tool_framing", json!("267"), false)] fn each_attribute_reads_its_own_type( #[case] key: &str, #[case] value: Value, diff --git a/litellm-rust/crates/model-catalog/src/wire/lax.rs b/litellm-rust/crates/model-catalog/src/wire/lax.rs index d514a6a82c1..d98a601e8c1 100644 --- a/litellm-rust/crates/model-catalog/src/wire/lax.rs +++ b/litellm-rust/crates/model-catalog/src/wire/lax.rs @@ -8,6 +8,8 @@ pub(crate) struct LaxI64; pub(crate) struct Count; +pub(crate) struct NumericCount; + impl<'de> DeserializeAs<'de, i64> for LaxI64 { fn deserialize_as>(deserializer: D) -> Result { deserializer.deserialize_any(Self) @@ -52,6 +54,34 @@ impl<'de> DeserializeAs<'de, u64> for Count { } } +impl<'de> DeserializeAs<'de, u64> for NumericCount { + fn deserialize_as>(deserializer: D) -> Result { + deserializer.deserialize_any(Self) + } +} + +impl Visitor<'_> for NumericCount { + type Value = u64; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("a whole number of zero or more") + } + + fn visit_u64(self, value: u64) -> Result { + Ok(value) + } + + fn visit_i64(self, value: i64) -> Result { + u64::try_from(value).map_err(|_| E::invalid_value(de::Unexpected::Signed(value), &self)) + } + + fn visit_f64(self, value: f64) -> Result { + integral_float(value) + .and_then(|whole| u64::try_from(whole).ok()) + .ok_or_else(|| E::invalid_value(de::Unexpected::Float(value), &self)) + } +} + fn integer_string(value: &str) -> Option { let integer = match value.split_once('.') { Some((integer, fraction)) => { @@ -99,6 +129,12 @@ mod tests { .map(DeserializeAsWrap::into_inner) } + fn numeric_count(value: Value) -> Option { + DeserializeAsWrap::::deserialize(value) + .ok() + .map(DeserializeAsWrap::into_inner) + } + use serde::Deserialize; #[rstest] @@ -125,4 +161,23 @@ mod tests { ) { assert_eq!(count(written), expected); } + + #[rstest] + #[case::integer(json!(214), Some(214))] + #[case::zero(json!(0), Some(0))] + #[case::whole_float(json!(214.0), Some(214))] + #[case::negative_zero(json!(-0.0), Some(0))] + #[case::above_i64(json!(u64::MAX), Some(u64::MAX))] + #[case::negative(json!(-1), None)] + #[case::negative_whole_float(json!(-1.0), None)] + #[case::fraction(json!(1.5), None)] + #[case::string(json!("214"), None)] + #[case::boolean(json!(true), None)] + #[case::list(json!([214]), None)] + fn numeric_counts_accept_only_whole_json_numbers( + #[case] written: Value, + #[case] expected: Option, + ) { + assert_eq!(numeric_count(written), expected); + } } From a2ab27373f7edec256ae21a3c12f386d2d5da681 Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Mon, 2 Nov 2026 10:38:35 -0800 Subject: [PATCH 12/13] fix(proxy): count the context_management edits /v1/messages applies on every bridged provider (#49234) * fix(proxy): count the context_management edits /v1/messages applies on every bridged provider * fix: apply context_management edits in count_tokens on Responses bridges, compaction blocks and Bedrock Claude The Responses bridge now applies the same edits inference applies for every provider, the chat bridge cuts the history at a client compaction block, and Bedrock Claude forwards the edits AWS accepts with their beta to the mantle and runtime count APIs instead of dropping them. * refactor(count_tokens): return count-view failures as values and raise them once in the caller * fix(bedrock): keep the runtime's 400 when mantle does not serve the model it retries * fix(count_tokens): forward context edits on Vertex and Azure Claude counts and keep Bedrock's mantle answer without edits * fix(proxy): count edited bodies on native routes without a count API, keep counts on editor failures * fix(proxy): keep the Bedrock pass-through count shape and provider 400s on empty bridged bodies The pass-through count picks Converse or InvokeModel from the content again, and the edited InvokeModel count wraps string turns in text blocks. A bridged view with no messages sends the body as before. Azure and Vertex load the beta config off the event loop before a count with edits, and a model with no deployment reads drop_params from the request and router defaults. * fix(bedrock): send an empty-string turn as is when a count forwards context edits Wrapping an empty string into a text block made AWS reject the count with 'text content blocks must be non-empty', while base and /v1/messages accept the empty string. Only non-empty string turns are wrapped now. * test(azure_ai): compare the count's betas with the betas Azure inference sends * test(azure_ai): record the Azure inference request with respx instead of appending to a list --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../context_management/inference_view.py | 138 ++++ .../anthropic/count_tokens/handler.py | 6 + .../anthropic/count_tokens/token_counter.py | 1 + .../anthropic/count_tokens/transformation.py | 17 +- litellm/llms/bedrock/common_utils.py | 7 + .../count_tokens/bedrock_token_counter.py | 68 +- .../bedrock/count_tokens/mantle_handler.py | 11 +- .../anthropic_claude_transformation.py | 8 +- litellm/llms/gemini/common_utils.py | 9 +- litellm/llms/vertex_ai/common_utils.py | 13 +- .../llms/vertex_ai/count_tokens/context.py | 129 +--- .../count_tokens/handler.py | 17 +- .../vertex_ai_partner_models/main.py | 11 +- litellm/messages/dispatch.py | 42 +- litellm/proxy/proxy_server.py | 190 +++++- .../context_management/test_inference_view.py | 239 +++++++ .../count_tokens/test_token_counter.py | 94 ++- ...est_bedrock_count_tokens_transformation.py | 18 + .../test_bedrock_token_counter.py | 99 +++ .../vertex_ai/count_tokens/test_context.py | 230 +------ .../vertex_ai/test_vertex_ai_common_utils.py | 99 ++- tests/unit/proxy/test_proxy_token_counter.py | 638 +++++++++++++++++- 22 files changed, 1691 insertions(+), 393 deletions(-) create mode 100644 litellm/llms/anthropic/messages/context_management/inference_view.py create mode 100644 tests/unit/llms/anthropic/messages/context_management/test_inference_view.py diff --git a/litellm/llms/anthropic/messages/context_management/inference_view.py b/litellm/llms/anthropic/messages/context_management/inference_view.py new file mode 100644 index 00000000000..004beebe838 --- /dev/null +++ b/litellm/llms/anthropic/messages/context_management/inference_view.py @@ -0,0 +1,138 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Final + +from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.types.llms.anthropic import AnthropicMessagesTextParam +from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat + + +@dataclass(frozen=True, slots=True) +class InferenceRequestView: + messages: Sequence[Mapping[str, object]] + system: object | None + + +@dataclass(frozen=True, slots=True) +class InferenceRequestViewError: + error_message: str + status_code: int + + +@dataclass(frozen=True, slots=True) +class CountContextEdits: + spec: JsonValue + model: str + tools: tuple[Mapping[str, object], ...] | None + drop_invalid_edits: bool = False + + +_CONTEXT_MANAGEMENT_SPEC: Final[TypeAdapter[Mapping[str, object] | tuple[Mapping[str, object], ...]]] = TypeAdapter( + Mapping[str, object] | tuple[Mapping[str, object], ...] +) +_SYSTEM_BLOCKS: Final = TypeAdapter(tuple[AnthropicMessagesTextParam, ...]) + + +def count_context_edits( + context_management: CountContextManagement | None, + tools: Sequence[Mapping[str, object]] | None, + litellm_params: Mapping[str, object], + fallback_model: str, +) -> CountContextEdits: + model: Final = litellm_params.get("model") + return CountContextEdits( + spec=None if context_management is None else context_management.spec, + model=model if isinstance(model, str) else fallback_model, + tools=None if tools is None else tuple(tools), + drop_invalid_edits=context_management is not None and context_management.drop_invalid_edits, + ) + + +def _typed_spec(spec: JsonValue) -> Mapping[str, object] | tuple[Mapping[str, object], ...] | None: + try: + return _CONTEXT_MANAGEMENT_SPEC.validate_python(spec) + except ValidationError: + return None + + +def _typed_system(system: object | None) -> str | tuple[AnthropicMessagesTextParam, ...] | None | ValidationError: + if system is None or isinstance(system, str): + return system + try: + return _SYSTEM_BLOCKS.validate_python(system) + except ValidationError as error: + return error + + +def _view_without_edits( + messages: Sequence[Mapping[str, object]], + system: object | None, + typed_system: str | Sequence[Mapping[str, object]] | None, +) -> InferenceRequestView: + from litellm.llms.anthropic.messages.context_management.editors.compact import ( + apply_client_compaction_block_history, + ) + + history: Final = apply_client_compaction_block_history( + messages=[dict(message) for message in messages], + system=typed_system + if typed_system is None or isinstance(typed_system, str) + else [dict(block) for block in typed_system], + ) + if history is None: + return InferenceRequestView(messages=messages, system=system) + return InferenceRequestView(messages=history.messages, system=history.system) + + +async def _anthropic_inference_view( + messages: Sequence[Mapping[str, object]], system: object | None, edits: CountContextEdits +) -> InferenceRequestView | InferenceRequestViewError: + from litellm.llms.anthropic.messages.adapters.handler import prepare_context_managed_request + from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError + from litellm.llms.anthropic.messages.native.utils import normalize_native_messages + + typed_system: Final = _typed_system(system) + if isinstance(typed_system, ValidationError): + return InferenceRequestView(messages=messages, system=system) + normalized: Final = normalize_native_messages(messages) + spec: Final = _typed_spec(edits.spec) + try: + prepared: Final = await prepare_context_managed_request( + model=edits.model, + messages=normalized, + tools=None if edits.tools is None else [dict(tool) for tool in edits.tools], + system=[dict(block) for block in typed_system] if isinstance(typed_system, tuple) else typed_system, + context_management_spec=( + None if spec is None else dict(spec) if isinstance(spec, Mapping) else [dict(edit) for edit in spec] + ), + litellm_metadata=None, + additional_drop_params=None, + llm_router=None, + count_only=True, + drop_invalid_edits=edits.drop_invalid_edits, + ) + except AnthropicContextManagementError as error: + if error.status_code >= 500: + return _view_without_edits(normalized, system, typed_system) + return InferenceRequestViewError(error_message=error.message, status_code=error.status_code) + if prepared is None: + return InferenceRequestView(messages=normalized, system=system) + return InferenceRequestView(messages=prepared.messages, system=prepared.system) + + +async def inference_request_view( + messages: Sequence[Mapping[str, object]], + system: object | None, + messages_format: TokenCountMessagesFormat, + edits: CountContextEdits, +) -> InferenceRequestView | InferenceRequestViewError: + """Apply the history rewrites and context edits /v1/messages runs before it converts a body for another provider.""" + match messages_format: + case "openai": + return InferenceRequestView(messages=messages, system=system) + case "anthropic": + return await _anthropic_inference_view(messages, system, edits) + case _: + assert_never(messages_format) diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index c324f0f0980..97ee44bd3bb 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -12,6 +12,7 @@ from pydantic import JsonValue import litellm from litellm._logging import verbose_logger +from litellm.anthropic_beta_headers_manager import aload_beta_headers_config from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.llms.anthropic.common_utils import AnthropicError @@ -39,6 +40,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): tools: list[dict[str, JsonValue]] | None = None, system: object = None, optional_params: Mapping[str, JsonValue] | None = None, + context_management: JsonValue = None, ) -> dict[str, object]: """ Handle a CountTokens request using httpx with Azure authentication. @@ -70,6 +72,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): tools=tools, system=system, optional_params=optional_params, + context_management=context_management, drop_params=drops_unsupported_params((litellm_params or {}).get("drop_params")), ) @@ -80,11 +83,14 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): verbose_logger.debug("Making request to: %s", endpoint_url) + if context_management is not None: + await aload_beta_headers_config() # Get required headers with Azure authentication headers: Final = self.get_required_headers( api_key=api_key, litellm_params=litellm_params, messages=messages, + context_management=context_management, ) if not {"x-api-key", "Authorization"} & headers.keys(): raise AnthropicError( diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index af7cf9d81ce..14f7d4bfd2b 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -135,6 +135,7 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): tools=tools, system=system, optional_params=count_options, + context_management=None if context_management is None else context_management.spec, ) if result is not None: diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py index c02c92a95d6..91ed40d9271 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -8,9 +8,16 @@ from collections.abc import Sequence from typing import Final from urllib.parse import urlsplit +from pydantic import JsonValue + +from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION from litellm.litellm_core_utils.url_utils import join_api_base -from litellm.llms.anthropic.common_utils import count_tokens_feature_betas, merge_anthropic_beta_headers +from litellm.llms.anthropic.common_utils import ( + context_management_betas, + count_tokens_feature_betas, + merge_anthropic_beta_headers, +) from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountTokensConfig, ) @@ -31,6 +38,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): api_key: str | None, litellm_params: dict[str, object] | None = None, messages: Sequence[object] = (), + context_management: JsonValue = None, ) -> dict[str, str]: """ Get the required headers for the Azure AI Anthropic CountTokens API. @@ -42,6 +50,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): api_key: The Azure AI API key litellm_params: Optional LiteLLM parameters for additional auth config messages: The request messages, read for betas the counted features need + context_management: The edit spec the count forwards, read for the betas its edits need Returns: Dictionary of required headers with both x-api-key and Azure authentication @@ -51,7 +60,11 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): "Content-Type": "application/json", "anthropic-version": "2023-06-01", "anthropic-beta": merge_anthropic_beta_headers( - ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, count_tokens_feature_betas("azure_ai", {}, messages) + ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, + ( + *count_tokens_feature_betas("azure_ai", {}, messages), + *filter_and_transform_beta_headers(list(context_management_betas(context_management)), "azure_ai"), + ), ), **({"x-api-key": api_key} if api_key else {}), } diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index d6f5d5b6bd7..dca5e1c7a7b 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -43,6 +43,7 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.request_metadata import bedrock_request_metadata_is_owned from litellm.secret_managers.main import get_secret, get_secret_str +from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams if TYPE_CHECKING: @@ -50,6 +51,12 @@ if TYPE_CHECKING: from litellm.types.llms.openai import AllMessageValues BEDROCK_INVOKE_TRACE_FIELD: Final = "amazon-bedrock-trace" +BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS: Final[Mapping[str, str]] = MappingProxyType( + { + "compact_20260112": ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value, + "clear_tool_uses_20250919": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value, + } +) _ERROR_REQUEST_URL: Final = "https://docs.litellm.ai/docs" diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 26bc5904e06..dec099f37c0 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -13,9 +13,14 @@ from typing_extensions import assert_never import litellm from litellm._logging import verbose_logger from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.llms.anthropic.common_utils import context_management_betas, context_management_with_edits from litellm.llms.anthropic.count_tokens.transformation import ChatCountTranslationError, anthropic_count_fields from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model +from litellm.llms.bedrock.common_utils import ( + BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS, + BedrockError, + get_bedrock_base_model, +) from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model @@ -74,6 +79,16 @@ def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: ) +def _mantle_retry_outcome( + runtime: CountTokensOutcome, mantle: CountTokensOutcome, carried_edits: bool +) -> CountTokensOutcome: + match runtime, mantle: + case CountTokensFailure(status_code=400), CountTokensFailure(status_code=404) if carried_edits: + return runtime + case _: + return mantle + + def _request_data( model: str, messages: Sequence[Mapping[str, object]], tools: Sequence[Mapping[str, object]] | None, system: object ) -> Mapping[str, object]: @@ -85,6 +100,33 @@ def _request_data( } +def _invoke_model_message(message: Mapping[str, object]) -> Mapping[str, object]: + content: Final = message.get("content") + return ( + {**message, "content": [{"type": "text", "text": content}]} if isinstance(content, str) and content else message + ) + + +def _invoke_supports_edit(edit: JsonValue) -> bool: + edit_type: Final = edit.get("type") if isinstance(edit, dict) else None + return isinstance(edit_type, str) and edit_type in BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS + + +def _invoke_context_management(context_management: CountContextManagement | None) -> Mapping[str, JsonValue]: + spec: Final = context_management_with_edits(None if context_management is None else context_management.spec) + edits: Final = spec.get("edits") if isinstance(spec, dict) else None + supported_edits: Final = tuple(filter(_invoke_supports_edit, edits if isinstance(edits, list) else ())) + if not isinstance(spec, dict) or not supported_edits: + return MappingProxyType({}) + supported_spec: Final[JsonValue] = {**spec, "edits": list(supported_edits)} + return MappingProxyType( + { + "context_management": supported_spec, + "anthropic_beta": list(context_management_betas(supported_spec)), + } + ) + + class BedrockTokenCounter(BaseTokenCounter): """Counts on the configured Bedrock endpoint, with Mantle retry when Runtime rejects a Claude model""" @@ -167,11 +209,19 @@ class BedrockTokenCounter(BaseTokenCounter): ) if isinstance(translated, ChatCountTranslationError): return translated.response(request_model, model_to_use, primary_tokenizer) - runtime_request: Final = ( - _request_data(model_to_use, messages, tools, system) - if translated is None - else _request_data(model_to_use, translated.messages, translated.tools, translated.system) + native_context_management: Final = _invoke_context_management( + context_management if is_mantle_claude_model(resolved_model) else None ) + counted_messages: Final = messages if translated is None else translated.messages + runtime_request: Final = { + **_request_data( + model_to_use, + tuple(map(_invoke_model_message, counted_messages)) if native_context_management else counted_messages, + tools if translated is None else translated.tools, + system if translated is None else translated.system, + ), + **native_context_management, + } mantle_request: Final = runtime_request if translated is None else {**runtime_request, **translated.options} primary: Final = ( @@ -184,8 +234,12 @@ class BedrockTokenCounter(BaseTokenCounter): ) ) outcome: Final = ( - await self._count_with( - self._mantle_handler, MANTLE_TOKENIZER_TYPE, mantle_request, litellm_params, resolved_model + _mantle_retry_outcome( + primary, + await self._count_with( + self._mantle_handler, MANTLE_TOKENIZER_TYPE, mantle_request, litellm_params, resolved_model + ), + carried_edits=bool(native_context_management), ) if not use_mantle and _runtime_rejected_claude_model(primary, resolved_model) else primary diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py index 5cd26235ba2..5e5c5516b94 100644 --- a/litellm/llms/bedrock/count_tokens/mantle_handler.py +++ b/litellm/llms/bedrock/count_tokens/mantle_handler.py @@ -8,7 +8,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.url_utils import join_api_base -from litellm.llms.anthropic.common_utils import tools_without_custom_field +from litellm.llms.anthropic.common_utils import context_management_betas, tools_without_custom_field from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES, AnthropicCountTokensConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url @@ -22,6 +22,7 @@ class MantleCountTokensRequest(TypedDict): messages: ReadOnly[list[dict[str, JsonValue]]] system: ReadOnly[NotRequired[JsonValue]] tools: ReadOnly[NotRequired[list[dict[str, JsonValue]]]] + context_management: ReadOnly[NotRequired[JsonValue]] _COUNT_REQUEST: Final = TypeAdapter(MantleCountTokensRequest) @@ -74,13 +75,19 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM): optional_params=_COUNT_OPTIONS.validate_python( {name: value for name, value in request_data.items() if name in COUNT_TOKEN_OPTION_NAMES} ), + context_management=request.get("context_management"), ) + betas: Final = context_management_betas(request.get("context_management")) verbose_logger.debug("Making bedrock-mantle count_tokens request to: %s", endpoint_url) api_key: Final = litellm_params.get("api_key") signed_headers, signed_body = await run_aws_signing( self._sign_request, service_name="bedrock", - headers={"Content-Type": "application/json", "anthropic-version": MANTLE_ANTHROPIC_VERSION}, + headers={ + "Content-Type": "application/json", + "anthropic-version": MANTLE_ANTHROPIC_VERSION, + **({"anthropic-beta": ",".join(betas)} if betas else {}), + }, optional_params=litellm_params, request_data=body, api_base=endpoint_url, diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py index e5d6a168d57..f358dbe96db 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py @@ -37,6 +37,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation bedrock_invoke_guardrail_headers, ) from litellm.llms.bedrock.common_utils import ( + BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS, BedrockError, apply_bedrock_invoke_structured_output, bedrock_supports_tool_search, @@ -499,12 +500,7 @@ class AmazonAnthropicClaudeMessagesConfig( # Bedrock InvokeModel DOES support ``clear_tool_uses_20250919`` under the # ``context-management-2025-06-27`` beta. AWS docs: # https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-tool-use.md - _BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: Mapping[str, str] = MappingProxyType( - { - "compact_20260112": ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value, - "clear_tool_uses_20250919": ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value, - } - ) + _BEDROCK_INVOKE_SUPPORTED_CONTEXT_MANAGEMENT_EDITS: Mapping[str, str] = BEDROCK_INVOKE_CONTEXT_MANAGEMENT_EDIT_BETAS @classmethod def _filter_context_management_for_bedrock_invoke( diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 88c6494e3f7..74cd6f22704 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -537,13 +537,16 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): ) -> TokenCountResponse | None: import copy + from litellm.llms.anthropic.messages.context_management.inference_view import ( + InferenceRequestViewError, + count_context_edits, + inference_request_view, + ) from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContextError, - count_context_edits, count_context_error_response, gemini_count_contents, - inference_request_view, malformed_tool_result_text_count_error, provider_error_response, unconvertible_messages_error, @@ -556,7 +559,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): messages_format, count_context_edits(context_management, tools, (deployment or {}).get("litellm_params", {}), model_to_use), ) - if isinstance(request_view, VertexGeminiCountContextError): + if isinstance(request_view, InferenceRequestViewError): return count_context_error_response(request_view, request_model, model_to_use) count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools, model=model_to_use) if isinstance(count_context, VertexGeminiCountContextError): diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index ad64d852ef9..52345e1ff40 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1458,6 +1458,7 @@ class VertexAITokenCounter(BaseTokenCounter): ) -> TokenCountResponse | None: import copy + from litellm.llms.anthropic.common_utils import context_management_with_edits from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import count_tokens_publisher from litellm.llms.vertex_ai.vertex_ai_partner_models.main import ( VertexAIError as PartnerVertexAIError, @@ -1539,6 +1540,9 @@ class VertexAITokenCounter(BaseTokenCounter): system=system, tools=tools, count_options=count_options, + context_management=( + None if context_management is None else context_management_with_edits(context_management.spec) + ), ) except (PartnerVertexAIError, httpx.HTTPStatusError) as e: status_code: Final = e.response.status_code @@ -1565,12 +1569,15 @@ class VertexAITokenCounter(BaseTokenCounter): original_response=result, ) else: + from litellm.llms.anthropic.messages.context_management.inference_view import ( + InferenceRequestViewError, + count_context_edits, + inference_request_view, + ) from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContextError, - count_context_edits, count_context_error_response, gemini_count_contents, - inference_request_view, malformed_tool_result_text_count_error, provider_error_response, unconvertible_messages_error, @@ -1584,7 +1591,7 @@ class VertexAITokenCounter(BaseTokenCounter): messages_format, count_context_edits(context_management, tools, count_tokens_params_request, model_to_use), ) - if isinstance(request_view, VertexGeminiCountContextError): + if isinstance(request_view, InferenceRequestViewError): return count_context_error_response(request_view, request_model, model_to_use) count_context: Final = vertex_gemini_count_context( system=request_view.system, tools=tools, model=model_to_use diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py index fc7d6545624..c29f2ecbcf8 100644 --- a/litellm/llms/vertex_ai/count_tokens/context.py +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -11,10 +11,11 @@ from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_ from litellm.types.llms.anthropic import AnthropicMessagesTextParam from litellm.types.llms.openai import AllMessageValues from litellm.types.llms.vertex_ai import ContentType -from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse +from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse if TYPE_CHECKING: from litellm.exceptions import APIError + from litellm.llms.anthropic.messages.context_management.inference_view import InferenceRequestViewError @dataclass(frozen=True, slots=True) @@ -152,7 +153,7 @@ def malformed_tool_result_text_count_error( def count_context_error_response( - error: VertexGeminiCountContextError, request_model: str, model_used: str + error: "VertexGeminiCountContextError | InferenceRequestViewError", request_model: str, model_used: str ) -> TokenCountResponse: return TokenCountResponse( total_tokens=0, @@ -178,130 +179,6 @@ def provider_error_response(error: "APIError", request_model: str, model_used: s ) -@dataclass(frozen=True, slots=True) -class InferenceRequestView: - messages: Sequence[Mapping[str, object]] - system: object | None - - -@dataclass(frozen=True, slots=True) -class CountContextEdits: - spec: JsonValue - model: str - tools: tuple[Mapping[str, object], ...] | None - drop_invalid_edits: bool = False - - -_CONTEXT_MANAGEMENT_SPEC: Final[TypeAdapter[Mapping[str, object] | tuple[Mapping[str, object], ...]]] = TypeAdapter( - Mapping[str, object] | tuple[Mapping[str, object], ...] -) - - -def count_context_edits( - context_management: CountContextManagement | None, - tools: Sequence[Mapping[str, object]] | None, - litellm_params: Mapping[str, object], - fallback_model: str, -) -> CountContextEdits: - model: Final = litellm_params.get("model") - return CountContextEdits( - spec=None if context_management is None else context_management.spec, - model=model if isinstance(model, str) else fallback_model, - tools=None if tools is None else tuple(tools), - drop_invalid_edits=context_management is not None and context_management.drop_invalid_edits, - ) - - -def _typed_spec(spec: JsonValue) -> Mapping[str, object] | tuple[Mapping[str, object], ...] | None: - try: - return _CONTEXT_MANAGEMENT_SPEC.validate_python(spec) - except ValidationError: - return None - - -def _typed_system( - system: object | None, -) -> str | tuple[AnthropicMessagesTextParam, ...] | None | VertexGeminiCountContextError: - if system is None or isinstance(system, str): - return system - try: - return _SYSTEM_BLOCKS.validate_python(system) - except ValidationError: - return _INVALID_SYSTEM - - -def _view_without_edits( - messages: Sequence[Mapping[str, object]], - system: object | None, - typed_system: str | Sequence[Mapping[str, object]] | None, -) -> InferenceRequestView: - from litellm.llms.anthropic.messages.context_management.editors.compact import ( - apply_client_compaction_block_history, - ) - - history: Final = apply_client_compaction_block_history( - messages=[dict(message) for message in messages], - system=typed_system - if typed_system is None or isinstance(typed_system, str) - else [dict(block) for block in typed_system], - ) - if history is None: - return InferenceRequestView(messages=messages, system=system) - return InferenceRequestView(messages=history.messages, system=history.system) - - -async def _anthropic_inference_view( - messages: Sequence[Mapping[str, object]], system: object | None, edits: CountContextEdits -) -> InferenceRequestView | VertexGeminiCountContextError: - from litellm.llms.anthropic.messages.adapters.handler import prepare_context_managed_request - from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError - from litellm.llms.anthropic.messages.native.utils import normalize_native_messages - - typed_system: Final = _typed_system(system) - if isinstance(typed_system, VertexGeminiCountContextError): - return typed_system - normalized: Final = normalize_native_messages(messages) - spec: Final = _typed_spec(edits.spec) - try: - prepared: Final = await prepare_context_managed_request( - model=edits.model, - messages=normalized, - tools=None if edits.tools is None else [dict(tool) for tool in edits.tools], - system=[dict(block) for block in typed_system] if isinstance(typed_system, tuple) else typed_system, - context_management_spec=( - None if spec is None else dict(spec) if isinstance(spec, Mapping) else [dict(edit) for edit in spec] - ), - litellm_metadata=None, - additional_drop_params=None, - llm_router=None, - count_only=True, - drop_invalid_edits=edits.drop_invalid_edits, - ) - except AnthropicContextManagementError as error: - if error.status_code >= 500: - return _view_without_edits(normalized, system, typed_system) - return VertexGeminiCountContextError(error_message=error.message, status_code=error.status_code) - if prepared is None: - return InferenceRequestView(messages=normalized, system=system) - return InferenceRequestView(messages=prepared.messages, system=prepared.system) - - -async def inference_request_view( - messages: Sequence[Mapping[str, object]], - system: object | None, - messages_format: TokenCountMessagesFormat, - edits: CountContextEdits, -) -> InferenceRequestView | VertexGeminiCountContextError: - """Apply the history rewrites and context edits /v1/messages runs before it converts a body for Gemini.""" - match messages_format: - case "openai": - return InferenceRequestView(messages=messages, system=system) - case "anthropic": - return await _anthropic_inference_view(messages, system, edits) - case _: - assert_never(messages_format) - - def _chat_messages( messages: Sequence[Mapping[str, object]], messages_format: TokenCountMessagesFormat, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py index e341c8f5f48..8d39ae76456 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py @@ -9,12 +9,14 @@ their respective publisher-specific count-tokens endpoints. from collections.abc import Mapping, Sequence from typing import Any, Final -from pydantic import TypeAdapter +from pydantic import JsonValue, TypeAdapter +from litellm.anthropic_beta_headers_manager import aload_beta_headers_config, filter_and_transform_beta_headers from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.common_utils import ( OBJECT_LIST_ADAPTER, + context_management_betas, count_tokens_feature_betas, is_claude_model_target, tools_without_custom_field, @@ -31,6 +33,7 @@ _COUNT_TOKENS_PUBLISHERS: Final = ( _REQUEST_MESSAGES: Final = TypeAdapter(Sequence[object]) +_CONTEXT_MANAGEMENT: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) def count_tokens_request_body( @@ -55,8 +58,14 @@ def count_tokens_publisher(model: str) -> str | None: def count_tokens_headers(access_token: str, request_data: Mapping[str, object]) -> Mapping[str, str]: authorization: Final = {"Authorization": f"Bearer {access_token}"} - betas: Final = count_tokens_feature_betas( - "vertex_ai", request_data, _REQUEST_MESSAGES.validate_python(request_data["messages"]) + betas: Final = ( + *count_tokens_feature_betas( + "vertex_ai", request_data, _REQUEST_MESSAGES.validate_python(request_data["messages"]) + ), + *filter_and_transform_beta_headers( + list(context_management_betas(_CONTEXT_MANAGEMENT.validate_python(request_data.get("context_management")))), + "vertex_ai", + ), ) return {**authorization, "anthropic-beta": ",".join(betas)} if betas else authorization @@ -193,6 +202,8 @@ class VertexAIPartnerModelsTokenCounter(VertexBase): api_base=litellm_params.get("api_base"), ) + if request_data.get("context_management") is not None: + await aload_beta_headers_config() headers: Final = dict(count_tokens_headers(access_token, request_data)) # Get async HTTP client diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index a3a3427670d..c298ca404e2 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -268,6 +268,7 @@ class VertexAIPartnerModels(VertexBase): system: object | None = None, tools: list[dict[str, object]] | None = None, count_options: Mapping[str, object] = MappingProxyType({}), + context_management: object = None, ): """ Count tokens for Vertex AI partner models (Anthropic Claude, Mistral, etc.) @@ -302,7 +303,15 @@ class VertexAIPartnerModels(VertexBase): "model": model, "messages": messages, **MappingProxyType( - {key: value for key, value in (("system", system), ("tools", tools)) if value is not None} + { + key: value + for key, value in ( + ("system", system), + ("tools", tools), + ("context_management", context_management), + ) + if value is not None + } ), **count_options, } diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index 35136c41dc2..5bf719c6179 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -75,9 +75,7 @@ def _public_request( return native_call(legacy, args, _native_kwargs(fields, kwargs)) -def _resolved(request: NativeCall) -> tuple[str, str | None]: - model: Final = str(request.resolved["model"]) - provider: Final = optional_str(request.resolved.get("custom_llm_provider")) +def _resolved_model(model: str, provider: str | None) -> tuple[str, str | None]: try: resolved_model, resolved_provider, _, _ = get_llm_provider(model, provider) except BadRequestError: @@ -85,21 +83,43 @@ def _resolved(request: NativeCall) -> tuple[str, str | None]: return resolved_model, resolved_provider -def _target(request: NativeCall, model: str, provider: str | None) -> Target: - requested_model: Final = str(request.resolved["model"]) - model_info: Final = request.resolved.get("model_info") - api_base: Final = request.resolved.get("api_base") - native_config: Final = main._native_messages_config( # pyright: ignore[reportPrivateUsage] # the handler's own passthrough lookup - model, provider, model_info, api_base if isinstance(api_base, str) else None - ) +def _resolved(request: NativeCall) -> tuple[str, str | None]: + return _resolved_model(str(request.resolved["model"]), optional_str(request.resolved.get("custom_llm_provider"))) + + +def _route_target( + requested_model: str, + model: str, + provider: str | None, + model_info: object, + api_base: str | None, + compacting: bool, +) -> Target: + native_config: Final = main._native_messages_config(model, provider, model_info, api_base) # pyright: ignore[reportPrivateUsage] # the handler's own passthrough lookup if native_config is not None: return Target.MESSAGES - compacting: Final = request.resolved.get("compaction") is not None if not compacting and main._should_route_to_responses_api(provider, requested_model, model, model_info): # pyright: ignore[reportPrivateUsage] # the handler's own translation choice return Target.RESPONSES return Target.CHAT_COMPLETIONS +def _target(request: NativeCall, model: str, provider: str | None) -> Target: + api_base: Final = request.resolved.get("api_base") + return _route_target( + str(request.resolved["model"]), + model, + provider, + request.resolved.get("model_info"), + api_base if isinstance(api_base, str) else None, + compacting=request.resolved.get("compaction") is not None, + ) + + +def messages_target(model: str, custom_llm_provider: str | None, model_info: object, api_base: str | None) -> Target: + resolved_model, resolved_provider = _resolved_model(model, custom_llm_provider) + return _route_target(model, resolved_model, resolved_provider, model_info, api_base, compacting=False) + + _FEATURES: Final[detectors.Detectors] = MappingProxyType( { **detectors.COMMON, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7eace4ade9b..fb08272e687 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -33,7 +33,7 @@ from collections.abc import ( from collections.abc import ( Set as AbstractSet, ) -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import datetime, timedelta, timezone from itertools import chain from types import MappingProxyType, TracebackType, UnionType @@ -199,6 +199,7 @@ from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( snapshot_tuning_baselines, tuning_limit_violation, ) +from litellm.rust_bridge.catalog import Target from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ( PRICING_OVERRIDES_KEY, @@ -389,6 +390,12 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( mask_sensitive_keys, ) from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact +from litellm.llms.anthropic.messages.context_management.inference_view import ( + InferenceRequestView, + InferenceRequestViewError, + count_context_edits, + inference_request_view, +) from litellm.llms.anthropic.messages.convertible_blocks import ( deployment_drops_params, drop_non_string_tool_result_texts, @@ -399,6 +406,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.openai.live import parse_live_session_start from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.messages.dispatch import messages_target from litellm.proxy._lazy_features import attach_lazy_features, reserve_lazy_slot from litellm.proxy._types import * from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolver, resolve_agent_access_group_ceiling @@ -14927,7 +14935,7 @@ async def _try_provider_token_count( deployment: Mapping[str, object] | None, request_model: str, tools: list | None = None, - system: str | None = None, + system: object | None = None, messages_format: TokenCountMessagesFormat = "openai", output_format: JsonValue = None, context_management: CountContextManagement | None = None, @@ -14997,6 +15005,7 @@ async def _provider_token_count( request_model: str, settings: TokenCountSettings, local_estimate_allowed: bool, + forwarded_edits: bool = False, ) -> Union["TokenCountResponse", ProviderCountFailure, None]: """Count with the provider, using the token_count_model in its region when one is set. @@ -15020,7 +15029,7 @@ async def _provider_token_count( pinned_count_model: Final = None if count_model is None or count_model.region is None else count_model region_hint: Final = ( None - if pinned_count_model is None + if pinned_count_model is None or (forwarded_edits and counted.status_code == 400) else f"This model counts tokens with {pinned_count_model.model} in {pinned_count_model.region}. " "If that region is not enabled for the AWS account, enable it or set token_count_model in model_info" ) @@ -15282,6 +15291,9 @@ class _TokenCountDeployment(TypedDict, total=False): _TOKEN_COUNT_DEPLOYMENT: Final = TypeAdapter(_TokenCountDeployment) +_COUNT_TOOLS: Final[TypeAdapter[tuple[Mapping[str, object], ...] | None]] = TypeAdapter( + tuple[Mapping[str, object], ...] | None +) _ROUTED_DEPLOYMENT: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) @@ -15351,12 +15363,15 @@ def _count_context_management( model, general_settings_view(), user_model if isinstance(user_model, str) else None ) router_defaults: Final = _router_default_litellm_params(llm_router) - sent_params: Final = tuple( - {**deployment["litellm_params"], **router_defaults, **request_litellm_params} + deployment_params: Final = tuple( + deployment["litellm_params"] for deployment in llm_router.registry().deployments( inference_model if isinstance(inference_model, str) else model, source="routable" ) ) + sent_params: Final = tuple( + {**params, **router_defaults, **request_litellm_params} for params in deployment_params or ({},) + ) if any(context_management_explicitly_dropped(params.get("additional_drop_params")) for params in sent_params): return None return CountContextManagement( @@ -15366,6 +15381,90 @@ def _count_context_management( ) +_GATEWAY_EDIT_ROUTES: Final = frozenset({Target.CHAT_COMPLETIONS, Target.RESPONSES}) + + +def _messages_route(deployment: _TokenCountDeployment, request_model: str) -> Target: + params: Final = deployment.get("litellm_params", {}) + return messages_target( + _non_empty_str(params.get("model")) or request_model, + _non_empty_str(params.get("custom_llm_provider")), + deployment.get("model_info"), + _deployment_api_base(deployment), + ) + + +def _rejected_view(status_code: int, message: str) -> ProxyException: + return ProviderCountFailure( + status_code=status_code, message=message, rejected_before_provider_call=True + ).to_proxy_exception() + + +async def _bridged_request_view( + messages: Sequence[Mapping[str, object]], + system: object, + tools: Sequence[Mapping[str, object]] | None, + messages_format: TokenCountMessagesFormat, + edits: CountContextManagement | None, + deployment: _TokenCountDeployment, + fallback_model: str, +) -> InferenceRequestView | InferenceRequestViewError | None: + view: Final = await inference_request_view( + messages, + system, + messages_format, + count_context_edits(edits, tools, deployment.get("litellm_params", {}), fallback_model), + ) + return None if isinstance(view, InferenceRequestView) and not view.messages else view + + +async def _responses_request_view( + messages: Sequence[Mapping[str, object]], + system: object, + tools: Sequence[Mapping[str, object]] | None, + edits: CountContextManagement, + deployment: _TokenCountDeployment, + fallback_model: str, +) -> InferenceRequestView | InferenceRequestViewError: + from litellm.llms.anthropic.messages.context_management import AnthropicContextManagementError + from litellm.llms.anthropic.messages.responses_adapters.handler import count_view_messages + + try: + edited: Final = await count_view_messages( + model=_non_empty_str(deployment.get("litellm_params", {}).get("model")) or fallback_model, + messages=messages, + tools=tools, + system=system, + context_management=edits, + ) + except AnthropicContextManagementError as error: + if error.status_code >= 500: + return InferenceRequestView(messages=messages, system=system) + return InferenceRequestViewError(error_message=error.message, status_code=error.status_code) + return InferenceRequestView(messages=edited, system=system) + + +async def _gateway_request_view( + route: Target, + messages: Sequence[Mapping[str, object]], + system: object, + tools: Sequence[Mapping[str, object]] | None, + messages_format: TokenCountMessagesFormat, + edits: CountContextManagement | None, + deployment: _TokenCountDeployment, + fallback_model: str, +) -> InferenceRequestView | InferenceRequestViewError | None: + match route: + case Target.CHAT_COMPLETIONS: + return await _bridged_request_view( + messages, system, tools, messages_format, edits, deployment, fallback_model + ) + case Target.RESPONSES if edits is not None: + return await _responses_request_view(messages, system, tools, edits, deployment, fallback_model) + case _: + return None + + def _responses_routing_kwargs(responses_request: Mapping[str, object] | None) -> Mapping[str, object]: previous_response_id: Final = None if responses_request is None else responses_request.get("previous_response_id") return {"previous_response_id": previous_response_id} if isinstance(previous_response_id, str) else {} @@ -15559,6 +15658,40 @@ async def _count_request_tokens( ) prompt_only: Final = bool(prompt) and conversation_empty request_messages: Final = _counted_request_messages(messages, messages_format, count_deployment) + local_count_error: Final = malformed_tool_result_count_error(request_messages, messages_format, count_deployment) + malformed_messages: Final = ( + None + if local_count_error is None + else ProxyException(message=local_count_error, type="token_counting_error", param="messages", code=400) + ) + context_edits: Final = _count_context_management(context_management, request.model, request_litellm_params) + route: Final = ( + _messages_route(deployment_view, request.model) + if request_messages is not None and messages_format == "anthropic" + else None + ) + if context_edits is not None and route in _GATEWAY_EDIT_ROUTES and malformed_messages is not None: + raise malformed_messages + edited_request: Final = ( + None + if route is None or request_messages is None + else await _gateway_request_view( + route, + request_messages, + system, + _COUNT_TOOLS.validate_python(tools), + messages_format, + context_edits, + deployment_view, + model_to_use, + ) + ) + if isinstance(edited_request, InferenceRequestViewError): + raise _rejected_view(edited_request.status_code, edited_request.error_message) + inference_messages: Final = ( + request_messages if edited_request is None else [dict(message) for message in edited_request.messages] + ) + inference_system: Final = system if edited_request is None else edited_request.system contents_estimable: Final = not contents or (bool(messages) and _estimable_contents(contents)) prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES routed_deployment: Final = _ROUTED_DEPLOYMENT.validate_python(deployment) @@ -15598,30 +15731,30 @@ async def _count_request_tokens( if chat_count is not None else [{"role": "user", "content": prompt}] if prompt_as_message - else request_messages + else inference_messages ), contents=None if prompt_as_message else contents, deployment=count_deployment, request_model=request.model, tools=tools if chat_count is None else [dict(tool) for tool in chat_count.tools or ()] or None, - system=system, + system=inference_system, messages_format=messages_format, output_format=output_format, - context_management=_count_context_management(context_management, request.model, request_litellm_params), + context_management=None if edited_request is not None else context_edits, chat_params=MappingProxyType({}) if chat_count is None else chat_count.params, ), model_to_use=model_to_use, request_model=request.model, settings=settings, local_estimate_allowed=litellm.disable_token_counter is not True and contents_estimable, + forwarded_edits=context_edits is not None and edited_request is None, ) ) if isinstance(ladder, TokenCountResponse): return ladder - local_count_error: Final = malformed_tool_result_count_error(request_messages, messages_format, count_deployment) - if local_count_error is not None: - raise ProxyException(message=local_count_error, type="token_counting_error", param="messages", code=400) + if malformed_messages is not None: + raise malformed_messages # Check if token counter is disabled before fallback if litellm.disable_token_counter is True: @@ -15632,6 +15765,27 @@ async def _count_request_tokens( code=503, ) + local_view: Final = ( + None + if route is not Target.MESSAGES or context_edits is None or request_messages is None + else await _bridged_request_view( + request_messages, + system, + _COUNT_TOOLS.validate_python(tools), + messages_format, + replace(context_edits, drop_invalid_edits=False), + deployment_view, + model_to_use, + ) + ) + if isinstance(local_view, InferenceRequestViewError): + raise _rejected_view(local_view.status_code, local_view.error_message) + counted_view: Final = edited_request if local_view is None else local_view + local_messages: Final = ( + request_messages if counted_view is None else [dict(message) for message in counted_view.messages] + ) + local_system: Final = system if counted_view is None else counted_view.system + # Default LiteLLM token counting custom_tokenizer: CustomHuggingfaceTokenizer | None = None if model_info is not None: @@ -15639,11 +15793,11 @@ async def _count_request_tokens( CustomHuggingfaceTokenizer | None, model_info.get("custom_tokenizer", None), ) - system_message: Final = _system_message(system) + system_message: Final = _system_message(local_system) prompt_turn: Final = prompt_only and not (system_message is None and not tools) typed_messages: Final = cast( # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes Sequence[AllMessageValues] | None, - [{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else request_messages), + [{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else local_messages), ) counted_messages: Final = ( typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages) @@ -15654,7 +15808,15 @@ async def _count_request_tokens( framing: Final = settings.framing if counted_messages is not None and settings.framing is not None else 0 kimi_k3_messages: Final = counted_messages if settings.tokenizer == "kimi_k3" else None exact: Final = ( - await _kimi_k3_count(request, messages_format, model_to_use, kimi_k3_messages, counted_tools) + await _kimi_k3_count( + request + if counted_view is None + else request.model_copy(update={"messages": local_messages, "system": local_system}), + messages_format, + model_to_use, + kimi_k3_messages, + counted_tools, + ) if kimi_k3_messages is not None and custom_tokenizer is None else None ) diff --git a/tests/unit/llms/anthropic/messages/context_management/test_inference_view.py b/tests/unit/llms/anthropic/messages/context_management/test_inference_view.py new file mode 100644 index 00000000000..8c2747592cc --- /dev/null +++ b/tests/unit/llms/anthropic/messages/context_management/test_inference_view.py @@ -0,0 +1,239 @@ +import logging +from typing import Final +from unittest.mock import patch + +import pytest +from pydantic import JsonValue, TypeAdapter + +import litellm +from litellm.llms.anthropic.messages.context_management.inference_view import ( + CountContextEdits, + InferenceRequestView, + InferenceRequestViewError, + inference_request_view, +) + +_NO_EDITS: Final = CountContextEdits(spec=None, model="gemini/gemini-2.5-flash", tools=None) + + +@pytest.mark.parametrize( + "system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system") +) +async def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None: + view: Final = await inference_request_view( + ( + {"role": "user", "content": "first question " * 50}, + {"role": "assistant", "content": [{"type": "compaction", "content": "Summary of the earlier turns."}]}, + {"role": "user", "content": "latest question"}, + ), + system, + "anthropic", + _NO_EDITS, + ) + + assert isinstance(view, InferenceRequestView) + assert view.messages == [{"role": "user", "content": "latest question"}] + assert "Summary of the earlier turns." in str(view.system) + assert "Be brief." in str(view.system) + + +async def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None: + snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." + view: Final = await inference_request_view( + ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": "", + "snippet": snippet, + } + ], + }, + ], + }, + ), + None, + "anthropic", + _NO_EDITS, + ) + + assert isinstance(view, InferenceRequestView) + assistant_blocks: Final = TypeAdapter(list[dict[str, str]]).validate_python(view.messages[1]["content"]) + assert [block["type"] for block in assistant_blocks] == ["text"] + assert snippet in assistant_blocks[0]["text"] + + +async def test_openai_view_is_the_request_unchanged() -> None: + messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},) + + view: Final = await inference_request_view(messages, "Be brief.", "openai", _NO_EDITS) + + assert isinstance(view, InferenceRequestView) + assert view.messages == messages + assert view.system == "Be brief." + + +_TOOL_OUTPUT: Final = "line of a long file that the agent read earlier. " * 180 +_READ_FILE_TOOL: Final[dict[str, object]] = { + "name": "read_file", + "description": "Read a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, +} +_CLEAR_EDIT: Final[JsonValue] = { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 500}, + "keep": {"type": "tool_uses", "value": 1}, +} +_CLEAR_OLD_TOOL_RESULTS: Final[JsonValue] = {"edits": [_CLEAR_EDIT]} + + +def _agent_turns(rounds: int, first_round: int = 0) -> tuple[dict[str, JsonValue], ...]: + return ( + *( + turn + for index in range(first_round, first_round + rounds) + for turn in ( + {"role": "user", "content": f"Read file {index}."}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}"}} + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": f"toolu_{index}", "content": _TOOL_OUTPUT}], + }, + ) + ), # comprehension-ok: flattens each round into its three turns + {"role": "user", "content": "Summarize what you read."}, + ) + + +def _tool_outputs_kept(view: InferenceRequestView) -> tuple[bool, ...]: + return tuple( + _TOOL_OUTPUT in str(message["content"]) for message in view.messages if "tool_result" in str(message["content"]) + ) + + +def _edits(spec: JsonValue) -> CountContextEdits: + return CountContextEdits(spec=spec, model="gemini/gemini-2.5-flash", tools=(_READ_FILE_TOOL,)) + + +async def test_anthropic_view_clears_the_tool_results_inference_clears() -> None: + view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) + + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (False, False, True) + + +async def test_anthropic_view_applies_clear_edits_without_calling_the_compaction_summary_model() -> None: + spec: Final[JsonValue] = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50000}}, _CLEAR_EDIT] + } + turns: Final = ( + {"role": "user", "content": "first question " * 30000}, + {"role": "assistant", "content": "first answer"}, + *_agent_turns(3), + ) + + with ( + patch( + "litellm.llms.anthropic.messages.context_management.editors.compact._read_summary_model_setting", + return_value="gemini/gemini-2.5-flash", + ), + patch.object(litellm, "acompletion", side_effect=AssertionError("count_tokens must not summarize")) as summary, + ): + view: Final = await inference_request_view(turns, None, "anthropic", _edits(spec)) + + assert summary.call_count == 0 + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (False, False, True) + + +async def test_anthropic_view_counts_the_unedited_body_without_an_error_log_when_an_edit_fails( + caplog: pytest.LogCaptureFixture, +) -> None: + with ( + patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")), + caplog.at_level(logging.DEBUG, logger="LiteLLM"), + ): + view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) + + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (True, True, True) + assert any("skipping edits" in record.getMessage() for record in caplog.records) + assert not any(record.levelno >= logging.ERROR for record in caplog.records) + + +async def test_anthropic_view_counts_the_turns_after_a_compaction_block_when_an_edit_fails() -> None: + turns: Final = ( + *_agent_turns(3)[:-1], + {"role": "assistant", "content": [{"type": "compaction", "content": "Read files 0 to 2."}]}, + *_agent_turns(2, first_round=3), + ) + + with patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")): + view: Final = await inference_request_view(turns, None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) + + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (True, True) + assert "Read files 0 to 2." in str(view.system) + + +async def test_anthropic_view_clears_only_the_turns_after_a_compaction_block_when_clear_runs_before_compact() -> None: + clear_before_compact: Final[JsonValue] = { + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "tool_uses", "value": 3}, + "keep": {"type": "tool_uses", "value": 1}, + }, + {"type": "compact_20260112"}, + ] + } + turns: Final = ( + *_agent_turns(3)[:-1], + {"role": "assistant", "content": [{"type": "compaction", "content": "Read files 0 to 2."}]}, + *_agent_turns(2, first_round=3), + ) + + with patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")): + view: Final = await inference_request_view(turns, None, "anthropic", _edits(clear_before_compact)) + + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (True, True) + assert "Read files 0 to 2." in str(view.system) + + +async def test_anthropic_view_leaves_a_system_that_is_not_text_blocks_and_its_turns_to_the_provider_check() -> None: + image_system: Final[JsonValue] = [{"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}] + + view: Final = await inference_request_view( + _agent_turns(3), image_system, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS) + ) + + assert isinstance(view, InferenceRequestView) + assert _tool_outputs_kept(view) == (True, True, True) + assert view.system == image_system + + +async def test_anthropic_view_rejects_a_compact_trigger_inference_rejects() -> None: + low_trigger: Final[JsonValue] = { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] + } + + view: Final = await inference_request_view(_agent_turns(1), None, "anthropic", _edits(low_trigger)) + + assert isinstance(view, InferenceRequestViewError) + assert view.status_code == 400 diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py index 7807975be24..faa33e559ca 100644 --- a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py +++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py @@ -1,5 +1,7 @@ import json import operator +import threading +from collections.abc import Mapping from typing import Final import httpx @@ -8,8 +10,9 @@ import respx from pydantic import ValidationError import litellm +from litellm import anthropic_beta_headers_manager from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import AzureAIAnthropicTokenCounter -from litellm.types.utils import TokenCountResponse +from litellm.types.utils import CountContextManagement, TokenCountResponse API_BASE: Final = "https://my-resource.services.ai.azure.com" COUNT_URL: Final = f"{API_BASE}/anthropic/v1/messages/count_tokens" @@ -268,3 +271,92 @@ async def test_count_sends_the_tool_change_beta_when_the_history_adds_a_tool(mon assert result is not None assert result.total_tokens == 541 assert route.calls.last.request.headers["anthropic-beta"] == "mid-conversation-tool-changes-2026-07-01,token-counting-2024-11-01" + + +async def _betas_azure_inference_sends(context_management: Mapping[str, object]) -> frozenset[str]: + with respx.mock: + route: Final = respx.post(f"{API_BASE}/anthropic/v1/messages").mock( + return_value=httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 2}, + }, + ) + ) + await litellm.anthropic.messages.acreate( + max_tokens=64, + messages=[{"role": "user", "content": "Summarize the thread."}], + model="azure_ai/claude-sonnet-5-5", + api_key="azure-key", + api_base=API_BASE, + context_management=context_management, + ) + return frozenset(filter(None, route.calls.last.request.headers.get("anthropic-beta", "").split(","))) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +@pytest.mark.parametrize("edit_type", ["clear_tool_uses_20250919", "compact_20260112"]) +async def test_count_forwards_the_context_edits_with_the_betas_azure_inference_sends( + monkeypatch: pytest.MonkeyPatch, edit_type: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + spec: Final = {"edits": [{"type": edit_type}]} + inference_betas: Final = await _betas_azure_inference_sends(spec) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 88})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-5-5", + messages=[{"role": "user", "content": "Summarize the thread."}], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}}, + request_model="azure_ai/claude-sonnet-5-5", + context_management=CountContextManagement(spec=spec, drop_invalid_edits=False), + ) + assert result is not None + assert result.total_tokens == 88 + assert json.loads(route.calls.last.request.content)["context_management"] == spec + assert frozenset(route.calls.last.request.headers["anthropic-beta"].split(",")) == inference_betas | { + "token-counting-2024-11-01" + } + + +class _FetchThreads: + seen: tuple[int, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.seen = (*self.seen, threading.get_ident()) + return httpx.Response(200, json={}) + + +@pytest.mark.asyncio +async def test_count_with_an_edit_fetches_the_remote_beta_config_off_the_event_loop( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False) + monkeypatch.setattr(anthropic_beta_headers_manager, "_BETA_HEADERS_CONFIG", None) + fetch_threads: Final = _FetchThreads() + with respx.mock: + respx.get(litellm.anthropic_beta_headers_url).mock(side_effect=fetch_threads) + respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 88})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-5-5", + messages=[{"role": "user", "content": "Summarize the thread."}], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}}, + request_model="azure_ai/claude-sonnet-5-5", + context_management=CountContextManagement( + spec={"edits": [{"type": "clear_tool_uses_20250919"}]}, drop_invalid_edits=False + ), + ) + assert result is not None + assert len(fetch_threads.seen) == 1 + assert threading.get_ident() not in fetch_threads.seen diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py index 8ac35caba9d..d5b57e42146 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py @@ -1,6 +1,8 @@ import base64 import json +import pytest + from litellm.llms.bedrock.count_tokens.transformation import ( DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS, BedrockCountTokensConfig, @@ -53,6 +55,22 @@ def test_detect_input_type_converse_blocks_route_to_converse(): assert config._detect_input_type(request) == "converse" +@pytest.mark.parametrize("content", [[{"text": "hi"}], "hi"]) +@pytest.mark.parametrize( + "context_management", [None, {"edits": [{"type": "clear_tool_uses_20250919"}]}] +) +def test_a_pass_through_body_with_context_management_keeps_its_converse_input(content, context_management): + request = { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "messages": [{"role": "user", "content": content}], + "context_management": context_management, + } + + transformed = BedrockCountTokensConfig().transform_anthropic_to_bedrock_count_tokens(request) + + assert transformed["input"]["converse"]["messages"] == [{"role": "user", "content": [{"text": "hi"}]}] + + def test_transform_to_invoke_model_format_base64_encodes_body(): """The CountTokens API expects invokeModel.body as a base64-encoded blob; Anthropic Messages bodies additionally need anthropic_version/max_tokens diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py index cec7950b859..30063936b2b 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py @@ -9,6 +9,7 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.types.utils import CountContextManagement RUNTIME_HOST: Final = "bedrock-runtime.us-east-1.amazonaws.com" MANTLE_HOST: Final = "bedrock-mantle.us-east-1.api.aws" @@ -144,6 +145,104 @@ async def test_mantle_failure_is_reported_with_its_status() -> None: assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) +@pytest.mark.asyncio +async def test_runtime_rejection_of_an_edit_stands_when_mantle_does_not_serve_the_model() -> None: + rejection: Final = "'claude-haiku-4-5-20251001' does not support the 'compact_20260112' context management strategy" + bedrock: Final = _Bedrock( + runtime=(400, {"message": rejection}), + mantle=(404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}}), + ) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model="claude-haiku-4-5", + context_management=CountContextManagement( + spec={"edits": [{"type": "compact_20260112"}]}, drop_invalid_edits=False + ), + ) + + assert result is not None + assert result.error is True + assert result.status_code == 400 + assert result.tokenizer_type == "bedrock_api" + assert rejection in (result.error_message or "") + assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) + + +@pytest.mark.asyncio +async def test_runtime_400_without_an_edit_still_reports_the_mantle_failure() -> None: + bedrock: Final = _Bedrock( + runtime=(400, {"message": "The provided model identifier is invalid."}), + mantle=(404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}}), + ) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model="claude-haiku-4-5", + ) + + assert result is not None + assert result.status_code == 404 + assert result.tokenizer_type == "bedrock_mantle_api" + assert "does not exist" in (result.error_message or "") + assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) + + +@pytest.mark.asyncio +async def test_plain_text_history_with_an_edit_is_counted_through_invoke_carrying_the_edit() -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 57}), mantle=(200, {"input_tokens": 1})) + spec: Final[dict[str, JsonValue]] = {"edits": [{"type": "clear_tool_uses_20250919"}]} + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-opus-4-8", + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-4-8", + context_management=CountContextManagement(spec=spec, drop_invalid_edits=False), + ) + + assert result is not None + assert result.total_tokens == 57 + assert bedrock.posted_hosts() == (RUNTIME_HOST,) + invoke_body: Final = _invoke_body(bedrock.requests[0]) + assert invoke_body["context_management"] == spec + assert invoke_body["anthropic_beta"] == ["context-management-2025-06-27"] + + +@pytest.mark.asyncio +async def test_an_empty_string_turn_with_an_edit_is_sent_as_an_empty_string_not_an_empty_text_block() -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 16}), mantle=(200, {"input_tokens": 1})) + spec: Final[dict[str, JsonValue]] = {"edits": [{"type": "clear_tool_uses_20250919"}]} + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-opus-4-8", + messages=[ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": ""}, + {"role": "user", "content": "third question here"}, + ], + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-4-8", + context_management=CountContextManagement(spec=spec, drop_invalid_edits=False), + ) + + assert result is not None + assert result.total_tokens == 16 + assert _invoke_body(bedrock.requests[0])["messages"] == [ + {"role": "user", "content": [{"type": "text", "text": "first"}]}, + {"role": "assistant", "content": ""}, + {"role": "user", "content": [{"type": "text", "text": "third question here"}]}, + ] + + @pytest.mark.asyncio async def test_claude_runtime_404_retries_with_system_and_tools_on_mantle() -> None: bedrock: Final = _Bedrock(runtime=(404, UNSUPPORTED), mantle=(200, {"input_tokens": 2177})) diff --git a/tests/unit/llms/vertex_ai/count_tokens/test_context.py b/tests/unit/llms/vertex_ai/count_tokens/test_context.py index 67733884e12..233e5d888fa 100644 --- a/tests/unit/llms/vertex_ai/count_tokens/test_context.py +++ b/tests/unit/llms/vertex_ai/count_tokens/test_context.py @@ -1,6 +1,5 @@ import asyncio import base64 -import logging import uuid from collections.abc import Generator from contextlib import contextmanager @@ -11,20 +10,22 @@ from typing import Final, Literal from unittest.mock import patch import pytest -from pydantic import JsonValue, TypeAdapter +from pydantic import JsonValue import litellm from litellm.llms.anthropic.messages.adapters.transformation import ( LiteLLMAnthropicMessagesAdapter, truncate_tool_name, ) -from litellm.llms.vertex_ai.count_tokens.context import ( +from litellm.llms.anthropic.messages.context_management.inference_view import ( CountContextEdits, InferenceRequestView, + inference_request_view, +) +from litellm.llms.vertex_ai.count_tokens.context import ( VertexGeminiCountContext, VertexGeminiCountContextError, gemini_count_contents, - inference_request_view, unconvertible_messages_error, vertex_gemini_count_context, ) @@ -424,73 +425,6 @@ async def test_image_that_cannot_be_fetched_returns_an_error_value( assert isinstance(contents, VertexGeminiCountContextError) -@pytest.mark.parametrize( - "system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system") -) -async def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None: - view: Final = await inference_request_view( - ( - {"role": "user", "content": "first question " * 50}, - {"role": "assistant", "content": [{"type": "compaction", "content": "Summary of the earlier turns."}]}, - {"role": "user", "content": "latest question"}, - ), - system, - "anthropic", - _NO_EDITS, - ) - - assert isinstance(view, InferenceRequestView) - assert view.messages == [{"role": "user", "content": "latest question"}] - assert "Summary of the earlier turns." in str(view.system) - assert "Be brief." in str(view.system) - - -async def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None: - snippet: Final = "LiteLLM ships a weekly release candidate and a stable line." - view: Final = await inference_request_view( - ( - {"role": "user", "content": "What is new?"}, - { - "role": "assistant", - "content": [ - {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, - { - "type": "web_search_tool_result", - "tool_use_id": "srvtoolu_01", - "content": [ - { - "type": "web_search_result", - "url": "https://docs.litellm.ai/release_notes", - "title": "Release notes", - "encrypted_content": "", - "snippet": snippet, - } - ], - }, - ], - }, - ), - None, - "anthropic", - _NO_EDITS, - ) - - assert isinstance(view, InferenceRequestView) - assistant_blocks: Final = TypeAdapter(list[dict[str, str]]).validate_python(view.messages[1]["content"]) - assert [block["type"] for block in assistant_blocks] == ["text"] - assert snippet in assistant_blocks[0]["text"] - - -async def test_openai_view_is_the_request_unchanged() -> None: - messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},) - - view: Final = await inference_request_view(messages, "Be brief.", "openai", _NO_EDITS) - - assert isinstance(view, InferenceRequestView) - assert view.messages == messages - assert view.system == "Be brief." - - async def test_gcs_media_type_lookup_uses_the_deployment_vertex_project() -> None: with patch( "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type", @@ -556,138 +490,30 @@ async def test_anthropic_replayed_web_search_result_is_counted_as_its_text( assert snippet in str(contents) -_TOOL_OUTPUT: Final = "line of a long file that the agent read earlier. " * 180 -_READ_FILE_TOOL: Final[dict[str, object]] = { - "name": "read_file", - "description": "Read a file", - "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, -} -_CLEAR_EDIT: Final[JsonValue] = { - "type": "clear_tool_uses_20250919", - "trigger": {"type": "input_tokens", "value": 500}, - "keep": {"type": "tool_uses", "value": 1}, -} -_CLEAR_OLD_TOOL_RESULTS: Final[JsonValue] = {"edits": [_CLEAR_EDIT]} - - -def _agent_turns(rounds: int, first_round: int = 0) -> tuple[dict[str, JsonValue], ...]: - return ( - *( - turn - for index in range(first_round, first_round + rounds) - for turn in ( - {"role": "user", "content": f"Read file {index}."}, - { - "role": "assistant", - "content": [ - {"type": "tool_use", "id": f"toolu_{index}", "name": "read_file", "input": {"path": f"{index}"}} - ], - }, - { - "role": "user", - "content": [{"type": "tool_result", "tool_use_id": f"toolu_{index}", "content": _TOOL_OUTPUT}], - }, - ) - ), # comprehension-ok: flattens each round into its three turns - {"role": "user", "content": "Summarize what you read."}, - ) - - -def _tool_outputs_kept(view: InferenceRequestView) -> tuple[bool, ...]: - return tuple( - _TOOL_OUTPUT in str(message["content"]) for message in view.messages if "tool_result" in str(message["content"]) - ) - - -def _edits(spec: JsonValue) -> CountContextEdits: - return CountContextEdits(spec=spec, model="gemini/gemini-2.5-flash", tools=(_READ_FILE_TOOL,)) - - -async def test_anthropic_view_clears_the_tool_results_inference_clears() -> None: - view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) - - assert isinstance(view, InferenceRequestView) - assert _tool_outputs_kept(view) == (False, False, True) - - -async def test_anthropic_view_applies_clear_edits_without_calling_the_compaction_summary_model() -> None: - spec: Final[JsonValue] = { - "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50000}}, _CLEAR_EDIT] - } - turns: Final = ( - {"role": "user", "content": "first question " * 30000}, - {"role": "assistant", "content": "first answer"}, - *_agent_turns(3), - ) - - with ( - patch( - "litellm.llms.anthropic.messages.context_management.editors.compact._read_summary_model_setting", - return_value="gemini/gemini-2.5-flash", - ), - patch.object(litellm, "acompletion", side_effect=AssertionError("count_tokens must not summarize")) as summary, - ): - view: Final = await inference_request_view(turns, None, "anthropic", _edits(spec)) - - assert summary.call_count == 0 - assert isinstance(view, InferenceRequestView) - assert _tool_outputs_kept(view) == (False, False, True) - - -async def test_anthropic_view_counts_the_unedited_body_without_an_error_log_when_an_edit_fails( - caplog: pytest.LogCaptureFixture, -) -> None: - with ( - patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")), - caplog.at_level(logging.DEBUG, logger="LiteLLM"), - ): - view: Final = await inference_request_view(_agent_turns(3), None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) - - assert isinstance(view, InferenceRequestView) - assert _tool_outputs_kept(view) == (True, True, True) - assert any("skipping edits" in record.getMessage() for record in caplog.records) - assert not any(record.levelno >= logging.ERROR for record in caplog.records) - - -async def test_anthropic_view_clears_only_the_turns_after_a_compaction_block_when_clear_runs_before_compact() -> None: - clear_before_compact: Final[JsonValue] = { - "edits": [ - { - "type": "clear_tool_uses_20250919", - "trigger": {"type": "tool_uses", "value": 3}, - "keep": {"type": "tool_uses", "value": 1}, - }, - {"type": "compact_20260112"}, - ] - } - turns: Final = ( - *_agent_turns(3)[:-1], - {"role": "assistant", "content": [{"type": "compaction", "content": "Read files 0 to 2."}]}, - *_agent_turns(2, first_round=3), - ) - - with patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer unavailable")): - view: Final = await inference_request_view(turns, None, "anthropic", _edits(clear_before_compact)) - - assert isinstance(view, InferenceRequestView) - assert _tool_outputs_kept(view) == (True, True) - assert "Read files 0 to 2." in str(view.system) - - -async def test_anthropic_view_rejects_a_compact_trigger_inference_rejects() -> None: - low_trigger: Final[JsonValue] = { - "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] - } - - view: Final = await inference_request_view(_agent_turns(1), None, "anthropic", _edits(low_trigger)) - - assert isinstance(view, VertexGeminiCountContextError) - assert view.status_code == 400 - - async def test_anthropic_view_names_the_block_gemini_cannot_count_when_an_edit_trigger_cannot_tokenize_it() -> None: + clear_old_tool_results: Final = CountContextEdits( + spec={ + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 500}, + "keep": {"type": "tool_uses", "value": 1}, + } + ] + }, + model="gemini/gemini-2.5-flash", + tools=({"name": "read_file", "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}},), + ) turns: Final = ( - *_agent_turns(1)[:-1], + {"role": "user", "content": "Read file 0."}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "toolu_0", "name": "read_file", "input": {"path": "0"}}], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_0", "content": "line of a long file. " * 180}], + }, { "role": "user", "content": [ @@ -701,7 +527,7 @@ async def test_anthropic_view_names_the_block_gemini_cannot_count_when_an_edit_t }, ) - view: Final = await inference_request_view(turns, None, "anthropic", _edits(_CLEAR_OLD_TOOL_RESULTS)) + view: Final = await inference_request_view(turns, None, "anthropic", clear_old_tool_results) assert isinstance(view, InferenceRequestView) blocks_error: Final = unconvertible_messages_error(view.messages, "anthropic") diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index 3b222be6e5e..32d68249858 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1,16 +1,19 @@ import copy +import threading import time from collections.abc import Iterator from contextlib import contextmanager from enum import Enum from typing import Final -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import respx from pydantic import JsonValue, TypeAdapter import litellm +from litellm import anthropic_beta_headers_manager from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler @@ -1826,6 +1829,100 @@ async def test_vertex_ai_token_counter_forwards_system_and_tools_to_partner_requ assert without_optional_fields.total_tokens == 37 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("edit_type", "expected_beta"), + [("clear_tool_uses_20250919", "context-management-2025-06-27"), ("compact_20260112", "compact-2026-01-12")], +) +async def test_vertex_claude_count_forwards_the_context_edits_and_beta_vertex_inference_sends( + edit_type: str, expected_beta: str, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens import handler + + class FakeResponse: + status_code = 200 + + def json(self) -> dict[str, int]: + return {"input_tokens": 21} + + class FakeHttpClient: + posted: tuple[tuple[dict[str, str], dict[str, object]], ...] = () + + async def post( + self, url: str, headers: dict[str, str], json: dict[str, object], timeout: float + ) -> FakeResponse: + self.posted = (*self.posted, (headers, json)) + return FakeResponse() + + fake_http_client: Final = FakeHttpClient() + spec: Final[dict[str, JsonValue]] = {"edits": [{"type": edit_type}]} + application_default_credentials: Final = MagicMock(token="fake-token", valid=True, expiry=None) + monkeypatch.setattr("google.auth.default", lambda *args, **kwargs: (application_default_credentials, "test-project")) + with patch.object(handler, "get_async_httpx_client", return_value=fake_http_client): + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="claude-opus-4-8", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-east5"}}, + messages_format="anthropic", + context_management=CountContextManagement(spec=spec, drop_invalid_edits=False), + ) + + assert result is not None + assert result.total_tokens == 21 + ((headers, body),) = fake_http_client.posted + assert body["context_management"] == spec + assert headers["anthropic-beta"] == expected_beta + + +@pytest.mark.asyncio +async def test_vertex_claude_count_with_an_edit_fetches_the_remote_beta_config_off_the_event_loop( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens import handler + + class FakeResponse: + status_code = 200 + + def json(self) -> dict[str, int]: + return {"input_tokens": 21} + + class FakeHttpClient: + async def post( + self, url: str, headers: dict[str, str], json: dict[str, object], timeout: float + ) -> FakeResponse: + return FakeResponse() + + class FetchThreads: + seen: tuple[int, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.seen = (*self.seen, threading.get_ident()) + return httpx.Response(200, json={}) + + monkeypatch.delenv("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS", raising=False) + monkeypatch.setattr(anthropic_beta_headers_manager, "_BETA_HEADERS_CONFIG", None) + fetch_threads: Final = FetchThreads() + application_default_credentials: Final = MagicMock(token="fake-token", valid=True, expiry=None) + monkeypatch.setattr("google.auth.default", lambda *args, **kwargs: (application_default_credentials, "test-project")) + with patch.object(handler, "get_async_httpx_client", return_value=FakeHttpClient()), respx.mock: + respx.get(litellm.anthropic_beta_headers_url).mock(side_effect=fetch_threads) + result: Final = await VertexAITokenCounter().count_tokens( + model_to_use="claude-opus-4-8", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-east5"}}, + messages_format="anthropic", + context_management=CountContextManagement( + spec={"edits": [{"type": "clear_tool_uses_20250919"}]}, drop_invalid_edits=False + ), + ) + + assert result is not None + assert len(fetch_threads.seen) == 1 + assert threading.get_ident() not in fetch_threads.seen + + _PARTNER_CHAT_TOOL_TURNS: Final[list[dict[str, object]]] = [ {"role": "user", "content": "Weather in SF?"}, { diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 68eb05e60dd..608efe20c50 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -25,7 +25,7 @@ import pytest import respx import tiktoken from dotenv import load_dotenv -from pydantic import BaseModel, ConfigDict, JsonValue +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter from tokenizers import Tokenizer from tokenizers.models import WordLevel from tokenizers.pre_tokenizers import Whitespace @@ -82,6 +82,7 @@ from litellm.proxy.proxy_server import ( ) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted +from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem from litellm.types.utils import ( ADDRESSED_RESPONSE_ID_FIELD, @@ -1741,7 +1742,10 @@ _OPUS_IN_US_EAST_1: Final = TokenCountModel(model="us.anthropic.claude-opus-5-5" async def _ladder( - counter: _ScriptedCounter, settings: TokenCountSettings, local_estimate_allowed: bool = True + counter: _ScriptedCounter, + settings: TokenCountSettings, + local_estimate_allowed: bool = True, + forwarded_edits: bool = False, ) -> TokenCountResponse | None: return await _provider_token_count( count_with=partial( @@ -1759,6 +1763,7 @@ async def _ladder( request_model=_COUNT_REQUEST.model, settings=settings, local_estimate_allowed=local_estimate_allowed, + forwarded_edits=forwarded_edits, ) @@ -1986,6 +1991,26 @@ async def test_configured_token_count_model_error_stands_even_with_framing( assert raised.value.message.count("counts tokens with us.anthropic.claude-opus-5-5 in us-east-1") == region_hints +@pytest.mark.parametrize(("status_code", "region_hints"), [(400, 0), (403, 1)]) +@pytest.mark.asyncio +async def test_a_pinned_token_count_model_rejecting_forwarded_edits_answers_without_the_region_hint( + status_code: int, region_hints: int +): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} + ) + + with pytest.raises(ProxyException) as raised: + await _ladder( + counter, + TokenCountSettings(framing=None, count_model=replace(_OPUS_IN_US_EAST_1, configured=True)), + forwarded_edits=True, + ) + + assert raised.value.code == str(status_code) + assert raised.value.message.count("counts tokens with us.anthropic.claude-opus-5-5 in us-east-1") == region_hints + + @pytest.mark.parametrize( ("count_model", "region_hints"), [(_OPUS_IN_US_EAST_1, frozenset({True})), (replace(_OPUS_IN_US_EAST_1, region=None), frozenset())], @@ -3172,15 +3197,14 @@ _READ_FILE_TOOL: Final[dict[str, JsonValue]] = { "name": "read_file", "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, } -_CLEAR_TOOL_USES: Final[dict[str, JsonValue]] = { - "edits": [ - { - "type": "clear_tool_uses_20250919", - "trigger": {"type": "input_tokens", "value": 500}, - "keep": {"type": "tool_uses", "value": 1}, - } - ] -} +_CLEAR_TOOL_USES_EDITS: Final[list[JsonValue]] = [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 500}, + "keep": {"type": "tool_uses", "value": 1}, + } +] +_CLEAR_TOOL_USES: Final[dict[str, JsonValue]] = {"edits": _CLEAR_TOOL_USES_EDITS} _INVALID_COMPACT_EDIT: Final[dict[str, JsonValue]] = { "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}] } @@ -3516,6 +3540,28 @@ async def test_gemini_messages_count_tokens_counts_the_unedited_body_when_drop_p assert sent_bodies[0].decode().count(_GEMINI_TOOL_OUTPUT) == 3 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("router_defaults", "request_options", "expected_status"), + (({}, {"drop_params": True}, None), ({"drop_params": True}, {}, None), ({}, {}, 400)), + ids=("request-drop-params", "router-default-drop-params", "no-drop-params"), +) +async def test_messages_count_tokens_for_a_model_with_no_deployment_drops_an_invalid_edit_only_under_drop_params( + monkeypatch: pytest.MonkeyPatch, + router_defaults: dict[str, JsonValue], + request_options: dict[str, JsonValue], + expected_status: int | None, +) -> None: + monkeypatch.setenv("GEMINI_API_KEY", "test-key") + with patch.object(litellm, "drop_params", False): + response, _ = await _count_through_the_messages_endpoint( + Router(model_list=[_gemini_deployment("gemini-agent", {})], default_litellm_params=router_defaults), + {**request_options, "model": "gemini/gemini-3.8-pro", "context_management": _INVALID_COMPACT_EDIT}, + ) + + assert (response.status_code if isinstance(response, JSONResponse) else None) == expected_status + + @pytest.mark.asyncio async def test_gemini_messages_count_tokens_rejects_an_invalid_edit_without_drop_params() -> None: with patch.object(litellm, "drop_params", False): @@ -3529,6 +3575,576 @@ async def test_gemini_messages_count_tokens_rejects_an_invalid_edit_without_drop assert sent_bodies == () +_MISTRAL_MODEL: Final = "mistral/mistral-large-latest" + + +def _mistral_router( + litellm_params: Mapping[str, JsonValue] = MappingProxyType({}), + router_defaults: Mapping[str, JsonValue] = MappingProxyType({}), +) -> Router: + return Router( + model_list=[ + { + "model_name": "mistral-agent", + "litellm_params": {"model": _MISTRAL_MODEL, "api_key": "test-key", **litellm_params}, + } + ], + default_litellm_params=dict(router_defaults), + ) + + +async def _messages_endpoint_count(router: Router, request_body: Mapping[str, JsonValue]) -> object: + import litellm.proxy.anthropic_endpoints.endpoints as anthropic_endpoints + import litellm.proxy.proxy_server as proxy_server + + mock_request: Final = MagicMock(spec=Request) + mock_request.headers = {} + with ( + patch.object(proxy_server, "llm_router", router), + patch.object(anthropic_endpoints, "read_request_body", AsyncMock(return_value=dict(request_body))), + ): + return await anthropic_count_tokens(mock_request, MagicMock()) + + +class _MessagesCountReply(BaseModel): + model_config = ConfigDict(frozen=True) + input_tokens: int + + +async def _input_tokens(router: Router, request_body: Mapping[str, JsonValue]) -> int: + return _MessagesCountReply.model_validate(await _messages_endpoint_count(router, request_body)).input_tokens + + +async def _messages_the_inference_polyfill_sends(model: str) -> list[dict[str, JsonValue]]: + from litellm.llms.anthropic.messages.adapters.handler import prepare_context_managed_request + + prepared: Final = await prepare_context_managed_request( + model=model, + messages=[dict(message) for message in _agent_turns_with_tool_results(3)], + tools=[dict(_READ_FILE_TOOL)], + system=None, + context_management_spec=dict(_CLEAR_TOOL_USES), + litellm_metadata=None, + additional_drop_params=None, + llm_router=None, + ) + assert prepared is not None + return TypeAdapter(list[dict[str, JsonValue]]).validate_python(prepared.messages) + + +@pytest.mark.asyncio +async def test_messages_count_tokens_counts_the_body_the_inference_polyfill_sends_on_a_bridged_provider() -> None: + router: Final = _mistral_router() + unedited_body: Final[dict[str, JsonValue]] = { + "model": "mistral-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + } + + edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES}) + unedited: Final = await _input_tokens(router, unedited_body) + polyfilled: Final = await _input_tokens( + router, {**unedited_body, "messages": await _messages_the_inference_polyfill_sends(_MISTRAL_MODEL)} + ) + + assert edited < unedited + assert edited == polyfilled + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployment_params", "router_defaults", "request_options"), + ( + ({"additional_drop_params": ["context_management"]}, {}, {}), + ({}, {"additional_drop_params": ["context_management"]}, {}), + ({}, {}, {"additional_drop_params": ["context_management"]}), + ), + ids=("deployment-drops", "router-default-drops", "request-drops"), +) +async def test_messages_count_tokens_on_a_bridged_provider_keeps_the_unedited_count_when_context_management_is_dropped( + deployment_params: dict[str, JsonValue], + router_defaults: dict[str, JsonValue], + request_options: dict[str, JsonValue], +) -> None: + router: Final = _mistral_router(deployment_params, router_defaults) + unedited_body: Final[dict[str, JsonValue]] = { + "model": "mistral-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + **request_options, + } + + edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES}) + + assert edited == await _input_tokens(router, unedited_body) + + +@pytest.mark.asyncio +async def test_messages_count_tokens_on_a_bridged_provider_rejects_an_invalid_edit_without_drop_params() -> None: + with patch.object(litellm, "drop_params", False): + response: Final = await _messages_endpoint_count( + _mistral_router(), + { + "model": "mistral-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "context_management": _INVALID_COMPACT_EDIT, + }, + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + + +@pytest.mark.asyncio +async def test_messages_count_tokens_on_a_bridged_provider_counts_without_an_invalid_edit_under_drop_params() -> None: + router: Final = _mistral_router() + unedited_body: Final[dict[str, JsonValue]] = { + "model": "mistral-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "drop_params": True, + } + + with patch.object(litellm, "drop_params", False): + edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _INVALID_COMPACT_EDIT}) + unedited: Final = await _input_tokens(router, unedited_body) + + assert edited == unedited + + +def _turns_compacted_by_the_client() -> list[dict[str, JsonValue]]: + return [ + *_agent_turns_with_tool_results(3)[:-1], + {"role": "assistant", "content": [{"type": "compaction", "content": "The agent read files 0, 1 and 2."}]}, + {"role": "user", "content": "Summarize what you read."}, + ] + + +@pytest.mark.asyncio +async def test_messages_count_tokens_on_a_bridged_provider_counts_only_what_follows_a_client_compaction_block() -> None: + from litellm.llms.anthropic.messages.adapters.handler import prepare_context_managed_request + + router: Final = _mistral_router() + sent: Final = await prepare_context_managed_request( + model=_MISTRAL_MODEL, + messages=[dict(message) for message in _turns_compacted_by_the_client()], + tools=[dict(_READ_FILE_TOOL)], + system=None, + context_management_spec=None, + litellm_metadata=None, + additional_drop_params=None, + llm_router=None, + ) + assert sent is not None + + compacted: Final = await _input_tokens( + router, {"model": "mistral-agent", "messages": _turns_compacted_by_the_client(), "tools": [_READ_FILE_TOOL]} + ) + uncompacted: Final = await _input_tokens( + router, {"model": "mistral-agent", "messages": _agent_turns_with_tool_results(3), "tools": [_READ_FILE_TOOL]} + ) + sent_body_count: Final = await _input_tokens( + router, + { + "model": "mistral-agent", + "messages": TypeAdapter(list[dict[str, JsonValue]]).validate_python(sent.messages), + "system": TypeAdapter(JsonValue).validate_python(sent.system), + "tools": [_READ_FILE_TOOL], + }, + ) + + assert compacted < uncompacted + assert compacted == sent_body_count + + +@pytest.mark.asyncio +async def test_messages_count_tokens_applies_the_edits_a_responses_api_bridge_polyfills_for_a_non_openai_provider() -> ( + None +): + perplexity_model: Final = "perplexity/perplexity/sonar" + router: Final = Router( + model_list=[{"model_name": "sonar-agent", "litellm_params": {"model": perplexity_model, "api_key": "test-key"}}] + ) + unedited_body: Final[dict[str, JsonValue]] = { + "model": "sonar-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + } + + edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES}) + unedited: Final = await _input_tokens(router, unedited_body) + polyfilled: Final = await _input_tokens( + router, {**unedited_body, "messages": await _messages_the_inference_polyfill_sends(perplexity_model)} + ) + + assert edited < unedited + assert edited == polyfilled + + +@pytest.mark.asyncio +async def test_messages_count_tokens_on_a_responses_api_bridge_counts_the_unedited_body_when_an_edit_fails() -> None: + from litellm.llms.anthropic.messages.responses_adapters import handler as responses_adapters_handler + + router: Final = Router( + model_list=[ + { + "model_name": "sonar-agent", + "litellm_params": {"model": "perplexity/perplexity/sonar", "api_key": "test-key"}, + } + ] + ) + unedited_body: Final[dict[str, JsonValue]] = { + "model": "sonar-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + } + + with patch.object( + responses_adapters_handler, "apply_context_management", AsyncMock(side_effect=RuntimeError("editor bug")) + ): + edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES}) + + assert edited == await _input_tokens(router, unedited_body) + + +@pytest.mark.asyncio +async def test_messages_count_tokens_on_a_responses_api_bridge_rejects_an_invalid_edit_without_drop_params() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "sonar-agent", + "litellm_params": {"model": "perplexity/perplexity/sonar", "api_key": "test-key"}, + } + ] + ) + invalid_clear_edit: Final[dict[str, JsonValue]] = { + "edits": [{"type": "clear_tool_uses_20250919", "clear_at_least": "all of them"}] + } + + with patch.object(litellm, "drop_params", False): + response: Final = await _messages_endpoint_count( + router, + { + "model": "sonar-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "context_management": invalid_clear_edit, + }, + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + message: Final = json.loads(bytes(response.body))["error"]["message"] + assert "context_management.edits.0.clear_tool_uses_20250919.clear_at_least" in message + + +class _CountTokensPeer: + def __init__(self, reply: Mapping[str, JsonValue]) -> None: + self._reply: Final = reply + self.requests: tuple[httpx.Request, ...] = () + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.requests = (*self.requests, request) + return httpx.Response(200, json=dict(self._reply), request=request) + + +class _BedrockInvokeModel(BaseModel): + model_config = ConfigDict(frozen=True) + body: str + + +class _BedrockCountTokensInput(BaseModel): + model_config = ConfigDict(frozen=True) + invoke_model: _BedrockInvokeModel = Field(alias="invokeModel") + + +class _BedrockCountTokensBody(BaseModel): + model_config = ConfigDict(frozen=True) + input: _BedrockCountTokensInput + + +def _bedrock_counted_request(request_body: bytes) -> str: + invoke_model: Final = _BedrockCountTokensBody.model_validate_json(request_body).input.invoke_model + return base64.b64decode(invoke_model.body).decode() + + +async def _count_on_a_bridged_provider_count_api( + request_body: Mapping[str, JsonValue], + monkeypatch: pytest.MonkeyPatch, + bedrock_model: str = "bedrock/us.amazon.nova-2-lite-v1:0", + peer_reply: Mapping[str, JsonValue] = MappingProxyType({"inputTokens": 57}), + model_info: Mapping[str, JsonValue] = MappingProxyType({}), +) -> tuple[object, tuple[httpx.Request, ...]]: + from litellm.llms.bedrock.count_tokens import handler as bedrock_count_tokens_handler + from litellm.llms.bedrock.count_tokens import mantle_handler as bedrock_mantle_count_tokens_handler + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + peer: Final = _CountTokensPeer(peer_reply) + router: Final = Router( + model_list=[ + { + "model_name": "nova-agent", + "litellm_params": { + "model": bedrock_model, + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-east-1", + }, + "model_info": dict(model_info), + } + ] + ) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer)) + with ( + patch.object(bedrock_count_tokens_handler, "get_async_httpx_client", return_value=client), + patch.object(bedrock_mantle_count_tokens_handler, "get_async_httpx_client", return_value=client), + ): + response: Final = await _messages_endpoint_count(router, {**request_body, "model": "nova-agent"}) + return response, peer.requests + + +@pytest.mark.asyncio +async def test_messages_count_tokens_sends_the_edited_messages_to_a_bridged_provider_count_api( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response, sent_bodies = await _count_on_a_bridged_provider_count_api( + { + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "context_management": _CLEAR_TOOL_USES, + }, + monkeypatch, + ) + + assert response == {"input_tokens": 57, "litellm_count_source": "provider"} + assert len(sent_bodies) == 1 + assert _bedrock_counted_request(sent_bodies[0].content).count(_GEMINI_TOOL_OUTPUT) == 1 + + +@pytest.mark.asyncio +async def test_messages_count_tokens_sends_a_body_with_only_empty_text_to_the_bridged_provider_count_api( + monkeypatch: pytest.MonkeyPatch, +) -> None: + empty_turn: Final[dict[str, JsonValue]] = {"role": "user", "content": [{"type": "text", "text": ""}]} + + response, sent_bodies = await _count_on_a_bridged_provider_count_api({"messages": [empty_turn]}, monkeypatch) + + assert response == {"input_tokens": 57, "litellm_count_source": "provider"} + assert len(sent_bodies) == 1 + + +_BEDROCK_CLAUDE_MODEL: Final = "bedrock/us.anthropic.claude-sonnet-5-5" +_CLEAR_THINKING_EDIT: Final[dict[str, JsonValue]] = {"type": "clear_thinking_20251015"} +_COUNTED_ON_MANTLE: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"token_count_model": {"model": "us.anthropic.claude-sonnet-5-5", "litellm_provider": "bedrock_mantle"}} +) +_COUNTED_ON_RUNTIME: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"token_count_model": {"model": "us.anthropic.claude-sonnet-5-5", "litellm_provider": "bedrock"}} +) + + +def _bedrock_claude_count_request(edits: list[JsonValue]) -> Mapping[str, JsonValue]: + return MappingProxyType( + { + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "context_management": {"edits": edits}, + } + ) + + +@pytest.mark.asyncio +async def test_messages_count_tokens_forwards_bedrock_supported_edits_and_their_beta_to_the_mantle_count_api( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response, sent = await _count_on_a_bridged_provider_count_api( + _bedrock_claude_count_request([*_CLEAR_TOOL_USES_EDITS, _CLEAR_THINKING_EDIT]), + monkeypatch, + _BEDROCK_CLAUDE_MODEL, + MappingProxyType({"input_tokens": 57, "context_management": {"original_input_tokens": 4100}}), + _COUNTED_ON_MANTLE, + ) + + assert isinstance(response, dict) + assert (response["input_tokens"], response["context_management"]) == (57, {"original_input_tokens": 4100}) + assert len(sent) == 1 + assert json.loads(sent[0].content)["context_management"] == _CLEAR_TOOL_USES + assert sent[0].headers["anthropic-beta"] == ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value + + +@pytest.mark.asyncio +async def test_messages_count_tokens_forwards_bedrock_supported_edits_and_their_beta_to_the_runtime_count_api( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response, sent = await _count_on_a_bridged_provider_count_api( + _bedrock_claude_count_request([*_CLEAR_TOOL_USES_EDITS, _CLEAR_THINKING_EDIT]), + monkeypatch, + _BEDROCK_CLAUDE_MODEL, + model_info=_COUNTED_ON_RUNTIME, + ) + + assert isinstance(response, dict) + assert response["input_tokens"] == 57 + assert len(sent) == 1 + invoke_body: Final = json.loads(_bedrock_counted_request(sent[0].content)) + assert (invoke_body["context_management"], invoke_body["anthropic_beta"]) == ( + _CLEAR_TOOL_USES, + [ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value], + ) + + +@pytest.mark.asyncio +async def test_messages_count_tokens_sends_no_context_management_to_bedrock_when_no_edit_is_bedrock_supported( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response, sent = await _count_on_a_bridged_provider_count_api( + _bedrock_claude_count_request([_CLEAR_THINKING_EDIT]), + monkeypatch, + _BEDROCK_CLAUDE_MODEL, + MappingProxyType({"input_tokens": 57}), + _COUNTED_ON_MANTLE, + ) + + assert isinstance(response, dict) + assert response["input_tokens"] == 57 + assert len(sent) == 1 + assert "context_management" not in json.loads(sent[0].content) + assert "anthropic-beta" not in sent[0].headers + + +@pytest.mark.asyncio +async def test_messages_count_tokens_rejects_a_non_string_tool_result_text_before_a_bridged_provider_applies_edits( + monkeypatch: pytest.MonkeyPatch, +) -> None: + response, sent_bodies = await _count_on_a_bridged_provider_count_api( + {"messages": _build_log_messages(5), "context_management": _CLEAR_TOOL_USES}, monkeypatch + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + assert "tool_result toolu_log content.1.text" in json.loads(bytes(response.body))["error"]["message"] + assert sent_bodies == () + + +@pytest.mark.asyncio +async def test_messages_count_tokens_applies_a_responses_api_deployment_edits_once() -> None: + peer: Final = _CountTokensPeer({"object": "response.input_tokens", "input_tokens": 71}) + router: Final = Router( + model_list=[{"model_name": "gpt-agent", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "test-key"}}] + ) + + with patch.object( + openai_count_tokens_http, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ): + response: Final = await _messages_endpoint_count( + router, + { + "model": "gpt-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "context_management": _CLEAR_TOOL_USES, + }, + ) + + assert response == {"input_tokens": 71, "litellm_count_source": "provider"} + assert len(peer.requests) == 1 + assert peer.requests[0].content.decode().count(_GEMINI_TOOL_OUTPUT) == 1 + + +_DATABRICKS_CLAUDE_MODEL: Final = "databricks/databricks-claude-sonnet-4-6" + + +def _databricks_claude_router() -> Router: + return Router( + model_list=[ + { + "model_name": "databricks-claude-agent", + "litellm_params": { + "model": _DATABRICKS_CLAUDE_MODEL, + "api_key": "test-key", + "api_base": "https://example.cloud.databricks.com/serving-endpoints", + }, + } + ] + ) + + +@pytest.mark.asyncio +async def test_messages_count_tokens_counts_the_edited_body_on_a_native_messages_deployment_without_a_count_api() -> ( + None +): + router: Final = _databricks_claude_router() + unedited_body: Final[dict[str, JsonValue]] = { + "model": "databricks-claude-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + } + + edited: Final = await _input_tokens(router, {**unedited_body, "context_management": _CLEAR_TOOL_USES}) + unedited: Final = await _input_tokens(router, unedited_body) + polyfilled: Final = await _input_tokens( + router, {**unedited_body, "messages": await _messages_the_inference_polyfill_sends(_DATABRICKS_CLAUDE_MODEL)} + ) + + assert edited < unedited + assert edited == polyfilled + + +@pytest.mark.asyncio +@pytest.mark.parametrize("drop_params", [False, True]) +async def test_messages_count_tokens_on_a_native_messages_deployment_without_a_count_api_rejects_an_invalid_edit( + drop_params: bool, +) -> None: + with patch.object(litellm, "drop_params", drop_params): + response: Final = await _messages_endpoint_count( + _databricks_claude_router(), + { + "model": "databricks-claude-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "context_management": _INVALID_COMPACT_EDIT, + }, + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + + +@pytest.mark.asyncio +async def test_messages_count_tokens_leaves_the_edits_to_a_native_messages_deployment() -> None: + peer: Final = _CountTokensPeer({"input_tokens": 63}) + router: Final = Router( + model_list=[ + { + "model_name": "claude-agent", + "litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": "test-key"}, + } + ] + ) + + with patch.object( + anthropic_count_tokens_handler, + "get_async_httpx_client", + return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)), + ): + response: Final = await _messages_endpoint_count( + router, + { + "model": "claude-agent", + "messages": _agent_turns_with_tool_results(3), + "tools": [_READ_FILE_TOOL], + "context_management": _CLEAR_TOOL_USES, + }, + ) + + assert response == {"input_tokens": 63, "litellm_count_source": "provider"} + assert len(peer.requests) == 1 + assert peer.requests[0].content.decode().count(_GEMINI_TOOL_OUTPUT) == 3 + + @pytest.mark.asyncio @pytest.mark.parametrize( "text", ["hello there", "The quick brown fox jumps over the lazy dog. " * 50], ids=["short", "long"] From f91774e2f9e47b87062445d755e2bfa2ef1f0fe8 Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Mon, 2 Nov 2026 10:46:30 -0800 Subject: [PATCH 13/13] fix(fireworks): keep developer items in place for models that accept a later system message (#49399) * fix(fireworks): keep developer items in place for models that accept a later system message * fix(fireworks_ai): keep developer items in place for short Fireworks model names and stop flagging gpt-oss-120b * fix(fireworks): drop the mid-conversation system flag from nemotron-lightning nemotron-lightning-3p5-30b-a3b followed a later instruction less often with the developer item kept in place than folded, so both of its cost map keys go back to folding. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../fireworks_ai/responses/transformation.py | 16 +++- ...odel_prices_and_context_window_backup.json | 31 ++++++++ model_prices_and_context_window.json | 31 ++++++++ ...t_fireworks_ai_responses_transformation.py | 78 +++++++++++++++++-- 4 files changed, 148 insertions(+), 8 deletions(-) diff --git a/litellm/llms/fireworks_ai/responses/transformation.py b/litellm/llms/fireworks_ai/responses/transformation.py index 463b9cbfc9a..d599f0b65f7 100644 --- a/litellm/llms/fireworks_ai/responses/transformation.py +++ b/litellm/llms/fireworks_ai/responses/transformation.py @@ -29,6 +29,7 @@ from litellm.utils import ( declared_mid_conversation_system_support, deployment_declared_mid_conversation_system_support, deployment_model_info, + get_bundled_model_cost_map, ) if TYPE_CHECKING: @@ -54,6 +55,15 @@ def _forwarded_user_params(litellm_params: GenericLiteLLMParams) -> Mapping[str, ) +def _bundled_cost_map_declares_mid_conversation_system(model: str) -> bool: + bundled_cost_map: Final[Mapping[str, object]] = get_bundled_model_cost_map() + names: Final = (model.removeprefix("fireworks_ai/"), resolve_fireworks_resource_name(model)) + return any( + deployment_declared_mid_conversation_system_support(bundled_cost_map.get(f"fireworks_ai/{name}")) is True + for name in names + ) + + class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig): @property def custom_llm_provider(self) -> LlmProviders: @@ -90,9 +100,13 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig): validated_input: Final = self._validate_input_param(input) model_info: Final = deployment_model_info(litellm_params) declared: Final = declared_mid_conversation_system_support(model, self.custom_llm_provider.value, model_info) + keeps_developer_items_in_place: Final = declared is True and ( + deployment_declared_mid_conversation_system_support(model_info) is True + or _bundled_cost_map_declares_mid_conversation_system(model) + ) folded: Final = ( None - if deployment_declared_mid_conversation_system_support(model_info) is True + if keeps_developer_items_in_place else instruction_items_folded( validated_input, response_api_optional_request_params, "user" if declared is False else "system" ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index f2a3bb1a47e..02bf7dedcef 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -64039,6 +64039,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64058,6 +64059,7 @@ "output_cost_per_token_priority": 2.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64100,6 +64102,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64137,6 +64140,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64156,6 +64160,7 @@ "output_cost_per_token_priority": 2.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64230,6 +64235,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64250,6 +64256,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64273,6 +64280,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64348,6 +64356,7 @@ "output_cost_per_token_priority": 3e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64384,6 +64393,7 @@ "output_cost_per_token_priority": 6.375e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_vision": true }, @@ -64420,6 +64430,7 @@ "output_cost_per_token_priority": 3e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64490,6 +64501,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64513,6 +64525,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67306,6 +67319,7 @@ "output_cost_per_token_priority": 5.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67324,6 +67338,7 @@ "output_cost_per_token_priority": 8.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67342,6 +67357,7 @@ "output_cost_per_token_priority": 5.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67360,6 +67376,7 @@ "output_cost_per_token_priority": 8.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67375,6 +67392,7 @@ "output_cost_per_token": 6.6e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67417,6 +67435,7 @@ "output_cost_per_token": 6.6e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67435,6 +67454,7 @@ "output_cost_per_token_priority": 6.25e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67453,6 +67473,7 @@ "output_cost_per_token_priority": 9.375e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67470,6 +67491,7 @@ "output_cost_per_token_priority": 6.25e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67488,6 +67510,7 @@ "output_cost_per_token_priority": 9.375e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67506,6 +67529,7 @@ "output_cost_per_token_priority": 4.05e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -73747,6 +73771,7 @@ "output_cost_per_token_priority": 1.875e-05, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -77249,6 +77274,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -77407,6 +77433,7 @@ "output_cost_per_token_priority": 9e-07, "source": "https://platform.minimax.io/docs/guides/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -81184,6 +81211,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81200,6 +81228,7 @@ "output_cost_per_token": 3.08e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81216,6 +81245,7 @@ "output_cost_per_token": 8.4e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81232,6 +81262,7 @@ "output_cost_per_token": 8.4e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index f2a3bb1a47e..02bf7dedcef 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -64039,6 +64039,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64058,6 +64059,7 @@ "output_cost_per_token_priority": 2.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64100,6 +64102,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64137,6 +64140,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64156,6 +64160,7 @@ "output_cost_per_token_priority": 2.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64230,6 +64235,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64250,6 +64256,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64273,6 +64280,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64348,6 +64356,7 @@ "output_cost_per_token_priority": 3e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64384,6 +64393,7 @@ "output_cost_per_token_priority": 6.375e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_vision": true }, @@ -64420,6 +64430,7 @@ "output_cost_per_token_priority": 3e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64490,6 +64501,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64513,6 +64525,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67306,6 +67319,7 @@ "output_cost_per_token_priority": 5.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67324,6 +67338,7 @@ "output_cost_per_token_priority": 8.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67342,6 +67357,7 @@ "output_cost_per_token_priority": 5.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67360,6 +67376,7 @@ "output_cost_per_token_priority": 8.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67375,6 +67392,7 @@ "output_cost_per_token": 6.6e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67417,6 +67435,7 @@ "output_cost_per_token": 6.6e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67435,6 +67454,7 @@ "output_cost_per_token_priority": 6.25e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67453,6 +67473,7 @@ "output_cost_per_token_priority": 9.375e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67470,6 +67491,7 @@ "output_cost_per_token_priority": 6.25e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67488,6 +67510,7 @@ "output_cost_per_token_priority": 9.375e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67506,6 +67529,7 @@ "output_cost_per_token_priority": 4.05e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -73747,6 +73771,7 @@ "output_cost_per_token_priority": 1.875e-05, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -77249,6 +77274,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -77407,6 +77433,7 @@ "output_cost_per_token_priority": 9e-07, "source": "https://platform.minimax.io/docs/guides/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -81184,6 +81211,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81200,6 +81228,7 @@ "output_cost_per_token": 3.08e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81216,6 +81245,7 @@ "output_cost_per_token": 8.4e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81232,6 +81262,7 @@ "output_cost_per_token": 8.4e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, diff --git a/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py index 7fc6ac2c436..6538ccf1cac 100644 --- a/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py +++ b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py @@ -172,10 +172,10 @@ def test_responses_call_forwards_previous_response_id_and_store() -> None: def test_responses_call_folds_developer_items_into_instructions() -> None: - client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b")) with patch(HTTPX_CLIENT_FACTORY, return_value=client): litellm.responses( - model="fireworks_ai/accounts/fireworks/models/kimi-k3", + model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b", input=[ {"role": "user", "content": "Hi there"}, {"role": "developer", "content": "Answer with exactly one word."}, @@ -214,6 +214,68 @@ def test_responses_call_keeps_developer_items_in_place_for_a_deployment_that_acc ) +@respx.mock +@pytest.mark.parametrize( + "model", + [ + "fireworks_ai/accounts/fireworks/models/kimi-k3", + "fireworks_ai/accounts/fireworks/models/glm-5p3", + "fireworks_ai/accounts/fireworks/routers/kimi-k3-fast", + "fireworks_ai/glm-5p3-us", + "fireworks_ai/models/kimi-k3", + "fireworks_ai/routers/kimi-k3-fast", + ], +) +def test_responses_call_keeps_developer_items_in_place_for_a_model_the_cost_map_declares_accepts_them( + model: str, +) -> None: + route: Final = _fireworks_responses_route(model.removeprefix("fireworks_ai/")) + litellm.responses( + model=model, + instructions="You are Codex", + input=[ + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ], + api_key="fw-test-key", + ) + body: Final = _routed_body(route) + assert body["instructions"] == "You are Codex" + assert tuple(body["input"]) == ( + {"role": "user", "content": "Hi there"}, + {"role": "system", "content": "Answer with exactly one word.", "type": "message"}, + {"role": "user", "content": "What is the capital of France?"}, + ) + + +@respx.mock +def test_responses_call_folds_developer_items_when_the_loaded_cost_map_turns_the_flag_off( + declare_mid_conversation_system: Callable[[str, bool], None], +) -> None: + declare_mid_conversation_system("fireworks_ai/accounts/fireworks/models/kimi-k3", False) + route: Final = _fireworks_responses_route("accounts/fireworks/models/kimi-k3") + litellm.responses( + model="fireworks_ai/accounts/fireworks/models/kimi-k3", + input=[ + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + {"type": "message", "role": "assistant", "content": "Paris"}, + {"role": "developer", "content": "Reply with PONG."}, + ], + api_key="fw-test-key", + ) + body: Final = _routed_body(route) + assert body["instructions"] == "Answer with exactly one word." + assert tuple(body["input"]) == ( + {"role": "user", "content": "Hi there"}, + {"role": "user", "content": "What is the capital of France?"}, + {"type": "message", "role": "assistant", "content": "Paris"}, + {"role": "user", "content": "Reply with PONG.", "type": "message"}, + ) + + @respx.mock @pytest.mark.parametrize( ("model_info", "closing_role"), @@ -222,9 +284,9 @@ def test_responses_call_keeps_developer_items_in_place_for_a_deployment_that_acc def test_responses_call_sends_the_closing_developer_item_as_a_user_turn_only_for_a_declared_system_first_deployment( model_info: Mapping[str, object] | None, closing_role: str ) -> None: - route: Final = _fireworks_responses_route("accounts/fireworks/models/kimi-k3") + route: Final = _fireworks_responses_route("accounts/fireworks/models/qwen3p8-2p4t-a95b") litellm.responses( - model="fireworks_ai/accounts/fireworks/models/kimi-k3", + model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b", input=[ {"role": "user", "content": "Hi there"}, {"role": "developer", "content": "Answer with exactly one word."}, @@ -402,7 +464,7 @@ def test_transform_request_forwards_non_string_instructions_and_input_untouched( developer_item: Final = {"role": "developer", "content": "Answer with exactly one word."} user_item: Final = {"role": "user", "content": "What is the capital of France?"} request: Final = FireworksAIResponsesAPIConfig().transform_responses_api_request( - model="accounts/fireworks/models/kimi-k3", + model="accounts/fireworks/models/qwen3p8-2p4t-a95b", input=cast(ResponseInputParam, [developer_item, user_item]), response_api_optional_request_params={"instructions": ["not", "a", "string"]}, litellm_params=GenericLiteLLMParams(), @@ -416,7 +478,7 @@ def test_transform_request_forwards_non_string_instructions_and_input_untouched( def test_responses_call_maps_pydantic_developer_items_and_replays_pydantic_output_items() -> None: - client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b")) pydantic_input: Final = cast( ResponseInputParam, [ @@ -435,7 +497,9 @@ def test_responses_call_maps_pydantic_developer_items_and_replays_pydantic_outpu ) with patch(HTTPX_CLIENT_FACTORY, return_value=client): litellm.responses( - model="fireworks_ai/accounts/fireworks/models/kimi-k3", input=pydantic_input, api_key="fw-test-key" + model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b", + input=pydantic_input, + api_key="fw-test-key", ) _, _, body = _sent_request(client) assert body["instructions"] == "Answer with exactly one word."