From f91774e2f9e47b87062445d755e2bfa2ef1f0fe8 Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Mon, 2 Nov 2026 10:46:30 -0800 Subject: [PATCH 01/79] fix(fireworks): keep developer items in place for models that accept a later system message (#49399) * fix(fireworks): keep developer items in place for models that accept a later system message * fix(fireworks_ai): keep developer items in place for short Fireworks model names and stop flagging gpt-oss-120b * fix(fireworks): drop the mid-conversation system flag from nemotron-lightning nemotron-lightning-3p5-30b-a3b followed a later instruction less often with the developer item kept in place than folded, so both of its cost map keys go back to folding. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../fireworks_ai/responses/transformation.py | 16 +++- ...odel_prices_and_context_window_backup.json | 31 ++++++++ model_prices_and_context_window.json | 31 ++++++++ ...t_fireworks_ai_responses_transformation.py | 78 +++++++++++++++++-- 4 files changed, 148 insertions(+), 8 deletions(-) diff --git a/litellm/llms/fireworks_ai/responses/transformation.py b/litellm/llms/fireworks_ai/responses/transformation.py index 463b9cbfc9a..d599f0b65f7 100644 --- a/litellm/llms/fireworks_ai/responses/transformation.py +++ b/litellm/llms/fireworks_ai/responses/transformation.py @@ -29,6 +29,7 @@ from litellm.utils import ( declared_mid_conversation_system_support, deployment_declared_mid_conversation_system_support, deployment_model_info, + get_bundled_model_cost_map, ) if TYPE_CHECKING: @@ -54,6 +55,15 @@ def _forwarded_user_params(litellm_params: GenericLiteLLMParams) -> Mapping[str, ) +def _bundled_cost_map_declares_mid_conversation_system(model: str) -> bool: + bundled_cost_map: Final[Mapping[str, object]] = get_bundled_model_cost_map() + names: Final = (model.removeprefix("fireworks_ai/"), resolve_fireworks_resource_name(model)) + return any( + deployment_declared_mid_conversation_system_support(bundled_cost_map.get(f"fireworks_ai/{name}")) is True + for name in names + ) + + class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig): @property def custom_llm_provider(self) -> LlmProviders: @@ -90,9 +100,13 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig): validated_input: Final = self._validate_input_param(input) model_info: Final = deployment_model_info(litellm_params) declared: Final = declared_mid_conversation_system_support(model, self.custom_llm_provider.value, model_info) + keeps_developer_items_in_place: Final = declared is True and ( + deployment_declared_mid_conversation_system_support(model_info) is True + or _bundled_cost_map_declares_mid_conversation_system(model) + ) folded: Final = ( None - if deployment_declared_mid_conversation_system_support(model_info) is True + if keeps_developer_items_in_place else instruction_items_folded( validated_input, response_api_optional_request_params, "user" if declared is False else "system" ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index f2a3bb1a47e..02bf7dedcef 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -64039,6 +64039,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64058,6 +64059,7 @@ "output_cost_per_token_priority": 2.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64100,6 +64102,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64137,6 +64140,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64156,6 +64160,7 @@ "output_cost_per_token_priority": 2.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64230,6 +64235,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64250,6 +64256,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64273,6 +64280,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64348,6 +64356,7 @@ "output_cost_per_token_priority": 3e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64384,6 +64393,7 @@ "output_cost_per_token_priority": 6.375e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_vision": true }, @@ -64420,6 +64430,7 @@ "output_cost_per_token_priority": 3e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64490,6 +64501,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64513,6 +64525,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67306,6 +67319,7 @@ "output_cost_per_token_priority": 5.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67324,6 +67338,7 @@ "output_cost_per_token_priority": 8.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67342,6 +67357,7 @@ "output_cost_per_token_priority": 5.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67360,6 +67376,7 @@ "output_cost_per_token_priority": 8.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67375,6 +67392,7 @@ "output_cost_per_token": 6.6e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67417,6 +67435,7 @@ "output_cost_per_token": 6.6e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67435,6 +67454,7 @@ "output_cost_per_token_priority": 6.25e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67453,6 +67473,7 @@ "output_cost_per_token_priority": 9.375e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67470,6 +67491,7 @@ "output_cost_per_token_priority": 6.25e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67488,6 +67510,7 @@ "output_cost_per_token_priority": 9.375e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67506,6 +67529,7 @@ "output_cost_per_token_priority": 4.05e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -73747,6 +73771,7 @@ "output_cost_per_token_priority": 1.875e-05, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -77249,6 +77274,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -77407,6 +77433,7 @@ "output_cost_per_token_priority": 9e-07, "source": "https://platform.minimax.io/docs/guides/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -81184,6 +81211,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81200,6 +81228,7 @@ "output_cost_per_token": 3.08e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81216,6 +81245,7 @@ "output_cost_per_token": 8.4e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81232,6 +81262,7 @@ "output_cost_per_token": 8.4e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index f2a3bb1a47e..02bf7dedcef 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -64039,6 +64039,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64058,6 +64059,7 @@ "output_cost_per_token_priority": 2.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64100,6 +64102,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64137,6 +64140,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64156,6 +64160,7 @@ "output_cost_per_token_priority": 2.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -64230,6 +64235,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64250,6 +64256,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64273,6 +64280,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64348,6 +64356,7 @@ "output_cost_per_token_priority": 3e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64384,6 +64393,7 @@ "output_cost_per_token_priority": 6.375e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_vision": true }, @@ -64420,6 +64430,7 @@ "output_cost_per_token_priority": 3e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64490,6 +64501,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -64513,6 +64525,7 @@ ], "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67306,6 +67319,7 @@ "output_cost_per_token_priority": 5.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67324,6 +67338,7 @@ "output_cost_per_token_priority": 8.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67342,6 +67357,7 @@ "output_cost_per_token_priority": 5.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67360,6 +67376,7 @@ "output_cost_per_token_priority": 8.25e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67375,6 +67392,7 @@ "output_cost_per_token": 6.6e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67417,6 +67435,7 @@ "output_cost_per_token": 6.6e-06, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -67435,6 +67454,7 @@ "output_cost_per_token_priority": 6.25e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67453,6 +67473,7 @@ "output_cost_per_token_priority": 9.375e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67470,6 +67491,7 @@ "output_cost_per_token_priority": 6.25e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67488,6 +67510,7 @@ "output_cost_per_token_priority": 9.375e-07, "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -67506,6 +67529,7 @@ "output_cost_per_token_priority": 4.05e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true @@ -73747,6 +73771,7 @@ "output_cost_per_token_priority": 1.875e-05, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -77249,6 +77274,7 @@ "output_cost_per_token_priority": 1.5e-06, "source": "https://api.fireworks.ai/v1/serverless/models?format=nested", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -77407,6 +77433,7 @@ "output_cost_per_token_priority": 9e-07, "source": "https://platform.minimax.io/docs/guides/pricing", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -81184,6 +81211,7 @@ ], "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81200,6 +81228,7 @@ "output_cost_per_token": 3.08e-06, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81216,6 +81245,7 @@ "output_cost_per_token": 8.4e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -81232,6 +81262,7 @@ "output_cost_per_token": 8.4e-07, "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, + "supports_mid_conversation_system": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, diff --git a/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py index 7fc6ac2c436..6538ccf1cac 100644 --- a/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py +++ b/tests/unit/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py @@ -172,10 +172,10 @@ def test_responses_call_forwards_previous_response_id_and_store() -> None: def test_responses_call_folds_developer_items_into_instructions() -> None: - client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b")) with patch(HTTPX_CLIENT_FACTORY, return_value=client): litellm.responses( - model="fireworks_ai/accounts/fireworks/models/kimi-k3", + model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b", input=[ {"role": "user", "content": "Hi there"}, {"role": "developer", "content": "Answer with exactly one word."}, @@ -214,6 +214,68 @@ def test_responses_call_keeps_developer_items_in_place_for_a_deployment_that_acc ) +@respx.mock +@pytest.mark.parametrize( + "model", + [ + "fireworks_ai/accounts/fireworks/models/kimi-k3", + "fireworks_ai/accounts/fireworks/models/glm-5p3", + "fireworks_ai/accounts/fireworks/routers/kimi-k3-fast", + "fireworks_ai/glm-5p3-us", + "fireworks_ai/models/kimi-k3", + "fireworks_ai/routers/kimi-k3-fast", + ], +) +def test_responses_call_keeps_developer_items_in_place_for_a_model_the_cost_map_declares_accepts_them( + model: str, +) -> None: + route: Final = _fireworks_responses_route(model.removeprefix("fireworks_ai/")) + litellm.responses( + model=model, + instructions="You are Codex", + input=[ + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + ], + api_key="fw-test-key", + ) + body: Final = _routed_body(route) + assert body["instructions"] == "You are Codex" + assert tuple(body["input"]) == ( + {"role": "user", "content": "Hi there"}, + {"role": "system", "content": "Answer with exactly one word.", "type": "message"}, + {"role": "user", "content": "What is the capital of France?"}, + ) + + +@respx.mock +def test_responses_call_folds_developer_items_when_the_loaded_cost_map_turns_the_flag_off( + declare_mid_conversation_system: Callable[[str, bool], None], +) -> None: + declare_mid_conversation_system("fireworks_ai/accounts/fireworks/models/kimi-k3", False) + route: Final = _fireworks_responses_route("accounts/fireworks/models/kimi-k3") + litellm.responses( + model="fireworks_ai/accounts/fireworks/models/kimi-k3", + input=[ + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": "What is the capital of France?"}, + {"type": "message", "role": "assistant", "content": "Paris"}, + {"role": "developer", "content": "Reply with PONG."}, + ], + api_key="fw-test-key", + ) + body: Final = _routed_body(route) + assert body["instructions"] == "Answer with exactly one word." + assert tuple(body["input"]) == ( + {"role": "user", "content": "Hi there"}, + {"role": "user", "content": "What is the capital of France?"}, + {"type": "message", "role": "assistant", "content": "Paris"}, + {"role": "user", "content": "Reply with PONG.", "type": "message"}, + ) + + @respx.mock @pytest.mark.parametrize( ("model_info", "closing_role"), @@ -222,9 +284,9 @@ def test_responses_call_keeps_developer_items_in_place_for_a_deployment_that_acc def test_responses_call_sends_the_closing_developer_item_as_a_user_turn_only_for_a_declared_system_first_deployment( model_info: Mapping[str, object] | None, closing_role: str ) -> None: - route: Final = _fireworks_responses_route("accounts/fireworks/models/kimi-k3") + route: Final = _fireworks_responses_route("accounts/fireworks/models/qwen3p8-2p4t-a95b") litellm.responses( - model="fireworks_ai/accounts/fireworks/models/kimi-k3", + model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b", input=[ {"role": "user", "content": "Hi there"}, {"role": "developer", "content": "Answer with exactly one word."}, @@ -402,7 +464,7 @@ def test_transform_request_forwards_non_string_instructions_and_input_untouched( developer_item: Final = {"role": "developer", "content": "Answer with exactly one word."} user_item: Final = {"role": "user", "content": "What is the capital of France?"} request: Final = FireworksAIResponsesAPIConfig().transform_responses_api_request( - model="accounts/fireworks/models/kimi-k3", + model="accounts/fireworks/models/qwen3p8-2p4t-a95b", input=cast(ResponseInputParam, [developer_item, user_item]), response_api_optional_request_params={"instructions": ["not", "a", "string"]}, litellm_params=GenericLiteLLMParams(), @@ -416,7 +478,7 @@ def test_transform_request_forwards_non_string_instructions_and_input_untouched( def test_responses_call_maps_pydantic_developer_items_and_replays_pydantic_output_items() -> None: - client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/qwen3p8-2p4t-a95b")) pydantic_input: Final = cast( ResponseInputParam, [ @@ -435,7 +497,9 @@ def test_responses_call_maps_pydantic_developer_items_and_replays_pydantic_outpu ) with patch(HTTPX_CLIENT_FACTORY, return_value=client): litellm.responses( - model="fireworks_ai/accounts/fireworks/models/kimi-k3", input=pydantic_input, api_key="fw-test-key" + model="fireworks_ai/accounts/fireworks/models/qwen3p8-2p4t-a95b", + input=pydantic_input, + api_key="fw-test-key", ) _, _, body = _sent_request(client) assert body["instructions"] == "Answer with exactly one word." From 508a98586eda49e57f2783741617185f67cbd608 Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Mon, 2 Nov 2026 11:17:40 -0800 Subject: [PATCH 02/79] fix(databricks): join plain text blocks for non-Claude json_object requests (#49385) * fix(databricks): join plain text blocks for non-Claude json_object requests * fix(databricks): join only the last message for json_object and drop cache_control under drop_params Databricks checks for "json" only in the last message, so earlier turns now go out as sent and only the last message, plus any system messages merged into it, is joined into a string. Under drop_params a text block's cache_control is dropped so that message can be joined too, and without it the request still goes out as a list. * fix(databricks): skip dropped assistant turns and blank parts when joining for json_object The json_object join picked the last message before the request transform drops bare assistant turns, so a trailing thinking-only or reasoning-only turn (a Responses input ending in a reasoning item included) left the real last user message as a list and Databricks still returned its 400. The pick now skips the turns that transform drops. The join also ran when only one part had text, which sent whitespace-only parts the existing empty block filter used to drop on requests that already worked. It now joins only two or more non-blank parts, so a request with one non-blank part goes out exactly as before. Without drop_params, a message with a message-level cache_control now goes out as sent instead of being rewritten into one cached block. * fix(databricks): join json_object text parts only on models Databricks hosts Databricks checks for the word "json" only on the models it hosts (databricks-* pay per token names and system.ai.* AI Gateway names). An external model endpoint skips that check and takes a list of text parts today, so its request now goes out as sent instead of joined. The drop_params read now goes through drops_unsupported_params, and a test pins that a text block with a truthy non-string text is kept. * fix(databricks): reject a text part holding a non-string locally with a clear 400 * refactor(databricks): reject a non-string text part at the transform_request boundary * fix(databricks): read pydantic message objects in the non-string text check --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../llms/databricks/chat/transformation.py | 167 +++++++++- litellm/llms/databricks/common_utils.py | 1 + .../providers/test_databricks_chat_wire.py | 169 +++++++++- .../test_databricks_chat_transformation.py | 289 +++++++++++++++++- 4 files changed, 607 insertions(+), 19 deletions(-) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 60081c38073..53674907341 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -4,6 +4,7 @@ Translates from OpenAI's `/v1/chat/completions` to Databricks' `/chat/completion import os from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence +from itertools import chain, takewhile from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload @@ -12,6 +13,7 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter from litellm.anthropic_beta_gated_fields import log_beta_gated_strips, strip_beta_gated_thinking from litellm.constants import RESPONSE_FORMAT_TOOL_NAME +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( handle_invalid_parallel_tool_calls, should_convert_tool_call_to_json_mode, @@ -65,7 +67,7 @@ from ...anthropic.chat.transformation import ( tools_without_unsupported_tool_search, ) from ...openai_like.chat.transformation import OpenAILikeChatConfig -from ..common_utils import DatabricksBase, DatabricksException +from ..common_utils import DATABRICKS_HOSTED_MODEL_PREFIXES, DatabricksBase, DatabricksException _RESPONSE_FORMAT_ADAPTER: ( # mutable-ok: mirrors the dict return contract of get_json_schema_from_pydantic_object Final[TypeAdapter[dict[str, object] | None]] @@ -81,15 +83,139 @@ def _is_bare_assistant_message(message_dict: Mapping[str, object]) -> bool: ) +_CONTENT_BLOCKS: Final = TypeAdapter(tuple[object, ...]) + + +def _content_blocks(message: object) -> tuple[object, ...]: + match message.model_dump(exclude_none=True) if isinstance(message, BaseModel) else message: + case {"content": list() as content}: + return _CONTENT_BLOCKS.validate_python(content) + case _: + return () + + +def _plain_text(block: object) -> str | None: + match block: + case {"type": "text", "text": str() as text} if block == {"type": "text", "text": text}: + return text + case _: + return None + + +def _plain_texts(message: Mapping[str, object]) -> tuple[str, ...]: + blocks: Final = _content_blocks(message) + texts: Final = tuple(text for block in blocks if (text := _plain_text(block)) is not None) + return texts if len(texts) == len(blocks) else () + + def _collapse_single_text_block(message: Mapping[str, object]) -> Mapping[str, object]: - content: Final = message.get("content") - match content: - case [{"type": "text", "text": str() as text}] if content == [{"type": "text", "text": text}]: + match _plain_texts(message): + case (text,): return {**message, "content": text} case _: return message +def _without_text_cache_control(block: object) -> object: + match block: + case {"type": "text", "text": str() as text, "cache_control": object() as cache_control} if block == { + "type": "text", + "text": text, + "cache_control": cache_control, + }: + return {"type": "text", "text": text} + case _: + return block + + +def _without_cache_control(message: AllMessageValues) -> AllMessageValues: + uncached: Final = {key: value for key, value in message.items() if key != "cache_control"} + match content := uncached.get("content"): + case list(): + blocks: Final = [_without_text_cache_control(block) for block in _CONTENT_BLOCKS.validate_python(content)] + return cast(AllMessageValues, {**uncached, "content": blocks}) # cast-ok: spread keeps the shape + case _: + return cast(AllMessageValues, uncached) # cast-ok: only cache_control was left out + + +def _has_string_content(message: Mapping[str, object]) -> bool: + return isinstance(message.get("content"), str) + + +def _is_system(message: Mapping[str, object]) -> bool: + return message.get("role") == "system" + + +def _json_mode_message(message: AllMessageValues, drop_params: bool) -> AllMessageValues: + candidate: Final = _without_cache_control(message) if drop_params else message + texts: Final = _plain_texts(candidate) + if "cache_control" in candidate or not (texts or _has_string_content(candidate)): + return message + if sum(1 for text in texts if text.strip()) < 2: + return candidate + return cast(AllMessageValues, {**candidate, "content": "".join(texts)}) # cast-ok: spread keeps the shape + + +def _is_dropped_on_send(message: Mapping[str, object]) -> bool: + sanitized: Final = {**message} + _sanitize_empty_content(sanitized) + return _is_bare_assistant_message(sanitized) + + +def _is_system_or_dropped(message: Mapping[str, object]) -> bool: + return _is_system(message) or _is_dropped_on_send(message) + + +def _json_checked_start(messages: Sequence[AllMessageValues]) -> int: + trailing: Final = tuple(takewhile(_is_system_or_dropped, reversed(messages))) + last_sent_offset: Final = 0 if any(_is_system(message) for message in trailing) else 1 + return max(len(messages) - len(trailing) - last_sent_offset, 0) + + +def _json_object_ready_messages( + model: str, + messages: list[AllMessageValues], # mutable-ok: inherited transform_request contract + response_format: object, + drop_params: bool, +) -> list[AllMessageValues]: # mutable-ok: inherited transform_request contract + match response_format: + case {"type": "json_object"} if "claude" not in model and model.startswith(DATABRICKS_HOSTED_MODEL_PREFIXES): + checked_start: Final = _json_checked_start(messages) + return [ + *messages[:checked_start], + *(_json_mode_message(message, drop_params) for message in messages[checked_start:]), + ] + case _: + return messages + + +def _is_empty_text_block(block: object) -> bool: + match block: + case {"type": "text", "text": str() as text}: + return not text.strip() + case {"type": "text", "text": object() as text}: + return not text + case {"type": "text"}: + return True + case _: + return False + + +def _non_string_text_type(block: object) -> str | None: + match block: + case {"type": "text", "text": str()}: + return None + case {"type": "text", "text": object() as text} if text: + return type(text).__name__ + case _: + return None + + +def _first_non_string_text_type(messages: Sequence[AllMessageValues]) -> str | None: + blocks: Final = chain.from_iterable(map(_content_blocks, messages)) + return next((type_name for block in blocks if (type_name := _non_string_text_type(block))), None) + + def _sanitize_empty_content(message_dict: dict[str, object]) -> None: """ Remove or filter content so empty text blocks are not sent. @@ -108,9 +234,7 @@ def _sanitize_empty_content(message_dict: dict[str, object]) -> None: message_dict.pop("content") return filtered: Final = [ - block - for block in content - if not (isinstance(block, dict) and block.get("type") == "text" and not (block.get("text") or "").strip()) + block for block in _CONTENT_BLOCKS.validate_python(content) if not _is_empty_text_block(block) ] if not filtered: message_dict.pop("content") @@ -535,6 +659,35 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): return False + def transform_request( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: inherited contract + optional_params: dict[str, object], # mutable-ok: inherited contract + litellm_params: dict[str, object], # mutable-ok: inherited contract + headers: dict[str, object], # mutable-ok: inherited contract + ) -> dict[str, object]: # mutable-ok: inherited contract + match _first_non_string_text_type(messages): + case str() as type_name: + raise DatabricksException( + status_code=400, + message=f"A text content block needs a string 'text', got {type_name}", + ) + case None: + pass + return super().transform_request( + model=model, + messages=_json_object_ready_messages( + model, + messages, + optional_params.get("response_format"), + drop_params=drops_unsupported_params(litellm_params.get("drop_params")), + ), + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + @overload def _transform_messages( self, messages: list[AllMessageValues], model: str, is_async: Literal[True] diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index b1562bef6b8..e8313d9a003 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -56,6 +56,7 @@ DATABRICKS_WORKSPACE_ROUTE_PREFIXES: Final = ( DATABRICKS_AI_GATEWAY_MESSAGES_PATH, ) DATABRICKS_NATIVE_MESSAGES_MODEL_PREFIXES: Final = ("databricks-claude-", "system.ai.claude-") +DATABRICKS_HOSTED_MODEL_PREFIXES: Final = ("databricks-", "system.ai.") DATABRICKS_ENDPOINT_PATHS: Final = MappingProxyType( { "chat_completions": "/chat/completions", diff --git a/tests/integration/providers/test_databricks_chat_wire.py b/tests/integration/providers/test_databricks_chat_wire.py index b93cebdbc4c..3d9d749eb04 100644 --- a/tests/integration/providers/test_databricks_chat_wire.py +++ b/tests/integration/providers/test_databricks_chat_wire.py @@ -353,9 +353,172 @@ def _shape_reply(identity: str, stream: bool) -> Reply: "/chat/completions", id="A5-responses-list-input", ), + pytest.param( + "json-object-two-blocks", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]} + ], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A6-chat-json-object-two-blocks-joined", + ), + pytest.param( + "json-object-two-blocks-stream", + "/v1/chat/completions", + "databricks-qwen35-122b-a10b", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]} + ], + "response_format": {"type": "json_object"}, + "stream": True, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A7-chat-stream-json-object-two-blocks-joined", + ), + pytest.param( + "json-object-history", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Code: A"}, {"type": "text", "text": "BC"}]}, + {"role": "assistant", "content": "Got it."}, + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}, + ], + "response_format": {"type": "json_object"}, + }, + [ + {"role": "user", "content": [{"type": "text", "text": "Code: A"}, {"type": "text", "text": "BC"}]}, + {"role": "assistant", "content": "Got it."}, + {"role": "user", "content": "Reply in JSON"}, + ], + False, + "/chat/completions", + id="A8-chat-json-object-earlier-turns-kept", + ), + pytest.param( + "json-object-trailing-system", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + {"role": "user", "content": "Say hi"}, + {"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]}, + {"role": "system", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}, + ], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": "Say hi"}, {"role": "system", "content": "Alpha rule.\n\nReply in JSON"}], + False, + "/chat/completions", + id="A8b-chat-json-object-trailing-system-keeps-break", + ), + pytest.param( + "json-object-drop-params-cache-control", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Reply"}, + {"type": "text", "text": " in JSON", "cache_control": {"type": "ephemeral"}}, + ], + } + ], + "response_format": {"type": "json_object"}, + "drop_params": True, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A8c-chat-json-object-drop-params-drops-cache-control", + ), + pytest.param( + "json-object-trailing-thinking-turn", + "/v1/chat/completions", + "databricks-meta-llama-3-3-70b-instruct", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}, + { + "role": "assistant", + "content": "", + "thinking_blocks": [{"type": "thinking", "thinking": "Plan.", "signature": "sig"}], + }, + ], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A8d-chat-json-object-trailing-thinking-turn-dropped", + ), + pytest.param( + "json-object-external-endpoint", + "/v1/chat/completions", + "openai-gpt-external", + { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]} + ], + "response_format": {"type": "json_object"}, + }, + [{"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}], + False, + "/chat/completions", + id="A8e-chat-json-object-external-endpoint-kept-as-sent", + ), + pytest.param( + "json-object-responses-two-parts", + "/v1/responses", + "databricks-qwen35-122b-a10b", + { + "input": [ + { + "role": "user", + "content": [{"type": "input_text", "text": "Reply"}, {"type": "input_text", "text": " in JSON"}], + } + ], + "text": {"format": {"type": "json_object"}}, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A9-responses-json-object-two-parts-joined", + ), + pytest.param( + "json-object-responses-trailing-reasoning", + "/v1/responses", + "databricks-qwen35-122b-a10b", + { + "input": [ + { + "role": "user", + "content": [{"type": "input_text", "text": "Reply"}, {"type": "input_text", "text": " in JSON"}], + }, + {"type": "reasoning", "id": "rs_1", "summary": [{"type": "summary_text", "text": "Plan."}]}, + ], + "text": {"format": {"type": "json_object"}}, + }, + [{"role": "user", "content": "Reply in JSON"}], + False, + "/chat/completions", + id="A9b-responses-json-object-trailing-reasoning-dropped", + ), ], ) -def test_databricks_non_claude_single_text_block_reaches_upstream_as_string( +def test_databricks_non_claude_plain_text_content_reaches_upstream_as_string( gateway: Gateway, case: str, endpoint: str, @@ -365,7 +528,7 @@ def test_databricks_non_claude_single_text_block_reaches_upstream_as_string( responses_mode: bool, upstream_path: str, ) -> None: - identity: Final = f"databricks-single-text-{case}-{uuid.uuid4().hex}" + identity: Final = f"databricks-plain-text-{case}-{uuid.uuid4().hex}" streaming: Final = request_body.get("stream") is True def respond(request: Request) -> Reply: @@ -1453,7 +1616,7 @@ def test_databricks_non_string_text_block_is_rejected_before_reaching_upstream(g ) received: Final = wire.drain() assert response.status_code == 400, response.text - assert "'int' object has no attribute 'strip'" in response.text + assert "A text content block needs a string 'text', got int" in response.text assert received == () diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index 7a6f9c7dc18..33c1e7a2e22 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -1,14 +1,14 @@ import copy import json from collections.abc import AsyncIterator, Iterator -from typing import Final +from typing import Final, cast from unittest.mock import MagicMock, patch import httpx import pytest import respx from openai.lib.streaming.chat import ChatCompletionStreamState -from openai.types.chat import ChatCompletionChunk +from openai.types.chat import ChatCompletionChunk, ChatCompletionMessage from pydantic import TypeAdapter from fastapi.testclient import TestClient from pydantic import BaseModel, ConfigDict @@ -21,6 +21,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.databricks.common_utils import DatabricksException from litellm.llms.databricks.chat.transformation import ( DatabricksChatResponseIterator, DatabricksConfig, @@ -539,13 +540,6 @@ def test_transform_request_preserves_unity_model_service_name(): [{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}], id="cache-control", ), - pytest.param( - "databricks-meta-llama-3-3-70b-instruct", - [{"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}], - {"response_format": {"type": "json_object"}}, - [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}], - id="multiple-text-blocks", - ), pytest.param( "databricks-meta-llama-3-3-70b-instruct", [ @@ -603,6 +597,283 @@ def test_transform_request_collapses_single_plain_text_block_for_non_claude_mode assert outgoing_messages == [{**messages[0], "content": expected_content}] +_TWO_TEXT_BLOCKS: Final = [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}] +_JSON_OBJECT: Final = {"response_format": {"type": "json_object"}} +_IMAGE_BLOCK: Final = {"type": "image_url", "image_url": {"url": "https://example.com/image.png"}} +_CACHED_BLOCK: Final = {"type": "text", "text": " in JSON", "cache_control": {"type": "ephemeral"}} + + +def _outgoing_messages( + model: str, messages: list[AllMessageValues], optional_params: dict[str, object], drop_params: bool = False +) -> list[dict[str, object]]: + request: Final = TypeAdapter(dict[str, object]).validate_python( + DatabricksConfig().transform_request( # pyright: ignore[reportUnknownMemberType] # inherited hook is untyped + model=model, + messages=messages, + optional_params=optional_params, + litellm_params={"drop_params": True} if drop_params else {}, + headers={}, + ) + ) + return TypeAdapter(list[dict[str, object]]).validate_python(request["messages"]) + + +@pytest.mark.parametrize( + "model,messages,optional_params,expected_content", + [ + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + "Reply in JSON", + id="json-object-two-blocks-joined", + ), + pytest.param( + "databricks-qwen35-122b-a10b", + [{"role": "user", "content": [{"type": "text", "text": "A"}, {"type": "text", "text": "B json"}] * 2}], + _JSON_OBJECT, + "AB jsonAB json", + id="json-object-four-blocks-joined-as-sent", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + {}, + _TWO_TEXT_BLOCKS, + id="no-response-format-stays-list", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + {"response_format": {"type": "json_schema", "json_schema": {"name": "r", "schema": {"type": "object"}}}}, + _TWO_TEXT_BLOCKS, + id="json-schema-stays-list", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": [{"type": "text", "text": "Reply"}, _CACHED_BLOCK]}], + _JSON_OBJECT, + [{"type": "text", "text": "Reply"}, _CACHED_BLOCK], + id="json-object-cache-control-stays-list", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": "Reply in JSON", "cache_control": {"type": "ephemeral"}}], + _JSON_OBJECT, + [{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}], + id="json-object-message-cache-control-stays-block", + ), + pytest.param( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}, _IMAGE_BLOCK]}], + _JSON_OBJECT, + [{"type": "text", "text": "Reply in JSON"}, _IMAGE_BLOCK], + id="json-object-text-and-image-stays-list", + ), + pytest.param( + "databricks-claude-sonnet-5", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + _TWO_TEXT_BLOCKS, + id="json-object-claude-stays-list", + ), + pytest.param( + "system.ai.deepseek-v4-flash-0731", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + "Reply in JSON", + id="json-object-ai-gateway-model-joined", + ), + pytest.param( + "openai-gpt-external", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + _TWO_TEXT_BLOCKS, + id="json-object-external-endpoint-stays-list", + ), + ], +) +def test_transform_request_joins_plain_text_blocks_only_for_hosted_non_claude_json_object( + model: str, + messages: list[AllMessageValues], + optional_params: dict[str, object], + expected_content: object, +) -> None: + assert _outgoing_messages(model, messages, optional_params) == [{"role": "user", "content": expected_content}] + + +@pytest.mark.parametrize("response_format", [None, {"type": "json_object"}], ids=["no-response-format", "json-object"]) +def test_transform_request_rejects_a_text_block_whose_text_is_a_truthy_non_string( + response_format: dict[str, str] | None, +) -> None: + with pytest.raises(DatabricksException) as error: + _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}, {"type": "text", "text": 7}]}], + {} if response_format is None else {"response_format": response_format}, + ) + + assert error.value.status_code == 400 + + +def test_transform_request_accepts_an_openai_message_object_in_history() -> None: + earlier_answer: Final = ChatCompletionMessage(role="assistant", content="Earlier answer") + + outgoing: Final = _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [ + {"role": "user", "content": "Earlier question"}, + cast(AllMessageValues, earlier_answer), # cast-ok: SDK callers pass response messages back as history + {"role": "user", "content": "Next question"}, + ], + {}, + ) + + assert [message.get("content") for message in outgoing] == ["Earlier question", "Earlier answer", "Next question"] + + +def test_transform_request_json_object_keeps_earlier_turns_as_sent() -> None: + earlier_turn: Final[AllMessageValues] = { + "role": "user", + "content": [{"type": "text", "text": "Code: A"}, {"type": "text", "text": "BC"}], + } + reply: Final[AllMessageValues] = {"role": "assistant", "content": "Got it."} + + assert _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [earlier_turn, reply, {"role": "user", "content": _TWO_TEXT_BLOCKS}], + _JSON_OBJECT, + ) == [earlier_turn, reply, {"role": "user", "content": "Reply in JSON"}] + + +def test_transform_request_json_object_joins_a_trailing_system_run_and_keeps_its_break() -> None: + messages: Final[list[AllMessageValues]] = [ + {"role": "user", "content": "Say hi"}, + {"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]}, + {"role": "system", "content": _TWO_TEXT_BLOCKS}, + ] + + assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", messages, _JSON_OBJECT) == [ + {"role": "user", "content": "Say hi"}, + {"role": "system", "content": "Alpha rule.\n\nReply in JSON"}, + ] + + +_CACHED_TEXT_AND_IMAGE: Final = [{**_CACHED_BLOCK, "text": "Reply in JSON"}, _IMAGE_BLOCK] + + +@pytest.mark.parametrize( + "message,expected", + [ + pytest.param( + {"role": "user", "content": [{"type": "text", "text": "Reply"}, _CACHED_BLOCK]}, + {"role": "user", "content": "Reply in JSON"}, + id="block-cache-control-dropped-and-joined", + ), + pytest.param( + {"role": "user", "content": "Reply in JSON", "cache_control": {"type": "ephemeral"}}, + {"role": "user", "content": "Reply in JSON"}, + id="message-cache-control-dropped", + ), + pytest.param( + {"role": "user", "content": _CACHED_TEXT_AND_IMAGE}, + {"role": "user", "content": _CACHED_TEXT_AND_IMAGE}, + id="image-keeps-cache-control", + ), + ], +) +def test_transform_request_json_object_drops_last_message_cache_control_under_drop_params( + message: AllMessageValues, expected: dict[str, object] +) -> None: + assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", [message], _JSON_OBJECT, drop_params=True) == [ + expected + ] + + +_THINKING_BLOCK: Final = {"type": "thinking", "thinking": "Plan the reply.", "signature": "sig"} + + +@pytest.mark.parametrize( + "dropped_turn", + [ + pytest.param({"role": "assistant", "content": "", "thinking_blocks": [_THINKING_BLOCK]}, id="thinking-only"), + pytest.param({"role": "assistant", "content": None, "reasoning_content": "Plan."}, id="reasoning-only"), + pytest.param({"role": "assistant", "content": [{"type": "text", "text": " "}]}, id="blank-text-part"), + ], +) +def test_transform_request_json_object_joins_the_message_before_a_dropped_assistant_turn( + dropped_turn: AllMessageValues, +) -> None: + assert _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": _TWO_TEXT_BLOCKS}, dropped_turn], + _JSON_OBJECT, + ) == [{"role": "user", "content": "Reply in JSON"}] + + +def test_transform_request_json_object_joins_a_system_run_split_by_a_dropped_assistant_turn() -> None: + earlier_turn: Final[AllMessageValues] = { + "role": "user", + "content": [{"type": "text", "text": "Say"}, {"type": "text", "text": " hi"}], + } + messages: Final[list[AllMessageValues]] = [ + earlier_turn, + {"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]}, + {"role": "assistant", "content": "", "thinking_blocks": [_THINKING_BLOCK]}, + {"role": "system", "content": _TWO_TEXT_BLOCKS}, + ] + + assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", messages, _JSON_OBJECT) == [ + earlier_turn, + {"role": "system", "content": "Alpha rule.\n\nReply in JSON"}, + ] + + +@pytest.mark.parametrize( + "content,expected_content", + [ + pytest.param( + [{"type": "text", "text": "Reply in JSON"}, {"type": "text", "text": "\n\n"}], + "Reply in JSON", + id="one-non-blank-part-goes-out-as-before", + ), + pytest.param( + [{"type": "text", "text": "Reply"}, {"type": "text", "text": " "}, {"type": "text", "text": "in JSON"}], + "Reply in JSON", + id="blank-separator-between-parts-kept", + ), + ], +) +def test_transform_request_json_object_joins_only_two_or_more_non_blank_parts( + content: list[object], expected_content: str +) -> None: + assert _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", [{"role": "user", "content": content}], _JSON_OBJECT + ) == [{"role": "user", "content": expected_content}] + + +def test_transform_request_json_object_keeps_message_cache_control_parts_as_sent_without_drop_params() -> None: + message: Final[AllMessageValues] = { + "role": "user", + "content": _TWO_TEXT_BLOCKS, + "cache_control": {"type": "ephemeral"}, + } + + assert _outgoing_messages("databricks-meta-llama-3-3-70b-instruct", [message], {}) == _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", [message], _JSON_OBJECT + ) + + +def test_transform_request_json_object_reads_the_global_drop_params(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", True) + + assert _outgoing_messages( + "databricks-meta-llama-3-3-70b-instruct", + [{"role": "user", "content": [{"type": "text", "text": "Reply"}, _CACHED_BLOCK]}], + _JSON_OBJECT, + ) == [{"role": "user", "content": "Reply in JSON"}] + + def test_transform_request_strips_thinking_blocks_and_reasoning_content(): """Regression for LIT-6762: replaying an assistant turn that litellm decorated with `thinking_blocks` / `reasoning_content` made Databricks 400 with From f8f12a270a0db490e75adfa35d9e8175727f8bf4 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 2 Nov 2026 11:23:11 -0800 Subject: [PATCH 03/79] test(integration): repair stale contract fixtures (#49700) --- .../test_response_cache_empty_output.py | 19 ++++++++--- .../compatibility/test_a2a_wire_versions.py | 25 +++++++++++--- .../cost_calculation/cases/azure.py | 1 + .../cost_calculation/cases/openai.py | 1 + .../test_pdf_document_count_tokens_wire.py | 6 ++-- .../test_prometheus_provider_auth_failure.py | 4 ++- .../test_device_code_login_guard_boot.py | 33 +++++-------------- .../test_device_code_login_guard_wire.py | 17 +++++----- ...anthropic_json_tool_cut_stream_sync_sdk.py | 14 +++++--- 9 files changed, 70 insertions(+), 50 deletions(-) diff --git a/tests/integration/caching/test_response_cache_empty_output.py b/tests/integration/caching/test_response_cache_empty_output.py index ce5a4d64ca7..7a20e712d4e 100644 --- a/tests/integration/caching/test_response_cache_empty_output.py +++ b/tests/integration/caching/test_response_cache_empty_output.py @@ -23,7 +23,7 @@ from typing import Final, TypeAlias, TypeVar import httpx import pytest from anthropic import Anthropic, AsyncAnthropic -from openai import AsyncOpenAI, OpenAI +from openai import AsyncOpenAI, InternalServerError, OpenAI from pydantic import JsonValue from tests.integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value @@ -420,14 +420,25 @@ def _streamed_chat(gateway: Gateway, model: str, text: str) -> _Answer: return ids, streamed, tuple(chunk.object for chunk in chunks) -def test_streamed_chat_without_choices_is_never_assembled_or_cached(gateway: Gateway) -> None: +def test_streamed_chat_without_choices_fails_and_is_never_cached(gateway: Gateway) -> None: upstream: Final = Upstream(gateway.upstream_url) + store: Final = redis_store() text: Final = prompt() with gateway.scenario() as scenario: handle: Final = scripted(scenario, CHAT_STREAM_WITHOUT_CHOICES) model: Final = scenario.model(api_base=handle.api_base()) - empty: Final = _never_replayed(partial(_streamed_chat, gateway, model, text), upstream, handle.scenario_id) - assert all(answer.text == "" for answer in empty), empty + stream: Final = partial(_streamed_chat, gateway, model, text) + for attempt in range(1, EMPTY_STREAMS + 1): + with pytest.raises(InternalServerError, match="Stream ended before a finish_reason was received"): + stream() + _assert_prompt_reached_upstream(upstream, handle.scenario_id, text, attempt) + rescript(handle, CHAT_STREAM) + real_ids, real_text, _ = stream() + assert real_text == "streamed response" and len(real_ids) == 1, (real_ids, real_text) + _assert_prompt_reached_upstream(upstream, handle.scenario_id, text, EMPTY_STREAMS + 1) + await_entry(store, next(iter(real_ids))) + replay: Final = _replayed(stream, upstream, handle.scenario_id, real_ids) + assert replay.text == real_text, replay def _streamed_text(gateway: Gateway, model: str, text: str) -> _Answer: diff --git a/tests/integration/compatibility/test_a2a_wire_versions.py b/tests/integration/compatibility/test_a2a_wire_versions.py index 4471322d1e9..b3ca0c99392 100644 --- a/tests/integration/compatibility/test_a2a_wire_versions.py +++ b/tests/integration/compatibility/test_a2a_wire_versions.py @@ -1023,9 +1023,18 @@ def test_a2a_task_methods_pass_through_verbatim_and_push_callbacks_are_validated {"taskId": "t1", "url": callback}, ): rejected = call("tasks/pushNotificationConfig/set", params) - assert (rejected.status_code, rejected.json()) == (400, {"detail": detail}), ( - f"{params}: {rejected.text}" - ) + assert (rejected.status_code, rejected.json()) == ( + 400, + { + "jsonrpc": "2.0", + "id": marker + "tasks/pushNotificationConfig/set", + "error": { + "code": -32000, + "message": detail, + "data": {"type": "invalid_request_error", "param": None, "code": "400"}, + }, + }, + ), f"{params}: {rejected.text}" assert tuple(item for item in wire.drain() if item.method == "POST") == () public: Final = {"taskId": "t1", "pushNotificationConfig": {"url": "https://1.1.1.1/callback"}} @@ -1118,7 +1127,15 @@ def test_a2a_v1_task_methods_and_extended_card_lower_to_0_3_upstream_and_answer_ rejected = call("CreateTaskPushNotificationConfig", params) assert (rejected.status_code, rejected.json()) == ( 400, - {"detail": "Push notification URL must use HTTPS"}, + { + "jsonrpc": "2.0", + "id": marker + "CreateTaskPushNotificationConfig", + "error": { + "code": -32000, + "message": "Push notification URL must use HTTPS", + "data": {"type": "invalid_request_error", "param": None, "code": "400"}, + }, + }, ), f"{params}: {rejected.text}" assert forwarded() == () diff --git a/tests/integration/cost_calculation/cases/azure.py b/tests/integration/cost_calculation/cases/azure.py index d9f17c12bda..6495199b8bb 100644 --- a/tests/integration/cost_calculation/cases/azure.py +++ b/tests/integration/cost_calculation/cases/azure.py @@ -1588,6 +1588,7 @@ AZURE_PINNED_GPT_5_4_MINI_STREAM: Final = CostTrackingTestCase( content_type="text/event-stream", frames=( 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1,"model":"$MODEL","choices":[{"index":0,"delta":{"role":"assistant","content":"ok"},"finish_reason":null}]}\n\n', + 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1,"model":"$MODEL","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n', 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1,"model":"$MODEL","choices":[],"usage":{"prompt_tokens":1840,"completion_tokens":412,"total_tokens":2252}}\n\n', "data: [DONE]\n\n", ), diff --git a/tests/integration/cost_calculation/cases/openai.py b/tests/integration/cost_calculation/cases/openai.py index bc7fb3486be..270bb6e105d 100644 --- a/tests/integration/cost_calculation/cases/openai.py +++ b/tests/integration/cost_calculation/cases/openai.py @@ -7952,6 +7952,7 @@ GPT_5_6_STREAM_USAGE_IN_EMPTY_CHOICES_CHUNK: Final = CostTrackingTestCase( frames=( 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1789788263,"model":"gpt-5.6","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}],"usage":null}', 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1789788263,"model":"gpt-5.6","choices":[{"index":0,"delta":{"content":"scripted answer"},"finish_reason":null}],"usage":null}', + 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1789788263,"model":"gpt-5.6","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":null}', 'data: {"id":"chatcmpl-$REQUEST_ID","object":"chat.completion.chunk","created":1789788263,"model":"gpt-5.6","choices":[],"usage":{"prompt_tokens":1840,"completion_tokens":412,"total_tokens":2252}}', "data: [DONE]", ), diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py index 8eeaab66553..210b754f626 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py @@ -199,14 +199,14 @@ def test_utils_token_counter_adds_title_and_context_once_per_document( @pytest.mark.parametrize("label", list(_UNREADABLE)) -def test_utils_token_counter_prices_unreadable_pdf_bytes_like_one_image( +def test_utils_token_counter_uses_the_same_bounded_fallback_for_unreadable_pdf_bytes( gateway: Gateway, provider: SharedProvider, label: str ) -> None: data: Final = _UNREADABLE[label] as_pdf: Final = _local(gateway, provider, [document(base64_source(data))]) - as_png: Final = _local(gateway, provider, [document(base64_source(data, media_type="image/png"))]) + fallback: Final = _local(gateway, provider, [document(base64_source(encoded(b"unreadable-document")))]) without: Final = _local(gateway, provider, []) - assert as_pdf == as_png, (as_pdf, as_png) + assert as_pdf == fallback, (as_pdf, fallback) assert 0 <= as_pdf - without < rendered_tokens((LETTER,)), (as_pdf, without) diff --git a/tests/integration/observability/test_prometheus_provider_auth_failure.py b/tests/integration/observability/test_prometheus_provider_auth_failure.py index 13d264ca6b0..de1f19bc523 100644 --- a/tests/integration/observability/test_prometheus_provider_auth_failure.py +++ b/tests/integration/observability/test_prometheus_provider_auth_failure.py @@ -995,7 +995,7 @@ def _gated(gateway: Gateway, cell: str, key: str | None) -> httpx.Response: def recorded_before(cell: str) -> list[tuple[str, str, str, float]]: - if cell in ("invalid-key", "expired-key"): + if cell in ("invalid-key", "missing-key", "expired-key", "blocked-key"): return [] return sorted([(FAILURES, "", "", 1.0), (TOTALS, "", "", 1.0), (LLM_FAILED, identity(cell), "", 1.0)]) @@ -1020,7 +1020,9 @@ def test_a_request_the_proxy_rejects_itself_never_marks_the_deployment_down_and_ assert response.status_code == (403 if cell == "no-model-access" else 401), response.text assert rig.upstream.drain() == () samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name(cell)) + assert values(samples, PROXY_FAILED, requested_model=name(cell)) == [1.0], recorded(samples) assert deployment_rows(samples, cell) == recorded_before(cell), recorded(samples, requested_model=name(cell)) + assert values(samples, STATE, model_id=identity(cell)) == [], recorded(samples) def test_a_key_rate_limit_rejected_at_the_proxy_is_still_counted_as_before(rig: Rig) -> None: diff --git a/tests/integration/providers/test_device_code_login_guard_boot.py b/tests/integration/providers/test_device_code_login_guard_boot.py index 8aba6841cf9..91f502179f3 100644 --- a/tests/integration/providers/test_device_code_login_guard_boot.py +++ b/tests/integration/providers/test_device_code_login_guard_boot.py @@ -28,7 +28,6 @@ from pydantic import JsonValue, TypeAdapter pytestmark: Final = [pytest.mark.timeout(900), pytest.mark.shares_worker_state] _STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") -_DROPPED: Final = re.compile(r"original model: (\S+), ignoring and continuing") _FIXED: Final[Mapping[str, str]] = MappingProxyType( {"chatgpt-fixed": "chatgpt/gpt-5.6-terra", "copilot-fixed": "github_copilot/gpt-5.4"} ) @@ -41,8 +40,6 @@ _CONTROL_CALLS_BEFORE_THE_RELOAD: Final = 1 _SHAPES: Final[Mapping[str, bool]] = MappingProxyType( {"docs-shapes": False, "docs-shapes-beside-a-universal-wildcard": True} ) -_NOT_FOUND: Final = "Invalid model name passed in model=" -_NO_HEALTHY: Final = "There are no healthy deployments for this model" _CONTROL: Final = "device-login-boot-control" _EXTRA: Final[Mapping[str, JsonValue]] = MappingProxyType({"num_retries": 0, "cache": {"no-cache": True}}) _CALL_ID: Final = "x-litellm-call-id" @@ -174,15 +171,11 @@ def _lists_login_models(listed: tuple[str, ...]) -> bool: ) -def _every_worker_lists( - owned: OwnedProxy, workers: tuple[int, ...], *, login_models: bool -) -> Mapping[int, tuple[str, ...]]: +def _every_worker_lists(owned: OwnedProxy, workers: tuple[int, ...]) -> Mapping[int, tuple[str, ...]]: def settled(by_worker: Mapping[int, tuple[str, ...]]) -> bool: if set(by_worker) != set(workers): return False - if login_models: - return all(_lists_login_models(listed) for listed in by_worker.values()) - return all(_login_names_listed(listed) == frozenset() for listed in by_worker.values()) + return all(_lists_login_models(listed) for listed in by_worker.values()) return eventually(lambda: _listing_by_worker(owned, workers), settled, seconds=90) @@ -212,16 +205,8 @@ def _reload(owned: OwnedProxy) -> None: assert response.status_code == 200, response.text -def _refusal_before_the_reload(name: str, universal: bool) -> str: - if name in _FIXED: - return _NO_HEALTHY - if not universal: - return _NOT_FOUND - return dl.CHATGPT_REFUSAL if name.startswith("chatgpt/") else dl.COPILOT_REFUSAL - - @pytest.mark.parametrize("universal", _SHAPES.values(), ids=_SHAPES.keys()) -def test_booting_without_tokens_drops_the_login_deployments_and_a_reload_after_mounting_brings_them_back( +def test_booting_without_tokens_lists_login_deployments_and_serves_only_after_tokens_are_mounted( gateway: Gateway, tmp_path: Path, universal: bool ) -> None: dirs: Final = _token_dirs(tmp_path) @@ -240,20 +225,18 @@ def test_booting_without_tokens_drops_the_login_deployments_and_a_reload_after_m workers=2, ) as owned: workers: Final = _workers(owned) - _every_worker_lists(owned, workers, login_models=False) + _every_worker_lists(owned, workers) boot_log: Final = owned.log.read_text() assert peers.auth_connections() == () assert dl.CHATGPT_USER_CODE not in boot_log and dl.COPILOT_USER_CODE not in boot_log, boot_log[-3000:] - dropped_models: Final = _DROPPED.findall(boot_log) - expected_drops: Final = frozenset((*_FIXED.values(), *_WILDCARDS)) - assert frozenset(dropped_models) == expected_drops, boot_log[-3000:] - assert all(dropped_models.count(model) >= len(workers) for model in expected_drops), boot_log[-3000:] with owned.gateway.scenario() as scenario: key: Final = scenario.key() for name in _LOGIN_NAMES: refused: Final = _chat(_base_url(owned), key, name, uuid.uuid4().hex) assert refused.status == 400, refused.text - assert _refusal_before_the_reload(name, universal) in refused.message(), refused.text + refusal: Final = dl.CHATGPT_REFUSAL if name.startswith("chatgpt") else dl.COPILOT_REFUSAL + assert refusal in refused.message(), refused.text + assert peers.api.drain() == () control: Final = _chat(_base_url(owned), key, _CONTROL, uuid.uuid4().hex) assert control.status == 200, control.text assert peers.auth_connections() == () @@ -261,7 +244,7 @@ def test_booting_without_tokens_drops_the_login_deployments_and_a_reload_after_m dirs.store_valid_tokens(peers.api.url) _reload(owned) - _every_worker_lists(owned, workers, login_models=True) + _every_worker_lists(owned, workers) for served_logins, name in enumerate(_LOGIN_NAMES, start=1): served: Final = _chat(_base_url(owned), key, name, uuid.uuid4().hex) assert served.status == 200, served.text diff --git a/tests/integration/providers/test_device_code_login_guard_wire.py b/tests/integration/providers/test_device_code_login_guard_wire.py index b4044674ed5..39e36fd95fa 100644 --- a/tests/integration/providers/test_device_code_login_guard_wire.py +++ b/tests/integration/providers/test_device_code_login_guard_wire.py @@ -746,10 +746,7 @@ def test_a_recent_device_code_request_does_not_hold_the_worker_for_the_rest_of_t assert scene.rig.peers.auth_connections() == () -_NOT_LIVE: Final = re.compile(r"saved to the database, but the model id\(s\) \['([0-9a-f-]+)'\] are not live") - - -def test_model_created_without_a_token_is_saved_not_live_and_serves_once_the_token_is_mounted(scene: _Scene) -> None: +def test_model_created_without_a_token_refuses_requests_and_serves_once_the_token_is_mounted(scene: _Scene) -> None: gateway: Final = scene.rig.owned.gateway name: Final = f"device-login-db-{uuid.uuid4().hex}" identity: Final = str(uuid.uuid4()) @@ -763,14 +760,18 @@ def test_model_created_without_a_token_is_saved_not_live_and_serves_once_the_tok }, ) try: - assert created.status_code == 500, created.text - not_live: Final = _NOT_LIVE.search(created.text) - assert not_live is not None and not_live.group(1) == identity, created.text + assert created.status_code == 200, created.text + assert created.json()["model_info"]["id"] == identity, created.text assert scene.rig.peers.auth_connections() == () - scene.store_valid_token("chatgpt") call: Final = _Call( scene.rig.base_url(), scene.key, name, "responses", False, "httpx", uuid.uuid4().hex, SimpleQueue() ) + refused: Final = _serve(call) + assert refused.status == 400, refused.text + assert dl.CHATGPT_REFUSAL in _error_message(refused.text), refused.text + assert scene.rig.peers.auth_connections() == () + assert scene.rig.peers.api.drain() == () + scene.store_valid_token("chatgpt") served: Final = eventually(lambda: _serve(call), lambda answer: answer.status == 200, seconds=90) assert set(rv.MARKER.findall(served.text)) == {call.marker}, served.text assert scene.rig.peers.auth_connections() == () diff --git a/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py b/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py index da091b1c12e..54a6df02edf 100644 --- a/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py +++ b/tests/integration/sdk/test_anthropic_json_tool_cut_stream_sync_sdk.py @@ -10,6 +10,7 @@ from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue import litellm +from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper _JSON_TOOL: Final = "json_tool_call" @@ -122,12 +123,10 @@ def _credentials(model: str) -> dict[str, str]: @pytest.mark.parametrize("model", [_ANTHROPIC_MODEL, _BEDROCK_INVOKE_MODEL], ids=["anthropic", "bedrock_invoke"]) @pytest.mark.parametrize( ("shape", "cut"), - [("plain", False), ("plain", True), ("wrapped", True)], - ids=["plain_complete", "plain_cut", "wrapped_cut"], + [("plain", False), ("wrapped", False), ("plain", True), ("wrapped", True)], + ids=["plain_complete", "wrapped_complete", "plain_cut", "wrapped_cut"], ) -def test_sync_stream_cut_before_the_json_tool_block_closes_still_yields_the_whole_answer( - model: str, shape: str, cut: bool -) -> None: +def test_sync_json_tool_stream_requires_a_provider_finish_reason(model: str, shape: str, cut: bool) -> None: marker_hex: Final = uuid.uuid4().hex with wire_server(_peer(marker_hex, shape, cut)) as wire: response: Final = litellm.completion( @@ -141,6 +140,11 @@ def test_sync_stream_cut_before_the_json_tool_block_closes_still_yields_the_whol **_credentials(model), ) assert isinstance(response, CustomStreamWrapper), type(response) + if cut: + with pytest.raises(MidStreamFallbackError, match="Stream ended before a finish_reason was received"): + tuple(response) + assert len(wire.drain()) == 1 + return text: Final = "".join(str(chunk.choices[0].delta.content or "") for chunk in tuple(response)) assert len(wire.drain()) == 1 assert json.loads(text) == _payload(marker_hex), text From 6be98d8c751f874aa038dbb556e8b12d647c9b16 Mon Sep 17 00:00:00 2001 From: zoveki <349028329+zoveki@users.noreply.github.com> Date: Mon, 2 Nov 2026 11:42:45 -0800 Subject: [PATCH 04/79] fix(bedrock)!: give Invoke and Mantle Claude the model's own default max_tokens (#49692) * fix(bedrock): resolve Invoke Claude default max_tokens from the base model id Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(bedrock): build Invoke Claude request params without mutating Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(bedrock): type tool-search normalizer return to clear unknown-type gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(bedrock): resolve Invoke and Mantle Claude default max_tokens from the deployment id, ARN base_model, and Mantle row * fix(bedrock): default Invoke Claude max_tokens from the deployment's own cap and honor AnthropicConfig Read the called deployment's model_info.max_output_tokens before the base model's cost row, skip the shared regional key the router fills from every deployment, and leave the default to litellm.AnthropicConfig(max_tokens) when it is set. * fix(bedrock): send the mapped default max_tokens on bedrock_mantle Claude chat * fix(bedrock): bound a global AnthropicConfig max_tokens by the deployment's own cap * test(bedrock): annotate the new Mantle default max_tokens tests with Final * refactor(bedrock): inject the AnthropicConfig max_tokens reader into the Invoke Claude config --------- Co-authored-by: mateo Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../anthropic_claude_transformation.py | 75 ++++++++++- .../bedrock/chat/mantle/transformation.py | 3 + .../test_bedrock_mantle_claude_chat_wire.py | 52 ++++---- ...mations_anthropic_claude_transformation.py | 118 ++++++++++++++++++ ...test_bedrock_chat_mantle_transformation.py | 19 +++ .../test_bedrock_mantle_transformation.py | 19 +++ 6 files changed, 259 insertions(+), 27 deletions(-) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py index e32df0bbdbb..e150791b09e 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py @@ -1,11 +1,13 @@ -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import TypeAdapter +import litellm from litellm.anthropic_beta_gated_fields import strip_beta_gated_request_body from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers +from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_mapping from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_anthropic_image_obj, @@ -25,7 +27,9 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation ) from litellm.llms.bedrock.common_utils import ( BEDROCK_INVOKE_TRACE_FIELD, + BedrockModelInfo, apply_bedrock_invoke_structured_output, + bedrock_capability_model, bedrock_chat_uses_native_structured_output, bedrock_cleaned_output_schema, bedrock_supports_tool_search, @@ -35,6 +39,7 @@ from litellm.llms.bedrock.common_utils import ( normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, requests_opt_in_native_output, + strip_bedrock_routing_prefix, strip_unsupported_bedrock_invoke_output_config_keys, tools_without_eager_input_streaming, ) @@ -59,6 +64,19 @@ _NON_DEFAULT_PARAMS: Final = TypeAdapter(Mapping[str, object]) _LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) +def _anthropic_config_max_tokens() -> int | None: + return litellm.AnthropicConfig.max_tokens + + +def _mapped_max_tokens(model: str) -> int | None: + row: Final = object_mapping(litellm.model_cost.get(model)) or {} # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # cost map rows are untyped JSON + match row.get("max_output_tokens"): + case int() as max_tokens: + return max_tokens + case _: + return None + + class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): """ Reference: @@ -79,9 +97,18 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): anthropic_version: str = "bedrock-2023-05-31" - def __init__(self, model_info: Mapping[str, object] | None = None) -> None: + def __init__( + self, + model_info: Mapping[str, object] | None = None, + read_configured_max_tokens: Callable[[], int | None] = _anthropic_config_max_tokens, + ) -> None: super().__init__() self._deployment_model_info: Final = model_info + self._read_configured_max_tokens: Final = read_configured_max_tokens + + @classmethod + def get_config(cls, *, model: str | None = None) -> dict[str, object]: # mutable-ok: overrides BaseConfig's dict + return {key: value for key, value in super().get_config(model=model).items() if key != "max_tokens"} @property def custom_llm_provider(self) -> str | None: @@ -241,6 +268,36 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): headers=headers, ) + def _default_max_tokens_models(self, model: str, litellm_params: Mapping[str, object]) -> tuple[str, ...]: + return ( + BedrockModelInfo.get_base_model( + bedrock_capability_model(strip_bedrock_routing_prefix(model), litellm_params.get("base_model")) + ), + ) + + def _default_max_tokens(self, model: str, litellm_params: Mapping[str, object]) -> int: + configured_max_tokens: Final = self._read_configured_max_tokens() + match deployment_model_info(litellm_params).get("max_output_tokens"): + case int() as deployment_max_tokens: + return ( + deployment_max_tokens + if configured_max_tokens is None + else min(deployment_max_tokens, configured_max_tokens) + ) + case _ if configured_max_tokens is not None: + return configured_max_tokens + case _: + return next( + ( + max_tokens + for max_tokens in map( + _mapped_max_tokens, self._default_max_tokens_models(model, litellm_params) + ) + if max_tokens is not None + ), + DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, + ) + def _build_bedrock_anthropic_request_base( self, model: str, @@ -257,13 +314,21 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): model=model, output_config=filtered_params["output_config"], ) - filtered_params = self._normalize_bedrock_tool_search_tools(filtered_params) + normalized_params: Final = self._normalize_bedrock_tool_search_tools(filtered_params) + default_max_tokens: Final = ( + None + if "max_tokens" in normalized_params + else self._default_max_tokens(model, _LITELLM_PARAMS.validate_python(litellm_params)) + ) + request_params: Final = ( + normalized_params if default_max_tokens is None else {**normalized_params, "max_tokens": default_max_tokens} + ) anthropic_request: Final = AnthropicConfig.transform_request( self, model=model, messages=messages, - optional_params=filtered_params, + optional_params=request_params, litellm_params=litellm_params, headers=headers, ) @@ -377,7 +442,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): "data": image_chunk["data"], } - def _normalize_bedrock_tool_search_tools(self, optional_params: dict) -> dict: + def _normalize_bedrock_tool_search_tools(self, optional_params: dict) -> dict[str, object]: """ Convert tool search entries to the format supported by the Bedrock Invoke API. """ diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index b5a352bab48..cbde6f918bc 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -54,6 +54,9 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): region=region, ) + def _default_max_tokens_models(self, model: str, litellm_params: Mapping[str, object]) -> tuple[str, ...]: + return (f"bedrock_mantle/{model}", *super()._default_max_tokens_models(model, litellm_params)) + def validate_environment( self, headers: dict, diff --git a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py index d4f9533fcdc..32b6c4f1690 100644 --- a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py @@ -70,10 +70,18 @@ def _chat_messages(marker: str) -> list[JsonValue]: return [{"role": "user", "content": _prompt(marker)}] -def _outbound(backend: str, messages: list[JsonValue], **fields: JsonValue) -> dict[str, JsonValue]: +def _default_max_tokens(gateway: Gateway, backend: str) -> int: + match gateway.get("/public/litellm_model_cost_map").get(f"bedrock_mantle/{backend}"): + case {"max_output_tokens": int() as mapped}: + return mapped + case _: + return DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS + + +def _outbound(gateway: Gateway, backend: str, messages: list[JsonValue], **fields: JsonValue) -> dict[str, JsonValue]: return { "messages": messages, - "max_tokens": DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, + "max_tokens": _default_max_tokens(gateway, backend), "anthropic_version": _BRIDGE_VERSION, "model": backend, **fields, @@ -290,7 +298,7 @@ def test_chat_completion_on_a_mantle_claude_id_posts_native_messages_and_spends_ gateway: Gateway, backend: str ) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(backend, [_user_turn(marker)]) + expected: Final = _outbound(gateway, backend, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker, backend))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{backend}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -314,7 +322,7 @@ def test_chat_stream_on_a_mantle_claude_id_relays_openai_chunks_from_the_native_ gateway: Gateway, stream_options: dict[str, JsonValue], usage_chunks: list[tuple[JsonValue, ...]] ) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], stream=True) with wire_server(_bearer_peer(expected, _text_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -355,7 +363,7 @@ def test_chat_request_fields_are_translated_to_the_native_messages_shape( outbound_fields: dict[str, JsonValue], ) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], **outbound_fields) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], **outbound_fields) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -377,7 +385,7 @@ def test_chat_tools_reach_mantle_as_input_schema_and_the_tool_use_block_returns_ marker: Final = uuid4().hex backend: Final = _mantle_claude_where(gateway, "supports_forced_tool_use", disabled=False) expected: Final = _outbound( - backend, [_user_turn(marker)], tools=[_WEATHER_TOOL_OUTBOUND], tool_choice=outbound_choice + gateway, backend, [_user_turn(marker)], tools=[_WEATHER_TOOL_OUTBOUND], tool_choice=outbound_choice ) reply: Final = _tool_reply(marker, backend, "get_weather") with wire_server(_bearer_peer(expected, reply)) as wire, gateway.scenario() as scenario: @@ -407,7 +415,7 @@ def test_chat_tools_reach_mantle_as_input_schema_and_the_tool_use_block_returns_ def test_chat_stream_relays_a_native_tool_use_block_as_tool_call_chunks(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], tools=[_WEATHER_TOOL_OUTBOUND], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], tools=[_WEATHER_TOOL_OUTBOUND], stream=True) with wire_server(_bearer_peer(expected, _tool_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -439,7 +447,7 @@ def test_chat_tool_result_turn_reaches_mantle_as_tool_use_and_tool_result_blocks marker: Final = uuid4().hex tool_call: Final = f"toolu_{marker}" expected: Final = _outbound( - _HAIKU, + gateway, _HAIKU, [ _user_turn(marker), { @@ -484,7 +492,7 @@ def test_chat_json_schema_response_format_is_sent_as_a_forced_tool_and_returned_ ) -> None: marker: Final = uuid4().hex expected: Final = _outbound( - _HAIKU, + gateway, _HAIKU, [_user_turn(marker)], tools=[{"name": "json_tool_call", "input_schema": _CITY_SCHEMA}], tool_choice={"name": "json_tool_call", "type": "tool"}, @@ -537,7 +545,7 @@ def _sigv4_peer(expected: Mapping[str, JsonValue], reply: Reply) -> Callable[[Re def test_chat_completion_signs_the_native_messages_request_with_the_deployment_aws_keys(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_sigv4_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model( model=f"bedrock_mantle/{_HAIKU}", @@ -556,7 +564,7 @@ def test_chat_completion_signs_the_native_messages_request_with_the_deployment_a def test_region_prefix_in_a_mantle_claude_id_is_routing_only_and_keeps_the_mantle_price(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/us-east-2/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -577,7 +585,7 @@ def test_region_prefix_in_a_mantle_claude_id_is_routing_only_and_keeps_the_mantl def test_openai_sdk_chat_completion_on_a_mantle_claude_id(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) with OpenAI( @@ -606,7 +614,7 @@ def test_openai_sdk_chat_completion_on_a_mantle_claude_id(gateway: Gateway) -> N def test_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], stream=True) with wire_server(_bearer_peer(expected, _text_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) with OpenAI( @@ -645,7 +653,7 @@ def test_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gateway) -> None: async def test_async_openai_sdk_chat_completion_on_a_mantle_claude_id(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) async with AsyncOpenAI( @@ -674,7 +682,7 @@ async def test_async_openai_sdk_chat_completion_on_a_mantle_claude_id(gateway: G async def test_async_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], stream=True) with wire_server(_bearer_peer(expected, _text_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) async with AsyncOpenAI( @@ -714,7 +722,7 @@ def test_identical_chat_request_is_served_from_the_response_cache_and_logged_as_ gateway: Gateway, ) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) request: Final[dict[str, JsonValue]] = {"model": model, "messages": _chat_messages(marker)} @@ -802,7 +810,7 @@ def test_open_weight_mantle_id_keeps_the_openai_compatible_chat_route_and_its_sa ) def test_api_base_with_a_route_suffix_still_posts_to_the_native_messages_path(gateway: Gateway, suffix: str) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url + suffix, api_key=_API_KEY) response: Final = gateway.request( @@ -814,7 +822,7 @@ def test_api_base_with_a_route_suffix_still_posts_to_the_native_messages_path(ga def test_responses_api_on_a_mantle_claude_id_bridges_to_native_messages(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request("POST", "/v1/responses", {"model": model, "input": _prompt(marker)}) @@ -839,7 +847,7 @@ def test_responses_api_on_a_mantle_claude_id_bridges_to_native_messages(gateway: def test_responses_api_stream_on_a_mantle_claude_id_bridges_to_the_native_stream(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], stream=True) with wire_server(_bearer_peer(expected, _text_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -996,7 +1004,7 @@ def test_unsupported_claude_param_is_dropped_under_drop_params( ) -> None: marker: Final = uuid4().hex backend: Final = _backend_without(gateway, flag) - expected: Final = _outbound(backend, [_user_turn(marker)], **outbound_fields) + expected: Final = _outbound(gateway, backend, [_user_turn(marker)], **outbound_fields) with wire_server(_bearer_peer(expected, _text_reply(marker, backend))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{backend}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -1114,7 +1122,7 @@ def _wildcard_prefix(gateway: Gateway, scenario: Scenario, api_base: str) -> str def test_uppercase_claude_id_through_a_mantle_wildcard_is_sent_to_native_messages(gateway: Gateway) -> None: marker: Final = uuid4().hex backend: Final = _HAIKU.upper() - expected: Final = _outbound(backend, [_user_turn(marker)]) + expected: Final = _outbound(gateway, backend, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker, backend))) as wire, gateway.scenario() as scenario: model: Final = f"{_wildcard_prefix(gateway, scenario, wire.url)}/{backend}" response: Final = gateway.request( @@ -1129,7 +1137,7 @@ def test_five_kilobyte_claude_id_through_a_mantle_wildcard_reaches_mantle_and_it ) -> None: marker: Final = uuid4().hex backend: Final = f"{_HAIKU}-{'x' * 5000}" - expected: Final = _outbound(backend, [_user_turn(marker)]) + expected: Final = _outbound(gateway, backend, [_user_turn(marker)]) unknown: Final = _error_reply(404, "not_found_error", "model: synthetic unknown model") with wire_server(_bearer_peer(expected, unknown)) as wire, gateway.scenario() as scenario: model: Final = f"{_wildcard_prefix(gateway, scenario, wire.url)}/{backend}" diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py index d40abb4860b..4850a402103 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py @@ -14,6 +14,7 @@ from pydantic import BaseModel, ConfigDict # Ensure the project root is on the import path so `litellm` can be imported when # tests are executed from any working directory. import litellm +from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) @@ -2010,6 +2011,123 @@ def test_invoke_uses_opt_in_native_output_is_false_when_native_output_did_not_co ) +BEDROCK_HAIKU_45_BASE_ID: Final = "anthropic.claude-haiku-4-5-20251001-v1:0" + + +@pytest.mark.parametrize( + "model", + [ + f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + f"invoke/{BEDROCK_HAIKU_45_BASE_ID}", + ], +) +def test_bedrock_invoke_claude_default_max_tokens_comes_from_base_model(local_model_cost_map, model: str): + mapped_cap: Final = litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) + assert mapped_cap != DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "the model's mapped cap must differ from the fallback" + + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == mapped_cap, f"default max_tokens should be the base model's cap, got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_explicit_max_tokens_is_kept(local_model_cost_map): + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model=f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + messages=[{"role": "user", "content": "hi"}], + optional_params={"max_tokens": 123}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == 123, f"explicit max_tokens must not be overridden, got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_default_max_tokens_takes_the_deployments_own_cap(local_model_cost_map): + deployment_cap: Final = 8192 + assert litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) != deployment_cap + + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model=f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={"model_info": {"max_output_tokens": deployment_cap}}, + headers={}, + ) + + assert body["max_tokens"] == deployment_cap, f"the deployment's own cap should win, got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_default_max_tokens_ignores_a_sibling_deployments_cap(local_model_cost_map): + regional_id: Final = f"us.{BEDROCK_HAIKU_45_BASE_ID}" + mapped_cap: Final = litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) + litellm.register_model(model_cost={regional_id: {"max_output_tokens": 50}}, persist_across_reloads=False) + + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model=f"invoke/{regional_id}", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == mapped_cap, f"a sibling's cap on the shared id leaked, got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_default_max_tokens_yields_to_anthropic_config(local_model_cost_map): + configured_max_tokens: Final = 1000 + assert litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) != configured_max_tokens + + body: Final = AmazonAnthropicClaudeConfig( + read_configured_max_tokens=lambda: configured_max_tokens + ).transform_request( + model=f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == configured_max_tokens, f"AnthropicConfig(max_tokens) was ignored, got {body['max_tokens']}" + + +@pytest.mark.parametrize(("configured_max_tokens", "deployment_cap"), [(100000, 50), (1000, 8192)]) +def test_bedrock_invoke_claude_default_max_tokens_is_the_lower_of_anthropic_config_and_the_deployments_cap( + local_model_cost_map, configured_max_tokens: int, deployment_cap: int +): + body: Final = AmazonAnthropicClaudeConfig( + read_configured_max_tokens=lambda: configured_max_tokens + ).transform_request( + model=f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={"model_info": {"max_output_tokens": deployment_cap}}, + headers={}, + ) + + assert body["max_tokens"] == min(configured_max_tokens, deployment_cap), f"got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_default_max_tokens_reads_an_inference_profile_arns_base_model(local_model_cost_map): + mapped_cap: Final = litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) + assert mapped_cap != DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "the model's mapped cap must differ from the fallback" + + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model="invoke/anthropic/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={"base_model": BEDROCK_HAIKU_45_BASE_ID}, + headers={}, + ) + + assert body["max_tokens"] == mapped_cap, f"an ARN should take its base_model's cap, got {body['max_tokens']}" + + @pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") @pytest.mark.parametrize("drop_params", [True, False]) def test_bedrock_chat_invoke_adds_the_binding_beta_so_block_binding_is_kept( diff --git a/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py b/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py index 154be9f0370..56d6b039788 100644 --- a/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py +++ b/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py @@ -2,6 +2,7 @@ import base64 import json import uuid from types import SimpleNamespace +from typing import Final import httpx import pytest @@ -243,3 +244,21 @@ async def test_mantle_claude_async_completion_streams_a_wrapped_json_schema_answ chunks = tuple([chunk async for chunk in stream]) assert _streamed_answer(chunks) == BARE_STREAMED_ANSWER + + +def test_mantle_claude_default_max_tokens_comes_from_the_mantle_row(local_model_cost_map): + from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS + from litellm.llms.bedrock.chat.mantle.transformation import AmazonMantleConfig + + mapped_cap: Final = litellm.get_max_tokens("bedrock_mantle/anthropic.claude-haiku-4-5") + assert mapped_cap != DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "the model's mapped cap must differ from the fallback" + + body: Final = AmazonMantleConfig().transform_request( + model="mantle/anthropic.claude-haiku-4-5", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == mapped_cap, f"Mantle should send its mapped cap, got {body['max_tokens']}" diff --git a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index f252d9886d3..6005acffc0a 100644 --- a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -1096,6 +1096,25 @@ class TestBedrockMantleClaudeChatRoute: ) assert response._hidden_params["response_cost"] != pytest.approx(_row_cost("anthropic.claude-opus-5-5", 10, 5)) + def test_claude_completion_without_max_tokens_sends_the_mantle_rows_output_limit(self, monkeypatch, local_cost_map): + from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + mapped_cap: Final = litellm.model_cost["bedrock_mantle/anthropic.claude-haiku-4-5"]["max_output_tokens"] + assert mapped_cap != DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "the model's mapped cap must differ from the fallback" + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "mantle-key") + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + handler: Final = Mock(side_effect=_anthropic_message) + + litellm.completion( + model="bedrock_mantle/anthropic.claude-haiku-4-5", + messages=[{"role": "user", "content": "hello"}], + aws_region_name="us-east-1", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))), + ) + + assert json.loads(handler.call_args.args[0].content)["max_tokens"] == mapped_cap + def test_claude_streaming_completion_uses_native_messages_endpoint(self, monkeypatch, local_cost_map): from litellm.llms.custom_httpx.http_handler import HTTPHandler From a4de7201449530811958862ee38e450d3b5d6a91 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Mon, 2 Nov 2026 12:37:35 -0800 Subject: [PATCH 05/79] fix(count_tokens): count grok-4-1, grok-4.3 and DeepSeek V4.1 tool requests the way their hosts bill them (#49074) * fix(count_tokens): count grok-4-1, grok-4.3 and DeepSeek V4.1 tool requests like their chat templates Render grok-4-1, grok-4.3 and DeepSeek V4.1 tool schemas the way each chat template does, and add a per-deployment token_count_tool_framing next to token_count_framing for the hidden tool preamble a host bills. The Azure AI Foundry grok-4-1, xAI grok-4.3 and Fireworks DeepSeek V4.1 cost map rows carry the measured framing, overridable in model_info * fix(count_tokens): count Azure grok-4-1 Responses requests with their own framing and use the tool framing in place of the base framing * fix(count_tokens): keep operator framing on tool and Responses counts, frame bridged Responses like chat, follow base_model for tool rendering, count grok tool_choice * fix(token_counter): count Azure Responses strict default, skip false strict and built-in-only tools, per-kind operator framing * fix(token_counter): count an explicit null strict like false on Azure Responses * fix(count_tokens): follow the chat completions bridge flag and Azure's strict default on azure/ deployments Keep Azure grok-4-1 chat tool counts at or over the bill by raising its tool framing to 302 * fix(token_counter): keep strict false for non-grok templates, coerce the bridge flag, frame DeepSeek API counts * fix(proxy): frame Responses counts for the api_base router defaults send inference to * fix(proxy): keep the api_base helper's name so main's /v1/messages count caller resolves --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../crates/model-catalog/src/model_info.rs | 3 + .../prompt_templates/chat_template_tools.py | 26 +- litellm/litellm_core_utils/token_counter.py | 25 +- ...odel_prices_and_context_window_backup.json | 16 +- .../common_utils/token_count_settings.py | 37 +- litellm/proxy/proxy_server.py | 108 +++- litellm/types/utils.py | 13 +- litellm/utils.py | 3 + model_prices_and_context_window.json | 16 +- .../test_chat_template_tool_count_wire.py | 10 +- ...ls_prompt_templates_chat_template_tools.py | 17 +- .../litellm_core_utils/test_token_counter.py | 64 ++- .../common_utils/test_token_count_settings.py | 101 ++++ .../test_model_management_endpoints.py | 3 + tests/unit/proxy/test_proxy_token_counter.py | 491 ++++++++++++++++++ .../unit/test_router_model_cost_isolation.py | 5 +- tests/unit/test_utils.py | 3 + 17 files changed, 904 insertions(+), 37 deletions(-) diff --git a/litellm-rust/crates/model-catalog/src/model_info.rs b/litellm-rust/crates/model-catalog/src/model_info.rs index 8f7549c1c6d..141702bf857 100644 --- a/litellm-rust/crates/model-catalog/src/model_info.rs +++ b/litellm-rust/crates/model-catalog/src/model_info.rs @@ -186,6 +186,9 @@ pub struct ModelInfoFields { pub supports_image_size: Option, pub supported_audio_formats: Option, pub token_count_framing: Option, + pub token_count_tool_framing: Option, + pub token_count_responses_framing: Option, + pub token_count_responses_tool_framing: Option, pub token_count_model: Option, pub token_count_tokenizer: Option, pub vertex_ai_audio_api: Option, diff --git a/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py b/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py index 60134456577..49ffae03c40 100644 --- a/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py +++ b/litellm/litellm_core_utils/prompt_templates/chat_template_tools.py @@ -10,12 +10,18 @@ class ChatTemplateToolSection: preamble_tokens: int implicit_system_turn_tokens: int render_tool_schemas: Callable[[Sequence[Mapping[str, object]]], str] + counts_tool_choice: bool = False + renders_strict_only_when_true: bool = False -def _deepseek_v4_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: +def _json_line_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: return "\n".join(json.dumps(function, ensure_ascii=False) for function in functions) +def _compact_json_line_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: + return "\n".join(json.dumps(function, ensure_ascii=False, separators=(",", ":")) for function in functions) + + def _qwen3_5_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: return "".join( "\n" + json.dumps({"type": "function", "function": function}, ensure_ascii=False) for function in functions @@ -23,13 +29,29 @@ def _qwen3_5_tool_schemas(functions: Sequence[Mapping[str, object]]) -> str: _TOOL_SECTIONS: Final = ( + ( + re.compile(r"deepseek-v4(?:[-._]|p)1(?!\d)", re.IGNORECASE), + ChatTemplateToolSection(211, 0, _json_line_tool_schemas), + ), ( re.compile( r"deepseek-v4(?![.\d]|p\d)|^(?:deepseek/)?deepseek-(?:chat|reasoner|flash)(?::free)?$", re.IGNORECASE ), - ChatTemplateToolSection(222, 0, _deepseek_v4_tool_schemas), + ChatTemplateToolSection(222, 0, _json_line_tool_schemas), ), (re.compile(r"qwen3[._]?5(?!\d)", re.IGNORECASE), ChatTemplateToolSection(196, 5, _qwen3_5_tool_schemas)), + ( + re.compile(r"grok-4[-.]1(?![.\d])", re.IGNORECASE), + ChatTemplateToolSection( + 0, 0, _compact_json_line_tool_schemas, counts_tool_choice=True, renders_strict_only_when_true=True + ), + ), + ( + re.compile(r"grok-4[-.]3(?![.\d])", re.IGNORECASE), + ChatTemplateToolSection( + 0, 0, _json_line_tool_schemas, counts_tool_choice=True, renders_strict_only_when_true=True + ), + ), ) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 323af72940a..df8038e756c 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -990,7 +990,11 @@ def _count_extra( tool_section (Optional[ChatTemplateToolSection]): How the model's chat template renders tools, if known. """ reply_primer_tokens: Final = 3 - tool_choice_tokens: Final = 0 if tool_section is not None else _count_tool_choice(count_function, tool_choice) + tool_choice_tokens: Final = ( + _count_tool_choice(count_function, tool_choice) + if tool_section is None or tool_section.counts_tool_choice + else 0 + ) return ( reply_primer_tokens + _count_tool_definitions(count_function, tools, includes_system_message, tool_section) @@ -1029,14 +1033,29 @@ def _count_tool_choice( return 0 +def _named_function_definitions(tools: Sequence[object]) -> Iterable[Mapping[str, object]]: + return (function for function in _function_definitions(_tools_sent_as_functions(tools)) if function.get("name")) + + +def sends_function_tools(tools: Sequence[object] | None) -> bool: + return tools is not None and any(_named_function_definitions(tools)) + + +def _without_false_strict(function: Mapping[str, object]) -> Mapping[str, object]: + return {key: value for key, value in function.items() if key != "strict" or value} + + def _count_chat_template_tools( count_function: TokenCounterFunction, tools: Sequence[object], includes_system_message: bool, tool_section: ChatTemplateToolSection, ) -> int: - functions: Final = tuple( - function for function in _function_definitions(_tools_sent_as_functions(tools)) if function.get("name") + named_functions: Final = tuple(_named_function_definitions(tools)) + functions: Final = ( + tuple(map(_without_false_strict, named_functions)) + if tool_section.renders_strict_only_when_true + else named_functions ) if not functions: return 0 diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 02bf7dedcef..e6f65d6b416 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -13478,6 +13478,10 @@ "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", + "token_count_framing": 149, + "token_count_tool_framing": 302, + "token_count_responses_framing": 145, + "token_count_responses_tool_framing": 100, "max_input_tokens": 128000, "max_output_tokens": 128000, "mode": "chat", @@ -54457,6 +54461,8 @@ "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", + "token_count_framing": 185, + "token_count_tool_framing": 214, "max_input_tokens": 1000000, "max_output_tokens": 1000000, "mode": "chat", @@ -62640,7 +62646,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 23, + "token_count_tool_framing": 12 }, "deepseek/deepseek-v4-flash-vision-exp": { "cache_creation_input_token_cost": 0.0, @@ -62746,7 +62754,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "token_count_framing": 76, + "token_count_tool_framing": 75 }, "tencent/deepseek-v4-pro": { "cache_creation_input_token_cost": 0.0, @@ -64032,6 +64042,7 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_priority": 3.75e-07, "litellm_provider": "fireworks_ai", + "token_count_framing": 23, "max_input_tokens": 1048576, "max_output_tokens": 393216, "mode": "chat", @@ -64133,6 +64144,7 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_priority": 3.75e-07, "litellm_provider": "fireworks_ai", + "token_count_framing": 23, "max_input_tokens": 1048576, "max_output_tokens": 393216, "mode": "chat", diff --git a/litellm/proxy/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py index dd04e531f2d..4e3059cfe26 100644 --- a/litellm/proxy/common_utils/token_count_settings.py +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -56,6 +56,22 @@ class TokenCountSettings: framing: int | None count_model: TokenCountModel | None tokenizer: TokenCountTokenizer | None = None + tool_framing: int | None = None + responses_framing: int | None = None + responses_tool_framing: int | None = None + + @property + def frames_counts(self) -> bool: + return any( + framing is not None + for framing in (self.framing, self.tool_framing, self.responses_framing, self.responses_tool_framing) + ) + + def hidden_prompt_tokens(self, carries_tools: bool, responses_api: bool) -> int: + chat: Final = self.tool_framing if carries_tools else None + responses: Final = self.responses_tool_framing if carries_tools else self.responses_framing + narrower: Final = responses if responses_api else chat + return (narrower if narrower is not None else self.framing) or 0 @dataclass(frozen=True, slots=True) @@ -105,6 +121,15 @@ def _framing(value: object) -> int | None: return None +def _first_framing(source: Mapping[str, object], keys: tuple[str, ...]) -> int | None: + return next((framing for framing in map(_framing, map(source.get, keys)) if framing is not None), None) + + +def _configured_framing(configured: Mapping[str, object], default: Mapping[str, object], *keys: str) -> int | None: + operator_sets_one: Final = _first_framing(configured, (*keys, "token_count_framing")) is not None + return _first_framing(configured if operator_sets_one else default, keys) + + def _non_blank(value: object) -> str | None: return value if isinstance(value, str) and value.strip() else None @@ -175,13 +200,21 @@ def _tokenizer(value: object) -> TokenCountTokenizer | None: def resolve_token_count_settings( configured: Mapping[str, object], default: Mapping[str, object], deployment_region: Callable[[], str] ) -> TokenCountSettings: - framing: Final = _framing(configured.get("token_count_framing")) return TokenCountSettings( - framing=framing if framing is not None else _framing(default.get("token_count_framing")), + framing=_configured_framing(configured, default, "token_count_framing"), count_model=_resolved_count_model( configured.get("token_count_model"), default.get("token_count_model"), deployment_region ), tokenizer=_tokenizer(default.get("token_count_tokenizer")), + tool_framing=_configured_framing(configured, default, "token_count_tool_framing"), + responses_framing=_configured_framing(configured, default, "token_count_responses_framing"), + responses_tool_framing=_configured_framing( + configured, + default, + "token_count_responses_tool_framing", + "token_count_responses_framing", + "token_count_tool_framing", + ), ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fb08272e687..0e35560f30f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -107,6 +107,7 @@ from litellm.litellm_core_utils.litellm_logging import ( initialize_sentry, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker +from litellm.litellm_core_utils.prompt_templates.chat_template_tools import chat_template_tool_section from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias from litellm.litellm_core_utils.routing_groups import parse_routing_groups from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -115,6 +116,7 @@ from litellm.litellm_core_utils.token_counter import ( UnsupportedContentTypeError, deployment_openai_model, offload_token_count, + sends_function_tools, ) from litellm.llms.bedrock.count_tokens.kimi_k3 import ( HUGGING_FACE_SOURCE, @@ -1007,6 +1009,7 @@ from litellm.proxy.video_endpoints.endpoints import router as video_router from litellm.repositories.base_repository import SupportsModelDump from litellm.repositories.credentials_repository import CredentialsRepository from litellm.repositories.prisma_protocols import TableActions +from litellm.responses.main import provider_stores_the_conversation from litellm.responses.streaming_iterator import drain_logging_tasks from litellm.router import ( AssistantsTypedDict, @@ -15037,7 +15040,7 @@ async def _provider_token_count( framing_can_stand_in: Final = (count_model is None or not count_model.configured) and ( failure.cannot_count_model or (pinned_count_model is not None and failure.denied_in_region) ) - if not (framing_can_stand_in and settings.framing is not None and local_estimate_allowed): + if not (framing_can_stand_in and settings.frames_counts and local_estimate_allowed): raise failure.to_proxy_exception() if pinned_count_model is not None: verbose_proxy_logger.warning( @@ -15237,6 +15240,7 @@ async def _tokenizer_or_provider_error( async def _local_token_estimate( model: str, + template_model: str, prompt: str | None, messages: Sequence[AllMessageValues] | None, tools: list[ChatCompletionToolParam] | None, @@ -15249,7 +15253,7 @@ async def _local_token_estimate( else await _tokenizer_or_provider_error(model, custom_tokenizer, provider_failure) ) count: Final = offload_token_count(litellm.token_counter)( - model=model, + model=template_model, text=prompt, messages=messages, tools=tools, @@ -15315,10 +15319,15 @@ def _deployment_aws_region(deployment: _TokenCountDeployment) -> str: ) +def _deployment_base_model(deployment: _TokenCountDeployment) -> str | None: + return _non_empty_str(deployment.get("model_info", {}).get("base_model")) or _non_empty_str( + deployment.get("litellm_params", {}).get("base_model") + ) + + def _token_count_settings_model(deployment: _TokenCountDeployment, request_model: str) -> tuple[str, str | None]: - model_info: Final = deployment.get("model_info", {}) params: Final = deployment.get("litellm_params", {}) - base_model: Final = _non_empty_str(model_info.get("base_model")) or _non_empty_str(params.get("base_model")) + base_model: Final = _deployment_base_model(deployment) if base_model is not None: return base_model, None model: Final = _non_empty_str(params.get("model")) @@ -15470,12 +15479,15 @@ def _responses_routing_kwargs(responses_request: Mapping[str, object] | None) -> return {"previous_response_id": previous_response_id} if isinstance(previous_response_id, str) else {} +def _inference_params( + litellm_params: Mapping[str, object], responses_request: Mapping[str, object] +) -> Mapping[str, object]: + return {**litellm_params, **_router_default_litellm_params(llm_router), **responses_request} + + def _count_drop_policy(deployment: Mapping[str, object] | None, responses_request: Mapping[str, object]) -> DropPolicy: litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params")) or {} - router_defaults: Final[Mapping[str, object]] = ( - {} if llm_router is None else _router_default_litellm_params(llm_router) - ) - return drop_policy({**litellm_params, **router_defaults, **responses_request}) + return drop_policy(_inference_params(litellm_params, responses_request)) def _native_responses_count( @@ -15498,8 +15510,54 @@ def _native_responses_count( ) -def _deployment_api_base(deployment: _TokenCountDeployment) -> str | None: - return _non_empty_str(deployment.get("litellm_params", {}).get("api_base")) +def _deployment_api_base( + deployment: _TokenCountDeployment, responses_request: Mapping[str, object] | None = None +) -> str | None: + return _non_empty_str( + _inference_params(deployment.get("litellm_params", {}), responses_request or {}).get("api_base") + ) + + +def _deployment_provider( + deployment: _TokenCountDeployment, request_model: str, responses_request: Mapping[str, object] | None +) -> str | None: + params: Final = deployment.get("litellm_params", {}) + try: + return litellm.get_llm_provider( + model=_non_empty_str(params.get("model")) or request_model, + custom_llm_provider=_non_empty_str(params.get("custom_llm_provider")), + api_base=_deployment_api_base(deployment, responses_request), + )[1] + except litellm.BadRequestError: + return None + + +def _serves_native_responses( + deployment: _TokenCountDeployment, request_model: str, responses_request: Mapping[str, object] +) -> bool: + params: Final = deployment.get("litellm_params", {}) + inference: Final = _inference_params(params, responses_request) + try: + return provider_stores_the_conversation( + _non_empty_str(params.get("model")) or request_model, + _non_empty_str(params.get("custom_llm_provider")), + { + "use_chat_completions_api": bool(inference.get("use_chat_completions_api")), + "model_info": dict(deployment.get("model_info", {})), + "api_base": _non_empty_str(inference.get("api_base")), + }, + ) + except litellm.BadRequestError: + return False + + +_AZURE_PROVIDERS: Final = frozenset((litellm.LlmProviders.AZURE.value, litellm.LlmProviders.AZURE_AI.value)) + + +def _with_strict_by_default(tools: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, object], ...]: + return tuple( + {**tool, "strict": True} if tool.get("type") == "function" and "strict" not in tool else tool for tool in tools + ) def _system_message(system: object) -> ChatCompletionSystemMessage | None: @@ -15704,7 +15762,7 @@ async def _count_request_tokens( else await chat_count_request( model_to_use, custom_llm_provider, - _deployment_api_base(deployment_view), + _deployment_api_base(deployment_view, routed_request), routed_request, _count_drop_policy(routed_deployment, routed_request), ) @@ -15802,10 +15860,31 @@ async def _count_request_tokens( counted_messages: Final = ( typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages) ) - counted_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats + native_responses: Final = responses_request is not None and _serves_native_responses( + deployment_view, request.model, responses_request + ) + sent_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None ) - framing: Final = settings.framing if counted_messages is not None and settings.framing is not None else 0 + openai_model: Final = deployment_openai_model(deployment_view) + template_model: Final = openai_model or _deployment_base_model(deployment_view) or model_to_use + tool_section: Final = chat_template_tool_section(template_model) + counted_tools: Final = ( + cast(list[ChatCompletionToolParam], list(_with_strict_by_default(sent_tools))) # cast-ok: strict added + if sent_tools is not None + and native_responses + and _deployment_provider(deployment_view, request.model, responses_request) in _AZURE_PROVIDERS + and tool_section is not None + and tool_section.renders_strict_only_when_true + else sent_tools + ) + framing: Final = ( + 0 + if counted_messages is None + else settings.hidden_prompt_tokens( + carries_tools=sends_function_tools(counted_tools), responses_api=native_responses + ) + ) kimi_k3_messages: Final = counted_messages if settings.tokenizer == "kimi_k3" else None exact: Final = ( await _kimi_k3_count( @@ -15829,7 +15908,8 @@ async def _count_request_tokens( count_source="local_estimate", ) estimate: Final = await _local_token_estimate( - model=deployment_openai_model(deployment_view) or model_to_use, + model=openai_model or model_to_use, + template_model=template_model, prompt=prompt if counted_messages is None else None, messages=counted_messages, tools=counted_tools, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index da1d99957a7..bf69445fdd1 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -246,6 +246,9 @@ class ProviderSpecificModelInfo(TypedDict, total=False): bedrock_converse_supports_strict_tools: bool | None supports_regex_lookaround: ReadOnly[bool | None] token_count_framing: ReadOnly[int | None] + token_count_tool_framing: ReadOnly[int | None] + token_count_responses_framing: ReadOnly[int | None] + token_count_responses_tool_framing: ReadOnly[int | None] token_count_model: ReadOnly[TokenCountModelInfo | None] token_count_tokenizer: ReadOnly[TokenCountTokenizer | None] @@ -4236,7 +4239,15 @@ DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset( } ) -DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS: Final[frozenset[str]] = frozenset({"token_count_framing", "token_count_model"}) +DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS: Final[frozenset[str]] = frozenset( + { + "token_count_framing", + "token_count_tool_framing", + "token_count_responses_framing", + "token_count_responses_tool_framing", + "token_count_model", + } +) SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = ( frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__) diff --git a/litellm/utils.py b/litellm/utils.py index b0826f84006..5f40b64e79f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7398,6 +7398,9 @@ def get_model_info_helper( supports_image_size=_model_info.get("supports_image_size", None), supported_audio_formats=_model_info.get("supported_audio_formats", None), token_count_framing=_model_info.get("token_count_framing", None), + token_count_tool_framing=_model_info.get("token_count_tool_framing", None), + token_count_responses_framing=_model_info.get("token_count_responses_framing", None), + token_count_responses_tool_framing=_model_info.get("token_count_responses_tool_framing", None), token_count_model=_model_info.get("token_count_model", None), token_count_tokenizer=_model_info.get("token_count_tokenizer", None), vertex_ai_audio_api=_model_info.get("vertex_ai_audio_api", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 02bf7dedcef..e6f65d6b416 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -13478,6 +13478,10 @@ "input_cost_per_token": 2e-07, "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", + "token_count_framing": 149, + "token_count_tool_framing": 302, + "token_count_responses_framing": 145, + "token_count_responses_tool_framing": 100, "max_input_tokens": 128000, "max_output_tokens": 128000, "mode": "chat", @@ -54457,6 +54461,8 @@ "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "xai", + "token_count_framing": 185, + "token_count_tool_framing": 214, "max_input_tokens": 1000000, "max_output_tokens": 1000000, "mode": "chat", @@ -62640,7 +62646,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "token_count_framing": 23, + "token_count_tool_framing": 12 }, "deepseek/deepseek-v4-flash-vision-exp": { "cache_creation_input_token_cost": 0.0, @@ -62746,7 +62754,9 @@ "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "token_count_framing": 76, + "token_count_tool_framing": 75 }, "tencent/deepseek-v4-pro": { "cache_creation_input_token_cost": 0.0, @@ -64032,6 +64042,7 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_priority": 3.75e-07, "litellm_provider": "fireworks_ai", + "token_count_framing": 23, "max_input_tokens": 1048576, "max_output_tokens": 393216, "mode": "chat", @@ -64133,6 +64144,7 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_priority": 3.75e-07, "litellm_provider": "fireworks_ai", + "token_count_framing": 23, "max_input_tokens": 1048576, "max_output_tokens": 393216, "mode": "chat", diff --git a/tests/integration/providers/test_chat_template_tool_count_wire.py b/tests/integration/providers/test_chat_template_tool_count_wire.py index a8b97fa732e..552df3a87bc 100644 --- a/tests/integration/providers/test_chat_template_tool_count_wire.py +++ b/tests/integration/providers/test_chat_template_tool_count_wire.py @@ -39,7 +39,7 @@ _API_KEY: Final = "scripted-provider-key" _DEEPSEEK_V4: Final = "azure_ai/DeepSeek-V4-Flash-0731" _DEEPSEEK_R1: Final = "azure_ai/DeepSeek-R1-0528" _QWEN3_5: Final = "databricks/databricks-qwen35-122b-a10b" -_GROK: Final = "azure_ai/grok-4-1" +_GROK_3: Final = "azure_ai/grok-3" _DEEPSEEK_V4_TOOL_PREAMBLE: Final = 222 _QWEN3_5_TOOL_PREAMBLE: Final = 196 _QWEN3_5_OPENED_SYSTEM_TURN: Final = 5 @@ -250,7 +250,7 @@ def test_messages_count_tokens_skips_web_search_and_counts_unmapped_tool_keys_fo def test_messages_count_tokens_keeps_openai_tool_framing_for_other_models(gateway: Gateway) -> None: with wire_server(_completion) as wire, gateway.scenario() as scenario: - model: Final = _deployment(scenario, "count", _GROK, wire.url) + model: Final = _deployment(scenario, "count", _GROK_3, wire.url) without_tools: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, ())) plain: Final = _count(gateway, _MESSAGES_COUNT, _MESSAGES_COUNT.body(model, (_ANTHROPIC_TOOL,))) deferred: Final = _count( @@ -428,7 +428,7 @@ def test_project_itpm_reserves_deepseek_v4_tools_the_way_its_chat_template_rende def test_project_itpm_keeps_openai_tool_framing_for_other_models(gateway: Gateway) -> None: chat: Final = _ENDPOINTS["chat-completions"] with wire_server(_completion) as wire, gateway.scenario() as scenario: - group: Final = _deployment(scenario, "grok-4-1", _GROK, wire.url) + group: Final = _deployment(scenario, "grok-3", _GROK_3, wire.url) response: Final = gateway.request( "POST", chat.path, @@ -504,8 +504,8 @@ def _agent_body(group: str, marker: str, file_body: str, trigger: int) -> dict[s @pytest.mark.parametrize( ("model", "files_kept"), - ((_DEEPSEEK_V4, (False, False, True)), (_GROK, (True, True, True))), - ids=("deepseek-v4", "grok-4-1"), + ((_DEEPSEEK_V4, (False, False, True)), (_GROK_3, (True, True, True))), + ids=("deepseek-v4", "grok-3"), ) def test_messages_context_management_trigger_counts_tools_the_way_the_models_chat_template_renders_them( gateway: Gateway, model: str, files_kept: tuple[bool, ...] diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py index 51249e01d4d..a8162860684 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_chat_template_tools.py @@ -9,7 +9,10 @@ _BARE_FUNCTION: Final = {"name": "g", "parameters": {}} # Tool schemas as rendered by render_tools in encoding/encoding_dsv4.py (deepseek-ai/DeepSeek-V4-Flash) and by -# chat_template.jinja (Qwen/Qwen3.5-122B-A10B), both read on Hugging Face on 2026-10-17 +# chat_template.jinja (Qwen/Qwen3.5-122B-A10B), both read on Hugging Face on 2026-10-17, and by chat_template.jinja +# (deepseek-ai/DeepSeek-V4.1-Flash) read on 2026-10-29. Grok's template is closed, so its two renderings are the ones +# whose token counts left a constant preamble across seven tool shapes billed by Azure AI Foundry +# grok-4-1-fast-reasoning and xAI grok-4.3 on 2026-10-29 @pytest.mark.parametrize( ("model", "rendered"), [ @@ -17,6 +20,18 @@ _BARE_FUNCTION: Final = {"name": "g", "parameters": {}} "DeepSeek-V4-Flash-0731", '{"name": "f", "description": "é", "parameters": {}}\n{"name": "g", "parameters": {}}', ), + ( + "DeepSeek-V4.1-Flash", + '{"name": "f", "description": "é", "parameters": {}}\n{"name": "g", "parameters": {}}', + ), + ( + "grok-4-1-fast-reasoning", + '{"name":"f","description":"é","parameters":{}}\n{"name":"g","parameters":{}}', + ), + ( + "grok-4.3", + '{"name": "f", "description": "é", "parameters": {}}\n{"name": "g", "parameters": {}}', + ), ( "databricks-qwen35-122b-a10b", '\n{"type": "function", "function": {"name": "f", "description": "é", "parameters": {}}}' diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 697ac76b4e5..8ef118154d1 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -638,8 +638,14 @@ def test_token_counter_tools_match_billed_prompt_for_json_tool_chat_templates( ("deepseek/deepseek-chat:free", "DeepSeek-V4-Flash-0731"), ("deepseek-reasoner", "DeepSeek-V4-Flash-0731"), ("deepseek/deepseek-flash", "DeepSeek-V4-Flash-0731"), - ("DeepSeek-V4.1-Flash", "deepseek-v3.2"), - ("accounts/fireworks/models/deepseek-v4p1-flash", "deepseek-v3.2"), + ("DeepSeek-V4.1-Flash", "accounts/fireworks/models/deepseek-v4p1-flash"), + ("databricks-deepseek-v4-1-flash", "accounts/fireworks/models/deepseek-v4p1-flash"), + ("deepseek-ai/deepseek-v4.1-flash-maas", "accounts/fireworks/models/deepseek-v4p1-flash"), + ("deepseek-v4p1-flash-1017", "accounts/fireworks/models/deepseek-v4p1-flash"), + ("grok-4.1-fast-non-reasoning", "grok-4-1-fast-reasoning"), + ("xai.grok-4.1-fast", "grok-4-1-fast-reasoning"), + ("grok-4.3-latest", "grok-4.3"), + ("xai.grok-4-3", "grok-4.3"), ("Qwen/Qwen3.5-122B-A10B", "databricks-qwen35-122b-a10b"), ("qwen3-235b-a22b", "mistral-large-latest"), ], @@ -704,6 +710,48 @@ def test_token_counter_tool_choice_adds_nothing_for_deepseek_v4(tool_choice: obj ) == token_counter_new(model="DeepSeek-V4-Flash-0731", messages=list(_USER_MESSAGES), tools=tools) +# A request with a one-function tool billed 283 prompt tokens on xAI grok-4.3 without tool_choice and 297 with a +# named one, and 357 and 358 on Azure AI Foundry grok-4-1-fast-reasoning, on 2026-10-29 +@pytest.mark.parametrize("model", ["grok-4.3", "grok-4-1-fast-reasoning"]) +def test_token_counter_counts_a_named_tool_choice_for_grok(model: str): + tools: Final = [_openai_tool(_WEATHER_FUNCTION)] + named: Final = {"type": "function", "function": {"name": "get_weather"}} + + assert token_counter_new( + model=model, messages=list(_USER_MESSAGES), tools=tools, tool_choice=named + ) > token_counter_new(model=model, messages=list(_USER_MESSAGES), tools=tools) + + +def _counted_with_strict(model: str, **strict: bool) -> int: + return token_counter_new( + model=model, messages=list(_USER_MESSAGES), tools=[_openai_tool({**_WEATHER_FUNCTION, **strict})] + ) + + +# A one-function tool billed the same prompt tokens with strict unset and false on 2026-10-29: 357 through Azure AI +# Foundry grok-4-1-fast-reasoning chat and 283 through xAI grok-4.3 +@pytest.mark.parametrize("model", ["grok-4-1-fast-reasoning", "grok-4.3"]) +def test_token_counter_renders_a_false_strict_like_an_unset_one_for_grok(model: str): + assert ( + _counted_with_strict(model, strict=False) + == _counted_with_strict(model) + < _counted_with_strict(model, strict=True) + ) + + +# A one-function tool billed 5 prompt tokens more with strict false than unset, the same as with strict true, on +# 2026-10-30: 310 against 305 through Baseten DeepSeek-V4.1-Flash and 309 against 304 through Together +@pytest.mark.parametrize( + "model", ["accounts/fireworks/models/deepseek-v4p1-flash", "deepseek-v4-flash", "qwen3.5-plus"] +) +def test_token_counter_renders_a_false_strict_like_a_true_one_for_other_chat_templates(model: str): + assert ( + _counted_with_strict(model, strict=False) + == _counted_with_strict(model, strict=True) + > _counted_with_strict(model) + ) + + @pytest.mark.parametrize( ("tool_choice", "expected_tool_choice"), [ @@ -759,9 +807,17 @@ def test_token_counter_counts_anthropic_tools_as_sent_to_chat_template_families( # Qwen3.5's chat_template.jinja opens a system turn for the tool schemas when the request has none, 5 tokens that a # request with a system message already pays, while DeepSeek V4's encoding_dsv4.py appends them to the first turn -# either way (both read on Hugging Face on 2026-10-17) +# either way (both read on Hugging Face on 2026-10-17). A system message cost the same with and without a tool on +# Fireworks deepseek-v4p1-flash, Azure AI Foundry grok-4-1-fast-reasoning and xAI grok-4.3 on 2026-10-29 @pytest.mark.parametrize( - ("model", "implicit_system_turn_tokens"), [("DeepSeek-V4-Flash-0731", 0), ("databricks-qwen35-122b-a10b", 5)] + ("model", "implicit_system_turn_tokens"), + [ + ("DeepSeek-V4-Flash-0731", 0), + ("databricks-qwen35-122b-a10b", 5), + ("accounts/fireworks/models/deepseek-v4p1-flash", 0), + ("grok-4-1-fast-reasoning", 0), + ("grok-4.3", 0), + ], ) def test_token_counter_counts_the_tools_system_turn_only_without_a_system_message( model: str, implicit_system_turn_tokens: int diff --git a/tests/unit/proxy/common_utils/test_token_count_settings.py b/tests/unit/proxy/common_utils/test_token_count_settings.py index 43edc7e449f..68910f52b02 100644 --- a/tests/unit/proxy/common_utils/test_token_count_settings.py +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -131,6 +131,107 @@ def test_shipped_australian_count_region_stays_with_an_australian_deployment() - assert settings.count_model == TokenCountModel(model="au.anthropic.claude-opus-5-5", region="ap-southeast-4") +_ABSENT: Final = object() + + +@pytest.mark.parametrize( + ("key", "field"), + [ + ("token_count_tool_framing", "tool_framing"), + ("token_count_responses_framing", "responses_framing"), + ("token_count_responses_tool_framing", "responses_tool_framing"), + ], +) +@pytest.mark.parametrize( + ("configured", "default", "expected"), + [ + (_ABSENT, 151, 151), + (40, 151, 40), + (0, 151, 0), + (40.0, _ABSENT, 40), + (-1, 151, 151), + (True, 151, 151), + ("40", _ABSENT, None), + ], +) +def test_each_request_kind_framing_reads_its_value_like_framing( + key: str, field: str, configured: object, default: object, expected: int | None +) -> None: + settings: Final = resolve_token_count_settings( + configured={} if configured is _ABSENT else {key: configured}, + default={} if default is _ABSENT else {key: default}, + deployment_region=_in_us_west_2, + ) + assert getattr(settings, field) == expected + + +def _hidden_prompts(settings: TokenCountSettings) -> tuple[int, int, int, int]: + return ( + settings.hidden_prompt_tokens(carries_tools=False, responses_api=False), + settings.hidden_prompt_tokens(carries_tools=True, responses_api=False), + settings.hidden_prompt_tokens(carries_tools=False, responses_api=True), + settings.hidden_prompt_tokens(carries_tools=True, responses_api=True), + ) + + +@pytest.mark.parametrize( + ("default", "expected"), + [ + ({}, (0, 0, 0, 0)), + ({"token_count_framing": 23}, (23, 23, 23, 23)), + ({"token_count_framing": 185, "token_count_tool_framing": 214}, (185, 214, 185, 214)), + ( + { + "token_count_framing": 149, + "token_count_tool_framing": 300, + "token_count_responses_framing": 145, + "token_count_responses_tool_framing": 104, + }, + (149, 300, 145, 104), + ), + ({"token_count_framing": 149, "token_count_tool_framing": 0}, (149, 0, 149, 0)), + ({"token_count_framing": 149, "token_count_responses_tool_framing": 0}, (149, 149, 149, 0)), + ({"token_count_framing": 149, "token_count_responses_framing": 145}, (149, 149, 145, 145)), + ], +) +def test_hidden_prompt_is_the_most_specific_framing_set_for_the_request_api_and_tools( + default: dict[str, object], expected: tuple[int, int, int, int] +) -> None: + settings: Final = resolve_token_count_settings(configured={}, default=default, deployment_region=_in_us_west_2) + assert _hidden_prompts(settings) == expected + + +_PER_API_COST_MAP: Final = { + "token_count_framing": 149, + "token_count_tool_framing": 300, + "token_count_responses_framing": 145, + "token_count_responses_tool_framing": 104, +} + + +@pytest.mark.parametrize( + ("configured", "expected"), + [ + ({}, (149, 300, 145, 104)), + ({"token_count_framing": 500}, (500, 500, 500, 500)), + ({"token_count_framing": 0}, (0, 0, 0, 0)), + ({"token_count_tool_framing": 250}, (149, 250, 145, 250)), + ({"token_count_framing": 500, "token_count_tool_framing": 250}, (500, 250, 500, 250)), + ({"token_count_responses_framing": 120}, (149, 300, 120, 120)), + ({"token_count_tool_framing": 250, "token_count_responses_framing": 120}, (149, 250, 120, 120)), + ({"token_count_responses_tool_framing": 90}, (149, 300, 145, 90)), + ({"token_count_framing": -1}, (149, 300, 145, 104)), + ], +) +def test_an_operator_framing_applies_to_every_request_kind_the_cost_map_splits_more_finely( + configured: dict[str, object], expected: tuple[int, int, int, int] +) -> None: + settings: Final = resolve_token_count_settings( + configured=configured, default=_PER_API_COST_MAP, deployment_region=_in_us_west_2 + ) + assert _hidden_prompts(settings) == expected + + def test_no_setting_anywhere_means_the_model_counts_itself_without_framing() -> None: assert resolve_token_count_settings( configured={}, default={}, deployment_region=_in_us_west_2 diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index ab6cfcfebfe..44c43713438 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -4896,6 +4896,9 @@ class TestModelInfoCostMapEchoFilter: [ ("bedrock/us.anthropic.claude-sonnet-5-5", "token_count_model", {}), ("bedrock/us.moonshotai.kimi-k3", "token_count_framing", 0), + ("azure_ai/grok-4-1-fast-reasoning", "token_count_tool_framing", 0), + ("azure_ai/grok-4-1-fast-reasoning", "token_count_responses_framing", 0), + ("azure_ai/grok-4-1-fast-reasoning", "token_count_responses_tool_framing", 0), ], ) def test_token_count_setting_is_stored_only_when_the_operator_changed_it( diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 608efe20c50..2f85d1c7516 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -1925,6 +1925,23 @@ async def test_uncountable_model_without_framing_raises(): assert raised.value.code == "404" +@pytest.mark.parametrize( + "settings", + [ + TokenCountSettings(framing=None, count_model=None, tool_framing=300), + TokenCountSettings(framing=None, count_model=None, responses_framing=145), + TokenCountSettings(framing=None, count_model=None, responses_tool_framing=100), + ], +) +@pytest.mark.asyncio +async def test_uncountable_model_with_only_a_request_kind_framing_falls_through_to_the_local_estimate( + settings: TokenCountSettings, +): + counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=404)}) + + assert isinstance(await _ladder(counter, settings), ProviderCountFailure) + + @pytest.mark.parametrize( ("count_model", "counted_as"), [ @@ -2770,6 +2787,480 @@ async def test_framing_comes_from_the_row_the_deployment_provider_or_base_model_ assert aliased.total_tokens == litellm.token_counter(model="gpt-5.4-mini", messages=messages) + framing +@pytest.mark.asyncio +async def test_a_tool_request_counts_its_framing_in_place_of_the_plain_one_per_api(monkeypatch: pytest.MonkeyPatch): + def deployment(name: str, model_info: dict[str, object]) -> dict[str, object]: + return { + "model_name": name, + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": f"sk-unused-{name}"}, + "model_info": model_info, + } + + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + deployment("plain", {}), + deployment("chat", {"token_count_framing": 7, "token_count_tool_framing": 50}), + deployment( + "per-api", + { + "token_count_framing": 7, + "token_count_tool_framing": 50, + "token_count_responses_framing": 5, + "token_count_responses_tool_framing": 2, + }, + ), + ] + ), + ) + messages: Final = [{"role": "user", "content": "What is the weather in Paris?"}] + tools: Final = [_KIMI_WEATHER] + + async def count(model: str, request_tools: list[dict[str, object]] | None, responses: bool) -> int: + response: Final = await count_request_tokens( + request=TokenCountRequest(model=model, messages=messages, tools=request_tools), + responses_request={"model": model, "input": "What is the weather in Paris?"} if responses else None, + ) + return response.total_tokens + + plain_with_tools: Final = await count("plain", tools, responses=False) + plain_without_tools: Final = await count("plain", None, responses=False) + + assert await count("plain", tools, responses=True) == plain_with_tools + assert await count("chat", tools, responses=False) == plain_with_tools + 50 + assert await count("chat", None, responses=False) == plain_without_tools + 7 + assert await count("chat", [], responses=False) == plain_without_tools + 7 + assert await count("chat", tools, responses=True) == plain_with_tools + 50 + assert await count("chat", None, responses=True) == plain_without_tools + 7 + assert await count("per-api", tools, responses=False) == plain_with_tools + 50 + assert await count("per-api", tools, responses=True) == plain_with_tools + 2 + assert await count("per-api", None, responses=True) == plain_without_tools + 5 + + +_GET_WEATHER: Final = { + "name": "get_weather", + "description": "Get the current weather for a city", + "input_schema": { + "type": "object", + "properties": {"city": {"type": "string", "description": "City name"}}, + "required": ["city"], + }, +} +_THREE_TOOLS: Final = ( + _GET_WEATHER, + { + "name": "get_time", + "description": "Get the local time in a timezone", + "input_schema": { + "type": "object", + "properties": {"timezone": {"type": "string", "description": "IANA timezone"}}, + "required": ["timezone"], + }, + }, + { + "name": "search_docs", + "description": "Search internal documents and return the top matches", + "input_schema": { + "type": "object", + "properties": {"query": {"type": "string"}, "limit": {"type": "integer", "description": "Max results"}}, + "required": ["query"], + }, + }, +) +_FORM_FIELDS: Final = tuple(f"field_{index}" for index in range(10)) +_FILL_FORM: Final = { + "name": "fill_form", + "description": "Fill a form", + "input_schema": { + "type": "object", + "properties": { + field: {"type": "string", "description": f"Value for field number {index}"} + for index, field in enumerate(_FORM_FIELDS) + }, + "required": list(_FORM_FIELDS), + }, +} +_WEATHER_SYSTEM: Final = "You are a helpful assistant that answers weather questions briefly." +_AZURE_GROK_4_1: Final = "azure_ai/grok-4-1-fast-reasoning" +_FIREWORKS_DEEPSEEK_V4_1: Final = "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash" +_XAI_GROK_4_3: Final = "xai/grok-4.3" +_NATIVE_RESPONSES_BASE: Final = { + _AZURE_GROK_4_1: "https://unused.services.ai.azure.com/models", + _FIREWORKS_DEEPSEEK_V4_1: "https://unused.invalid", + _XAI_GROK_4_3: "https://unused.invalid", +} + + +def _responses_function(tool: Mapping[str, object], **strict: bool | None) -> dict[str, object]: + return { + "type": "function", + "name": tool["name"], + "description": tool["description"], + "parameters": tool["input_schema"], + **strict, + } + + +def _get_weather_copies(count: int) -> tuple[dict[str, object], ...]: + return tuple({**_GET_WEATHER, "name": f"get_weather_{index}"} for index in range(count)) + + +# usage.input_tokens plus cache_read_input_tokens that POST /v1/messages billed for these requests with the user turn +# "What is the weather in Paris?" through Azure AI Foundry grok-4-1-fast-reasoning, Fireworks deepseek-v4p1-flash and +# xAI grok-4.3 on 2026-10-29, and usage.prompt_tokens that POST /v1/chat/completions billed through DeepSeek's API on +# 2026-10-30 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "system", "tools", "billed_prompt_tokens"), + [ + (_AZURE_GROK_4_1, None, (), 163), + (_AZURE_GROK_4_1, None, (_GET_WEATHER,), 357), + (_AZURE_GROK_4_1, None, _THREE_TOOLS, 445), + (_AZURE_GROK_4_1, None, (_FILL_FORM,), 550), + (_AZURE_GROK_4_1, _WEATHER_SYSTEM, (_GET_WEATHER,), 368), + (_FIREWORKS_DEEPSEEK_V4_1, None, (), 37), + (_FIREWORKS_DEEPSEEK_V4_1, None, (_GET_WEATHER,), 304), + (_FIREWORKS_DEEPSEEK_V4_1, None, _THREE_TOOLS, 430), + (_FIREWORKS_DEEPSEEK_V4_1, None, (_FILL_FORM,), 550), + (_FIREWORKS_DEEPSEEK_V4_1, _WEATHER_SYSTEM, (_GET_WEATHER,), 315), + (_XAI_GROK_4_3, None, (), 199), + (_XAI_GROK_4_3, None, (_GET_WEATHER,), 283), + (_XAI_GROK_4_3, None, _THREE_TOOLS, 402), + (_XAI_GROK_4_3, None, (_FILL_FORM,), 530), + (_XAI_GROK_4_3, _WEATHER_SYSTEM, (_GET_WEATHER,), 297), + ("deepseek/deepseek-v4-flash", None, (), 37), + ("deepseek/deepseek-v4-flash", None, (_GET_WEATHER,), 304), + ("deepseek/deepseek-v4-flash", None, _THREE_TOOLS, 430), + ("deepseek/deepseek-v4-flash", _WEATHER_SYSTEM, (_GET_WEATHER,), 315), + ("deepseek/deepseek-v4-pro", None, (), 90), + ("deepseek/deepseek-v4-pro", None, (_GET_WEATHER,), 367), + ("deepseek/deepseek-v4-pro", None, _THREE_TOOLS, 493), + ("deepseek/deepseek-v4-pro", _WEATHER_SYSTEM, (_GET_WEATHER,), 378), + ], +) +async def test_local_estimate_matches_the_billed_prompt_with_the_shipped_framing( + monkeypatch: pytest.MonkeyPatch, + model: str, + system: str | None, + tools: tuple[dict[str, object], ...], + billed_prompt_tokens: int, +): + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "billed", + "litellm_params": {"model": model, "api_key": "unused", "api_base": "https://unused.invalid"}, + } + ] + ), + ) + + counted: Final = await token_counter( + request=TokenCountRequest( + model="billed", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=list(tools) or None, + system=system, + ) + ) + + assert counted.count_source == "local_estimate" + assert abs(counted.total_tokens - billed_prompt_tokens) <= billed_prompt_tokens * 0.05 + + +# usage.input_tokens that POST /v1/responses billed for these requests with the input "What is the weather in Paris?" +# through Azure AI Foundry grok-4-1-fast-reasoning, Fireworks deepseek-v4p1-flash and xAI grok-4.3 on 2026-10-29 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model", "system", "tools", "billed_prompt_tokens"), + [ + (_AZURE_GROK_4_1, None, (), 159), + (_AZURE_GROK_4_1, None, (_GET_WEATHER,), 159), + (_AZURE_GROK_4_1, None, _THREE_TOOLS, 257), + (_AZURE_GROK_4_1, None, (_FILL_FORM,), 352), + (_AZURE_GROK_4_1, _WEATHER_SYSTEM, (_GET_WEATHER,), 170), + (_FIREWORKS_DEEPSEEK_V4_1, None, (_GET_WEATHER,), 304), + (_XAI_GROK_4_3, None, (_GET_WEATHER,), 283), + ], +) +async def test_responses_estimate_matches_the_billed_responses_prompt_with_the_shipped_framing( + monkeypatch: pytest.MonkeyPatch, + model: str, + system: str | None, + tools: tuple[dict[str, object], ...], + billed_prompt_tokens: int, +): + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "billed", + "litellm_params": {"model": model, "api_key": "unused", "api_base": _NATIVE_RESPONSES_BASE[model]}, + } + ] + ), + ) + + counted: Final = await count_request_tokens( + request=TokenCountRequest( + model="billed", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=[_responses_function(tool) for tool in tools] or None, + system=system, + ), + responses_request={"model": "billed", "input": "What is the weather in Paris?"}, + ) + + assert counted.count_source == "local_estimate" + assert abs(counted.total_tokens - billed_prompt_tokens) <= billed_prompt_tokens * 0.05 + + +def _single_deployment_router(litellm_params: dict[str, object], model_info: dict[str, object]) -> Router: + return Router( + model_list=[ + { + "model_name": "counted", + "litellm_params": {"api_key": "unused", **litellm_params}, + "model_info": model_info, + } + ] + ) + + +async def _count_with_router( + monkeypatch: pytest.MonkeyPatch, router: Router, tools: tuple[dict[str, object], ...], responses: bool +) -> int: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + counted: Final = await count_request_tokens( + request=TokenCountRequest( + model="counted", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=list(tools) or None, + ), + responses_request={"model": "counted", "input": "What is the weather in Paris?"} if responses else None, + ) + return counted.total_tokens + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tools", [(), (_GET_WEATHER,), _THREE_TOOLS, (_FILL_FORM,)]) +@pytest.mark.parametrize("responses", [False, True]) +async def test_a_custom_deployment_name_counts_like_the_base_model_it_names( + monkeypatch: pytest.MonkeyPatch, tools: tuple[dict[str, object], ...], responses: bool +): + api_base: Final = _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1] + custom_name: Final = _single_deployment_router( + {"model": "azure_ai/prod-grok", "api_base": api_base}, {"base_model": _AZURE_GROK_4_1} + ) + row_name: Final = _single_deployment_router({"model": _AZURE_GROK_4_1, "api_base": api_base}, {}) + + assert await _count_with_router(monkeypatch, custom_name, tools, responses) == await _count_with_router( + monkeypatch, row_name, tools, responses + ) + + +# usage.input_tokens that POST /v1/responses billed through Azure AI Foundry grok-4-1-fast-reasoning on 2026-10-29 for +# the input "What is the weather in Paris?" with these function tools, strict left unset or sent as false or null, +# the same on the .services.ai.azure.com and .openai.azure.com hosts +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tools", "strict", "billed_prompt_tokens"), + [ + ((_GET_WEATHER,), {}, 159), + ((_GET_WEATHER,), {"strict": False}, 155), + ((_GET_WEATHER,), {"strict": None}, 155), + (_THREE_TOOLS, {}, 257), + (_get_weather_copies(5), {}, 345), + (_get_weather_copies(10), {}, 575), + (_get_weather_copies(10), {"strict": False}, 535), + (_get_weather_copies(10), {"strict": None}, 535), + ], +) +@pytest.mark.parametrize( + ("litellm_params", "model_info"), + [ + ({"model": _AZURE_GROK_4_1, "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1]}, {}), + ( + {"model": "azure/grok-4-1-fast-reasoning", "api_base": "https://unused.openai.azure.com/openai/v1"}, + {"base_model": _AZURE_GROK_4_1}, + ), + ], +) +async def test_azure_responses_estimate_stays_at_or_just_over_the_billed_prompt_as_function_tools_are_added( + monkeypatch: pytest.MonkeyPatch, + tools: tuple[dict[str, object], ...], + strict: dict[str, bool | None], + billed_prompt_tokens: int, + litellm_params: dict[str, object], + model_info: dict[str, object], +): + router: Final = _single_deployment_router(litellm_params, model_info) + + counted: Final = await _count_with_router( + monkeypatch, router, tuple(_responses_function(tool, **strict) for tool in tools), responses=True + ) + + assert billed_prompt_tokens <= counted <= billed_prompt_tokens * 1.02 + + +# POST /v1/responses billed the same usage.input_tokens with strict unset and false on 2026-10-29, 283 for one tool +# through xAI grok-4.3 +@pytest.mark.asyncio +async def test_a_responses_host_that_bills_an_unset_strict_like_false_counts_them_alike( + monkeypatch: pytest.MonkeyPatch, +): + router: Final = _single_deployment_router( + {"model": _XAI_GROK_4_3, "api_base": _NATIVE_RESPONSES_BASE[_XAI_GROK_4_3]}, {} + ) + + unset: Final = await _count_with_router(monkeypatch, router, (_responses_function(_GET_WEATHER),), responses=True) + + assert unset == await _count_with_router( + monkeypatch, router, (_responses_function(_GET_WEATHER, strict=False),), responses=True + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("built_in_tool", "responses"), + [({"type": "web_search"}, True), ({"type": "web_search_20250305", "name": "web_search"}, False)], +) +async def test_a_request_whose_only_tools_are_built_in_counts_like_one_without_tools( + monkeypatch: pytest.MonkeyPatch, built_in_tool: dict[str, object], responses: bool +): + router: Final = _single_deployment_router( + {"model": _AZURE_GROK_4_1, "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1]}, {} + ) + + assert await _count_with_router(monkeypatch, router, (built_in_tool,), responses) == await _count_with_router( + monkeypatch, router, (), responses + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tools", [(), (_GET_WEATHER,), _THREE_TOOLS]) +async def test_a_responses_count_the_proxy_bridges_to_chat_completions_frames_like_the_chat_count( + monkeypatch: pytest.MonkeyPatch, tools: tuple[dict[str, object], ...] +): + bridged: Final = _single_deployment_router( + {"model": _AZURE_GROK_4_1, "api_base": "https://unused.cognitiveservices.azure.com/models"}, {} + ) + native: Final = _single_deployment_router( + {"model": _AZURE_GROK_4_1, "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1]}, {} + ) + chat_count: Final = await _count_with_router(monkeypatch, bridged, tools, responses=False) + + assert await _count_with_router(monkeypatch, bridged, tools, responses=True) == chat_count + assert await _count_with_router(monkeypatch, native, tools, responses=True) != chat_count + + +@pytest.mark.asyncio +@pytest.mark.parametrize("from_router_defaults", [False, True]) +async def test_a_responses_count_that_asks_for_the_chat_completions_bridge_frames_like_the_chat_count( + monkeypatch: pytest.MonkeyPatch, from_router_defaults: bool +): + router: Final = Router( + model_list=[ + { + "model_name": "counted", + "litellm_params": { + "api_key": "unused", + "model": _AZURE_GROK_4_1, + "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1], + }, + "model_info": {}, + } + ], + default_litellm_params={"use_chat_completions_api": True} if from_router_defaults else None, + ) + chat_count: Final = await _count_with_router(monkeypatch, router, (_GET_WEATHER,), responses=False) + + counted: Final = await count_request_tokens( + request=TokenCountRequest( + model="counted", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=[_GET_WEATHER], + ), + responses_request={ + "model": "counted", + "input": "What is the weather in Paris?", + **({} if from_router_defaults else {"use_chat_completions_api": True}), + }, + ) + + assert counted.total_tokens == chat_count + + +@pytest.mark.asyncio +async def test_a_responses_count_frames_for_the_api_base_router_defaults_send_inference_to( + monkeypatch: pytest.MonkeyPatch, +): + router: Final = Router( + model_list=[ + { + "model_name": "counted", + "litellm_params": { + "api_key": "unused", + "model": _AZURE_GROK_4_1, + "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1], + }, + "model_info": {}, + } + ], + default_litellm_params={"api_base": "https://unused.cognitiveservices.azure.com/models"}, + ) + chat_count: Final = await _count_with_router(monkeypatch, router, (_GET_WEATHER,), responses=False) + + assert await _count_with_router(monkeypatch, router, (_GET_WEATHER,), responses=True) == chat_count + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_chat_completions_api", ["maybe", {"a": 1}]) +async def test_a_responses_count_with_a_malformed_bridge_flag_reads_it_the_way_the_responses_route_does( + monkeypatch: pytest.MonkeyPatch, use_chat_completions_api: object +): + router: Final = _single_deployment_router( + {"model": _AZURE_GROK_4_1, "api_base": _NATIVE_RESPONSES_BASE[_AZURE_GROK_4_1]}, {} + ) + chat_count: Final = await _count_with_router(monkeypatch, router, (_GET_WEATHER,), responses=False) + + counted: Final = await count_request_tokens( + request=TokenCountRequest( + model="counted", + messages=[{"role": "user", "content": "What is the weather in Paris?"}], + tools=[_GET_WEATHER], + ), + responses_request={ + "model": "counted", + "input": "What is the weather in Paris?", + "use_chat_completions_api": use_chat_completions_api, + }, + ) + + assert counted.total_tokens == chat_count + + +@pytest.mark.asyncio +async def test_an_azure_responses_count_leaves_strict_unset_for_templates_not_measured_there( + monkeypatch: pytest.MonkeyPatch, +): + router: Final = _single_deployment_router( + {"model": "azure_ai/deepseek-v4.1-flash", "api_base": "https://unused.services.ai.azure.com/models"}, {} + ) + + async def counted(**strict: bool) -> int: + return await _count_with_router(monkeypatch, router, (_responses_function(_GET_WEATHER, **strict),), True) + + assert await counted() < await counted(strict=True) + + _KIMI_WEATHER: Final = { "type": "function", "function": { diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 91f0b2339c7..447a4dfb81d 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -672,7 +672,7 @@ def test_pdf_opt_out_on_one_deployment_does_not_block_its_sibling(respx_mock: re def test_token_count_settings_stay_on_the_deployment_that_set_them() -> None: - """A deployment's ``token_count_model`` and ``token_count_framing`` must not land on + """A deployment's ``token_count_model`` and framing settings must not land on the shared ``{provider}/{model}`` key, or a sibling deployment of that model would count with them.""" from litellm.proxy.common_utils.token_count_settings import cost_map_entry @@ -681,6 +681,9 @@ def test_token_count_settings_stay_on_the_deployment_that_set_them() -> None: configured: Final = { "token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-west-2"}, "token_count_framing": 7, + "token_count_tool_framing": 3, + "token_count_responses_framing": 5, + "token_count_responses_tool_framing": 2, } builtin: Final = {field: cost_map_entry(backend_model).get(field) for field in configured} assert all(builtin[field] != value for field, value in configured.items()), builtin diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 5cf872f4ea7..61825e463f8 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1178,6 +1178,9 @@ _COST_MAP_ROW_SCHEMA: Final = { "supports_vision": {"type": "boolean"}, "supports_web_search": {"type": "boolean"}, "token_count_framing": {"type": "integer"}, + "token_count_tool_framing": {"type": "integer"}, + "token_count_responses_framing": {"type": "integer"}, + "token_count_responses_tool_framing": {"type": "integer"}, "token_count_model": { "type": "object", "properties": {"model": {"type": "string"}, "region": {"type": "string"}}, From 05b76878b05fda8fbd687099659d538eb2e60525 Mon Sep 17 00:00:00 2001 From: jesus-berri Date: Mon, 2 Nov 2026 12:42:23 -0800 Subject: [PATCH 06/79] fix(cohere): map Cohere v2 finish_reason on non-streaming chat responses (#49529) * fix(cohere): map Cohere v2 finish_reason on non-streaming chat responses * fix(cohere): default a missing Cohere v2 finish_reason to stop instead of raising A 200 body without finish_reason used to return stop and started raising KeyError after the mapping change, so fall back to COMPLETE --- litellm/llms/cohere/chat/v2_transformation.py | 4 ++ .../cohere/chat/test_cohere_transformation.py | 48 +++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/litellm/llms/cohere/chat/v2_transformation.py b/litellm/llms/cohere/chat/v2_transformation.py index cf033e10b41..141a174e558 100644 --- a/litellm/llms/cohere/chat/v2_transformation.py +++ b/litellm/llms/cohere/chat/v2_transformation.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx import litellm +from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.llms.cohere import CohereV2ChatResponse @@ -248,6 +249,9 @@ class CohereV2ChatConfig(OpenAIGPTConfig): prompt_tokens: Final = token_usage.get("input_tokens", 0) completion_tokens: Final = token_usage.get("output_tokens", 0) + model_response.choices[0].finish_reason = map_finish_reason( # rebind-ok: ModelResponse is filled in place + cohere_v2_chat_response.get("finish_reason", "COMPLETE") + ) model_response.created = int(time.time()) model_response.model = model usage: Final = Usage( diff --git a/tests/unit/llms/cohere/chat/test_cohere_transformation.py b/tests/unit/llms/cohere/chat/test_cohere_transformation.py index d32b3bcbeca..fcb3d3f1cc0 100644 --- a/tests/unit/llms/cohere/chat/test_cohere_transformation.py +++ b/tests/unit/llms/cohere/chat/test_cohere_transformation.py @@ -576,6 +576,54 @@ def test_cohere_v2_tool_calling(respx_mock: respx.MockRouter) -> None: assert [(call.id, call.function.name, call.function.arguments) for call in tool_calls] == [ ("get_weather_k9q2xw1n", "get_weather", '{"location":"New York, NY"}') ] + assert response.choices[0].finish_reason == "tool_calls" + + +@pytest.mark.usefixtures("_cohere_httpx_transport") +@pytest.mark.respx(assert_all_called=True) +@pytest.mark.parametrize( + ("cohere_finish_reason", "expected_finish_reason"), + [("COMPLETE", "stop"), ("MAX_TOKENS", "length"), ("STOP_SEQUENCE", "stop")], +) +def test_cohere_v2_maps_finish_reason( + respx_mock: respx.MockRouter, cohere_finish_reason: str, expected_finish_reason: str +) -> None: + respx_mock.post(COHERE_V2_CHAT_URL).mock( + return_value=httpx.Response( + 200, + json=_cohere_v2_response( + {"content": [{"type": "text", "text": "Hello"}]}, + finish_reason=cohere_finish_reason, + ), + ) + ) + + response: Final = litellm.completion( + model="cohere_chat/v2/command-a-03-2025", + messages=[{"role": "user", "content": "Say hello"}], + api_key="cohere-test-key", + ) + + assert response.choices[0].finish_reason == expected_finish_reason + + +@pytest.mark.usefixtures("_cohere_httpx_transport") +@pytest.mark.respx(assert_all_called=True) +def test_cohere_v2_missing_finish_reason_defaults_to_stop(respx_mock: respx.MockRouter) -> None: + body_without_finish_reason: Final = { + key: value + for key, value in _cohere_v2_response({"content": [{"type": "text", "text": "Hello"}]}).items() + if key != "finish_reason" + } + respx_mock.post(COHERE_V2_CHAT_URL).mock(return_value=httpx.Response(200, json=body_without_finish_reason)) + + response: Final = litellm.completion( + model="cohere_chat/v2/command-a-03-2025", + messages=[{"role": "user", "content": "Say hello"}], + api_key="cohere-test-key", + ) + + assert (response.choices[0].message.content, response.choices[0].finish_reason) == ("Hello", "stop") @pytest.mark.usefixtures("_cohere_httpx_transport") From 044bd824c023b1ec30c7804b688f6ea6e37db8aa Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 2 Nov 2026 12:43:15 -0800 Subject: [PATCH 07/79] test(integration): align remaining release qualification checks (#49714) * test(integration): align remaining release qualification checks * test(integration): scope offline resolution to the restricted runner --- .circleci/scripts/run_integration.sh | 2 ++ .../test_cache_hit_guardrail_metrics.py | 23 ++++++++++------- .../pricing/test_price_data_reload_prices.py | 2 +- .../providers/test_decisions_wire.py | 25 ++++++++++++++++--- 4 files changed, 39 insertions(+), 13 deletions(-) diff --git a/.circleci/scripts/run_integration.sh b/.circleci/scripts/run_integration.sh index 1380b07de53..fb91bcce77f 100644 --- a/.circleci/scripts/run_integration.sh +++ b/.circleci/scripts/run_integration.sh @@ -128,6 +128,8 @@ fi sudo iptables -L integration_only -n -v -x > "$results/egress-guard.txt" awk '$3 == "REJECT" && $1 > 0 { rejected=1 } END { exit !rejected }' "$results/egress-guard.txt" +export UV_OFFLINE=true + if [ "$mode" = parity ]; then export INTEGRATION_WORKERS=1 else diff --git a/tests/integration/observability/test_cache_hit_guardrail_metrics.py b/tests/integration/observability/test_cache_hit_guardrail_metrics.py index c52649c46df..13cd2133d4f 100644 --- a/tests/integration/observability/test_cache_hit_guardrail_metrics.py +++ b/tests/integration/observability/test_cache_hit_guardrail_metrics.py @@ -253,6 +253,16 @@ def _count(samples: tuple, name: str, model_id: str) -> float: ) +def _deployment_states(rig: Rig) -> tuple[float, ...]: + return tuple( + sample.value + for sample in _metric_samples(rig.candidate, rig.model_name) + if sample.name == DEPLOYMENT_STATE + and sample.labels.get("model_id") == rig.deployment_id + and sample.labels.get("api_base") == rig.provider.url + "/v1" + ) + + def _populated_failures(samples: tuple, rig: Rig, api_provider: str) -> float: return float( sum( @@ -558,7 +568,8 @@ def test_cache_hit_post_call_guardrail_outage_keeps_deployment_labels( ) assert reject.status_code >= 400, reject.text assert rig.provider.received.qsize() == 1, rig.provider.drain() - _expect_metrics(rig, 0, 0, pf_populated=1, pf_blank=0) + _expect_metrics(rig, 1, 0) + _assert_spend(_call_id(reject), rig) def test_two_identical_cache_hit_rejects_increment_populated_series(gateway: Gateway, tmp_path: Path) -> None: @@ -635,6 +646,7 @@ def test_cache_hit_reject_deployment_metric_set_diff(gateway: Gateway, tmp_path: text: Final = "cache hit control e4 " + marker warm: Final = rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None)) assert warm.status_code == 200, warm.text + eventually(lambda: _deployment_states(rig), lambda states: states == (0.0,)) reject: Final = rig.candidate.request( "POST", "/v1/chat/completions", _chat_body(rig.model_name, text, rig.guardrail_name) ) @@ -642,11 +654,4 @@ def test_cache_hit_reject_deployment_metric_set_diff(gateway: Gateway, tmp_path: samples: Final = _expect_metrics(rig, 1, 0) blank: Final = tuple(sample for sample in samples if sample.labels.get("model_id") == "") assert blank == (), blank - states: Final = tuple( - sample.value - for sample in samples - if sample.name == DEPLOYMENT_STATE - and sample.labels.get("model_id") == rig.deployment_id - and sample.labels.get("api_base") == "" - ) - assert states == (1.0,), states + eventually(lambda: _deployment_states(rig), lambda states: states == (1.0,)) diff --git a/tests/integration/pricing/test_price_data_reload_prices.py b/tests/integration/pricing/test_price_data_reload_prices.py index 0aa3a3bcba2..a91f5cba1ad 100644 --- a/tests/integration/pricing/test_price_data_reload_prices.py +++ b/tests/integration/pricing/test_price_data_reload_prices.py @@ -284,7 +284,7 @@ def _charged_as_listed(proxy: Gateway, deployment: _Deployment, scale: float) -> assert _listed_rates(proxy, deployment.name) == _expected_listing(rates), deployment.name request_id, charged = _billed(proxy.client, proxy.key, deployment) expected: Final = _expected_cost(rates, deployment.usage) - assert charged == (None if rates is None else pytest.approx(expected, rel=1e-6)), deployment.name + assert charged == pytest.approx(expected, rel=1e-6), deployment.name return _Charge(request_id=request_id, expected=expected) diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index 5af1532e9ed..4e091460667 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -262,11 +262,28 @@ def _upstream_calls(gateway: Gateway, handle: ScenarioHandle) -> list[dict[str, return _calls_to(_observed_requests(gateway), handle) -def _sdk_inference_call(gateway: Gateway, handle: ScenarioHandle, provider: _Provider) -> dict[str, JsonValue]: +def _sdk_inference_call( + gateway: Gateway, + handle: ScenarioHandle, + provider: _Provider, + *, + metadata_authorization: str | None = None, +) -> dict[str, JsonValue]: observed: Final = _upstream_calls(gateway, handle) inference_path: Final = f"/{handle.scenario_id}/v1/systemone" (call,) = tuple(request for request in observed if request["path"] == inference_path) - metadata_calls: Final = tuple(request for request in observed if request["path"] != inference_path) + environment_metadata: Final = { + "path": f"/{handle.scenario_id}/api/show", + "authorization": metadata_authorization, + "body": {"name": provider.body_model}, + "method": "POST", + "api_key": "", + } + metadata_calls: Final = tuple( + request + for request in observed + if request["path"] != inference_path and (metadata_authorization is None or request != environment_metadata) + ) assert metadata_calls in ( (), ( @@ -556,7 +573,9 @@ def test_ollama_decisions_sdk_respects_environment_and_explicit_auth( api_key=api_key, ) assert response.model_dump(mode="json") == _answer_body(provider) - call: Final = _sdk_inference_call(gateway, handle, provider) + call: Final = _sdk_inference_call( + gateway, handle, provider, metadata_authorization=authorization if environment_base else None + ) assert (call["path"], call["authorization"], call["body"]) == ( f"/{handle.scenario_id}/v1/systemone", authorization, From 1563776cd82f2d9a7516ab275133b22b05375dec Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 2 Nov 2026 12:56:53 -0800 Subject: [PATCH 08/79] fix(prometheus): retain deployment attribution on proxy failures (#49717) --- litellm/integrations/prometheus.py | 12 ++++-- .../integrations/test_prometheus_labels.py | 42 +++++++++++++++++++ 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 42fe0837c29..1c9a443b546 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -3036,7 +3036,7 @@ class PrometheusLogger(CustomLogger): async def async_post_call_failure_hook( self, - request_data: dict, + request_data: dict[str, object], original_exception: Exception, user_api_key_dict: UserAPIKeyAuth, traceback_str: str | None = None, @@ -3059,10 +3059,14 @@ class PrometheusLogger(CustomLogger): try: _tags: Final = StandardLoggingPayloadSetup.get_request_tags( litellm_params=request_data, - proxy_server_request=request_data.get("proxy_server_request", {}), + proxy_server_request=_logged_fields(request_data.get("proxy_server_request")), + ) + _metadata: Final = _logged_fields(request_data.get("metadata")) + model_id: Final = ( + _logged_fields(_metadata.get("model_info")).get("id") + or _logged_fields(request_data.get("model_info")).get("id") + or _logged_fields(request_data.get("standard_logging_object")).get("model_id") ) - _metadata: Final = request_data.get("metadata", {}) or {} - model_id: Final = _metadata.get("model_info", {}).get("id") or request_data.get("model_info", {}).get("id") rate_limit_category, rate_limit_type = self._extract_rate_limit_labels(original_exception) api_provider: Final = self._extract_api_provider_from_request_data( request_data diff --git a/tests/unit/integrations/test_prometheus_labels.py b/tests/unit/integrations/test_prometheus_labels.py index 878f19c66f8..04eb7c3a5a4 100644 --- a/tests/unit/integrations/test_prometheus_labels.py +++ b/tests/unit/integrations/test_prometheus_labels.py @@ -5,6 +5,7 @@ Unit tests for prometheus metric labels configuration from typing import Final, Literal import pytest +from pydantic import JsonValue import litellm from litellm.types.integrations.prometheus import ( @@ -802,6 +803,47 @@ async def test_failure_hook_emits_api_provider_value_on_failed_requests_metric() _clear_prometheus_registry() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("metadata_id", "request_id", "logged_id", "expected"), + ( + (None, None, "logged-deployment", "logged-deployment"), + ("", None, "logged-deployment", "logged-deployment"), + ("metadata-deployment", "request-deployment", "logged-deployment", "metadata-deployment"), + (None, "request-deployment", "logged-deployment", "request-deployment"), + (None, None, None, "None"), + ), +) +async def test_failure_counters_keep_logged_deployment_attribution( + metadata_id: str | None, request_id: str | None, logged_id: str | None, expected: str +) -> None: + from fastapi import HTTPException + + from litellm.integrations.prometheus import PrometheusLogger + from litellm.proxy._types import UserAPIKeyAuth + + request_data: Final[dict[str, JsonValue]] = { + "model": "metric-test-alias", + "metadata": {"model_info": {"id": metadata_id}}, + "model_info": {"id": request_id}, + "standard_logging_object": {"model_id": logged_id, "custom_llm_provider": "openai"}, + } + _clear_prometheus_registry() + try: + await PrometheusLogger().async_post_call_failure_hook( + request_data=request_data, + original_exception=HTTPException(status_code=400, detail="guardrail rejected"), + user_api_key_dict=UserAPIKeyAuth(token="metric-test-token"), + ) + for metric_name in ("litellm_proxy_failed_requests_metric_total", "litellm_proxy_total_requests_metric_total"): + samples: Final = _collected_samples(metric_name) + assert [(sample.labels["model_id"], sample.value) for sample in samples] == [ + (expected, 1.0) + ] + finally: + _clear_prometheus_registry() + + @pytest.mark.asyncio async def test_failure_hook_labels_a_cli_session_with_the_per_user_alias_not_the_login_token(): from litellm.integrations.prometheus import PrometheusLogger From f106a09102b66d187a27376071a41c3a7d0128bd Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 2 Nov 2026 12:57:56 -0800 Subject: [PATCH 09/79] test(integration): isolate stranger batch authorization from the batch-cost poller (#49522) * test(integration): align stream and A2A contracts and isolate batch authorization Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(sdk): keep EOF assertion compatible with pytest lint Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(sdk): collect cut stream content without mutation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../security/test_files_batches_access.py | 23 +++++++++++-------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/tests/integration/security/test_files_batches_access.py b/tests/integration/security/test_files_batches_access.py index 119055c64c7..93ccc742d58 100644 --- a/tests/integration/security/test_files_batches_access.py +++ b/tests/integration/security/test_files_batches_access.py @@ -25,7 +25,7 @@ from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa from integration._support.client import Gateway, Scenario, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import owned_proxy +from integration._support.process import OwnedState, owned_proxy from integration._support.tls import server_context, write_self_signed_cert from integration._support.wire import Reply, Request, Wire, wire_server from openai import OpenAI, PermissionDeniedError @@ -250,15 +250,20 @@ def test_stranger_cannot_read_download_or_delete_another_teams_managed_file(gate assert _drained_other_than_model_list_probes(wire) == () -def test_stranger_cannot_retrieve_or_cancel_another_teams_managed_batch(gateway: Gateway) -> None: +@pytest.mark.timeout(180) +def test_stranger_cannot_retrieve_or_cancel_another_teams_managed_batch(gateway: Gateway, tmp_path: Path) -> None: key: Final = "provider-key-" + uuid.uuid4().hex[:8] - with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario: + with ( + owned_proxy(gateway, tmp_path, {"PROXY_BATCH_POLLING_ENABLED": "false"}, state=OwnedState.PRIVATE) as candidate, + wire_server(_openai_backend()) as wire, + candidate.scenario() as scenario, + ): model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key) - _wait_until_every_worker_serves(gateway, model) + _wait_until_every_worker_serves(candidate, model) owner: Final = _member(scenario, model) stranger: Final = _member(scenario, model) - uploaded: Final = gateway.request_multipart( + uploaded: Final = candidate.request_multipart( "/v1/files", {"purpose": "batch", "target_model_names": model}, {"file": (INPUT_FILENAME, _jsonl(model), "application/jsonl")}, @@ -266,7 +271,7 @@ def test_stranger_cannot_retrieve_or_cancel_another_teams_managed_batch(gateway: ) assert uploaded.status_code == 200, uploaded.text managed_input: Final = string_value(_json(uploaded)["id"]) - created: Final = gateway.request( + created: Final = candidate.request( "POST", "/v1/batches", {"input_file_id": managed_input, "endpoint": "/v1/chat/completions", "completion_window": "24h"}, @@ -277,13 +282,13 @@ def test_stranger_cannot_retrieve_or_cancel_another_teams_managed_batch(gateway: wire.drain() expected_message: Final = f"User {stranger.user} does not have access to the object {unified_batch}" - stranger_get: Final = gateway.request("GET", f"/v1/batches/{unified_batch}", key=stranger.key) + stranger_get: Final = candidate.request("GET", f"/v1/batches/{unified_batch}", key=stranger.key) assert _error_message(stranger_get, 403) == expected_message, stranger_get.text - stranger_cancel: Final = gateway.request("POST", f"/v1/batches/{unified_batch}/cancel", {}, key=stranger.key) + stranger_cancel: Final = candidate.request("POST", f"/v1/batches/{unified_batch}/cancel", {}, key=stranger.key) assert _error_message(stranger_cancel, 403) == expected_message, stranger_cancel.text assert _drained_other_than_model_list_probes(wire) == (), "stranger batch calls must not reach the provider" - owner_cancel: Final = gateway.request("POST", f"/v1/batches/{unified_batch}/cancel", {}, key=owner.key) + owner_cancel: Final = candidate.request("POST", f"/v1/batches/{unified_batch}/cancel", {}, key=owner.key) assert owner_cancel.status_code == 200, owner_cancel.text assert string_value(_json(owner_cancel)["id"]) == unified_batch, owner_cancel.text (cancel_request,) = _drained_other_than_model_list_probes(wire) From 61da6fcb4dc02cd58e39cb8c57ab3c99bea39a4f Mon Sep 17 00:00:00 2001 From: simonp9 <349053457+simonp9@users.noreply.github.com> Date: Mon, 2 Nov 2026 13:13:03 -0800 Subject: [PATCH 10/79] fix(mcp)!: answer 400 for a hosted MCP tool a translated route cannot run, drop it under drop_params (#48931) * fix(bedrock, gemini, deepseek, databricks, gigachat, anthropic): answer 400 for a hosted MCP tool a route cannot run, drop it under drop_params A `type: "mcp"` tool sent to a route that rebuilds tools for its provider (Bedrock Converse and Nova, Bedrock Invoke Claude, Vertex AI Claude, Vertex AI and Google AI Studio Gemini, DeepSeek, Databricks Claude, GigaChat) was dropped silently and the request still answered 200. These routes now answer UnsupportedParamsError (400) naming the model, route and server, and pointing to the LiteLLM MCP gateway. Under drop_params the tool is dropped with a warning, other tools are kept, and no empty tools field is sent. Anthropic, Azure AI Claude and Claude Platform on AWS keep mapping the tool to the MCP connector. Native Responses providers and raw OpenAI-compatible forwarders are unchanged. * fix(mcp): check hosted MCP tools once per request and keep base tool_choice handling Each route now runs one request-level step before its param mapping. When the drop leaves no tools, the request sends no tools, tool_choice or parallel_tool_calls. Converse and GigaChat go back to their base tool_choice and empty-tools handling. * refactor(deepseek): hold the sent tools in a tuple * fix(mcp): keep tool_choice for legacy functions when every hosted MCP tool drops * fix(mcp): answer 400 for MCP servers a route cannot run on Messages, Ollama and the Responses playground /v1/messages requests with mcp_servers on a chat-adapter or Responses-bridge route now answer 400 naming each server, or drop it with a warning under drop_params, instead of being forwarded or silently ignored. Ollama generate no longer pastes a hosted MCP tool (headers included) into the prompt, and ollama_chat filters hosted MCP tools the same way. The Admin UI Responses playground now sends litellm_proxy/mcp references like the other endpoints. * fix(ollama): leave JSON mode off when drop_params drops every hosted MCP tool Ollama generate set format json for any request with tools, even after the hosted MCP check dropped all of them, so a plain answer came back as JSON or as a call to a tool the model never had. The check now runs before the JSON mode switch, which stays on only while a tool is left for the prompt. * fix(messages): honor additional_drop_params for mcp_servers and read the beta map only for hosted MCP tools * fix(mcp_security): accept gateway server names and toolsets the gateway resolves The guardrail compared litellm_proxy/mcp/ refs against registry keys, which are server ids, so a saved server referenced by name, alias, or toolset name was blocked although the gateway routes it. It now resolves each ref the way the gateway does and still blocks unknown names. * fix(tool_fields): keep MCP drop warnings on one log line Server labels, model names, and routes come from the caller, so a line break in one could forge extra log lines. The drop warnings now strip CR and LF from each value. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../prompt_templates/factory.py | 8 +- litellm/litellm_core_utils/tool_fields.py | 151 +++++++++++++++++- litellm/llms/anthropic/chat/transformation.py | 28 +++- .../anthropic/messages/adapters/handler.py | 39 ++++- .../messages/responses_adapters/handler.py | 17 +- .../bedrock/chat/converse_transformation.py | 16 +- .../bedrock/claude_platform/transformation.py | 3 + .../llms/databricks/chat/transformation.py | 26 +-- litellm/llms/deepseek/chat/transformation.py | 23 +-- litellm/llms/gemini/chat/transformation.py | 4 +- litellm/llms/gigachat/chat/transformation.py | 6 +- litellm/llms/ollama/chat/transformation.py | 13 +- litellm/llms/opencode/chat/transformation.py | 3 +- .../vertex_and_google_ai_studio_gemini.py | 14 +- .../mcp_security/mcp_security_guardrail.py | 19 ++- litellm/types/llms/anthropic.py | 1 + litellm/utils.py | 31 +++- .../litellm_core_utils/test_tool_fields.py | 96 +++++++++++ .../test_anthropic_chat_transformation.py | 79 +++++++++ .../test_handler_output_config_passthrough.py | 116 +++++++++++++- .../test_responses_adapters_handler.py | 51 ++++++ .../chat/test_converse_transformation.py | 64 ++++++++ .../test_databricks_chat_transformation.py | 68 ++++++++ .../chat/test_deepseek_chat_transformation.py | 45 ++++++ .../llms/gigachat/chat/test_transformation.py | 55 +++++++ .../ollama/test_ollama_chat_transformation.py | 67 ++++++++ .../test_ollama_completion_transformation.py | 71 ++++++++ .../llms/opencode/chat/test_transformation.py | 11 ++ ...test_vertex_and_google_ai_studio_gemini.py | 84 ++++++++++ .../guardrail_hooks/test_mcp_security.py | 142 ++++++++-------- .../test_litellm_completion_responses.py | 23 +++ .../llm_calls/responses_api.test.tsx | 44 ++++- .../components/llm_calls/responses_api.tsx | 58 ++----- 33 files changed, 1291 insertions(+), 185 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 98348e434b2..ab1d2dea74e 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -34,6 +34,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import split_leading_system_run, string_field, ) +from litellm.litellm_core_utils.tool_fields import tools_without_hosted_mcp from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client, header_value from litellm.types.files import get_file_extension_from_mime_type @@ -5916,8 +5917,13 @@ def function_call_prompt( llm_provider: str, drop_params: bool, ) -> list[dict[str, object]]: + prompt_functions: Final = tools_without_hosted_mcp( + functions, model=model, route=llm_provider, drop_params=drop_params + ) + if not prompt_functions: + return messages header: Final = """To call a function, reply with JSON ONLY in this format {"name": "function_name", "arguments":{"argument_name": "argument_value"}}. Once a function result answers the request, reply to the user in plain text instead of calling a function again. The following functions are available to you:""" - function_prompt: Final = header + "".join(f"\n{function}\n" for function in functions) + function_prompt: Final = header + "".join(f"\n{function}\n" for function in prompt_functions) if not any(_is_system_message(message) for message in messages): return [*messages, {"role": "system", "content": function_prompt}] return [ diff --git a/litellm/litellm_core_utils/tool_fields.py b/litellm/litellm_core_utils/tool_fields.py index f90c8f5d174..5ee96ea9133 100644 --- a/litellm/litellm_core_utils/tool_fields.py +++ b/litellm/litellm_core_utils/tool_fields.py @@ -1,13 +1,16 @@ -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence +from copy import copy from dataclasses import dataclass from functools import reduce -from typing import Annotated, Final, TypeAlias +from typing import Annotated, Final, TypeAlias, TypeVar from pydantic import BeforeValidator, ConfigDict, Json, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm -from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm._logging import verbose_logger +from litellm.exceptions import UnsupportedParamsError +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, normalize_drop_params ROUTES_WITH_TOOLS: Final = frozenset( { @@ -162,3 +165,145 @@ def resolve_deployment_tool_fields( if "proxy_server_request" not in kwargs or route not in ROUTES_WITH_TOOLS: return kwargs return resolve_tool_fields(kwargs, lambda: (drop_policy(kwargs),)) + + +_ToolT: Final = TypeVar("_ToolT") + +HOSTED_MCP_DROP_WARNING: Final = ( + "Dropping the hosted MCP tool(s) %s for model=%s on %s (drop_params=True) because this route cannot run a " + "hosted MCP tool." +) + + +def _hosted_mcp_server_label(tool: object) -> str | None: + match tool: + case {"type": "mcp", "server_label": str() as server_label}: + return server_label + case {"type": "mcp"}: + return "unlabeled" + case _: + return None + + +def _drop_params_hint(dropped: str) -> str: + return ( + f"To drop {dropped} instead, set `litellm.drop_params=True` or for proxy:\n\n" + "`litellm_settings:\n drop_params: true`\n" + ) + + +def _one_log_line(text: str) -> str: + return text.replace("\r", "").replace("\n", "") + + +def _unrunnable_mcp_error(message: str, model: str, route: str) -> UnsupportedParamsError: + return UnsupportedParamsError(message=message, model=model, llm_provider=route, status_code=400) + + +def tools_without_hosted_mcp(tools: Sequence[_ToolT], model: str, route: str, drop_params: bool) -> Sequence[_ToolT]: + server_labels: Final = tuple(label for label in map(_hosted_mcp_server_label, tools) if label is not None) + if not server_labels: + return tools + named_servers: Final = ", ".join(f"`{label}`" for label in server_labels) + if not drops_unsupported_params(drop_params): + raise _unrunnable_mcp_error( + f'{model} on {route} cannot run the hosted MCP tool {named_servers} (`type: "mcp"`). To have LiteLLM ' + "run the MCP server for this model, register it on the LiteLLM MCP gateway and send " + '`server_url: "litellm_proxy/mcp/"`. ' + _drop_params_hint("hosted MCP tools"), + model=model, + route=route, + ) + verbose_logger.warning( + HOSTED_MCP_DROP_WARNING, _one_log_line(named_servers), _one_log_line(model), _one_log_line(route) + ) + return [tool for tool in tools if _hosted_mcp_server_label(tool) is None] + + +MCP_SERVERS_DROP_WARNING: Final = ( + "Dropping mcp_servers %s and their mcp_toolset tools for model=%s on %s (drop_params=True) because LiteLLM " + "cannot pass mcp_servers to this route." +) + +ROUTES_WITHOUT_MCP_SERVERS: Final = frozenset( + {"bedrock", "databricks", "deepseek", "gemini", "gigachat", "ollama", "ollama_chat", "vertex_ai", "vertex_ai_beta"} +) + + +def _mcp_server_name(server: object) -> str: + match server: + case {"name": str() as name}: + return name + case _: + return "unnamed" + + +def _mcp_server_names(mcp_servers: object) -> tuple[str, ...]: + if not mcp_servers: + return () + if not isinstance(mcp_servers, list): + return ("unnamed",) + return tuple(map(_mcp_server_name, _TOOLS.validate_python(mcp_servers))) + + +def _is_mcp_toolset(tool: object) -> bool: + match tool: + case {"type": "mcp_toolset"}: + return True + case _: + return False + + +def drops_mcp_servers(request_kwargs: Mapping[str, object]) -> bool: + additional_drop_params: Final = request_kwargs.get("additional_drop_params") + return drops_unsupported_params(request_kwargs.get("drop_params")) or ( + isinstance(additional_drop_params, list) and "mcp_servers" in additional_drop_params + ) + + +def tools_without_mcp_servers( + mcp_servers: object, tools: Sequence[_ToolT], model: str, route: str, drop_params: bool +) -> Sequence[_ToolT]: + server_names: Final = _mcp_server_names(mcp_servers) + if not server_names: + return tools + named_servers: Final = ", ".join(f"`{name}`" for name in server_names) + if not drops_unsupported_params(drop_params): + raise _unrunnable_mcp_error( + f"LiteLLM cannot pass `mcp_servers` to {model} on {route}, so the MCP server {named_servers} would not " + 'run. To have LiteLLM run it for this model, register it on the LiteLLM MCP gateway and send a `type: "mcp"` ' + 'tool with `server_url: "litellm_proxy/mcp/"`. ' + _drop_params_hint("`mcp_servers`"), + model=model, + route=route, + ) + verbose_logger.warning( + MCP_SERVERS_DROP_WARNING, _one_log_line(named_servers), _one_log_line(model), _one_log_line(route) + ) + return [tool for tool in tools if not _is_mcp_toolset(tool)] + + +_TOOL_SELECTION_PARAMS: Final = frozenset({"tools", "tool_choice", "parallel_tool_calls"}) +_ParamsT: Final = TypeVar("_ParamsT", bound=dict[str, object]) + + +def tool_params_without_hosted_mcp( + params: _ParamsT, model: str, route: str, drop_params: bool, runs_hosted_mcp: Callable[[], bool] | None = None +) -> _ParamsT: + tools: Final = params.get("tools") + if not isinstance(tools, list) or not tools: + return params + sent_tools: Final = _TOOLS.validate_python(tools) + if all(_hosted_mcp_server_label(tool) is None for tool in sent_tools) or ( + runs_hosted_mcp is not None and runs_hosted_mcp() + ): + return params + kept_tools: Final = tools_without_hosted_mcp(tools=sent_tools, model=model, route=route, drop_params=drop_params) + if len(kept_tools) == len(sent_tools): + return params + kept_params: Final = copy(params) + if kept_tools: + kept_params["tools"] = list(kept_tools) + return kept_params + dropped_keys: Final = ("tools",) if params.get("functions") else _TOOL_SELECTION_PARAMS + for key in dropped_keys: + _ = kept_params.pop(key, None) + return kept_params diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index f198d76f72a..01a4c508cf5 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -13,6 +13,7 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm +from litellm.anthropic_beta_headers_manager import get_unsupported_headers from litellm.constants import ( ANTHROPIC_BILLING_METADATA_PREFIX, ANTHROPIC_MIN_THINKING_BUDGET_TOKENS, @@ -56,6 +57,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import split_leading_system_run, ) from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_non_string_reasoning_effort +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.llms.base_llm.base_utils import ( ResponseFormatSchema, as_json_object, @@ -1315,6 +1317,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): *(tool for tool in tools[first_search_tool_index + 1 :] if not _is_openai_tool_search_tool(tool)), ] + def runs_hosted_mcp(self) -> bool: + return ANTHROPIC_BETA_HEADER_VALUES.MCP_CLIENT_2025_04_04.value not in get_unsupported_headers( + self._resolved_provider + ) + def _tools_param_with_openai_tool_search_mapped( self, non_default_params: Mapping[str, object], model: str, drop_params: bool ) -> Mapping[str, object]: @@ -1903,7 +1910,14 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): model: str, drop_params: bool, ) -> dict: - is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params=non_default_params) + request_params: Final = tool_params_without_hosted_mcp( + non_default_params, + model=model, + route=self._resolved_provider, + drop_params=drop_params, + runs_hosted_mcp=self.runs_hosted_mcp, + ) + is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params=request_params) # NB: ``map_openai_params`` deliberately does NOT sanitize tool names # here. Names are the *original* OpenAI names at this stage, and must @@ -1916,7 +1930,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # via ``data = {**optional_params}`` and would 400 with # ``Extra inputs are not permitted``). - for param, value in non_default_params.items(): + for param, value in request_params.items(): if param == "max_tokens" or param == "max_completion_tokens": optional_params["max_tokens"] = value if isinstance(value, int) else max(1, int(round(value))) elif param == "tools": @@ -1930,8 +1944,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): optional_params["mcp_servers"] = mcp_servers elif param == "tool_choice" or param == "parallel_tool_calls": _tool_choice: AnthropicMessagesToolChoice | None = self._map_tool_choice( - tool_choice=non_default_params.get("tool_choice"), - parallel_tool_use=non_default_params.get("parallel_tool_calls"), + tool_choice=request_params.get("tool_choice"), + parallel_tool_use=request_params.get("parallel_tool_calls"), ) if _tool_choice is not None: @@ -1957,7 +1971,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if self.uses_native_structured_output_for_schema( model, self._extract_json_schema_from_response_format(value), - REQUEST_PARAMS.validate_python(non_default_params), + REQUEST_PARAMS.validate_python(request_params), ): _output_format = self.map_response_format_to_anthropic_output_format(value) if _output_format is not None: @@ -1993,7 +2007,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # `thinking={type: enabled, budget_tokens}` interface a # pre-4.6 model actually supports instead of forwarding a # shape the model will reject. - max_tokens = non_default_params.get("max_completion_tokens") or non_default_params.get("max_tokens") + max_tokens = request_params.get("max_completion_tokens") or request_params.get("max_tokens") legacy_thinking = AnthropicConfig._map_reasoning_effort( reasoning_effort="medium", model=model, @@ -2087,7 +2101,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ## handle thinking tokens self.update_optional_params_with_thinking_tokens( - non_default_params=non_default_params, optional_params=optional_params + non_default_params=request_params, optional_params=optional_params ) return optional_params diff --git a/litellm/llms/anthropic/messages/adapters/handler.py b/litellm/llms/anthropic/messages/adapters/handler.py index d13fd4f47ff..31a4e9d9d73 100644 --- a/litellm/llms/anthropic/messages/adapters/handler.py +++ b/litellm/llms/anthropic/messages/adapters/handler.py @@ -15,6 +15,11 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +from litellm.litellm_core_utils.tool_fields import ( + ROUTES_WITHOUT_MCP_SERVERS, + drops_mcp_servers, + tools_without_mcp_servers, +) from litellm.llms.anthropic.messages.adapters.transformation import ( AnthropicAdapter, ) @@ -610,6 +615,23 @@ class LiteLLMMessagesToCompletionTransformationHandler: Logging as LiteLLMLoggingObject, ) + provider_kwargs: Final = extra_kwargs or {} + custom_llm_provider: Final = provider_kwargs.get("custom_llm_provider") + target_provider: Final = custom_llm_provider if isinstance(custom_llm_provider, str) else None + drop_params: Final = drops_unsupported_params(provider_kwargs.get("drop_params")) + strips_mcp_servers: Final = target_provider in ROUTES_WITHOUT_MCP_SERVERS + request_tools: Final = ( + tools_without_mcp_servers( + provider_kwargs.get("mcp_servers"), + tools or (), + model=model, + route=target_provider, + drop_params=drops_mcp_servers(provider_kwargs), + ) + if strips_mcp_servers + else tools + ) + request_data: Final[dict[str, object]] = { "model": model, "messages": messages, @@ -626,10 +648,10 @@ class LiteLLMMessagesToCompletionTransformationHandler: request_data["temperature"] = temperature if thinking: request_data["thinking"] = thinking - if tool_choice: + if tool_choice and (request_tools or not tools): request_data["tool_choice"] = tool_choice - if tools: - request_data["tools"] = tools + if request_tools: + request_data["tools"] = request_tools if top_k is not None: request_data["top_k"] = top_k if top_p is not None: @@ -647,16 +669,15 @@ class LiteLLMMessagesToCompletionTransformationHandler: if "output_config" in extra_kwargs: request_data["output_config"] = extra_kwargs["output_config"] - custom_llm_provider: Final = extra_kwargs.get("custom_llm_provider") api_base: Final = extra_kwargs.get("api_base") ( openai_request, tool_name_mapping, ) = ANTHROPIC_ADAPTER.translate_completion_input_params_with_tool_mapping( request_data, - custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None, + custom_llm_provider=target_provider, api_base=api_base if isinstance(api_base, str) else None, - drop_params=drops_unsupported_params(extra_kwargs.get("drop_params")), + drop_params=drop_params, ) if openai_request is None: @@ -684,7 +705,11 @@ class LiteLLMMessagesToCompletionTransformationHandler: # Maintainability: when adding a new Anthropic-only request param to # ``AnthropicMessagesRequestOptionalParams``, also extend # ``ANTHROPIC_ONLY_REQUEST_KEYS`` here so it doesn't silently leak. - excluded_keys: Final = ANTHROPIC_ONLY_REQUEST_KEYS | CONVERTED_STREAM_KEYS | {"anthropic_messages"} + excluded_keys: Final = ( + ANTHROPIC_ONLY_REQUEST_KEYS + | CONVERTED_STREAM_KEYS + | ({"anthropic_messages", "mcp_servers"} if strips_mcp_servers else {"anthropic_messages"}) + ) # NOTE: extra_kwargs was already coerced from None to {} at the top of # this method (line ~220). It is guaranteed to be a dict here. for key, value in extra_kwargs.items(): diff --git a/litellm/llms/anthropic/messages/responses_adapters/handler.py b/litellm/llms/anthropic/messages/responses_adapters/handler.py index 08b00dd5fcf..8b638f2c03a 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/handler.py +++ b/litellm/llms/anthropic/messages/responses_adapters/handler.py @@ -24,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( prompt_cache_breakpoint_with_readable_mode, with_prompt_cache_breakpoint, ) +from litellm.litellm_core_utils.tool_fields import drops_mcp_servers, tools_without_mcp_servers from litellm.responses.utils import auto_reasoning_summary from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS from litellm.types.llms.anthropic import ( @@ -272,6 +273,14 @@ def _build_responses_kwargs( """ forwarded_kwargs: Final = _forwarded_kwargs(extra_kwargs) reject_malformed_tool_result_text(messages, model, forwarded_kwargs) + target_provider: Final = forwarded_kwargs.get("custom_llm_provider") + request_tools: Final = tools_without_mcp_servers( + forwarded_kwargs.get("mcp_servers"), + tools or (), + model=model, + route=target_provider if isinstance(target_provider, str) else "this route", + drop_params=drops_mcp_servers(forwarded_kwargs), + ) block_rewrite: Final = _prompt_cache_block_rewrite(model, forwarded_kwargs) request_messages: Final[AnthropicRequestMessages] = ( [ @@ -304,10 +313,10 @@ def _build_responses_kwargs( request_data["temperature"] = temperature if thinking: request_data["thinking"] = thinking - if tool_choice: + if tool_choice and (request_tools or not tools): request_data["tool_choice"] = tool_choice - if tools: - request_data["tools"] = tools + if request_tools: + request_data["tools"] = list(request_tools) if top_p is not None: request_data["top_p"] = top_p if output_format: @@ -347,7 +356,7 @@ def _build_responses_kwargs( responses_kwargs["stream"] = True # Forward litellm-specific kwargs (api_key, api_base, logging obj, etc.) - excluded: Final = CONVERTED_STREAM_KEYS | {"anthropic_messages"} + excluded: Final = CONVERTED_STREAM_KEYS | {"anthropic_messages", "mcp_servers"} for key, value in forwarded_kwargs.items(): if key == "litellm_logging_obj" and value is not None: from litellm.litellm_core_utils.litellm_logging import ( diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 21af97e523e..9a345ccdc9d 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -52,6 +52,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import parts_of, ) from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.llms.anthropic.chat.transformation import ( DROP_UNSUPPORTED_ADAPTIVE_THINKING_WARNING, DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING, @@ -1068,20 +1069,23 @@ class AmazonConverseConfig(BaseConfig): drop_params: bool, model_info: Mapping[str, object] | None = None, ) -> dict: - is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params) + request_params: Final = tool_params_without_hosted_mcp( + non_default_params, model=model, route="bedrock", drop_params=drop_params + ) + is_thinking_enabled: Final = self.is_thinking_enabled(request_params) base_model: Final = BedrockModelInfo.get_base_model(model) drop_thinking_param: Final = self._is_deepseek_model(model=model, base_model=base_model) drop_reasoning_effort_param: Final = self._model_rejects_reasoning_effort_param( model=model, base_model=base_model ) - for param, value in non_default_params.items(): + for param, value in request_params.items(): if param == "response_format" and isinstance(value, dict): optional_params = self._translate_response_format_param( value=value, model=model, optional_params=optional_params, - non_default_params=non_default_params, + non_default_params=request_params, is_thinking_enabled=is_thinking_enabled, model_info=model_info, ) @@ -1135,7 +1139,7 @@ class AmazonConverseConfig(BaseConfig): ), ), model=model, - non_default_params=non_default_params, + non_default_params=request_params, optional_params=optional_params, ) if param == "tool_choice": @@ -1160,7 +1164,7 @@ class AmazonConverseConfig(BaseConfig): and value.get("type") == "adaptive" and not AnthropicConfig.is_adaptive_thinking_model(model, "bedrock") ): - max_tokens = non_default_params.get("max_completion_tokens") or non_default_params.get("max_tokens") + max_tokens = request_params.get("max_completion_tokens") or request_params.get("max_tokens") legacy_thinking = AnthropicConfig.map_reasoning_effort( reasoning_effort="medium", model=model, @@ -1230,7 +1234,7 @@ class AmazonConverseConfig(BaseConfig): # Nova 2 handles token budgeting differently through reasoningConfig if "gpt-oss" not in model and not self._is_nova_2_model(model): self.update_optional_params_with_thinking_tokens( - non_default_params=non_default_params, optional_params=optional_params + non_default_params=request_params, optional_params=optional_params ) final_is_thinking_enabled: Final = self.is_thinking_enabled(optional_params) diff --git a/litellm/llms/bedrock/claude_platform/transformation.py b/litellm/llms/bedrock/claude_platform/transformation.py index a5ecc37d394..6efc4b506e8 100644 --- a/litellm/llms/bedrock/claude_platform/transformation.py +++ b/litellm/llms/bedrock/claude_platform/transformation.py @@ -30,6 +30,9 @@ class BedrockClaudePlatformConfig(BedrockClaudePlatformMixin, AnthropicConfig): def should_strip_billing_metadata(self) -> bool: return True + def runs_hosted_mcp(self) -> bool: + return True + def validate_environment( self, headers: dict, diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 53674907341..b5393fa1603 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_name_from_message, strip_prefill_prefix_field, ) +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.base_utils import ( @@ -467,7 +468,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): use_ai_gateway: Final = self.uses_ai_gateway(model) return join_api_base(self._get_api_base(api_base, use_ai_gateway=use_ai_gateway), "chat/completions") - def get_supported_openai_params(self, model: str | None = None) -> list: + def get_supported_openai_params(self, model: str | None = None) -> list[str]: return [ "stream", "stop", @@ -573,8 +574,13 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): drop_params: bool, replace_max_completion_tokens_with_max_tokens: bool = True, ) -> dict: - is_thinking_enabled: Final = self.is_thinking_enabled(non_default_params) - mapped_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params) + request_params: Final = ( + tool_params_without_hosted_mcp(non_default_params, model=model, route="databricks", drop_params=drop_params) + if "claude" in model + else non_default_params + ) + is_thinking_enabled: Final = self.is_thinking_enabled(request_params) + mapped_params: Final = super().map_openai_params(request_params, optional_params, model, drop_params) if "claude" in model: AnthropicConfig.translate_legacy_thinking_for_adaptive_model( model=model, optional_params=mapped_params, custom_llm_provider="databricks", drop_params=drop_params @@ -583,16 +589,16 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): mapped_params["tools"] = self._map_openai_to_dbrx_tool( model=model, tools=mapped_params["tools"], drop_params=drop_params ) - if "max_completion_tokens" in non_default_params and replace_max_completion_tokens_with_max_tokens: - mapped_params["max_tokens"] = non_default_params[ + if "max_completion_tokens" in request_params and replace_max_completion_tokens_with_max_tokens: + mapped_params["max_tokens"] = request_params[ "max_completion_tokens" ] # most openai-compatible providers support 'max_tokens' not 'max_completion_tokens' mapped_params.pop("max_completion_tokens", None) - if "response_format" in non_default_params and "claude" in model: + if "response_format" in request_params and "claude" in model: _tool: Final = self.map_response_format_to_databricks_tool( model, - non_default_params["response_format"], + request_params["response_format"], mapped_params, is_thinking_enabled, ) @@ -610,8 +616,8 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): "response_format", None ) # unsupported for claude models - if json_schema -> convert to tool call - if "reasoning_effort" in non_default_params and self._uses_anthropic_thinking_param(model): - reasoning_effort_value: Final = non_default_params.get("reasoning_effort") + if "reasoning_effort" in request_params and self._uses_anthropic_thinking_param(model): + reasoning_effort_value: Final = request_params.get("reasoning_effort") mapped_thinking: Final = AnthropicConfig._map_reasoning_effort( reasoning_effort=reasoning_effort_value, model=model, @@ -637,7 +643,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): optional_params.pop("reasoning_effort", None) ## handle thinking tokens self.update_optional_params_with_thinking_tokens( - non_default_params=non_default_params, optional_params=mapped_params + non_default_params=request_params, optional_params=mapped_params ) gated_thinking: Final = strip_beta_gated_thinking( mapped_params.get("thinking"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # mapped_params is the base config's bare dict diff --git a/litellm/llms/deepseek/chat/transformation.py b/litellm/llms/deepseek/chat/transformation.py index 5b7c2a2c410..9c222f03c8c 100644 --- a/litellm/llms/deepseek/chat/transformation.py +++ b/litellm/llms/deepseek/chat/transformation.py @@ -12,6 +12,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_search_results_text, ) from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -33,17 +34,15 @@ def _warn_on_dropped_budget(thinking: object, model: str) -> None: class DeepSeekChatConfig(OpenAIGPTConfig): - def get_supported_openai_params(self, model: str) -> list: + def get_supported_openai_params(self, model: str) -> list[str]: """ DeepSeek reasoner models support thinking parameter. """ - params: Final = super().get_supported_openai_params(model) - params.extend(["thinking", "reasoning_effort"]) - return params + return [*super().get_supported_openai_params(model), "thinking", "reasoning_effort"] def map_openai_params( self, - non_default_params: dict, + non_default_params: dict[str, object], optional_params: dict, model: str, drop_params: bool, @@ -59,7 +58,12 @@ class DeepSeekChatConfig(OpenAIGPTConfig): Reference: https://api-docs.deepseek.com/guides/thinking_mode """ # Let parent handle standard params first - optional_params = super().map_openai_params(non_default_params, optional_params, model, drop_params) + optional_params = super().map_openai_params( + tool_params_without_hosted_mcp(non_default_params, model=model, route="deepseek", drop_params=drop_params), + optional_params, + model, + drop_params, + ) # Pop thinking/reasoning_effort from optional_params first (parent may have added them) # Then re-add only if valid for DeepSeek @@ -265,6 +269,7 @@ class DeepSeekChatConfig(OpenAIGPTConfig): tools: Final = optional_params.get("tools") if not isinstance(tools, list) or not tools: return optional_params + sent_tools: Final[tuple[object, ...]] = tuple(tools) def _is_function_tool(tool: object) -> bool: return isinstance(tool, dict) and tool.get("type") == "function" @@ -289,14 +294,14 @@ class DeepSeekChatConfig(OpenAIGPTConfig): name: Final = function.get("name") return isinstance(name, str) and name in function_tool_names - function_tools: Final = [tool for tool in tools if _is_function_tool(tool)] - if len(function_tools) == len(tools): + function_tools: Final = [tool for tool in sent_tools if _is_function_tool(tool)] + if len(function_tools) == len(sent_tools): return optional_params dropped_types: Final = sorted( { str(tool.get("type")) if isinstance(tool, dict) else type(tool).__name__ - for tool in tools + for tool in sent_tools if not _is_function_tool(tool) } ) diff --git a/litellm/llms/gemini/chat/transformation.py b/litellm/llms/gemini/chat/transformation.py index 936919af83d..1539f13455f 100644 --- a/litellm/llms/gemini/chat/transformation.py +++ b/litellm/llms/gemini/chat/transformation.py @@ -1,5 +1,5 @@ from collections.abc import Mapping -from typing import Final, cast +from typing import ClassVar, Final, cast import litellm from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -49,6 +49,8 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): Note: Please make sure to modify the default parameters as required for your use case. """ + _route: ClassVar[str] = "gemini" + temperature: float | None = None max_output_tokens: int | None = None top_p: float | None = None diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index 81a186eadf4..7bc783b960f 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -17,6 +17,7 @@ import httpx from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.gigachat.utils import convert_usage, get_api_base @@ -169,7 +170,10 @@ class GigaChatConfig(BaseConfig): drop_params: bool, ) -> dict: # mutable-ok: base class contract returns dict """Map OpenAI parameters to GigaChat parameters.""" - for param, value in non_default_params.items(): + request_params: Final = tool_params_without_hosted_mcp( + dict(non_default_params), model=model, route="gigachat", drop_params=drop_params + ) + for param, value in request_params.items(): if param == "stream": optional_params["stream"] = value elif param == "temperature": diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 168b562b358..59020b058a2 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -4,7 +4,7 @@ from collections.abc import AsyncIterator, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final, cast from httpx._models import Headers, Response -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter import litellm from litellm._uuid import uuid @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_reasoning_content, ) from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException @@ -40,6 +41,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_REQUEST_PARAMS: Final = TypeAdapter(dict[str, object]) + class OllamaChatConfig(BaseConfig): """ @@ -152,7 +155,13 @@ class OllamaChatConfig(BaseConfig): model: str, drop_params: bool, ) -> dict: - for param, value in non_default_params.items(): + request_params: Final = tool_params_without_hosted_mcp( + _REQUEST_PARAMS.validate_python(non_default_params), + model=model, + route="ollama_chat", + drop_params=drop_params, + ) + for param, value in request_params.items(): if param == "max_tokens" or param == "max_completion_tokens": optional_params["num_predict"] = value if param == "stream": diff --git a/litellm/llms/opencode/chat/transformation.py b/litellm/llms/opencode/chat/transformation.py index dcb30b5835c..7e16f93d7cb 100644 --- a/litellm/llms/opencode/chat/transformation.py +++ b/litellm/llms/opencode/chat/transformation.py @@ -1,5 +1,5 @@ from collections.abc import AsyncIterator, Iterator, Mapping, Sequence -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, ClassVar, Final import httpx from pydantic import TypeAdapter @@ -159,6 +159,7 @@ class OpenCodeGoMessagesChatConfig(_GoReportedCostMixin, AnthropicConfig): class OpenCodeGoGeminiChatConfig(_GoReportedCostMixin, GoogleAIStudioGeminiConfig): go_usage_key: str = "usageMetadata" + _route: ClassVar[str] = "opencode_go" @property def supports_stream_param_in_request_body(self) -> bool: diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index d0710339464..9f5bec37ac0 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -9,7 +9,7 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence from copy import deepcopy from functools import lru_cache, partial from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, Union, cast, get_args import httpx from pydantic import JsonValue, TypeAdapter @@ -34,6 +34,7 @@ from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAcc from litellm.litellm_core_utils.model_param_helper import embeds_thought_signature_in_tool_call_id from litellm.litellm_core_utils.prompt_templates.factory import encode_tool_call_id_with_signature from litellm.litellm_core_utils.reasoning_effort_utils import warn_dropped_reasoning +from litellm.litellm_core_utils.tool_fields import tool_params_without_hosted_mcp from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -614,6 +615,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): Note: Please make sure to modify the default parameters as required for your use case. """ + _route: ClassVar[str] = "vertex_ai" + temperature: float | None = None max_output_tokens: int | None = None top_p: float | None = None @@ -1667,9 +1670,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): drop_params: bool, allow_integer_enums: bool = True, ) -> dict: - self._apply_include_server_side_tool_invocations(non_default_params, optional_params) + request_params: Final = tool_params_without_hosted_mcp( + non_default_params, model=model, route=self._route, drop_params=drop_params + ) + self._apply_include_server_side_tool_invocations(request_params, optional_params) supports_sampling_params: Final = self.supports_sampling_params(model) - for param, value in non_default_params.items(): + for param, value in request_params.items(): if param in ("temperature", "top_p", "top_k") and not supports_sampling_params: continue if param == "temperature": @@ -1733,7 +1739,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if _tool_choice_value is not None: optional_params["tool_choice"] = _tool_choice_value elif param == "parallel_tool_calls": - tools_list = non_default_params.get("tools", non_default_params.get("functions")) + tools_list = request_params.get("tools", request_params.get("functions")) num_tools = len(tools_list) if isinstance(tools_list, list) else 0 # Gemini does not support parallel_tool_calls=False with multiple # tools. Drop the param instead of failing — Responses API clients diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py index 9c0a574ae3b..b4652e010ff 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_security/mcp_security_guardrail.py @@ -107,9 +107,20 @@ class MCPSecurityGuardrail(CustomGuardrail): from litellm.proxy.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) + from litellm.proxy.proxy_server import prisma_client async with global_mcp_server_manager.catalog.operation(): - registry: Final = global_mcp_server_manager.get_registry() - registered_names: Final = set(registry.keys()) - - return frozenset(requested_servers - registered_names) + not_servers: Final = frozenset( + name + for name in requested_servers + if global_mcp_server_manager.get_mcp_server_by_identifier(name) is None + ) + if prisma_client is None: + return not_servers + return frozenset( + [ + name + for name in not_servers + if await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) is None + ] + ) diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index dc35af81091..5737d717661 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -849,6 +849,7 @@ class ANTHROPIC_BETA_HEADER_VALUES(str, Enum): ADVISOR_TOOL_2026_03_01 = "advisor-tool-2026-03-01" PER_TURN_CONTROL_2026_07_01 = "per-turn-control-2026-07-01" DANGEROUS_TOOL_USE_2026_09_03 = "dangerous-tool-use-2026-09-03" + MCP_CLIENT_2025_04_04 = "mcp-client-2025-04-04" # Tool search beta header constant (for Anthropic direct API and Microsoft Foundry) diff --git a/litellm/utils.py b/litellm/utils.py index 5f40b64e79f..f03b1eb7ad9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -98,7 +98,11 @@ from litellm.litellm_core_utils.model_retirement import is_hidden_as_retired from litellm.litellm_core_utils.offline_mode import is_offline_mode from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens -from litellm.litellm_core_utils.tool_fields import MalformedField, resolve_deployment_tool_fields +from litellm.litellm_core_utils.tool_fields import ( + MalformedField, + resolve_deployment_tool_fields, + tools_without_hosted_mcp, +) from litellm.rust_bridge import tokenizer as tokenizer_dispatch from litellm.rust_bridge.catalog import decision from litellm.rust_bridge.configuration import Decision @@ -5186,7 +5190,12 @@ def remove_sensitive_keys_from_dict(d: dict) -> dict: return d -def pre_process_optional_params(passed_params: dict, non_default_params: dict, custom_llm_provider: str) -> dict: +_TOOL_LIST: Final = TypeAdapter(list[object]) + + +def pre_process_optional_params( + passed_params: dict, non_default_params: dict, custom_llm_provider: str, model: str, drop_params: bool | None +) -> dict: """For .completion(), preprocess optional params""" optional_params: dict = {} @@ -5243,14 +5252,20 @@ def pre_process_optional_params(passed_params: dict, non_default_params: dict, c and custom_llm_provider not in litellm.openai_compatible_providers ): if custom_llm_provider == "ollama": + if "tools" in non_default_params: + sent_tools: Final = _TOOL_LIST.validate_python(non_default_params.pop("tools")) + non_default_params.pop("tool_choice", None) # causes ollama requests to hang + prompt_tools: Final = tools_without_hosted_mcp( + sent_tools, model=model, route=custom_llm_provider, drop_params=bool(drop_params) + ) + if sent_tools and not prompt_tools: + return optional_params + optional_params["functions_unsupported_model"] = prompt_tools + elif "functions" in non_default_params: + optional_params["functions_unsupported_model"] = non_default_params.pop("functions") # ollama actually supports json output optional_params["format"] = "json" litellm.add_function_to_prompt = True # so that main.py adds the function call to the prompt - if "tools" in non_default_params: - optional_params["functions_unsupported_model"] = non_default_params.pop("tools") - non_default_params.pop("tool_choice", None) # causes ollama requests to hang - elif "functions" in non_default_params: - optional_params["functions_unsupported_model"] = non_default_params.pop("functions") elif litellm.add_function_to_prompt: # if user opts to add it to prompt instead optional_params["functions_unsupported_model"] = non_default_params.pop( "tools", non_default_params.pop("functions", None) @@ -5385,6 +5400,8 @@ def get_optional_params( passed_params=passed_params, non_default_params=non_default_params, custom_llm_provider=custom_llm_provider, + model=model, + drop_params=drop_params, ) def _check_valid_arg(supported_params: list[str]): diff --git a/tests/unit/litellm_core_utils/test_tool_fields.py b/tests/unit/litellm_core_utils/test_tool_fields.py index 23d7de98fbc..a99278d18dd 100644 --- a/tests/unit/litellm_core_utils/test_tool_fields.py +++ b/tests/unit/litellm_core_utils/test_tool_fields.py @@ -1,4 +1,5 @@ import datetime +import logging from collections.abc import Callable, Mapping from typing import Final @@ -6,7 +7,13 @@ import pytest import litellm from litellm import Router +from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.tool_fields import ( + tool_params_without_hosted_mcp, + tools_without_hosted_mcp, + tools_without_mcp_servers, +) from litellm.utils import async_pre_call_deployment_hook _GOOD_TOOL: Final = {"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}} @@ -185,3 +192,92 @@ async def test_tools_core_accepts_reach_the_deployment(tools: object) -> None: sent: Final = await async_pre_call_deployment_hook(dict(_proxy_request("keeps", tools=tools)), "acompletion") assert sent["tools"] is tools + + +def test_a_hosted_mcp_error_names_every_server_and_marks_an_unlabeled_one(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError) as raised: + tools_without_hosted_mcp( + tools=[ + {"type": "mcp", "server_label": "deepwiki", "server_url": "https://mcp.deepwiki.com/mcp"}, + {"type": "mcp", "server_url": "https://mcp.example.com/mcp"}, + ], + model="gemini-3.8-flash", + route="gemini", + drop_params=False, + ) + assert "the hosted MCP tool `deepwiki`, `unlabeled` " in str(raised.value) + + +_HOSTED_MCP_TOOL: Final = {"type": "mcp", "server_label": "deepwiki", "server_url": "https://mcp.deepwiki.com/mcp"} + + +def test_dropping_a_hosted_mcp_tool_keeps_the_other_tool_params_and_leaves_the_caller_dict_alone() -> None: + tool_choice: Final = {"type": "function", "function": {"name": "get_weather"}} + sent: Final = { + "tools": [_HOSTED_MCP_TOOL, _GOOD_TOOL], + "tool_choice": tool_choice, + "parallel_tool_calls": False, + "temperature": 0.2, + } + params: Final = tool_params_without_hosted_mcp(sent, model="deepseek-v4-flash", route="deepseek", drop_params=True) + + assert params == { + "tools": [_GOOD_TOOL], + "tool_choice": tool_choice, + "parallel_tool_calls": False, + "temperature": 0.2, + } + assert sent["tools"] == [_HOSTED_MCP_TOOL, _GOOD_TOOL] + + +def test_dropping_every_hosted_mcp_tool_clears_the_tool_selection_params() -> None: + params: Final = tool_params_without_hosted_mcp( + {"tools": [_HOSTED_MCP_TOOL], "tool_choice": "required", "parallel_tool_calls": True, "temperature": 0.2}, + model="deepseek-v4-flash", + route="deepseek", + drop_params=True, + ) + + assert params == {"temperature": 0.2} + + +class _Records(logging.Handler): + def __init__(self) -> None: + super().__init__(level=logging.WARNING) + self.messages: Final[list[str]] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.messages.append(record.getMessage()) + + +@pytest.mark.parametrize( + "drop", + [ + lambda model: tools_without_hosted_mcp( + [{"type": "mcp", "server_label": "wiki\r\nFORGED", "server_url": "https://mcp.example.com/mcp"}], + model=model, + route="deepseek", + drop_params=True, + ), + lambda model: tools_without_mcp_servers( + [{"type": "url", "name": "wiki\r\nFORGED", "url": "https://mcp.example.com/mcp"}], + tools=[], + model=model, + route="bedrock", + drop_params=True, + ), + ], + ids=["hosted_mcp_tool", "mcp_servers"], +) +def test_a_drop_warning_stays_on_one_log_line(drop: Callable[[str], object]) -> None: + records: Final = _Records() + verbose_logger.addHandler(records) + try: + drop("deepseek-v4-flash\nFORGED level=CRITICAL") + finally: + verbose_logger.removeHandler(records) + + assert len(records.messages) == 1 + assert "\n" not in records.messages[0] and "\r" not in records.messages[0] + assert "FORGED level=CRITICAL" in records.messages[0] 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 63a304ab77d..395255f0608 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -10014,6 +10014,85 @@ def test_transform_request_tool_config_point_without_tools_drops_the_key(): assert litellm_params["metadata"] == {} +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} +_ROUTES_WITHOUT_THE_MCP_CONNECTOR: Final = ( + pytest.param(AmazonAnthropicClaudeConfig, "bedrock", id="bedrock_invoke"), + pytest.param(VertexAIAnthropicConfig, "vertex_ai", id="vertex_ai"), +) + + +@pytest.mark.parametrize("config_class", [AnthropicConfig, AzureAnthropicConfig, BedrockClaudePlatformConfig]) +def test_routes_with_the_mcp_connector_send_a_hosted_mcp_tool_as_an_mcp_server(monkeypatch, config_class): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = config_class().map_openai_params( + non_default_params={"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, + optional_params={}, + model="claude-haiku-5-5", + drop_params=False, + ) + assert [server["url"] for server in mapped["mcp_servers"]] == [_HOSTED_MCP_TOOL["server_url"]] + assert [tool["name"] for tool in mapped["tools"]] == ["get_weather"] + + +@pytest.mark.parametrize(("config_class", "route"), _ROUTES_WITHOUT_THE_MCP_CONNECTOR) +def test_a_hosted_mcp_tool_answers_400_on_claude_routes_without_the_mcp_connector(monkeypatch, config_class, route): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + config_class().map_openai_params( + non_default_params={"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, + optional_params={}, + model="claude-haiku-5-5", + drop_params=False, + ) + assert raised.value.status_code == 400 + assert f"claude-haiku-5-5 on {route} " in str(raised.value) + assert "`deepwiki`" in str(raised.value) + + +@pytest.mark.parametrize(("config_class", "route"), _ROUTES_WITHOUT_THE_MCP_CONNECTOR) +def test_claude_routes_without_the_mcp_connector_drop_a_hosted_mcp_tool_under_drop_params( + monkeypatch, caplog, config_class, route +): + monkeypatch.setattr(litellm, "drop_params", False) + with caplog.at_level("WARNING"): + mapped: Final = config_class().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL]}, + optional_params={}, + model="claude-haiku-5-5", + drop_params=True, + ) + assert [tool["name"] for tool in mapped["tools"]] == ["get_weather"] + assert "mcp_servers" not in mapped + assert any( + "`deepwiki`" in record.getMessage() and f" on {route} (" in record.getMessage() for record in caplog.records + ) + + +@pytest.mark.parametrize(("config_class", "route"), _ROUTES_WITHOUT_THE_MCP_CONNECTOR) +def test_claude_routes_without_the_mcp_connector_send_no_tools_when_drop_params_drops_every_one( + monkeypatch, config_class, route +): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = config_class().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL], "tool_choice": "required"}, + optional_params={}, + model="claude-haiku-5-5", + drop_params=True, + ) + assert "tools" not in mapped + assert "mcp_servers" not in mapped + assert "tool_choice" not in mapped + + _KEPT_SEARCH_CALL: Final = { "role": "assistant", "content": None, diff --git a/tests/unit/llms/anthropic/messages/adapters/test_handler_output_config_passthrough.py b/tests/unit/llms/anthropic/messages/adapters/test_handler_output_config_passthrough.py index c34cb686928..12769e78329 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_handler_output_config_passthrough.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_handler_output_config_passthrough.py @@ -27,11 +27,14 @@ Tests cover (consolidating PRs #23706 and #22727): import os import sys +from collections.abc import Mapping from typing import Final from unittest.mock import MagicMock, patch import pytest +import litellm + # Anchor sys.path to this file's location — not the working-directory-relative # pattern Greptile flagged on PR #23706. Resolves correctly regardless of # where pytest is invoked from. @@ -67,8 +70,8 @@ def _call_prepare(extra_kwargs, model="gpt-4o", output_format=None, **overrides) system=None, temperature=None, thinking=None, - tool_choice=None, - tools=None, + tool_choice=overrides.get("tool_choice"), + tools=overrides.get("tools"), top_k=None, top_p=None, output_format=output_format, @@ -263,3 +266,112 @@ async def test_messages_layer_converted_stream_flag_still_gets_a_message_back(co ) assert response["content"][0]["text"] == "hi" + + +_DEEPWIKI_SERVER: Final = {"type": "url", "url": "https://mcp.deepwiki.com/mcp", "name": "deepwiki"} + + +@pytest.mark.parametrize( + ("mcp_servers", "named"), + [ + ([_DEEPWIKI_SERVER], "`deepwiki`"), + ([_DEEPWIKI_SERVER, {"type": "url", "url": "https://mcp.example.com/mcp"}], "`deepwiki`, `unnamed`"), + ({"name": "deepwiki"}, "`unnamed`"), + ], + ids=["named", "one-unnamed", "not-a-list"], +) +def test_mcp_servers_a_chat_route_cannot_run_answer_400_naming_each_server( + monkeypatch: pytest.MonkeyPatch, mcp_servers: object, named: str +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError) as raised: + _call_prepare( + extra_kwargs={"mcp_servers": mcp_servers, "custom_llm_provider": "gemini"}, model="gemini-3.8-flash" + ) + + assert raised.value.status_code == 400 + assert f"to gemini-3.8-flash on gemini, so the MCP server {named} would not run" in str(raised.value) + + +@pytest.mark.parametrize( + ("route", "forwards"), + [("litellm_proxy", True), ("hosted_vllm", True), ("bedrock", False), ("vertex_ai", False)], +) +def test_mcp_servers_reach_routes_that_carry_them_and_answer_400_elsewhere( + monkeypatch: pytest.MonkeyPatch, route: str, forwards: bool +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + extra_kwargs: Final = {"mcp_servers": [_DEEPWIKI_SERVER], "custom_llm_provider": route} + if not forwards: + with pytest.raises(litellm.UnsupportedParamsError, match=f"on {route}, so the MCP server `deepwiki`"): + _call_prepare(extra_kwargs=extra_kwargs, model="some-model") + return + + completion_kwargs, _ = _call_prepare(extra_kwargs=extra_kwargs, model="some-model") + + assert completion_kwargs["mcp_servers"] == [_DEEPWIKI_SERVER] + + +_DEEPWIKI_TOOLSET: Final = {"type": "mcp_toolset", "mcp_server_name": "deepwiki"} +_WEATHER_TOOL: Final = { + "name": "get_weather", + "description": "Get the weather for a city", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}}, +} + + +@pytest.mark.parametrize( + ("tools", "kept_tool_names"), + [([_DEEPWIKI_TOOLSET, _WEATHER_TOOL], ["get_weather"]), ([_DEEPWIKI_TOOLSET], None)], + ids=["with-function", "toolset-only"], +) +def test_drop_params_drops_the_mcp_toolset_tools_of_dropped_mcp_servers( + monkeypatch: pytest.MonkeyPatch, tools: list[dict[str, object]], kept_tool_names: list[str] | None +) -> None: + monkeypatch.setattr(litellm, "drop_params", True) + completion_kwargs, _ = _call_prepare( + extra_kwargs={"mcp_servers": [_DEEPWIKI_SERVER], "custom_llm_provider": "gemini"}, + model="gemini-3.8-flash", + tools=tools, + tool_choice={"type": "auto"}, + ) + + sent_tools: Final = completion_kwargs.get("tools") + sent_names: Final = [tool["function"]["name"] for tool in sent_tools] if sent_tools else None + assert sent_names == kept_tool_names + assert ("tool_choice" in completion_kwargs) == bool(kept_tool_names) + + +_DROP_FORMS: Final[Mapping[str, Mapping[str, object]]] = { + "request": {"drop_params": True}, + "global": {}, + "additional_drop_params": {"additional_drop_params": ["mcp_servers"]}, +} + + +@pytest.mark.parametrize("drop_params", list(_DROP_FORMS)) +def test_drop_params_drops_mcp_servers_instead_of_forwarding_them( + monkeypatch: pytest.MonkeyPatch, drop_params: str +) -> None: + monkeypatch.setattr(litellm, "drop_params", drop_params == "global") + completion_kwargs, _ = _call_prepare( + extra_kwargs={ + "mcp_servers": [_DEEPWIKI_SERVER], + "custom_llm_provider": "gemini", + **_DROP_FORMS[drop_params], + }, + model="gemini-3.8-flash", + ) + + assert "mcp_servers" not in completion_kwargs + assert completion_kwargs["messages"] == MESSAGES + + +@pytest.mark.parametrize("mcp_servers", [None, []], ids=["absent", "empty"]) +def test_no_mcp_servers_needs_no_drop_params(monkeypatch: pytest.MonkeyPatch, mcp_servers: object) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + completion_kwargs, _ = _call_prepare( + extra_kwargs={"mcp_servers": mcp_servers, "custom_llm_provider": "gemini"}, model="gemini-3.8-flash" + ) + + assert "mcp_servers" not in completion_kwargs diff --git a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py index 382917fc756..c06c9d721ac 100644 --- a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_handler.py @@ -764,6 +764,57 @@ async def test_recomputed_cost_of_a_served_bridge_response_prices_the_reasoning_ assert logging_obj.cost_breakdown["reasoning_cost"] == pytest.approx(expected_reasoning) +_DEEPWIKI_SERVER: Final = {"type": "url", "url": "https://mcp.deepwiki.com/mcp", "name": "deepwiki"} + + +def test_mcp_servers_the_responses_bridge_cannot_run_answer_400(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError) as rejected: + _build_responses_kwargs( + model="gpt-5.6", + extra_kwargs={**OPENAI_API, "mcp_servers": [_DEEPWIKI_SERVER]}, + max_tokens=1024, + messages=MESSAGES, + ) + + assert rejected.value.status_code == 400 + assert "to gpt-5.6 on openai, so the MCP server `deepwiki` would not run" in rejected.value.message + + +@pytest.mark.parametrize( + "drop_form", + [{"drop_params": True}, {"additional_drop_params": ["mcp_servers"]}], + ids=["drop_params", "additional_drop_params"], +) +def test_drop_params_keeps_mcp_servers_off_the_responses_request( + monkeypatch: pytest.MonkeyPatch, drop_form: Mapping[str, object] +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + responses_kwargs: Final = _build_responses_kwargs( + model="gpt-5.6", + extra_kwargs={**OPENAI_API, "mcp_servers": [_DEEPWIKI_SERVER], **drop_form}, + max_tokens=1024, + messages=MESSAGES, + ) + + assert "mcp_servers" not in responses_kwargs + assert responses_kwargs["input"] + + +def test_drop_params_drops_the_mcp_toolset_tools_of_dropped_mcp_servers(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + weather_tool: Final = {"name": "get_weather", "input_schema": {"type": "object", "properties": {}}} + responses_kwargs: Final = _build_responses_kwargs( + model="gpt-5.6", + extra_kwargs={**OPENAI_API, "mcp_servers": [_DEEPWIKI_SERVER], "drop_params": True}, + max_tokens=1024, + messages=MESSAGES, + tools=[{"type": "mcp_toolset", "mcp_server_name": "deepwiki"}, weather_tool], + ) + + assert [tool["name"] for tool in responses_kwargs["tools"]] == ["get_weather"] + + TEXT_RESPONSE: Final = { "id": "resp_cleared", "object": "response", diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 9ab7060d766..44d3427be91 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -11097,6 +11097,70 @@ def test_a_deployment_with_a_batch_control_plane_endpoint_keeps_it_out_of_chat_r assert "aws_bedrock_control_plane_endpoint" not in route.calls[0].request.content.decode() +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} + + +def _converse_request_for_tools( + model: str, non_default_params: Mapping[str, object], drop_params: bool +) -> Mapping[str, object]: + config: Final = AmazonConverseConfig() + return config._transform_request_helper( + model=model, + system_content_blocks=[], + optional_params=config.map_openai_params( + non_default_params=dict(non_default_params), optional_params={}, model=model, drop_params=drop_params + ), + messages=[{"role": "user", "content": "hi"}], + ) + + +@pytest.mark.parametrize( + "model", ["us.openai.gpt-oss-120b-1:0", "us.anthropic.claude-haiku-5-5", "us.amazon.nova-2-lite-v1:0"] +) +def test_a_hosted_mcp_tool_answers_400_on_converse_without_drop_params(monkeypatch, model): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + _converse_request_for_tools(model, {"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, drop_params=False) + assert raised.value.status_code == 400 + assert model in str(raised.value) + assert "`deepwiki`" in str(raised.value) + assert 'server_url: "litellm_proxy/mcp/"' in str(raised.value) + assert "drop_params: true" in str(raised.value) + + +def test_converse_drops_a_hosted_mcp_tool_under_drop_params_and_keeps_the_function_tool(monkeypatch, caplog): + monkeypatch.setattr(litellm, "drop_params", False) + model: Final = "us.openai.gpt-oss-120b-1:0" + with caplog.at_level("WARNING"): + request: Final = _converse_request_for_tools( + model, {"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL]}, drop_params=True + ) + assert [tool["toolSpec"]["name"] for tool in request["toolConfig"]["tools"]] == ["get_weather"] + assert any( + "`deepwiki`" in record.getMessage() and f"model={model} on bedrock (" in record.getMessage() + for record in caplog.records + ) + + +def test_converse_sends_no_tools_when_drop_params_drops_every_hosted_mcp_tool(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + request: Final = _converse_request_for_tools( + "us.openai.gpt-oss-120b-1:0", {"tools": [_HOSTED_MCP_TOOL], "tool_choice": "auto"}, drop_params=True + ) + assert "toolConfig" not in request + assert "toolChoice" not in json.dumps(request) + assert "tool_choice" not in json.dumps(request) + + def _warnings_naming(caplog: pytest.LogCaptureFixture, param: str) -> tuple[str, ...]: return tuple( r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index 33c1e7a2e22..d2b78235c29 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -3046,6 +3046,74 @@ def test_claude_json_schema_stream_split_across_json_tools_keeps_the_non_stream_ assert [choice.finish_reason for choice in choices if choice.finish_reason] == [answer.finish_reason] +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} + + +def test_databricks_claude_answers_400_for_a_hosted_mcp_tool_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + DatabricksConfig().map_openai_params( + non_default_params={"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, + optional_params={}, + model="databricks-claude-opus-5-5", + drop_params=False, + ) + assert raised.value.status_code == 400 + assert "databricks-claude-opus-5-5 on databricks " in str(raised.value) + assert "`deepwiki`" in str(raised.value) + + +def test_databricks_claude_drops_a_hosted_mcp_tool_under_drop_params_and_keeps_the_function_tool(monkeypatch, caplog): + monkeypatch.setattr(litellm, "drop_params", False) + with caplog.at_level("WARNING"): + mapped: Final = DatabricksConfig().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL]}, + optional_params={}, + model="databricks-claude-opus-5-5", + drop_params=True, + ) + assert [tool["function"]["name"] for tool in mapped["tools"]] == ["get_weather"] + assert any( + "`deepwiki`" in record.getMessage() and " on databricks (" in record.getMessage() for record in caplog.records + ) + + +def test_databricks_claude_sends_no_tools_when_drop_params_drops_every_hosted_mcp_tool(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = DatabricksConfig().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL], "tool_choice": "required"}, + optional_params={}, + model="databricks-claude-opus-5-5", + drop_params=True, + ) + assert "tools" not in mapped + assert "tool_choice" not in mapped + + +def test_databricks_claude_still_sends_max_completion_tokens_as_max_tokens_when_dropping_a_hosted_mcp_tool( + monkeypatch, +): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = DatabricksConfig().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL], "max_completion_tokens": 64}, + optional_params={}, + model="databricks-claude-opus-5-5", + drop_params=True, + ) + assert mapped["max_tokens"] == 64 + assert "max_completion_tokens" not in mapped + assert [tool["function"]["name"] for tool in mapped["tools"]] == ["get_weather"] + + @pytest.mark.parametrize( ("model", "base_model"), [("databricks-claude-haiku-5-5", None), ("support-bot-endpoint", "databricks/databricks-claude-haiku-5-5")], diff --git a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py index 56af47aeebe..160ed3950b3 100644 --- a/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py +++ b/tests/unit/llms/deepseek/chat/test_deepseek_chat_transformation.py @@ -955,6 +955,51 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str): ) +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} + + +def test_a_hosted_mcp_tool_answers_400_on_deepseek_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + DeepSeekChatConfig().map_openai_params( + {"tools": [_function_tool("shell"), _HOSTED_MCP_TOOL]}, {}, "deepseek-v4-flash", False + ) + assert raised.value.status_code == 400 + assert "deepseek-v4-flash on deepseek " in str(raised.value) + assert "`deepwiki`" in str(raised.value) + + +def test_deepseek_drops_a_hosted_mcp_tool_under_drop_params_and_keeps_the_function_tool(monkeypatch, caplog): + monkeypatch.setattr(litellm, "drop_params", False) + with caplog.at_level("WARNING"): + result: Final = DeepSeekChatConfig().map_openai_params( + {"tools": [_HOSTED_MCP_TOOL, _function_tool("shell")], "tool_choice": "auto"}, {}, "deepseek-v4-flash", True + ) + assert result["tools"] == [_function_tool("shell")] + assert result["tool_choice"] == "auto" + assert any( + "`deepwiki`" in record.getMessage() and " on deepseek (" in record.getMessage() for record in caplog.records + ) + + +def test_deepseek_sends_no_tools_when_drop_params_drops_every_hosted_mcp_tool(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + result: Final = DeepSeekChatConfig().map_openai_params( + {"tools": [_HOSTED_MCP_TOOL], "tool_choice": "required", "parallel_tool_calls": True}, + {}, + "deepseek-v4-flash", + True, + ) + assert "tools" not in result + assert "tool_choice" not in result + assert "parallel_tool_calls" not in result + + def _budget_warnings(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: return tuple( r.getMessage() diff --git a/tests/unit/llms/gigachat/chat/test_transformation.py b/tests/unit/llms/gigachat/chat/test_transformation.py index 9116a4d1979..4d8a916ea87 100644 --- a/tests/unit/llms/gigachat/chat/test_transformation.py +++ b/tests/unit/llms/gigachat/chat/test_transformation.py @@ -7,10 +7,12 @@ Run with: pytest tests/unit/llms/gigachat/chat/test_transformation.py -v import asyncio import importlib +from typing import Final import pytest import litellm +from litellm.llms.gigachat.chat.transformation import GigaChatConfig from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -575,3 +577,56 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), "cohere_key": getattr(litellm, "cohere_key", None), } + + +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} + + +def test_a_hosted_mcp_tool_answers_400_on_gigachat_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + GigaChatConfig().map_openai_params( + non_default_params={"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, + optional_params={}, + model="GigaChat-2-Max", + drop_params=False, + ) + assert raised.value.status_code == 400 + assert "GigaChat-2-Max on gigachat " in str(raised.value) + assert "`deepwiki`" in str(raised.value) + + +def test_gigachat_drops_a_hosted_mcp_tool_under_drop_params_and_keeps_the_function_tool(monkeypatch, caplog): + monkeypatch.setattr(litellm, "drop_params", False) + with caplog.at_level("WARNING"): + mapped: Final = GigaChatConfig().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL]}, + optional_params={}, + model="GigaChat-2-Max", + drop_params=True, + ) + assert [function["name"] for function in mapped["functions"]] == ["get_weather"] + assert any( + "`deepwiki`" in record.getMessage() and " on gigachat (" in record.getMessage() for record in caplog.records + ) + + +def test_gigachat_sends_no_functions_when_drop_params_drops_every_hosted_mcp_tool(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = GigaChatConfig().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL], "tool_choice": "auto"}, + optional_params={}, + model="GigaChat-2-Max", + drop_params=True, + ) + assert "functions" not in mapped + assert "function_call" not in mapped diff --git a/tests/unit/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py index 6a80f29ac43..98ca5aee407 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -6,6 +6,7 @@ import sys from typing import Final, cast import pytest +import respx from pydantic import BaseModel sys.path.insert( @@ -1211,6 +1212,72 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str): ) +_DEEPWIKI_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}}, +} + + +def _ollama_chat_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post("http://ollama.test:11434/api/chat").respond( + json={ + "model": "qwen3:0.6b", + "message": {"role": "assistant", "content": "Sunny"}, + "done": True, + "prompt_eval_count": 5, + "eval_count": 3, + } + ) + + +def test_ollama_chat_answers_400_for_a_hosted_mcp_tool_it_cannot_run( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + route: Final = _ollama_chat_route(respx_mock) + + with pytest.raises(litellm.UnsupportedParamsError) as raised: + litellm.completion( + model="ollama_chat/qwen3:0.6b", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=[_DEEPWIKI_MCP_TOOL, _WEATHER_TOOL], + api_base="http://ollama.test:11434", + ) + + assert raised.value.status_code == 400 + assert "qwen3:0.6b on ollama_chat cannot run the hosted MCP tool `deepwiki`" in str(raised.value) + assert not route.called + + +@pytest.mark.parametrize( + ("tools", "sent_tools"), + [([_DEEPWIKI_MCP_TOOL, _WEATHER_TOOL], [_WEATHER_TOOL]), ([_DEEPWIKI_MCP_TOOL], None)], + ids=["beside-a-function", "alone"], +) +def test_ollama_chat_drop_params_sends_only_the_tools_ollama_can_call( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, tools: list[object], sent_tools: object +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + route: Final = _ollama_chat_route(respx_mock) + + response: Final = litellm.completion( + model="ollama_chat/qwen3:0.6b", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=tools, + api_base="http://ollama.test:11434", + drop_params=True, + ) + + assert json.loads(route.calls.last.request.content).get("tools") == sent_tools + assert response.choices[0].message.content == "Sunny" + + def _warnings_naming(caplog: pytest.LogCaptureFixture, param: str) -> tuple[str, ...]: return tuple( r.getMessage() for r in caplog.records if r.levelno == logging.WARNING and f"Dropping {param}=" in r.getMessage() diff --git a/tests/unit/llms/ollama/test_ollama_completion_transformation.py b/tests/unit/llms/ollama/test_ollama_completion_transformation.py index edea39bb10c..26addf1b24c 100644 --- a/tests/unit/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_completion_transformation.py @@ -993,6 +993,77 @@ def test_get_complete_url_ignores_a_trailing_slash_on_api_base(api_base: str): == "http://gateway.example.com/ollama/api/generate" ) + +_DEEPWIKI_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_TOOL: Final = { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, +} + + +def _ollama_generate_route(respx_mock: respx.MockRouter) -> respx.Route: + return respx_mock.post("http://ollama.test:11434/api/generate").respond( + json={"model": "llama3.2", "response": "Sunny", "done": True, "prompt_eval_count": 5, "eval_count": 3} + ) + + +def test_ollama_answers_400_for_a_hosted_mcp_tool_instead_of_pasting_it_into_the_prompt( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr(litellm, "add_function_to_prompt", False) + route: Final = _ollama_generate_route(respx_mock) + + with pytest.raises(litellm.UnsupportedParamsError) as raised: + litellm.completion( + model="ollama/llama3.2", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=[_DEEPWIKI_MCP_TOOL, _WEATHER_TOOL], + api_base="http://ollama.test:11434", + ) + + assert raised.value.status_code == 400 + assert "cannot run the hosted MCP tool `deepwiki`" in str(raised.value) + assert not route.called + + +@pytest.mark.parametrize( + ("tools", "prompt_names_weather"), + [([_DEEPWIKI_MCP_TOOL, _WEATHER_TOOL], True), ([_DEEPWIKI_MCP_TOOL], False)], + ids=["beside-a-function", "alone"], +) +def test_ollama_drop_params_keeps_a_hosted_mcp_tool_out_of_the_prompt( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, tools: list[object], prompt_names_weather: bool +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setattr(litellm, "add_function_to_prompt", False) + route: Final = _ollama_generate_route(respx_mock) + + litellm.completion( + model="ollama/llama3.2", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=tools, + api_base="http://ollama.test:11434", + drop_params=True, + ) + + body: Final = json.loads(route.calls.last.request.content) + prompt: Final = body["prompt"] + assert (body.get("format") == "json") is prompt_names_weather + assert "deepwiki" not in prompt + assert ("get_weather" in prompt) is prompt_names_weather + assert prompt.startswith("### User:\nweather in Paris?") + assert prompt_names_weather or prompt == "### User:\nweather in Paris?\n\n" + + def _reasoning_effort_warnings(caplog: pytest.LogCaptureFixture) -> tuple[str, ...]: return tuple( r.getMessage() diff --git a/tests/unit/llms/opencode/chat/test_transformation.py b/tests/unit/llms/opencode/chat/test_transformation.py index 0a5e49ca1b2..91361a3d81d 100644 --- a/tests/unit/llms/opencode/chat/test_transformation.py +++ b/tests/unit/llms/opencode/chat/test_transformation.py @@ -354,6 +354,17 @@ def test_gemini_custom_base_normalizes_version_and_models(go_catalog: None, suff assert response.choices[0].message.content == "hello" +def test_gemini_hosted_mcp_error_names_the_opencode_go_route(go_catalog: None) -> None: + with pytest.raises(litellm.UnsupportedParamsError, match="fixture-gemini on opencode_go cannot run"): + litellm.completion( + model="opencode_go/fixture-gemini", + api_key="fake-go-key", + api_base="https://gateway.example/google", + messages=[{"role": "user", "content": "hello"}], + tools=[{"type": "mcp", "server_label": "deepwiki", "server_url": "https://mcp.deepwiki.com/mcp"}], + ) + + def test_gemini_sync_stream_preserves_tool_calls_and_usage(go_catalog: None) -> None: chunk: Final = { "candidates": [ diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index ef71f09fc7c..2eecb83e2df 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -10861,6 +10861,90 @@ async def test_gemini_sends_a_boolean_for_a_gateway_default_written_as_text_or_n assert sent["includeThoughts"] is expected +_HOSTED_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", +} +_WEATHER_FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} +_GEMINI_ROUTES: Final = ( + pytest.param(VertexGeminiConfig, "vertex_ai", id="vertex_ai"), + pytest.param(GoogleAIStudioGeminiConfig, "gemini", id="gemini"), +) + + +@pytest.mark.parametrize(("config_class", "route"), _GEMINI_ROUTES) +def test_a_hosted_mcp_tool_answers_400_on_gemini_without_drop_params(monkeypatch, config_class, route): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError) as raised: + config_class().map_openai_params( + non_default_params={"tools": [_WEATHER_FUNCTION_TOOL, _HOSTED_MCP_TOOL]}, + optional_params={}, + model="gemini-3.8-flash", + drop_params=False, + ) + assert raised.value.status_code == 400 + assert f"gemini-3.8-flash on {route} " in str(raised.value) + assert "`deepwiki`" in str(raised.value) + assert "litellm_proxy/mcp/" in str(raised.value) + + +@pytest.mark.parametrize(("config_class", "route"), _GEMINI_ROUTES) +def test_gemini_drops_a_hosted_mcp_tool_under_drop_params_and_keeps_the_function_tool( + monkeypatch, caplog, config_class, route +): + monkeypatch.setattr(litellm, "drop_params", False) + with caplog.at_level("WARNING"): + mapped: Final = config_class().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL, _WEATHER_FUNCTION_TOOL]}, + optional_params={}, + model="gemini-3.8-flash", + drop_params=True, + ) + assert "deepwiki" not in json.dumps(mapped) + assert "get_weather" in json.dumps(mapped["tools"]) + assert any( + "`deepwiki`" in record.getMessage() and f" on {route} (" in record.getMessage() for record in caplog.records + ) + + +@pytest.mark.parametrize(("config_class", "route"), _GEMINI_ROUTES) +def test_gemini_sends_no_tools_when_drop_params_drops_every_hosted_mcp_tool(monkeypatch, config_class, route): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = config_class().map_openai_params( + non_default_params={"tools": [_HOSTED_MCP_TOOL], "tool_choice": "required"}, + optional_params={}, + model="gemini-3.8-flash", + drop_params=True, + ) + assert "tools" not in mapped + assert "tool_choice" not in mapped + + +@pytest.mark.parametrize(("config_class", "route"), _GEMINI_ROUTES) +def test_gemini_keeps_tool_choice_none_for_legacy_functions_after_dropping_every_hosted_mcp_tool( + monkeypatch, config_class, route +): + monkeypatch.setattr(litellm, "drop_params", False) + mapped: Final = config_class().map_openai_params( + non_default_params={ + "tools": [_HOSTED_MCP_TOOL], + "functions": [_WEATHER_FUNCTION_TOOL["function"]], + "tool_choice": "none", + }, + optional_params={}, + model="gemini-3.8-flash", + drop_params=True, + ) + assert "deepwiki" not in json.dumps(mapped) + assert "get_weather" in json.dumps(mapped["tools"]) + assert mapped["tool_choice"] == {"functionCallingConfig": {"mode": "NONE"}} + + class _BodyInChunks(httpx.SyncByteStream, httpx.AsyncByteStream): def __init__(self, chunks: tuple[bytes, ...]) -> None: self._chunks: Final = chunks diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_mcp_security.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_mcp_security.py index 0e6670c26e7..43238e37f63 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_mcp_security.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_mcp_security.py @@ -6,19 +6,24 @@ and allows requests with only registered servers. Covers both /chat/completions and /responses API paths (same pre_call_hook logic, different call_type). """ -from typing import Literal -from unittest.mock import MagicMock, patch +from typing import Final, Literal +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException import litellm +from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.mcp_security import initialize_guardrail from litellm.proxy.guardrails.guardrail_hooks.mcp_security.mcp_security_guardrail import ( MCPSecurityGuardrail, ) +from litellm.proxy.mcp_server.mcp_server_manager import MCPServerManager from litellm.types.guardrails import Guardrail, GuardrailEventHooks, LitellmParams +from litellm.types.mcp import MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer +from litellm.types.mcp_server.mcp_toolset import MCPToolset @pytest.fixture @@ -60,93 +65,67 @@ class TestExtractMCPServerNames: assert names == set() +def _manager_with_saved_server() -> MCPServerManager: + manager: Final = MCPServerManager() + server: Final = MCPServer( + server_id="3f6b2c1e-saved-zapier", + name="zapier", + server_name="zapier", + alias="zap", + transport=MCPTransport.http, + url="https://upstream.example/mcp", + ) + manager.registry[server.server_id] = server + return manager + + class TestMCPSecurityGuardrailPreCall: @pytest.mark.asyncio - @patch( - "litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager" - ) - async def test_blocks_unregistered_server_chat_completions( - self, mock_manager, guardrail - ): - """Simulates /chat/completions with an unregistered MCP server.""" - mock_manager.get_registry.return_value = {"zapier": MagicMock()} - + @pytest.mark.parametrize("call_type", ["acompletion", "aresponses"]) + async def test_blocks_unregistered_server(self, guardrail, call_type): data = { "tools": [ {"type": "mcp", "server_url": "litellm_proxy/mcp/zapier"}, {"type": "mcp", "server_url": "litellm_proxy/mcp/evil_server"}, ], "model": "gpt-4", - "messages": [{"role": "user", "content": "hi"}], "guardrails": ["test-mcp-security"], } - with pytest.raises(HTTPException) as exc_info: + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", _manager_with_saved_server()), + pytest.raises(HTTPException) as exc_info, + ): await guardrail.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth(), cache=MagicMock(), data=data, - call_type="acompletion", + call_type=call_type, ) assert exc_info.value.status_code == 400 - assert "evil_server" in str(exc_info.value.detail) + assert exc_info.value.detail["unregistered_servers"] == ["evil_server"] @pytest.mark.asyncio - @patch( - "litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager" - ) - async def test_blocks_unregistered_server_responses_api( - self, mock_manager, guardrail - ): - """Simulates /responses with an unregistered MCP server.""" - mock_manager.get_registry.return_value = {"github": MagicMock()} - + @pytest.mark.parametrize("call_type", ["acompletion", "aresponses"]) + @pytest.mark.parametrize("reference", ["zapier", "zap", "3f6b2c1e-saved-zapier"]) + async def test_allows_saved_server_by_any_name_the_gateway_routes(self, guardrail, call_type, reference): data = { - "tools": [ - {"type": "mcp", "server_url": "litellm_proxy/mcp/unknown_server"}, - ], - "model": "gpt-4o", - "input": "What can you do?", - "guardrails": ["test-mcp-security"], - } - - with pytest.raises(HTTPException) as exc_info: - await guardrail.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth(), - cache=MagicMock(), - data=data, - call_type="aresponses", - ) - assert exc_info.value.status_code == 400 - assert "unknown_server" in str(exc_info.value.detail) - - @pytest.mark.asyncio - @patch( - "litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager" - ) - async def test_allows_registered_servers(self, mock_manager, guardrail): - """All MCP servers are registered - request passes through.""" - mock_manager.get_registry.return_value = { - "zapier": MagicMock(), - "github": MagicMock(), - } - - data = { - "tools": [ - {"type": "mcp", "server_url": "litellm_proxy/mcp/zapier"}, - {"type": "mcp", "server_url": "litellm_proxy/mcp/github"}, - ], + "tools": [{"type": "mcp", "server_url": f"litellm_proxy/mcp/{reference}"}], "model": "gpt-4", - "messages": [{"role": "user", "content": "hi"}], "guardrails": ["test-mcp-security"], } - result = await guardrail.async_pre_call_hook( - user_api_key_dict=UserAPIKeyAuth(), - cache=MagicMock(), - data=data, - call_type="acompletion", - ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", _manager_with_saved_server()), + ): + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(), + data=data, + call_type=call_type, + ) assert result == data @pytest.mark.asyncio @@ -222,11 +201,13 @@ async def test_guardrail_observes_saved_server_creation_and_deletion_on_another_ prisma.writer_db = prisma.db prisma.db.litellm_mcpservertable.find_many = AsyncMock(side_effect=([row], [])) prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + prisma.db.litellm_mcptoolsettable.find_first = AsyncMock(return_value=None) data = {"tools": [{"type": "mcp", "server_url": "litellm_proxy/mcp/peer-server"}], "guardrails": ["test-mcp-security"]} with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", manager), ): result = await guardrail.async_pre_call_hook(UserAPIKeyAuth(), MagicMock(), data, call_type) @@ -235,3 +216,34 @@ async def test_guardrail_observes_saved_server_creation_and_deletion_on_another_ await guardrail.async_pre_call_hook(UserAPIKeyAuth(), MagicMock(), data, call_type) assert exc.value.status_code == 400 assert exc.value.detail["unregistered_servers"] == ["peer-server"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["acompletion", "aresponses"]) +async def test_allows_saved_toolset_by_name_and_blocks_unknown_name(guardrail, call_type): + saved_toolset: Final = MCPToolset(toolset_id="7c1d-research", toolset_name="research") + + async def find_toolset(*, where: dict[str, str]) -> MCPToolset | None: + return saved_toolset if where["toolset_name"] == saved_toolset.toolset_name else None + + prisma: Final = MagicMock() + prisma.writer_db = prisma.db + prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + prisma.db.litellm_mcptoolsettable.find_first = AsyncMock(side_effect=find_toolset) + data: Final = { + "tools": [ + {"type": "mcp", "server_url": "litellm_proxy/mcp/research"}, + {"type": "mcp", "server_url": "litellm_proxy/mcp/evil_server"}, + ], + "guardrails": ["test-mcp-security"], + } + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()), + patch("litellm.proxy.mcp_server.mcp_server_manager.global_mcp_server_manager", MCPServerManager()), + pytest.raises(HTTPException) as exc, + ): + await guardrail.async_pre_call_hook(UserAPIKeyAuth(), MagicMock(), data, call_type) + assert exc.value.detail["unregistered_servers"] == ["evil_server"] diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index b4cd606374a..d5619515491 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -7192,6 +7192,29 @@ async def test_bridge_partial_response_on_disconnect_bills_only_the_input_since_ ) == await _input_tokens_billed_on_disconnect(since_compaction) +def test_a_hosted_mcp_tool_sent_to_responses_on_bedrock_converse_answers_400(monkeypatch) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError) as raised: + litellm.responses( + model="bedrock/converse/us.openai.gpt-oss-120b-1:0", + input="What does the deepwiki MCP server say about litellm?", + tools=[ + { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", + } + ], + aws_access_key_id="AKIAHOSTEDMCP", + aws_secret_access_key="hosted-mcp-secret", + aws_region_name="us-east-1", + ) + assert raised.value.status_code == 400 + assert "`deepwiki`" in str(raised.value) + assert "litellm_proxy/mcp/" in str(raised.value) + + _ANTHROPIC_MODEL: Final = "anthropic/claude-sonnet-5-5" _WEATHER_ASK: Final = {"role": "user", "content": "Weather in Paris?"} diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx index f9c5e3b0181..9285a49c93f 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx @@ -207,6 +207,46 @@ describe("responses_api", () => { expect(onTimingData).not.toHaveBeenCalled(); }); + it.each([ + [["__all__"], [], [{ server_label: "litellm", server_url: "litellm_proxy/mcp" }]], + [ + ["server-1"], + [{ server_id: "server-1", server_name: "my server" }], + [{ server_label: "my server", server_url: "litellm_proxy/mcp/my server" }], + ], + ])( + "sends MCP selection %j as gateway references the proxy runs for any model", + async (selected, servers, expected) => { + await makeOpenAIResponsesRequest( + messages, + mockUpdateTextUI, + "gemini-pro", + "test-token", + undefined, // tags + undefined, // signal + undefined, // onReasoningContent + undefined, // onTimingData + undefined, // onUsageData + undefined, // traceId + undefined, // vector_store_ids + undefined, // guardrails + undefined, // policies + selected, + undefined, // previousResponseId + undefined, // onResponseId + undefined, // onMCPEvent + undefined, // codeInterpreterEnabled + undefined, // onCodeInterpreterResult + undefined, // customBaseUrl + servers as never, + ); + + expect(mockResponsesCreate.mock.calls[0][0].tools).toEqual( + expected.map((tool) => ({ type: "mcp", require_approval: "never", ...tool })), + ); + }, + ); + it("should report total latency in both streaming and non-streaming modes", async () => { const onTotalLatency = vi.fn(); const callWithStreaming = (streamingEnabled: boolean) => @@ -475,14 +515,14 @@ describe("responses_api", () => { { type: "mcp", server_label: "Alpha", - server_url: "https://backend.example.com/mcp/Alpha", + server_url: "litellm_proxy/mcp/Alpha", require_approval: "never", allowed_tools: ["toolA"], }, { type: "mcp", server_label: "Beta", - server_url: "https://backend.example.com/mcp/Beta", + server_url: "litellm_proxy/mcp/Beta", require_approval: "never", allowed_tools: ["toolB", "toolC"], }, diff --git a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx index f25ac939331..ed1aeb25074 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx @@ -7,11 +7,12 @@ import type { } from "openai/resources/responses/responses"; import { MessageType } from "../chat_ui/types"; import { TokenUsage } from "../chat_ui/ResponseMetrics"; -import { getGatewayBaseUrl, getProxyBaseUrl } from "@/components/networking"; +import { getGatewayBaseUrl } from "@/components/networking"; import { toast } from "@/lib/toast"; import { extractPromptCacheTokens } from "@/utils/promptCacheUsage"; import { parseUsageCost } from "./usage_cost"; import { buildPlaygroundHeaders, type CustomHeaders } from "./request_headers"; +import { buildMcpToolBlocks, type BuildMcpToolBlocksArgs } from "./mcp_tool_blocks"; import type { MCPEvent } from "@/components/mcp_tools/types"; import { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; import { @@ -132,7 +133,6 @@ export async function makeOpenAIResponsesRequest( } const proxyBaseUrl = customBaseUrl || getGatewayBaseUrl(); - const mcpBaseUrl = customBaseUrl || getProxyBaseUrl(); const headers = buildPlaygroundHeaders(tags, customHeaders); const client = new openai.OpenAI({ @@ -178,53 +178,13 @@ export async function makeOpenAIResponsesRequest( }); // Build tools array - const tools: any[] = []; - - // Add MCP servers if selected - if (selectedMCPServers && selectedMCPServers.length > 0) { - if (selectedMCPServers.includes("__all__")) { - // All MCP Servers selected - const allServersTool = { - type: "mcp", - server_label: "litellm", - server_url: `${mcpBaseUrl}/mcp`, - require_approval: "never", - }; - tools.push(allServersTool); - } else { - // Individual servers/toolsets selected - create one entry per item - selectedMCPServers.forEach((serverId) => { - if (serverId.startsWith("toolset:")) { - // Toolset: same /{name}/mcp pattern as individual servers - const toolsetId = serverId.slice("toolset:".length); - const toolset = mcpToolsets?.find((t) => t.toolset_id === toolsetId); - const toolsetName = toolset?.toolset_name || toolsetId; - const toolsetTool = { - type: "mcp", - server_label: toolsetName, - server_url: `${mcpBaseUrl}/mcp/${encodeURIComponent(toolsetName)}`, - require_approval: "never", - }; - tools.push(toolsetTool); - } else { - const server = mcpServers?.find((s) => s.server_id === serverId); - // Use server_name for both routing and labelling. server_name is the - // unique registered identifier; aliases can collide across servers. - const routeName = server?.server_name || serverId; - const allowedTools = mcpServerToolRestrictions?.[serverId] || []; - - const serverTool = { - type: "mcp", - server_label: routeName, // unique per request — collisions cause silent tool-routing failures - server_url: `${mcpBaseUrl}/mcp/${encodeURIComponent(routeName)}`, - require_approval: "never", - ...(allowedTools.length > 0 ? { allowed_tools: allowedTools } : {}), - }; - tools.push(serverTool); - } - }); - } - } + const mcpSelection: BuildMcpToolBlocksArgs = { + selectedMCPServers, + mcpServers, + mcpToolsets, + mcpServerToolRestrictions, + }; + const tools: any[] = buildMcpToolBlocks(mcpSelection); // Add code_interpreter tool if enabled (OpenAI auto-creates container) if (codeInterpreterEnabled) { From b3c397e964fd528b4043203d127441b84bc694b1 Mon Sep 17 00:00:00 2001 From: "Prabal Pratap Singh (ZeroX)" Date: Tue, 3 Nov 2026 02:56:29 +0530 Subject: [PATCH 11/79] feat(hosted_vllm)!: keep vLLM's stop_reason and token_ids on chat choices and stream chunks (#48393) * feat(hosted_vllm): keep vLLM's stop_reason and token_ids on chat choices Non-streamed hosted_vllm chat completions moved vLLM's choice-level stop_reason and token_ids into provider_specific_fields, so clients written against vLLM could not find them where vLLM puts them. They are now also set on the choice itself, and the copy in provider_specific_fields stays for existing readers. Fixes #47945 * feat(hosted_vllm): keep stop_reason and token_ids on streamed chunks and proxy stream output Streamed chat finish chunks and streamed /v1/completions finish chunks now carry vLLM's stop_reason. Each stream chunk keeps the token_ids vLLM sent with it, including held-back empty chunks, and the proxy no longer strips those chunks as empty usage chunks or drops the fields in its fast serializer. * fix(hosted_vllm): keep the upstream choice index on token-id-only stream chunks * fix(hosted_vllm): keep vLLM choice fields for null, cache hits, n>1, and streamed completions Response conversion keeps stop_reason and token_ids on the choice for any OpenAI compatible response, cache replays carry them, finish fields are recorded per choice index, text completion streams keep token_ids, and the proxy passes only these two keys through, nulls included. * fix(hosted_vllm): send empty text on token-id chunks, keep fields on their own choice, and keep null stop_reason on slow-path streams Token-id-only chunks now carry an empty content or text string, a finish chunk only takes stop_reason and token_ids from its own upstream choice, content chunks no longer repeat stop_reason, the slow stream serializer keeps a null stop_reason or token_ids, and fields cached under provider_specific_fields come back on the choice * fix(hosted_vllm): keep stop_reason on the last content chunk, keep both fields in built stream responses, and key the cache on return_token_ids * fix(hosted_vllm): keep another choice's final token_ids and stop_reason on their own stream chunk With n above 1, a finish-only vLLM chunk for a choice other than the one the single finish chunk carries was dropped or taken as choice 0's finish, losing that choice's last token_ids. It now goes out as an id-only chunk on its own choice index. * fix(hosted_vllm): send the first choice's finish chunk before another choice's later token-id chunks * fix(streaming): keep each streamed text completion choice on its own index * fix(streaming): finish each streamed text completion choice on its own index and null token_ids in redacted logs * fix(proxy): drop token_ids from guarded streams, redact choice-level provider fields, and keep one message_delta per Anthropic stream * fix(proxy): skip token-id-only chunks on guarded streams and put stop_reason on the final text chunk * fix(proxy): drop token_ids from the logged and cached response of a guarded stream * fix(proxy): drop token_ids under per-chunk stream checks and redact text-completion stream choices * fix(caching): key return_token_ids and continuation_token sent in extra_body * fix(streaming): drop token ids from streamed chunks while post_call_rules run vLLM holds text back behind a possible stop string but sends its token ids at once, so an SDK stream sent ids before litellm.post_call_rules checked their text. While those rules are set, streamed choices drop token_ids and chunks that carry only ids are skipped, as the proxy already does ahead of a per-chunk check * fix(streaming): keep vLLM fields when a stream's only choice index is not 0 --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/caching/caching.py | 21 +- litellm/constants.py | 3 + litellm/litellm_core_utils/litellm_logging.py | 1 + .../convert_dict_to_response.py | 43 ++- litellm/litellm_core_utils/redact_messages.py | 30 +- .../litellm_core_utils/stream_usage_chunks.py | 8 + .../streaming_chunk_builder_utils.py | 46 ++- .../litellm_core_utils/streaming_handler.py | 196 ++++++++++- .../messages/adapters/streaming_iterator.py | 15 +- litellm/proxy/proxy_server.py | 47 +++ litellm/proxy/utils.py | 86 ++++- .../text_completions/streaming_iterator.py | 12 +- tests/unit/caching/test_caching.py | 24 ++ .../test_convert_dict_to_response.py | 67 ++++ .../test_redact_messages.py | 77 +++++ .../test_streaming_chunk_builder_utils.py | 89 +++++ .../test_streaming_handler.py | 327 ++++++++++++++++++ .../test_streaming_iterator_combined_chunk.py | 14 + .../test_hosted_vllm_chat_transformation.py | 61 +++- .../proxy_server/test_streaming_helpers.py | 100 +++++- .../test_model_dump_with_preserved_fields.py | 28 ++ .../test_proxy_logging_hook_detection.py | 235 +++++++++++++ ...test_proxy_server_endpoints_and_startup.py | 17 +- .../test_streaming_iterator.py | 148 +++++++- 24 files changed, 1656 insertions(+), 39 deletions(-) diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 91cf9bb2cfe..3b7ea66ad6b 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -45,13 +45,27 @@ from .redis_semantic_cache import RedisSemanticCache from .s3_cache import S3Cache RESPONSE_CACHE_TARGET: Final = "llm_response" -_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token"}) +_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token", "return_token_ids"}) _DECODED_LITERAL: Final = TypeAdapter(object) _REDIS_STARTUP_NODES_ADAPTER: Final = TypeAdapter( list[dict[str, object]], config=ConfigDict(strict=True, hide_input_in_errors=True) ) +def _keyed_provider_params_in_extra_body(kwargs: Mapping[str, object]) -> str: + raw_extra_body: Final = kwargs.get("extra_body") + if not isinstance(raw_extra_body, dict): + return "" + extra_body: Final = cast( # cast-ok: isinstance leaves the key and value types unknown + Mapping[str, object], raw_extra_body + ) + return "".join( + f"extra_body.{param}: {extra_body[param]}" + for param in sorted(_ALWAYS_KEYED_PROVIDER_PARAMS) + if extra_body.get(param) + ) + + @contextmanager def response_cache_phase(operation: Literal["get", "set"]) -> Generator[None]: """The ``cache.get llm_response`` / ``cache.set llm_response`` span a response-cache read or write runs @@ -429,7 +443,10 @@ class Cache: "" if embeds_thought_signature_in_tool_call_id(kwargs) else "thought_signature_in_tool_call_id: False" ) hashed_cache_key = Cache._get_hashed_cache_key( - cache_key + reasoning_summary_material + tool_call_id_shape_material + cache_key + + _keyed_provider_params_in_extra_body(kwargs) + + reasoning_summary_material + + tool_call_id_shape_material ) hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs) verbose_logger.debug( diff --git a/litellm/constants.py b/litellm/constants.py index cd7c0176428..14a2a9f9eca 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -699,6 +699,9 @@ STREAM_SSE_DONE_STRING: Final[str] = "[DONE]" STREAM_SSE_DATA_PREFIX: Final[str] = "data: " STREAM_SSE_KEEPALIVE_PING_CHUNK: Final[str] = 'event: ping\ndata: {"type": "ping"}\n\n' STREAM_SSE_KEEPALIVE_PING_BYTES: Final[bytes] = STREAM_SSE_KEEPALIVE_PING_CHUNK.encode("utf-8") +STREAM_FINISH_CHOICE_FIELDS: Final = ("stop_reason",) +STREAM_CHUNK_CHOICE_FIELDS: Final = ("token_ids",) +VLLM_CHOICE_FIELDS: Final = STREAM_FINISH_CHOICE_FIELDS + STREAM_CHUNK_CHOICE_FIELDS ### SPEND TRACKING ### DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND: Final = get_env_float( "DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND", 0.001400 diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 639003c9a6a..5f909f639c4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -918,6 +918,7 @@ class Logging(LiteLLMLoggingBaseClass): # post_call guardrails have run; the @client decorator then stores the # enqueue closure here instead of firing it immediately. self.defer_async_logging: bool = False + self.drops_stream_token_ids: bool = False self.enqueue_deferred_logging: Callable[[], None] | None = None self._async_success_scheduled: bool = False self._async_success_enqueued: bool = False diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 58fe5b4297f..6bb036e9080 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -3,7 +3,8 @@ import json import re import time import traceback -from collections.abc import Generator, Iterator, Mapping, Sequence +from collections.abc import Generator, Iterable, Iterator, Mapping, Sequence +from types import MappingProxyType from typing import Final, Literal, cast from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config @@ -11,7 +12,12 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger -from litellm.constants import RESPONSE_FORMAT_TOOL_NAME +from litellm.constants import ( + RESPONSE_FORMAT_TOOL_NAME, + STREAM_CHUNK_CHOICE_FIELDS, + STREAM_FINISH_CHOICE_FIELDS, + VLLM_CHOICE_FIELDS, +) from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_reasoning_content, @@ -170,6 +176,28 @@ def _split_assembled_content_for_replay(content: str | None) -> list[str]: return _REPLAY_CONTENT_SLICE_RE.findall(content) +def _drop_choice_fields(choice: StreamingChoices, fields: Iterable[str]) -> None: + for field in fields: + if hasattr(choice, field): + delattr(choice, field) + + +def _vllm_choice_fields(choice: Mapping[str, object]) -> Mapping[str, object]: + cached_fields: Final = _choice_provider_specific_fields(choice) + return MappingProxyType( + { + field: choice[field] if field in choice else cached_fields[field] + for field in VLLM_CHOICE_FIELDS + if field in choice or field in cached_fields + } + ) + + +def _choice_provider_specific_fields(choice: Mapping[str, object]) -> Mapping[str, object]: + fields: Final = choice.get("provider_specific_fields") + return _FIELD_MAP.validate_python(fields) if isinstance(fields, dict) else MappingProxyType({}) + + _REPLAY_REASONING_FIELDS: Final = frozenset({"reasoning_content", "thinking_blocks", "reasoning_items"}) _REPLAY_LEAD_FIELDS: Final = _REPLAY_REASONING_FIELDS | {"role", "provider_specific_fields"} _REPLAY_TOOL_FIELDS: Final = frozenset({"tool_calls", "function_call"}) @@ -220,8 +248,9 @@ def _replay_frame( choice.finish_reason = finish_reason if not keeps_choice_metadata: choice.logprobs = None - if hasattr(choice, "enhancements"): - del choice.enhancements + _drop_choice_fields(choice, ("enhancements", *STREAM_CHUNK_CHOICE_FIELDS)) + if finish_reason is None: + _drop_choice_fields(choice, STREAM_FINISH_CHOICE_FIELDS) return frame @@ -325,7 +354,9 @@ async def convert_to_streaming_response_async( logprobs = choice.get("logprobs", None) - choice = StreamingChoices(finish_reason=finish_reason, index=idx, delta=delta, logprobs=logprobs) + choice = StreamingChoices( + finish_reason=finish_reason, index=idx, delta=delta, logprobs=logprobs, **_vllm_choice_fields(choice) + ) choice_list.append(choice) model_response_object.choices = choice_list @@ -388,6 +419,7 @@ def convert_to_streaming_response( delta=delta, logprobs=logprobs, enhancements=enhancements, + **_vllm_choice_fields(choice), ) choice_list.append(choice) @@ -872,6 +904,7 @@ def convert_to_model_response_object( logprobs=logprobs, enhancements=enhancements, provider_specific_fields=provider_specific_fields, + **_vllm_choice_fields(choice), ) if provider_sent_logprobs and logprobs is None: choice.logprobs = None diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 981859c620b..2f8ecdcc9b9 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -196,6 +196,9 @@ def _redact_function_call(function_call) -> None: def _redact_choice_content(choice): """Helper to redact content in a choice (message or delta).""" + if isinstance(choice, (litellm.Choices, litellm.utils.StreamingChoices, litellm.utils.TextChoices)): + drop_token_ids(choice) + _redact_provider_specific_fields(choice) if isinstance(choice, litellm.Choices): if choice.message.content is not None: choice.message.content = REDACTED_BY_LITELLM @@ -216,6 +219,9 @@ def _redact_choice_content(choice): _redact_provider_specific_fields(choice.delta) _redact_tool_calls(getattr(choice.delta, "tool_calls", None)) _redact_function_call(getattr(choice.delta, "function_call", None)) + elif isinstance(choice, litellm.utils.TextChoices): + _redact_field_when_present(choice, "text") + _redact_field_when_present(choice, "reasoning_content") def _as_mapping(value: object) -> Mapping[str, object] | None: @@ -227,14 +233,27 @@ def _redact_provider_specific_fields(message: object) -> None: setattr(message, "provider_specific_fields", _redacted_provider_specific_fields(fields)) # noqa: B010 # typed as object +def _drop_provider_specific_token_ids(holder: object) -> None: + fields: Final = _as_mapping(getattr(holder, "provider_specific_fields", None)) + if fields is not None and fields.get("token_ids") is not None: + setattr(holder, "provider_specific_fields", {**fields, "token_ids": None}) # noqa: B010 # typed as object + + +def drop_token_ids(choice: object) -> None: + if getattr(choice, "token_ids", None) is not None: + setattr(choice, "token_ids", None) # noqa: B010 # pydantic extra, not a declared field + for holder in (choice, getattr(choice, "delta", None), getattr(choice, "message", None)): + _drop_provider_specific_token_ids(holder) + + def _redact_provider_specific_fields_dict(choice: Mapping[str, object]) -> None: - for message in (choice.get("message"), choice.get("delta")): + for holder in (choice, choice.get("message"), choice.get("delta")): if ( - (readable := _as_mapping(message)) is not None + (readable := _as_mapping(holder)) is not None and (fields := _as_mapping(readable.get("provider_specific_fields"))) is not None - and isinstance(message, dict) + and isinstance(holder, dict) ): - message["provider_specific_fields"] = _redacted_provider_specific_fields(fields) + holder["provider_specific_fields"] = _redacted_provider_specific_fields(fields) def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[str, object]: @@ -251,6 +270,7 @@ def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[ "compaction_blocks", "compaction_start", "compaction_delta", + "token_ids", ) if key in fields }, @@ -422,6 +442,8 @@ def _redact_model_response_dict_choices(choices, redacted_str: str): for choice in choices: if isinstance(choice, dict): _redact_provider_specific_fields_dict(choice) + if choice.get("token_ids") is not None: + choice["token_ids"] = None if "message" in choice and isinstance(choice["message"], dict): if choice["message"].get("content") is not None: choice["message"]["content"] = redacted_str diff --git a/litellm/litellm_core_utils/stream_usage_chunks.py b/litellm/litellm_core_utils/stream_usage_chunks.py index eda57c7ea12..d073ddc2724 100644 --- a/litellm/litellm_core_utils/stream_usage_chunks.py +++ b/litellm/litellm_core_utils/stream_usage_chunks.py @@ -1,6 +1,7 @@ from collections.abc import Iterable from typing import Final +from litellm.constants import VLLM_CHOICE_FIELDS from litellm.types.utils import ModelResponseStream, StreamingChoices @@ -9,12 +10,19 @@ def is_empty_streaming_choice(choice: StreamingChoices) -> bool: return False if getattr(choice, "logprobs", None) is not None: return False + if not _all_values_none(_vllm_choice_values(choice)): + return False delta: Final = choice.delta if hasattr(choice, "delta") else None if delta is None: return True return _all_values_none(delta.model_dump().values()) +def _vllm_choice_values(choice: StreamingChoices) -> Iterable[object]: + choice_extras: Final = choice.model_extra or {} + return (choice_extras.get(field) for field in VLLM_CHOICE_FIELDS) + + def _all_values_none(values: Iterable[object]) -> bool: return all(value is None for value in values) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 7191b2ce8e8..c87357122ce 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -12,6 +12,7 @@ from typing_extensions import ReadOnly, Required import litellm from litellm._logging import verbose_logger +from litellm.constants import VLLM_CHOICE_FIELDS from litellm.litellm_core_utils.core_helpers import ( PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS, UPSTREAM_USAGE_COST, @@ -31,6 +32,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_content_from_model_response, # pyright: ignore[reportUnknownVariableType] # response param accepts an untyped dict ) +from litellm.litellm_core_utils.redact_messages import drop_token_ids from litellm.litellm_core_utils.token_counter import usage_tokenizer_model from litellm.types.llms.openai import ( AllMessageValues, @@ -434,6 +436,8 @@ class ChunkProcessor: chunks, unreported="tool_calls" if _carries_tool_calls(chunks) else "stop" ), message={"role": ChunkProcessor._get_role_from_chunks(chunks), "content": ""}, + provider_specific_fields=None, + **_streamed_vllm_choice_fields(chain.from_iterable(chunk.get("choices") or () for chunk in chunks)), ) @staticmethod @@ -504,6 +508,7 @@ class ChunkProcessor: "index": 0, "message": {"role": role, "content": ""}, "finish_reason": finish_reason, + **_streamed_vllm_choice_fields(_first_choices(chunks)), } ] ), @@ -1371,6 +1376,7 @@ def _rebuilt_text_choice(index: int, choices: Sequence[TextChoices]) -> TextChoi index=index, logprobs=_merged_logprobs(choices), finish_reason=next((choice.finish_reason for choice in reversed(choices) if choice.finish_reason), None), + **_streamed_vllm_choice_fields(choices), ) @@ -1504,6 +1510,31 @@ def _field_of(value: object, name: str) -> object: return getattr(value, name, None) +def _first_choices(chunks: Iterable[object]) -> tuple[object, ...]: + return tuple(next(iter(object_items(_field_of(chunk, "choices"))), None) for chunk in chunks) + + +def _choice_extras(choice: object) -> Mapping[str, object]: + if isinstance(choice, BaseModel): + return choice.model_extra or {} + return object_mapping(choice) or {} + + +def _streamed_vllm_choice_fields(choices: Iterable[object]) -> Mapping[str, object]: + extras: Final = tuple(map(_choice_extras, choices)) + token_id_runs: Final = tuple(object_items(fields.get("token_ids")) for fields in extras) + stop_reasons: Final = tuple(reason for fields in extras if (reason := fields.get("stop_reason")) is not None) + has_token_ids: Final = any(fields.get("token_ids") is not None for fields in extras) + streamed_values: Final[Mapping[str, object]] = MappingProxyType( + { + "token_ids": list(chain.from_iterable(token_id_runs)) if has_token_ids else None, + "stop_reason": stop_reasons[-1] if stop_reasons else None, + } + ) + present_fields: Final = frozenset(chain.from_iterable(fields.keys() for fields in extras)) + return MappingProxyType({field: streamed_values[field] for field in VLLM_CHOICE_FIELDS if field in present_fields}) + + def _delta_reasoning_items(delta: object) -> Sequence[object]: items: Final = _field_of(delta, "reasoning_items") return cast(Sequence[object], items) if isinstance(items, list) else () # cast-ok: items forwarded verbatim @@ -1776,6 +1807,19 @@ def stream_chunk_builder( end_time: object = None, logging_obj: Optional["Logging"] = None, count_prompt_tokens: Callable[[], int] | None = None, +) -> ModelResponse | TextCompletionResponse | None: + response: Final = _assembled_stream_response(chunks, messages, logging_obj, count_prompt_tokens) + if response is not None and getattr(logging_obj, "drops_stream_token_ids", False) is True: + for choice in response.choices: + drop_token_ids(choice) + return response + + +def _assembled_stream_response( + chunks: Sequence[object], + messages: Sequence[AllMessageValues | Message] | None, + logging_obj: Optional["Logging"], + count_prompt_tokens: Callable[[], int] | None, ) -> ModelResponse | TextCompletionResponse | None: try: if chunks is None: @@ -1788,7 +1832,7 @@ def stream_chunk_builder( if not chunks: return None - processor: Final = ChunkProcessor(chunks, messages) + processor: Final = ChunkProcessor(list(chunks), messages) chunks = processor.chunks ### BASE-CASE ### diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 9bb5b761fd5..a9ef06c3b4f 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -23,6 +23,7 @@ from typing_extensions import NotRequired, TypedDict import litellm from litellm import verbose_logger from litellm._uuid import uuid +from litellm.constants import STREAM_CHUNK_CHOICE_FIELDS, STREAM_FINISH_CHOICE_FIELDS from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import ( UPSTREAM_USAGE_COST, @@ -37,10 +38,11 @@ from litellm.litellm_core_utils.hidden_params import ( set_hidden_param, stored_attribute, ) +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_items from litellm.litellm_core_utils.model_response_utils import ( is_model_response_stream_empty, ) -from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject +from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject, drop_token_ids from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer from litellm.types.llms.base import LiteLLMBaseModel @@ -137,6 +139,70 @@ def _prepend_upstream_choice_extras( choice.__pydantic_fields_set__.update(upstream_extras) +_NO_UPSTREAM_CHOICE_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) +_NO_CHOICE_FIELDS_BY_INDEX: Final[Mapping[int, Mapping[str, object]]] = MappingProxyType({}) + + +def _upstream_first_choice(upstream_chunk: object) -> BaseModel | None: + match getattr(upstream_chunk, "choices", None): + case [BaseModel() as first_choice, *_]: + return first_choice + case _: + return None + + +def _upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]: + first_choice: Final = _upstream_first_choice(upstream_chunk) + if first_choice is None: + return _NO_UPSTREAM_CHOICE_FIELDS + upstream_extras: Final[Mapping[str, object]] = first_choice.model_extra or {} + return MappingProxyType({field: upstream_extras[field] for field in fields if field in upstream_extras}) + + +def _upstream_choice_index(upstream_chunk: object) -> int: + upstream_index: Final[object] = getattr(_upstream_first_choice(upstream_chunk), "index", 0) + return upstream_index if isinstance(upstream_index, int) else 0 + + +def _response_choice_index(response_obj: Mapping[str, object]) -> int: + choice_index: Final = response_obj.get("choice_index") + if isinstance(choice_index, int): + return choice_index + return _upstream_choice_index(response_obj.get("original_chunk")) + + +def _non_null_upstream_choice_fields(upstream_chunk: object, fields: Sequence[str]) -> Mapping[str, object]: + first_choice: Final = _upstream_first_choice(upstream_chunk) + upstream_extras: Final = first_choice.model_extra if first_choice is not None else None + if not upstream_extras: + return _NO_UPSTREAM_CHOICE_FIELDS + return MappingProxyType( + {field: upstream_extras[field] for field in fields if upstream_extras.get(field) is not None} + ) + + +def _apply_choice_fields(choice: StreamingChoices, fields: Mapping[str, object]) -> None: + for field, value in fields.items(): + setattr(choice, field, value) + + +def _choice_fields_chunk( + model_response: ModelResponseStream, upstream_index: int, fields: Mapping[str, object] +) -> ModelResponseStream: + model_response.choices[0].delta = Delta(content="") # rebind-ok: chunks are built in place here + model_response.choices[0].index = upstream_index # rebind-ok: chunks are built in place here + _apply_choice_fields(model_response.choices[0], fields) + return model_response + + +def _without_token_ids_under_post_call_rules(chunk: ModelResponseStream | None) -> ModelResponseStream | None: + if chunk is None or not litellm.post_call_rules: + return chunk + for choice in chunk.choices: + drop_token_ids(choice) + return chunk + + def _holds_only_scalars(fields: Iterable[tuple[str, object]]) -> bool: return all(value is None or isinstance(value, (str, int, float, bool)) for _, value in fields) @@ -378,6 +444,8 @@ class CustomStreamWrapper: self._provider_response_model: str | None = None self._upstream_provider: str | None = None self.received_finish_reason: str | None = None + self.received_finish_choice_fields: Mapping[int, Mapping[str, object]] = _NO_CHOICE_FIELDS_BY_INDEX + self.carried_chunk_choice_fields: Mapping[int, Mapping[str, object]] = _NO_CHOICE_FIELDS_BY_INDEX self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream self.received_any_finish_reason = False # a later chunk can reset intermittent_finish_reason self.special_tokens = [ @@ -823,6 +891,7 @@ class CustomStreamWrapper: "logprobs": logprobs, "original_chunk": str_line, "usage": usage, + "chunk_choice_fields": _non_null_upstream_choice_fields(str_line, STREAM_CHUNK_CHOICE_FIELDS), } except Exception as e: raise e @@ -847,7 +916,7 @@ class CustomStreamWrapper: except Exception as e: raise e - def handle_openai_text_completion_chunk(self, chunk): + def handle_openai_text_completion_chunk(self, chunk: object): try: text = "" is_finished = False @@ -860,11 +929,15 @@ class CustomStreamWrapper: is_finished = True finish_reason = choices[0].finish_reason usage = getattr(chunk, "usage", None) + choice_index: Final = _upstream_choice_index(chunk) return { "text": text, "is_finished": is_finished, "finish_reason": finish_reason, "usage": usage, + "chunk_choice_fields": _non_null_upstream_choice_fields(chunk, STREAM_CHUNK_CHOICE_FIELDS), + "choice_index": choice_index, + "sibling_finish_reason": finish_reason if choice_index != 0 else None, } except Exception as e: @@ -1183,6 +1256,13 @@ class CustomStreamWrapper: ) is_chunk_non_empty: Final = self.is_chunk_non_empty(completion_obj, model_response, response_obj) + upstream_chunk_choice_fields: Final[Mapping[str, object]] = ( + _NO_UPSTREAM_CHOICE_FIELDS + if litellm.post_call_rules + else response_obj.get("chunk_choice_fields", _NO_UPSTREAM_CHOICE_FIELDS) + ) + chunk_choice_fields: Final = self._take_carried_chunk_choice_fields(response_obj, upstream_chunk_choice_fields) + sibling_finish_reason: Final[str | None] = response_obj.get("sibling_finish_reason") if is_chunk_non_empty: # cannot set content of an OpenAI Object to be an empty string self.raise_on_model_repetition() @@ -1270,16 +1350,40 @@ class CustomStreamWrapper: if response_obj.get("provider_specific_fields") is not None: completion_obj["provider_specific_fields"] = response_obj["provider_specific_fields"] model_response.choices[0].delta = Delta(**completion_obj) - _index: Final[int | None] = completion_obj.get("index") + _index: Final[int | None] = completion_obj.get("index", response_obj.get("choice_index")) if _index is not None: model_response.choices[0].index = _index + _apply_choice_fields(model_response.choices[0], chunk_choice_fields) + if sibling_finish_reason is not None: + self._finish_sibling_choice(model_response.choices[0], sibling_finish_reason) + elif response_obj.get("is_finished"): + self._apply_finish_choice_fields(model_response.choices[0]) self._optional_combine_thinking_block_in_choices(model_response=model_response) return model_response else: return + elif sibling_finish_reason is not None: + sibling_chunk: Final = _choice_fields_chunk( + model_response, _response_choice_index(response_obj), chunk_choice_fields + ) + self._finish_sibling_choice(sibling_chunk.choices[0], sibling_finish_reason) + return sibling_chunk elif self.received_finish_reason is not None: + chunk_upstream_index: Final = _response_choice_index(response_obj) + if chunk_choice_fields and self._streams_past_the_pending_finish( + model_response, response_obj, chunk_upstream_index + ): + self.carried_chunk_choice_fields = MappingProxyType( + {**self.carried_chunk_choice_fields, chunk_upstream_index: chunk_choice_fields} + ) + else: + another_choice_fields_chunk: Final = self._another_choice_fields_chunk( + model_response, response_obj, chunk_choice_fields + ) + if another_choice_fields_chunk is not None: + return another_choice_fields_chunk if self.sent_last_chunk is True: # Bedrock returns the guardrail trace in the last chunk - we want to return this here if self.custom_llm_provider == "bedrock" and "trace" in model_response: @@ -1317,12 +1421,20 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = self._normalize_tool_call_finish_reason( map_finish_reason(finish_reason=self.received_finish_reason) ) # ensure consistent output to openai + self._apply_finish_choice_fields(model_response.choices[0]) + if ( + response_obj.get("is_finished") + and _response_choice_index(response_obj) == model_response.choices[0].index + ): + _apply_choice_fields(model_response.choices[0], chunk_choice_fields) self.sent_last_chunk = True return model_response elif self._has_special_delta_content(model_response): return self._handle_special_delta_content(model_response) + elif chunk_choice_fields: + return _choice_fields_chunk(model_response, _response_choice_index(response_obj), chunk_choice_fields) else: if hasattr(model_response, "usage"): self.chunks.append(model_response) @@ -1545,6 +1657,8 @@ class CustomStreamWrapper: completion_obj["content"] = response_obj["text"] print_verbose(f"completion obj content: {completion_obj['content']}") if response_obj["is_finished"]: + self._record_finish_choice_fields(chunk) + if response_obj["is_finished"] and response_obj["sibling_finish_reason"] is None: self.received_finish_reason = response_obj["finish_reason"] if response_obj["usage"] is not None: _text_completion_usage: Final[Usage] = response_obj["usage"] @@ -1608,6 +1722,7 @@ class CustomStreamWrapper: self.system_fingerprint = cached_chunk.system_fingerprint if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] + self._record_finish_choice_fields(cached_chunk) else: # openai / azure chat model self.uses_openai_chat_stream = True if isinstance(chunk, ChatCompletionChunk): @@ -1637,6 +1752,7 @@ class CustomStreamWrapper: f"{self.custom_llm_provider} raised a streaming error - finish_reason: error, no content string given. Received Chunk={response_obj}" ) self.received_finish_reason = response_obj["finish_reason"] + self._record_finish_choice_fields(chunk) if response_obj.get("original_chunk", None) is not None: if hasattr(response_obj["original_chunk"], "id"): model_response = self.set_model_id(response_obj["original_chunk"].id, model_response) @@ -1704,7 +1820,7 @@ class CustomStreamWrapper: completion_obj=completion_obj, ) if isinstance(dispatch_result, _ProviderChunkEarlyReturn): - return dispatch_result.value + return _without_token_ids_under_post_call_rules(dispatch_result.value) response_obj = dispatch_result.response_obj model_response.model = self._served_openrouter_model() or self.model @@ -1794,7 +1910,7 @@ class CustomStreamWrapper: response_obj=response_obj, plain_delta_created=plain_delta_created, ) - return result + return _without_token_ids_under_post_call_rules(result) except StopIteration: raise StopIteration @@ -1987,6 +2103,75 @@ class CustomStreamWrapper: return "tool_calls" return finish_reason + def _record_finish_choice_fields(self, upstream_chunk: object) -> None: + self.received_finish_choice_fields = MappingProxyType( + { + **self.received_finish_choice_fields, + _upstream_choice_index(upstream_chunk): _upstream_choice_fields( + upstream_chunk, STREAM_FINISH_CHOICE_FIELDS + ), + } + ) + + def _apply_finish_choice_fields(self, choice: StreamingChoices) -> None: + _apply_choice_fields(choice, self.received_finish_choice_fields.get(choice.index, _NO_UPSTREAM_CHOICE_FIELDS)) + + def _finish_sibling_choice(self, choice: StreamingChoices, finish_reason: str) -> None: + choice.finish_reason = map_finish_reason(finish_reason) # rebind-ok: chunks are built in place here + self._apply_finish_choice_fields(choice) + + def _streams_past_the_pending_finish( + self, model_response: ModelResponseStream, response_obj: Mapping[str, object], upstream_index: int + ) -> bool: + finishing_index: Final = model_response.choices[0].index + return ( + not self.sent_last_chunk + and not response_obj.get("is_finished") + and upstream_index != finishing_index + and finishing_index in self.received_finish_choice_fields + ) + + def _take_carried_chunk_choice_fields( + self, response_obj: Mapping[str, object], chunk_choice_fields: Mapping[str, object] + ) -> Mapping[str, object]: + if not self.carried_chunk_choice_fields: + return chunk_choice_fields + upstream_index: Final = _response_choice_index(response_obj) + carried: Final = self.carried_chunk_choice_fields.get(upstream_index) + if carried is None: + return chunk_choice_fields + self.carried_chunk_choice_fields = MappingProxyType( + {index: fields for index, fields in self.carried_chunk_choice_fields.items() if index != upstream_index} + ) + return MappingProxyType( + { + field: [*object_items(carried.get(field)), *object_items(chunk_choice_fields.get(field))] + for field in STREAM_CHUNK_CHOICE_FIELDS + if field in carried or field in chunk_choice_fields + } + ) + + def _another_choice_fields_chunk( + self, + model_response: ModelResponseStream, + response_obj: Mapping[str, object], + chunk_choice_fields: Mapping[str, object], + ) -> ModelResponseStream | None: + upstream_index: Final = _response_choice_index(response_obj) + is_finished: Final = bool(response_obj.get("is_finished")) + if is_finished and upstream_index == model_response.choices[0].index and not self.sent_last_chunk: + return None + finish_fields: Final = ( + self.received_finish_choice_fields.get(upstream_index, _NO_UPSTREAM_CHOICE_FIELDS) + if is_finished + else _NO_UPSTREAM_CHOICE_FIELDS + ) + if not chunk_choice_fields and all(value is None for value in finish_fields.values()): + return None + return _choice_fields_chunk( + model_response, upstream_index, MappingProxyType({**finish_fields, **chunk_choice_fields}) + ) + def finish_reason_handler(self): model_response: Final = self.model_response_creator( hidden_params=None if self.received_model_id is None else {"received_model_id": self.received_model_id} @@ -2005,6 +2190,7 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason = self._normalize_tool_call_finish_reason( _finish_reason if _finish_reason is not None else "stop" ) + self._apply_finish_choice_fields(model_response.choices[0]) close_tag: Final = self._close_open_think_block() if close_tag is not None: model_response.choices[0].delta.content = close_tag diff --git a/litellm/llms/anthropic/messages/adapters/streaming_iterator.py b/litellm/llms/anthropic/messages/adapters/streaming_iterator.py index 61d23c2de0d..c41deb00f1d 100644 --- a/litellm/llms/anthropic/messages/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/messages/adapters/streaming_iterator.py @@ -493,10 +493,23 @@ class _CombinedChunkSplitter: for position, block in enumerate(blocks) ) + @staticmethod + def _without_sibling_finish(chunk: "ModelResponseStream") -> "ModelResponseStream": + choices: Final = _optional_attr_sequence(chunk, "choices") + if ( + len(choices) != 1 + or _optional_attr(choices[0], "index") in (None, 0) + or _optional_attr(choices[0], "finish_reason") is None + ): + return chunk + unfinished: Final = copy.deepcopy(chunk) + unfinished.choices[0].finish_reason = None + return unfinished + def _expand(self, chunk: "ModelResponseStream") -> "tuple[ModelResponseStream, ...]": split_chunks: Final = tuple( split_chunk - for combined_chunk in self._split(chunk) + for combined_chunk in self._split(self._without_sibling_finish(chunk)) for split_chunk in self._split_by_payload_kind(combined_chunk) # comprehension-ok: flattens nested splits ) block_chunks: Final = tuple(itertools.chain.from_iterable(map(self._split_thinking_blocks, split_chunks))) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0e35560f30f..807f9a0860d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -90,6 +90,7 @@ from litellm.constants import ( LITELLM_UI_SESSION_DURATION, OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS, RUNTIME_UPDATABLE_ROUTER_SETTINGS, + VLLM_CHOICE_FIELDS, ) from litellm.litellm_core_utils.access_windows import access_windows_config_error from litellm.litellm_core_utils.add_retry_fallback_headers import ( @@ -10226,6 +10227,11 @@ def _restamp_nested_response_model_in_sse_line(line: bytes, target_model: str) - return _SSE_DATA_FIELD + b" " + serialized_event.encode() + line_ending +def _vllm_choice_fields(choice: BaseModel) -> Mapping[str, object]: + choice_extras: Final = choice.model_extra or {} + return MappingProxyType({field: choice_extras[field] for field in VLLM_CHOICE_FIELDS if field in choice_extras}) + + def _json_encodes(value: object) -> bool: try: to_jsonable_python(value) @@ -10281,11 +10287,52 @@ def _serialize_chat_chunk(chunk: ModelResponseStream) -> bytes: def _serialize_streaming_chunk(chunk: BaseModel) -> bytes: + choices: Final = _streaming_chunk_choices(chunk) + null_choice_fields: Final = tuple(_null_vllm_choice_fields(choice) for choice in choices) + if any(null_choice_fields): + serialized_chunk: Final = _serialize_keeping_null_choice_fields(chunk, choices, null_choice_fields) + if serialized_chunk is not None: + return serialized_chunk if isinstance(chunk, ModelResponseStream): return _serialize_chat_chunk(chunk) return chunk.model_dump_json(exclude_none=True, exclude_unset=True).encode() +def _serialize_keeping_null_choice_fields( + chunk: BaseModel, choices: Sequence[BaseModel], null_choice_fields: Sequence[Mapping[str, None]] +) -> bytes | None: + try: + return json.dumps( + { + **_dump_set_fields(chunk, exclude=frozenset({"choices"})), + "choices": [ + {**_dump_set_fields(choice), **nulls} + for choice, nulls in zip(choices, null_choice_fields, strict=True) + ], + }, + separators=(",", ":"), + ensure_ascii=False, + ).encode() + except ValueError: + return None + + +def _streaming_chunk_choices(chunk: BaseModel) -> Sequence[BaseModel]: + match chunk: + case ModelResponseStream() | TextCompletionResponse(): + return chunk.choices + case _: + return () + + +def _null_vllm_choice_fields(choice: BaseModel) -> Mapping[str, None]: + return MappingProxyType({field: None for field, value in _vllm_choice_fields(choice).items() if value is None}) + + +def _dump_set_fields(model: BaseModel, exclude: AbstractSet[str] = frozenset()) -> Mapping[str, object]: + return model.model_dump(mode="json", exclude_none=True, exclude_unset=True, exclude=set(exclude)) + + async def _apply_streaming_chunk_hooks( *, chunk: Any, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a1b36e5b7a1..0668c882a67 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -68,6 +68,7 @@ from litellm.constants import ( SPEND_LOG_QUEUE_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS, + VLLM_CHOICE_FIELDS, ) from litellm.cost_calculator import BaseTokenUsageProcessor from litellm.litellm_core_utils.bug_report import ( @@ -170,6 +171,7 @@ from litellm.litellm_core_utils.litellm_logging import ( deployment_pricing_model_info, ) from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.litellm_core_utils.redact_messages import drop_token_ids from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -179,6 +181,7 @@ from litellm.litellm_core_utils.served_output_texts import ( ) from litellm.litellm_core_utils.token_counter import offload_token_count, usage_tokenizer_model from litellm.llms import load_guardrail_translation_mappings +from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( AlertType, @@ -707,6 +710,56 @@ def _pipeline_step_streams_live(step: "PipelineStep") -> bool: return not unified_guardrail.resolve_streaming_flag(callback, "streaming_buffer_until_moderated", True) +def _drop_stream_chunk_token_ids(chunk: object) -> None: + for choice in stream_item_items(chunk, "choices"): + drop_token_ids(choice) + + +def _drop_built_stream_token_ids(request_data: Mapping[str, object]) -> None: + logging_obj: Final = request_data.get("litellm_logging_obj") + if isinstance(logging_obj, Logging): + logging_obj.drops_stream_token_ids = True + + +def _drop_token_ids_ahead_of_chunk_check( + callback: CustomLogger, chunk: object, request_data: Mapping[str, object] +) -> None: + if _overrides_hook(callback, "async_post_call_streaming_hook"): + _drop_stream_chunk_token_ids(chunk) + _drop_built_stream_token_ids(request_data) + + +_PLAIN_STREAM_FIELDS: Final = TypeAdapter(dict[str, object]) +_EMPTY_STREAM_VALUES: Final[tuple[object, ...]] = (None, "", [], {}) + + +def _holds_stream_payload(value: object) -> bool: + if isinstance(value, dict): + nested: Final = _PLAIN_STREAM_FIELDS.validate_python(value).values() + return any(item not in _EMPTY_STREAM_VALUES for item in nested) + return value not in _EMPTY_STREAM_VALUES + + +def _choice_carries_only_token_ids(choice: object) -> bool: + if stream_item_field(choice, "token_ids") is None or not isinstance(choice, BaseModel): + return False + fields: Final = _PLAIN_STREAM_FIELDS.validate_python(choice.model_dump()) + return not any(_holds_stream_payload(value) for key, value in fields.items() if key not in ("index", "token_ids")) + + +def _is_token_id_only_chunk(chunk: object) -> bool: + choices: Final = stream_item_items(chunk, "choices") + return bool(choices) and all(_choice_carries_only_token_ids(choice) for choice in choices) + + +async def _without_token_id_only_chunks(response: AsyncIterable[object]) -> AsyncGenerator[object, None]: + async for chunk in response: + if _is_token_id_only_chunk(chunk): + _drop_stream_chunk_token_ids(chunk) + continue + yield chunk + + def _pipelines_stream_live(pipelines: Sequence[tuple[str, "GuardrailPipeline"]]) -> bool: return all(_pipeline_step_streams_live(step) for step in _pipeline_steps(pipelines)) @@ -4710,6 +4763,7 @@ class ProxyLogging: else: _callback = callback if _callback is not None and isinstance(_callback, CustomLogger): + _drop_token_ids_ahead_of_chunk_check(_callback, response, data) if str_so_far is not None: complete_response = str_so_far + response_str else: @@ -4731,7 +4785,7 @@ class ProxyLogging: self, response, user_api_key_dict: UserAPIKeyAuth, - request_data: dict, + request_data: dict[str, object], ): """ Allow user to modify outgoing streaming data -> Given a whole response iterator. @@ -4765,20 +4819,26 @@ class ProxyLogging: # Merge model-level guardrails before checking which guardrails to run request_data = check_and_merge_model_level_guardrails(data=request_data, llm_router=llm_router) - current_response = response stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict) pipeline_gated_names: Final = _pipeline_step_guardrail_names(post_call_pipelines) + running_overrides: Final = tuple( + (callback, kind) + for callback, kind in caps.iterator_overrides + if not isinstance(callback, CustomGuardrail) + or ( + callback.guardrail_name not in pipeline_gated_names + and callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is True + ) + ) + drops_token_ids: Final = bool(post_call_pipelines) or any( + isinstance(callback, CustomGuardrail) for callback, _ in running_overrides + ) + if drops_token_ids: + _drop_built_stream_token_ids(request_data) + current_response = _without_token_id_only_chunks(response) if drops_token_ids else response guarded_layers: Final[list[AsyncGenerator[object, None]]] = [] # mutable-ok: closed on disconnect - for resolved_callback, kind in caps.iterator_overrides: - if isinstance(resolved_callback, CustomGuardrail): - if resolved_callback.guardrail_name in pipeline_gated_names: - continue - if ( - resolved_callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) - is not True - ): - continue + for resolved_callback, kind in running_overrides: effective_kind = ( "apply_guardrail" if ( @@ -4836,6 +4896,8 @@ class ProxyLogging: try: async for chunk in current_response: served_chunks.append(chunk) + if drops_token_ids: + _drop_stream_chunk_token_ids(served_chunks[-1]) yield chunk except (GeneratorExit, asyncio.CancelledError): await ProxyLogging._close_guarded_layers(guarded_layers) @@ -10203,7 +10265,7 @@ _PRESERVED_NONE_FIELDS: Final[list[tuple[str, str]]] = [ ("message", "refusal"), ("delta", "content"), # null in streaming chunks ] -_PRESERVED_NONE_CHOICE_FIELDS: Final = ("logprobs",) +_PRESERVED_NONE_CHOICE_FIELDS: Final = ("logprobs", *VLLM_CHOICE_FIELDS) _PRESERVED_NONE_RESPONSE_FIELDS: Final = ("system_fingerprint",) diff --git a/litellm/text_completions/streaming_iterator.py b/litellm/text_completions/streaming_iterator.py index 61e47e9d7af..6bba56a5afe 100644 --- a/litellm/text_completions/streaming_iterator.py +++ b/litellm/text_completions/streaming_iterator.py @@ -1,5 +1,6 @@ from typing import Final +from litellm.constants import VLLM_CHOICE_FIELDS from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact from litellm.types.utils import Choices, ModelResponse, ModelResponseStream, TextChoices, TextCompletionResponse @@ -37,16 +38,19 @@ class TextCompletionStreamWrapper: response["object"] = "text_completion" response["created"] = chunk.get("created", None) response["model"] = chunk.get("model", None) - text_choices: Final = TextChoices() if isinstance(chunk, Choices): # chunk should always be of type StreamingChoices raise Exception if chunk["choices"]: - delta: Final = chunk["choices"][0]["delta"] + chat_choice: Final = chunk["choices"][0] + text_choices: Final = TextChoices( + **{field: chat_choice[field] for field in VLLM_CHOICE_FIELDS if field in chat_choice} + ) + delta: Final = chat_choice["delta"] text_choices["text"] = delta["content"] text_choices["reasoning_content"] = delta.get("reasoning_content") - text_choices["index"] = chunk["choices"][0]["index"] - text_choices["finish_reason"] = chunk["choices"][0]["finish_reason"] + text_choices["index"] = chat_choice["index"] + text_choices["finish_reason"] = chat_choice["finish_reason"] response["choices"] = [text_choices] else: response["choices"] = [] diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index 0f6ef2304fb..0911d5ef888 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -472,6 +472,30 @@ def test_cache_key_tells_apart_requests_that_differ_only_in_continuation_token() assert cache.get_cache_key(**request, continuation_token="") == first_segment_key +def test_cache_key_tells_apart_requests_that_ask_for_token_ids() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + request: Final = {"model": "hosted_vllm/qwen3", "messages": [{"role": "user", "content": "count"}]} + + plain_key: Final = cache.get_cache_key(**request) + + assert cache.get_cache_key(**request, return_token_ids=True) != plain_key + assert cache.get_cache_key(**request, return_token_ids=False) == plain_key + assert cache.get_cache_key(**request, return_token_ids=None) == plain_key + + +def test_cache_key_tells_apart_requests_that_ask_for_token_ids_through_extra_body() -> None: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + request: Final = {"model": "hosted_vllm/qwen3", "messages": [{"role": "user", "content": "count"}]} + + plain_key: Final = cache.get_cache_key(**request) + other_extra_body_key: Final = cache.get_cache_key(**request, extra_body={"top_k": 5}) + + assert cache.get_cache_key(**request, extra_body={"return_token_ids": True}) != plain_key + assert cache.get_cache_key(**request, extra_body={"top_k": 5, "return_token_ids": True}) != other_extra_body_key + assert cache.get_cache_key(**request, extra_body={"return_token_ids": False}) == plain_key + assert cache.get_cache_key(**request, extra_body={"continuation_token": "abc"}) != plain_key + + def test_cache_key_tells_apart_requests_that_differ_only_in_tool_call_id_shape() -> None: cache: Final = Cache(type=LiteLLMCacheType.LOCAL) request: Final = { diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index d0f5bfc6b9f..3e197740e98 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -3089,6 +3089,73 @@ def test_convert_to_model_response_object_function_output(): assert result._hidden_params is not None +@pytest.mark.parametrize( + "vllm_choice_fields", + [{"stop_reason": ", 5", "token_ids": [16, 11, 220]}, {"stop_reason": 151645}, {"stop_reason": None, "token_ids": None}], + ids=["stop-string-with-token-ids", "stop-token", "natural-stop"], +) +def test_vllm_choice_fields_stay_on_the_choice_through_a_cache_round_trip( + vllm_choice_fields: dict[str, object], +) -> None: + upstream: Final = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "Qwen/Qwen3-0.6B", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "1, 2, 3, 4"}, + "finish_reason": "stop", + **vllm_choice_fields, + } + ], + } + response: Final = convert_to_model_response_object(response_object=upstream, model_response_object=ModelResponse()) + assert isinstance(response, ModelResponse) + cache_hit: Final = convert_to_model_response_object( + response_object=json.loads(response.model_dump_json()), model_response_object=ModelResponse() + ) + assert isinstance(cache_hit, ModelResponse) + + dumped_choices: Final = [json.loads(rebuilt.model_dump_json())["choices"][0] for rebuilt in (response, cache_hit)] + assert [{field: choice.get(field, "missing") for field in vllm_choice_fields} for choice in dumped_choices] == [ + vllm_choice_fields, + vllm_choice_fields, + ] + + +@pytest.mark.parametrize( + "vllm_choice_fields", + [{"stop_reason": ", 5", "token_ids": [16, 11, 220]}, {"stop_reason": None, "token_ids": None}], + ids=["stop-string-with-token-ids", "natural-stop"], +) +def test_vllm_choice_fields_cached_only_under_provider_specific_fields_come_back_on_the_choice( + vllm_choice_fields: dict[str, object], +) -> None: + cached_before_upgrade: Final = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "Qwen/Qwen3-0.6B", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "1, 2, 3, 4"}, + "finish_reason": "stop", + "provider_specific_fields": vllm_choice_fields, + } + ], + } + cache_hit: Final = convert_to_model_response_object( + response_object=cached_before_upgrade, model_response_object=ModelResponse() + ) + assert isinstance(cache_hit, ModelResponse) + + dumped_choice: Final = json.loads(cache_hit.model_dump_json())["choices"][0] + assert {field: dumped_choice.get(field, "missing") for field in vllm_choice_fields} == vllm_choice_fields + + @pytest.mark.parametrize( "reported_cost", [ diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 9e385c588ae..0aee61d8365 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -561,6 +561,83 @@ class TestPerformRedaction: {"text": "redacted-by-litellm", "finish_reason": "length", "index": 1}, ] + def test_redacts_token_ids_that_decode_to_the_completion(self): + details = { + "stream": True, + "complete_streaming_response": litellm.ModelResponse( + choices=[ + litellm.Choices(message=litellm.Message(content="secret", role="assistant"), token_ids=[20, 11]) + ] + ), + "standard_logging_object": { + "response": { + "choices": [ + { + "token_ids": [20, 11], + "message": {"content": "secret", "provider_specific_fields": {"token_ids": [20, 11]}}, + }, + {"text": " Paris.", "token_ids": [20, 11], "index": 1}, + ] + } + }, + } + + perform_redaction(details, None) + + logged_choices = details["standard_logging_object"]["response"]["choices"] + assert [choice["token_ids"] for choice in logged_choices] == [None, None] + assert logged_choices[0]["message"]["provider_specific_fields"]["token_ids"] is None + assert details["complete_streaming_response"].choices[0].get("token_ids") is None + + def test_redacts_text_completion_stream_text_and_the_token_ids_that_decode_to_it(self): + details = { + "stream": True, + "async_complete_streaming_response": litellm.TextCompletionResponse( + choices=[ + litellm.utils.TextChoices( + index=0, text=" Paris.", reasoning_content="thinking", token_ids=[20, 11], finish_reason="stop" + ) + ] + ), + } + + perform_redaction(details, None) + + logged_choice = details["async_complete_streaming_response"].choices[0] + assert (logged_choice.text, logged_choice.reasoning_content) == ("redacted-by-litellm", "redacted-by-litellm") + assert logged_choice.get("token_ids") is None + assert logged_choice.finish_reason == "stop" + + def test_redacts_token_ids_a_cache_hit_keeps_on_the_choice_provider_specific_fields(self): + result = litellm.ModelResponse( + choices=[ + litellm.Choices( + message=litellm.Message(content="secret", role="assistant"), + provider_specific_fields={"token_ids": [20, 11], "stop_reason": "END"}, + ) + ] + ) + details = { + "standard_logging_object": { + "response": { + "choices": [ + { + "message": {"content": "secret"}, + "provider_specific_fields": {"token_ids": [20, 11], "stop_reason": "END"}, + } + ] + } + } + } + + redacted = perform_redaction(details, result) + + logged_fields = details["standard_logging_object"]["response"]["choices"][0]["provider_specific_fields"] + assert logged_fields == {"token_ids": None, "stop_reason": "END"} + assert redacted.choices[0].provider_specific_fields["token_ids"] is None + assert redacted.choices[0].provider_specific_fields["stop_reason"] == "END" + assert result.choices[0].provider_specific_fields["token_ids"] == [20, 11] + def test_redacts_object_choices_inside_model_response_dict(self): result = { "choices": [ diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index bbf1897b3a7..527c83d5760 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -2757,6 +2757,95 @@ def test_stream_chunk_builder_carries_the_bedrock_guardrail_trace_and_records_it assert entry["guardrail_cost"] == pytest.approx(2 * 0.00015) +def _vllm_chat_chunk(index: int, content: str, finish_reason: str | None = None, **fields: object) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-1", + created=1, + model="qwen3", + choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason, **fields)], + ) + + +@pytest.mark.parametrize("finish_content", ["", "4"], ids=["finish-only", "content-then-finish"]) +def test_stream_chunk_builder_keeps_the_vllm_stop_reason_and_every_token_id_of_each_choice( + finish_content: str, +) -> None: + built: Final = stream_chunk_builder( + chunks=[ + _vllm_chat_chunk(0, "1,", stop_reason=None, token_ids=[16, 11]), + _vllm_chat_chunk(1, "7,", stop_reason=None, token_ids=[22, 11]), + _vllm_chat_chunk(0, "2,3", stop_reason=None, token_ids=[17, 11, 18]), + _vllm_chat_chunk(0, finish_content, "stop", stop_reason=",5", token_ids=[19]), + ] + ) + + assert built is not None + assert [ + {field: choice[field] for field in ("index", "stop_reason", "token_ids")} + for choice in built.model_dump()["choices"] + ] == [ + {"index": 0, "stop_reason": ",5", "token_ids": [16, 11, 17, 11, 18, 19]}, + {"index": 1, "stop_reason": None, "token_ids": [22, 11]}, + ] + + +def test_stream_chunk_builder_keeps_the_vllm_fields_of_a_sole_choice_streamed_on_a_nonzero_index() -> None: + built: Final = stream_chunk_builder( + chunks=[ + _vllm_chat_chunk(1, "7,", stop_reason=None, token_ids=[22, 11]), + _vllm_chat_chunk(1, "8", "stop", stop_reason=",9", token_ids=[23]), + ] + ) + + assert built is not None + assert {field: built.model_dump()["choices"][0].get(field) for field in ("stop_reason", "token_ids")} == { + "stop_reason": ",9", + "token_ids": [22, 11, 23], + } + + +def test_stream_chunk_builder_keeps_null_vllm_fields_and_adds_none_for_other_providers() -> None: + vllm_built: Final = stream_chunk_builder( + chunks=[ + _vllm_chat_chunk(0, "hi", stop_reason=None, token_ids=None), + _vllm_chat_chunk(0, "", "stop", stop_reason=None, token_ids=None), + ] + ) + other_built: Final = stream_chunk_builder( + chunks=[_vllm_chat_chunk(0, "hi"), _vllm_chat_chunk(0, "", "stop")] + ) + + assert vllm_built is not None and other_built is not None + assert { + field: vllm_built.model_dump()["choices"][0][field] for field in ("stop_reason", "token_ids") + } == {"stop_reason": None, "token_ids": None} + assert {"stop_reason", "token_ids"}.isdisjoint(other_built.model_dump()["choices"][0]) + + +def test_stream_chunk_builder_text_completion_keeps_each_choices_vllm_stop_reason_and_token_ids() -> None: + def text_chunk(index: int, text: str, **fields: object) -> TextCompletionResponse: + return TextCompletionResponse( + id="cmpl-1", + object="text_completion", + created=1, + model="qwen3", + choices=[TextChoices(index=index, text=text, **fields)], + ) + + built: Final = stream_chunk_builder_text_completion( + chunks=[ + text_chunk(0, "1,", token_ids=[16, 11], stop_reason=None), + text_chunk(1, "7,", token_ids=[22, 11], stop_reason=None), + text_chunk(0, "2", token_ids=[17], stop_reason=", 3"), + text_chunk(1, "8", token_ids=[23], stop_reason=151645), + ] + ) + + assert [ + (choice["index"], choice["stop_reason"], choice["token_ids"]) for choice in built.model_dump()["choices"] + ] == [(0, ", 3", [16, 11, 17]), (1, 151645, [22, 11, 23])] + + @pytest.mark.parametrize( ("base_model", "counted_as"), [("azure/gpt-6.1-sol", "azure/gpt-6.1-sol"), ("claude-sonnet-4-5", "my-gpt-6")] ) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index da88f1358ad..cf8de9c40f5 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -10,6 +10,7 @@ import threading import time from collections.abc import Mapping, Sequence from dataclasses import dataclass +from itertools import chain from pathlib import Path from types import CodeType, FrameType, MappingProxyType, TracebackType from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast @@ -27,6 +28,10 @@ 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.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response, + convert_to_streaming_response_async, +) from litellm.litellm_core_utils.streaming_handler import ( CustomStreamWrapper, _ProviderChunkEarlyReturn, @@ -34,6 +39,7 @@ from litellm.litellm_core_utils.streaming_handler import ( ) from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicModelResponseIterator from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator as GeminiModelResponseIterator, @@ -6971,6 +6977,327 @@ def test_chunk_creator_keeps_upstream_index_on_single_choice_chunk(logging_obj: assert result.choices[0].delta.content == "hello" +_TOKEN_IDS: Final = TypeAdapter(list[int] | None) +_CHOICE_INDEX: Final = TypeAdapter(int) +_DELTA: Final = TypeAdapter(Mapping[str, object]) + + +def _joined_token_ids(choices: Sequence[Mapping[str, object]]) -> tuple[int, ...]: + return tuple(chain.from_iterable(_TOKEN_IDS.validate_python(choice.get("token_ids")) or () for choice in choices)) + + +def _hosted_vllm_stream( + logging_obj: Logging, upstream_choices: Sequence[Mapping[str, object]] +) -> CustomStreamWrapper: + sse_lines: Final = ( + *( + "data: " + json.dumps({"id": "chatcmpl-1", "object": "chat.completion.chunk", "model": "m", "choices": [c]}) + for c in upstream_choices + ), + "data: [DONE]", + ) + completion_stream: Final = HostedVLLMChatConfig().get_model_response_iterator( + streaming_response=iter(sse_lines), sync_stream=True, json_mode=False + ) + return CustomStreamWrapper( + completion_stream=completion_stream, model="m", custom_llm_provider="hosted_vllm", logging_obj=logging_obj + ) + + +def _hosted_vllm_stream_choices( + logging_obj: Logging, upstream_choices: Sequence[Mapping[str, object]] +) -> tuple[dict[str, object], ...]: + return tuple(chunk.choices[0].model_dump() for chunk in _hosted_vllm_stream(logging_obj, upstream_choices)) + + +@pytest.mark.parametrize("stop_reason", [",5", 151645]) +@pytest.mark.parametrize("final_content", [",3", ""], ids=["content-then-finish", "finish-only"]) +def test_hosted_vllm_stream_keeps_stop_reason_and_every_token_id( + logging_obj: Logging, final_content: str, stop_reason: str | int +) -> None: + emitted: Final = _hosted_vllm_stream_choices( + logging_obj, + ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]}, + {"index": 0, "delta": {"content": "1,2"}, "finish_reason": None, "token_ids": [11, 17]}, + { + "index": 0, + "delta": {"content": final_content}, + "finish_reason": "stop", + "stop_reason": stop_reason, + "token_ids": [20], + }, + ), + ) + + assert [choice.get("stop_reason") for choice in emitted if choice["finish_reason"] is not None] == [stop_reason] + assert [ + (_DELTA.validate_python(choice["delta"]).get("content"), choice["stop_reason"]) + for choice in emitted + if choice["finish_reason"] is None and choice.get("stop_reason") is not None + ] == ([(final_content, stop_reason)] if final_content else []) + assert _joined_token_ids(emitted) == (16, 11, 17, 20) + + +def test_hosted_vllm_stream_sends_empty_content_on_token_id_only_chunks(logging_obj: Logging) -> None: + emitted: Final = _hosted_vllm_stream_choices( + logging_obj, + ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [11]}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [220]}, + {"index": 0, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [16]}, + ), + ) + + assert [ + (_DELTA.validate_python(choice["delta"]).get("content"), choice.get("token_ids")) + for choice in emitted + if choice.get("token_ids") + ] == [("", [11]), ("", [220]), ("1", [16])] + + +def test_hosted_vllm_stream_keeps_held_back_token_ids_on_their_own_choice(logging_obj: Logging) -> None: + emitted: Final = _hosted_vllm_stream_choices( + logging_obj, + ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 1, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 1, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]}, + {"index": 0, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [17]}, + {"index": 1, "delta": {"content": "2"}, "finish_reason": None, "token_ids": [18]}, + ), + ) + + assert [(choice["index"], choice.get("token_ids")) for choice in emitted if choice.get("token_ids")] == [ + (1, [16]), + (0, [17]), + (1, [18]), + ] + + +def _rejects_a_two(text: str) -> bool: + return "2" not in text + + +def _choices_sent_before_rule_failure(stream: CustomStreamWrapper) -> Iterator[dict[str, object]]: + try: + yield from (chunk.choices[0].model_dump() for chunk in stream) + except MidStreamFallbackError: + return + + +def test_hosted_vllm_stream_sends_no_token_ids_ahead_of_a_failing_post_call_rule( + logging_obj: Logging, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "post_call_rules", [_rejects_a_two]) + stream: Final = _hosted_vllm_stream( + logging_obj, + ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]}, + {"index": 0, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [11, 17]}, + {"index": 0, "delta": {"content": ",2"}, "finish_reason": "stop", "stop_reason": None, "token_ids": [20]}, + ), + ) + + assert [ + (_DELTA.validate_python(choice["delta"]).get("content"), choice.get("token_ids")) + for choice in _choices_sent_before_rule_failure(stream) + ] == [("", None), ("1", None)] + + +def _held_back_stop_string_stream(token_ids: Mapping[str, None]) -> tuple[Mapping[str, object], ...]: + return ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, **token_ids}, + {"index": 0, "delta": {"content": "1,2"}, "finish_reason": None, **token_ids}, + {"index": 0, "delta": {"content": ""}, "finish_reason": "stop", "stop_reason": ",3", **token_ids}, + ) + + +def test_hosted_vllm_stream_with_null_token_ids_emits_the_same_chunks_as_without_them(logging_obj: Logging) -> None: + with_null_ids: Final = _hosted_vllm_stream_choices(logging_obj, _held_back_stop_string_stream({"token_ids": None})) + without_ids: Final = _hosted_vllm_stream_choices(logging_obj, _held_back_stop_string_stream({})) + + assert [(c["delta"], c["finish_reason"]) for c in with_null_ids] == [ + (c["delta"], c["finish_reason"]) for c in without_ids + ] + + +_TWO_CHOICE_STOP_REASONS: Final[Mapping[int, str | int]] = MappingProxyType({0: ", 3", 1: 151645}) +_TWO_CHOICE_TOKEN_IDS: Final[Mapping[int, tuple[int, int]]] = MappingProxyType({0: (16, 17), 1: (26, 27)}) + + +def _two_choice_stream(first_finished: int, finish_contents: tuple[str, str]) -> tuple[Mapping[str, object], ...]: + return ( + *({"index": index, "delta": {"role": "assistant", "content": ""}, "finish_reason": None} for index in (0, 1)), + *( + {"index": index, "delta": {"content": "1"}, "finish_reason": None, "token_ids": [ids[0]]} + for index, ids in _TWO_CHOICE_TOKEN_IDS.items() + ), + *( + { + "index": index, + "delta": {"content": content}, + "finish_reason": "stop", + "stop_reason": _TWO_CHOICE_STOP_REASONS[index], + "token_ids": [_TWO_CHOICE_TOKEN_IDS[index][1]], + } + for index, content in zip((first_finished, 1 - first_finished), finish_contents) + ), + ) + + +def _carries_another_choices_fields(choice: Mapping[str, object]) -> bool: + own_index: Final = _CHOICE_INDEX.validate_python(choice["index"]) + token_ids: Final = _TOKEN_IDS.validate_python(choice.get("token_ids")) or () + return not set(token_ids) <= set(_TWO_CHOICE_TOKEN_IDS[own_index]) or choice.get("stop_reason") not in ( + None, + _TWO_CHOICE_STOP_REASONS[own_index], + ) + + +@pytest.mark.parametrize("first_finished", [0, 1]) +@pytest.mark.parametrize( + "finish_contents", + [("", ""), ("2", ""), ("2", "2")], + ids=["both-finish-only", "content-then-finish-only", "both-content-then-finish"], +) +def test_hosted_vllm_stream_with_two_choices_keeps_ids_and_stop_reason_on_their_own_choice( + logging_obj: Logging, first_finished: int, finish_contents: tuple[str, str] +) -> None: + emitted: Final = _hosted_vllm_stream_choices(logging_obj, _two_choice_stream(first_finished, finish_contents)) + + assert [ + (choice["index"], choice.get("token_ids"), choice.get("stop_reason")) + for choice in emitted + if _carries_another_choices_fields(choice) + ] == [] + + +@pytest.mark.parametrize("first_finished", [0, 1]) +@pytest.mark.parametrize( + "finish_contents", [("2", ""), ("2", "2")], ids=["content-then-finish-only", "both-content-then-finish"] +) +def test_hosted_vllm_stream_with_two_choices_keeps_each_stop_reason_on_its_final_content_chunk( + logging_obj: Logging, first_finished: int, finish_contents: tuple[str, str] +) -> None: + emitted: Final = _hosted_vllm_stream_choices(logging_obj, _two_choice_stream(first_finished, finish_contents)) + content_chunk_stop_reasons: Final = frozenset( + (choice["index"], choice["stop_reason"]) + for choice in emitted + if _DELTA.validate_python(choice["delta"]).get("content") and choice.get("stop_reason") is not None + ) + finished_with_content: Final = frozenset( + (index, _TWO_CHOICE_STOP_REASONS[index]) + for index, content in zip((first_finished, 1 - first_finished), finish_contents) + if content + ) + + assert content_chunk_stop_reasons == finished_with_content + + +@pytest.mark.parametrize("first_finished", [0, 1]) +@pytest.mark.parametrize( + "finish_contents", + [("", ""), ("2", ""), ("", "2"), ("2", "2")], + ids=["both-finish-only", "content-then-finish-only", "finish-only-then-content", "both-content-then-finish"], +) +def test_hosted_vllm_stream_with_two_choices_delivers_every_token_id_and_stop_reason_of_each_choice( + logging_obj: Logging, first_finished: int, finish_contents: tuple[str, str] +) -> None: + emitted: Final = _hosted_vllm_stream_choices(logging_obj, _two_choice_stream(first_finished, finish_contents)) + + assert { + index: _joined_token_ids([choice for choice in emitted if choice["index"] == index]) + for index in _TWO_CHOICE_TOKEN_IDS + } == _TWO_CHOICE_TOKEN_IDS + assert frozenset( + (choice["index"], choice["stop_reason"]) for choice in emitted if choice.get("stop_reason") is not None + ) == frozenset(_TWO_CHOICE_STOP_REASONS.items()) + + +def test_hosted_vllm_stream_with_two_choices_finishes_the_first_choice_before_the_other_choice_streams_on( + logging_obj: Logging, +) -> None: + emitted: Final = _hosted_vllm_stream_choices( + logging_obj, + ( + *({"index": index, "delta": {"role": "assistant", "content": ""}, "finish_reason": None} for index in (0, 1)), + {"index": 0, "delta": {"content": "1"}, "finish_reason": "stop", "stop_reason": ",2", "token_ids": [16]}, + {"index": 1, "delta": {"content": ""}, "finish_reason": None, "token_ids": [26]}, + {"index": 1, "delta": {"content": ""}, "finish_reason": None, "token_ids": [11]}, + {"index": 1, "delta": {"content": "1,"}, "finish_reason": None, "token_ids": [27]}, + {"index": 1, "delta": {"content": "3"}, "finish_reason": "stop", "stop_reason": 151645, "token_ids": [28]}, + ), + ) + first_choice_finish: Final = next( + position for position, choice in enumerate(emitted) if choice["index"] == 0 and choice["finish_reason"] + ) + + assert emitted[first_choice_finish]["stop_reason"] == ",2" + assert _joined_token_ids([choice for choice in emitted[first_choice_finish:] if choice["index"] == 1]) == ( + 26, + 11, + 27, + 28, + ) + + +_CACHED_VLLM_TOKEN_IDS: Final = (16, 11, 220, 17, 11, 220, 18, 11, 220, 19) + + +def _cached_vllm_completion() -> dict[str, object]: + return { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "Qwen/Qwen3-0.6B", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "1, 2, 3, 4"}, + "finish_reason": "stop", + "stop_reason": ", 5", + "token_ids": list(_CACHED_VLLM_TOKEN_IDS), + } + ], + } + + +async def _replayed_cache_hit_choices(logging_obj: Logging, replay: str) -> tuple[dict[str, object], ...]: + if replay == "async": + async_wrapper: Final = CustomStreamWrapper( + completion_stream=convert_to_streaming_response_async(_cached_vllm_completion()), + model="m", + custom_llm_provider="cached_response", + logging_obj=logging_obj, + ) + return tuple([chunk.choices[0].model_dump() async for chunk in async_wrapper if chunk.choices]) + sync_wrapper: Final = CustomStreamWrapper( + completion_stream=convert_to_streaming_response(_cached_vllm_completion()), + model="m", + custom_llm_provider="cached_response", + logging_obj=logging_obj, + ) + return tuple(chunk.choices[0].model_dump() for chunk in sync_wrapper if chunk.choices) + + +@pytest.mark.parametrize("replay", ["sync", "async"]) +@pytest.mark.asyncio +async def test_a_streamed_cache_hit_keeps_the_vllm_stop_reason_and_every_token_id( + logging_obj: Logging, replay: str +) -> None: + emitted: Final = await _replayed_cache_hit_choices(logging_obj, replay) + + assert [ + (choice["finish_reason"], choice["stop_reason"]) for choice in emitted if choice.get("stop_reason") is not None + ] == [("stop", ", 5")] + assert _joined_token_ids(emitted) == _CACHED_VLLM_TOKEN_IDS + + def _later_chunk_outputs( custom_llm_provider: str, later_chunks: Sequence[ChatCompletionChunk], logging_obj: Logging ) -> tuple[ModelResponseStream, ...]: diff --git a/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_combined_chunk.py b/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_combined_chunk.py index d4911427e73..a57011b6410 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_combined_chunk.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_streaming_iterator_combined_chunk.py @@ -204,6 +204,20 @@ def test_splitter_splits_combined_chunk_into_content_then_finish(): assert finish_chunk.choices[0].delta.content is None +def test_a_finish_on_another_choice_index_does_not_end_the_message(): + chunks: Final = [ + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="Hello"))]), + ModelResponseStream(choices=[StreamingChoices(index=1, delta=Delta(content="Hi"))]), + ModelResponseStream(choices=[StreamingChoices(index=1, delta=Delta(content=""), finish_reason="length")]), + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")]), + ] + + events: Final = list(AnthropicStreamWrapper(completion_stream=iter(chunks), model="qwen3")) + + stop_reasons: Final = [event["delta"]["stop_reason"] for event in events if event.get("type") == "message_delta"] + assert stop_reasons == ["end_turn"] + + def test_is_combined_false_when_choices_empty(): """A metadata-only chunk with no choices is never treated as combined.""" assert _CombinedChunkSplitter._is_combined(SimpleNamespace(choices=[])) is False diff --git a/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py b/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py index fc5ae3631ff..94f1873faba 100644 --- a/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py +++ b/tests/unit/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py @@ -3,14 +3,15 @@ import json from typing import Final from unittest.mock import MagicMock, patch +import httpx import pytest - from litellm.constants import ( DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, ) from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig +from litellm.types.utils import ModelResponse def test_hosted_vllm_chat_transformation_file_url(): @@ -401,6 +402,64 @@ def test_hosted_vllm_custom_tools_use_top_level_input_schema(): assert tools[0]["function"]["parameters"] == input_schema +def _transform_vllm_chat_response(choice_extras: dict[str, object]) -> ModelResponse: + payload: Final = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "glm", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "1, 2, 3, 4"}, + "logprobs": None, + "finish_reason": "stop", + **choice_extras, + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, + } + raw_response: Final = httpx.Response( + 200, json=payload, request=httpx.Request("POST", "http://vllm.local/v1/chat/completions") + ) + return HostedVLLMChatConfig().transform_response( + model="hosted_vllm/glm", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +@pytest.mark.parametrize( + "choice_extras", + [ + {"stop_reason": ", 5", "token_ids": [16, 11, 220]}, + {"stop_reason": 128009}, + {"stop_reason": None}, + ], +) +def test_hosted_vllm_keeps_vllm_choice_fields_on_the_choice(choice_extras: dict[str, object]): + response: Final = _transform_vllm_chat_response(choice_extras) + + dumped_choice: Final[dict[str, dict[str, object]]] = json.loads(response.model_dump_json())["choices"][0] + for field, value in choice_extras.items(): + assert dumped_choice[field] == value + assert dumped_choice["provider_specific_fields"][field] == value + + +def test_hosted_vllm_does_not_invent_vllm_choice_fields(): + response: Final = _transform_vllm_chat_response({}) + + dumped_choice: Final[dict[str, object]] = json.loads(response.model_dump_json())["choices"][0] + assert "stop_reason" not in dumped_choice + assert "token_ids" not in dumped_choice + + def _strict_weather_tool() -> dict[str, object]: return { "type": "function", diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index c7870dec17a..80bed119e79 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -18,7 +18,7 @@ import asyncio import hashlib import json import logging -from collections.abc import AsyncIterator, Callable +from collections.abc import AsyncIterator, Callable, Mapping from types import MappingProxyType from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock @@ -62,6 +62,8 @@ from litellm.types.utils import ( Function, ModelResponseStream, StreamingChoices, + TextChoices, + TextCompletionResponse, Usage, ) from tests._support.clock_loops import AutojumpClockLoop @@ -553,6 +555,102 @@ def test_serialize_streaming_chunk_simple_payload(): } +@pytest.mark.parametrize( + ("choice_extras", "serialized_extras"), + [ + ({"stop_reason": ",5", "token_ids": [20]}, {"stop_reason": ",5", "token_ids": [20]}), + ({"finish_reason": "stop", "stop_reason": None}, {"finish_reason": "stop", "stop_reason": None}), + ], + ids=["vllm-fields", "natural-stop"], +) +def test_serialize_streaming_chunk_keeps_the_vllm_choice_fields( + choice_extras: dict[str, object], serialized_extras: dict[str, object] +): + chunk: Final = ModelResponseStream( + id="chatcmpl-test", + choices=[StreamingChoices(index=0, delta=Delta(content="1,2,3,4", role="assistant"), **choice_extras)], + created=0, + model="gpt-4", + object="chat.completion.chunk", + ) + + assert json.loads(_serialize_streaming_chunk(chunk))["choices"] == [ + {"index": 0, "delta": {"role": "assistant", "content": "1,2,3,4"}, **serialized_extras} + ] + + +def _text_completion_chunk(choice_extras: Mapping[str, object]) -> TextCompletionResponse: + return TextCompletionResponse( + id="cmpl-1", + choices=[TextChoices(finish_reason="length", index=0, text="", **choice_extras)], + created=1, + model="Qwen/Qwen3-0.6B", + ) + + +def _chat_chunk_with_fingerprint(choice_extras: Mapping[str, object]) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-1", + choices=[StreamingChoices(finish_reason="length", index=0, delta=Delta(content=None), **choice_extras)], + created=1, + model="Qwen/Qwen3-0.6B", + object="chat.completion.chunk", + system_fingerprint="fp_1", + ) + + +_VLLM_SHAPED_CHUNKS: Final = pytest.mark.parametrize( + "build_chunk", [_text_completion_chunk, _chat_chunk_with_fingerprint], ids=["text-completion", "chat-fingerprint"] +) + + +@_VLLM_SHAPED_CHUNKS +@pytest.mark.parametrize( + "vllm_choice_fields", + [{"stop_reason": None}, {"stop_reason": None, "token_ids": [21]}, {"stop_reason": None, "token_ids": None}], + ids=["null-stop-reason", "null-stop-reason-with-token-ids", "null-stop-reason-and-token-ids"], +) +def test_serialize_streaming_chunk_keeps_null_vllm_choice_fields( + build_chunk: Callable[[Mapping[str, object]], TextCompletionResponse | ModelResponseStream], + vllm_choice_fields: dict[str, object], +) -> None: + chunk: Final = build_chunk(vllm_choice_fields) + plain_chunk: Final = json.loads(build_chunk({}).model_dump_json(exclude_none=True, exclude_unset=True)) + + assert json.loads(_serialize_streaming_chunk(chunk)) == { + **plain_chunk, + "choices": [{**plain_chunk["choices"][0], **vllm_choice_fields}], + } + + +@_VLLM_SHAPED_CHUNKS +@pytest.mark.parametrize("choice_extras", [{}, {"stop_reason": ",5", "token_ids": [20]}], ids=["no-vllm", "set-vllm"]) +def test_serialize_streaming_chunk_without_null_vllm_fields_matches_model_dump_json( + build_chunk: Callable[[Mapping[str, object]], TextCompletionResponse | ModelResponseStream], + choice_extras: dict[str, object], +) -> None: + chunk: Final = build_chunk(choice_extras) + + assert _serialize_streaming_chunk(chunk) == chunk.model_dump_json(exclude_none=True, exclude_unset=True).encode() + + +def test_serialize_streaming_chunk_keeps_streaming_a_null_stop_reason_chunk_with_a_value_json_cannot_encode() -> None: + chunk: Final = ModelResponseStream( + id="chatcmpl-vllm", + choices=[StreamingChoices(index=0, delta=Delta(content="a"), stop_reason=None, token_ids=[21])], + created=0, + model="Qwen/Qwen3-0.6B", + object="chat.completion.chunk", + ) + setattr(chunk, "trace_chunk", hashlib.sha256(b"trace").digest()) + + payload: Final = json.loads(_serialize_streaming_chunk(chunk)) + + assert payload["choices"][0]["delta"] == {"content": "a"} + assert payload["choices"][0]["token_ids"] == [21] + assert "trace_chunk" not in payload + + def test_serialize_streaming_chunk_invalid_input_raises_attribute_error(): """The helper is typed as ``BaseModel`` — handing it a plain dict trips the attribute-access path (no ``model_dump_json``).""" diff --git a/tests/unit/proxy/test_model_dump_with_preserved_fields.py b/tests/unit/proxy/test_model_dump_with_preserved_fields.py index 31531dccd16..92c7e58126b 100644 --- a/tests/unit/proxy/test_model_dump_with_preserved_fields.py +++ b/tests/unit/proxy/test_model_dump_with_preserved_fields.py @@ -545,3 +545,31 @@ def test_provider_refusal_string_is_preserved_on_message() -> None: "object": "chat.completion", } assert model_response.choices[0].message.refusal == "I can't help with that." + + +def test_vllm_null_stop_reason_is_kept_only_on_the_choice_that_carries_it() -> None: + model_response: Final = convert_to_model_response_object( + response_object={ + "id": "chatcmpl-vllm-natural-stop", + "created": 4, + "model": "Qwen/Qwen3-0.6B", + "object": "chat.completion", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": {"role": "assistant", "content": "1, 2"}, + "stop_reason": None, + }, + {"finish_reason": "stop", "index": 1, "message": {"role": "assistant", "content": "1, 2"}}, + ], + }, + model_response_object=ModelResponse(), + ) + + dumped_choices: Final = model_dump_with_preserved_fields(model_response, exclude_unset=True)["choices"] + + assert [(choice["index"], "stop_reason" in choice, choice.get("stop_reason")) for choice in dumped_choices] == [ + (0, True, None), + (1, False, None), + ] diff --git a/tests/unit/proxy/test_proxy_logging_hook_detection.py b/tests/unit/proxy/test_proxy_logging_hook_detection.py index 182ef833069..afcb89010fe 100644 --- a/tests/unit/proxy/test_proxy_logging_hook_detection.py +++ b/tests/unit/proxy/test_proxy_logging_hook_detection.py @@ -347,6 +347,241 @@ async def test_post_call_stream_guardrail_keeps_own_iterator_on_chat_completions assert delivered_text != "" +async def _served_token_ids(monkeypatch, guardrail_default_on: bool) -> tuple[object, ...]: + from litellm.caching.caching import DualCache + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + guardrail = ContentFilterGuardrail( + guardrail_name="output-filter", + blocked_words=[BlockedWord(keyword="zebra", action=ContentFilterAction("MASK"))], + event_hook="post_call", + default_on=guardrail_default_on, + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + def chunk(content: str, token_ids: list[int], finish_reason: str | None = None) -> ModelResponseStream: + streamed = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=content), finish_reason=finish_reason)] + ) + setattr(streamed.choices[0], "token_ids", token_ids) + return streamed + + async def fake_stream(): + yield chunk("the zebra runs", [1, 2, 3]) + yield chunk("", [4], finish_reason="stop") + + served = [ + served_chunk + async for served_chunk in ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=fake_stream(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/chat/completions"), + request_data={"model": "gpt-4o-mini", "metadata": {}}, + ) + ] + return tuple(getattr(served_chunk.choices[0], "token_ids", None) for served_chunk in served) + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_drops_token_ids_that_decode_to_the_unmasked_text(monkeypatch): + served_token_ids = await _served_token_ids(monkeypatch, guardrail_default_on=True) + + assert served_token_ids + assert all(token_ids is None for token_ids in served_token_ids) + + +@pytest.mark.asyncio +async def test_post_call_stream_keeps_token_ids_when_no_guardrail_runs_on_the_request(monkeypatch): + served_token_ids = await _served_token_ids(monkeypatch, guardrail_default_on=False) + + assert served_token_ids == ([1, 2, 3], [4]) + + +class _ChecksEachChunkText(CustomGuardrail): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + return None + + +async def _per_chunk_hook_token_ids(monkeypatch, guardrail_default_on: bool) -> tuple[object, bool]: + import time + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + monkeypatch.setattr( + litellm, + "callbacks", + [_ChecksEachChunkText(guardrail_name="chunk-check", event_hook="post_call", default_on=guardrail_default_on)], + ) + logging_obj = Logging( + model="m", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="per-chunk-call", + function_id="per-chunk-call", + ) + held_back_text_chunk = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=""))]) + setattr(held_back_text_chunk.choices[0], "token_ids", [220, 19, 11]) + + served = await ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_hook( + data={"model": "m", "metadata": {}, "litellm_logging_obj": logging_obj}, + response=held_back_text_chunk, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/chat/completions"), + ) + return getattr(served.choices[0], "token_ids", None), logging_obj.drops_stream_token_ids + + +@pytest.mark.asyncio +async def test_per_chunk_stream_check_drops_token_ids_served_ahead_of_the_text_it_checks(monkeypatch): + assert await _per_chunk_hook_token_ids(monkeypatch, guardrail_default_on=True) == (None, True) + + +@pytest.mark.asyncio +async def test_per_chunk_stream_keeps_token_ids_when_its_check_does_not_run_on_the_request(monkeypatch): + assert await _per_chunk_hook_token_ids(monkeypatch, guardrail_default_on=False) == ([220, 19, 11], False) + + +class _ScansWholeStreamFirst(CustomGuardrail): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + for chunk in [chunk async for chunk in response]: + yield chunk + + +async def _logged_stream_token_ids(monkeypatch, guardrail_default_on: bool) -> object: + import asyncio + import json + import time + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig + + guardrail = _ScansWholeStreamFirst( + guardrail_name="whole-stream-scan", event_hook="post_call", default_on=guardrail_default_on + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + logging_obj = Logging( + model="m", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="token-ids-call", + function_id="token-ids-call", + ) + logged = asyncio.get_running_loop().create_future() + + async def capture_logged_response(complete_streaming_response, cache_hit): + logged.set_result(complete_streaming_response) + + logging_obj._on_deferred_stream_complete = capture_logged_response + + async def sse_lines(): + for choice in ( + {"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}, + {"index": 0, "delta": {"content": ""}, "finish_reason": None, "token_ids": [16]}, + {"index": 0, "delta": {"content": "1,2"}, "finish_reason": None, "token_ids": [11, 17]}, + {"index": 0, "delta": {"content": ""}, "finish_reason": "stop", "token_ids": [20]}, + ): + yield "data: " + json.dumps( + {"id": "chatcmpl-1", "object": "chat.completion.chunk", "model": "m", "choices": [choice]} + ) + yield "data: [DONE]" + + stream = CustomStreamWrapper( + completion_stream=HostedVLLMChatConfig().get_model_response_iterator( + streaming_response=sse_lines(), sync_stream=False, json_mode=False + ), + model="m", + custom_llm_provider="hosted_vllm", + logging_obj=logging_obj, + ) + async for _ in ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=stream, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/chat/completions"), + request_data={"model": "m", "metadata": {}, "litellm_logging_obj": logging_obj}, + ): + pass + return (await asyncio.wait_for(logged, timeout=5)).choices[0].token_ids + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_drops_token_ids_from_the_response_it_logs_and_caches(monkeypatch): + assert await _logged_stream_token_ids(monkeypatch, guardrail_default_on=True) is None + + +@pytest.mark.asyncio +async def test_post_call_stream_logs_every_token_id_when_no_guardrail_runs_on_the_request(monkeypatch): + assert await _logged_stream_token_ids(monkeypatch, guardrail_default_on=False) == [16, 11, 17, 20] + + +def _completion_chunk(text: str, token_ids: list[int], finish_reason: str | None = None) -> object: + from litellm.types.utils import TextChoices, TextCompletionResponse + + return TextCompletionResponse( + object="text_completion", + choices=[TextChoices(index=0, text=text, finish_reason=finish_reason, token_ids=token_ids)], + ) + + +def _guarded_completions_stream(monkeypatch, action: str, held_back: tuple[object, ...]): + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction + + monkeypatch.setattr( + litellm, + "callbacks", + [ + ContentFilterGuardrail( + guardrail_name="output-filter", + blocked_words=[BlockedWord(keyword="7", action=ContentFilterAction(action))], + event_hook="post_call", + default_on=True, + ) + ], + ) + + async def upstream(): + for chunk in held_back: + yield chunk + yield _completion_chunk(", 6, 7", [21, 11, 22], finish_reason="stop") + + return ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook( + response=upstream(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/v1/completions"), + request_data={"model": "qwen3", "metadata": {}}, + ) + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_blocks_completions_before_serving_token_id_only_chunks(monkeypatch): + stream = _guarded_completions_stream( + monkeypatch, "BLOCK", (_completion_chunk("", [11]), _completion_chunk("", [220])) + ) + + with pytest.raises(HTTPException): + await stream.__anext__() + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_serves_no_chunk_left_empty_by_the_token_id_drop(monkeypatch): + held_back = (_completion_chunk("", [11]), _completion_chunk("", [220])) + + served = [chunk async for chunk in _guarded_completions_stream(monkeypatch, "MASK", held_back)] + + assert [(chunk.choices[0].text, chunk.choices[0].token_ids) for chunk in served] == [ + (", 6, [KEYWORD_REDACTED]", None) + ] + assert [chunk.choices[0].token_ids for chunk in held_back] == [None, None] + + @pytest.mark.asyncio async def test_post_call_stream_records_masked_text_for_deferred_logging(monkeypatch): from litellm.caching.caching import DualCache diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 3647146eb4e..76d94628d81 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -13900,7 +13900,7 @@ def _stream_usage_generator_chunks(): def test_is_injected_stream_usage_artifact(): from litellm.litellm_core_utils.stream_usage_chunks import is_injected_stream_usage_artifact - from litellm.types.utils import ModelResponseStream, Usage + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage content_chunk, finish_chunk, empty_choices_usage_chunk = _stream_usage_test_chunks() assert is_injected_stream_usage_artifact(empty_choices_usage_chunk) is True @@ -13923,6 +13923,21 @@ def test_is_injected_stream_usage_artifact(): assert is_injected_stream_usage_artifact({"usage": {"prompt_tokens": 1}}) is False + held_back_token_ids_chunk: Final = ModelResponseStream( + model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), token_ids=[16])] + ) + assert is_injected_stream_usage_artifact(held_back_token_ids_chunk) is False + + null_token_ids_chunk: Final = ModelResponseStream( + model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), token_ids=None)] + ) + assert is_injected_stream_usage_artifact(null_token_ids_chunk) is True + + unrelated_extra_chunk: Final = ModelResponseStream( + model="qwen3", choices=[StreamingChoices(delta=Delta(content=None), text="")] + ) + assert is_injected_stream_usage_artifact(unrelated_extra_chunk) is True + async def _collect_async_data_generator_frames(request_data: dict) -> list: from litellm.proxy.proxy_server import async_data_generator diff --git a/tests/unit/text_completions/test_streaming_iterator.py b/tests/unit/text_completions/test_streaming_iterator.py index e5459a5a0cd..571451164d9 100644 --- a/tests/unit/text_completions/test_streaming_iterator.py +++ b/tests/unit/text_completions/test_streaming_iterator.py @@ -1,11 +1,15 @@ -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Iterator, Sequence +from itertools import chain from typing import Final import pytest +from openai.types.completion import Completion from litellm.exceptions import APIConnectionError +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper -from litellm.types.utils import Choices, Delta, ModelResponseStream, StreamingChoices, Usage +from litellm.types.utils import Choices, Delta, ModelResponseStream, StreamingChoices, TextCompletionResponse, Usage def test_reasoning_content_preserved_in_text_completion_wrapper(): @@ -85,6 +89,146 @@ async def test_text_completion_stream_usage_is_opt_in( ] +def _completion_chunk(text: str, finish_reason: str | None, **choice_extras: object) -> Completion: + return Completion.construct( + id="cmpl-1", + created=1, + model="m", + object="text_completion", + choices=[{"index": 0, "text": text, "finish_reason": finish_reason, "logprobs": None, **choice_extras}], + ) + + +def _text_completion_stream(upstream: Sequence[Completion]) -> list[TextCompletionResponse]: + logging_obj: Final = Logging( + model="m", + messages=[{"role": "user", "content": "1, 2, 3, 4"}], + stream=True, + call_type="text_completion", + start_time=0.0, + litellm_call_id="call-1", + function_id="function-1", + ) + chat_stream: Final = CustomStreamWrapper( + completion_stream=iter(upstream), + model="m", + custom_llm_provider="text-completion-openai", + logging_obj=logging_obj, + ) + return list(TextCompletionStreamWrapper(completion_stream=chat_stream, model="m")) + + +@pytest.mark.parametrize("final_text", [", 5, 6", ""], ids=["text-then-finish", "finish-only"]) +def test_text_completion_stream_keeps_stop_reason_on_the_finish_chunk(final_text: str) -> None: + responses: Final = _text_completion_stream( + (_completion_chunk(", 5", None), _completion_chunk(final_text, "stop", stop_reason=", 7")) + ) + + assert responses[-1].choices[0].model_dump(exclude_none=True) == { + "finish_reason": "stop", + "index": 0, + "stop_reason": ", 7", + } + + +def test_text_completion_stream_puts_stop_reason_on_the_final_text_chunk() -> None: + responses: Final = _text_completion_stream( + (_completion_chunk(", 5", None), _completion_chunk(", 6", "stop", stop_reason=", 7")) + ) + + assert [(response.choices[0].get("text"), response.choices[0].get("stop_reason")) for response in responses] == [ + (", 5", None), + (", 6", ", 7"), + (None, ", 7"), + ] + + +@pytest.mark.parametrize("final_text", [", 6", ""], ids=["text-then-finish", "finish-only"]) +def test_text_completion_stream_keeps_held_back_token_ids_and_the_stop_reason(final_text: str) -> None: + responses: Final = _text_completion_stream( + ( + _completion_chunk("", None, token_ids=[11]), + _completion_chunk("", None, token_ids=[220]), + _completion_chunk(", 5", None, token_ids=[20]), + _completion_chunk(final_text, "stop", stop_reason=", 7", token_ids=[11, 220, 21]), + ) + ) + choices: Final = [response.choices[0] for response in responses] + + assert list(chain.from_iterable(choice.get("token_ids") or () for choice in choices)) == [11, 220, 20, 11, 220, 21] + assert [choice.get("stop_reason") for choice in choices if choice.finish_reason is not None] == [", 7"] + + +def test_text_completion_stream_sends_empty_text_on_token_id_only_chunks() -> None: + responses: Final = _text_completion_stream( + ( + _completion_chunk("", None, token_ids=[11]), + _completion_chunk("", None, token_ids=[220]), + _completion_chunk(", 5", None, token_ids=[20]), + ) + ) + + choices: Final = [response.choices[0] for response in responses] + + assert [(choice.text, choice.get("token_ids")) for choice in choices if choice.get("token_ids")] == [ + ("", [11]), + ("", [220]), + (", 5", [20]), + ] + + +def test_text_completion_stream_keeps_each_choices_text_and_token_ids_on_its_own_index() -> None: + responses: Final = _text_completion_stream( + ( + _completion_chunk(" 5", None, index=0, token_ids=[20]), + _completion_chunk(" 7", None, index=1, token_ids=[22]), + _completion_chunk(", 6", None, index=0, token_ids=[11, 21]), + _completion_chunk("", None, index=1, token_ids=[11]), + _completion_chunk("", "stop", index=0, stop_reason=", 9", token_ids=[11, 220]), + _completion_chunk(", 8", "length", index=1, token_ids=[220, 23]), + ) + ) + choices: Final = tuple(chain.from_iterable(response.choices for response in responses)) + + assert { + index: ( + "".join(choice.text or "" for choice in choices if choice.index == index), + list(chain.from_iterable(choice.get("token_ids") or () for choice in choices if choice.index == index)), + ) + for index in (0, 1) + } == {0: (" 5, 6", [20, 11, 21, 11, 220]), 1: (" 7, 8", [22, 11, 220, 23])} + + +@pytest.mark.parametrize( + "upstream", + [ + ( + _completion_chunk(" a", None, index=0), + _completion_chunk(" b", None, index=1), + _completion_chunk("", "length", index=2), + _completion_chunk(" red", "stop", index=1, stop_reason=","), + _completion_chunk(" s", "stop", index=0, stop_reason=","), + ), + ( + _completion_chunk(" a", None, index=0), + _completion_chunk(" b", None, index=1), + _completion_chunk(" s", "stop", index=0, stop_reason=","), + _completion_chunk("", "length", index=2), + _completion_chunk(" red", "stop", index=1, stop_reason=","), + ), + ], + ids=["other-choices-finish-first", "choice-0-finishes-first"], +) +def test_text_completion_stream_finishes_each_choice_on_its_own_index(upstream: Sequence[Completion]) -> None: + finished: Final = [ + (choice.index, choice.finish_reason, choice.get("stop_reason")) + for choice in chain.from_iterable(response.choices for response in _text_completion_stream(upstream)) + if choice.finish_reason is not None + ] + + assert sorted(finished) == [(0, "stop", ","), (1, "stop", ","), (2, "length", None)] + + def test_text_completion_conversion_rejects_non_streaming_choices() -> None: chunk: Final = Choices() wrapper: Final = TextCompletionStreamWrapper(completion_stream=None, model="test-model") From c720b23238ab12a574617e6d4d5d8c71cb6c697c Mon Sep 17 00:00:00 2001 From: tessellark <349031000+tessellark@users.noreply.github.com> Date: Mon, 2 Nov 2026 13:27:58 -0800 Subject: [PATCH 12/79] feat(proxy)!: honor model_info.max_input_tokens in the Codex catalog for stock models (#49701) * feat(proxy): honor model_info.max_input_tokens in the Codex catalog for stock models A Codex model catalog entry copied from Codex's stock entry (a stock slug, or an alias whose upstream is one) now takes the max_input_tokens every routable deployment of the group sets as its context_window, capped at the stock entry's max_context_window like Codex's own model_context_window. With nothing set, the stock window stays. * feat(proxy): honor the smallest max_input_tokens a group's deployments set in the Codex catalog A deployment without a limit counts as Codex's own window, so a limit on only some deployments can lower the window but never raise it * fix(proxy): skip a zero or negative max_input_tokens when sizing a Codex catalog window * test(proxy): stub the max_input_tokens group field and type the Codex window tests * fix(proxy): skip a max_input_tokens a saved /model/info response copied from the cost map --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/client/cli/commands/agents.py | 7 +- .../proxy/common_utils/codex_model_catalog.py | 52 +++++-- litellm/router_api.py | 18 ++- litellm/router_python_core.py | 32 ++++ .../compatibility/test_codex_model_catalog.py | 35 +++++ .../common_utils/test_codex_model_catalog.py | 145 +++++++++++++++++- .../proxy/proxy_server/test_routes_models.py | 1 + tests/unit/proxy/router_entry_point_stubs.py | 2 + 8 files changed, 274 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 12d97ef6cce..b70432d1e72 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -770,9 +770,10 @@ def codex_model_catalog( A model Codex does not know gets the fallback entry, with the same base instructions Codex itself uses so the agent never runs without a system prompt, and offers the reasoning efforts `groups` gives its model group - that the installed Codex, at `client_version`, can parse. Apart from aliases, the catalog is - the one `/v1/models?client_version=...` serves, built from the installed Codex's own entries - instead of the proxy's vendored ones. + that the installed Codex, at `client_version`, can parse. Apart from aliases, and the window the + deployments' `model_info.max_input_tokens` gives a model Codex knows, the catalog is the one + `/v1/models?client_version=...` serves, built from the installed Codex's own entries instead of + the proxy's vendored ones. """ rows: Final = tuple(_codex_row(m, upstreams.get(m.id), groups) for m in models) body: Final = codex_models_response_json( diff --git a/litellm/proxy/common_utils/codex_model_catalog.py b/litellm/proxy/common_utils/codex_model_catalog.py index 63fb27d62cf..7b48aa07365 100644 --- a/litellm/proxy/common_utils/codex_model_catalog.py +++ b/litellm/proxy/common_utils/codex_model_catalog.py @@ -7,8 +7,10 @@ never the OpenAI list shape. This module builds that body from the rows ``/v1/mo lists: a model Codex knows keeps Codex's own entry for it, any other model gets the fallback entry Codex itself uses for an unknown slug, and ``model_info.service_tiers`` becomes the entry's ``service_tiers``, which Codex turns into slash commands (``/ultrafast``) that send ``service_tier`` -upstream. A fallback entry offers the model group's ``supported_reasoning_efforts`` as its reasoning -levels, trimmed to the ones the requesting Codex release can parse. +upstream. A ``model_info.max_input_tokens`` the operator sets replaces a known model's +``context_window``, up to the most Codex lets its own ``model_context_window`` raise that model to. +A fallback entry offers the model group's ``supported_reasoning_efforts`` as its reasoning levels, +trimmed to the ones the requesting Codex release can parse. The stock entries are Codex 0.159.3's bundled catalog, vendored unchanged from https://github.com/openai/codex/blob/rust-v0.159.3/codex-rs/models-manager/models.json (Apache-2.0) @@ -158,6 +160,7 @@ class CodexStockModel(LiteLLMBaseModel): supports_search_tool: bool = False supported_reasoning_levels: tuple[CodexReasoningLevel, ...] = () context_window: int | None = None + max_context_window: int | None = None class CodexStockCatalog(LiteLLMBaseModel): @@ -173,8 +176,9 @@ class CodexCatalogRow: """What the catalog needs to know about one listed model: its public id, the listing's mode and input limit, the upstream model behind it, the reasoning effort its deployments pin, and the `model_info` values the operator set; `service_tiers` holds one raw `model_info.service_tiers` - value per deployment behind the id, None where a deployment sets none. `supported_reasoning_efforts` - is the model group's, None when unknown.""" + value per deployment behind the id, None where a deployment sets none, and + `configured_max_input_tokens` one `model_info.max_input_tokens` per deployment the same way. + `supported_reasoning_efforts` is the model group's, None when unknown.""" id: str mode: str | None = None @@ -185,6 +189,7 @@ class CodexCatalogRow: service_tiers: tuple[object, ...] = () supported_reasoning_efforts: tuple[str, ...] | None = None search_tool_support: tuple[bool | None, ...] = () + configured_max_input_tokens: tuple[int | None, ...] = () @cache @@ -333,6 +338,25 @@ def _alias_context_window(row: CodexCatalogRow, stock_model: CodexStockModel) -> return MappingProxyType({"context_window": min(stock_model.context_window, row.max_input_tokens)}) +def _configured_window(row: CodexCatalogRow, stock_model: CodexStockModel) -> Mapping[str, int]: + """The `context_window` the operator's `model_info.max_input_tokens` gives a known model's entry, + or nothing, which keeps `stock_model`'s own. + + A request to the row can route to any of its deployments, so the window is the smallest one of them + takes: its `max_input_tokens` where the operator set a positive one, Codex's own window where not. A + limit set on only some deployments can therefore lower the window but never raise it. The window never + goes past `stock_model`'s `max_context_window`, the most Codex lets its own `model_context_window` + setting raise that model to.""" + per_deployment: Final = row.configured_max_input_tokens + declared: Final = tuple(limit for limit in per_deployment if limit is not None and limit > 0) + unset_window: Final = () if len(declared) == len(per_deployment) else (stock_model.context_window,) + if not declared or None in unset_window: + return MappingProxyType({}) + smallest: Final = min(limit for limit in (*declared, *unset_window) if limit is not None) + ceiling: Final = stock_model.max_context_window + return MappingProxyType({"context_window": smallest if ceiling is None else min(smallest, ceiling)}) + + def _alias_reasoning_level(row: CodexCatalogRow, stock_model: CodexStockModel) -> str | None: efforts: Final = frozenset(level.effort for level in stock_model.supported_reasoning_levels) return row.reasoning_effort if row.reasoning_effort in efforts else None @@ -398,6 +422,7 @@ def _entry( "upgrade": upgrade, **({} if search_tool is None else {"supports_search_tool": search_tool}), **(_deployment_defaults(row, stock_model) if is_alias and alias_defaults == "deployment" else {}), + **_configured_window(row, stock_model), } if configured is None: return stock_model.model_copy(update=listing) @@ -473,10 +498,12 @@ def codex_models_response_json( byte limit silently and keeps its bundled one, so the body holds only the entries that fit, the ones offering a service tier taken first and the rest in listing order, an entry too large for the bytes left passed over, every kept entry at its listing position; `byte_limit=None` - keeps every entry. A row listed under another name than the Codex model its `upstream_model` - names (an alias) takes that model's entry; `alias_defaults="deployment"` drops the defaults - Codex picks for that model in favor of what the alias's deployments pin, while `"codex"` keeps - them, as Codex would use them for the model's own name. `client_version` is the requesting Codex's version, which decides the reasoning + keeps every entry. A row whose deployments set `model_info.max_input_tokens` gets the smallest + window any of them takes as its known model's `context_window`, capped at that model's + `max_context_window`. A row listed under another name than the Codex model its `upstream_model` names (an alias) takes that model's + entry; `alias_defaults="deployment"` drops the defaults Codex picks for that model in favor of + what the alias's deployments pin, while `"codex"` keeps them, as Codex would use them for the + model's own name. `client_version` is the requesting Codex's version, which decides the reasoning levels a fallback entry may offer; None counts as a version too old for custom levels. """ known: Final = bundled_codex_models() if stock is None else stock @@ -519,6 +546,9 @@ def _catalog_row( search_tool_support=( registry.resolve(lookup_id, kind="search_tool_support", team_id=team_id) if registry is not None else () ), + configured_max_input_tokens=( + registry.group(lookup_id, field="max_input_tokens", team_id=team_id) if registry is not None else () + ), ) @@ -532,9 +562,9 @@ def codex_catalog_rows( """`rows` joined with the router's configured metadata, looked up by each entry's internal id so team-scoped rows resolve the way the Anthropic listing's display names do, a `model_group_alias` reading its target's deployments under its own name; `team_id` is the requesting key's team, so the - upstream model that picks Codex's stock entry and the tiers are read only off the deployments its - requests can route to, and the reasoning efforts are the ones `/model_group/info` reports for the - entry. Display names are read under `label_entries`' ids when given.""" + upstream model that picks Codex's stock entry, the tiers, and the configured input limits are read + only off the deployments its requests can route to, and the reasoning efforts are the ones + `/model_group/info` reports for the entry. Display names are read under `label_entries`' ids when given.""" lookup_ids: Final = MappingProxyType(dict(entries)) label_ids: Final = lookup_ids if label_entries is None else MappingProxyType(dict(label_entries)) return tuple( diff --git a/litellm/router_api.py b/litellm/router_api.py index c013a99573f..3007b5345c7 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -155,6 +155,7 @@ GroupField: TypeAlias = Literal[ "display_name", "mode", "service_tiers", + "max_input_tokens", "token_limits", "listing_info", ] @@ -855,6 +856,16 @@ class RouterRegistry: team_id: str | None = None, ) -> tuple[object, ...]: ... + @overload + def group( + self, + name: str, + *, + field: Literal["max_input_tokens"], + include_hidden: bool = False, + team_id: str | None = None, + ) -> tuple[int | None, ...]: ... + @overload def group( self, @@ -882,9 +893,7 @@ class RouterRegistry: field: GroupField = "info", include_hidden: bool = False, team_id: str | None = None, - ) -> ( - "ModelGroupInfo | DeploymentModelListingInfo | tuple[int | None, int | None] | tuple[object, ...] | str | None" - ): + ) -> "ModelGroupInfo | DeploymentModelListingInfo | tuple[int | None, ...] | tuple[object, ...] | str | None": """ What the router knows about a model group. ``field``: @@ -894,6 +903,7 @@ class RouterRegistry: (were ``get_configured_display_name``, ``get_configured_mode``, ``get_configured_token_limits``, ``get_model_listing_info``) - ``service_tiers``: (was ``get_configured_service_tiers``) + - ``max_input_tokens``: each routable deployment's own ``model_info.max_input_tokens`` """ r: Final = self._router if field == "info": @@ -906,6 +916,8 @@ class RouterRegistry: return r._core_get_configured_mode(name) if field == "service_tiers": return r._core_get_configured_service_tiers(name, team_id=team_id) + if field == "max_input_tokens": + return r._core_get_configured_max_input_tokens(name, team_id=team_id) if field == "token_limits": return r._core_get_configured_token_limits(name) if field == "listing_info": diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index c63e538c998..98c9e641278 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -141,6 +141,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import ( is_registered_custom_provider, resolve_model_and_provider_for_metadata, ) +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.get_retry_from_policy import ( get_num_retries_from_retry_policy as _get_num_retries_from_retry_policy, ) @@ -475,6 +476,7 @@ from litellm.types.router import ( from litellm.types.services import ServiceTypes from litellm.types.utils import ( AUTOROUTER_CLASSIFIER_CALL_ORIGIN, + COST_MAP_LOOKUP_KEY, PROMPT_QUOTING_ROUTING_DECISION_FIELDS, CustomPricingLiteLLMParams, GenericBudgetConfigType, @@ -486,6 +488,7 @@ from litellm.types.utils import ( StandardLoggingRoutingDecision, Usage, all_litellm_params, + echoed_cost_map_fields, is_server_derived_pricing_key, shared_backend_model_info, ) @@ -772,6 +775,17 @@ def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, obj return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {}) +def _operator_max_input_tokens(deployment: Mapping[str, object]) -> int | None: + model_info: Final = _configured_model_info(deployment) + catalog_key: Final = model_info.get(COST_MAP_LOOKUP_KEY) + catalog_entry: Final = ( + GetModelCostMap.loaded_model_cost_map().get(catalog_key) if isinstance(catalog_key, str) else None + ) + if catalog_entry is not None and "max_input_tokens" in echoed_cost_map_fields(model_info, catalog_entry): + return None + return coerce_token_limit(model_info.get("max_input_tokens")) + + def _search_tool_support(deployment: DeploymentTypedDict) -> bool | None: from litellm.responses.litellm_completion_transformation.tool_search import chat_config_builds_anthropic_request @@ -12693,6 +12707,24 @@ class PythonCore(DeprecatedRouterMembers): for deployment in self._routable_deployments(model_name, team_id) ) + def _core_get_configured_max_input_tokens( + self, model_name: str, team_id: str | None = None + ) -> tuple[int | None, ...]: + """ + Return the max_input_tokens each deployment a request from team_id to + model_name can route to (see _routable_deployments) sets in its own + model_info, in model_list order, None for a deployment that sets none + or a malformed one. Unlike get_model_listing_info, neither the cost map + nor a discovered limit is read, so only an operator's own value counts, + and a value a saved /model/info response copied from the cost map is + skipped too. + + Returns an empty tuple for wildcard-expanded or unknown names. + """ + return tuple( + _operator_max_input_tokens(deployment) for deployment in self._routable_deployments(model_name, team_id) + ) + def _core_get_search_tool_support(self, model_name: str, team_id: str | None = None) -> tuple[bool | None, ...]: """ Whether each deployment a request from team_id to model_name can route diff --git a/tests/integration/compatibility/test_codex_model_catalog.py b/tests/integration/compatibility/test_codex_model_catalog.py index a3ba032e612..fbf05b0f526 100644 --- a/tests/integration/compatibility/test_codex_model_catalog.py +++ b/tests/integration/compatibility/test_codex_model_catalog.py @@ -351,6 +351,41 @@ def test_stock_model_keeps_codex_entry(gateway: Gateway) -> None: assert catalog[name] == expected, catalog[name] +def test_stock_model_alias_takes_configured_max_input_tokens(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + stock: Final = _bundled_stock()["gpt-6-sol"] + stock_window: Final = stock["context_window"] + ceiling: Final = stock["max_context_window"] + assert isinstance(stock_window, int) and isinstance(ceiling, int) and stock_window < ceiling, stock + raised: Final = _model_name("codex-raised-window") + lowered: Final = _model_name("codex-lowered-window") + half_set: Final = _model_name("codex-half-set-window") + half_lowered: Final = _model_name("codex-half-lowered-window") + unset: Final = _model_name("codex-unset-window") + names: Final = (raised, lowered, half_set, half_lowered, unset) + _new_model(scenario, raised, model="openai/gpt-6-sol", model_info={"max_input_tokens": ceiling + 50000}) + _new_model(scenario, lowered, model="openai/gpt-6-sol", model_info={"max_input_tokens": stock_window // 2}) + _new_model(scenario, half_set, model="openai/gpt-6-sol", model_info={"max_input_tokens": ceiling}) + _new_model(scenario, half_set, model="openai/gpt-6-sol") + _new_model(scenario, half_lowered, model="openai/gpt-6-sol", model_info={"max_input_tokens": stock_window // 3}) + _new_model(scenario, half_lowered, model="openai/gpt-6-sol") + _new_model(scenario, unset, model="openai/gpt-6-sol") + expected: Final = { + raised: ceiling, + lowered: stock_window // 2, + half_set: stock_window, + half_lowered: stock_window // 3, + unset: stock_window, + } + windows: Final = _settled( + lambda: { + slug: entry["context_window"] for slug, entry in _by_slug(_catalog(gateway)).items() if slug in names + }, + lambda value: value == expected, + ) + assert windows == expected, windows + + def test_stock_model_configured_tiers_override(gateway: Gateway) -> None: with gateway.scenario() as scenario: stock: Final = _bundled_stock()["gpt-6-sol"] diff --git a/tests/unit/proxy/common_utils/test_codex_model_catalog.py b/tests/unit/proxy/common_utils/test_codex_model_catalog.py index 509303d9635..e3d941a6aa4 100644 --- a/tests/unit/proxy/common_utils/test_codex_model_catalog.py +++ b/tests/unit/proxy/common_utils/test_codex_model_catalog.py @@ -1,9 +1,12 @@ import json +from collections.abc import Mapping from types import MappingProxyType +from typing import Final, Literal import pytest import litellm +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.proxy.common_utils.codex_model_catalog import ( CODEX_CATALOG_BYTE_LIMIT, CodexCatalogRow, @@ -233,6 +236,140 @@ def test_alias_with_deployment_defaults_caps_codexs_window_at_its_listed_max_inp assert json.loads(body.json)["models"][0]["context_window"] == expected +_WIDE_STOCK = MappingProxyType( + { + "gpt-6-sol": _stock("gpt-6-sol").model_copy(update={"context_window": 272000, "max_context_window": 872000}), + "gpt-5.5": _stock("gpt-5.5").model_copy(update={"context_window": 272000}), + } +) + + +@pytest.mark.parametrize( + ("configured", "expected"), + [ + ((500000,), 500000), + ((100000,), 100000), + ((1050000,), 872000), + ((600000, 400000), 400000), + ((600000, None), 272000), + ((100000, None), 100000), + ((None,), 272000), + ((), 272000), + ((0,), 272000), + ((-5, 400000), 272000), + ((0, 100000), 100000), + ], +) +@pytest.mark.parametrize("alias_defaults", ["codex", "deployment"]) +def test_a_known_model_and_its_alias_take_the_smallest_window_any_deployment_takes( + configured: tuple[int | None, ...], expected: int, alias_defaults: Literal["codex", "deployment"] +) -> None: + rows: Final = ( + _row("gpt-6-sol", max_input_tokens=1050000, configured_max_input_tokens=configured), + _row( + "sol-1m", + upstream_model="openai/gpt-6-sol", + max_input_tokens=1050000, + configured_max_input_tokens=configured, + ), + ) + + body: Final = codex_models_response_json( + rows, stock=_WIDE_STOCK, instructions=_PROMPT, alias_defaults=alias_defaults + ) + + assert {entry["slug"]: entry["context_window"] for entry in json.loads(body.json)["models"]} == { + "gpt-6-sol": expected, + "sol-1m": expected, + } + + +def test_a_configured_window_goes_past_codexs_own_when_the_stock_entry_sets_no_ceiling() -> None: + (entry,) = json.loads( + codex_models_response_json( + (_row("gpt-5.5", configured_max_input_tokens=(1050000,)),), stock=_WIDE_STOCK, instructions=_PROMPT + ).json + )["models"] + + assert entry["context_window"] == 1050000 + + +def test_a_deployment_without_a_limit_keeps_a_windowless_stock_entry_as_codex_ships_it() -> None: + windowless: Final = MappingProxyType({"gpt-5.5": _stock("gpt-5.5").model_copy(update={"context_window": None})}) + + (entry,) = json.loads( + codex_models_response_json( + (_row("gpt-5.5", configured_max_input_tokens=(1050000, None)),), stock=windowless, instructions=_PROMPT + ).json + )["models"] + + assert entry["context_window"] is None + + +def test_model_list_body_sizes_a_known_models_window_off_the_deployments_the_key_team_routes_to() -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra"}, + "model_info": {"max_input_tokens": 922000}, + }, + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://team-2.example"}, + "model_info": {"team_id": "team-2", "max_input_tokens": 400000}, + }, + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://team-3.example"}, + "model_info": {"team_id": "team-3"}, + }, + ], + model_group_alias={"astra-1m": "gpt-6-astra"}, + ) + listing: Final = tuple( + {"id": name, "object": "model", "created": 0, "owned_by": "openai", "mode": "chat"} + for name in ("gpt-6-astra", "astra-1m") + ) + entries: Final = (("gpt-6-astra", "gpt-6-astra"), ("astra-1m", "astra-1m")) + stock: Final = bundled_codex_models()["gpt-6-astra"] + assert stock.context_window is not None and stock.max_context_window is not None + assert stock.context_window < 400000 < stock.max_context_window < 922000 + + def windows(team_id: str) -> dict[str, int | None]: + models: Final = json.loads(codex_model_list_body(listing, entries, router, team_id))["models"] + return {entry["slug"]: entry["context_window"] for entry in models} + + assert windows("team-1") == {"gpt-6-astra": stock.max_context_window, "astra-1m": stock.max_context_window} + assert windows("team-2") == {"gpt-6-astra": 400000, "astra-1m": 400000} + assert windows("team-3") == {"gpt-6-astra": stock.context_window, "astra-1m": stock.context_window} + + +def test_model_list_body_skips_a_limit_a_saved_model_info_response_copied_from_the_cost_map() -> None: + catalog_limit: Final = GetModelCostMap.loaded_model_cost_map()["gpt-6-astra"]["max_input_tokens"] + stock: Final = bundled_codex_models()["gpt-6-astra"] + assert isinstance(catalog_limit, int) and stock.max_context_window is not None + assert stock.context_window is not None and stock.context_window < 600000 < stock.max_context_window + + def window(model_info: Mapping[str, object]) -> int | None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra"}, + "model_info": dict(model_info), + } + ] + ) + listing: Final = ({"id": "gpt-6-astra", "object": "model", "created": 0, "owned_by": "openai", "mode": "chat"},) + (entry,) = json.loads(codex_model_list_body(listing, (("gpt-6-astra", "gpt-6-astra"),), router))["models"] + return entry["context_window"] + + assert window({"key": "gpt-6-astra", "max_input_tokens": catalog_limit}) == stock.context_window + assert window({"key": "gpt-6-astra", "max_input_tokens": 600000}) == 600000 + assert window({"max_input_tokens": catalog_limit}) == min(catalog_limit, stock.max_context_window) + + def test_alias_with_deployment_defaults_drops_codexs_tier_and_upgrade_while_the_native_name_keeps_them(): rows = (_row("gpt-6-sol"), _row("team-sol", upstream_model="gpt-6-sol"), _row("team-55", upstream_model="gpt-5.5")) codex = {entry["slug"]: entry for entry in _models(*rows)} @@ -599,9 +736,14 @@ def test_catalog_rows_read_the_router_by_each_entry_lookup_id(): "model_name_team-1_c0ffee" ).supported_reasoning_efforts, search_tool_support=(None, None), + configured_max_input_tokens=(None, None), ), CodexCatalogRow( - id="plain", upstream_model="openai/some-unmapped-model", service_tiers=(None,), search_tool_support=(None,) + id="plain", + upstream_model="openai/some-unmapped-model", + service_tiers=(None,), + search_tool_support=(None,), + configured_max_input_tokens=(None,), ), ) @@ -631,6 +773,7 @@ def test_catalog_rows_read_an_alias_off_its_target_under_the_alias_name(): service_tiers=(["ultrafast"],), supported_reasoning_efforts=router.get_model_group_info("gpt-6-astra").supported_reasoning_efforts, search_tool_support=(None,), + configured_max_input_tokens=(None,), ) (entry,) = json.loads(codex_model_list_body(listing, (("gpt-6", "gpt-6"),), router))["models"] assert (entry["slug"], entry["display_name"], [tier["id"] for tier in entry["service_tiers"]]) == ( diff --git a/tests/unit/proxy/proxy_server/test_routes_models.py b/tests/unit/proxy/proxy_server/test_routes_models.py index 874f471f3e7..ff74cd3fa1d 100644 --- a/tests/unit/proxy/proxy_server/test_routes_models.py +++ b/tests/unit/proxy/proxy_server/test_routes_models.py @@ -58,6 +58,7 @@ def patched_models(monkeypatch): router.get_routable_upstream_model = MagicMock(return_value="gpt-4") router.get_configured_display_name = MagicMock(return_value=None) router.get_configured_service_tiers = MagicMock(return_value=()) + router.get_configured_max_input_tokens = MagicMock(return_value=()) router.get_model_group_info = MagicMock(return_value=None) router.get_search_tool_support = MagicMock(return_value=()) route_entry_points_to_stubs(router) diff --git a/tests/unit/proxy/router_entry_point_stubs.py b/tests/unit/proxy/router_entry_point_stubs.py index 6912049d225..51b5bf13f16 100644 --- a/tests/unit/proxy/router_entry_point_stubs.py +++ b/tests/unit/proxy/router_entry_point_stubs.py @@ -123,6 +123,8 @@ def route_entry_points_to_stubs(router: _RouterDouble) -> _RouterDouble: return router.get_configured_display_name(name) if field == "service_tiers": return router.get_configured_service_tiers(name, **_given(team_id=team_id)) + if field == "max_input_tokens": + return router.get_configured_max_input_tokens(name, **_given(team_id=team_id)) if field == "listing_info": return router.get_model_listing_info(name) if field == "mode": From b4a80d175c6f8b1a14bfc22eca934a936a0b8a46 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 2 Nov 2026 13:33:05 -0800 Subject: [PATCH 13/79] ci: remove workflows whose secrets are dead (#49722) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/deploy-dev-on-main.yml | 30 -------- .github/workflows/weekly_load_anomaly.yml | 85 ----------------------- tests/e2e/AGENTS.md | 2 +- 3 files changed, 1 insertion(+), 116 deletions(-) delete mode 100644 .github/workflows/deploy-dev-on-main.yml delete mode 100644 .github/workflows/weekly_load_anomaly.yml diff --git a/.github/workflows/deploy-dev-on-main.yml b/.github/workflows/deploy-dev-on-main.yml deleted file mode 100644 index 1208a2dbbdb..00000000000 --- a/.github/workflows/deploy-dev-on-main.yml +++ /dev/null @@ -1,30 +0,0 @@ -name: Deploy main to gateway-dev - -on: - push: - branches: [main] - -permissions: {} - -concurrency: - group: deploy-dev-dispatch - cancel-in-progress: true - -jobs: - dispatch: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - timeout-minutes: 5 - steps: - - uses: actions/create-github-app-token@fee1f7d63c2ff003460e3d139729b119787bc349 # v2.2.2 - id: token - with: - app-id: ${{ secrets.LITELLM_OPS_DISPATCH_APP_ID }} - private-key: ${{ secrets.LITELLM_OPS_DISPATCH_APP_PRIVATE_KEY }} - owner: BerriAI - repositories: litellm-ops - permission-actions: write - - name: Dispatch the litellm-ops dev deploy - env: - GH_TOKEN: ${{ steps.token.outputs.token }} - run: gh api --method POST repos/BerriAI/litellm-ops/actions/workflows/deploy-dev.yml/dispatches -f ref=main diff --git a/.github/workflows/weekly_load_anomaly.yml b/.github/workflows/weekly_load_anomaly.yml deleted file mode 100644 index 0ba4df89cc4..00000000000 --- a/.github/workflows/weekly_load_anomaly.yml +++ /dev/null @@ -1,85 +0,0 @@ -name: "Weekly Load Anomaly Check" - -on: - schedule: - - cron: "0 12 * * 6" - workflow_dispatch: - -permissions: - contents: read - -jobs: - weekly-load-anomaly: - if: github.event_name != 'schedule' || github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - timeout-minutes: 45 - services: - postgres: - image: postgres:16.6 - env: - POSTGRES_USER: llmproxy - POSTGRES_PASSWORD: dbpassword9090 - POSTGRES_DB: litellm - ports: - - 5432:5432 - options: >- - --health-cmd "pg_isready -U llmproxy" - --health-interval 5s - --health-timeout 5s - --health-retries 10 - env: - DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm - LITELLM_MASTER_KEY: sk-weekly-anomaly-check - ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} - AWS_BEARER_TOKEN_BEDROCK: ${{ secrets.AWS_BEARER_TOKEN_BEDROCK }} - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Set up uv - uses: ./.github/actions/setup-uv-with-retries - with: - version: "0.10.9" - - - name: Cache the Rust build - uses: ./.github/actions/cache-cargo-build - - - name: Install dependencies - run: | - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy - - - name: Cache Prisma binaries - uses: ./.github/actions/cache-prisma-binaries - - - name: Generate Prisma client - run: | - uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - - - name: Start the proxy - run: | - nohup uv run --no-sync litellm --config tests/e2e/perf/weekly_anomaly_config.yml --port 4000 > proxy.log 2>&1 & - for _ in $(seq 1 90); do - if curl -fs http://localhost:4000/health/liveliness > /dev/null; then - exit 0 - fi - sleep 2 - done - echo "proxy never became live" - tail -n 100 proxy.log - exit 1 - - - name: Run the weekly session anomaly test - env: - E2E_WEEKLY_ANOMALY: "1" - run: | - uv run --no-sync pytest tests/e2e/perf/test_weekly_session_anomaly_e2e.py -v --tb=short -rA - - - name: Show proxy log on failure - if: failure() - run: tail -n 300 proxy.log diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index cf1e41a59a6..496a85688b3 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -44,7 +44,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a loose ceiling over the measured batches nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) -- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares, or a `Ramp` that steps the rate up through the run, built with `ramped()`, whose plan carries each request's send second and whose run reports `sustained_rps`, the last step the pod held, scored in `perf/setup/ramp.py`, with every step's latency, completed rate, non-200s, CPU per request, Redis latency and Postgres connections cut from the samplers' time-stamped `Series` (`perf/setup/series.py`) to the step's window and emitted as `perf_ramp_step__` JUnit properties for the dashboard) and `Limit`s (one metric at one percentile under one ceiling, or over one floor for `sustained_rps`). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; the same overhead as the gateway itself reports it, read from the `x-litellm-overhead-duration-ms` header per request (`litellm_overhead_seconds`, gated next to `gateway_added_seconds`, the outside reading of the same thing, so the two can be charted together) and the gateway's own end-to-end time from `x-litellm-response-duration-ms` (`litellm_response_seconds`), the run mean of the gateway's `litellm_redis_latency` histogram from a `/metrics` scrape before and after the run (`telemetry.py`, `prometheus_redis_latency_mean_seconds`), and the database's own view of the run from the same `pg_stat_statements` rows and a once-a-second `pg_stat_activity` sample: execution seconds per wall second (`postgres_active_sessions`, the average-active-sessions number RDS reports as DB load), the slowest statement's mean (`postgres_slowest_statement_seconds`, named in the detail) and the peak client connection count (`postgres_peak_connections`), each under a ceiling set from the first local runs and due a retune once the lane has its ten-run baseline; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` +- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, run by hand since its weekly workflow never had working Bedrock credentials), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares, or a `Ramp` that steps the rate up through the run, built with `ramped()`, whose plan carries each request's send second and whose run reports `sustained_rps`, the last step the pod held, scored in `perf/setup/ramp.py`, with every step's latency, completed rate, non-200s, CPU per request, Redis latency and Postgres connections cut from the samplers' time-stamped `Series` (`perf/setup/series.py`) to the step's window and emitted as `perf_ramp_step__` JUnit properties for the dashboard) and `Limit`s (one metric at one percentile under one ceiling, or over one floor for `sustained_rps`). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; the same overhead as the gateway itself reports it, read from the `x-litellm-overhead-duration-ms` header per request (`litellm_overhead_seconds`, gated next to `gateway_added_seconds`, the outside reading of the same thing, so the two can be charted together) and the gateway's own end-to-end time from `x-litellm-response-duration-ms` (`litellm_response_seconds`), the run mean of the gateway's `litellm_redis_latency` histogram from a `/metrics` scrape before and after the run (`telemetry.py`, `prometheus_redis_latency_mean_seconds`), and the database's own view of the run from the same `pg_stat_statements` rows and a once-a-second `pg_stat_activity` sample: execution seconds per wall second (`postgres_active_sessions`, the average-active-sessions number RDS reports as DB load), the slowest statement's mean (`postgres_slowest_statement_seconds`, named in the detail) and the peak client connection count (`postgres_peak_connections`), each under a ceiling set from the first local runs and due a retune once the lane has its ten-run baseline; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at, see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) - `gateway/` - proxy configuration only (`litellm-config.yml`). No tests From 0ea7e96a664297b0db2a2e8e952efbe38216ffc1 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 2 Nov 2026 13:44:45 -0800 Subject: [PATCH 14/79] test(integration): repair agent reads and bulk deletion setup (#49719) --- .../management/test_agent_management.py | 12 +++++++----- .../test_team_delete_large_membership.py | 15 +++++---------- 2 files changed, 12 insertions(+), 15 deletions(-) diff --git a/tests/integration/management/test_agent_management.py b/tests/integration/management/test_agent_management.py index 5c0fdc68e09..a4c2650ffdd 100644 --- a/tests/integration/management/test_agent_management.py +++ b/tests/integration/management/test_agent_management.py @@ -5,12 +5,14 @@ from typing import Final from integration._support.client import Gateway, Scenario from integration._support.database import read_rows from integration._support.wire import Reply, Request, wire_server -from pydantic import BaseModel, JsonValue +from pydantic import BaseModel, ConfigDict, JsonValue _REDACTED: Final = "REDACTED_BY_LITELM" class _AgentRead(BaseModel): + model_config = ConfigDict(frozen=True) + agent_id: str agent_name: str litellm_params: dict[str, JsonValue] | None @@ -84,7 +86,8 @@ def test_agent_reads_redact_credentials_filter_by_key_permission_and_keep_the_st visible: Final = _register_agent(gateway, scenario, marker + "-visible", wire.url, secret) hidden: Final = _register_agent(gateway, scenario, marker + "-hidden", wire.url, secret) restricted: Final = scenario.key(object_permission={"agents": [visible]}) - redacted_params: Final = {"api_key": _REDACTED, "cost_per_query": 0.1, "is_public": False} + stored_params: Final = {"api_key": _REDACTED, "cost_per_query": 0.1} + redacted_params: Final = {**stored_params, "is_public": False} listed: Final = gateway.client.get("/v1/agents", headers={"Authorization": f"Bearer {restricted}"}) assert listed.status_code == 200, listed.text @@ -98,9 +101,8 @@ def test_agent_reads_redact_credentials_filter_by_key_permission_and_keep_the_st assert fetched.status_code == 200, fetched.text assert secret not in fetched.text, fetched.text read: Final = _AgentRead.model_validate_json(fetched.content) - assert (read.agent_id, read.litellm_params, read.static_headers) == (visible, redacted_params, None), ( - fetched.text - ) + assert (read.agent_id, read.static_headers) == (visible, None), fetched.text + assert read.litellm_params in (stored_params, redacted_params), fetched.text refused: Final = gateway.client.get(f"/v1/agents/{hidden}", headers={"Authorization": f"Bearer {restricted}"}) assert (refused.status_code, refused.json()) == ( diff --git a/tests/integration/management/test_team_delete_large_membership.py b/tests/integration/management/test_team_delete_large_membership.py index 8b3c6b06bba..2a8508ff3c9 100644 --- a/tests/integration/management/test_team_delete_large_membership.py +++ b/tests/integration/management/test_team_delete_large_membership.py @@ -24,6 +24,7 @@ import yaml from pydantic import JsonValue from redis import Redis +from litellm.constants import TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS from tests.integration._support.client import ( Gateway, Scenario, @@ -166,19 +167,11 @@ def _remove_team_by_sql(team_id: str) -> None: assert _team_rows(team_id) == [] -def _remove_team_if_present(gateway: Gateway, team_id: str) -> None: - """Cleanup for teams the test deletes itself: the API delete first, SQL for anything it leaves.""" - if not _team_rows(team_id): - return - gateway.request("POST", "/team/delete", {"team_ids": [team_id]}) - _remove_team_by_sql(team_id) - - def _create_team(scenario: Scenario, **fields: JsonValue) -> str: """A team the test deletes itself; cleanup only removes it if the test left it behind.""" created: Final = scenario.gateway.post("/team/new", {"team_alias": f"integration-{uuid.uuid4().hex}", **fields}) team_id: Final = string_value(created["team_id"]) - scenario.cleanups.callback(_remove_team_if_present, scenario.gateway, team_id) + scenario.cleanups.callback(_remove_team_by_sql, team_id) return team_id @@ -783,7 +776,9 @@ def test_deleting_a_team_with_more_keys_than_the_publish_cap_revokes_every_key_o keys: Final = _seeded_team_keys(team, KEYS_WELL_PAST_THE_PUBLISH_CAP) assert _statuses(peer, keys) == {200: len(keys)} - response: Final = _delete_teams(gateway, [team]) + response: Final = _post_with_timeout( + gateway, "/team/delete", {"team_ids": [team]}, timeout=2 * TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS + ) assert response.status_code == 200, response.text assert read_rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE team_id = %s', (team,)) == [] From 3af23024bc5e1aaf36d57b6008402f2949e3fb23 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 2 Nov 2026 13:50:56 -0800 Subject: [PATCH 15/79] test(e2e): share the coding-CLI harness between suites (#49370) * test(e2e): share the CLI version rule, env lockdown and rate limiter in tests/e2e/cli_harness Moves the npm release rule (now taking the package name), the env allowlist plus fresh HOME, and the per-upstream rate limiter out of claude_code so the Codex suite can reuse them. Claude Code behavior is unchanged, its CI, cron and Dockerfile paths now call the shared resolver with @anthropic-ai/claude-code * test(e2e): freeze the packument map and note the rate limiter self-test * test(e2e): race real processes on one rate limiter bucket so a missing lock fails Eight spawned processes start together and drain a bucket that never refills. With the flock they get exactly its burst between them, without it every attempt is granted * test(e2e): keep the old Claude Code version resolver path for the Buildkite runner images The runner Dockerfiles still run tests/e2e/claude_code/pr_gate_version_resolver.py and silently install @latest when it is missing, which drops the three day review window. The old path now forwards to cli_harness/npm_version_resolver.py --- .github/workflows/compat-matrix-image.yml | 4 +- .github/workflows/test-linting.yml | 2 +- .../check_e2e_no_raw_requests.py | 4 +- tests/e2e/AGENTS.md | 3 +- tests/e2e/claude_code/cli_driver.py | 139 ++----- tests/e2e/claude_code/cron_vm/README.md | 2 +- tests/e2e/claude_code/cron_vm/run_daily.sh | 4 +- tests/e2e/claude_code/http_probe.py | 10 +- .../claude_code/pr_gate_version_resolver.py | 157 +------- tests/e2e/claude_code/rate_limiter.py | 378 ------------------ tests/e2e/claude_code/run_compat.sh | 2 +- tests/e2e/cli_harness/__init__.py | 0 tests/e2e/cli_harness/lockdown.py | 49 +++ tests/e2e/cli_harness/npm_version_resolver.py | 92 +++++ tests/e2e/cli_harness/rate_limiter.py | 151 +++++++ tests/e2e_harness/AGENTS.md | 2 +- .../test_pr_gate_version_resolver.py | 60 --- .../claude_code/test_request_determinism.py | 2 +- .../e2e_harness/cli_harness/test_lockdown.py | 32 ++ .../cli_harness/test_npm_version_resolver.py | 63 +++ .../cli_harness/test_rate_limiter.py | 78 ++++ 21 files changed, 534 insertions(+), 700 deletions(-) delete mode 100644 tests/e2e/claude_code/rate_limiter.py create mode 100644 tests/e2e/cli_harness/__init__.py create mode 100644 tests/e2e/cli_harness/lockdown.py create mode 100644 tests/e2e/cli_harness/npm_version_resolver.py create mode 100644 tests/e2e/cli_harness/rate_limiter.py delete mode 100644 tests/e2e_harness/claude_code/test_pr_gate_version_resolver.py create mode 100644 tests/e2e_harness/cli_harness/test_lockdown.py create mode 100644 tests/e2e_harness/cli_harness/test_npm_version_resolver.py create mode 100644 tests/e2e_harness/cli_harness/test_rate_limiter.py diff --git a/.github/workflows/compat-matrix-image.yml b/.github/workflows/compat-matrix-image.yml index f08792c904c..ddd2aa1a50b 100644 --- a/.github/workflows/compat-matrix-image.yml +++ b/.github/workflows/compat-matrix-image.yml @@ -4,7 +4,7 @@ on: pull_request: paths: - tests/e2e/claude_code/cron_vm/** - - tests/e2e/claude_code/pr_gate_version_resolver.py + - tests/e2e/cli_harness/npm_version_resolver.py - .github/workflows/compat-matrix-image.yml workflow_dispatch: @@ -36,7 +36,7 @@ jobs: whoami gh --version uv --version - version="$(uv run --no-project --python 3.12 python /opt/litellm/tests/e2e/claude_code/pr_gate_version_resolver.py)" + version="$(uv run --no-project --python 3.12 python /opt/litellm/tests/e2e/cli_harness/npm_version_resolver.py @anthropic-ai/claude-code)" /opt/litellm/tests/e2e/claude_code/cron_vm/install_claude_code.sh "${version}" /tmp/claude-cli /tmp/claude-cli/claude --version ' diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index c3388430589..cbeea71ad0f 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -208,7 +208,7 @@ jobs: exit 0 fi retry() { "$@" || { sleep 15; "$@"; } || { sleep 30; "$@"; }; } - CLAUDE_VERSION="$(retry uv run --no-sync python tests/e2e/claude_code/pr_gate_version_resolver.py)" + CLAUDE_VERSION="$(retry uv run --no-sync python tests/e2e/cli_harness/npm_version_resolver.py @anthropic-ai/claude-code)" tests/e2e/claude_code/cron_vm/install_claude_code.sh "$CLAUDE_VERSION" "$RUNNER_TEMP/claude-cli" PATH="$RUNNER_TEMP/claude-cli:$PATH" uv run --no-sync pytest -q tests/e2e_harness uv run --no-sync pytest --tb=short -q --noconftest -o addopts= -o "pythonpath=tests/e2e tests/e2e_harness" -p no:rerunfailures \ diff --git a/tests/code_coverage_tests/check_e2e_no_raw_requests.py b/tests/code_coverage_tests/check_e2e_no_raw_requests.py index bea177d3b7d..94cba6334e0 100644 --- a/tests/code_coverage_tests/check_e2e_no_raw_requests.py +++ b/tests/code_coverage_tests/check_e2e_no_raw_requests.py @@ -3,7 +3,7 @@ raw HTTP client imports (requests, urllib.request, httpx, aiohttp, http.client) banned in suite code, and tests/e2e_harness, which tests that harness, is held to the same ban. Importing requests' exception types for catching is fine anywhere; a small allowlist grandfathers the files that legitimately make raw calls (the transport -itself, the root conftest liveness probe, the claude_code version resolver's constant +itself, the root conftest liveness probe, the cli_harness npm version resolver's constant registry URL fetch, and the mcp OAuth client, whose httpx client is the object the official mcp SDK's streamable_http_client requires and so cannot go through the sync requests transport). Referenced by tests/e2e/AGENTS.md.""" @@ -24,7 +24,7 @@ BANNED_MODULES = ("requests", "urllib.request", "http.client", "httpx", "aiohttp ALLOWED_RAW_CLIENT_FILES = { "e2e/e2e_http.py": ("requests",), "e2e/conftest.py": ("requests",), - "e2e/claude_code/pr_gate_version_resolver.py": ("urllib.request",), + "e2e/cli_harness/npm_version_resolver.py": ("urllib.request",), } EXCEPTION_ONLY_NAMES = frozenset({"RequestException", "ConnectionError", "Timeout", "HTTPError"}) diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index 496a85688b3..9e2c628028e 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -48,7 +48,8 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at, see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) - `gateway/` - proxy configuration only (`litellm-config.yml`). No tests -- `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture. Ships its own driver/builder/publisher, covered by `tests/e2e_harness/claude_code/`. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`). The CLI-driving path stays bespoke +- `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture. Ships its own driver/builder/publisher, covered by `tests/e2e_harness/claude_code/`. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`). The CLI-driving path stays bespoke. `pr_gate_version_resolver.py` only forwards to `cli_harness/npm_version_resolver.py`, because the Buildkite runner images still run that path and install @latest without it +- `cli_harness/` - no tests, only the plumbing every suite that drives a real coding-agent CLI shares: `npm_version_resolver.py` picks the CLI release to install (the newest npm release at least three days old, run by CI as `python tests/e2e/cli_harness/npm_version_resolver.py `, stdlib only), `lockdown.py` gives the CLI subprocess an allowlisted env and a fresh empty HOME so it never sees the runner's credentials or dotfiles, and `rate_limiter.py` is the cross-process token bucket per upstream that keeps parallel CLI runs under one request rate. Covered by `tests/e2e_harness/cli_harness/`. Code only one CLI suite needs stays in that suite - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream, see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport. The Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json`. Its specs follow the dashboard's rule against testing what shadcn and Base UI already guarantee, in `ui/litellm-dashboard/AGENTS.md` ## MCP suite: real Datadog only diff --git a/tests/e2e/claude_code/cli_driver.py b/tests/e2e/claude_code/cli_driver.py index 6f7a7b39a39..599cfe6aa75 100644 --- a/tests/e2e/claude_code/cli_driver.py +++ b/tests/e2e/claude_code/cli_driver.py @@ -16,23 +16,18 @@ from __future__ import annotations import json import os import re -import shutil import subprocess import sys import tempfile import time from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field -from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union +from typing import Any, Callable, Dict, Final, List, Mapping, Optional, Sequence, Tuple, Union +from cli_harness.lockdown import isolated_home, locked_down_env +from cli_harness.rate_limiter import RateBucket, RateLimiter, get_default_limiter from e2e_metadata import step -from claude_code.rate_limiter import ( - RateLimiter, - get_default_limiter, - infer_provider, -) - CLAUDE_CLI_DEFAULT = "claude" # 120s is fine for a single isolated CLI call against an unloaded # upstream, but the matrix run launches up to 75 concurrent calls and @@ -77,61 +72,36 @@ Anthropic / Bedrock / Vertex) to reset. Both knobs are env-tunable so a matrix run can trade wall time for resilience without code edits; retries=0 disables the behavior entirely.""" -# Env vars the `claude` Node CLI legitimately needs to function: -# locating its own binary + node, basic locale/terminal plumbing. -# Deliberately excludes every credential-bearing var that the -# surrounding CI job sets for the proxy (ANTHROPIC_API_KEY, AWS_*, -# AZURE_*, VERTEXAI_CREDENTIALS, GITHUB_TOKEN, OPENAI_API_KEY, -# DATABASE_URL, ...). The CLI talks to the proxy via the explicit -# ANTHROPIC_BASE_URL/ANTHROPIC_AUTH_TOKEN we set below — it has no -# business reading the proxy's upstream credentials, and a -# compromised CLI release shouldn't be able to exfiltrate them out -# of the CI environment. -# -# `HOME` is intentionally NOT in this list: see `_make_isolated_home` -# below. We give the CLI a fresh empty per-invocation HOME so a -# compromised claude package or a model-directed `Read` tool call -# can't reach files like `~/.config/gh/hosts.yml`, `~/.ssh/id_rsa`, -# or `~/.bash_history` on the cron VM (and on the CircleCI executor -# the same isolation prevents accidentally exposing checkout-adjacent -# files even though the runner home is ephemeral there). -_CLI_ENV_ALLOWLIST: tuple = ( - "PATH", - "USER", - "LOGNAME", - "SHELL", - "TERM", - "TMPDIR", - "LANG", - "LC_ALL", - "LC_CTYPE", - "NODE_PATH", - "NVM_DIR", - "NVM_BIN", -) +def infer_provider(model: str) -> RateBucket: + """Map a model alias to its compat-matrix rate bucket. -def _make_isolated_home() -> str: - """Create a fresh empty HOME directory for a single `claude` subprocess. - - The CLI needs *a* writable HOME (it caches per-session state under - `$HOME/.claude/projects//`), but it has no legitimate need - for the *user's* HOME. Handing it the real one means a compromised - `@anthropic-ai/claude-code` release, or a model-directed `Read` - tool call during a PDF/vision cell, can read host files like - `~/.config/gh/hosts.yml` (GitHub CLI host token), `~/.ssh/`, - `~/.bash_history`, or any other dotfile under the runtime user's - home. On the cron VM the runtime user is a real interactive - account (`mateo`) with a populated home directory, so this is a - real exfiltration surface. - - The directory is created under `tempfile.gettempdir()` (which is - `/tmp` on Linux; under systemd's `PrivateTmp=true` that's a - per-service tmpfs that the service user can't otherwise reach). - Caller is responsible for `shutil.rmtree`-ing it after the - subprocess exits. + The matrix column layout is fixed; aliases registered in the proxy + encode the provider via a suffix (`-bedrock-converse`, + `-bedrock-invoke`, `-azure`, `-vertex`, `-openai`, `-azure-openai`, + `-bedrock-mantle`) or its absence (Anthropic). Order matters: + `-azure-openai` also ends with `-openai`, and the bedrock suffixes + all contain `bedrock`, so the more-specific suffixes are tested + first. """ - return tempfile.mkdtemp(prefix="claude-cli-home-") + if not model: + raise ValueError("model must be a non-empty string") + lower: Final = model.lower() + if lower.endswith("-azure-openai"): + return RateBucket.AZURE_OPENAI + if lower.endswith("-openai"): + return RateBucket.OPENAI + if lower.endswith("-bedrock-mantle"): + return RateBucket.BEDROCK_MANTLE + if lower.endswith("-bedrock-converse"): + return RateBucket.BEDROCK_CONVERSE + if lower.endswith("-bedrock-invoke"): + return RateBucket.BEDROCK_INVOKE + if lower.endswith("-azure"): + return RateBucket.AZURE + if lower.endswith("-vertex"): + return RateBucket.VERTEX_AI + return RateBucket.ANTHROPIC _FIXED_CLI_USER_ID = "0" * 64 @@ -176,7 +146,7 @@ def _stable_cli_state() -> Tuple[str, str]: Pinning both makes the bodies repeatable across builds. This narrows what survives rather than widening it: HOME stays fresh and - empty per invocation, so the isolation `_make_isolated_home` describes is + empty per invocation, so the isolation `cli_harness.lockdown` describes is unchanged, and the CLI's own state no longer outlives the pod either. The working directory is deliberately not the checkout, so a model-directed `Read` sees an empty directory instead of the repository. @@ -291,38 +261,26 @@ def run_claude( cmd.append("--") cmd.append(prompt) - # Build a minimal env for the CLI subprocess: only the allowlisted - # process-runtime vars from os.environ, plus the explicit proxy - # creds, plus any caller-supplied overrides. See _CLI_ENV_ALLOWLIST - # above for the security rationale. - env: Dict[str, str] = { - key: os.environ[key] for key in _CLI_ENV_ALLOWLIST if key in os.environ - } - env["ANTHROPIC_BASE_URL"] = base_url - env["ANTHROPIC_AUTH_TOKEN"] = api_key - env["DISABLE_AUTOUPDATER"] = "1" - # Hand the CLI a fresh empty HOME so a compromised claude package - # or a model-directed Read tool call can't see the runtime user's - # real dotfiles. Created here, removed in the `finally` below - # regardless of how the subprocess exits. - isolated_home = _make_isolated_home() - env["HOME"] = isolated_home - config_dir, workspace = _stable_cli_state() - env["CLAUDE_CONFIG_DIR"] = config_dir - if extra_env: - env.update(extra_env) - # Throttle by provider *before* launching the CLI. Doing this here # (rather than per-test) means every code path that lands on # `run_claude` is rate-limited automatically — including # `run_claude_models_parallel`, which is the hot path during the # full matrix run. limiter = rate_limiter if rate_limiter is not None else get_default_limiter() - provider = infer_provider(model) - limiter.acquire(provider) - run_fn = runner or subprocess.run - try: + with isolated_home("claude-cli-home-") as home: + config_dir, workspace = _stable_cli_state() + env = locked_down_env( + home=home, + overrides={ + "ANTHROPIC_BASE_URL": base_url, + "ANTHROPIC_AUTH_TOKEN": api_key, + "DISABLE_AUTOUPDATER": "1", + "CLAUDE_CONFIG_DIR": config_dir, + **(extra_env or {}), + }, + ) + limiter.acquire(infer_provider(model)) try: completed = run_fn( cmd, @@ -340,15 +298,6 @@ def run_claude( ) from exc except subprocess.TimeoutExpired as exc: raise ClaudeCLIError(f"claude CLI timed out after {timeout}s") from exc - finally: - # Best-effort cleanup. If the subprocess wrote a `.claude/` - # session dir under the isolated HOME, we remove it here so - # parallel matrix runs don't accumulate per-call tmpdirs. - # `ignore_errors=True` because rmtree races with any - # not-yet-reaped child (SIGTERM'd `claude` on a host-side - # timeout) are benign — the next matrix run starts from a - # fresh tmpdir anyway. - shutil.rmtree(isolated_home, ignore_errors=True) events = _parse_stream_json(completed.stdout or "") text = _extract_assistant_text(events) diff --git a/tests/e2e/claude_code/cron_vm/README.md b/tests/e2e/claude_code/cron_vm/README.md index 4c16b6e237a..7451e86e09b 100644 --- a/tests/e2e/claude_code/cron_vm/README.md +++ b/tests/e2e/claude_code/cron_vm/README.md @@ -52,7 +52,7 @@ than as a GitHub Action or on a dedicated VM. Trade-offs: the stable venv doesn't install) is never loaded. The tag's own `tests/e2e/` is deliberately not used. 3. **Resolves and installs the Claude Code CLI under test**: - `pr_gate_version_resolver.py` (run on the venv, the image has no + `tests/e2e/cli_harness/npm_version_resolver.py` (run on the venv, the image has no Python of its own) picks the newest `@anthropic-ai/claude-code` npm release published at least 3 days ago, the same buffer the PR gate uses, unless `CLAUDE_CODE_VERSION` pins one, and diff --git a/tests/e2e/claude_code/cron_vm/run_daily.sh b/tests/e2e/claude_code/cron_vm/run_daily.sh index 17b628b8c78..749d5166a34 100755 --- a/tests/e2e/claude_code/cron_vm/run_daily.sh +++ b/tests/e2e/claude_code/cron_vm/run_daily.sh @@ -8,7 +8,7 @@ # Releases API. # 2. Update a long-lived worktree at $WORKTREE to that tag and `uv sync` it. # 3. Resolve the Claude Code CLI version under test (the newest npm -# release at least 3 days old, via pr_gate_version_resolver.py, or +# release at least 3 days old, via cli_harness/npm_version_resolver.py, or # $CLAUDE_CODE_VERSION when set) and download that release's # linux-x64 binary into the run's scratch dir, checksum-verified # against the vendor's release manifest (install_claude_code.sh). @@ -329,7 +329,7 @@ PROXY_CONFIG="${WORKTREE}/tests/e2e/claude_code/test_config.yaml" if [[ -z "${CLAUDE_CODE_VERSION}" ]]; then CLAUDE_CODE_VERSION="$( cd "${WORKTREE}" \ - && "${WORKTREE_UV}" run --no-sync python "${POPULATOR_DIR}/../pr_gate_version_resolver.py" + && "${WORKTREE_UV}" run --no-sync python "${POPULATOR_DIR}/../../cli_harness/npm_version_resolver.py" @anthropic-ai/claude-code )" || die "could not resolve the Claude Code version to test" log "resolved claude code: ${CLAUDE_CODE_VERSION}" else diff --git a/tests/e2e/claude_code/http_probe.py b/tests/e2e/claude_code/http_probe.py index a6c587f98cf..390ef4c24d4 100644 --- a/tests/e2e/claude_code/http_probe.py +++ b/tests/e2e/claude_code/http_probe.py @@ -31,8 +31,7 @@ from __future__ import annotations from dataclasses import dataclass from typing import TYPE_CHECKING -from pydantic import BaseModel - +from cli_harness.rate_limiter import RateLimiter, get_default_limiter from e2e_http import ( NetworkError, RateLimitedError, @@ -59,12 +58,9 @@ from models import ( JsonSchemaProperty, ToolInputSchema, ) +from pydantic import BaseModel -from claude_code.rate_limiter import ( - RateLimiter, - get_default_limiter, - infer_provider, -) +from claude_code.cli_driver import infer_provider if TYPE_CHECKING: from proxy_client import ProxyClient diff --git a/tests/e2e/claude_code/pr_gate_version_resolver.py b/tests/e2e/claude_code/pr_gate_version_resolver.py index 756dafb7d73..72392f8a0ae 100644 --- a/tests/e2e/claude_code/pr_gate_version_resolver.py +++ b/tests/e2e/claude_code/pr_gate_version_resolver.py @@ -1,155 +1,16 @@ -"""Claude Code PR-Gate Version Resolver. - -Resolves the `@anthropic-ai/claude-code` npm version that the PR-gate CI -job installs. Selects the newest version (by publish timestamp) whose -publish timestamp is at least 3 days old. The 3-day window is a security -review buffer — see PRD #26476, "Version resolvers". - -Two surfaces: - -- ``resolve_pr_gate_version(...)`` — the importable function. Accepts - pre-fetched npm metadata (for unit tests) or a custom ``fetcher`` - callable. The default fetcher hits the public npm registry. -- ``python -m claude_code.pr_gate_version_resolver`` — prints the - resolved version string to stdout, suitable for piping into a shell - ``$(...)`` substitution inside the CircleCI job. - -The CLI form is what CircleCI runs at job start; engineers reading the -job log can see the selected version on a single line above the -``npm install -g`` step (acceptance criterion: "the selected Claude -Code version is logged in the CI output"). -""" +"""The path the Buildkite runner images run (`python3 tests/e2e/claude_code/pr_gate_version_resolver.py`) to pick the +Claude Code release they install. They fall back to @latest when this file is missing, so keep it until they call +`cli_harness/npm_version_resolver.py` themselves.""" from __future__ import annotations -import json +import subprocess import sys -import urllib.request -from datetime import datetime, timedelta, timezone -from typing import Callable, Mapping, Optional - -PACKAGE_NAME = "@anthropic-ai/claude-code" -NPM_REGISTRY_URL = "https://registry.npmjs.org/{package}" -DEFAULT_MIN_AGE = timedelta(days=3) -DEFAULT_FETCH_TIMEOUT_SECONDS = 30 - -# npm's `time` map mixes per-version timestamps with these meta keys. -_TIME_META_KEYS = frozenset({"created", "modified"}) - - -class NoEligibleVersionError(RuntimeError): - """Raised when no version in the npm metadata satisfies the min-age cutoff.""" - - -def _parse_npm_timestamp(value: str) -> datetime: - """Parse the ISO-8601 timestamps npm emits (always UTC, may use ``Z``).""" - if value.endswith("Z"): - value = value[:-1] + "+00:00" - parsed = datetime.fromisoformat(value) - if parsed.tzinfo is None: - parsed = parsed.replace(tzinfo=timezone.utc) - return parsed - - -def _default_fetcher(package_name: str) -> dict: - """Fetch the npm packument for ``package_name`` over HTTPS. - - Uses urllib (stdlib) so this module has no extra dependencies in the - CI environment. Returns the raw JSON dict. - """ - # urllib.parse.quote would encode the leading '@' / '/' which the - # npm registry expects literally; do a minimal hand-roll instead. - url = NPM_REGISTRY_URL.format(package=package_name.replace("/", "%2F")) - req = urllib.request.Request(url, headers={"Accept": "application/json"}) - with urllib.request.urlopen( # noqa: S310 — registry URL is constant - req, timeout=DEFAULT_FETCH_TIMEOUT_SECONDS - ) as response: - body = response.read().decode("utf-8") - return json.loads(body) - - -def resolve_pr_gate_version( - *, - metadata: Optional[Mapping] = None, - fetcher: Optional[Callable[[str], Mapping]] = None, - as_of: Optional[datetime] = None, - min_age: timedelta = DEFAULT_MIN_AGE, - package_name: str = PACKAGE_NAME, -) -> str: - """Return the newest npm version of ``package_name`` published >= ``min_age`` ago. - - "Newest" means newest by **publish time**, not semver string order — - if a patch lands on an older major after a newer release, the - patched line is the eligible one. A version npm has unpublished keeps - its ``time`` entry but drops out of ``versions``, so only versions - still present in ``versions`` are candidates. - - Args: - metadata: Pre-fetched npm packument (skips the HTTP call). Useful - for unit tests. - fetcher: Callable taking a package name and returning the - packument. Defaults to a stdlib HTTPS fetcher. - as_of: The clock used to decide whether a version is "old - enough". Defaults to ``datetime.now(timezone.utc)``. - min_age: Minimum publish age. Defaults to 3 days. - package_name: Defaults to ``@anthropic-ai/claude-code``. - - Raises: - NoEligibleVersionError: when no version in the registry meets - the age cutoff. - """ - if metadata is None: - fetch = fetcher or _default_fetcher - metadata = fetch(package_name) - - times = metadata.get("time") or {} - versions = metadata.get("versions") or {} - if as_of is None: - as_of = datetime.now(timezone.utc) - cutoff = as_of - min_age - - eligible: list[tuple[datetime, str]] = [] - for version, raw_ts in times.items(): - if version in _TIME_META_KEYS: - continue - if version not in versions: - continue - if not isinstance(raw_ts, str): - continue - if "-" in version: - continue - published = _parse_npm_timestamp(raw_ts) - if published <= cutoff: - eligible.append((published, version)) - - if not eligible: - raise NoEligibleVersionError( - f"no version of {package_name} is at least {min_age} old " - f"as of {as_of.isoformat()}" - ) - - eligible.sort(key=lambda pair: pair[0], reverse=True) - return eligible[0][1] - - -def _main(argv: list[str]) -> int: - """Print the resolved version to stdout. Exit code 0 on success. - - Stderr carries the human-readable announcement so the version can be - captured cleanly with ``$(python -m ...)`` in shell. - """ - try: - version = resolve_pr_gate_version() - except Exception as exc: # noqa: BLE001 — CLI surface, want everything - print(f"pr_gate_version_resolver: {exc}", file=sys.stderr) # noqa: T201 - return 1 - print( # noqa: T201 - f"pr_gate_version_resolver: selected {PACKAGE_NAME}@{version}", - file=sys.stderr, - ) - print(version) # noqa: T201 - return 0 +from pathlib import Path +from typing import Final +_RESOLVER: Final = Path(__file__).resolve().parent.parent / "cli_harness" / "npm_version_resolver.py" +_PACKAGE: Final = "@anthropic-ai/claude-code" if __name__ == "__main__": - raise SystemExit(_main(sys.argv[1:])) + raise SystemExit(subprocess.run((sys.executable, "-I", str(_RESOLVER), _PACKAGE), check=False).returncode) diff --git a/tests/e2e/claude_code/rate_limiter.py b/tests/e2e/claude_code/rate_limiter.py deleted file mode 100644 index 5818338ff7f..00000000000 --- a/tests/e2e/claude_code/rate_limiter.py +++ /dev/null @@ -1,378 +0,0 @@ -"""Cross-process token-bucket rate limiter for the Claude Code compat suite. - -The compat matrix runs 75 live `claude` CLI invocations (25 cells × 3 -Claude tiers per cell). When pytest-xdist fans these out across worker -processes, each worker would maintain its own in-memory rate limiter -and the *aggregate* request rate hitting any one upstream provider -would be `workers × per-worker rate` — exactly the situation that -trips Anthropic / Azure / Bedrock / Vertex 429s in the middle of a -matrix run and silently flips green cells red. - -The fix is a **shared, cross-process** token bucket per provider, -backed by a small JSON state file and an OS-level `flock`. Each -`run_claude` invocation acquires one token (sleeping if the bucket is -empty) before launching the CLI; refills happen lazily based on wall -time, so workers can be killed and restarted without losing or -double-spending budget. - -Configuration is driven entirely by environment variables so a -binary-search workflow can shift per-provider rates without code -edits: - - LITELLM_COMPAT_RATE_ANTHROPIC (req/s, default 5.0) - LITELLM_COMPAT_RATE_AZURE (req/s, default 5.0) - LITELLM_COMPAT_RATE_VERTEX_AI (req/s, default 5.0) - LITELLM_COMPAT_RATE_BEDROCK_CONVERSE (req/s, default 5.0) - LITELLM_COMPAT_RATE_BEDROCK_INVOKE (req/s, default 5.0) - LITELLM_COMPAT_RATE_OPENAI (req/s, default 5.0) - LITELLM_COMPAT_RATE_AZURE_OPENAI (req/s, default 5.0) - LITELLM_COMPAT_RATE_BEDROCK_MANTLE (req/s, default 5.0) - LITELLM_COMPAT_RATE_BURST (per-bucket burst override; - default = rate) - LITELLM_COMPAT_RATE_STATE_DIR (state file directory; - default = $TMPDIR/litellm-claude-compat-ratelimit) - -A rate of 0 (or any non-positive value) disables throttling for that -provider — useful when you trust the upstream to handle the burst or -when running unit-test-shaped workloads that never actually hit the -network. - -The provider id is inferred from the model id by `infer_provider`, -mirroring the matrix's column layout (`anthropic`, `azure`, -`vertex_ai`, `bedrock_converse`, `bedrock_invoke`, `openai`, -`azure_openai`, `bedrock_mantle`). The `vertex_ai_gpt` matrix column -has no bucket: its cells are static not_applicable stubs that never -reach the network. -""" - -from __future__ import annotations - -import contextlib -import json -import math -import os -import tempfile -import time -from dataclasses import dataclass -from pathlib import Path -from typing import Dict, Iterator, Mapping, Optional - -# `fcntl` is POSIX-only; the suite is Linux/macOS only, so we don't -# attempt a Windows fallback. Importing at module load fails fast on -# the (currently non-existent) Windows runner so we don't silently -# degrade to no-locking behavior. -import fcntl - -PROVIDER_ANTHROPIC = "anthropic" -PROVIDER_AZURE = "azure" -PROVIDER_VERTEX_AI = "vertex_ai" -PROVIDER_BEDROCK_CONVERSE = "bedrock_converse" -PROVIDER_BEDROCK_INVOKE = "bedrock_invoke" -PROVIDER_OPENAI = "openai" -PROVIDER_AZURE_OPENAI = "azure_openai" -PROVIDER_BEDROCK_MANTLE = "bedrock_mantle" - -ALL_PROVIDERS = ( - PROVIDER_ANTHROPIC, - PROVIDER_AZURE, - PROVIDER_VERTEX_AI, - PROVIDER_BEDROCK_CONVERSE, - PROVIDER_BEDROCK_INVOKE, - PROVIDER_OPENAI, - PROVIDER_AZURE_OPENAI, - PROVIDER_BEDROCK_MANTLE, -) - -DEFAULT_RATE = 5.0 # req/s per provider, conservative starting point -RATE_ENV_PREFIX = "LITELLM_COMPAT_RATE_" -BURST_ENV = "LITELLM_COMPAT_RATE_BURST" -STATE_DIR_ENV = "LITELLM_COMPAT_RATE_STATE_DIR" -DEFAULT_STATE_DIR_NAME = "litellm-claude-compat-ratelimit" - - -def infer_provider(model: str) -> str: - """Map a model alias to its compat-matrix provider id. - - The matrix column layout is fixed; aliases registered in the proxy - encode the provider via a suffix (`-bedrock-converse`, - `-bedrock-invoke`, `-azure`, `-vertex`, `-openai`, `-azure-openai`, - `-bedrock-mantle`) or its absence (Anthropic). Order matters: - `-azure-openai` also ends with `-openai`, and the bedrock suffixes - all contain `bedrock`, so the more-specific suffixes are tested - first. - """ - if not model: - raise ValueError("model must be a non-empty string") - lower = model.lower() - if lower.endswith("-azure-openai"): - return PROVIDER_AZURE_OPENAI - if lower.endswith("-openai"): - return PROVIDER_OPENAI - if lower.endswith("-bedrock-mantle"): - return PROVIDER_BEDROCK_MANTLE - if lower.endswith("-bedrock-converse"): - return PROVIDER_BEDROCK_CONVERSE - if lower.endswith("-bedrock-invoke"): - return PROVIDER_BEDROCK_INVOKE - if lower.endswith("-azure"): - return PROVIDER_AZURE - if lower.endswith("-vertex"): - return PROVIDER_VERTEX_AI - return PROVIDER_ANTHROPIC - - -@dataclass(frozen=True) -class ProviderConfig: - """Static config snapshot for one provider's bucket. - - Captured up front (rather than re-read per acquire) so the limiter's - behavior in a single process is stable even if env vars are mutated - mid-run. A fresh `RateLimiter` picks up env changes on construction. - """ - - rate_per_sec: float - burst: float - - @property - def enabled(self) -> bool: - return self.rate_per_sec > 0 and self.burst > 0 - - -def load_config( - env: Optional[Mapping[str, str]] = None, -) -> Dict[str, ProviderConfig]: - """Build a {provider: ProviderConfig} from env, applying defaults. - - Parsing failures fall back to the default rate rather than - crashing the test session — a typo in `LITELLM_COMPAT_RATE_AZURE` - should not silently disable throttling, but it also shouldn't - abort 75 live tests with a `ValueError` ten minutes in. - """ - src = env if env is not None else os.environ - - def _as_float(value: Optional[str], default: float) -> float: - if value is None or value == "": - return default - try: - return float(value) - except (TypeError, ValueError): - return default - - burst_override = _as_float(src.get(BURST_ENV), -1.0) - - out: Dict[str, ProviderConfig] = {} - for provider in ALL_PROVIDERS: - env_key = RATE_ENV_PREFIX + provider.upper() - rate = _as_float(src.get(env_key), DEFAULT_RATE) - burst = burst_override if burst_override > 0 else max(rate, 1.0) - out[provider] = ProviderConfig(rate_per_sec=rate, burst=burst) - return out - - -def _state_dir(env: Optional[Mapping[str, str]] = None) -> Path: - """Resolve the directory holding per-provider state files. - - A user-supplied `LITELLM_COMPAT_RATE_STATE_DIR` wins for tests and - container environments where `$TMPDIR` may be ephemeral or shared - in surprising ways. The directory is created lazily with - `parents=True, exist_ok=True` so first-run setup needs no - fixture wiring. - """ - src = env if env is not None else os.environ - explicit = src.get(STATE_DIR_ENV) - if explicit: - return Path(explicit) - return Path(tempfile.gettempdir()) / DEFAULT_STATE_DIR_NAME - - -class RateLimiter: - """Cross-process token bucket per provider. - - Each `acquire(provider)` call: - 1. Opens (creating if needed) `/.json`. - 2. Holds an exclusive `flock` while reading + updating the - {tokens, last_refill} state. - 3. Refills tokens based on `now - last_refill`, capped at burst. - 4. If tokens >= 1, subtracts one and returns immediately. - 5. Otherwise, computes the wall-time delay needed to earn a - single token at the configured rate, releases the lock, and - sleeps. After sleeping it retries — staying under the lock - while sleeping would serialize all workers behind whichever - one held it longest. - - This bounds the *aggregate* req/s seen by the upstream, regardless - of how many xdist workers, threads, or processes are concurrently - running tests against the same provider. - - A `_clock` / `_sleep` injection seam keeps the unit tests fast and - deterministic; production callers should never override either. - """ - - def __init__( - self, - config: Optional[Mapping[str, ProviderConfig]] = None, - state_dir: Optional[Path] = None, - clock: Optional[callable] = None, - sleep: Optional[callable] = None, - ) -> None: - self._config = dict(config) if config is not None else load_config() - self._state_dir = Path(state_dir) if state_dir is not None else _state_dir() - self._clock = clock or time.monotonic - self._sleep = sleep or time.sleep - # `_state_dir.mkdir` once on construction is fine; concurrent - # workers all racing to create the same directory is benign. - self._state_dir.mkdir(parents=True, exist_ok=True) - - def acquire(self, provider: str) -> float: - """Block until one token is available for `provider`. - - Returns the cumulative wall-time spent waiting (0.0 when the - bucket had budget and we returned immediately). Callers can - log this to attribute slow cells to throttling vs. upstream - latency — same role `DriverResult.duration_ms` plays for the - actual CLI invocation. - """ - cfg = self._config.get(provider) - if cfg is None or not cfg.enabled: - return 0.0 - - path = self._state_path(provider) - total_waited = 0.0 - while True: - now = self._clock() - sleep_for = self._try_consume(path, cfg, now) - if sleep_for <= 0: - return total_waited - self._sleep(sleep_for) - total_waited += sleep_for - - def _state_path(self, provider: str) -> Path: - return self._state_dir / f"{provider}.json" - - def _try_consume(self, path: Path, cfg: ProviderConfig, now: float) -> float: - """Atomically refill and try to take one token. - - Returns 0.0 if a token was consumed, or a positive sleep - duration (seconds) if the caller must wait before retrying. - - We hold an exclusive `flock` only across the read-modify-write - of the JSON state — never across a `sleep` — so workers don't - serialize while one of them is parked. - """ - # `os.open` + `os.O_CREAT | os.O_RDWR` gives us a fd we can - # both lock and read/write through. Opening with `"a+"` then - # seeking is equivalent but uglier; this version is closer to - # the canonical flock recipe. - fd = os.open(path, os.O_CREAT | os.O_RDWR, 0o600) - try: - fcntl.flock(fd, fcntl.LOCK_EX) - try: - tokens, last_refill = self._read_state(fd, cfg, now) - tokens, last_refill = self._refill(tokens, last_refill, now, cfg) - if tokens >= 1.0: - tokens -= 1.0 - self._write_state(fd, tokens, last_refill) - return 0.0 - # Not enough budget. Persist the refilled state so a - # subsequent caller doesn't have to redo the math, then - # release the lock and tell the caller how long to - # sleep before retrying. - self._write_state(fd, tokens, last_refill) - deficit = 1.0 - tokens - # `deficit / rate` seconds will earn exactly enough - # for one token. Add a tiny safety margin so we don't - # wake up nanoseconds early and spin. - return (deficit / cfg.rate_per_sec) + 1e-3 - finally: - fcntl.flock(fd, fcntl.LOCK_UN) - finally: - os.close(fd) - - @staticmethod - def _read_state(fd: int, cfg: ProviderConfig, now: float) -> tuple: - """Read {tokens, last_refill} from `fd`, defaulting to a full - bucket on a missing/empty/corrupt file. - - New files start full so the very first request never waits; - corrupt files are treated like new files because the - alternative — refusing to run — is worse than briefly - over-spending one bucket's worth of budget. - """ - os.lseek(fd, 0, os.SEEK_SET) - raw = os.read(fd, 4096).decode("utf-8") - if not raw.strip(): - return cfg.burst, now - try: - obj = json.loads(raw) - tokens = float(obj.get("tokens", cfg.burst)) - last_refill = float(obj.get("last_refill", now)) - return tokens, last_refill - except (ValueError, TypeError): - return cfg.burst, now - - @staticmethod - def _refill( - tokens: float, last_refill: float, now: float, cfg: ProviderConfig - ) -> tuple: - """Apply elapsed time to the bucket, capped at burst. - - Negative elapsed (clock went backward, e.g. across host - sleep/resume or a manually-tweaked monotonic mock) is clamped - to zero so we never *remove* tokens. - """ - elapsed = max(0.0, now - last_refill) - tokens = min(cfg.burst, tokens + elapsed * cfg.rate_per_sec) - return tokens, now - - @staticmethod - def _write_state(fd: int, tokens: float, last_refill: float) -> None: - payload = json.dumps({"tokens": tokens, "last_refill": last_refill}).encode( - "utf-8" - ) - os.lseek(fd, 0, os.SEEK_SET) - os.ftruncate(fd, 0) - os.write(fd, payload) - - -# A single process-wide instance is all we need: provider-keyed state -# is stored in files, so multiple `RateLimiter` instances would just -# duplicate the in-process bookkeeping. We expose a getter rather than -# the instance directly so unit tests can install a custom limiter -# scoped to a tmp directory without monkeypatching globals. -_default: Optional[RateLimiter] = None - - -def get_default_limiter() -> RateLimiter: - global _default - if _default is None: - _default = RateLimiter() - return _default - - -def reset_default_limiter() -> None: - """Drop the cached default limiter; the next `get_default_limiter` - call rebuilds it from the current environment. - - Useful between unit tests that patch env vars: without this they'd - keep reading the stale config snapshot from the first call. - """ - global _default - _default = None - - -@contextlib.contextmanager -def use_limiter(limiter: RateLimiter) -> Iterator[RateLimiter]: - """Temporarily install `limiter` as the process default. - - The driver's `run_claude` calls `get_default_limiter()`; tests that - want a controlled tmp-dir-backed limiter use this contextmanager - to swap one in without touching env vars or the on-disk default - state. - """ - global _default - previous = _default - _default = limiter - try: - yield limiter - finally: - _default = previous diff --git a/tests/e2e/claude_code/run_compat.sh b/tests/e2e/claude_code/run_compat.sh index c42416283e7..e3410391661 100755 --- a/tests/e2e/claude_code/run_compat.sh +++ b/tests/e2e/claude_code/run_compat.sh @@ -1,7 +1,7 @@ #!/usr/bin/env bash # Run the Claude Code compat matrix end-to-end against a live LiteLLM # proxy, with per-provider rate limits applied via the cross-process -# token bucket in `tests/e2e/claude_code/rate_limiter.py`. +# token bucket in `tests/e2e/cli_harness/rate_limiter.py`. # # Designed for binary-searching the ideal X / Y / Z req/s per provider: # 1. Pick an initial rate (e.g. 5/s for everyone). diff --git a/tests/e2e/cli_harness/__init__.py b/tests/e2e/cli_harness/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/e2e/cli_harness/lockdown.py b/tests/e2e/cli_harness/lockdown.py new file mode 100644 index 00000000000..1d430529851 --- /dev/null +++ b/tests/e2e/cli_harness/lockdown.py @@ -0,0 +1,49 @@ +"""What a coding-agent CLI under test may see of the host. + +The CI job holds the proxy's upstream credentials (provider keys, AWS, Azure, the database URL, GitHub tokens) in its +environment, and the cron VM runs as a real user with a populated home. A compromised CLI release, or a model told to +read a file, must reach neither, so the CLI gets only the allowlisted runtime variables plus what its suite passes in, +and a fresh empty HOME that is deleted when it exits. +""" + +from __future__ import annotations + +import os +import shutil +import tempfile +from collections.abc import Generator, Mapping +from contextlib import contextmanager +from pathlib import Path +from types import MappingProxyType +from typing import Final + +CLI_ENV_ALLOWLIST: Final = ( + "PATH", + "USER", + "LOGNAME", + "SHELL", + "TERM", + "TMPDIR", + "LANG", + "LC_ALL", + "LC_CTYPE", + "NODE_PATH", + "NVM_DIR", + "NVM_BIN", +) + + +@contextmanager +def isolated_home(prefix: str) -> Generator[Path]: + home: Final = Path(tempfile.mkdtemp(prefix=prefix)) + try: + yield home + finally: + shutil.rmtree(home, ignore_errors=True) + + +def locked_down_env( + *, home: Path, overrides: Mapping[str, str], source: Mapping[str, str] = os.environ +) -> Mapping[str, str]: + allowed: Final = {name: source[name] for name in CLI_ENV_ALLOWLIST if name in source} + return MappingProxyType({**allowed, "HOME": str(home), **overrides}) diff --git a/tests/e2e/cli_harness/npm_version_resolver.py b/tests/e2e/cli_harness/npm_version_resolver.py new file mode 100644 index 00000000000..3c3f9b89ce5 --- /dev/null +++ b/tests/e2e/cli_harness/npm_version_resolver.py @@ -0,0 +1,92 @@ +"""The npm release of a coding-agent CLI that CI installs: the newest one published at least three days ago. + +The three days are a security review buffer (PRD #26476, "Version resolvers"), so a release that gets pulled or +patched within days never drives a run. Newest means by publish time, not semver order. + +`python tests/e2e/cli_harness/npm_version_resolver.py ` prints the version on stdout for a shell +`$(...)`. It is stdlib only because the cron image runs it without the project's dependencies. +""" + +from __future__ import annotations + +import json +import sys +import urllib.parse +import urllib.request +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from types import MappingProxyType +from typing import BinaryIO, Final, cast + +MIN_AGE: Final = timedelta(days=3) +_REGISTRY_URL: Final = "https://registry.npmjs.org/{package}" +_FETCH_TIMEOUT_SECONDS: Final = 30 + + +@dataclass(frozen=True, slots=True) +class Packument: + published_at: Mapping[str, datetime] + versions: frozenset[str] + + @classmethod + def from_json(cls, raw: object) -> Packument: + document: Final = _as_mapping(raw) + times: Final = _as_mapping(document.get("time")) + return cls( + published_at=MappingProxyType({version: _parse_utc(stamp) for version, stamp in _string_values(times)}), + versions=frozenset(_as_mapping(document.get("versions"))), + ) + + +def newest_version_at_least(packument: Packument, *, min_age: timedelta, as_of: datetime) -> str | None: + cutoff: Final = as_of - min_age + eligible: Final = ( + (published, version) + for version, published in packument.published_at.items() + if version in packument.versions and "-" not in version and published <= cutoff + ) + newest: Final = max(eligible, default=None) + return None if newest is None else newest[1] + + +def fetch_packument(package: str) -> Packument: + url: Final = _REGISTRY_URL.format(package=urllib.parse.quote(package, safe="@")) + request: Final = urllib.request.Request(url, headers={"Accept": "application/json"}) + with cast(BinaryIO, urllib.request.urlopen(request, timeout=_FETCH_TIMEOUT_SECONDS)) as response: + return Packument.from_json(cast(object, json.loads(response.read()))) + + +def _as_mapping(raw: object) -> Mapping[str, object]: + if not isinstance(raw, dict): + return {} + return {str(key): value for key, value in cast(dict[object, object], raw).items()} + + +def _string_values(raw: Mapping[str, object]) -> tuple[tuple[str, str], ...]: + return tuple((key, value) for key, value in raw.items() if isinstance(value, str)) + + +def _parse_utc(stamp: str) -> datetime: + parsed: Final = datetime.fromisoformat(stamp.replace("Z", "+00:00")) + return parsed if parsed.tzinfo is not None else parsed.replace(tzinfo=timezone.utc) + + +def _main(argv: list[str]) -> int: + if len(argv) != 1: + print("usage: npm_version_resolver.py ", file=sys.stderr) # noqa: T201 # CLI surface + return 2 + package: Final = argv[0] + version: Final = newest_version_at_least( + fetch_packument(package), min_age=MIN_AGE, as_of=datetime.now(timezone.utc) + ) + if version is None: + print(f"npm_version_resolver: no release of {package} is at least {MIN_AGE} old", file=sys.stderr) # noqa: T201 # CLI surface + return 1 + print(f"npm_version_resolver: selected {package}@{version}", file=sys.stderr) # noqa: T201 # CLI surface + print(version) # noqa: T201 # CLI surface + return 0 + + +if __name__ == "__main__": + raise SystemExit(_main(sys.argv[1:])) diff --git a/tests/e2e/cli_harness/rate_limiter.py b/tests/e2e/cli_harness/rate_limiter.py new file mode 100644 index 00000000000..bd8119ffe87 --- /dev/null +++ b/tests/e2e/cli_harness/rate_limiter.py @@ -0,0 +1,151 @@ +"""A token bucket per upstream, shared by every process on the host, so parallel CLI suites and xdist workers +together stay under one request rate per provider instead of one rate each. + +Each bucket is a small JSON file read and written under an exclusive `flock`, in `LITELLM_COMPAT_RATE_STATE_DIR` +(default: a fixed folder under `$TMPDIR`). `LITELLM_COMPAT_RATE_` sets a bucket's req/s (default 5, 0 or less +turns it off) and `LITELLM_COMPAT_RATE_BURST` overrides every bucket's burst (default: its rate). +""" + +from __future__ import annotations + +import fcntl +import os +import tempfile +import time +from collections.abc import Callable, Iterator, Mapping +from dataclasses import dataclass, field +from enum import StrEnum +from functools import cache +from pathlib import Path +from types import MappingProxyType +from typing import Final + +from pydantic import BaseModel, ConfigDict, ValidationError + +DEFAULT_RATE: Final = 5.0 +RATE_ENV_PREFIX: Final = "LITELLM_COMPAT_RATE_" +BURST_ENV: Final = "LITELLM_COMPAT_RATE_BURST" +STATE_DIR_ENV: Final = "LITELLM_COMPAT_RATE_STATE_DIR" +DEFAULT_STATE_DIR_NAME: Final = "litellm-claude-compat-ratelimit" +_WAKE_MARGIN_SECONDS: Final = 1e-3 + + +class RateBucket(StrEnum): + ANTHROPIC = "anthropic" + AZURE = "azure" + VERTEX_AI = "vertex_ai" + BEDROCK_CONVERSE = "bedrock_converse" + BEDROCK_INVOKE = "bedrock_invoke" + OPENAI = "openai" + AZURE_OPENAI = "azure_openai" + BEDROCK_MANTLE = "bedrock_mantle" + + +@dataclass(frozen=True, slots=True) +class ProviderConfig: + rate_per_sec: float + burst: float + + @property + def enabled(self) -> bool: + return self.rate_per_sec > 0 and self.burst > 0 + + +class _BucketState(BaseModel): + model_config = ConfigDict(frozen=True) + + tokens: float + last_refill: float + + +def load_config(env: Mapping[str, str] = os.environ) -> Mapping[RateBucket, ProviderConfig]: + burst_override: Final = _as_float(env.get(BURST_ENV), -1.0) + return MappingProxyType( + { + bucket: _bucket_config(_as_float(env.get(_rate_env(bucket)), DEFAULT_RATE), burst_override) + for bucket in RateBucket + } + ) + + +def default_state_dir(env: Mapping[str, str] = os.environ) -> Path: + explicit: Final = env.get(STATE_DIR_ENV) + return Path(explicit) if explicit else Path(tempfile.gettempdir()) / DEFAULT_STATE_DIR_NAME + + +@dataclass(frozen=True, slots=True) +class RateLimiter: + config: Mapping[RateBucket, ProviderConfig] = field(default_factory=load_config) + state_dir: Path = field(default_factory=default_state_dir) + clock: Callable[[], float] = time.monotonic + sleep: Callable[[float], None] = time.sleep + + def __post_init__(self) -> None: + self.state_dir.mkdir(parents=True, exist_ok=True) + + def acquire(self, bucket: RateBucket) -> float: + cfg: Final = self.config.get(bucket) + if cfg is None or not cfg.enabled: + return 0.0 + return sum(self._waits(self.state_dir / f"{bucket.value}.json", cfg)) + + def _waits(self, path: Path, cfg: ProviderConfig) -> Iterator[float]: + while (wait := _try_consume(path, cfg, self.clock())) > 0: + self.sleep(wait) + yield wait + + +@cache +def get_default_limiter() -> RateLimiter: + return RateLimiter() + + +def _rate_env(bucket: RateBucket) -> str: + return RATE_ENV_PREFIX + bucket.value.upper() + + +def _bucket_config(rate: float, burst_override: float) -> ProviderConfig: + return ProviderConfig(rate_per_sec=rate, burst=burst_override if burst_override > 0 else max(rate, 1.0)) + + +def _as_float(value: str | None, default: float) -> float: + if not value: + return default + try: + return float(value) + except ValueError: + return default + + +def _try_consume(path: Path, cfg: ProviderConfig, now: float) -> float: + fd: Final = os.open(path, os.O_CREAT | os.O_RDWR, 0o600) + try: + fcntl.flock(fd, fcntl.LOCK_EX) + state: Final = _refilled(_read_state(fd, cfg, now), now, cfg) + if state.tokens >= 1.0: + _write_state(fd, _BucketState(tokens=state.tokens - 1.0, last_refill=state.last_refill)) + return 0.0 + _write_state(fd, state) + return (1.0 - state.tokens) / cfg.rate_per_sec + _WAKE_MARGIN_SECONDS + finally: + os.close(fd) + + +def _read_state(fd: int, cfg: ProviderConfig, now: float) -> _BucketState: + os.lseek(fd, 0, os.SEEK_SET) + raw: Final = os.read(fd, 4096) + try: + return _BucketState.model_validate_json(raw) + except ValidationError: + return _BucketState(tokens=cfg.burst, last_refill=now) + + +def _refilled(state: _BucketState, now: float, cfg: ProviderConfig) -> _BucketState: + elapsed: Final = max(0.0, now - state.last_refill) + return _BucketState(tokens=min(cfg.burst, state.tokens + elapsed * cfg.rate_per_sec), last_refill=now) + + +def _write_state(fd: int, state: _BucketState) -> None: + os.lseek(fd, 0, os.SEEK_SET) + os.ftruncate(fd, 0) + os.write(fd, state.model_dump_json().encode("utf-8")) diff --git a/tests/e2e_harness/AGENTS.md b/tests/e2e_harness/AGENTS.md index 65a3b32df96..8528424166c 100644 --- a/tests/e2e_harness/AGENTS.md +++ b/tests/e2e_harness/AGENTS.md @@ -2,7 +2,7 @@ Tests of the harness under `tests/e2e/` (the transport, the clients, the fixture bundle and replay edge, the stack lock, the IdP launcher, the coverage collector, the JUnit properties and the Claude Code driver), not of the product. They live outside `tests/e2e/` because the Buildkite e2e run copies that folder into the runner image and runs every test in it, so a harness test in there counts as a product test in the nightly numbers. Nothing here needs a proxy, provider keys or the network -The layout mirrors `tests/e2e/`: `test_e2e_http.py` covers `tests/e2e/e2e_http.py`, `logging/test_datadog_reader.py` covers `tests/e2e/logging/datadog_reader.py`, and `claude_code/` covers the driver, builder, probe and version resolver. Put a new harness test under the folder that mirrors the suite folder whose module it covers +The layout mirrors `tests/e2e/`: `test_e2e_http.py` covers `tests/e2e/e2e_http.py`, `logging/test_datadog_reader.py` covers `tests/e2e/logging/datadog_reader.py`, `claude_code/` covers the driver, builder and probe, and `cli_harness/` covers the version resolver, the CLI env lockdown and the rate limiter. Put a new harness test under the folder that mirrors the suite folder whose module it covers Run them from the repo root. `e2e_config` reads `LITELLM_MASTER_KEY` at import and any value will do, the CI lane sets a dummy: diff --git a/tests/e2e_harness/claude_code/test_pr_gate_version_resolver.py b/tests/e2e_harness/claude_code/test_pr_gate_version_resolver.py deleted file mode 100644 index 0ed7e2bf083..00000000000 --- a/tests/e2e_harness/claude_code/test_pr_gate_version_resolver.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Unit tests for the Claude Code PR-gate version resolver. - -Markerless harness tests: they feed the resolver a hand-built packument and a -fixed clock, so they run without a proxy, never reach the npm registry, and -carry no `e2e` marker. -""" - -from __future__ import annotations - -from datetime import datetime, timezone -from typing import Final, Mapping - -import pytest - -from claude_code.pr_gate_version_resolver import NoEligibleVersionError, resolve_pr_gate_version - -NOW: Final = datetime(2026, 4, 25, 12, 0, tzinfo=timezone.utc) -INSIDE_THE_2_1_88_WINDOW: Final = datetime(2026, 4, 3, 12, 0, tzinfo=timezone.utc) - - -def _packument(times: Mapping[str, str], unpublished: frozenset[str] = frozenset()) -> dict[str, object]: - return { - "name": "@anthropic-ai/claude-code", - "time": {"created": "2024-01-01T00:00:00.000Z", "modified": "2026-04-25T00:00:00.000Z", **times}, - "versions": {version: {"version": version} for version in times if version not in unpublished}, - } - - -def test_skips_a_version_npm_has_unpublished() -> None: - metadata: Final = _packument( - { - "2.1.87": "2026-03-28T20:00:00.000Z", - "2.1.88": "2026-03-30T22:36:48.424Z", - "2.1.89": "2026-03-31T23:32:40.000Z", - }, - unpublished=frozenset({"2.1.88"}), - ) - assert resolve_pr_gate_version(metadata=metadata, as_of=INSIDE_THE_2_1_88_WINDOW) == "2.1.87" - - -def test_raises_when_the_only_old_enough_version_is_unpublished() -> None: - metadata: Final = _packument( - {"2.1.88": "2026-03-30T22:36:48.424Z", "2.1.89": "2026-03-31T23:32:40.000Z"}, - unpublished=frozenset({"2.1.88"}), - ) - with pytest.raises(NoEligibleVersionError): - resolve_pr_gate_version(metadata=metadata, as_of=INSIDE_THE_2_1_88_WINDOW) - - -def test_picks_the_newest_published_version_at_least_min_age_old() -> None: - metadata: Final = _packument( - { - "2.1.118": "2026-04-15T10:00:00.000Z", - "2.1.119": "2026-04-21T10:00:00.000Z", - "2.2.0-rc.1": "2026-04-22T10:00:00.000Z", - "2.1.120": "2026-04-23T10:00:00.000Z", - "2.1.121": "2026-04-25T11:00:00.000Z", - } - ) - assert resolve_pr_gate_version(metadata=metadata, as_of=NOW) == "2.1.119" diff --git a/tests/e2e_harness/claude_code/test_request_determinism.py b/tests/e2e_harness/claude_code/test_request_determinism.py index b7d330b7da6..c6edb02d1c5 100644 --- a/tests/e2e_harness/claude_code/test_request_determinism.py +++ b/tests/e2e_harness/claude_code/test_request_determinism.py @@ -34,7 +34,7 @@ from typing import List, Tuple import pytest from claude_code.cli_driver import _FIXED_CLI_USER_ID, _seed_cli_identity, _stable_cli_state, run_claude -from claude_code.rate_limiter import RateLimiter +from cli_harness.rate_limiter import RateLimiter pytestmark = pytest.mark.cli_determinism diff --git a/tests/e2e_harness/cli_harness/test_lockdown.py b/tests/e2e_harness/cli_harness/test_lockdown.py new file mode 100644 index 00000000000..fcb24f4a830 --- /dev/null +++ b/tests/e2e_harness/cli_harness/test_lockdown.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Final + +from cli_harness.lockdown import isolated_home, locked_down_env + +RUNNER_ENV: Final = { + "PATH": "/usr/bin", + "LANG": "C.UTF-8", + "HOME": "/home/runner", + "OPENAI_API_KEY": "sk-runner-secret", + "AWS_SECRET_ACCESS_KEY": "runner-aws-secret", + "DATABASE_URL": "postgresql://runner", +} + + +def test_cli_sees_runtime_vars_and_its_overrides_but_no_runner_credentials(tmp_path: Path) -> None: + env: Final = locked_down_env(home=tmp_path, overrides={"OPENAI_BASE_URL": "http://gateway"}, source=RUNNER_ENV) + assert dict(env) == { + "PATH": "/usr/bin", + "LANG": "C.UTF-8", + "HOME": str(tmp_path), + "OPENAI_BASE_URL": "http://gateway", + } + + +def test_isolated_home_starts_empty_and_is_gone_after_the_cli_exits() -> None: + with isolated_home("cli-harness-test-") as home: + assert home.is_dir() and not any(home.iterdir()) + (home / ".config").mkdir() + assert not home.exists() diff --git a/tests/e2e_harness/cli_harness/test_npm_version_resolver.py b/tests/e2e_harness/cli_harness/test_npm_version_resolver.py new file mode 100644 index 00000000000..3a84ccd361a --- /dev/null +++ b/tests/e2e_harness/cli_harness/test_npm_version_resolver.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +from typing import Final + +from cli_harness.npm_version_resolver import MIN_AGE, Packument, newest_version_at_least + +NOW: Final = datetime(2026, 4, 25, 12, 0, tzinfo=timezone.utc) +INSIDE_THE_2_1_88_WINDOW: Final = datetime(2026, 4, 3, 12, 0, tzinfo=timezone.utc) + + +def _packument(times: Mapping[str, str], unpublished: frozenset[str] = frozenset()) -> Packument: + return Packument.from_json( + { + "name": "@scope/some-cli", + "time": {"created": "2024-01-01T00:00:00.000Z", "modified": "2026-04-25T00:00:00.000Z", **times}, + "versions": {version: {"version": version} for version in times if version not in unpublished}, + } + ) + + +def test_skips_a_version_npm_has_unpublished() -> None: + packument: Final = _packument( + { + "2.1.87": "2026-03-28T20:00:00.000Z", + "2.1.88": "2026-03-30T22:36:48.424Z", + "2.1.89": "2026-03-31T23:32:40.000Z", + }, + unpublished=frozenset({"2.1.88"}), + ) + assert newest_version_at_least(packument, min_age=MIN_AGE, as_of=INSIDE_THE_2_1_88_WINDOW) == "2.1.87" + + +def test_finds_nothing_when_the_only_old_enough_version_is_unpublished() -> None: + packument: Final = _packument( + {"2.1.88": "2026-03-30T22:36:48.424Z", "2.1.89": "2026-03-31T23:32:40.000Z"}, + unpublished=frozenset({"2.1.88"}), + ) + assert newest_version_at_least(packument, min_age=MIN_AGE, as_of=INSIDE_THE_2_1_88_WINDOW) is None + + +def test_picks_the_newest_stable_version_at_least_min_age_old() -> None: + packument: Final = _packument( + { + "2.1.118": "2026-04-15T10:00:00.000Z", + "2.1.119": "2026-04-21T10:00:00.000Z", + "2.2.0-rc.1": "2026-04-22T10:00:00.000Z", + "2.1.120": "2026-04-23T10:00:00.000Z", + "2.1.121": "2026-04-25T11:00:00.000Z", + } + ) + assert newest_version_at_least(packument, min_age=MIN_AGE, as_of=NOW) == "2.1.119" + + +def test_newest_means_latest_published_not_highest_semver() -> None: + packument: Final = _packument( + { + "3.0.0": "2026-04-01T10:00:00.000Z", + "2.9.9": "2026-04-10T10:00:00.000Z", + } + ) + assert newest_version_at_least(packument, min_age=timedelta(days=3), as_of=NOW) == "2.9.9" diff --git a/tests/e2e_harness/cli_harness/test_rate_limiter.py b/tests/e2e_harness/cli_harness/test_rate_limiter.py new file mode 100644 index 00000000000..78ad3e92450 --- /dev/null +++ b/tests/e2e_harness/cli_harness/test_rate_limiter.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +import multiprocessing +from collections.abc import Mapping +from dataclasses import dataclass +from multiprocessing.synchronize import Barrier +from pathlib import Path +from types import MappingProxyType +from typing import Final, NoReturn + +from cli_harness.rate_limiter import ProviderConfig, RateBucket, RateLimiter + +_RACERS: Final = 8 +_ATTEMPTS_PER_RACER: Final = 100 +_BURST: Final = 400 +_NO_REFILL: Final[Mapping[RateBucket, ProviderConfig]] = MappingProxyType( + {RateBucket.OPENAI: ProviderConfig(rate_per_sec=1e-9, burst=_BURST)} +) + + +@dataclass +class _FakeTime: + now: float = 0.0 + + def clock(self) -> float: + return self.now + + def sleep(self, seconds: float) -> None: + self.now += seconds + + +def test_two_processes_on_one_host_share_a_single_budget_per_bucket(tmp_path: Path) -> None: + fake_time: Final = _FakeTime() + config: Final = {RateBucket.OPENAI: ProviderConfig(rate_per_sec=2.0, burst=1.0)} + first_process: Final = RateLimiter(config, tmp_path, fake_time.clock, fake_time.sleep) + second_process: Final = RateLimiter(config, tmp_path, fake_time.clock, fake_time.sleep) + + assert first_process.acquire(RateBucket.OPENAI) == 0.0 + assert 0.5 <= second_process.acquire(RateBucket.OPENAI) < 0.51 + assert second_process.acquire(RateBucket.BEDROCK_MANTLE) == 0.0 + + +def test_processes_racing_for_one_bucket_never_get_more_than_its_burst(tmp_path: Path) -> None: + grants: Final = tmp_path / "grants" + + _race(_RACERS, state_dir=tmp_path / "state", grants=grants) + + assert len(grants.read_bytes()) == _BURST + + +def _race(racers: int, *, state_dir: Path, grants: Path) -> None: + context: Final = multiprocessing.get_context("spawn") + start_line: Final = context.Barrier(racers) + processes: Final = tuple( + context.Process(target=_take_tokens_until_told_to_wait, args=(state_dir, grants, start_line)) + for _ in range(racers) + ) + for process in processes: + process.start() + for process in processes: + process.join() + + +def _take_tokens_until_told_to_wait(state_dir: Path, grants: Path, start_line: Barrier) -> None: + limiter: Final = RateLimiter(_NO_REFILL, state_dir, _frozen_clock, _exit_instead_of_sleeping) + with grants.open("ab", buffering=0) as log: + _ = start_line.wait() + for _ in range(_ATTEMPTS_PER_RACER): + _ = limiter.acquire(RateBucket.OPENAI) + _ = log.write(b".") + + +def _frozen_clock() -> float: + return 0.0 + + +def _exit_instead_of_sleeping(_seconds: float) -> NoReturn: + raise SystemExit(0) From 4ede2c012f15bad9df94ed08a0688c8a744c38b3 Mon Sep 17 00:00:00 2001 From: jesus-berri Date: Mon, 2 Nov 2026 13:53:58 -0800 Subject: [PATCH 16/79] feat(model_info): add supports_continuation capability flag for Gemini long decode continuation (#49500) * feat(model_info): add supports_continuation capability flag for Gemini long decode continuation * fix(router): keep a deployment's supports_continuation override on that deployment A supports_continuation value set in one deployment's model_info was written onto the shared provider/model cost map key, so every other deployment of that model reported it too. List it with the other per-deployment capability flags, and cite Google's long decode page next to the cost map assertion * fix(cost-map): flag vertex_ai/gemini-4-argon-cyber-permissive with supports_continuation Google's Vertex Gemini 4 Argon guide gives gemini-4-argon-cyber-permissive the same long decode limit as gemini-4-argon, and LiteLLM already continues its long requests, so the cost map now says so in both copies. * fix(router): count only a boolean true supports_continuation when building a model group --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/__init__.py | 1 + litellm/_lazy_imports_registry.py | 2 + ...odel_prices_and_context_window_backup.json | 3 ++ litellm/router_python_core.py | 3 ++ litellm/types/router.py | 1 + litellm/types/utils.py | 2 + litellm/utils.py | 10 +++++ model_prices_and_context_window.json | 3 ++ .../test_model_prices_and_context_window.py | 20 +++++++++ .../test_router/test_router_helper_utils.py | 45 +++++++++++++++++++ .../unit/test_router_model_cost_isolation.py | 5 ++- tests/unit/test_utils.py | 29 ++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 +++ 13 files changed, 128 insertions(+), 1 deletion(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 0e5101b49c3..baee4e9f7ba 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -2374,6 +2374,7 @@ if TYPE_CHECKING: supports_web_search: Callable[..., bool] supports_tool_search: Callable[..., bool] supports_url_context: Callable[..., bool] + supports_continuation: Callable[..., bool] supports_response_schema: Callable[..., bool] supports_parallel_function_calling: Callable[..., bool] supports_vision: Callable[..., bool] diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index de290070ed3..c9949613303 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -33,6 +33,7 @@ UTILS_NAMES: Final = ( "supports_web_search", "supports_tool_search", "supports_url_context", + "supports_continuation", "supports_response_schema", "supports_parallel_function_calling", "supports_vision", @@ -488,6 +489,7 @@ _UTILS_IMPORT_MAP: Final = { "supports_web_search": (".utils", "supports_web_search"), "supports_tool_search": (".utils", "supports_tool_search"), "supports_url_context": (".utils", "supports_url_context"), + "supports_continuation": (".utils", "supports_continuation"), "supports_response_schema": (".utils", "supports_response_schema"), "supports_parallel_function_calling": ( ".utils", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e6f65d6b416..8bfb8949890 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -76813,6 +76813,7 @@ "supports_audio_input": true, "supports_audio_output": false, "supports_computer_use": true, + "supports_continuation": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -76872,6 +76873,7 @@ ], "supports_audio_input": true, "supports_computer_use": true, + "supports_continuation": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -76929,6 +76931,7 @@ "text" ], "supports_audio_input": true, + "supports_continuation": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 98c9e641278..2235abc9897 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -13247,6 +13247,7 @@ class PythonCore(DeprecatedRouterMembers): "providers": [llm_provider], "description": model_group_description, **model_info, + "supports_continuation": model_info.get("supports_continuation") is True, "supports_fast_mode": True, "supported_reasoning_efforts": None, "supported_endpoints": None, @@ -13314,6 +13315,8 @@ class PythonCore(DeprecatedRouterMembers): and model_info["supports_url_context"] is True ): model_group_info.supports_url_context = True + if model_info.get("supports_continuation") is True: + model_group_info.supports_continuation = True if model_info.get("supports_reasoning", None) is not None and model_info["supports_reasoning"] is True: model_group_info.supports_reasoning = True diff --git a/litellm/types/router.py b/litellm/types/router.py index 41a986837da..4e3e4d92b3c 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -1061,6 +1061,7 @@ class ModelGroupInfo(LiteLLMBaseModel): supports_vision: bool = Field(default=False) supports_web_search: bool = Field(default=False) supports_url_context: bool = Field(default=False) + supports_continuation: bool = Field(default=False) supports_reasoning: bool = Field(default=False) supports_function_calling: bool = Field(default=False) supports_fast_mode: bool = Field(default=False) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index bf69445fdd1..eeb0956a06b 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -227,6 +227,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_tool_search: bool | None supports_mid_conversation_system: bool | None supports_url_context: bool | None + supports_continuation: ReadOnly[bool | None] supports_none_reasoning_effort: bool | None supports_minimal_reasoning_effort: bool | None supports_low_reasoning_effort: bool | None @@ -4234,6 +4235,7 @@ DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset( "supports_assistant_prefill", "supports_bedrock_native_structured_output", "supports_pdf_input", + "supports_continuation", "reasoning_modes", "reasoning_contexts", } diff --git a/litellm/utils.py b/litellm/utils.py index f03b1eb7ad9..55423a42260 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3187,6 +3187,14 @@ def supports_url_context(model: str, custom_llm_provider: str | None = None) -> ) +def supports_continuation(model: str, custom_llm_provider: str | None = None) -> bool: + return supports_factory( + model=model, + custom_llm_provider=custom_llm_provider, + key="supports_continuation", + ) + + def supports_native_streaming(model: str, custom_llm_provider: str | None) -> bool: """ Check if the given model supports native streaming and return a boolean value. @@ -7377,6 +7385,7 @@ def get_model_info_helper( ), supports_web_search=_model_info.get("supports_web_search", None), supports_url_context=_model_info.get("supports_url_context", None), + supports_continuation=_model_info.get("supports_continuation", None), supports_reasoning=_model_info.get("supports_reasoning", None), supports_adaptive_thinking=_model_info.get("supports_adaptive_thinking", None), supports_legacy_thinking=_model_info.get("supports_legacy_thinking", None), @@ -7536,6 +7545,7 @@ def get_model_info( supports_pdf_input: Optional[bool] supports_web_search: Optional[bool] supports_url_context: Optional[bool] + supports_continuation: Optional[bool] supports_reasoning: Optional[bool] Raises: Exception: If the model is not mapped yet. diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e6f65d6b416..8bfb8949890 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -76813,6 +76813,7 @@ "supports_audio_input": true, "supports_audio_output": false, "supports_computer_use": true, + "supports_continuation": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -76872,6 +76873,7 @@ ], "supports_audio_input": true, "supports_computer_use": true, + "supports_continuation": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -76929,6 +76931,7 @@ "text" ], "supports_audio_input": true, + "supports_continuation": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index e4a2beb3a86..4e8ede5c446 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -845,6 +845,26 @@ def _search_free_requests_rows(path: Path) -> Mapping[str, dict[str, JsonValue]] ) +def _supports_continuation_rows(path: Path) -> frozenset[str]: + prices: Final[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]).validate_json(path.read_bytes()) + return frozenset( + key for key, row in prices.items() if isinstance(row, dict) and row.get("supports_continuation") is True + ) + + +def test_gemini_4_argon_google_rows_carry_supports_continuation_in_both_cost_map_copies() -> None: + main_rows: Final = _supports_continuation_rows(PRICES_PATH) + # long decode models: https://ai.google.dev/gemini-api/docs/generate-content/long-decode and + # https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/guides/gemini-4-argon (2026-11-02) + long_decode_rows: Final = { + "gemini/gemini-4-argon", + "vertex_ai/gemini-4-argon", + "vertex_ai/gemini-4-argon-cyber-permissive", + } + assert long_decode_rows <= main_rows, sorted(main_rows) + assert _supports_continuation_rows(BACKUP_PRICES_PATH) == main_rows + + @pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) def test_search_free_requests_rows_are_grounding_rows_with_one_allowance_per_group(path: Path) -> None: rows = _search_free_requests_rows(path) diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 4964dea498f..b46a57ad5f1 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -1497,6 +1497,51 @@ def test_set_model_group_info(model_list, user_facing_model_group_name): assert resp.model_group == user_facing_model_group_name +def test_set_model_group_info_reports_supports_continuation_when_any_deployment_has_it(): + router: Final = Router( + model_list=[ + { + "model_name": "long-decode", + "litellm_params": {"model": "openai/gpt-4.1", "api_key": "fake"}, + }, + { + "model_name": "long-decode", + "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "fake"}, + "model_info": {"supports_continuation": True}, + }, + { + "model_name": "long-decode", + "litellm_params": {"model": "openai/gpt-4.1-nano", "api_key": "fake"}, + }, + { + "model_name": "plain", + "litellm_params": {"model": "openai/gpt-4.1", "api_key": "fake"}, + }, + ] + ) + long_decode: Final = router._set_model_group_info( + model_group="long-decode", user_facing_model_group_name="long-decode" + ) + plain: Final = router._set_model_group_info(model_group="plain", user_facing_model_group_name="plain") + assert long_decode is not None and long_decode.supports_continuation is True + assert plain is not None and plain.supports_continuation is False + + +@pytest.mark.parametrize("value", ["maybe", "", ["x"], "true", 1]) +def test_set_model_group_info_counts_only_a_boolean_true_supports_continuation(value: object) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "long-decode", + "litellm_params": {"model": "openai/gpt-4.1", "api_key": "fake"}, + "model_info": {"supports_continuation": value}, + }, + ] + ) + group: Final = router._set_model_group_info(model_group="long-decode", user_facing_model_group_name="long-decode") + assert group is not None and group.supports_continuation is False + + @pytest.mark.asyncio async def test_set_response_headers(model_list): """Test if the 'set_response_headers' function is working correctly""" diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 447a4dfb81d..96519a0acbe 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -585,7 +585,10 @@ def test_should_not_pollute_shared_key_with_custom_nonzero_pricing(): ) -@pytest.mark.parametrize("flag", ["supports_regex_lookaround", "supports_assistant_prefill", "supports_sampling_params"]) +@pytest.mark.parametrize( + "flag", + ["supports_regex_lookaround", "supports_assistant_prefill", "supports_sampling_params", "supports_continuation"], +) def test_capability_flag_stays_on_the_deployment_that_set_it(flag: str) -> None: """A deployment's capability override must not land on the shared ``{provider}/{model}`` key, or every sibling deployment of that model would inherit it.""" diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 61825e463f8..a9b56568afb 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -476,6 +476,34 @@ def test_supports_video_input_reads_its_own_cost_map_flag( assert litellm.get_model_info(model)["supports_video_input"] is video_input +@pytest.mark.parametrize("continuation", [True, False]) +def test_supports_continuation_reads_its_own_cost_map_flag( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, continuation: bool +) -> None: + model: Final = "gemini/continuation-flag-model" + monkeypatch.setitem( + litellm.model_cost, + model, + { + "litellm_provider": "gemini", + "mode": "chat", + "supports_continuation": continuation, + "supports_url_context": not continuation, + }, + ) + assert litellm.supports_continuation(model) is continuation + assert litellm.get_model_info(model).get("supports_continuation") is continuation + + +def test_supports_continuation_is_unset_when_the_cost_map_omits_it( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + model: Final = "gemini/no-continuation-flag-model" + monkeypatch.setitem(litellm.model_cost, model, {"litellm_provider": "gemini", "mode": "chat"}) + assert litellm.supports_continuation(model) is False + assert litellm.get_model_info(model).get("supports_continuation") is None + + def test_get_optional_params_image_gen(): from litellm.llms.azure.image_generation import AzureGPTImageGenerationConfig @@ -1192,6 +1220,7 @@ _COST_MAP_ROW_SCHEMA: Final = { "supports_bedrock_runtime_chat_completions_inline_reasoning": {"type": "boolean"}, "supports_bedrock_converse_tool_result_images": {"type": "boolean"}, "supports_url_context": {"type": "boolean"}, + "supports_continuation": {"type": "boolean"}, "supports_multimodal": {"type": "boolean"}, "uses_embed_content": {"type": "boolean"}, "supports_reasoning": {"type": "boolean"}, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 96d4c0e1bc4..3aa7c888e98 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -38948,6 +38948,11 @@ export interface components { * @default false */ supports_background_chat: boolean; + /** + * Supports Continuation + * @default false + */ + supports_continuation: boolean; /** * Supports Fast Mode * @default false From 250af8f027dbfaa9a7d2e1867cf53254056de60f Mon Sep 17 00:00:00 2001 From: simonp9 <349053457+simonp9@users.noreply.github.com> Date: Mon, 2 Nov 2026 14:13:17 -0800 Subject: [PATCH 17/79] fix(spend)!: save streamed /v1/responses spend rows under the encrypted resp_ id the client received (#49493) * fix(responses): log the encrypted response id on streamed /v1/responses spend rows With response id security on, the stream hook encrypts the resp_ id on every frame, but logging had already received the SDK's own copy of the terminal response, so the spend row, the standard logging payload and callbacks got the managed id (or the chatcmpl id on a bridged stream) the client never saw. The hook now records the exact id it advertises on the request's logging object, and the logging layer stamps that id on the logged streamed response * fix(responses): return the id the client sent on GET and cancel, and keep read rows off the create row's id Every GET /v1/responses/{id} re-encrypted the response id, so a client that continued a bridged conversation with the id a GET returned referenced a ciphertext no spend row was filed under. The follow-up lost turn 1 and opened a new session. GET and cancel now answer with the exact id the client sent when it is an id this proxy issued and it decodes to the response the provider returned. A create with previous_response_id still gets a fresh id for its own response, and a raw provider id still gets a managed id Because a read now carries the create's id, a GET or cancel spend row is filed as _retrieve_ or _cancel_. That keeps one row per read, as before, and a read row can never take the create row's request_id, so the create row survives whichever row is written first and a follow-up replays turn 1 once. The background settlement row keeps _background_cost, and the cost poll's own status checks keep the response id * fix(responses): log a fresh id for each read and cancel, echo the held id only to its owner A GET or cancel of a stored response now logs a new encrypted id of its own, so callbacks and spend rows never reuse the create's id, while the owner of the id still gets back the exact id they sent. An admin reading someone else's response gets an id bound to the admin, as before. The spend row id suffix for reads and cancels is gone, so with response id security off the rows and session history match the default branch. A streamed GET follows the same rules. * fix(responses): echo the sent id to every reader that passes the response access check A GET, streamed GET or cancel of a response now answers with the id the caller sent whenever the caller passed the access check, not only when the caller is the user or team the id was issued to. An admin, the master key or a teammate who read someone else's streamed response got a fresh id back, and the session lookup could not reach the first turn from it, so a bridged follow-up through that id lost the conversation history Every read and cancel still logs a fresh id bound to the reader, never the create's id * fix(responses): keep a read or cancel row in session history only when it brings new output A GET, streamed GET or cancel sent with a session header joins the create's session, and its stored response repeats the create's output. The replay then sent the same assistant turn once per read. Session history now keeps an aget_responses or acancel_responses row only when its stored Responses output has an item id that no earlier kept row carried. A polled background response keeps the first poll that returns the finished output, since its queued create row has none, and later polls of the same output are dropped. Every other row keeps today's rule * fix(responses): replay each history item once, echo the sent id on delete, keep background cost files apart A create row now owns its output items in session history, and a read or cancel row adds only items no create and no earlier read carried, in one linear pass, so a background response polled while it grows, or a read that ends before its streamed create, replays the answer once and after its prompt DELETE /v1/responses/{id} answers with the id the caller sent, on a copy, and logs a fresh encrypted id like a read does Azure Data Lake Storage per-event files and unbatched GCS objects are named after the event's spend row id, so a streamed background create and its cost event, which now share the client's id, keep separate files The integration tests that decoded the managed id now match the exact id the client received * fix(responses): keep main's answer for a raw-id delete and main's storage names outside background cost events A DELETE by a raw provider id returns the provider's id as main did, and logging still records a fresh encrypted id. A streamed GET matches the held id by upstream response id, like a plain GET Only a background response's cost event gets the _background_cost suffix on its Azure Data Lake Storage file and unbatched GCS object, so batch retrieves and events with no id keep main's names An item a session read replays takes the newest read's copy at the first read's place, so an item a poll caught half written replays finished Tests pin the client's id over the router's continued response id, and the storage test no longer reads the clock after the upload * fix(spend): give background cost events their own OpenMeter and DynamoDB ids, replay id-less read items A streamed background create and its later cost event share the client's id, so OpenMeter dropped the second event and DynamoDB overwrote the create's item. Both now key the cost event as _background_cost, like its spend row and storage file. Session history again replays a read's output items that carry no id, as main did. * fix(integrations): keep OpenMeter and DynamoDB logging working without the proxy extras * fix(integrations): give background cost events their own ClickHouse request id and DeepEval span id A streamed background create and its later cost event now share the id the client received, so ClickHouse and DeepEval saw both events under one id. Route both through logged_event_key, the way the Azure storage, GCS, OpenMeter and DynamoDB loggers already do, so the cost event's ClickHouse request_id matches its spend row's `_background_cost` and its DeepEval span gets its own uuid. ClickHouse response_id stays the raw provider id. * test(deepeval): annotate the background cost span test's locals with Final * test(responses): run the DELETE wire test that pins the echoed client id DELETE /v1/responses/{id} now answers with the id the caller sent, so the skip that marked it as a known bug no longer applies. The test passes on the /v1, /openai/v1 and bare routes. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../azure_storage/azure_storage.py | 50 +- .../clickhouse/clickhouse_spend_logger.py | 7 +- litellm/integrations/deepeval/deepeval.py | 5 +- litellm/integrations/dynamodb.py | 18 +- litellm/integrations/gcs_bucket/gcs_bucket.py | 37 +- litellm/integrations/openmeter.py | 11 +- .../internal_call_metadata.py | 27 + litellm/litellm_core_utils/litellm_logging.py | 26 +- litellm/proxy/hooks/responses_id_security.py | 132 +++- .../spend_tracking/spend_tracking_utils.py | 33 +- .../session_handler.py | 98 ++- .../test_responses_retrieve_e2e.py | 3 - .../test_openai_sdk_client_paths.py | 15 +- .../test_langfuse_otel_metadata_chaos.py | 6 +- .../test_langfuse_otel_metadata_clients.py | 9 +- .../test_otel_conversation_id.py | 44 +- .../observability/test_signoz_delivery.py | 37 +- .../test_bedrock_gpt_responses_native_wire.py | 6 +- .../test_bedrock_mantle_output_limit_wire.py | 2 +- .../test_bedrock_model_discovery_wire.py | 9 +- ..._bedrock_runtime_chat_completions_chaos.py | 6 +- .../test_responses_lifecycle_wire.py | 4 - .../test_responses_stream_spend_row_id.py | 465 +++++++++++++ .../azure_storage/test_azure_storage.py | 181 ++++- .../test_clickhouse_spend_logger.py | 43 +- tests/unit/integrations/conftest.py | 36 +- .../gcs_bucket/test_gcs_bucket.py | 101 ++- tests/unit/integrations/test_deepeval.py | 35 + tests/unit/integrations/test_dynamodb.py | 52 ++ tests/unit/integrations/test_openmeter.py | 44 ++ .../test_litellm_logging.py | 34 + .../test_spend_tracking_utils.py | 17 + .../test_session_handler.py | 122 ++++ tests/unit/test_responses_id_security.py | 618 ++++++++++++++---- 34 files changed, 1981 insertions(+), 352 deletions(-) create mode 100644 tests/integration/spend/test_responses_stream_spend_row_id.py create mode 100644 tests/unit/integrations/test_dynamodb.py diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py index 37b4d2d3937..2ff523571f2 100644 --- a/litellm/integrations/azure_storage/azure_storage.py +++ b/litellm/integrations/azure_storage/azure_storage.py @@ -1,7 +1,8 @@ import asyncio import os import time -from collections.abc import Callable +from collections.abc import Callable, Mapping +from dataclasses import dataclass from datetime import datetime, timedelta, timezone from functools import cache from typing import Final @@ -15,6 +16,7 @@ from litellm.constants import ( AZURE_STORAGE_MSFT_VERSION, ) from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.azure.common_utils import get_azure_ad_token_from_entra_id from litellm.llms.custom_httpx.http_handler import ( @@ -47,6 +49,20 @@ def adls_safe_file_name(payload_id: str | None) -> str: return f"{name.translate(_ADLS_SAFE_NAME)}.json" +@dataclass(frozen=True, slots=True) +class AzureStorageLogQueueItem: + payload: StandardLoggingPayload + file_name: str + + +def _queued_log(kwargs: Mapping[str, object], payload: StandardLoggingPayload) -> AzureStorageLogQueueItem: + payload_id: Final = payload.get("id") + return AzureStorageLogQueueItem( + payload=payload, + file_name=adls_safe_file_name(payload_id and logged_event_key(kwargs, payload_id)), + ) + + @cache def _cached_credential_chain_token_provider() -> Callable[[], str]: return get_azure_ad_token_provider( @@ -99,7 +115,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() - self.log_queue: list[StandardLoggingPayload] = [] + self.log_queue: list[AzureStorageLogQueueItem] = [] super().__init__(**kwargs, flush_lock=self.flush_lock) except Exception as e: verbose_logger.exception( @@ -115,7 +131,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): def azure_storage_blob_endpoint(self) -> str: return f"https://{self.azure_storage_account_name}.blob.{self.azure_storage_endpoint_suffix}" - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_success_event(self, kwargs, response_obj: object, start_time: datetime, end_time: datetime): """ Async Log success events to Azure Blob Storage @@ -133,12 +149,12 @@ class AzureBlobStorageLogger(CustomBatchLogger): if standard_logging_payload is None: raise ValueError("standard_logging_payload is not set") - self.log_queue.append(standard_logging_payload) + self.log_queue.append(_queued_log(kwargs, standard_logging_payload)) except Exception as e: verbose_logger.exception("AzureBlobStorageLogger Layer Error - %s", e) - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_failure_event(self, kwargs, response_obj: object, start_time: datetime, end_time: datetime): """ Async Log failure events to Azure Blob Storage @@ -156,7 +172,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): if standard_logging_payload is None: raise ValueError("standard_logging_payload is not set") - self.log_queue.append(standard_logging_payload) + self.log_queue.append(_queued_log(kwargs, standard_logging_payload)) except Exception as e: verbose_logger.exception("AzureBlobStorageLogger Layer Error - %s", e) @@ -195,13 +211,15 @@ class AzureBlobStorageLogger(CustomBatchLogger): len(self.log_queue), ) - for payload in self.log_queue: - await self.async_upload_payload_to_azure_blob_storage(payload=payload) + for item in self.log_queue: + await self.async_upload_payload_to_azure_blob_storage(payload=item.payload, file_name=item.file_name) except Exception as e: verbose_logger.exception("AzureBlobStorageLogger Error sending batch API - %s", e) - async def async_upload_payload_to_azure_blob_storage(self, payload: StandardLoggingPayload): + async def async_upload_payload_to_azure_blob_storage( + self, payload: StandardLoggingPayload, file_name: str | None = None + ): """ Uploads the payload to Azure Blob Storage using a 3-step process: 1. Create file resource @@ -210,14 +228,14 @@ class AzureBlobStorageLogger(CustomBatchLogger): """ try: if self.azure_storage_account_key: - await self.upload_to_azure_data_lake_with_azure_account_key(payload=payload) + await self.upload_to_azure_data_lake_with_azure_account_key(payload=payload, file_name=file_name) else: # Get a valid token instead of always requesting a new one await self.set_valid_azure_ad_token() async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) json_payload: Final = safe_dumps(payload) + "\n" # Add newline for each log entry payload_bytes: Final = json_payload.encode("utf-8") - filename: Final = adls_safe_file_name(payload.get("id")) + filename: Final = file_name or adls_safe_file_name(payload.get("id")) base_url = f"{self.azure_storage_dfs_endpoint}/{self.azure_storage_file_system}/{filename}" # Execute the 3-step upload process @@ -377,7 +395,9 @@ class AzureBlobStorageLogger(CustomBatchLogger): self._service_client_timeout = self._clock() + _DEFAULT_TTL_FOR_HTTPX_CLIENTS return self._service_client - async def upload_to_azure_data_lake_with_azure_account_key(self, payload: StandardLoggingPayload): + async def upload_to_azure_data_lake_with_azure_account_key( + self, payload: StandardLoggingPayload, file_name: str | None = None + ): """ Uploads the payload to Azure Data Lake using the Azure SDK @@ -386,11 +406,11 @@ class AzureBlobStorageLogger(CustomBatchLogger): try: today: Final = datetime.now().strftime("%Y-%m-%d") - file_name: Final = adls_safe_file_name(payload.get("id")) + resolved_file_name: Final = file_name or adls_safe_file_name(payload.get("id")) await self._upload_bytes_with_account_key( - file_path=f"{today}/{file_name}", content=safe_dumps(payload).encode("utf-8") + file_path=f"{today}/{resolved_file_name}", content=safe_dumps(payload).encode("utf-8") ) - verbose_logger.debug("Successfully uploaded and wrote to %s/%s", today, file_name) + verbose_logger.debug("Successfully uploaded and wrote to %s/%s", today, resolved_file_name) except Exception as e: verbose_logger.exception("Error occurred: %s", e) diff --git a/litellm/integrations/clickhouse/clickhouse_spend_logger.py b/litellm/integrations/clickhouse/clickhouse_spend_logger.py index 1d01965cd2c..c8453456052 100644 --- a/litellm/integrations/clickhouse/clickhouse_spend_logger.py +++ b/litellm/integrations/clickhouse/clickhouse_spend_logger.py @@ -19,6 +19,7 @@ from litellm._logging import verbose_logger from litellm.integrations.clickhouse.clickhouse_batch_logger import ClickHouseBatchLogger from litellm.integrations.clickhouse.context import is_lens_analysis from litellm.integrations.clickhouse.schema import SPEND_LOGS_TABLE +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.litellm_core_utils.llm_response_utils.get_headers import get_provider_request_id from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.tracing.types import SpendLogPayload, SpendLogRecord @@ -134,7 +135,7 @@ def spend_log_row_from_payload( usage: Mapping[str, Any] = metadata.get("usage_object") or hidden_params.get("usage_object") or MappingProxyType({}) cache_read_tokens, cache_write_tokens = _cache_tokens(usage) trace_id, span_id = _find_traceparent(metadata, kwargs) - request_id = str(payload.get("id") or "") + response_id: Final = str(payload.get("id") or "") redact = litellm.turn_off_message_logging is True track_end_user: Final = not litellm.disable_end_user_cost_tracking completion_start_ms = _to_ms(payload.get("completionStartTime")) @@ -157,8 +158,8 @@ def spend_log_row_from_payload( ) ) return SpendLogRecord( - request_id=request_id, - response_id=strip_cache_hit_suffix(request_id), + request_id=logged_event_key(kwargs, response_id), + response_id=strip_cache_hit_suffix(response_id), provider_request_id=( get_provider_request_id(kwargs.get("response_headers")) or get_provider_request_id(hidden_params.get("additional_headers")) diff --git a/litellm/integrations/deepeval/deepeval.py b/litellm/integrations/deepeval/deepeval.py index 71d149e4014..e38b987d46f 100644 --- a/litellm/integrations/deepeval/deepeval.py +++ b/litellm/integrations/deepeval/deepeval.py @@ -15,6 +15,7 @@ from litellm.integrations.deepeval.utils import ( to_zod_compatible_iso, validate_environment, ) +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key # This file includes the custom callbacks for LiteLLM Proxy @@ -109,8 +110,10 @@ class DeepEvalLogger(CustomLogger): ) else: output = str(standard_logging_object.get("error_string", "")) + logged_id: Final = standard_logging_object.get("id", uuid.uuid4()) + span_id: Final = logged_event_key(kwargs, logged_id) if isinstance(logged_id, str) else logged_id return BaseApiSpan( - uuid=standard_logging_object.get("id", uuid.uuid4()), + uuid=span_id, name=("litellm_success_callback" if is_success else "litellm_failure_callback"), status=(TraceSpanApiStatus.SUCCESS if is_success else TraceSpanApiStatus.ERRORED), type=SpanApiType.LLM, diff --git a/litellm/integrations/dynamodb.py b/litellm/integrations/dynamodb.py index 6569e70329b..8672eb00ce5 100644 --- a/litellm/integrations/dynamodb.py +++ b/litellm/integrations/dynamodb.py @@ -9,6 +9,8 @@ from typing import Final, Protocol import litellm from litellm._uuid import uuid +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key class _DynamoTable(Protocol): @@ -22,11 +24,13 @@ class _DynamoResource(Protocol): class DyanmoDBLogger: # Class variables or attributes - def __init__(self): + def __init__(self, dynamodb: _DynamoResource | None = None): # Instance variables import boto3 - self.dynamodb: Final[_DynamoResource] = boto3.resource("dynamodb", region_name=os.environ["AWS_REGION_NAME"]) + self.dynamodb: Final[_DynamoResource] = dynamodb or boto3.resource( + "dynamodb", region_name=os.environ["AWS_REGION_NAME"] + ) if litellm.dynamodb_table_name is None: raise ValueError( "LiteLLM Error, trying to use DynamoDB but not table name passed. Create a table and set `litellm.dynamodb_table_name=`" @@ -45,23 +49,23 @@ class DyanmoDBLogger: _async_log_event = async_log_event - def log_event(self, kwargs, response_obj, start_time, end_time, print_verbose): + def log_event(self, kwargs: Mapping[str, object], response_obj, start_time, end_time, print_verbose): try: print_verbose(f"DynamoDB Logging - Enters logging function for model {kwargs}") # construct payload to send to DynamoDB # follows the same params as langfuse.py - litellm_params: Final = kwargs.get("litellm_params", {}) - metadata: Final = litellm_params.get("metadata", {}) or {} # if litellm_params['metadata'] == None + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} + metadata: Final = as_str_mapping(litellm_params.get("metadata")) or {} messages: Final = kwargs.get("messages") optional_params: Final = kwargs.get("optional_params", {}) call_type: Final = kwargs.get("call_type", "litellm.completion") usage: Final = response_obj["usage"] - id: Final = response_obj.get("id", str(uuid.uuid4())) + response_id: Final = response_obj.get("id", str(uuid.uuid4())) # Build the initial payload payload: Final[dict[str, object]] = { - "id": id, + "id": logged_event_key(kwargs, response_id) if isinstance(response_id, str) else response_id, "call_type": call_type, "startTime": start_time, "endTime": end_time, diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 99cd8b31b63..7057e719764 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.cloud_storage_security import ( sanitize_cloud_object_component, ) from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus from litellm.types.integrations.gcs_bucket import * from litellm.types.utils import StandardBatchLineItemsPayload, StandardLoggingPayload @@ -296,30 +297,28 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): else: await self._send_individual_logs(items_to_process) - def _get_object_name(self, kwargs: dict, logging_payload: StandardLoggingPayload, response_obj: Any) -> str: + def _get_object_name( + self, kwargs: Mapping[str, object], logging_payload: StandardLoggingPayload, response_obj: Any + ) -> str: """ Get the object name to use for the current payload """ current_date: Final = self._get_object_date_from_datetime(datetime.now(timezone.utc)) + metadata: Final = as_str_mapping((as_str_mapping(kwargs.get("litellm_params")) or {}).get("metadata")) or {} + gcs_log_id: Final = metadata.get("gcs_log_id") + safe_log_id: Final = sanitize_cloud_object_component( + gcs_log_id if isinstance(gcs_log_id, str) else None, fallback="" + ) + if safe_log_id: + return f"{current_date}/custom-{uuid.uuid4().hex}-{safe_log_id}" if logging_payload.get("error_str", None) is not None: - object_name = self._generate_failure_object_name( - request_date_str=current_date, - ) - else: - object_name = self._generate_success_object_name( - request_date_str=current_date, - response_id=response_obj.get("id", ""), - ) - - # used for testing - _litellm_params: Final = kwargs.get("litellm_params", None) or {} - _metadata: Final = _litellm_params.get("metadata", None) or {} - if "gcs_log_id" in _metadata: - safe_log_id: Final = sanitize_cloud_object_component(_metadata.get("gcs_log_id"), fallback="") - if safe_log_id: - object_name = f"{current_date}/custom-{uuid.uuid4().hex}-{safe_log_id}" - - return object_name + return self._generate_failure_object_name(request_date_str=current_date) + raw_response_id: Final[object] = response_obj.get("id", "") + response_id: Final = raw_response_id if isinstance(raw_response_id, str) else "" + return self._generate_success_object_name( + request_date_str=current_date, + response_id=logged_event_key(kwargs, response_id), + ) async def get_request_response_payload( self, diff --git a/litellm/integrations/openmeter.py b/litellm/integrations/openmeter.py index db2fe386dec..918e849a086 100644 --- a/litellm/integrations/openmeter.py +++ b/litellm/integrations/openmeter.py @@ -3,12 +3,15 @@ import json import os +from collections.abc import Mapping from typing import Final import httpx import litellm from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.internal_call_metadata import logged_event_key from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, get_async_httpx_client, @@ -48,7 +51,7 @@ class OpenMeterLogger(CustomLogger): if len(missing_keys) > 0: raise Exception(f"Missing keys={missing_keys} in environment.") - def _common_logic(self, kwargs: dict, response_obj): + def _common_logic(self, kwargs: Mapping[str, object], response_obj): call_id: Final = response_obj.get("id", kwargs.get("litellm_call_id")) dt: Final = get_utc_datetime().isoformat() cost: Final = kwargs.get("response_cost", None) @@ -74,8 +77,8 @@ class OpenMeterLogger(CustomLogger): # If no user provided directly, try to get it from token user_id if user_param is None: # Check if user_id is available from the API key metadata - litellm_params: Final = kwargs.get("litellm_params", {}) - metadata: Final = litellm_params.get("metadata", {}) + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} + metadata: Final = as_str_mapping(litellm_params.get("metadata")) or {} user_api_key_user_id: Final = metadata.get("user_api_key_user_id", None) if user_api_key_user_id is not None: @@ -89,7 +92,7 @@ class OpenMeterLogger(CustomLogger): return { "specversion": "1.0", "type": os.getenv("OPENMETER_EVENT_TYPE", "litellm_tokens"), - "id": call_id, + "id": logged_event_key(kwargs, call_id) if isinstance(call_id, str) else call_id, "time": dt, "subject": subject, "source": "litellm-proxy", diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index 453d4358b35..138a01c6524 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -26,7 +26,12 @@ from litellm.constants import ( MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, NON_INFERENCE_CALL_TYPES, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.initialize_dynamic_callback_params import initialize_standard_callback_dynamic_params +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) from litellm.types.utils import ( BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN, @@ -35,6 +40,8 @@ from litellm.types.utils import ( BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"}) +BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX: Final = "_background_cost" + ITPM_RESERVED_KEY: Final = "_litellm_itpm_reserved" OTPM_RESERVED_KEY: Final = "_litellm_otpm_reserved" ITPM_CACHE_KEY: Final = "_litellm_itpm_cache_key" @@ -151,6 +158,26 @@ def is_unbilled_non_inference_call_from_params( return is_unbilled_non_inference_call(call_type, metadata, response) +def settles_background_response(call_type: str, metadata: Mapping[str, object], kwargs: Mapping[str, object]) -> bool: + return kwargs.get(SETTLES_BACKGROUND_RESPONSE_KWARG) is BackgroundResponseSettlement.SETTLES or ( + call_type in NON_INFERENCE_CALL_TYPES + and metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + ) + + +def logged_event_key(kwargs: Mapping[str, object], event_id: str) -> str: + call_type: Final = kwargs.get("call_type") + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} + metadata: Final = ( + as_str_mapping(litellm_params.get("litellm_metadata")) or as_str_mapping(litellm_params.get("metadata")) or {} + ) + if not event_id or not settles_background_response( + call_type if isinstance(call_type, str) else "", metadata, kwargs + ): + return event_id + return f"{event_id}{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}" + + def sanitize_user_api_key_auth(auth: object) -> object: """Copy of the auth object with its budget reservation removed; the cost callback falls back to reading the reservation from inside the auth object.""" diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5f909f639c4..37657a94f16 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -872,6 +872,7 @@ class Logging(LiteLLMLoggingBaseClass): # The model the proxy shows the client on streamed chunks. The logged streamed response carries it # once that response is priced, the same way a non-streamed response is logged self.client_facing_stream_model: str | None = None + self.client_facing_response_id: str | None = None self.zero_cost_warned: bool = False self.llm_caching_handler: LLMCachingHandler | None = None @@ -3238,6 +3239,23 @@ class Logging(LiteLLMLoggingBaseClass): return response return response.model_copy(update={"id": response_id}) + def _with_client_facing_response_id( + self, + response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse, + ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse: + response_id: Final = self.client_facing_response_id + if response_id is None or getattr(response, "id", None) == response_id: + return response + return response.model_copy(update={"id": response_id}) + + def _client_facing_streaming_response( + self, + response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse, + ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse: + return self._with_client_facing_response_id( + self._with_continued_response_id(self._with_client_facing_stream_model(response)) + ) + def _success_handler_helper_fn( self, result: object = None, @@ -3641,9 +3659,7 @@ class Logging(LiteLLMLoggingBaseClass): ) self._add_continued_partial_spend(complete_streaming_response) self._merge_hidden_params_from_response_into_metadata(complete_streaming_response) - logged_streaming_response: Final = self._with_continued_response_id( - self._with_client_facing_stream_model(complete_streaming_response) - ) + logged_streaming_response: Final = self._client_facing_streaming_response(complete_streaming_response) self.model_call_details["complete_streaming_response"] = logged_streaming_response self.capture_usage_before_logging(logged_streaming_response, start_time, end_time) ## STANDARDIZED LOGGING PAYLOAD @@ -4240,9 +4256,7 @@ class Logging(LiteLLMLoggingBaseClass): await self._prepare_baseline_cache_estimate(complete_streaming_response) - logged_streaming_response: Final = self._with_continued_response_id( - self._with_client_facing_stream_model(complete_streaming_response) - ) + logged_streaming_response: Final = self._client_facing_streaming_response(complete_streaming_response) self.model_call_details["async_complete_streaming_response"] = logged_streaming_response ## STANDARDIZED LOGGING PAYLOAD diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index 771d68a0ea0..61ab1eb4771 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -7,22 +7,25 @@ instead of writing immediately on each request. from collections.abc import AsyncGenerator, Callable, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, TypeVar from fastapi import HTTPException from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import LitellmUserRoles from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( BaseLiteLLMOpenAIResponseObject, ResponsesAPIResponse, ) +from litellm.types.responses.main import DeleteResponseResult from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, CallTypesLiteral, LLMResponseTypes, SpecialEnums if TYPE_CHECKING: @@ -67,6 +70,8 @@ def _proxy_signing_key() -> str | None: _RESPONSE_PAYLOAD_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_ResponseObjectT = TypeVar("_ResponseObjectT", bound=BaseLiteLLMOpenAIResponseObject) +_AddressedResultT = TypeVar("_AddressedResultT", bound=ResponsesAPIResponse | DeleteResponseResult) def _response_payload(response_obj: object) -> Mapping[str, object] | None: @@ -77,9 +82,9 @@ def _response_payload(response_obj: object) -> Mapping[str, object] | None: def _rewrite_advertised_id( - event: BaseLiteLLMOpenAIResponseObject, + event: _ResponseObjectT, rewrite: Callable[[str], str], -) -> BaseLiteLLMOpenAIResponseObject: +) -> _ResponseObjectT: event_id: Final = getattr(event, "id", None) if isinstance(event_id, str) and event_id.startswith("resp_"): setattr(event, "id", rewrite(event_id)) @@ -116,6 +121,35 @@ def _streams_response_ids(request_route: str | None) -> bool: return canonical in _RESPONSES_API_ROUTES or (parent in _RESPONSES_API_ROUTES and response_id != "") +_NO_HELD_IDS: Final[Mapping[str, str]] = MappingProxyType({}) + + +def _upstream_response_id(response_id: str) -> str: + return ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(response_id) + + +def _names_the_same_response(response_id: str | None, original_response_id: str) -> bool: + return response_id is not None and _upstream_response_id(response_id) == _upstream_response_id(original_response_id) + + +def _held_id_naming(response_id: str | None, held_ids: Mapping[str, str]) -> str | None: + return next( + ( + held_id + for original_response_id, held_id in held_ids.items() + if _names_the_same_response(response_id, original_response_id) + ), + None, + ) + + +def _id_sent_back(response: ResponsesAPIResponse | DeleteResponseResult, held_ids: Mapping[str, str]) -> str | None: + held_id: Final = _held_id_naming(response.id, held_ids) + if held_id is not None: + return held_id + return response.id if isinstance(response, DeleteResponseResult) else None + + class ResponsesIDSecurity(CustomLogger): def __init__( self, @@ -275,10 +309,18 @@ class ResponsesIDSecurity(CustomLogger): def _encrypt_response_id( self, - response: BaseLiteLLMOpenAIResponseObject, + response: _ResponseObjectT, user_api_key_dict: "UserAPIKeyAuth", request_cache: dict[str, str] | None = None, - ) -> BaseLiteLLMOpenAIResponseObject: + ) -> _ResponseObjectT: + encrypt: Final = self._response_id_encrypter(user_api_key_dict, request_cache) + return response if encrypt is None else _rewrite_advertised_id(response, encrypt) + + def _response_id_encrypter( + self, + user_api_key_dict: "UserAPIKeyAuth", + request_cache: dict[str, str] | None, + ) -> Callable[[str], str] | None: signing_key: Final = self._get_signing_key() if signing_key is None: verbose_proxy_logger.debug( @@ -287,7 +329,7 @@ class ResponsesIDSecurity(CustomLogger): "Skipping response ID encryption. " "See: https://docs.litellm.ai/docs/proxy/prod#5-set-litellm-salt-key" ) - return response + return None def encrypt(original_id: str) -> str: cached: Final = request_cache.get(original_id) if request_cache is not None else None @@ -304,47 +346,73 @@ class ResponsesIDSecurity(CustomLogger): request_cache[original_id] = encrypted_id return encrypted_id - return _rewrite_advertised_id(response, encrypt) + return encrypt + + def _streamed_response_id_advertiser( + self, + user_api_key_dict: "UserAPIKeyAuth", + request_data: Mapping[str, object], + ) -> Callable[[str], str] | None: + if not _streams_response_ids(user_api_key_dict.request_route): + return None + if self._general_settings_reader().get("disable_responses_id_security", False): + return None + encrypt: Final = self._response_id_encrypter(user_api_key_dict, request_cache={}) + if encrypt is None: + return None + logging_obj: Final = request_data.get("litellm_logging_obj") + held_ids: Final = self._ids_the_caller_holds(request_data.get(ADDRESSED_RESPONSE_ID_FIELD)) + + def advertise(original_id: str) -> str: + logged_id: Final = encrypt(original_id) + if isinstance(logging_obj, LiteLLMLoggingObj): + logging_obj.client_facing_response_id = logged_id + return _held_id_naming(original_id, held_ids) or logged_id + + return advertise + + def _ids_the_caller_holds(self, addressed_id: object) -> Mapping[str, str]: + if not isinstance(addressed_id, str) or not self._is_encrypted_response_id(addressed_id): + return _NO_HELD_IDS + original_response_id, _, _ = self._decrypt_response_id(addressed_id) + return MappingProxyType({original_response_id: addressed_id}) + + def _advertised_response( + self, + response: _AddressedResultT, + user_api_key_dict: "UserAPIKeyAuth", + request_data: Mapping[str, object], + ) -> _AddressedResultT: + sent_back_id: Final = _id_sent_back( + response, self._ids_the_caller_holds(request_data.get(ADDRESSED_RESPONSE_ID_FIELD)) + ) + logged_response: Final = self._encrypt_response_id(response, user_api_key_dict, request_cache=None) + return logged_response if sent_back_id is None else logged_response.model_copy(update={"id": sent_back_id}) async def async_post_call_success_hook( self, - data: dict, + data: dict[str, object], user_api_key_dict: "UserAPIKeyAuth", response: LLMResponseTypes, ) -> LLMResponseTypes: - """ - Queue response IDs for batch processing instead of writing directly to DB. - - This method adds response IDs to an in-memory queue, which are then - batch-processed by the DBSpendUpdateWriter during regular database update cycles. - """ general_settings: Final = self._general_settings_reader() if general_settings.get("disable_responses_id_security", False): return response - if isinstance(response, ResponsesAPIResponse): - response = cast( - ResponsesAPIResponse, - self._encrypt_response_id(response, user_api_key_dict, request_cache=None), - ) + if isinstance(response, ResponsesAPIResponse | DeleteResponseResult): + return self._advertised_response(response, user_api_key_dict, data) return response async def async_post_call_streaming_iterator_hook( - self, user_api_key_dict: "UserAPIKeyAuth", response: Any, request_data: dict + self, user_api_key_dict: "UserAPIKeyAuth", response: Any, request_data: dict[str, object] ) -> AsyncGenerator[BaseLiteLLMOpenAIResponseObject, None]: - general_settings: Final = self._general_settings_reader() - - # Create a request-scoped cache for consistent encryption across streaming chunks. - request_encryption_cache: Final[dict[str, str]] = {} - + advertise: Final = self._streamed_response_id_advertiser(user_api_key_dict, request_data) async for chunk in response: - if ( - isinstance(chunk, BaseLiteLLMOpenAIResponseObject) - and _streams_response_ids(user_api_key_dict.request_route) - and not general_settings.get("disable_responses_id_security", False) - ): - chunk = self._encrypt_response_id(chunk, user_api_key_dict, request_encryption_cache) - yield chunk + yield ( + _rewrite_advertised_id(chunk, advertise) + if advertise is not None and isinstance(chunk, BaseLiteLLMOpenAIResponseObject) + else chunk + ) def decrypted_response_id(response_id: str) -> str: diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index ba5f873c4e9..8b0ff718fec 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -20,7 +20,6 @@ from litellm.constants import ( CALLER_SESSION_ID_METADATA_KEY, CLI_SESSION_KEY_PREFIX, EMPTY_MAPPING, - INTERNAL_CALL_ORIGIN_METADATA_KEY, KUBERNETES_POD_ROUTING_KEY, LITELLM_PROXY_MASTER_KEY_ALIAS, LITELLM_TRUNCATED_PAYLOAD_FIELD, @@ -30,7 +29,6 @@ from litellm.constants import ( MAX_SPEND_LOG_MODEL_NAME_LENGTH, MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB, MCP_SPEND_LOG_MODEL_PREFIX, - NON_INFERENCE_CALL_TYPES, REDACTED_BY_LITELM_STRING, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL, @@ -53,7 +51,11 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider from litellm.litellm_core_utils.hidden_params import served_from_cache -from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call +from litellm.litellm_core_utils.internal_call_metadata import ( + BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX, + is_unbilled_non_inference_call, + settles_background_response, +) from litellm.litellm_core_utils.litellm_logging import ( StandardLoggingPayloadSetup, UsageObservation, @@ -79,13 +81,8 @@ from litellm.responses.litellm_completion_transformation.compaction_items import summary_only_compaction_content, with_compaction_summary, ) -from litellm.types.passthrough_endpoints.pass_through_endpoints import ( - SETTLES_BACKGROUND_RESPONSE_KWARG, - BackgroundResponseSettlement, -) from litellm.types.router import DeploymentTypedDict, LiteLLM_Params, SilentExperimentFlag from litellm.types.utils import ( - BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, PROMPT_CARRYING_GUARDRAIL_FIELDS, AzureSpillover, CallTypes, @@ -347,24 +344,21 @@ def _get_spend_logs_metadata( BATCH_COST_REQUEST_ID_SUFFIX: Final = "_batch_cost" -BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX: Final = "_background_cost" -def _cost_row_request_id_suffix( - call_type: str, metadata: Mapping[str, object], settles_background_response: bool -) -> str: +def _cost_row_request_id_suffix(call_type: str, metadata: Mapping[str, object], kwargs: Mapping[str, object]) -> str: if call_type == CallTypes.aretrieve_batch.value: return BATCH_COST_REQUEST_ID_SUFFIX - if settles_background_response or ( - call_type in NON_INFERENCE_CALL_TYPES - and metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN - ): + if settles_background_response(call_type, metadata, kwargs): return BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX return "" def get_spend_logs_id( - call_type: str, response_obj: dict, kwargs: Mapping[str, object], metadata: Mapping[str, object] = EMPTY_MAPPING + call_type: str, + response_obj: Mapping[str, object], + kwargs: Mapping[str, object], + metadata: Mapping[str, object] = EMPTY_MAPPING, ) -> str | None: standard_logging_payload = kwargs.get("standard_logging_object") candidate_ids: Final = ( @@ -377,10 +371,7 @@ def get_spend_logs_id( ) if resolved_id is None: return None - settles_background_response: Final = ( - kwargs.get(SETTLES_BACKGROUND_RESPONSE_KWARG) is BackgroundResponseSettlement.SETTLES - ) - return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata, settles_background_response)}" + return f"{resolved_id}{_cost_row_request_id_suffix(call_type, metadata, kwargs)}" _MISSING_ATTRIBUTE: Final = object() diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 7cf163e7ce8..7e6ee84e2b2 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -1,14 +1,16 @@ import asyncio import json -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from itertools import chain +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER from litellm.litellm_core_utils.cold_storage_handler import ColdStorageHandler +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.responses.litellm_completion_transformation.client_tool_search import ( bridge_client_tool_search, client_tool_search_loaded_tools, @@ -69,6 +71,90 @@ def _is_a_conversation_turn(spend_log: "SpendLogsPayload") -> bool: return spend_log.get("call_type") not in _MCP_TOOL_CALL_TYPES and not _is_a_failed_turn(spend_log) +_RESPONSE_READ_CALL_TYPES: Final = frozenset(("aget_responses", "acancel_responses")) + + +def _output_item_id(item: object) -> str | None: + match item: + case {"id": str() as item_id}: + return item_id + case _: + return None + + +def _is_a_response_read(spend_log: "SpendLogsPayload") -> bool: + return "call_type" in spend_log and spend_log["call_type"] in _RESPONSE_READ_CALL_TYPES + + +def _stored_response(spend_log: "SpendLogsPayload") -> object: + return spend_log.get("response") + + +def _output_items(spend_log: "SpendLogsPayload") -> Iterator[tuple[str | None, object]]: + output: Final = _stored_responses_api_output(_stored_response(spend_log)) or [] + return zip(map(_output_item_id, output), output) + + +def _identified_output_items(spend_log: "SpendLogsPayload") -> Iterator[tuple[str, object]]: + return ((item_id, item) for item_id, item in _output_items(spend_log) if item_id is not None) + + +def _output_item_ids(spend_log: "SpendLogsPayload") -> frozenset[str]: + return frozenset(item_id for item_id, _ in _identified_output_items(spend_log)) + + +def _first_index_holding_each_item(indexed_turns: Iterable[tuple[int, "SpendLogsPayload"]]) -> Mapping[str, int]: + return dict( + chain.from_iterable( + ((item_id, index) for item_id in _output_item_ids(turn)) for index, turn in reversed(tuple(indexed_turns)) + ) + ) + + +def _turn_replaying_each_item(turns: Sequence["SpendLogsPayload"]) -> Mapping[str, int]: + indexed_turns: Final = tuple(enumerate(turns)) + first_read_holding: Final = _first_index_holding_each_item( + (index, turn) for index, turn in indexed_turns if _is_a_response_read(turn) + ) + first_other_turn_holding: Final = _first_index_holding_each_item( + (index, turn) for index, turn in indexed_turns if not _is_a_response_read(turn) + ) + return MappingProxyType({**first_read_holding, **first_other_turn_holding}) + + +def _latest_read_copy_of_each_item(turns: Sequence["SpendLogsPayload"]) -> Mapping[str, object]: + return MappingProxyType( + dict(chain.from_iterable(_identified_output_items(turn) for turn in turns if _is_a_response_read(turn))) + ) + + +def _with_only_new_output( + index: int, turn: "SpendLogsPayload", turn_replaying: Mapping[str, int], latest_copy: Mapping[str, object] +) -> "SpendLogsPayload | None": + if not _is_a_response_read(turn): + return turn + stored_response: Final = as_str_mapping(_stored_response(turn)) + new_output: Final = [ + item if item_id is None else latest_copy[item_id] + for item_id, item in _output_items(turn) + if item_id is None or turn_replaying.get(item_id) == index + ] + if stored_response is None or not new_output: + return None + replayed_turn: Final[SpendLogsPayload] = {**turn, "response": {**stored_response, "output": new_output}} + return replayed_turn + + +def _conversation_turns(spend_logs: Sequence["SpendLogsPayload"]) -> tuple["SpendLogsPayload", ...]: + turns: Final = tuple(filter(_is_a_conversation_turn, spend_logs)) + turn_replaying: Final = _turn_replaying_each_item(turns) + latest_copy: Final = _latest_read_copy_of_each_item(turns) + replayed: Final = ( + _with_only_new_output(index, turn, turn_replaying, latest_copy) for index, turn in enumerate(turns) + ) + return tuple(turn for turn in replayed if turn is not None) + + _MCP_AUTO_EXECUTE_OUTPUT_TYPES: Final = frozenset(("mcp_tools_fetched", "tool_execution_results")) @@ -101,7 +187,7 @@ class StoredTurns: def _compaction_point( index: int, spend_log: "SpendLogsPayload", proxy_server_request: Mapping[str, object] | None ) -> _ReplayStart | None: - if response_holds_compaction_block(spend_log.get("response")): + if response_holds_compaction_block(_stored_response(spend_log)): return _ReplayStart(index=index, at_output=True) if proxy_server_request is not None and holds_compaction_summary(proxy_server_request.get("input")): return _ReplayStart(index=index, at_output=False) @@ -271,12 +357,10 @@ class ResponsesSessionHandler: previous_response_id: str, client_response_ids: tuple[str, ...] = (), ) -> StoredTurns: - spend_logs: Final = tuple( - spend_log - for spend_log in await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + spend_logs: Final = _conversation_turns( + await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( previous_response_id, client_response_ids ) - if _is_a_conversation_turn(spend_log) ) return StoredTurns( spend_logs=spend_logs, @@ -378,7 +462,7 @@ class ResponsesSessionHandler: ############################################################ # Add Output messages for this Spend Log ############################################################ - _response_output: Final[object] = spend_log.get("response", "{}") + _response_output: Final = _stored_response(spend_log) if (responses_api_output := _stored_responses_api_output(_response_output)) is not None: chat_completion_message_history.extend( LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( diff --git a/tests/e2e/llm_translation/test_responses_retrieve_e2e.py b/tests/e2e/llm_translation/test_responses_retrieve_e2e.py index 4a08259c028..0e954af0952 100644 --- a/tests/e2e/llm_translation/test_responses_retrieve_e2e.py +++ b/tests/e2e/llm_translation/test_responses_retrieve_e2e.py @@ -69,9 +69,6 @@ def _retrieve_response(proxy: ProxyClient, key: str, response_id: str) -> Respon class TestResponsesRetrieve: - @pytest.mark.skip( - reason="stage red: product gap (LIT-5446), retrieve returns a different id than the stored response (non-idempotent response-id re-encryption)" - ) @pytest.mark.covers("llm.responses.openai.basic.nonstream.works") @meta( Subject( diff --git a/tests/integration/compatibility/test_openai_sdk_client_paths.py b/tests/integration/compatibility/test_openai_sdk_client_paths.py index f8d7449038a..d6b6d1d460d 100644 --- a/tests/integration/compatibility/test_openai_sdk_client_paths.py +++ b/tests/integration/compatibility/test_openai_sdk_client_paths.py @@ -630,19 +630,6 @@ def _spend_row(identity: str) -> Mapping[str, JsonValue]: return rows[0] -def _spend_row_keyed_by_any(identities: frozenset[str]) -> Mapping[str, JsonValue]: - keys: Final = tuple(sorted(identities)) - placeholders: Final = ", ".join("%s" for _ in keys) - rows: Final = eventually( - lambda: read_rows( - f'SELECT request_id, call_type, model FROM "LiteLLM_SpendLogs" WHERE request_id IN ({placeholders})', keys - ), - lambda rows: len(rows) == 1, - seconds=70, - ) - return rows[0] - - def _spend_request_ids(identities: tuple[str, str, str]) -> frozenset[str]: rows: Final = eventually( lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id IN (%s, %s, %s)', identities), @@ -815,7 +802,7 @@ def test_r07_responses_stream_through_the_async_openai_sdk(gateway: Gateway) -> served: Final = _served(wire) assert _routes(served) == (("POST", "/v1/responses"),), served assert _body(served[0])["stream"] is True, served[0].body - assert _spend_row_keyed_by_any(response_identities(last.response.id))["model"] == DEPLOYMENT_MODEL + assert _spend_row(last.response.id)["model"] == DEPLOYMENT_MODEL def test_r08_completions_non_stream_through_the_sync_openai_sdk(gateway: Gateway) -> None: diff --git a/tests/integration/observability/test_langfuse_otel_metadata_chaos.py b/tests/integration/observability/test_langfuse_otel_metadata_chaos.py index bb386f6d41e..15bc076b8ba 100644 --- a/tests/integration/observability/test_langfuse_otel_metadata_chaos.py +++ b/tests/integration/observability/test_langfuse_otel_metadata_chaos.py @@ -107,9 +107,11 @@ def test_langfuse_otel_sink_outage_mid_burst_lands_every_response_id_once(gatewa ) payloads: Final = tuple(json.loads(line) for line in events) found: Final = next( - payload.get("response_id") or (payload.get("message") or {}).get("id") or payload.get("id") + (payload.get("response") or {}).get("id") + or (payload.get("message") or {}).get("id") + or payload.get("id") for payload in payloads - if payload.get("response_id") or payload.get("message") or payload.get("id") + if payload.get("response") or payload.get("message") or payload.get("id") ) return str(found) response: Final = rig.candidate.client.post(path, json=body, headers={"Authorization": f"Bearer {key}"}) diff --git a/tests/integration/observability/test_langfuse_otel_metadata_clients.py b/tests/integration/observability/test_langfuse_otel_metadata_clients.py index c2e0573b48a..fd2b9758449 100644 --- a/tests/integration/observability/test_langfuse_otel_metadata_clients.py +++ b/tests/integration/observability/test_langfuse_otel_metadata_clients.py @@ -162,8 +162,13 @@ def test_langfuse_otel_metadata_reaches_langfuse_for_responses_sse_stream(gatewa for line in stream.iter_lines() if line.startswith("data: ") and line != "data: [DONE]" ) - response_id: Final = next( - event["response_id"] for event in events if event.get("type") == "response.output_text.delta" + advertised: Final = frozenset( + event["response"]["id"] + for event in events + if event.get("type") in ("response.created", "response.completed") ) + assert len(advertised) == 1, events + response_id: Final = next(iter(advertised)) attrs: Final = exported_llm_span(rig.collector, response_id) + assert attrs["gen_ai.response.id"] == response_id, attrs assert_identity(attrs, chat_identity_fields(identity)) diff --git a/tests/integration/observability/test_otel_conversation_id.py b/tests/integration/observability/test_otel_conversation_id.py index d7f4de2961b..cd095991e53 100644 --- a/tests/integration/observability/test_otel_conversation_id.py +++ b/tests/integration/observability/test_otel_conversation_id.py @@ -1,5 +1,4 @@ import asyncio -import base64 import json import re import signal @@ -118,17 +117,6 @@ def _responses_reply(identity: str, stream: bool) -> Reply: ) -def _decoded_responses_id(identity: str) -> str: - try: - return base64.b64decode(identity.removeprefix("resp_").encode()).decode() - except (ValueError, UnicodeDecodeError): - return identity - - -def _canonical_id(identity: str) -> str: - return _decoded_responses_id(identity).rpartition("response_id:")[2] - - def _sse_events(text: str) -> tuple[dict[str, JsonValue], ...]: return tuple( json.loads(line[6:]) for line in text.splitlines() if line.startswith("data: ") and line != "data: [DONE]" @@ -170,8 +158,7 @@ class Collector: return tuple( attributes for attributes in self.attributes() - if isinstance(logged := attributes.get("gen_ai.response.id", {}).get("stringValue"), str) - and _canonical_id(logged) == _canonical_id(response_id) + if attributes.get("gen_ai.response.id", {}).get("stringValue") == response_id ) def conversation_ids(self, response_id: str) -> tuple[str | None, ...]: @@ -183,10 +170,6 @@ class Collector: def single_span(self, response_id: str) -> str | None: return eventually(lambda: self.conversation_ids(response_id), lambda values: len(values) == 1, seconds=30)[0] - def logged_id(self, response_id: str) -> str: - spans: Final = eventually(lambda: self.spans(response_id), lambda values: len(values) == 1, seconds=30) - return str(spans[0]["gen_ai.response.id"]["stringValue"]) - @pytest.fixture(scope="session") def collector() -> Iterator[Collector]: @@ -437,12 +420,13 @@ def test_responses_stream_raw_metadata_session_id_lands_as_conversation_id(rig: completed: Final = tuple(event for event in events if event["type"] == "response.completed") assert len(completed) == 1, events assert completed[0]["response"]["output"][0]["id"] == f"msg_resp_{marker}", completed - assert str(completed[0]["response"]["id"]).startswith("resp_"), completed + client_id: Final = str(completed[0]["response"]["id"]) + assert client_id.startswith("resp_"), completed assert rig.upstream_bodies(marker) == ( {"model": "gpt-4o-mini", "input": marker, "metadata": {"session_id": session}, "stream": True}, ) - assert rig.sink.single_span(f"resp_{marker}") == session - assert rig.spend_session(rig.sink.logged_id(f"resp_{marker}")) == session + assert rig.sink.single_span(client_id) == session + assert rig.spend_session(client_id) == session @pytest.mark.covers("other.observability.otel.conversation_id_from_metadata_chat_raw") @@ -689,15 +673,11 @@ def _burst(rig: Rig, count: int, session_for: Mapping[int, str]) -> tuple[tuple[ return tuple(pool.map(one, range(count))) -def _is_encrypted_responses_id(identity: str) -> bool: - return identity.startswith("resp_") and _decoded_responses_id(identity) == identity - - def _landed(rig: Rig, expected: Mapping[str, str]) -> dict[str, tuple[str, ...]]: spans: Final = rig.sink.attributes() return { session: tuple( - _canonical_id(str(attributes["gen_ai.response.id"]["stringValue"])) + str(attributes["gen_ai.response.id"]["stringValue"]) for attributes in spans if attributes.get(CONVERSATION, {}).get("stringValue") == session and "gen_ai.response.id" in attributes ) @@ -707,21 +687,13 @@ def _landed(rig: Rig, expected: Mapping[str, str]) -> dict[str, tuple[str, ...]] def _assert_exactly_once(rig: Rig, expected: Mapping[str, str], landed: Mapping[str, tuple[str, ...]]) -> None: spend: Final = eventually( - lambda: { - session: tuple(_canonical_id(identity) for identity in rig.spend_request_ids(session)) - for session in expected.values() - }, + lambda: {session: rig.spend_request_ids(session) for session in expected.values()}, lambda rows: all(len(values) >= 1 for values in rows.values()), seconds=70, ) assert landed == spend, (landed, spend) assert all(len(values) == 1 for values in landed.values()), landed - caller_visible: Final = { - session: (_canonical_id(identity),) - for identity, session in expected.items() - if not _is_encrypted_responses_id(identity) - } - assert {session: landed[session] for session in caller_visible} == caller_visible, landed + assert landed == {session: (identity,) for identity, session in expected.items()}, landed @pytest.mark.covers("other.observability.otel.conversation_id_sink_outage_recovers_exactly_once") diff --git a/tests/integration/observability/test_signoz_delivery.py b/tests/integration/observability/test_signoz_delivery.py index e25126e4fa8..93539d7641b 100644 --- a/tests/integration/observability/test_signoz_delivery.py +++ b/tests/integration/observability/test_signoz_delivery.py @@ -1,5 +1,4 @@ import asyncio -import base64 import json import re import signal @@ -9,7 +8,6 @@ from collections import deque from collections.abc import Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass -from functools import partial from pathlib import Path from typing import Final @@ -137,17 +135,6 @@ def _upstream(request: Request) -> Reply: return _chat_reply(f"chatcmpl-{marker}", stream) -def _decoded_responses_id(identity: str) -> str: - try: - return base64.b64decode(identity.removeprefix("resp_").encode()).decode() - except (ValueError, UnicodeDecodeError): - return identity - - -def _canonical_id(identity: str) -> str: - return _decoded_responses_id(identity).rpartition("response_id:")[2] - - def _sse_events(text: str) -> tuple[dict[str, JsonValue], ...]: return tuple( object_value(JSON.validate_json(line[6:])) @@ -232,8 +219,7 @@ class Collector: return tuple( span for span in self.spans() - if RESPONSE_ID in span.attributes - and _canonical_id(span.attributes[RESPONSE_ID]) == _canonical_id(response_id) + if RESPONSE_ID in span.attributes and span.attributes[RESPONSE_ID] == response_id ) def single_span(self, response_id: str, *, elsewhere: "Collector | None" = None) -> Span: @@ -249,11 +235,8 @@ class Collector: def landed(self, response_ids: Sequence[str]) -> dict[str, int]: spans: Final = self.spans() return { - _canonical_id(identity): sum( - 1 - for span in spans - if RESPONSE_ID in span.attributes - and _canonical_id(span.attributes[RESPONSE_ID]) == _canonical_id(identity) + identity: sum( + 1 for span in spans if RESPONSE_ID in span.attributes and span.attributes[RESPONSE_ID] == identity ) for identity in response_ids } @@ -448,7 +431,7 @@ def _assert_operator_span(rig: Rig, response_id: str, marker: str) -> Span: span: Final = rig.sink.single_span(response_id) assert span.target == "/v1/traces", span assert span.ingestion_key == OPERATOR_KEY, span - assert rig.tenant_sink.landed((response_id,)) == {_canonical_id(response_id): 0} + assert rig.tenant_sink.landed((response_id,)) == {response_id: 0} bodies: Final = rig.upstream_bodies(marker) assert len(bodies) == 1, bodies assert "signoz" not in json.dumps(bodies[0]).replace(marker, ""), bodies[0] @@ -522,7 +505,7 @@ def test_responses_stream_raw_httpx_span_lands_once_after_the_stream_is_consumed marker: Final = _marker() response: Final = rig.proxy.request("POST", "/v1/responses", {"model": rig.model, "input": marker, "stream": True}) assert response.status_code == 200, response.text - _assert_operator_span(rig, _responses_id(response, marker), marker) + _assert_operator_span(rig, _responses_id(response), marker) def test_v2_flag_on_still_delivers_the_operator_span_with_the_ingestion_key(v2_rig: Rig) -> None: @@ -821,7 +804,7 @@ def _chat_id(response: httpx.Response) -> str: return next(iter(identities)) -def _responses_id(response: httpx.Response, marker: str) -> str: +def _responses_id(response: httpx.Response) -> str: if not _is_event_stream(response): return _body_id(response) completed: Final = tuple( @@ -830,7 +813,7 @@ def _responses_id(response: httpx.Response, marker: str) -> str: if event.get("type") == "response.completed" ) assert len(completed) == 1 and completed[0].startswith("resp_"), response.text - return f"resp_{marker}" + return completed[0] def _message_id(response: httpx.Response) -> str: @@ -852,7 +835,7 @@ def _burst(rig: Rig, count: int) -> tuple[tuple[int, str, str | None], ...]: stream: Final = index % 2 == 0 path, body, identity_of = ( ("/v1/chat/completions", {"model": rig.model, "messages": [{"role": "user", "content": marker}]}, _chat_id), - ("/v1/responses", {"model": rig.model, "input": marker}, partial(_responses_id, marker=marker)), + ("/v1/responses", {"model": rig.model, "input": marker}, _responses_id), ( "/v1/messages", {"model": rig.model, "max_tokens": 16, "messages": [{"role": "user", "content": marker}]}, @@ -874,8 +857,8 @@ def _assert_exactly_once(rig: Rig, identities: Sequence[str]) -> None: landed: Final = eventually( lambda: rig.sink.landed(identities), lambda seen: all(count >= 1 for count in seen.values()), seconds=80 ) - assert landed == {_canonical_id(identity): 1 for identity in identities}, landed - charged: Final = frozenset(identity for identity in identities if not identity.startswith("resp_")) + assert landed == {identity: 1 for identity in identities}, landed + charged: Final = frozenset(identities) spend: Final = eventually( lambda: read_rows( 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s::text[])', diff --git a/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py b/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py index 3d6eb7fcaff..0788b9ac85a 100644 --- a/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py +++ b/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py @@ -62,12 +62,12 @@ def _body(request: Request) -> dict[str, JsonValue]: # ciphertext the caller received, because the spend row id is read from response_obj["id"] before the # ResponsesIDSecurity hook rewrites it in place; the row is looked up under both ids until that ordering is fixed on # main -def _spend_row(client_id: str, issued_id: str) -> dict[str, JsonValue]: +def _spend_row(*request_ids: str) -> dict[str, JsonValue]: rows: Final = eventually( lambda: read_rows( 'SELECT model_group, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" ' "WHERE request_id = ANY(%s)", - ([client_id, issued_id],), # pyright: ignore[reportArgumentType] # psycopg adapts the list to a text array + (list(request_ids),), # pyright: ignore[reportArgumentType] # psycopg adapts the list to a text array ), lambda found: len(found) == 1, seconds=70, @@ -120,4 +120,4 @@ async def test_async_openai_sdk_responses_stream_is_served_by_the_native_respons assert issued.upstream == f"resp_upstream_{marker}", completed.response.id request: Final = _native_request(wire) assert _body(request) == {"model": GPT, "input": _prompt(marker), "stream": True}, request.body - assert _spend_row(completed.response.id, issued.issued) == _success_row(model) + assert _spend_row(completed.response.id) == _success_row(model) diff --git a/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py b/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py index 89ad226761e..5aede070896 100644 --- a/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_output_limit_wire.py @@ -516,7 +516,7 @@ def test_responses_max_output_tokens_is_forwarded_natively_on_both_legs(gateway: for attempt in attempts: assert attempt.body["max_output_tokens"] == 1000 and _limits(attempt.body) == (None, None), attempt.body _assert_logged(issued, model, "success") - _assert_logged(_managed_response_id(streamed_issued), model, "success") + _assert_logged(streamed_issued, model, "success") def test_region_prefixed_deployment_on_a_v1_base_sends_the_bare_model_with_max_completion_tokens( diff --git a/tests/integration/providers/test_bedrock_model_discovery_wire.py b/tests/integration/providers/test_bedrock_model_discovery_wire.py index 3bf4583e736..2871e51de5e 100644 --- a/tests/integration/providers/test_bedrock_model_discovery_wire.py +++ b/tests/integration/providers/test_bedrock_model_discovery_wire.py @@ -301,11 +301,7 @@ def _spend_row_id(rows: Sequence[Mapping[str, JsonValue]], tag: str) -> str: def _landed_once(served: tuple[tuple[str, str, bool], ...], tags: tuple[str, ...]) -> None: - """One spend row per tagged call, carrying the id the client saw. - - FIXME: a streamed /v1/responses row carries the managed `resp_` id LiteLLM built before the proxy encrypted the - advertised one, so that row is matched by tag alone until the spend log reads the id the client received. - """ + """One spend row per tagged call, carrying the id the client saw.""" rows: Final = eventually( lambda: read_rows( 'SELECT request_id, request_tags FROM "LiteLLM_SpendLogs" WHERE request_tags ?| %s', @@ -316,9 +312,6 @@ def _landed_once(served: tuple[tuple[str, str, bool], ...], tags: tuple[str, ... ) for (identity, path, stream), tag in zip(served, tags, strict=True): row_id: Final = _spend_row_id(rows, tag) - if path == "/v1/responses" and stream: - assert row_id.startswith("resp_"), (row_id, identity) - continue assert row_id == identity, (path, stream, row_id, identity) diff --git a/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py b/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py index 3cadc1b5ccf..42950fb18c3 100644 --- a/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py +++ b/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py @@ -172,7 +172,11 @@ def _upstream_id_inside(row_id: str) -> str | None: def _row_belongs_to(row_id: str, served: _Served) -> bool: if row_id == _response_id(served): return True - return served.call.endpoint == "responses" and _upstream_id_inside(row_id) == f"resp_upstream_{served.call.marker}" + return ( + served.call.endpoint == "responses" + and not served.call.stream + and _upstream_id_inside(row_id) == f"resp_upstream_{served.call.marker}" + ) def _assert_each_success_landed_once(rows: list[dict[str, JsonValue]], served: tuple[_Served, ...]) -> None: diff --git a/tests/integration/providers/test_responses_lifecycle_wire.py b/tests/integration/providers/test_responses_lifecycle_wire.py index 02269885b44..ac5b4dc6d5c 100644 --- a/tests/integration/providers/test_responses_lifecycle_wire.py +++ b/tests/integration/providers/test_responses_lifecycle_wire.py @@ -203,9 +203,6 @@ def test_retrieve_routes_to_the_deployment_that_created_the_response( @pytest.mark.parametrize("prefix", _ALIASES) def test_retrieve_returns_the_id_the_client_holds(gateway: Gateway, prefix: str) -> None: - pytest.skip( - "BUG: GET /v1/responses/{id} returns a re-encrypted id that differs byte-for-byte from the id the client sent" - ) raw_id: Final = f"resp_retrieve_id_{uuid.uuid4().hex}" def respond(request: Request) -> Reply: @@ -697,7 +694,6 @@ def test_delete_returns_the_client_held_id_not_the_provider_id( gateway: Gateway, prefix: str, ) -> None: - pytest.skip("BUG: DELETE /v1/responses/{id} returns the raw provider id instead of the proxy-issued id") raw_id: Final = f"resp_delete_id_{uuid.uuid4().hex}" def respond(request: Request) -> Reply: diff --git a/tests/integration/spend/test_responses_stream_spend_row_id.py b/tests/integration/spend/test_responses_stream_spend_row_id.py new file mode 100644 index 00000000000..58ed9f7785f --- /dev/null +++ b/tests/integration/spend/test_responses_stream_spend_row_id.py @@ -0,0 +1,465 @@ +import json +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +from hashlib import sha256 +from pathlib import Path +from types import MappingProxyType +from typing import Final +from uuid import uuid4 + +import httpx +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, list_value, object_value, string_value +from integration._support.process import owned_proxy +from integration._support.responses_stream import frame, response_object +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_NATIVE_MODEL: Final = "openai/gpt-5.6" +_BRIDGED_MODEL: Final = "deepseek/deepseek-v4-pro" +_ANSWER: Final = "streamed answer" +_USAGE: Final[dict[str, JsonValue]] = {"input_tokens": 30, "output_tokens": 5, "total_tokens": 35} +_JSON_LIST: Final = TypeAdapter(list[dict[str, JsonValue]]) + + +def _sse(payload: dict[str, JsonValue]) -> bytes: + return f"data: {json.dumps(payload)}\n\n".encode() + + +def _native_completed(upstream_id: str) -> dict[str, JsonValue]: + message: Final[dict[str, JsonValue]] = { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": _ANSWER, "annotations": []}], + } + return response_object(upstream_id, "completed", output=[message], usage=_USAGE) + + +def _native_stream(upstream_id: str) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=( + frame( + { + "type": "response.created", + "sequence_number": 0, + "response": response_object(upstream_id, "in_progress"), + } + ), + frame( + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": _ANSWER, + } + ), + frame({"type": "response.completed", "sequence_number": 2, "response": _native_completed(upstream_id)}), + ), + ) + + +def _chat_stream(upstream_id: str, model: str) -> Reply: + def chunk(delta: dict[str, JsonValue], finish_reason: str | None) -> dict[str, JsonValue]: + return { + "id": upstream_id, + "object": "chat.completion.chunk", + "created": 1, + "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + + return Reply( + content_type="text/event-stream", + chunks=( + _sse(chunk({"role": "assistant", "content": _ANSWER}, None)), + _sse( + { + **chunk({}, "stop"), + "usage": {"prompt_tokens": 30, "completion_tokens": 5, "total_tokens": 35}, + } + ), + b"data: [DONE]\n\n", + ), + ) + + +def _respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + if request.method == "GET": + stored_id: Final = request.target.removeprefix("/v1/responses/") + assert stored_id.startswith("resp_upstream-"), request.target + return Reply(body=json.dumps(_native_completed(stored_id)).encode()) + if request.target.endswith("/cancel"): + cancelled_id: Final = request.target.removeprefix("/v1/responses/").removesuffix("/cancel") + assert cancelled_id.startswith("resp_upstream-"), request.target + return Reply(body=json.dumps(response_object(cancelled_id, "cancelled")).encode()) + body: Final = JSON_OBJECT.validate_json(request.body) + upstream_id: Final = f"upstream-{uuid4().hex}" + if request.target == "/v1/responses" and body.get("stream") is not True: + return Reply(body=json.dumps(_native_completed(f"resp_{upstream_id}")).encode()) + assert body.get("stream") is True, body + if request.target == "/v1/responses": + return _native_stream(f"resp_{upstream_id}") + assert request.target == "/v1/chat/completions", request.target + return _chat_stream(f"chatcmpl-{upstream_id}", string_value(body["model"])) + + +def _terminal_response_id( + gateway: Gateway, key: str, body: dict[str, JsonValue], headers: Mapping[str, str] | None = None +) -> str: + response: Final = gateway.request("POST", "/v1/responses", body, key=key, headers=headers) + assert response.status_code == 200, response.text + frames: Final = tuple( + JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + assert frames and frames[-1]["type"] == "response.completed", response.text + advertised: Final = tuple( + dict.fromkeys(string_value(object_value(event["response"])["id"]) for event in frames if "response" in event) + ) + assert len(advertised) == 1, advertised + return advertised[0] + + +def _created_response_id( + gateway: Gateway, key: str, body: dict[str, JsonValue], headers: Mapping[str, str] | None = None +) -> str: + if body.get("stream") is True: + return _terminal_response_id(gateway, key, body, headers) + response: Final = gateway.request("POST", "/v1/responses", body, key=key, headers=headers) + assert response.status_code == 200, response.text + return string_value(JSON_OBJECT.validate_json(response.content)["id"]) + + +def _read_and_cancel( + gateway: Gateway, key: str, response_id: str, headers: Mapping[str, str] | None = None +) -> tuple[httpx.Response, ...]: + responses: Final = tuple( + gateway.request(method, path, key=key, headers=headers) + for method, path in ( + ("GET", f"/v1/responses/{response_id}"), + ("GET", f"/v1/responses/{response_id}"), + ("POST", f"/v1/responses/{response_id}/cancel"), + ) + ) + assert [response.status_code for response in responses] == [200, 200, 200], [r.text for r in responses] + return responses + + +def _spend_logs_lookup(gateway: Gateway, request_id: str) -> list[JsonValue]: + response: Final = gateway.request("GET", "/spend/logs", params={"request_id": request_id}) + assert response.status_code == 200, response.text + return [row["request_id"] for row in _JSON_LIST.validate_json(response.content)] + + +def _request_ids_for_key(gateway: Gateway, key: str, count: int) -> list[str]: + rows: Final = eventually( + lambda: gateway.rows( + 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "endTime"', + (sha256(key.encode()).hexdigest(),), + ), + lambda found: len(found) >= count, + seconds=70, + ) + return [string_value(row["request_id"]) for row in rows] + + +def _logs_page_search(gateway: Gateway, client_id: str) -> list[str]: + now: Final = datetime.now(timezone.utc) + page: Final = gateway.get( + "/spend/logs/ui", + params={ + "search": client_id, + "start_date": (now - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"), + "end_date": (now + timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"), + }, + ) + return [string_value(object_value(row)["request_id"]) for row in list_value(page["data"])] + + +@pytest.mark.timeout(120) +@pytest.mark.parametrize("litellm_model", [_NATIVE_MODEL, _BRIDGED_MODEL], ids=["native", "bridge"]) +def test_streamed_responses_spend_row_is_filed_under_the_id_the_client_saw( + gateway: Gateway, litellm_model: str +) -> None: + with wire_server(_respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=litellm_model, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[model]) + + client_id: Final = _terminal_response_id( + gateway, key, {"model": model, "input": f"marker-{uuid4().hex}", "stream": True} + ) + + assert client_id.startswith("resp_"), client_id + assert _request_ids_for_key(gateway, key, 1) == [client_id] + assert _logs_page_search(gateway, client_id) == [client_id] + assert _spend_logs_lookup(gateway, client_id) == [client_id] + + +def _proxy_config( + directory: Path, + general_settings: Mapping[str, JsonValue], + litellm_settings: Mapping[str, JsonValue] = MappingProxyType({}), +) -> Path: + stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + path: Final = directory / "responses-spend.yaml" + path.write_text( + yaml.safe_dump( + { + **stock, + "general_settings": {**object_value(stock["general_settings"]), **general_settings}, + "litellm_settings": {**object_value(stock["litellm_settings"]), **litellm_settings}, + } + ) + ) + return path + + +def _prompt_storing_config(directory: Path) -> Path: + return _proxy_config(directory, {"store_prompts_in_spend_logs": True}) + + +def test_a_follow_up_turn_on_a_streamed_bridged_response_replays_the_first_turn( + gateway: Gateway, tmp_path: Path +) -> None: + first_prompt: Final = f"first turn {uuid4().hex}" + with ( + wire_server(_respond) as wire, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storing_config(tmp_path)) as storing, + storing.scenario() as scenario, + ): + model: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[model]) + first_id: Final = _terminal_response_id(storing, key, {"model": model, "input": first_prompt, "stream": True}) + assert _request_ids_for_key(storing, key, 1) == [first_id] + wire.drain() + + second_id: Final = _terminal_response_id( + storing, + key, + {"model": model, "input": "second turn", "previous_response_id": first_id, "stream": True}, + ) + + replayed: Final = _JSON_LIST.validate_python(JSON_OBJECT.validate_json(wire.drain()[-1].body)["messages"]) + assert [(message["role"], message["content"]) for message in replayed] == [ + ("user", first_prompt), + ("assistant", _ANSWER), + ("user", "second turn"), + ], replayed + assert _request_ids_for_key(storing, key, 2) == [first_id, second_id] + + +@pytest.mark.timeout(180) +@pytest.mark.parametrize("stream_first_turn", [True, False], ids=["streamed-first-turn", "non-streamed-first-turn"]) +def test_a_follow_up_through_the_id_a_retrieve_returned_replays_the_first_turn_once( + gateway: Gateway, tmp_path: Path, stream_first_turn: bool +) -> None: + first_prompt: Final = f"first turn {uuid4().hex}" + with ( + wire_server(_respond) as wire, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storing_config(tmp_path)) as storing, + storing.scenario() as scenario, + ): + native: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + bridged: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[native, bridged]) + first_id: Final = _created_response_id( + storing, key, {"model": native, "input": first_prompt, "stream": stream_first_turn} + ) + + retrieved: Final = storing.request("GET", f"/v1/responses/{first_id}", key=key) + + assert retrieved.status_code == 200, retrieved.text + retrieved_id: Final = string_value(JSON_OBJECT.validate_json(retrieved.content)["id"]) + assert retrieved_id == first_id + stored: Final = _request_ids_for_key(storing, key, 2) + assert stored.count(first_id) == 1, stored + assert len(set(stored)) == 2, stored + assert _spend_logs_lookup(storing, first_id) == [first_id] + wire.drain() + + _terminal_response_id( + storing, + key, + {"model": bridged, "input": "second turn", "previous_response_id": retrieved_id, "stream": True}, + ) + + replayed: Final = _JSON_LIST.validate_python(JSON_OBJECT.validate_json(wire.drain()[-1].body)["messages"]) + assert [(message["role"], message["content"]) for message in replayed] == [ + ("user", first_prompt), + ("assistant", _ANSWER), + ("user", "second turn"), + ], replayed + + +def _replayed_turns(wire_body: bytes) -> list[tuple[JsonValue, JsonValue]]: + replayed: Final = _JSON_LIST.validate_python(JSON_OBJECT.validate_json(wire_body)["messages"]) + return [(message["role"], message["content"]) for message in replayed] + + +@pytest.mark.timeout(180) +@pytest.mark.parametrize("reader", ["proxy-admin", "master-key"]) +def test_a_reader_other_than_the_creator_continues_through_the_id_its_retrieve_returned( + gateway: Gateway, tmp_path: Path, reader: str +) -> None: + first_prompt: Final = f"first turn {uuid4().hex}" + with ( + wire_server(_respond) as wire, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storing_config(tmp_path)) as storing, + storing.scenario() as scenario, + ): + native: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + bridged: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + creator: Final = scenario.key(user_id=scenario.user(), models=[native, bridged]) + reader_key: Final = ( + scenario.key(user_id=scenario.user(user_role="proxy_admin"), models=[native, bridged]) + if reader == "proxy-admin" + else storing.key + ) + first_id: Final = _terminal_response_id( + storing, creator, {"model": native, "input": first_prompt, "stream": True} + ) + assert _request_ids_for_key(storing, creator, 1) == [first_id] + + retrieved: Final = storing.request("GET", f"/v1/responses/{first_id}", key=reader_key) + + assert retrieved.status_code == 200, retrieved.text + retrieved_id: Final = string_value(JSON_OBJECT.validate_json(retrieved.content)["id"]) + wire.drain() + _terminal_response_id( + storing, + reader_key, + {"model": bridged, "input": "second turn", "previous_response_id": retrieved_id, "stream": True}, + ) + assert _replayed_turns(wire.drain()[-1].body) == [ + ("user", first_prompt), + ("assistant", _ANSWER), + ("user", "second turn"), + ] + assert retrieved_id == first_id + + +@pytest.mark.timeout(180) +def test_reads_and_cancels_sent_with_a_session_header_add_no_replayed_turns(gateway: Gateway, tmp_path: Path) -> None: + first_prompt: Final = f"first turn {uuid4().hex}" + session: Final = {"x-litellm-session-id": f"session-{uuid4().hex}"} + with ( + wire_server(_respond) as wire, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storing_config(tmp_path)) as storing, + storing.scenario() as scenario, + ): + native: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + bridged: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[native, bridged]) + first_id: Final = _created_response_id(storing, key, {"model": native, "input": first_prompt}, session) + + _read_and_cancel(storing, key, first_id, session) + + assert len(_request_ids_for_key(storing, key, 4)) == 4 + wire.drain() + _terminal_response_id( + storing, + key, + {"model": bridged, "input": "second turn", "previous_response_id": first_id, "stream": True}, + session, + ) + assert _replayed_turns(wire.drain()[-1].body) == [ + ("user", first_prompt), + ("assistant", _ANSWER), + ("user", "second turn"), + ] + + +@pytest.mark.timeout(180) +def test_each_read_and_cancel_reaches_callbacks_and_spend_logs_under_an_id_of_its_own( + gateway: Gateway, tmp_path: Path +) -> None: + logged_call_types: Final = frozenset({"aresponses", "aget_responses", "acancel_responses"}) + with ( + wire_server(_respond) as wire, + wire_server(lambda _: Reply()) as sink, + owned_proxy( + gateway, + tmp_path, + {"GENERIC_LOGGER_ENDPOINT": sink.url}, + config=_proxy_config(tmp_path, {}, {"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}), + ) as logging_proxy, + logging_proxy.scenario() as scenario, + ): + model: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[model]) + first_id: Final = _created_response_id(logging_proxy, key, {"model": model, "input": f"marker-{uuid4().hex}"}) + + reads: Final = _read_and_cancel(logging_proxy, key, first_id) + + assert [string_value(JSON_OBJECT.validate_json(read.content)["id"]) for read in reads] == [first_id] * 3 + batches: Final[list[bytes]] = [] # mutable-ok: drain() consumes, later polls keep earlier batches + + def delivered() -> tuple[tuple[str, str], ...]: + batches.extend(batch.body for batch in sink.drain()) + return tuple( + (string_value(event["call_type"]), string_value(event["id"])) + for event in (object_value(item) for batch in batches for item in _JSON_LIST.validate_json(batch)) + if event["call_type"] in logged_call_types + ) + + events: Final = eventually(delivered, lambda found: len(found) >= 4, seconds=30) + read_ids: Final = [logged_id for call_type, logged_id in events if call_type != "aresponses"] + assert [logged_id for call_type, logged_id in events if call_type == "aresponses"] == [first_id], events + assert len(read_ids) == 3, events + assert first_id not in read_ids, events + assert len(set(read_ids)) == 3, events + assert sorted(_request_ids_for_key(logging_proxy, key, 4)) == sorted([first_id, *read_ids]) + + +@pytest.mark.timeout(240) +def test_with_id_security_off_reads_and_cancels_add_no_rows_and_no_replayed_turns( + gateway: Gateway, tmp_path: Path +) -> None: + first_prompt: Final = f"first turn {uuid4().hex}" + session: Final = {"x-litellm-session-id": f"session-{uuid4().hex}"} + with ( + wire_server(_respond) as wire, + owned_proxy( + gateway, + tmp_path, + {}, + config=_proxy_config( + tmp_path, {"store_prompts_in_spend_logs": True, "disable_responses_id_security": True} + ), + ) as insecure, + insecure.scenario() as scenario, + ): + native: Final = scenario.model(model=_NATIVE_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + bridged: Final = scenario.model(model=_BRIDGED_MODEL, api_base=f"{wire.url}/v1", api_key="synthetic-key") + key: Final = scenario.key(models=[native, bridged]) + first_id: Final = _created_response_id(insecure, key, {"model": native, "input": first_prompt}, session) + + reads: Final = _read_and_cancel(insecure, key, first_id, session) + + assert [string_value(JSON_OBJECT.validate_json(read.content)["id"]) for read in reads] == [first_id] * 3 + later_id: Final = _created_response_id(insecure, key, {"model": native, "input": "a later request"}) + assert _request_ids_for_key(insecure, key, 2) == [first_id, later_id] + wire.drain() + + _terminal_response_id( + insecure, + key, + {"model": bridged, "input": "second turn", "previous_response_id": first_id, "stream": True}, + session, + ) + + replayed: Final = _JSON_LIST.validate_python(JSON_OBJECT.validate_json(wire.drain()[-1].body)["messages"]) + assert [(message["role"], message["content"]) for message in replayed] == [ + ("user", first_prompt), + ("assistant", _ANSWER), + ("user", "second turn"), + ], replayed diff --git a/tests/unit/integrations/azure_storage/test_azure_storage.py b/tests/unit/integrations/azure_storage/test_azure_storage.py index 0f0ce24ff84..ee12291e7a0 100644 --- a/tests/unit/integrations/azure_storage/test_azure_storage.py +++ b/tests/unit/integrations/azure_storage/test_azure_storage.py @@ -5,19 +5,28 @@ import re import sys import threading from datetime import datetime, timedelta, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest -from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS +from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS, INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.azure_storage import azure_storage as azure_storage_module from litellm.integrations.azure_storage.azure_storage import ( AzureBlobStorageLogger, _cached_credential_chain_token_provider, adls_safe_file_name, ) +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) from litellm.types.secret_managers.get_azure_ad_token_provider import AzureCredentialType -from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload +from litellm.types.utils import ( + BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, + StandardAuditLogPayload, + StandardLoggingPayload, +) @pytest.fixture @@ -648,3 +657,171 @@ async def test_audit_log_is_uploaded_to_a_dated_audit_logs_path_with_an_account_ assert json.loads(uploaded) == audit_log assert file_client.append_data.await_args.kwargs == {"data": uploaded, "offset": 0, "length": len(uploaded)} file_client.flush_data.assert_awaited_once_with(position=len(uploaded), offset=0) + + +BACKGROUND_RESPONSE_ID: Final = "resp_YWJjZA==" +BACKGROUND_COST_POLL_KWARGS: Final = { + "call_type": "aget_responses", + "litellm_params": { + "metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}, + }, +} +PASSTHROUGH_SETTLEMENT_KWARGS: Final = { + "call_type": "pass_through_endpoint", + "litellm_params": {"metadata": {}}, + SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES, +} + + +def _background_create_cost_and_other_events( + cost_kwargs: dict[str, object], +) -> tuple[tuple[dict[str, object], StandardLoggingPayload], ...]: + create: Final = StandardLoggingPayload( + id=BACKGROUND_RESPONSE_ID, + call_type="aresponses", + messages=[{"role": "user", "content": "background create input"}], + ) + cost: Final = StandardLoggingPayload(id=BACKGROUND_RESPONSE_ID, call_type="aget_responses", response_cost=0.42) + ordinary: Final = StandardLoggingPayload(id="chatcmpl-ordinary", call_type="acompletion") + batch_retrieve: Final = StandardLoggingPayload(id="batch_abc123", call_type="aretrieve_batch", response_cost=0.07) + return ( + ({"call_type": "aresponses", "litellm_params": {"metadata": {}}, "standard_logging_object": create}, create), + ({**cost_kwargs, "standard_logging_object": cost}, cost), + ( + {"call_type": "acompletion", "litellm_params": {"metadata": {}}, "standard_logging_object": ordinary}, + ordinary, + ), + ( + { + "call_type": "aretrieve_batch", + "litellm_params": {"metadata": {}}, + "standard_logging_object": batch_retrieve, + }, + batch_retrieve, + ), + ) + + +def _expected_files( + events: tuple[tuple[dict[str, object], StandardLoggingPayload], ...], +) -> dict[str, StandardLoggingPayload]: + create, cost, ordinary, batch_retrieve = (payload for _, payload in events) + return { + "resp_YWJjZA.json": create, + "resp_YWJjZA_background_cost.json": cost, + "chatcmpl-ordinary.json": ordinary, + "batch_abc123.json": batch_retrieve, + } + + +async def _log_and_flush( + logger: AzureBlobStorageLogger, events: tuple[tuple[dict[str, object], StandardLoggingPayload], ...] +) -> None: + for kwargs, payload in events: + await logger.async_log_success_event(kwargs, {"id": payload["id"]}, None, None) + await logger.async_send_batch() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "cost_kwargs", + ( + pytest.param(BACKGROUND_COST_POLL_KWARGS, id="cost-poll"), + pytest.param(PASSTHROUGH_SETTLEMENT_KWARGS, id="passthrough-settlement"), + ), +) +async def test_background_cost_event_keeps_its_own_file_beside_the_create_on_the_entra_path( + mock_env_vars, monkeypatch, cost_kwargs +): + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + events: Final = _background_create_cost_and_other_events(cost_kwargs) + with patch( # test-quality-ok: REST client is created inside the method; assert emitted requests + "litellm.integrations.azure_storage.azure_storage.get_async_httpx_client" + ) as mock_get_client: + mock_http_client = AsyncMock() + mock_http_client.put.return_value = MagicMock() + mock_http_client.patch.return_value = MagicMock() + mock_get_client.return_value = mock_http_client + + logger = AzureBlobStorageLogger() + logger.azure_auth_token = "mock-azure-ad-token" + logger.token_expiry = None + await _log_and_flush(logger, events) + + base_url: Final = "https://test-account.dfs.core.windows.net/test-container/" + appends: Final = [call for call in mock_http_client.patch.call_args_list if "action=append" in call.args[0]] + uploaded: Final = { + call.args[0].removeprefix(base_url).split("?")[0]: json.loads(call.kwargs["data"]) for call in appends + } + assert uploaded == _expected_files(events), ( + "the create, its background cost event and every other event must each keep their own file" + ) + + +@pytest.mark.asyncio +async def test_failure_event_is_uploaded_to_a_file_named_after_its_id(mock_env_vars, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + failure: Final = StandardLoggingPayload(id="chatcmpl-failed", call_type="acompletion", error_str="provider failure") + with patch( # test-quality-ok: REST client is created inside the method; assert emitted requests + "litellm.integrations.azure_storage.azure_storage.get_async_httpx_client" + ) as mock_get_client: + mock_http_client = AsyncMock() + mock_http_client.put.return_value = MagicMock() + mock_http_client.patch.return_value = MagicMock() + mock_get_client.return_value = mock_http_client + + logger = AzureBlobStorageLogger() + logger.azure_auth_token = "mock-azure-ad-token" + logger.token_expiry = None + await logger.async_log_failure_event( + {"call_type": "acompletion", "litellm_params": {"metadata": {}}, "standard_logging_object": failure}, + None, + None, + None, + ) + await logger.async_send_batch() + + base_url: Final = "https://test-account.dfs.core.windows.net/test-container/" + appends: Final = [call for call in mock_http_client.patch.call_args_list if "action=append" in call.args[0]] + assert {call.args[0].removeprefix(base_url).split("?")[0]: json.loads(call.kwargs["data"]) for call in appends} == { + "chatcmpl-failed.json": failure + } + + +@pytest.mark.asyncio +async def test_background_cost_event_keeps_its_own_file_beside_the_create_with_an_account_key( + workload_identity_env_vars, monkeypatch +): + monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_KEY", "dGVzdC1rZXk=") + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr(azure_storage_module, "datetime", _FrozenDatetime) + events: Final = _background_create_cost_and_other_events(BACKGROUND_COST_POLL_KWARGS) + + uploaded: Final[dict[str, StandardLoggingPayload]] = {} + + def file_client_for(file_path: str) -> MagicMock: + file_client = MagicMock() + file_client.create_file = AsyncMock() + file_client.flush_data = AsyncMock() + + async def append_data(data: bytes, offset: int, length: int) -> None: + uploaded[file_path] = json.loads(data) + + file_client.append_data = append_data + return file_client + + file_system_client = MagicMock() + file_system_client.get_file_client = MagicMock(side_effect=file_client_for) + service_client = MagicMock() + service_client.get_file_system_client = MagicMock(return_value=file_system_client) + fake_aio_module = MagicMock() + fake_aio_module.DataLakeServiceClient = MagicMock(return_value=service_client) + + with patch.dict(sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}): + logger = AzureBlobStorageLogger() + await _log_and_flush(logger, events) + + today: Final = _FrozenDatetime.now().strftime("%Y-%m-%d") + assert uploaded == {f"{today}/{name}": payload for name, payload in _expected_files(events).items()}, ( + "the create, its background cost event and every other event must each keep their own file" + ) diff --git a/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py b/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py index f94d7bbee37..42e69f95bb4 100644 --- a/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py +++ b/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py @@ -20,10 +20,16 @@ from litellm.integrations.clickhouse.clickhouse_spend_logger import ( strip_cache_hit_suffix, ) from litellm.integrations.clickhouse.context import lens_analysis +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.litellm_core_utils.internal_call_metadata import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX from litellm.litellm_core_utils import litellm_logging from litellm.litellm_core_utils.secret_redaction import REDACTED from litellm.tracing.types import SpendLogRecord -from litellm.types.utils import StandardLoggingPayload +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, StandardLoggingPayload _JSON_OBJECT_ADAPTER: Final = TypeAdapter(Mapping[str, JsonValue]) @@ -349,6 +355,41 @@ def test_cache_hit_id_is_stripped_for_response_id(): assert strip_cache_hit_suffix("chatcmpl-xyz") == "chatcmpl-xyz" +@pytest.mark.parametrize( + "cost_kwargs", + ( + pytest.param( + { + "call_type": "aget_responses", + "litellm_params": { + "metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN} + }, + }, + id="cost-poll", + ), + pytest.param( + { + "call_type": "pass_through_endpoint", + "litellm_params": {"metadata": {}}, + SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES, + }, + id="passthrough-settlement", + ), + ), +) +def test_background_cost_row_keeps_its_own_request_id_beside_its_create(cost_kwargs: Mapping[str, object]) -> None: + payload: Final = cast(StandardLoggingPayload, _payload(id="resp_YWJjZA", call_type="aresponses")) + create_kwargs: Final[Mapping[str, object]] = {"call_type": "aresponses", "litellm_params": {"metadata": {}}} + + rows: Final = [spend_log_row_from_payload(payload, kwargs) for kwargs in (create_kwargs, cost_kwargs)] + + assert [row["request_id"] for row in rows] == [ + "resp_YWJjZA", + f"resp_YWJjZA{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}", + ] + assert [row["response_id"] for row in rows] == ["resp_YWJjZA", "resp_YWJjZA"] + + def test_parse_traceparent_valid_missing_malformed(): assert parse_traceparent(TRACEPARENT) == (TRACE_ID, SPAN_ID) assert parse_traceparent(None) == ("", "") diff --git a/tests/unit/integrations/conftest.py b/tests/unit/integrations/conftest.py index 48a01ed8d48..9278a280fe7 100644 --- a/tests/unit/integrations/conftest.py +++ b/tests/unit/integrations/conftest.py @@ -1,14 +1,18 @@ import functools import http.server +import importlib.abc import ipaddress import os import queue import ssl +import sys import threading -from collections.abc import Iterator +from collections.abc import Iterator, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone +from importlib.machinery import ModuleSpec from pathlib import Path +from types import ModuleType from typing import Final import pytest @@ -75,6 +79,36 @@ def write_self_signed_cert(directory: Path, stem: str) -> tuple[Path, Path]: return certificate_path, key_path +PROXY_ONLY_PACKAGES: Final = frozenset({"fastapi", "starlette"}) + + +class _ProxyOnlyPackageBlocker(importlib.abc.MetaPathFinder): + def find_spec( + self, + fullname: str, + path: Sequence[str] | None = None, + target: ModuleType | None = None, + ) -> ModuleSpec | None: + if fullname.partition(".")[0] in PROXY_ONLY_PACKAGES: + raise ModuleNotFoundError(f"No module named {fullname!r}") + return None + + +def _needs_proxy_extras(module_name: str) -> bool: + return ( + module_name.partition(".")[0] in PROXY_ONLY_PACKAGES + or module_name == "litellm.proxy" + or module_name.startswith("litellm.proxy.") + ) + + +@pytest.fixture +def sdk_only_install(monkeypatch: pytest.MonkeyPatch) -> None: + for module_name in tuple(filter(_needs_proxy_extras, sys.modules)): + monkeypatch.delitem(sys.modules, module_name) + monkeypatch.setattr(sys, "meta_path", [_ProxyOnlyPackageBlocker(), *sys.meta_path]) + + @pytest.fixture(autouse=True) def restore_process_environment() -> Iterator[None]: original: Final = dict(os.environ) diff --git a/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py b/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py index 6f143a244f6..c29985b1e41 100644 --- a/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py +++ b/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py @@ -13,9 +13,11 @@ import litellm import pytest import respx +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN mock_response_data: Final = { @@ -172,9 +174,7 @@ async def test_get_payload_not_found(monkeypatch): assert payload is None -def _gcs_logger_with_upload_status( - monkeypatch: pytest.MonkeyPatch, status_code: int -) -> tuple[GCSBucketLogger, Mock]: +def _gcs_logger_with_upload_status(monkeypatch: pytest.MonkeyPatch, status_code: int) -> tuple[GCSBucketLogger, Mock]: monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) monkeypatch.setattr(google.auth, "default", _google_default_credentials) monkeypatch.setenv("GCS_FLUSH_INTERVAL", "3600") @@ -394,3 +394,98 @@ async def test_batch_line_items_are_serialized_off_the_event_loop(monkeypatch: p (thread,) = threads assert thread is not threading.main_thread() + + +@pytest.mark.asyncio +async def test_unbatched_background_cost_event_keeps_its_own_object_readable_by_its_spend_log_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("GCS_USE_BATCHED_LOGGING", "false") + logger, _ = _gcs_logger_with_upload_status(monkeypatch, 200) + object_prefix: Final = "/storage/v1/b/test-bucket/o/" + stored: Final[dict[str, object]] = {} + + def bucket(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + stored[request.url.params["name"]] = json.loads(request.content) + return httpx.Response(200, json={"name": request.url.params["name"]}) + name: Final = request.url.path.removeprefix(object_prefix) + if name not in stored: + return httpx.Response(404, text="No such object") + return httpx.Response(200, content=json.dumps(stored[name]).encode("utf-8")) + + logger.async_httpx_client = AsyncHTTPHandler(transport=httpx.MockTransport(bucket)) + response_id: Final = "resp_YmFja2dyb3VuZA" + create: Final = { + "id": response_id, + "messages": [{"role": "user", "content": "background create input"}], + "error_str": None, + } + cost: Final = {"id": response_id, "response_cost": 0.42, "error_str": None} + ordinary: Final = {"id": "chatcmpl-ordinary", "error_str": None} + batch_retrieve: Final = {"id": "batch_abc123", "response_cost": 0.07, "error_str": None} + events: Final = ( + ({"call_type": "aresponses", "litellm_params": {"metadata": {}}}, create), + ( + { + "call_type": "aget_responses", + "litellm_params": { + "metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN} + }, + }, + cost, + ), + ({"call_type": "acompletion", "litellm_params": {"metadata": {}}}, ordinary), + ({"call_type": "aretrieve_batch", "litellm_params": {"metadata": {}}}, batch_retrieve), + ) + + for kwargs, payload in events: + await logger.async_log_success_event( + {**kwargs, "standard_logging_object": payload}, {"id": payload["id"]}, None, None + ) + await logger.flush_queue() + + assert {name.split("/", maxsplit=1)[1]: payload for name, payload in stored.items()} == { + response_id: create, + f"{response_id}_background_cost": cost, + "chatcmpl-ordinary": ordinary, + "batch_abc123": batch_retrieve, + }, "the create, its background cost event and every other event must each keep their own object" + now: Final = datetime.now(timezone.utc) + assert await logger.get_request_response_payload(f"{response_id}_background_cost", now, None) == cost + assert await logger.get_request_response_payload(response_id, now, None) == create + + +@pytest.mark.asyncio +async def test_unbatched_gcs_log_id_still_names_a_background_cost_event_object(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("GCS_USE_BATCHED_LOGGING", "false") + logger, _ = _gcs_logger_with_upload_status(monkeypatch, 200) + stored: Final[dict[str, object]] = {} + + def bucket(request: httpx.Request) -> httpx.Response: + stored[request.url.params["name"]] = json.loads(request.content) + return httpx.Response(200, json={"name": request.url.params["name"]}) + + logger.async_httpx_client = AsyncHTTPHandler(transport=httpx.MockTransport(bucket)) + cost: Final = {"id": "resp_YmFja2dyb3VuZA", "response_cost": 0.42, "error_str": None} + await logger.async_log_success_event( + { + "call_type": "aget_responses", + "litellm_params": { + "metadata": { + INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, + "gcs_log_id": "qa-run-7", + } + }, + "standard_logging_object": cost, + }, + {"id": cost["id"]}, + None, + None, + ) + await logger.flush_queue() + + ((name, payload),) = stored.items() + object_name: Final = name.split("/", maxsplit=1)[1] + assert object_name.startswith("custom-") and object_name.endswith("-qa-run-7"), object_name + assert payload == cost diff --git a/tests/unit/integrations/test_deepeval.py b/tests/unit/integrations/test_deepeval.py index 6dd3aef0a03..4523a9c61d2 100644 --- a/tests/unit/integrations/test_deepeval.py +++ b/tests/unit/integrations/test_deepeval.py @@ -2,11 +2,15 @@ import os import unittest from litellm._uuid import uuid from datetime import datetime, timezone +from typing import Final from unittest.mock import MagicMock, patch from litellm.integrations.deepeval.api import Endpoints, HttpMethods from litellm.integrations.deepeval.deepeval import DeepEvalLogger from litellm.integrations.deepeval.types import SpanApiType, TraceSpanApiStatus +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.litellm_core_utils.internal_call_metadata import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN class TestDeepEvalLogger(unittest.TestCase): @@ -94,6 +98,37 @@ class TestDeepEvalLogger(unittest.TestCase): self.assertEqual(llm_span["inputTokenCount"], 10) self.assertEqual(llm_span["outputTokenCount"], 20) + def test_background_cost_span_keeps_its_own_uuid_beside_its_create(self) -> None: + standard_logging_object: Final = { + "id": "resp_YWJjZA", + "trace_id": self.trace_id, + "model": self.model, + "response": {"usage": {}, "choices": [{"message": {"content": "OK"}}]}, + } + create_kwargs: Final = { + "input": self.input_str, + "call_type": "aresponses", + "litellm_params": {"metadata": {}}, + "standard_logging_object": standard_logging_object, + } + cost_kwargs: Final = { + "input": self.input_str, + "call_type": "aget_responses", + "litellm_params": { + "metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN} + }, + "standard_logging_object": standard_logging_object, + } + + for kwargs in (create_kwargs, cost_kwargs): + self.logger.log_success_event(kwargs, self.mock_response_obj, self.start_time, self.end_time) + + span_uuids: Final = [ + call.kwargs["body"]["llmSpans"][0]["uuid"] + for call in self.mock_api_instance.send_request.call_args_list + ] + self.assertEqual(span_uuids, ["resp_YWJjZA", f"resp_YWJjZA{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}"]) + def test_log_failure_event(self): error_message = "This is an error." kwargs = { diff --git a/tests/unit/integrations/test_dynamodb.py b/tests/unit/integrations/test_dynamodb.py new file mode 100644 index 00000000000..d139c49a49e --- /dev/null +++ b/tests/unit/integrations/test_dynamodb.py @@ -0,0 +1,52 @@ +from collections.abc import Mapping +from datetime import datetime, timezone +from typing import Final +from unittest.mock import MagicMock + +import pytest + +import litellm +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.integrations.dynamodb import DyanmoDBLogger +from litellm.litellm_core_utils.internal_call_metadata import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + +BACKGROUND_COST_POLL_KWARGS: Final[Mapping[str, object]] = { + "call_type": "aget_responses", + "litellm_params": {"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}}, +} +PASSTHROUGH_SETTLEMENT_KWARGS: Final[Mapping[str, object]] = { + "call_type": "pass_through_endpoint", + "litellm_params": {"metadata": {}}, + SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES, +} + + +@pytest.mark.parametrize( + "cost_kwargs", + ( + pytest.param(BACKGROUND_COST_POLL_KWARGS, id="cost-poll"), + pytest.param(PASSTHROUGH_SETTLEMENT_KWARGS, id="passthrough-settlement"), + ), +) +@pytest.mark.usefixtures("sdk_only_install") +def test_background_cost_event_keeps_its_own_item_beside_its_create( + monkeypatch: pytest.MonkeyPatch, cost_kwargs: Mapping[str, object] +) -> None: + monkeypatch.setattr(litellm, "dynamodb_table_name", "litellm-logs") + dynamodb: Final = MagicMock() + logger: Final = DyanmoDBLogger(dynamodb=dynamodb) + response_obj: Final = {"id": "resp_YWJjZA", "usage": {"total_tokens": 3}} + now: Final = datetime.now(timezone.utc) + create_kwargs: Final[Mapping[str, object]] = {"call_type": "aresponses", "litellm_params": {"metadata": {}}} + + for kwargs in (create_kwargs, cost_kwargs): + logger.log_event(kwargs, response_obj, now, now, print_verbose=print) + + dynamodb.Table.assert_called_with("litellm-logs") + written_ids: Final = [call.kwargs["Item"]["id"] for call in dynamodb.Table.return_value.put_item.call_args_list] + assert written_ids == ["resp_YWJjZA", f"resp_YWJjZA{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}"] diff --git a/tests/unit/integrations/test_openmeter.py b/tests/unit/integrations/test_openmeter.py index 2d09e1572db..fc4e2fcd7a3 100644 --- a/tests/unit/integrations/test_openmeter.py +++ b/tests/unit/integrations/test_openmeter.py @@ -1,11 +1,30 @@ import json import os +from collections.abc import Mapping +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest import litellm +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.openmeter import OpenMeterLogger +from litellm.litellm_core_utils.internal_call_metadata import BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX +from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + SETTLES_BACKGROUND_RESPONSE_KWARG, + BackgroundResponseSettlement, +) +from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + +BACKGROUND_COST_POLL_KWARGS: Final[Mapping[str, object]] = { + "call_type": "aget_responses", + "litellm_params": {"metadata": {INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}}, +} +PASSTHROUGH_SETTLEMENT_KWARGS: Final[Mapping[str, object]] = { + "call_type": "pass_through_endpoint", + "litellm_params": {"metadata": {}}, + SETTLES_BACKGROUND_RESPONSE_KWARG: BackgroundResponseSettlement.SETTLES, +} class TestOpenMeterIntegration: @@ -36,6 +55,31 @@ class TestOpenMeterIntegration: with pytest.raises(Exception, match=r"Missing keys.*OPENMETER_API_KEY"): OpenMeterLogger() + @pytest.mark.parametrize( + "cost_kwargs", + ( + pytest.param(BACKGROUND_COST_POLL_KWARGS, id="cost-poll"), + pytest.param(PASSTHROUGH_SETTLEMENT_KWARGS, id="passthrough-settlement"), + ), + ) + @pytest.mark.usefixtures("sdk_only_install") + def test_background_cost_event_gets_its_own_event_id_beside_its_create( + self, cost_kwargs: Mapping[str, object] + ) -> None: + logger: Final = OpenMeterLogger() + response_obj: Final = {"id": "resp_YWJjZA"} + create_kwargs: Final[Mapping[str, object]] = { + "user": "user-1", + "call_type": "aresponses", + "litellm_params": {"metadata": {}}, + } + + create_event: Final = logger._common_logic(create_kwargs, response_obj) + cost_event: Final = logger._common_logic({**cost_kwargs, "user": "user-1", "response_cost": 0.42}, response_obj) + + assert create_event["id"] == "resp_YWJjZA" + assert cost_event["id"] == f"resp_YWJjZA{BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX}" + def test_common_logic_with_string_user(self): """Test that _common_logic correctly handles string user parameter""" logger = OpenMeterLogger() diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index f1854563cd5..13e181b2b92 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -11090,6 +11090,40 @@ def test_get_assembled_streaming_response_bills_a_provider_reported_usage_cost() assert logging_obj.response_cost_calculator(result=assembled) == 0.0042 +@pytest.mark.parametrize("continued_response_id", ["resp_continuedafteracutstream", None]) +@pytest.mark.parametrize("client_facing_response_id", ["resp_encryptedidtheclientsaw", None]) +def test_sync_success_handler_logs_a_stream_under_the_id_the_client_saw( + client_facing_response_id, continued_response_id +): + class Recorder(CustomLogger): + def __init__(self): + super().__init__() + self.logged = [] + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + self.logged.append((response_obj.id, kwargs["standard_logging_object"]["id"])) + + recorder: Final = Recorder() + logging_obj: Final = _responses_stream_logging_obj() + logging_obj.dynamic_success_callbacks = [recorder] + logging_obj.client_facing_response_id = client_facing_response_id + if continued_response_id is not None: + logging_obj.record_continued_partial( + Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), 0.0, continued_response_id + ) + now: Final = datetime.datetime.now() + + logging_obj.success_handler( + result=_completed_responses_event(ResponseAPIUsage(input_tokens=3, output_tokens=2, total_tokens=5)), + start_time=now, + end_time=now, + cache_hit=False, + ) + + expected_id: Final = client_facing_response_id or continued_response_id or "resp-1" + assert recorder.logged == [(expected_id, expected_id)] + + def test_response_cost_calculator_prices_terminal_responses_event_from_its_response(): logging_obj: Final = _responses_stream_logging_obj() inner_response: Final = ResponsesAPIResponse( diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index f52f40ad55d..f73d33cc93e 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -5150,6 +5150,23 @@ def test_a_background_response_billing_event_writes_a_row_apart_from_the_read_th assert request_ids == ("resp_shared", "resp_shared_background_cost") +@pytest.mark.parametrize("call_type", ["aget_responses", "acancel_responses"]) +def test_a_read_or_cancel_row_carries_the_id_its_callbacks_logged(call_type: str): + row: Final = get_logging_payload( + kwargs={ + "call_type": call_type, + "model": "gpt-5.6", + "litellm_call_id": "call-read", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + response_obj={"id": "resp_logged_by_callbacks", "object": "response", "status": "completed"}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert row["request_id"] == "resp_logged_by_callbacks" + + @pytest.mark.parametrize("forged_marker", [True, BackgroundResponseSettlement.SETTLES.value]) def test_a_request_body_cannot_move_its_spend_row_onto_a_background_settlement_id(forged_marker): """A failed call logs its raw request body as kwargs, so a body field named like the settlement diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index f1d841a8e5e..9a8fa19e2e4 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -1083,6 +1083,128 @@ async def test_previous_response_history_starts_at_a_compaction_item_in_a_stored ] +def _stored_response_turn( + request_id: str, call_type: str, status: str, output: list[dict[str, object]], prompt: str | None = None +) -> dict[str, object]: + return { + "request_id": request_id, + "call_type": call_type, + "session_id": "session-with-reads", + "proxy_server_request": json.dumps({} if prompt is None else {"input": [{"role": "user", "content": prompt}]}), + "response": {"id": request_id, "object": "response", "status": status, "output": output}, + } + + +def _identified_output_message(item_id: str, text: str) -> dict[str, object]: + return {**_output_message(text), "id": item_id} + + +def _replayed_roles_and_texts(session: ChatCompletionSession) -> list[tuple[object, object]]: + return [ + (role, _text_of(content)) for role, content in (_role_and_content(message) for message in session["messages"]) + ] + + +@pytest.mark.asyncio +async def test_reads_and_cancels_of_a_response_the_session_already_holds_replay_no_turns(): + answer: Final = [_identified_output_message("msg_answer", "OK")] + spend_logs: Final = ( + _stored_response_turn("resp_create", "aresponses", "completed", answer, prompt="Reply with OK"), + _stored_response_turn("resp_read", "aget_responses", "completed", answer), + _stored_response_turn("resp_streamed_read", "aget_responses", "completed", answer), + _stored_response_turn("resp_cancel", "acancel_responses", "completed", answer), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_create") + + assert _replayed_roles_and_texts(session) == [("user", "Reply with OK"), ("assistant", "OK")] + + +@pytest.mark.asyncio +async def test_a_polled_background_response_replays_its_finished_output_once(): + answer: Final = [_identified_output_message("msg_answer", "One, two, three")] + spend_logs: Final = ( + _stored_response_turn("resp_queued", "aresponses", "queued", [], prompt="Count to three"), + _stored_response_turn("resp_poll_running", "aget_responses", "in_progress", []), + _stored_response_turn("resp_poll_done", "aget_responses", "completed", answer), + _stored_response_turn("resp_poll_again", "aget_responses", "completed", answer), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_queued") + + assert _replayed_roles_and_texts(session) == [("user", "Count to three"), ("assistant", "One, two, three")] + + +@pytest.mark.asyncio +async def test_a_background_response_that_grows_between_polls_replays_each_output_item_once(): + first: Final = _identified_output_message("msg_first", "One") + second: Final = _identified_output_message("msg_second", "Two") + spend_logs: Final = ( + _stored_response_turn("resp_queued", "aresponses", "queued", [], prompt="Count to two"), + _stored_response_turn("resp_poll_partial", "aget_responses", "in_progress", [first]), + _stored_response_turn("resp_poll_done", "aget_responses", "completed", [first, second]), + _stored_response_turn("resp_poll_again", "aget_responses", "completed", [first, second]), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_queued") + + assert _replayed_roles_and_texts(session) == [("user", "Count to two"), ("assistant", "One"), ("assistant", "Two")] + + +@pytest.mark.asyncio +async def test_an_item_a_poll_caught_half_written_replays_its_finished_text_once(): + spend_logs: Final = ( + _stored_response_turn("resp_queued", "aresponses", "queued", [], prompt="Count to three"), + _stored_response_turn( + "resp_poll_partial", "aget_responses", "in_progress", [_identified_output_message("msg_answer", "One, t")] + ), + _stored_response_turn( + "resp_poll_done", + "aget_responses", + "completed", + [_identified_output_message("msg_answer", "One, two, three"), _identified_output_message("msg_more", "Done")], + ), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_queued") + + assert _replayed_roles_and_texts(session) == [ + ("user", "Count to three"), + ("assistant", "One, two, three"), + ("assistant", "Done"), + ] + + +@pytest.mark.asyncio +async def test_a_polled_output_item_without_an_id_still_replays_in_order(): + spend_logs: Final = ( + _stored_response_turn("resp_queued", "aresponses", "queued", [], prompt="Count to two"), + _stored_response_turn( + "resp_poll_done", + "aget_responses", + "completed", + [_output_message("One"), _identified_output_message("msg_two", "Two")], + ), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_queued") + + assert _replayed_roles_and_texts(session) == [("user", "Count to two"), ("assistant", "One"), ("assistant", "Two")] + + +@pytest.mark.asyncio +async def test_a_read_that_ends_before_its_streamed_create_replays_the_answer_once_after_the_prompt(): + answer: Final = [_identified_output_message("msg_answer", "OK")] + spend_logs: Final = ( + _stored_response_turn("resp_read_mid_stream", "aget_responses", "completed", answer), + _stored_response_turn("resp_create", "aresponses", "completed", answer, prompt="Say OK"), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_create") + + assert _replayed_roles_and_texts(session) == [("user", "Say OK"), ("assistant", "OK")] + + def _tool_call_ids(message: object) -> tuple[str, ...]: tool_calls: Final = message.get("tool_calls") if isinstance(message, dict) else getattr(message, "tool_calls", None) return tuple(call["id"] if isinstance(call, dict) else call.id for call in tool_calls or ()) diff --git a/tests/unit/test_responses_id_security.py b/tests/unit/test_responses_id_security.py index 7e74bfd0f17..8dc50fa33c7 100644 --- a/tests/unit/test_responses_id_security.py +++ b/tests/unit/test_responses_id_security.py @@ -4,23 +4,33 @@ Tests for ResponsesIDSecurity hook. Tests the security hook that prevents user B from seeing response from user A. """ +import datetime +from collections.abc import Callable, Mapping from unittest.mock import MagicMock, patch import pytest from fastapi import HTTPException +from litellm.caching.dual_cache import DualCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.responses_id_security import ( ResponsesIDSecurity, _streams_response_ids, ) +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( GenericEvent, + ResponseAPIUsage, ResponseCompletedEvent, ResponseCreatedEvent, ResponsesAPIResponse, ResponsesAPIStreamEvents, ) -from litellm.types.utils import SpecialEnums +from litellm.types.responses.main import DeleteResponseResult +from litellm.types.utils import CallTypesLiteral, Choices, Message, ModelResponse, SpecialEnums, Usage @pytest.fixture @@ -55,12 +65,12 @@ class TestIsEncryptedResponseId: import litellm.proxy.hooks.responses_id_security as responses_module with patch.object(responses_module, "decrypt_value_helper") as mock_decrypt: - mock_decrypt.return_value = f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}response_id:resp_123;user_id:user-456" - - result = responses_id_security._is_encrypted_response_id( - "resp_encrypted_value" + mock_decrypt.return_value = ( + f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}response_id:resp_123;user_id:user-456" ) + result = responses_id_security._is_encrypted_response_id("resp_encrypted_value") + assert result is True mock_decrypt.assert_called_once() @@ -88,9 +98,7 @@ class TestDecryptResponseId: with patch.object(responses_module, "decrypt_value_helper") as mock_decrypt: mock_decrypt.return_value = f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}response_id:resp_original_123;user_id:user-456;team_id:team-789" - original_id, user_id, team_id = responses_id_security._decrypt_response_id( - "resp_encrypted_value" - ) + original_id, user_id, team_id = responses_id_security._decrypt_response_id("resp_encrypted_value") assert original_id == "resp_original_123" assert user_id == "user-456" @@ -104,9 +112,7 @@ class TestDecryptResponseId: with patch.object(responses_module, "decrypt_value_helper") as mock_decrypt: mock_decrypt.return_value = None - original_id, user_id, team_id = responses_id_security._decrypt_response_id( - "resp_plain_value" - ) + original_id, user_id, team_id = responses_id_security._decrypt_response_id("resp_plain_value") assert original_id == "resp_plain_value" assert user_id is None @@ -116,9 +122,7 @@ class TestDecryptResponseId: class TestCheckUserAccessToResponseId: """Test check_user_access_to_response_id function""" - def test_check_user_access_same_user( - self, responses_id_security, mock_user_api_key_dict - ): + def test_check_user_access_same_user(self, responses_id_security, mock_user_api_key_dict): """Test that same user has access to their response ID""" result = responses_id_security.check_user_access_to_response_id( response_id_user_id="test-user-123", @@ -128,9 +132,7 @@ class TestCheckUserAccessToResponseId: assert result is True - def test_check_user_access_different_user_raises_exception( - self, responses_id_security, mock_user_api_key_dict - ): + def test_check_user_access_different_user_raises_exception(self, responses_id_security, mock_user_api_key_dict): """Test that different user is denied access to response ID""" with patch("litellm.proxy.proxy_server.general_settings", {}): with pytest.raises(HTTPException) as exc_info: @@ -143,9 +145,7 @@ class TestCheckUserAccessToResponseId: assert exc_info.value.status_code == 403 assert "Forbidden" in exc_info.value.detail - def test_check_user_access_different_team_raises_exception( - self, responses_id_security, mock_user_api_key_dict - ): + def test_check_user_access_different_team_raises_exception(self, responses_id_security, mock_user_api_key_dict): """Test that different team is denied access to response ID""" with patch("litellm.proxy.proxy_server.general_settings", {}): with pytest.raises(HTTPException) as exc_info: @@ -158,9 +158,7 @@ class TestCheckUserAccessToResponseId: assert exc_info.value.status_code == 403 assert "Forbidden" in exc_info.value.detail - def test_check_user_access_team_a_to_team_b_without_user_id( - self, responses_id_security - ): + def test_check_user_access_team_a_to_team_b_without_user_id(self, responses_id_security): """Test that key from team A (without user_id) cannot access response from team B (without user_id)""" # Create a mock user from team A without user_id mock_auth_team_a = MagicMock() @@ -179,9 +177,7 @@ class TestCheckUserAccessToResponseId: assert exc_info.value.status_code == 403 assert "team" in exc_info.value.detail.lower() - def test_check_user_access_team_a_to_team_b_with_user_id( - self, responses_id_security - ): + def test_check_user_access_team_a_to_team_b_with_user_id(self, responses_id_security): """Test that key from team A (without user_id) cannot access response from team B (with user_id)""" # Create a mock user from team A without user_id mock_auth_team_a = MagicMock() @@ -217,9 +213,7 @@ class TestCheckUserAccessToResponseId: assert result is True - def test_check_user_access_admin_can_access_any_response( - self, responses_id_security - ): + def test_check_user_access_admin_can_access_any_response(self, responses_id_security): """Test that proxy admin can access any response ID""" from litellm.proxy._types import LitellmUserRoles @@ -238,9 +232,7 @@ class TestCheckUserAccessToResponseId: assert result is True - def test_check_user_access_security_disabled( - self, responses_id_security, mock_user_api_key_dict - ): + def test_check_user_access_security_disabled(self, responses_id_security, mock_user_api_key_dict): """Test that when security is disabled, any user can access any response""" with patch( "litellm.proxy.proxy_server.general_settings", @@ -266,9 +258,7 @@ class TestAsyncPreCallHook: """Test pre-call hook decrypts previous_response_id for aresponses call""" data = {"previous_response_id": "resp_encrypted_value"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -284,15 +274,11 @@ class TestAsyncPreCallHook: assert result["previous_response_id"] == "resp_original_123" @pytest.mark.asyncio - async def test_async_pre_call_hook_aget_responses( - self, responses_id_security, mock_user_api_key_dict, mock_cache - ): + async def test_async_pre_call_hook_aget_responses(self, responses_id_security, mock_user_api_key_dict, mock_cache): """Test pre-call hook decrypts response_id for aget_responses call""" data = {"response_id": "resp_encrypted_456"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -308,9 +294,7 @@ class TestAsyncPreCallHook: assert result["response_id"] == "resp_original_456" @pytest.mark.asyncio - async def test_async_pre_call_hook_team_a_accessing_team_b_response( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_team_a_accessing_team_b_response(self, responses_id_security, mock_cache): """Test pre-call hook prevents team A from accessing team B response""" # Create a mock user from team A mock_auth_team_a = MagicMock() @@ -320,9 +304,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_b"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -341,9 +323,7 @@ class TestAsyncPreCallHook: assert "team" in exc_info.value.detail.lower() @pytest.mark.asyncio - async def test_async_pre_call_hook_team_a_accessing_team_b_with_user( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_team_a_accessing_team_b_with_user(self, responses_id_security, mock_cache): """Test pre-call hook prevents team A (no user) from accessing team B response (with user)""" # Create a mock user from team A without user_id mock_auth_team_a = MagicMock() @@ -353,9 +333,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_b_with_user"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -375,9 +353,7 @@ class TestAsyncPreCallHook: assert "forbidden" in exc_info.value.detail.lower() @pytest.mark.asyncio - async def test_async_pre_call_hook_same_team_access( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_same_team_access(self, responses_id_security, mock_cache): """Test pre-call hook allows team A to access their own team's response""" # Create a mock user from team A mock_auth_team_a = MagicMock() @@ -387,9 +363,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_a"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -405,9 +379,7 @@ class TestAsyncPreCallHook: assert result["response_id"] == "resp_original_team_a" @pytest.mark.asyncio - async def test_async_pre_call_hook_adelete_responses_team_security( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_adelete_responses_team_security(self, responses_id_security, mock_cache): """Test pre-call hook prevents team A from deleting team B's response""" # Create a mock user from team A mock_auth_team_a = MagicMock() @@ -417,9 +389,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_b"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -438,9 +408,7 @@ class TestAsyncPreCallHook: assert "team" in exc_info.value.detail.lower() @pytest.mark.asyncio - async def test_async_pre_call_hook_acancel_responses_team_security( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_acancel_responses_team_security(self, responses_id_security, mock_cache): """Test pre-call hook prevents team A from canceling team B's response""" # Create a mock user from team A mock_auth_team_a = MagicMock() @@ -450,9 +418,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_b"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -470,16 +436,13 @@ class TestAsyncPreCallHook: assert exc_info.value.status_code == 403 assert "team" in exc_info.value.detail.lower() - @pytest.mark.asyncio async def test_async_pre_call_hook_alist_input_items_decrypts_response_id( self, responses_id_security, mock_user_api_key_dict, mock_cache ): data = {"response_id": "resp_encrypted_789"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -496,9 +459,7 @@ class TestAsyncPreCallHook: assert result["response_id"] == "resp_original_789" @pytest.mark.asyncio - async def test_async_pre_call_hook_alist_input_items_team_security( - self, responses_id_security, mock_cache - ): + async def test_async_pre_call_hook_alist_input_items_team_security(self, responses_id_security, mock_cache): mock_auth_team_a = MagicMock() mock_auth_team_a.user_id = None mock_auth_team_a.team_id = "team-a" @@ -506,9 +467,7 @@ class TestAsyncPreCallHook: data = {"response_id": "resp_encrypted_team_b"} - with patch.object( - responses_id_security, "_is_encrypted_response_id", return_value=True - ): + with patch.object(responses_id_security, "_is_encrypted_response_id", return_value=True): with patch.object( responses_id_security, "_decrypt_response_id", @@ -620,9 +579,7 @@ class TestAsyncPostCallStreamingIteratorHook: "/openai/v1/responses/resp_rawprovider123", ], ) - async def test_streamed_id_encrypted_on_all_responses_routes( - self, responses_id_security, route, monkeypatch - ): + async def test_streamed_id_encrypted_on_all_responses_routes(self, responses_id_security, route, monkeypatch): streamed_id = await self._drain_streamed_id(responses_id_security, route, monkeypatch) assert streamed_id != "resp_rawprovider123" @@ -634,12 +591,8 @@ class TestAsyncPostCallStreamingIteratorHook: ) @pytest.mark.asyncio - async def test_streamed_id_untouched_on_non_responses_route( - self, responses_id_security, monkeypatch - ): - streamed_id = await self._drain_streamed_id( - responses_id_security, "/chat/completions", monkeypatch - ) + async def test_streamed_id_untouched_on_non_responses_route(self, responses_id_security, monkeypatch): + streamed_id = await self._drain_streamed_id(responses_id_security, "/chat/completions", monkeypatch) assert streamed_id == "resp_rawprovider123" assert not responses_id_security._is_encrypted_response_id(streamed_id) @@ -698,12 +651,8 @@ class TestStreamedGenericEventIdEncryption: def _advertised_ids(events): nested = (getattr(event, "response", None) for event in events) return [ - payload["id"] if isinstance(payload, dict) else payload.id - for payload in nested - if payload is not None - ] + [ - event.id for event in events if isinstance(getattr(event, "id", None), str) - ] + payload["id"] if isinstance(payload, dict) else payload.id for payload in nested if payload is not None + ] + [event.id for event in events if isinstance(getattr(event, "id", None), str)] async def _drain(self, responses_id_security, monkeypatch): monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-abcdefghij") @@ -723,9 +672,7 @@ class TestStreamedGenericEventIdEncryption: ] @pytest.mark.asyncio - async def test_every_event_advertises_the_same_encrypted_id( - self, responses_id_security, monkeypatch - ): + async def test_every_event_advertises_the_same_encrypted_id(self, responses_id_security, monkeypatch): events = await self._drain(responses_id_security, monkeypatch) advertised = self._advertised_ids(events) @@ -742,22 +689,14 @@ class TestStreamedGenericEventIdEncryption: ) @pytest.mark.asyncio - async def test_raw_provider_id_never_reaches_the_client( - self, responses_id_security, monkeypatch - ): + async def test_raw_provider_id_never_reaches_the_client(self, responses_id_security, monkeypatch): events = await self._drain(responses_id_security, monkeypatch) - assert [self.RAW_ID in event.model_dump_json() for event in events] == [ - False - ] * len(events) + assert [self.RAW_ID in event.model_dump_json() for event in events] == [False] * len(events) @pytest.mark.asyncio - async def test_sibling_fields_survive_the_rewrite( - self, responses_id_security, monkeypatch - ): - _, queued, keepalive, later, _ = await self._drain( - responses_id_security, monkeypatch - ) + async def test_sibling_fields_survive_the_rewrite(self, responses_id_security, monkeypatch): + _, queued, keepalive, later, _ = await self._drain(responses_id_security, monkeypatch) assert queued.response["status"] == "queued" assert later.response["status"] == "in_progress" @@ -769,27 +708,19 @@ class TestAsyncPostCallSuccessHook: """Test async_post_call_success_hook function""" @pytest.mark.asyncio - async def test_async_post_call_success_hook_encrypts_response( - self, responses_id_security, mock_user_api_key_dict - ): + async def test_async_post_call_success_hook_encrypts_response(self, responses_id_security, mock_user_api_key_dict): """Test post-call hook encrypts ResponsesAPIResponse""" - mock_response = ResponsesAPIResponse( - id="resp_789", created_at=1234567890, output=[], status="completed" - ) + mock_response = ResponsesAPIResponse(id="resp_789", created_at=1234567890, output=[], status="completed") data = {} - with patch.object( - responses_id_security, "_encrypt_response_id", return_value=mock_response - ) as mock_encrypt: + with patch.object(responses_id_security, "_encrypt_response_id", return_value=mock_response) as mock_encrypt: result = await responses_id_security.async_post_call_success_hook( data=data, user_api_key_dict=mock_user_api_key_dict, response=mock_response, ) - mock_encrypt.assert_called_once_with( - mock_response, mock_user_api_key_dict, request_cache=None - ) + mock_encrypt.assert_called_once_with(mock_response, mock_user_api_key_dict, request_cache=None) assert result == mock_response @pytest.mark.asyncio @@ -846,9 +777,7 @@ def _auth(user_id="owner-user", team_id="owner-team", user_role=None): def _issue_managed_id(hook, owner, provider_response_id=_FABRICATED_PROVIDER_RESPONSE_ID): """Mint an id exactly the way the proxy hands one to a client on create.""" issued = hook._encrypt_response_id( - ResponsesAPIResponse( - id=provider_response_id, created_at=1234567890, output=[], status="completed" - ), + ResponsesAPIResponse(id=provider_response_id, created_at=1234567890, output=[], status="completed"), owner, ) return issued.id @@ -1080,9 +1009,7 @@ class TestClientSuppliedRetainedIdCannotBypassAuthorization: assert exc_info.value.status_code == 403 @pytest.mark.asyncio - async def test_stranger_forging_their_own_id_never_reaches_someone_elses_response( - self, mock_cache, salt_key_env - ): + async def test_stranger_forging_their_own_id_never_reaches_someone_elses_response(self, mock_cache, salt_key_env): hook = _hook() stranger = _auth(user_id="stranger-user", team_id="stranger-team") stranger_id = _issue_managed_id(hook, stranger, provider_response_id="resp_strangerownprovideridcccccccc") @@ -1098,3 +1025,436 @@ class TestClientSuppliedRetainedIdCannotBypassAuthorization: assert result["response_id"] == "resp_strangerownprovideridcccccccc" assert result["response_id"] != victim_provider_id + + +def _admin(): + from litellm.proxy._types import LitellmUserRoles + + return _auth(user_id="admin-user", team_id="admin-team", user_role=LitellmUserRoles.PROXY_ADMIN) + + +def _provider_answer(response_id: str) -> ResponsesAPIResponse: + return ResponsesAPIResponse(id=response_id, created_at=1234567890, output=[], status="completed") + + +class TestReadsAnswerWithTheIdTheClientHolds: + """A follow-up looks its history up by the id the client holds, so a read must not hand out a new one.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", ["aget_responses", "acancel_responses"]) + @pytest.mark.parametrize( + "issuer, reader", + [ + (_auth(), _auth()), + (_auth(team_id=None), _auth(team_id="another-team")), + (_auth(user_id=None), _auth(user_id=None)), + (_auth(user_id=None), _auth(user_id="teammate-user")), + (_auth(), _admin()), + (_auth(user_id=None), _auth(user_id=None, team_id="admin-team", user_role=_admin().user_role)), + ], + ids=[ + "same-key", + "same-user-new-team", + "same-team-no-user", + "teammate-reading-a-team-key-id", + "admin-user", + "admin-without-user", + ], + ) + async def test_a_caller_with_access_gets_back_the_exact_id_they_sent( + self, mock_cache, salt_key_env, call_type, issuer, reader + ): + hook = _hook() + issued_id = _issue_managed_id(hook, issuer) + data = await hook.async_pre_call_hook( + user_api_key_dict=reader, cache=mock_cache, data={"response_id": issued_id}, call_type=call_type + ) + + result = await hook.async_post_call_success_hook( + data=data, user_api_key_dict=reader, response=_provider_answer(data["response_id"]) + ) + + assert result.id == issued_id + + @pytest.mark.asyncio + async def test_a_follow_up_gets_an_id_for_its_own_response_never_the_previous_one(self, mock_cache, salt_key_env): + hook = _hook() + owner = _auth() + previous_id = _issue_managed_id(hook, owner) + data = await hook.async_pre_call_hook( + user_api_key_dict=owner, + cache=mock_cache, + data={"model": "gpt-fake", "previous_response_id": previous_id}, + call_type="aresponses", + ) + + result = await hook.async_post_call_success_hook( + data=data, user_api_key_dict=owner, response=_provider_answer("resp_followupprovideridzzzzzzzzzzzzzzz") + ) + + assert result.id != previous_id + assert hook._decrypt_response_id(result.id) == ( + "resp_followupprovideridzzzzzzzzzzzzzzz", + "owner-user", + "owner-team", + ) + + @pytest.mark.asyncio + async def test_an_admin_read_by_a_raw_provider_id_gets_a_managed_id_back(self, mock_cache, salt_key_env): + hook = _hook() + admin = _admin() + data = await hook.async_pre_call_hook( + user_api_key_dict=admin, + cache=mock_cache, + data={"response_id": _FABRICATED_UNMANAGED_ID}, + call_type="aget_responses", + ) + + result = await hook.async_post_call_success_hook( + data=data, user_api_key_dict=admin, response=_provider_answer(_FABRICATED_UNMANAGED_ID) + ) + + assert hook._is_encrypted_response_id(result.id) + assert hook._decrypt_response_id(result.id) == (_FABRICATED_UNMANAGED_ID, "admin-user", "admin-team") + + +_LOGGED_MANAGED_ID = "resp_managedidlitellmbuiltbeforethehookran" + + +def _streamed_response(response_id): + return ResponsesAPIResponse( + id=response_id, + created_at=0, + model="gpt-5.1", + object="response", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + usage=ResponseAPIUsage(input_tokens=3, output_tokens=2, total_tokens=5), + ) + + +def _upstream_frames(): + return ( + ResponseCreatedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_CREATED, response=_streamed_response(_LOGGED_MANAGED_ID) + ), + ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=_streamed_response(_LOGGED_MANAGED_ID) + ), + ) + + +def _native_logged_result(): + return ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=_streamed_response(_LOGGED_MANAGED_ID) + ) + + +def _bridged_logged_result(): + return ModelResponse( + id="chatcmpl-bridgedchatcompletionid", + model="gpt-5.1", + choices=[Choices(message=Message(role="assistant", content="hi"), finish_reason="stop", index=0)], + usage=Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5), + ) + + +class _SpendRowRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.spend_row_ids = [] + self.standard_logging_ids = [] + self.callback_response_ids = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.spend_row_ids.append(get_logging_payload(kwargs, response_obj, start_time, end_time)["request_id"]) + self.callback_response_ids.append(response_obj.id) + if "standard_logging_object" in kwargs: + self.standard_logging_ids.append(kwargs["standard_logging_object"]["id"]) + + +async def _stream_logged_mid_stream( + hook: ResponsesIDSecurity, + route: str, + logged_result: ResponseCompletedEvent | ModelResponse, + caller: UserAPIKeyAuth | None = None, + request_data: Mapping[str, object] | None = None, +) -> tuple[str, _SpendRowRecorder]: + streaming_caller = caller if caller is not None else _auth() + recorder = _SpendRowRecorder() + logging_obj = LiteLLMLoggingObj( + model="gpt-5.1", + messages=[], + stream=True, + call_type="aresponses", + start_time=datetime.datetime.now(), + litellm_call_id="streamed-responses-call", + function_id="streamed-responses-fn", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="gpt-5.1", user="", optional_params={}, litellm_params={"api_base": ""} + ) + upstream = _upstream_frames() + + async def frames(): + for frame in upstream: + yield frame + + client_stream = hook.async_post_call_streaming_iterator_hook( + user_api_key_dict=streaming_caller.model_copy(update={"request_route": route}), + response=frames(), + request_data={**(request_data or {}), "litellm_logging_obj": logging_obj}, + ) + first_frame = await anext(client_stream) + now = datetime.datetime.now() + await logging_obj.async_success_handler(result=logged_result, start_time=now, end_time=now, cache_hit=False) + client_frames = [first_frame, *[frame async for frame in client_stream]] + return client_frames[-1].response.id, recorder + + +class TestStreamedSpendRowCarriesTheIdTheClientSaw: + @pytest.mark.asyncio + @pytest.mark.parametrize("logged_result", [_native_logged_result, _bridged_logged_result], ids=["native", "bridge"]) + @pytest.mark.parametrize("route", ["/v1/responses", "/responses", "/openai/v1/responses"]) + async def test_spend_row_and_callbacks_log_the_encrypted_id_of_the_terminal_frame( + self, salt_key_env, route, logged_result + ): + hook = _hook() + + client_id, recorder = await _stream_logged_mid_stream(hook, route, logged_result()) + + assert hook._decrypt_response_id(client_id) == (_LOGGED_MANAGED_ID, "owner-user", "owner-team") + assert recorder.spend_row_ids == [client_id] + assert recorder.standard_logging_ids == [client_id] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "hook, route", + [ + (_hook(), "/chat/completions"), + (_hook(general_settings={"disable_responses_id_security": True}), "/v1/responses"), + (_hook(signing_key=None), "/v1/responses"), + ], + ids=["not-a-create-route", "security-disabled", "no-signing-key"], + ) + async def test_stream_the_hook_leaves_alone_keeps_its_logged_id(self, salt_key_env, hook, route): + client_id, recorder = await _stream_logged_mid_stream(hook, route, _native_logged_result()) + + assert client_id == _LOGGED_MANAGED_ID + assert recorder.spend_row_ids == [_LOGGED_MANAGED_ID] + assert recorder.standard_logging_ids == [_LOGGED_MANAGED_ID] + + +class TestStreamedReadAnswersWithTheIdTheClientHolds: + @pytest.mark.asyncio + @pytest.mark.parametrize("route_prefix", ["/v1/responses", "/responses", "/openai/v1/responses"]) + async def test_the_owner_sees_their_id_while_callbacks_log_one_of_the_reads_own( + self, mock_cache, salt_key_env, route_prefix + ): + hook = _hook() + owner = _auth() + issued_id = _issue_managed_id(hook, owner, provider_response_id=_LOGGED_MANAGED_ID) + data = await hook.async_pre_call_hook( + user_api_key_dict=owner, cache=mock_cache, data={"response_id": issued_id}, call_type="aget_responses" + ) + + client_id, recorder = await _stream_logged_mid_stream( + hook, f"{route_prefix}/{issued_id}", _native_logged_result(), caller=owner, request_data=data + ) + + assert client_id == issued_id + assert recorder.spend_row_ids == recorder.standard_logging_ids + assert len(recorder.standard_logging_ids) == 1 + assert recorder.standard_logging_ids[0] != issued_id + assert hook._decrypt_response_id(recorder.standard_logging_ids[0]) == ( + _LOGGED_MANAGED_ID, + "owner-user", + "owner-team", + ) + + @pytest.mark.asyncio + async def test_an_admin_streamed_read_of_someone_elses_response_sees_the_id_sent_and_logs_one_of_its_own( + self, mock_cache, salt_key_env + ): + hook = _hook() + admin = _admin() + issued_id = _issue_managed_id(hook, _auth(), provider_response_id=_LOGGED_MANAGED_ID) + data = await hook.async_pre_call_hook( + user_api_key_dict=admin, cache=mock_cache, data={"response_id": issued_id}, call_type="aget_responses" + ) + + client_id, recorder = await _stream_logged_mid_stream( + hook, f"/v1/responses/{issued_id}", _native_logged_result(), caller=admin, request_data=data + ) + + assert client_id == issued_id + assert recorder.spend_row_ids == recorder.standard_logging_ids + assert len(recorder.standard_logging_ids) == 1 + assert hook._decrypt_response_id(recorder.standard_logging_ids[0]) == ( + _LOGGED_MANAGED_ID, + "admin-user", + "admin-team", + ) + + + @pytest.mark.asyncio + async def test_a_frame_naming_the_held_response_under_another_wrapping_still_gets_the_id_sent( + self, mock_cache, salt_key_env + ): + hook = _hook() + owner = _auth() + issued_id = _issue_managed_id(hook, owner, provider_response_id=_routed_response_id(_LOGGED_MANAGED_ID)) + data = await hook.async_pre_call_hook( + user_api_key_dict=owner, cache=mock_cache, data={"response_id": issued_id}, call_type="aget_responses" + ) + + client_id, _ = await _stream_logged_mid_stream( + hook, f"/v1/responses/{issued_id}", _native_logged_result(), caller=owner, request_data=data + ) + + assert client_id == issued_id + + +async def _read_logged_by_callbacks( + hook: ResponsesIDSecurity, + caller: UserAPIKeyAuth, + issued_id: str, + call_type: CallTypesLiteral, + cache: DualCache, + answer: Callable[[str], ResponsesAPIResponse | DeleteResponseResult] = _provider_answer, +) -> tuple[str, _SpendRowRecorder]: + recorder = _SpendRowRecorder() + now = datetime.datetime.now() + logging_obj = LiteLLMLoggingObj( + model="gpt-5.1", + messages=[], + stream=False, + call_type=call_type, + start_time=now, + litellm_call_id=f"{call_type}-call", + function_id=f"{call_type}-fn", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="gpt-5.1", user="", optional_params={}, litellm_params={"api_base": ""} + ) + data = await hook.async_pre_call_hook( + user_api_key_dict=caller, cache=cache, data={"response_id": issued_id}, call_type=call_type + ) + provider_answer = answer(data["response_id"]) + + client_response = await hook.async_post_call_success_hook( + data=data, user_api_key_dict=caller, response=provider_answer + ) + await logging_obj.async_success_handler(result=provider_answer, start_time=now, end_time=now, cache_hit=False) + return client_response.id, recorder + + +class TestEveryReadLogsAnIdOfItsOwn: + @pytest.mark.asyncio + @pytest.mark.parametrize("call_type", ["aget_responses", "acancel_responses"]) + @pytest.mark.parametrize("reader", [_auth(), _admin()], ids=["owner", "admin"]) + async def test_callbacks_and_the_spend_row_log_a_new_id_per_read_never_the_one_the_client_holds( + self, mock_cache, salt_key_env, call_type, reader + ): + hook = _hook() + issued_id = _issue_managed_id(hook, _auth()) + + reads = [await _read_logged_by_callbacks(hook, reader, issued_id, call_type, mock_cache) for _ in range(2)] + + assert [client_id for client_id, _ in reads] == [issued_id, issued_id] + logged_ids = [id_ for _, recorder in reads for id_ in recorder.standard_logging_ids] + assert [recorder.spend_row_ids for _, recorder in reads] == [[logged_id] for logged_id in logged_ids] + assert len(set(logged_ids)) == 2 + assert issued_id not in logged_ids + assert {hook._decrypt_response_id(logged_id) for logged_id in logged_ids} == { + (_FABRICATED_PROVIDER_RESPONSE_ID, reader.user_id, reader.team_id) + } + + +_RAW_UPSTREAM_RESPONSE_ID = "resp_06de59f3upstreamiddddddddddddddddd" + + +def _routed_response_id(upstream_response_id: str) -> str: + return ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( + _provider_answer(upstream_response_id), "openai", {"model_info": {"id": "deployment-1"}} + ).id + + +def _deletion_of(response_id: str) -> DeleteResponseResult: + return DeleteResponseResult( + id=ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(response_id), + object="response.deleted", + deleted=True, + ) + + +class TestDeleteAnswersWithTheIdTheClientHolds: + @pytest.mark.asyncio + @pytest.mark.parametrize("reader", [_auth(), _admin()], ids=["owner", "admin"]) + @pytest.mark.parametrize( + "stored_id", + [_routed_response_id(_RAW_UPSTREAM_RESPONSE_ID), _RAW_UPSTREAM_RESPONSE_ID], + ids=["routed-deployment", "plain-id"], + ) + async def test_a_caller_with_access_gets_back_the_exact_id_they_sent( + self, mock_cache, salt_key_env, reader, stored_id + ): + hook = _hook() + issued_id = _issue_managed_id(hook, _auth(), provider_response_id=stored_id) + data = await hook.async_pre_call_hook( + user_api_key_dict=reader, cache=mock_cache, data={"response_id": issued_id}, call_type="adelete_responses" + ) + + result = await hook.async_post_call_success_hook( + data=data, user_api_key_dict=reader, response=_deletion_of(data["response_id"]) + ) + + assert (result.id, result.object, result.deleted) == (issued_id, "response.deleted", True) + + @pytest.mark.asyncio + @pytest.mark.parametrize("reader", [_auth(), _admin()], ids=["owner", "admin"]) + async def test_callbacks_and_the_spend_row_log_an_id_of_its_own_never_the_one_the_client_holds( + self, mock_cache, salt_key_env, reader + ): + hook = _hook() + issued_id = _issue_managed_id( + hook, _auth(), provider_response_id=_routed_response_id(_RAW_UPSTREAM_RESPONSE_ID) + ) + + client_id, recorder = await _read_logged_by_callbacks( + hook, reader, issued_id, "adelete_responses", mock_cache, answer=_deletion_of + ) + + assert client_id == issued_id + assert recorder.spend_row_ids == recorder.callback_response_ids + assert len(recorder.callback_response_ids) == 1 + assert recorder.callback_response_ids[0] != issued_id + assert hook._decrypt_response_id(recorder.callback_response_ids[0]) == ( + _RAW_UPSTREAM_RESPONSE_ID, + reader.user_id, + reader.team_id, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "sent_id", + [_routed_response_id(_RAW_UPSTREAM_RESPONSE_ID), _RAW_UPSTREAM_RESPONSE_ID], + ids=["routed-deployment", "plain-id"], + ) + async def test_an_admin_delete_by_a_raw_provider_id_gets_the_provider_id_back_and_logs_an_id_of_its_own( + self, mock_cache, salt_key_env, sent_id + ): + hook = _hook() + + client_id, recorder = await _read_logged_by_callbacks( + hook, _admin(), sent_id, "adelete_responses", mock_cache, answer=_deletion_of + ) + + assert client_id == _RAW_UPSTREAM_RESPONSE_ID + assert [hook._decrypt_response_id(logged_id) for logged_id in recorder.callback_response_ids] == [ + (_RAW_UPSTREAM_RESPONSE_ID, "admin-user", "admin-team") + ] From cc142a819132f78b35e71f3893e9660d84ec42a4 Mon Sep 17 00:00:00 2001 From: jesus-berri Date: Mon, 2 Nov 2026 14:17:10 -0800 Subject: [PATCH 18/79] fix(cost): add deepgram/streaming/nova-3-medical and update stale Nova-3 streaming rates (#49733) The /deepgram/v1/listen WebSocket refuses nova-3-medical with 1008 because the cost map has no streaming row for it. Deepgram's pricing page does not list Nova-3 Medical, so the new row uses the $0.0077/min streaming launch price from Deepgram's 2025-03-03 announcement deepgram/streaming/nova-3 and nova-3-multilingual were $0.0048 and $0.0058 per minute against $0.0077 and $0.0092 on deepgram.com/pricing as of 2026-11-02, so they are updated too Resolves LIT-10393 --- ...odel_prices_and_context_window_backup.json | 29 ++++++++++++++----- model_prices_and_context_window.json | 29 ++++++++++++++----- .../deepgram/test_deepgram_common_utils.py | 6 +++- 3 files changed, 49 insertions(+), 15 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8bfb8949890..517aead5d29 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23523,12 +23523,12 @@ ] }, "deepgram/streaming/nova-3": { - "input_cost_per_second": 8e-05, + "input_cost_per_second": 0.00012833, "litellm_provider": "deepgram", "metadata": { - "calculation": "$0.0048/60 seconds = $0.00008000 per second", + "calculation": "$0.0077/60 seconds = $0.00012833 per second", "note": "Nova-3 monolingual streaming, pay as you go", - "original_pricing_per_minute": 0.0048 + "original_pricing_per_minute": 0.0077 }, "mode": "audio_transcription", "output_cost_per_second": 0.0, @@ -23537,13 +23537,28 @@ "/v1/listen" ] }, - "deepgram/streaming/nova-3-multilingual": { - "input_cost_per_second": 9.667e-05, + "deepgram/streaming/nova-3-medical": { + "input_cost_per_second": 0.00012833, "litellm_provider": "deepgram", "metadata": { - "calculation": "$0.0058/60 seconds = $0.00009667 per second", + "calculation": "$0.0077/60 seconds = $0.00012833 per second", + "note": "Nova-3 Medical streaming, launch price from Deepgram's 2025-03-03 announcement. deepgram.com/pricing does not list Nova-3 Medical, and this matches its current Nova-3 monolingual streaming pay as you go rate", + "original_pricing_per_minute": 0.0077 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://www.businesswire.com/news/home/20250303470439/en", + "supported_endpoints": [ + "/v1/listen" + ] + }, + "deepgram/streaming/nova-3-multilingual": { + "input_cost_per_second": 0.00015333, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0092/60 seconds = $0.00015333 per second", "note": "Nova-3 multilingual (language=multi) streaming, pay as you go", - "original_pricing_per_minute": 0.0058 + "original_pricing_per_minute": 0.0092 }, "mode": "audio_transcription", "output_cost_per_second": 0.0, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8bfb8949890..517aead5d29 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23523,12 +23523,12 @@ ] }, "deepgram/streaming/nova-3": { - "input_cost_per_second": 8e-05, + "input_cost_per_second": 0.00012833, "litellm_provider": "deepgram", "metadata": { - "calculation": "$0.0048/60 seconds = $0.00008000 per second", + "calculation": "$0.0077/60 seconds = $0.00012833 per second", "note": "Nova-3 monolingual streaming, pay as you go", - "original_pricing_per_minute": 0.0048 + "original_pricing_per_minute": 0.0077 }, "mode": "audio_transcription", "output_cost_per_second": 0.0, @@ -23537,13 +23537,28 @@ "/v1/listen" ] }, - "deepgram/streaming/nova-3-multilingual": { - "input_cost_per_second": 9.667e-05, + "deepgram/streaming/nova-3-medical": { + "input_cost_per_second": 0.00012833, "litellm_provider": "deepgram", "metadata": { - "calculation": "$0.0058/60 seconds = $0.00009667 per second", + "calculation": "$0.0077/60 seconds = $0.00012833 per second", + "note": "Nova-3 Medical streaming, launch price from Deepgram's 2025-03-03 announcement. deepgram.com/pricing does not list Nova-3 Medical, and this matches its current Nova-3 monolingual streaming pay as you go rate", + "original_pricing_per_minute": 0.0077 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://www.businesswire.com/news/home/20250303470439/en", + "supported_endpoints": [ + "/v1/listen" + ] + }, + "deepgram/streaming/nova-3-multilingual": { + "input_cost_per_second": 0.00015333, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0092/60 seconds = $0.00015333 per second", "note": "Nova-3 multilingual (language=multi) streaming, pay as you go", - "original_pricing_per_minute": 0.0058 + "original_pricing_per_minute": 0.0092 }, "mode": "audio_transcription", "output_cost_per_second": 0.0, diff --git a/tests/unit/llms/deepgram/test_deepgram_common_utils.py b/tests/unit/llms/deepgram/test_deepgram_common_utils.py index 4c1b4603076..a1672e965bb 100644 --- a/tests/unit/llms/deepgram/test_deepgram_common_utils.py +++ b/tests/unit/llms/deepgram/test_deepgram_common_utils.py @@ -285,6 +285,9 @@ NOVA_2_URL: Final = "wss://api.deepgram.com/v1/listen?model=nova-2" [ pytest.param(NOVA_3_URL, (), True, id="streaming entry present"), pytest.param(f"{NOVA_3_URL}&language=multi", (), True, id="multilingual entry present"), + pytest.param( + "wss://api.deepgram.com/v1/listen?model=nova-3-medical", (), True, id="medical streaming entry present" + ), pytest.param(NOVA_2_URL, (), False, id="only the pre-recorded entry"), pytest.param(f"{NOVA_2_URL}&language=multi", ("deepgram/streaming/nova-2",), False, id="needs multilingual"), pytest.param("wss://api.deepgram.com/v1/listen?model=nova-99-unmapped", (), False, id="nothing priced"), @@ -295,7 +298,8 @@ NOVA_2_URL: Final = "wss://api.deepgram.com/v1/listen?model=nova-2" def test_deepgram_listen_is_priced( monkeypatch: pytest.MonkeyPatch, upstream_url: str, extra_rows: tuple[str, ...], expected: bool ): - """The bundled map prices only nova-3 for streaming; nova-2 has a pre-recorded row, which must never count.""" + """The bundled map prices nova-3 and nova-3-medical for streaming; nova-2 has only a pre-recorded row, which must + never count.""" monkeypatch.delitem(litellm.model_cost, "deepgram/streaming/nova-2", raising=False) assert "deepgram/nova-2" in litellm.model_cost for row in extra_rows: From 10b3deae0d058d4da0f71d5d4642082bc5102743 Mon Sep 17 00:00:00 2001 From: Daniel van Strien Date: Mon, 2 Nov 2026 22:24:02 +0000 Subject: [PATCH 19/79] build(deps): allow huggingface-hub 2.x (#49056) Widen the huggingface-hub cap from <2.0 to <3.0. The locked version stays at 1.32.0; only the matching specifier line in uv.lock changes. Claude-Session: https://claude.ai/code/session_016isSPojcgVpmxSRshXQEPw Co-authored-by: Claude Opus 5.5 Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5c23a64ebb3..7a653b28f81 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ "tiktoken>=0.12.0,<1.0; python_version == '3.14'", "tiktoken>=0.14.0,<1.0; python_version >= '3.15'", "tokenizers>=0.21.0,<1.0", - "huggingface-hub>=0.34.0,<2.0", + "huggingface-hub>=0.34.0,<3.0", "click>=8.0.0,<9.0", "jinja2>=3.1.6,<4.0", "aiohttp>=3.14.2,<4.0", diff --git a/uv.lock b/uv.lock index 19d3a1ffd4f..e5948530f17 100644 --- a/uv.lock +++ b/uv.lock @@ -5262,7 +5262,7 @@ requires-dist = [ { name = "httpx", extras = ["http2"], specifier = ">=0.28.0,<1.0" }, { name = "httpx2", marker = "extra == 'mcp'", specifier = ">=2.10.0,<3" }, { name = "httpx2", marker = "extra == 'proxy'", specifier = ">=2.10.0,<3" }, - { name = "huggingface-hub", specifier = ">=0.34.0,<2.0" }, + { name = "huggingface-hub", specifier = ">=0.34.0,<3.0" }, { name = "importlib-metadata", specifier = ">=8.0.0,<9.0" }, { name = "inquirerpy", marker = "extra == 'cli'", specifier = ">=0.3.4,<1.0" }, { name = "inquirerpy", marker = "extra == 'proxy'", specifier = ">=0.3.4,<1.0" }, From 523b87ec718cc733363ac0d66312730c51a23a27 Mon Sep 17 00:00:00 2001 From: Etienne Chabert Date: Mon, 2 Nov 2026 23:41:59 +0100 Subject: [PATCH 20/79] fix(proxy): charge budget windows when the cached key or team expired before the request ended (#47608) * fix(proxy): charge budget windows when the cached key or team expired before the request ended The post-call spend update reads a key's or team's budget_limits from user_api_key_cache and returned early on a miss, skipping both the window counter increment and the LiteLLM_BudgetWindowSpend row. The cached object expires after user_api_key_cache_ttl (60s by default), so any request that ran longer than that without another request refreshing the entry was counted in the key's lifetime spend but never in its budget windows. With budget reservation disabled nothing else charges the windows, so keys kept running past their window caps On a miss the update now reads budget_limits straight from the key or team row. It does not re-cache the object, since a write from the spend path could race a /key/delete and put a just-deleted key back into the auth cache * test(proxy): assert the budget window fallback never re-caches the key or team * fix(proxy): read a missing key or team row once per cache TTL in the budget window fallback The master key and UI session tokens are logged under names no key row carries, so every one of their requests missed the auth cache and read the database. A missing row is now remembered for the cache TTL in this process only. A failed read is not remembered, and a row that exists is read fresh on every miss so a key update is never served stale. * fix(proxy): skip the budget window row read for a key or team with no windows The fallback remembered only a missing row, so a key or team that exists but holds no budget_limits was read from the database on every request that outlived the key cache. It now remembers both cases under one cache key built in user_api_key_cache.py, for the same TTL as the auth cache * fix(proxy): keep no-window markers in their own cache and clear them when an update adds windows Markers for keys and teams with no budget windows lived in the shared 200-entry auth cache, so a burst of late costs could evict cached teams and users. They now live in a dedicated in-memory cache on UserApiKeyCache, reached by local and broadcast evictions. /key/update and /team/update now evict the marker on every worker, so a late cost after an update charges the windows the update added. * test(proxy): give the spend counter tests the proxy's real key cache and type the window spend helper * fix(proxy): keep a row read that an update outran from restoring the no-window marker * fix(proxy): clear finished no-window reads and evict a marker-shaped user id from both caches * fix(proxy): apply other spend counters before a window row read and let no-window notes follow the key cache settings --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/caching/dual_cache.py | 2 +- litellm/proxy/auth/auth_checks.py | 12 + .../auth_cache_invalidation_pubsub.py | 2 +- .../proxy/common_utils/user_api_key_cache.py | 22 ++ .../management_endpoints/team_endpoints.py | 4 + litellm/proxy/proxy_server.py | 257 +++++++++------ .../test_auth_cache_invalidation_pubsub.py | 29 +- .../common_utils/test_user_api_key_cache.py | 37 ++- .../test_key_management_endpoints.py | 13 +- .../test_team_endpoints.py | 8 +- tests/unit/proxy/test_budget_reservation.py | 2 +- ...test_proxy_server_endpoints_and_startup.py | 307 +++++++++++++++++- 12 files changed, 575 insertions(+), 120 deletions(-) diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index e6214a3650c..c82d25c9425 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -673,7 +673,7 @@ class DualCache(BaseCache): if self.redis_cache is not None: self.redis_cache.flush_cache() - def delete_cache(self, key): + def delete_cache(self, key: str) -> None: """ Delete a key from the cache """ diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index f5d3a80ee91..fdc8aa82268 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -125,6 +125,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( NO_TEAM_MEMBERSHIP_SENTINEL, TAG_REGISTRY_OVERFLOW_SENTINEL, UserApiKeyCache, + budget_windows_absent_cache_key, end_user_cache_key, end_user_restricted_registry_cache_key, get_management_object_ttl, @@ -3471,11 +3472,22 @@ async def delete_cache_key_object( ) await publish_auth_cache_invalidation(cache_key=key) + await evict_budget_windows_absent_marker( + entity_type=Litellm_EntityType.KEY, entity_id=hashed_token, user_api_key_cache=user_api_key_cache + ) _delete_cache_key_object: Final = delete_cache_key_object +async def evict_budget_windows_absent_marker( + entity_type: Litellm_EntityType, entity_id: str, user_api_key_cache: UserApiKeyCache +) -> None: + marker: Final = budget_windows_absent_cache_key(entity_type.value, entity_id) + evict_local(cache_keys=(marker,), user_api_key_cache=user_api_key_cache) + await publish_auth_cache_invalidation(cache_key=marker) + + async def delete_cache_key_objects( hashed_tokens: Sequence[str], user_api_key_cache: UserApiKeyCache, diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index b2640d2a765..edc8048ee73 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -183,7 +183,7 @@ def evict_local(cache_keys: Sequence[str], user_api_key_cache: "UserApiKeyCache" def _drop_from_memory(cache_key: str, user_api_key_cache: "UserApiKeyCache") -> None: - user_api_key_cache.in_memory_cache_for(cache_key).delete_cache(cache_key) + user_api_key_cache.delete_from_memory(cache_key) @with_service_target(AUTH_OBJECTS_TARGET) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index e64fd755ef9..a2272c3f73d 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -24,6 +24,7 @@ T = TypeVar("T", bound=BaseModel) from litellm.constants import AUTH_OBJECTS_TARGET as AUTH_OBJECTS_TARGET # noqa: E402, PLC0414 # public re-export _HASHED_TOKEN_CACHE_KEY: Final = re.compile(r"[0-9a-f]{64}") +_BUDGET_WINDOWS_ABSENT_PREFIX: Final = "budget_windows_absent:" def is_user_key_cache_key(key: str) -> bool: @@ -70,6 +71,7 @@ class UserApiKeyCache(DualCache): default_in_memory_ttl: float | None = None, default_redis_ttl: float | None = None, key_object_in_memory_cache: InMemoryCache | None = None, + budget_windows_absent_in_memory_cache: InMemoryCache | None = None, ) -> None: super().__init__( in_memory_cache=in_memory_cache, @@ -83,19 +85,33 @@ class UserApiKeyCache(DualCache): default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl, ) + self.budget_windows_absent_cache: Final = DualCache( + in_memory_cache=budget_windows_absent_in_memory_cache + or InMemoryCache(default_ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL), + default_in_memory_ttl=default_in_memory_ttl, + ) def in_memory_cache_for(self, key: str) -> InMemoryCache: return self.key_object_cache.in_memory_cache if is_user_key_cache_key(key) else self.in_memory_cache + def delete_from_memory(self, key: str) -> None: + self.in_memory_cache_for(key).delete_cache(key) + if key.startswith(_BUDGET_WINDOWS_ABSENT_PREFIX): + self.budget_windows_absent_cache.delete_cache(key) + def update_cache_ttl(self, default_in_memory_ttl: float | None, default_redis_ttl: float | None) -> None: super().update_cache_ttl(default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl) self.key_object_cache.update_cache_ttl( default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl ) + self.budget_windows_absent_cache.update_cache_ttl( + default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=None + ) def update_in_memory_max_size(self, max_size: int | None) -> None: super().update_in_memory_max_size(max_size) self.key_object_cache.update_in_memory_max_size(max_size) + self.budget_windows_absent_cache.update_in_memory_max_size(max_size) def attach_redis_cache( self, redis_cache: RedisCache | None = None, *, default_redis_ttl: float | None = None @@ -272,6 +288,7 @@ class UserApiKeyCache(DualCache): def flush_cache(self) -> None: super().flush_cache() self.key_object_cache.in_memory_cache.flush_cache() + self.budget_windows_absent_cache.flush_cache() async def async_set_cache_pipeline( self, cache_list: Sequence[tuple[str, object]], local_only: bool = False, **kwargs: object @@ -400,6 +417,11 @@ def team_membership_reservation_cache_key(user_id: str, team_id: str) -> str: NO_TEAM_MEMBERSHIP_SENTINEL: Final = "__no_team_membership__" +def budget_windows_absent_cache_key(entity_type: str, entity_id: str) -> str: + """Cache key marking a key or team that holds no budget windows, so the spend fallback skips its row read.""" + return f"{_BUDGET_WINDOWS_ABSENT_PREFIX}{entity_type}:{entity_id}" + + def get_management_object_ttl(cache: DualCache) -> float: """ In-memory TTL for management-object cache writes (keys, teams, users, budgets, ...). diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index b712ba42a6a..f505406fe0e 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -109,6 +109,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo can_org_access_model, delete_cache_key_objects, delete_cache_team_object, + evict_budget_windows_absent_marker, get_jwt_key_mapping_cache_keys_for_tokens, get_org_object, get_team_membership, @@ -612,6 +613,9 @@ async def refresh_cached_team( await publish_auth_cache_invalidation(cache_key=f"team_id:{team_row.team_id}") if team_table.team_alias is not None: await publish_auth_cache_invalidation(cache_key=f"team_alias:{team_table.team_alias}") + await evict_budget_windows_absent_marker( + entity_type=Litellm_EntityType.TEAM, entity_id=team_row.team_id, user_api_key_cache=user_api_key_cache + ) _refresh_cached_team: Final = refresh_cached_team diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 807f9a0860d..22e5bd6b3bd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -26,6 +26,7 @@ from collections.abc import ( Collection, Container, Iterable, + Iterator, Mapping, MutableMapping, Sequence, @@ -588,6 +589,7 @@ from litellm.proxy.common_utils.token_count_log import ( ) from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, + budget_windows_absent_cache_key, end_user_cache_key, get_management_object_ttl, model_access_group_cache_key, @@ -3636,6 +3638,52 @@ def _cached_budget_limits(cached: object) -> object: return _UNTYPED_VALUE.validate_python(json.loads(stored)) if isinstance(stored, str) else stored +def _has_budget_windows(budget_limits: object) -> bool: + return isinstance(budget_limits, list) and budget_limits != [] + + +def _budget_windows_absent_marker(absent_key: str) -> object: + return _UNTYPED_VALUE.validate_python(user_api_key_cache.budget_windows_absent_cache.get_cache(key=absent_key)) + + +def _finish_pending_read(absent_key: str, pending_read: str, remember_absent: bool) -> None: + if _budget_windows_absent_marker(absent_key) != pending_read: + return + if remember_absent: + user_api_key_cache.budget_windows_absent_cache.set_cache(key=absent_key, value=True) + return + user_api_key_cache.budget_windows_absent_cache.delete_cache(absent_key) + + +async def _stored_budget_limits(entity_type: Litellm_EntityType, entity_id: str) -> object: + """Read only: caching the row here could race a delete and put a just-deleted key back into the auth cache.""" + if prisma_client is None: + return None + absent_key: Final = budget_windows_absent_cache_key(entity_type.value, entity_id) + if _budget_windows_absent_marker(absent_key) is True: + return None + pending_read: Final = f"reading:{secrets.token_hex(8)}" + user_api_key_cache.budget_windows_absent_cache.set_cache(key=absent_key, value=pending_read) + try: + row: Final[SupportsModelDump | None] = ( + await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": entity_id}) + if entity_type == Litellm_EntityType.KEY + else await TeamRepository(prisma_client).table.find_unique(where={"team_id": entity_id}) + ) + except Exception as e: # noqa: BLE001 # spend tracking must never fail the cost callback + verbose_proxy_logger.warning( + "Unable to read budget windows for %s=%s; this request's window spend is not counted: %s", + entity_type.value, + entity_id, + e, + ) + _finish_pending_read(absent_key, pending_read, remember_absent=False) + return None + budget_limits: Final[object] = None if row is None else _cached_budget_limits(row) + _finish_pending_read(absent_key, pending_read, remember_absent=not _has_budget_windows(budget_limits)) + return budget_limits + + async def _increment_spend_counters_batched( token: str | None, team_id: str | None, @@ -3666,7 +3714,68 @@ async def _increment_spend_counters_batched( cost: Final[float] = response_cost - async def _key_scope(key_token: str) -> tuple[PendingSpendIncrement | BaseException, ...]: + async def _window_increment( + entity_type: Litellm_EntityType, entity_id: str, window: object + ) -> PendingSpendIncrement | None: + duration: Final = ( + window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) + ) + window_reset_at: Final = ( + window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) + ) + window_counter: Final = f"spend:{entity_type.value}:{entity_id}:window:{duration}" + window_start: Final = get_budget_window_start(window) + pending_window: Final = ( + await _prepare_window_spend_counter_increment( + counter_key=window_counter, + entity_type=entity_type.value.capitalize(), + entity_id=entity_id, + window_duration=duration, + window_start=window_start, + increment=cost, + ) + if window_counter not in reserved_counter_keys + else None + ) + await _enqueue_window_spend_row_update( + entity_type=entity_type, + entity_id=entity_id, + reset_at=window_reset_at, + window_duration=duration, + window_start=window_start, + increment=cost, + request_started_at=request_started_at, + ) + return pending_window + + async def _window_increments( + entity_type: Litellm_EntityType, entity_id: str, budget_limits: object + ) -> tuple[PendingSpendIncrement | BaseException, ...]: + if not isinstance(budget_limits, list): + return () + window_pending: Final = await asyncio.gather( + *(_window_increment(entity_type, entity_id, window) for window in budget_limits), return_exceptions=True + ) + return tuple(item for item in window_pending if item is not None) + + async def _cached_window_increments( + entity_type: Litellm_EntityType, entity_id: str, cache_key: str + ) -> tuple[_SpendScopeItem, ...]: + cached: Final[object] = _UNTYPED_VALUE.validate_python(await user_api_key_cache.async_get_cache(key=cache_key)) + if cached is None: + return (_UncachedBudgetWindows(entity_type=entity_type, entity_id=entity_id),) + return await _window_increments(entity_type, entity_id, _cached_budget_limits(cached)) + + async def _stored_window_increments( + uncached: _UncachedBudgetWindows, + ) -> tuple[PendingSpendIncrement | BaseException, ...]: + return await _window_increments( + uncached.entity_type, + uncached.entity_id, + await _stored_budget_limits(uncached.entity_type, uncached.entity_id), + ) + + async def _key_scope(key_token: str) -> tuple[_SpendScopeItem, ...]: # key_token arrives pre-hashed from metadata["user_api_key"] (auth flow # hashes raw "sk-..." keys before they reach the callback). The # startswith("sk-") check is a safety net matching update_cache — @@ -3689,52 +3798,9 @@ async def _increment_spend_counters_batched( ) ) - async def _key_window_increment(window: object) -> PendingSpendIncrement | None: - duration = ( - window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) - ) - key_window_reset_at = ( - window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) - ) - key_window_counter: Final = f"spend:key:{hashed_token}:window:{duration}" - key_window_start = get_budget_window_start(window) - pending_window: Final = ( - await _prepare_window_spend_counter_increment( - counter_key=key_window_counter, - entity_type="Key", - entity_id=hashed_token, - window_duration=duration, - window_start=key_window_start, - increment=cost, - ) - if key_window_counter not in reserved_counter_keys - else None - ) - await _enqueue_window_spend_row_update( - entity_type=Litellm_EntityType.KEY, - entity_id=hashed_token, - reset_at=key_window_reset_at, - window_duration=duration, - window_start=key_window_start, - increment=cost, - request_started_at=request_started_at, - ) - return pending_window + return key_pending + await _cached_window_increments(Litellm_EntityType.KEY, hashed_token, hashed_token) - key_obj: Final[object] = _UNTYPED_VALUE.validate_python( - await user_api_key_cache.async_get_cache(key=hashed_token) - ) - if key_obj is None: - return key_pending - key_budget_limits: Final = _cached_budget_limits(key_obj) - if not isinstance(key_budget_limits, list): - return key_pending - window_pending: Final = await asyncio.gather( - *(_key_window_increment(window) for window in key_budget_limits), return_exceptions=True - ) - return key_pending + tuple(item for item in window_pending if item is not None) - - async def _team_scope(scope_team_id: str) -> tuple[PendingSpendIncrement | BaseException, ...]: + async def _team_scope(scope_team_id: str) -> tuple[_SpendScopeItem, ...]: team_counter_key: Final = f"spend:team:{scope_team_id}" team_pending: Final[tuple[PendingSpendIncrement, ...]] = ( () @@ -3748,50 +3814,9 @@ async def _increment_spend_counters_batched( ) ) - async def _team_window_increment(window: object) -> PendingSpendIncrement | None: - duration = ( - window["budget_duration"] if isinstance(window, dict) else getattr(window, "budget_duration", None) - ) - team_window_reset_at = ( - window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None) - ) - team_window_counter: Final = f"spend:team:{scope_team_id}:window:{duration}" - team_window_start = get_budget_window_start(window) - pending_window: Final = ( - await _prepare_window_spend_counter_increment( - counter_key=team_window_counter, - entity_type="Team", - entity_id=scope_team_id, - window_duration=duration, - window_start=team_window_start, - increment=cost, - ) - if team_window_counter not in reserved_counter_keys - else None - ) - await _enqueue_window_spend_row_update( - entity_type=Litellm_EntityType.TEAM, - entity_id=scope_team_id, - reset_at=team_window_reset_at, - window_duration=duration, - window_start=team_window_start, - increment=cost, - request_started_at=request_started_at, - ) - return pending_window - - team_obj: Final[object] = _UNTYPED_VALUE.validate_python( - await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}") + return team_pending + await _cached_window_increments( + Litellm_EntityType.TEAM, scope_team_id, f"team_id:{scope_team_id}" ) - if team_obj is None: - return team_pending - team_budget_limits: Final = _cached_budget_limits(team_obj) - if not isinstance(team_budget_limits, list): - return team_pending - window_pending: Final = await asyncio.gather( - *(_team_window_increment(window) for window in team_budget_limits), return_exceptions=True - ) - return team_pending + tuple(item for item in window_pending if item is not None) async def _team_member_scope( scope_user_id: str, scope_team_id: str @@ -3862,22 +3887,30 @@ async def _increment_spend_counters_batched( # return_exceptions so a failing scope does not leave its siblings running # as orphaned tasks that race the caller's reservation-counter invalidation; # all scopes settle, then the first error propagates as before. - scope_results: Final = await asyncio.gather(*scope_coros, return_exceptions=True) - scope_errors: Final = tuple( - item - for scope in scope_results - for item in (scope if isinstance(scope, tuple) else (scope,)) - if isinstance(item, BaseException) + scope_items: Final = tuple(_flatten_spend_scopes(await asyncio.gather(*scope_coros, return_exceptions=True))) + uncached_windows: Final = tuple(item for item in scope_items if isinstance(item, _UncachedBudgetWindows)) + await _apply_spend_counter_increments( + pending=reservation_update.pending + + tuple(item for item in scope_items if isinstance(item, PendingSpendIncrement)) ) - pending: Final = tuple( - item - for scope in scope_results - if not isinstance(scope, BaseException) - for item in scope - if not isinstance(item, BaseException) + stored_window_items: Final = ( + tuple( + _flatten_spend_scopes( + await asyncio.gather( + *(_stored_window_increments(uncached) for uncached in uncached_windows), return_exceptions=True + ) + ) + ) + if uncached_windows + else () ) - await _apply_spend_counter_increments(pending=reservation_update.pending + pending) + stored_window_pending: Final = tuple( + item for item in stored_window_items if isinstance(item, PendingSpendIncrement) + ) + if stored_window_pending: + await _apply_spend_counter_increments(pending=stored_window_pending) stamp_budget_reservation_actual_cost(budget_reservation=budget_reservation, actual_cost=response_cost) + scope_errors: Final = tuple(item for item in scope_items + stored_window_items if isinstance(item, BaseException)) if scope_errors: raise scope_errors[0] @@ -3885,6 +3918,25 @@ async def _increment_spend_counters_batched( budget_reservation["finalized"] = True +@dataclass(frozen=True, slots=True) +class _UncachedBudgetWindows: + entity_type: Litellm_EntityType + entity_id: str + + +_SpendScopeItem: TypeAlias = PendingSpendIncrement | BaseException | _UncachedBudgetWindows + + +def _flatten_spend_scopes( + scope_results: Sequence[tuple[_SpendScopeItem, ...] | BaseException], +) -> Iterator[_SpendScopeItem]: + for scope in scope_results: + if isinstance(scope, BaseException): + yield scope + continue + yield from scope + + @dataclass(frozen=True, slots=True) class _ReservationCounterUpdate: """The reserved counters the direct increment must skip, and the adjustments that settle them on the actual @@ -14840,6 +14892,7 @@ from litellm.repositories.table_repositories import ( ) from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import VerificationTokenRepository from litellm.responses.litellm_completion_transformation.count_request import ( ChatCountBridgeError, chat_count_request, diff --git a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py index 694b817ab86..2dc98b5ca6b 100644 --- a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py +++ b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -3,6 +3,7 @@ import hashlib import json import time from collections.abc import Iterable +from typing import Final from unittest.mock import patch import pytest @@ -19,7 +20,11 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( evict_and_broadcast, publish_auth_cache_invalidation, ) -from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache +from litellm.proxy.common_utils.user_api_key_cache import ( + AUTH_OBJECTS_TARGET, + UserApiKeyCache, + budget_windows_absent_cache_key, +) class _RecordingRedisClient(Redis): @@ -243,6 +248,28 @@ async def test_subscriber_deletes_key_object_partition_entry_on_message() -> Non assert cache.get_cache(hashed_token, model_type=UserAPIKeyAuth) is None +@pytest.mark.asyncio +async def test_subscriber_evicts_a_user_whose_id_looks_like_a_no_window_marker_along_with_the_marker() -> None: + user_id: Final = budget_windows_absent_cache_key("key", "alice") + cache: Final = UserApiKeyCache() + cache.set_cache(user_id, {"user_id": user_id, "models": ["gpt-4o"]}) + cache.budget_windows_absent_cache.set_cache(key=user_id, value=True) + + pubsub: Final = _QueuePubSub(initial_messages=[_invalidation_message(user_id)]) + subscriber: Final = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])), + user_api_key_cache=cache, + ) + subscriber.start() + try: + assert await asyncio.wait_for(pubsub.drained.wait(), timeout=5) + finally: + await subscriber.stop() + + assert cache.get_cache(user_id) is None + assert cache.budget_windows_absent_cache.get_cache(key=user_id) is None + + @pytest.mark.asyncio async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message() -> None: """ diff --git a/tests/unit/proxy/common_utils/test_user_api_key_cache.py b/tests/unit/proxy/common_utils/test_user_api_key_cache.py index 09523cd1901..577e5e5a389 100644 --- a/tests/unit/proxy/common_utils/test_user_api_key_cache.py +++ b/tests/unit/proxy/common_utils/test_user_api_key_cache.py @@ -1,6 +1,6 @@ import hashlib import json -from typing import Any +from typing import Any, Final import pytest @@ -11,6 +11,7 @@ from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, + budget_windows_absent_cache_key, end_user_cache_key, get_management_object_ttl, is_user_key_cache_key, @@ -32,6 +33,17 @@ class CapturingInMemoryCache(InMemoryCache): super().set_cache(key, value, **kwargs) +class _ManualClock: + def __init__(self) -> None: + self.seconds = 1000.0 + + def __call__(self) -> float: + return self.seconds + + def advance(self, seconds: float) -> None: + self.seconds += seconds + + class FakeRedisCache(RedisCache): """ In-memory fake that enforces the UserApiKeyCache Redis payload contract. @@ -432,6 +444,29 @@ class TestUserKeyObjectPartition: assert first_key.token == tokens[0] assert cache.get_cache(end_user_cache_key("u0")) == {"user_id": "u0"} + def test_update_in_memory_max_size_applies_to_no_window_notes(self): + cache = UserApiKeyCache() + cache.update_in_memory_max_size(1) + + cache.budget_windows_absent_cache.set_cache(key=budget_windows_absent_cache_key("key", "first"), value=True) + cache.budget_windows_absent_cache.set_cache(key=budget_windows_absent_cache_key("key", "second"), value=True) + + assert cache.budget_windows_absent_cache.get_cache(key=budget_windows_absent_cache_key("key", "first")) is None + assert cache.budget_windows_absent_cache.get_cache(key=budget_windows_absent_cache_key("key", "second")) is True + + def test_update_cache_ttl_applies_to_no_window_notes(self): + clock: Final = _ManualClock() + cache: Final = UserApiKeyCache( + default_in_memory_ttl=60, budget_windows_absent_in_memory_cache=InMemoryCache(clock=clock) + ) + cache.update_cache_ttl(default_in_memory_ttl=5, default_redis_ttl=5) + note: Final = budget_windows_absent_cache_key("key", "no-windows") + + cache.budget_windows_absent_cache.set_cache(key=note, value=True) + clock.advance(6) + + assert cache.budget_windows_absent_cache.get_cache(key=note) is None + def test_update_in_memory_max_size_none_resets_key_object_partition_to_default(self): cache = UserApiKeyCache(key_object_in_memory_cache=InMemoryCache(max_size_in_memory=1)) cache.update_in_memory_max_size(None) diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 0198f582be7..b453f985512 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -6,7 +6,7 @@ from contextlib import ExitStack from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch from uuid import UUID, uuid4 import pytest @@ -51,7 +51,11 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, project_cache_key +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + budget_windows_absent_cache_key, + project_cache_key, +) from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker from litellm.proxy.management_endpoints.key_management_endpoints import ( KeyProjectBindingError, @@ -8911,7 +8915,10 @@ async def test_delete_cache_key_object_broadcasts_invalidation(monkeypatch): # Real, observable state: the cache object itself no longer holds the entry. assert real_user_api_key_cache.get_cache(key="hashed-broadcast-key") is None - mock_publish.assert_awaited_once_with(cache_key="hashed-broadcast-key") + assert mock_publish.await_args_list == [ + call(cache_key="hashed-broadcast-key"), + call(cache_key=budget_windows_absent_cache_key("key", "hashed-broadcast-key")), + ] @pytest.mark.asyncio diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index e33743bb149..2f4dad0ccd7 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -2684,7 +2684,7 @@ async def test_team_block_unblock_reach_warmed_auth_on_every_worker( from litellm.proxy._types import BlockTeamRequest from litellm.proxy.auth.auth_checks import get_team_object from litellm.proxy.auth.team_grants import team_model_aliases - from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, budget_windows_absent_cache_key from litellm.proxy.management_endpoints import team_endpoints team_id: Final = "team-kill-switch" @@ -2721,7 +2721,11 @@ async def test_team_block_unblock_reach_warmed_auth_on_every_worker( seen_by_auth: Final = await get_team_object(team_id=team_id, prisma_client=prisma_client, user_api_key_cache=cache) assert seen_by_auth.blocked is blocked assert team_model_aliases(seen_by_auth) == {"fast": "gpt-4o"} - assert redis.broadcast_keys == [f"team_id:{team_id}", "team_alias:kill-switch"] + assert redis.broadcast_keys == [ + f"team_id:{team_id}", + "team_alias:kill-switch", + budget_windows_absent_cache_key("team", team_id), + ] assert (response.blocked, response.litellm_model_table, response.object_permission) == (blocked, None, None) diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index 535c3da8c3a..6ce8d400daf 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -66,7 +66,7 @@ def spend_counter_state(): original_prisma_client = ps.prisma_client counter_cache = DualCache() - key_cache = DualCache() + key_cache = UserApiKeyCache() ps.spend_counter_cache = counter_cache ps.user_api_key_cache = key_cache ps.prisma_client = None diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 76d94628d81..d09008d62bc 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -1,5 +1,6 @@ import asyncio import contextlib +import dataclasses import importlib import json import logging @@ -10,9 +11,10 @@ import subprocess import sys import time import types -from collections.abc import Callable +from collections.abc import Callable, Iterator, Mapping, Sequence from datetime import datetime, timedelta, timezone from pathlib import Path +from types import MappingProxyType from typing import Final from unittest import mock from unittest.mock import AsyncMock, MagicMock, call, create_autospec, mock_open, patch @@ -48,6 +50,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.login_throttle import LoginThrottle from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendUpdateQueue from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash from litellm.proxy.proxy_server import app, aws_error_type_exception_handler, initialize, openai_exception_handler from litellm.utils import _invalidate_model_cost_lowercase_map @@ -14075,17 +14078,18 @@ def test_startup_is_silent_when_mock_testing_params_disabled(caplog): @contextlib.contextmanager -def _window_spend_enqueue_env(cached_objects: dict): +def _window_spend_enqueue_env( + cached_objects: Mapping[str, object], prisma_client: MagicMock | None = None +) -> Iterator[WindowSpendUpdateQueue]: """Point increment_spend_counters at throwaway caches and a real WindowSpendUpdateQueue, and hand back the queue to inspect.""" import litellm.proxy.proxy_server as ps from litellm.caching.dual_cache import DualCache - from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( - WindowSpendUpdateQueue, - ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache - user_api_key_cache = MagicMock() - user_api_key_cache.async_get_cache = AsyncMock(side_effect=lambda key, **_: cached_objects.get(key)) + user_api_key_cache: Final = UserApiKeyCache() + for key, value in cached_objects.items(): + user_api_key_cache.set_cache(key=key, value=value) queue = WindowSpendUpdateQueue() proxy_logging_obj = MagicMock() @@ -14099,7 +14103,7 @@ def _window_spend_enqueue_env(cached_objects: dict): ) ps.user_api_key_cache = user_api_key_cache ps.spend_counter_cache = DualCache() - ps.prisma_client = None + ps.prisma_client = prisma_client ps.proxy_logging_obj = proxy_logging_obj try: yield queue @@ -14287,6 +14291,293 @@ async def test_team_window_spend_row_carries_the_request_start_time(): assert enqueued[0]["started_at"] == "2026-08-10T12:00:00.500000" +_WINDOW_RESET_AT: Final = "2026-08-17T00:00:00+00:00" + + +def _window(budget_duration: str, max_budget: float) -> dict[str, object]: + return {"budget_duration": budget_duration, "max_budget": max_budget, "reset_at": _WINDOW_RESET_AT} + + +@dataclasses.dataclass(frozen=True, slots=True) +class _StoredRow: + budget_limits: list[dict[str, object]] | None + + +def _prisma_with_rows( + keys: Mapping[str, _StoredRow] = MappingProxyType({}), + teams: Mapping[str, _StoredRow] = MappingProxyType({}), +) -> MagicMock: + """Key and team tables that answer find_unique from the given rows.""" + client: Final = MagicMock() + client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=lambda where: keys.get(where["token"])) + client.db.litellm_teamtable.find_unique = AsyncMock(side_effect=lambda where: teams.get(where["team_id"])) + return client + + +def _charged_windows(enqueued: Sequence[Mapping[str, object]]) -> list[tuple[object, object, object]]: + return [(item["entity_id"], item["window_duration"], item["spend"]) for item in enqueued] + + +@pytest.mark.asyncio +async def test_key_window_is_charged_when_the_cached_key_expired_before_the_request_ended(): + """A stream that outlives user_api_key_cache_ttl finds no cached key when its + cost lands. Its spend still belongs to the key's budget windows.""" + import litellm.proxy.proxy_server as ps + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows(keys={"hashed-token": _StoredRow([_window("30d", 100.0)])}) + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) + enqueued: Final = await _drain(queue) + assert await ps.user_api_key_cache.async_get_cache(key="hashed-token") is None + + assert [item["entity_type"] for item in enqueued] == ["key"] + assert _charged_windows(enqueued) == [("hashed-token", "30d", pytest.approx(0.25))] + + +@pytest.mark.asyncio +async def test_team_window_is_charged_when_the_cached_team_expired_before_the_request_ended(): + import litellm.proxy.proxy_server as ps + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows(teams={"team-1": _StoredRow([_window("7d", 50.0)])}) + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + await increment_spend_counters(token=None, team_id="team-1", user_id=None, response_cost=1.5) + enqueued: Final = await _drain(queue) + assert await ps.user_api_key_cache.async_get_cache(key="team_id:team-1") is None + + assert [item["entity_type"] for item in enqueued] == ["team"] + assert _charged_windows(enqueued) == [("team-1", "7d", pytest.approx(1.5))] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "token, keys", + [ + pytest.param("litellm_proxy_master_key", {}, id="no-key-row"), + pytest.param("hashed-token", {"hashed-token": _StoredRow(budget_limits=None)}, id="key-row-without-windows"), + ], +) +async def test_a_token_without_budget_windows_reads_the_database_once_across_requests( + token: str, keys: Mapping[str, _StoredRow] +): + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows(keys=keys) + find_unique: Final = prisma_client.db.litellm_verificationtoken.find_unique + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + await increment_spend_counters(token=token, team_id=None, user_id=None, response_cost=0.25) + reads_after_first_request: Final = find_unique.await_count + for _ in range(3): + await increment_spend_counters(token=token, team_id=None, user_id=None, response_cost=0.25) + enqueued: Final = await _drain(queue) + + assert enqueued == [] + assert reads_after_first_request > 0 + assert find_unique.await_count == reads_after_first_request + + +@pytest.mark.asyncio +async def test_every_cache_miss_charges_the_windows_the_key_row_holds_at_that_moment(): + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows(keys={"hashed-token": _StoredRow([_window("30d", 100.0)])}) + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=_StoredRow([_window("7d", 5.0)]) + ) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.5) + enqueued: Final = await _drain(queue) + + assert sorted(_charged_windows(enqueued), key=lambda charge: str(charge[1])) == [ + ("hashed-token", "30d", pytest.approx(0.25)), + ("hashed-token", "7d", pytest.approx(0.5)), + ] + + +@pytest.mark.asyncio +async def test_a_failed_key_row_read_skips_that_request_and_the_next_one_reads_again(): + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows() + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + side_effect=ConnectionError("database is unreachable") + ) + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) + after_failure: Final = await _drain(queue) + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=_StoredRow([_window("7d", 5.0)]) + ) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.5) + after_recovery: Final = await _drain(queue) + + assert after_failure == [] + assert _charged_windows(after_recovery) == [("hashed-token", "7d", pytest.approx(0.5))] + + +@pytest.mark.asyncio +async def test_a_late_request_after_a_key_update_charges_windows_the_key_did_not_have_before(): + import litellm.proxy.proxy_server as ps + from litellm.proxy.auth.auth_checks import delete_cache_key_object + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows(keys={"hashed-token": _StoredRow(budget_limits=None)}) + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=_StoredRow([_window("7d", 5.0)]) + ) + await delete_cache_key_object( + hashed_token="hashed-token", user_api_key_cache=ps.user_api_key_cache, proxy_logging_obj=None + ) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.5) + enqueued: Final = await _drain(queue) + + assert _charged_windows(enqueued) == [("hashed-token", "7d", pytest.approx(0.5))] + + +@pytest.mark.asyncio +async def test_a_key_update_that_lands_during_the_row_read_does_not_hide_the_windows_it_adds(): + import litellm.proxy.proxy_server as ps + from litellm.proxy.auth.auth_checks import delete_cache_key_object + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows() + updated_row: Final = _StoredRow([_window("7d", 5.0)]) + + async def read_the_row_from_before_the_update(where: Mapping[str, str]) -> _StoredRow: + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=updated_row) + await delete_cache_key_object( + hashed_token=where["token"], user_api_key_cache=ps.user_api_key_cache, proxy_logging_obj=None + ) + return _StoredRow(budget_limits=None) + + prisma_client.db.litellm_verificationtoken.find_unique = read_the_row_from_before_the_update + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + ps.spend_counter_cache.set_cache(key="spend:key:hashed-token", value=0.0) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) + _ = await _drain(queue) + await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.5) + after_the_update: Final = await _drain(queue) + + assert _charged_windows(after_the_update) == [("hashed-token", "7d", pytest.approx(0.5))] + + +@pytest.mark.asyncio +async def test_late_requests_from_many_keys_with_windows_do_not_push_out_a_known_no_window_key(): + import litellm.proxy.proxy_server as ps + from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY + from litellm.proxy.proxy_server import increment_spend_counters + + windowed_tokens: Final = tuple(f"windowed-{index}" for index in range(DEFAULT_MAX_SIZE_IN_MEMORY)) + rows: Final = { + "no-windows": _StoredRow(budget_limits=None), + **{token: _StoredRow([_window("7d", 5.0)]) for token in windowed_tokens}, + } + prisma_client: Final = _prisma_with_rows(keys=rows) + find_unique: Final = prisma_client.db.litellm_verificationtoken.find_unique + + def no_window_key_reads() -> int: + return sum(1 for call in find_unique.await_args_list if call.kwargs["where"]["token"] == "no-windows") + + with _window_spend_enqueue_env({}, prisma_client=prisma_client): + await increment_spend_counters(token="no-windows", team_id=None, user_id=None, response_cost=0.1) + reads_after_the_first_request: Final = no_window_key_reads() + for token in windowed_tokens: + await increment_spend_counters(token=token, team_id=None, user_id=None, response_cost=0.1) + ps.spend_counter_cache.set_cache(key="spend:key:no-windows", value=0.1) + await increment_spend_counters(token="no-windows", team_id=None, user_id=None, response_cost=0.1) + + assert no_window_key_reads() == reads_after_the_first_request + + +@pytest.mark.asyncio +async def test_a_slow_window_row_read_does_not_hold_back_the_key_spend_counter(): + import litellm.proxy.proxy_server as ps + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows() + read_started: Final = asyncio.Event() + finish_read: Final = asyncio.Event() + + async def slow_row_read(where: Mapping[str, str]) -> _StoredRow: + read_started.set() + await finish_read.wait() + return _StoredRow([_window("7d", 5.0)]) + + prisma_client.db.litellm_verificationtoken.find_unique = slow_row_read + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + ps.spend_counter_cache.set_cache(key="spend:key:hashed-token", value=0.0) + ps.spend_counter_cache.set_cache(key="spend:key:hashed-token:window:7d", value=1.0) + charge: Final = asyncio.create_task( + increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25) + ) + assert await asyncio.wait_for(read_started.wait(), timeout=5) + key_spend_while_the_row_read_waits: Final = ps.spend_counter_cache.get_cache(key="spend:key:hashed-token") + finish_read.set() + await charge + window_spend: Final = ps.spend_counter_cache.get_cache(key="spend:key:hashed-token:window:7d") + charged: Final = await _drain(queue) + + assert key_spend_while_the_row_read_waits == pytest.approx(0.25) + assert window_spend == pytest.approx(1.25) + assert _charged_windows(charged) == [("hashed-token", "7d", pytest.approx(0.25))] + + +@pytest.mark.asyncio +async def test_a_late_request_that_misses_the_updated_team_charges_windows_the_team_did_not_have_before(): + """Another worker's /team/update only evicts this worker's cached team, so its next late cost reads the row.""" + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_team + from litellm.proxy.proxy_server import increment_spend_counters + + prisma_client: Final = _prisma_with_rows(teams={"team-1": _StoredRow(budget_limits=None)}) + updated_row: Final = _StoredRow([_window("7d", 50.0)]) + + with _window_spend_enqueue_env({}, prisma_client=prisma_client) as queue: + await increment_spend_counters(token=None, team_id="team-1", user_id=None, response_cost=1.5) + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=updated_row) + await refresh_cached_team( + team_row=LiteLLM_TeamTable(team_id="team-1", budget_limits=updated_row.budget_limits), + user_api_key_cache=ps.user_api_key_cache, + proxy_logging_obj=ps.proxy_logging_obj, + ) + ps.user_api_key_cache.delete_cache(key="team_id:team-1") + await increment_spend_counters(token=None, team_id="team-1", user_id=None, response_cost=2.0) + enqueued: Final = await _drain(queue) + + assert _charged_windows(enqueued) == [("team-1", "7d", pytest.approx(2.0))] + + +@pytest.mark.asyncio +async def test_late_costs_for_many_tokens_without_windows_leave_the_cached_team_in_place(): + from litellm.proxy.proxy_server import increment_spend_counters + + team_obj: Final = MagicMock() + team_obj.budget_limits = [_window("7d", 50.0)] + + with _window_spend_enqueue_env({"team_id:team-1": team_obj}, prisma_client=_prisma_with_rows()) as queue: + for token_number in range(250): + await increment_spend_counters( + token=f"token-{token_number}", team_id=None, user_id=None, response_cost=0.01 + ) + await increment_spend_counters(token=None, team_id="team-1", user_id=None, response_cost=1.5) + enqueued: Final = await _drain(queue) + + assert _charged_windows(enqueued) == [("team-1", "7d", pytest.approx(1.5))] + + def _mock_startup_prisma_client(health_check_error=None, connect_error=None): client = MagicMock() client.connect = AsyncMock(side_effect=connect_error) From 7fe6b1de34a5374b7283e5af29ad3dcc68b13f20 Mon Sep 17 00:00:00 2001 From: paveln83 <349044532+paveln83@users.noreply.github.com> Date: Mon, 2 Nov 2026 14:58:46 -0800 Subject: [PATCH 21/79] fix(proxy)!: ignore stored general_settings keys no settings API manages (#48967) Apply database general settings only for managed keys, warn about ignored copies, and allow admins to delete stale stored keys Move intended database-only settings into config.yaml before upgrading, including authentication and access restrictions --- .../proxy/config_resolvers/settings_store.py | 19 +- litellm/proxy/proxy_server.py | 108 ++++++--- .../configuration/test_effective_settings.py | 56 ++++- .../config_resolvers/test_settings_store.py | 31 +++ .../proxy/proxy_server/test_proxy_config.py | 134 ++++++++++-- ...test_proxy_server_endpoints_and_startup.py | 206 +++++++++++++++++- 6 files changed, 494 insertions(+), 60 deletions(-) diff --git a/litellm/proxy/config_resolvers/settings_store.py b/litellm/proxy/config_resolvers/settings_store.py index 70486be0068..3be66821915 100644 --- a/litellm/proxy/config_resolvers/settings_store.py +++ b/litellm/proxy/config_resolvers/settings_store.py @@ -42,10 +42,12 @@ _EMPTY_ROWS: Final[Mapping[DbRow, Mapping[str, JsonValue]]] = MappingProxyType({ class SettingsStore(MutableMapping[str, JsonValue]): - def __init__(self, section: Section) -> None: + def __init__(self, section: Section, *, managed_db_keys: frozenset[str] | None = None) -> None: self._section: Final = section + self._managed_db_keys: Final = managed_db_keys self._yaml_values: Mapping[str, JsonValue] = _EMPTY_VALUES self._database_rows: Mapping[DbRow, Mapping[str, JsonValue]] = _EMPTY_ROWS + self._unmanaged_db_keys: frozenset[str] = frozenset() self._runtime_values: Mapping[str, JsonValue] = _EMPTY_VALUES self._deleted_runtime_keys: frozenset[str] = frozenset() @@ -77,15 +79,24 @@ class SettingsStore(MutableMapping[str, JsonValue]): return self.owned_by_config(key) and self._db_value_is_shadowed(key) def apply_db_row(self, row: DbRow, db_row: Mapping[str, JsonValue]) -> None: + applied_row: Final = MappingProxyType({key: db_row[key] for key in db_row if self._db_manages(row, key)}) previous_row: Final = self._database_rows.get(row, _EMPTY_VALUES) changed: Final = frozenset( key - for key in (*previous_row, *db_row) - if previous_row.get(key, ABSENT) != db_row.get(key, ABSENT) # pyright: ignore[reportUnknownArgumentType] # JsonValue vs Absent compare + for key in (*previous_row, *applied_row) + if previous_row.get(key, ABSENT) != applied_row.get(key, ABSENT) # pyright: ignore[reportUnknownArgumentType] # JsonValue vs Absent compare ) - self._database_rows = MappingProxyType({**self._database_rows, row: MappingProxyType(dict(db_row))}) + self._database_rows = MappingProxyType({**self._database_rows, row: applied_row}) + if row == self._section: + self._unmanaged_db_keys = frozenset(key for key in db_row if key not in applied_row) self._clear_runtime_keys(changed) + def unmanaged_db_keys(self) -> tuple[str, ...]: + return tuple(sorted(self._unmanaged_db_keys)) + + def _db_manages(self, row: DbRow, key: str) -> bool: + return row != self._section or self._managed_db_keys is None or key in self._managed_db_keys + def resolved(self) -> Mapping[str, JsonValue]: return MappingProxyType(dict(self)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 22e5bd6b3bd..33dcf978854 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5722,6 +5722,14 @@ def _reload_settings_store(section: Section, store: SettingsStore, section_confi store["pass_through_endpoints"] = serving_pass_throughs +def _general_settings_row_is_readable(row: _ConfigParamValueRow | None) -> bool: + return row is None or row.param_value is None or isinstance(row.param_value, Mapping) + + +def _db_managed_general_settings_keys() -> frozenset[str]: + return frozenset((*ConfigGeneralSettings.model_fields, *_CLEANUP_SCHEDULE_KEYS, "allowed_ips", "email_settings")) + + def _bind_general_settings_store(settings: SettingsStore) -> None: global general_settings general_settings = settings # pyright: ignore[reportAssignmentType] # legacy global accepts mappings @@ -5872,7 +5880,9 @@ class ProxyConfig: # whether an existing request predates the prices it just fetched, and re-serving one # costs a single fetch where skipping one leaves it priced wrong indefinitely self.model_cost_map_applied_revision: int = 0 - self.settings: Final[SettingsStore] = SettingsStore("general_settings") + self.settings: Final[SettingsStore] = SettingsStore( + "general_settings", managed_db_keys=_db_managed_general_settings_keys() + ) self.router_settings: Final[SettingsStore] = SettingsStore("router_settings") self.litellm_settings: Final[SettingsStore] = SettingsStore("litellm_settings") self.environment_variables: Final[SettingsStore] = SettingsStore("environment_variables") @@ -5883,6 +5893,8 @@ class ProxyConfig: self._db_environment_keys: frozenset[str] = frozenset() self._sso_environment_removals: frozenset[str] = frozenset() self._warned_shadowed_keys: frozenset[tuple[Section, str]] = frozenset() + self._warned_unmanaged_keys: frozenset[str] = frozenset() + self._warned_unreadable_db_general_settings: bool = False self._db_config_callback_entries: Mapping[tuple[str, str], _DbCallbackRegistration] = MappingProxyType({}) self._settings_stores: Final[Mapping[Section, SettingsStore]] = MappingProxyType( { @@ -8406,15 +8418,20 @@ class ProxyConfig: general_settings.get("maximum_spend_logs_retention_interval", "1d"), ) - async def _update_general_settings(self, db_general_settings: Mapping[str, SettingsJsonValue] | None) -> None: + async def _update_general_settings(self, db_general_settings: Mapping[str, SettingsJsonValue] | str | None) -> None: global general_settings if db_general_settings is None: return + if not isinstance(db_general_settings, Mapping): + self._warn_about_unreadable_db_general_settings() + return + self._warned_unreadable_db_general_settings = False if not isinstance(general_settings, SettingsStore): self.settings.load_yaml(_as_settings_mapping(general_settings)) cache_size_was_db: Final = self.settings.source("user_api_key_cache_max_size") == "db" previous_cleanup_schedule: Final = self._resolved_cleanup_schedule() self.settings.apply_db_row("general_settings", db_general_settings) + self._warn_about_unmanaged_db_settings() _bind_general_settings_store(self.settings) await self._apply_general_settings_side_effects( db_general_settings, @@ -8632,6 +8649,7 @@ class ProxyConfig: ) self._warn_about_shadowed_db_settings() + self._warn_about_unmanaged_db_settings() return self._config_with_resolved_settings(config) async def _apply_db_environment_variables_row( @@ -8690,6 +8708,30 @@ class ProxyConfig: ) self._warned_shadowed_keys = shadowed + def _warn_about_unmanaged_db_settings(self) -> None: + unmanaged: Final = frozenset( + key for key in self.settings.unmanaged_db_keys() if key not in _GENERAL_SETTINGS_UI_LITELLM_FIELDS + ) + for key in sorted(unmanaged - self._warned_unmanaged_keys): + verbose_proxy_logger.warning( + "general_settings.%s is stored in the database, but the proxy does not apply that stored value, so it " + "is ignored. %s", + key, + "The config file's value applies. Remove the stored copy with POST /config/field/delete." + if self.settings.owned_by_config(key) + else "Set it in the config file to use it, or remove it with POST /config/field/delete.", + ) + self._warned_unmanaged_keys = unmanaged + + def _warn_about_unreadable_db_general_settings(self) -> None: + if self._warned_unreadable_db_general_settings: + return + verbose_proxy_logger.warning( + "general_settings is stored in the database as something other than a JSON object, so the proxy " + "ignores every stored general setting." + ) + self._warned_unreadable_db_general_settings = True + def _prepared_db_settings_values(self, section: Section, value: object) -> Mapping[str, SettingsJsonValue]: scrubbed: Final = _scrub_db_overlay_remote_module_loads(section=section, db_value=value) return _as_settings_mapping(scrubbed) @@ -8857,7 +8899,10 @@ class ProxyConfig: ) # initialize vector stores, guardrails, etc. table in db - await self._init_non_llm_objects_in_db(prisma_client=prisma_client) + await self._init_non_llm_objects_in_db( + prisma_client=prisma_client, + refresh_pass_through_endpoints=_general_settings_row_is_readable(db_general_settings), + ) except Exception as e: verbose_proxy_logger.exception("litellm.proxy.proxy_server.py::ProxyConfig:add_deployment - %s", e) @@ -8926,7 +8971,9 @@ class ProxyConfig: except Exception as e: # noqa: BLE001 # best-effort: a failing stop must not break proxy shutdown verbose_proxy_logger.error("Error stopping auth cache invalidation subscriber: %s", e) - async def _init_non_llm_objects_in_db(self, prisma_client: PrismaClient): + async def _init_non_llm_objects_in_db( + self, prisma_client: PrismaClient, *, refresh_pass_through_endpoints: bool = True + ): """ Use this to read non-llm objects from the db and initialize them @@ -8950,7 +8997,7 @@ class ProxyConfig: if self._should_load_db_object(object_type="agents"): await self._init_agents_in_db(prisma_client=prisma_client) - if self._should_load_db_object(object_type="pass_through_endpoints"): + if refresh_pass_through_endpoints and self._should_load_db_object(object_type="pass_through_endpoints"): await self._init_pass_through_endpoints_in_db() if self._should_load_db_object(object_type="prompts"): @@ -21111,53 +21158,56 @@ async def delete_config_general_settings( if data.field_name in _GENERAL_SETTINGS_UI_LITELLM_FIELDS: return await _reset_general_settings_ui_litellm_field(data.field_name, user_api_key_dict) - if data.field_name not in ConfigGeneralSettings.model_fields: + config_table: Final[TableActions[_ConfigParamValueRow]] = _config_param_table(prisma_client) + stored_row: Final = await config_table.find_first(where={"param_name": "general_settings"}) + stored: Final = ( + _SETTINGS_MAPPING.validate_python(stored_row.param_value) + if stored_row is not None and isinstance(stored_row.param_value, Mapping) + else None + ) + + if not _general_settings_row_is_readable(stored_row): + raise HTTPException( + status_code=400, + detail={ + "error": "general_settings is stored in the database as something other than a JSON object, so no " + "field can be removed from it" + }, + ) + + if data.field_name not in ConfigGeneralSettings.model_fields and (stored is None or data.field_name not in stored): raise HTTPException( status_code=400, detail={"error": f"Invalid field={data.field_name} passed in."}, ) - proxy_config.reject_config_owned_deletes(section_name="general_settings", keys=(data.field_name,)) + if data.field_name in _db_managed_general_settings_keys(): + proxy_config.reject_config_owned_deletes(section_name="general_settings", keys=(data.field_name,)) - ## get general settings from db - db_general_settings: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first( - where={"param_name": "general_settings"} - ) - ### pop the value - - if db_general_settings is None or db_general_settings.param_value is None: + if stored is None: raise HTTPException( status_code=400, detail={"error": f"Field name={data.field_name} not in config"}, ) - else: - general_settings = dict(db_general_settings.param_value) - - before_general_settings: Final = copy.deepcopy(general_settings) - - ## update db - - general_settings.pop(data.field_name, None) + remaining: Final = {key: value for key, value in stored.items() if key != data.field_name} response: Final = await ConfigRepository(prisma_client).table.upsert( where={"param_name": "general_settings"}, data={ "create": { "param_name": "general_settings", - "param_value": json.dumps(general_settings), + "param_value": json.dumps(remaining), }, - "update": {"param_value": json.dumps(general_settings)}, + "update": {"param_value": json.dumps(remaining)}, }, ) await invalidate_config_param("general_settings") - proxy_config.settings.apply_db_row("general_settings", general_settings) + proxy_config.settings.apply_db_row("general_settings", remaining) if is_resource_list("general_settings", data.field_name): - stored_endpoints: Final = general_settings.get("pass_through_endpoints") + stored_endpoints: Final = remaining.get("pass_through_endpoints") await proxy_config.serve_pass_through_endpoints(stored_endpoints if isinstance(stored_endpoints, list) else ()) asyncio.create_task( - create_config_audit_log( - "general_settings", "deleted", before_general_settings, general_settings, user_api_key_dict - ) + create_config_audit_log("general_settings", "deleted", dict(stored), remaining, user_api_key_dict) ) return response diff --git a/tests/integration/configuration/test_effective_settings.py b/tests/integration/configuration/test_effective_settings.py index 07477f1f048..1b0086c619a 100644 --- a/tests/integration/configuration/test_effective_settings.py +++ b/tests/integration/configuration/test_effective_settings.py @@ -20,7 +20,7 @@ from pydantic import BaseModel, JsonValue, TypeAdapter from tests.integration._support.client import Gateway, eventually, object_value, string_value from tests.integration._support.database import read_rows, scratch_database -from tests.integration._support.process import owned_proxy +from tests.integration._support.process import owned_proxy, owned_proxy_process from tests.integration._support.wire import Reply, Request, Wire, wire_server pytestmark: Final = pytest.mark.shares_worker_state @@ -1701,3 +1701,57 @@ def test_db_environment_variables_saved_from_a_session_outside_the_tables_schema assert isinstance(first, int) and isinstance(second, int) assert 0 < first < second + + +def test_a_stored_general_setting_no_settings_api_manages_stays_off_and_the_rest_of_the_row_applies( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + alias: Final = "stale-database-setting-" + uuid.uuid4().hex + with scratch_database() as database_url, wire_server(lambda _request: _completion_reply("pong")) as provider: + monkeypatch.setenv("DATABASE_URL", database_url) + _push_schema() + with psycopg.connect(database_url, autocommit=True) as writer: + writer.execute( + 'INSERT INTO "LiteLLM_Config" (param_name, param_value) VALUES (%s, %s)', + ("general_settings", Jsonb({"enforce_user_param": True, "store_prompts_in_spend_logs": True})), + ) + config: Final = tmp_path / "stale-database-setting.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": alias, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": provider.url + "/v1", + "api_key": "integration-provider-key", + }, + } + ], + "general_settings": {"master_key": gateway.key, "store_model_in_db": True}, + } + ) + ) + with owned_proxy_process( + gateway, + tmp_path / "stale-database-setting", + {"DATABASE_URL": database_url, "DISABLE_SCHEMA_UPDATE": "true"}, + config=config, + database_setup=(), + ) as owned: + answer: Final = owned.gateway.request( + "POST", "/v1/chat/completions", {"model": alias, "messages": [{"role": "user", "content": "ping"}]} + ) + listed: Final = owned.gateway.request("GET", "/config/list", params={"config_type": "general_settings"}) + log: Final = owned.log.read_text() + + assert answer.status_code == 200, answer.text + assert listed.status_code == 200, listed.text + stored_prompts: Final = next( + entry + for entry in TypeAdapter(list[dict[str, JsonValue]]).validate_json(listed.content) + if entry["field_name"] == "store_prompts_in_spend_logs" + ) + assert (stored_prompts["field_value"], stored_prompts["source"]) == (True, "db") + assert "general_settings.enforce_user_param is stored in the database, but the proxy does not apply" in log diff --git a/tests/unit/proxy/config_resolvers/test_settings_store.py b/tests/unit/proxy/config_resolvers/test_settings_store.py index 7b2cd404b46..76850825937 100644 --- a/tests/unit/proxy/config_resolvers/test_settings_store.py +++ b/tests/unit/proxy/config_resolvers/test_settings_store.py @@ -471,3 +471,34 @@ def test_settings_store_truthiness_matches_emptiness() -> None: assert bool(store) is True del store["max_parallel_requests"] assert bool(store) is False + + +def test_settings_store_ignores_stored_keys_outside_its_managed_db_keys_and_still_applies_them_from_yaml() -> None: + store: Final = SettingsStore("general_settings", managed_db_keys=frozenset({"max_parallel_requests"})) + stored_row: Final[dict[str, JsonValue]] = {"enforce_user_param": True, "max_parallel_requests": 5} + + store.apply_db_row("general_settings", stored_row) + + assert dict(store) == {"max_parallel_requests": 5} + assert store.source("enforce_user_param") == "unset" + assert store.unmanaged_db_keys() == ("enforce_user_param",) + + store.load_yaml({"enforce_user_param": True}) + store.apply_db_row("general_settings", stored_row) + + assert dict(store) == {"enforce_user_param": True, "max_parallel_requests": 5} + assert store.source("enforce_user_param") == "config" + assert store.unmanaged_db_keys() == ("enforce_user_param",) + + +def test_settings_store_managed_db_keys_leave_other_rows_and_unrestricted_stores_whole() -> None: + restricted: Final = SettingsStore("general_settings", managed_db_keys=frozenset()) + unrestricted: Final = SettingsStore("general_settings") + + restricted.apply_db_row("ui_settings", {"disable_custom_api_keys": True}) + unrestricted.apply_db_row("general_settings", {"enforce_user_param": True}) + + assert dict(restricted) == {"disable_custom_api_keys": True} + assert restricted.unmanaged_db_keys() == () + assert dict(unrestricted) == {"enforce_user_param": True} + assert unrestricted.unmanaged_db_keys() == () diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index d9ff7fef9f9..a78209e16fe 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -102,7 +102,9 @@ def test_proxy_config_upserts_existing_deployment_without_duplicating_it(monkeyp @pytest.mark.asyncio -async def test_proxy_config_preserves_models_when_config_read_fails(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +async def test_proxy_config_preserves_models_when_config_read_fails( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: from litellm.proxy import proxy_server router = litellm.Router(model_list=[_deployment("config-deployment").to_json(exclude_none=True)]) @@ -135,6 +137,8 @@ async def test_proxy_config_keeps_database_models_and_deletes_stale_models( assert desired == frozenset({"kept-deployment"}) assert router.get_model_ids() == ["kept-deployment"] + + from tests._master_key import MASTER_KEY @@ -150,7 +154,6 @@ async def test_proxy_config_loads_lens_store_from_yaml(tmp_path, monkeypatch) -> assert is_lens_tracing_enabled(settings["tracing"], {}) is True - @pytest.mark.asyncio @pytest.mark.parametrize("shutdown_error", [False, True]) async def test_tracing_config_automatically_exports_spend_without_a_storage_dependency( @@ -176,7 +179,9 @@ async def test_tracing_config_automatically_exports_spend_without_a_storage_depe outcome: Final = pytest.raises(RuntimeError, match="shutdown failure") if shutdown_error else nullcontext() with outcome: async with manage_tracing(enabled=True, client_factory=client): - logger: Final = next(callback for callback in litellm._async_success_callback if isinstance(callback, LensExporter)) + logger: Final = next( + callback for callback in litellm._async_success_callback if isinstance(callback, LensExporter) + ) await logger.async_log_success_event( {"standard_logging_object": {"id": "response-1", "response_cost": 0.25}}, None, None, None ) @@ -2693,7 +2698,6 @@ async def test_load_config_rejects_non_boolean_deny_by_default( await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) - @pytest.mark.asyncio @pytest.mark.parametrize( "yaml_value", ["files-api-2025-04-14", "[files-api-2025-04-14, 7]"], ids=["bare-string", "non-string-entry"] @@ -2713,6 +2717,7 @@ async def test_load_config_rejects_a_malformed_vertex_drop_anthropic_betas_list( with pytest.raises(ValidationError, match="vertex_ai_pass_through_drop_anthropic_betas"): await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) + @pytest.mark.asyncio async def test_ProxyConfig_load_config_resolves_router_settings_plugins(tmp_path, monkeypatch): """Regression: router_settings.plugins dotted-path strings must be resolved to @@ -4470,6 +4475,86 @@ async def test_ProxyConfig__update_general_settings_updates_health_check_retenti reschedule.assert_awaited_once() +@pytest.mark.asyncio +async def test_ProxyConfig__update_general_settings_ignores_a_stored_key_no_settings_api_manages_and_warns_once( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + pc: Final = ProxyConfig() + stored_row: Final = {"enforce_user_param": True, "store_prompts_in_spend_logs": True} + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await pc._update_general_settings(stored_row) + await pc._update_general_settings(stored_row) + from litellm.proxy import proxy_server + + assert proxy_server.general_settings["store_prompts_in_spend_logs"] is True + assert "enforce_user_param" not in proxy_server.general_settings + warnings: Final = [record.getMessage() for record in caplog.records if "enforce_user_param" in record.getMessage()] + assert warnings == [ + "general_settings.enforce_user_param is stored in the database, but the proxy does not apply that stored " + "value, so it is ignored. Set it in the config file to use it, or remove it with POST /config/field/delete." + ] + + +@pytest.mark.asyncio +async def test_ProxyConfig__update_general_settings_warns_about_an_ignored_stored_key_the_config_file_also_sets( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + pc: Final = ProxyConfig() + pc._load_yaml_settings_stores({"general_settings": {"enforce_user_param": False}}) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await pc._update_general_settings({"enforce_user_param": True}) + from litellm.proxy import proxy_server + + assert proxy_server.general_settings["enforce_user_param"] is False + warnings: Final = [record.getMessage() for record in caplog.records if "enforce_user_param" in record.getMessage()] + assert warnings == [ + "general_settings.enforce_user_param is stored in the database, but the proxy does not apply that stored " + "value, so it is ignored. The config file's value applies. Remove the stored copy with POST " + "/config/field/delete." + ] + + +@pytest.mark.asyncio +async def test_ProxyConfig__update_general_settings_does_not_warn_about_a_stored_copy_of_an_admin_ui_field( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + pc: Final = ProxyConfig() + pc._load_yaml_settings_stores({"general_settings": {}}) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await pc._update_general_settings({"budget_rollover": True, "enforce_user_param": True}) + + warned: Final = [record.getMessage().split(" ")[0] for record in caplog.records if "is ignored" in record.getMessage()] + assert warned == ["general_settings.enforce_user_param"] + + +@pytest.mark.asyncio +async def test_ProxyConfig__update_general_settings_warns_once_each_time_the_stored_row_stops_being_an_object( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + pc: Final = ProxyConfig() + pc._load_yaml_settings_stores({"general_settings": {}}) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await pc._update_general_settings(None) + await pc._update_general_settings('{"store_prompts_in_spend_logs": true}') + await pc._update_general_settings('{"store_prompts_in_spend_logs": true}') + await pc._update_general_settings({"store_prompts_in_spend_logs": True}) + await pc._update_general_settings('{"store_prompts_in_spend_logs": true}') + + warnings: Final = [record.getMessage() for record in caplog.records if "JSON object" in record.getMessage()] + assert warnings == [ + "general_settings is stored in the database as something other than a JSON object, so the proxy ignores " + "every stored general setting." + ] * 2 + + def _paused_scheduler(monkeypatch): from apscheduler.schedulers.asyncio import AsyncIOScheduler @@ -4918,9 +5003,6 @@ async def test_ProxyConfig__update_general_settings_none_input_noop(): # None input returns early. result = await pc._update_general_settings(db_general_settings=None) assert result is None - # Error-style: dict() will fail on non-mapping non-None input. - with pytest.raises(TypeError): - await pc._update_general_settings(db_general_settings=12345) # type: ignore[arg-type] @pytest.mark.asyncio @@ -5080,7 +5162,9 @@ async def test_ProxyConfig__update_config_from_db_keeps_keys_the_config_file_omi @pytest.mark.asyncio @pytest.mark.parametrize("ui_settings_already_synced", [False, True]) -async def test_ProxyConfig_add_deployment_continues_after_null_pass_through_endpoints(monkeypatch, ui_settings_already_synced): +async def test_ProxyConfig_add_deployment_continues_after_null_pass_through_endpoints( + monkeypatch: pytest.MonkeyPatch, ui_settings_already_synced: bool +) -> None: from litellm.proxy import proxy_server pc = ProxyConfig() @@ -5106,6 +5190,32 @@ async def test_ProxyConfig_add_deployment_continues_after_null_pass_through_endp assert settings_refresh.await_count == (0 if ui_settings_already_synced else 1) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "stored_row, rebuilds_routes", + [({"pass_through_endpoints": []}, True), ('{"pass_through_endpoints": []}', False), ([], False)], + ids=["object", "json-string", "array"], +) +async def test_ProxyConfig_add_deployment_rebuilds_pass_through_routes_only_from_an_object_general_settings_row( + monkeypatch: pytest.MonkeyPatch, stored_row: object, rebuilds_routes: bool +) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.pass_through_endpoints import pass_through_endpoints + + pc: Final = ProxyConfig() + rebuild_routes: Final = AsyncMock() + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "prefetch_config_params", AsyncMock()) + monkeypatch.setattr(proxy_server, "get_config_param", AsyncMock(return_value=SimpleNamespace(param_value=stored_row))) + monkeypatch.setattr(pass_through_endpoints, "initialize_pass_through_endpoints_in_db", rebuild_routes) + monkeypatch.setattr(pc, "_should_load_db_object", lambda *, object_type: object_type == "pass_through_endpoints") + monkeypatch.setattr(pc, "get_credentials", AsyncMock()) + + await pc.add_deployment(prisma_client=MagicMock(), proxy_logging_obj=MagicMock(), ui_settings_already_synced=True) + + assert rebuild_routes.await_count == (1 if rebuilds_routes else 0) + + # --------------------------------------------------------------------------- # ProxyConfig._update_config_from_db # --------------------------------------------------------------------------- @@ -5285,9 +5395,7 @@ def revision_test_environment(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv(key) -def _environment_variables_row( - values: Mapping[str, str | None] | None, revision: int | None -) -> SimpleNamespace | None: +def _environment_variables_row(values: Mapping[str, str | None] | None, revision: int | None) -> SimpleNamespace | None: if values is None: return None return SimpleNamespace( @@ -5473,9 +5581,7 @@ async def test_environment_variables_row_read_before_an_overlapping_newer_sync_i await newer_sync_finished.wait() older_sync: Final = asyncio.create_task( - _sync_from( - pc, _config_database(read_older_row, AsyncMock(return_value=older_row), wait_for_the_newer_sync) - ) + _sync_from(pc, _config_database(read_older_row, AsyncMock(return_value=older_row), wait_for_the_newer_sync)) ) await older_environment_row_read.wait() await _sync_from(pc, _rows_database(newer_row, newer_row)) diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index d09008d62bc..487638d6919 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -1807,9 +1807,7 @@ class _DatabaseWithOneStoredCredential: @pytest.mark.asyncio -@pytest.mark.parametrize( - "encrypted_with, asks_to_migrate", [(PUBLICLY_KNOWN_KEY, True), ("sk-some-other-key", False)] -) +@pytest.mark.parametrize("encrypted_with, asks_to_migrate", [(PUBLICLY_KNOWN_KEY, True), ("sk-some-other-key", False)]) async def test_proxy_startup_asks_to_migrate_only_when_the_database_holds_values_under_the_unsafe_key( monkeypatch, tmp_path, encrypted_with, asks_to_migrate ): @@ -1927,9 +1925,7 @@ async def test_proxy_startup_names_the_config_file_that_set_the_unsafe_key(monke from litellm.proxy.auth.master_key_boot_check import UnsafeMasterKeyError from litellm.proxy.proxy_server import proxy_startup_event - config_path, announced = _boot_with_general_settings( - monkeypatch, tmp_path, {"master_key": PUBLICLY_KNOWN_KEY} - ) + config_path, announced = _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": PUBLICLY_KNOWN_KEY}) with pytest.raises(UnsafeMasterKeyError): async with proxy_startup_event(FastAPI()): @@ -3810,6 +3806,24 @@ async def test_db_reload_keeps_the_resolved_value_of_a_config_owned_env_referenc assert proxy_server_module.store_model_in_db is True +@pytest.mark.asyncio +async def test_db_reload_of_a_json_string_general_settings_row_reports_no_ignored_keys( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server as proxy_server_module + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setattr(proxy_server_module, "user_api_key_cache", MagicMock(), raising=False) + proxy_config: Final = ProxyConfig() + proxy_config.settings.load_yaml({}) + monkeypatch.setattr(proxy_server_module, "general_settings", proxy_config.settings, raising=False) + + await proxy_config._update_general_settings('{"budget_rollover": true, "master_key": "sk-stored"}') + + assert proxy_config.settings.unmanaged_db_keys() == () + assert proxy_config.settings.source("master_key") != "db" + + def test_max_ui_session_budget_default_is_one_dollar(): """LIT-4662: the dashboard session budget default is a product decision; the old 0.25 default locked admins out of auto router Test Connection and the @@ -7298,7 +7312,7 @@ def test_deep_merge_dicts_skips_none_and_empty_lists(monkeypatch): db_param_value = { "max_parallel_requests": None, "allowed_models": [], - "new_key": "new_value", + "store_prompts_in_spend_logs": True, "nested": { "key1": "updated_value1", "key3": "value3", @@ -7311,7 +7325,7 @@ def test_deep_merge_dicts_skips_none_and_empty_lists(monkeypatch): assert result["max_parallel_requests"] == 10 assert result["allowed_models"] == ["gpt-3.5-turbo", "gpt-4"] - assert result["new_key"] == "new_value" + assert result["store_prompts_in_spend_logs"] is True assert result["nested"] == {"key1": "value1", "key2": "value2"} @@ -8397,10 +8411,18 @@ async def test_update_general_settings_db_pass_through_endpoint_overrides_yaml_e request.headers = {} request.query_params = {} - settings: Final = patch("litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]}) # test-quality-ok: the method reads this module global; no injection seam - yaml_endpoints: Final = patch("litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint]) # test-quality-ok: module global holding the YAML endpoints the fix merges in - initialize: Final = patch("litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock()) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here - master_key: Final = patch("litellm.proxy.proxy_server.master_key", MASTER_KEY) # test-quality-ok: a set master key is what makes a missing Authorization header a 401 + settings: Final = patch( + "litellm.proxy.proxy_server.general_settings", {"pass_through_endpoints": [yaml_endpoint]} + ) # test-quality-ok: the method reads this module global; no injection seam + yaml_endpoints: Final = patch( + "litellm.proxy.proxy_server.config_passthrough_endpoints", [yaml_endpoint] + ) # test-quality-ok: module global holding the YAML endpoints the fix merges in + initialize: Final = patch( + "litellm.proxy.proxy_server.initialize_pass_through_endpoints", AsyncMock() + ) # test-quality-ok: route registration needs the FastAPI app; auth is the observable here + master_key: Final = patch( + "litellm.proxy.proxy_server.master_key", MASTER_KEY + ) # test-quality-ok: a set master key is what makes a missing Authorization header a 401 with settings, yaml_endpoints, initialize, master_key: await ProxyConfig()._update_general_settings(db_general_settings={"pass_through_endpoints": [db_endpoint]}) @@ -16135,6 +16157,166 @@ async def test_delete_config_general_settings_still_removes_a_key_the_database_o assert "max_request_size_mb" not in pc.settings +@pytest.mark.asyncio +async def test_delete_config_general_settings_removes_a_stored_key_no_settings_api_manages( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy._types import ConfigFieldDelete + from litellm.proxy.proxy_server import ProxyConfig, delete_config_general_settings + + pc: Final = ProxyConfig() + pc._load_yaml_settings_stores({"general_settings": {}}) + stored_row: Final = {"enforce_user_param": True, "store_prompts_in_spend_logs": True} + pc.settings.apply_db_row("general_settings", stored_row) + fake_prisma: Final = _fake_prisma_with_config(stored_row) + monkeypatch.setattr(proxy_server_module, "proxy_config", pc) + monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) + + admin: Final = UserAPIKeyAuth(api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + await delete_config_general_settings( + data=ConfigFieldDelete(field_name="enforce_user_param", config_type="general_settings"), + user_api_key_dict=admin, + ) + + written: Final = json.loads(fake_prisma.db.litellm_config.upsert.await_args.kwargs["data"]["update"]["param_value"]) + assert written == {"store_prompts_in_spend_logs": True} + assert pc.settings.unmanaged_db_keys() == () + + +@pytest.mark.asyncio +async def test_delete_config_general_settings_stops_serving_the_removed_pass_through_routes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastapi import FastAPI + from starlette.routing import Route + + from litellm.proxy._types import ConfigFieldDelete + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import _registered_pass_through_routes + from litellm.proxy.proxy_server import ProxyConfig, delete_config_general_settings + from litellm.types.passthrough_endpoints.pass_through_endpoints import LITELLM_PASS_THROUGH_ENDPOINT_MARKER + + endpoint: Final = {"id": "stored-route", "path": "/stored-route", "target": "http://example.com"} + stored_row: Final = {"pass_through_endpoints": [endpoint], "store_prompts_in_spend_logs": True} + app: Final = FastAPI() + monkeypatch.setattr(proxy_server_module, "app", app) + monkeypatch.setattr(proxy_server_module, "premium_user", True) + monkeypatch.setattr(proxy_server_module, "config_passthrough_endpoints", None) + pc: Final = ProxyConfig() + pc._load_yaml_settings_stores({"general_settings": {}}) + pc.settings.apply_db_row("general_settings", stored_row) + monkeypatch.setattr(proxy_server_module, "proxy_config", pc) + monkeypatch.setattr(proxy_server_module, "prisma_client", _fake_prisma_with_config(stored_row)) + + def served_paths() -> frozenset[str]: + return frozenset( + route.path + for route in app.router.routes + if isinstance(route, Route) and getattr(route.endpoint, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, False) is True + ) + + saved_registry: Final = dict(_registered_pass_through_routes) + try: + _registered_pass_through_routes.clear() + await pc.serve_pass_through_endpoints([endpoint]) + served_before: Final = served_paths() + await delete_config_general_settings( + data=ConfigFieldDelete(field_name="pass_through_endpoints", config_type="general_settings"), + user_api_key_dict=UserAPIKeyAuth( + api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN + ), + ) + served_after: Final = served_paths() + finally: + _registered_pass_through_routes.clear() + _registered_pass_through_routes.update(saved_registry) + + assert "/stored-route" in served_before + assert "/stored-route" not in served_after + + +@pytest.mark.asyncio +async def test_delete_config_general_settings_removes_an_ignored_stored_key_the_config_file_also_sets( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy._types import ConfigFieldDelete + from litellm.proxy.proxy_server import ProxyConfig, delete_config_general_settings + + pc: Final = ProxyConfig() + pc._load_yaml_settings_stores({"general_settings": {"enforce_user_param": False}}) + stored_row: Final = {"enforce_user_param": True, "store_prompts_in_spend_logs": True} + pc.settings.apply_db_row("general_settings", stored_row) + fake_prisma: Final = _fake_prisma_with_config(stored_row) + monkeypatch.setattr(proxy_server_module, "proxy_config", pc) + monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) + + admin: Final = UserAPIKeyAuth(api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + await delete_config_general_settings( + data=ConfigFieldDelete(field_name="enforce_user_param", config_type="general_settings"), + user_api_key_dict=admin, + ) + + written: Final = json.loads(fake_prisma.db.litellm_config.upsert.await_args.kwargs["data"]["update"]["param_value"]) + assert written == {"store_prompts_in_spend_logs": True} + assert pc.settings["enforce_user_param"] is False + assert pc.settings.unmanaged_db_keys() == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stored_row", ['{"max_request_size_mb": 42, "disable_custom_api_keys": true}', "not json"]) +@pytest.mark.parametrize("field_name", ["max_request_size_mb", "disable_custom_api_keys", "not_a_setting"]) +async def test_delete_config_general_settings_leaves_a_string_row_untouched( + monkeypatch: pytest.MonkeyPatch, stored_row: str, field_name: str +) -> None: + from litellm.proxy._types import ConfigFieldDelete + from litellm.proxy.proxy_server import ProxyConfig, delete_config_general_settings + + pc: Final = ProxyConfig() + pc._load_yaml_settings_stores({"general_settings": {}}) + fake_prisma: Final = _fake_prisma_with_config(stored_row) + monkeypatch.setattr(proxy_server_module, "proxy_config", pc) + monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) + + admin: Final = UserAPIKeyAuth(api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + with pytest.raises(HTTPException) as refused: + await delete_config_general_settings( + data=ConfigFieldDelete(field_name=field_name, config_type="general_settings"), + user_api_key_dict=admin, + ) + + assert (refused.value.status_code, refused.value.detail) == ( + 400, + { + "error": "general_settings is stored in the database as something other than a JSON object, so no field " + "can be removed from it" + }, + ) + fake_prisma.db.litellm_config.upsert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delete_config_general_settings_refuses_a_key_neither_managed_nor_stored( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy._types import ConfigFieldDelete + from litellm.proxy.proxy_server import ProxyConfig, delete_config_general_settings + + pc: Final = ProxyConfig() + pc._load_yaml_settings_stores({"general_settings": {}}) + fake_prisma: Final = _fake_prisma_with_config({"store_prompts_in_spend_logs": True}) + monkeypatch.setattr(proxy_server_module, "proxy_config", pc) + monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) + + admin: Final = UserAPIKeyAuth(api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + with pytest.raises(HTTPException) as refused: + await delete_config_general_settings( + data=ConfigFieldDelete(field_name="enforce_user_param", config_type="general_settings"), + user_api_key_dict=admin, + ) + + assert refused.value.status_code == 400 + fake_prisma.db.litellm_config.upsert.assert_not_awaited() + + @pytest.mark.asyncio async def test_config_field_info_reports_the_declared_value_of_a_config_owned_secret(monkeypatch): from litellm.proxy.proxy_server import ProxyConfig, get_config_general_settings From a2d27881e87145058d652a481db1862acc4d5b94 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 2 Nov 2026 15:08:10 -0800 Subject: [PATCH 22/79] fix(bedrock): reserve DeepSeek V3.2 output capacity and date GovCloud West Llama 3 (#49732) * fix(bedrock): set the regional DeepSeek V3.2 context to the model card 164K and date GovCloud West Llama 3 * fix(bedrock): reserve DeepSeek V3.2 output capacity and date GovCloud West Llama 3 --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 38 ++++++++++--------- model_prices_and_context_window.json | 38 ++++++++++--------- 2 files changed, 40 insertions(+), 36 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 517aead5d29..7a430148a53 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -13835,7 +13835,7 @@ "input_cost_per_token_flex": 3.7e-07, "input_cost_per_token_priority": 1.295e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, @@ -13849,7 +13849,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/ap-northeast-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -14006,7 +14006,7 @@ "input_cost_per_token_flex": 3.7e-07, "input_cost_per_token_priority": 1.295e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, @@ -14020,7 +14020,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/ap-south-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -14144,12 +14144,12 @@ "input_cost_per_token": 7.4e-07, "input_cost_per_token_flex": 3.7e-07, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, "output_cost_per_token_flex": 1.11e-06, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.html", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md", "supports_audio_input": false, "supports_function_calling": true, "supports_native_structured_output": true, @@ -14254,7 +14254,7 @@ "input_cost_per_token_flex": 3.7e-07, "input_cost_per_token_priority": 1.295e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, @@ -14268,7 +14268,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/eu-north-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -14790,7 +14790,7 @@ "input_cost_per_token_flex": 3.7e-07, "input_cost_per_token_priority": 1.295e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, @@ -14804,7 +14804,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/sa-east-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -15039,7 +15039,7 @@ "input_cost_per_token_flex": 3.1e-07, "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, @@ -15053,7 +15053,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/us-east-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -15160,7 +15160,7 @@ "input_cost_per_token_flex": 3.1e-07, "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, @@ -15174,7 +15174,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/us-east-2/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -15556,6 +15556,7 @@ "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md" }, "bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0": { + "deprecation_date": "2027-03-30", "input_cost_per_token": 2.65e-06, "litellm_provider": "bedrock", "max_input_tokens": 8000, @@ -15566,6 +15567,7 @@ "supports_pdf_input": true }, "bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0": { + "deprecation_date": "2027-03-30", "input_cost_per_token": 3e-07, "litellm_provider": "bedrock", "max_input_tokens": 8000, @@ -15712,7 +15714,7 @@ "input_cost_per_token_flex": 3.1e-07, "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, @@ -15726,7 +15728,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/us-west-2/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -24538,7 +24540,7 @@ "input_cost_per_token_flex": 3.1e-07, "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock_converse", - "max_input_tokens": 164000, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, @@ -24549,7 +24551,7 @@ "supports_reasoning": true, "supports_native_structured_output": true, "supports_tool_choice": true, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 517aead5d29..7a430148a53 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -13835,7 +13835,7 @@ "input_cost_per_token_flex": 3.7e-07, "input_cost_per_token_priority": 1.295e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, @@ -13849,7 +13849,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/ap-northeast-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -14006,7 +14006,7 @@ "input_cost_per_token_flex": 3.7e-07, "input_cost_per_token_priority": 1.295e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, @@ -14020,7 +14020,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/ap-south-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -14144,12 +14144,12 @@ "input_cost_per_token": 7.4e-07, "input_cost_per_token_flex": 3.7e-07, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, "output_cost_per_token_flex": 1.11e-06, - "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.html", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md", "supports_audio_input": false, "supports_function_calling": true, "supports_native_structured_output": true, @@ -14254,7 +14254,7 @@ "input_cost_per_token_flex": 3.7e-07, "input_cost_per_token_priority": 1.295e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, @@ -14268,7 +14268,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/eu-north-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -14790,7 +14790,7 @@ "input_cost_per_token_flex": 3.7e-07, "input_cost_per_token_priority": 1.295e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 2.22e-06, @@ -14804,7 +14804,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/sa-east-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -15039,7 +15039,7 @@ "input_cost_per_token_flex": 3.1e-07, "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, @@ -15053,7 +15053,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/us-east-1/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -15160,7 +15160,7 @@ "input_cost_per_token_flex": 3.1e-07, "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, @@ -15174,7 +15174,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/us-east-2/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -15556,6 +15556,7 @@ "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-sonnet-4-5.md" }, "bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0": { + "deprecation_date": "2027-03-30", "input_cost_per_token": 2.65e-06, "litellm_provider": "bedrock", "max_input_tokens": 8000, @@ -15566,6 +15567,7 @@ "supports_pdf_input": true }, "bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0": { + "deprecation_date": "2027-03-30", "input_cost_per_token": 3e-07, "litellm_provider": "bedrock", "max_input_tokens": 8000, @@ -15712,7 +15714,7 @@ "input_cost_per_token_flex": 3.1e-07, "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock", - "max_input_tokens": 163840, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, @@ -15726,7 +15728,7 @@ "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/index.json" + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md" }, "bedrock/us-west-2/minimax.minimax-m2.1": { "deprecation_date": "2027-03-30", @@ -24538,7 +24540,7 @@ "input_cost_per_token_flex": 3.1e-07, "input_cost_per_token_priority": 1.085e-06, "litellm_provider": "bedrock_converse", - "max_input_tokens": 164000, + "max_input_tokens": 156000, "max_output_tokens": 8000, "mode": "chat", "output_cost_per_token": 1.85e-06, @@ -24549,7 +24551,7 @@ "supports_reasoning": true, "supports_native_structured_output": true, "supports_tool_choice": true, - "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-deepseek-deepseek-v3-2.md", "supports_audio_input": false, "supports_response_schema": true, "supports_vision": false From 4e268c2c552eb26bbbfe37fd2c11342c94148073 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 2 Nov 2026 23:40:32 +0000 Subject: [PATCH 23/79] test(gemini): assert vertex lists on content chunks (#49734) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/providers/test_gemini_long_decode_wire.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/integration/providers/test_gemini_long_decode_wire.py b/tests/integration/providers/test_gemini_long_decode_wire.py index 5ecb5e2a942..90e7616e26c 100644 --- a/tests/integration/providers/test_gemini_long_decode_wire.py +++ b/tests/integration/providers/test_gemini_long_decode_wire.py @@ -380,7 +380,9 @@ def test_cut_chat_stream_hands_its_token_on_a_chunk_with_empty_vertex_ai_lists_a chunks: Final = payloads(post(gateway, "chat", request_body("chat", model, True)).text, True) carriers: Final = [chunk for chunk in chunks if "continuation_token" in json.dumps(chunk)] assert [_vertex_lists(chunk) for chunk in carriers] == [dict(_EMPTY_VERTEX_LISTS)], chunks - assert [_vertex_lists(chunk) for chunk in chunks if chunk not in carriers] == [{}] * (len(chunks) - 1), chunks + assert [_vertex_lists(chunk) for chunk in chunks] == [ + dict(_EMPTY_VERTEX_LISTS) if at(chunk, "choices", 0, "delta", "content") else {} for chunk in chunks + ], chunks metered: Final = [chunk for chunk in chunks if isinstance(chunk.get("usage"), dict)] assert [at(chunk, "usage", "completion_tokens_details") for chunk in metered] == [ {"text_tokens": CUT.output_tokens} From 01d6812f30c2ccee636b82aa27f8e8307400a222 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Mon, 2 Nov 2026 15:41:16 -0800 Subject: [PATCH 24/79] fix(model-catalog): declare the thinking_on_by_default key (#49720) * fix(model-catalog): declare the thinking_on_by_default key The cost map guard runs main's catalog on a PR's cost map, so the key has to be declared on main before a PR can set it on the Claude 5 rows * chore(model-catalog): keep the attribute list sorted and document thinking_on_by_default --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm-rust/crates/model-catalog/src/wire/attributes.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/litellm-rust/crates/model-catalog/src/wire/attributes.rs b/litellm-rust/crates/model-catalog/src/wire/attributes.rs index 6ac7b4bb304..307f25f0bda 100644 --- a/litellm-rust/crates/model-catalog/src/wire/attributes.rs +++ b/litellm-rust/crates/model-catalog/src/wire/attributes.rs @@ -104,6 +104,8 @@ attributes! { /// Cloud regions the model is available in ('global' or region ids). supported_regions: Vec, thinking_always_on: bool, + /// Whether the model thinks when the request carries no thinking param. + thinking_on_by_default: bool, token_count_framing: u64 as NumericCount, token_count_model: TokenCountModel, token_count_responses_framing: u64 as NumericCount, @@ -301,6 +303,8 @@ mod tests { #[case::known_mode("mode", json!("chat"), true)] #[case::unknown_mode("mode", json!("chatt"), false)] #[case::negative_dbu("input_dbu_cost_per_token", json!(-1), false)] + #[case::thinking_on_by_default("thinking_on_by_default", json!(true), true)] + #[case::thinking_on_by_default_not_a_flag("thinking_on_by_default", json!("yes"), false)] #[case::reasoning_contexts("reasoning_contexts", json!(["reasoning", "non_reasoning"]), true)] #[case::reasoning_modes("reasoning_modes", json!(["none", "adaptive"]), true)] #[case::reasoning_context_not_a_list("reasoning_contexts", json!("reasoning"), false)] From 9601740d0590f945b9c225ea42380da0aa082deb Mon Sep 17 00:00:00 2001 From: shrey Date: Mon, 2 Nov 2026 15:53:39 -0800 Subject: [PATCH 25/79] fix(ui): make the off state of switches visible in light theme (#49718) The unchecked track used --input, which is nearly white in the light theme, so an off switch read as a white knob on a white page. The light unchecked track now uses --muted-foreground, and dark mode keeps its own override --- ui/litellm-dashboard/src/components/ui/switch.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/ui/switch.tsx b/ui/litellm-dashboard/src/components/ui/switch.tsx index 5d7171965d4..31233e049c5 100644 --- a/ui/litellm-dashboard/src/components/ui/switch.tsx +++ b/ui/litellm-dashboard/src/components/ui/switch.tsx @@ -16,7 +16,7 @@ function Switch({ data-slot="switch" data-size={size} className={cn( - "peer group/switch relative inline-flex shrink-0 items-center rounded-full border border-transparent shadow-xs transition-all outline-none after:absolute after:-inset-x-3 after:-inset-y-2 focus-visible:border-ring focus-visible:ring-3 focus-visible:ring-ring/50 aria-invalid:border-destructive aria-invalid:ring-3 aria-invalid:ring-destructive/20 data-[size=default]:h-[18.4px] data-[size=default]:w-[32px] data-[size=sm]:h-[14px] data-[size=sm]:w-[24px] dark:aria-invalid:border-destructive/50 dark:aria-invalid:ring-destructive/40 data-checked:bg-primary data-unchecked:bg-input dark:data-unchecked:bg-input/80 data-disabled:cursor-not-allowed data-disabled:opacity-50", + "peer group/switch relative inline-flex shrink-0 items-center rounded-full border border-transparent shadow-xs transition-all outline-none after:absolute after:-inset-x-3 after:-inset-y-2 focus-visible:border-ring focus-visible:ring-3 focus-visible:ring-ring/50 aria-invalid:border-destructive aria-invalid:ring-3 aria-invalid:ring-destructive/20 data-[size=default]:h-[18.4px] data-[size=default]:w-[32px] data-[size=sm]:h-[14px] data-[size=sm]:w-[24px] dark:aria-invalid:border-destructive/50 dark:aria-invalid:ring-destructive/40 data-checked:bg-primary data-unchecked:bg-muted-foreground dark:data-unchecked:bg-input/80 data-disabled:cursor-not-allowed data-disabled:opacity-50", className, )} {...props} From a5fe912ecf63ab73efddb0d59dca451d81839a57 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 2 Nov 2026 23:55:10 +0000 Subject: [PATCH 26/79] fix(proxy): preserve DB settings during config reload (#49735) * fix(proxy): preserve DB settings during config reload Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): mark reload settings variable as final --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 17 ++++++++++++++++- .../proxy/proxy_server/test_proxy_config.py | 11 +++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 33dcf978854..91bbc418ffd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5717,7 +5717,6 @@ def _with_config_file_pass_through_endpoints( def _reload_settings_store(section: Section, store: SettingsStore, section_config: object) -> None: serving_pass_throughs: Final = store.get("pass_through_endpoints") store.load_yaml(_as_settings_mapping(section_config)) - store.apply_db_row(section, _EMPTY_SETTINGS_MAPPING) if is_resource_list(section, "pass_through_endpoints") and serving_pass_throughs is not None: store["pass_through_endpoints"] = serving_pass_throughs @@ -6540,6 +6539,9 @@ class ProxyConfig: prisma_client=prisma_client, store_model_in_db=store_model_in_db, ) + else: + for section, store in self._settings_stores.items(): + store.apply_db_row(section, _EMPTY_SETTINGS_MAPPING) config = self._config_with_resolved_settings(config) @@ -8623,12 +8625,24 @@ class ProxyConfig: environment_read: Final = next(self._environment_read_numbers) sections: Final = tuple(self._settings_stores) responses: Final = await asyncio.gather(*(get_config_param(prisma_client, section) for section in sections)) + general_settings_response: Final = responses[sections.index("general_settings")] + general_settings_value: Final[Mapping[str, SettingsJsonValue]] = ( + self._prepared_db_settings_values("general_settings", general_settings_response.param_value) + if general_settings_response is not None and general_settings_response.param_value is not None + else _EMPTY_SETTINGS_MAPPING + ) + self._settings_stores["general_settings"].apply_db_row("general_settings", general_settings_value) + response: _ConfigParamValueRow | None for section, response in zip(sections, responses): + if section == "general_settings": + continue if response is None: + self._settings_stores[section].apply_db_row(section, _EMPTY_SETTINGS_MAPPING) continue param_value: JsonValue = getattr(response, "param_value", None) if param_value is None: + self._settings_stores[section].apply_db_row(section, _EMPTY_SETTINGS_MAPPING) continue verbose_proxy_logger.debug( @@ -8639,6 +8653,7 @@ class ProxyConfig: if section == "litellm_settings": self._apply_litellm_settings_db_values(self._prepared_db_settings_values(section, param_value)) elif section == "environment_variables": + self._settings_stores[section].apply_db_row(section, _EMPTY_SETTINGS_MAPPING) await self._apply_db_environment_variables_row( prisma_client, response, environment_before_db, environment_read ) diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index a78209e16fe..55226fc629b 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -32,11 +32,13 @@ import litellm from litellm._logging import turn_on_json from litellm.proxy._types import CommonProxyErrors, ConfigGeneralSettings from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper +from litellm.proxy.config_resolvers.settings_store import SettingsStore from litellm.router_api import ApplyResult, RouterSettings from litellm.types.router import Deployment, LiteLLM_Params from litellm.proxy.proxy_server import ( ProxyConfig, _is_remote_module_url, + _reload_settings_store, _scrub_db_overlay_remote_module_loads, _scrub_guardrail_inner, resolve_complexity_router_plugins, @@ -70,6 +72,15 @@ def _db_model(model_id: str, model_name: str = "gpt-3.5-turbo") -> SimpleNamespa ) +def test_yaml_settings_reload_keeps_db_general_settings_until_sync() -> None: + settings: Final = SettingsStore("general_settings") + settings.apply_db_row("general_settings", {"max_file_downloads_per_minute": 2}) + + _reload_settings_store("general_settings", settings, {}) + + assert settings["max_file_downloads_per_minute"] == 2 + + @pytest.mark.asyncio async def test_proxy_config_adds_and_deletes_stale_deployment(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy import proxy_server From 1bbfb3a02f8375e9ccdf5285497e30ea16491f31 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 2 Nov 2026 15:56:37 -0800 Subject: [PATCH 27/79] refactor(ui): move MCP text onto i18n (#49473) Route the page-authored English on /mcp-servers, /mcp/oauth/callback and the shared MCP permission, server selector and BYOK components through t() and Trans, with 1,094 messages added to en.json under mcpServers, mcpServerManagement, mcpTools and mcpOAuthCallback. The rendered English is unchanged. --- .../_components/AwsSigV4Fields.tsx | 273 +-- .../_components/CreateMCPServer.tsx | 134 +- .../_components/DcrBridgeToggle.tsx | 6 +- .../_components/EnvVarsSection.tsx | 73 +- .../_components/IdJagFormFields.tsx | 88 +- .../_components/ImportMCPServers.tsx | 25 +- .../_components/MCPGatewaySessionsTab.tsx | 146 +- .../_components/MCPLogoSelector.tsx | 19 +- .../_components/MCPNetworkSettings.tsx | 115 +- .../_components/MCPPermissionManagement.tsx | 118 +- .../mcp-servers/_components/MCPServerCard.tsx | 103 +- .../MCPServerUserCredentialsPanel.tsx | 77 +- .../_components/MCPStandardsSettings.tsx | 57 +- .../_components/MCPSubmissionsTab.tsx | 153 +- .../_components/MCPToolsetTableColumns.tsx | 55 +- .../_components/MCPToolsetsTab.tsx | 160 +- .../_components/OAuthFormFields.tsx | 204 +- .../_components/OpenAPIFormSection.tsx | 8 +- .../_components/OpenAPIQuickPicker.tsx | 10 +- .../_components/OpenApiByokFields.tsx | 23 +- .../PassthroughAuthorizeSection.tsx | 84 +- .../_components/StdioAvailability.tsx | 58 +- .../_components/StdioConfiguration.tsx | 10 +- .../TokenEndpointAuthMethodField.tsx | 83 +- .../_components/TokenExchangeFormFields.tsx | 74 +- .../_components/ToolArgumentsForm.tsx | 68 +- .../mcp-servers/_components/ToolTestPanel.tsx | 76 +- .../_components/TruePassthroughWarning.tsx | 11 +- .../_components/UpstreamTokenHeaderField.tsx | 46 +- .../_components/UserEnvVarsModal.tsx | 67 +- .../_components/duplicateServerCheck.ts | 3 +- .../_components/editServerPayload.ts | 18 +- .../_components/editToolPreview.ts | 6 +- .../_components/importConnectorConfig.ts | 21 +- .../mcp-servers/_components/mcpFieldRules.ts | 7 +- .../mcp-servers/_components/mcp_connect.tsx | 245 ++- .../_components/mcp_connection_status.tsx | 42 +- .../mcp-servers/_components/mcp_discovery.tsx | 28 +- .../_components/mcp_server_cost_config.tsx | 49 +- .../_components/mcp_server_cost_display.tsx | 24 +- .../_components/mcp_server_edit.tsx | 194 +- .../_components/mcp_server_view.tsx | 122 +- .../mcp-servers/_components/mcp_servers.tsx | 121 +- .../_components/mcp_tool_configuration.tsx | 90 +- .../mcp-servers/_components/mcp_tools.tsx | 66 +- .../_components/toolCallArguments.ts | 22 +- .../mcp-servers/_components/utils.tsx | 34 +- .../src/app/mcp/oauth/callback/page.tsx | 15 +- .../MCPSemanticFilterSettings.tsx | 82 +- .../MCPSemanticFilterTestPanel.tsx | 67 +- .../semanticFilterTestUtils.ts | 46 +- .../MCPToolSearchSettings.tsx | 51 +- .../MCPServerSelector.tsx | 24 +- .../MCPToolPermissions.tsx | 91 +- .../mcp_tools/ByokCredentialModal.tsx | 64 +- .../mcp_tools/McpCrudPermissionPanel.tsx | 27 +- .../src/hooks/useMcpOAuthFlow.tsx | 18 +- .../src/hooks/useTestMCPConnection.tsx | 4 +- .../src/hooks/useToolsOAuthFlow.tsx | 8 +- ui/litellm-dashboard/src/i18n/locales/en.json | 1768 +++++++++++++++++ 60 files changed, 4151 insertions(+), 1630 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/AwsSigV4Fields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/AwsSigV4Fields.tsx index 5f64d003c11..b153d0a6456 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/AwsSigV4Fields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/AwsSigV4Fields.tsx @@ -1,5 +1,6 @@ import { Info } from "lucide-react"; import React from "react"; +import { useTranslation } from "react-i18next"; import { SimpleTooltip } from "@/components/ui/tooltip"; import { MountedFormField } from "@/components/common_components/MountedFormField"; @@ -22,134 +23,150 @@ const FieldLabel: React.FC<{ label: string; tooltip: string }> = ({ label, toolt const ACCESS_KEY_PATH = ["credentials", "aws_access_key_id"] as const; const SECRET_KEY_PATH = ["credentials", "aws_secret_access_key"] as const; -const AwsSigV4Fields: React.FC = () => ( - <> -

- For MCP servers hosted on AWS Bedrock AgentCore.{" "} - { + const { t } = useTranslation(); + return ( + <> +

+ {t("mcpServers.awsSigV4Fields.intro")}{" "} + + {t("mcpServers.awsSigV4Fields.viewDocs")} + +

+ + } + name={["credentials", "aws_region_name"]} + required + rules={{ validate: { required: requiredRule(t("mcpServers.awsSigV4Fields.region.required")) } }} > - View docs → - -

- } - name={["credentials", "aws_region_name"]} - required - rules={{ validate: { required: requiredRule("AWS region is required for SigV4 auth") } }} - > - {(control) => } - - - } - name={["credentials", "aws_service_name"]} - > - {(control) => } - - - } - name={ACCESS_KEY_PATH} - rules={{ - deps: ["credentials.aws_secret_access_key"], - validate: { - pairedWithSecret: requiredWhenSiblingSet( - SECRET_KEY_PATH, - "Access Key ID is required when Secret Access Key is provided", - ), - }, - }} - > - {(control) => ( - - )} - - - } - name={SECRET_KEY_PATH} - rules={{ - deps: ["credentials.aws_access_key_id"], - validate: { - pairedWithAccessKey: requiredWhenSiblingSet( - ACCESS_KEY_PATH, - "Secret Access Key is required when Access Key ID is provided", - ), - }, - }} - > - {(control) => ( - - )} - - } - name={["credentials", "aws_session_token"]} - > - {(control) => ( - - )} - - - } - name={["credentials", "aws_role_name"]} - > - {(control) => ( - - )} - - - } - name={["credentials", "aws_session_name"]} - > - {(control) => ( - - )} - - -); + {(control) => } + + + } + name={["credentials", "aws_service_name"]} + > + {(control) => } + + + } + name={ACCESS_KEY_PATH} + rules={{ + deps: ["credentials.aws_secret_access_key"], + validate: { + pairedWithSecret: requiredWhenSiblingSet( + SECRET_KEY_PATH, + t("mcpServers.awsSigV4Fields.accessKeyId.requiredWithSecret"), + ), + }, + }} + > + {(control) => ( + + )} + + + } + name={SECRET_KEY_PATH} + rules={{ + deps: ["credentials.aws_access_key_id"], + validate: { + pairedWithAccessKey: requiredWhenSiblingSet( + ACCESS_KEY_PATH, + t("mcpServers.awsSigV4Fields.secretAccessKey.requiredWithAccessKey"), + ), + }, + }} + > + {(control) => ( + + )} + + + } + name={["credentials", "aws_session_token"]} + > + {(control) => ( + + )} + + + } + name={["credentials", "aws_role_name"]} + > + {(control) => ( + + )} + + + } + name={["credentials", "aws_session_name"]} + > + {(control) => ( + + )} + + + ); +}; export default AwsSigV4Fields; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.tsx index 07618abb9a4..14a6e3d941f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.tsx @@ -1,5 +1,7 @@ +import type { TFunction } from "i18next"; import React, { useState } from "react"; import { FormProvider, useForm, useWatch } from "react-hook-form"; +import { useTranslation } from "react-i18next"; import { ChevronDown, Info } from "lucide-react"; import { Button } from "@/components/ui/button"; import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; @@ -35,7 +37,7 @@ import { buildCreateServerPayload, reduceStaticHeaders, } from "./createServerPayload"; -import { DUPLICATE_IDENTIFIER_MESSAGE, findDuplicateMcpServer, mcpSubmitErrorReason } from "./duplicateServerCheck"; +import { DUPLICATE_IDENTIFIER_MESSAGE_KEY, findDuplicateMcpServer, mcpSubmitErrorReason } from "./duplicateServerCheck"; import { readCreateUiSnapshot, writeCreateUiSnapshot } from "./createOAuthUiState"; import AwsSigV4Fields from "./AwsSigV4Fields"; import OpenApiByokFields from "./OpenApiByokFields"; @@ -86,14 +88,14 @@ interface CreateMCPServerProps { stdioEnabled?: boolean; } -const payloadErrorMessage = (result: Exclude): string => { +const payloadErrorMessage = (result: Exclude, t: TFunction): string => { switch (result.kind) { case "invalid_tool_display_name": - return `Tool display name "${result.displayName}" is invalid. Only letters, digits, underscores, and hyphens are allowed (no spaces).`; + return t("mcpServers.createMcpServer.payloadErrors.invalidToolDisplayName", { displayName: result.displayName }); case "invalid_stdio_json": - return "Invalid JSON in stdio configuration"; + return t("mcpServers.createMcpServer.payloadErrors.invalidStdioJson"); case "invalid_token_validation_json": - return "Invalid JSON in Token Validation Rules"; + return t("mcpServers.createMcpServer.payloadErrors.invalidTokenValidationJson"); } }; @@ -117,6 +119,7 @@ const CreateMCPServer: React.FC = ({ onBackToDiscovery, stdioEnabled = true, }) => { + const { t } = useTranslation(); const form = useForm({ mode: "onChange", defaultValues: CREATE_DEFAULTS }); const registry = useMountRegistry(); const [isLoading, setIsLoading] = useState(false); @@ -255,9 +258,7 @@ const CreateMCPServer: React.FC = ({ // which would persist it as server-level credentials on the created server row. Mirrors the // edit form's onTokenReceived early return. setAuthorizedIdentity(getOAuthAuthorizationIdentity(allFieldsValue(form))); - toast.success( - "Token held for this browser session. Tools can now be previewed and configured; the token is not saved to LiteLLM.", - ); + toast.success(t("mcpServers.createMcpServer.toasts.browserTokenHeld")); return; } @@ -283,7 +284,7 @@ const CreateMCPServer: React.FC = ({ // its own credential write. setAuthorizedIdentity(getOAuthAuthorizationIdentity(allFieldsValue(form))); - toast.success("OAuth authorization successful! Please click 'Create MCP Server' to save the configuration."); + toast.success(t("mcpServers.createMcpServer.toasts.oauthAuthorized")); }, onBeforeRedirect: persistCreateUiState, flowSource: "create", @@ -425,8 +426,8 @@ const CreateMCPServer: React.FC = ({ typeof values.alias === "string" ? values.alias : undefined, ); if (duplicate) { - form.setError(duplicate.field, { type: "duplicate", message: DUPLICATE_IDENTIFIER_MESSAGE }); - toast.fromError(DUPLICATE_IDENTIFIER_MESSAGE); + form.setError(duplicate.field, { type: "duplicate", message: t(DUPLICATE_IDENTIFIER_MESSAGE_KEY) }); + toast.fromError(t(DUPLICATE_IDENTIFIER_MESSAGE_KEY)); return; } const built = buildCreateServerPayload(values, { @@ -440,7 +441,7 @@ const CreateMCPServer: React.FC = ({ dcrClient: dcrClientRef.current, }); if (built.kind !== "ok") { - toast.fromError(payloadErrorMessage(built)); + toast.fromError(payloadErrorMessage(built, t)); return; } const payload = built.payload; @@ -482,10 +483,10 @@ const CreateMCPServer: React.FC = ({ } if (isAdmin) { - toast.success("MCP Server created successfully"); + toast.success(t("mcpServers.createMcpServer.toasts.created")); } else { - toast.success("MCP Server submitted for admin review", { - description: "Once an admin approves it, the server will appear in your MCP Servers list.", + toast.success(t("mcpServers.createMcpServer.toasts.submittedForReview"), { + description: t("mcpServers.createMcpServer.toasts.submittedForReviewDescription"), }); } form.reset(CREATE_DEFAULTS); @@ -500,7 +501,11 @@ const CreateMCPServer: React.FC = ({ } } catch (error) { const reason = mcpSubmitErrorReason(error); - toast.fromError(isAdmin ? `Error creating MCP Server: ${reason}` : `Error submitting MCP Server: ${reason}`); + toast.fromError( + isAdmin + ? t("mcpServers.createMcpServer.toasts.createFailed", { reason }) + : t("mcpServers.createMcpServer.toasts.submitFailed", { reason }), + ); } finally { setIsLoading(false); } @@ -640,9 +645,11 @@ const CreateMCPServer: React.FC = ({ ← )} - MCP Logo + {t("mcpServers.createMcpServer.logoAlt")} - {isAdmin ? "Add New MCP Server" : "Submit MCP Server for Review"} + {isAdmin + ? t("mcpServers.createMcpServer.title.admin") + : t("mcpServers.createMcpServer.title.submitForReview")} @@ -653,27 +660,26 @@ const CreateMCPServer: React.FC = ({
{!isAdmin && (
- Your submission will be sent for admin review. Once approved, the server will appear in your MCP - Servers list. The request must be made with a team-scoped API key. + {t("mcpServers.createMcpServer.reviewNotice")}
)}
- MCP Server Name - + {t("mcpServers.createMcpServer.serverName.label")} + } name="server_name" - rules={{ validate: validatorRules({ validator: (_, value) => validateMCPServerName(value) }) }} + rules={{ validate: validatorRules({ validator: (_, value) => validateMCPServerName(value, t) }) }} > {(control) => ( )} @@ -682,19 +688,19 @@ const CreateMCPServer: React.FC = ({ - Alias - + {t("mcpServers.createMcpServer.alias.label")} + } name="alias" - rules={{ validate: validatorRules({ validator: (_, value) => validateMCPServerName(value) }) }} + rules={{ validate: validatorRules({ validator: (_, value) => validateMCPServerName(value, t) }) }} > {(control) => ( { control.onChange(event); @@ -705,13 +711,17 @@ const CreateMCPServer: React.FC = ({ Description} + label={ + + {t("mcpServers.createMcpServer.description.label")} + + } name="description" > {(control) => ( )} @@ -720,7 +730,11 @@ const CreateMCPServer: React.FC = ({ GitHub / Source URL} + label={ + + {t("mcpServers.createMcpServer.sourceUrl.label")} + + } name="source_url" > {(control) => ( @@ -733,10 +747,14 @@ const CreateMCPServer: React.FC = ({ Transport Type} + label={ + + {t("mcpServers.createMcpServer.transport.label")} + + } name="transport" required - rules={{ validate: { required: requiredRule("Please select a transport type") } }} + rules={{ validate: { required: requiredRule(t("mcpServers.createMcpServer.transport.required")) } }} > {(control) => ( (control)} items={AUTH_TYPE_ITEMS}> - + {AUTH_TYPE_ITEMS.map((item) => ( @@ -886,8 +912,8 @@ const CreateMCPServer: React.FC = ({ - Authentication Value - + {t("mcpServers.createMcpServer.auth.valueLabel")} + @@ -895,14 +921,14 @@ const CreateMCPServer: React.FC = ({ name={["credentials", "auth_value"]} rules={{ validate: { - notWhitespace: notOnlyWhitespace("Authentication value cannot be empty whitespace"), + notWhitespace: notOnlyWhitespace(t("mcpServers.createMcpServer.auth.valueWhitespace")), }, }} > {(control) => ( )} @@ -1000,11 +1026,13 @@ const CreateMCPServer: React.FC = ({
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/DcrBridgeToggle.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/DcrBridgeToggle.tsx index 893981fa9c6..48a8f4d7a64 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/DcrBridgeToggle.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/DcrBridgeToggle.tsx @@ -1,5 +1,6 @@ import { Info } from "lucide-react"; import React from "react"; +import { useTranslation } from "react-i18next"; import { SimpleTooltip } from "@/components/ui/tooltip"; import { Switch } from "@/components/ui/switch"; @@ -23,13 +24,14 @@ export default function DcrBridgeToggle({ authType?: string | null; initialChecked?: boolean; }) { + const { t } = useTranslation(); if (!isClientForwardedTokenMode(authType)) return null; return ( - Gateway-hosted sign-in (DCR bridge) - + {t("mcpServers.dcrBridgeToggle.label")} + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/EnvVarsSection.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/EnvVarsSection.tsx index 378410e4790..a804d4697c4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/EnvVarsSection.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/EnvVarsSection.tsx @@ -1,5 +1,7 @@ +import type { ParseKeys } from "i18next"; import { CircleMinus, Info, Plus } from "lucide-react"; import React from "react"; +import { Trans, useTranslation } from "react-i18next"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { SimpleTooltip } from "@/components/ui/tooltip"; import { Button } from "@/components/ui/button"; @@ -16,11 +18,13 @@ import { requiredRule } from "@/components/common_components/formRules"; import { matchesPattern, selectControl, selectTriggerControl, textControl } from "./mcpFieldRules"; import { listControl } from "./mcpFormStore"; -const SCOPE_OPTIONS = [ - { value: "global", label: "Instance" }, - { value: "user", label: "Per-user" }, +const SCOPE_OPTIONS: ReadonlyArray<{ value: string; labelKey: ParseKeys }> = [ + { value: "global", labelKey: "mcpServers.envVarsSection.scopes.instance" }, + { value: "user", labelKey: "mcpServers.envVarsSection.scopes.perUser" }, ]; +const VARIABLE_SYNTAX = "${VAR_NAME}"; + const VARIABLE_NAME_PATTERN = /^[A-Za-z_][A-Za-z0-9_]*$/; /** @@ -34,31 +38,34 @@ const VARIABLE_NAME_PATTERN = /^[A-Za-z_][A-Za-z0-9_]*$/; * The parent form reads the ``env_vars`` field from the form values. */ const EnvVarsSection: React.FC = () => { + const { t } = useTranslation(); const { control } = useFormContext(); const { fields, append, remove } = useFieldArray({ control: listControl(control), name: "env_vars" }); useMountedName("env_vars"); + const scopeOptions = SCOPE_OPTIONS.map((option) => ({ value: option.value, label: t(option.labelKey) })); return (
- Variables + {t("mcpServers.envVarsSection.title")} - Define variables you can interpolate in Static Headers or Authentication using{" "} - {"${VAR_NAME}"}.
- Instance: admin-defined value used for every user. -
- Per-user: each user supplies their own value (e.g. personal credentials) via the MCP Gateway - dashboard. - + , b: }} + /> } >
- Reference these in Static Headers or Authentication as {"${VAR_NAME}"}. For example:{" "} + }} + />{" "} {"${DB_PROTOCOL}://${CORP_USERNAME}:${CORP_PASSWORD}@${DB_HOSTNAME}"} @@ -67,9 +74,9 @@ const EnvVarsSection: React.FC = () => {
{fields.length > 0 && (
-
Variable Name
-
Value / Description
-
Scope
+
{t("mcpServers.envVarsSection.columns.name")}
+
{t("mcpServers.envVarsSection.columns.value")}
+
{t("mcpServers.envVarsSection.columns.scope")}
)} @@ -80,16 +87,17 @@ const EnvVarsSection: React.FC = () => { className="mb-0 flex-1" rules={{ validate: { - required: requiredRule("Variable name is required"), - pattern: matchesPattern( - VARIABLE_NAME_PATTERN, - "Use letters, digits, underscores; cannot start with a digit.", - ), + required: requiredRule(t("mcpServers.envVarsSection.name.required")), + pattern: matchesPattern(VARIABLE_NAME_PATTERN, t("mcpServers.envVarsSection.name.invalid")), }, }} > {(control) => ( - + )}
@@ -97,12 +105,12 @@ const EnvVarsSection: React.FC = () => {
{(control) => ( - (control)} items={scopeOptions}> - {SCOPE_OPTIONS.map((option) => ( + {scopeOptions.map((option) => ( {option.label} @@ -121,7 +129,7 @@ const EnvVarsSection: React.FC = () => { ))}
@@ -132,6 +140,7 @@ const EnvVarsSection: React.FC = () => { // vars the value comes from each user later, so the column instead captures an // optional description that the per-user fill-in modal shows as a hint. const ScopedValueOrDescription: React.FC<{ index: number }> = ({ index }) => { + const { t } = useTranslation(); const isPerUser = useWatch({ name: `env_vars.${index}.scope` }) === "user"; if (isPerUser) { return ( @@ -139,16 +148,16 @@ const ScopedValueOrDescription: React.FC<{ index: number }> = ({ index }) => { {(control) => ( - + - Hint + {t("mcpServers.envVarsSection.hint.label")} @@ -158,7 +167,13 @@ const ScopedValueOrDescription: React.FC<{ index: number }> = ({ index }) => { } return ( - {(control) => } + {(control) => ( + + )} ); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx index 70d1bc40c18..98bb3254ba5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx @@ -1,5 +1,6 @@ import { Info } from "lucide-react"; import React from "react"; +import { useTranslation } from "react-i18next"; import { SimpleTooltip } from "@/components/ui/tooltip"; import { MountedFormField } from "@/components/common_components/MountedFormField"; @@ -27,9 +28,10 @@ const FieldLabel: React.FC<{ label: string; tooltip: string }> = ({ label, toolt ); const PRIVATE_KEY_PATH = ["credentials", "client_private_key"] as const; +const PEM_HEADER = "-----BEGIN PRIVATE KEY-----"; const IdJagFormFields: React.FC = ({ isEditing = false }) => { - const placeholderSuffix = isEditing ? " (leave blank to keep existing)" : ""; + const { t } = useTranslation(); const requiredWhenCreating = (message: string) => isEditing ? undefined : { validate: { required: requiredRule(message) } }; @@ -38,13 +40,13 @@ const IdJagFormFields: React.FC = ({ isEditing = false }) } name="token_exchange_endpoint" required={!isEditing} - rules={requiredWhenCreating("The org token endpoint is required for ID-JAG")} + rules={requiredWhenCreating(t("mcpServers.idJagFormFields.orgTokenEndpoint.required"))} > {(control) => ( = ({ isEditing = false }) } name={["credentials", "id_jag_resource_token_endpoint"]} required={!isEditing} - rules={requiredWhenCreating("The resource token endpoint is required for ID-JAG")} + rules={requiredWhenCreating(t("mcpServers.idJagFormFields.resourceTokenEndpoint.required"))} > {(control) => ( = ({ isEditing = false }) )} } + label={ + + } name={["credentials", "client_id"]} required={!isEditing} - rules={requiredWhenCreating("Client ID is required for ID-JAG")} + rules={requiredWhenCreating(t("mcpServers.idJagFormFields.clientId.required"))} > {(control) => ( )} @@ -90,8 +101,8 @@ const IdJagFormFields: React.FC = ({ isEditing = false }) } name={["credentials", "client_secret"]} @@ -103,7 +114,7 @@ const IdJagFormFields: React.FC = ({ isEditing = false }) validate: { secretOrPrivateKey: requiredUnlessSiblingSet( PRIVATE_KEY_PATH, - "Provide either a client secret or a client private key", + t("mcpServers.idJagFormFields.clientSecret.secretOrPrivateKey"), ), }, } @@ -112,7 +123,11 @@ const IdJagFormFields: React.FC = ({ isEditing = false }) {(control) => ( )} @@ -120,8 +135,8 @@ const IdJagFormFields: React.FC = ({ isEditing = false }) } name={PRIVATE_KEY_PATH} @@ -130,7 +145,11 @@ const IdJagFormFields: React.FC = ({ isEditing = false })