From e0d903b971fe36d57d12f5185b153c44b063afff Mon Sep 17 00:00:00 2001 From: Avani Prajapati <143805019+Avani-prajapati@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:06:19 +0530 Subject: [PATCH 1/3] fix(proxy)!: drop router_settings_override before routing and apply it on every branch (#47257) * fix(proxy): drop router_settings_override before routing and apply it on every branch route_request only popped router_settings_override inside the branch that runs after the api_key/api_base, comma-separated batch and user_config branches. Requests taking one of those earlier branches forwarded the field to the provider, which Anthropic rejects with a 400, and silently lost the key/team num_retries, fallbacks and timeout. The override is now popped and merged into the request before any branch is chosen, keeping the existing precedence where request values win and None values are skipped Fixes #43582 * fix(proxy): apply key/team fallbacks when the request carries api_key or api_base The previous commit skipped fallbacks, context_window_fallbacks and content_policy_fallbacks from router_settings_override whenever the request had api_key or api_base. That broke the scenario from #43582, where a trusted pre-call hook sets api_key and the key/team router_settings still need to apply. All per-request router settings are merged again on every branch, with None values skipped and request values still winning The tests now check that fallbacks run on the api_key and api_base paths and that the provider never receives router_settings_override * fix(proxy): answer 400 for a non-object router_settings_override and treat null as absent * fix(proxy): drop a malformed router_settings_override under drop_params * fix(proxy): route key and team router settings like any request, honor user_config and global drop_params * fix(proxy): keep Router-only override settings out of direct provider calls Calls that route_request sends straight to litellm (evals, endpoints that need no model, pass_through_all_models, the video fallback and the no-router paths) now leave out context_window_fallbacks, content_policy_fallbacks and model_group_retry_policy merged from key or team router settings. Only the Router reads them, and providers reject them as unknown parameters. * test(proxy): drop a comment from the router settings override test * fix(proxy): keep key router settings on the model read-through retry route_request now merges router_settings_override once, before the first attempt. The retry after a model read-through used to pop the merged enable_tag_filtering again and recompute an empty Router-only set, so a model created on another instance was served without the key's tag filtering * test(proxy): check the read-through retry keeps key tag filtering through a real router The regression test now routes through a real Router with two tagged deployments from the database instead of replacing Router.acompletion. Before the fix the retry picks the lower-order untagged-for-request deployment, after it the tagged one * fix(proxy): keep key Router-format fallbacks out of direct provider calls Key and team fallbacks use the Router's model-group mapping. litellm.acompletion reads each mapping as a fallback config, so a direct call retried the same model with the mapping merged into the provider body. The override's fallbacks now stay out of direct calls like the other Router-only settings, while num_retries and timeout still pass --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/route_llm_request.py | 124 ++++-- tests/unit/proxy/test_route_llm_request.py | 460 ++++++++++++++++++++- 2 files changed, 541 insertions(+), 43 deletions(-) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 8bab7215b17..e16e14c8d54 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -2,13 +2,14 @@ import asyncio from collections.abc import Mapping from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Annotated, Any, Final, Literal import httpx from fastapi import HTTPException, status -from pydantic import TypeAdapter, ValidationError +from pydantic import BeforeValidator, TypeAdapter, ValidationError import litellm +from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.router_utils.common_utils import is_proxy_admin_request @@ -38,6 +39,21 @@ else: LitellmRouter = Any +_PER_REQUEST_ROUTER_SETTINGS: Final = ( + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "num_retries", + "timeout", + "model_group_retry_policy", + "routing_strategy", + "enable_tag_filtering", +) +_ROUTER_ONLY_SETTINGS: Final = frozenset( + ("fallbacks", "context_window_fallbacks", "content_policy_fallbacks", "model_group_retry_policy") +) + + def _route_user_config_request(data: dict, route_type: str): """Route a request using the user-provided router config.""" router_config: Final = data.pop("user_config") @@ -204,6 +220,18 @@ REQUIRED_ONE_OF_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, str]]] = Map ) JSON_OBJECT_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) +_OPTIONAL_JSON_OBJECT_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter( + Mapping[str, object] | None +) + + +def _none_unless_json_object(value: object) -> object: + return value if isinstance(value, Mapping) else None + + +_JSON_OBJECT_OR_DROPPED_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter( + Annotated[Mapping[str, object] | None, BeforeValidator(_none_unless_json_object)] +) class ProxyMissingRequiredParamError(ProxyException): @@ -596,6 +624,43 @@ RouteType = Literal[ ] +def _merge_router_settings_override( + data: dict[str, object], # mutable-ok: route_request mutates the shared request body in place +) -> frozenset[str]: + data.pop("enable_tag_filtering", None) + drop_params: Final = litellm.drop_params is True or normalize_drop_params(data.get("drop_params")) is True + router_settings_override_adapter: Final = ( + _JSON_OBJECT_OR_DROPPED_ADAPTER if drop_params else _OPTIONAL_JSON_OBJECT_ADAPTER + ) + try: + router_settings_override: Final = router_settings_override_adapter.validate_python( + data.pop("router_settings_override", None) + ) + except ValidationError: + raise ProxyException( + message="router_settings_override must be a JSON object.", + type="invalid_request_error", + param="router_settings_override", + code=status.HTTP_400_BAD_REQUEST, + ) from None + if "user_config" in data and data["user_config"] is None: + data.pop("user_config") + user_config: Final = _JSON_OBJECT_OR_DROPPED_ADAPTER.validate_python(data.get("user_config")) + router_settings_from_override: Final = MappingProxyType( + { + key: router_settings_override[key] + for key in _PER_REQUEST_ROUTER_SETTINGS + if router_settings_override.get(key) is not None + and key not in data + and (user_config is None or key not in user_config) + } + if router_settings_override is not None + else {} + ) + data.update(router_settings_from_override) + return _ROUTER_ONLY_SETTINGS.intersection(router_settings_from_override) + + async def route_request( data: dict, llm_router: LitellmRouter | None, @@ -606,12 +671,14 @@ async def route_request( """ Common helper to route the request """ + router_only_settings_from_override: Final = _merge_router_settings_override(data) # pyright: ignore[reportUnknownArgumentType] # route_request takes an untyped dict body try: return await _route_request_single_attempt( data=data, llm_router=llm_router, user_model=user_model, route_type=route_type, + router_only_settings_from_override=router_only_settings_from_override, user_api_key_dict=user_api_key_dict, ) except (ProxyModelNotFoundError, ProxyMissingParamWithoutLoadedModelError) as e: @@ -634,6 +701,7 @@ async def route_request( llm_router=proxy_server.llm_router, user_model=user_model, route_type=route_type, + router_only_settings_from_override=router_only_settings_from_override, user_api_key_dict=user_api_key_dict, ) @@ -643,6 +711,7 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr llm_router: LitellmRouter | None, user_model: str | None, route_type: RouteType, + router_only_settings_from_override: frozenset[str], user_api_key_dict: UserAPIKeyAuth | None = None, ): raise_if_required_body_param_missing(route_type=route_type, data=data, llm_router=llm_router) @@ -651,9 +720,14 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr raise_if_mock_testing_params_disallowed(data, allowed=mock_testing_params_allowed()) - data.pop("enable_tag_filtering", None) - if "user_config" in data and data["user_config"] is None: - data.pop("user_config") # pyright: ignore[reportUnknownMemberType] # route_request takes an untyped dict body + def direct_litellm_kwargs() -> Mapping[str, object]: + return MappingProxyType( + { + key: value + for key, value in JSON_OBJECT_ADAPTER.validate_python(data).items() + if key not in router_only_settings_from_override + } + ) team_id: Final = get_team_id_from_data(data) router_model_names: Final = llm_router.model_names if llm_router is not None else [] @@ -668,7 +742,7 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr if llm_router is not None: return getattr(llm_router, f"{route_type}")(**data) else: - return getattr(litellm, f"{route_type}")(**data) + return getattr(litellm, f"{route_type}")(**direct_litellm_kwargs()) elif ( route_type == "acompletion" @@ -687,34 +761,6 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr elif "user_config" in data: return _route_user_config_request(data, route_type) - elif "router_settings_override" in data: - # Apply per-request router settings overrides from key/team config - # Instead of creating a new Router (expensive), merge settings into kwargs - # The Router already supports per-request overrides for these settings - override_settings: Final = data.pop("router_settings_override") - - # Settings that the Router accepts as per-request kwargs - # These override the global router settings for this specific request - per_request_settings: Final = [ - "fallbacks", - "context_window_fallbacks", - "content_policy_fallbacks", - "num_retries", - "timeout", - "model_group_retry_policy", - "routing_strategy", - "enable_tag_filtering", - ] - - for key in per_request_settings: - if override_settings.get(key) is not None and key not in data: - data[key] = override_settings[key] - - # Use main router with overridden kwargs - if llm_router is not None: - return getattr(llm_router, f"{route_type}")(**data) - else: - return getattr(litellm, f"{route_type}")(**data) elif llm_router is not None: _raise_if_model_fully_blocked(llm_router=llm_router, model_name=data.get("model"), team_id=team_id) # Evals API: always route to litellm directly (not through router) @@ -755,7 +801,7 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr # If we can't get deployment creds, continue without them pass - return getattr(litellm, f"{route_type}")(**data) + return getattr(litellm, f"{route_type}")(**direct_litellm_kwargs()) # Skip model-based routing for container operations if route_type in [ "acreate_container", @@ -806,7 +852,7 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr "aingest", ] and (data.get("model") is None or data.get("model") == ""): # These endpoints don't need a model, use custom_llm_provider directly - return getattr(litellm, f"{route_type}")(**data) + return getattr(litellm, f"{route_type}")(**direct_litellm_kwargs()) if "model" not in data: raise ProxyMissingRequiredParamError( @@ -845,7 +891,7 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr # Check wildcards before checking deployment_names # Priority: 1. Exact model_name match, 2. Wildcard match, 3. deployment_names match if llm_router.router_general_settings.pass_through_all_models: - return getattr(litellm, f"{route_type}")(**data) + return getattr(litellm, f"{route_type}")(**direct_litellm_kwargs()) elif llm_router.default_deployment is not None or len(llm_router.pattern_router.patterns) > 0: return getattr(llm_router, f"{route_type}")(**data) elif data["model"] in llm_router.deployment_names: @@ -897,7 +943,7 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr return getattr(llm_router, f"{route_type}")(**data) except Exception: # If router fails (e.g., model not found in router), fall back to direct call - return getattr(litellm, f"{route_type}")(**data) + return getattr(litellm, f"{route_type}")(**direct_litellm_kwargs()) elif _is_a2a_agent_model(data.get("model", "")): from litellm.proxy.agent_endpoints.a2a_routing import ( route_a2a_agent_request, @@ -909,7 +955,7 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr # Fall through to raise exception below if result is None elif user_model is not None or route_type == "allm_passthrough_route": - return getattr(litellm, f"{route_type}")(**data) + return getattr(litellm, f"{route_type}")(**direct_litellm_kwargs()) # if no route found then it's a bad request route_name: Final = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type) diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 4e4e2f23950..8928b79bcd4 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1,4 +1,6 @@ +import contextlib import httpx +import json import litellm import pytest import respx @@ -10,7 +12,8 @@ from unittest.mock import MagicMock from fastapi import HTTPException -from litellm.proxy.route_llm_request import ProxyModelNotFoundError, route_request +from litellm.proxy.route_llm_request import JSON_OBJECT_ADAPTER, ProxyModelNotFoundError, route_request +from litellm.types.router import RouterGeneralSettings @pytest.mark.parametrize( @@ -390,7 +393,6 @@ async def test_route_request_with_router_settings_override(): "timeout": 30, "model_group_retry_policy": {"gpt-3.5-turbo": {"RateLimitErrorRetries": 3}}, "routing_strategy": "least-busy", - # This setting should be ignored (not in per_request_settings list) "model_group_alias": {"alias": "real_model"}, }, } @@ -437,12 +439,12 @@ async def test_route_request_with_router_settings_override_no_router(): litellm.acompletion = mock_completion try: - response = await route_request(data, None, None, "acompletion") + response = await route_request(data, None, "gpt-3.5-turbo", "acompletion") assert response == "success" # Verify litellm.acompletion was called with merged settings call_kwargs = mock_completion.call_args[1] - assert call_kwargs["fallbacks"] == [{"gpt-3.5-turbo": ["gpt-4"]}] + assert "fallbacks" not in call_kwargs assert call_kwargs["num_retries"] == 3 finally: litellm.acompletion = original_acompletion @@ -1039,6 +1041,408 @@ async def test_route_request_override_enable_tag_filtering_beats_body_value(): assert call_kwargs["enable_tag_filtering"] is True +_PROXY_DEPLOYMENT_BASE: Final = "https://proxy-deployment.example/v1" +_HELLO_MESSAGES: Final = ({"role": "user", "content": "Hello"},) + + +def _chat_completion_response(model: str) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + +def _proxy_router_with_two_model_groups() -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": f"openai/{model}", + "api_key": "sk-proxy", + "api_base": _PROXY_DEPLOYMENT_BASE, + }, + } + for model in ("gpt-3.5-turbo", "gpt-4") + ], + num_retries=0, + ) + + +@pytest.mark.parametrize( + "extra_body, upstream_base", + [ + pytest.param({"api_key": "sk-from-pre-call-hook"}, _PROXY_DEPLOYMENT_BASE, id="clientside-api-key"), + pytest.param( + {"api_base": "https://clientside.example/v1"}, "https://clientside.example/v1", id="clientside-api-base" + ), + pytest.param({"model": "gpt-3.5-turbo,gpt-4"}, _PROXY_DEPLOYMENT_BASE, id="comma-separated-batch"), + pytest.param({"drop_params": True}, _PROXY_DEPLOYMENT_BASE, id="drop-params"), + pytest.param( + { + "user_config": { + "model_list": [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "openai/gpt-3.5-turbo", + "api_key": "sk-user-config", + "api_base": "https://user-config.example/v1", + }, + } + ], + } + }, + "https://user-config.example/v1", + id="user-config", + ), + ], +) +@pytest.mark.asyncio +async def test_route_request_never_forwards_router_settings_override( + monkeypatch: pytest.MonkeyPatch, extra_body: dict[str, object], upstream_base: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + data: Final = { + "model": "gpt-3.5-turbo", + "messages": list(_HELLO_MESSAGES), + "router_settings_override": {"num_retries": 2, "timeout": None}, + **extra_body, + } + + with respx.mock(assert_all_called=False) as respx_mock: + primary: Final = respx_mock.post(f"{upstream_base}/chat/completions", json__model="gpt-3.5-turbo").mock( + side_effect=[ + httpx.Response(500, json={"error": {"message": "boom"}}), + _chat_completion_response("gpt-3.5-turbo"), + ] + ) + respx_mock.post(f"{upstream_base}/chat/completions", json__model="gpt-4").mock( + return_value=_chat_completion_response("gpt-4") + ) + response: Final = await (await route_request(data, _proxy_router_with_two_model_groups(), None, "acompletion")) + + primary_response: Final = response[0] if isinstance(response, list) else response + sent_bodies: Final = [json.loads(call.request.content) for call in primary.calls] + assert not [body for body in sent_bodies if "router_settings_override" in body], sent_bodies + assert primary_response._hidden_params["additional_headers"]["x-litellm-max-retries"] == 2, ( + primary_response._hidden_params + ) + + +@pytest.mark.parametrize("extra_body", [{}, {"api_key": "sk-from-pre-call-hook"}], ids=["router", "clientside-api-key"]) +@pytest.mark.asyncio +async def test_route_request_treats_null_router_settings_override_as_absent( + monkeypatch: pytest.MonkeyPatch, extra_body: dict[str, str] +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + data: Final = { + "model": "gpt-3.5-turbo", + "messages": list(_HELLO_MESSAGES), + "router_settings_override": None, + **extra_body, + } + + with respx.mock(assert_all_called=False) as respx_mock: + upstream: Final = respx_mock.post(f"{_PROXY_DEPLOYMENT_BASE}/chat/completions").mock( + return_value=_chat_completion_response("gpt-3.5-turbo") + ) + response: Final = await (await route_request(data, _proxy_router_with_two_model_groups(), None, "acompletion")) + + assert response.choices[0].message.content == "hi", response + assert "router_settings_override" not in json.loads(upstream.calls.last.request.content) + + +@pytest.mark.parametrize( + "global_drop_params, extra_body", + [ + pytest.param(False, {}, id="drop-params-off"), + pytest.param(False, {"drop_params": False}, id="request-drop-params-false"), + ], +) +@pytest.mark.parametrize("override", ["fallbacks", ["num_retries"], 3]) +@pytest.mark.asyncio +async def test_route_request_rejects_non_object_router_settings_override( + monkeypatch: pytest.MonkeyPatch, global_drop_params: bool, extra_body: dict[str, bool | str], override: object +) -> None: + from litellm.proxy._types import ProxyException + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "drop_params", global_drop_params) + data: Final = { + "model": "gpt-3.5-turbo", + "messages": list(_HELLO_MESSAGES), + "router_settings_override": override, + **extra_body, + } + + with respx.mock(assert_all_called=False) as respx_mock: + upstream: Final = respx_mock.post(f"{_PROXY_DEPLOYMENT_BASE}/chat/completions").mock( + return_value=_chat_completion_response("gpt-3.5-turbo") + ) + with pytest.raises(ProxyException) as exc_info: + await route_request(data, _proxy_router_with_two_model_groups(), None, "acompletion") + + assert exc_info.value.code == "400" + assert exc_info.value.param == "router_settings_override" + assert upstream.call_count == 0 + + +@pytest.mark.parametrize( + "global_drop_params, extra_body", + [ + pytest.param(False, {"drop_params": True}, id="request-drop-params"), + pytest.param(True, {}, id="global-drop-params"), + pytest.param(True, {"drop_params": False}, id="global-drop-params-with-request-false"), + pytest.param(False, {"drop_params": "true"}, id="request-drop-params-string"), + ], +) +@pytest.mark.parametrize("override", ["fallbacks", ["num_retries"], 3]) +@pytest.mark.asyncio +async def test_route_request_drops_non_object_router_settings_override_under_drop_params( + monkeypatch: pytest.MonkeyPatch, global_drop_params: bool, extra_body: dict[str, bool | str], override: object +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "drop_params", global_drop_params) + data: Final = { + "model": "gpt-3.5-turbo", + "messages": list(_HELLO_MESSAGES), + "router_settings_override": override, + **extra_body, + } + + with respx.mock(assert_all_called=False) as respx_mock: + upstream: Final = respx_mock.post(f"{_PROXY_DEPLOYMENT_BASE}/chat/completions").mock( + return_value=_chat_completion_response("gpt-3.5-turbo") + ) + response: Final = await (await route_request(data, _proxy_router_with_two_model_groups(), None, "acompletion")) + + assert response.choices[0].message.content == "hi", response + assert upstream.call_count == 1 + assert "router_settings_override" not in json.loads(upstream.calls.last.request.content) + + +@pytest.mark.parametrize( + "user_config_settings, expected_upstream_calls", + [ + pytest.param({}, 2, id="key-setting-fills-unset-user-config-setting"), + pytest.param({"num_retries": 0}, 1, id="user-config-setting-beats-key-setting"), + ], +) +@pytest.mark.asyncio +async def test_route_request_user_config_router_settings_beat_key_settings( + monkeypatch: pytest.MonkeyPatch, user_config_settings: dict[str, int], expected_upstream_calls: int +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + user_config_base: Final = "https://user-config.example/v1" + data: Final = { + "model": "gpt-3.5-turbo", + "messages": list(_HELLO_MESSAGES), + "router_settings_override": {"num_retries": 1}, + "user_config": { + **user_config_settings, + "model_list": [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "openai/gpt-3.5-turbo", + "api_key": "sk-user-config", + "api_base": user_config_base, + }, + } + ], + }, + } + + with respx.mock(assert_all_called=False) as respx_mock: + upstream: Final = respx_mock.post(f"{user_config_base}/chat/completions").mock( + side_effect=[ + httpx.Response(500, json={"error": {"message": "boom"}}), + _chat_completion_response("gpt-3.5-turbo"), + ] + ) + with contextlib.suppress(litellm.InternalServerError): + await (await route_request(data, _proxy_router_with_two_model_groups(), None, "acompletion")) + + assert upstream.call_count == expected_upstream_calls + + +@pytest.mark.parametrize( + "extra_body", + [ + pytest.param({}, id="no-key-settings"), + pytest.param({"router_settings_override": {"num_retries": 2}}, id="key-settings"), + ], +) +@pytest.mark.asyncio +async def test_route_request_refuses_blocked_model_whether_or_not_key_has_router_settings( + extra_body: dict[str, dict[str, int]], +) -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "sk-proxy"}, + "model_info": {"blocked": True}, + } + ], + ) + data: Final = {"model": "gpt-3.5-turbo", "messages": list(_HELLO_MESSAGES), **extra_body} + + with pytest.raises(litellm.PermissionDeniedError): + await route_request(data, router, None, "acompletion") + + +@pytest.mark.parametrize( + "credential", + [ + pytest.param({"api_key": "sk-from-pre-call-hook"}, id="clientside-api-key"), + pytest.param({"api_base": _PROXY_DEPLOYMENT_BASE}, id="clientside-api-base"), + ], +) +@pytest.mark.parametrize( + "fallback_field, upstream_error", + [ + pytest.param("fallbacks", httpx.Response(500, json={"error": {"message": "boom"}}), id="fallbacks"), + pytest.param( + "context_window_fallbacks", + httpx.Response( + 400, + json={ + "error": { + "message": "This model's maximum context length is 16385 tokens", + "code": "context_length_exceeded", + } + }, + ), + id="context-window-fallbacks", + ), + pytest.param( + "content_policy_fallbacks", + httpx.Response( + 400, + json={ + "error": { + "message": "Your request was rejected by the content filter", + "type": "invalid_request_error", + "code": "content_policy_violation", + } + }, + ), + id="content-policy-fallbacks", + ), + ], +) +@pytest.mark.asyncio +async def test_route_request_applies_override_fallbacks_when_request_carries_credentials( + monkeypatch: pytest.MonkeyPatch, + credential: dict[str, str], + fallback_field: str, + upstream_error: httpx.Response, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + data: Final = { + "model": "gpt-3.5-turbo", + "messages": list(_HELLO_MESSAGES), + "router_settings_override": {fallback_field: [{"gpt-3.5-turbo": ["gpt-4"]}]}, + **credential, + } + + with respx.mock(assert_all_called=False) as respx_mock: + respx_mock.post(f"{_PROXY_DEPLOYMENT_BASE}/chat/completions", json__model="gpt-3.5-turbo").mock( + return_value=upstream_error + ) + fallback: Final = respx_mock.post(f"{_PROXY_DEPLOYMENT_BASE}/chat/completions", json__model="gpt-4").mock( + return_value=_chat_completion_response("gpt-4") + ) + response: Final = await (await route_request(data, _proxy_router_with_two_model_groups(), None, "acompletion")) + + assert fallback.call_count == 1, [str(call.request.url) for call in respx_mock.calls] + assert "router_settings_override" not in json.loads(fallback.calls.last.request.content) + assert response.choices[0].message.content == "hi", response + + +@pytest.mark.parametrize( + "router_only_setting", + [ + pytest.param({"model_group_retry_policy": {"gpt-4": {"TimeoutErrorRetries": 1}}}, id="model-group-retry-policy"), + pytest.param({"context_window_fallbacks": [{"gpt-4": ["gpt-3.5-turbo"]}]}, id="context-window-fallbacks"), + pytest.param({"content_policy_fallbacks": [{"gpt-4": ["gpt-3.5-turbo"]}]}, id="content-policy-fallbacks"), + ], +) +@pytest.mark.asyncio +async def test_route_request_keeps_router_only_override_settings_out_of_pass_through_provider_body( + monkeypatch: pytest.MonkeyPatch, router_only_setting: dict[str, object] +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("OPENAI_API_KEY", "sk-pass-through") + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4", "api_key": "sk-proxy", "api_base": _PROXY_DEPLOYMENT_BASE}, + } + ], + router_general_settings=RouterGeneralSettings(pass_through_all_models=True), + ) + data: Final = { + "model": "openai/gpt-pass-through", + "messages": list(_HELLO_MESSAGES), + "router_settings_override": router_only_setting, + } + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=_chat_completion_response("gpt-pass-through") + ) + await (await route_request(data, router, None, "acompletion")) + + assert set(router_only_setting).isdisjoint(JSON_OBJECT_ADAPTER.validate_json(upstream.calls.last.request.content)) + + +@pytest.mark.asyncio +async def test_route_request_keeps_key_fallbacks_out_of_pass_through_provider_body( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("OPENAI_API_KEY", "sk-pass-through") + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4", "api_key": "sk-proxy", "api_base": _PROXY_DEPLOYMENT_BASE}, + } + ], + router_general_settings=RouterGeneralSettings(pass_through_all_models=True), + ) + data: Final = { + "model": "openai/gpt-pass-through", + "messages": list(_HELLO_MESSAGES), + "router_settings_override": {"fallbacks": [{"gpt-4": ["gpt-4"]}], "num_retries": 0}, + } + + with respx.mock(assert_all_called=True) as respx_mock: + upstream: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response(500, json={"error": {"message": "upstream down"}}) + ) + with pytest.raises(litellm.InternalServerError): + await (await route_request(data, router, None, "acompletion")) + + provider_bodies: Final = tuple(JSON_OBJECT_ADAPTER.validate_json(call.request.content) for call in upstream.calls) + assert provider_bodies and all("gpt-4" not in body for body in provider_bodies), provider_bodies + + @pytest.mark.parametrize( "route_type, param, route", [ @@ -1524,6 +1928,54 @@ async def test_route_request_read_through_recovers_model_created_on_sibling_repl assert table.find_many_wheres[0] == {"model_name": model_name} +def _tagged_db_model_row(model_name: str, tag: str, order: int) -> SimpleNamespace: + return SimpleNamespace( + model_id=f"{model_name}-{tag}-id", + model_name=model_name, + litellm_params={ + "model": "openai/gpt-4o", + "api_key": "fake", + "mock_response": f"{tag}-deployment", + "tags": [tag], + "order": order, + }, + model_info={}, + blocked=False, + ) + + +@pytest.mark.asyncio +async def test_route_request_read_through_retry_keeps_key_tag_filtering(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm.proxy.proxy_server as proxy_server + + model_name: Final = "sibling-replica-model-for-key-with-tag-filtering" + router: Final = litellm.Router( + model_list=[{"model_name": "some-other-model", "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}}] + ) + fake_prisma, _ = _fake_prisma_client_with_models( + [_tagged_db_model_row(model_name, "free", order=1), _tagged_db_model_row(model_name, "paid", order=2)] + ) + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", router) + + response: Final = await ( + await route_request( + { + "model": model_name, + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"tags": ["paid"]}, + "router_settings_override": {"enable_tag_filtering": True}, + }, + router, + None, + "acompletion", + ) + ) + + assert response.choices[0].message.content == "paid-deployment" + + @pytest.mark.asyncio async def test_route_request_unknown_model_raises_and_hits_db_once_within_ttl(monkeypatch): import litellm From 270d108aceb9850b6e924406c13a4bd6dbc7cd94 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 08:38:40 +0000 Subject: [PATCH 2/3] fix(anthropic): json_schema envelope edge cases on the Anthropic json tool path (#47478) * test(anthropic): cover json_schema envelope edge cases on the Anthropic json tool path * fix(anthropic): unwrap required-only keys, walk past tool-named envelopes, keep non-ASCII text, flush cut streams, check Vertex output_format * fix(anthropic): send held json tool text on cut Bedrock Invoke streams and drop a non-object output_format under drop_params A Bedrock Invoke stream that ends before the json tool block closes now sends the held text and closing brace, as the direct Anthropic stream already did. A non-object output_format on the Anthropic messages route and its providers is dropped under drop_params and otherwise forwarded, so the provider answers its own 400. Vertex Claude names the proxy setting in its 400 and reads drop_params through the shared helper. Adds integration cases for the json_schema envelope and the cut stream. * fix(anthropic): drop a non-object output_format on chat, unwrap a whole envelope on a cut stream, escape lone surrogates Chat completions to Anthropic, Azure AI, Bedrock Invoke and Vertex now drop a non-object output_format under drop_params, matching /v1/messages. A stream cut after a complete held envelope is unwrapped like a normal end. A json_schema answer holding a lone surrogate keeps its \u escape so the response can still be encoded * fix(anthropic): join a split answer when the stream ends in its last block A cut during the second json tool block now runs the held block through the same flush as a block stop, so it joins the streamed keys or the earlier split blocks instead of sending two objects back to back. The Vertex output_format fold drops its unreachable drop_params branch, since both parent configs drop a non-object value first. * test(vertex): expect the LiteLLM 400 for a non-object output_format on /v1/messages --- litellm/llms/anthropic/chat/handler.py | 34 ++++- litellm/llms/anthropic/chat/transformation.py | 15 +- .../pass_through/messages/transformation.py | 9 ++ litellm/llms/base_llm/base_utils.py | 50 +++--- litellm/llms/bedrock/chat/invoke_handler.py | 12 ++ .../anthropic/output_params_utils.py | 22 ++- ...vertex_anthropic_structured_output_wire.py | 29 ++-- .../response_format/__init__.py | 0 ...t_completions_response_format_anthropic.py | 123 +++++++++++++++ ...st_chat_completions_streaming_anthropic.py | 144 +++++++++++++++++- .../chat/test_anthropic_chat_handler.py | 110 +++++++++++++ .../test_anthropic_chat_transformation.py | 75 +++++++++ ...t_anthropic_messages_structured_outputs.py | 50 ++++++ ...azure_anthropic_messages_transformation.py | 45 ++++++ .../llms/bedrock/chat/test_invoke_handler.py | 81 ++++++++++ ...artner_models_anthropic_messages_config.py | 39 +++++ ...partner_models_anthropic_transformation.py | 38 +++++ 17 files changed, 829 insertions(+), 47 deletions(-) create mode 100644 tests/integration/translation/chat_completions/response_format/__init__.py create mode 100644 tests/integration/translation/chat_completions/response_format/test_chat_completions_response_format_anthropic.py diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 718536120b6..f0143f64be8 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -28,6 +28,7 @@ from litellm.llms.base_llm.base_utils import ( ResponseFormatSchema, names_a_response_format_property, partial_response_format_answer, + response_format_answer_json, response_format_envelope_contents, split_response_format_answer, ) @@ -104,17 +105,17 @@ def _key_after_a_string_verdict(string_value_start: str, schema: ResponseFormatS next_key: Final = _STRING_VALUE_THEN_KEY.match(string_value_start) if next_key is None: return "undecided" - return "pass_through" if next_key.group(1) in schema.property_names else "hold" + return "pass_through" if next_key.group(1) in schema.declared_names else "hold" def _is_a_label(text: str, schema: ResponseFormatSchema, tool_name: str) -> bool: - return text in schema.property_names or text == tool_name + return text in schema.declared_names or text == tool_name def _label_verdict(label: str, next_key: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict: if _is_a_label(label, schema, tool_name): return "hold" - return "pass_through" if next_key in schema.property_names else "undecided" + return "pass_through" if next_key in schema.declared_names else "undecided" def _leading_labels_verdict(string_value_start: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict: @@ -159,9 +160,9 @@ def _prefix_verdict(buffered_json: str, schema: ResponseFormatSchema, tool_name: return "hold" if first_value is None and _NAME_VALUE_PREFIX.match(buffered_json): return "undecided" - if first_key.group(1) == "values" and "values" not in schema.property_names: + if first_key.group(1) == "values" and "values" not in schema.declared_names: return "hold" - if first_key.group(1) in schema.property_names or not schema.property_names: + if first_key.group(1) in schema.declared_names or not schema.declared_names: return "pass_through" return _first_value_verdict(buffered_json[first_key.end() :], schema, tool_name) @@ -172,13 +173,13 @@ def _held_answer(held: str, schema: ResponseFormatSchema) -> str: except (json.JSONDecodeError, RecursionError): return held contents: Final = response_format_envelope_contents(payload, schema, RESPONSE_FORMAT_TOOL_NAME) - return held if contents is None else json.dumps(contents) + return held if contents is None else response_format_answer_json(contents) def _split_answer_text(blocks: tuple[str, ...], schema: ResponseFormatSchema) -> str: answer: Final = split_response_format_answer(blocks, schema) if answer is not None: - return json.dumps(answer) + return response_format_answer_json(answer) return "".join(_held_answer(block, schema) for block in blocks) @@ -1283,7 +1284,7 @@ class ModelResponseIterator: return None self._streamed_answer_keys = streamed | frozenset(fields) separator: Final = ", " if streamed else "" - return f"{separator}{json.dumps(fields)[1:-1]}" if fields else "" + return f"{separator}{response_format_answer_json(fields)[1:-1]}" if fields else "" def _flush_response_format_buffer(self) -> str: schema: Final = self.response_format_schema @@ -1314,6 +1315,17 @@ class ModelResponseIterator: return "" return _split_answer_text(blocks, self.response_format_schema) + def cut_stream_chunk(self) -> ModelResponseStream | None: + unsent: Final = self._flush_response_format_buffer() if self._response_format_buffer.getvalue() else "" + text: Final = unsent + self._release_answer_close() + self._drain_split_answer_blocks() + if not text: + return None + return ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=text))], + id=self.response_id, + model=self.served_model, + ) + def _handle_message_delta( self, chunk: dict ) -> tuple[str, Usage | None, dict[str, Any] | None, Mapping[str, object]]: @@ -1428,6 +1440,9 @@ class ModelResponseIterator: result = self._handle_accumulated_json_chunk(data_str="", is_final=True) if result is not None: return result + cut_stream_chunk = self.cut_stream_chunk() + if cut_stream_chunk is not None: + return cut_stream_chunk raise StopIteration except ValueError as e: raise RuntimeError(f"Error receiving chunk from stream: {e}") @@ -1474,6 +1489,9 @@ class ModelResponseIterator: result = self._handle_accumulated_json_chunk(data_str="", is_final=True) if result is not None: return result + cut_stream_chunk = self.cut_stream_chunk() + if cut_stream_chunk is not None: + return cut_stream_chunk raise StopAsyncIteration except ValueError as e: raise RuntimeError(f"Error receiving chunk from stream: {e}") diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 9624dbb582b..867dd935a03 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -23,7 +23,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, RESPONSE_FORMAT_TOOL_NAME, ) -from litellm.litellm_core_utils.core_helpers import map_finish_reason +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, map_finish_reason from litellm.litellm_core_utils.prompt_templates.common_utils import ( DEFS_MAX_INLINED_BYTES, is_object_schema, @@ -50,6 +50,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import from litellm.llms.base_llm.base_utils import ( ResponseFormatSchema, as_json_object, + response_format_answer_json, response_format_schema, split_response_format_answer, type_to_response_format_param, @@ -163,6 +164,7 @@ _ANTHROPIC_TOOL_NAME_MAX_LEN: Final = 128 _DEFER_LOADING_ADAPTER: Final = TypeAdapter(bool) _TOOL_FIELDS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _TOOLS_ADAPTER: Final = TypeAdapter(list[object]) +_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _NATIVE_OUTPUT_STRING_FORMATS: Final = frozenset( ("date-time", "time", "date", "duration", "email", "hostname", "uri", "ipv4", "ipv6", "uuid") ) @@ -2436,6 +2438,13 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): drop_params=litellm_params.get("drop_params") is True, output_key="top_k", ) + output_format: Final = request_params.get("output_format") + if ( + output_format is not None + and not isinstance(output_format, Mapping) + and drops_unsupported_params(_LITELLM_PARAMS_ADAPTER.validate_python(litellm_params).get("drop_params")) + ): + request_params.pop("output_format") data: Final = { "model": model, @@ -3166,7 +3175,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) ) if split_answer is not None: - return litellm.Message(content=json.dumps(split_answer)) + return litellm.Message(content=response_format_answer_json(split_answer)) try: payload: Final = _DECODED_JSON.validate_python(json.loads(json_mode_content_str)) except (json.JSONDecodeError, RecursionError): @@ -3174,7 +3183,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): return litellm.Message(content=json_mode_content_str) if response_format_schema is not None: return litellm.Message( - content=json.dumps( + content=response_format_answer_json( unwrap_response_format_envelope(payload, response_format_schema, RESPONSE_FORMAT_TOOL_NAME) ) ) diff --git a/litellm/llms/anthropic/pass_through/messages/transformation.py b/litellm/llms/anthropic/pass_through/messages/transformation.py index 0e8c1b49e8e..a71cd72b768 100644 --- a/litellm/llms/anthropic/pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/pass_through/messages/transformation.py @@ -7,6 +7,7 @@ import litellm from litellm._logging import verbose_logger from litellm.anthropic_beta_headers_manager import get_unsupported_headers, is_beta_header_supported_by_bundled_config from litellm.exceptions import AuthenticationError +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.pass_through.convertible_blocks import drop_non_string_tool_result_texts from litellm.llms.anthropic.pass_through.safeguards import proxy_serves_safeguards @@ -627,6 +628,14 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): ): anthropic_messages_optional_request_params.pop("safeguards") # pyright: ignore[reportUnknownMemberType] # the request params arrive as a bare dict + output_format: Final[object] = anthropic_messages_optional_request_params.get("output_format") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the request params arrive as a bare dict + if ( + output_format is not None + and not isinstance(output_format, Mapping) + and drops_unsupported_params(litellm_params.drop_params) + ): + anthropic_messages_optional_request_params.pop("output_format") # pyright: ignore[reportUnknownMemberType] # the request params arrive as a bare dict + self._translate_reasoning_effort_to_anthropic( model=model, optional_params=anthropic_messages_optional_request_params, diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index 5ed73e13188..640930a780d 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -5,6 +5,7 @@ Utility functions for base LLM classes. import copy import itertools import json +import re from abc import ABC, abstractmethod from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass @@ -172,7 +173,7 @@ _convert_tool_response_to_message = convert_tool_response_to_message @dataclass(frozen=True, slots=True) class ResponseFormatSchema: - property_names: frozenset[str] + declared_names: frozenset[str] required_names: frozenset[str] = frozenset() allows_undeclared_keys: bool = False @@ -216,9 +217,10 @@ def response_format_schema(value: object) -> ResponseFormatSchema | None: properties: Final = as_json_object(document.get("properties", {})) if properties is None: return None + required_names: Final = _names(document.get("required", ())) return ResponseFormatSchema( - property_names=frozenset(properties), - required_names=_names(document.get("required", ())), + declared_names=frozenset(properties) | required_names, + required_names=required_names, allows_undeclared_keys=_allows_undeclared_keys(document), ) @@ -244,8 +246,9 @@ def _json_object_or_encoded_object(value: JsonValue) -> Mapping[str, JsonValue] def _label_keys(payload_object: Mapping[str, JsonValue], schema: ResponseFormatSchema) -> frozenset[str]: if not schema.required_names: return frozenset() - declared: Final = schema.property_names | schema.required_names - return frozenset(key for key, value in payload_object.items() if isinstance(value, str) and key not in declared) + return frozenset( + key for key, value in payload_object.items() if isinstance(value, str) and key not in schema.declared_names + ) def _envelope_layer( @@ -256,7 +259,7 @@ def _envelope_layer( ) label_keys: Final = name_keys or _label_keys(payload_object, schema) wrapper_keys: Final = tuple(key for key in payload_object if key not in label_keys) - if len(wrapper_keys) != 1 or wrapper_keys[0] in schema.property_names: + if len(wrapper_keys) != 1 or wrapper_keys[0] in schema.declared_names: return None wrapped: Final = payload_object[wrapper_keys[0]] inner: Final = _json_object_or_encoded_object(wrapped) if name_keys else as_json_object(wrapped) @@ -275,15 +278,16 @@ def _envelope_layers( def _schema_object_inside_envelope( payload_object: Mapping[str, JsonValue], schema: ResponseFormatSchema, tool_name: str ) -> Mapping[str, JsonValue] | None: - layers: Final = _envelope_layers(payload_object, schema, tool_name) - candidates: Final = layers if schema.property_names else itertools.islice(layers, 1) - return next((layer.inner for layer in candidates if _holds_the_schema_object(layer, schema)), None) + walk: Final = _envelope_layers(payload_object, schema, tool_name) + layers: Final = tuple(walk if schema.declared_names else itertools.islice(walk, 1)) + schema_object: Final = next((layer.inner for layer in layers if _shares_a_schema_key(layer.inner, schema)), None) + if schema_object is not None: + return schema_object + return next((layer.inner for layer in layers if layer.names_the_tool), None) -def _holds_the_schema_object(layer: _EnvelopeLayer, schema: ResponseFormatSchema) -> bool: - if layer.names_the_tool: - return True - return not schema.accepts_a_lone_undeclared_key and not schema.property_names.isdisjoint(layer.inner) +def _shares_a_schema_key(candidate: Mapping[str, JsonValue], schema: ResponseFormatSchema) -> bool: + return not schema.accepts_a_lone_undeclared_key and not schema.declared_names.isdisjoint(candidate) def response_format_envelope_contents(payload: object, schema: ResponseFormatSchema, tool_name: str) -> object | None: @@ -293,7 +297,15 @@ def response_format_envelope_contents(payload: object, schema: ResponseFormatSch schema_object: Final = _schema_object_inside_envelope(payload_object, schema, tool_name) if schema_object is not None: return schema_object - return None if "values" in schema.property_names else payload_object.get("values") + return None if "values" in schema.declared_names else payload_object.get("values") + + +_LONE_SURROGATE: Final = re.compile("[\ud800-\udfff]") + + +def response_format_answer_json(value: object) -> str: + text: Final = json.dumps(value, ensure_ascii=False) + return json.dumps(value) if _LONE_SURROGATE.search(text) else text def unwrap_response_format_envelope(payload: object, schema: ResponseFormatSchema, tool_name: str) -> object: @@ -305,7 +317,7 @@ def _field_names( blocks: tuple[Mapping[str, JsonValue], ...], key: str, schema: ResponseFormatSchema ) -> tuple[str, ...] | None: values: Final = (block[key] for block in blocks) - names: Final = tuple(value for value in values if isinstance(value, str) and value in schema.property_names) + names: Final = tuple(value for value in values if isinstance(value, str) and value in schema.declared_names) return names if len(frozenset(names)) == len(blocks) else None @@ -324,7 +336,7 @@ def _split_answer_fields( blocks: tuple[Mapping[str, JsonValue], ...], schema: ResponseFormatSchema ) -> Mapping[str, JsonValue] | None: keys: Final = frozenset(blocks[0]) - if not keys.isdisjoint(schema.property_names) or any(frozenset(block) != keys for block in blocks): + if not keys.isdisjoint(schema.declared_names) or any(frozenset(block) != keys for block in blocks): return None named: Final = tuple((key, names) for key in keys if (names := _field_names(blocks, key, schema)) is not None) if len(named) != 1: @@ -340,7 +352,7 @@ def _merged_partial_answers( blocks: tuple[Mapping[str, JsonValue], ...], schema: ResponseFormatSchema ) -> Mapping[str, JsonValue] | None: keys: Final = frozenset(itertools.chain.from_iterable(blocks)) - if not keys <= schema.property_names or len(keys) != sum(len(block) for block in blocks): + if not keys <= schema.declared_names or len(keys) != sum(len(block) for block in blocks): return None return dict(itertools.chain.from_iterable(block.items() for block in blocks)) @@ -358,14 +370,14 @@ def split_response_format_answer( def partial_response_format_answer(encoded_block: str, schema: ResponseFormatSchema) -> Mapping[str, JsonValue] | None: block: Final = _json_object_or_encoded_object(encoded_block) - return block if block is not None and frozenset(block) <= schema.property_names else None + return block if block is not None and frozenset(block) <= schema.declared_names else None def names_a_response_format_property(encoded_block: str, schema: ResponseFormatSchema) -> bool: block: Final = _json_object_or_encoded_object(encoded_block) if block is None: return False - return any(isinstance(value, str) and value in schema.property_names for value in block.values()) + return any(isinstance(value, str) and value in schema.declared_names for value in block.values()) def _dict_to_response_format_helper(response_format: dict, ref_template: str | None = None) -> dict: diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 7553dacd35d..3d4821fcaa2 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -857,6 +857,9 @@ class AWSEventStreamDecoder: undecoded_stream_error: Final = tally.undecoded_stream_error(response_headers) if undecoded_stream_error is not None: raise undecoded_stream_error + stream_end_chunk: Final = self._stream_end_chunk() + if stream_end_chunk is not None: + yield stream_end_chunk async def aiter_bytes( self, iterator: AsyncIterator[bytes], *, response_headers: Mapping[str, str] | None = None @@ -878,6 +881,12 @@ class AWSEventStreamDecoder: undecoded_stream_error: Final = tally.undecoded_stream_error(response_headers) if undecoded_stream_error is not None: raise undecoded_stream_error + stream_end_chunk: Final = self._stream_end_chunk() + if stream_end_chunk is not None: + yield stream_end_chunk + + def _stream_end_chunk(self) -> ModelResponseStream | None: + return None def _response_stream_shape(self) -> "Shape": response_stream_shape: Final = get_bedrock_response_stream_shape() @@ -956,6 +965,9 @@ class AmazonAnthropicClaudeStreamDecoder(AWSEventStreamDecoder): def _chunk_parser(self, chunk_data: dict) -> ModelResponseStream: return self.anthropic_model_response_iterator.chunk_parser(chunk=chunk_data) + def _stream_end_chunk(self) -> ModelResponseStream | None: + return self.anthropic_model_response_iterator.cut_stream_chunk() + class AmazonDeepSeekR1StreamDecoder(AWSEventStreamDecoder): def __init__( diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/output_params_utils.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/output_params_utils.py index d5950d49d2b..db27427a5be 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/output_params_utils.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/output_params_utils.py @@ -45,11 +45,22 @@ def _as_output_config(value: object) -> Mapping[str, object] | None: def _fold_output_format_into_output_config( - output_config: Mapping[str, object] | None, output_format: object + output_config: Mapping[str, object] | None, output_format: object, model: str ) -> Mapping[str, object] | None: if output_format is None: return output_config - return {**(output_config or {}), "format": output_format} + output_format_object: Final = _as_output_config(output_format) + if output_format_object is not None: + return {**(output_config or {}), "format": output_format_object} + from litellm.exceptions import UnsupportedParamsError + + raise UnsupportedParamsError( + message=( + f"{model} on vertex_ai needs `output_format` to be a JSON object, got {type(output_format).__name__}. " + "To drop it, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n" + ), + status_code=400, + ) def sanitize_vertex_anthropic_output_params(data: dict, model: str) -> None: @@ -67,13 +78,16 @@ def sanitize_vertex_anthropic_output_params(data: dict, model: str) -> None: body has no empty dict. * ``output_format`` is moved into ``output_config.format``, because Vertex AI Claude rejects the top-level field as deprecated. When both - are set, ``output_format`` wins, as it does in the Anthropic adapter. + are set, a dict ``output_format`` wins, as it does in the Anthropic + adapter. A non-dict ``output_format`` raises a 400 + ``UnsupportedParamsError``, since the parent config already dropped it + under ``drop_params``. * Non-dict values for ``output_config`` are dropped to avoid sending malformed payloads downstream. """ request: Final[dict[str, object]] = data # mutable-ok: the caller's request body is sanitized in place output_config: Final = _fold_output_format_into_output_config( - _as_output_config(request.get("output_config")), request.pop("output_format", None) + _as_output_config(request.get("output_config")), request.pop("output_format", None), model ) if output_config is None: request.pop("output_config", None) diff --git a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py index cb295815ae2..2918f364c56 100644 --- a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py +++ b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py @@ -1819,25 +1819,32 @@ def test_vertex_chat_without_response_format_sends_no_tool_and_no_output_config( assert not {"tools", "output_config", "output_format"} & sent.body.keys(), sent.body -@pytest.mark.parametrize( - "output_format", - [5, [], {}, "x" * 5000], - ids=["int", "list", "empty", "5kb"], -) -def test_vertex_messages_malformed_output_format_reaches_the_caller_as_the_peer_400( +def test_vertex_messages_empty_output_format_reaches_the_caller_as_the_peer_400(gateway: Gateway) -> None: + marker, marker_hex = _marker("native") + with wire_server(_peer) as wire, gateway.scenario() as scenario: + model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) + response: Final = gateway.request("POST", "/v1/messages", _messages_body(model, marker, output_format={})) + assert response.status_code == 400, response.text + assert _BAD_FORMAT in response.text, response.text + sent: Final = _one_delivery(wire, marker_hex) + assert "output_format" not in sent.body, sent.body + assert object_value(sent.body["output_config"])["format"] == {}, sent.body + assert gateway.request("GET", "/health/liveliness").status_code == 200 + + +@pytest.mark.parametrize("output_format", [5, [], "x" * 5000], ids=["int", "list", "5kb"]) +def test_vertex_messages_non_object_output_format_is_refused_without_reaching_the_peer( gateway: Gateway, output_format: JsonValue ) -> None: - marker, marker_hex = _marker("native") + marker, _ = _marker("native") with wire_server(_peer) as wire, gateway.scenario() as scenario: model: Final = _vertex_deployment(gateway, scenario, wire.url, _NATIVE_MODEL) response: Final = gateway.request( "POST", "/v1/messages", _messages_body(model, marker, output_format=output_format) ) assert response.status_code == 400, response.text - assert _BAD_FORMAT in response.text, response.text - sent: Final = _one_delivery(wire, marker_hex) - assert "output_format" not in sent.body, sent.body - assert object_value(sent.body["output_config"])["format"] == output_format, sent.body + assert "needs `output_format` to be a JSON object" in response.text, response.text + assert _delivered(wire) == () assert gateway.request("GET", "/health/liveliness").status_code == 200 diff --git a/tests/integration/translation/chat_completions/response_format/__init__.py b/tests/integration/translation/chat_completions/response_format/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/integration/translation/chat_completions/response_format/test_chat_completions_response_format_anthropic.py b/tests/integration/translation/chat_completions/response_format/test_chat_completions_response_format_anthropic.py new file mode 100644 index 00000000000..07e4358ff97 --- /dev/null +++ b/tests/integration/translation/chat_completions/response_format/test_chat_completions_response_format_anthropic.py @@ -0,0 +1,123 @@ +from dataclasses import replace +from typing import Final +from unittest.mock import ANY + +import pytest +from integration._support.client import Gateway +from integration._support.provider import SharedProvider +from integration.translation.case import TranslationTestCase +from integration.translation.chat_completions.bases.anthropic import CLAUDE_HAIKU_5_5_TEST_CASE +from integration.translation.runner import assert_translation + +PLACE_SCHEMA: Final = { + "type": "object", + "properties": {"city": {"type": "string"}, "country": {"type": "string"}}, + "required": ["city", "country"], + "additionalProperties": True, +} +PLACE_PROMPT: Final = "Which city is the capital of France? Answer as JSON with city and country." +PLACE_TOOL_USE: Final = { + "type": "tool_use", + "id": "toolu_01Sr1iptyZnsy1eNxabMZejF", + "name": "json_tool_call", + "input": {"city": "Paris", "country": "France"}, + "caller": {"type": "direct"}, +} + +CLAUDE_HAIKU_5_5_JSON_SCHEMA_TEST_CASE: Final = replace( + CLAUDE_HAIKU_5_5_TEST_CASE, + scenario="json_schema", + litellm_request={ + "model": "anthropic/claude-haiku-5-5", + "max_tokens": 128, + "messages": [{"role": "user", "content": PLACE_PROMPT}], + "response_format": {"type": "json_schema", "json_schema": {"name": "place", "schema": PLACE_SCHEMA}}, + "cache": {"no-cache": True}, + }, + expected_provider_request={ + "model": "claude-haiku-5-5", + "messages": [{"role": "user", "content": [{"type": "text", "text": PLACE_PROMPT}]}], + "max_tokens": 128, + "tool_choice": {"name": "json_tool_call", "type": "tool"}, + "tools": [{"name": "json_tool_call", "input_schema": PLACE_SCHEMA}], + }, + mock_provider_response={ + "model": "claude-haiku-5-5", + "id": "msg_011CgCoP2UR3vcM5ZWmEBd8D", + "type": "message", + "role": "assistant", + "content": [PLACE_TOOL_USE], + "container": None, + "stop_reason": "tool_use", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 528, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 55, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + "speed": "standard", + }, + "diagnostics": None, + }, + expected_litellm_response={ + "id": ANY, + "created": ANY, + "object": "chat.completion", + "model": "anthropic/claude-haiku-5-5", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": {"content": '{"city": "Paris", "country": "France"}', "role": "assistant"}, + } + ], + "usage": { + "completion_tokens": 55, + "prompt_tokens": 528, + "total_tokens": 583, + "completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 55}, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 528, + "cache_write_tokens": 0, + "cache_creation_tokens": 0, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "inference_geo": "global", + "service_tier": "standard", + "speed": "standard", + }, + }, +) + +CLAUDE_HAIKU_5_5_JSON_SCHEMA_TOOL_NAMED_ENVELOPE_TEST_CASE: Final = replace( + CLAUDE_HAIKU_5_5_JSON_SCHEMA_TEST_CASE, + scenario="json_schema_tool_named_envelope", + mock_provider_response={ + **CLAUDE_HAIKU_5_5_JSON_SCHEMA_TEST_CASE.mock_provider_response, + "content": [ + { + **PLACE_TOOL_USE, + "input": {"name": "json_tool_call", "arguments": {"input": {"city": "Paris", "country": "France"}}}, + } + ], + }, +) + + +@pytest.mark.parametrize( + "case", + [CLAUDE_HAIKU_5_5_JSON_SCHEMA_TEST_CASE, CLAUDE_HAIKU_5_5_JSON_SCHEMA_TOOL_NAMED_ENVELOPE_TEST_CASE], + ids=lambda case: case.id, +) +def test_chat_completions_response_format_anthropic( + case: TranslationTestCase, gateway: Gateway, provider: SharedProvider +) -> None: + assert_translation(case, gateway, provider) diff --git a/tests/integration/translation/chat_completions/streaming/test_chat_completions_streaming_anthropic.py b/tests/integration/translation/chat_completions/streaming/test_chat_completions_streaming_anthropic.py index ccba51cff23..b2c0d39f9aa 100644 --- a/tests/integration/translation/chat_completions/streaming/test_chat_completions_streaming_anthropic.py +++ b/tests/integration/translation/chat_completions/streaming/test_chat_completions_streaming_anthropic.py @@ -5,8 +5,12 @@ import pytest from integration._support.client import Gateway from integration._support.provider import SharedProvider from integration.translation.case import SSEFrame, StreamingTranslationTestCase -from integration.translation.chat_completions.bases.anthropic import CLAUDE_SONNET_5_5_TEST_CASE +from integration.translation.chat_completions.bases.anthropic import ( + CLAUDE_HAIKU_5_5_TEST_CASE, + CLAUDE_SONNET_5_5_TEST_CASE, +) from integration.translation.runner import assert_streaming_translation +from pydantic import JsonValue CLAUDE_SONNET_5_5_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( scenario="streaming", @@ -128,7 +132,143 @@ CLAUDE_SONNET_5_5_STREAMING_TEST_CASE: Final = StreamingTranslationTestCase( ) -@pytest.mark.parametrize("case", [CLAUDE_SONNET_5_5_STREAMING_TEST_CASE], ids=lambda case: case.id) +PLACE_SCHEMA: Final = { + "type": "object", + "properties": {"city": {"type": "string"}, "country": {"type": "string"}}, + "required": ["city", "country"], + "additionalProperties": True, +} +PLACE_PROMPT: Final = "Which city is the capital of France? Answer as JSON with city and country." + + +def _place_delta(partial_json: str) -> SSEFrame: + return SSEFrame( + event="content_block_delta", + data={ + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": partial_json}, + }, + ) + + +def _place_chunk(delta: dict[str, JsonValue], *, finish_reason: str | None = None) -> SSEFrame: + finish: Final = {} if finish_reason is None else {"finish_reason": finish_reason} + return SSEFrame( + data={ + "id": ANY, + "object": "chat.completion.chunk", + "created": ANY, + "model": "anthropic/claude-haiku-5-5", + "choices": [{"index": 0, "delta": delta, **finish}], + } + ) + + +CLAUDE_HAIKU_5_5_JSON_SCHEMA_CUT_STREAM_TEST_CASE: Final = StreamingTranslationTestCase( + scenario="json_schema_cut_stream", + litellm_endpoint=CLAUDE_HAIKU_5_5_TEST_CASE.litellm_endpoint, + litellm_request={ + "model": "anthropic/claude-haiku-5-5", + "max_tokens": 128, + "messages": [{"role": "user", "content": PLACE_PROMPT}], + "response_format": {"type": "json_schema", "json_schema": {"name": "place", "schema": PLACE_SCHEMA}}, + "cache": {"no-cache": True}, + "stream": True, + }, + expected_provider_endpoint=CLAUDE_HAIKU_5_5_TEST_CASE.expected_provider_endpoint, + expected_provider_headers=CLAUDE_HAIKU_5_5_TEST_CASE.expected_provider_headers, + expected_provider_request={ + "model": "claude-haiku-5-5", + "messages": [{"role": "user", "content": [{"type": "text", "text": PLACE_PROMPT}]}], + "max_tokens": 128, + "tool_choice": {"name": "json_tool_call", "type": "tool"}, + "tools": [{"name": "json_tool_call", "input_schema": PLACE_SCHEMA}], + "stream": True, + }, + mock_provider_frames=( + SSEFrame( + event="message_start", + data={ + "type": "message_start", + "message": { + "model": "claude-haiku-5-5", + "id": "msg_011CgCoP5RkexwE6N7PRZEca", + "type": "message", + "role": "assistant", + "content": [], + "container": None, + "stop_reason": None, + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 528, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 16, + "service_tier": "standard", + "inference_geo": "global", + "speed": "standard", + }, + "diagnostics": None, + }, + }, + ), + SSEFrame( + event="content_block_start", + data={ + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "tool_use", + "id": "toolu_01B4QSMQtft8TA3Rn6XFziWe", + "name": "json_tool_call", + "input": {}, + "caller": {"type": "direct"}, + }, + }, + ), + SSEFrame(event="ping", data={"type": "ping"}), + _place_delta(""), + _place_delta('{"city"'), + _place_delta(': "Pari'), + _place_delta('s"'), + _place_delta(', "coun'), + _place_delta('try": "'), + _place_delta('France"}'), + ), + expected_litellm_frames=( + _place_chunk({"role": "assistant", "content": '{"city": "Pari'}), + _place_chunk({"content": 's"'}), + _place_chunk({"content": ', "coun'}), + _place_chunk({"content": 'try": "'}), + _place_chunk({"content": 'France"'}), + _place_chunk({"content": "}"}), + _place_chunk({}, finish_reason="stop"), + SSEFrame(data="[DONE]"), + ), + expected_litellm_response={ + "id": ANY, + "object": "chat.completion", + "created": ANY, + "model": "anthropic/claude-haiku-5-5", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": '{"city": "Paris", "country": "France"}'}, + "finish_reason": "stop", + } + ], + }, +) + + +@pytest.mark.parametrize( + "case", + [CLAUDE_SONNET_5_5_STREAMING_TEST_CASE, CLAUDE_HAIKU_5_5_JSON_SCHEMA_CUT_STREAM_TEST_CASE], + ids=lambda case: case.id, +) def test_chat_completions_streaming_anthropic( case: StreamingTranslationTestCase, gateway: Gateway, provider: SharedProvider ) -> None: diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py index a2d9f321122..1f4a5aaa256 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -3108,6 +3108,92 @@ def test_streamed_json_mode_answer_with_a_foreign_name_matches_the_non_streamed_ assert unwrap_response_format_envelope(answer, schema, RESPONSE_FORMAT_TOOL_NAME) == answer +def test_streamed_json_mode_tool_name_envelope_around_an_input_wrapper_is_unwrapped_to_the_answer(): + envelope = json.dumps({"name": RESPONSE_FORMAT_TOOL_NAME, "arguments": {"input": _STREAMED_RISK_OBJECT}}) + + contents = _streamed_json_mode_contents((envelope[:9], envelope[9:40], envelope[40:])) + + assert json.loads("".join(contents)) == _STREAMED_RISK_OBJECT + + +def test_streamed_json_mode_unwrapped_answer_keeps_non_ascii_text(): + answer = {"tier": 3, "summary": "Hôtel Lutetia, 東京"} + envelope = json.dumps({"input": answer}, ensure_ascii=False) + + contents = _streamed_json_mode_contents((envelope[:9], envelope[9:30], envelope[30:])) + + assert "Hôtel Lutetia, 東京" in "".join(contents) + assert json.loads("".join(contents)) == answer + + +def test_streamed_json_mode_unwrapped_answer_with_a_lone_surrogate_stays_utf8_encodable(): + answer = {"tier": 3, "summary": "smile \ud83d end"} + envelope = json.dumps({"input": answer}) + + contents = _streamed_json_mode_contents((envelope[:9], envelope[9:30], envelope[30:])) + + assert json.loads("".join(contents).encode("utf-8")) == answer + + +_CUT_STREAM_FRAGMENTS: Final = ( + pytest.param( + json.dumps({"name": RESPONSE_FORMAT_TOOL_NAME, "parameters": _STREAMED_RISK_OBJECT})[:-1], + id="held_envelope", + ), + pytest.param(json.dumps(_STREAMED_RISK_OBJECT), id="streamed_answer"), +) + + +def _cut_stream_lines(sent: str) -> tuple[bytes, ...]: + events: Final = _json_tool_events((sent[:9], sent[9:30], sent[30:])) + return tuple(f"data: {json.dumps(event)}\n".encode() for event in events[:-2]) + + +def _cut_stream_iterator(sync_stream: bool, streaming_response: object) -> ModelResponseIterator: + return ModelResponseIterator( + streaming_response=streaming_response, + sync_stream=sync_stream, + json_mode=True, + response_format_schema=response_format_schema(_STREAMED_RISK_SCHEMA), + ) + + +@pytest.mark.parametrize("sent", _CUT_STREAM_FRAGMENTS) +def test_streamed_json_mode_text_is_delivered_when_the_stream_ends_before_the_block_stops(sent: str): + iterator = _cut_stream_iterator(True, iter(_cut_stream_lines(sent))) + + assert "".join(chunk.choices[0].delta.content or "" for chunk in iterator) == sent + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sent", _CUT_STREAM_FRAGMENTS) +async def test_async_streamed_json_mode_text_is_delivered_when_the_stream_ends_before_the_block_stops(sent: str): + async def lines(): + for line in _cut_stream_lines(sent): + yield line + + iterator = _cut_stream_iterator(False, lines()) + + assert "".join([chunk.choices[0].delta.content or "" async for chunk in iterator]) == sent + + +def test_streamed_json_mode_complete_envelope_is_unwrapped_when_the_stream_ends_before_the_block_stops(): + envelope = json.dumps({"name": RESPONSE_FORMAT_TOOL_NAME, "parameters": _STREAMED_RISK_OBJECT}) + + iterator = _cut_stream_iterator(True, iter(_cut_stream_lines(envelope))) + + assert json.loads("".join(chunk.choices[0].delta.content or "" for chunk in iterator)) == _STREAMED_RISK_OBJECT + + +def test_streamed_json_mode_answer_keyed_by_a_required_only_name_is_kept_whole(): + required_only_schema = {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["payload"]} + encoded = json.dumps({"payload": {"id": "abc", "rows": 3}}) + + contents = _streamed_json_mode_contents((encoded[:12], encoded[12:24], encoded[24:]), required_only_schema) + + assert "".join(contents) == encoded + + def _json_tool_block_events(index: int, encoded: str) -> tuple[dict[str, object], ...]: return ( { @@ -3202,6 +3288,30 @@ def test_streamed_json_mode_answer_split_into_one_partial_object_per_field_arriv assert contents[3] == blocks[0][:-1] +@pytest.mark.parametrize( + "blocks", + [ + pytest.param((_STREAMED_TIER_FIELD, _STREAMED_SUMMARY_FIELD), id="one_block_per_field"), + pytest.param( + tuple(json.dumps({name: value}) for name, value in _STREAMED_RISK_OBJECT.items()), + id="one_partial_object_per_field", + ), + ], +) +def test_streamed_json_mode_split_answer_arrives_as_one_object_when_the_stream_ends_in_its_last_block( + blocks: tuple[str, ...], +): + message_start, *_ = _json_tool_events(()) + events = ( + message_start, + *chain.from_iterable(_json_tool_block_events(index, encoded) for index, encoded in enumerate(blocks)), + )[:-1] + + iterator = _cut_stream_iterator(True, iter(tuple(f"data: {json.dumps(event)}\n".encode() for event in events))) + + assert json.loads("".join(chunk.choices[0].delta.content or "" for chunk in iterator)) == _STREAMED_RISK_OBJECT + + def test_streamed_json_mode_empty_object_then_field_blocks_arrives_as_one_object(): blocks = ("{}", *(json.dumps({name: value}) for name, value in _STREAMED_RISK_OBJECT.items())) 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 ef8447712d6..32f18b952ac 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -2251,6 +2251,38 @@ def test_effort_with_other_features(): assert "thinking" in result +_OBJECT_OUTPUT_FORMAT: Final = {"type": "json_schema", "schema": {"type": "object", "properties": {}}} + + +@pytest.mark.parametrize( + ("output_format", "request_drop_params", "global_drop_params", "forwarded"), + [ + pytest.param("json", True, False, False, id="string_dropped_under_request_drop_params"), + pytest.param("json", None, True, False, id="string_dropped_under_global_drop_params"), + pytest.param("json", None, False, True, id="string_forwarded_without_drop_params"), + pytest.param(_OBJECT_OUTPUT_FORMAT, True, False, True, id="object_kept_under_drop_params"), + ], +) +def test_anthropic_chat_drops_only_a_non_object_output_format_and_only_under_drop_params( + output_format: object, + request_drop_params: bool | None, + global_drop_params: bool, + forwarded: bool, + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "drop_params", global_drop_params) + + result = AnthropicConfig().transform_request( + model="claude-haiku-4-5", + messages=[{"role": "user", "content": "Hello"}], + optional_params={"output_format": output_format}, + litellm_params={"drop_params": request_drop_params}, + headers={}, + ) + + assert result.get("output_format") == (output_format if forwarded else None) + + def test_anthropic_drop_params_strips_output_config_for_pre_4_5_models(): """``drop_params=True`` strips unsupported ``output_config`` for pre-4.5 models.""" config = AnthropicConfig() @@ -7882,6 +7914,24 @@ def _json_tool_call(arguments: object) -> list[ChatCompletionToolCallChunk]: {"summary": "likes tea"}, id="tool_name_envelope_on_a_schema_allowing_extra_keys", ), + pytest.param( + {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["payload"]}, + {"payload": {"id": "abc", "rows": 3}}, + {"payload": {"id": "abc", "rows": 3}}, + id="key_listed_only_in_required", + ), + pytest.param( + _RISK_SCHEMA, + {"name": RESPONSE_FORMAT_TOOL_NAME, "arguments": {"input": _RISK_OBJECT}}, + _RISK_OBJECT, + id="tool_name_envelope_around_an_input_wrapper", + ), + pytest.param( + _PER_USER_SCHEMA, + {"name": RESPONSE_FORMAT_TOOL_NAME, "parameters": _PER_USER_ANSWER}, + _PER_USER_ANSWER, + id="tool_name_envelope_around_an_answer_with_one_allowed_key", + ), ], ) def test_convert_tool_response_to_message_unwraps_only_an_unambiguous_schema_envelope( @@ -7896,6 +7946,31 @@ def test_convert_tool_response_to_message_unwraps_only_an_unambiguous_schema_env assert json.loads(message.content) == expected_content +def test_convert_tool_response_to_message_keeps_non_ascii_text_in_an_unwrapped_answer(): + answer = {"tier": 3, "summary": "Hôtel Lutetia, 東京"} + + message = AnthropicConfig.convert_tool_response_to_message( + tool_calls=_json_tool_call({"input": answer}), + response_format_schema=_schema(_RISK_SCHEMA), + ) + + assert message is not None + assert "Hôtel Lutetia, 東京" in message.content + assert json.loads(message.content) == answer + + +def test_convert_tool_response_to_message_answer_with_a_lone_surrogate_stays_utf8_encodable(): + answer = {"tier": 3, "summary": "smile \ud83d end"} + + message = AnthropicConfig.convert_tool_response_to_message( + tool_calls=_json_tool_call({"input": answer}), + response_format_schema=_schema(_RISK_SCHEMA), + ) + + assert message is not None + assert json.loads(message.content.encode("utf-8")) == answer + + _TOO_DEEP_TO_DECODE: Final = "[" * 100_000 + "]" * 100_000 diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_messages_structured_outputs.py b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_messages_structured_outputs.py index 670e3de5110..db3ff58d80a 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_messages_structured_outputs.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_messages_structured_outputs.py @@ -2,10 +2,16 @@ Tests for structured outputs support in Anthropic /v1/messages endpoint. """ +from collections.abc import Mapping +from typing import Final + import pytest + +import litellm from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) +from litellm.types.router import GenericLiteLLMParams def test_output_format_supported_and_transforms_correctly(): @@ -109,3 +115,47 @@ def test_output_format_works_with_bedrock_and_azure(): headers={}, ) assert "output_format" in azure_result + + +_OBJECT_OUTPUT_FORMAT: Final = {"type": "json_schema", "schema": {"type": "object", "properties": {}}} + + +def _request_with_output_format( + output_format: object, litellm_params: GenericLiteLLMParams +) -> Mapping[str, object]: + return AnthropicMessagesConfig().transform_anthropic_messages_request( + model="claude-haiku-5-5", + messages=[{"role": "user", "content": "What is 6 times 7?"}], + anthropic_messages_optional_request_params={"max_tokens": 1024, "output_format": output_format}, + litellm_params=litellm_params, + headers={}, + ) + + +@pytest.mark.parametrize("output_format", ["json", ["json_schema"], 42]) +def test_malformed_output_format_is_dropped_under_drop_params(output_format: object) -> None: + result: Final = _request_with_output_format(output_format, GenericLiteLLMParams(drop_params=True)) + + assert "output_format" not in result + + +def test_malformed_output_format_is_dropped_under_the_global_drop_params_flag( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", True) + + result: Final = _request_with_output_format("json", GenericLiteLLMParams()) + + assert "output_format" not in result + + +def test_object_output_format_is_kept_under_drop_params() -> None: + result: Final = _request_with_output_format(_OBJECT_OUTPUT_FORMAT, GenericLiteLLMParams(drop_params=True)) + + assert result["output_format"] == _OBJECT_OUTPUT_FORMAT + + +def test_malformed_output_format_is_forwarded_for_the_provider_to_reject_without_drop_params() -> None: + result: Final = _request_with_output_format("json", GenericLiteLLMParams()) + + assert result["output_format"] == "json" diff --git a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py index 176dcfc86c9..38483b8a1fa 100644 --- a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py +++ b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py @@ -7,6 +7,7 @@ sys.path.insert( 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) ) +from collections.abc import Mapping from typing import Final from unittest.mock import patch @@ -551,3 +552,47 @@ async def test_outgoing_azure_request_sends_safeguards_with_the_dangerous_tool_u sent: Final = received[0] assert "dangerous-tool-use-2026-09-03" in sent.headers["anthropic-beta"].split(","), dict(sent.headers) assert json.loads(sent.content)["safeguards"] == list(AUTO_MODE_SAFEGUARDS) + + +async def _azure_request_body_for_output_format(output_format: object, drop_params: bool) -> Mapping[str, object]: + received: Final[list[httpx.Request]] = [] + + def foundry_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return _foundry_answer(request) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(foundry_answers)) + try: + await litellm.anthropic.messages.acreate( + max_tokens=256, + messages=[{"role": "user", "content": "What is 6 times 7?"}], + model="azure_ai/claude-sonnet-4-6", + api_key="test-api-key", + api_base="https://test-resource.services.ai.azure.com", + output_format=output_format, + drop_params=drop_params, + client=upstream, + ) + finally: + await upstream.close() + + assert len(received) == 1, received + return json.loads(received[0].content) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +async def test_outgoing_azure_request_drops_a_malformed_output_format_under_drop_params() -> None: + body: Final = await _azure_request_body_for_output_format("json", drop_params=True) + + assert "output_format" not in body + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +async def test_outgoing_azure_request_keeps_an_object_output_format_under_drop_params() -> None: + output_format: Final = {"type": "json_schema", "schema": {"type": "object", "properties": {}}} + + body: Final = await _azure_request_body_for_output_format(output_format, drop_params=True) + + assert body["output_format"] == output_format diff --git a/tests/unit/llms/bedrock/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py index 069438e25c0..f3ceac1d93c 100644 --- a/tests/unit/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/unit/llms/bedrock/chat/test_invoke_handler.py @@ -1880,3 +1880,84 @@ async def test_converse_stream_thinking_merges_into_content_with_think_tags() -> "Hi I am Anthropic, I am a helpful assistant" ) assert [getattr(chunk.choices[0].delta, "reasoning_content", None) for chunk in chunks] == [None] * len(chunks) + + +_PLACE_JSON_MODE_REQUEST: Final = json.dumps( + { + "tools": [ + { + "name": "json_tool_call", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string"}, "country": {"type": "string"}}, + "required": ["city", "country"], + "additionalProperties": True, + }, + } + ] + } +) + +_CUT_JSON_TOOL_STREAM: Final = b"".join( + _bedrock_event_stream_frame(event) + for event in ( + { + "type": "message_start", + "message": { + "id": "msg_cut_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 528, "output_tokens": 1}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "tool_use", "id": "toolu_cut_1", "name": "json_tool_call", "input": {}}, + }, + {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"city"'}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ': "Paris", '}}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '"country": "France"}'}, + }, + ) +) + + +def _cut_json_tool_stream_decoder() -> AWSEventStreamDecoder: + return _stream_decoder( + "anthropic", + model="anthropic.claude-haiku-5-5", + json_mode=True, + sync_stream=True, + response_format_schema=response_format_schema_from_request(_PLACE_JSON_MODE_REQUEST, json_mode=True), + ) + + +def _streamed_content(chunks: Sequence[object]) -> str: + return "".join( + chunk.choices[0].delta.content or "" + for chunk in chunks + if isinstance(chunk, ModelResponseStream) and chunk.choices + ) + + +def test_anthropic_invoke_stream_cut_before_the_json_tool_block_ends_sends_the_held_close() -> None: + chunks: Final = list(_cut_json_tool_stream_decoder().iter_bytes(iter([_CUT_JSON_TOOL_STREAM]))) + + assert json.loads(_streamed_content(chunks)) == {"city": "Paris", "country": "France"} + + +@pytest.mark.asyncio +async def test_async_anthropic_invoke_stream_cut_before_the_json_tool_block_ends_sends_the_held_close() -> None: + async def _chunks() -> AsyncIterator[bytes]: + yield _CUT_JSON_TOOL_STREAM + + chunks: Final = [chunk async for chunk in _cut_json_tool_stream_decoder().aiter_bytes(_chunks())] + + assert json.loads(_streamed_content(chunks)) == {"city": "Paris", "country": "France"} diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index 113ddcfe2a9..ae9b9684f2d 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -1,6 +1,7 @@ import copy import json import os +from collections.abc import Mapping from typing import Final from unittest.mock import MagicMock, patch @@ -537,6 +538,44 @@ def test_messages_request_moves_output_format_into_output_config(): assert result["output_config"] == {"effort": "high", "format": output_format} +_CALLER_OUTPUT_CONFIG: Final = { + "effort": "high", + "format": {"type": "json_schema", "schema": {"type": "object", "properties": {"answer": {"type": "integer"}}}}, +} + + +def _messages_request_with_output_format(output_format: object, drop_params: bool) -> Mapping[str, object]: + return VertexAIPartnerModelsAnthropicMessagesConfig().transform_anthropic_messages_request( + model="claude-opus-4-6", + messages=[{"role": "user", "content": "What is 6 times 7?"}], + anthropic_messages_optional_request_params={ + "max_tokens": 1024, + "output_format": output_format, + "output_config": copy.deepcopy(_CALLER_OUTPUT_CONFIG), + }, + litellm_params=GenericLiteLLMParams(drop_params=drop_params), + headers={}, + ) + + +@pytest.mark.parametrize("output_format", ["json", ["json_schema"], 42]) +def test_messages_request_drops_a_malformed_output_format_under_drop_params(output_format: object): + result = _messages_request_with_output_format(output_format, drop_params=True) + + assert "output_format" not in result + assert result["output_config"] == _CALLER_OUTPUT_CONFIG + + +@pytest.mark.parametrize("output_format", ["json", ["json_schema"], 42]) +def test_messages_request_rejects_a_malformed_output_format_without_drop_params(output_format: object): + import litellm + + with pytest.raises(litellm.UnsupportedParamsError) as error: + _messages_request_with_output_format(output_format, drop_params=False) + + assert error.value.status_code == 400 + + def test_provider_config_manager_reuses_vertex_anthropic_messages_config_instance(): """ Regression test: repeated provider config lookups for the same Vertex Claude model 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 ace24c8a3f4..21bda6f0345 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 @@ -918,6 +918,44 @@ def test_sanitize_folds_output_format_into_output_config(model: str, data: dict, assert request["output_config"] == expected_output_config +def _vertex_chat_request_with_output_format( + output_format: object, litellm_params: Mapping[str, object] +) -> Mapping[str, object]: + return VertexAIAnthropicConfig().transform_request( + model="claude-opus-4-6", + messages=[{"role": "user", "content": "What is 6 times 7?"}], + optional_params={"max_tokens": 64, "output_format": output_format}, + litellm_params=dict(litellm_params), + headers={}, + ) + + +def test_vertex_chat_request_drops_a_malformed_output_format_under_per_request_drop_params(): + request: Final = _vertex_chat_request_with_output_format("json", {"drop_params": True}) + + assert "output_format" not in request + assert "output_config" not in request + + +def test_vertex_chat_request_drops_a_malformed_output_format_under_the_global_drop_params_flag( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "drop_params", True) + + request: Final = _vertex_chat_request_with_output_format("json", {}) + + assert "output_format" not in request + assert "output_config" not in request + + +def test_vertex_chat_request_rejects_a_malformed_output_format_without_drop_params_naming_the_proxy_setting(): + with pytest.raises(litellm.UnsupportedParamsError) as error: + _vertex_chat_request_with_output_format("json", {}) + + assert error.value.status_code == 400 + assert "litellm_settings:\n drop_params: true" in str(error.value) + + def test_vertex_ai_native_response_format_reaches_vertex_as_output_config_format(local_model_cost_map): config = VertexAIAnthropicConfig() schema: Final = {"type": "object", "properties": {"result": {"type": "string"}}, "required": ["result"]} From 66cd1a651647639f298af0fc9554ecc9d9cd212b Mon Sep 17 00:00:00 2001 From: Zuraiz Anjum Date: Tue, 20 Oct 2026 13:50:29 +0500 Subject: [PATCH 3/3] fix(streaming): restore the token estimate for streams that send no usage (#44230) * fix(streaming): restore the token estimate for streams that send no usage A stream without stream_options got a made-up Usage(0, 0) written into its final chunk's hidden params. Since #42323 the chunk builder treats that as a provider-reported zero, so these streams were logged and billed at 0 tokens and $0 instead of the tokenizer estimate. Only attach the summed usage when some chunk actually carried usage. * test(streaming): use a fixed start time and check the estimated cost * fix(streaming): keep reported usage on the synthesized finish chunk and test it * test(streaming): cover explicit zero usage reported by the provider * fix(streaming): set the token estimate on a copy of the final chunk and ignore usage with no counts * fix(streaming): keep final chunk edits in the logged response when the provider sends no usage * fix(streaming): mark the unreported usage placeholder so every chunk rebuild falls back to the estimate * test(streaming): wait on the success callback event instead of sleep-polling --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../streaming_chunk_builder_utils.py | 3 + .../litellm_core_utils/streaming_handler.py | 142 +++++------ litellm/types/utils.py | 4 + .../test_streaming_handler.py | 235 ++++++++++++++++++ .../proxy/test_common_request_processing.py | 60 ++++- 5 files changed, 372 insertions(+), 72 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index b42210e2793..39f675099d7 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -46,6 +46,7 @@ from litellm.types.utils import ( StreamingChoices, TextChoices, TextCompletionResponse, + UnreportedUsage, Usage, ) from litellm.utils import print_verbose, token_counter @@ -876,6 +877,8 @@ class ChunkProcessor: if isinstance(usage_chunk, dict): return Usage(**usage_chunk) + if isinstance(usage_chunk, UnreportedUsage): + return None if usage_chunk is None or isinstance(usage_chunk, Usage): return usage_chunk return Usage(**usage_chunk.model_dump()) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index a627b26b9af..6ec426d6da5 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -24,7 +24,12 @@ import litellm from litellm import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.asyncify import asyncify -from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, get_hidden_params, stored_attribute +from litellm.litellm_core_utils.hidden_params import ( + HIDDEN_PARAMS_ATTR, + get_hidden_params, + set_hidden_param, + stored_attribute, +) from litellm.litellm_core_utils.model_response_utils import ( is_model_response_stream_empty, ) @@ -43,6 +48,8 @@ from litellm.types.utils import ( ModelResponseStream, PromptTokensDetailsWrapper, StreamingChoices, + TextCompletionResponse, + UnreportedUsage, Usage, ) from litellm.types.utils import GenericStreamingChunk as GChunk @@ -62,6 +69,7 @@ FUNCTION_CALL_ATTRIBUTE: Final = "function_call" _SYNC_ITER_EXHAUSTED: Final = object() _GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__) +_CompleteStreamingResponse = ModelResponse | TextCompletionResponse | ModelResponseStream | None _USAGE_COST_HEADER_PROVIDERS: Final[frozenset[str]] = frozenset({LlmProviders.OPENROUTER.value}) @@ -335,7 +343,8 @@ class CustomStreamWrapper: self.sent_stream_usage = False self.send_stream_usage = True if self.check_send_stream_usage(self.stream_options) else False self.tool_call = False - self.chunks: list = [] # keep track of the returned chunks - used for calculating the input/output tokens for stream options + # keep track of the returned chunks - used for calculating the input/output tokens for stream options + self.chunks: list[ModelResponseStream] = [] self._repeated_messages_count = 1 self.is_function_call = self.check_is_function_call(logging_obj=logging_obj) self.created: int | None = None @@ -544,7 +553,7 @@ class CustomStreamWrapper: last_content: Final = self.chunks[-1].choices[0].delta.content if ( - last_content is None or not isinstance(last_content, str) or len(last_content) <= 2 + not isinstance(last_content, str) or len(last_content) <= 2 ): # ignore empty content - https://github.com/BerriAI/litellm/issues/5158#issuecomment-2287156946 self._repeated_messages_count = 1 return @@ -1927,6 +1936,51 @@ class CustomStreamWrapper: response.hidden_params["additional_headers"] = {} response.hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = _cost + def _builder_failure_response(self, error: Exception) -> ModelResponseStream | None: + verbose_logger.warning( + "stream_chunk_builder raised at end-of-stream (%s); logging best-effort usage from chunks.", + str(error), + ) + try: + return self.model_response_creator(chunk={"usage": calculate_total_usage(chunks=self.chunks)}) + except Exception: + return None + + def _build_complete_streaming_response(self) -> _CompleteStreamingResponse: + try: + return litellm.stream_chunk_builder( + chunks=self.chunks, + messages=self.messages, + logging_obj=self.logging_obj, + count_prompt_tokens=self.count_prompt_tokens, + ) + except Exception as e: + return self._builder_failure_response(e) + + async def _async_build_complete_streaming_response(self) -> _CompleteStreamingResponse: + try: + return await asyncify(litellm.stream_chunk_builder)( + chunks=self.chunks, + messages=self.messages, + logging_obj=self.logging_obj, + count_prompt_tokens=self.count_prompt_tokens, + ) + except Exception as e: + return self._builder_failure_response(e) + + def _set_final_chunk_usage(self, chunk: ModelResponseStream) -> None: + reported_usage: Final = _reported_total_usage(chunks=self.chunks) + set_hidden_param( + chunk, + "usage", + UnreportedUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0) + if reported_usage is None + else reported_usage, + ) + self._last_returned_hidden_params = cast( # cast-ok: preserve dynamic mapping behavior + dict[str, object], getattr(chunk, HIDDEN_PARAMS_ATTR) + ) + def __next__(self) -> "ModelResponseStream": if self._closed: raise StopIteration @@ -2007,12 +2061,7 @@ class CustomStreamWrapper: continue # add usage as hidden param if self.sent_last_chunk is True and self.stream_options is None: - usage = calculate_total_usage(chunks=self.chunks) - response_hidden_params_for_final_chunk = cast( # cast-ok: preserve dynamic mapping behavior - dict[str, object], getattr(response, HIDDEN_PARAMS_ATTR) - ) - response_hidden_params_for_final_chunk["usage"] = usage - self._last_returned_hidden_params = response_hidden_params_for_final_chunk + self._set_final_chunk_usage(response) # Add MCP metadata to final chunk if present response = self._add_mcp_metadata_to_final_chunk(response) # RETURN RESULT @@ -2020,30 +2069,7 @@ class CustomStreamWrapper: except StopIteration: if self.sent_last_chunk is True: - try: - complete_streaming_response = litellm.stream_chunk_builder( - chunks=self.chunks, - messages=self.messages, - logging_obj=self.logging_obj, - count_prompt_tokens=self.count_prompt_tokens, - ) - except Exception as e: - # stream_chunk_builder can re-raise (as APIError) on large agentic - # streams. The raise originates inside this except-StopIteration block, - # so the sibling `except Exception` below does not catch it; it would - # escape __next__ and drop the request from SpendLogs. Recover - # best-effort usage from the raw chunks so cost is still tracked - verbose_logger.warning( - "stream_chunk_builder raised at end-of-stream (%s); logging best-effort usage from chunks.", - str(e), - ) - try: - complete_streaming_response = self.model_response_creator( - chunk={"usage": calculate_total_usage(chunks=self.chunks)} - ) - except Exception: - complete_streaming_response = None - + complete_streaming_response: Final = self._build_complete_streaming_response() response = self.model_response_creator() if complete_streaming_response is not None: self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) @@ -2107,9 +2133,8 @@ class CustomStreamWrapper: # last content chunk (stripped there), but never add the synthetic "stop" used when it sent none. if self.received_finish_reason is not None or self.intermittent_finish_reason is not None: self.chunks.append(processed_chunk) - if self.stream_options is None: # add usage as hidden param - usage = calculate_total_usage(chunks=self.chunks) - processed_chunk.hidden_params["usage"] = usage + if self.stream_options is None: + self._set_final_chunk_usage(processed_chunk) ## LOGGING executor.submit( self.run_success_logging_and_cache_storage, @@ -2224,12 +2249,7 @@ class CustomStreamWrapper: # add usage as hidden param if self.sent_last_chunk is True and self.stream_options is None: - usage = calculate_total_usage(chunks=self.chunks) - processed_chunk_hidden_params = cast( # cast-ok: preserve dynamic mapping behavior - dict[str, object], getattr(processed_chunk, HIDDEN_PARAMS_ATTR) - ) - processed_chunk_hidden_params["usage"] = usage - self._last_returned_hidden_params = processed_chunk_hidden_params + self._set_final_chunk_usage(processed_chunk) # Call post-call streaming deployment hook for final chunk if self.sent_last_chunk is True: @@ -2285,29 +2305,7 @@ class CustomStreamWrapper: async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream": if self.sent_last_chunk is True: - # log the final chunk with accurate streaming values - try: - complete_streaming_response = await asyncify(litellm.stream_chunk_builder)( - chunks=self.chunks, - messages=self.messages, - logging_obj=self.logging_obj, - count_prompt_tokens=self.count_prompt_tokens, - ) - except Exception as e: - # see sync __next__: a raise from stream_chunk_builder inside this - # except handler escapes __anext__ and drops the request from SpendLogs. - # Recover best-effort usage from the raw chunks so cost is still tracked - verbose_logger.warning( - "stream_chunk_builder raised at end-of-stream (%s); logging best-effort usage from chunks.", - str(e), - ) - try: - complete_streaming_response = self.model_response_creator( - chunk={"usage": calculate_total_usage(chunks=self.chunks)} - ) - except Exception: - complete_streaming_response = None - + complete_streaming_response: Final = await self._async_build_complete_streaming_response() response: Final = self.model_response_creator() if complete_streaming_response is not None: self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) @@ -2384,11 +2382,7 @@ class CustomStreamWrapper: if self.received_finish_reason is not None or self.intermittent_finish_reason is not None: self.chunks.append(processed_chunk) if self.stream_options is None: - usage: Final = calculate_total_usage(chunks=self.chunks) - processed_chunk_hidden_params: Final = cast( # cast-ok: preserve dynamic mapping behavior - dict[str, object], getattr(processed_chunk, HIDDEN_PARAMS_ATTR) - ) - processed_chunk_hidden_params["usage"] = usage + self._set_final_chunk_usage(processed_chunk) # see sync __next__'s sibling branch: deliberately do NOT restore # here - this chunk is still this call's own data, and restoring # before returning it would corrupt the caller's own log @@ -2659,7 +2653,13 @@ def _coerce_token_details( return details_type(**(raw if isinstance(raw, dict) else raw.model_dump())) -def calculate_total_usage(chunks: list[ModelResponse]) -> Usage: +def _reported_total_usage(chunks: Sequence[ModelResponse | ModelResponseStream]) -> Usage | None: + if not any("usage" in chunk and chunk["usage"] is not None for chunk in chunks): + return None + return calculate_total_usage(chunks=chunks) + + +def calculate_total_usage(chunks: Sequence[ModelResponse | ModelResponseStream]) -> Usage: """Assume most recent usage chunk has total usage uptil then.""" from litellm.litellm_core_utils.streaming_chunk_builder_utils import ( attach_cache_creation_token_details, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 576c4f093d3..2721a0c1cc1 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2062,6 +2062,10 @@ class Usage(SafeAttributeModel, CompletionUsage): setattr(self, key, value) +class UnreportedUsage(Usage): + pass + + class StreamingChoices(OpenAIObject): def __init__( self, diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 691442e80d0..68aabc3c01e 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -1,4 +1,5 @@ import asyncio +import datetime import json import sys import threading @@ -18,6 +19,7 @@ import litellm from litellm import verbose_logger from litellm._logging import session_id_var, trace_id_var from litellm.exceptions import MidStreamFallbackError +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import ( @@ -26,6 +28,7 @@ from litellm.litellm_core_utils.streaming_handler import ( _ProviderChunkParsed, ) from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response from litellm.types.utils import ( CompletionTokensDetailsWrapper, Delta, @@ -2820,6 +2823,238 @@ def test_usage_chunk_after_finish_reason_updates_hidden_params(logging_obj): ), f"Expected completion_tokens=135 from provider, got {hidden_usage.completion_tokens}" +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("ends_with_finish_chunk", [True, False]) +@pytest.mark.asyncio +async def test_stream_without_provider_usage_falls_back_to_the_token_estimate( + sync_mode: bool, ends_with_finish_chunk: bool +) -> None: + model: Final = "gpt-4o" + messages: Final = [{"role": "user", "content": "Write a long sentence about a fox. " * 30}] + text: Final = "The quick brown fox jumps over the lazy dog. " * 20 + content_chunk: Final = ModelResponseStream( + id="chatcmpl-no-usage", + created=1741037890, + model=model, + choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content=text))], + ) + finish_chunk: Final = ModelResponseStream( + id="chatcmpl-no-usage", + created=1741037890, + model=model, + choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")], + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator( + model_responses=[content_chunk, finish_chunk] if ends_with_finish_chunk else [content_chunk] + ), + model=model, + custom_llm_provider="openai", + logging_obj=Logging( + model=model, + messages=messages, + stream=True, + call_type="completion", + start_time=datetime.datetime(2025, 3, 3, 21, 38, 10), + litellm_call_id="no-usage-call", + function_id="no-usage-fn", + ), + stream_options=None, + ) + + yielded: Final = ( + [(chunk, chunk._hidden_params.get("usage")) for chunk in wrapper] + if sync_mode + else [(chunk, chunk._hidden_params.get("usage")) async for chunk in wrapper] + ) + collected: Final = [chunk for chunk, _ in yielded] + + expected: Final = ( + litellm.token_counter(model=model, messages=messages), + litellm.token_counter(model=model, text=text, count_response_tokens=True), + ) + assembled: Final = litellm.stream_chunk_builder(chunks=collected, messages=messages) + assert assembled is not None + assert (assembled.usage.prompt_tokens, assembled.usage.completion_tokens) == expected + last_chunk, usage_when_yielded = yielded[-1] + assert usage_when_yielded is not None + usage_after_drain: Final = last_chunk._hidden_params["usage"] + assert (usage_after_drain.prompt_tokens, usage_after_drain.completion_tokens) == expected + + rates: Final = litellm.model_cost[model] + spend: Final = expected[0] * rates["input_cost_per_token"] + expected[1] * rates["output_cost_per_token"] + assert litellm.completion_cost(completion_response=assembled) == pytest.approx(spend) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_stream_with_usage_after_the_finish_chunk_bills_only_the_reported_usage(sync_mode: bool) -> None: + model: Final = "gpt-4o" + messages: Final = [{"role": "user", "content": "Write a long sentence about a fox. " * 30}] + text: Final = "The quick brown fox jumps over the lazy dog. " * 20 + reported: Final = (6, 20) + chunks: Final = [ + ModelResponseStream(id="cmpl-late-usage", created=1741037890, model=model, choices=choices, usage=usage) + for choices, usage in ( + ([StreamingChoices(index=0, delta=Delta(role="assistant", content=text))], None), + ([StreamingChoices(index=0, delta=Delta(content=""), finish_reason="length")], None), + ([], Usage(prompt_tokens=reported[0], completion_tokens=reported[1], total_tokens=sum(reported))), + ) + ] + wrapper: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=chunks), + model=model, + custom_llm_provider="openai", + logging_obj=Logging( + model=model, + messages=messages, + stream=True, + call_type="completion", + start_time=datetime.datetime(2025, 3, 3, 21, 38, 10), + litellm_call_id="late-usage-call", + function_id="late-usage-fn", + ), + stream_options=None, + ) + + collected: Final = [chunk for chunk in wrapper] if sync_mode else [chunk async for chunk in wrapper] + + usage_after_drain: Final = collected[-1]._hidden_params["usage"] + assert (usage_after_drain.prompt_tokens, usage_after_drain.completion_tokens) == reported + assembled: Final = litellm.stream_chunk_builder(chunks=collected, messages=messages) + assert assembled is not None + rates: Final = litellm.model_cost[model] + spend: Final = reported[0] * rates["input_cost_per_token"] + reported[1] * rates["output_cost_per_token"] + assert litellm.completion_cost(completion_response=assembled) == pytest.approx(spend) + + +class _LoggedResponseRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.responses: list[object] = [] + self.logged: Final = threading.Event() + + def log_success_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.responses.append(response_obj) + self.logged.set() + + async def async_log_success_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + self.responses.append(response_obj) + self.logged.set() + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_stream_without_provider_usage_logs_the_estimate_with_the_final_chunk_mcp_metadata( + sync_mode: bool, +) -> None: + model: Final = "gpt-4o" + messages: Final = [{"role": "user", "content": "Write a long sentence about a fox. " * 30}] + text: Final = "The quick brown fox jumps over the lazy dog. " * 20 + tool_calls: Final = [{"name": "lookup_fox", "arguments": "{}"}] + recorder: Final = _LoggedResponseRecorder() + logging_obj: Final = Logging( + model=model, + messages=messages, + stream=True, + call_type="completion" if sync_mode else "acompletion", + start_time=datetime.datetime(2025, 3, 3, 21, 38, 10), + litellm_call_id="mcp-no-usage-call", + function_id="mcp-no-usage-fn", + dynamic_success_callbacks=[recorder], + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model=model, user=None, optional_params={}, litellm_params={"litellm_call_id": "mcp-no-usage-call"} + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator( + model_responses=[ + ModelResponseStream(id="chatcmpl-mcp", created=1741037890, model=model, choices=choices) + for choices in ( + [StreamingChoices(index=0, delta=Delta(role="assistant", content=text))], + [StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")], + ) + ] + ), + model=model, + custom_llm_provider="openai", + logging_obj=logging_obj, + stream_options=None, + ) + _add_mcp_metadata_to_response(wrapper, openai_tools=None, tool_calls=tool_calls) + + _ = [chunk for chunk in wrapper] if sync_mode else [chunk async for chunk in wrapper] + + assert await asyncio.to_thread(recorder.logged.wait, 5) + assert len(recorder.responses) == 1 + logged: Final = recorder.responses[0] + assert isinstance(logged, ModelResponse) + assert logged.choices[0].message.provider_specific_fields == {"mcp_tool_calls": tool_calls} + assert (logged.usage.prompt_tokens, logged.usage.completion_tokens) == ( + litellm.token_counter(model=model, messages=messages), + litellm.token_counter(model=model, text=text, count_response_tokens=True), + ) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("ends_with_finish_chunk", [True, False]) +@pytest.mark.parametrize("reported_tokens", [(30, 12), (0, 0)]) +@pytest.mark.asyncio +async def test_stream_with_provider_usage_hands_that_usage_to_the_last_chunk( + logging_obj: Logging, sync_mode: bool, ends_with_finish_chunk: bool, reported_tokens: tuple[int, int] +) -> None: + model: Final = "gpt-4o" + reported: Final = Usage( + prompt_tokens=reported_tokens[0], + completion_tokens=reported_tokens[1], + total_tokens=sum(reported_tokens), + ) + content_chunk: Final = ModelResponseStream( + id="chatcmpl-with-usage", + created=1741037890, + model=model, + choices=[ + StreamingChoices( + index=0, + delta=Delta(role="assistant", content="Hello there"), + finish_reason=None if ends_with_finish_chunk else "stop", + ) + ], + **({} if ends_with_finish_chunk else {"usage": reported}), + ) + finish_chunk: Final = ModelResponseStream( + id="chatcmpl-with-usage", + created=1741037890, + model=model, + choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")], + usage=reported, + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator( + model_responses=[content_chunk, finish_chunk] if ends_with_finish_chunk else [content_chunk] + ), + model=model, + custom_llm_provider="openai", + logging_obj=logging_obj, + stream_options=None, + ) + + yielded: Final = ( + [(chunk, chunk._hidden_params.get("usage")) for chunk in wrapper] + if sync_mode + else [(chunk, chunk._hidden_params.get("usage")) async for chunk in wrapper] + ) + + last_chunk, usage_when_yielded = yielded[-1] + assert usage_when_yielded is not None + assert (usage_when_yielded.prompt_tokens, usage_when_yielded.completion_tokens) == reported_tokens + usage_after_drain: Final = last_chunk._hidden_params["usage"] + assert (usage_after_drain.prompt_tokens, usage_after_drain.completion_tokens) == reported_tokens + + @pytest.mark.asyncio async def test_custom_stream_wrapper_aclose(): """Test that aclose() delegates to the underlying completion_stream's aclose()""" diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 679e1bfd85a..f837dc3fb5f 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -25,7 +25,7 @@ from litellm.litellm_core_utils.bug_report import ( build_bug_report, ) from litellm.types.router import RouterRateLimitError, RouterRateLimitErrorBasic -from litellm.types.utils import CachingDetails +from litellm.types.utils import CachingDetails, Delta, ModelResponseStream, StreamingChoices from litellm.constants import ( CLIENT_REQUESTED_MODEL_SCOPE_KEY, MAX_LITELLM_CALL_ID_LENGTH, @@ -76,6 +76,8 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( ) from litellm.proxy.utils import ProxyLogging from litellm.router import Router +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.utils import ModelResponseListIterator from litellm.router_utils.add_retry_fallback_headers import prepare_response_for_header_attachment @@ -8480,9 +8482,11 @@ class _RecordingSuccessLogger(CustomLogger): def __init__(self): super().__init__() self.success_events = [] + self.logged: Final = asyncio.Event() async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): self.success_events.append({"kwargs": kwargs, "response_obj": response_obj}) + self.logged.set() class TestStreamingClientDisconnectBilling: @@ -8593,6 +8597,60 @@ class TestStreamingClientDisconnectBilling: assert getattr(partial_response, "service_tier") == "priority" assert partial_response.usage.total_tokens > 0 + @pytest.mark.asyncio + async def test_disconnect_after_finish_chunk_bills_the_estimate_when_the_provider_sent_no_usage(self): + model: Final = "gpt-4o" + messages: Final = [{"role": "user", "content": "Write a long sentence about a fox. " * 30}] + text: Final = "The quick brown fox jumps over the lazy dog. " * 20 + recorder: Final = _RecordingSuccessLogger() + logging_obj: Final = LiteLLMLoggingObj( + model=model, + messages=messages, + stream=True, + call_type="acompletion", + start_time=datetime.datetime(2025, 3, 3, 21, 38, 10), + litellm_call_id="no-usage-disconnect-call", + function_id="no-usage-disconnect-fn", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model=model, user=None, optional_params={}, litellm_params={"litellm_call_id": "no-usage-disconnect-call"} + ) + response: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator( + model_responses=[ + ModelResponseStream(id="chatcmpl-no-usage", created=1741037890, model=model, choices=choices) + for choices in ( + [StreamingChoices(index=0, delta=Delta(role="assistant", content=text))], + [StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")], + ) + ] + ), + model=model, + custom_llm_provider="openai", + logging_obj=logging_obj, + stream_options=None, + ) + stream_iter: Final = response.__aiter__() + first_chunk: Final = await stream_iter.__anext__() + finish_chunk: Final = await stream_iter.__anext__() + assert first_chunk.choices[0].finish_reason is None + assert finish_chunk.choices[0].finish_reason == "stop" + + billed: Final = await _bill_partial_streamed_spend_on_disconnect( + {"litellm_logging_obj": logging_obj}, + response, + ) + await asyncio.wait_for(recorder.logged.wait(), timeout=5) + + assert billed is True + assert len(recorder.success_events) == 1 + usage: Final = recorder.success_events[0]["response_obj"].usage + assert (usage.prompt_tokens, usage.completion_tokens) == ( + litellm.token_counter(model=model, messages=messages), + litellm.token_counter(model=model, text=text, count_response_tokens=True), + ) + @pytest.mark.asyncio async def test_completed_stream_does_not_double_bill_on_late_disconnect(self): recorder = _RecordingSuccessLogger()