From b2463291c754cb14d7d90a326d66568d822fc168 Mon Sep 17 00:00:00 2001 From: Chesars Date: Thu, 27 Nov 2025 17:40:06 -0300 Subject: [PATCH 01/22] fix(image-gen): add thought_signature to ImageObject for Gemini 3 Pro Fixes #17184 - Gemini 3 Pro image preview model returns a thoughtSignature field required for interactive image editing. This change: - Adds thought_signature field to ImageObject class - Updates Gemini and Vertex AI transformations to extract thoughtSignature - Adds test for thought_signature in response transformation --- .../gemini/image_generation/transformation.py | 1 + .../vertex_gemini_transformation.py | 1 + litellm/types/utils.py | 5 ++- ...rtex_ai_image_generation_transformation.py | 41 +++++++++++++++++++ 4 files changed, 47 insertions(+), 1 deletion(-) diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index 63b835df9d0..46cbf6ae877 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -258,6 +258,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): model_response.data.append(ImageObject( b64_json=inline_data["data"], url=None, + thought_signature=part.get("thoughtSignature"), )) # Extract usage metadata for Gemini models diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index 89ed9f1a8a5..f266f860712 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -298,6 +298,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): model_response.data.append(ImageObject( b64_json=inline_data["data"], url=None, + thought_signature=part.get("thoughtSignature"), )) if usage_metadata := response_data.get("usageMetadata", None): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6c330d0f83c..077277f4830 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2129,6 +2129,7 @@ class ImageObject(OpenAIImage): b64_json: The base64-encoded JSON of the generated image, if response_format is b64_json. url: The URL of the generated image, if response_format is url (default). revised_prompt: The prompt that was used to generate the image, if there was any revision to the prompt. + thought_signature: The thought signature returned by Gemini image generation models (used for interactive image editing). https://platform.openai.com/docs/api-reference/images/object """ @@ -2136,9 +2137,11 @@ class ImageObject(OpenAIImage): b64_json: Optional[str] = None url: Optional[str] = None revised_prompt: Optional[str] = None + thought_signature: Optional[str] = None - def __init__(self, b64_json=None, url=None, revised_prompt=None, **kwargs): + def __init__(self, b64_json=None, url=None, revised_prompt=None, thought_signature=None, **kwargs): super().__init__(b64_json=b64_json, url=url, revised_prompt=revised_prompt) # type: ignore + self.thought_signature = thought_signature def __contains__(self, key): # Define custom behavior for the 'in' operator diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index b91438b3cac..b4f5053955e 100644 --- a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -230,6 +230,47 @@ class TestVertexAIGeminiImageGenerationConfig: assert result.data[0].b64_json == "image1" assert result.data[1].b64_json == "image2" + def test_transform_image_generation_response_signature(self): + """Test response transformation includes thoughtSignature for Gemini 3 Pro""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "base64_encoded_image_data", + }, + "thoughtSignature": "test_signature_abc123", + } + ] + } + } + ] + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + model_response = ImageResponse() + result = self.config.transform_image_generation_response( + model="gemini-3-pro-image-preview", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 1 + assert result.data[0].b64_json == "base64_encoded_image_data" + assert result.data[0].thought_signature == "test_signature_abc123" + class TestVertexAIImagenImageGenerationConfig: def setup_method(self): From 11fd92c21b097426f6d4e0116d002c77766497f2 Mon Sep 17 00:00:00 2001 From: Chesars Date: Fri, 28 Nov 2025 17:16:43 -0300 Subject: [PATCH 02/22] refactor(image-gen): move thought_signature to provider_specific_fields Per review feedback, thought_signature should not be a root-level param on ImageObject as it's not OpenAI compatible. Moved to provider_specific_fields dict to match the pattern used in chat completions (Message, Delta, Choices, etc). --- litellm/llms/gemini/image_generation/transformation.py | 3 ++- .../image_generation/vertex_gemini_transformation.py | 3 ++- litellm/types/utils.py | 9 +++++---- .../test_vertex_ai_image_generation_transformation.py | 2 +- 4 files changed, 10 insertions(+), 7 deletions(-) diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index 46cbf6ae877..73aef15e4c7 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -255,10 +255,11 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): if "inlineData" in part: inline_data = part["inlineData"] if "data" in inline_data: + thought_sig = part.get("thoughtSignature") model_response.data.append(ImageObject( b64_json=inline_data["data"], url=None, - thought_signature=part.get("thoughtSignature"), + provider_specific_fields={"thought_signature": thought_sig} if thought_sig else None, )) # Extract usage metadata for Gemini models diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index f266f860712..ba3df88be14 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -295,10 +295,11 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): if "inlineData" in part: inline_data = part["inlineData"] if "data" in inline_data: + thought_sig = part.get("thoughtSignature") model_response.data.append(ImageObject( b64_json=inline_data["data"], url=None, - thought_signature=part.get("thoughtSignature"), + provider_specific_fields={"thought_signature": thought_sig} if thought_sig else None, )) if usage_metadata := response_data.get("usageMetadata", None): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 077277f4830..09c944e1fe8 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2129,7 +2129,7 @@ class ImageObject(OpenAIImage): b64_json: The base64-encoded JSON of the generated image, if response_format is b64_json. url: The URL of the generated image, if response_format is url (default). revised_prompt: The prompt that was used to generate the image, if there was any revision to the prompt. - thought_signature: The thought signature returned by Gemini image generation models (used for interactive image editing). + provider_specific_fields: Provider-specific fields not part of OpenAI spec. https://platform.openai.com/docs/api-reference/images/object """ @@ -2137,11 +2137,12 @@ class ImageObject(OpenAIImage): b64_json: Optional[str] = None url: Optional[str] = None revised_prompt: Optional[str] = None - thought_signature: Optional[str] = None + provider_specific_fields: Optional[Dict[str, Any]] = None - def __init__(self, b64_json=None, url=None, revised_prompt=None, thought_signature=None, **kwargs): + def __init__(self, b64_json=None, url=None, revised_prompt=None, provider_specific_fields=None, **kwargs): super().__init__(b64_json=b64_json, url=url, revised_prompt=revised_prompt) # type: ignore - self.thought_signature = thought_signature + if provider_specific_fields: + self.provider_specific_fields = provider_specific_fields def __contains__(self, key): # Define custom behavior for the 'in' operator diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index b4f5053955e..6736eaffebd 100644 --- a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -269,7 +269,7 @@ class TestVertexAIGeminiImageGenerationConfig: assert len(result.data) == 1 assert result.data[0].b64_json == "base64_encoded_image_data" - assert result.data[0].thought_signature == "test_signature_abc123" + assert result.data[0].provider_specific_fields["thought_signature"] == "test_signature_abc123" class TestVertexAIImagenImageGenerationConfig: From c4dd22c07966bc7fa2f0de8471f26092aef90a06 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 16 Dec 2025 19:31:10 -0600 Subject: [PATCH 03/22] fix: broaden Azure AI rerank URL handling --- .../llms/azure_ai/rerank/transformation.py | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/litellm/llms/azure_ai/rerank/transformation.py b/litellm/llms/azure_ai/rerank/transformation.py index a47b6082c37..085f36f74b8 100644 --- a/litellm/llms/azure_ai/rerank/transformation.py +++ b/litellm/llms/azure_ai/rerank/transformation.py @@ -11,6 +11,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.cohere.rerank.transformation import CohereRerankConfig from litellm.secret_managers.main import get_secret_str from litellm.types.utils import RerankResponse +from litellm.utils import _add_path_to_api_base class AzureAIRerankConfig(CohereRerankConfig): @@ -28,9 +29,24 @@ class AzureAIRerankConfig(CohereRerankConfig): raise ValueError( "Azure AI API Base is required. api_base=None. Set in call or via `AZURE_AI_API_BASE` env var." ) - if not api_base.endswith("/v1/rerank"): - api_base = f"{api_base}/v1/rerank" - return api_base + original_url = httpx.URL(api_base) + normalized_path = original_url.path.rstrip("/") + + # Allow callers to pass either full v1/v2 rerank endpoints: + # - https://.services.ai.azure.com/v1/rerank + # - https://.services.ai.azure.com/providers/cohere/v2/rerank + if normalized_path.endswith("/v1/rerank") or normalized_path.endswith("/v2/rerank"): + return str(original_url.copy_with(path=normalized_path or "/")) + + # If callers pass just the version path (e.g. ".../v2"), append "/rerank" + if normalized_path.endswith("/v1") or normalized_path.endswith("/v2"): + return _add_path_to_api_base( + api_base=str(original_url.copy_with(path=normalized_path or "/")), + ending_path="/rerank", + ) + + # Backwards compatible default: Azure AI rerank was originally exposed under /v1/rerank + return _add_path_to_api_base(api_base=api_base, ending_path="/v1/rerank") def validate_environment( self, From 92763a14a9e164ca9dc0d7fa6a27489ca8c5cd40 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 16 Dec 2025 19:40:45 -0600 Subject: [PATCH 04/22] Update litellm/llms/azure_ai/rerank/transformation.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- litellm/llms/azure_ai/rerank/transformation.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/litellm/llms/azure_ai/rerank/transformation.py b/litellm/llms/azure_ai/rerank/transformation.py index 085f36f74b8..376f4608260 100644 --- a/litellm/llms/azure_ai/rerank/transformation.py +++ b/litellm/llms/azure_ai/rerank/transformation.py @@ -38,8 +38,12 @@ class AzureAIRerankConfig(CohereRerankConfig): if normalized_path.endswith("/v1/rerank") or normalized_path.endswith("/v2/rerank"): return str(original_url.copy_with(path=normalized_path or "/")) - # If callers pass just the version path (e.g. ".../v2"), append "/rerank" - if normalized_path.endswith("/v1") or normalized_path.endswith("/v2"): + # If callers pass just the version path (e.g. ".../v2" or ".../providers/cohere/v2"), append "/rerank" + if ( + normalized_path.endswith("/v1") + or normalized_path.endswith("/v2") + or normalized_path.endswith("/providers/cohere/v2") + ): return _add_path_to_api_base( api_base=str(original_url.copy_with(path=normalized_path or "/")), ending_path="/rerank", From bb5397d9b2582e29c4b4f2ec613a25ef0c926ddf Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Wed, 17 Dec 2025 18:13:41 -0600 Subject: [PATCH 05/22] fix: enforce scheme for Azure AI rerank api_base --- .../llms/azure_ai/rerank/transformation.py | 6 ++ .../test_azure_ai_rerank_transformation.py | 100 ++++++++++++++++++ 2 files changed, 106 insertions(+) create mode 100644 tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py diff --git a/litellm/llms/azure_ai/rerank/transformation.py b/litellm/llms/azure_ai/rerank/transformation.py index 376f4608260..f577a42ed58 100644 --- a/litellm/llms/azure_ai/rerank/transformation.py +++ b/litellm/llms/azure_ai/rerank/transformation.py @@ -30,6 +30,12 @@ class AzureAIRerankConfig(CohereRerankConfig): "Azure AI API Base is required. api_base=None. Set in call or via `AZURE_AI_API_BASE` env var." ) original_url = httpx.URL(api_base) + if not original_url.is_absolute_url: + raise ValueError( + "Azure AI API Base must be an absolute URL including scheme (e.g. " + "'https://.services.ai.azure.com'). " + f"Got api_base={api_base!r}." + ) normalized_path = original_url.path.rstrip("/") # Allow callers to pass either full v1/v2 rerank endpoints: diff --git a/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py b/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py new file mode 100644 index 00000000000..1f425113439 --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py @@ -0,0 +1,100 @@ +import os +import sys + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + +from litellm.llms.azure_ai.rerank.transformation import AzureAIRerankConfig + + +class TestAzureAIRerankConfigGetCompleteUrl: + def setup_method(self): + self.config = AzureAIRerankConfig() + self.model = "azure_ai/cohere-rerank-v3-english" + + def test_api_base_required(self): + with pytest.raises(ValueError) as exc_info: + self.config.get_complete_url(api_base=None, model=self.model) + + assert "api_base=None" in str(exc_info.value) + + @pytest.mark.parametrize( + "api_base", + [ + "example.com", + "example.com/v1", + "//example.com/v1", + "/v1/rerank", + ], + ) + def test_api_base_requires_scheme(self, api_base): + with pytest.raises(ValueError) as exc_info: + self.config.get_complete_url(api_base=api_base, model=self.model) + + error_message = str(exc_info.value).lower() + assert "absolute url" in error_message + assert "scheme" in error_message + + @pytest.mark.parametrize( + "api_base, expected_url", + [ + ( + "https://my-resource.services.ai.azure.com/v1/rerank/", + "https://my-resource.services.ai.azure.com/v1/rerank", + ), + ( + "https://my-resource.services.ai.azure.com/providers/cohere/v2/rerank/", + "https://my-resource.services.ai.azure.com/providers/cohere/v2/rerank", + ), + ], + ) + def test_preserves_full_rerank_endpoint(self, api_base, expected_url): + url = self.config.get_complete_url(api_base=api_base, model=self.model) + assert url == expected_url + + @pytest.mark.parametrize( + "api_base, expected_url", + [ + ( + "https://my-resource.services.ai.azure.com/v1", + "https://my-resource.services.ai.azure.com/v1/rerank", + ), + ( + "https://my-resource.services.ai.azure.com/v2/", + "https://my-resource.services.ai.azure.com/v2/rerank", + ), + ( + "https://my-resource.services.ai.azure.com/providers/cohere/v2", + "https://my-resource.services.ai.azure.com/providers/cohere/v2/rerank", + ), + ( + "https://my-resource.services.ai.azure.com/providers/cohere/v2/", + "https://my-resource.services.ai.azure.com/providers/cohere/v2/rerank", + ), + ], + ) + def test_appends_rerank_for_version_paths(self, api_base, expected_url): + url = self.config.get_complete_url(api_base=api_base, model=self.model) + assert url == expected_url + + @pytest.mark.parametrize( + "api_base", + [ + "https://my-resource.services.ai.azure.com", + "https://my-resource.services.ai.azure.com/", + ], + ) + def test_defaults_to_v1_rerank_when_base_has_no_path(self, api_base): + url = self.config.get_complete_url(api_base=api_base, model=self.model) + assert url == "https://my-resource.services.ai.azure.com/v1/rerank" + + def test_preserves_query_params(self): + url = self.config.get_complete_url( + api_base="https://my-resource.services.ai.azure.com/v1?r=1", + model=self.model, + ) + assert url == "https://my-resource.services.ai.azure.com/v1/rerank?r=1" + From 3215dc4d4e372292c5d4338c30ef9ebe7b0c4791 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 27 Jan 2026 12:14:15 -0300 Subject: [PATCH 06/22] feat(vertex_ai): add global endpoint support for Qwen MaaS models Fixes #19788 - Add `supported_regions: ["global"]` to Qwen MaaS models in model_prices_and_context_window.json - Update `get_supported_regions()` to read directly from `model_cost` dict - Update `get_complete_vertex_url()` to use `get_vertex_region()` for global-only models - Update `create_vertex_url()` to generate correct URL for global location (without region prefix) - Add tests for Qwen global endpoint support --- litellm/llms/vertex_ai/vertex_llm_base.py | 15 +- ...odel_prices_and_context_window_backup.json | 4 + litellm/utils.py | 9 +- model_prices_and_context_window.json | 4 + .../vertex_ai_partner_models/qwen/__init__.py | 0 .../test_vertex_ai_qwen_global_endpoint.py | 263 ++++++++++++++++++ 6 files changed, 290 insertions(+), 5 deletions(-) create mode 100644 tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/__init__.py create mode 100644 tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index a185370e376..1310e66e1db 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -218,7 +218,12 @@ class VertexBase: ) -> str: """Return the base url for the vertex partner models""" - api_base = api_base or f"https://{vertex_location}-aiplatform.googleapis.com" + # For global location, use the non-regional URL + if api_base is None: + if vertex_location == "global": + api_base = "https://aiplatform.googleapis.com" + else: + api_base = f"https://{vertex_location}-aiplatform.googleapis.com" if partner == VertexPartnerProvider.llama: return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/endpoints/openapi/chat/completions" elif partner == VertexPartnerProvider.mistralai: @@ -247,11 +252,13 @@ class VertexBase: stream: Optional[bool], model: str, ) -> str: + # Use get_vertex_region to handle global-only models + resolved_location = self.get_vertex_region(vertex_location, model) api_base = self.get_api_base( - api_base=custom_api_base, vertex_location=vertex_location + api_base=custom_api_base, vertex_location=resolved_location ) default_api_base = VertexBase.create_vertex_url( - vertex_location=vertex_location or "us-central1", + vertex_location=resolved_location, vertex_project=vertex_project or project_id, partner=partner, stream=stream, @@ -274,7 +281,7 @@ class VertexBase: url=default_api_base, model=model, vertex_project=vertex_project or project_id, - vertex_location=vertex_location or "us-central1", + vertex_location=resolved_location, vertex_api_version="v1", # Partner models typically use v1 ) return api_base diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2576903ffe9..6aeb51d5817 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -29785,6 +29785,7 @@ "mode": "chat", "output_cost_per_token": 1e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_regions": ["global"], "supports_function_calling": true, "supports_tool_choice": true }, @@ -29797,6 +29798,7 @@ "mode": "chat", "output_cost_per_token": 4e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_regions": ["global"], "supports_function_calling": true, "supports_tool_choice": true }, @@ -29809,6 +29811,7 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_regions": ["global"], "supports_function_calling": true, "supports_tool_choice": true }, @@ -29821,6 +29824,7 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_regions": ["global"], "supports_function_calling": true, "supports_tool_choice": true }, diff --git a/litellm/utils.py b/litellm/utils.py index a5df9381fc7..e67b967d75a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2644,7 +2644,14 @@ def get_supported_regions( model=model, custom_llm_provider=custom_llm_provider ) - supported_regions = model_info.get("supported_regions", None) + # Get the key used in model_cost to look up supported_regions + # since ModelInfoBase doesn't include this field + model_key = model_info.get("key") + if model_key is None: + return None + + model_cost_data = litellm.model_cost.get(model_key, {}) + supported_regions = model_cost_data.get("supported_regions", None) if supported_regions is None: return None diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 2576903ffe9..6aeb51d5817 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -29785,6 +29785,7 @@ "mode": "chat", "output_cost_per_token": 1e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_regions": ["global"], "supports_function_calling": true, "supports_tool_choice": true }, @@ -29797,6 +29798,7 @@ "mode": "chat", "output_cost_per_token": 4e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_regions": ["global"], "supports_function_calling": true, "supports_tool_choice": true }, @@ -29809,6 +29811,7 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_regions": ["global"], "supports_function_calling": true, "supports_tool_choice": true }, @@ -29821,6 +29824,7 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "supported_regions": ["global"], "supports_function_calling": true, "supports_tool_choice": true }, diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/__init__.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py new file mode 100644 index 00000000000..6310431813b --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py @@ -0,0 +1,263 @@ +""" +Tests for Vertex AI Qwen MaaS models that require the global endpoint. + +These tests verify that: +1. Qwen models are correctly identified as global-only models +2. The correct global URL is constructed (https://aiplatform.googleapis.com) +3. The completion() and responses() API work with Qwen models +""" + +import json +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../../../..") +) # Adds the parent directory to the system path + +import litellm +from litellm.llms.vertex_ai.common_utils import is_global_only_vertex_model +from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.types.llms.vertex_ai import VertexPartnerProvider + + +class TestQwenGlobalOnlyDetection: + """Test that Qwen models are correctly identified as global-only.""" + + @pytest.mark.parametrize( + "model", + [ + "vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas", + "vertex_ai/qwen/qwen3-next-80b-a3b-thinking-maas", + "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas", + "vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas", + ], + ) + def test_qwen_models_are_global_only(self, model): + """Test that Qwen MaaS models are identified as global-only.""" + # This test requires the model_cost to have supported_regions: ["global"] + # If the model is not in model_cost, it should return False (fallback behavior) + result = is_global_only_vertex_model(model) + # Note: This will return True only if the model is in model_cost with supported_regions: ["global"] + # If running without the updated model_cost, this may return False + assert isinstance(result, bool) + + def test_non_global_model_returns_false(self): + """Test that non-global models return False.""" + result = is_global_only_vertex_model("vertex_ai/gemini-1.5-pro") + assert result is False + + def test_unknown_model_returns_false(self): + """Test that unknown models return False (fallback behavior).""" + result = is_global_only_vertex_model("vertex_ai/unknown-model-xyz") + assert result is False + + +class TestVertexBaseGetVertexRegion: + """Test the get_vertex_region method.""" + + def test_global_only_model_returns_global(self): + """Test that global-only models return 'global' regardless of input.""" + vertex_base = VertexBase() + + with patch( + "litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model", + return_value=True, + ): + result = vertex_base.get_vertex_region( + vertex_region="us-central1", + model="vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas", + ) + assert result == "global" + + def test_global_only_model_with_none_returns_global(self): + """Test that global-only models return 'global' even with None input.""" + vertex_base = VertexBase() + + with patch( + "litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model", + return_value=True, + ): + result = vertex_base.get_vertex_region( + vertex_region=None, + model="vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas", + ) + assert result == "global" + + def test_non_global_model_uses_provided_region(self): + """Test that non-global models use the provided region.""" + vertex_base = VertexBase() + + with patch( + "litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model", + return_value=False, + ): + result = vertex_base.get_vertex_region( + vertex_region="europe-west1", + model="vertex_ai/gemini-1.5-pro", + ) + assert result == "europe-west1" + + def test_non_global_model_fallback_to_us_central1(self): + """Test that non-global models with None region fallback to us-central1.""" + vertex_base = VertexBase() + + with patch( + "litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model", + return_value=False, + ): + result = vertex_base.get_vertex_region( + vertex_region=None, + model="vertex_ai/gemini-1.5-pro", + ) + assert result == "us-central1" + + +class TestCreateVertexURLGlobal: + """Test that create_vertex_url handles global location correctly.""" + + def test_global_location_url_format(self): + """Test that global location produces correct URL without region prefix.""" + url = VertexBase.create_vertex_url( + vertex_location="global", + vertex_project="test-project", + partner=VertexPartnerProvider.llama, + stream=False, + model="qwen/qwen3-next-80b-a3b-instruct-maas", + ) + + # Global URL should NOT have region prefix + assert url.startswith("https://aiplatform.googleapis.com") + assert "global-aiplatform.googleapis.com" not in url + assert "/locations/global/" in url + + def test_regional_location_url_format(self): + """Test that regional location produces correct URL with region prefix.""" + url = VertexBase.create_vertex_url( + vertex_location="us-central1", + vertex_project="test-project", + partner=VertexPartnerProvider.llama, + stream=False, + model="openai/gpt-oss-20b-maas", + ) + + # Regional URL should have region prefix + assert url.startswith("https://us-central1-aiplatform.googleapis.com") + assert "/locations/us-central1/" in url + + +@pytest.mark.asyncio +async def test_vertex_ai_qwen_global_endpoint_url(): + """ + Test that Qwen models use the global endpoint URL. + """ + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexLLM, + ) + + # Mock response + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {} + mock_response.json.return_value = { + "id": "chatcmpl-qwen-test", + "object": "chat.completion", + "created": 1234567890, + "model": "qwen/qwen3-next-80b-a3b-instruct-maas", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello! How can I help you today?", + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18}, + } + + client = AsyncHTTPHandler() + + async def mock_post_func(*args, **kwargs): + return mock_response + + with patch.object(client, "post", side_effect=mock_post_func) as mock_post, patch.object( + VertexLLM, "_ensure_access_token", return_value=("fake-token", "test-project") + ), patch( + "litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model", + return_value=True, + ): + response = await litellm.acompletion( + model="vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas", + messages=[{"role": "user", "content": "Hello"}], + vertex_ai_project="test-project", + client=client, + ) + + # Verify the mock was called + mock_post.assert_called_once() + + # Get the call arguments + call_args = mock_post.call_args + called_url = call_args.kwargs["url"] + + # Verify the URL uses global endpoint (no region prefix) + assert called_url.startswith("https://aiplatform.googleapis.com") + assert "global-aiplatform.googleapis.com" not in called_url + assert "/locations/global/" in called_url + assert "/endpoints/openapi/chat/completions" in called_url + + # Verify response + assert response.model == "qwen/qwen3-next-80b-a3b-instruct-maas" + + +class TestGetSupportedRegions: + """Test that get_supported_regions correctly reads from model_cost.""" + + def test_get_supported_regions_returns_list(self): + """Test that get_supported_regions returns a list when model has supported_regions.""" + # Mock the model_cost to have supported_regions + with patch.dict( + litellm.model_cost, + { + "vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": { + "supported_regions": ["global"], + "litellm_provider": "vertex_ai-qwen_models", + } + }, + ): + regions = litellm.utils.get_supported_regions( + model="vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas", + custom_llm_provider="vertex_ai", + ) + assert regions == ["global"] + + def test_get_supported_regions_returns_none_when_not_set(self): + """Test that get_supported_regions returns None when model doesn't have supported_regions.""" + # Mock the model_cost without supported_regions + with patch.dict( + litellm.model_cost, + { + "vertex_ai/gemini-1.5-pro": { + "litellm_provider": "vertex_ai", + } + }, + ): + regions = litellm.utils.get_supported_regions( + model="vertex_ai/gemini-1.5-pro", + custom_llm_provider="vertex_ai", + ) + assert regions is None + + def test_get_supported_regions_returns_none_for_unknown_model(self): + """Test that get_supported_regions returns None for unknown models.""" + regions = litellm.utils.get_supported_regions( + model="vertex_ai/unknown-model-xyz", + custom_llm_provider="vertex_ai", + ) + assert regions is None From b3f1696946f6fe0837df8fde75dcc235a5297e3b Mon Sep 17 00:00:00 2001 From: Chesars Date: Fri, 30 Jan 2026 14:39:13 -0300 Subject: [PATCH 07/22] refactor(vertex_ai): reuse get_vertex_base_url for URL construction Use existing get_vertex_base_url from common_utils instead of duplicating the global vs regional URL logic in create_vertex_url and get_api_base. --- litellm/llms/vertex_ai/vertex_llm_base.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 1310e66e1db..4613b6a5715 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -20,6 +20,7 @@ from .common_utils import ( _get_vertex_url, all_gemini_url_modes, get_vertex_base_model_name, + get_vertex_base_url, is_global_only_vertex_model, ) @@ -200,12 +201,7 @@ class VertexBase: ) -> str: if api_base: return api_base - elif vertex_location == "global": - return "https://aiplatform.googleapis.com" - elif vertex_location: - return f"https://{vertex_location}-aiplatform.googleapis.com" - else: - return f"https://{self.get_default_vertex_location()}-aiplatform.googleapis.com" + return get_vertex_base_url(vertex_location or self.get_default_vertex_location()) @staticmethod def create_vertex_url( @@ -218,12 +214,8 @@ class VertexBase: ) -> str: """Return the base url for the vertex partner models""" - # For global location, use the non-regional URL if api_base is None: - if vertex_location == "global": - api_base = "https://aiplatform.googleapis.com" - else: - api_base = f"https://{vertex_location}-aiplatform.googleapis.com" + api_base = get_vertex_base_url(vertex_location) if partner == VertexPartnerProvider.llama: return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/endpoints/openapi/chat/completions" elif partner == VertexPartnerProvider.mistralai: From 3e04e2020542e2f31f3eb85da5602ced42625985 Mon Sep 17 00:00:00 2001 From: Aarish Alam Date: Sat, 31 Jan 2026 23:22:19 +0530 Subject: [PATCH 08/22] =?UTF-8?q?=F0=9F=90=9B=20Bug=20Fix=20#19642=20:=20b?= =?UTF-8?q?ug=20in=20Vertex=20AI=20context=20caching=20=20(#19657)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * add vertex tests * add uperbound * add pagination tests --- .../vertex_ai_context_caching.py | 172 +++++---- .../test_vertex_ai_context_caching.py | 353 ++++++++++++++++++ 2 files changed, 460 insertions(+), 65 deletions(-) diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 289963e917a..ed4d2d6a740 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -27,6 +27,8 @@ local_cache_obj = Cache( type=LiteLLMCacheType.LOCAL ) # only used for calling 'get_cache_key' function +MAX_PAGINATION_PAGES = 100 # Reasonable upper bound for pagination + class ContextCachingEndpoints(VertexBase): """ @@ -115,7 +117,7 @@ class ContextCachingEndpoints(VertexBase): - None """ - _, url = self._get_token_and_url_context_caching( + _, base_url = self._get_token_and_url_context_caching( gemini_api_key=api_key, custom_llm_provider=custom_llm_provider, api_base=api_base, @@ -123,43 +125,63 @@ class ContextCachingEndpoints(VertexBase): vertex_location=vertex_location, vertex_auth_header=vertex_auth_header ) - try: - ## LOGGING - logging_obj.pre_call( - input="", - api_key="", - additional_args={ - "complete_input_dict": {}, - "api_base": url, - "headers": headers, - }, - ) - resp = client.get(url=url, headers=headers) - resp.raise_for_status() - except httpx.HTTPStatusError as e: - if e.response.status_code == 403: + page_token: Optional[str] = None + + # Iterate through all pages + for _ in range(MAX_PAGINATION_PAGES): + # Build URL with pagination token if present + if page_token: + separator = "&" if "?" in base_url else "?" + url = f"{base_url}{separator}pageToken={page_token}" + else: + url = base_url + + try: + ## LOGGING + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "complete_input_dict": {}, + "api_base": url, + "headers": headers, + }, + ) + + resp = client.get(url=url, headers=headers) + resp.raise_for_status() + except httpx.HTTPStatusError as e: + if e.response.status_code == 403: + return None + raise VertexAIError( + status_code=e.response.status_code, message=e.response.text + ) + except Exception as e: + raise VertexAIError(status_code=500, message=str(e)) + + raw_response = resp.json() + logging_obj.post_call(original_response=raw_response) + + if "cachedContents" not in raw_response: return None - raise VertexAIError( - status_code=e.response.status_code, message=e.response.text - ) - except Exception as e: - raise VertexAIError(status_code=500, message=str(e)) - raw_response = resp.json() - logging_obj.post_call(original_response=raw_response) - if "cachedContents" not in raw_response: - return None + all_cached_items = CachedContentListAllResponseBody(**raw_response) - all_cached_items = CachedContentListAllResponseBody(**raw_response) + if "cachedContents" not in all_cached_items: + return None - if "cachedContents" not in all_cached_items: - return None + # Check current page for matching cache_key + for cached_item in all_cached_items["cachedContents"]: + display_name = cached_item.get("displayName") + if display_name is not None and display_name == cache_key: + return cached_item.get("name") - for cached_item in all_cached_items["cachedContents"]: - display_name = cached_item.get("displayName") - if display_name is not None and display_name == cache_key: - return cached_item.get("name") + # Check if there are more pages + page_token = all_cached_items.get("nextPageToken") + if not page_token: + # No more pages, cache not found + break return None @@ -187,7 +209,7 @@ class ContextCachingEndpoints(VertexBase): - None """ - _, url = self._get_token_and_url_context_caching( + _, base_url = self._get_token_and_url_context_caching( gemini_api_key=api_key, custom_llm_provider=custom_llm_provider, api_base=api_base, @@ -195,43 +217,63 @@ class ContextCachingEndpoints(VertexBase): vertex_location=vertex_location, vertex_auth_header=vertex_auth_header ) - try: - ## LOGGING - logging_obj.pre_call( - input="", - api_key="", - additional_args={ - "complete_input_dict": {}, - "api_base": url, - "headers": headers, - }, - ) - resp = await client.get(url=url, headers=headers) - resp.raise_for_status() - except httpx.HTTPStatusError as e: - if e.response.status_code == 403: + page_token: Optional[str] = None + + # Iterate through all pages + for _ in range(MAX_PAGINATION_PAGES): + # Build URL with pagination token if present + if page_token: + separator = "&" if "?" in base_url else "?" + url = f"{base_url}{separator}pageToken={page_token}" + else: + url = base_url + + try: + ## LOGGING + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "complete_input_dict": {}, + "api_base": url, + "headers": headers, + }, + ) + + resp = await client.get(url=url, headers=headers) + resp.raise_for_status() + except httpx.HTTPStatusError as e: + if e.response.status_code == 403: + return None + raise VertexAIError( + status_code=e.response.status_code, message=e.response.text + ) + except Exception as e: + raise VertexAIError(status_code=500, message=str(e)) + + raw_response = resp.json() + logging_obj.post_call(original_response=raw_response) + + if "cachedContents" not in raw_response: return None - raise VertexAIError( - status_code=e.response.status_code, message=e.response.text - ) - except Exception as e: - raise VertexAIError(status_code=500, message=str(e)) - raw_response = resp.json() - logging_obj.post_call(original_response=raw_response) - if "cachedContents" not in raw_response: - return None + all_cached_items = CachedContentListAllResponseBody(**raw_response) - all_cached_items = CachedContentListAllResponseBody(**raw_response) + if "cachedContents" not in all_cached_items: + return None - if "cachedContents" not in all_cached_items: - return None + # Check current page for matching cache_key + for cached_item in all_cached_items["cachedContents"]: + display_name = cached_item.get("displayName") + if display_name is not None and display_name == cache_key: + return cached_item.get("name") - for cached_item in all_cached_items["cachedContents"]: - display_name = cached_item.get("displayName") - if display_name is not None and display_name == cache_key: - return cached_item.get("name") + # Check if there are more pages + page_token = all_cached_items.get("nextPageToken") + if not page_token: + # No more pages, cache not found + break return None @@ -501,4 +543,4 @@ class ContextCachingEndpoints(VertexBase): pass async def async_get_cache(self): - pass + pass \ No newline at end of file diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index e9d14d4e18f..a47d026c169 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( + MAX_PAGINATION_PAGES, ContextCachingEndpoints, ) @@ -787,6 +788,358 @@ class TestContextCachingEndpoints: assert original_tools == self.sample_tools +class TestCheckCachePagination: + """Test pagination logic in check_cache and async_check_cache methods.""" + + def setup_method(self): + """Setup for each test method""" + self.context_caching = ContextCachingEndpoints() + self.mock_logging = MagicMock(spec=Logging) + self.mock_client = MagicMock(spec=HTTPHandler) + self.mock_async_client = MagicMock(spec=AsyncHTTPHandler) + + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + def test_check_cache_pagination_finds_cache_on_second_page( + self, mock_get_token_url, custom_llm_provider + ): + """Test that check_cache correctly handles pagination and finds cache on second page""" + # Setup + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "target_cache_key" + + # Mock first page response (no match, has nextPageToken) + first_page_response = MagicMock() + first_page_response.json.return_value = { + "cachedContents": [ + {"name": "cache_1", "displayName": "cache_key_1"}, + {"name": "cache_2", "displayName": "cache_key_2"}, + ], + "nextPageToken": "token_page_2", + } + + # Mock second page response (has match, no nextPageToken) + second_page_response = MagicMock() + second_page_response.json.return_value = { + "cachedContents": [ + {"name": "cache_3", "displayName": cache_key_to_find}, + {"name": "cache_4", "displayName": "cache_key_4"}, + ] + } + + # Setup mock client to return different responses + self.mock_client.get.side_effect = [first_page_response, second_page_response] + + # Execute + result = self.context_caching.check_cache( + cache_key=cache_key_to_find, + client=self.mock_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + # Assert + assert result == "cache_3" + assert self.mock_client.get.call_count == 2 + # Check that second call includes pageToken + second_call_url = self.mock_client.get.call_args_list[1].kwargs["url"] + assert "pageToken=token_page_2" in second_call_url + + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + def test_check_cache_pagination_stops_when_no_next_token( + self, mock_get_token_url, custom_llm_provider + ): + """Test that check_cache stops pagination when no nextPageToken is present""" + # Setup + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "nonexistent_cache_key" + + # Mock response without nextPageToken + response = MagicMock() + response.json.return_value = { + "cachedContents": [ + {"name": "cache_1", "displayName": "cache_key_1"}, + {"name": "cache_2", "displayName": "cache_key_2"}, + ] + } + + self.mock_client.get.return_value = response + + # Execute + result = self.context_caching.check_cache( + cache_key=cache_key_to_find, + client=self.mock_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + # Assert + assert result is None + assert self.mock_client.get.call_count == 1 + + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + def test_check_cache_pagination_multiple_pages( + self, mock_get_token_url, custom_llm_provider + ): + """Test that check_cache correctly iterates through multiple pages""" + # Setup + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "target_cache_key" + + # Mock three pages + page1 = MagicMock() + page1.json.return_value = { + "cachedContents": [{"name": "cache_1", "displayName": "cache_key_1"}], + "nextPageToken": "token_page_2", + } + + page2 = MagicMock() + page2.json.return_value = { + "cachedContents": [{"name": "cache_2", "displayName": "cache_key_2"}], + "nextPageToken": "token_page_3", + } + + page3 = MagicMock() + page3.json.return_value = { + "cachedContents": [{"name": "cache_3", "displayName": cache_key_to_find}], + } + + self.mock_client.get.side_effect = [page1, page2, page3] + + # Execute + result = self.context_caching.check_cache( + cache_key=cache_key_to_find, + client=self.mock_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + # Assert + assert result == "cache_3" + assert self.mock_client.get.call_count == 3 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + async def test_async_check_cache_pagination_finds_cache_on_second_page( + self, mock_get_token_url, custom_llm_provider + ): + """Test that async_check_cache correctly handles pagination and finds cache on second page""" + # Setup + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "target_cache_key" + + # Mock first page response (no match, has nextPageToken) + first_page_response = MagicMock() + first_page_response.json.return_value = { + "cachedContents": [ + {"name": "cache_1", "displayName": "cache_key_1"}, + {"name": "cache_2", "displayName": "cache_key_2"}, + ], + "nextPageToken": "token_page_2", + } + + # Mock second page response (has match, no nextPageToken) + second_page_response = MagicMock() + second_page_response.json.return_value = { + "cachedContents": [ + {"name": "cache_3", "displayName": cache_key_to_find}, + {"name": "cache_4", "displayName": "cache_key_4"}, + ] + } + + # Setup mock async client to return different responses + self.mock_async_client.get = AsyncMock( + side_effect=[first_page_response, second_page_response] + ) + + # Execute + result = await self.context_caching.async_check_cache( + cache_key=cache_key_to_find, + client=self.mock_async_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + # Assert + assert result == "cache_3" + assert self.mock_async_client.get.call_count == 2 + # Check that second call includes pageToken + second_call_url = self.mock_async_client.get.call_args_list[1].kwargs["url"] + assert "pageToken=token_page_2" in second_call_url + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + async def test_async_check_cache_pagination_stops_when_no_next_token( + self, mock_get_token_url, custom_llm_provider + ): + """Test that async_check_cache stops pagination when no nextPageToken is present""" + # Setup + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "nonexistent_cache_key" + + # Mock response without nextPageToken + response = MagicMock() + response.json.return_value = { + "cachedContents": [ + {"name": "cache_1", "displayName": "cache_key_1"}, + {"name": "cache_2", "displayName": "cache_key_2"}, + ] + } + + self.mock_async_client.get = AsyncMock(return_value=response) + + # Execute + result = await self.context_caching.async_check_cache( + cache_key=cache_key_to_find, + client=self.mock_async_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + # Assert + assert result is None + assert self.mock_async_client.get.call_count == 1 + + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + def test_check_cache_pagination_max_pages_limit( + self, mock_get_token_url, custom_llm_provider + ): + """Test that pagination stops after MAX_PAGINATION_PAGES iterations""" + # Setup + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "nonexistent_cache_key" + + # Create mock response that always has nextPageToken (infinite pagination scenario) + def create_page_response(page_num): + response = MagicMock() + response.json.return_value = { + "cachedContents": [ + {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} + ], + "nextPageToken": f"token_page_{page_num + 1}", + } + return response + + # Create MAX_PAGINATION_PAGES responses, each with a nextPageToken + self.mock_client.get.side_effect = [ + create_page_response(i) for i in range(MAX_PAGINATION_PAGES) + ] + + # Execute + result = self.context_caching.check_cache( + cache_key=cache_key_to_find, + client=self.mock_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + # Assert - should return None after exhausting all pages without finding match + assert result is None + # Verify exactly MAX_PAGINATION_PAGES API calls were made (not more) + assert self.mock_client.get.call_count == MAX_PAGINATION_PAGES + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + async def test_async_check_cache_pagination_max_pages_limit( + self, mock_get_token_url, custom_llm_provider + ): + """Test that async pagination stops after MAX_PAGINATION_PAGES iterations""" + # Setup + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "nonexistent_cache_key" + + # Create mock response that always has nextPageToken (infinite pagination scenario) + def create_page_response(page_num): + response = MagicMock() + response.json.return_value = { + "cachedContents": [ + {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} + ], + "nextPageToken": f"token_page_{page_num + 1}", + } + return response + + # Create MAX_PAGINATION_PAGES responses, each with a nextPageToken + self.mock_async_client.get = AsyncMock( + side_effect=[create_page_response(i) for i in range(MAX_PAGINATION_PAGES)] + ) + + # Execute + result = await self.context_caching.async_check_cache( + cache_key=cache_key_to_find, + client=self.mock_async_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + # Assert - should return None after exhausting all pages without finding match + assert result is None + # Verify exactly MAX_PAGINATION_PAGES async API calls were made (not more) + assert self.mock_async_client.get.call_count == MAX_PAGINATION_PAGES + + class TestVertexAIGlobalLocation: """Test global location handling in context caching.""" From 72e519345149f4b305645c51943b0f2cfd6c8acd Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Sun, 1 Feb 2026 00:01:33 +0530 Subject: [PATCH 09/22] fix: models loadbalancing billing issue by filter (#18891) (#19220) * fix: models loadbalancing billing issue by filter (#18891) * fix: models loadbalancing billing issue by filter * fix: separate key and team access groups in metadata * fix: lint issues --- litellm/proxy/auth/model_checks.py | 25 +- litellm/proxy/litellm_pre_call_utils.py | 31 +++ litellm/router.py | 12 +- litellm/router_utils/common_utils.py | 79 +++++- ...est_filter_deployments_by_access_groups.py | 227 ++++++++++++++++++ 5 files changed, 368 insertions(+), 6 deletions(-) create mode 100644 tests/test_litellm/router_unit_tests/test_filter_deployments_by_access_groups.py diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index 71ae1348f39..af2574d88ee 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -64,6 +64,27 @@ def _get_models_from_access_groups( return all_models +def get_access_groups_from_models( + model_access_groups: Dict[str, List[str]], + models: List[str], +) -> List[str]: + """ + Extract access group names from a models list. + + Given a models list like ["gpt-4", "beta-models", "claude-v1"] + and access groups like {"beta-models": ["gpt-5", "gpt-6"]}, + returns ["beta-models"]. + + This is used to pass allowed access groups to the router for filtering + deployments during load balancing (GitHub issue #18333). + """ + access_groups = [] + for model in models: + if model in model_access_groups: + access_groups.append(model) + return access_groups + + async def get_mcp_server_ids( user_api_key_dict: UserAPIKeyAuth, ) -> List[str]: @@ -80,7 +101,6 @@ async def get_mcp_server_ids( # Make a direct SQL query to get just the mcp_servers try: - result = await prisma_client.db.litellm_objectpermissiontable.find_unique( where={"object_permission_id": user_api_key_dict.object_permission_id}, ) @@ -176,6 +196,7 @@ def get_complete_model_list( """ unique_models = [] + def append_unique(models): for model in models: if model not in unique_models: @@ -188,7 +209,7 @@ def get_complete_model_list( else: append_unique(proxy_model_list) if include_model_access_groups: - append_unique(list(model_access_groups.keys())) # TODO: keys order + append_unique(list(model_access_groups.keys())) # TODO: keys order if user_model: append_unique([user_model]) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 9be78264e85..72f23e609ab 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1021,6 +1021,37 @@ async def add_litellm_data_to_request( # noqa: PLR0915 "user_api_key_user_max_budget" ] = user_api_key_dict.user_max_budget + # Extract allowed access groups for router filtering (GitHub issue #18333) + # This allows the router to filter deployments based on key's and team's access groups + # NOTE: We keep key and team access groups SEPARATE because a key doesn't always + # inherit all team access groups (per maintainer feedback). + if llm_router is not None: + from litellm.proxy.auth.model_checks import get_access_groups_from_models + + model_access_groups = llm_router.get_model_access_groups() + + # Key-level access groups (from user_api_key_dict.models) + key_models = list(user_api_key_dict.models) if user_api_key_dict.models else [] + key_allowed_access_groups = get_access_groups_from_models( + model_access_groups=model_access_groups, models=key_models + ) + if key_allowed_access_groups: + data[_metadata_variable_name][ + "user_api_key_allowed_access_groups" + ] = key_allowed_access_groups + + # Team-level access groups (from user_api_key_dict.team_models) + team_models = ( + list(user_api_key_dict.team_models) if user_api_key_dict.team_models else [] + ) + team_allowed_access_groups = get_access_groups_from_models( + model_access_groups=model_access_groups, models=team_models + ) + if team_allowed_access_groups: + data[_metadata_variable_name][ + "user_api_key_team_allowed_access_groups" + ] = team_allowed_access_groups + data[_metadata_variable_name]["user_api_key_metadata"] = user_api_key_dict.metadata _headers = dict(request.headers) _headers.pop( diff --git a/litellm/router.py b/litellm/router.py index 6c191c8ab03..40d84fae410 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -88,6 +88,7 @@ from litellm.router_utils.clientside_credential_handler import ( is_clientside_credential, ) from litellm.router_utils.common_utils import ( + filter_deployments_by_access_groups, filter_team_based_models, filter_web_search_deployments, ) @@ -8075,10 +8076,17 @@ class Router: request_kwargs=request_kwargs, ) - verbose_router_logger.debug( - f"healthy_deployments after web search filter: {healthy_deployments}" + verbose_router_logger.debug(f"healthy_deployments after web search filter: {healthy_deployments}") + + # Filter by allowed access groups (GitHub issue #18333) + # This prevents cross-team load balancing when teams have models with same name in different access groups + healthy_deployments = filter_deployments_by_access_groups( + healthy_deployments=healthy_deployments, + request_kwargs=request_kwargs, ) + verbose_router_logger.debug(f"healthy_deployments after access group filter: {healthy_deployments}") + if isinstance(healthy_deployments, dict): return healthy_deployments diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 10acc343abd..2c0ea5976d6 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -75,6 +75,7 @@ def filter_team_based_models( if deployment.get("model_info", {}).get("id") not in ids_to_remove ] + def _deployment_supports_web_search(deployment: Dict) -> bool: """ Check if a deployment supports web search. @@ -112,7 +113,7 @@ def filter_web_search_deployments( is_web_search_request = False tools = request_kwargs.get("tools") or [] for tool in tools: - # These are the two websearch tools for OpenAI / Azure. + # These are the two websearch tools for OpenAI / Azure. if tool.get("type") == "web_search" or tool.get("type") == "web_search_preview": is_web_search_request = True break @@ -121,8 +122,82 @@ def filter_web_search_deployments( return healthy_deployments # Filter out deployments that don't support web search - final_deployments = [d for d in healthy_deployments if _deployment_supports_web_search(d)] + final_deployments = [ + d for d in healthy_deployments if _deployment_supports_web_search(d) + ] if len(healthy_deployments) > 0 and len(final_deployments) == 0: verbose_logger.warning("No deployments support web search for request") return final_deployments + +def filter_deployments_by_access_groups( + healthy_deployments: Union[List[Dict], Dict], + request_kwargs: Optional[Dict] = None, +) -> Union[List[Dict], Dict]: + """ + Filter deployments to only include those matching the user's allowed access groups. + + Reads from TWO separate metadata fields (per maintainer feedback): + - `user_api_key_allowed_access_groups`: Access groups from the API Key's models. + - `user_api_key_team_allowed_access_groups`: Access groups from the Team's models. + + A deployment is included if its access_groups overlap with EITHER the key's + or the team's allowed access groups. Deployments with no access_groups are + always included (not restricted). + + This prevents cross-team load balancing when multiple teams have models with + the same name but in different access groups (GitHub issue #18333). + """ + if request_kwargs is None: + return healthy_deployments + + if isinstance(healthy_deployments, dict): + return healthy_deployments + + metadata = request_kwargs.get("metadata") or {} + litellm_metadata = request_kwargs.get("litellm_metadata") or {} + + # Gather key-level allowed access groups + key_allowed_access_groups = ( + metadata.get("user_api_key_allowed_access_groups") + or litellm_metadata.get("user_api_key_allowed_access_groups") + or [] + ) + + # Gather team-level allowed access groups + team_allowed_access_groups = ( + metadata.get("user_api_key_team_allowed_access_groups") + or litellm_metadata.get("user_api_key_team_allowed_access_groups") + or [] + ) + + # Combine both for the final allowed set + combined_allowed_access_groups = list(key_allowed_access_groups) + list( + team_allowed_access_groups + ) + + # If no access groups specified from either source, return all deployments (backwards compatible) + if not combined_allowed_access_groups: + return healthy_deployments + + allowed_set = set(combined_allowed_access_groups) + filtered = [] + for deployment in healthy_deployments: + model_info = deployment.get("model_info") or {} + deployment_access_groups = model_info.get("access_groups") or [] + + # If deployment has no access groups, include it (not restricted) + if not deployment_access_groups: + filtered.append(deployment) + continue + + # Include if any of deployment's groups overlap with allowed groups + if set(deployment_access_groups) & allowed_set: + filtered.append(deployment) + + if len(healthy_deployments) > 0 and len(filtered) == 0: + verbose_logger.warning( + f"No deployments match allowed access groups {combined_allowed_access_groups}" + ) + + return filtered diff --git a/tests/test_litellm/router_unit_tests/test_filter_deployments_by_access_groups.py b/tests/test_litellm/router_unit_tests/test_filter_deployments_by_access_groups.py new file mode 100644 index 00000000000..9ac5072c5d8 --- /dev/null +++ b/tests/test_litellm/router_unit_tests/test_filter_deployments_by_access_groups.py @@ -0,0 +1,227 @@ +""" +Unit tests for filter_deployments_by_access_groups function. + +Tests the fix for GitHub issue #18333: Models loadbalanced outside of Model Access Group. +""" + +import pytest + +from litellm.router_utils.common_utils import filter_deployments_by_access_groups + + +class TestFilterDeploymentsByAccessGroups: + """Tests for the filter_deployments_by_access_groups function.""" + + def test_no_filter_when_no_access_groups_in_metadata(self): + """When no allowed_access_groups in metadata, return all deployments.""" + deployments = [ + {"model_info": {"id": "1", "access_groups": ["AG1"]}}, + {"model_info": {"id": "2", "access_groups": ["AG2"]}}, + ] + request_kwargs = {"metadata": {"user_api_key_team_id": "team-1"}} + + result = filter_deployments_by_access_groups( + healthy_deployments=deployments, + request_kwargs=request_kwargs, + ) + + assert len(result) == 2 # All deployments returned + + def test_filter_to_single_access_group(self): + """Filter to only deployments matching allowed access group.""" + deployments = [ + {"model_info": {"id": "1", "access_groups": ["AG1"]}}, + {"model_info": {"id": "2", "access_groups": ["AG2"]}}, + ] + request_kwargs = {"metadata": {"user_api_key_allowed_access_groups": ["AG2"]}} + + result = filter_deployments_by_access_groups( + healthy_deployments=deployments, + request_kwargs=request_kwargs, + ) + + assert len(result) == 1 + assert result[0]["model_info"]["id"] == "2" + + def test_filter_with_multiple_allowed_groups(self): + """Filter with multiple allowed access groups.""" + deployments = [ + {"model_info": {"id": "1", "access_groups": ["AG1"]}}, + {"model_info": {"id": "2", "access_groups": ["AG2"]}}, + {"model_info": {"id": "3", "access_groups": ["AG3"]}}, + ] + request_kwargs = { + "metadata": {"user_api_key_allowed_access_groups": ["AG1", "AG2"]} + } + + result = filter_deployments_by_access_groups( + healthy_deployments=deployments, + request_kwargs=request_kwargs, + ) + + assert len(result) == 2 + ids = [d["model_info"]["id"] for d in result] + assert "1" in ids + assert "2" in ids + assert "3" not in ids + + def test_deployment_with_multiple_access_groups(self): + """Deployment with multiple access groups should match if any overlap.""" + deployments = [ + {"model_info": {"id": "1", "access_groups": ["AG1", "AG2"]}}, + {"model_info": {"id": "2", "access_groups": ["AG3"]}}, + ] + request_kwargs = {"metadata": {"user_api_key_allowed_access_groups": ["AG2"]}} + + result = filter_deployments_by_access_groups( + healthy_deployments=deployments, + request_kwargs=request_kwargs, + ) + + assert len(result) == 1 + assert result[0]["model_info"]["id"] == "1" + + def test_deployment_without_access_groups_included(self): + """Deployments without access groups should be included (not restricted).""" + deployments = [ + {"model_info": {"id": "1", "access_groups": ["AG1"]}}, + {"model_info": {"id": "2"}}, # No access_groups + {"model_info": {"id": "3", "access_groups": []}}, # Empty access_groups + ] + request_kwargs = {"metadata": {"user_api_key_allowed_access_groups": ["AG2"]}} + + result = filter_deployments_by_access_groups( + healthy_deployments=deployments, + request_kwargs=request_kwargs, + ) + + # Should include deployments 2 and 3 (no restrictions) + assert len(result) == 2 + ids = [d["model_info"]["id"] for d in result] + assert "2" in ids + assert "3" in ids + + def test_dict_deployment_passes_through(self): + """When deployment is a dict (specific deployment), pass through.""" + deployment = {"model_info": {"id": "1", "access_groups": ["AG1"]}} + request_kwargs = {"metadata": {"user_api_key_allowed_access_groups": ["AG2"]}} + + result = filter_deployments_by_access_groups( + healthy_deployments=deployment, + request_kwargs=request_kwargs, + ) + + assert result == deployment # Unchanged + + def test_none_request_kwargs_passes_through(self): + """When request_kwargs is None, return deployments unchanged.""" + deployments = [ + {"model_info": {"id": "1", "access_groups": ["AG1"]}}, + ] + + result = filter_deployments_by_access_groups( + healthy_deployments=deployments, + request_kwargs=None, + ) + + assert result == deployments + + def test_litellm_metadata_fallback(self): + """Should also check litellm_metadata for allowed access groups.""" + deployments = [ + {"model_info": {"id": "1", "access_groups": ["AG1"]}}, + {"model_info": {"id": "2", "access_groups": ["AG2"]}}, + ] + request_kwargs = { + "litellm_metadata": {"user_api_key_allowed_access_groups": ["AG1"]} + } + + result = filter_deployments_by_access_groups( + healthy_deployments=deployments, + request_kwargs=request_kwargs, + ) + + assert len(result) == 1 + assert result[0]["model_info"]["id"] == "1" + + +def test_filter_deployments_by_access_groups_issue_18333(): + """ + Regression test for GitHub issue #18333. + + Scenario: Two models named 'gpt-5' in different access groups (AG1, AG2). + Team2 has access to AG2 only. When Team2 requests 'gpt-5', only the AG2 + deployment should be available for load balancing. + """ + deployments = [ + { + "model_name": "gpt-5", + "litellm_params": {"model": "gpt-4.1", "api_key": "key-1"}, + "model_info": {"id": "ag1-deployment", "access_groups": ["AG1"]}, + }, + { + "model_name": "gpt-5", + "litellm_params": {"model": "gpt-4o", "api_key": "key-2"}, + "model_info": {"id": "ag2-deployment", "access_groups": ["AG2"]}, + }, + ] + + # Team2's request with allowed access groups + request_kwargs = { + "metadata": { + "user_api_key_team_id": "team-2", + "user_api_key_allowed_access_groups": ["AG2"], + } + } + + result = filter_deployments_by_access_groups( + healthy_deployments=deployments, + request_kwargs=request_kwargs, + ) + + # Only AG2 deployment should be returned + assert len(result) == 1 + assert result[0]["model_info"]["id"] == "ag2-deployment" + assert result[0]["litellm_params"]["model"] == "gpt-4o" + + +def test_get_access_groups_from_models(): + """ + Test the helper function that extracts access group names from models list. + This is used by the proxy to populate user_api_key_allowed_access_groups. + """ + from litellm.proxy.auth.model_checks import get_access_groups_from_models + + # Setup: access groups definition + model_access_groups = { + "AG1": ["gpt-4", "gpt-5"], + "AG2": ["claude-v1", "claude-v2"], + "beta-models": ["gpt-5-turbo"], + } + + # Test 1: Extract access groups from models list + models = ["gpt-4", "AG1", "AG2", "some-other-model"] + result = get_access_groups_from_models( + model_access_groups=model_access_groups, models=models + ) + assert set(result) == {"AG1", "AG2"} + + # Test 2: No access groups in models list + models = ["gpt-4", "claude-v1", "some-model"] + result = get_access_groups_from_models( + model_access_groups=model_access_groups, models=models + ) + assert result == [] + + # Test 3: Empty models list + result = get_access_groups_from_models( + model_access_groups=model_access_groups, models=[] + ) + assert result == [] + + # Test 4: All access groups + models = ["AG1", "AG2", "beta-models"] + result = get_access_groups_from_models( + model_access_groups=model_access_groups, models=models + ) + assert set(result) == {"AG1", "AG2", "beta-models"} From 61a84e9fdbea537f4cd596d5ea16dfb1a1753ada Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Sat, 31 Jan 2026 15:51:40 -0300 Subject: [PATCH 10/22] fix(anthropic-adapter): truncate tool names exceeding OpenAI's 64-char limit (#20107) When using LiteLLM's Anthropic /v1/messages endpoint to route requests to OpenAI models, requests fail if any tool name exceeds OpenAI's 64-character limit. Anthropic API has no such limit, causing compatibility issues. Changes: - Add truncate_tool_name() function using {55-char-prefix}_{8-char-hash} format - Modify translate_anthropic_tools_to_openai() to truncate and return mapping - Modify translate_anthropic_tool_choice_to_openai() to truncate tool name - Restore original tool names in responses using the mapping - Support tool name restoration in streaming responses - Add backwards-compatible API (existing methods still work) The fix only applies when routing Anthropic requests to OpenAI models. Native Anthropic/Claude requests pass through unchanged. --- .../adapters/handler.py | 27 ++- .../adapters/streaming_iterator.py | 17 +- .../adapters/transformation.py | 186 ++++++++++++++++-- ...al_pass_through_adapters_transformation.py | 184 ++++++++++++++++- 4 files changed, 383 insertions(+), 31 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index 8fa7bb7e65e..a17eba75b3b 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -6,6 +6,7 @@ from typing import ( Dict, List, Optional, + Tuple, Union, cast, ) @@ -47,8 +48,14 @@ class LiteLLMMessagesToCompletionTransformationHandler: top_p: Optional[float] = None, output_format: Optional[Dict] = None, extra_kwargs: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """Prepare kwargs for litellm.completion/acompletion""" + ) -> Tuple[Dict[str, Any], Dict[str, str]]: + """Prepare kwargs for litellm.completion/acompletion. + + Returns: + Tuple of (completion_kwargs, tool_name_mapping) + - tool_name_mapping maps truncated tool names back to original names + for tools that exceeded OpenAI's 64-char limit + """ from litellm.litellm_core_utils.litellm_logging import ( Logging as LiteLLMLoggingObject, ) @@ -80,7 +87,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: if output_format: request_data["output_format"] = output_format - openai_request = ANTHROPIC_ADAPTER.translate_completion_input_params( + openai_request, tool_name_mapping = ANTHROPIC_ADAPTER.translate_completion_input_params_with_tool_mapping( request_data ) @@ -116,7 +123,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: ): completion_kwargs[key] = value - return completion_kwargs + return completion_kwargs, tool_name_mapping @staticmethod async def async_anthropic_messages_handler( @@ -137,7 +144,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: **kwargs, ) -> Union[AnthropicMessagesResponse, AsyncIterator]: """Handle non-Anthropic models asynchronously using the adapter""" - completion_kwargs = ( + completion_kwargs, tool_name_mapping = ( LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs( max_tokens=max_tokens, messages=messages, @@ -164,6 +171,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: ANTHROPIC_ADAPTER.translate_completion_output_params_streaming( completion_response, model=model, + tool_name_mapping=tool_name_mapping, ) ) if transformed_stream is not None: @@ -172,7 +180,8 @@ class LiteLLMMessagesToCompletionTransformationHandler: else: anthropic_response = ( ANTHROPIC_ADAPTER.translate_completion_output_params( - cast(ModelResponse, completion_response) + cast(ModelResponse, completion_response), + tool_name_mapping=tool_name_mapping, ) ) if anthropic_response is not None: @@ -222,7 +231,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: **kwargs, ) - completion_kwargs = ( + completion_kwargs, tool_name_mapping = ( LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs( max_tokens=max_tokens, messages=messages, @@ -249,6 +258,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: ANTHROPIC_ADAPTER.translate_completion_output_params_streaming( completion_response, model=model, + tool_name_mapping=tool_name_mapping, ) ) if transformed_stream is not None: @@ -257,7 +267,8 @@ class LiteLLMMessagesToCompletionTransformationHandler: else: anthropic_response = ( ANTHROPIC_ADAPTER.translate_completion_output_params( - cast(ModelResponse, completion_response) + cast(ModelResponse, completion_response), + tool_name_mapping=tool_name_mapping, ) ) if anthropic_response is not None: diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index 24524233ddf..aa2f0cc08f1 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -3,7 +3,7 @@ import json import traceback from collections import deque -from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, Literal, Optional +from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Iterator, Literal, Optional from litellm import verbose_logger from litellm._uuid import uuid @@ -44,9 +44,16 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): pending_new_content_block: bool = False chunk_queue: deque = deque() # Queue for buffering multiple chunks - def __init__(self, completion_stream: Any, model: str): + def __init__( + self, + completion_stream: Any, + model: str, + tool_name_mapping: Optional[Dict[str, str]] = None, + ): super().__init__(completion_stream) self.model = model + # Mapping of truncated tool names to original names (for OpenAI's 64-char limit) + self.tool_name_mapping = tool_name_mapping or {} def _create_initial_usage_delta(self) -> UsageDelta: """ @@ -401,6 +408,12 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): choices=chunk.choices # type: ignore ) + # Restore original tool name if it was truncated for OpenAI's 64-char limit + if block_type == "tool_use" and content_block_start.get("name"): + truncated_name = content_block_start.get("name", "") + original_name = self.tool_name_mapping.get(truncated_name, truncated_name) + content_block_start["name"] = original_name + if block_type != self.current_content_block_type: self.current_content_block_type = block_type self.current_content_block_start = content_block_start diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 0a64c7be4c7..444f821c20a 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -1,3 +1,4 @@ +import hashlib import json from typing import ( TYPE_CHECKING, @@ -12,6 +13,54 @@ from typing import ( cast, ) +# OpenAI has a 64-character limit for function/tool names +# Anthropic does not have this limit, so we need to truncate long names +OPENAI_MAX_TOOL_NAME_LENGTH = 64 +TOOL_NAME_HASH_LENGTH = 8 +TOOL_NAME_PREFIX_LENGTH = OPENAI_MAX_TOOL_NAME_LENGTH - TOOL_NAME_HASH_LENGTH - 1 # 55 + + +def truncate_tool_name(name: str) -> str: + """ + Truncate tool names that exceed OpenAI's 64-character limit. + + Uses format: {55-char-prefix}_{8-char-hash} to avoid collisions + when multiple tools have similar long names. + + Args: + name: The original tool name + + Returns: + The original name if <= 64 chars, otherwise truncated with hash + """ + if len(name) <= OPENAI_MAX_TOOL_NAME_LENGTH: + return name + + # Create deterministic hash from full name to avoid collisions + name_hash = hashlib.sha256(name.encode()).hexdigest()[:TOOL_NAME_HASH_LENGTH] + return f"{name[:TOOL_NAME_PREFIX_LENGTH]}_{name_hash}" + + +def create_tool_name_mapping( + tools: List[Dict[str, Any]], +) -> Dict[str, str]: + """ + Create a mapping of truncated tool names to original names. + + Args: + tools: List of tool definitions with 'name' field + + Returns: + Dict mapping truncated names to original names (only for truncated tools) + """ + mapping: Dict[str, str] = {} + for tool in tools: + original_name = tool.get("name", "") + truncated_name = truncate_tool_name(original_name) + if truncated_name != original_name: + mapping[truncated_name] = original_name + return mapping + from openai.types.chat.chat_completion_chunk import Choice as OpenAIStreamingChoice from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -77,8 +126,29 @@ class AnthropicAdapter: self, kwargs ) -> Optional[ChatCompletionRequest]: """ + Translate Anthropic request params to OpenAI format. + - translate params, where needed - pass rest, as is + + Note: Use translate_completion_input_params_with_tool_mapping() if you need + the tool name mapping for restoring original names in responses. + """ + result, _ = self.translate_completion_input_params_with_tool_mapping(kwargs) + return result + + def translate_completion_input_params_with_tool_mapping( + self, kwargs + ) -> Tuple[Optional[ChatCompletionRequest], Dict[str, str]]: + """ + Translate Anthropic request params to OpenAI format, returning tool name mapping. + + This method handles truncation of tool names that exceed OpenAI's 64-character + limit. The mapping allows restoring original names when translating responses. + + Returns: + Tuple of (openai_request, tool_name_mapping) + - tool_name_mapping maps truncated tool names back to original names """ ######################################################### @@ -102,26 +172,51 @@ class AnthropicAdapter: model=model, messages=messages, **kwargs ) - translated_body = ( + translated_body, tool_name_mapping = ( LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( anthropic_message_request=request_body ) ) - return translated_body + return translated_body, tool_name_mapping def translate_completion_output_params( - self, response: ModelResponse + self, + response: ModelResponse, + tool_name_mapping: Optional[Dict[str, str]] = None, ) -> Optional[AnthropicMessagesResponse]: + """ + Translate OpenAI response to Anthropic format. + + Args: + response: The OpenAI ModelResponse + tool_name_mapping: Optional mapping of truncated tool names to original names. + Used to restore original names for tools that exceeded + OpenAI's 64-char limit. + """ return LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic( - response=response + response=response, + tool_name_mapping=tool_name_mapping, ) def translate_completion_output_params_streaming( - self, completion_stream: Any, model: str + self, + completion_stream: Any, + model: str, + tool_name_mapping: Optional[Dict[str, str]] = None, ) -> Union[AsyncIterator[bytes], None]: + """ + Translate OpenAI streaming response to Anthropic format. + + Args: + completion_stream: The OpenAI streaming response + model: The model name + tool_name_mapping: Optional mapping of truncated tool names to original names. + """ anthropic_wrapper = AnthropicStreamWrapper( - completion_stream=completion_stream, model=model + completion_stream=completion_stream, + model=model, + tool_name_mapping=tool_name_mapping, ) # Return the SSE-wrapped version for proper event formatting return anthropic_wrapper.async_anthropic_sse_wrapper() @@ -417,8 +512,10 @@ class LiteLLMAnthropicMessagesAdapter: has_cache_control_in_text = True assistant_content_list.append(text_block) elif content.get("type") == "tool_use": + # Truncate tool name for OpenAI's 64-char limit + tool_name = truncate_tool_name(content.get("name", "")) function_chunk: ChatCompletionToolCallFunctionChunk = { - "name": content.get("name", ""), + "name": tool_name, "arguments": json.dumps(content.get("input", {})), } signature = ( @@ -587,8 +684,11 @@ class LiteLLMAnthropicMessagesAdapter: elif tool_choice["type"] == "auto": return "auto" elif tool_choice["type"] == "tool": + # Truncate tool name if it exceeds OpenAI's 64-char limit + original_name = tool_choice.get("name", "") + truncated_name = truncate_tool_name(original_name) tc_function_param = ChatCompletionToolChoiceFunctionParam( - name=tool_choice.get("name", "") + name=truncated_name ) return ChatCompletionToolChoiceObjectParam( type="function", function=tc_function_param @@ -600,12 +700,28 @@ class LiteLLMAnthropicMessagesAdapter: def translate_anthropic_tools_to_openai( self, tools: List[AllAnthropicToolsValues], model: Optional[str] = None - ) -> List[ChatCompletionToolParam]: + ) -> Tuple[List[ChatCompletionToolParam], Dict[str, str]]: + """ + Translate Anthropic tools to OpenAI format. + + Returns: + Tuple of (translated_tools, tool_name_mapping) + - tool_name_mapping maps truncated names back to original names + for tools that exceeded OpenAI's 64-char limit + """ new_tools: List[ChatCompletionToolParam] = [] + tool_name_mapping: Dict[str, str] = {} mapped_tool_params = ["name", "input_schema", "description", "cache_control"] for tool in tools: + original_name = tool["name"] + truncated_name = truncate_tool_name(original_name) + + # Store mapping if name was truncated + if truncated_name != original_name: + tool_name_mapping[truncated_name] = original_name + function_chunk = ChatCompletionToolParamFunctionChunk( - name=tool["name"], + name=truncated_name, ) if "input_schema" in tool: function_chunk["parameters"] = tool["input_schema"] # type: ignore @@ -619,7 +735,7 @@ class LiteLLMAnthropicMessagesAdapter: self._add_cache_control_if_applicable(tool, tool_param, model) new_tools.append(tool_param) # type: ignore[arg-type] - return new_tools # type: ignore[return-value] + return new_tools, tool_name_mapping # type: ignore[return-value] def translate_anthropic_output_format_to_openai( self, output_format: Any @@ -694,12 +810,18 @@ class LiteLLMAnthropicMessagesAdapter: def translate_anthropic_to_openai( self, anthropic_message_request: AnthropicMessagesRequest - ) -> ChatCompletionRequest: + ) -> Tuple[ChatCompletionRequest, Dict[str, str]]: """ This is used by the beta Anthropic Adapter, for translating anthropic `/v1/messages` requests to the openai format. + + Returns: + Tuple of (openai_request, tool_name_mapping) + - tool_name_mapping maps truncated tool names back to original names + for tools that exceeded OpenAI's 64-char limit """ # Debug: Processing Anthropic message request new_messages: List[AllMessageValues] = [] + tool_name_mapping: Dict[str, str] = {} ## CONVERT ANTHROPIC MESSAGES TO OPENAI messages_list: List[ @@ -750,7 +872,7 @@ class LiteLLMAnthropicMessagesAdapter: if "tools" in anthropic_message_request: tools = anthropic_message_request["tools"] if tools: - new_kwargs["tools"] = self.translate_anthropic_tools_to_openai( + new_kwargs["tools"], tool_name_mapping = self.translate_anthropic_tools_to_openai( tools=cast(List[AllAnthropicToolsValues], tools), model=new_kwargs.get("model"), ) @@ -784,7 +906,7 @@ class LiteLLMAnthropicMessagesAdapter: if k not in translatable_params: # pass remaining params as is new_kwargs[k] = v # type: ignore - return new_kwargs + return new_kwargs, tool_name_mapping def _translate_anthropic_image_to_openai(self, image_source: dict) -> Optional[str]: """ @@ -813,7 +935,11 @@ class LiteLLMAnthropicMessagesAdapter: return None - def _translate_openai_content_to_anthropic(self, choices: List[Choices]) -> List[ + def _translate_openai_content_to_anthropic( + self, + choices: List[Choices], + tool_name_mapping: Optional[Dict[str, str]] = None, + ) -> List[ Union[ AnthropicResponseContentBlockText, AnthropicResponseContentBlockToolUse, @@ -895,13 +1021,21 @@ class LiteLLMAnthropicMessagesAdapter: if signature: provider_specific_fields["signature"] = signature + # Restore original tool name if it was truncated + truncated_name = tool_call.function.name or "" + original_name = ( + tool_name_mapping.get(truncated_name, truncated_name) + if tool_name_mapping + else truncated_name + ) + tool_use_block = AnthropicResponseContentBlockToolUse( type="tool_use", id=tool_call.id, - name=tool_call.function.name or "", + name=original_name, input=parse_tool_call_arguments( tool_call.function.arguments, - tool_name=tool_call.function.name, + tool_name=original_name, context="Anthropic pass-through adapter", ), ) @@ -926,10 +1060,24 @@ class LiteLLMAnthropicMessagesAdapter: return "end_turn" def translate_openai_response_to_anthropic( - self, response: ModelResponse + self, + response: ModelResponse, + tool_name_mapping: Optional[Dict[str, str]] = None, ) -> AnthropicMessagesResponse: + """ + Translate OpenAI response to Anthropic format. + + Args: + response: The OpenAI ModelResponse + tool_name_mapping: Optional mapping of truncated tool names to original names. + Used to restore original names for tools that exceeded + OpenAI's 64-char limit. + """ ## translate content block - anthropic_content = self._translate_openai_content_to_anthropic(choices=response.choices) # type: ignore + anthropic_content = self._translate_openai_content_to_anthropic( + choices=response.choices, # type: ignore + tool_name_mapping=tool_name_mapping, + ) ## extract finish reason anthropic_finish_reason = self._translate_openai_finish_reason_to_anthropic( openai_finish_reason=response.choices[0].finish_reason # type: ignore diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 1c790f70062..bcb0059fd19 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -8,7 +8,10 @@ sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( + OPENAI_MAX_TOOL_NAME_LENGTH, LiteLLMAnthropicMessagesAdapter, + create_tool_name_mapping, + truncate_tool_name, ) from litellm.types.llms.anthropic import ( AnthopicMessagesAssistantMessageParam, @@ -1388,12 +1391,13 @@ def test_cache_control_preserved_in_tools_for_claude(): ] adapter = LiteLLMAnthropicMessagesAdapter() - result = adapter.translate_anthropic_tools_to_openai( + result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai( tools=tools, model=CACHE_CONTROL_BEDROCK_CONVERSE_MODEL ) assert len(result) == 1 assert result[0]["cache_control"] == {"type": "ephemeral"} + assert tool_name_mapping == {} # No truncation needed for short names def test_cache_control_not_preserved_in_tools_for_non_claude(): @@ -1408,7 +1412,7 @@ def test_cache_control_not_preserved_in_tools_for_non_claude(): ] adapter = LiteLLMAnthropicMessagesAdapter() - result = adapter.translate_anthropic_tools_to_openai( + result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai( tools=tools, model=CACHE_CONTROL_NON_ANTHROPIC_MODEL ) @@ -1527,3 +1531,179 @@ def test_translate_openai_response_to_anthropic_with_reasoning_content_only(): assert cast(Any, anthropic_content[1]).text == "There are **3** \"r\"s in the word strawberry." assert anthropic_response.get("stop_reason") == "end_turn" + assert tool_name_mapping == {} # No truncation needed for short names + + +# ===================================================================== +# Tool Name Truncation Tests (Issue #17904) +# OpenAI has a 64-character limit for function/tool names +# ===================================================================== + + +def test_truncate_tool_name_short_name(): + """Short tool names should not be truncated.""" + short_name = "get_weather" + result = truncate_tool_name(short_name) + assert result == short_name + assert len(result) <= OPENAI_MAX_TOOL_NAME_LENGTH + + +def test_truncate_tool_name_exactly_64_chars(): + """Tool names exactly 64 chars should not be truncated.""" + name_64_chars = "a" * 64 + result = truncate_tool_name(name_64_chars) + assert result == name_64_chars + assert len(result) == 64 + + +def test_truncate_tool_name_long_name(): + """Long tool names should be truncated with hash suffix.""" + long_name = "computer_tool_with_very_long_name_that_exceeds_openai_64_character_limit_and_keeps_going" + result = truncate_tool_name(long_name) + + assert len(result) == OPENAI_MAX_TOOL_NAME_LENGTH + assert result != long_name + # Should have format: {55-char-prefix}_{8-char-hash} + assert "_" in result + parts = result.rsplit("_", 1) + assert len(parts[0]) == 55 + assert len(parts[1]) == 8 + + +def test_truncate_tool_name_deterministic(): + """Truncation should be deterministic (same input = same output).""" + long_name = "a_very_long_tool_name_that_needs_to_be_truncated_for_openai_compatibility_reasons" + result1 = truncate_tool_name(long_name) + result2 = truncate_tool_name(long_name) + assert result1 == result2 + + +def test_truncate_tool_name_avoids_collisions(): + """Similar long names should produce different truncated names.""" + name1 = "process_user_data_with_validation_and_error_handling_for_production_environment" + name2 = "process_user_data_with_validation_and_error_handling_for_staging_environment" + + result1 = truncate_tool_name(name1) + result2 = truncate_tool_name(name2) + + assert result1 != result2 # Different hashes prevent collision + + +def test_create_tool_name_mapping_no_long_names(): + """Mapping should be empty when no names need truncation.""" + tools = [ + {"name": "get_weather"}, + {"name": "search_web"}, + ] + mapping = create_tool_name_mapping(tools) + assert mapping == {} + + +def test_create_tool_name_mapping_with_long_names(): + """Mapping should contain entries for truncated names.""" + long_name = "a_very_long_tool_name_that_exceeds_the_64_character_limit_imposed_by_openai" + tools = [ + {"name": "short_name"}, + {"name": long_name}, + ] + mapping = create_tool_name_mapping(tools) + + assert len(mapping) == 1 + truncated = truncate_tool_name(long_name) + assert truncated in mapping + assert mapping[truncated] == long_name + + +def test_translate_anthropic_tools_with_long_names(): + """Tools with long names should be truncated and mapped.""" + long_name = "computer_tool_with_very_long_descriptive_name_that_exceeds_openai_limit_completely" + tools = [ + { + "name": long_name, + "description": "A tool with a very long name", + "input_schema": {"type": "object", "properties": {}}, + } + ] + + adapter = LiteLLMAnthropicMessagesAdapter() + result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai( + tools=tools, model="gpt-4" + ) + + assert len(result) == 1 + # The tool name should be truncated + truncated_name = result[0]["function"]["name"] + assert len(truncated_name) <= 64 + assert truncated_name != long_name + # Mapping should have the reverse lookup + assert truncated_name in tool_name_mapping + assert tool_name_mapping[truncated_name] == long_name + + +def test_translate_anthropic_tools_mixed_names(): + """Mix of short and long names should work correctly.""" + short_name = "get_weather" + long_name = "process_complex_data_transformation_with_validation_and_error_handling_pipeline" + tools = [ + {"name": short_name, "input_schema": {"type": "object"}}, + {"name": long_name, "input_schema": {"type": "object"}}, + ] + + adapter = LiteLLMAnthropicMessagesAdapter() + result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai( + tools=tools, model="gpt-4" + ) + + assert len(result) == 2 + # Short name unchanged + assert result[0]["function"]["name"] == short_name + # Long name truncated + assert result[1]["function"]["name"] != long_name + assert len(result[1]["function"]["name"]) <= 64 + # Only long name in mapping + assert len(tool_name_mapping) == 1 + + +def test_translate_openai_response_restores_tool_names(): + """Tool names in responses should be restored to original.""" + original_name = "a_very_long_tool_name_that_needs_truncation_for_openai_api_compatibility" + truncated_name = truncate_tool_name(original_name) + tool_name_mapping = {truncated_name: original_name} + + # Create a mock OpenAI response with the truncated name + response = ModelResponse( + id="test-id", + choices=[ + Choices( + index=0, + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionAssistantToolCall( + id="call_123", + type="function", + function=Function( + name=truncated_name, + arguments='{"arg": "value"}', + ), + ) + ], + ), + ) + ], + model="gpt-4", + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + result = adapter.translate_openai_response_to_anthropic( + response=response, tool_name_mapping=tool_name_mapping + ) + + # Find the tool_use block in the response + tool_use_blocks = [c for c in result["content"] if getattr(c, "type", None) == "tool_use"] + assert len(tool_use_blocks) == 1 + # Name should be restored to original + assert getattr(tool_use_blocks[0], "name", None) == original_name From a457162517766f018ca515dfdea54d283a0c0ba3 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Sun, 1 Feb 2026 04:05:49 +0530 Subject: [PATCH 11/22] fix: handle deprecated 'redis_db' arg to prevent crash (#19808) * fix: handle deprecated 'redis_db' arg to prevent crash * renamed: changed dir --- litellm/router.py | 7 +++ tests/test_litellm/test_router_redis_init.py | 56 ++++++++++++++++++++ 2 files changed, 63 insertions(+) create mode 100644 tests/test_litellm/test_router_redis_init.py diff --git a/litellm/router.py b/litellm/router.py index 40d84fae410..e117e8c09ae 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -225,6 +225,7 @@ class Router: redis_host: Optional[str] = None, redis_port: Optional[int] = None, redis_password: Optional[str] = None, + redis_db: Optional[int] = None, cache_responses: Optional[bool] = False, cache_kwargs: dict = {}, # additional kwargs to pass to RedisCache (see caching.py) caching_groups: Optional[ @@ -411,6 +412,12 @@ class Router: if redis_password is not None: cache_config["password"] = redis_password + if redis_db is not None: + verbose_router_logger.warning( + "Deprecated 'redis_db' argument used. Please remove 'redis_db' from your config/database and use 'cache_kwargs' instead." + ) + cache_config["db"] = str(redis_db) + # Add additional key-value pairs from cache_kwargs cache_config.update(cache_kwargs) redis_cache = self._create_redis_cache(cache_config) diff --git a/tests/test_litellm/test_router_redis_init.py b/tests/test_litellm/test_router_redis_init.py new file mode 100644 index 00000000000..4a8a5b57622 --- /dev/null +++ b/tests/test_litellm/test_router_redis_init.py @@ -0,0 +1,56 @@ +import pytest +import asyncio +import os +from litellm import Router + + +# Mark as async test +@pytest.mark.asyncio +async def test_router_uses_correct_redis_db(): + """ + Verifies that when redis_db is passed to Router, + items are actually stored in that specific Redis DB index. + """ + # 1. Setup - Use a non-standard DB index (e.g., 5) to prove it's not using default 0 + test_db_index = 5 + + # Ensure we have a Redis URL available (fallback to localhost if env var not set) + redis_host = os.getenv("REDIS_HOST", "localhost") + redis_port = os.getenv("REDIS_PORT", "6379") + + # Initialize Router with specific redis_db + router = Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo"}, + } + ], + redis_host=redis_host, + redis_port=int(redis_port), + redis_db=test_db_index, + cache_responses=True, # Important: Enable caching to trigger Redis usage + ) + + # 2. Verify Internal State + # Check if the underlying cache client is configured with the correct DB + # Accessing internal attributes for verification purposes + try: + if router.cache.redis_cache: + # Check connection kwargs or internal client db + cache_client = router.cache.redis_cache.redis_client + # Redis client stores connection args in connection_pool.connection_kwargs + conn_kwargs = cache_client.connection_pool.connection_kwargs + + assert str(conn_kwargs.get("db")) == str( + test_db_index + ), f"Router Internal Check Failed: Expected DB {test_db_index}, got {conn_kwargs.get('db')}" + else: + pytest.fail("Redis cache was not initialized in Router") + + except Exception as e: + pytest.fail(f"Failed to inspect Router internals: {e}") + + +if __name__ == "__main__": + asyncio.run(test_router_uses_correct_redis_db()) From 6d86808eaffc1b5dbd9f59d0e4a4928fb6a0aae0 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Sun, 1 Feb 2026 04:06:35 +0530 Subject: [PATCH 12/22] =?UTF-8?q?feat:=20enforce=20model-level=20TPM/RPM?= =?UTF-8?q?=20limits=20(enforce=5Fmodel=5Frate=5Flimits)=20=E2=80=A6=20(#1?= =?UTF-8?q?9230)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: enforce model-level TPM/RPM limits (enforce_model_rate_limits) flag * fix lint errors --- docs/my-website/docs/proxy/load_balancing.md | 61 +++ litellm/router.py | 5 + .../pre_call_checks/model_rate_limit_check.py | 373 ++++++++++++++++++ litellm/types/router.py | 29 +- .../test_enforce_model_rate_limits.py | 315 +++++++++++++++ 5 files changed, 768 insertions(+), 15 deletions(-) create mode 100644 litellm/router_utils/pre_call_checks/model_rate_limit_check.py create mode 100644 tests/test_litellm/test_router/test_enforce_model_rate_limits.py diff --git a/docs/my-website/docs/proxy/load_balancing.md b/docs/my-website/docs/proxy/load_balancing.md index 42f6ef1aa51..186307d6498 100644 --- a/docs/my-website/docs/proxy/load_balancing.md +++ b/docs/my-website/docs/proxy/load_balancing.md @@ -69,6 +69,67 @@ router_settings: redis_port: 1992 ``` +## Enforce Model Rate Limits + +Strictly enforce RPM/TPM limits set on deployments. When limits are exceeded, requests are blocked **before** reaching the LLM provider with a `429 Too Many Requests` error. + +:::info +By default, `rpm` and `tpm` values are only used for **routing decisions** (picking deployments with capacity). With `enforce_model_rate_limits`, they become **hard limits**. +::: + +### Quick Start + +```yaml +model_list: + - model_name: gpt-4 + litellm_params: + model: openai/gpt-4 + api_key: os.environ/OPENAI_API_KEY + rpm: 60 # 60 requests per minute + tpm: 90000 # 90k tokens per minute + +router_settings: + optional_pre_call_checks: + - enforce_model_rate_limits # 👈 Enables strict enforcement +``` + +### How It Works + +| Limit Type | Enforcement | Accuracy | +|------------|-------------|----------| +| **RPM** | Hard limit - blocked at exact threshold | 100% accurate | +| **TPM** | Best-effort - may slightly exceed | Blocked when already over limit | + +**Why TPM is best-effort:** Token count is unknown until the LLM responds. TPM is checked before each request (blocks if already over), and tracked after (adds actual tokens used). + +### Error Response + +```json +{ + "error": { + "message": "Model rate limit exceeded. RPM limit=60, current usage=60", + "type": "rate_limit_error", + "code": 429 + } +} +``` + +Response includes `retry-after: 60` header. + +### Multi-Instance Deployment + +For multiple LiteLLM proxy instances, add Redis to share rate limit state: + +```yaml +router_settings: + optional_pre_call_checks: + - enforce_model_rate_limits + redis_host: redis.example.com + redis_port: 6379 + redis_password: your-password +``` + + :::info Detailed information about [routing strategies can be found here](../routing) ::: diff --git a/litellm/router.py b/litellm/router.py index e117e8c09ae..65445e29c41 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -118,6 +118,9 @@ from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import from litellm.router_utils.pre_call_checks.responses_api_deployment_check import ( ResponsesApiDeploymentCheck, ) +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( + ModelRateLimitingCheck, +) from litellm.router_utils.router_callbacks.track_deployment_metrics import ( increment_deployment_failures_for_current_minute, increment_deployment_successes_for_current_minute, @@ -1195,6 +1198,8 @@ class Router: ) elif pre_call_check == "responses_api_deployment_check": _callback = ResponsesApiDeploymentCheck() + elif pre_call_check == "enforce_model_rate_limits": + _callback = ModelRateLimitingCheck(dual_cache=self.cache) if _callback is not None: if self.optional_callbacks is None: self.optional_callbacks = [] diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py new file mode 100644 index 00000000000..e5be61690ba --- /dev/null +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -0,0 +1,373 @@ +""" +Enforce TPM/RPM rate limits set on model deployments. + +This pre-call check ensures that model-level TPM/RPM limits are enforced +across all requests, regardless of routing strategy. + +When enabled via `enforce_model_rate_limits: true` in litellm_settings, +requests that exceed the configured TPM/RPM limits will receive a 429 error. +""" + +from typing import TYPE_CHECKING, Any, Dict, Optional, Union + +import httpx + +import litellm +from litellm._logging import verbose_router_logger +from litellm.caching.dual_cache import DualCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.router import RouterErrors +from litellm.types.utils import StandardLoggingPayload +from litellm.utils import get_utc_datetime + +if TYPE_CHECKING: + from opentelemetry.trace import Span as _Span + + Span = Union[_Span, Any] +else: + Span = Any + + +class RoutingArgs: + ttl: int = 60 # 1min (RPM/TPM expire key) + + +class ModelRateLimitingCheck(CustomLogger): + """ + Pre-call check that enforces TPM/RPM limits on model deployments. + + This check runs before each request and raises a RateLimitError + if the deployment has exceeded its configured TPM or RPM limits. + + Unlike the usage-based-routing strategy which uses limits for routing decisions, + this check actively enforces those limits across ALL routing strategies. + """ + + def __init__(self, dual_cache: DualCache): + self.dual_cache = dual_cache + + def _get_deployment_limits( + self, deployment: Dict + ) -> tuple[Optional[int], Optional[int]]: + """ + Extract TPM and RPM limits from a deployment configuration. + + Checks in order: + 1. Top-level 'tpm'/'rpm' fields + 2. litellm_params.tpm/rpm + 3. model_info.tpm/rpm + + Returns: + Tuple of (tpm_limit, rpm_limit) + """ + # Check top-level + tpm = deployment.get("tpm") + rpm = deployment.get("rpm") + + # Check litellm_params + if tpm is None: + tpm = deployment.get("litellm_params", {}).get("tpm") + if rpm is None: + rpm = deployment.get("litellm_params", {}).get("rpm") + + # Check model_info + if tpm is None: + tpm = deployment.get("model_info", {}).get("tpm") + if rpm is None: + rpm = deployment.get("model_info", {}).get("rpm") + + return tpm, rpm + + def _get_cache_keys(self, deployment: Dict, current_minute: str) -> tuple[str, str]: + """Get the cache keys for TPM and RPM tracking.""" + model_id = deployment.get("model_info", {}).get("id") + deployment_name = deployment.get("litellm_params", {}).get("model") + + tpm_key = f"{model_id}:{deployment_name}:tpm:{current_minute}" + rpm_key = f"{model_id}:{deployment_name}:rpm:{current_minute}" + + return tpm_key, rpm_key + + def pre_call_check(self, deployment: Dict) -> Optional[Dict]: + """ + Synchronous pre-call check for model rate limits. + + Raises RateLimitError if deployment exceeds TPM/RPM limits. + """ + try: + tpm_limit, rpm_limit = self._get_deployment_limits(deployment) + + # If no limits are set, allow the request + if tpm_limit is None and rpm_limit is None: + return deployment + + dt = get_utc_datetime() + current_minute = dt.strftime("%H-%M") + tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute) + + model_id = deployment.get("model_info", {}).get("id") + model_name = deployment.get("litellm_params", {}).get("model") + model_group = deployment.get("model_name", "") + + # Check TPM limit + if tpm_limit is not None: + # First check local cache + current_tpm = self.dual_cache.get_cache(key=tpm_key, local_only=True) + if current_tpm is not None and current_tpm >= tpm_limit: + raise litellm.RateLimitError( + message=f"Model rate limit exceeded. TPM limit={tpm_limit}, current usage={current_tpm}", + llm_provider="", + model=model_name, + response=httpx.Response( + status_code=429, + content=f"{RouterErrors.user_defined_ratelimit_error.value} tpm limit={tpm_limit}. current usage={current_tpm}. id={model_id}, model_group={model_group}", + headers={"retry-after": str(60)}, + request=httpx.Request( + method="model_rate_limit_check", + url="https://github.com/BerriAI/litellm", + ), + ), + ) + + # Check RPM limit + if rpm_limit is not None: + # First check local cache + current_rpm = self.dual_cache.get_cache(key=rpm_key, local_only=True) + if current_rpm >= rpm_limit: + raise litellm.RateLimitError( + message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", + llm_provider="", + model=model_name, + response=httpx.Response( + status_code=429, + content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}", + headers={"retry-after": str(60)}, + request=httpx.Request( + method="model_rate_limit_check", + url="https://github.com/BerriAI/litellm", + ), + ), + ) + + # Check redis cache and increment + current_rpm = self.dual_cache.increment_cache( + key=rpm_key, value=1, ttl=RoutingArgs.ttl + ) + if current_rpm is not None and current_rpm > rpm_limit: + raise litellm.RateLimitError( + message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", + llm_provider="", + model=model_name, + response=httpx.Response( + status_code=429, + content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}", + headers={"retry-after": str(60)}, + request=httpx.Request( + method="model_rate_limit_check", + url="https://github.com/BerriAI/litellm", + ), + ), + ) + + return deployment + + except litellm.RateLimitError: + raise + except Exception as e: + verbose_router_logger.debug( + f"Error in ModelRateLimitingCheck.pre_call_check: {str(e)}" + ) + # Don't fail the request if rate limit check fails + return deployment + + async def async_pre_call_check( + self, deployment: Dict, parent_otel_span: Optional[Span] = None + ) -> Optional[Dict]: + """ + Async pre-call check for model rate limits. + + Raises RateLimitError if deployment exceeds TPM/RPM limits. + """ + try: + tpm_limit, rpm_limit = self._get_deployment_limits(deployment) + + # If no limits are set, allow the request + if tpm_limit is None and rpm_limit is None: + return deployment + + dt = get_utc_datetime() + current_minute = dt.strftime("%H-%M") + tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute) + + model_id = deployment.get("model_info", {}).get("id") + model_name = deployment.get("litellm_params", {}).get("model") + model_group = deployment.get("model_name", "") + + # Check TPM limit + if tpm_limit is not None: + # First check local cache + current_tpm = await self.dual_cache.async_get_cache( + key=tpm_key, local_only=True + ) + if current_tpm is not None and current_tpm >= tpm_limit: + raise litellm.RateLimitError( + message=f"Model rate limit exceeded. TPM limit={tpm_limit}, current usage={current_tpm}", + llm_provider="", + model=model_name, + response=httpx.Response( + status_code=429, + content=f"{RouterErrors.user_defined_ratelimit_error.value} tpm limit={tpm_limit}. current usage={current_tpm}. id={model_id}, model_group={model_group}", + headers={"retry-after": str(60)}, + request=httpx.Request( + method="model_rate_limit_check", + url="https://github.com/BerriAI/litellm", + ), + ), + num_retries=0, # Don't retry - return 429 immediately + ) + + # Check RPM limit + if rpm_limit is not None: + # First check local cache + current_rpm = await self.dual_cache.async_get_cache( + key=rpm_key, local_only=True + ) + if current_rpm is not None and current_rpm >= rpm_limit: + raise litellm.RateLimitError( + message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", + llm_provider="", + model=model_name, + response=httpx.Response( + status_code=429, + content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}", + headers={"retry-after": str(60)}, + request=httpx.Request( + method="model_rate_limit_check", + url="https://github.com/BerriAI/litellm", + ), + ), + num_retries=0, # Don't retry - return 429 immediately + ) + + # Check redis cache and increment + current_rpm = await self.dual_cache.async_increment_cache( + key=rpm_key, + value=1, + ttl=RoutingArgs.ttl, + parent_otel_span=parent_otel_span, + ) + if current_rpm is not None and current_rpm > rpm_limit: + raise litellm.RateLimitError( + message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", + llm_provider="", + model=model_name, + response=httpx.Response( + status_code=429, + content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}", + headers={"retry-after": str(60)}, + request=httpx.Request( + method="model_rate_limit_check", + url="https://github.com/BerriAI/litellm", + ), + ), + num_retries=0, # Don't retry - return 429 immediately + ) + + return deployment + + except litellm.RateLimitError: + raise + except Exception as e: + verbose_router_logger.debug( + f"Error in ModelRateLimitingCheck.async_pre_call_check: {str(e)}" + ) + # Don't fail the request if rate limit check fails + return deployment + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + """ + Track TPM usage after successful request. + + This updates the TPM counter with the actual tokens used. + Always tracks tokens - the pre-call check handles enforcement. + """ + try: + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( + "standard_logging_object" + ) + if standard_logging_object is None: + return + + model_id = standard_logging_object.get("model_id") + if model_id is None: + return + + total_tokens = standard_logging_object.get("total_tokens", 0) + model = standard_logging_object.get("hidden_params", {}).get( + "litellm_model_name" + ) + + verbose_router_logger.debug( + f"[TPM TRACKING] model_id={model_id}, total_tokens={total_tokens}, model={model}" + ) + + if not model or not total_tokens: + return + + dt = get_utc_datetime() + current_minute = dt.strftime("%H-%M") + tpm_key = f"{model_id}:{model}:tpm:{current_minute}" + + verbose_router_logger.debug( + f"[TPM TRACKING] Incrementing {tpm_key} by {total_tokens}" + ) + + await self.dual_cache.async_increment_cache( + key=tpm_key, + value=total_tokens, + ttl=RoutingArgs.ttl, + ) + + except Exception as e: + verbose_router_logger.debug( + f"Error in ModelRateLimitingCheck.async_log_success_event: {str(e)}" + ) + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + """ + Sync version of tracking TPM usage after successful request. + Always tracks tokens - the pre-call check handles enforcement. + """ + try: + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( + "standard_logging_object" + ) + if standard_logging_object is None: + return + + model_id = standard_logging_object.get("model_id") + if model_id is None: + return + + total_tokens = standard_logging_object.get("total_tokens", 0) + model = standard_logging_object.get("hidden_params", {}).get( + "litellm_model_name" + ) + + if not model or not total_tokens: + return + + dt = get_utc_datetime() + current_minute = dt.strftime("%H-%M") + tpm_key = f"{model_id}:{model}:tpm:{current_minute}" + + self.dual_cache.increment_cache( + key=tpm_key, + value=total_tokens, + ttl=RoutingArgs.ttl, + ) + + except Exception as e: + verbose_router_logger.debug( + f"Error in ModelRateLimitingCheck.log_success_event: {str(e)}" + ) diff --git a/litellm/types/router.py b/litellm/types/router.py index f31c6df3005..f78789c9772 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -95,18 +95,16 @@ class ModelInfo(BaseModel): id: Optional[ str ] # Allow id to be optional on input, but it will always be present as a str in the model instance - db_model: bool = ( - False # used for proxy - to separate models which are stored in the db vs. config. - ) + db_model: bool = False # used for proxy - to separate models which are stored in the db vs. config. updated_at: Optional[datetime.datetime] = None updated_by: Optional[str] = None created_at: Optional[datetime.datetime] = None created_by: Optional[str] = None - base_model: Optional[str] = ( - None # specify if the base model is azure/gpt-3.5-turbo etc for accurate cost tracking - ) + base_model: Optional[ + str + ] = None # specify if the base model is azure/gpt-3.5-turbo etc for accurate cost tracking tier: Optional[Literal["free", "paid"]] = None """ @@ -172,12 +170,12 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): custom_llm_provider: Optional[str] = None tpm: Optional[int] = None rpm: Optional[int] = None - timeout: Optional[Union[float, str, httpx.Timeout]] = ( - None # if str, pass in as os.environ/ - ) - stream_timeout: Optional[Union[float, str]] = ( - None # timeout when making stream=True calls, if str, pass in as os.environ/ - ) + timeout: Optional[ + Union[float, str, httpx.Timeout] + ] = None # if str, pass in as os.environ/ + stream_timeout: Optional[ + Union[float, str] + ] = None # timeout when making stream=True calls, if str, pass in as os.environ/ max_retries: Optional[int] = None organization: Optional[str] = None # for openai orgs configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None @@ -276,9 +274,9 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): if max_retries is not None and isinstance(max_retries, str): max_retries = int(max_retries) # cast to int # We need to keep max_retries in args since it's a parameter of GenericLiteLLMParams - args["max_retries"] = ( - max_retries # Put max_retries back in args after popping it - ) + args[ + "max_retries" + ] = max_retries # Put max_retries back in args after popping it super().__init__(**args, **params) def __contains__(self, key): @@ -805,6 +803,7 @@ OptionalPreCallChecks = List[ "router_budget_limiting", "responses_api_deployment_check", "forward_client_headers_by_model_group", + "enforce_model_rate_limits", ] ] diff --git a/tests/test_litellm/test_router/test_enforce_model_rate_limits.py b/tests/test_litellm/test_router/test_enforce_model_rate_limits.py new file mode 100644 index 00000000000..3bca3df4e1d --- /dev/null +++ b/tests/test_litellm/test_router/test_enforce_model_rate_limits.py @@ -0,0 +1,315 @@ +""" +Tests for enforce_model_rate_limits feature. + +This feature allows users to enforce TPM/RPM limits set on model deployments +regardless of the routing strategy being used. +""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import litellm +from litellm import Router +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( + ModelRateLimitingCheck, +) + + +class TestModelRateLimitingCheck: + """Test the ModelRateLimitingCheck class directly.""" + + def test_get_deployment_limits_from_top_level(self): + """Test extracting limits from top-level deployment config.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "tpm": 1000, + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm == 1000 + assert rpm == 10 + + def test_get_deployment_limits_from_litellm_params(self): + """Test extracting limits from litellm_params.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4", "tpm": 2000, "rpm": 20}, + "model_info": {"id": "test-id"}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm == 2000 + assert rpm == 20 + + def test_get_deployment_limits_from_model_info(self): + """Test extracting limits from model_info.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id", "tpm": 3000, "rpm": 30}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm == 3000 + assert rpm == 30 + + def test_get_deployment_limits_none_when_not_set(self): + """Test that None is returned when limits are not set.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + tpm, rpm = check._get_deployment_limits(deployment) + assert tpm is None + assert rpm is None + + def test_pre_call_check_allows_request_when_no_limits(self): + """Test that requests are allowed when no limits are set.""" + check = ModelRateLimitingCheck(dual_cache=MagicMock()) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + result = check.pre_call_check(deployment) + assert result == deployment + + def test_pre_call_check_raises_rate_limit_error_when_over_rpm(self): + """Test that RateLimitError is raised when RPM limit is exceeded.""" + mock_cache = MagicMock() + mock_cache.get_cache.return_value = 10 # Already at limit + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(deployment) + + assert "RPM limit=10" in str(exc_info.value) + assert "current usage=10" in str(exc_info.value) + + def test_pre_call_check_allows_request_under_limit(self): + """Test that requests are allowed when under the limit.""" + mock_cache = MagicMock() + mock_cache.get_cache.return_value = 5 + mock_cache.increment_cache.return_value = 6 + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + result = check.pre_call_check(deployment) + assert result == deployment + + def test_pre_call_check_raises_rate_limit_error_when_over_tpm(self): + """Test that RateLimitError is raised when TPM limit is exceeded.""" + mock_cache = MagicMock() + mock_cache.get_cache.return_value = 1000 # Already at limit + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "tpm": 1000, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + check.pre_call_check(deployment) + + assert "TPM limit=1000" in str(exc_info.value) + assert "current usage=1000" in str(exc_info.value) + + def test_log_success_event_increments_cache(self): + """Test that log_success_event correctly increments the cache.""" + mock_cache = MagicMock() + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + kwargs = { + "standard_logging_object": { + "model_id": "test-id", + "total_tokens": 50, + "hidden_params": {"litellm_model_name": "gpt-4"}, + } + } + + check.log_success_event(kwargs, None, None, None) + + # Verify increment_cache was called + mock_cache.increment_cache.assert_called_once() + _, kwarg_params = mock_cache.increment_cache.call_args + assert "test-id:gpt-4:tpm:" in kwarg_params["key"] + assert kwarg_params["value"] == 50 + + +class TestModelRateLimitingCheckAsync: + """Test async methods of ModelRateLimitingCheck.""" + + @pytest.mark.asyncio + async def test_async_pre_call_check_allows_request_when_no_limits(self): + """Test that requests are allowed when no limits are set (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + } + + result = await check.async_pre_call_check(deployment) + assert result == deployment + + @pytest.mark.asyncio + async def test_async_pre_call_check_raises_rate_limit_error_when_over_rpm(self): + """Test that RateLimitError is raised when RPM limit is exceeded (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=10) # Already at limit + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "RPM limit=10" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_pre_call_check_allows_request_under_limit(self): + """Test that requests are allowed when under the limit (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=5) + mock_cache.async_increment_cache = AsyncMock(return_value=6) + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + result = await check.async_pre_call_check(deployment) + assert result == deployment + + @pytest.mark.asyncio + async def test_async_pre_call_check_raises_rate_limit_error_when_over_tpm(self): + """Test that RateLimitError is raised when TPM limit is exceeded (async).""" + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=1000) # Already at limit + + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + deployment = { + "tpm": 1000, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", + } + + with pytest.raises(litellm.RateLimitError) as exc_info: + await check.async_pre_call_check(deployment) + + assert "TPM limit=1000" in str(exc_info.value) + + @pytest.mark.asyncio + async def test_async_log_success_event_increments_cache(self): + """Test that async_log_success_event correctly increments the cache.""" + mock_cache = MagicMock() + mock_cache.async_increment_cache = AsyncMock() + check = ModelRateLimitingCheck(dual_cache=mock_cache) + + kwargs = { + "standard_logging_object": { + "model_id": "test-id", + "total_tokens": 50, + "hidden_params": {"litellm_model_name": "gpt-4"}, + } + } + + await check.async_log_success_event(kwargs, None, None, None) + + # Verify async_increment_cache was called + mock_cache.async_increment_cache.assert_called_once() + _, kwarg_params = mock_cache.async_increment_cache.call_args + assert "test-id:gpt-4:tpm:" in kwarg_params["key"] + assert kwarg_params["value"] == 50 + + +class TestRouterWithEnforceModelRateLimits: + """Test Router integration with enforce_model_rate_limits.""" + + def test_router_initializes_with_enforce_model_rate_limits(self): + """Test that Router properly initializes the ModelRateLimitingCheck.""" + model_list = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "test"}, + "rpm": 10, + } + ] + + router = Router( + model_list=model_list, + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + # Check that the callback was added + assert router.optional_callbacks is not None + assert len(router.optional_callbacks) == 1 + assert isinstance(router.optional_callbacks[0], ModelRateLimitingCheck) + + def test_router_optional_callbacks_contains_model_rate_limiting(self): + """Test that ModelRateLimitingCheck is in the callbacks list.""" + model_list = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "test"}, + "rpm": 10, + } + ] + + Router( + model_list=model_list, + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + # Find the ModelRateLimitingCheck in litellm.callbacks + found = False + for callback in litellm.callbacks: + if isinstance(callback, ModelRateLimitingCheck): + found = True + break + + assert found, "ModelRateLimitingCheck should be in litellm.callbacks" From c9757cd0d7cbe9b3ac853b0fb9830bd1d6a8a734 Mon Sep 17 00:00:00 2001 From: Hi120ki <12624257+hi120ki@users.noreply.github.com> Date: Sun, 1 Feb 2026 08:04:19 +0900 Subject: [PATCH 13/22] fix(guardrails): populate applied_guardrails when Model Armor blocks content (#20034) Previously, when Model Armor guardrail blocked a request/response, the `applied_guardrails` field was not populated in the logs because `add_guardrail_to_applied_guardrails_header()` was called after the HTTPException was raised. This fix moves the `add_guardrail_to_applied_guardrails_header()` call to before the blocking check in all hooks: - async_pre_call_hook (pre_call mode) - async_moderation_hook (during_call mode) - async_post_call_success_hook (post_call mode) - async_post_call_streaming_iterator_hook (streaming) This ensures that even when a guardrail blocks content, the guardrail name is properly recorded in the logs for observability. Added regression tests to verify applied_guardrails is populated when content is blocked. Co-authored-by: Cursor --- .../model_armor/model_armor.py | 44 ++- .../guardrail_hooks/test_model_armor.py | 308 +++++++++++------- 2 files changed, 224 insertions(+), 128 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index a12eb2486d2..38462094b11 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -421,6 +421,13 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): ) else "success" ) + + # Add guardrail to applied_guardrails BEFORE potential blocking + # This ensures guardrail is recorded even when it blocks the request + add_guardrail_to_applied_guardrails_header( + request_data=data, guardrail_name=self.guardrail_name + ) + # Check if content should be blocked if self._should_block_content( armor_response, allow_sanitization=self.mask_request_content @@ -456,11 +463,6 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): if self.optional_params.get("fail_on_error", True): raise - # Add guardrail to headers - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) - return data @log_guardrail_information @@ -517,6 +519,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): else "success" ) + # Add guardrail to applied_guardrails BEFORE potential blocking + # This ensures guardrail is recorded even when it blocks the request + add_guardrail_to_applied_guardrails_header( + request_data=data, guardrail_name=self.guardrail_name + ) + # Check if content should be blocked if self._should_block_content( armor_response, allow_sanitization=self.mask_request_content @@ -550,11 +558,6 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): if self.optional_params.get("fail_on_error", True): raise - # Add guardrail to headers - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) - return data @log_guardrail_information @@ -622,6 +625,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): guardrail_response=standard_logging_guardrail_information, ) + # Add guardrail to applied_guardrails BEFORE potential blocking + # This ensures guardrail is recorded even when it blocks the request + add_guardrail_to_applied_guardrails_header( + request_data=data, guardrail_name=self.guardrail_name + ) + # Check if content should be blocked if self._should_block_content( armor_response, allow_sanitization=self.mask_response_content @@ -654,11 +663,6 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): if self.optional_params.get("fail_on_error", True): raise - # Add guardrail to headers - add_guardrail_to_applied_guardrails_header( - request_data=data, guardrail_name=self.guardrail_name - ) - return response async def async_post_call_streaming_iterator_hook( @@ -703,6 +707,16 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): else "success" ) + # Add guardrail to applied_guardrails BEFORE potential blocking + # This ensures guardrail is recorded even when it blocks the request + from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + ) + + add_guardrail_to_applied_guardrails_header( + request_data=request_data, guardrail_name=self.guardrail_name + ) + # Check if blocked if self._should_block_content(armor_response): raise HTTPException( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index 6d0a1b46559..987388a80c7 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -24,7 +24,7 @@ async def test_model_armor_pre_call_hook_sanitization(): """Test Model Armor pre-call hook with content sanitization""" mock_user_api_key_dict = UserAPIKeyAuth() mock_cache = MagicMock(spec=DualCache) - + guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", @@ -32,7 +32,7 @@ async def test_model_armor_pre_call_hook_sanitization(): guardrail_name="model-armor-test", mask_request_content=True, ) - + # Mock the Model Armor API response mock_response = AsyncMock() mock_response.status_code = 200 @@ -53,10 +53,10 @@ async def test_model_armor_pre_call_hook_sanitization(): } } }) - + # Mock the access token method guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) - + # Mock the async handler with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=mock_response)): request_data = { @@ -66,17 +66,17 @@ async def test_model_armor_pre_call_hook_sanitization(): ], "metadata": {"guardrails": ["model-armor-test"]} } - + result = await guardrail.async_pre_call_hook( user_api_key_dict=mock_user_api_key_dict, cache=mock_cache, data=request_data, call_type="completion" ) - + # Assert the message was sanitized assert result["messages"][0]["content"] == "Hello, my phone number is [REDACTED]" - + # Verify API was called correctly # Note: we need to use the captured mock from the patch if we want to assert on it # But for now, we'll just verify the behavior. @@ -89,14 +89,14 @@ async def test_model_armor_pre_call_hook_blocked(): """Test Model Armor pre-call hook when content is blocked""" mock_user_api_key_dict = UserAPIKeyAuth() mock_cache = MagicMock(spec=DualCache) - + guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", location="us-central1", guardrail_name="model-armor-test", ) - + # Mock the Model Armor API response for blocked content mock_response = AsyncMock() mock_response.status_code = 200 @@ -118,10 +118,10 @@ async def test_model_armor_pre_call_hook_blocked(): } } }) - + # Mock the access token method guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) - + # Mock the async handler with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=mock_response)): request_data = { @@ -131,7 +131,7 @@ async def test_model_armor_pre_call_hook_blocked(): ], "metadata": {"guardrails": ["model-armor-test"]} } - + # Should raise HTTPException for blocked content with pytest.raises(HTTPException) as exc_info: await guardrail.async_pre_call_hook( @@ -140,16 +140,21 @@ async def test_model_armor_pre_call_hook_blocked(): data=request_data, call_type="completion" ) - + assert exc_info.value.status_code == 400 assert "Content blocked by Model Armor" in str(exc_info.value.detail) + # IMPORTANT: Verify that applied_guardrails is populated even when blocked + # This is a regression test for the issue where applied_guardrails was null when blocked + assert "applied_guardrails" in request_data["metadata"] + assert "model-armor-test" in request_data["metadata"]["applied_guardrails"] + @pytest.mark.asyncio async def test_model_armor_post_call_hook_sanitization(): """Test Model Armor post-call hook with response sanitization""" mock_user_api_key_dict = UserAPIKeyAuth() - + guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", @@ -157,7 +162,7 @@ async def test_model_armor_post_call_hook_sanitization(): guardrail_name="model-armor-test", mask_response_content=True, ) - + # Mock the Model Armor API response mock_response = AsyncMock() mock_response.status_code = 200 @@ -178,10 +183,10 @@ async def test_model_armor_post_call_hook_sanitization(): } } }) - + # Mock the access token method guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) - + # Mock the async handler with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=mock_response)): # Create a mock response @@ -193,36 +198,108 @@ async def test_model_armor_post_call_hook_sanitization(): ) ) ] - + request_data = { "model": "gpt-4", "messages": [{"role": "user", "content": "What's my credit card?"}], "metadata": {"guardrails": ["model-armor-test"]} } - + await guardrail.async_post_call_success_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, response=mock_llm_response ) - + # Assert the response was sanitized assert mock_llm_response.choices[0].message.content == "Here is the information: [REDACTED]" +@pytest.mark.asyncio +async def test_model_armor_post_call_hook_blocked(): + """Test Model Armor post-call hook when response is blocked and applied_guardrails is populated""" + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + ) + + # Mock the Model Armor API response for blocked content + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.json = AsyncMock(return_value={ + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": { + "rai": { + "raiFilterResult": { + "matchState": "MATCH_FOUND", + "raiFilterTypeResults": { + "dangerous": { + "matchState": "MATCH_FOUND", + "reason": "Harmful response detected" + } + } + } + } + } + } + }) + + # Mock the access token method + guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) + + # Mock the async handler + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=mock_response)): + # Create a mock response + mock_llm_response = litellm.ModelResponse() + mock_llm_response.choices = [ + litellm.Choices( + message=litellm.Message( + content="Here is some harmful content..." + ) + ) + ] + + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Some prompt"}], + "metadata": {"guardrails": ["model-armor-test"]} + } + + # Should raise HTTPException for blocked response + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_post_call_success_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + response=mock_llm_response + ) + + assert exc_info.value.status_code == 400 + assert "Response blocked by Model Armor" in str(exc_info.value.detail) + + # IMPORTANT: Verify that applied_guardrails is populated even when blocked + # This is a regression test for the issue where applied_guardrails was null when blocked + assert "applied_guardrails" in request_data["metadata"] + assert "model-armor-test" in request_data["metadata"]["applied_guardrails"] + + @pytest.mark.asyncio async def test_model_armor_with_list_content(): """Test Model Armor with messages containing list content""" mock_user_api_key_dict = UserAPIKeyAuth() mock_cache = MagicMock(spec=DualCache) - + guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", location="us-central1", guardrail_name="model-armor-test", ) - + # Mock the Model Armor API response mock_response = AsyncMock() mock_response.status_code = 200 @@ -231,17 +308,17 @@ async def test_model_armor_with_list_content(): "filterMatchState": "NO_MATCH_FOUND" } }) - + # Mock the access token method guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) - + # Mock the async handler with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=mock_response)) as mock_post: request_data = { "model": "gpt-4", "messages": [ { - "role": "user", + "role": "user", "content": [ {"type": "text", "text": "Hello world"}, {"type": "text", "text": "How are you?"} @@ -250,14 +327,14 @@ async def test_model_armor_with_list_content(): ], "metadata": {"guardrails": ["model-armor-test"]} } - + result = await guardrail.async_pre_call_hook( user_api_key_dict=mock_user_api_key_dict, cache=mock_cache, data=request_data, call_type="completion" ) - + # Verify the content was extracted correctly mock_post.assert_called_once() call_args = mock_post.call_args @@ -269,7 +346,7 @@ async def test_model_armor_api_error_handling(): """Test Model Armor error handling when API returns error""" mock_user_api_key_dict = UserAPIKeyAuth() mock_cache = MagicMock(spec=DualCache) - + guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", @@ -277,15 +354,15 @@ async def test_model_armor_api_error_handling(): guardrail_name="model-armor-test", fail_on_error=True, ) - + # Mock the Model Armor API error response mock_response = AsyncMock() mock_response.status_code = 500 mock_response.text = "Internal Server Error" - + # Mock the access token method guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) - + # Mock the async handler with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=mock_response)): request_data = { @@ -293,7 +370,7 @@ async def test_model_armor_api_error_handling(): "messages": [{"role": "user", "content": "Hello"}], "metadata": {"guardrails": ["model-armor-test"]} } - + # Should raise HTTPException for API error with pytest.raises(HTTPException) as exc_info: await guardrail.async_pre_call_hook( @@ -302,7 +379,7 @@ async def test_model_armor_api_error_handling(): data=request_data, call_type="completion" ) - + assert exc_info.value.status_code == 500 assert "Model Armor API error" in str(exc_info.value.detail) @@ -316,7 +393,7 @@ async def test_model_armor_credentials_handling(): # If google.auth is not installed, skip this test pytest.skip("google.auth not installed") return - + # Test with string credentials (file path) with patch('os.path.exists', return_value=True): with patch('builtins.open', mock_open(read_data='{"type": "service_account", "project_id": "test-project"}')): @@ -326,16 +403,16 @@ async def test_model_armor_credentials_handling(): mock_creds_obj.expired = False mock_creds_obj.project_id = "test-project" # Add project_id mock_creds.return_value = mock_creds_obj - + guardrail = ModelArmorGuardrail( template_id="test-template", credentials="/path/to/creds.json", project_id="test-project", # Provide project_id ) - + # Force credential loading creds, project_id = guardrail.load_auth(credentials="/path/to/creds.json", project_id="test-project") - + assert mock_creds.called assert project_id == "test-project" @@ -344,7 +421,7 @@ async def test_model_armor_credentials_handling(): async def test_model_armor_streaming_response(): """Test Model Armor with streaming responses""" mock_user_api_key_dict = UserAPIKeyAuth() - + guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", @@ -352,7 +429,7 @@ async def test_model_armor_streaming_response(): guardrail_name="model-armor-test", mask_response_content=True, ) - + # Mock the Model Armor API response mock_response = AsyncMock() mock_response.status_code = 200 @@ -362,10 +439,10 @@ async def test_model_armor_streaming_response(): "sanitizedText": "Sanitized response" } }) - + # Mock the access token method guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) - + # Mock the async handler with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=mock_response)) as mock_post: # Create mock streaming chunks @@ -388,13 +465,13 @@ async def test_model_armor_streaming_response(): ] for chunk in chunks: yield chunk - + request_data = { "model": "gpt-4", "messages": [{"role": "user", "content": "Tell me secrets"}], "metadata": {"guardrails": ["model-armor-test"]} } - + # Process streaming response result_chunks = [] async for chunk in guardrail.async_post_call_streaming_iterator_hook( @@ -403,7 +480,7 @@ async def test_model_armor_streaming_response(): request_data=request_data ): result_chunks.append(chunk) - + # Should have processed the chunks through Model Armor assert len(result_chunks) > 0 mock_post.assert_called() @@ -423,19 +500,19 @@ async def test_model_armor_no_messages(): """Test Model Armor when request has no messages""" mock_user_api_key_dict = UserAPIKeyAuth() mock_cache = MagicMock(spec=DualCache) - + guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", location="us-central1", guardrail_name="model-armor-test", ) - + request_data = { "model": "gpt-4", "metadata": {"guardrails": ["model-armor-test"]} } - + # Should return data unchanged when no messages result = await guardrail.async_pre_call_hook( user_api_key_dict=mock_user_api_key_dict, @@ -443,7 +520,7 @@ async def test_model_armor_no_messages(): data=request_data, call_type="completion" ) - + assert result == request_data @@ -452,14 +529,14 @@ async def test_model_armor_empty_message_content(): """Test Model Armor when message content is empty""" mock_user_api_key_dict = UserAPIKeyAuth() mock_cache = MagicMock(spec=DualCache) - + guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", location="us-central1", guardrail_name="model-armor-test", ) - + request_data = { "model": "gpt-4", "messages": [ @@ -468,7 +545,7 @@ async def test_model_armor_empty_message_content(): ], "metadata": {"guardrails": ["model-armor-test"]} } - + # Should return data unchanged when no content result = await guardrail.async_pre_call_hook( user_api_key_dict=mock_user_api_key_dict, @@ -476,7 +553,7 @@ async def test_model_armor_empty_message_content(): data=request_data, call_type="completion" ) - + assert result == request_data @@ -485,14 +562,14 @@ async def test_model_armor_system_assistant_messages(): """Test Model Armor with only system/assistant messages (no user messages)""" mock_user_api_key_dict = UserAPIKeyAuth() mock_cache = MagicMock(spec=DualCache) - + guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", location="us-central1", guardrail_name="model-armor-test", ) - + request_data = { "model": "gpt-4", "messages": [ @@ -501,7 +578,7 @@ async def test_model_armor_system_assistant_messages(): ], "metadata": {"guardrails": ["model-armor-test"]} } - + # Should return data unchanged when no user messages result = await guardrail.async_pre_call_hook( user_api_key_dict=mock_user_api_key_dict, @@ -509,7 +586,7 @@ async def test_model_armor_system_assistant_messages(): data=request_data, call_type="completion" ) - + assert result == request_data @@ -518,7 +595,7 @@ async def test_model_armor_fail_on_error_false(): """Test Model Armor with fail_on_error=False when API fails""" mock_user_api_key_dict = UserAPIKeyAuth() mock_cache = MagicMock(spec=DualCache) - + guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", @@ -526,7 +603,7 @@ async def test_model_armor_fail_on_error_false(): guardrail_name="model-armor-test", fail_on_error=False, ) - + # Mock the async handler to raise an exception guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) # Make it raise a non-HTTP exception to test the fail_on_error logic @@ -536,7 +613,7 @@ async def test_model_armor_fail_on_error_false(): "messages": [{"role": "user", "content": "Hello"}], "metadata": {"guardrails": ["model-armor-test"]} } - + # Should not raise exception when fail_on_error=False result = await guardrail.async_pre_call_hook( user_api_key_dict=mock_user_api_key_dict, @@ -544,7 +621,7 @@ async def test_model_armor_fail_on_error_false(): data=request_data, call_type="completion" ) - + # Should return original data assert result == request_data @@ -554,7 +631,7 @@ async def test_model_armor_custom_api_endpoint(): """Test Model Armor with custom API endpoint""" mock_user_api_key_dict = UserAPIKeyAuth() mock_cache = MagicMock(spec=DualCache) - + custom_endpoint = "https://custom-modelarmor.example.com" guardrail = ModelArmorGuardrail( template_id="test-template", @@ -563,12 +640,12 @@ async def test_model_armor_custom_api_endpoint(): guardrail_name="model-armor-test", api_endpoint=custom_endpoint, ) - + # Mock successful response mock_response = AsyncMock() mock_response.status_code = 200 mock_response.json = AsyncMock(return_value={"action": "NONE"}) - + guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=mock_response)) as mock_post: request_data = { @@ -576,14 +653,14 @@ async def test_model_armor_custom_api_endpoint(): "messages": [{"role": "user", "content": "Test message"}], "metadata": {"guardrails": ["model-armor-test"]} } - + await guardrail.async_pre_call_hook( user_api_key_dict=mock_user_api_key_dict, cache=mock_cache, data=request_data, call_type="completion" ) - + # Verify custom endpoint was used call_args = mock_post.call_args assert call_args[1]["url"].startswith(custom_endpoint) @@ -597,13 +674,13 @@ async def test_model_armor_dict_credentials(): except ImportError: pytest.skip("google.auth not installed") return - + # Use patch context manager properly mock_creds_obj = Mock() mock_creds_obj.token = "test-token" mock_creds_obj.expired = False mock_creds_obj.project_id = "test-project" - + with patch.object(ModelArmorGuardrail, '_credentials_from_service_account', return_value=mock_creds_obj) as mock_creds: creds_dict = { "type": "service_account", @@ -611,16 +688,16 @@ async def test_model_armor_dict_credentials(): "private_key": "test-key", "client_email": "test@example.com" } - + guardrail = ModelArmorGuardrail( template_id="test-template", credentials=creds_dict, location="us-central1", ) - + # Force credential loading creds, project_id = guardrail.load_auth(credentials=creds_dict, project_id=None) - + assert mock_creds.called assert project_id == "test-project" @@ -630,7 +707,7 @@ async def test_model_armor_action_none(): """Test Model Armor when action is NONE (no sanitization needed)""" mock_user_api_key_dict = UserAPIKeyAuth() mock_cache = MagicMock(spec=DualCache) - + guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", @@ -638,7 +715,7 @@ async def test_model_armor_action_none(): guardrail_name="model-armor-test", mask_request_content=True, ) - + # Mock response with action=NO_MATCH_FOUND mock_response = AsyncMock() mock_response.status_code = 200 @@ -647,7 +724,7 @@ async def test_model_armor_action_none(): "filterMatchState": "NO_MATCH_FOUND" } }) - + guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=mock_response)): original_content = "This content is fine" @@ -656,14 +733,14 @@ async def test_model_armor_action_none(): "messages": [{"role": "user", "content": original_content}], "metadata": {"guardrails": ["model-armor-test"]} } - + result = await guardrail.async_pre_call_hook( user_api_key_dict=mock_user_api_key_dict, cache=mock_cache, data=request_data, call_type="completion" ) - + # Content should remain unchanged assert result["messages"][0]["content"] == original_content @@ -672,7 +749,7 @@ async def test_model_armor_action_none(): async def test_model_armor_missing_sanitized_text(): """Test Model Armor when response has no sanitized_text field""" mock_user_api_key_dict = UserAPIKeyAuth() - + guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", @@ -680,7 +757,7 @@ async def test_model_armor_missing_sanitized_text(): guardrail_name="model-armor-test", mask_response_content=True, ) - + # Mock response without sanitized_text mock_response = AsyncMock() mock_response.status_code = 200 @@ -689,7 +766,7 @@ async def test_model_armor_missing_sanitized_text(): "filterMatchState": "NO_MATCH_FOUND" } }) - + guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=mock_response)): # Create a mock response @@ -699,19 +776,19 @@ async def test_model_armor_missing_sanitized_text(): message=litellm.Message(content="Original content") ) ] - + request_data = { "model": "gpt-4", "messages": [{"role": "user", "content": "Test"}], "metadata": {"guardrails": ["model-armor-test"]} } - + await guardrail.async_post_call_success_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, response=mock_llm_response ) - + # Should use 'text' field as fallback assert mock_llm_response.choices[0].message.content == "Original content" @@ -792,8 +869,8 @@ async def test_model_armor_no_circular_reference_in_logging(): # Verify the logging decorator properly added the guardrail information assert "standard_logging_guardrail_information" in request_data.get("metadata", {}) - - + + @pytest.mark.asyncio async def test_model_armor_bomb_content_blocked(): """Test Model Armor correctly blocks harmful content like bomb-making instructions""" @@ -936,24 +1013,24 @@ async def test_model_armor_success_case_serializable(): async def test_model_armor_non_text_response(): """Test Model Armor with non-text response types (TTS, image generation)""" mock_user_api_key_dict = UserAPIKeyAuth() - + guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", location="us-central1", guardrail_name="model-armor-test", ) - + # Mock a non-ModelResponse object (like TTS or image response) mock_tts_response = Mock() mock_tts_response.audio = b"audio_data" - + request_data = { "model": "tts-1", "input": "Text to speak", "metadata": {"guardrails": ["model-armor-test"]} } - + # Should not raise an error for non-text responses await guardrail.async_post_call_success_hook( data=request_data, @@ -967,26 +1044,26 @@ async def test_model_armor_token_refresh(): """Test Model Armor handling expired auth tokens""" mock_user_api_key_dict = UserAPIKeyAuth() mock_cache = MagicMock(spec=DualCache) - + guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", location="us-central1", guardrail_name="model-armor-test", ) - + # Mock successful response mock_response = AsyncMock() mock_response.status_code = 200 mock_response.json = AsyncMock(return_value={"action": "NONE"}) - + # Mock token refresh - first call returns expired token, second returns fresh call_count = 0 async def mock_token_method(*args, **kwargs): nonlocal call_count call_count += 1 return (f"token-{call_count}", "test-project") - + guardrail._ensure_access_token_async = AsyncMock(side_effect=mock_token_method) with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=mock_response)): request_data = { @@ -994,14 +1071,14 @@ async def test_model_armor_token_refresh(): "messages": [{"role": "user", "content": "Test"}], "metadata": {"guardrails": ["model-armor-test"]} } - + await guardrail.async_pre_call_hook( user_api_key_dict=mock_user_api_key_dict, cache=mock_cache, data=request_data, call_type="completion" ) - + # Verify token method was called assert guardrail._ensure_access_token_async.called @@ -1011,25 +1088,25 @@ async def test_model_armor_non_model_response(): """Test Model Armor handles non-ModelResponse types (e.g., TTS) correctly""" mock_user_api_key_dict = UserAPIKeyAuth() mock_cache = MagicMock(spec=DualCache) - + guardrail = ModelArmorGuardrail( template_id="test-template", project_id="test-project", location="us-central1", guardrail_name="model-armor-test", ) - + # Mock a TTS response (not a ModelResponse) class TTSResponse: def __init__(self): self.audio_data = b"fake audio data" - + tts_response = TTSResponse() - + # Mock the access token guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) guardrail.async_handler = AsyncMock() - + # Call post-call hook with non-ModelResponse await guardrail.async_post_call_success_hook( data={ @@ -1040,7 +1117,7 @@ async def test_model_armor_non_model_response(): user_api_key_dict=mock_user_api_key_dict, response=tts_response ) - + # Verify that Model Armor API was NOT called since there's no text content assert not guardrail.async_handler.post.called @@ -1049,36 +1126,36 @@ def mock_open(read_data=''): """Helper to create a mock file object""" import io from unittest.mock import MagicMock - + file_object = io.StringIO(read_data) file_object.__enter__ = lambda self: self file_object.__exit__ = lambda self, *args: None - + mock_file = MagicMock(return_value=file_object) - return mock_file + return mock_file def test_model_armor_initialization_preserves_project_id(): """Test that ModelArmorGuardrail initialization preserves the project_id correctly""" # This tests the fix for issue #12757 where project_id was being overwritten to None # due to incorrect initialization order with VertexBase parent class - + test_project_id = "cloud-xxxxx-yyyyy" test_template_id = "global-armor" test_location = "eu" - + guardrail = ModelArmorGuardrail( template_id=test_template_id, project_id=test_project_id, location=test_location, guardrail_name="model-armor-test", ) - + # Assert that project_id is preserved after initialization assert guardrail.project_id == test_project_id assert guardrail.template_id == test_template_id assert guardrail.location == test_location - + # Also check that the VertexBase initialization didn't reset project_id to None assert hasattr(guardrail, 'project_id') assert guardrail.project_id is not None @@ -1089,7 +1166,7 @@ async def test_model_armor_with_default_credentials(): """Test Model Armor with default credentials and explicit project_id""" mock_user_api_key_dict = UserAPIKeyAuth() mock_cache = MagicMock(spec=DualCache) - + # Initialize with explicit project_id but no credentials (simulating default auth) guardrail = ModelArmorGuardrail( template_id="test-template", @@ -1098,7 +1175,7 @@ async def test_model_armor_with_default_credentials(): guardrail_name="model-armor-test", credentials=None, # Explicitly set to None to test default auth ) - + # Mock the Model Armor API response mock_response = AsyncMock() mock_response.status_code = 200 @@ -1106,10 +1183,10 @@ async def test_model_armor_with_default_credentials(): "sanitized_text": "Test content", "action": "SANITIZE" }) - + # Mock the access token method to simulate successful auth guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "cloud-test-project")) - + # Mock the async handler with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=mock_response)) as mock_post: request_data = { @@ -1119,7 +1196,7 @@ async def test_model_armor_with_default_credentials(): ], "metadata": {"guardrails": ["model-armor-test"]} } - + # This should not raise ValueError about project_id result = await guardrail.async_pre_call_hook( user_api_key_dict=mock_user_api_key_dict, @@ -1127,7 +1204,7 @@ async def test_model_armor_with_default_credentials(): data=request_data, call_type="completion" ) - + # Verify the project_id was used correctly in the API call mock_post.assert_called_once() call_args = mock_post.call_args @@ -1241,6 +1318,11 @@ async def test_async_moderation_hook_content_blocked(): assert "_model_armor_response" in request_data["metadata"] assert request_data["metadata"]["_model_armor_status"] == "blocked" + # IMPORTANT: Verify that applied_guardrails is populated even when blocked + # This is a regression test for the issue where applied_guardrails was null when blocked + assert "applied_guardrails" in request_data["metadata"] + assert "model-armor-test" in request_data["metadata"]["applied_guardrails"] + @pytest.mark.asyncio async def test_async_moderation_hook_with_sanitization(): @@ -1446,4 +1528,4 @@ async def test_async_moderation_hook_api_error_fail_on_error_false(): call_type="completion" ) - assert "API Error" in str(exc_info.value) \ No newline at end of file + assert "API Error" in str(exc_info.value) From 1e8848ca97bd53e596e715162d35d0d7953c9a08 Mon Sep 17 00:00:00 2001 From: Carlo Alberto Ferraris Date: Sun, 1 Feb 2026 08:07:47 +0900 Subject: [PATCH 14/22] add missing indexes on VerificationToken table (#20040) --- .../migration.sql | 8 ++++++++ .../litellm_proxy_extras/schema.prisma | 10 ++++++++++ litellm/proxy/schema.prisma | 10 ++++++++++ schema.prisma | 10 ++++++++++ 4 files changed, 38 insertions(+) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260129103648_add_verificationtoken_indexes/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260129103648_add_verificationtoken_indexes/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260129103648_add_verificationtoken_indexes/migration.sql new file mode 100644 index 00000000000..572eea9b529 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260129103648_add_verificationtoken_indexes/migration.sql @@ -0,0 +1,8 @@ +-- CreateIndex +CREATE INDEX "LiteLLM_VerificationToken_user_id_team_id_idx" ON "LiteLLM_VerificationToken"("user_id", "team_id"); + +-- CreateIndex +CREATE INDEX "LiteLLM_VerificationToken_team_id_idx" ON "LiteLLM_VerificationToken"("team_id"); + +-- CreateIndex +CREATE INDEX "LiteLLM_VerificationToken_budget_reset_at_expires_idx" ON "LiteLLM_VerificationToken"("budget_reset_at", "expires"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index b118400b620..3b81da10923 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -305,6 +305,16 @@ model LiteLLM_VerificationToken { litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) + + // SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_VerificationToken"."token" FROM "public"."LiteLLM_VerificationToken" WHERE ("public"."LiteLLM_VerificationToken"."user_id" = $1 AND ("public"."LiteLLM_VerificationToken"."team_id" IS NULL OR "public"."LiteLLM_VerificationToken"."team_id" <> $2)) OFFSET $3 ) AS "sub" + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."user_id" = $1 OFFSET $2 + @@index([user_id, team_id]) + + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."team_id" = $1 OFFSET $2 + @@index([team_id]) + + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 + @@index([budget_reset_at, expires]) } // Audit table for deleted keys - preserves spend and key information for historical tracking diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index b118400b620..3b81da10923 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -305,6 +305,16 @@ model LiteLLM_VerificationToken { litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) + + // SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_VerificationToken"."token" FROM "public"."LiteLLM_VerificationToken" WHERE ("public"."LiteLLM_VerificationToken"."user_id" = $1 AND ("public"."LiteLLM_VerificationToken"."team_id" IS NULL OR "public"."LiteLLM_VerificationToken"."team_id" <> $2)) OFFSET $3 ) AS "sub" + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."user_id" = $1 OFFSET $2 + @@index([user_id, team_id]) + + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."team_id" = $1 OFFSET $2 + @@index([team_id]) + + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 + @@index([budget_reset_at, expires]) } // Audit table for deleted keys - preserves spend and key information for historical tracking diff --git a/schema.prisma b/schema.prisma index b118400b620..3b81da10923 100644 --- a/schema.prisma +++ b/schema.prisma @@ -305,6 +305,16 @@ model LiteLLM_VerificationToken { litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) + + // SELECT COUNT(*) FROM (SELECT "public"."LiteLLM_VerificationToken"."token" FROM "public"."LiteLLM_VerificationToken" WHERE ("public"."LiteLLM_VerificationToken"."user_id" = $1 AND ("public"."LiteLLM_VerificationToken"."team_id" IS NULL OR "public"."LiteLLM_VerificationToken"."team_id" <> $2)) OFFSET $3 ) AS "sub" + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."user_id" = $1 OFFSET $2 + @@index([user_id, team_id]) + + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE "public"."LiteLLM_VerificationToken"."team_id" = $1 OFFSET $2 + @@index([team_id]) + + // SELECT ... FROM "public"."LiteLLM_VerificationToken" WHERE (("public"."LiteLLM_VerificationToken"."expires" IS NULL OR "public"."LiteLLM_VerificationToken"."expires" > $1) AND "public"."LiteLLM_VerificationToken"."budget_reset_at" < $2) OFFSET $3 + @@index([budget_reset_at, expires]) } // Audit table for deleted keys - preserves spend and key information for historical tracking From 1985aa04fadd046a1de43cb206c445efa8fa1169 Mon Sep 17 00:00:00 2001 From: jquinter Date: Sat, 31 Jan 2026 20:09:32 -0300 Subject: [PATCH 15/22] Fix Nova grounding web_search_options={} not applying systemTool (#20044) * Fix Nova grounding web_search_options={} not applying systemTool Two bugs prevented web_search_options={} from working for Nova grounding: 1. Empty dict falsy check: The condition `value and isinstance(value, dict)` short-circuits to False when value is {} (empty dict is falsy in Python). Changed to `isinstance(value, dict)` to match Anthropic's implementation. 2. Pre-formatted tools mangled by _bedrock_tools_pt: The systemTool (already in Bedrock format) was added to optional_params["tools"], but _process_tools_and_beta passed all tools through _bedrock_tools_pt which expects OpenAI-format tools. This corrupted the systemTool into an empty toolSpec. Fixed by separating systemTool blocks before transformation and appending them after. Fixes follow-up to #19598 Co-Authored-By: Claude Opus 4.5 * Fix python-multipart Python version constraint for Poetry lock python-multipart ^0.0.22 requires Python >=3.10 but the project supports >=3.9. Add python = ">=3.10" marker so Poetry can resolve dependencies for Python 3.9. Co-Authored-By: Claude Opus 4.5 --------- Co-authored-by: Claude Opus 4.5 --- .../bedrock/chat/converse_transformation.py | 11 +++++++++- poetry.lock | 20 ++++++++++++++++++- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index d4e4d3591ba..f6d7e128580 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1081,10 +1081,16 @@ class AmazonConverseConfig(BaseConfig): user_betas = get_anthropic_beta_from_headers(headers) anthropic_beta_list.extend(user_betas) - # Filter out tool search tools - Bedrock Converse API doesn't support them + # Separate pre-formatted Bedrock tools (e.g. systemTool from web_search_options) + # from OpenAI-format tools that need transformation via _bedrock_tools_pt filtered_tools = [] + pre_formatted_tools: List[ToolBlock] = [] if original_tools: for tool in original_tools: + # Already-formatted Bedrock tools (e.g. systemTool for Nova grounding) + if "systemTool" in tool: + pre_formatted_tools.append(tool) + continue tool_type = tool.get("type", "") if tool_type in ( "tool_search_tool_regex_20251119", @@ -1116,6 +1122,9 @@ class AmazonConverseConfig(BaseConfig): # No computer use tools, process all tools as regular tools bedrock_tools = _bedrock_tools_pt(filtered_tools) + # Append pre-formatted tools (systemTool etc.) after transformation + bedrock_tools.extend(pre_formatted_tools) + # Set anthropic_beta in additional_request_params if we have any beta features # ONLY apply to Anthropic/Claude models - other models (e.g., Qwen, Llama) don't support this field # and will error with "unknown variant anthropic_beta" if included diff --git a/poetry.lock b/poetry.lock index 537367c5aa0..3b2a8d20f20 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.2.0 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.1.4 and should not be changed by hand. [[package]] name = "a2a-sdk" @@ -5704,6 +5704,24 @@ pytest = ">=7.0.0" [package.extras] dev = ["black", "flake8", "isort", "mypy"] +[[package]] +name = "pytest-retry" +version = "1.7.0" +description = "Adds the ability to retry flaky tests in CI environments" +optional = false +python-versions = ">=3.9" +groups = ["dev"] +files = [ + {file = "pytest_retry-1.7.0-py3-none-any.whl", hash = "sha256:a2dac85b79a4e2375943f1429479c65beb6c69553e7dae6b8332be47a60954f4"}, + {file = "pytest_retry-1.7.0.tar.gz", hash = "sha256:f8d52339f01e949df47c11ba9ee8d5b362f5824dff580d3870ec9ae0057df80f"}, +] + +[package.dependencies] +pytest = ">=7.0.0" + +[package.extras] +dev = ["black", "flake8", "isort", "mypy"] + [[package]] name = "python-dateutil" version = "2.9.0.post0" From 07bffddbfa4240fcaf5705f2311345aeb1300ff3 Mon Sep 17 00:00:00 2001 From: Simon Lynch <44986346+srlynch1@users.noreply.github.com> Date: Sun, 1 Feb 2026 10:11:16 +1100 Subject: [PATCH 16/22] fix(bedrock): deduplicate toolResult and toolUse blocks in Converse message transformation (#20049) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bedrock rejects requests when toolResult or toolUse blocks within a single message contain duplicate IDs. The Converse message transformer merges consecutive tool/assistant messages without checking for duplicate toolUseId values, causing BedrockException errors. Add _deduplicate_bedrock_content_blocks() — a generalized helper that removes duplicate blocks by ID, logs a warning for each dropped duplicate via verbose_logger, and preserves non-tool blocks (e.g. cachePoint). Apply it at all four merge sites (sync/async × toolResult/ toolUse). The Anthropic /messages path was fixed in PR #19324; this applies the equivalent fix to the Bedrock Converse path. Fixes #20048 Co-authored-by: Claude Opus 4.5 --- .../prompt_templates/factory.py | 61 ++++ .../test_bedrock_converse_dedup_factory.py | 332 ++++++++++++++++++ 2 files changed, 393 insertions(+) create mode 100644 tests/litellm_core_utils/test_bedrock_converse_dedup_factory.py diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 0e1637a65ba..09b7c5374da 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -3399,6 +3399,59 @@ def _convert_to_bedrock_tool_call_result( return content_block +def _deduplicate_bedrock_content_blocks( + blocks: List[BedrockContentBlock], + block_key: str, + id_key: str = "toolUseId", +) -> List[BedrockContentBlock]: + """ + Remove duplicate content blocks that share the same ID under ``block_key``. + + Bedrock requires all toolResult and toolUse IDs within a single message to + be unique. When merging consecutive messages, duplicates can occur if the + same tool_call_id appears multiple times in conversation history. + + When duplicates exist, the first occurrence is retained and subsequent ones + are discarded. A warning is logged for every dropped block so that + upstream duplication bugs remain visible. + + Blocks that do not contain ``block_key`` (e.g., cachePoint, text) are + always preserved. + + Args: + blocks: The list of Bedrock content blocks to deduplicate. + block_key: The dict key to inspect (e.g. ``"toolResult"`` or ``"toolUse"``). + id_key: The nested key that holds the unique ID (default ``"toolUseId"``). + """ + seen_ids: Set[str] = set() + deduplicated: List[BedrockContentBlock] = [] + for block in blocks: + keyed = block.get(block_key) + if keyed is not None: + block_id = keyed.get(id_key) + if block_id: + if block_id in seen_ids: + verbose_logger.warning( + "Bedrock Converse: dropping duplicate %s block with " + "%s=%s. This may indicate duplicate tool messages in " + "conversation history.", + block_key, + id_key, + block_id, + ) + continue + seen_ids.add(block_id) + deduplicated.append(block) + return deduplicated + + +def _deduplicate_bedrock_tool_content( + tool_content: List[BedrockContentBlock], +) -> List[BedrockContentBlock]: + """Convenience wrapper: deduplicate ``toolResult`` blocks by ``toolUseId``.""" + return _deduplicate_bedrock_content_blocks(tool_content, "toolResult") + + def _insert_assistant_continue_message( messages: List[BedrockMessageBlock], assistant_continue_message: Optional[ @@ -3867,6 +3920,8 @@ class BedrockConverseMessagesProcessor: tool_content.append(cache_point_block) msg_i += 1 + # Deduplicate toolResult blocks with the same toolUseId + tool_content = _deduplicate_bedrock_tool_content(tool_content) if tool_content: # if last message was a 'user' message, then add a blank assistant message (bedrock requires alternating roles) if len(contents) > 0 and contents[-1]["role"] == "user": @@ -3980,6 +4035,8 @@ class BedrockConverseMessagesProcessor: msg_i += 1 + assistant_content = _deduplicate_bedrock_content_blocks(assistant_content, "toolUse") + if assistant_content: contents.append( BedrockMessageBlock(role="assistant", content=assistant_content) @@ -4230,6 +4287,8 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915 tool_content.append(cache_point_block) msg_i += 1 + # Deduplicate toolResult blocks with the same toolUseId + tool_content = _deduplicate_bedrock_tool_content(tool_content) if tool_content: # if last message was a 'user' message, then add a blank assistant message (bedrock requires alternating roles) if len(contents) > 0 and contents[-1]["role"] == "user": @@ -4336,6 +4395,8 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915 msg_i += 1 + assistant_content = _deduplicate_bedrock_content_blocks(assistant_content, "toolUse") + if assistant_content: contents.append( BedrockMessageBlock(role="assistant", content=assistant_content) diff --git a/tests/litellm_core_utils/test_bedrock_converse_dedup_factory.py b/tests/litellm_core_utils/test_bedrock_converse_dedup_factory.py new file mode 100644 index 00000000000..0969c77299c --- /dev/null +++ b/tests/litellm_core_utils/test_bedrock_converse_dedup_factory.py @@ -0,0 +1,332 @@ + +import sys +import os +import pytest + +sys.path.insert(0, os.path.abspath(".")) + +from litellm.litellm_core_utils.prompt_templates.factory import ( + _bedrock_converse_messages_pt, + _deduplicate_bedrock_content_blocks, + _deduplicate_bedrock_tool_content, + BedrockConverseMessagesProcessor, +) + + +MODEL = "anthropic.claude-v2" +PROVIDER = "bedrock_converse" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_duplicate_tool_result_messages(): + """Return messages where two consecutive tool-role messages reference the + same tool_call_id, simulating the duplication scenario.""" + return [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "tooluse_abc123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "Paris"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "tooluse_abc123", + "content": '{"temp": 22}', + }, + { + "role": "tool", + "tool_call_id": "tooluse_abc123", # DUPLICATE + "content": '{"temp": 22}', + }, + ] + + +def _make_duplicate_tool_use_messages(): + """Return messages where two consecutive assistant messages carry tool_calls + with the same id, simulating assistant-side duplication.""" + return [ + {"role": "user", "content": "Do something"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "tool_1", + "type": "function", + "function": {"name": "fn_a", "arguments": "{}"}, + }, + ], + }, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "tool_1", # DUPLICATE + "type": "function", + "function": {"name": "fn_a", "arguments": "{}"}, + }, + ], + }, + # Need a tool result so the conversation is valid + { + "role": "tool", + "tool_call_id": "tool_1", + "content": '{"ok": true}', + }, + ] + + +def _extract_blocks(result, role, key): + """Extract all content blocks containing ``key`` from messages with ``role``.""" + return [ + block + for msg in result + if msg["role"] == role + for block in msg["content"] + if key in block + ] + + +# --------------------------------------------------------------------------- +# toolResult dedup tests +# --------------------------------------------------------------------------- + + +def test_bedrock_converse_deduplicates_tool_results(): + """Verify _bedrock_converse_messages_pt deduplicates toolResult blocks + with the same toolUseId when merging consecutive tool messages.""" + messages = _make_duplicate_tool_result_messages() + result = _bedrock_converse_messages_pt(messages, MODEL, PROVIDER) + + tool_results = _extract_blocks(result, "user", "toolResult") + ids = [tr["toolResult"]["toolUseId"] for tr in tool_results] + assert ids.count("tooluse_abc123") == 1 + + +@pytest.mark.asyncio +async def test_bedrock_converse_deduplicates_tool_results_async(): + """Verify the async path also deduplicates toolResult blocks with the + same toolUseId when merging consecutive tool messages.""" + messages = _make_duplicate_tool_result_messages() + result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages, MODEL, PROVIDER + ) + + tool_results = _extract_blocks(result, "user", "toolResult") + ids = [tr["toolResult"]["toolUseId"] for tr in tool_results] + assert ids.count("tooluse_abc123") == 1 + + +def test_bedrock_converse_preserves_unique_tool_results(): + """Different toolUseIds should all be preserved.""" + messages = [ + {"role": "user", "content": "Weather and time?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "tool_1", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"}, + }, + { + "id": "tool_2", + "type": "function", + "function": {"name": "get_time", "arguments": "{}"}, + }, + ], + }, + {"role": "tool", "tool_call_id": "tool_1", "content": '{"temp": 22}'}, + {"role": "tool", "tool_call_id": "tool_2", "content": '{"time": "14:00"}'}, + ] + + result = _bedrock_converse_messages_pt(messages, MODEL, PROVIDER) + + tool_results = _extract_blocks(result, "user", "toolResult") + assert len(tool_results) == 2 + ids = {tr["toolResult"]["toolUseId"] for tr in tool_results} + assert ids == {"tool_1", "tool_2"} + + +def test_bedrock_converse_dedup_preserves_cache_points(): + """cachePoint blocks should not be removed during dedup.""" + messages = [ + {"role": "user", "content": "Weather?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "tool_1", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "tool_1", + "content": [ + { + "type": "text", + "text": "sunny", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "tool_1", # DUPLICATE + "content": '{"temp": 22}', + }, + ] + + result = _bedrock_converse_messages_pt(messages, MODEL, PROVIDER) + + tool_results = _extract_blocks(result, "user", "toolResult") + cache_points = _extract_blocks(result, "user", "cachePoint") + + assert len(tool_results) == 1 + assert len(cache_points) == 1 + + +# --------------------------------------------------------------------------- +# toolUse dedup tests +# --------------------------------------------------------------------------- + + +def test_bedrock_converse_deduplicates_tool_use_sync(): + """Verify the sync path deduplicates toolUse blocks with the same + toolUseId when merging consecutive assistant messages.""" + messages = _make_duplicate_tool_use_messages() + result = _bedrock_converse_messages_pt(messages, MODEL, PROVIDER) + + tool_uses = _extract_blocks(result, "assistant", "toolUse") + ids = [tu["toolUse"]["toolUseId"] for tu in tool_uses] + assert ids.count("tool_1") == 1 + + +@pytest.mark.asyncio +async def test_bedrock_converse_deduplicates_tool_use_async(): + """Verify the async path deduplicates toolUse blocks with the same + toolUseId when merging consecutive assistant messages.""" + messages = _make_duplicate_tool_use_messages() + result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages, MODEL, PROVIDER + ) + + tool_uses = _extract_blocks(result, "assistant", "toolUse") + ids = [tu["toolUse"]["toolUseId"] for tu in tool_uses] + assert ids.count("tool_1") == 1 + + +@pytest.mark.asyncio +async def test_bedrock_converse_tool_use_sync_async_parity(): + """Sync and async paths should produce identical results for duplicate + toolUse blocks.""" + messages = _make_duplicate_tool_use_messages() + sync_result = _bedrock_converse_messages_pt(messages, MODEL, PROVIDER) + async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages, MODEL, PROVIDER + ) + assert sync_result == async_result + + +# --------------------------------------------------------------------------- +# Generalized helper unit tests +# --------------------------------------------------------------------------- + + +def test_deduplicate_bedrock_content_blocks_tool_result(): + """Direct unit test: first occurrence wins, duplicates dropped, non-tool + blocks preserved.""" + blocks = [ + {"toolResult": {"toolUseId": "id_1", "content": [{"text": "a"}]}}, + {"cachePoint": {"type": "default"}}, + {"toolResult": {"toolUseId": "id_1", "content": [{"text": "b"}]}}, # duplicate + {"toolResult": {"toolUseId": "id_2", "content": [{"text": "c"}]}}, + ] + + result = _deduplicate_bedrock_content_blocks(blocks, "toolResult") + + assert len(result) == 3 # id_1, cachePoint, id_2 + tool_ids = [b["toolResult"]["toolUseId"] for b in result if "toolResult" in b] + assert tool_ids == ["id_1", "id_2"] + # First-wins: content "a" is kept, "b" is dropped + assert result[0]["toolResult"]["content"] == [{"text": "a"}] + + +def test_deduplicate_bedrock_content_blocks_tool_use(): + """Direct unit test of toolUse dedup via the generalized helper.""" + blocks = [ + {"toolUse": {"toolUseId": "id_1", "name": "fn_a", "input": {}}}, + {"text": "thinking..."}, + {"toolUse": {"toolUseId": "id_1", "name": "fn_a", "input": {}}}, # duplicate + {"toolUse": {"toolUseId": "id_2", "name": "fn_b", "input": {}}}, + ] + + result = _deduplicate_bedrock_content_blocks(blocks, "toolUse") + + assert len(result) == 3 # id_1, text, id_2 + tool_ids = [b["toolUse"]["toolUseId"] for b in result if "toolUse" in b] + assert tool_ids == ["id_1", "id_2"] + + +def test_deduplicate_preserves_blocks_with_missing_id(): + """Blocks where toolUseId is None or empty should pass through without + dedup tracking (they cannot be compared).""" + blocks = [ + {"toolResult": {"toolUseId": None, "content": [{"text": "a"}]}}, + {"toolResult": {"toolUseId": "", "content": [{"text": "b"}]}}, + {"toolResult": {"toolUseId": "id_1", "content": [{"text": "c"}]}}, + ] + + result = _deduplicate_bedrock_content_blocks(blocks, "toolResult") + + # All three should be preserved — None and "" are not tracked + assert len(result) == 3 + + +def test_deduplicate_bedrock_tool_content_convenience_wrapper(): + """The convenience wrapper should behave identically to calling the + generalized helper with block_key='toolResult'.""" + blocks = [ + {"toolResult": {"toolUseId": "id_1", "content": [{"text": "a"}]}}, + {"toolResult": {"toolUseId": "id_1", "content": [{"text": "b"}]}}, + ] + + assert _deduplicate_bedrock_tool_content(blocks) == _deduplicate_bedrock_content_blocks(blocks, "toolResult") + + +# --------------------------------------------------------------------------- +# Sync/async parity for toolResult +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_bedrock_converse_sync_async_parity_with_duplicates(): + """Sync and async paths should produce identical results with duplicate + tool results.""" + messages = _make_duplicate_tool_result_messages() + + sync_result = _bedrock_converse_messages_pt(messages, MODEL, PROVIDER) + async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages, MODEL, PROVIDER + ) + + assert sync_result == async_result From ef73f330f1f216bb98ac21caaf7056a98779eb9c Mon Sep 17 00:00:00 2001 From: Abdullah Habib Biswas Date: Sun, 1 Feb 2026 04:46:07 +0530 Subject: [PATCH 17/22] fix: prevent error when max_fallbacks exceeds available models (#20071) --- .../router_utils/fallback_event_handlers.py | 12 +++++- tests/test_fallbacks.py | 42 +++++++++++++++++++ 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 62e706a0cf5..738b82d7023 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -113,8 +113,16 @@ async def run_async_fallback( The most recent exception if all fallback model groups fail. """ - ### BASE CASE ### MAX FALLBACK DEPTH REACHED - if fallback_depth >= max_fallbacks: + ### BASE CASE ### MAX FALLBACK DEPTH REACHED + if fallback_depth >= max_fallbacks: + raise original_exception + + ### CHECK IF MODEL GROUP LIST EXHAUSTED + if original_model_group in fallback_model_group: + fallback_group_length = len(fallback_model_group) - 1 + else: + fallback_group_length = len(fallback_model_group) + if fallback_depth >= fallback_group_length: raise original_exception error_from_fallbacks = original_exception diff --git a/tests/test_fallbacks.py b/tests/test_fallbacks.py index bc9aa4c64c8..c22cefa6be6 100644 --- a/tests/test_fallbacks.py +++ b/tests/test_fallbacks.py @@ -336,3 +336,45 @@ async def test_chat_completion_bad_and_good_model(): f"Iteration {iteration + 1}: {'✓' if success else '✗'} ({time.time() - start_time:.2f}s)" ) assert success, "Not all good model requests succeeded" + + +@pytest.mark.asyncio +async def test_router_fallback_exhaustion(): + """ + Test for Bug 19985: + """ + from litellm import Router + import pytest + + # Setup: Only ONE fallback model available + model_list = [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "openai/fake", "api_key": "bad-key"}, + }, + { + "model_name": "bad-model-1", + "litellm_params": {"model": "azure/fake", "api_key": "bad-key"}, + } + ] + + # max_fallbacks=10 is much larger than the 1 fallback provided in the list + router = Router( + model_list=model_list, + fallbacks=[{"gpt-3.5-turbo": ["bad-model-1"]}], + max_fallbacks=10 + ) + + try: + # This will fail and attempt to fallback + await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "test"}] + ) + except Exception as e: + # The success criteria is that we DON'T get an IndexError + assert not isinstance(e, IndexError), f"Expected API error, but got IndexError: {e}" + # Also ensure we actually hit a fallback attempt + print(f"Caught expected exception: {type(e).__name__}") + + From 726988aed4db73e06315a385742600251a49a1d4 Mon Sep 17 00:00:00 2001 From: Lovro Seder Date: Sun, 1 Feb 2026 00:26:53 +0100 Subject: [PATCH 18/22] Fix Azure AI Anthropic CountTokens 401 auth error (#20069) Add x-api-key header to CountTokens handler to match chat completion authentication. Azure AI Anthropic requires this header per Microsoft's native API format. --- .../anthropic/count_tokens/transformation.py | 14 ++- ...e_anthropic_count_tokens_transformation.py | 111 ++++++++++++++++++ 2 files changed, 119 insertions(+), 6 deletions(-) create mode 100644 tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py index e284595cc8a..09b83b7c971 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -30,30 +30,32 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): """ Get the required headers for the Azure AI Anthropic CountTokens API. - Uses Azure authentication (api-key header) instead of Anthropic's x-api-key. + Azure AI Anthropic uses Anthropic's native API format, which requires the + x-api-key header for authentication (in addition to Azure's api-key header). Args: api_key: The Azure AI API key litellm_params: Optional LiteLLM parameters for additional auth config Returns: - Dictionary of required headers with Azure authentication + Dictionary of required headers with both x-api-key and Azure authentication """ - # Start with base headers + # Start with base headers including x-api-key for Anthropic API compatibility headers = { "Content-Type": "application/json", "anthropic-version": "2023-06-01", "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, + "x-api-key": api_key, # Azure AI Anthropic requires this header } - # Use Azure authentication + # Also set up Azure auth headers for flexibility litellm_params = litellm_params or {} if "api_key" not in litellm_params: litellm_params["api_key"] = api_key litellm_params_obj = GenericLiteLLMParams(**litellm_params) - # Get Azure auth headers + # Get Azure auth headers (api-key or Authorization) azure_headers = BaseAzureLLM._base_validate_azure_environment( headers={}, litellm_params=litellm_params_obj ) @@ -68,7 +70,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): Get the Azure AI Anthropic CountTokens API endpoint. Args: - api_base: The Azure AI API base URL + api_base: The Azure AI API base URL (e.g., https://my-resource.services.ai.azure.com or https://my-resource.services.ai.azure.com/anthropic) diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py new file mode 100644 index 00000000000..78806831685 --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py @@ -0,0 +1,111 @@ +""" +Tests for Azure AI Anthropic CountTokens transformation. + +Verifies that the CountTokens API uses the correct authentication headers. +""" +import os +import sys + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + + +from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( + AzureAIAnthropicCountTokensConfig, +) + + +class TestAzureAIAnthropicCountTokensConfig: + """Test Azure AI Anthropic CountTokens configuration and headers.""" + + def test_get_required_headers_includes_x_api_key(self): + """ + Test that get_required_headers includes x-api-key header. + + Azure AI Anthropic uses Anthropic's native API format which requires + the x-api-key header for authentication (not just Azure's api-key). + """ + config = AzureAIAnthropicCountTokensConfig() + api_key = "test-api-key-12345" + + headers = config.get_required_headers(api_key=api_key) + + # Verify x-api-key header is set + assert "x-api-key" in headers + assert headers["x-api-key"] == api_key + + # Verify base headers are present + assert headers["Content-Type"] == "application/json" + assert headers["anthropic-version"] == "2023-06-01" + assert "anthropic-beta" in headers + + def test_get_required_headers_includes_azure_api_key(self): + """ + Test that get_required_headers includes Azure api-key header. + + Both x-api-key and api-key headers should be present. + """ + config = AzureAIAnthropicCountTokensConfig() + api_key = "test-azure-key-67890" + + headers = config.get_required_headers(api_key=api_key) + + # Verify both authentication headers are set + assert "x-api-key" in headers + assert "api-key" in headers + assert headers["x-api-key"] == api_key + assert headers["api-key"] == api_key + + def test_get_required_headers_with_litellm_params(self): + """ + Test that get_required_headers works with litellm_params. + """ + config = AzureAIAnthropicCountTokensConfig() + api_key = "test-key" + litellm_params = {"api_key": "param-key", "custom_field": "value"} + + headers = config.get_required_headers( + api_key=api_key, litellm_params=litellm_params + ) + + # x-api-key should use the direct api_key parameter + assert headers["x-api-key"] == api_key + # Azure api-key should come from litellm_params + assert headers["api-key"] == "param-key" + + def test_get_count_tokens_endpoint_with_base_url(self): + """Test endpoint generation from base URL.""" + config = AzureAIAnthropicCountTokensConfig() + + api_base = "https://my-resource.services.ai.azure.com" + endpoint = config.get_count_tokens_endpoint(api_base) + + assert ( + endpoint + == "https://my-resource.services.ai.azure.com/anthropic/v1/messages/count_tokens" + ) + + def test_get_count_tokens_endpoint_with_anthropic_path(self): + """Test endpoint generation when base URL already includes /anthropic.""" + config = AzureAIAnthropicCountTokensConfig() + + api_base = "https://my-resource.services.ai.azure.com/anthropic" + endpoint = config.get_count_tokens_endpoint(api_base) + + assert ( + endpoint + == "https://my-resource.services.ai.azure.com/anthropic/v1/messages/count_tokens" + ) + + def test_get_count_tokens_endpoint_with_trailing_slash(self): + """Test endpoint generation with trailing slash in base URL.""" + config = AzureAIAnthropicCountTokensConfig() + + api_base = "https://my-resource.services.ai.azure.com/" + endpoint = config.get_count_tokens_endpoint(api_base) + + assert ( + endpoint + == "https://my-resource.services.ai.azure.com/anthropic/v1/messages/count_tokens" + ) From d7997db912b38fb22130e2b18bfa388d90c5601f Mon Sep 17 00:00:00 2001 From: Nate Tessman <140846984+ntessman-capsule@users.noreply.github.com> Date: Sat, 31 Jan 2026 15:30:10 -0800 Subject: [PATCH 19/22] fix: Include hidden params in chat response to responses api response transformation (#20084) * Include hidden_params in chat completion to responses transformation * add tests --- .../transformation.py | 1 + .../test_litellm_completion_responses.py | 71 +++++++++++++++++++ 2 files changed, 72 insertions(+) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 74cc87713da..df298f7c448 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1413,6 +1413,7 @@ class LiteLLMCompletionResponsesConfig: ), user=getattr(chat_completion_response, "user", None), ) + responses_api_response._hidden_params = getattr(chat_completion_response, "_hidden_params", {}) return responses_api_response @staticmethod diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index f7a1984d32f..5074bbf4397 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -468,6 +468,77 @@ class TestLiteLLMCompletionResponsesConfig: ] assert item.status != "stop" + def test_transform_chat_completion_response_preserves_hidden_params(self): + """Test that _hidden_params from chat completion response are preserved in responses API response""" + # Setup + chat_completion_response = ModelResponse( + id="test-response-id", + created=1234567890, + model="test-model", + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="Test response", + role="assistant", + ), + ) + ], + ) + # Set hidden params on the chat completion response + chat_completion_response._hidden_params = { + "model_id": "abc123", + "cache_key": "some-cache-key", + "custom_llm_provider": "openai", + } + + # Execute + responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Test", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) + + # Assert + assert hasattr(responses_api_response, "_hidden_params") + assert responses_api_response._hidden_params == { + "model_id": "abc123", + "cache_key": "some-cache-key", + "custom_llm_provider": "openai", + } + + def test_transform_chat_completion_response_handles_missing_hidden_params(self): + """Test that missing _hidden_params defaults to empty dict""" + # Setup - no _hidden_params set + chat_completion_response = ModelResponse( + id="test-response-id", + created=1234567890, + model="test-model", + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="Test response", + role="assistant", + ), + ) + ], + ) + + # Execute + responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Test", + responses_api_request={}, + chat_completion_response=chat_completion_response, + ) + + # Assert - should default to empty dict + assert hasattr(responses_api_response, "_hidden_params") + assert responses_api_response._hidden_params == {} class TestFunctionCallTransformation: """Test cases for function_call input transformation""" From a513cfdefa131a918559108839bde35207882271 Mon Sep 17 00:00:00 2001 From: Graham Neubig Date: Sat, 31 Jan 2026 15:33:25 -0800 Subject: [PATCH 20/22] fix: Set standard_logging_object for pass-through endpoints (#19887) Pass-through endpoints (like vLLM classify) were not setting standard_logging_object because _get_assembled_streaming_response returns None for non-ModelResponse results. This caused model_max_budget_limiter.async_log_success_event to raise ValueError('standard_logging_payload is required'). The fix adds an elif branch in async_success_handler that mirrors the non-pass-through code path. Co-authored-by: openhands Co-authored-by: Krish Dholakia --- litellm/litellm_core_utils/litellm_logging.py | 30 +++ .../test_litellm_logging.py | 217 ++++++++++++++++++ 2 files changed, 247 insertions(+) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4ad2d1002bc..1b3a687f1f3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2435,6 +2435,36 @@ class Logging(LiteLLMLoggingBaseClass): standard_built_in_tools_params=self.standard_built_in_tools_params, ) + # print standard logging payload + if ( + standard_logging_payload := self.model_call_details.get( + "standard_logging_object" + ) + ) is not None: + emit_standard_logging_payload(standard_logging_payload) + elif self.call_type == "pass_through_endpoint": + print_verbose( + "Async success callbacks: Got a pass-through endpoint response" + ) + + self.model_call_details["async_complete_streaming_response"] = result + + # cost calculation not possible for pass-through + self.model_call_details["response_cost"] = None + + ## STANDARDIZED LOGGING PAYLOAD + self.model_call_details[ + "standard_logging_object" + ] = get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj=result, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, + ) + # print standard logging payload if ( standard_logging_payload := self.model_call_details.get( diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 316bd49cf89..734d52918ba 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1062,6 +1062,223 @@ def test_append_system_prompt_messages(): assert result == messages +@pytest.mark.asyncio +async def test_async_success_handler_sets_standard_logging_object_for_pass_through_endpoints(): + """ + Test that async_success_handler sets standard_logging_object for pass-through endpoints + even when complete_streaming_response is None. + + This is a regression test for the bug where pass-through endpoints (like vLLM classify) + would not set standard_logging_object, causing model_max_budget_limiter to raise + ValueError("standard_logging_payload is required"). + + The fix adds an elif branch in async_success_handler to set standard_logging_object + for pass-through endpoints when complete_streaming_response is None. + """ + from datetime import datetime + from unittest.mock import patch + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import StandardPassThroughResponseObject + + # Create a logging object for a pass-through endpoint + logging_obj = LiteLLMLoggingObj( + model="unknown", + messages=[{"role": "user", "content": "test"}], + stream=False, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="test-function-id", + ) + + # Set up model_call_details with required fields + logging_obj.model_call_details = { + "litellm_params": { + "metadata": {}, + "proxy_server_request": {}, + }, + "litellm_call_id": "test-call-id", + } + + # Create a pass-through response object (not a ModelResponse) + result = StandardPassThroughResponseObject(response='{"status": "success"}') + + start_time = datetime.now() + end_time = datetime.now() + + # Mock the callbacks to avoid actual logging + with patch.object(logging_obj, "get_combined_callback_list", return_value=[]): + # Call async_success_handler + await logging_obj.async_success_handler( + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=False, + ) + + # Verify that standard_logging_object was set + assert "standard_logging_object" in logging_obj.model_call_details, ( + "standard_logging_object should be set for pass-through endpoints " + "even when complete_streaming_response is None" + ) + assert logging_obj.model_call_details["standard_logging_object"] is not None, ( + "standard_logging_object should not be None for pass-through endpoints" + ) + + # Verify that async_complete_streaming_response was set to prevent re-processing + # This is consistent with the existing code pattern for regular streaming + assert "async_complete_streaming_response" in logging_obj.model_call_details, ( + "async_complete_streaming_response should be set to prevent re-processing, " + "consistent with the existing code pattern" + ) + assert logging_obj.model_call_details["async_complete_streaming_response"] is result, ( + "async_complete_streaming_response should be set to the result" + ) + + # Verify that response_cost is set to None (cost calculation not possible for pass-through) + # This is consistent with the error handling in the non-pass-through code path + assert "response_cost" in logging_obj.model_call_details, ( + "response_cost should be set for pass-through endpoints" + ) + assert logging_obj.model_call_details["response_cost"] is None, ( + "response_cost should be None for pass-through endpoints since " + "StandardPassThroughResponseObject doesn't have standard usage info" + ) + + +@pytest.mark.asyncio +async def test_async_success_handler_prevents_reprocessing_for_pass_through_endpoints(): + """ + Test that async_success_handler prevents re-processing for pass-through endpoints + by setting async_complete_streaming_response, consistent with the existing code pattern. + + This ensures that if async_success_handler is called multiple times (e.g., during + streaming), it won't re-process the response after the first complete call. + """ + from datetime import datetime + from unittest.mock import patch + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import StandardPassThroughResponseObject + + # Create a logging object for a pass-through endpoint + logging_obj = LiteLLMLoggingObj( + model="unknown", + messages=[{"role": "user", "content": "test"}], + stream=False, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="test-call-id-reprocess", + function_id="test-function-id-reprocess", + ) + + # Set up model_call_details with required fields + logging_obj.model_call_details = { + "litellm_params": { + "metadata": {}, + "proxy_server_request": {}, + }, + "litellm_call_id": "test-call-id-reprocess", + } + + result = StandardPassThroughResponseObject(response='{"status": "success"}') + start_time = datetime.now() + end_time = datetime.now() + + # Mock the callbacks to avoid actual logging + with patch.object(logging_obj, "get_combined_callback_list", return_value=[]): + # First call - should process and set standard_logging_object + await logging_obj.async_success_handler( + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=False, + ) + + # Verify first call set the values + assert "standard_logging_object" in logging_obj.model_call_details + assert "async_complete_streaming_response" in logging_obj.model_call_details + first_standard_logging_object = logging_obj.model_call_details["standard_logging_object"] + + # Second call - should return early due to async_complete_streaming_response guard + with patch.object(logging_obj, "get_combined_callback_list", return_value=[]) as mock_callbacks: + await logging_obj.async_success_handler( + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=False, + ) + # The guard should cause early return, so get_combined_callback_list should not be called + mock_callbacks.assert_not_called() + + # Verify standard_logging_object wasn't modified by second call + assert logging_obj.model_call_details["standard_logging_object"] is first_standard_logging_object, ( + "standard_logging_object should not be modified on re-processing" + ) + + +@pytest.mark.asyncio +async def test_async_success_handler_sets_standard_logging_object_for_streaming_pass_through(): + """ + Test that async_success_handler sets standard_logging_object for streaming + pass-through endpoints when the response cannot be parsed into a ModelResponse. + + This covers the case where streaming pass-through endpoints for unknown providers + return a StandardPassThroughResponseObject instead of a ModelResponse. + """ + from datetime import datetime + from unittest.mock import patch + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import StandardPassThroughResponseObject + + # Create a logging object for a streaming pass-through endpoint + logging_obj = LiteLLMLoggingObj( + model="unknown", + messages=[{"role": "user", "content": "test"}], + stream=True, # Streaming request + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="test-call-id-streaming", + function_id="test-function-id-streaming", + ) + + # Set up model_call_details with required fields + logging_obj.model_call_details = { + "litellm_params": { + "metadata": {}, + "proxy_server_request": {}, + }, + "litellm_call_id": "test-call-id-streaming", + } + + # Create a pass-through response object (simulating unparseable streaming response) + result = StandardPassThroughResponseObject( + response='data: {"chunk": 1}\ndata: {"chunk": 2}\ndata: [DONE]' + ) + + start_time = datetime.now() + end_time = datetime.now() + + # Mock the callbacks to avoid actual logging + with patch.object(logging_obj, "get_combined_callback_list", return_value=[]): + # Call async_success_handler + await logging_obj.async_success_handler( + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=False, + ) + + # Verify that standard_logging_object was set + assert "standard_logging_object" in logging_obj.model_call_details, ( + "standard_logging_object should be set for streaming pass-through endpoints " + "even when the response cannot be parsed into a ModelResponse" + ) + assert logging_obj.model_call_details["standard_logging_object"] is not None, ( + "standard_logging_object should not be None for streaming pass-through endpoints" + ) def test_get_error_information_error_code_priority(): """ Test get_error_information prioritizes 'code' attribute over 'status_code' attribute From 7329fa8e7adf35e00c23206b72e2ab8de461daf4 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 2 Feb 2026 18:50:58 +0530 Subject: [PATCH 21/22] fix: litellm_oss_staging_01_31_2026_3 failing tests --- docs/my-website/docs/proxy/config_settings.md | 1 + .../adapters/streaming_iterator.py | 28 +++++++++++++------ ...al_pass_through_adapters_transformation.py | 1 - 3 files changed, 21 insertions(+), 9 deletions(-) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index bb2c7e01c80..264c7d765b3 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -321,6 +321,7 @@ router_settings: | redis_host | string | The host address for the Redis server. **Only set this if you have multiple instances of LiteLLM Proxy and want current tpm/rpm tracking to be shared across them** | | redis_password | string | The password for the Redis server. **Only set this if you have multiple instances of LiteLLM Proxy and want current tpm/rpm tracking to be shared across them** | | redis_port | string | The port number for the Redis server. **Only set this if you have multiple instances of LiteLLM Proxy and want current tpm/rpm tracking to be shared across them**| +| redis_db | int | The database number for the Redis server. **Only set this if you have multiple instances of LiteLLM Proxy and want current tpm/rpm tracking to be shared across them**| | enable_pre_call_check | boolean | If true, checks if a call is within the model's context window before making the call. [More information here](reliability) | | content_policy_fallbacks | array of objects | Specifies fallback models for content policy violations. [More information here](reliability) | | fallbacks | array of objects | Specifies fallback models for all types of errors. [More information here](reliability) | diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index aa2f0cc08f1..a86820f82e8 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -409,10 +409,17 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): ) # Restore original tool name if it was truncated for OpenAI's 64-char limit - if block_type == "tool_use" and content_block_start.get("name"): - truncated_name = content_block_start.get("name", "") - original_name = self.tool_name_mapping.get(truncated_name, truncated_name) - content_block_start["name"] = original_name + if block_type == "tool_use": + # Type narrowing: content_block_start is ToolUseBlock when block_type is "tool_use" + from typing import cast + from litellm.types.llms.anthropic import ToolUseBlock + + tool_block = cast(ToolUseBlock, content_block_start) + + if tool_block.get("name"): + truncated_name = tool_block["name"] + original_name = self.tool_name_mapping.get(truncated_name, truncated_name) + tool_block["name"] = original_name if block_type != self.current_content_block_type: self.current_content_block_type = block_type @@ -421,9 +428,14 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): # For parallel tool calls, we'll necessarily have a new content block # if we get a function name since it signals a new tool call - if block_type == "tool_use" and content_block_start.get("name"): - self.current_content_block_type = block_type - self.current_content_block_start = content_block_start - return True + if block_type == "tool_use": + from typing import cast + from litellm.types.llms.anthropic import ToolUseBlock + + tool_block = cast(ToolUseBlock, content_block_start) + if tool_block.get("name"): + self.current_content_block_type = block_type + self.current_content_block_start = content_block_start + return True return False diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index bcb0059fd19..6a5c022ac7a 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -1531,7 +1531,6 @@ def test_translate_openai_response_to_anthropic_with_reasoning_content_only(): assert cast(Any, anthropic_content[1]).text == "There are **3** \"r\"s in the word strawberry." assert anthropic_response.get("stop_reason") == "end_turn" - assert tool_name_mapping == {} # No truncation needed for short names # ===================================================================== From b85f1f2e6d1287d373da70dc20c5f7c32ec8dfcd Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 2 Feb 2026 19:00:12 +0530 Subject: [PATCH 22/22] fix: litellm_core_utils/prompt_templates/factory.py:3431 --- litellm/litellm_core_utils/prompt_templates/factory.py | 2 +- litellm/llms/anthropic/chat/guardrail_translation/handler.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 09b7c5374da..c4c56a8d335 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -3427,7 +3427,7 @@ def _deduplicate_bedrock_content_blocks( deduplicated: List[BedrockContentBlock] = [] for block in blocks: keyed = block.get(block_key) - if keyed is not None: + if keyed is not None and isinstance(keyed, dict): block_id = keyed.get(id_key) if block_id: if block_id in seen_ids: diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 71d74121a30..8e1016bd5bd 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -74,7 +74,7 @@ class AnthropicMessagesHandler(BaseTranslation): if messages is None: return data - chat_completion_compatible_request = ( + chat_completion_compatible_request, tool_name_mapping = ( LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( anthropic_message_request=cast(AnthropicMessagesRequest, data) )