From 81fefc69c9b77470d94e2247e910380dd0972810 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Benedikt=20=C3=93skarsson?= Date: Thu, 8 Jan 2026 00:44:59 +0000 Subject: [PATCH 01/58] fix(bedrock): handle thinking with tool calls for Claude 4 models --- .../bedrock/chat/converse_transformation.py | 38 +- litellm/llms/bedrock/chat/invoke_handler.py | 36 +- litellm/utils.py | 58 ++- .../chat/test_converse_transformation.py | 486 +++++++++++------- 4 files changed, 398 insertions(+), 220 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 13dbec3952a..9935cce68dc 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -53,7 +53,12 @@ from litellm.types.utils import ( PromptTokensDetailsWrapper, Usage, ) -from litellm.utils import add_dummy_tool, has_tool_call_blocks, supports_reasoning +from litellm.utils import ( + add_dummy_tool, + has_tool_call_blocks, + last_assistant_with_tool_calls_has_no_thinking_blocks, + supports_reasoning, +) from ..common_utils import ( BedrockError, @@ -729,7 +734,7 @@ class AmazonConverseConfig(BaseConfig): return optional_params """ - Follow similar approach to anthropic - translate to a single tool call. + Follow similar approach to anthropic - translate to a single tool call. When using tools in this way: - https://docs.anthropic.com/en/docs/build-with-claude/tool-use#json-mode - You usually want to provide a single tool @@ -912,16 +917,16 @@ class AmazonConverseConfig(BaseConfig): inference_params = { k: v for k, v in inference_params.items() if k in total_supported_params } - + # Only set the topK value in for models that support it additional_request_params.update( self._handle_top_k_value(model, inference_params) ) - + # Filter out internal/MCP-related parameters that shouldn't be sent to the API # These are LiteLLM internal parameters, not API parameters additional_request_params = filter_internal_params(additional_request_params) - + # Filter out non-serializable objects (exceptions, callables, logging objects, etc.) # from additional_request_params to prevent JSON serialization errors # This filters: Exception objects, callable objects (functions), Logging objects, etc. @@ -1021,9 +1026,24 @@ class AmazonConverseConfig(BaseConfig): llm_provider="bedrock", ) + # Drop thinking param if thinking is enabled but thinking_blocks are missing + # This prevents the error: "Expected thinking or redacted_thinking, but found tool_use" + # Related issues: https://github.com/BerriAI/litellm/issues/14194 + if ( + optional_params.get("thinking") is not None + and messages is not None + and last_assistant_with_tool_calls_has_no_thinking_blocks(messages) + ): + if litellm.modify_params: + optional_params.pop("thinking", None) + litellm.verbose_logger.warning( + "Dropping 'thinking' param because the last assistant message with tool_calls " + "has no thinking_blocks. The model won't use extended thinking for this turn." + ) + # Prepare and separate parameters - inference_params, additional_request_params, request_metadata = ( - self._prepare_request_params(optional_params, model) + inference_params, additional_request_params, request_metadata = self._prepare_request_params( + optional_params, model ) original_tools = inference_params.pop("tools", []) @@ -1410,11 +1430,11 @@ class AmazonConverseConfig(BaseConfig): ) """ - Bedrock Response Object has optional message block + Bedrock Response Object has optional message block completion_response["output"].get("message", None) - A message block looks like this (Example 1): + A message block looks like this (Example 1): "output": { "message": { "role": "assistant", diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 49292545208..ea612da52c6 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -374,6 +374,29 @@ class BedrockLLM(BaseAWSLLM): def __init__(self) -> None: super().__init__() + @staticmethod + def is_claude_messages_api_model(model: str) -> bool: + """ + Check if the model uses the Claude Messages API (Claude 3+). + + Handles: + - Regional prefixes: eu.anthropic.claude-*, us.anthropic.claude-* + - Claude 3 models: claude-3-haiku, claude-3-sonnet, claude-3-opus, claude-3-5-*, claude-3-7-* + - Claude 4 models: claude-opus-4, claude-sonnet-4, claude-haiku-4 + """ + # Normalize model string to lowercase for matching + model_lower = model.lower() + + # Claude 3+ indicators (all use Messages API) + messages_api_indicators = [ + "claude-3", # Claude 3.x models + "claude-opus-4", # Claude Opus 4 + "claude-sonnet-4", # Claude Sonnet 4 + "claude-haiku-4", # Claude Haiku 4 + ] + + return any(indicator in model_lower for indicator in messages_api_indicators) + def convert_messages_to_prompt( self, model, messages, provider, custom_prompt_dict ) -> Tuple[str, Optional[list]]: @@ -465,7 +488,7 @@ class BedrockLLM(BaseAWSLLM): completion_response["generations"][0]["finish_reason"] ) elif provider == "anthropic": - if model.startswith("anthropic.claude-3"): + if self.is_claude_messages_api_model(model): json_schemas: dict = {} _is_function_call = False ## Handle Tool Calling @@ -595,13 +618,12 @@ class BedrockLLM(BaseAWSLLM): outputText = choice["message"].get("content") elif "text" in choice: # fallback for completion format outputText = choice["text"] - # Set finish reason if "finish_reason" in choice: model_response.choices[0].finish_reason = map_finish_reason( choice["finish_reason"] ) - + # Set usage if available if "usage" in completion_response: usage = completion_response["usage"] @@ -842,7 +864,7 @@ class BedrockLLM(BaseAWSLLM): ] = True # cohere requires stream = True in inference params data = json.dumps({"prompt": prompt, **inference_params}) elif provider == "anthropic": - if model.startswith("anthropic.claude-3"): + if self.is_claude_messages_api_model(model): # Separate system prompt from rest of message system_prompt_idx: list[int] = [] system_messages: list[str] = [] @@ -940,13 +962,13 @@ class BedrockLLM(BaseAWSLLM): # Use AmazonBedrockOpenAIConfig for proper OpenAI transformation openai_config = AmazonBedrockOpenAIConfig() supported_params = openai_config.get_supported_openai_params(model=model) - + # Filter to only supported OpenAI params filtered_params = { - k: v for k, v in inference_params.items() + k: v for k, v in inference_params.items() if k in supported_params } - + # OpenAI uses messages format, not prompt data = json.dumps({"messages": messages, **filtered_params}) else: diff --git a/litellm/utils.py b/litellm/utils.py index fbbaa94f7a1..5011b45c0f7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -621,7 +621,7 @@ def load_credentials_from_list(kwargs: dict): """ # Access CredentialAccessor via module to trigger lazy loading if needed CredentialAccessor = getattr(sys.modules[__name__], 'CredentialAccessor') - + credential_name = kwargs.get("litellm_credential_name") if credential_name and litellm.credential_list: credential_accessor = CredentialAccessor.get_credential_values(credential_name) @@ -648,7 +648,7 @@ def _is_gemini_model(model: Optional[str], custom_llm_provider: Optional[str]) - if custom_llm_provider in ["vertex_ai", "vertex_ai_beta"]: return model is not None and "gemini" in model.lower() return True - + # Check if model name contains gemini return model is not None and "gemini" in model.lower() @@ -670,7 +670,7 @@ def _process_assistant_message_tool_calls( """ role = msg_copy.get("role") tool_calls = msg_copy.get("tool_calls") - + if role == "assistant" and isinstance(tool_calls, list): new_tool_calls = [] for tc in tool_calls: @@ -683,17 +683,17 @@ def _process_assistant_message_tool_calls( else: new_tool_calls.append(tc) continue - + # Remove thought signature from ID if present if isinstance(tc_dict.get("id"), str): if thought_signature_separator in tc_dict["id"]: tc_dict["id"] = _remove_thought_signature_from_id( tc_dict["id"], thought_signature_separator ) - + new_tool_calls.append(tc_dict) msg_copy["tool_calls"] = new_tool_calls - + return msg_copy @@ -708,7 +708,7 @@ def _process_tool_message_id(msg_copy: dict, thought_signature_separator: str) - msg_copy["tool_call_id"] = _remove_thought_signature_from_id( msg_copy["tool_call_id"], thought_signature_separator ) - + return msg_copy @@ -719,7 +719,7 @@ def _remove_thought_signatures_from_messages( Remove thought signatures from tool call IDs in all messages. """ processed_messages = [] - + for msg in messages: # Handle Pydantic models (convert to dict) if hasattr(msg, "model_dump"): @@ -730,17 +730,17 @@ def _remove_thought_signatures_from_messages( # Unknown type, keep as is processed_messages.append(msg) continue - + # Process assistant messages with tool_calls msg_dict = _process_assistant_message_tool_calls( msg_dict, thought_signature_separator ) - + # Process tool messages with tool_call_id msg_dict = _process_tool_message_id(msg_dict, thought_signature_separator) - + processed_messages.append(msg_dict) - + return processed_messages @@ -960,7 +960,7 @@ def function_setup( # noqa: PLR0915 input=buffer.getvalue(), model=model, ) - + ### REMOVE THOUGHT SIGNATURES FROM TOOL CALL IDS FOR NON-GEMINI MODELS ### # Gemini models embed thought signatures in tool call IDs. When sending # messages with tool calls to non-Gemini providers, we need to remove these @@ -976,7 +976,7 @@ def function_setup( # noqa: PLR0915 # Get custom_llm_provider to determine target provider custom_llm_provider = kwargs.get("custom_llm_provider") - + # If custom_llm_provider not in kwargs, try to determine it from the model if not custom_llm_provider and model: try: @@ -987,18 +987,18 @@ def function_setup( # noqa: PLR0915 except Exception: # If we can't determine the provider, skip this processing pass - + # Only process if target is NOT a Gemini model if not _is_gemini_model(model, custom_llm_provider): verbose_logger.debug( "Removing thought signatures from tool call IDs for non-Gemini model" ) - + # Process messages to remove thought signatures processed_messages = _remove_thought_signatures_from_messages( messages, THOUGHT_SIGNATURE_SEPARATOR ) - + # Update messages in kwargs or args if "messages" in kwargs: kwargs["messages"] = processed_messages @@ -2977,7 +2977,7 @@ def get_optional_params_embeddings( # noqa: PLR0915 ): # Lazy load get_supported_openai_params get_supported_openai_params = getattr(sys.modules[__name__], 'get_supported_openai_params') - + # retrieve all parameters passed to the function passed_params = locals() custom_llm_provider = passed_params.pop("custom_llm_provider", None) @@ -4063,7 +4063,14 @@ def get_optional_params( # noqa: PLR0915 ), ) elif "anthropic" in bedrock_base_model and bedrock_route == "invoke": - if bedrock_base_model.startswith("anthropic.claude-3"): + # Check for Claude 3+ models (Messages API) including regional prefixes and Claude 4 + # Models like eu.anthropic.claude-opus-4-5, us.anthropic.claude-3-5-sonnet, etc. + bedrock_base_model_lower = bedrock_base_model.lower() + is_messages_api_model = any( + indicator in bedrock_base_model_lower + for indicator in ["claude-3", "claude-opus-4", "claude-sonnet-4", "claude-haiku-4"] + ) + if is_messages_api_model: optional_params = ( litellm.AmazonAnthropicClaudeConfig().map_openai_params( non_default_params=non_default_params, @@ -6911,7 +6918,7 @@ def get_valid_models( # init litellm_params ################################# from litellm.types.router import LiteLLM_Params - + if litellm_params is None: litellm_params = LiteLLM_Params(model="") if api_key is not None: @@ -7513,7 +7520,7 @@ class ProviderConfigManager: return litellm.IBMWatsonXAIConfig() elif litellm.LlmProviders.EMPOWER == provider: return litellm.EmpowerChatConfig() - elif litellm.LlmProviders.MINIMAX == provider: + elif litellm.LlmProviders.MINIMAX == provider: return litellm.MinimaxChatConfig() elif litellm.LlmProviders.GITHUB == provider: return litellm.GithubChatConfig() @@ -8314,8 +8321,7 @@ class ProviderConfigManager: from litellm.llms.vertex_ai.ocr.common_utils import get_vertex_ai_ocr_config return get_vertex_ai_ocr_config(model=model) - - MistralOCRConfig = getattr(sys.modules[__name__], 'MistralOCRConfig') +MistralOCRConfig = getattr(sys.modules[__name__], 'MistralOCRConfig') PROVIDER_TO_CONFIG_MAP = { litellm.LlmProviders.MISTRAL: MistralOCRConfig, } @@ -8752,12 +8758,12 @@ def __getattr__(name: str) -> Any: """Lazy import handler for utils module with cached registry for improved performance.""" # Use cached registry from _lazy_imports instead of importing tuples every time from litellm._lazy_imports import _get_lazy_import_registry - + registry = _get_lazy_import_registry() - + # Check if name is in registry and call the cached handler function if name in registry: handler_func = registry[name] return handler_func(name) - + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index e603f94ab87..473847c04fd 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -275,10 +275,10 @@ def test_get_supported_openai_params(): def test_get_supported_openai_params_bedrock_converse(): """ - Test that all documented bedrock converse models have the same set of supported openai params when using + Test that all documented bedrock converse models have the same set of supported openai params when using `bedrock/converse/` or `bedrock/` prefix. - Note: This test is critical for routing, if we ever remove `litellm.BEDROCK_CONVERSE_MODELS`, + Note: This test is critical for routing, if we ever remove `litellm.BEDROCK_CONVERSE_MODELS`, please update this test to read `bedrock_converse` models from the model cost map. """ for model in litellm.BEDROCK_CONVERSE_MODELS: @@ -380,7 +380,7 @@ def test_transform_response_with_computer_use_tool(): @property def text(self): return json.dumps(response_json) - + config = AmazonConverseConfig() model_response = ModelResponse() optional_params = { @@ -471,7 +471,7 @@ def test_transform_response_with_bash_tool(): @property def text(self): return json.dumps(response_json) - + config = AmazonConverseConfig() model_response = ModelResponse() optional_params = { @@ -525,7 +525,7 @@ def test_transform_response_with_structured_response_being_called(): "toolUseId": "tooluse_456", "name": "json_tool_call", "input": { - "Current_Temperature": 62, + "Current_Temperature": 62, "Weather_Explanation": "San Francisco typically has mild, cool weather year-round due to its coastal location and marine influence. The city is known for its fog, moderate temperatures, and relatively stable climate with little seasonal variation."}, } } @@ -550,51 +550,51 @@ def test_transform_response_with_structured_response_being_called(): @property def text(self): return json.dumps(response_json) - + config = AmazonConverseConfig() model_response = ModelResponse() optional_params = { "json_mode": True, "tools": [ { - 'type': 'function', + 'type': 'function', 'function': { - 'name': 'get_weather', - 'description': 'Get the current weather in a given location', + 'name': 'get_weather', + 'description': 'Get the current weather in a given location', 'parameters': { - 'type': 'object', + 'type': 'object', 'properties': { 'location': { - 'type': 'string', + 'type': 'string', 'description': 'The city and state, e.g. San Francisco, CA' - }, + }, 'unit': { - 'type': 'string', + 'type': 'string', 'enum': ['celsius', 'fahrenheit'] } - }, + }, 'required': ['location'] } } - }, + }, { - 'type': 'function', + 'type': 'function', 'function': { - 'name': 'json_tool_call', + 'name': 'json_tool_call', 'parameters': { - '$schema': 'http://json-schema.org/draft-07/schema#', - 'type': 'object', - 'required': ['Weather_Explanation', 'Current_Temperature'], + '$schema': 'http://json-schema.org/draft-07/schema#', + 'type': 'object', + 'required': ['Weather_Explanation', 'Current_Temperature'], 'properties': { 'Weather_Explanation': { - 'type': ['string', 'null'], + 'type': ['string', 'null'], 'description': '1-2 sentences explaining the weather in the location' - }, + }, 'Current_Temperature': { - 'type': ['number', 'null'], + 'type': ['number', 'null'], 'description': 'Current temperature in the location' } - }, + }, 'additionalProperties': False } } @@ -629,36 +629,36 @@ def test_transform_response_with_structured_response_calling_tool(): response_json = { "metrics": { "latencyMs": 1148 - }, + }, "output": { - "message": + "message": { "content": [ { "text": "I\'ll check the current weather in San Francisco for you." - }, + }, { "toolUse": { "input": { "location": "San Francisco, CA", "unit": "celsius" - }, - "name": "get_weather", + }, + "name": "get_weather", "toolUseId": "tooluse_oKk__QrqSUmufMw3Q7vGaQ" } } - ], + ], "role": "assistant" } - }, - "stopReason": "tool_use", + }, + "stopReason": "tool_use", "usage": { - "cacheReadInputTokenCount": 0, - "cacheReadInputTokens": 0, - "cacheWriteInputTokenCount": 0, - "cacheWriteInputTokens": 0, - "inputTokens": 534, - "outputTokens": 69, + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 534, + "outputTokens": 69, "totalTokens": 603 } } @@ -669,51 +669,51 @@ def test_transform_response_with_structured_response_calling_tool(): @property def text(self): return json.dumps(response_json) - + config = AmazonConverseConfig() model_response = ModelResponse() optional_params = { "json_mode": True, "tools": [ { - 'type': 'function', + 'type': 'function', 'function': { - 'name': 'get_weather', - 'description': 'Get the current weather in a given location', + 'name': 'get_weather', + 'description': 'Get the current weather in a given location', 'parameters': { - 'type': 'object', + 'type': 'object', 'properties': { 'location': { - 'type': 'string', + 'type': 'string', 'description': 'The city and state, e.g. San Francisco, CA' - }, + }, 'unit': { - 'type': 'string', + 'type': 'string', 'enum': ['celsius', 'fahrenheit'] } - }, + }, 'required': ['location'] } } - }, + }, { - 'type': 'function', + 'type': 'function', 'function': { - 'name': 'json_tool_call', + 'name': 'json_tool_call', 'parameters': { - '$schema': 'http://json-schema.org/draft-07/schema#', - 'type': 'object', - 'required': ['Weather_Explanation', 'Current_Temperature'], + '$schema': 'http://json-schema.org/draft-07/schema#', + 'type': 'object', + 'required': ['Weather_Explanation', 'Current_Temperature'], 'properties': { 'Weather_Explanation': { - 'type': ['string', 'null'], + 'type': ['string', 'null'], 'description': '1-2 sentences explaining the weather in the location' - }, + }, 'Current_Temperature': { - 'type': ['number', 'null'], + 'type': ['number', 'null'], 'description': 'Current temperature in the location' } - }, + }, 'additionalProperties': False } } @@ -743,7 +743,7 @@ def test_transform_response_with_structured_response_calling_tool(): @pytest.mark.asyncio async def test_bedrock_bash_tool_acompletion(): """Test Bedrock with bash tool for ls command using acompletion.""" - + # Test with bash tool instead of computer tool tools = [ { @@ -751,14 +751,14 @@ async def test_bedrock_bash_tool_acompletion(): "name": "bash", } ] - + messages = [ { - "role": "user", + "role": "user", "content": "run ls command and find all python files" } ] - + try: response = await litellm.acompletion( model="bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0", @@ -771,13 +771,13 @@ async def test_bedrock_bash_tool_acompletion(): assert False, "Expected authentication error but got successful response" except Exception as e: error_str = str(e).lower() - + # Check if it's an expected authentication/credentials error auth_error_indicators = [ - "credentials", "authentication", "unauthorized", "access denied", + "credentials", "authentication", "unauthorized", "access denied", "aws", "region", "profile", "token", "invalid", "signature" ] - + if any(auth_error in error_str for auth_error in auth_error_indicators): # This is expected - request formatting succeeded, auth failed as expected assert True @@ -789,7 +789,7 @@ async def test_bedrock_bash_tool_acompletion(): @pytest.mark.asyncio async def test_bedrock_computer_use_acompletion(): """Test Bedrock computer use with acompletion function.""" - + # Test with computer use tool tools = [ { @@ -800,10 +800,10 @@ async def test_bedrock_computer_use_acompletion(): "display_number": 0, } ] - + messages = [ { - "role": "user", + "role": "user", "content": [ { "type": "text", @@ -818,7 +818,7 @@ async def test_bedrock_computer_use_acompletion(): ] } ] - + try: response = await litellm.acompletion( model="bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0", @@ -831,13 +831,13 @@ async def test_bedrock_computer_use_acompletion(): assert False, "Expected authentication error but got successful response" except Exception as e: error_str = str(e).lower() - + # Check if it's an expected authentication/credentials error auth_error_indicators = [ - "credentials", "authentication", "unauthorized", "access denied", + "credentials", "authentication", "unauthorized", "access denied", "aws", "region", "profile", "token", "invalid", "signature" ] - + if any(auth_error in error_str for auth_error in auth_error_indicators): # This is expected - request formatting succeeded, auth failed as expected assert True @@ -849,9 +849,9 @@ async def test_bedrock_computer_use_acompletion(): @pytest.mark.asyncio async def test_transformation_directly(): """Test the transformation directly to verify the request structure.""" - + config = AmazonConverseConfig() - + tools = [ { "type": "computer_20241022", @@ -865,14 +865,14 @@ async def test_transformation_directly(): "name": "bash", } ] - + messages = [ { "role": "user", "content": "run ls command and find all python files" } ] - + # Transform request request_data = config.transform_request( model="anthropic.claude-3-5-sonnet-20241022-v2:0", @@ -881,19 +881,19 @@ async def test_transformation_directly(): litellm_params={}, headers={} ) - + # Verify the structure assert "additionalModelRequestFields" in request_data additional_fields = request_data["additionalModelRequestFields"] - + # Check that anthropic_beta is set correctly for computer use assert "anthropic_beta" in additional_fields assert additional_fields["anthropic_beta"] == ["computer-use-2024-10-22"] - + # Check that tools are present assert "tools" in additional_fields assert len(additional_fields["tools"]) == 2 - + # Verify tool types tool_types = [tool.get("type") for tool in additional_fields["tools"]] assert "computer_20241022" in tool_types @@ -933,7 +933,7 @@ def test_transform_request_helper_includes_anthropic_beta_and_tools_bash(): def test_transform_request_with_multiple_tools(): """Test transformation with multiple tools including computer, bash, and function tools.""" config = AmazonConverseConfig() - + # Use the exact payload from the user's error tools = [ { @@ -974,14 +974,14 @@ def test_transform_request_with_multiple_tools(): } } ] - + messages = [ { "role": "user", "content": "run ls command and find all python files" } ] - + # Transform request request_data = config.transform_request( model="anthropic.claude-3-5-sonnet-20241022-v2:0", @@ -990,25 +990,25 @@ def test_transform_request_with_multiple_tools(): litellm_params={}, headers={} ) - + # Verify the structure assert "additionalModelRequestFields" in request_data additional_fields = request_data["additionalModelRequestFields"] - + # Check that anthropic_beta is set correctly for computer use assert "anthropic_beta" in additional_fields assert additional_fields["anthropic_beta"] == ["computer-use-2024-10-22"] - + # Check that tools are present assert "tools" in additional_fields assert len(additional_fields["tools"]) == 3 # computer, bash, text_editor tools - + # Verify tool types tool_types = [tool.get("type") for tool in additional_fields["tools"]] assert "computer_20241022" in tool_types assert "bash_20241022" in tool_types assert "text_editor_20241022" in tool_types - + # Function tools are processed separately and not included in computer use tools # They would be in toolConfig if present @@ -1016,7 +1016,7 @@ def test_transform_request_with_multiple_tools(): def test_transform_request_with_computer_tool_only(): """Test transformation with only computer tool.""" config = AmazonConverseConfig() - + tools = [ { "type": "computer_20241022", @@ -1026,10 +1026,10 @@ def test_transform_request_with_computer_tool_only(): "display_number": 0, } ] - + messages = [ { - "role": "user", + "role": "user", "content": [ { "type": "text", @@ -1044,7 +1044,7 @@ def test_transform_request_with_computer_tool_only(): ] } ] - + # Transform request request_data = config.transform_request( model="anthropic.claude-3-5-sonnet-20241022-v2:0", @@ -1053,15 +1053,15 @@ def test_transform_request_with_computer_tool_only(): litellm_params={}, headers={} ) - + # Verify the structure assert "additionalModelRequestFields" in request_data additional_fields = request_data["additionalModelRequestFields"] - + # Check that anthropic_beta is set correctly for computer use assert "anthropic_beta" in additional_fields assert additional_fields["anthropic_beta"] == ["computer-use-2024-10-22"] - + # Check that tools are present assert "tools" in additional_fields assert len(additional_fields["tools"]) == 1 @@ -1071,21 +1071,21 @@ def test_transform_request_with_computer_tool_only(): def test_transform_request_with_bash_tool_only(): """Test transformation with only bash tool.""" config = AmazonConverseConfig() - + tools = [ { "type": "bash_20241022", "name": "bash", } ] - + messages = [ { - "role": "user", + "role": "user", "content": "run ls command and find all python files" } ] - + # Transform request request_data = config.transform_request( model="anthropic.claude-3-5-sonnet-20241022-v2:0", @@ -1094,15 +1094,15 @@ def test_transform_request_with_bash_tool_only(): litellm_params={}, headers={} ) - + # Verify the structure assert "additionalModelRequestFields" in request_data additional_fields = request_data["additionalModelRequestFields"] - + # Check that anthropic_beta is set correctly for computer use assert "anthropic_beta" in additional_fields assert additional_fields["anthropic_beta"] == ["computer-use-2024-10-22"] - + # Check that tools are present assert "tools" in additional_fields assert len(additional_fields["tools"]) == 1 @@ -1112,21 +1112,21 @@ def test_transform_request_with_bash_tool_only(): def test_transform_request_with_text_editor_tool(): """Test transformation with text editor tool.""" config = AmazonConverseConfig() - + tools = [ { "type": "text_editor_20241022", "name": "str_replace_editor", } ] - + messages = [ { "role": "user", "content": "Edit this text file" } ] - + # Transform request request_data = config.transform_request( model="anthropic.claude-3-5-sonnet-20241022-v2:0", @@ -1135,15 +1135,15 @@ def test_transform_request_with_text_editor_tool(): litellm_params={}, headers={} ) - + # Verify the structure assert "additionalModelRequestFields" in request_data additional_fields = request_data["additionalModelRequestFields"] - + # Check that anthropic_beta is set correctly for computer use assert "anthropic_beta" in additional_fields assert additional_fields["anthropic_beta"] == ["computer-use-2024-10-22"] - + # Check that tools are present assert "tools" in additional_fields assert len(additional_fields["tools"]) == 1 @@ -1153,7 +1153,7 @@ def test_transform_request_with_text_editor_tool(): def test_transform_request_with_function_tool(): """Test transformation with function tool.""" config = AmazonConverseConfig() - + tools = [ { "type": "function", @@ -1174,14 +1174,14 @@ def test_transform_request_with_function_tool(): } } ] - + messages = [ { "role": "user", "content": "What's the weather like in San Francisco?" } ] - + # Transform request request_data = config.transform_request( model="anthropic.claude-3-5-sonnet-20241022-v2:0", @@ -1190,11 +1190,11 @@ def test_transform_request_with_function_tool(): litellm_params={}, headers={} ) - + # Verify the structure assert "additionalModelRequestFields" in request_data additional_fields = request_data["additionalModelRequestFields"] - + # Function tools are not computer use tools, so they don't get anthropic_beta # They are processed through the regular tool config assert "toolConfig" in request_data @@ -1206,7 +1206,7 @@ def test_transform_request_with_function_tool(): def test_map_openai_params_with_response_format(): """Test map_openai_params with response_format.""" config = AmazonConverseConfig() - + tools = [ { "type": "function", @@ -1277,12 +1277,12 @@ async def test_assistant_message_cache_control(): messages = [ {"role": "user", "content": "Hello"}, { - "role": "assistant", + "role": "assistant", "content": "Hi there!", "cache_control": {"type": "ephemeral"} } ] - + result = _bedrock_converse_messages_pt( messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", @@ -1294,7 +1294,7 @@ async def test_assistant_message_cache_control(): model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) - + assert result == async_result async_result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( @@ -1302,14 +1302,14 @@ async def test_assistant_message_cache_control(): model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) - + assert result == async_result - + # Should have user message and assistant message assert len(result) == 2 assert result[0]["role"] == "user" assert result[1]["role"] == "assistant" - + # Assistant message should have text content and cachePoint assistant_content = result[1]["content"] assert len(assistant_content) == 2 @@ -1325,7 +1325,7 @@ async def test_assistant_message_list_content_cache_control(): BedrockConverseMessagesProcessor, _bedrock_converse_messages_pt, ) - + messages = [ {"role": "user", "content": "Hello"}, { @@ -1339,7 +1339,7 @@ async def test_assistant_message_list_content_cache_control(): ] } ] - + result = _bedrock_converse_messages_pt( messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", @@ -1351,9 +1351,9 @@ async def test_assistant_message_list_content_cache_control(): model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) - + assert result == async_result - + # Assistant message should have text content and cachePoint assistant_content = result[1]["content"] assert len(assistant_content) == 2 @@ -1369,7 +1369,7 @@ async def test_tool_message_cache_control(): BedrockConverseMessagesProcessor, _bedrock_converse_messages_pt, ) - + messages = [ {"role": "user", "content": "What's the weather?"}, { @@ -1395,7 +1395,7 @@ async def test_tool_message_cache_control(): ] } ] - + result = _bedrock_converse_messages_pt( messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", @@ -1407,20 +1407,20 @@ async def test_tool_message_cache_control(): model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) - + assert result == async_result - + # Should have user, assistant, and user (tool results) messages assert len(result) == 3 - + # Last message should contain tool result and cachePoint tool_message_content = result[2]["content"] assert len(tool_message_content) == 2 - + # First should be tool result assert "toolResult" in tool_message_content[0] assert tool_message_content[0]["toolResult"]["content"][0]["text"] == "Weather data: sunny, 25°C" - + # Second should be cachePoint assert "cachePoint" in tool_message_content[1] assert tool_message_content[1]["cachePoint"]["type"] == "default" @@ -1433,7 +1433,7 @@ async def test_tool_message_string_content_cache_control(): BedrockConverseMessagesProcessor, _bedrock_converse_messages_pt, ) - + messages = [ {"role": "user", "content": "What's the weather?"}, { @@ -1442,7 +1442,7 @@ async def test_tool_message_string_content_cache_control(): "tool_calls": [ { "id": "call_123", - "type": "function", + "type": "function", "function": {"name": "get_weather", "arguments": "{}"} } ] @@ -1454,7 +1454,7 @@ async def test_tool_message_string_content_cache_control(): "cache_control": {"type": "ephemeral"} } ] - + result = _bedrock_converse_messages_pt( messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", @@ -1466,17 +1466,17 @@ async def test_tool_message_string_content_cache_control(): model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) - + assert result == async_result - + # Last message should contain tool result and cachePoint tool_message_content = result[2]["content"] assert len(tool_message_content) == 2 - + # First should be tool result assert "toolResult" in tool_message_content[0] assert tool_message_content[0]["toolResult"]["content"][0]["text"] == "Weather: sunny, 25°C" - + # Second should be cachePoint assert "cachePoint" in tool_message_content[1] assert tool_message_content[1]["cachePoint"]["type"] == "default" @@ -1489,7 +1489,7 @@ async def test_assistant_tool_calls_cache_control(): BedrockConverseMessagesProcessor, _bedrock_converse_messages_pt, ) - + messages = [ {"role": "user", "content": "Calculate 2+2"}, { @@ -1505,7 +1505,7 @@ async def test_assistant_tool_calls_cache_control(): ] } ] - + result = _bedrock_converse_messages_pt( messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", @@ -1517,18 +1517,18 @@ async def test_assistant_tool_calls_cache_control(): model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) - + assert result == async_result - + # Assistant message should have tool use and cachePoint assistant_content = result[1]["content"] assert len(assistant_content) == 2 - + # First should be tool use assert "toolUse" in assistant_content[0] assert assistant_content[0]["toolUse"]["name"] == "calc" assert assistant_content[0]["toolUse"]["toolUseId"] == "call_proxy_123" - + # Second should be cachePoint assert "cachePoint" in assistant_content[1] assert assistant_content[1]["cachePoint"]["type"] == "default" @@ -1541,7 +1541,7 @@ async def test_multiple_tool_calls_with_mixed_cache_control(): BedrockConverseMessagesProcessor, _bedrock_converse_messages_pt, ) - + messages = [ {"role": "user", "content": "Do multiple calculations"}, { @@ -1563,7 +1563,7 @@ async def test_multiple_tool_calls_with_mixed_cache_control(): ] } ] - + result = _bedrock_converse_messages_pt( messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", @@ -1575,21 +1575,21 @@ async def test_multiple_tool_calls_with_mixed_cache_control(): model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) - + assert result == async_result - + # Assistant message should have: toolUse1, cachePoint, toolUse2 assistant_content = result[1]["content"] assert len(assistant_content) == 3 - + # First tool use with cache assert "toolUse" in assistant_content[0] assert assistant_content[0]["toolUse"]["toolUseId"] == "call_1" - + # Cache point for first tool assert "cachePoint" in assistant_content[1] assert assistant_content[1]["cachePoint"]["type"] == "default" - + # Second tool use without cache assert "toolUse" in assistant_content[2] assert assistant_content[2]["toolUse"]["toolUseId"] == "call_2" @@ -1602,7 +1602,7 @@ async def test_no_cache_control_no_cache_point(): BedrockConverseMessagesProcessor, _bedrock_converse_messages_pt, ) - + messages = [ {"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi there!"}, # No cache_control @@ -1612,7 +1612,7 @@ async def test_no_cache_control_no_cache_point(): "content": "Tool result" # No cache_control } ] - + result = _bedrock_converse_messages_pt( messages=messages, model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", @@ -1624,14 +1624,14 @@ async def test_no_cache_control_no_cache_point(): model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) - + assert result == async_result - + # Assistant message should only have text content, no cachePoint assistant_content = result[1]["content"] assert len(assistant_content) == 1 assert assistant_content[0]["text"] == "Hi there!" - + # Tool message should only have tool result, no cachePoint tool_content = result[2]["content"] assert len(tool_content) == 1 @@ -1867,11 +1867,11 @@ def test_guarded_text_with_tool_calls(): # First should be regular text assert "text" in content[0] assert content[0]["text"] == "What's the weather?" - + # Second should be guardContent assert "guardContent" in content[1] assert content[1]["guardContent"]["text"]["text"] == "Please be careful with sensitive information" - + # Other messages should not have guardContent for i in range(1, 3): content = result[i]["content"] @@ -2115,7 +2115,7 @@ def test_auto_convert_in_full_transformation(): # Verify the transformation worked assert "messages" in result assert len(result["messages"]) == 1 - + # The message should have guardContent message = result["messages"][0] assert "content" in message @@ -2626,79 +2626,79 @@ def test_empty_assistant_message_handling(): {"role": "assistant", "content": ""}, # Empty content {"role": "user", "content": "How are you?"} ] - + # Enable modify_params to prevent consecutive user message merging original_modify_params = litellm.modify_params litellm.modify_params = True - + try: result = _bedrock_converse_messages_pt( messages=messages, model="anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) - + # Should have 3 messages: user, assistant (with placeholder), user assert len(result) == 3 assert result[0]["role"] == "user" assert result[1]["role"] == "assistant" assert result[2]["role"] == "user" - + # Assistant message should have placeholder text instead of empty content assert len(result[1]["content"]) == 1 assert result[1]["content"][0]["text"] == "Please continue." - + # Test case 2: Whitespace-only content messages = [ {"role": "user", "content": "Hello"}, {"role": "assistant", "content": " "}, # Whitespace-only content {"role": "user", "content": "How are you?"} ] - + result = _bedrock_converse_messages_pt( messages=messages, model="anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) - + # Assistant message should have placeholder text instead of whitespace assert len(result[1]["content"]) == 1 assert result[1]["content"][0]["text"] == "Please continue." - + # Test case 3: Empty list content messages = [ {"role": "user", "content": "Hello"}, {"role": "assistant", "content": [{"type": "text", "text": ""}]}, # Empty text in list {"role": "user", "content": "How are you?"} ] - + result = _bedrock_converse_messages_pt( messages=messages, model="anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) - + # Assistant message should have placeholder text instead of empty text assert len(result[1]["content"]) == 1 assert result[1]["content"][0]["text"] == "Please continue." - + # Test case 4: Normal content should not be affected messages = [ {"role": "user", "content": "Hello"}, {"role": "assistant", "content": "I'm doing well, thank you!"}, # Normal content {"role": "user", "content": "How are you?"} ] - + result = _bedrock_converse_messages_pt( messages=messages, model="anthropic.claude-3-5-sonnet-20240620-v1:0", llm_provider="bedrock_converse" ) - + # Assistant message should keep original content assert len(result[1]["content"]) == 1 assert result[1]["content"][0]["text"] == "I'm doing well, thank you!" - + finally: # Restore original modify_params setting litellm.modify_params = original_modify_params @@ -2707,31 +2707,161 @@ def test_empty_assistant_message_handling(): def test_is_nova_lite_2_model(): """Test the _is_nova_lite_2_model() method for detecting Nova 2 models.""" config = AmazonConverseConfig() - + # Test with amazon.nova-2-lite-v1:0 assert config._is_nova_lite_2_model("amazon.nova-2-lite-v1:0") is True - + # Test with regional variants assert config._is_nova_lite_2_model("us.amazon.nova-2-lite-v1:0") is True assert config._is_nova_lite_2_model("eu.amazon.nova-2-lite-v1:0") is True assert config._is_nova_lite_2_model("apac.amazon.nova-2-lite-v1:0") is True - + # Test with other Nova 2 variants (pro, micro) assert config._is_nova_lite_2_model("amazon.nova-pro-1-5-v1:0") is False assert config._is_nova_lite_2_model("amazon.nova-micro-1-5-v1:0") is False assert config._is_nova_lite_2_model("us.amazon.nova-pro-1-5-v1:0") is False assert config._is_nova_lite_2_model("eu.amazon.nova-micro-1-5-v1:0") is False - + # Test with non-Nova-1.5 lite models (should return False) assert config._is_nova_lite_2_model("amazon.nova-lite-v1:0") is False assert config._is_nova_lite_2_model("amazon.nova-pro-v1:0") is False assert config._is_nova_lite_2_model("amazon.nova-micro-v1:0") is False - + # Test with Nova v1:0 models (should return False) assert config._is_nova_lite_2_model("us.amazon.nova-lite-v1:0") is False assert config._is_nova_lite_2_model("eu.amazon.nova-pro-v1:0") is False - + # Test with completely different models (should return False) assert config._is_nova_lite_2_model("anthropic.claude-3-5-sonnet-20240620-v1:0") is False assert config._is_nova_lite_2_model("meta.llama3-70b-instruct-v1:0") is False assert config._is_nova_lite_2_model("mistral.mistral-7b-instruct-v0:2") is False + +def test_drop_thinking_param_when_thinking_blocks_missing(): + """ + Test that thinking param is dropped when modify_params=True and + thinking_blocks are missing from assistant message with tool_calls. + + This prevents the Anthropic/Bedrock error: + "Expected thinking or redacted_thinking, but found tool_use" + + Related issue: https://github.com/BerriAI/litellm/issues/14194 + """ + from litellm.utils import last_assistant_with_tool_calls_has_no_thinking_blocks + + # Save original modify_params setting + original_modify_params = litellm.modify_params + + try: + # Test case 1: thinking should be dropped when modify_params=True + # and assistant message has tool_calls but no thinking_blocks + litellm.modify_params = True + + messages_without_thinking_blocks = [ + {"role": "user", "content": "Search for weather"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": {"name": "search", "arguments": "{}"}, + } + ], + # No thinking_blocks - simulates OpenAI-compatible client + }, + {"role": "tool", "content": "Weather is sunny", "tool_call_id": "call_123"}, + ] + + optional_params = {"thinking": {"type": "enabled", "budget_tokens": 1000}} + + # Verify the condition is detected + assert last_assistant_with_tool_calls_has_no_thinking_blocks( + messages_without_thinking_blocks + ), "Should detect missing thinking_blocks" + + # Simulate what _transform_request_helper does + if ( + optional_params.get("thinking") is not None + and messages_without_thinking_blocks is not None + and last_assistant_with_tool_calls_has_no_thinking_blocks( + messages_without_thinking_blocks + ) + ): + if litellm.modify_params: + optional_params.pop("thinking", None) + + assert "thinking" not in optional_params, ( + "thinking param should be dropped when modify_params=True " + "and thinking_blocks are missing" + ) + + # Test case 2: thinking should NOT be dropped when thinking_blocks are present + messages_with_thinking_blocks = [ + {"role": "user", "content": "Search for weather"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": {"name": "search", "arguments": "{}"}, + } + ], + "thinking_blocks": [ + {"type": "thinking", "thinking": "Let me search for weather..."} + ], + }, + {"role": "tool", "content": "Weather is sunny", "tool_call_id": "call_123"}, + ] + + optional_params_with_thinking = { + "thinking": {"type": "enabled", "budget_tokens": 1000} + } + + # Verify the condition is NOT detected when thinking_blocks are present + assert not last_assistant_with_tool_calls_has_no_thinking_blocks( + messages_with_thinking_blocks + ), "Should NOT detect missing thinking_blocks when they are present" + + # Simulate what _transform_request_helper does + if ( + optional_params_with_thinking.get("thinking") is not None + and messages_with_thinking_blocks is not None + and last_assistant_with_tool_calls_has_no_thinking_blocks( + messages_with_thinking_blocks + ) + ): + if litellm.modify_params: + optional_params_with_thinking.pop("thinking", None) + + assert "thinking" in optional_params_with_thinking, ( + "thinking param should NOT be dropped when thinking_blocks are present" + ) + + # Test case 3: thinking should NOT be dropped when modify_params=False + litellm.modify_params = False + + optional_params_no_modify = { + "thinking": {"type": "enabled", "budget_tokens": 1000} + } + + # Simulate what _transform_request_helper does + if ( + optional_params_no_modify.get("thinking") is not None + and messages_without_thinking_blocks is not None + and last_assistant_with_tool_calls_has_no_thinking_blocks( + messages_without_thinking_blocks + ) + ): + if litellm.modify_params: + optional_params_no_modify.pop("thinking", None) + + assert "thinking" in optional_params_no_modify, ( + "thinking param should NOT be dropped when modify_params=False" + ) + + finally: + # Restore original modify_params setting + litellm.modify_params = original_modify_params From d3fdac84682e4d972d3c170925a5225224886135 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Benedikt=20=C3=93skarsson?= Date: Thu, 8 Jan 2026 02:04:52 +0000 Subject: [PATCH 02/58] chore: fix formatting error --- litellm/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 5011b45c0f7..e088b35037a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8321,7 +8321,8 @@ class ProviderConfigManager: from litellm.llms.vertex_ai.ocr.common_utils import get_vertex_ai_ocr_config return get_vertex_ai_ocr_config(model=model) -MistralOCRConfig = getattr(sys.modules[__name__], 'MistralOCRConfig') + + MistralOCRConfig = getattr(sys.modules[__name__], 'MistralOCRConfig') PROVIDER_TO_CONFIG_MAP = { litellm.LlmProviders.MISTRAL: MistralOCRConfig, } From 5e1a5a8949a31bc3b39f4d689c76e9a423a80740 Mon Sep 17 00:00:00 2001 From: stiyyagura0901 Date: Sat, 17 Jan 2026 10:15:28 -0500 Subject: [PATCH 03/58] fix: UI dashboard respects custom authentication header override Fix UI dashboard to use custom authentication header name configured in general settings instead of hardcoded 'Authorization' header. - Replace hardcoded Authorization headers with getGlobalLitellmHeaderName() in all hook files - Replace hardcoded Authorization headers with globalLitellmHeaderName in networking.tsx - Update test files to properly mock getGlobalLitellmHeaderName - Fixes issue where UI becomes unusable when custom auth header is configured --- .../hooks/cloudzero/useCloudZeroCreate.ts | 4 +-- .../hooks/cloudzero/useCloudZeroDryRun.ts | 4 +-- .../hooks/cloudzero/useCloudZeroExport.ts | 4 +-- .../hooks/cloudzero/useCloudZeroSettings.ts | 8 +++--- .../hooks/router/useRouterFields.test.ts | 1 + .../hooks/router/useRouterFields.ts | 4 +-- .../pricing_calculator/use_cost_estimate.ts | 4 +-- .../use_multi_cost_estimate.ts | 4 +-- .../use_discount_config.ts | 6 ++-- .../CostTrackingSettings/use_margin_config.ts | 6 ++-- .../src/components/cloudzero_export_modal.tsx | 7 +++-- .../src/components/cost_tracking_settings.tsx | 6 ++-- .../guardrails/edit_guardrail_form.tsx | 4 +-- .../src/components/networking.tsx | 28 +++++++++---------- .../chat_ui/CodeInterpreterOutput.test.tsx | 1 + .../chat_ui/CodeInterpreterOutput.tsx | 6 ++-- .../playground/llm_calls/a2a_send_message.tsx | 6 ++-- .../llm_calls/embeddings_api.test.tsx | 1 + .../playground/llm_calls/embeddings_api.tsx | 4 +-- .../playground/llm_calls/fetch_agents.tsx | 4 +-- .../conversation_panel/useConversation.ts | 4 +-- .../src/components/ui_theme_settings.tsx | 8 +++--- 22 files changed, 64 insertions(+), 60 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroCreate.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroCreate.ts index e1263903622..29c5ae3a0f0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroCreate.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroCreate.ts @@ -1,4 +1,4 @@ -import { getProxyBaseUrl } from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; import { useMutation } from "@tanstack/react-query"; interface CreateParams { @@ -18,7 +18,7 @@ const performCloudZeroCreate = async (accessToken: string, params: CreateParams) const response = await fetch(url, { method: "POST", headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, body: JSON.stringify({ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroDryRun.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroDryRun.ts index 1ed8a141603..e00fd2a6375 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroDryRun.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroDryRun.ts @@ -1,4 +1,4 @@ -import { getProxyBaseUrl } from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; import { useMutation } from "@tanstack/react-query"; interface DryRunParams { @@ -16,7 +16,7 @@ const performCloudZeroDryRun = async (accessToken: string, params: DryRunParams const response = await fetch(url, { method: "POST", headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, body: JSON.stringify({ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroExport.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroExport.ts index 47d559b20d2..ba9a013ecc6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroExport.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroExport.ts @@ -1,4 +1,4 @@ -import { getProxyBaseUrl } from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; import { useMutation } from "@tanstack/react-query"; interface ExportParams { @@ -16,7 +16,7 @@ const performCloudZeroExport = async (accessToken: string, params: ExportParams const response = await fetch(url, { method: "POST", headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, body: JSON.stringify({ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroSettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroSettings.ts index 96f5ab2f944..7fd61077385 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroSettings.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/cloudzero/useCloudZeroSettings.ts @@ -1,5 +1,5 @@ import { CloudZeroSettings } from "@/components/CloudZeroCostTracking/types"; -import { getProxyBaseUrl } from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; @@ -12,7 +12,7 @@ const getCloudZeroSettings = async (accessToken: string): Promise ({ proxyBaseUrl: null, + getGlobalLitellmHeaderName: vi.fn(() => "Authorization"), })); // Mock useAuthorized hook diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/router/useRouterFields.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/router/useRouterFields.ts index 589508c5ddb..5fa22b41096 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/router/useRouterFields.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/router/useRouterFields.ts @@ -1,7 +1,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useQuery, UseQueryResult } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; -import { proxyBaseUrl } from "@/components/networking"; +import { proxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; export interface RouterSettingsField { field_name: string; @@ -39,7 +39,7 @@ const getRouterFields = async (accessToken: string): Promise = ({ isOpen, onC const response = await fetch("/cloudzero/settings", { method: "GET", headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, }); @@ -88,7 +89,7 @@ const CloudZeroExportModal: React.FC = ({ isOpen, onC const response = await fetch(endpoint, { method, headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, body: JSON.stringify(payload), @@ -128,7 +129,7 @@ const CloudZeroExportModal: React.FC = ({ isOpen, onC const response = await fetch("/cloudzero/export", { method: "POST", headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, body: JSON.stringify({ diff --git a/ui/litellm-dashboard/src/components/cost_tracking_settings.tsx b/ui/litellm-dashboard/src/components/cost_tracking_settings.tsx index bab61da6775..75d650306ee 100644 --- a/ui/litellm-dashboard/src/components/cost_tracking_settings.tsx +++ b/ui/litellm-dashboard/src/components/cost_tracking_settings.tsx @@ -15,7 +15,7 @@ import { Col, Subtitle, } from "@tremor/react"; -import { getProxyBaseUrl } from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; import NotificationsManager from "./molecules/notifications_manager"; interface CostTrackingSettingsProps { @@ -56,7 +56,7 @@ const CostTrackingSettings: React.FC = ({ const response = await fetch(url, { method: "GET", headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, }); @@ -86,7 +86,7 @@ const CostTrackingSettings: React.FC = ({ const response = await fetch(url, { method: "PATCH", headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, body: JSON.stringify(discountConfig), diff --git a/ui/litellm-dashboard/src/components/guardrails/edit_guardrail_form.tsx b/ui/litellm-dashboard/src/components/guardrails/edit_guardrail_form.tsx index caa197e0cd5..a2cc3ad41dd 100644 --- a/ui/litellm-dashboard/src/components/guardrails/edit_guardrail_form.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/edit_guardrail_form.tsx @@ -2,7 +2,7 @@ import React, { useState, useEffect } from "react"; import { Form, Typography, Select, Input, Switch, Modal } from "antd"; import { Button, TextInput } from "@tremor/react"; import { guardrail_provider_map, guardrailLogoMap, getGuardrailProviders } from "./guardrail_info_helpers"; -import { getGuardrailUISettings } from "../networking"; +import { getGuardrailUISettings, getGlobalLitellmHeaderName } from "../networking"; import PiiConfiguration from "./pii_configuration"; import NotificationsManager from "../molecules/notifications_manager"; @@ -183,7 +183,7 @@ const EditGuardrailForm: React.FC = ({ const response = await fetch(url, { method: "PUT", headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, body: JSON.stringify(guardrailData), diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 82894dfb0e2..cc434adee1a 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -6226,7 +6226,7 @@ export const tagCreateCall = async (accessToken: string, formValues: TagNewReque method: "POST", headers: { "Content-Type": "application/json", - Authorization: `Bearer ${accessToken}`, + [globalLitellmHeaderName]: `Bearer ${accessToken}`, }, body: JSON.stringify(formValues), }); @@ -6252,7 +6252,7 @@ export const tagUpdateCall = async (accessToken: string, formValues: TagUpdateRe method: "POST", headers: { "Content-Type": "application/json", - Authorization: `Bearer ${accessToken}`, + [globalLitellmHeaderName]: `Bearer ${accessToken}`, }, body: JSON.stringify(formValues), }); @@ -6278,7 +6278,7 @@ export const tagInfoCall = async (accessToken: string, tagNames: string[]): Prom method: "POST", headers: { "Content-Type": "application/json", - Authorization: `Bearer ${accessToken}`, + [globalLitellmHeaderName]: `Bearer ${accessToken}`, }, body: JSON.stringify({ names: tagNames }), }); @@ -6304,7 +6304,7 @@ export const tagListCall = async (accessToken: string): Promise const response = await fetch(url, { method: "GET", headers: { - Authorization: `Bearer ${accessToken}`, + [globalLitellmHeaderName]: `Bearer ${accessToken}`, }, }); @@ -6330,7 +6330,7 @@ export const tagDeleteCall = async (accessToken: string, tagName: string): Promi method: "POST", headers: { "Content-Type": "application/json", - Authorization: `Bearer ${accessToken}`, + [globalLitellmHeaderName]: `Bearer ${accessToken}`, }, body: JSON.stringify({ name: tagName }), }); @@ -6422,7 +6422,7 @@ export const getTeamPermissionsCall = async (accessToken: string, teamId: string method: "GET", headers: { "Content-Type": "application/json", - Authorization: `Bearer ${accessToken}`, + [globalLitellmHeaderName]: `Bearer ${accessToken}`, }, }); @@ -6450,7 +6450,7 @@ export const teamPermissionsUpdateCall = async (accessToken: string, teamId: str method: "POST", headers: { "Content-Type": "application/json", - Authorization: `Bearer ${accessToken}`, + [globalLitellmHeaderName]: `Bearer ${accessToken}`, }, body: JSON.stringify({ team_id: teamId, @@ -6514,7 +6514,7 @@ export const vectorStoreCreateCall = async (accessToken: string, formValues: Rec method: "POST", headers: { "Content-Type": "application/json", - Authorization: `Bearer ${accessToken}`, + [globalLitellmHeaderName]: `Bearer ${accessToken}`, }, body: JSON.stringify(formValues), }); @@ -6543,7 +6543,7 @@ export const vectorStoreListCall = async ( method: "GET", headers: { "Content-Type": "application/json", - Authorization: `Bearer ${accessToken}`, + [globalLitellmHeaderName]: `Bearer ${accessToken}`, }, }); @@ -6567,7 +6567,7 @@ export const vectorStoreDeleteCall = async (accessToken: string, vectorStoreId: method: "POST", headers: { "Content-Type": "application/json", - Authorization: `Bearer ${accessToken}`, + [globalLitellmHeaderName]: `Bearer ${accessToken}`, }, body: JSON.stringify({ vector_store_id: vectorStoreId }), }); @@ -6592,7 +6592,7 @@ export const vectorStoreInfoCall = async (accessToken: string, vectorStoreId: st method: "POST", headers: { "Content-Type": "application/json", - Authorization: `Bearer ${accessToken}`, + [globalLitellmHeaderName]: `Bearer ${accessToken}`, }, body: JSON.stringify({ vector_store_id: vectorStoreId }), }); @@ -6617,7 +6617,7 @@ export const vectorStoreUpdateCall = async (accessToken: string, formValues: Rec method: "POST", headers: { "Content-Type": "application/json", - Authorization: `Bearer ${accessToken}`, + [globalLitellmHeaderName]: `Bearer ${accessToken}`, }, body: JSON.stringify(formValues), }); @@ -7738,7 +7738,7 @@ export const vectorStoreSearchCall = async ( const response = await fetch(url, { method: "POST", headers: { - Authorization: `Bearer ${accessToken}`, + [globalLitellmHeaderName]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, body: JSON.stringify({ @@ -7771,7 +7771,7 @@ export const searchToolQueryCall = async ( const response = await fetch(url, { method: "POST", headers: { - Authorization: `Bearer ${accessToken}`, + [globalLitellmHeaderName]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, body: JSON.stringify({ diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterOutput.test.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterOutput.test.tsx index d87a74f4641..a42bd42ed81 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterOutput.test.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterOutput.test.tsx @@ -5,6 +5,7 @@ import CodeInterpreterOutput from "./CodeInterpreterOutput"; vi.mock("@/components/networking", () => ({ getProxyBaseUrl: vi.fn(() => "https://example.com"), + getGlobalLitellmHeaderName: vi.fn(() => "Authorization"), })); global.fetch = vi.fn(); diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterOutput.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterOutput.tsx index 5625d0cbf4b..b72e55feaea 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterOutput.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterOutput.tsx @@ -9,7 +9,7 @@ import { } from "@ant-design/icons"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import { coy } from "react-syntax-highlighter/dist/esm/styles/prism"; -import { getProxyBaseUrl } from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; interface ContainerFileCitation { type: "container_file_citation"; @@ -55,7 +55,7 @@ const CodeInterpreterOutput: React.FC = ({ `${proxyBaseUrl}/v1/containers/${annotation.container_id}/files/${annotation.file_id}/content`, { headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, }, } ); @@ -90,7 +90,7 @@ const CodeInterpreterOutput: React.FC = ({ `${proxyBaseUrl}/v1/containers/${annotation.container_id}/files/${annotation.file_id}/content`, { headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, }, } ); diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/a2a_send_message.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/a2a_send_message.tsx index 0654db32920..b5cfc820f36 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/a2a_send_message.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/a2a_send_message.tsx @@ -2,7 +2,7 @@ // A2A Protocol (JSON-RPC 2.0) implementation for sending messages to agents import { v4 as uuidv4 } from "uuid"; -import { getProxyBaseUrl } from "../../networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "../../networking"; import { A2ATaskMetadata } from "../chat_ui/types"; interface A2AMessagePart { @@ -143,7 +143,7 @@ export const makeA2ASendMessageRequest = async ( const response = await fetch(url, { method: "POST", headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, body: JSON.stringify(jsonRpcRequest), @@ -276,7 +276,7 @@ export const makeA2AStreamMessageRequest = async ( const response = await fetch(url, { method: "POST", headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, body: JSON.stringify(jsonRpcRequest), diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/embeddings_api.test.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/embeddings_api.test.tsx index f24a50074c6..19a79f3f9be 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/embeddings_api.test.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/embeddings_api.test.tsx @@ -3,6 +3,7 @@ import { makeOpenAIEmbeddingsRequest } from "./embeddings_api"; vi.mock("@/components/networking", () => ({ getProxyBaseUrl: vi.fn(() => "https://example.com"), + getGlobalLitellmHeaderName: vi.fn(() => "Authorization"), })); describe("embeddings_api", () => { diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/embeddings_api.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/embeddings_api.tsx index 84192a1d866..cc3c0f81dc5 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/embeddings_api.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/embeddings_api.tsx @@ -1,5 +1,5 @@ import NotificationManager from "@/components/molecules/notifications_manager"; -import { getProxyBaseUrl } from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; export async function makeOpenAIEmbeddingsRequest( input: string, @@ -34,7 +34,7 @@ export async function makeOpenAIEmbeddingsRequest( method: "POST", headers: { "Content-Type": "application/json", - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, ...headers, }, body: JSON.stringify({ diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/fetch_agents.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/fetch_agents.tsx index 889b012c5f7..f0f283c5252 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/fetch_agents.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/fetch_agents.tsx @@ -1,6 +1,6 @@ // fetch_agents.tsx -import { getProxyBaseUrl } from "../../networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "../../networking"; export interface Agent { agent_id: string; @@ -27,7 +27,7 @@ export const fetchAvailableAgents = async ( const response = await fetch(url, { method: "GET", headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, }); diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/useConversation.ts b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/useConversation.ts index 4372622232e..7b595c777c1 100644 --- a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/useConversation.ts +++ b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/useConversation.ts @@ -3,7 +3,7 @@ import NotificationsManager from "../../../molecules/notifications_manager"; import { TokenUsage } from "../../../playground/chat_ui/ResponseMetrics"; import { Message } from "./types"; import { convertToDotPrompt, extractVariables } from "../utils"; -import { getProxyBaseUrl } from "../../../networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "../../../networking"; export const useConversation = (prompt: any, accessToken: string | null) => { const [isLoading, setIsLoading] = useState(false); @@ -91,7 +91,7 @@ export const useConversation = (prompt: any, accessToken: string | null) => { const response = await fetch(`${proxyBaseUrl}/prompts/test`, { method: "POST", headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, body: JSON.stringify(requestBody), diff --git a/ui/litellm-dashboard/src/components/ui_theme_settings.tsx b/ui/litellm-dashboard/src/components/ui_theme_settings.tsx index 83c5cbb5c15..636163d1d70 100644 --- a/ui/litellm-dashboard/src/components/ui_theme_settings.tsx +++ b/ui/litellm-dashboard/src/components/ui_theme_settings.tsx @@ -1,7 +1,7 @@ import React, { useState, useEffect } from "react"; import { Card, Title, Text, TextInput, Button } from "@tremor/react"; import { useTheme } from "@/contexts/ThemeContext"; -import { getProxyBaseUrl } from "@/components/networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; import NotificationsManager from "./molecules/notifications_manager"; interface UIThemeSettingsProps { @@ -29,7 +29,7 @@ const UIThemeSettings: React.FC = ({ userID, userRole, acc const response = await fetch(url, { method: "GET", headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, }); @@ -53,7 +53,7 @@ const UIThemeSettings: React.FC = ({ userID, userRole, acc const response = await fetch(url, { method: "PATCH", headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, body: JSON.stringify({ @@ -87,7 +87,7 @@ const UIThemeSettings: React.FC = ({ userID, userRole, acc const response = await fetch(url, { method: "PATCH", headers: { - Authorization: `Bearer ${accessToken}`, + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, body: JSON.stringify({ From 6cd4b3603fb4b9273f3d4cd4af40b04123ce8dbf Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Mon, 19 Jan 2026 18:48:35 +0530 Subject: [PATCH 04/58] fix(router): prevent retrying 4xx client errors (#19275) --- litellm/router.py | 6 + tests/local_testing/test_router_retries.py | 137 +++++++++++++++++++-- 2 files changed, 130 insertions(+), 13 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 0b07d5ed8c1..5a19d5a9a44 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4865,6 +4865,12 @@ class Router: ): raise error + status_code = getattr(error, "status_code", None) + if status_code is not None and not litellm._should_retry(status_code): + # 401/403 are special cases - allow retry if multiple deployments exist (handled below) + if status_code not in (401, 403): + raise error + if isinstance(error, litellm.NotFoundError): raise error # Error we should only retry if there are other deployments diff --git a/tests/local_testing/test_router_retries.py b/tests/local_testing/test_router_retries.py index 65539cfeee0..1d72d5914c5 100644 --- a/tests/local_testing/test_router_retries.py +++ b/tests/local_testing/test_router_retries.py @@ -676,6 +676,85 @@ def test_no_retry_for_not_found_error_404(): print("got exception", e) +def test_no_retry_for_bad_request_error_400(): + """ + Test that 400 BadRequestError is NOT retried, even if healthy deployments exist. + This tests the fix for GitHub issue #19216. + """ + healthy_deployments = ["deployment1", "deployment2"] # Multiple healthy deployments + + router = Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "azure/gpt-4.1-mini", + "api_key": os.getenv("AZURE_API_KEY"), + "api_version": os.getenv("AZURE_API_VERSION"), + "api_base": os.getenv("AZURE_API_BASE"), + }, + } + ] + ) + + # Act & Assert + error = litellm.BadRequestError( + message="400 Invalid request parameters", + model="gpt-3.5-turbo", + llm_provider="azure", + ) + try: + response = router.should_retry_this_error( + error=error, healthy_deployments=healthy_deployments + ) + pytest.fail( + "Should have raised BadRequestError - 400 errors should never be retried" + ) + except litellm.BadRequestError as e: + print("Correctly raised BadRequestError without retry:", e) + + +def test_no_retry_for_unprocessable_entity_error_422(): + """ + Test that 422 UnprocessableEntityError is NOT retried, even if healthy deployments exist. + """ + healthy_deployments = ["deployment1", "deployment2"] # Multiple healthy deployments + + router = Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "azure/gpt-4.1-mini", + "api_key": os.getenv("AZURE_API_KEY"), + "api_version": os.getenv("AZURE_API_VERSION"), + "api_base": os.getenv("AZURE_API_BASE"), + }, + } + ] + ) + + # Act & Assert + error = litellm.UnprocessableEntityError( + message="422 Unprocessable Entity", + model="gpt-3.5-turbo", + llm_provider="azure", + response=httpx.Response( + status_code=422, + request=httpx.Request(method="POST", url="https://api.openai.com/v1"), + ), + ) + try: + response = router.should_retry_this_error( + error=error, healthy_deployments=healthy_deployments + ) + pytest.fail( + "Should have raised UnprocessableEntityError - 422 errors should never be retried" + ) + except litellm.UnprocessableEntityError as e: + print("Correctly raised UnprocessableEntityError without retry:", e) + + internal_server_error = litellm.InternalServerError( message="internal server error", model="gpt-12", @@ -809,7 +888,7 @@ async def test_router_timeout_model_specific_and_global(): async def test_router_retry_num_retries_tracking(): """ Test that num_retries attribute is correctly set on exceptions when all retries are exhausted. - + This verifies the fix for the bug where num_retries was incorrectly set to current_attempt (0-indexed) instead of the actual number of retries attempted. """ @@ -837,8 +916,17 @@ async def test_router_retry_num_retries_tracking(): ) with patch.object(router, "make_call", side_effect=mock_make_call): - with patch.object(router, "_async_get_healthy_deployments", return_value=([{"model_info": {"id": "test-id"}}], [{"model_info": {"id": "test-id"}}])): - with patch.object(router, "_time_to_sleep_before_retry", return_value=0.01): # Fast retries for testing + with patch.object( + router, + "_async_get_healthy_deployments", + return_value=( + [{"model_info": {"id": "test-id"}}], + [{"model_info": {"id": "test-id"}}], + ), + ): + with patch.object( + router, "_time_to_sleep_before_retry", return_value=0.01 + ): # Fast retries for testing try: await router.acompletion( model="gpt-3.5-turbo", @@ -847,15 +935,27 @@ async def test_router_retry_num_retries_tracking(): pytest.fail("Expected exception to be raised") except litellm.RateLimitError as e: # Verify num_retries is correctly set to 3 (not 2, which would be current_attempt) - assert hasattr(e, "num_retries"), "Exception should have num_retries attribute" - assert hasattr(e, "max_retries"), "Exception should have max_retries attribute" - assert e.num_retries == 3, f"Expected num_retries to be 3, got {e.num_retries}" - assert e.max_retries == 3, f"Expected max_retries to be 3, got {e.max_retries}" - + assert hasattr( + e, "num_retries" + ), "Exception should have num_retries attribute" + assert hasattr( + e, "max_retries" + ), "Exception should have max_retries attribute" + assert ( + e.num_retries == 3 + ), f"Expected num_retries to be 3, got {e.num_retries}" + assert ( + e.max_retries == 3 + ), f"Expected max_retries to be 3, got {e.max_retries}" + # Verify the error message includes correct retry information error_str = str(e) - assert "LiteLLM Retried: 3 times" in error_str, f"Error message should indicate 3 retries: {error_str}" - assert "LiteLLM Max Retries: 3" in error_str, f"Error message should show max retries: {error_str}" + assert ( + "LiteLLM Retried: 3 times" in error_str + ), f"Error message should indicate 3 retries: {error_str}" + assert ( + "LiteLLM Max Retries: 3" in error_str + ), f"Error message should show max retries: {error_str}" @pytest.mark.asyncio @@ -887,7 +987,14 @@ async def test_router_retry_num_retries_single_retry(): ) with patch.object(router, "make_call", side_effect=mock_make_call): - with patch.object(router, "_async_get_healthy_deployments", return_value=([{"model_info": {"id": "test-id"}}], [{"model_info": {"id": "test-id"}}])): + with patch.object( + router, + "_async_get_healthy_deployments", + return_value=( + [{"model_info": {"id": "test-id"}}], + [{"model_info": {"id": "test-id"}}], + ), + ): with patch.object(router, "_time_to_sleep_before_retry", return_value=0.01): try: await router.acompletion( @@ -897,5 +1004,9 @@ async def test_router_retry_num_retries_single_retry(): pytest.fail("Expected exception to be raised") except litellm.Timeout as e: # With num_retries=1, we should attempt 1 retry - assert e.num_retries == 1, f"Expected num_retries to be 1, got {e.num_retries}" - assert e.max_retries == 1, f"Expected max_retries to be 1, got {e.max_retries}" \ No newline at end of file + assert ( + e.num_retries == 1 + ), f"Expected num_retries to be 1, got {e.num_retries}" + assert ( + e.max_retries == 1 + ), f"Expected max_retries to be 1, got {e.max_retries}" From 5db0e3289ad0f69e3dd327c0ba4b1b1dc8fa27f8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=B3n=20Levy?= Date: Mon, 19 Jan 2026 13:20:24 +0000 Subject: [PATCH 05/58] fix(agentcore): simplify agentcore streaming (#17141) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(agentcore): simplify agentcore streaming * fix(agentcore): move CustomStreamWrapper import to module level The deferred imports inside streaming methods caused initialization delays during health check requests, leading to timeouts in ECS deployments. - Move CustomStreamWrapper import to module-level (line 19) - Remove deferred imports from get_sync_custom_stream_wrapper (line 588) - Remove deferred import from get_async_custom_stream_wrapper (line 747) - Remove from TYPE_CHECKING block to use actual import This ensures the import happens at module load time rather than during first request processing, preventing health check endpoint blocking. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude * test(agentcore): ensure sync response * chore: upgrade boto3 to 1.40.76 in pyproject.toml * chore: added taplo.toml * fix(types): correct annotation type hint for MyPy compatibility Update _convert_annotations_to_chat_format return type from Dict[str, Any] to ChatCompletionAnnotation TypedDict to match the Message class's expected type signature. Co-Authored-By: Claude --------- Co-authored-by: Claude Co-authored-by: Benedikt Óskarsson --- .../transformation.py | 4 +- .../bedrock/chat/agentcore/sse_iterator.py | 252 ---------------- .../bedrock/chat/agentcore/transformation.py | 276 +++++++++++++---- pyproject.toml | 2 +- requirements.txt | 4 +- taplo.toml | 23 ++ .../llm_translation/test_bedrock_agentcore.py | 281 +++++++++++++++++- 7 files changed, 526 insertions(+), 316 deletions(-) delete mode 100644 litellm/llms/bedrock/chat/agentcore/sse_iterator.py create mode 100644 taplo.toml diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index af8185aa215..c98799d2542 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -779,10 +779,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): @staticmethod def _convert_annotations_to_chat_format( annotations: Optional[List[Any]], - ) -> Optional[List["ChatCompletionAnnotation"]]: + ) -> Optional[List[ChatCompletionAnnotation]]: """ Convert annotations from Responses API to Chat Completions format. - + Annotations are already in compatible format between both APIs, so we just need to convert Pydantic models to dicts. """ diff --git a/litellm/llms/bedrock/chat/agentcore/sse_iterator.py b/litellm/llms/bedrock/chat/agentcore/sse_iterator.py deleted file mode 100644 index 90c5ada769f..00000000000 --- a/litellm/llms/bedrock/chat/agentcore/sse_iterator.py +++ /dev/null @@ -1,252 +0,0 @@ -""" -SSE Stream Iterator for Bedrock AgentCore. - -Handles Server-Sent Events (SSE) streaming responses from AgentCore. -""" - -import json -from typing import TYPE_CHECKING, Any, Optional - -import httpx - -from litellm._logging import verbose_logger -from litellm._uuid import uuid -from litellm.types.llms.bedrock_agentcore import AgentCoreUsage -from litellm.types.utils import Delta, ModelResponse, StreamingChoices, Usage - -if TYPE_CHECKING: - pass - - -class AgentCoreSSEStreamIterator: - """ - Iterator for AgentCore SSE streaming responses. - Supports both sync and async iteration. - - CRITICAL: The line iterators are created lazily on first access and reused. - We must NOT create new iterators in __aiter__/__iter__ because - CustomStreamWrapper calls __aiter__ on every call to its __anext__, - which would create new iterators and cause StreamConsumed errors. - """ - - def __init__(self, response: httpx.Response, model: str): - self.response = response - self.model = model - self.finished = False - self._sync_iter: Any = None - self._async_iter: Any = None - self._sync_iter_initialized = False - self._async_iter_initialized = False - - def __iter__(self): - """Initialize sync iteration - create iterator lazily on first call only.""" - if not self._sync_iter_initialized: - self._sync_iter = iter(self.response.iter_lines()) - self._sync_iter_initialized = True - return self - - def __aiter__(self): - """Initialize async iteration - create iterator lazily on first call only.""" - if not self._async_iter_initialized: - self._async_iter = self.response.aiter_lines().__aiter__() - self._async_iter_initialized = True - return self - - def _parse_sse_line(self, line: str) -> Optional[ModelResponse]: - """ - Parse a single SSE line and return a ModelResponse chunk if applicable. - - AgentCore SSE format: - - data: {"event": {"contentBlockDelta": {"delta": {"text": "..."}}}} - - data: {"event": {"metadata": {"usage": {...}}}} - - data: {"message": {...}} - """ - line = line.strip() - if not line or not line.startswith("data:"): - return None - - json_str = line[5:].strip() - if not json_str: - return None - - try: - data = json.loads(json_str) - - # Skip non-dict data (some lines contain Python repr strings) - if not isinstance(data, dict): - return None - - # Process content delta events - if "event" in data and isinstance(data["event"], dict): - event_payload = data["event"] - content_block_delta = event_payload.get("contentBlockDelta") - - if content_block_delta: - delta = content_block_delta.get("delta", {}) - text = delta.get("text", "") - - if text: - # Return chunk with text - chunk = ModelResponse( - id=f"chatcmpl-{uuid.uuid4()}", - created=0, - model=self.model, - object="chat.completion.chunk", - ) - - chunk.choices = [ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta(content=text, role="assistant"), - ) - ] - - return chunk - - # Check for metadata/usage - this signals the end - metadata = event_payload.get("metadata") - if metadata and "usage" in metadata: - chunk = ModelResponse( - id=f"chatcmpl-{uuid.uuid4()}", - created=0, - model=self.model, - object="chat.completion.chunk", - ) - - chunk.choices = [ - StreamingChoices( - finish_reason="stop", - index=0, - delta=Delta(), - ) - ] - - usage_data: AgentCoreUsage = metadata["usage"] # type: ignore - setattr( - chunk, - "usage", - Usage( - prompt_tokens=usage_data.get("inputTokens", 0), - completion_tokens=usage_data.get("outputTokens", 0), - total_tokens=usage_data.get("totalTokens", 0), - ), - ) - - self.finished = True - return chunk - - # Check for final message (alternative finish signal) - if "message" in data and isinstance(data["message"], dict): - if not self.finished: - chunk = ModelResponse( - id=f"chatcmpl-{uuid.uuid4()}", - created=0, - model=self.model, - object="chat.completion.chunk", - ) - - chunk.choices = [ - StreamingChoices( - finish_reason="stop", - index=0, - delta=Delta(), - ) - ] - - self.finished = True - return chunk - - except json.JSONDecodeError: - verbose_logger.debug(f"Skipping non-JSON SSE line: {line[:100]}") - - return None - - def _create_final_chunk(self) -> ModelResponse: - """Create a final chunk to signal stream completion.""" - chunk = ModelResponse( - id=f"chatcmpl-{uuid.uuid4()}", - created=0, - model=self.model, - object="chat.completion.chunk", - ) - - chunk.choices = [ - StreamingChoices( - finish_reason="stop", - index=0, - delta=Delta(), - ) - ] - - return chunk - - def __next__(self) -> ModelResponse: - """ - Sync iteration - parse SSE events and yield ModelResponse chunks. - - Uses next() on the stored iterator to properly resume between calls. - """ - try: - if self._sync_iter is None: - raise StopIteration - - # Keep getting lines until we have a result to return - while True: - try: - line = next(self._sync_iter) - except StopIteration: - # Stream ended - send final chunk if not already finished - if not self.finished: - self.finished = True - return self._create_final_chunk() - raise - - result = self._parse_sse_line(line) - if result is not None: - return result - - except StopIteration: - raise - except httpx.StreamConsumed: - raise StopIteration - except httpx.StreamClosed: - raise StopIteration - except Exception as e: - verbose_logger.error(f"Error in AgentCore SSE stream: {str(e)}") - raise StopIteration - - async def __anext__(self) -> ModelResponse: - """ - Async iteration - parse SSE events and yield ModelResponse chunks. - - Uses __anext__() on the stored iterator to properly resume between calls. - """ - try: - if self._async_iter is None: - raise StopAsyncIteration - - # Keep getting lines until we have a result to return - while True: - try: - line = await self._async_iter.__anext__() - except StopAsyncIteration: - # Stream ended - send final chunk if not already finished - if not self.finished: - self.finished = True - return self._create_final_chunk() - raise - - result = self._parse_sse_line(line) - if result is not None: - return result - - except StopAsyncIteration: - raise - except httpx.StreamConsumed: - raise StopAsyncIteration - except httpx.StreamClosed: - raise StopAsyncIteration - except Exception as e: - verbose_logger.error(f"Error in AgentCore SSE stream: {str(e)}") - raise StopAsyncIteration diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 7c65cad94df..94e845e3095 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -5,6 +5,7 @@ https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agentcore_InvokeAgen """ import json +from collections.abc import AsyncGenerator from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast from urllib.parse import quote @@ -15,9 +16,9 @@ from litellm._uuid import uuid from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM -from litellm.llms.bedrock.chat.agentcore.sse_iterator import AgentCoreSSEStreamIterator from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.llms.bedrock_agentcore import ( AgentCoreMessage, @@ -25,19 +26,17 @@ from litellm.types.llms.bedrock_agentcore import ( AgentCoreUsage, ) from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import Choices, Message, ModelResponse, Usage +from litellm.types.utils import Choices, Delta, Message, ModelResponse, StreamingChoices, Usage if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler - from litellm.utils import CustomStreamWrapper LiteLLMLoggingObj = _LiteLLMLoggingObj else: LiteLLMLoggingObj = Any HTTPHandler = Any AsyncHTTPHandler = Any - CustomStreamWrapper = Any class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): @@ -116,7 +115,8 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): fake_stream: Optional[bool] = None, ) -> Tuple[dict, Optional[bytes]]: # Check if api_key (bearer token) is provided for Cognito authentication - jwt_token = optional_params.get("api_key") + # Priority: api_key parameter first, then optional_params + jwt_token = api_key or optional_params.get("api_key") if jwt_token: verbose_logger.debug( f"AgentCore: Using Bearer token authentication (Cognito/JWT) - token: {jwt_token[:50]}..." @@ -437,22 +437,104 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): content=content, usage=usage_data, final_message=final_message ) - def get_streaming_response( + def _stream_agentcore_response_sync( self, + response: httpx.Response, model: str, - raw_response: httpx.Response, - ) -> AgentCoreSSEStreamIterator: + ): """ - Return a streaming iterator for SSE responses. - - Args: - model: The model name - raw_response: Raw HTTP response with streaming data - - Returns: - AgentCoreSSEStreamIterator: Iterator that yields ModelResponse chunks + Internal sync generator that parses SSE and yields ModelResponse chunks. """ - return AgentCoreSSEStreamIterator(response=raw_response, model=model) + buffer = "" + for text_chunk in response.iter_text(): + buffer += text_chunk + + # Process complete lines + while '\n' in buffer: + line, buffer = buffer.split('\n', 1) + line = line.strip() + + if not line or not line.startswith('data:'): + continue + + json_str = line[5:].strip() + if not json_str: + continue + + try: + data_obj = json.loads(json_str) + if not isinstance(data_obj, dict): + continue + + # Process contentBlockDelta events + if "event" in data_obj and isinstance(data_obj["event"], dict): + event_payload = data_obj["event"] + content_block_delta = event_payload.get("contentBlockDelta") + + if content_block_delta: + delta = content_block_delta.get("delta", {}) + text = delta.get("text", "") + + if text: + chunk = ModelResponse( + id=f"chatcmpl-{uuid.uuid4()}", + created=0, + model=model, + object="chat.completion.chunk", + ) + chunk.choices = [ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content=text, role="assistant"), + ) + ] + yield chunk + + # Process metadata/usage + metadata = event_payload.get("metadata") + if metadata and "usage" in metadata: + chunk = ModelResponse( + id=f"chatcmpl-{uuid.uuid4()}", + created=0, + model=model, + object="chat.completion.chunk", + ) + chunk.choices = [ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(), + ) + ] + usage_data: AgentCoreUsage = metadata["usage"] # type: ignore + setattr(chunk, "usage", Usage( + prompt_tokens=usage_data.get("inputTokens", 0), + completion_tokens=usage_data.get("outputTokens", 0), + total_tokens=usage_data.get("totalTokens", 0), + )) + yield chunk + + # Process final message + if "message" in data_obj and isinstance(data_obj["message"], dict): + chunk = ModelResponse( + id=f"chatcmpl-{uuid.uuid4()}", + created=0, + model=model, + object="chat.completion.chunk", + ) + chunk.choices = [ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(), + ) + ] + yield chunk + + except json.JSONDecodeError: + verbose_logger.debug(f"Skipping non-JSON SSE line: {line[:100]}") + continue def get_sync_custom_stream_wrapper( self, @@ -466,17 +548,14 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): client: Optional[Union[HTTPHandler, "AsyncHTTPHandler"]] = None, json_mode: Optional[bool] = None, signed_json_body: Optional[bytes] = None, - ) -> CustomStreamWrapper: + ) -> "CustomStreamWrapper": """ - Get a CustomStreamWrapper for synchronous streaming. - - This is called when stream=True is passed to completion(). + Simplified sync streaming - returns a generator that yields ModelResponse chunks. """ from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, _get_httpx_client, ) - from litellm.utils import CustomStreamWrapper if client is None or not isinstance(client, HTTPHandler): client = _get_httpx_client(params={}) @@ -488,7 +567,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): api_base, headers=headers, data=signed_json_body if signed_json_body else json.dumps(data), - stream=True, # THIS IS KEY - tells httpx to not buffer + stream=True, logging_obj=logging_obj, ) @@ -497,18 +576,6 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): status_code=response.status_code, message=str(response.read()) ) - # Create iterator for SSE stream - completion_stream = self.get_streaming_response( - model=model, raw_response=response - ) - - streaming_response = CustomStreamWrapper( - completion_stream=completion_stream, - model=model, - custom_llm_provider=custom_llm_provider, - logging_obj=logging_obj, - ) - # LOGGING logging_obj.post_call( input=messages, @@ -517,7 +584,112 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): additional_args={"complete_input_dict": data}, ) - return streaming_response + # Wrap the generator in CustomStreamWrapper + return CustomStreamWrapper( + completion_stream=self._stream_agentcore_response_sync(response, model), + model=model, + custom_llm_provider="bedrock", + logging_obj=logging_obj, + ) + + async def _stream_agentcore_response( + self, + response: httpx.Response, + model: str, + ) -> AsyncGenerator[ModelResponse, None]: + """ + Internal async generator that parses SSE and yields ModelResponse chunks. + """ + buffer = "" + async for text_chunk in response.aiter_text(): + buffer += text_chunk + + # Process complete lines + while '\n' in buffer: + line, buffer = buffer.split('\n', 1) + line = line.strip() + + if not line or not line.startswith('data:'): + continue + + json_str = line[5:].strip() + if not json_str: + continue + + try: + data_obj = json.loads(json_str) + if not isinstance(data_obj, dict): + continue + + # Process contentBlockDelta events + if "event" in data_obj and isinstance(data_obj["event"], dict): + event_payload = data_obj["event"] + content_block_delta = event_payload.get("contentBlockDelta") + + if content_block_delta: + delta = content_block_delta.get("delta", {}) + text = delta.get("text", "") + + if text: + chunk = ModelResponse( + id=f"chatcmpl-{uuid.uuid4()}", + created=0, + model=model, + object="chat.completion.chunk", + ) + chunk.choices = [ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content=text, role="assistant"), + ) + ] + yield chunk + + # Process metadata/usage + metadata = event_payload.get("metadata") + if metadata and "usage" in metadata: + chunk = ModelResponse( + id=f"chatcmpl-{uuid.uuid4()}", + created=0, + model=model, + object="chat.completion.chunk", + ) + chunk.choices = [ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(), + ) + ] + usage_data: AgentCoreUsage = metadata["usage"] # type: ignore + setattr(chunk, "usage", Usage( + prompt_tokens=usage_data.get("inputTokens", 0), + completion_tokens=usage_data.get("outputTokens", 0), + total_tokens=usage_data.get("totalTokens", 0), + )) + yield chunk + + # Process final message + if "message" in data_obj and isinstance(data_obj["message"], dict): + chunk = ModelResponse( + id=f"chatcmpl-{uuid.uuid4()}", + created=0, + model=model, + object="chat.completion.chunk", + ) + chunk.choices = [ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(), + ) + ] + yield chunk + + except json.JSONDecodeError: + verbose_logger.debug(f"Skipping non-JSON SSE line: {line[:100]}") + continue async def get_async_custom_stream_wrapper( self, @@ -531,17 +703,14 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): client: Optional["AsyncHTTPHandler"] = None, json_mode: Optional[bool] = None, signed_json_body: Optional[bytes] = None, - ) -> CustomStreamWrapper: + ) -> "CustomStreamWrapper": """ - Get a CustomStreamWrapper for asynchronous streaming. - - This is called when stream=True is passed to acompletion(). + Simplified async streaming - returns an async generator that yields ModelResponse chunks. """ from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) - from litellm.utils import CustomStreamWrapper if client is None or not isinstance(client, AsyncHTTPHandler): client = get_async_httpx_client( @@ -555,7 +724,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): api_base, headers=headers, data=signed_json_body if signed_json_body else json.dumps(data), - stream=True, # THIS IS KEY - tells httpx to not buffer + stream=True, logging_obj=logging_obj, ) @@ -564,18 +733,6 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): status_code=response.status_code, message=str(await response.aread()) ) - # Create iterator for SSE stream - completion_stream = self.get_streaming_response( - model=model, raw_response=response - ) - - streaming_response = CustomStreamWrapper( - completion_stream=completion_stream, - model=model, - custom_llm_provider=custom_llm_provider, - logging_obj=logging_obj, - ) - # LOGGING logging_obj.post_call( input=messages, @@ -584,7 +741,13 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): additional_args={"complete_input_dict": data}, ) - return streaming_response + # Wrap the async generator in CustomStreamWrapper + return CustomStreamWrapper( + completion_stream=self._stream_agentcore_response(response, model), + model=model, + custom_llm_provider="bedrock", + logging_obj=logging_obj, + ) @property def has_custom_stream_wrapper(self) -> bool: @@ -692,4 +855,5 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): stream: Optional[bool], custom_llm_provider: Optional[str] = None, ) -> bool: - return True + # AgentCore supports true streaming - don't buffer + return False diff --git a/pyproject.toml b/pyproject.toml index d6242b27786..fbb3aaf9a52 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,7 +56,7 @@ google-cloud-iam = {version = "^2.19.1", optional = true} resend = {version = ">=0.8.0", optional = true} pynacl = {version = "^1.5.0", optional = true} websockets = {version = "^15.0.1", optional = true} -boto3 = {version = "1.36.0", optional = true} +boto3 = { version = "1.40.76", optional = true } redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"} mcp = {version = "^1.21.2", optional = true, python = ">=3.10"} litellm-proxy-extras = {version = "0.4.23", optional = true} diff --git a/requirements.txt b/requirements.txt index a49a94ca274..1b5cae72013 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,7 +10,7 @@ uvicorn==0.31.1 # server dep gunicorn==23.0.0 # server dep fastuuid==0.13.5 # for uuid4 uvloop==0.21.0 # uvicorn dep, gives us much better performance under load -boto3==1.36.0 # aws bedrock/sagemaker calls +boto3==1.40.53 # aws bedrock/sagemaker calls (has bedrock-agentcore-control, compatible with aioboto3) redis==5.2.1 # redis caching prisma==0.11.0 # for db nodejs-wheel-binaries==24.12.0 ## required by prisma for migrations, prevents runtime download (updated from nodejs-bin for security fixes) @@ -58,8 +58,8 @@ tokenizers==0.20.2 # for calculating usage click==8.1.7 # for proxy cli rich==13.7.1 # for litellm proxy cli jinja2==3.1.6 # for prompt templates +aioboto3==15.5.0 # for async sagemaker calls (updated to match boto3 1.40.73) aiohttp==3.13.3 # for network calls -aioboto3==13.4.0 # for async sagemaker calls tenacity==8.5.0 # for retrying requests, when litellm.num_retries set pydantic>=2.11,<3 # proxy + openai req. + mcp jsonschema>=4.23.0,<5.0.0 # validating json schema - aligned with openapi-core + mcp diff --git a/taplo.toml b/taplo.toml new file mode 100644 index 00000000000..e5065a69616 --- /dev/null +++ b/taplo.toml @@ -0,0 +1,23 @@ +[formatting] + +# Keep your table/key order as written (project, project.scripts, dependency-groups, tool.uv, ...) +reorder_keys = false + +# Force arrays to stay multiline once expanded +array_auto_expand = true +array_auto_collapse = false + +# Keep nice spacing inside arrays and inline tables +compact_arrays = false +compact_inline_tables = true + +# Don’t align `=` vertically (matches your example) +align_entries = true + +# Reasonable defaults +align_comments = true +trailing_newline = true +reorder_arrays = true +reorder_inline_tables = false +allowed_blank_lines = 2 +crlf = false diff --git a/tests/llm_translation/test_bedrock_agentcore.py b/tests/llm_translation/test_bedrock_agentcore.py index 3afb01482ac..f00ceb8eab4 100644 --- a/tests/llm_translation/test_bedrock_agentcore.py +++ b/tests/llm_translation/test_bedrock_agentcore.py @@ -12,10 +12,9 @@ sys.path.insert( ) import litellm -from unittest.mock import MagicMock, patch -import pytest - +from unittest.mock import MagicMock, Mock, patch import pytest +import httpx @pytest.mark.parametrize( "model", [ @@ -367,3 +366,279 @@ def test_bedrock_agentcore_without_api_key_uses_sigv4(): assert "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id" in headers assert headers["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] == "sigv4-test-session" + +def test_agentcore_parse_json_response(): + """ + Unit test for JSON response parsing (non-streaming) + Verifies that content-type: application/json responses are parsed correctly + """ + from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig + + config = AmazonAgentCoreConfig() + + # Create a mock JSON response + mock_response = Mock(spec=httpx.Response) + mock_response.headers = {"content-type": "application/json"} + mock_response.json.return_value = { + "result": { + "role": "assistant", + "content": [{"text": "Hello from JSON response"}] + } + } + + # Parse the response + parsed = config._get_parsed_response(mock_response) + + # Verify content extraction + assert parsed["content"] == "Hello from JSON response" + # JSON responses don't include usage data + assert parsed["usage"] is None + # Final message should be the result object + assert parsed["final_message"] == mock_response.json.return_value["result"] + + +def test_agentcore_parse_sse_response(): + """ + Unit test for SSE response parsing (streaming response consumed as text) + Verifies that text/event-stream responses are parsed correctly + """ + from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig + + config = AmazonAgentCoreConfig() + + # Create a mock SSE response with multiple events + sse_data = """data: {"event":{"contentBlockDelta":{"delta":{"text":"Hello "}}}} + +data: {"event":{"contentBlockDelta":{"delta":{"text":"from SSE"}}}} + +data: {"event":{"metadata":{"usage":{"inputTokens":10,"outputTokens":5,"totalTokens":15}}}} + +data: {"message":{"role":"assistant","content":[{"text":"Hello from SSE"}]}} +""" + + mock_response = Mock(spec=httpx.Response) + mock_response.headers = {"content-type": "text/event-stream"} + mock_response.text = sse_data + + # Parse the response + parsed = config._get_parsed_response(mock_response) + + # Verify content extraction from final message + assert parsed["content"] == "Hello from SSE" + # SSE responses can include usage data + assert parsed["usage"] is not None + assert parsed["usage"]["inputTokens"] == 10 + assert parsed["usage"]["outputTokens"] == 5 + assert parsed["usage"]["totalTokens"] == 15 + # Final message should be present + assert parsed["final_message"] is not None + assert parsed["final_message"]["role"] == "assistant" + + +def test_agentcore_parse_sse_response_without_final_message(): + """ + Unit test for SSE response parsing when only deltas are present (no final message) + """ + from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig + + config = AmazonAgentCoreConfig() + + # Create a mock SSE response with only content deltas + sse_data = """data: {"event":{"contentBlockDelta":{"delta":{"text":"First "}}}} + +data: {"event":{"contentBlockDelta":{"delta":{"text":"second "}}}} + +data: {"event":{"contentBlockDelta":{"delta":{"text":"third"}}}} +""" + + mock_response = Mock(spec=httpx.Response) + mock_response.headers = {"content-type": "text/event-stream"} + mock_response.text = sse_data + + # Parse the response + parsed = config._get_parsed_response(mock_response) + + # Content should be concatenated from deltas + assert parsed["content"] == "First second third" + # No final message + assert parsed["final_message"] is None + + +def test_agentcore_transform_response_json(): + """ + Integration test for transform_response with JSON response + Verifies end-to-end transformation of JSON responses to ModelResponse + """ + from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig + from litellm.types.utils import ModelResponse + + config = AmazonAgentCoreConfig() + + # Create mock JSON response + mock_response = Mock(spec=httpx.Response) + mock_response.headers = {"content-type": "application/json"} + mock_response.json.return_value = { + "result": { + "role": "assistant", + "content": [{"text": "Response from transform_response"}] + } + } + mock_response.status_code = 200 + + # Create model response + model_response = ModelResponse() + + # Mock logging object + mock_logging = MagicMock() + + # Transform the response + result = config.transform_response( + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/test", + raw_response=mock_response, + model_response=model_response, + logging_obj=mock_logging, + request_data={}, + messages=[{"role": "user", "content": "test"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + # Verify ModelResponse structure + assert len(result.choices) == 1 + assert result.choices[0].message.content == "Response from transform_response" + assert result.choices[0].message.role == "assistant" + assert result.choices[0].finish_reason == "stop" + assert result.choices[0].index == 0 + + +def test_agentcore_transform_response_sse(): + """ + Integration test for transform_response with SSE response + Verifies end-to-end transformation of SSE responses to ModelResponse + """ + from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig + from litellm.types.utils import ModelResponse + + config = AmazonAgentCoreConfig() + + # Create mock SSE response + sse_data = """data: {"event":{"contentBlockDelta":{"delta":{"text":"SSE "}}}} + +data: {"event":{"contentBlockDelta":{"delta":{"text":"response"}}}} + +data: {"event":{"metadata":{"usage":{"inputTokens":20,"outputTokens":10,"totalTokens":30}}}} + +data: {"message":{"role":"assistant","content":[{"text":"SSE response"}]}} +""" + + mock_response = Mock(spec=httpx.Response) + mock_response.headers = {"content-type": "text/event-stream"} + mock_response.text = sse_data + mock_response.status_code = 200 + + # Create model response + model_response = ModelResponse() + + # Mock logging object + mock_logging = MagicMock() + + # Transform the response + result = config.transform_response( + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/test", + raw_response=mock_response, + model_response=model_response, + logging_obj=mock_logging, + request_data={}, + messages=[{"role": "user", "content": "test"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + # Verify ModelResponse structure + assert len(result.choices) == 1 + assert result.choices[0].message.content == "SSE response" + assert result.choices[0].message.role == "assistant" + assert result.choices[0].finish_reason == "stop" + + # Verify usage data from SSE metadata + assert hasattr(result, "usage") + assert result.usage.prompt_tokens == 20 + assert result.usage.completion_tokens == 10 + assert result.usage.total_tokens == 30 + + +def test_agentcore_synchronous_non_streaming_response(): + """ + Test that synchronous (non-streaming) AgentCore calls still work correctly + after streaming simplification changes. + + This test verifies: + 1. Synchronous completion calls work (stream=False or no stream param) + 2. Response is properly parsed and returned as ModelResponse + 3. Content is extracted correctly + 4. Usage data is calculated when not provided by API + + This is a regression test for the streaming simplification changes + to ensure we didn't break the non-streaming code path. + """ + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + litellm._turn_on_debug() + client = HTTPHandler() + + # Mock a JSON response (typical for synchronous AgentCore calls) + mock_json_response = { + "result": { + "role": "assistant", + "content": [{"text": "This is a synchronous response from AgentCore."}] + } + } + + # Create a mock response object + mock_response = Mock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_response.json.return_value = mock_json_response + + with patch.object(client, "post", return_value=mock_response) as mock_post: + # Make a synchronous (non-streaming) completion call + response = litellm.completion( + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", + messages=[ + { + "role": "user", + "content": "Test synchronous response", + } + ], + stream=False, # Explicitly disable streaming + client=client, + ) + + # Verify the response structure + assert response is not None + assert hasattr(response, "choices") + assert len(response.choices) > 0 + + # Verify content + message = response.choices[0].message + assert message is not None + assert message.content == "This is a synchronous response from AgentCore." + assert message.role == "assistant" + + # Verify completion metadata + assert response.choices[0].finish_reason == "stop" + assert response.choices[0].index == 0 + + # Verify usage data exists (either from API or calculated) + assert hasattr(response, "usage") + assert response.usage is not None + assert response.usage.prompt_tokens > 0 + assert response.usage.completion_tokens > 0 + assert response.usage.total_tokens > 0 + + print(f"Synchronous response: {response}") + print(f"Content: {message.content}") + print(f"Usage: prompt={response.usage.prompt_tokens}, completion={response.usage.completion_tokens}, total={response.usage.total_tokens}") + From 32562708ee0546551d0284910124d2d618778328 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Mon, 19 Jan 2026 19:06:54 +0530 Subject: [PATCH 06/58] fix(proxy): add /a2a/{agent_id}/.well-known/agent-card.json to agent_routes allowlist (#19277) --- litellm/proxy/_types.py | 90 +++++++++++++++++++++-------------------- 1 file changed, 46 insertions(+), 44 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f717cca9b9f..2fd03ac2128 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -360,7 +360,6 @@ class LiteLLMRoutes(enum.Enum): # OCR "/ocr", "/v1/ocr", - # containers API "/containers", "/v1/containers", @@ -421,6 +420,7 @@ class LiteLLMRoutes(enum.Enum): "/a2a/{agent_id}", "/a2a/{agent_id}/message/send", "/a2a/{agent_id}/message/stream", + "/a2a/{agent_id}/.well-known/agent-card.json", ] google_routes = [ @@ -836,9 +836,9 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} - model_max_budget: Optional[dict] = ( - {} - ) # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} + model_max_budget: Optional[ + dict + ] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} model_config = ConfigDict(protected_namespaces=()) model_rpm_limit: Optional[dict] = None @@ -1372,12 +1372,12 @@ class NewCustomerRequest(BudgetNewRequest): blocked: bool = False # allow/disallow requests for this end-user budget_id: Optional[str] = None # give either a budget_id or max_budget spend: Optional[float] = None - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model @model_validator(mode="before") @classmethod @@ -1399,12 +1399,12 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): blocked: bool = False # allow/disallow requests for this end-user max_budget: Optional[float] = None budget_id: Optional[str] = None # give either a budget_id or max_budget - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model class DeleteCustomerRequest(LiteLLMPydanticObjectBase): @@ -1490,15 +1490,15 @@ class NewTeamRequest(TeamBase): ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm model_tpm_limit: Optional[Dict[str, int]] = None - team_member_budget: Optional[float] = ( - None # allow user to set a budget for all team members - ) - team_member_rpm_limit: Optional[int] = ( - None # allow user to set RPM limit for all team members - ) - team_member_tpm_limit: Optional[int] = ( - None # allow user to set TPM limit for all team members - ) + team_member_budget: Optional[ + float + ] = None # allow user to set a budget for all team members + team_member_rpm_limit: Optional[ + int + ] = None # allow user to set RPM limit for all team members + team_member_tpm_limit: Optional[ + int + ] = None # allow user to set TPM limit for all team members team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m" allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None @@ -1586,9 +1586,9 @@ class BlockKeyRequest(LiteLLMPydanticObjectBase): class AddTeamCallback(LiteLLMPydanticObjectBase): callback_name: str - callback_type: Optional[Literal["success", "failure", "success_and_failure"]] = ( - "success_and_failure" - ) + callback_type: Optional[ + Literal["success", "failure", "success_and_failure"] + ] = "success_and_failure" callback_vars: Dict[str, str] @model_validator(mode="before") @@ -1916,9 +1916,9 @@ class ConfigList(LiteLLMPydanticObjectBase): stored_in_db: Optional[bool] field_default_value: Any premium_field: bool = False - nested_fields: Optional[List[FieldDetail]] = ( - None # For nested dictionary or Pydantic fields - ) + nested_fields: Optional[ + List[FieldDetail] + ] = None # For nested dictionary or Pydantic fields class UserHeaderMapping(LiteLLMPydanticObjectBase): @@ -2127,7 +2127,9 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): rotation_interval: Optional[str] = None # How often to rotate (e.g., "30d", "90d") last_rotation_at: Optional[datetime] = None # When this key was last rotated key_rotation_at: Optional[datetime] = None # When this key should next be rotated - router_settings: Optional[Dict] = None # Router settings for this key (Key > Team > Global precedence) + router_settings: Optional[ + Dict + ] = None # Router settings for this key (Key > Team > Global precedence) model_config = ConfigDict(protected_namespaces=()) @@ -2332,9 +2334,9 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): budget_id: Optional[str] = None created_at: datetime updated_at: datetime - user: Optional[Any] = ( - None # You might want to replace 'Any' with a more specific type if available - ) + user: Optional[ + Any + ] = None # You might want to replace 'Any' with a more specific type if available litellm_budget_table: Optional[LiteLLM_BudgetTable] = None model_config = ConfigDict(protected_namespaces=()) @@ -3306,9 +3308,9 @@ class TeamModelDeleteRequest(BaseModel): # Organization Member Requests class OrganizationMemberAddRequest(OrgMemberAddRequest): organization_id: str - max_budget_in_organization: Optional[float] = ( - None # Users max budget within the organization - ) + max_budget_in_organization: Optional[ + float + ] = None # Users max budget within the organization class OrganizationMemberDeleteRequest(MemberDeleteRequest): @@ -3523,9 +3525,9 @@ class ProviderBudgetResponse(LiteLLMPydanticObjectBase): Maps provider names to their budget configs. """ - providers: Dict[str, ProviderBudgetResponseObject] = ( - {} - ) # Dictionary mapping provider names to their budget configurations + providers: Dict[ + str, ProviderBudgetResponseObject + ] = {} # Dictionary mapping provider names to their budget configurations class ProxyStateVariables(TypedDict): @@ -3668,9 +3670,9 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): enforce_rbac: bool = False roles_jwt_field: Optional[str] = None # v2 on role mappings role_mappings: Optional[List[RoleMapping]] = None - object_id_jwt_field: Optional[str] = ( - None # can be either user / team, inferred from the role mapping - ) + object_id_jwt_field: Optional[ + str + ] = None # can be either user / team, inferred from the role mapping scope_mappings: Optional[List[ScopeMapping]] = None enforce_scope_based_access: bool = False enforce_team_based_model_access: bool = False From 29adf3431371c611d9e31d4a8c3ad2bc5a75de68 Mon Sep 17 00:00:00 2001 From: Manuel Schweigert Date: Mon, 19 Jan 2026 14:37:45 +0100 Subject: [PATCH 07/58] Add ChatGPT subscription support and responses bridge (#19030) * Add ChatGPT subscription support and responses bridge * Fix typing import for responses bridge * Guard device code timestamp parsing * add /v1/messages endpoint to chatgpt model --- docs/my-website/docs/providers/chatgpt.md | 84 ++++ docs/my-website/sidebars.js | 1 + litellm/__init__.py | 6 + litellm/_lazy_imports_registry.py | 4 + .../handler.py | 105 ++++- litellm/constants.py | 2 + .../get_llm_provider_logic.py | 8 + litellm/llms/chatgpt/authenticator.py | 388 ++++++++++++++++++ litellm/llms/chatgpt/chat/transformation.py | 75 ++++ litellm/llms/chatgpt/common_utils.py | 301 ++++++++++++++ .../llms/chatgpt/responses/transformation.py | 191 +++++++++ ...odel_prices_and_context_window_backup.json | 58 +++ litellm/types/utils.py | 1 + litellm/utils.py | 3 + model_prices_and_context_window.json | 57 +++ provider_endpoints_support.json | 18 + .../test_chatgpt_responses_transformation.py | 125 ++++++ .../chatgpt/test_chatgpt_authenticator.py | 68 +++ .../test_responses_api_bridge_non_stream.py | 44 ++ 19 files changed, 1536 insertions(+), 3 deletions(-) create mode 100644 docs/my-website/docs/providers/chatgpt.md create mode 100644 litellm/llms/chatgpt/authenticator.py create mode 100644 litellm/llms/chatgpt/chat/transformation.py create mode 100644 litellm/llms/chatgpt/common_utils.py create mode 100644 litellm/llms/chatgpt/responses/transformation.py create mode 100644 tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py create mode 100644 tests/test_litellm/llms/chatgpt/test_chatgpt_authenticator.py create mode 100644 tests/test_litellm/test_responses_api_bridge_non_stream.py diff --git a/docs/my-website/docs/providers/chatgpt.md b/docs/my-website/docs/providers/chatgpt.md new file mode 100644 index 00000000000..156bbf99df6 --- /dev/null +++ b/docs/my-website/docs/providers/chatgpt.md @@ -0,0 +1,84 @@ +# ChatGPT Subscription + +Use ChatGPT Pro/Max subscription models through LiteLLM with OAuth device flow authentication. + +| Property | Details | +|-------|-------| +| Description | ChatGPT subscription access (Codex + GPT-5.2 family) via ChatGPT backend API | +| Provider Route on LiteLLM | `chatgpt/` | +| Supported Endpoints | `/responses`, `/chat/completions` (bridged to Responses for supported models) | +| API Reference | https://chatgpt.com | + +ChatGPT subscription access is native to the Responses API. Chat Completions requests are bridged to Responses for supported models (for example `chatgpt/gpt-5.2`). + +Notes: +- The ChatGPT subscription backend rejects token limit fields (`max_tokens`, `max_output_tokens`, `max_completion_tokens`) and `metadata`. LiteLLM strips these fields for this provider. +- `/v1/chat/completions` honors `stream`. When `stream` is false (default), LiteLLM aggregates the Responses stream into a single JSON response. + +## Authentication + +ChatGPT subscription access uses an OAuth device code flow: + +1. LiteLLM prints a device code and verification URL +2. Open the URL, sign in, and enter the code +3. Tokens are stored locally for reuse + +## Usage - LiteLLM Python SDK + +### Responses (recommended for Codex models) + +```python showLineNumbers title="ChatGPT Responses" +import litellm + +response = litellm.responses( + model="chatgpt/gpt-5.2-codex", + input="Write a Python hello world" +) + +print(response) +``` + +### Chat Completions (bridged to Responses) + +```python showLineNumbers title="ChatGPT Chat Completions" +import litellm + +response = litellm.completion( + model="chatgpt/gpt-5.2", + messages=[{"role": "user", "content": "Write a Python hello world"}] +) + +print(response) +``` + +## Usage - LiteLLM Proxy + +```yaml showLineNumbers title="config.yaml" +model_list: + - model_name: chatgpt/gpt-5.2 + model_info: + mode: responses + litellm_params: + model: chatgpt/gpt-5.2 + - model_name: chatgpt/gpt-5.2-codex + model_info: + mode: responses + litellm_params: + model: chatgpt/gpt-5.2-codex +``` + +```bash showLineNumbers title="Start LiteLLM Proxy" +litellm --config config.yaml +``` + +## Configuration + +### Environment Variables + +- `CHATGPT_TOKEN_DIR`: Custom token storage directory +- `CHATGPT_AUTH_FILE`: Auth file name (default: `auth.json`) +- `CHATGPT_API_BASE`: Override API base (default: `https://chatgpt.com/backend-api/codex`) +- `OPENAI_CHATGPT_API_BASE`: Alias for `CHATGPT_API_BASE` +- `CHATGPT_ORIGINATOR`: Override the `originator` header value +- `CHATGPT_USER_AGENT`: Override the `User-Agent` header value +- `CHATGPT_USER_AGENT_SUFFIX`: Optional suffix appended to the `User-Agent` header diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 102e3dfe1c5..d1abe45332f 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -717,6 +717,7 @@ const sidebars = { "providers/galadriel", "providers/github", "providers/github_copilot", + "providers/chatgpt", "providers/gradient_ai", "providers/groq", "providers/helicone", diff --git a/litellm/__init__.py b/litellm/__init__.py index 9eb3f075d5e..83cbcfda3f1 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -557,6 +557,7 @@ docker_model_runner_models: Set = set() amazon_nova_models: Set = set() stability_models: Set = set() github_copilot_models: Set = set() +chatgpt_models: Set = set() minimax_models: Set = set() aws_polly_models: Set = set() gigachat_models: Set = set() @@ -812,6 +813,8 @@ def add_known_models(): stability_models.add(key) elif value.get("litellm_provider") == "github_copilot": github_copilot_models.add(key) + elif value.get("litellm_provider") == "chatgpt": + chatgpt_models.add(key) elif value.get("litellm_provider") == "minimax": minimax_models.add(key) elif value.get("litellm_provider") == "aws_polly": @@ -1025,6 +1028,7 @@ models_by_provider: dict = { "amazon_nova": amazon_nova_models, "stability": stability_models, "github_copilot": github_copilot_models, + "chatgpt": chatgpt_models, "minimax": minimax_models, "aws_polly": aws_polly_models, "gigachat": gigachat_models, @@ -1458,6 +1462,8 @@ if TYPE_CHECKING: from .llms.github_copilot.chat.transformation import GithubCopilotConfig as GithubCopilotConfig from .llms.github_copilot.responses.transformation import GithubCopilotResponsesAPIConfig as GithubCopilotResponsesAPIConfig from .llms.github_copilot.embedding.transformation import GithubCopilotEmbeddingConfig as GithubCopilotEmbeddingConfig + from .llms.chatgpt.chat.transformation import ChatGPTConfig as ChatGPTConfig + from .llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig as ChatGPTResponsesAPIConfig from .llms.gigachat.chat.transformation import GigaChatConfig as GigaChatConfig from .llms.gigachat.embedding.transformation import GigaChatEmbeddingConfig as GigaChatEmbeddingConfig from .llms.nebius.chat.transformation import NebiusConfig as NebiusConfig diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index f37c4dc6d04..46abde4eb5e 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -253,6 +253,8 @@ LLM_CONFIG_NAMES = ( "IBMWatsonXAudioTranscriptionConfig", "GithubCopilotConfig", "GithubCopilotResponsesAPIConfig", + "ChatGPTConfig", + "ChatGPTResponsesAPIConfig", "ManusResponsesAPIConfig", "GithubCopilotEmbeddingConfig", "NebiusConfig", @@ -648,6 +650,8 @@ _LLM_CONFIGS_IMPORT_MAP = { "GithubCopilotConfig": (".llms.github_copilot.chat.transformation", "GithubCopilotConfig"), "GithubCopilotResponsesAPIConfig": (".llms.github_copilot.responses.transformation", "GithubCopilotResponsesAPIConfig"), "GithubCopilotEmbeddingConfig": (".llms.github_copilot.embedding.transformation", "GithubCopilotEmbeddingConfig"), + "ChatGPTConfig": (".llms.chatgpt.chat.transformation", "ChatGPTConfig"), + "ChatGPTResponsesAPIConfig": (".llms.chatgpt.responses.transformation", "ChatGPTResponsesAPIConfig"), "NebiusConfig": (".llms.nebius.chat.transformation", "NebiusConfig"), "WandbConfig": (".llms.wandb.chat.transformation", "WandbConfig"), "GigaChatConfig": (".llms.gigachat.chat.transformation", "GigaChatConfig"), diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index 6ec49ce0620..5c051797e8b 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -2,10 +2,12 @@ Handler for transforming /chat/completions api requests to litellm.responses requests """ -from typing import TYPE_CHECKING, Any, Coroutine, Union +from typing import TYPE_CHECKING, Any, Coroutine, Optional, Union from typing_extensions import TypedDict +from litellm.types.llms.openai import ResponsesAPIResponse + if TYPE_CHECKING: from litellm import CustomStreamWrapper, LiteLLMLoggingObj, ModelResponse @@ -28,6 +30,71 @@ class ResponsesToCompletionBridgeHandler: super().__init__() self.transformation_handler = LiteLLMResponsesTransformationHandler() + @staticmethod + def _resolve_stream_flag(optional_params: dict, litellm_params: dict) -> bool: + stream = optional_params.get("stream") + if stream is None: + stream = litellm_params.get("stream", False) + return bool(stream) + + @staticmethod + def _coerce_response_object( + response_obj: Any, + hidden_params: Optional[dict], + ) -> "ResponsesAPIResponse": + if isinstance(response_obj, ResponsesAPIResponse): + response = response_obj + elif isinstance(response_obj, dict): + try: + response = ResponsesAPIResponse(**response_obj) + except Exception: + response = ResponsesAPIResponse.model_construct(**response_obj) + else: + raise ValueError("Unexpected responses stream payload") + + if hidden_params: + existing = getattr(response, "_hidden_params", None) + if not isinstance(existing, dict) or not existing: + setattr(response, "_hidden_params", dict(hidden_params)) + else: + for key, value in hidden_params.items(): + existing.setdefault(key, value) + return response + + def _collect_response_from_stream( + self, stream_iter: Any + ) -> "ResponsesAPIResponse": + for _ in stream_iter: + pass + + completed = getattr(stream_iter, "completed_response", None) + response_obj = getattr(completed, "response", None) if completed else None + if response_obj is None: + raise ValueError("Stream ended without a completed response") + + hidden_params = getattr(stream_iter, "_hidden_params", None) + response = self._coerce_response_object(response_obj, hidden_params) + if not isinstance(response, ResponsesAPIResponse): + raise ValueError("Stream completed response is invalid") + return response + + async def _collect_response_from_stream_async( + self, stream_iter: Any + ) -> "ResponsesAPIResponse": + async for _ in stream_iter: + pass + + completed = getattr(stream_iter, "completed_response", None) + response_obj = getattr(completed, "response", None) if completed else None + if response_obj is None: + raise ValueError("Stream ended without a completed response") + + hidden_params = getattr(stream_iter, "_hidden_params", None) + response = self._coerce_response_object(response_obj, hidden_params) + if not isinstance(response, ResponsesAPIResponse): + raise ValueError("Stream completed response is invalid") + return response + def validate_input_kwargs( self, kwargs: dict ) -> ResponsesToCompletionBridgeHandlerInputKwargs: @@ -87,7 +154,6 @@ class ResponsesToCompletionBridgeHandler: from litellm import responses from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - from litellm.types.llms.openai import ResponsesAPIResponse validated_kwargs = self.validate_input_kwargs(kwargs) model = validated_kwargs["model"] @@ -113,6 +179,7 @@ class ResponsesToCompletionBridgeHandler: **request_data, ) + stream = self._resolve_stream_flag(optional_params, litellm_params) if isinstance(result, ResponsesAPIResponse): return self.transformation_handler.transform_response( model=model, @@ -127,6 +194,21 @@ class ResponsesToCompletionBridgeHandler: api_key=kwargs.get("api_key"), json_mode=kwargs.get("json_mode"), ) + elif not stream: + responses_api_response = self._collect_response_from_stream(result) + return self.transformation_handler.transform_response( + model=model, + raw_response=responses_api_response, + model_response=model_response, + logging_obj=logging_obj, + request_data=request_data, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + encoding=kwargs.get("encoding"), + api_key=kwargs.get("api_key"), + json_mode=kwargs.get("json_mode"), + ) else: completion_stream = self.transformation_handler.get_model_response_iterator( streaming_response=result, # type: ignore @@ -146,7 +228,6 @@ class ResponsesToCompletionBridgeHandler: ) -> Union["ModelResponse", "CustomStreamWrapper"]: from litellm import aresponses from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - from litellm.types.llms.openai import ResponsesAPIResponse validated_kwargs = self.validate_input_kwargs(kwargs) model = validated_kwargs["model"] @@ -175,6 +256,7 @@ class ResponsesToCompletionBridgeHandler: aresponses=True, ) + stream = self._resolve_stream_flag(optional_params, litellm_params) if isinstance(result, ResponsesAPIResponse): return self.transformation_handler.transform_response( model=model, @@ -189,6 +271,23 @@ class ResponsesToCompletionBridgeHandler: api_key=kwargs.get("api_key"), json_mode=kwargs.get("json_mode"), ) + elif not stream: + responses_api_response = await self._collect_response_from_stream_async( + result + ) + return self.transformation_handler.transform_response( + model=model, + raw_response=responses_api_response, + model_response=model_response, + logging_obj=logging_obj, + request_data=request_data, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + encoding=kwargs.get("encoding"), + api_key=kwargs.get("api_key"), + json_mode=kwargs.get("json_mode"), + ) else: completion_stream = self.transformation_handler.get_model_response_iterator( streaming_response=result, # type: ignore diff --git a/litellm/constants.py b/litellm/constants.py index 3bdd943481e..f21e7517482 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -415,6 +415,7 @@ LITELLM_CHAT_PROVIDERS = [ "galadriel", "gradient_ai", "github_copilot", # GitHub Copilot Chat API + "chatgpt", # ChatGPT subscription API "novita", "meta_llama", "featherless_ai", @@ -613,6 +614,7 @@ openai_compatible_providers: List = [ "lm_studio", "galadriel", "github_copilot", # GitHub Copilot Chat API + "chatgpt", # ChatGPT subscription API "novita", "meta_llama", "publicai", # PublicAI - JSON-configured provider diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 21d69177336..807b66faec1 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -768,6 +768,14 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 ) = litellm.GithubCopilotConfig()._get_openai_compatible_provider_info( model, api_base, api_key, custom_llm_provider ) + elif custom_llm_provider == "chatgpt": + ( + api_base, + dynamic_api_key, + custom_llm_provider, + ) = litellm.ChatGPTConfig()._get_openai_compatible_provider_info( + model, api_base, api_key, custom_llm_provider + ) elif custom_llm_provider == "novita": api_base = ( api_base diff --git a/litellm/llms/chatgpt/authenticator.py b/litellm/llms/chatgpt/authenticator.py new file mode 100644 index 00000000000..ff053730c35 --- /dev/null +++ b/litellm/llms/chatgpt/authenticator.py @@ -0,0 +1,388 @@ +import base64 +import json +import os +import time +from typing import Any, Dict, Optional + +import httpx + +from litellm._logging import verbose_logger +from litellm.llms.custom_httpx.http_handler import _get_httpx_client + +from .common_utils import ( + CHATGPT_API_BASE, + CHATGPT_AUTH_BASE, + CHATGPT_CLIENT_ID, + CHATGPT_DEVICE_CODE_URL, + CHATGPT_DEVICE_TOKEN_URL, + CHATGPT_DEVICE_VERIFY_URL, + CHATGPT_OAUTH_TOKEN_URL, + GetAccessTokenError, + GetDeviceCodeError, + RefreshAccessTokenError, +) + +TOKEN_EXPIRY_SKEW_SECONDS = 60 +DEVICE_CODE_TIMEOUT_SECONDS = 15 * 60 +DEVICE_CODE_COOLDOWN_SECONDS = 5 * 60 +DEVICE_CODE_POLL_SLEEP_SECONDS = 5 + + +class Authenticator: + def __init__(self) -> None: + self.token_dir = os.getenv( + "CHATGPT_TOKEN_DIR", + os.path.expanduser("~/.config/litellm/chatgpt"), + ) + self.auth_file = os.path.join( + self.token_dir, os.getenv("CHATGPT_AUTH_FILE", "auth.json") + ) + self._ensure_token_dir() + + def get_api_base(self) -> str: + return ( + os.getenv("CHATGPT_API_BASE") + or os.getenv("OPENAI_CHATGPT_API_BASE") + or CHATGPT_API_BASE + ) + + def get_access_token(self) -> str: + auth_data = self._read_auth_file() + if auth_data: + access_token = auth_data.get("access_token") + if access_token and not self._is_token_expired(auth_data, access_token): + return access_token + refresh_token = auth_data.get("refresh_token") + if refresh_token: + try: + refreshed = self._refresh_tokens(refresh_token) + return refreshed["access_token"] + except RefreshAccessTokenError as exc: + verbose_logger.warning( + "ChatGPT refresh token failed, re-login required: %s", exc + ) + + cooldown_remaining = self._get_device_code_cooldown_remaining(auth_data) + if cooldown_remaining > 0: + token = self._wait_for_access_token(cooldown_remaining) + if token: + return token + + tokens = self._login_device_code() + return tokens["access_token"] + + def get_account_id(self) -> Optional[str]: + auth_data = self._read_auth_file() + if not auth_data: + return None + account_id = auth_data.get("account_id") + if account_id: + return account_id + id_token = auth_data.get("id_token") + access_token = auth_data.get("access_token") + derived = self._extract_account_id(id_token or access_token) + if derived: + auth_data["account_id"] = derived + self._write_auth_file(auth_data) + return derived + + def _ensure_token_dir(self) -> None: + if not os.path.exists(self.token_dir): + os.makedirs(self.token_dir, exist_ok=True) + + def _read_auth_file(self) -> Optional[Dict[str, Any]]: + try: + with open(self.auth_file, "r") as f: + return json.load(f) + except IOError: + return None + except json.JSONDecodeError as exc: + verbose_logger.warning("Invalid ChatGPT auth file: %s", exc) + return None + + def _write_auth_file(self, data: Dict[str, Any]) -> None: + try: + with open(self.auth_file, "w") as f: + json.dump(data, f) + except IOError as exc: + verbose_logger.error("Failed to write ChatGPT auth file: %s", exc) + + def _is_token_expired(self, auth_data: Dict[str, Any], access_token: str) -> bool: + expires_at = auth_data.get("expires_at") + if expires_at is None: + expires_at = self._get_expires_at(access_token) + if expires_at: + auth_data["expires_at"] = expires_at + self._write_auth_file(auth_data) + if expires_at is None: + return True + return time.time() >= float(expires_at) - TOKEN_EXPIRY_SKEW_SECONDS + + def _get_expires_at(self, token: str) -> Optional[int]: + claims = self._decode_jwt_claims(token) + exp = claims.get("exp") + if isinstance(exp, (int, float)): + return int(exp) + return None + + def _decode_jwt_claims(self, token: str) -> Dict[str, Any]: + try: + parts = token.split(".") + if len(parts) < 2: + return {} + payload_b64 = parts[1] + payload_b64 += "=" * (-len(payload_b64) % 4) + payload_bytes = base64.urlsafe_b64decode(payload_b64) + return json.loads(payload_bytes.decode("utf-8")) + except Exception: + return {} + + def _extract_account_id(self, token: Optional[str]) -> Optional[str]: + if not token: + return None + claims = self._decode_jwt_claims(token) + auth_claims = claims.get("https://api.openai.com/auth") + if isinstance(auth_claims, dict): + account_id = auth_claims.get("chatgpt_account_id") + if isinstance(account_id, str) and account_id: + return account_id + return None + + def _login_device_code(self) -> Dict[str, str]: + cooldown_remaining = self._get_device_code_cooldown_remaining( + self._read_auth_file() + ) + if cooldown_remaining > 0: + token = self._wait_for_access_token(cooldown_remaining) + if token: + return {"access_token": token} + + device_code = self._request_device_code() + self._record_device_code_request() + print( # noqa: T201 + "Sign in with ChatGPT using device code:\n" + f"1) Visit {CHATGPT_DEVICE_VERIFY_URL}\n" + f"2) Enter code: {device_code['user_code']}\n" + "Device codes are a common phishing target. Never share this code.", + flush=True, + ) + auth_code = self._poll_for_authorization_code(device_code) + tokens = self._exchange_code_for_tokens(auth_code) + auth_data = self._build_auth_record(tokens) + self._write_auth_file(auth_data) + return tokens + + def _request_device_code(self) -> Dict[str, str]: + try: + client = _get_httpx_client() + resp = client.post( + CHATGPT_DEVICE_CODE_URL, + json={"client_id": CHATGPT_CLIENT_ID}, + ) + resp.raise_for_status() + data = resp.json() + except httpx.HTTPStatusError as exc: + raise GetDeviceCodeError( + message=f"Failed to request device code: {exc}", + status_code=exc.response.status_code, + ) + except Exception as exc: + raise GetDeviceCodeError( + message=f"Failed to request device code: {exc}", + status_code=400, + ) + + device_auth_id = data.get("device_auth_id") + user_code = data.get("user_code") or data.get("usercode") + interval = data.get("interval") + if not device_auth_id or not user_code: + raise GetDeviceCodeError( + message=f"Device code response missing fields: {data}", + status_code=400, + ) + return { + "device_auth_id": device_auth_id, + "user_code": user_code, + "interval": str(interval or "5"), + } + + def _poll_for_authorization_code(self, device_code: Dict[str, str]) -> Dict[str, str]: + client = _get_httpx_client() + interval = int(device_code.get("interval", "5")) + start_time = time.time() + while time.time() - start_time < DEVICE_CODE_TIMEOUT_SECONDS: + try: + resp = client.post( + CHATGPT_DEVICE_TOKEN_URL, + json={ + "device_auth_id": device_code["device_auth_id"], + "user_code": device_code["user_code"], + }, + ) + if resp.status_code == 200: + data = resp.json() + if all( + key in data + for key in ( + "authorization_code", + "code_challenge", + "code_verifier", + ) + ): + return data + if resp.status_code in (403, 404): + time.sleep(max(interval, DEVICE_CODE_POLL_SLEEP_SECONDS)) + continue + resp.raise_for_status() + except httpx.HTTPStatusError as exc: + status_code = exc.response.status_code if exc.response else None + if status_code in (403, 404): + time.sleep(max(interval, DEVICE_CODE_POLL_SLEEP_SECONDS)) + continue + raise GetAccessTokenError( + message=f"Polling failed: {exc}", + status_code=exc.response.status_code, + ) + except Exception as exc: + raise GetAccessTokenError( + message=f"Polling failed: {exc}", + status_code=400, + ) + time.sleep(max(interval, DEVICE_CODE_POLL_SLEEP_SECONDS)) + + raise GetAccessTokenError( + message="Timed out waiting for device authorization", + status_code=408, + ) + + def _exchange_code_for_tokens(self, code_data: Dict[str, str]) -> Dict[str, str]: + try: + client = _get_httpx_client() + redirect_uri = f"{CHATGPT_AUTH_BASE}/deviceauth/callback" + body = ( + "grant_type=authorization_code" + f"&code={code_data['authorization_code']}" + f"&redirect_uri={redirect_uri}" + f"&client_id={CHATGPT_CLIENT_ID}" + f"&code_verifier={code_data['code_verifier']}" + ) + resp = client.post( + CHATGPT_OAUTH_TOKEN_URL, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + content=body, + ) + resp.raise_for_status() + data = resp.json() + except httpx.HTTPStatusError as exc: + raise GetAccessTokenError( + message=f"Token exchange failed: {exc}", + status_code=exc.response.status_code, + ) + except Exception as exc: + raise GetAccessTokenError( + message=f"Token exchange failed: {exc}", + status_code=400, + ) + + if not all(key in data for key in ("access_token", "refresh_token", "id_token")): + raise GetAccessTokenError( + message=f"Token exchange response missing fields: {data}", + status_code=400, + ) + return { + "access_token": data["access_token"], + "refresh_token": data["refresh_token"], + "id_token": data["id_token"], + } + + def _refresh_tokens(self, refresh_token: str) -> Dict[str, str]: + try: + client = _get_httpx_client() + resp = client.post( + CHATGPT_OAUTH_TOKEN_URL, + json={ + "client_id": CHATGPT_CLIENT_ID, + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "scope": "openid profile email", + }, + ) + resp.raise_for_status() + data = resp.json() + except httpx.HTTPStatusError as exc: + raise RefreshAccessTokenError( + message=f"Refresh token failed: {exc}", + status_code=exc.response.status_code, + ) + except Exception as exc: + raise RefreshAccessTokenError( + message=f"Refresh token failed: {exc}", + status_code=400, + ) + + access_token = data.get("access_token") + id_token = data.get("id_token") + if not access_token or not id_token: + raise RefreshAccessTokenError( + message=f"Refresh response missing fields: {data}", + status_code=400, + ) + + refreshed = { + "access_token": access_token, + "refresh_token": data.get("refresh_token", refresh_token), + "id_token": id_token, + } + auth_data = self._build_auth_record(refreshed) + self._write_auth_file(auth_data) + return refreshed + + def _build_auth_record(self, tokens: Dict[str, str]) -> Dict[str, Any]: + access_token = tokens.get("access_token") + id_token = tokens.get("id_token") + expires_at = self._get_expires_at(access_token) if access_token else None + account_id = self._extract_account_id(id_token or access_token) + return { + "access_token": access_token, + "refresh_token": tokens.get("refresh_token"), + "id_token": id_token, + "expires_at": expires_at, + "account_id": account_id, + } + + def _get_device_code_cooldown_remaining( + self, auth_data: Optional[Dict[str, Any]] + ) -> float: + if not auth_data: + return 0.0 + requested_at = auth_data.get("device_code_requested_at") + if not isinstance(requested_at, (int, float, str)): + return 0.0 + try: + requested_at = float(requested_at) + except (TypeError, ValueError): + return 0.0 + elapsed = time.time() - requested_at + remaining = DEVICE_CODE_COOLDOWN_SECONDS - elapsed + return max(0.0, remaining) + + def _record_device_code_request(self) -> None: + auth_data = self._read_auth_file() or {} + auth_data["device_code_requested_at"] = time.time() + self._write_auth_file(auth_data) + + def _wait_for_access_token(self, timeout_seconds: float) -> Optional[str]: + deadline = time.time() + timeout_seconds + while time.time() < deadline: + auth_data = self._read_auth_file() + if auth_data: + access_token = auth_data.get("access_token") + if access_token and not self._is_token_expired( + auth_data, access_token + ): + return access_token + sleep_for = min(DEVICE_CODE_POLL_SLEEP_SECONDS, max(0.0, deadline - time.time())) + if sleep_for <= 0: + break + time.sleep(sleep_for) + return None diff --git a/litellm/llms/chatgpt/chat/transformation.py b/litellm/llms/chatgpt/chat/transformation.py new file mode 100644 index 00000000000..2db5eb3c58d --- /dev/null +++ b/litellm/llms/chatgpt/chat/transformation.py @@ -0,0 +1,75 @@ +from typing import List, Optional, Tuple + +from litellm.exceptions import AuthenticationError +from litellm.llms.openai.openai import OpenAIConfig +from litellm.types.llms.openai import AllMessageValues + +from ..authenticator import Authenticator +from ..common_utils import ( + GetAccessTokenError, + ensure_chatgpt_session_id, + get_chatgpt_default_headers, +) + + +class ChatGPTConfig(OpenAIConfig): + def __init__( + self, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + custom_llm_provider: str = "openai", + ) -> None: + super().__init__() + self.authenticator = Authenticator() + + def _get_openai_compatible_provider_info( + self, + model: str, + api_base: Optional[str], + api_key: Optional[str], + custom_llm_provider: str, + ) -> Tuple[Optional[str], Optional[str], str]: + dynamic_api_base = self.authenticator.get_api_base() + try: + dynamic_api_key = self.authenticator.get_access_token() + except GetAccessTokenError as e: + raise AuthenticationError( + model=model, + llm_provider=custom_llm_provider, + message=str(e), + ) + return dynamic_api_base, dynamic_api_key, custom_llm_provider + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + validated_headers = super().validate_environment( + headers, model, messages, optional_params, litellm_params, api_key, api_base + ) + + account_id = self.authenticator.get_account_id() + session_id = ensure_chatgpt_session_id(litellm_params) + default_headers = get_chatgpt_default_headers( + api_key or "", account_id, session_id + ) + return {**default_headers, **validated_headers} + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + optional_params = super().map_openai_params( + non_default_params, optional_params, model, drop_params + ) + optional_params.setdefault("stream", False) + return optional_params diff --git a/litellm/llms/chatgpt/common_utils.py b/litellm/llms/chatgpt/common_utils.py new file mode 100644 index 00000000000..d80487cde24 --- /dev/null +++ b/litellm/llms/chatgpt/common_utils.py @@ -0,0 +1,301 @@ +""" +Constants and helpers for ChatGPT subscription OAuth. +""" +import os +import platform +from typing import Any, Optional, Union +from uuid import uuid4 + +import httpx + +from litellm.llms.base_llm.chat.transformation import BaseLLMException + +# OAuth + API constants (derived from openai/codex) +CHATGPT_AUTH_BASE = "https://auth.openai.com" +CHATGPT_DEVICE_CODE_URL = f"{CHATGPT_AUTH_BASE}/api/accounts/deviceauth/usercode" +CHATGPT_DEVICE_TOKEN_URL = f"{CHATGPT_AUTH_BASE}/api/accounts/deviceauth/token" +CHATGPT_OAUTH_TOKEN_URL = f"{CHATGPT_AUTH_BASE}/oauth/token" +CHATGPT_DEVICE_VERIFY_URL = f"{CHATGPT_AUTH_BASE}/codex/device" +CHATGPT_API_BASE = "https://chatgpt.com/backend-api/codex" +CHATGPT_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann" + +DEFAULT_ORIGINATOR = "codex_cli_rs" +DEFAULT_USER_AGENT = "codex_cli_rs/0.0.0 (Unknown 0; unknown) unknown" +CHATGPT_DEFAULT_INSTRUCTIONS = """You are Codex, based on GPT-5. You are running as a coding agent in the Codex CLI on a user's computer. + +## General + +- When searching for text or files, prefer using `rg` or `rg --files` respectively because `rg` is much faster than alternatives like `grep`. (If the `rg` command is not found, then use alternatives.) + +## Editing constraints + +- Default to ASCII when editing or creating files. Only introduce non-ASCII or other Unicode characters when there is a clear justification and the file already uses them. +- Add succinct code comments that explain what is going on if code is not self-explanatory. You should not add comments like "Assigns the value to the variable", but a brief comment might be useful ahead of a complex code block that the user would otherwise have to spend time parsing out. Usage of these comments should be rare. +- Try to use apply_patch for single file edits, but it is fine to explore other options to make the edit if it does not work well. Do not use apply_patch for changes that are auto-generated (i.e. generating package.json or running a lint or format command like gofmt) or when scripting is more efficient (such as search and replacing a string across a codebase). +- You may be in a dirty git worktree. + * NEVER revert existing changes you did not make unless explicitly requested, since these changes were made by the user. + * If asked to make a commit or code edits and there are unrelated changes to your work or changes that you didn't make in those files, don't revert those changes. + * If the changes are in files you've touched recently, you should read carefully and understand how you can work with the changes rather than reverting them. + * If the changes are in unrelated files, just ignore them and don't revert them. +- Do not amend a commit unless explicitly requested to do so. +- While you are working, you might notice unexpected changes that you didn't make. If this happens, STOP IMMEDIATELY and ask the user how they would like to proceed. +- **NEVER** use destructive commands like `git reset --hard` or `git checkout --` unless specifically requested or approved by the user. + +## Plan tool + +When using the planning tool: +- Skip using the planning tool for straightforward tasks (roughly the easiest 25%). +- Do not make single-step plans. +- When you made a plan, update it after having performed one of the sub-tasks that you shared on the plan. + +## Special user requests + +- If the user makes a simple request (such as asking for the time) which you can fulfill by running a terminal command (such as `date`), you should do so. +- If the user asks for a "review", default to a code review mindset: prioritise identifying bugs, risks, behavioural regressions, and missing tests. Findings must be the primary focus of the response - keep summaries or overviews brief and only after enumerating the issues. Present findings first (ordered by severity with file/line references), follow with open questions or assumptions, and offer a change-summary only as a secondary detail. If no findings are discovered, state that explicitly and mention any residual risks or testing gaps. + +## Frontend tasks +When doing frontend design tasks, avoid collapsing into "AI slop" or safe, average-looking layouts. +Aim for interfaces that feel intentional, bold, and a bit surprising. +- Typography: Use expressive, purposeful fonts and avoid default stacks (Inter, Roboto, Arial, system). +- Color & Look: Choose a clear visual direction; define CSS variables; avoid purple-on-white defaults. No purple bias or dark mode bias. +- Motion: Use a few meaningful animations (page-load, staggered reveals) instead of generic micro-motions. +- Background: Don't rely on flat, single-color backgrounds; use gradients, shapes, or subtle patterns to build atmosphere. +- Overall: Avoid boilerplate layouts and interchangeable UI patterns. Vary themes, type families, and visual languages across outputs. +- Ensure the page loads properly on both desktop and mobile + +Exception: If working within an existing website or design system, preserve the established patterns, structure, and visual language. + +## Presenting your work and final message + +You are producing plain text that will later be styled by the CLI. Follow these rules exactly. Formatting should make results easy to scan, but not feel mechanical. Use judgment to decide how much structure adds value. + +- Default: be very concise; friendly coding teammate tone. +- Ask only when needed; suggest ideas; mirror the user's style. +- For substantial work, summarize clearly; follow final-answer formatting. +- Skip heavy formatting for simple confirmations. +- Don't dump large files you've written; reference paths only. +- No "save/copy this file" - User is on the same machine. +- Offer logical next steps (tests, commits, build) briefly; add verify steps if you couldn't do something. +- For code changes: + * Lead with a quick explanation of the change, and then give more details on the context covering where and why a change was made. Do not start this explanation with "summary", just jump right in. + * If there are natural next steps the user may want to take, suggest them at the end of your response. Do not make suggestions if there are no natural next steps. + * When suggesting multiple options, use numeric lists for the suggestions so the user can quickly respond with a single number. +- The user does not command execution outputs. When asked to show the output of a command (e.g. `git show`), relay the important details in your answer or summarize the key lines so the user understands the result. + +### Final answer structure and style guidelines + +- Plain text; CLI handles styling. Use structure only when it helps scanability. +- Headers: optional; short Title Case (1-3 words) wrapped in **...**; no blank line before the first bullet; add only if they truly help. +- Bullets: use - ; merge related points; keep to one line when possible; 4-6 per list ordered by importance; keep phrasing consistent. +- Monospace: backticks for commands/paths/env vars/code ids and inline examples; use for literal keyword bullets; never combine with **. +- Code samples or multi-line snippets should be wrapped in fenced code blocks; include an info string as often as possible. +- Structure: group related bullets; order sections general -> specific -> supporting; for subsections, start with a bolded keyword bullet, then items; match complexity to the task. +- Tone: collaborative, concise, factual; present tense, active voice; self-contained; no "above/below"; parallel wording. +- Don'ts: no nested bullets/hierarchies; no ANSI codes; don't cram unrelated keywords; keep keyword lists short--wrap/reformat if long; avoid naming formatting styles in answers. +- Adaptation: code explanations -> precise, structured with code refs; simple tasks -> lead with outcome; big changes -> logical walkthrough + rationale + next actions; casual one-offs -> plain sentences, no headers/bullets. +- File References: When referencing files in your response follow the below rules: + * Use inline code to make file paths clickable. + * Each reference should have a stand alone path. Even if it's the same file. + * Accepted: absolute, workspace-relative, a/ or b/ diff prefixes, or bare filename/suffix. + * Optionally include line/column (1-based): :line[:column] or #Lline[Ccolumn] (column defaults to 1). + * Do not use URIs like file://, vscode://, or https://. + * Do not provide range of lines + * Examples: src/app.ts, src/app.ts:42, b/server/index.js#L10, C:\\repo\\project\\main.rs:12:5 +""" + + +class ChatGPTAuthError(BaseLLMException): + def __init__( + self, + status_code, + message, + request: Optional[httpx.Request] = None, + response: Optional[httpx.Response] = None, + headers: Optional[Union[httpx.Headers, dict]] = None, + body: Optional[dict] = None, + ): + super().__init__( + status_code=status_code, + message=message, + request=request, + response=response, + headers=headers, + body=body, + ) + + +class GetDeviceCodeError(ChatGPTAuthError): + pass + + +class GetAccessTokenError(ChatGPTAuthError): + pass + + +class RefreshAccessTokenError(ChatGPTAuthError): + pass + + +def _safe_header_value(value: str) -> str: + if not value: + return "" + return "".join(ch if 32 <= ord(ch) <= 126 else "_" for ch in value) + + +def _sanitize_user_agent_token(value: str) -> str: + if not value: + return "" + return "".join( + ch if (ch.isalnum() or ch in "-_./") else "_" for ch in value + ) + + +def _terminal_user_agent() -> str: + term_program = os.getenv("TERM_PROGRAM") + if term_program: + version = os.getenv("TERM_PROGRAM_VERSION") + token = f"{term_program}/{version}" if version else term_program + return _sanitize_user_agent_token(token) or "unknown" + + wezterm_version = os.getenv("WEZTERM_VERSION") + if wezterm_version is not None: + token = ( + f"WezTerm/{wezterm_version}" if wezterm_version else "WezTerm" + ) + return _sanitize_user_agent_token(token) or "WezTerm" + + if ( + os.getenv("ITERM_SESSION_ID") + or os.getenv("ITERM_PROFILE") + or os.getenv("ITERM_PROFILE_NAME") + ): + return "iTerm.app" + + if os.getenv("TERM_SESSION_ID"): + return "Apple_Terminal" + + if os.getenv("KITTY_WINDOW_ID") or "kitty" in (os.getenv("TERM") or ""): + return "kitty" + + if os.getenv("ALACRITTY_SOCKET") or os.getenv("TERM") == "alacritty": + return "Alacritty" + + konsole_version = os.getenv("KONSOLE_VERSION") + if konsole_version is not None: + token = ( + f"Konsole/{konsole_version}" if konsole_version else "Konsole" + ) + return _sanitize_user_agent_token(token) or "Konsole" + + if os.getenv("GNOME_TERMINAL_SCREEN"): + return "gnome-terminal" + + vte_version = os.getenv("VTE_VERSION") + if vte_version is not None: + token = f"VTE/{vte_version}" if vte_version else "VTE" + return _sanitize_user_agent_token(token) or "VTE" + + if os.getenv("WT_SESSION"): + return "WindowsTerminal" + + term = os.getenv("TERM") + if term: + return _sanitize_user_agent_token(term) or "unknown" + + return "unknown" + + +def _get_litellm_version() -> str: + try: + from importlib.metadata import version + + return version("litellm") + except Exception: + return "0.0.0" + + +def get_chatgpt_originator() -> str: + originator = os.getenv("CHATGPT_ORIGINATOR") or DEFAULT_ORIGINATOR + return _safe_header_value(originator) or DEFAULT_ORIGINATOR + + +def get_chatgpt_user_agent(originator: str) -> str: + override = os.getenv("CHATGPT_USER_AGENT") + if override: + return _safe_header_value(override) or DEFAULT_USER_AGENT + version = _get_litellm_version() + os_type = platform.system() or "Unknown" + os_version = platform.release() or "0" + arch = platform.machine() or "unknown" + terminal_ua = _terminal_user_agent() + suffix = os.getenv("CHATGPT_USER_AGENT_SUFFIX", "").strip() + suffix = f" ({suffix})" if suffix else "" + candidate = ( + f"{originator}/{version} ({os_type} {os_version}; {arch}) {terminal_ua}{suffix}" + ) + return _safe_header_value(candidate) or DEFAULT_USER_AGENT + + +def get_chatgpt_default_headers( + access_token: str, + account_id: Optional[str], + session_id: Optional[str] = None, +) -> dict: + originator = get_chatgpt_originator() + user_agent = get_chatgpt_user_agent(originator) + headers = { + "Authorization": f"Bearer {access_token}", + "content-type": "application/json", + "accept": "text/event-stream", + "originator": originator, + "user-agent": user_agent, + } + if session_id: + headers["session_id"] = session_id + if account_id: + headers["ChatGPT-Account-Id"] = account_id + return headers + + +def get_chatgpt_default_instructions() -> str: + return os.getenv("CHATGPT_DEFAULT_INSTRUCTIONS") or CHATGPT_DEFAULT_INSTRUCTIONS + + +def _normalize_litellm_params(litellm_params: Optional[Any]) -> dict: + if litellm_params is None: + return {} + if isinstance(litellm_params, dict): + return litellm_params + if hasattr(litellm_params, "model_dump"): + try: + return litellm_params.model_dump() + except Exception: + return {} + if hasattr(litellm_params, "dict"): + try: + return litellm_params.dict() + except Exception: + return {} + return {} + + +def get_chatgpt_session_id(litellm_params: Optional[Any]) -> Optional[str]: + params = _normalize_litellm_params(litellm_params) + for key in ("litellm_session_id", "session_id"): + value = params.get(key) + if value: + return str(value) + metadata = params.get("metadata") + if isinstance(metadata, dict): + value = metadata.get("session_id") + if value: + return str(value) + for key in ("litellm_trace_id", "litellm_call_id"): + value = params.get(key) + if value: + return str(value) + return None + + +def ensure_chatgpt_session_id(litellm_params: Optional[Any]) -> str: + return get_chatgpt_session_id(litellm_params) or str(uuid4()) diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py new file mode 100644 index 00000000000..0ce24f63a89 --- /dev/null +++ b/litellm/llms/chatgpt/responses/transformation.py @@ -0,0 +1,191 @@ +import json +from typing import Any, Optional + +from litellm.exceptions import AuthenticationError +from litellm.constants import STREAM_SSE_DONE_STRING +from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.llms.openai.common_utils import OpenAIError +from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + _safe_convert_created_field, +) +from litellm.types.llms.openai import ( + ResponsesAPIResponse, + ResponsesAPIStreamEvents, +) +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.utils import CustomStreamWrapper + +from ..authenticator import Authenticator +from ..common_utils import ( + CHATGPT_API_BASE, + GetAccessTokenError, + ensure_chatgpt_session_id, + get_chatgpt_default_headers, + get_chatgpt_default_instructions, +) + + +class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): + def __init__(self) -> None: + super().__init__() + self.authenticator = Authenticator() + + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.CHATGPT + + def validate_environment( + self, + headers: dict, + model: str, + litellm_params: Optional[GenericLiteLLMParams], + ) -> dict: + try: + access_token = self.authenticator.get_access_token() + except GetAccessTokenError as e: + raise AuthenticationError( + model=model, + llm_provider="chatgpt", + message=str(e), + ) + + account_id = self.authenticator.get_account_id() + session_id = ensure_chatgpt_session_id(litellm_params) + default_headers = get_chatgpt_default_headers( + access_token, account_id, session_id + ) + return {**default_headers, **headers} + + def transform_responses_api_request( + self, + model: str, + input: Any, + response_api_optional_request_params: dict, + litellm_params: GenericLiteLLMParams, + headers: dict, + ) -> dict: + request = super().transform_responses_api_request( + model, + input, + response_api_optional_request_params, + litellm_params, + headers, + ) + request.pop("max_output_tokens", None) + request.pop("max_tokens", None) + request.pop("max_completion_tokens", None) + request.pop("metadata", None) + base_instructions = get_chatgpt_default_instructions() + existing_instructions = request.get("instructions") + if existing_instructions: + if base_instructions not in existing_instructions: + request["instructions"] = ( + f"{base_instructions}\n\n{existing_instructions}" + ) + else: + request["instructions"] = base_instructions + request["store"] = False + request["stream"] = True + include = list(request.get("include") or []) + if "reasoning.encrypted_content" not in include: + include.append("reasoning.encrypted_content") + request["include"] = include + return request + + def transform_response_api_response( + self, + model: str, + raw_response: Any, + logging_obj: Any, + ): + content_type = (raw_response.headers or {}).get("content-type", "") + body_text = raw_response.text or "" + if "text/event-stream" not in content_type.lower(): + trimmed_body = body_text.lstrip() + if not ( + trimmed_body.startswith("event:") + or trimmed_body.startswith("data:") + or "\nevent:" in body_text + or "\ndata:" in body_text + ): + return super().transform_response_api_response( + model=model, + raw_response=raw_response, + logging_obj=logging_obj, + ) + + logging_obj.post_call( + original_response=raw_response.text, + additional_args={"complete_input_dict": {}}, + ) + + completed_response = None + error_message = None + for chunk in body_text.splitlines(): + stripped_chunk = CustomStreamWrapper._strip_sse_data_from_chunk(chunk) + if not stripped_chunk: + continue + stripped_chunk = stripped_chunk.strip() + if not stripped_chunk: + continue + if stripped_chunk == STREAM_SSE_DONE_STRING: + break + try: + parsed_chunk = json.loads(stripped_chunk) + except json.JSONDecodeError: + continue + if not isinstance(parsed_chunk, dict): + continue + event_type = parsed_chunk.get("type") + if event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED: + response_payload = parsed_chunk.get("response") + if isinstance(response_payload, dict): + response_payload = dict(response_payload) + if "created_at" in response_payload: + response_payload["created_at"] = _safe_convert_created_field( + response_payload["created_at"] + ) + try: + completed_response = ResponsesAPIResponse(**response_payload) + except Exception: + completed_response = ResponsesAPIResponse.model_construct( + **response_payload + ) + break + if event_type in ( + ResponsesAPIStreamEvents.RESPONSE_FAILED, + ResponsesAPIStreamEvents.ERROR, + ): + error_obj = parsed_chunk.get("error") or ( + parsed_chunk.get("response") or {} + ).get("error") + if error_obj is not None: + if isinstance(error_obj, dict): + error_message = error_obj.get("message") or str(error_obj) + else: + error_message = str(error_obj) + + if completed_response is None: + raise OpenAIError( + message=error_message or raw_response.text, + status_code=raw_response.status_code, + ) + + raw_headers = dict(raw_response.headers) + processed_headers = process_response_headers(raw_headers) + if not hasattr(completed_response, "_hidden_params"): + setattr(completed_response, "_hidden_params", {}) + completed_response._hidden_params["additional_headers"] = processed_headers + completed_response._hidden_params["headers"] = raw_headers + return completed_response + + def get_complete_url( + self, + api_base: Optional[str], + litellm_params: dict, + ) -> str: + api_base = api_base or self.authenticator.get_api_base() or CHATGPT_API_BASE + api_base = api_base.rstrip("/") + return f"{api_base}/responses" diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 470d598a25f..cd4187ff779 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -15932,6 +15932,64 @@ "max_tokens": 8191, "mode": "embedding" }, + "chatgpt/gpt-5.2-codex": { + "litellm_provider": "chatgpt", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.1-codex-max": { + "litellm_provider": "chatgpt", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.1-codex-mini": { + "litellm_provider": "chatgpt", + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.2": { + "litellm_provider": "chatgpt", + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/messages", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, "gigachat/GigaChat-2-Lite": { "input_cost_per_token": 0.0, "litellm_provider": "gigachat", diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 8c30e4d7e80..fd83b7e79f0 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2958,6 +2958,7 @@ GenericBudgetConfigType = Dict[str, BudgetConfig] class LlmProviders(str, Enum): OPENAI = "openai" + CHATGPT = "chatgpt" OPENAI_LIKE = "openai_like" # embedding only JINA_AI = "jina_ai" XAI = "xai" diff --git a/litellm/utils.py b/litellm/utils.py index ac194e4f339..f5b5c50468f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7659,6 +7659,7 @@ class ProviderConfigManager: LlmProviders.GITHUB: (lambda: litellm.GithubChatConfig(), False), LlmProviders.COMPACTIFAI: (lambda: litellm.CompactifAIChatConfig(), False), LlmProviders.GITHUB_COPILOT: (lambda: litellm.GithubCopilotConfig(), False), + LlmProviders.CHATGPT: (lambda: litellm.ChatGPTConfig(), False), LlmProviders.GIGACHAT: (lambda: litellm.GigaChatConfig(), False), LlmProviders.RAGFLOW: (lambda: litellm.RAGFlowConfig(), False), LlmProviders.CUSTOM: (lambda: litellm.OpenAILikeChatConfig(), False), @@ -8028,6 +8029,8 @@ class ProviderConfigManager: return litellm.XAIResponsesAPIConfig() elif litellm.LlmProviders.GITHUB_COPILOT == provider: return litellm.GithubCopilotResponsesAPIConfig() + elif litellm.LlmProviders.CHATGPT == provider: + return litellm.ChatGPTResponsesAPIConfig() elif litellm.LlmProviders.LITELLM_PROXY == provider: return litellm.LiteLLMProxyResponsesAPIConfig() elif litellm.LlmProviders.MANUS == provider: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 470d598a25f..c846edb85cb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -15932,6 +15932,63 @@ "max_tokens": 8191, "mode": "embedding" }, + "chatgpt/gpt-5.2-codex": { + "litellm_provider": "chatgpt", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.2": { + "litellm_provider": "chatgpt", + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.1-codex-max": { + "litellm_provider": "chatgpt", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, + "chatgpt/gpt-5.1-codex-mini": { + "litellm_provider": "chatgpt", + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "responses", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, "gigachat/GigaChat-2-Lite": { "input_cost_per_token": 0.0, "litellm_provider": "gigachat", diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index db11ec8a0de..b7892a67a17 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -919,6 +919,24 @@ "interactions": true } }, + "chatgpt": { + "display_name": "ChatGPT Subscription (`chatgpt`)", + "url": "https://docs.litellm.ai/docs/providers/chatgpt", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false, + "interactions": false + } + }, "github": { "display_name": "GitHub Models (`github`)", "url": "https://docs.litellm.ai/docs/providers/github", diff --git a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py new file mode 100644 index 00000000000..bec748d8dc8 --- /dev/null +++ b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -0,0 +1,125 @@ +""" +Tests for ChatGPT subscription Responses API transformation + +Source: litellm/llms/chatgpt/responses/transformation.py +""" +import json +import os +import sys +from unittest.mock import MagicMock, patch + +import httpx + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager +from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPIConfig + + +class TestChatGPTResponsesAPITransformation: + def test_chatgpt_provider_config_registration(self): + config = ProviderConfigManager.get_provider_responses_api_config( + model="chatgpt/gpt-5.2", + provider=LlmProviders.CHATGPT, + ) + + assert config is not None + assert isinstance(config, ChatGPTResponsesAPIConfig) + assert config.custom_llm_provider == LlmProviders.CHATGPT + + @patch("litellm.llms.chatgpt.responses.transformation.Authenticator") + def test_chatgpt_responses_endpoint_url(self, mock_authenticator_class): + mock_auth_instance = MagicMock() + mock_auth_instance.get_api_base.return_value = "https://chatgpt.example.com" + mock_authenticator_class.return_value = mock_auth_instance + + config = ChatGPTResponsesAPIConfig() + + url = config.get_complete_url(api_base=None, litellm_params={}) + assert url == "https://chatgpt.example.com/responses" + + custom_url = config.get_complete_url( + api_base="https://custom.chatgpt.com", litellm_params={} + ) + assert custom_url == "https://custom.chatgpt.com/responses" + + url_with_slash = config.get_complete_url( + api_base="https://chatgpt.example.com/", litellm_params={} + ) + assert url_with_slash == "https://chatgpt.example.com/responses" + + @patch("litellm.llms.chatgpt.responses.transformation.Authenticator") + def test_validate_environment_headers(self, mock_authenticator_class): + mock_auth_instance = MagicMock() + mock_auth_instance.get_access_token.return_value = "access-123" + mock_auth_instance.get_account_id.return_value = "acct-123" + mock_authenticator_class.return_value = mock_auth_instance + + config = ChatGPTResponsesAPIConfig() + litellm_params = GenericLiteLLMParams(litellm_session_id="session-123") + headers = config.validate_environment( + headers={"originator": "custom-origin"}, + model="gpt-5.2", + litellm_params=litellm_params, + ) + + assert headers["Authorization"] == "Bearer access-123" + assert headers["ChatGPT-Account-Id"] == "acct-123" + assert headers["originator"] == "custom-origin" + assert headers["content-type"] == "application/json" + assert headers["accept"] == "text/event-stream" + assert headers["session_id"] == "session-123" + + def test_chatgpt_forces_streaming_and_reasoning_include(self): + config = ChatGPTResponsesAPIConfig() + request = config.transform_responses_api_request( + model="chatgpt/gpt-5.2-codex", + input="hi", + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert request["stream"] is True + assert "reasoning.encrypted_content" in request["include"] + assert request["instructions"].startswith( + "You are Codex, based on GPT-5." + ) + + def test_chatgpt_non_stream_sse_response_parsing(self): + config = ChatGPTResponsesAPIConfig() + response_payload = { + "id": "resp_test", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.2-codex", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Hello!"}], + } + ], + } + sse_body = "\n".join( + [ + f"data: {json.dumps({'type': 'response.completed', 'response': response_payload})}", + "data: [DONE]", + "", + ] + ) + raw_response = httpx.Response( + 200, headers={"content-type": "text/event-stream"}, text=sse_body + ) + logging_obj = MagicMock() + + parsed = config.transform_response_api_response( + model="chatgpt/gpt-5.2-codex", + raw_response=raw_response, + logging_obj=logging_obj, + ) + + assert parsed.output_text == "Hello!" diff --git a/tests/test_litellm/llms/chatgpt/test_chatgpt_authenticator.py b/tests/test_litellm/llms/chatgpt/test_chatgpt_authenticator.py new file mode 100644 index 00000000000..5a4b58e159a --- /dev/null +++ b/tests/test_litellm/llms/chatgpt/test_chatgpt_authenticator.py @@ -0,0 +1,68 @@ +import base64 +import json +import time +from unittest.mock import mock_open, patch + +import pytest + +from litellm.llms.chatgpt.authenticator import Authenticator + + +def _make_jwt(payload: dict) -> str: + header = {"alg": "none", "typ": "JWT"} + + def _b64(obj: dict) -> str: + raw = json.dumps(obj, separators=(",", ":")).encode("utf-8") + return base64.urlsafe_b64encode(raw).decode("utf-8").rstrip("=") + + return f"{_b64(header)}.{_b64(payload)}." + + +class TestChatGPTAuthenticator: + @pytest.fixture + def authenticator(self): + with patch("os.path.exists", return_value=True): + return Authenticator() + + def test_get_access_token_from_file(self, authenticator): + future_time = time.time() + 3600 + auth_data = json.dumps({"access_token": "token-123", "expires_at": future_time}) + + with patch("builtins.open", mock_open(read_data=auth_data)): + token = authenticator.get_access_token() + assert token == "token-123" + + def test_get_access_token_refresh(self, authenticator): + past_time = time.time() - 10 + auth_data = json.dumps( + { + "access_token": "token-old", + "refresh_token": "refresh-123", + "expires_at": past_time, + } + ) + refreshed = { + "access_token": "token-new", + "refresh_token": "refresh-123", + "id_token": "id-123", + } + + with patch("builtins.open", mock_open(read_data=auth_data)), patch.object( + authenticator, "_refresh_tokens", return_value=refreshed + ): + token = authenticator.get_access_token() + assert token == "token-new" + + def test_get_account_id_from_id_token(self, authenticator): + id_token = _make_jwt( + {"https://api.openai.com/auth": {"chatgpt_account_id": "acct-123"}} + ) + auth_data = json.dumps({"id_token": id_token}) + + with patch("builtins.open", mock_open(read_data=auth_data)), patch.object( + authenticator, "_write_auth_file" + ) as mock_write: + account_id = authenticator.get_account_id() + assert account_id == "acct-123" + mock_write.assert_called_once() + assert mock_write.call_args[0][0]["account_id"] == "acct-123" diff --git a/tests/test_litellm/test_responses_api_bridge_non_stream.py b/tests/test_litellm/test_responses_api_bridge_non_stream.py new file mode 100644 index 00000000000..8f72debfc5d --- /dev/null +++ b/tests/test_litellm/test_responses_api_bridge_non_stream.py @@ -0,0 +1,44 @@ +from litellm.completion_extras.litellm_responses_transformation.handler import ( + ResponsesToCompletionBridgeHandler, +) +from litellm.types.llms.openai import ResponsesAPIResponse + + +class _CompletedEvent: + def __init__(self, response): + self.response = response + + +class _FakeResponsesStream: + def __init__(self, response): + self._emitted = False + self._response = response + self.completed_response = None + self._hidden_params = {"headers": {"x-test": "1"}} + + def __iter__(self): + return self + + def __next__(self): + if not self._emitted: + self._emitted = True + self.completed_response = _CompletedEvent(self._response) + return {"type": "response.completed"} + raise StopIteration + + +def test_should_collect_response_from_stream(): + handler = ResponsesToCompletionBridgeHandler() + response = ResponsesAPIResponse.model_construct( + id="resp-1", + created_at=0, + output=[], + object="response", + model="gpt-5.2", + ) + stream = _FakeResponsesStream(response) + + collected = handler._collect_response_from_stream(stream) + + assert collected.id == "resp-1" + assert collected._hidden_params.get("headers") == {"x-test": "1"} From 07fbd77c9164cbefd5d8af583a236f773f224427 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Mon, 19 Jan 2026 19:23:23 +0530 Subject: [PATCH 08/58] fix(logging): prevent duplicate StandardLoggingPayload logs (#19325) --- litellm/litellm_core_utils/litellm_logging.py | 66 +++++++++++++------ 1 file changed, 46 insertions(+), 20 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index bc5faf962c2..9cef21bfdf6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1897,6 +1897,14 @@ class Logging(LiteLLMLoggingBaseClass): status="success", standard_built_in_tools_params=self.standard_built_in_tools_params, ) + if ( + standard_logging_payload := self.model_call_details.get( + "standard_logging_object" + ) + ) is not None: + # Only emit for sync requests (async_success_handler handles async) + if is_sync_request: + emit_standard_logging_payload(standard_logging_payload) callbacks = self.get_combined_callback_list( dynamic_success_callbacks=self.dynamic_success_callbacks, global_callbacks=litellm.success_callback, @@ -2190,10 +2198,7 @@ class Logging(LiteLLMLoggingBaseClass): print_verbose=print_verbose, ) - if ( - callback == "openmeter" - and is_sync_request - ): + if callback == "openmeter" and is_sync_request: global openMeterLogger if openMeterLogger is None: print_verbose("Instantiates openmeter client") @@ -2405,6 +2410,14 @@ class Logging(LiteLLMLoggingBaseClass): 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( + "standard_logging_object" + ) + ) is not None: + emit_standard_logging_payload(standard_logging_payload) callbacks = self.get_combined_callback_list( dynamic_success_callbacks=self.dynamic_async_success_callbacks, global_callbacks=litellm._async_success_callback, @@ -2799,8 +2812,7 @@ class Logging(LiteLLMLoggingBaseClass): callback_func=callback, ) if ( - isinstance(callback, CustomLogger) - and is_sync_request + isinstance(callback, CustomLogger) and is_sync_request ): # custom logger class callback.log_failure_event( start_time=start_time, @@ -3743,10 +3755,13 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 OpenTelemetry, OpenTelemetryConfig, ) - logfire_base_url = os.getenv("LOGFIRE_BASE_URL", "https://logfire-api.pydantic.dev") + + logfire_base_url = os.getenv( + "LOGFIRE_BASE_URL", "https://logfire-api.pydantic.dev" + ) otel_config = OpenTelemetryConfig( exporter="otlp_http", - endpoint = f"{logfire_base_url.rstrip('/')}/v1/traces", + endpoint=f"{logfire_base_url.rstrip('/')}/v1/traces", headers=f"Authorization={os.getenv('LOGFIRE_TOKEN')}", ) for callback in _in_memory_loggers: @@ -4342,32 +4357,38 @@ class StandardLoggingPayloadSetup: def merge_litellm_metadata(litellm_params: dict) -> dict: """ Merge both litellm_metadata and metadata from litellm_params. - + litellm_metadata contains model-related fields, metadata contains user API key fields. We need both for complete standard logging payload. - + Args: litellm_params: Dictionary containing metadata and litellm_metadata - + Returns: dict: Merged metadata with user API key fields taking precedence """ merged_metadata: dict = {} - + # Start with metadata (user API key fields) - but skip non-serializable objects - if litellm_params.get("metadata") and isinstance(litellm_params.get("metadata"), dict): + if litellm_params.get("metadata") and isinstance( + litellm_params.get("metadata"), dict + ): for key, value in litellm_params["metadata"].items(): # Skip non-serializable objects like UserAPIKeyAuth if key == "user_api_key_auth": continue merged_metadata[key] = value - + # Then merge litellm_metadata (model-related fields) - this will NOT overwrite existing keys - if litellm_params.get("litellm_metadata") and isinstance(litellm_params.get("litellm_metadata"), dict): + if litellm_params.get("litellm_metadata") and isinstance( + litellm_params.get("litellm_metadata"), dict + ): for key, value in litellm_params["litellm_metadata"].items(): - if key not in merged_metadata: # Don't overwrite existing keys from metadata + if ( + key not in merged_metadata + ): # Don't overwrite existing keys from metadata merged_metadata[key] = value - + return merged_metadata @staticmethod @@ -4810,7 +4831,9 @@ class StandardLoggingPayloadSetup: """ Extract additional header tags for spend tracking based on config. """ - extra_headers: List[str] = getattr(litellm, "extra_spend_tag_headers", None) or [] + extra_headers: List[str] = ( + getattr(litellm, "extra_spend_tag_headers", None) or [] + ) if not extra_headers: return None @@ -4959,7 +4982,9 @@ def get_standard_logging_object_payload( proxy_server_request = litellm_params.get("proxy_server_request") or {} # Merge both litellm_metadata and metadata to get complete metadata - metadata: dict = StandardLoggingPayloadSetup.merge_litellm_metadata(litellm_params) + metadata: dict = StandardLoggingPayloadSetup.merge_litellm_metadata( + litellm_params + ) completion_start_time = kwargs.get("completion_start_time", end_time) call_type = kwargs.get("call_type") @@ -5129,7 +5154,8 @@ def get_standard_logging_object_payload( standard_built_in_tools_params=standard_built_in_tools_params, ) - emit_standard_logging_payload(payload) + # emit_standard_logging_payload(payload) - Moved to success_handler to prevent double emitting + return payload except Exception as e: verbose_logger.exception( From fe92f4af9cb8b5e22c923d41902a4b9ca64acefa Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Mon, 19 Jan 2026 19:23:47 +0530 Subject: [PATCH 09/58] fix(langfuse_otel): ignore service logs and fix callback shadowing (#19298) * fix(langfuse_otel): ignore service logs and fix callback shadowing * add test cases for service logger --- litellm/_service_logger.py | 39 +++++---- .../integrations/langfuse/langfuse_otel.py | 67 ++++++++++++--- tests/test_service_logger_otel.py | 85 +++++++++++++++++++ 3 files changed, 164 insertions(+), 27 deletions(-) create mode 100644 tests/test_service_logger_otel.py diff --git a/litellm/_service_logger.py b/litellm/_service_logger.py index 3128f02f409..b67d0d86063 100644 --- a/litellm/_service_logger.py +++ b/litellm/_service_logger.py @@ -145,16 +145,19 @@ class ServiceLogging(CustomLogger): event_metadata=event_metadata, ) elif callback == "otel" or isinstance(callback, OpenTelemetry): - from litellm.proxy.proxy_server import open_telemetry_logger + _otel_logger_to_use: Optional[OpenTelemetry] = None + if isinstance(callback, OpenTelemetry): + _otel_logger_to_use = callback + else: + from litellm.proxy.proxy_server import open_telemetry_logger - await self.init_otel_logger_if_none() + if open_telemetry_logger is not None and isinstance( + open_telemetry_logger, OpenTelemetry + ): + _otel_logger_to_use = open_telemetry_logger - if ( - parent_otel_span is not None - and open_telemetry_logger is not None - and isinstance(open_telemetry_logger, OpenTelemetry) - ): - await self.otel_logger.async_service_success_hook( + if _otel_logger_to_use is not None and parent_otel_span is not None: + await _otel_logger_to_use.async_service_success_hook( payload=payload, parent_otel_span=parent_otel_span, start_time=start_time, @@ -253,20 +256,24 @@ class ServiceLogging(CustomLogger): event_metadata=event_metadata, ) elif callback == "otel" or isinstance(callback, OpenTelemetry): - from litellm.proxy.proxy_server import open_telemetry_logger + _otel_logger_to_use: Optional[OpenTelemetry] = None + if isinstance(callback, OpenTelemetry): + _otel_logger_to_use = callback + else: + from litellm.proxy.proxy_server import open_telemetry_logger - await self.init_otel_logger_if_none() + if open_telemetry_logger is not None and isinstance( + open_telemetry_logger, OpenTelemetry + ): + _otel_logger_to_use = open_telemetry_logger if not isinstance(error, str): error = str(error) - if ( - parent_otel_span is not None - and open_telemetry_logger is not None - and isinstance(open_telemetry_logger, OpenTelemetry) - ): - await self.otel_logger.async_service_success_hook( + if _otel_logger_to_use is not None and parent_otel_span is not None: + await _otel_logger_to_use.async_service_failure_hook( payload=payload, + error=error, parent_otel_span=parent_otel_span, start_time=start_time, end_time=end_time, diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index 6992ea17cc8..08493a0e8ec 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/litellm/integrations/langfuse/langfuse_otel.py @@ -156,7 +156,11 @@ class LangfuseOtelLogger(OpenTelemetry): "arguments": arguments_obj, } transformed_tool_calls.append(langfuse_tool_call) - safe_set_attribute(span, LangfuseSpanAttributes.OBSERVATION_OUTPUT.value, safe_dumps(transformed_tool_calls)) + safe_set_attribute( + span, + LangfuseSpanAttributes.OBSERVATION_OUTPUT.value, + safe_dumps(transformed_tool_calls), + ) else: output_data = {} if message.get("role"): @@ -164,7 +168,11 @@ class LangfuseOtelLogger(OpenTelemetry): if message.get("content") is not None: output_data["content"] = message.get("content") if output_data: - safe_set_attribute(span, LangfuseSpanAttributes.OBSERVATION_OUTPUT.value, safe_dumps(output_data)) + safe_set_attribute( + span, + LangfuseSpanAttributes.OBSERVATION_OUTPUT.value, + safe_dumps(output_data), + ) output = response_obj.get("output", []) if output: @@ -175,15 +183,28 @@ class LangfuseOtelLogger(OpenTelemetry): if item_type == "reasoning" and hasattr(item, "summary"): for summary in item.summary: if hasattr(summary, "text"): - output_items_data.append({"role": "reasoning_summary", "content": summary.text}) + output_items_data.append( + { + "role": "reasoning_summary", + "content": summary.text, + } + ) elif item_type == "message": - output_items_data.append({ - "role": getattr(item, "role", "assistant"), - "content": getattr(getattr(item, "content", [{}])[0], "text", "") - }) + output_items_data.append( + { + "role": getattr(item, "role", "assistant"), + "content": getattr( + getattr(item, "content", [{}])[0], "text", "" + ), + } + ) elif item_type == "function_call": arguments_str = getattr(item, "arguments", "{}") - arguments_obj = json.loads(arguments_str) if isinstance(arguments_str, str) else arguments_str + arguments_obj = ( + json.loads(arguments_str) + if isinstance(arguments_str, str) + else arguments_str + ) langfuse_tool_call = { "id": getattr(item, "id", ""), "name": getattr(item, "name", ""), @@ -193,7 +214,11 @@ class LangfuseOtelLogger(OpenTelemetry): } output_items_data.append(langfuse_tool_call) if output_items_data: - safe_set_attribute(span, LangfuseSpanAttributes.OBSERVATION_OUTPUT.value, safe_dumps(output_items_data)) + safe_set_attribute( + span, + LangfuseSpanAttributes.OBSERVATION_OUTPUT.value, + safe_dumps(output_items_data), + ) @staticmethod def _set_langfuse_specific_attributes(span: Span, kwargs, response_obj): @@ -210,14 +235,22 @@ class LangfuseOtelLogger(OpenTelemetry): langfuse_environment = os.environ.get("LANGFUSE_TRACING_ENVIRONMENT") if langfuse_environment: - safe_set_attribute(span, LangfuseSpanAttributes.LANGFUSE_ENVIRONMENT.value, langfuse_environment) + safe_set_attribute( + span, + LangfuseSpanAttributes.LANGFUSE_ENVIRONMENT.value, + langfuse_environment, + ) metadata = LangfuseOtelLogger._extract_langfuse_metadata(kwargs) LangfuseOtelLogger._set_metadata_attributes(span=span, metadata=metadata) messages = kwargs.get("messages") if messages: - safe_set_attribute(span, LangfuseSpanAttributes.OBSERVATION_INPUT.value, safe_dumps(messages)) + safe_set_attribute( + span, + LangfuseSpanAttributes.OBSERVATION_INPUT.value, + safe_dumps(messages), + ) LangfuseOtelLogger._set_observation_output(span=span, response_obj=response_obj) @@ -319,3 +352,15 @@ class LangfuseOtelLogger(OpenTelemetry): dynamic_headers["Authorization"] = auth_header return dynamic_headers + + async def async_service_success_hook(self, *args, **kwargs): + """ + Langfuse should not receive service success logs. + """ + pass + + async def async_service_failure_hook(self, *args, **kwargs): + """ + Langfuse should not receive service failure logs. + """ + pass diff --git a/tests/test_service_logger_otel.py b/tests/test_service_logger_otel.py new file mode 100644 index 00000000000..5cb21dadeae --- /dev/null +++ b/tests/test_service_logger_otel.py @@ -0,0 +1,85 @@ +import os +import sys +import unittest +from unittest.mock import patch, AsyncMock, MagicMock + +# Add the project root to sys.path +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) + +import litellm +from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger +from litellm.integrations.opentelemetry import OpenTelemetry +from litellm.types.services import ServiceTypes +from litellm._service_logger import ServiceLogging + + +class TestServiceLoggerOTEL(unittest.IsolatedAsyncioTestCase): + def setUp(self): + # Reset callbacks before each test + litellm.service_callback = [] + os.environ["LANGFUSE_PUBLIC_KEY"] = "pk-lf-123" + os.environ["LANGFUSE_SECRET_KEY"] = "sk-lf-123" + + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_tracing") + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_metrics") + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_logs") + async def test_langfuse_otel_ignores_service_logs( + self, mock_logs, mock_metrics, mock_tracing + ): + """ + Test that LangfuseOtelLogger overrides the service logging hooks with 'pass'. + """ + logger = LangfuseOtelLogger() + + # Verify hooks are overriden + self.assertEqual( + logger.async_service_success_hook.__qualname__, + "LangfuseOtelLogger.async_service_success_hook", + ) + self.assertEqual( + logger.async_service_failure_hook.__qualname__, + "LangfuseOtelLogger.async_service_failure_hook", + ) + + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_tracing") + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_metrics") + @patch("litellm.integrations.opentelemetry.OpenTelemetry._init_logs") + async def test_service_logging_shadowing_fix( + self, mock_logs, mock_metrics, mock_tracing + ): + """ + Test the architectural fix: multiple OTEL loggers should receive logs independently. + """ + # 1. Initialize two loggers + langfuse_logger = LangfuseOtelLogger() + otel_logger = OpenTelemetry() + + # 2. Setup service_callback list + litellm.service_callback = [langfuse_logger, otel_logger] + + service_logging = ServiceLogging() + + # 3. Mock the base OpenTelemetry hook + with patch.object( + OpenTelemetry, "async_service_success_hook", new_callable=AsyncMock + ) as mock_base_hook: + # Trigger a service event + await service_logging.async_service_success_hook( + service=ServiceTypes.DB, + call_type="success", + duration=0.1, + parent_otel_span=MagicMock(), + start_time=0.0, + end_time=1.0, + ) + + # The architectural fix ensures we call each correctly. + self.assertEqual( + mock_base_hook.call_count, + 1, + "Generic OTEL logger should have received the log exactly once.", + ) + + +if __name__ == "__main__": + unittest.main() From 98e87c3e67ab36948020f769da6eea87fd3ae2c3 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Mon, 19 Jan 2026 19:27:24 +0530 Subject: [PATCH 10/58] feat: Add Redis-based migration lock with bug fixes (#19261) --- .../litellm_proxy_extras/utils.py | 250 ++++++++++++++++-- litellm/proxy/db/prisma_client.py | 13 +- .../test_litellm_proxy_extras_utils.py | 214 ++++++++++++++- 3 files changed, 452 insertions(+), 25 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 7ffbe95be13..4e95267d8ab 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -11,6 +11,11 @@ from typing import Optional from litellm_proxy_extras._logging import logger +try: + from litellm.caching.redis_cache import RedisCache +except ImportError: + RedisCache = None # type: ignore + def str_to_bool(value: Optional[str]) -> bool: if value is None: @@ -18,6 +23,154 @@ def str_to_bool(value: Optional[str]) -> bool: return value.lower() in ("true", "1", "t", "y", "yes") +class MigrationLockManager: + """Redis-based lock manager for database migrations. + + Prevents concurrent Prisma migrations in multi-pod deployments by using + a distributed lock. Only one pod can hold the lock and run migrations at a time. + """ + + MIGRATION_LOCK_KEY = "migration_lock" + LOCK_TTL_SECONDS = 300 # 5 minutes TTL + + def __init__(self, redis_cache: Optional["RedisCache"] = None): + """Initialize the migration lock manager. + + Args: + redis_cache: Optional RedisCache instance for distributed locking. + If None, migrations run without lock protection (single instance mode). + """ + self.redis_cache = redis_cache + self.lock_acquired = False + self.pod_id = f"pod_{os.getpid()}_{int(time.time())}" + + def _get_redis_lock_key(self) -> str: + """Get Redis lock key for migration.""" + return f"migration_lock:{self.MIGRATION_LOCK_KEY}" + + def acquire_lock(self) -> bool: + """Acquire migration lock using Redis SET NX. + + Returns: + bool: True if lock acquired, False otherwise. + """ + if self.redis_cache is None: + logger.warning( + "Redis cache is not available, running migration without lock protection" + ) + self.lock_acquired = True + return True + + try: + lock_key = self._get_redis_lock_key() + + # FIX: Use native Redis SET NX instead of RedisCache.set_cache() + # Original bug: set_cache() doesn't support nx parameter and returns None + # Fixed: Use redis_client.set() directly which returns True/False + acquired = self.redis_cache.redis_client.set( + name=lock_key, + value=self.pod_id, + nx=True, # Only set if key doesn't exist + ex=self.LOCK_TTL_SECONDS, # Set expiration time + ) + + if acquired: + self.lock_acquired = True + logger.info(f"Migration lock acquired by pod {self.pod_id}") + return True + else: + logger.info("Migration lock is already held by another pod") + return False + + except Exception as e: + logger.warning(f"Failed to acquire migration lock: {e}") + return False + + def wait_for_lock_release( + self, check_interval: int = 5, max_wait: int = 300 + ) -> bool: + """Wait for another process to release the lock. + + Args: + check_interval: Seconds to wait between lock acquisition attempts. + max_wait: Maximum seconds to wait for lock release. + + Returns: + bool: True if lock acquired after waiting, False if timeout. + """ + if self.redis_cache is None: + logger.warning("Redis cache is not available, cannot wait for lock") + return False + + logger.info(f"Waiting for migration lock to be released (max {max_wait}s)...") + start_time = time.time() + + while time.time() - start_time < max_wait: + # Try to acquire lock using the public acquire_lock method + if self.acquire_lock(): + logger.info( + f"Migration lock acquired after waiting by pod {self.pod_id}" + ) + return True + + time.sleep(check_interval) + + logger.warning(f"Failed to acquire migration lock within {max_wait} seconds") + return False + + def release_lock(self): + """Release migration lock atomically using Lua script. + + FIX: Use Lua script for atomic compare-and-delete to prevent race conditions. + Original bug: Non-atomic GET then DELETE allows another pod to acquire lock + between the GET and DELETE operations. + """ + if not self.lock_acquired or self.redis_cache is None: + return + + try: + lock_key = self._get_redis_lock_key() + + # FIX: Use Lua script for atomic compare-and-delete + # This prevents race condition where: + # 1. Pod A reads lock value (sees its own pod_id) + # 2. Lock TTL expires + # 3. Pod B acquires lock + # 4. Pod A deletes lock (deletes Pod B's lock!) + lua_script = """ + if redis.call("GET", KEYS[1]) == ARGV[1] then + return redis.call("DEL", KEYS[1]) + else + return 0 + end + """ + + result = self.redis_cache.redis_client.eval( + lua_script, + 1, # Number of keys + lock_key, # KEYS[1] + self.pod_id, # ARGV[1] + ) + + if result == 1: + logger.info(f"Migration lock released by pod {self.pod_id}") + else: + logger.warning(f"Pod {self.pod_id} cannot release lock (not owner)") + + except Exception as e: + logger.warning(f"Failed to release migration lock: {e}") + finally: + self.lock_acquired = False + + def __enter__(self): + """Context manager entry - acquire lock when entering with statement.""" + self.acquire_lock() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit - release lock when exiting with statement.""" + self.release_lock() + def _get_prisma_env() -> dict: """Get environment variables for Prisma, handling offline mode if configured.""" @@ -25,7 +178,9 @@ def _get_prisma_env() -> dict: if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")): # These env vars prevent Prisma from attempting downloads prisma_env["NPM_CONFIG_PREFER_OFFLINE"] = "true" - prisma_env["NPM_CONFIG_CACHE"] = os.getenv("NPM_CONFIG_CACHE", "/app/.cache/npm") + prisma_env["NPM_CONFIG_CACHE"] = os.getenv( + "NPM_CONFIG_CACHE", "/app/.cache/npm" + ) return prisma_env @@ -34,29 +189,28 @@ def _get_prisma_command() -> str: if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")): # Primary location where Prisma Python package installs the CLI default_cli_path = "/app/.cache/prisma-python/binaries/node_modules/.bin/prisma" - + # Check if custom path is provided (for flexibility) custom_cli_path = os.getenv("PRISMA_CLI_PATH") if custom_cli_path and os.path.exists(custom_cli_path): logger.info(f"Using custom Prisma CLI at {custom_cli_path}") return custom_cli_path - + # Check the default location if os.path.exists(default_cli_path): logger.info(f"Using cached Prisma CLI at {default_cli_path}") return default_cli_path - + # If not found, log warning and fall back logger.warning( f"Prisma CLI not found at {default_cli_path}. " "Falling back to Python wrapper (may attempt downloads)" ) - + # Fall back to the Python wrapper (will work in online mode) return "prisma" - class ProxyExtrasDBManager: @staticmethod def _get_prisma_dir() -> str: @@ -119,7 +273,7 @@ class ProxyExtrasDBManager: stdout=open(migration_file, "w"), check=True, timeout=30, - env=prisma_env + env=prisma_env, ) # 3. Mark the migration as applied since it represents current state @@ -134,7 +288,7 @@ class ProxyExtrasDBManager: ], check=True, timeout=30, - env=prisma_env + env=prisma_env, ) return True @@ -159,14 +313,20 @@ class ProxyExtrasDBManager: @staticmethod def _roll_back_migration(migration_name: str): """Mark a specific migration as rolled back""" - # Set up environment for offline mode if configured + # Set up environment for offline mode if configured prisma_env = _get_prisma_env() subprocess.run( - [_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name], + [ + _get_prisma_command(), + "migrate", + "resolve", + "--rolled-back", + migration_name, + ], timeout=60, check=True, capture_output=True, - env=prisma_env + env=prisma_env, ) @staticmethod @@ -178,7 +338,7 @@ class ProxyExtrasDBManager: timeout=60, check=True, capture_output=True, - env=prisma_env + env=prisma_env, ) @staticmethod @@ -248,7 +408,7 @@ class ProxyExtrasDBManager: if not database_url: logger.error("DATABASE_URL not set") return - + diff_dir = ( Path(migrations_dir) / "migrations" @@ -283,7 +443,7 @@ class ProxyExtrasDBManager: check=True, timeout=60, stdout=f, - env=_get_prisma_env() + env=_get_prisma_env(), ) except subprocess.CalledProcessError as e: logger.warning(f"Failed to generate migration diff: {e.stderr}") @@ -313,7 +473,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env() + env=_get_prisma_env(), ) logger.info(f"prisma db execute stdout: {result.stdout}") logger.info("✅ Migration diff applied successfully") @@ -331,12 +491,18 @@ class ProxyExtrasDBManager: try: logger.info(f"Resolving migration: {migration_name}") subprocess.run( - [_get_prisma_command(), "migrate", "resolve", "--applied", migration_name], + [ + _get_prisma_command(), + "migrate", + "resolve", + "--applied", + migration_name, + ], timeout=60, check=True, capture_output=True, text=True, - env=_get_prisma_env() + env=_get_prisma_env(), ) logger.debug(f"Resolved migration: {migration_name}") except subprocess.CalledProcessError as e: @@ -346,19 +512,57 @@ class ProxyExtrasDBManager: ) @staticmethod - def setup_database(use_migrate: bool = False) -> bool: + def setup_database( + use_migrate: bool = False, redis_cache: Optional["RedisCache"] = None + ) -> bool: """ - Set up the database using either prisma migrate or prisma db push - Uses migrations from litellm-proxy-extras package + Set up the database using either prisma migrate or prisma db push. + Uses migrations from litellm-proxy-extras package. + In multi-instance environment, use redis lock to prevent concurrent execution. Args: - schema_path (str): Path to the Prisma schema file use_migrate (bool): Whether to use prisma migrate instead of db push + redis_cache: Redis cache instance for distributed locking Returns: bool: True if setup was successful, False otherwise """ schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma" + + database_url = os.getenv("DATABASE_URL") + if not database_url: + logger.error("DATABASE_URL environment variable is not set") + return False + + # Use MigrationLockManager to prevent concurrent migration execution + with MigrationLockManager(redis_cache) as lock_manager: + # Lock is already acquired in __enter__, check if it was successful + if not lock_manager.lock_acquired: + # Cannot acquire lock, another process is running migration + logger.info( + "Another pod is running migration, waiting for completion..." + ) + + # Wait for other process to complete migration + if not lock_manager.wait_for_lock_release(): + logger.error("Failed to acquire migration lock after waiting") + return False + + # Successfully acquired lock, proceed with migration + logger.info("Acquired migration lock, proceeding with migration") + return ProxyExtrasDBManager._execute_migration(use_migrate, schema_path) + + @staticmethod + def _execute_migration(use_migrate: bool, schema_path: str) -> bool: + """Execute the actual migration. + + Args: + use_migrate: Whether to use prisma migrate instead of db push + schema_path: Path to the Prisma schema file + + Returns: + bool: True if migration was successful, False otherwise + """ for attempt in range(4): original_dir = os.getcwd() migrations_dir = ProxyExtrasDBManager._get_prisma_dir() @@ -375,7 +579,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env() + env=_get_prisma_env(), ) logger.info(f"prisma migrate deploy stdout: {result.stdout}") @@ -413,7 +617,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env() + env=_get_prisma_env(), ) logger.info( f"✅ Migration {failed_migration} marked as rolled back... retrying" diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index c9c0cfe8f68..65fed92340b 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -383,7 +383,18 @@ class PrismaManager: prisma_dir = PrismaManager._get_prisma_dir() - return ProxyExtrasDBManager.setup_database(use_migrate=use_migrate) + # Import redis_usage_cache for distributed locking + try: + from litellm.proxy.proxy_server import redis_usage_cache + except ImportError: + verbose_proxy_logger.warning( + "Could not import redis_usage_cache, migrations will run without distributed locking" + ) + redis_usage_cache = None + + return ProxyExtrasDBManager.setup_database( + use_migrate=use_migrate, redis_cache=redis_usage_cache + ) else: # Use prisma db push with increased timeout subprocess.run( diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 5714cd5c487..cff3a88ab6c 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -1,11 +1,13 @@ import os import sys +import time +from unittest.mock import Mock, patch sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path -from litellm_proxy_extras.utils import ProxyExtrasDBManager +from litellm_proxy_extras.utils import ProxyExtrasDBManager, MigrationLockManager def test_custom_prisma_dir(monkeypatch): @@ -125,3 +127,213 @@ class TestErrorClassificationPriority: error_message = "connection timeout" assert ProxyExtrasDBManager._is_permission_error(error_message) is False assert ProxyExtrasDBManager._is_idempotent_error(error_message) is False + + +class TestMigrationLockManager: + """Test cases for MigrationLockManager""" + + def test_acquire_lock_success(self): + """Test successful lock acquisition using Redis SET NX""" + mock_redis = Mock() + # FIX VERIFICATION: redis_client.set() returns True for successful SET NX + mock_redis.redis_client.set.return_value = True + + lock_manager = MigrationLockManager(mock_redis) + result = lock_manager.acquire_lock() + + assert result is True + assert lock_manager.lock_acquired is True + # Verify SET NX was called with correct parameters + mock_redis.redis_client.set.assert_called_once() + call_args = mock_redis.redis_client.set.call_args + assert call_args[1]["nx"] is True # SET NX parameter + assert call_args[1]["ex"] == 300 # TTL + + def test_acquire_lock_failure(self): + """Test lock acquisition failure when lock is already held""" + mock_redis = Mock() + # FIX VERIFICATION: redis_client.set() returns False when key exists + mock_redis.redis_client.set.return_value = False + + lock_manager = MigrationLockManager(mock_redis) + result = lock_manager.acquire_lock() + + assert result is False + assert lock_manager.lock_acquired is False + + def test_acquire_lock_without_redis(self): + """Test lock acquisition without Redis (graceful fallback)""" + lock_manager = MigrationLockManager(redis_cache=None) + result = lock_manager.acquire_lock() + + assert result is True + assert lock_manager.lock_acquired is True + + def test_release_lock_success(self): + """Test successful lock release using Lua script""" + mock_redis = Mock() + # FIX VERIFICATION: Lua script returns 1 for successful delete + mock_redis.redis_client.eval.return_value = 1 + + lock_manager = MigrationLockManager(mock_redis) + lock_manager.pod_id = "pod_123_456" + lock_manager.lock_acquired = True + + lock_manager.release_lock() + + # Verify Lua script was called for atomic compare-and-delete + mock_redis.redis_client.eval.assert_called_once() + call_args = mock_redis.redis_client.eval.call_args + # Verify Lua script contains atomic compare-and-delete logic + lua_script = call_args[0][0] + assert "GET" in lua_script + assert "DEL" in lua_script + assert lock_manager.lock_acquired is False + + def test_release_lock_wrong_owner(self): + """Test releasing lock when not the owner""" + mock_redis = Mock() + # FIX VERIFICATION: Lua script returns 0 when ownership check fails + mock_redis.redis_client.eval.return_value = 0 + + lock_manager = MigrationLockManager(mock_redis) + lock_manager.pod_id = "pod_123_456" + lock_manager.lock_acquired = True + + lock_manager.release_lock() + + mock_redis.redis_client.eval.assert_called_once() + assert lock_manager.lock_acquired is False + + def test_context_manager(self): + """Test MigrationLockManager as context manager""" + mock_redis = Mock() + mock_redis.redis_client.set.return_value = True + mock_redis.redis_client.eval.return_value = 1 + + lock_manager = MigrationLockManager(mock_redis) + lock_manager.pod_id = "pod_123_456" + + with lock_manager: + assert lock_manager.lock_acquired is True + + # Should call release_lock when exiting context + mock_redis.redis_client.eval.assert_called_once() + + def test_wait_for_lock_release_success(self): + """Test waiting for lock release and acquiring it""" + mock_redis = Mock() + # First call fails, second call succeeds + mock_redis.redis_client.set.side_effect = [False, True] + + lock_manager = MigrationLockManager(mock_redis) + result = lock_manager.wait_for_lock_release(check_interval=0.1, max_wait=1) + + assert result is True + assert lock_manager.lock_acquired is True + assert mock_redis.redis_client.set.call_count == 2 + + def test_wait_for_lock_release_timeout(self): + """Test timeout when waiting for lock release""" + mock_redis = Mock() + # Always fails to acquire lock + mock_redis.redis_client.set.return_value = False + + lock_manager = MigrationLockManager(mock_redis) + start_time = time.time() + result = lock_manager.wait_for_lock_release(check_interval=0.1, max_wait=0.5) + end_time = time.time() + + assert result is False + assert lock_manager.lock_acquired is False + assert end_time - start_time >= 0.5 # Should wait at least max_wait time + assert mock_redis.redis_client.set.call_count > 1 # Multiple attempts + + +class TestProxyExtrasDBManagerMigrationLock: + """Test cases for ProxyExtrasDBManager with migration locking""" + + @patch("litellm_proxy_extras.utils.ProxyExtrasDBManager._execute_migration") + def test_setup_database_with_redis_lock_success( + self, mock_execute_migration, monkeypatch + ): + """Test successful database setup with Redis lock""" + monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test") + mock_execute_migration.return_value = True + + # Mock Redis cache + mock_redis = Mock() + mock_redis.redis_client.set.return_value = True # Lock acquired + mock_redis.redis_client.eval.return_value = 1 # Lock released + + result = ProxyExtrasDBManager.setup_database( + use_migrate=True, redis_cache=mock_redis + ) + + assert result is True + mock_execute_migration.assert_called_once() + # Verify lock was acquired + mock_redis.redis_client.set.assert_called_once() + + @patch("litellm_proxy_extras.utils.ProxyExtrasDBManager._execute_migration") + def test_setup_database_with_redis_lock_wait_and_acquire( + self, mock_execute_migration, monkeypatch + ): + """Test database setup when lock is held, then acquired after waiting""" + monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test") + mock_execute_migration.return_value = True + + # Mock Redis cache - first call fails, second call succeeds + mock_redis = Mock() + mock_redis.redis_client.set.side_effect = [False, True] + mock_redis.redis_client.eval.return_value = 1 + + result = ProxyExtrasDBManager.setup_database( + use_migrate=True, redis_cache=mock_redis + ) + + assert result is True + mock_execute_migration.assert_called_once() + # Should have tried to acquire lock twice + assert mock_redis.redis_client.set.call_count == 2 + + @patch("litellm_proxy_extras.utils.ProxyExtrasDBManager._execute_migration") + def test_setup_database_without_redis(self, mock_execute_migration, monkeypatch): + """Test database setup without Redis cache (single instance mode)""" + monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test") + mock_execute_migration.return_value = True + + result = ProxyExtrasDBManager.setup_database(use_migrate=True, redis_cache=None) + + assert result is True + mock_execute_migration.assert_called_once() + + def test_setup_database_no_database_url(self): + """Test database setup without DATABASE_URL""" + with patch.dict(os.environ, {}, clear=True): + result = ProxyExtrasDBManager.setup_database( + use_migrate=True, redis_cache=None + ) + + assert result is False + + @patch("litellm_proxy_extras.utils.ProxyExtrasDBManager._execute_migration") + def test_setup_database_lock_timeout(self, mock_execute_migration, monkeypatch): + """Test database setup when lock acquisition times out""" + monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test") + + # Mock Redis cache - always fails to acquire lock + mock_redis = Mock() + mock_redis.redis_client.set.return_value = False + + # Patch wait_for_lock_release to simulate timeout + with patch.object(MigrationLockManager, "wait_for_lock_release") as mock_wait: + mock_wait.return_value = False # Simulate timeout + + result = ProxyExtrasDBManager.setup_database( + use_migrate=True, redis_cache=mock_redis + ) + + assert result is False + mock_execute_migration.assert_not_called() # Should not run migration + mock_wait.assert_called_once() From 1dc2d2ddacc4bf54335a8023ecd4645858f9b3e6 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Mon, 19 Jan 2026 19:30:23 +0530 Subject: [PATCH 11/58] fix(utils.py): correctly extract messages from google genai contents (#19156) * fix(utils.py): correctly extract messages from google genai contents * refactor use shared utilities --- litellm/utils.py | 392 ++++++++++++------ .../test_google_api_endpoints.py | 257 ++++++------ 2 files changed, 392 insertions(+), 257 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index f5b5c50468f..9198b8e2ecc 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -87,6 +87,7 @@ def _get_cached_custom_logger(): global _CustomLogger if _CustomLogger is None: from litellm.integrations.custom_logger import CustomLogger + _CustomLogger = CustomLogger return _CustomLogger @@ -100,6 +101,7 @@ def _get_cached_custom_guardrail(): global _CustomGuardrail if _CustomGuardrail is None: from litellm.integrations.custom_guardrail import CustomGuardrail + _CustomGuardrail = CustomGuardrail return _CustomGuardrail @@ -113,6 +115,7 @@ def _get_cached_caching_handler_response(): global _CachingHandlerResponse if _CachingHandlerResponse is None: from litellm.caching.caching_handler import CachingHandlerResponse + _CachingHandlerResponse = CachingHandlerResponse return _CachingHandlerResponse @@ -126,6 +129,7 @@ def _get_cached_llm_caching_handler(): global _LLMCachingHandler if _LLMCachingHandler is None: from litellm.caching.caching_handler import LLMCachingHandler + _LLMCachingHandler = LLMCachingHandler return _LLMCachingHandler @@ -144,9 +148,11 @@ def _get_cached_audio_utils(): global _audio_utils_module if _audio_utils_module is None: import litellm.litellm_core_utils.audio_utils.utils + _audio_utils_module = litellm.litellm_core_utils.audio_utils.utils return _audio_utils_module + from litellm.types.llms.openai import ( AllMessageValues, AllPromptValues, @@ -203,10 +209,6 @@ from litellm.types.utils import ( # Thank you users! We ❤️ you! - Krrish & Ishaan - - - - try: # Python 3.9+ with resources.files("litellm.litellm_core_utils.tokenizers").joinpath( @@ -250,10 +252,14 @@ from litellm.llms.base_llm.base_utils import ( if TYPE_CHECKING: # Heavy types that are only needed for type checking; avoid importing # their modules at runtime during `litellm` import. - from litellm.caching.caching_handler import CachingHandlerResponse, LLMCachingHandler + from litellm.caching.caching_handler import ( + CachingHandlerResponse, + LLMCachingHandler, + ) from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.files.transformation import BaseFilesConfig from litellm.proxy._types import AllowedModelRegion + # Type stubs for lazy-loaded functions to help mypy understand their types # These imports allow mypy to understand the types when these are accessed via __getattr__ from litellm.litellm_core_utils.exception_mapping_utils import exception_type @@ -288,10 +294,13 @@ if TYPE_CHECKING: ) from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig from litellm.llms.base_llm.search.transformation import BaseSearchConfig - from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig + from litellm.llms.base_llm.text_to_speech.transformation import ( + BaseTextToSpeechConfig, + ) from litellm.llms.bedrock.common_utils import BedrockModelInfo from litellm.llms.cohere.common_utils import CohereModelInfo from litellm.llms.mistral.ocr.transformation import MistralOCRConfig + # Type stubs for lazy-loaded functions and classes from litellm.litellm_core_utils.cached_imports import ( get_coroutine_checker, @@ -335,6 +344,7 @@ if TYPE_CHECKING: reset_retry_policy, ) from litellm.secret_managers.main import get_secret + # Type stubs for lazy-loaded config classes and types from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.containers.transformation import BaseContainerConfig @@ -618,8 +628,8 @@ def load_credentials_from_list(kwargs: dict): Updates kwargs with the credentials if credential_name in kwarg """ # Access CredentialAccessor via module to trigger lazy loading if needed - CredentialAccessor = getattr(sys.modules[__name__], 'CredentialAccessor') - + CredentialAccessor = getattr(sys.modules[__name__], "CredentialAccessor") + credential_name = kwargs.get("litellm_credential_name") if credential_name and litellm.credential_list: credential_accessor = CredentialAccessor.get_credential_values(credential_name) @@ -646,7 +656,7 @@ def _is_gemini_model(model: Optional[str], custom_llm_provider: Optional[str]) - if custom_llm_provider in ["vertex_ai", "vertex_ai_beta"]: return model is not None and "gemini" in model.lower() return True - + # Check if model name contains gemini return model is not None and "gemini" in model.lower() @@ -668,7 +678,7 @@ def _process_assistant_message_tool_calls( """ role = msg_copy.get("role") tool_calls = msg_copy.get("tool_calls") - + if role == "assistant" and isinstance(tool_calls, list): new_tool_calls = [] for tc in tool_calls: @@ -681,17 +691,17 @@ def _process_assistant_message_tool_calls( else: new_tool_calls.append(tc) continue - + # Remove thought signature from ID if present if isinstance(tc_dict.get("id"), str): if thought_signature_separator in tc_dict["id"]: tc_dict["id"] = _remove_thought_signature_from_id( tc_dict["id"], thought_signature_separator ) - + new_tool_calls.append(tc_dict) msg_copy["tool_calls"] = new_tool_calls - + return msg_copy @@ -699,14 +709,12 @@ def _process_tool_message_id(msg_copy: dict, thought_signature_separator: str) - """ Process tool message to remove thought signature from tool_call_id. """ - if msg_copy.get("role") == "tool" and isinstance( - msg_copy.get("tool_call_id"), str - ): + if msg_copy.get("role") == "tool" and isinstance(msg_copy.get("tool_call_id"), str): if thought_signature_separator in msg_copy["tool_call_id"]: msg_copy["tool_call_id"] = _remove_thought_signature_from_id( msg_copy["tool_call_id"], thought_signature_separator ) - + return msg_copy @@ -717,7 +725,7 @@ def _remove_thought_signatures_from_messages( Remove thought signatures from tool call IDs in all messages. """ processed_messages = [] - + for msg in messages: # Handle Pydantic models (convert to dict) if hasattr(msg, "model_dump"): @@ -728,17 +736,17 @@ def _remove_thought_signatures_from_messages( # Unknown type, keep as is processed_messages.append(msg) continue - + # Process assistant messages with tool_calls msg_dict = _process_assistant_message_tool_calls( msg_dict, thought_signature_separator ) - + # Process tool messages with tool_call_id msg_dict = _process_tool_message_id(msg_dict, thought_signature_separator) - + processed_messages.append(msg_dict) - + return processed_messages @@ -763,12 +771,12 @@ def function_setup( # noqa: PLR0915 function_id: Optional[str] = kwargs["id"] if "id" in kwargs else None ## LAZY LOAD COROUTINE CHECKER ## - get_coroutine_checker = getattr(sys.modules[__name__], 'get_coroutine_checker') + get_coroutine_checker = getattr(sys.modules[__name__], "get_coroutine_checker") ## DYNAMIC CALLBACKS ## - dynamic_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]] = ( - kwargs.pop("callbacks", None) - ) + dynamic_callbacks: Optional[ + List[Union[str, Callable, "CustomLogger"]] + ] = kwargs.pop("callbacks", None) all_callbacks = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks) if len(all_callbacks) > 0: @@ -811,7 +819,7 @@ def function_setup( # noqa: PLR0915 + litellm.failure_callback ) ) - get_set_callbacks = getattr(sys.modules[__name__], 'get_set_callbacks') + get_set_callbacks = getattr(sys.modules[__name__], "get_set_callbacks") get_set_callbacks()(callback_list=callback_list, function_id=function_id) ## ASYNC CALLBACKS if len(litellm.input_callback) > 0: @@ -939,7 +947,7 @@ def function_setup( # noqa: PLR0915 elif kwargs.get("messages", None): messages = kwargs["messages"] ### PRE-CALL RULES ### - Rules = getattr(sys.modules[__name__], 'Rules') + Rules = getattr(sys.modules[__name__], "Rules") if ( Rules.has_pre_call_rules() and isinstance(messages, list) @@ -947,7 +955,6 @@ def function_setup( # noqa: PLR0915 and isinstance(messages[0], dict) and "content" in messages[0] ): - buffer = StringIO() for m in messages: content = m.get("content", "") @@ -958,7 +965,7 @@ def function_setup( # noqa: PLR0915 input=buffer.getvalue(), model=model, ) - + ### REMOVE THOUGHT SIGNATURES FROM TOOL CALL IDS FOR NON-GEMINI MODELS ### # Gemini models embed thought signatures in tool call IDs. When sending # messages with tool calls to non-Gemini providers, we need to remove these @@ -974,7 +981,7 @@ def function_setup( # noqa: PLR0915 # Get custom_llm_provider to determine target provider custom_llm_provider = kwargs.get("custom_llm_provider") - + # If custom_llm_provider not in kwargs, try to determine it from the model if not custom_llm_provider and model: try: @@ -985,18 +992,18 @@ def function_setup( # noqa: PLR0915 except Exception: # If we can't determine the provider, skip this processing pass - + # Only process if target is NOT a Gemini model if not _is_gemini_model(model, custom_llm_provider): verbose_logger.debug( "Removing thought signatures from tool call IDs for non-Gemini model" ) - + # Process messages to remove thought signatures processed_messages = _remove_thought_signatures_from_messages( messages, THOUGHT_SIGNATURE_SEPARATOR ) - + # Update messages in kwargs or args if "messages" in kwargs: kwargs["messages"] = processed_messages @@ -1041,9 +1048,7 @@ def function_setup( # noqa: PLR0915 _file_obj: FileTypes = args[1] if len(args) > 1 else kwargs["file"] # Lazy import audio_utils.utils only when needed for transcription calls audio_utils = _get_cached_audio_utils() - file_checksum = audio_utils.get_audio_file_content_hash( - file_obj=_file_obj - ) + file_checksum = audio_utils.get_audio_file_content_hash(file_obj=_file_obj) if "metadata" in kwargs: kwargs["metadata"]["file_checksum"] = file_checksum else: @@ -1064,6 +1069,42 @@ def function_setup( # noqa: PLR0915 else kwargs.get("input") or kwargs.get("messages", "default-message-value") ) + elif ( + call_type == CallTypes.generate_content.value + or call_type == CallTypes.agenerate_content.value + or call_type == CallTypes.generate_content_stream.value + or call_type == CallTypes.agenerate_content_stream.value + ): + try: + from litellm.google_genai.adapters.transformation import ( + GoogleGenAIAdapter, + ) + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + get_last_user_message, + ) + + contents_param = args[1] if len(args) > 1 else kwargs.get("contents") + model_param = args[0] if len(args) > 0 else kwargs.get("model", "") + + if contents_param: + adapter = GoogleGenAIAdapter() + transformed = adapter.translate_generate_content_to_completion( + model=model_param, + contents=contents_param, + config=kwargs.get("config"), + ) + transformed_messages = transformed.get("messages", []) + messages = ( + get_last_user_message(transformed_messages) + or "default-message-value" + ) + else: + messages = "default-message-value" + except Exception as e: + verbose_logger.debug( + f"Error extracting messages from Google contents: {str(e)}" + ) + messages = "default-message-value" else: messages = "default-message-value" stream = False @@ -1072,7 +1113,9 @@ def function_setup( # noqa: PLR0915 call_type=call_type, ): stream = True - get_litellm_logging_class = getattr(sys.modules[__name__], 'get_litellm_logging_class') + get_litellm_logging_class = getattr( + sys.modules[__name__], "get_litellm_logging_class" + ) logging_obj = get_litellm_logging_class()( # Victim for object pool model=model, # type: ignore messages=messages, @@ -1156,8 +1199,10 @@ def _get_wrapper_num_retries( if num_retries is None: num_retries = litellm.num_retries if kwargs.get("retry_policy", None): - get_num_retries_from_retry_policy = getattr(sys.modules[__name__], 'get_num_retries_from_retry_policy') - reset_retry_policy = getattr(sys.modules[__name__], 'reset_retry_policy') + get_num_retries_from_retry_policy = getattr( + sys.modules[__name__], "get_num_retries_from_retry_policy" + ) + reset_retry_policy = getattr(sys.modules[__name__], "reset_retry_policy") retry_policy_num_retries = get_num_retries_from_retry_policy( exception=exception, retry_policy=kwargs.get("retry_policy"), @@ -1185,7 +1230,7 @@ def _get_wrapper_timeout( def check_coroutine(value) -> bool: - get_coroutine_checker = getattr(sys.modules[__name__], 'get_coroutine_checker') + get_coroutine_checker = getattr(sys.modules[__name__], "get_coroutine_checker") return get_coroutine_checker().is_async_callable(value) @@ -1344,7 +1389,7 @@ def post_call_processing( def client(original_function): # noqa: PLR0915 - Rules = getattr(sys.modules[__name__], 'Rules') + Rules = getattr(sys.modules[__name__], "Rules") rules_obj = Rules() @wraps(original_function) @@ -1530,7 +1575,9 @@ def client(original_function): # noqa: PLR0915 ) else: # RETURN RESULT - update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') + update_response_metadata = getattr( + sys.modules[__name__], "update_response_metadata" + ) update_response_metadata( result=result, logging_obj=logging_obj, @@ -1574,7 +1621,7 @@ def client(original_function): # noqa: PLR0915 # Copy the current context to propagate it to the background thread # This is essential for OpenTelemetry span context propagation ctx = contextvars.copy_context() - executor = getattr(sys.modules[__name__], 'executor') + executor = getattr(sys.modules[__name__], "executor") executor.submit( ctx.run, logging_obj.success_handler, @@ -1583,7 +1630,9 @@ def client(original_function): # noqa: PLR0915 end_time, ) # RETURN RESULT - update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') + update_response_metadata = getattr( + sys.modules[__name__], "update_response_metadata" + ) update_response_metadata( result=result, logging_obj=logging_obj, @@ -1600,15 +1649,19 @@ def client(original_function): # noqa: PLR0915 kwargs.get("num_retries", None) or litellm.num_retries or None ) if kwargs.get("retry_policy", None): - get_num_retries_from_retry_policy = getattr(sys.modules[__name__], 'get_num_retries_from_retry_policy') - reset_retry_policy = getattr(sys.modules[__name__], 'reset_retry_policy') + get_num_retries_from_retry_policy = getattr( + sys.modules[__name__], "get_num_retries_from_retry_policy" + ) + reset_retry_policy = getattr( + sys.modules[__name__], "reset_retry_policy" + ) num_retries = get_num_retries_from_retry_policy( exception=e, retry_policy=kwargs.get("retry_policy"), ) - kwargs["retry_policy"] = ( - reset_retry_policy() - ) # prevent infinite loops + kwargs[ + "retry_policy" + ] = reset_retry_policy() # prevent infinite loops litellm.num_retries = ( None # set retries to None to prevent infinite loops ) @@ -1645,15 +1698,19 @@ def client(original_function): # noqa: PLR0915 kwargs.get("num_retries", None) or litellm.num_retries or None ) if kwargs.get("retry_policy", None): - get_num_retries_from_retry_policy = getattr(sys.modules[__name__], 'get_num_retries_from_retry_policy') - reset_retry_policy = getattr(sys.modules[__name__], 'reset_retry_policy') + get_num_retries_from_retry_policy = getattr( + sys.modules[__name__], "get_num_retries_from_retry_policy" + ) + reset_retry_policy = getattr( + sys.modules[__name__], "reset_retry_policy" + ) num_retries = get_num_retries_from_retry_policy( exception=e, retry_policy=kwargs.get("retry_policy"), ) - kwargs["retry_policy"] = ( - reset_retry_policy() - ) # prevent infinite loops + kwargs[ + "retry_policy" + ] = reset_retry_policy() # prevent infinite loops litellm.num_retries = ( None # set retries to None to prevent infinite loops ) @@ -1804,7 +1861,9 @@ def client(original_function): # noqa: PLR0915 chunks, messages=kwargs.get("messages", None) ) else: - update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') + update_response_metadata = getattr( + sys.modules[__name__], "update_response_metadata" + ) update_response_metadata( result=result, logging_obj=logging_obj, @@ -1869,7 +1928,9 @@ def client(original_function): # noqa: PLR0915 end_time=end_time, ) - update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') + update_response_metadata = getattr( + sys.modules[__name__], "update_response_metadata" + ) update_response_metadata( result=result, logging_obj=logging_obj, @@ -1969,7 +2030,7 @@ def client(original_function): # noqa: PLR0915 setattr(e, "timeout", timeout) raise e - get_coroutine_checker = getattr(sys.modules[__name__], 'get_coroutine_checker') + get_coroutine_checker = getattr(sys.modules[__name__], "get_coroutine_checker") is_coroutine = get_coroutine_checker().is_async_callable(original_function) # Return the appropriate wrapper based on the original function type @@ -2330,7 +2391,7 @@ def supports_response_schema( """ ## GET LLM PROVIDER ## try: - get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') + get_llm_provider = getattr(sys.modules[__name__], "get_llm_provider") model, custom_llm_provider, _, _ = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider ) @@ -2694,10 +2755,10 @@ def register_model(model_cost: Union[str, dict]): # noqa: PLR0915 ## override / add new keys to the existing model cost dictionary updated_dictionary = _update_dictionary(existing_model, value) litellm.model_cost.setdefault(model_cost_key, {}).update(updated_dictionary) - + # Invalidate case-insensitive lookup map since model_cost was modified _invalidate_model_cost_lowercase_map() - + verbose_logger.debug( f"added/updated model={model_cost_key} in litellm.model_cost: {model_cost_key}" ) @@ -3034,8 +3095,10 @@ def get_optional_params_embeddings( # noqa: PLR0915 **kwargs, ): # Lazy load get_supported_openai_params - get_supported_openai_params = getattr(sys.modules[__name__], 'get_supported_openai_params') - + get_supported_openai_params = getattr( + sys.modules[__name__], "get_supported_openai_params" + ) + # retrieve all parameters passed to the function passed_params = locals() custom_llm_provider = passed_params.pop("custom_llm_provider", None) @@ -3308,8 +3371,8 @@ def get_optional_params_embeddings( # noqa: PLR0915 ) elif custom_llm_provider == "ollama": - if 'dimensions' in non_default_params: - optional_params['dimensions']=non_default_params.pop('dimensions') + if "dimensions" in non_default_params: + optional_params["dimensions"] = non_default_params.pop("dimensions") if len(non_default_params.keys()) > 0: if ( litellm.drop_params is True or drop_params is True @@ -3575,10 +3638,10 @@ def pre_process_non_default_params( if "response_format" in non_default_params: if provider_config is not None: - non_default_params["response_format"] = ( - provider_config.get_json_schema_from_pydantic_object( - response_format=non_default_params["response_format"] - ) + non_default_params[ + "response_format" + ] = provider_config.get_json_schema_from_pydantic_object( + response_format=non_default_params["response_format"] ) else: non_default_params["response_format"] = type_to_response_format_param( @@ -3707,16 +3770,16 @@ def pre_process_optional_params( True # so that main.py adds the function call to the prompt ) if "tools" in non_default_params: - optional_params["functions_unsupported_model"] = ( - non_default_params.pop("tools") - ) + optional_params[ + "functions_unsupported_model" + ] = non_default_params.pop("tools") non_default_params.pop( "tool_choice", None ) # causes ollama requests to hang elif "functions" in non_default_params: - optional_params["functions_unsupported_model"] = ( - non_default_params.pop("functions") - ) + optional_params[ + "functions_unsupported_model" + ] = non_default_params.pop("functions") elif ( litellm.add_function_to_prompt ): # if user opts to add it to prompt instead @@ -3840,7 +3903,9 @@ def get_optional_params( # noqa: PLR0915 message=f"{custom_llm_provider} does not support parameters: {list(unsupported_params.keys())}, for model={model}. To drop these, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n. \n If you want to use these params dynamically send allowed_openai_params={list(unsupported_params.keys())} in your request.", ) - get_supported_openai_params = getattr(sys.modules[__name__], 'get_supported_openai_params') + get_supported_openai_params = getattr( + sys.modules[__name__], "get_supported_openai_params" + ) supported_params = get_supported_openai_params( model=model, custom_llm_provider=custom_llm_provider ) @@ -4095,7 +4160,7 @@ def get_optional_params( # noqa: PLR0915 ), ) elif custom_llm_provider == "bedrock": - BedrockModelInfo = getattr(sys.modules[__name__], 'BedrockModelInfo') + BedrockModelInfo = getattr(sys.modules[__name__], "BedrockModelInfo") bedrock_route = BedrockModelInfo.get_bedrock_route(model) bedrock_base_model = BedrockModelInfo.get_base_model(model) if bedrock_route == "converse" or bedrock_route == "converse_like": @@ -4520,8 +4585,8 @@ def get_optional_params( # noqa: PLR0915 # Apply nested drops from additional_drop_params if additional_drop_params: - is_nested_path = getattr(sys.modules[__name__], 'is_nested_path') - delete_nested_value = getattr(sys.modules[__name__], 'delete_nested_value') + is_nested_path = getattr(sys.modules[__name__], "is_nested_path") + delete_nested_value = getattr(sys.modules[__name__], "delete_nested_value") nested_paths = [p for p in additional_drop_params if is_nested_path(p)] for path in nested_paths: optional_params = delete_nested_value(optional_params, path) @@ -4571,7 +4636,9 @@ def add_provider_specific_params_to_optional_params( else: processed_extra_body = initial_extra_body - _ensure_extra_body_is_safe = getattr(sys.modules[__name__], '_ensure_extra_body_is_safe') + _ensure_extra_body_is_safe = getattr( + sys.modules[__name__], "_ensure_extra_body_is_safe" + ) optional_params["extra_body"] = _ensure_extra_body_is_safe( extra_body=processed_extra_body ) @@ -4862,9 +4929,9 @@ def get_response_string(response_obj: Union[ModelResponse, ModelResponseStream]) return delta if isinstance(delta, str) else "" # Handle standard ModelResponse and ModelResponseStream - _choices: Union[List[Union[Choices, StreamingChoices]], List[StreamingChoices]] = ( - response_obj.choices - ) + _choices: Union[ + List[Union[Choices, StreamingChoices]], List[StreamingChoices] + ] = response_obj.choices # Use list accumulation to avoid O(n^2) string concatenation across choices response_parts: List[str] = [] @@ -4982,7 +5049,7 @@ def get_max_tokens(model: str) -> Optional[int]: return litellm.model_cost[model]["max_output_tokens"] elif "max_tokens" in litellm.model_cost[model]: return litellm.model_cost[model]["max_tokens"] - get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') + get_llm_provider = getattr(sys.modules[__name__], "get_llm_provider") model, custom_llm_provider, _, _ = get_llm_provider(model=model) if custom_llm_provider == "huggingface": max_tokens = _get_max_position_embeddings(model_name=model) @@ -5058,7 +5125,7 @@ _model_cost_lowercase_map: Optional[Dict[str, str]] = None def _invalidate_model_cost_lowercase_map() -> None: """Invalidate the case-insensitive lookup map for model_cost. - + Call this whenever litellm.model_cost is modified to ensure the map is rebuilt. """ global _model_cost_lowercase_map @@ -5067,7 +5134,7 @@ def _invalidate_model_cost_lowercase_map() -> None: def _rebuild_model_cost_lowercase_map() -> Dict[str, str]: """Rebuild the case-insensitive lookup map from the current model_cost. - + Returns: The rebuilt map (guaranteed to be not None). """ @@ -5081,9 +5148,9 @@ def _handle_stale_map_entry_rebuild( ) -> Optional[str]: """ Handle stale _model_cost_lowercase_map entry (key was popped). - + Rebuilds the map and retries the lookup. - + Returns: The matched key if found after rebuild, None otherwise. """ @@ -5100,9 +5167,9 @@ def _handle_new_key_with_scan( ) -> Optional[str]: """ Handle new key added to model_cost without invalidating _model_cost_lowercase_map. - + Scans model_cost for case-insensitive match and rebuilds the map if found. - + Returns: The matched key if found, None otherwise. """ @@ -5117,20 +5184,20 @@ def _handle_new_key_with_scan( def _get_model_cost_key(potential_key: str) -> Optional[str]: """ Get the actual key from model_cost, with case-insensitive fallback. - + WARNING: Only O(1) lookup operations are acceptable. O(n) lookups will cause severe CPU overhead. This function is called frequently during router operations. - + ALLOWED HELPER FUNCTIONS (conditionally called, O(n) operations are acceptable): - _rebuild_model_cost_lowercase_map: Rebuilds the lookup map (only when map is None) - _handle_stale_map_entry_rebuild: Rebuilds map when stale entry detected (rare case) - + If you need to add a new helper function with O(n) operations that is conditionally called and confirmed not to cause performance issues, add it to the allowed_helpers list in: tests/code_coverage_tests/check_get_model_cost_key_performance.py """ global _model_cost_lowercase_map - + # Exact match (O(1)) if potential_key in litellm.model_cost: return potential_key @@ -5138,20 +5205,20 @@ def _get_model_cost_key(potential_key: str) -> Optional[str]: # Case-insensitive lookup via map (O(1)) if _model_cost_lowercase_map is None: _model_cost_lowercase_map = _rebuild_model_cost_lowercase_map() - + potential_key_lower = potential_key.lower() matched_key = _model_cost_lowercase_map.get(potential_key_lower) - + # Verify key exists (O(1) - handles model_cost.pop() case) if matched_key is not None and matched_key in litellm.model_cost: return matched_key - + # Rebuild map if stale entry detected (O(n) rebuild, but only when stale entry found) if matched_key is not None: matched_key = _handle_stale_map_entry_rebuild(potential_key_lower) if matched_key is not None: return matched_key - + return None @@ -5183,9 +5250,12 @@ def _check_provider_match(model_info: dict, custom_llm_provider: Optional[str]) custom_llm_provider == "litellm_proxy" ): # litellm_proxy is a special case, it's not a provider, it's a proxy for the provider return True - elif custom_llm_provider == "azure_ai" and model_info["litellm_provider"] in ("azure", "openai"): - # Azure AI also works with azure models - # as a last attempt if the model is not on Azure AI, Azure then fallback to OpenAI cost + elif custom_llm_provider == "azure_ai" and model_info["litellm_provider"] in ( + "azure", + "openai", + ): + # Azure AI also works with azure models + # as a last attempt if the model is not on Azure AI, Azure then fallback to OpenAI cost # tracking the cost is better than attributing 0 cost to it. return True else: @@ -5211,7 +5281,7 @@ def _get_potential_model_names( if custom_llm_provider is None: # Get custom_llm_provider try: - get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') + get_llm_provider = getattr(sys.modules[__name__], "get_llm_provider") split_model, custom_llm_provider, _, _ = get_llm_provider(model=model) except Exception: split_model = model @@ -5941,7 +6011,7 @@ def validate_environment( # noqa: PLR0915 } ## EXTRACT LLM PROVIDER - if model name provided try: - get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') + get_llm_provider = getattr(sys.modules[__name__], "get_llm_provider") _, custom_llm_provider, _, _ = get_llm_provider(model=model) except Exception: custom_llm_provider = None @@ -6504,7 +6574,7 @@ def register_prompt_template( complete_model = model potential_models = [complete_model] try: - get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') + get_llm_provider = getattr(sys.modules[__name__], "get_llm_provider") model = get_llm_provider(model=model)[0] potential_models.append(model) except Exception: @@ -6590,7 +6660,7 @@ class TextCompletionStreamWrapper: except StopIteration: raise StopIteration except Exception as e: - exception_type = getattr(sys.modules[__name__], 'exception_type') + exception_type = getattr(sys.modules[__name__], "exception_type") raise exception_type( model=self.model, custom_llm_provider=self.custom_llm_provider or "", @@ -7084,7 +7154,7 @@ def get_valid_models( # init litellm_params ################################# from litellm.types.router import LiteLLM_Params - + if litellm_params is None: litellm_params = LiteLLM_Params(model="") if api_key is not None: @@ -7214,14 +7284,20 @@ def _get_base_model_from_metadata(model_call_details=None): return _base_model metadata = litellm_params.get("metadata", {}) - _get_base_model_from_litellm_call_metadata = getattr(sys.modules[__name__], '_get_base_model_from_litellm_call_metadata') - base_model_from_metadata = _get_base_model_from_litellm_call_metadata(metadata=metadata) + _get_base_model_from_litellm_call_metadata = getattr( + sys.modules[__name__], "_get_base_model_from_litellm_call_metadata" + ) + base_model_from_metadata = _get_base_model_from_litellm_call_metadata( + metadata=metadata + ) if base_model_from_metadata is not None: return base_model_from_metadata # Also check litellm_metadata (used by Responses API and other generic API calls) litellm_metadata = litellm_params.get("litellm_metadata", {}) - _get_base_model_from_litellm_call_metadata = getattr(sys.modules[__name__], '_get_base_model_from_litellm_call_metadata') + _get_base_model_from_litellm_call_metadata = getattr( + sys.modules[__name__], "_get_base_model_from_litellm_call_metadata" + ) return _get_base_model_from_litellm_call_metadata(metadata=litellm_metadata) return None @@ -7618,7 +7694,7 @@ class ProviderConfigManager: @staticmethod def _build_provider_config_map() -> dict[LlmProviders, tuple[Callable, bool]]: """Build the provider-to-config mapping dictionary. - + Returns a dict mapping provider to (factory_function, needs_model_parameter). This avoids expensive inspect.signature() calls at runtime. """ @@ -7627,12 +7703,30 @@ class ProviderConfigManager: # Format: (factory_function, needs_model_parameter: bool) LlmProviders.OPENAI: (lambda: litellm.OpenAIGPTConfig(), False), LlmProviders.ANTHROPIC: (lambda: litellm.AnthropicConfig(), False), - LlmProviders.AZURE: (lambda model: ProviderConfigManager._get_azure_config(model), True), - LlmProviders.AZURE_AI: (lambda model: ProviderConfigManager._get_azure_ai_config(model), True), - LlmProviders.VERTEX_AI: (lambda model: ProviderConfigManager._get_vertex_ai_config(model), True), - LlmProviders.BEDROCK: (lambda model: ProviderConfigManager._get_bedrock_config(model), True), - LlmProviders.COHERE: (lambda model: ProviderConfigManager._get_cohere_config(model), True), - LlmProviders.COHERE_CHAT: (lambda model: ProviderConfigManager._get_cohere_config(model), True), + LlmProviders.AZURE: ( + lambda model: ProviderConfigManager._get_azure_config(model), + True, + ), + LlmProviders.AZURE_AI: ( + lambda model: ProviderConfigManager._get_azure_ai_config(model), + True, + ), + LlmProviders.VERTEX_AI: ( + lambda model: ProviderConfigManager._get_vertex_ai_config(model), + True, + ), + LlmProviders.BEDROCK: ( + lambda model: ProviderConfigManager._get_bedrock_config(model), + True, + ), + LlmProviders.COHERE: ( + lambda model: ProviderConfigManager._get_cohere_config(model), + True, + ), + LlmProviders.COHERE_CHAT: ( + lambda model: ProviderConfigManager._get_cohere_config(model), + True, + ), # Simple provider mappings (no model parameter needed) LlmProviders.DEEPSEEK: (lambda: litellm.DeepSeekChatConfig(), False), LlmProviders.GROQ: (lambda: litellm.GroqChatConfig(), False), @@ -7642,7 +7736,10 @@ class ProviderConfigManager: LlmProviders.ZAI: (lambda: litellm.ZAIChatConfig(), False), LlmProviders.LAMBDA_AI: (lambda: litellm.LambdaAIChatConfig(), False), LlmProviders.LLAMA: (lambda: litellm.LlamaAPIConfig(), False), - LlmProviders.TEXT_COMPLETION_OPENAI: (lambda: litellm.OpenAITextCompletionConfig(), False), + LlmProviders.TEXT_COMPLETION_OPENAI: ( + lambda: litellm.OpenAITextCompletionConfig(), + False, + ), LlmProviders.SNOWFLAKE: (lambda: litellm.SnowflakeConfig(), False), LlmProviders.CLARIFAI: (lambda: litellm.ClarifaiConfig(), False), LlmProviders.ANTHROPIC_TEXT: (lambda: litellm.AnthropicTextConfig(), False), @@ -7665,7 +7762,10 @@ class ProviderConfigManager: LlmProviders.CUSTOM: (lambda: litellm.OpenAILikeChatConfig(), False), LlmProviders.CUSTOM_OPENAI: (lambda: litellm.OpenAILikeChatConfig(), False), LlmProviders.OPENAI_LIKE: (lambda: litellm.OpenAILikeChatConfig(), False), - LlmProviders.AIOHTTP_OPENAI: (lambda: litellm.AiohttpOpenAIChatConfig(), False), + LlmProviders.AIOHTTP_OPENAI: ( + lambda: litellm.AiohttpOpenAIChatConfig(), + False, + ), LlmProviders.HOSTED_VLLM: (lambda: litellm.HostedVLLMChatConfig(), False), LlmProviders.LLAMAFILE: (lambda: litellm.LlamafileChatConfig(), False), LlmProviders.LM_STUDIO: (lambda: litellm.LMStudioChatConfig(), False), @@ -7674,7 +7774,10 @@ class ProviderConfigManager: LlmProviders.HUGGINGFACE: (lambda: litellm.HuggingFaceChatConfig(), False), LlmProviders.TOGETHER_AI: (lambda: litellm.TogetherAIConfig(), False), LlmProviders.OPENROUTER: (lambda: litellm.OpenrouterConfig(), False), - LlmProviders.VERCEL_AI_GATEWAY: (lambda: litellm.VercelAIGatewayConfig(), False), + LlmProviders.VERCEL_AI_GATEWAY: ( + lambda: litellm.VercelAIGatewayConfig(), + False, + ), LlmProviders.COMETAPI: (lambda: litellm.CometAPIConfig(), False), LlmProviders.DATAROBOT: (lambda: litellm.DataRobotConfig(), False), LlmProviders.GEMINI: (lambda: litellm.GoogleAIStudioGeminiConfig(), False), @@ -7692,7 +7795,10 @@ class ProviderConfigManager: LlmProviders.CEREBRAS: (lambda: litellm.CerebrasConfig(), False), LlmProviders.BASETEN: (lambda: litellm.BasetenConfig(), False), LlmProviders.VOLCENGINE: (lambda: litellm.VolcEngineConfig(), False), - LlmProviders.TEXT_COMPLETION_CODESTRAL: (lambda: litellm.CodestralTextCompletionConfig(), False), + LlmProviders.TEXT_COMPLETION_CODESTRAL: ( + lambda: litellm.CodestralTextCompletionConfig(), + False, + ), LlmProviders.SAMBANOVA: (lambda: litellm.SambanovaConfig(), False), LlmProviders.MARITALK: (lambda: litellm.MaritalkConfig(), False), LlmProviders.VLLM: (lambda: litellm.VLLMConfig(), False), @@ -7700,17 +7806,26 @@ class ProviderConfigManager: LlmProviders.PREDIBASE: (lambda: litellm.PredibaseConfig(), False), LlmProviders.TRITON: (lambda: litellm.TritonConfig(), False), LlmProviders.PETALS: (lambda: litellm.PetalsConfig(), False), - LlmProviders.SAP_GENERATIVE_AI_HUB: (lambda: litellm.GenAIHubOrchestrationConfig(), False), + LlmProviders.SAP_GENERATIVE_AI_HUB: ( + lambda: litellm.GenAIHubOrchestrationConfig(), + False, + ), LlmProviders.FEATHERLESS_AI: (lambda: litellm.FeatherlessAIConfig(), False), LlmProviders.NOVITA: (lambda: litellm.NovitaConfig(), False), LlmProviders.NEBIUS: (lambda: litellm.NebiusConfig(), False), LlmProviders.WANDB: (lambda: litellm.WandbConfig(), False), LlmProviders.DASHSCOPE: (lambda: litellm.DashScopeChatConfig(), False), LlmProviders.MOONSHOT: (lambda: litellm.MoonshotChatConfig(), False), - LlmProviders.DOCKER_MODEL_RUNNER: (lambda: litellm.DockerModelRunnerChatConfig(), False), + LlmProviders.DOCKER_MODEL_RUNNER: ( + lambda: litellm.DockerModelRunnerChatConfig(), + False, + ), LlmProviders.V0: (lambda: litellm.V0ChatConfig(), False), LlmProviders.MORPH: (lambda: litellm.MorphChatConfig(), False), - LlmProviders.LITELLM_PROXY: (lambda: litellm.LiteLLMProxyChatConfig(), False), + LlmProviders.LITELLM_PROXY: ( + lambda: litellm.LiteLLMProxyChatConfig(), + False, + ), LlmProviders.GRADIENT_AI: (lambda: litellm.GradientAIConfig(), False), LlmProviders.NSCALE: (lambda: litellm.NscaleConfig(), False), LlmProviders.HEROKU: (lambda: litellm.HerokuChatConfig(), False), @@ -7718,7 +7833,10 @@ class ProviderConfigManager: LlmProviders.HYPERBOLIC: (lambda: litellm.HyperbolicChatConfig(), False), LlmProviders.OVHCLOUD: (lambda: litellm.OVHCloudChatConfig(), False), LlmProviders.AMAZON_NOVA: (lambda: litellm.AmazonNovaChatConfig(), False), - LlmProviders.LANGGRAPH: (lambda: ProviderConfigManager._get_langgraph_config(), False), + LlmProviders.LANGGRAPH: ( + lambda: ProviderConfigManager._get_langgraph_config(), + False, + ), } @staticmethod @@ -7748,6 +7866,7 @@ class ProviderConfigManager: from litellm.llms.vertex_ai.vertex_ai_partner_models.gpt_oss.transformation import ( VertexAIGPTOSSTransformation, ) + return VertexAIGPTOSSTransformation() elif model in litellm.vertex_mistral_models: if "codestral" in model: @@ -7762,12 +7881,13 @@ class ProviderConfigManager: def _get_bedrock_config(model: str) -> BaseConfig: """Get Bedrock config based on model.""" from litellm.llms.bedrock.common_utils import get_bedrock_chat_config + return get_bedrock_chat_config(model=model) @staticmethod def _get_cohere_config(model: str) -> BaseConfig: """Get Cohere config based on route.""" - CohereModelInfo = getattr(sys.modules[__name__], 'CohereModelInfo') + CohereModelInfo = getattr(sys.modules[__name__], "CohereModelInfo") route = CohereModelInfo.get_cohere_route(model) if route == "v2": return litellm.CohereV2ChatConfig() @@ -7777,6 +7897,7 @@ class ProviderConfigManager: def _get_langgraph_config() -> BaseConfig: """Get LangGraph config.""" from litellm.llms.langgraph.chat.transformation import LangGraphConfig + return LangGraphConfig() @staticmethod @@ -7785,7 +7906,7 @@ class ProviderConfigManager: ) -> Optional[BaseConfig]: """ Returns the provider config for a given provider. - + Uses O(1) dictionary lookup for fast provider resolution. """ # Check JSON providers FIRST (these override standard mappings) @@ -7807,7 +7928,9 @@ class ProviderConfigManager: # Initialize provider config map lazily (avoids circular imports) if ProviderConfigManager._PROVIDER_CONFIG_MAP is None: - ProviderConfigManager._PROVIDER_CONFIG_MAP = ProviderConfigManager._build_provider_config_map() + ProviderConfigManager._PROVIDER_CONFIG_MAP = ( + ProviderConfigManager._build_provider_config_map() + ) # O(1) dictionary lookup config_entry = ProviderConfigManager._PROVIDER_CONFIG_MAP.get(provider) @@ -7877,6 +8000,7 @@ class ProviderConfigManager: from litellm.llms.openrouter.embedding.transformation import ( OpenrouterEmbeddingConfig, ) + return OpenrouterEmbeddingConfig() elif litellm.LlmProviders.GIGACHAT == provider: return litellm.GigaChatEmbeddingConfig() @@ -8490,8 +8614,8 @@ class ProviderConfigManager: from litellm.llms.vertex_ai.ocr.common_utils import get_vertex_ai_ocr_config return get_vertex_ai_ocr_config(model=model) - - MistralOCRConfig = getattr(sys.modules[__name__], 'MistralOCRConfig') + + MistralOCRConfig = getattr(sys.modules[__name__], "MistralOCRConfig") PROVIDER_TO_CONFIG_MAP = { litellm.LlmProviders.MISTRAL: MistralOCRConfig, } @@ -8632,7 +8756,9 @@ def get_end_user_id_for_cost_tracking( service_type: "litellm_logging" or "prometheus" - used to allow prometheus only disable cost tracking. """ - get_litellm_metadata_from_kwargs = getattr(sys.modules[__name__], 'get_litellm_metadata_from_kwargs') + get_litellm_metadata_from_kwargs = getattr( + sys.modules[__name__], "get_litellm_metadata_from_kwargs" + ) _metadata = cast( dict, get_litellm_metadata_from_kwargs(dict(litellm_params=litellm_params)) ) @@ -8928,12 +9054,12 @@ def __getattr__(name: str) -> Any: """Lazy import handler for utils module with cached registry for improved performance.""" # Use cached registry from _lazy_imports instead of importing tuples every time from litellm._lazy_imports import _get_lazy_import_registry - + registry = _get_lazy_import_registry() - + # Check if name is in registry and call the cached handler function if name in registry: handler_func = registry[name] return handler_func(name) - + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py index 26c60b2c303..2f2eaa905be 100644 --- a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py @@ -26,30 +26,28 @@ def test_google_generate_content_endpoint(): from litellm.proxy.google_endpoints.endpoints import router as google_router except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - + # Create a FastAPI app and include the router (required for FastAPI 0.120+) app = FastAPI() app.include_router(google_router) - + # Create a test client client = TestClient(app) - + # Mock the router's agenerate_content method with patch("litellm.proxy.proxy_server.llm_router") as mock_router: mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - + # Send a request to the endpoint response = client.post( "/v1beta/models/test-model:generateContent", - json={ - "contents": [{"role": "user", "parts": [{"text": "Hello"}]}] - } + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, ) - + # Verify the response assert response.status_code == 200 assert response.json() == {"test": "response"} - + # Verify that agenerate_content was called mock_router.agenerate_content.assert_called_once() @@ -64,40 +62,42 @@ def test_google_stream_generate_content_endpoint(): from litellm.proxy.google_endpoints.endpoints import router as google_router except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - + # Create a FastAPI app and include the router (required for FastAPI 0.120+) app = FastAPI() app.include_router(google_router) - + # Create a test client client = TestClient(app) - + # Mock the router's agenerate_content_stream method to return a stream async def mock_stream_generator(): yield 'data: {"test": "stream_chunk_1"}\n\n' yield 'data: {"test": "stream_chunk_2"}\n\n' yield "data: [DONE]\n\n" - + with patch("litellm.proxy.proxy_server.llm_router") as mock_router: - mock_router.agenerate_content_stream = AsyncMock(return_value=mock_stream_generator()) - + mock_router.agenerate_content_stream = AsyncMock( + return_value=mock_stream_generator() + ) + # Send a request to the endpoint response = client.post( "/v1beta/models/test-model:streamGenerateContent", - json={ - "contents": [{"role": "user", "parts": [{"text": "Hello"}]}] - } + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, ) - + # Verify the response assert response.status_code == 200 - + # Verify that agenerate_content_stream was called with correct parameters mock_router.agenerate_content_stream.assert_called_once() call_args = mock_router.agenerate_content_stream.call_args assert call_args[1]["stream"] is True assert call_args[1]["model"] == "test-model" - assert call_args[1]["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] + assert call_args[1]["contents"] == [ + {"role": "user", "parts": [{"text": "Hello"}]} + ] def test_google_generate_content_with_cost_tracking_metadata(): @@ -110,25 +110,28 @@ def test_google_generate_content_with_cost_tracking_metadata(): from litellm.proxy.google_endpoints.endpoints import router as google_router except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - + # Create a FastAPI app and include the router (required for FastAPI 0.120+) app = FastAPI() app.include_router(google_router) - + # Create a test client client = TestClient(app) - + # Mock all required proxy server dependencies - with patch("litellm.proxy.proxy_server.llm_router") as mock_router, \ - patch("litellm.proxy.proxy_server.general_settings", {}), \ - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, \ - patch("litellm.proxy.proxy_server.version", "1.0.0"), \ - patch("litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request") as mock_add_data: - + with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch( + "litellm.proxy.proxy_server.general_settings", {} + ), patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, patch( + "litellm.proxy.proxy_server.version", "1.0.0" + ), patch( + "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" + ) as mock_add_data: mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - + # Mock add_litellm_data_to_request to return data with metadata - async def mock_add_litellm_data(data, request, user_api_key_dict, proxy_config, general_settings, version): + async def mock_add_litellm_data( + data, request, user_api_key_dict, proxy_config, general_settings, version + ): # Simulate adding user metadata data["litellm_metadata"] = { "user_api_key_user_id": "test-user-id", @@ -136,29 +139,27 @@ def test_google_generate_content_with_cost_tracking_metadata(): "user_api_key": "hashed-key", } return data - + mock_add_data.side_effect = mock_add_litellm_data - + # Send a request to the endpoint response = client.post( "/v1beta/models/test-model:generateContent", - json={ - "contents": [{"role": "user", "parts": [{"text": "Hello"}]}] - }, - headers={"Authorization": "Bearer sk-test-key"} + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, + headers={"Authorization": "Bearer sk-test-key"}, ) - + # Verify the response assert response.status_code == 200 - + # Verify that add_litellm_data_to_request was called mock_add_data.assert_called_once() - + # Verify that agenerate_content was called with metadata mock_router.agenerate_content.assert_called_once() call_args = mock_router.agenerate_content.call_args called_data = call_args[1] - + # Verify that litellm_metadata exists and contains user information assert "litellm_metadata" in called_data assert called_data["litellm_metadata"]["user_api_key_user_id"] == "test-user-id" @@ -174,30 +175,33 @@ def test_google_stream_generate_content_with_cost_tracking_metadata(): from litellm.proxy.google_endpoints.endpoints import router as google_router except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - + # Create a FastAPI app and include the router (required for FastAPI 0.120+) app = FastAPI() app.include_router(google_router) - + # Create a test client client = TestClient(app) - + # Mock the router's agenerate_content_stream method to return a stream mock_stream = AsyncMock() mock_stream.__aiter__ = lambda self: mock_stream mock_stream.__anext__.side_effect = StopAsyncIteration - + # Mock all required proxy server dependencies - with patch("litellm.proxy.proxy_server.llm_router") as mock_router, \ - patch("litellm.proxy.proxy_server.general_settings", {}), \ - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, \ - patch("litellm.proxy.proxy_server.version", "1.0.0"), \ - patch("litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request") as mock_add_data: - + with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch( + "litellm.proxy.proxy_server.general_settings", {} + ), patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, patch( + "litellm.proxy.proxy_server.version", "1.0.0" + ), patch( + "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" + ) as mock_add_data: mock_router.agenerate_content_stream = AsyncMock(return_value=mock_stream) - + # Mock add_litellm_data_to_request to return data with metadata - async def mock_add_litellm_data(data, request, user_api_key_dict, proxy_config, general_settings, version): + async def mock_add_litellm_data( + data, request, user_api_key_dict, proxy_config, general_settings, version + ): # Simulate adding user metadata data["litellm_metadata"] = { "user_api_key_user_id": "test-user-id", @@ -205,29 +209,27 @@ def test_google_stream_generate_content_with_cost_tracking_metadata(): "user_api_key": "hashed-key", } return data - + mock_add_data.side_effect = mock_add_litellm_data - + # Send a request to the endpoint response = client.post( "/v1beta/models/test-model:streamGenerateContent", - json={ - "contents": [{"role": "user", "parts": [{"text": "Hello"}]}] - }, - headers={"Authorization": "Bearer sk-test-key"} + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, + headers={"Authorization": "Bearer sk-test-key"}, ) - + # Verify the response assert response.status_code == 200 - + # Verify that add_litellm_data_to_request was called mock_add_data.assert_called_once() - + # Verify that agenerate_content_stream was called with metadata mock_router.agenerate_content_stream.assert_called_once() call_args = mock_router.agenerate_content_stream.call_args called_data = call_args[1] - + # Verify that litellm_metadata exists and contains user information assert "litellm_metadata" in called_data assert called_data["litellm_metadata"]["user_api_key_user_id"] == "test-user-id" @@ -239,7 +241,7 @@ def test_google_stream_generate_content_with_cost_tracking_metadata(): def test_google_generate_content_with_system_instruction(): """ Test that systemInstruction is correctly passed through from the endpoint to the router. - + This test verifies the fix for systemInstruction being dropped when forwarding requests to Vertex AI through the Google GenAI endpoint. """ @@ -250,62 +252,63 @@ def test_google_generate_content_with_system_instruction(): from litellm.proxy.google_endpoints.endpoints import router as google_router except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - + # Create a FastAPI app and include the router app = FastAPI() app.include_router(google_router) - + # Create a test client client = TestClient(app) - + # Mock all required proxy server dependencies - with patch("litellm.proxy.proxy_server.llm_router") as mock_router, \ - patch("litellm.proxy.proxy_server.general_settings", {}), \ - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, \ - patch("litellm.proxy.proxy_server.version", "1.0.0"), \ - patch("litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request") as mock_add_data: - + with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch( + "litellm.proxy.proxy_server.general_settings", {} + ), patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, patch( + "litellm.proxy.proxy_server.version", "1.0.0" + ), patch( + "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" + ) as mock_add_data: mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - + # Mock add_litellm_data_to_request to pass through data unchanged - async def mock_add_litellm_data(data, request, user_api_key_dict, proxy_config, general_settings, version): + async def mock_add_litellm_data( + data, request, user_api_key_dict, proxy_config, general_settings, version + ): return data - + mock_add_data.side_effect = mock_add_litellm_data - + # Define the systemInstruction to test - system_instruction = { - "parts": [{"text": "Your name is Doodle."}] - } - + system_instruction = {"parts": [{"text": "Your name is Doodle."}]} + # Send a request with systemInstruction response = client.post( "/v1beta/models/gemini-2.5-pro:generateContent", json={ "systemInstruction": system_instruction, "contents": [ - { - "parts": [{"text": "What is your name?"}], - "role": "user" - } - ] + {"parts": [{"text": "What is your name?"}], "role": "user"} + ], }, - headers={"Authorization": "Bearer sk-test-key"} + headers={"Authorization": "Bearer sk-test-key"}, ) - + # Verify the response assert response.status_code == 200 - + # Verify that agenerate_content was called mock_router.agenerate_content.assert_called_once() call_args = mock_router.agenerate_content.call_args called_data = call_args[1] - + # Verify that systemInstruction is present in the call arguments assert "systemInstruction" in called_data assert called_data["systemInstruction"] == system_instruction - assert called_data["systemInstruction"]["parts"][0]["text"] == "Your name is Doodle." - + assert ( + called_data["systemInstruction"]["parts"][0]["text"] + == "Your name is Doodle." + ) + # Verify contents are also present assert "contents" in called_data assert len(called_data["contents"]) == 1 @@ -315,7 +318,7 @@ def test_google_generate_content_with_system_instruction(): def test_google_generate_content_with_image_config(): """ Test that imageConfig is correctly passed through from generationConfig to the router. - + This test verifies that imageConfig parameters (aspectRatio, imageSize) are preserved when forwarding requests to Google GenAI through the endpoint. """ @@ -326,69 +329,75 @@ def test_google_generate_content_with_image_config(): from litellm.proxy.google_endpoints.endpoints import router as google_router except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - + # Create a FastAPI app and include the router app = FastAPI() app.include_router(google_router) - + # Create a test client client = TestClient(app) - + # Mock all required proxy server dependencies - with patch("litellm.proxy.proxy_server.llm_router") as mock_router, \ - patch("litellm.proxy.proxy_server.general_settings", {}), \ - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, \ - patch("litellm.proxy.proxy_server.version", "1.0.0"), \ - patch("litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request") as mock_add_data: - + with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch( + "litellm.proxy.proxy_server.general_settings", {} + ), patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, patch( + "litellm.proxy.proxy_server.version", "1.0.0" + ), patch( + "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" + ) as mock_add_data: mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - + # Mock add_litellm_data_to_request to pass through data unchanged - async def mock_add_litellm_data(data, request, user_api_key_dict, proxy_config, general_settings, version): + async def mock_add_litellm_data( + data, request, user_api_key_dict, proxy_config, general_settings, version + ): return data - + mock_add_data.side_effect = mock_add_litellm_data - + # Send a request with generationConfig containing imageConfig response = client.post( "/v1beta/models/gemini-3-pro-image-preview:generateContent", json={ - "contents": [{ - "role": "user", - "parts": [{"text": "Create a vibrant infographic about photosynthesis"}] - }], + "contents": [ + { + "role": "user", + "parts": [ + { + "text": "Create a vibrant infographic about photosynthesis" + } + ], + } + ], "generationConfig": { "responseModalities": ["TEXT", "IMAGE"], - "imageConfig": { - "aspectRatio": "9:16", - "imageSize": "4K" - } - } + "imageConfig": {"aspectRatio": "9:16", "imageSize": "4K"}, + }, }, - headers={"Authorization": "Bearer sk-test-key"} + headers={"Authorization": "Bearer sk-test-key"}, ) - + # Verify the response assert response.status_code == 200 - + # Verify that agenerate_content was called mock_router.agenerate_content.assert_called_once() call_args = mock_router.agenerate_content.call_args called_data = call_args[1] - + # Verify that config is present in the call arguments assert "config" in called_data - + # Verify that imageConfig is preserved in the config assert "imageConfig" in called_data["config"] assert called_data["config"]["imageConfig"]["aspectRatio"] == "9:16" assert called_data["config"]["imageConfig"]["imageSize"] == "4K" - + # Verify that responseModalities is also preserved assert "responseModalities" in called_data["config"] assert called_data["config"]["responseModalities"] == ["TEXT", "IMAGE"] - + # Verify contents are also present assert "contents" in called_data assert len(called_data["contents"]) == 1 - assert called_data["contents"][0]["role"] == "user" \ No newline at end of file + assert called_data["contents"][0]["role"] == "user" From fc9988b686ba23fd8da6c469e77380f620e87622 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Mon, 19 Jan 2026 19:48:02 +0530 Subject: [PATCH 12/58] fix/bedrock-inconsistent-postcall-hook (#19151) * fix/bedrock-inconsistent-postcall-hook * Add condition check to avoid multiple validation --- .../guardrail_hooks/bedrock_guardrails.py | 110 +++++++++++++----- 1 file changed, 83 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 59032189438..d3e8a74945f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -893,15 +893,53 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return ######################################################### - ########## 1. Make parallel Bedrock API requests ########## + ########## 1. Make Bedrock API requests ########## ######################################################### + # Import asyncio for parallel execution + import asyncio + + # Determine if INPUT validation is needed in post_call + # Skip INPUT validation if pre_call or during_call is already enabled + # (to avoid redundant validation - those hooks would have already validated INPUT) + should_validate_input = not ( + self._event_hook_is_event_type(GuardrailEventHooks.pre_call) + or self._event_hook_is_event_type(GuardrailEventHooks.during_call) + ) + output_content_bedrock: Optional[Union[BedrockGuardrailResponse, str]] = None - try: - output_content_bedrock = await self.make_bedrock_api_request( + + if should_validate_input: + # Prepare input messages (with optional filtering for latest role message) + input_filter = self._prepare_guardrail_messages_for_role( + messages=new_messages + ) + input_messages = input_filter.payload_messages or new_messages + + # Create tasks for parallel execution of both INPUT and OUTPUT validation + input_task = self.make_bedrock_api_request( + source="INPUT", + messages=input_messages, + request_data=data, + ) + output_task = self.make_bedrock_api_request( source="OUTPUT", response=response, request_data=data ) - except GuardrailInterventionNormalStringError as e: - output_content_bedrock = e.message + + # Execute both requests in parallel + try: + _, output_content_bedrock = await asyncio.gather( + input_task, output_task + ) + except GuardrailInterventionNormalStringError as e: + output_content_bedrock = e.message + else: + # Only run OUTPUT validation (INPUT was already validated in pre_call or during_call) + try: + output_content_bedrock = await self.make_bedrock_api_request( + source="OUTPUT", response=response, request_data=data + ) + except GuardrailInterventionNormalStringError as e: + output_content_bedrock = e.message ######################################################### ########## 2. Apply masking to response with output guardrail response ########## @@ -910,7 +948,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): response = self.create_guardrail_blocked_response( response=output_content_bedrock ) - else: + elif output_content_bedrock is not None: self._apply_masking_to_response( response=response, bedrock_guardrail_response=output_content_bedrock, @@ -997,36 +1035,54 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) if isinstance(assembled_model_response, ModelResponse): #################################################################### - ########## 1. Make parallel Bedrock Apply Guardrail API requests ########## + ########## 1. Make Bedrock Apply Guardrail API requests ########## # Bedrock will raise an exception if this violates the guardrail policy ################################################################### - # Create tasks for parallel execution - input_filter = self._prepare_guardrail_messages_for_role( - messages=request_data.get("messages") + # Determine if INPUT validation is needed in post_call + # Skip INPUT validation if pre_call or during_call is already enabled + # (to avoid redundant validation - those hooks would have already validated INPUT) + should_validate_input = not ( + self._event_hook_is_event_type(GuardrailEventHooks.pre_call) + or self._event_hook_is_event_type(GuardrailEventHooks.during_call) ) - input_messages = input_filter.payload_messages or request_data.get( - "messages" - ) - input_task = self.make_bedrock_api_request( - source="INPUT", - messages=input_messages, - request_data=request_data, - ) # Only input messages + output_guardrail_response: Optional[ Union[BedrockGuardrailResponse, str] ] = None - output_task = self.make_bedrock_api_request( - source="OUTPUT", response=assembled_model_response - ) # Only response - # Execute both requests in parallel - try: - _, output_guardrail_response = await asyncio.gather( - input_task, output_task + if should_validate_input: + # Create tasks for parallel execution + input_filter = self._prepare_guardrail_messages_for_role( + messages=request_data.get("messages") ) - except GuardrailInterventionNormalStringError as e: - output_guardrail_response = e.message + input_messages = input_filter.payload_messages or request_data.get( + "messages" + ) + input_task = self.make_bedrock_api_request( + source="INPUT", + messages=input_messages, + request_data=request_data, + ) # Only input messages + output_task = self.make_bedrock_api_request( + source="OUTPUT", response=assembled_model_response + ) # Only response + + # Execute both requests in parallel + try: + _, output_guardrail_response = await asyncio.gather( + input_task, output_task + ) + except GuardrailInterventionNormalStringError as e: + output_guardrail_response = e.message + else: + # Only run OUTPUT validation (INPUT was already validated in pre_call or during_call) + try: + output_guardrail_response = await self.make_bedrock_api_request( + source="OUTPUT", response=assembled_model_response + ) + except GuardrailInterventionNormalStringError as e: + output_guardrail_response = e.message ######################################################################### ########## 2. Apply masking to response with output guardrail response ########## From 0367e9c9f11b7f8f68166a7d1bb68c56523a8776 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Benedikt=20=C3=93skarsson?= Date: Mon, 19 Jan 2026 14:50:57 +0000 Subject: [PATCH 13/58] fix: pr ammends. --- litellm/llms/bedrock/chat/converse_transformation.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index fd8c0d75473..cb26a22edfa 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -55,6 +55,7 @@ from litellm.types.utils import ( ) from litellm.utils import ( add_dummy_tool, + any_assistant_message_has_thinking_blocks, has_tool_call_blocks, last_assistant_with_tool_calls_has_no_thinking_blocks, supports_reasoning, @@ -1077,11 +1078,15 @@ class AmazonConverseConfig(BaseConfig): # Drop thinking param if thinking is enabled but thinking_blocks are missing # This prevents the error: "Expected thinking or redacted_thinking, but found tool_use" + # + # IMPORTANT: Only drop thinking if NO assistant messages have thinking_blocks. + # If any message has thinking_blocks, we must keep thinking enabled, otherwise # Related issues: https://github.com/BerriAI/litellm/issues/14194 if ( optional_params.get("thinking") is not None and messages is not None and last_assistant_with_tool_calls_has_no_thinking_blocks(messages) + and not any_assistant_message_has_thinking_blocks(messages) ): if litellm.modify_params: optional_params.pop("thinking", None) From 1678f621db8305c7e5f9366f5ee46c7f4edb9a47 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Mon, 19 Jan 2026 23:57:01 +0530 Subject: [PATCH 14/58] feat: add retry_delay, exponential_backoff, and jitter to completion() (#19371) --- docs/my-website/docs/completion/input.md | 6 + .../docs/completion/reliable_completions.md | 30 +++ .../litellm_core_utils/get_litellm_params.py | 6 + litellm/main.py | 181 +++++++++++++++--- litellm/types/utils.py | 8 +- 5 files changed, 204 insertions(+), 27 deletions(-) diff --git a/docs/my-website/docs/completion/input.md b/docs/my-website/docs/completion/input.md index 2f6da4bedcd..1272c7eb161 100644 --- a/docs/my-website/docs/completion/input.md +++ b/docs/my-website/docs/completion/input.md @@ -264,6 +264,12 @@ messages=[{"role": "user", "content": [ - `num_retries`: *int (optional)* - The number of times to retry the API call if an APIError, TimeoutError or ServiceUnavailableError occurs +- `retry_delay`: *float (optional)* - Time in seconds to wait between retries. + +- `exponential_backoff`: *float (optional)* - If true, wait time doubles after each failure (1s, 2s, 4s...) + +- `jitter`: *float (optional)* - If true, adds randomness to the wait time to prevent thundering herd. + - `context_window_fallback_dict`: *dict (optional)* - A mapping of model to use if call fails due to context window error - `fallbacks`: *list (optional)* - A list of model names + params to be used, in case the initial call fails diff --git a/docs/my-website/docs/completion/reliable_completions.md b/docs/my-website/docs/completion/reliable_completions.md index f38917fe53d..431df66c9dc 100644 --- a/docs/my-website/docs/completion/reliable_completions.md +++ b/docs/my-website/docs/completion/reliable_completions.md @@ -31,6 +31,36 @@ response = completion( ) ``` +### Configurable Retries (Exponential Backoff, Jitter) + +You can also specify the `retry_delay` and `exponential_backoff` or `jitter` to the completion call. + +* `retry_delay`: (float) Time in seconds to wait between retries. +* `exponential_backoff`: (bool) If true, wait time doubles after each failure (1s, 2s, 4s...). +* `jitter`: (bool) If true, adds randomness to the wait time to prevent thundering herd. + +```python +import litellm + +# Exponential Backoff +response = litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hi"}], + num_retries=3, + retry_delay=1.0, # Start with 1s wait + exponential_backoff=True # Wait 1s, 2s, 4s... +) + +# Jitter +response = litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hi"}], + num_retries=3, + retry_delay=1.0, + jitter=True # Randomize wait times +) +``` + ## Fallbacks (SDK) :::info diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 0d35cfa3140..e9740ed1da6 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -64,6 +64,9 @@ def get_litellm_params( api_version: Optional[str] = None, max_retries: Optional[int] = None, litellm_request_debug: Optional[bool] = None, + retry_delay: Optional[float] = None, + exponential_backoff: Optional[bool] = None, + jitter: Optional[bool] = None, **kwargs, ) -> dict: litellm_params = { @@ -116,6 +119,9 @@ def get_litellm_params( "azure_password": kwargs.get("azure_password"), "azure_scope": kwargs.get("azure_scope"), "max_retries": max_retries, + "retry_delay": retry_delay, + "exponential_backoff": exponential_backoff, + "jitter": jitter, "timeout": kwargs.get("timeout"), "bucket_name": kwargs.get("bucket_name"), "vertex_credentials": kwargs.get("vertex_credentials"), diff --git a/litellm/main.py b/litellm/main.py index 969cf55a3d6..08384960b16 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -415,6 +415,11 @@ async def acompletion( web_search_options: Optional[OpenAIWebSearchOptions] = None, # Session management shared_session: Optional["ClientSession"] = None, + # Retry params + retry_delay: Optional[float] = None, + exponential_backoff: Optional[bool] = None, + jitter: Optional[bool] = None, + num_retries: Optional[int] = None, **kwargs, ) -> Union[ModelResponse, CustomStreamWrapper]: """ @@ -460,6 +465,16 @@ async def acompletion( - The `completion` function is called using `run_in_executor` to execute synchronously in the event loop. - If `stream` is True, the function returns an async generator that yields completion lines. """ + if _update_kwargs_with_retry_params( + retry_delay, + exponential_backoff, + jitter, + num_retries, + locals().get("max_retries"), + kwargs, + ): + return await acompletion_with_retries(model=model, messages=messages, **kwargs) + fallbacks = kwargs.get("fallbacks", None) mock_timeout = kwargs.get("mock_timeout", None) @@ -988,8 +1003,32 @@ def _drop_input_examples_from_tools( return cleaned_tools -@tracer.wrap() -@client +def _update_kwargs_with_retry_params( + retry_delay: Optional[float], + exponential_backoff: Optional[bool], + jitter: Optional[bool], + num_retries: Optional[int], + max_retries: Optional[int], + kwargs: dict, +) -> bool: + """ + Updates kwargs with retry parameters if any are provided. + Returns True if retry logic should be triggered, False otherwise. + """ + if retry_delay is not None or exponential_backoff is not None or jitter is not None: + kwargs["retry_delay"] = retry_delay + kwargs["exponential_backoff"] = exponential_backoff + kwargs["jitter"] = jitter + + if num_retries is not None: + kwargs["num_retries"] = num_retries + elif max_retries is not None and "num_retries" not in kwargs: + kwargs["num_retries"] = max_retries + + return True + return False + + def completion( # type: ignore # noqa: PLR0915 model: str, # Optional OpenAI params: see https://platform.openai.com/docs/api-reference/chat/create @@ -1039,6 +1078,11 @@ def completion( # type: ignore # noqa: PLR0915 thinking: Optional[AnthropicThinkingParam] = None, # Session management shared_session: Optional["ClientSession"] = None, + # Retry params + retry_delay: Optional[float] = None, + exponential_backoff: Optional[bool] = None, + jitter: Optional[bool] = None, + num_retries: Optional[int] = None, **kwargs, ) -> Union[ModelResponse, CustomStreamWrapper]: """ @@ -1086,6 +1130,20 @@ def completion( # type: ignore # noqa: PLR0915 - It supports various optional parameters for customizing the completion behavior. - If 'mock_response' is provided, a mock completion response is returned for testing or debugging. """ + if _update_kwargs_with_retry_params( + retry_delay, + exponential_backoff, + jitter, + num_retries, + locals().get("max_retries"), + kwargs, + ): + # check if this is an async call (acompletion=True in kwargs) + if kwargs.get("acompletion", False) is True: + return acompletion_with_retries(model=model, messages=messages, **kwargs) + + return completion_with_retries(model=model, messages=messages, **kwargs) + ### VALIDATE Request ### if model is None: raise ValueError("model param not passed in.") @@ -1111,7 +1169,9 @@ def completion( # type: ignore # noqa: PLR0915 # Check if MCP tools are present (following responses pattern) # Cast tools to Optional[Iterable[ToolParam]] for type checking tools_for_mcp = cast(Optional[Iterable[ToolParam]], tools) - if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools_for_mcp): + if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway( + tools=tools_for_mcp + ): # Return coroutine - acompletion will await it # completion() can return a coroutine when MCP tools are present, which acompletion() awaits return acompletion_with_mcp( # type: ignore[return-value] @@ -2337,11 +2397,7 @@ def completion( # type: ignore # noqa: PLR0915 input=messages, api_key=api_key, original_response=response ) elif custom_llm_provider == "minimax": - api_key = ( - api_key - or get_secret_str("MINIMAX_API_KEY") - or litellm.api_key - ) + api_key = api_key or get_secret_str("MINIMAX_API_KEY") or litellm.api_key api_base = ( api_base @@ -2389,7 +2445,9 @@ def completion( # type: ignore # noqa: PLR0915 or custom_llm_provider == "wandb" or custom_llm_provider == "clarifai" or custom_llm_provider in litellm.openai_compatible_providers - or JSONProviderRegistry.exists(custom_llm_provider) # JSON-configured providers + or JSONProviderRegistry.exists( + custom_llm_provider + ) # JSON-configured providers or "ft:gpt-3.5-turbo" in model # finetune gpt-3.5-turbo ): # allow user to make an openai call with a custom base # note: if a user sets a custom base - we should ensure this works @@ -4229,24 +4287,51 @@ def completion_with_retries(*args, **kwargs): retry_strategy: Literal["exponential_backoff_retry", "constant_retry"] = kwargs.pop( "retry_strategy", "constant_retry" ) # type: ignore + retry_delay = kwargs.pop("retry_delay", None) + exponential_backoff = kwargs.pop("exponential_backoff", False) + jitter = kwargs.pop("jitter", False) + original_function = kwargs.pop("original_function", completion) - if retry_strategy == "exponential_backoff_retry": + + # +1 because stop_after_attempt includes the initial attempt + stop_after = tenacity.stop_after_attempt(num_retries + 1) + + if retry_strategy == "exponential_backoff_retry" or exponential_backoff: + # Defaults for exponential backoff + multiplier = 1 + min_wait = 0 + if retry_delay is not None: + multiplier = retry_delay + min_wait = retry_delay + + wait_args = {"multiplier": multiplier, "min": min_wait, "max": 10} + + if jitter: + wait_strategy = tenacity.wait_random_exponential(**wait_args) + else: + wait_strategy = tenacity.wait_exponential(**wait_args) + retryer = tenacity.Retrying( - wait=tenacity.wait_exponential(multiplier=1, max=10), - stop=tenacity.stop_after_attempt(num_retries), + wait=wait_strategy, + stop=stop_after, reraise=True, ) else: + wait_strategy = tenacity.wait_none() + if retry_delay: + wait_strategy = tenacity.wait_fixed(retry_delay) + retryer = tenacity.Retrying( - stop=tenacity.stop_after_attempt(num_retries), reraise=True + wait=wait_strategy, + stop=stop_after, + reraise=True, ) return retryer(original_function, *args, **kwargs) async def acompletion_with_retries(*args, **kwargs): """ - [DEPRECATED]. Use 'acompletion' or router.acompletion instead! - Executes a litellm.completion() with 3 retries + Executes a litellm.completion() with retries. """ try: import tenacity @@ -4259,18 +4344,65 @@ async def acompletion_with_retries(*args, **kwargs): kwargs["max_retries"] = 0 kwargs["num_retries"] = 0 retry_strategy = kwargs.pop("retry_strategy", "constant_retry") + retry_delay = kwargs.pop("retry_delay", None) + exponential_backoff = kwargs.pop("exponential_backoff", False) + jitter = kwargs.pop("jitter", False) original_function = kwargs.pop("original_function", completion) - if retry_strategy == "exponential_backoff_retry": + + # +1 because stop_after_attempt includes the initial attempt + stop_after = tenacity.stop_after_attempt(num_retries + 1) + + # If the original function is completion but we are doing async retries + # we need to ensure it's treated as an async function if it returns a coroutine + # or wraps it. + # However, since we are in acompletion_with_retries, we expect to be waiting. + # If original_function is completion(acompletion=True), it returns a coro. + # Tenacity AsyncRetrying expects the function to be awaitable or return awaitable? + # Actually AsyncRetrying works with async def functions. + # If original_function is sync (but returns coro), we might need a wrapper. + + async def _async_original_function(*args, **kwargs): + # Ensure we await the result if it is a coroutine + result = original_function(*args, **kwargs) + if asyncio.iscoroutine(result): + return await result + return result + + if retry_strategy == "exponential_backoff_retry" or exponential_backoff: + # Defaults for exponential backoff + multiplier = 1 + min_wait = 0 + if retry_delay is not None: + multiplier = retry_delay + min_wait = retry_delay + + wait_args = {"multiplier": multiplier, "min": min_wait, "max": 10} + + if jitter: + # Use wait_random_exponential if available or combine + # Using wait_exponential + wait_random is one way, or wait_random_exponential + # tenacity.wait_random_exponential(multiplier=1, max=10) + wait_strategy = tenacity.wait_random_exponential(**wait_args) + else: + wait_strategy = tenacity.wait_exponential(**wait_args) + retryer = tenacity.AsyncRetrying( - wait=tenacity.wait_exponential(multiplier=1, max=10), - stop=tenacity.stop_after_attempt(num_retries), + wait=wait_strategy, + stop=stop_after, reraise=True, ) else: + wait_strategy = tenacity.wait_none() + if retry_delay: + wait_strategy = tenacity.wait_fixed(retry_delay) + retryer = tenacity.AsyncRetrying( - stop=tenacity.stop_after_attempt(num_retries), reraise=True + wait=wait_strategy, + stop=stop_after, + reraise=True, ) - return await retryer(original_function, *args, **kwargs) + + return await retryer(_async_original_function, *args, **kwargs) def responses_with_retries(*args, **kwargs): @@ -4700,7 +4832,7 @@ def embedding( # noqa: PLR0915 if headers is not None and headers != {}: optional_params["extra_headers"] = headers - + if encoding_format is not None: optional_params["encoding_format"] = encoding_format else: @@ -6735,9 +6867,7 @@ def speech( # noqa: PLR0915 if text_to_speech_provider_config is None: text_to_speech_provider_config = MinimaxTextToSpeechConfig() - minimax_config = cast( - MinimaxTextToSpeechConfig, text_to_speech_provider_config - ) + minimax_config = cast(MinimaxTextToSpeechConfig, text_to_speech_provider_config) if api_base is not None: litellm_params_dict["api_base"] = api_base @@ -6877,7 +7007,7 @@ async def ahealth_check( custom_llm_provider_from_params = model_params.get("custom_llm_provider", None) api_base_from_params = model_params.get("api_base", None) api_key_from_params = model_params.get("api_key", None) - + model, custom_llm_provider, _, _ = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider_from_params, @@ -7251,6 +7381,7 @@ def __getattr__(name: str) -> Any: _encoding = tiktoken.get_encoding("cl100k_base") # Cache it in the module's __dict__ for subsequent accesses import sys + sys.modules[__name__].__dict__["encoding"] = _encoding global _encoding_cache _encoding_cache = _encoding diff --git a/litellm/types/utils.py b/litellm/types/utils.py index fd83b7e79f0..e1f1c1c4b5b 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -112,13 +112,14 @@ class SearchContextCostPerQuery(TypedDict, total=False): class AgenticLoopParams(TypedDict, total=False): """ Parameters passed to agentic loop hooks (e.g., WebSearch interception). - + Stored in logging_obj.model_call_details["agentic_loop_params"] to provide agentic hooks with the original request context needed for follow-up calls. """ + model: str """The model string with provider prefix (e.g., 'bedrock/invoke/...')""" - + custom_llm_provider: str """The LLM provider name (e.g., 'bedrock', 'anthropic')""" @@ -2910,6 +2911,9 @@ all_litellm_params = ( "shared_session", "search_tool_name", "order", + "retry_delay", + "exponential_backoff", + "jitter", ] + list(StandardCallbackDynamicParams.__annotations__.keys()) + list(CustomPricingLiteLLMParams.model_fields.keys()) From 4ad5de10cb68d9e19fef6069495dbea4c158a181 Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Mon, 19 Jan 2026 15:37:41 -0300 Subject: [PATCH 15/58] fix(realtime): disable SSL for ws:// WebSocket connections (#19345) When using http:// api_base (converted to ws://), the websockets library throws "ssl argument is incompatible with a ws:// URI". Only pass SSL context for secure wss:// connections. Co-authored-by: Krish Dholakia --- litellm/llms/openai/realtime/handler.py | 4 +- .../realtime/test_openai_realtime_handler.py | 58 +++++++++++++++++++ 2 files changed, 61 insertions(+), 1 deletion(-) diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index 6ab43ab31e4..fd04ac4d458 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -56,7 +56,9 @@ class OpenAIRealtime(OpenAIChatCompletion): url = self._construct_url(api_base, query_params) try: - ssl_context = get_shared_realtime_ssl_context() + # Only use SSL context for secure websocket connections (wss://) + # websockets library doesn't accept ssl argument for ws:// URIs + ssl_context = None if url.startswith("ws://") else get_shared_realtime_ssl_context() # Log a masked request preview consistent with other endpoints. logging_obj.pre_call( input=None, diff --git a/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py b/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py index eb9f8027761..87923a8093b 100644 --- a/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py +++ b/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py @@ -265,3 +265,61 @@ async def test_async_realtime_uses_max_size_parameter(): mock_realtime_streaming.assert_called_once() mock_streaming_instance.bidirectional_forward.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_async_realtime_ws_url_has_no_ssl(): + """ + Test that when using http:// api_base (converted to ws://), the ssl argument + is set to None. The websockets library doesn't accept ssl argument for ws:// URIs. + + This verifies the fix for: https://github.com/BerriAI/litellm/issues/19222 + """ + from litellm.llms.openai.realtime.handler import OpenAIRealtime + from litellm.types.realtime import RealtimeQueryParams + + handler = OpenAIRealtime() + api_base = "http://localhost:8113" # Non-SSL local server + api_key = "test-key" + model = "test-model" + query_params: RealtimeQueryParams = {"model": model} + + dummy_websocket = AsyncMock() + dummy_logging_obj = MagicMock() + mock_backend_ws = AsyncMock() + + class DummyAsyncContextManager: + def __init__(self, value): + self.value = value + async def __aenter__(self): + return self.value + async def __aexit__(self, exc_type, exc, tb): + return None + + with patch("websockets.connect", return_value=DummyAsyncContextManager(mock_backend_ws)) as mock_ws_connect, \ + patch("litellm.llms.openai.realtime.handler.RealTimeStreaming") as mock_realtime_streaming: + + mock_streaming_instance = MagicMock() + mock_realtime_streaming.return_value = mock_streaming_instance + mock_streaming_instance.bidirectional_forward = AsyncMock() + + await handler.async_realtime( + model=model, + websocket=dummy_websocket, + logging_obj=dummy_logging_obj, + api_base=api_base, + api_key=api_key, + query_params=query_params, + ) + + # Verify websockets.connect was called + mock_ws_connect.assert_called_once() + called_url = mock_ws_connect.call_args[0][0] + called_kwargs = mock_ws_connect.call_args[1] + + # Verify URL was converted from http:// to ws:// + assert called_url.startswith("ws://localhost:8113/v1/realtime?") + assert f"model={model}" in called_url + + # Verify ssl is None for ws:// URLs (the fix for issue #19222) + assert called_kwargs["ssl"] is None From 57b1d99b4456c5aa159ed53fae543deeda2b7131 Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Mon, 19 Jan 2026 15:44:38 -0300 Subject: [PATCH 16/58] feat(azure): add support for Azure OpenAI v1 API (#19313) * feat(azure): add support for Azure OpenAI v1 API When api_version is 'v1', 'latest', or 'preview', use the standard OpenAI client instead of AzureOpenAI client with base_url pointing to /openai/v1/ endpoint. This follows Microsoft's documentation for the new v1 API format: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#api-specs Changes: - Add OpenAI/AsyncOpenAI imports to common_utils.py and azure.py - Modify get_azure_openai_client() to detect v1 API versions and create appropriate client type - Update isinstance checks and type hints to accept both client types - Add unit tests for v1 API client creation * fix(azure): fix MyPy type errors for v1 API support - Add type: ignore for AsyncOpenAI constructor - Update type hints in files/handler.py and batches/handler.py - Add OpenAI/AsyncOpenAI to Union types for client parameters - Update isinstance checks to include OpenAI/AsyncOpenAI * fix(azure): update type hints in files and batches handlers for v1 API Update async method signatures to accept Union[AsyncAzureOpenAI, AsyncOpenAI] to fix mypy errors when using v1 API client. --- litellm/llms/azure/azure.py | 42 ++++--- litellm/llms/azure/batches/handler.py | 36 +++--- litellm/llms/azure/common_utils.py | 47 +++++-- litellm/llms/azure/files/handler.py | 46 +++---- .../llms/azure/test_azure_common_utils.py | 116 ++++++++++++++++++ 5 files changed, 217 insertions(+), 70 deletions(-) diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 3ef0186ba0e..dced06b6b3a 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -4,7 +4,13 @@ import time from typing import Any, Callable, Coroutine, Dict, List, Optional, Union import httpx # type: ignore -from openai import APITimeoutError, AsyncAzureOpenAI, AzureOpenAI +from openai import ( + APITimeoutError, + AsyncAzureOpenAI, + AsyncOpenAI, + AzureOpenAI, + OpenAI, +) import litellm from litellm.constants import AZURE_OPERATION_POLLING_TIMEOUT, DEFAULT_MAX_RETRIES @@ -128,7 +134,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): def make_sync_azure_openai_chat_completion_request( self, - azure_client: AzureOpenAI, + azure_client: Union[AzureOpenAI, OpenAI], data: dict, timeout: Union[float, httpx.Timeout], ): @@ -151,7 +157,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): @track_llm_api_timing() async def make_azure_openai_chat_completion_request( self, - azure_client: AsyncAzureOpenAI, + azure_client: Union[AsyncAzureOpenAI, AsyncOpenAI], data: dict, timeout: Union[float, httpx.Timeout], logging_obj: LiteLLMLoggingObj, @@ -328,10 +334,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): _is_async=False, litellm_params=litellm_params, ) - if not isinstance(azure_client, AzureOpenAI): + if not isinstance(azure_client, (AzureOpenAI, OpenAI)): raise AzureOpenAIError( status_code=500, - message="azure_client is not an instance of AzureOpenAI", + message="azure_client is not an instance of AzureOpenAI or OpenAI", ) headers, response = self.make_sync_azure_openai_chat_completion_request( @@ -401,8 +407,8 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): _is_async=True, litellm_params=litellm_params, ) - if not isinstance(azure_client, AsyncAzureOpenAI): - raise ValueError("Azure client is not an instance of AsyncAzureOpenAI") + if not isinstance(azure_client, (AsyncAzureOpenAI, AsyncOpenAI)): + raise ValueError("Azure client is not an instance of AsyncAzureOpenAI or AsyncOpenAI") ## LOGGING logging_obj.pre_call( input=data["messages"], @@ -412,7 +418,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): "api_key": api_key, "azure_ad_token": azure_ad_token, }, - "api_base": azure_client._base_url._uri_reference, + "api_base": api_base, "acompletion": True, "complete_input_dict": data, }, @@ -520,10 +526,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): _is_async=False, litellm_params=litellm_params, ) - if not isinstance(azure_client, AzureOpenAI): + if not isinstance(azure_client, (AzureOpenAI, OpenAI)): raise AzureOpenAIError( status_code=500, - message="azure_client is not an instance of AzureOpenAI", + message="azure_client is not an instance of AzureOpenAI or OpenAI", ) ## LOGGING logging_obj.pre_call( @@ -534,7 +540,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): "api_key": api_key, "azure_ad_token": azure_ad_token, }, - "api_base": azure_client._base_url._uri_reference, + "api_base": api_base, "acompletion": True, "complete_input_dict": data, }, @@ -578,8 +584,8 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): _is_async=True, litellm_params=litellm_params, ) - if not isinstance(azure_client, AsyncAzureOpenAI): - raise ValueError("Azure client is not an instance of AsyncAzureOpenAI") + if not isinstance(azure_client, (AsyncAzureOpenAI, AsyncOpenAI)): + raise ValueError("Azure client is not an instance of AsyncAzureOpenAI or AsyncOpenAI") ## LOGGING logging_obj.pre_call( @@ -590,7 +596,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): "api_key": api_key, "azure_ad_token": azure_ad_token, }, - "api_base": azure_client._base_url._uri_reference, + "api_base": api_base, "acompletion": True, "complete_input_dict": data, }, @@ -657,8 +663,8 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): client=client, litellm_params=litellm_params, ) - if not isinstance(openai_aclient, AsyncAzureOpenAI): - raise ValueError("Azure client is not an instance of AsyncAzureOpenAI") + if not isinstance(openai_aclient, (AsyncAzureOpenAI, AsyncOpenAI)): + raise ValueError("Azure client is not an instance of AsyncAzureOpenAI or AsyncOpenAI") raw_response = await openai_aclient.embeddings.with_raw_response.create( **data, timeout=timeout @@ -776,10 +782,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): client=client, litellm_params=litellm_params, ) - if not isinstance(azure_client, AzureOpenAI): + if not isinstance(azure_client, (AzureOpenAI, OpenAI)): raise AzureOpenAIError( status_code=500, - message="azure_client is not an instance of AzureOpenAI", + message="azure_client is not an instance of AzureOpenAI or OpenAI", ) ## COMPLETION CALL diff --git a/litellm/llms/azure/batches/handler.py b/litellm/llms/azure/batches/handler.py index 7fc6388ba87..3996cb808e4 100644 --- a/litellm/llms/azure/batches/handler.py +++ b/litellm/llms/azure/batches/handler.py @@ -6,6 +6,8 @@ from typing import Any, Coroutine, Optional, Union, cast import httpx +from openai import AsyncOpenAI, OpenAI + from litellm.llms.azure.azure import AsyncAzureOpenAI, AzureOpenAI from litellm.types.llms.openai import ( Batch, @@ -33,7 +35,7 @@ class AzureBatchesAPI(BaseAzureLLM): async def acreate_batch( self, create_batch_data: CreateBatchRequest, - azure_client: AsyncAzureOpenAI, + azure_client: Union[AsyncAzureOpenAI, AsyncOpenAI], ) -> LiteLLMBatch: response = await azure_client.batches.create(**create_batch_data) return LiteLLMBatch(**response.model_dump()) @@ -47,11 +49,11 @@ class AzureBatchesAPI(BaseAzureLLM): api_version: Optional[str], timeout: Union[float, httpx.Timeout], max_retries: Optional[int], - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, ) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]: azure_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI] + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] ] = self.get_azure_openai_client( api_key=api_key, api_base=api_base, @@ -66,20 +68,20 @@ class AzureBatchesAPI(BaseAzureLLM): ) if _is_async is True: - if not isinstance(azure_client, AsyncAzureOpenAI): + if not isinstance(azure_client, (AsyncAzureOpenAI, AsyncOpenAI)): raise ValueError( "OpenAI client is not an instance of AsyncOpenAI. Make sure you passed an AsyncOpenAI client." ) return self.acreate_batch( # type: ignore create_batch_data=create_batch_data, azure_client=azure_client ) - response = cast(AzureOpenAI, azure_client).batches.create(**create_batch_data) + response = cast(Union[AzureOpenAI, OpenAI], azure_client).batches.create(**create_batch_data) return LiteLLMBatch(**response.model_dump()) async def aretrieve_batch( self, retrieve_batch_data: RetrieveBatchRequest, - client: AsyncAzureOpenAI, + client: Union[AsyncAzureOpenAI, AsyncOpenAI], ) -> LiteLLMBatch: response = await client.batches.retrieve(**retrieve_batch_data) return LiteLLMBatch(**response.model_dump()) @@ -93,11 +95,11 @@ class AzureBatchesAPI(BaseAzureLLM): api_version: Optional[str], timeout: Union[float, httpx.Timeout], max_retries: Optional[int], - client: Optional[AzureOpenAI] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, ): azure_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI] + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] ] = self.get_azure_openai_client( api_key=api_key, api_base=api_base, @@ -112,14 +114,14 @@ class AzureBatchesAPI(BaseAzureLLM): ) if _is_async is True: - if not isinstance(azure_client, AsyncAzureOpenAI): + if not isinstance(azure_client, (AsyncAzureOpenAI, AsyncOpenAI)): raise ValueError( "OpenAI client is not an instance of AsyncOpenAI. Make sure you passed an AsyncOpenAI client." ) return self.aretrieve_batch( # type: ignore retrieve_batch_data=retrieve_batch_data, client=azure_client ) - response = cast(AzureOpenAI, azure_client).batches.retrieve( + response = cast(Union[AzureOpenAI, OpenAI], azure_client).batches.retrieve( **retrieve_batch_data ) return LiteLLMBatch(**response.model_dump()) @@ -127,7 +129,7 @@ class AzureBatchesAPI(BaseAzureLLM): async def acancel_batch( self, cancel_batch_data: CancelBatchRequest, - client: AsyncAzureOpenAI, + client: Union[AsyncAzureOpenAI, AsyncOpenAI], ) -> Batch: response = await client.batches.cancel(**cancel_batch_data) return response @@ -141,11 +143,11 @@ class AzureBatchesAPI(BaseAzureLLM): api_version: Optional[str], timeout: Union[float, httpx.Timeout], max_retries: Optional[int], - client: Optional[AzureOpenAI] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, ): azure_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI] + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] ] = self.get_azure_openai_client( api_key=api_key, api_base=api_base, @@ -163,7 +165,7 @@ class AzureBatchesAPI(BaseAzureLLM): async def alist_batches( self, - client: AsyncAzureOpenAI, + client: Union[AsyncAzureOpenAI, AsyncOpenAI], after: Optional[str] = None, limit: Optional[int] = None, ): @@ -180,11 +182,11 @@ class AzureBatchesAPI(BaseAzureLLM): max_retries: Optional[int], after: Optional[str] = None, limit: Optional[int] = None, - client: Optional[AzureOpenAI] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, ): azure_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI] + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] ] = self.get_azure_openai_client( api_key=api_key, api_base=api_base, @@ -199,7 +201,7 @@ class AzureBatchesAPI(BaseAzureLLM): ) if _is_async is True: - if not isinstance(azure_client, AsyncAzureOpenAI): + if not isinstance(azure_client, (AsyncAzureOpenAI, AsyncOpenAI)): raise ValueError( "OpenAI client is not an instance of AsyncOpenAI. Make sure you passed an AsyncOpenAI client." ) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 85596a628da..25b218fca8c 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -3,7 +3,7 @@ import os from typing import Any, Callable, Dict, Literal, Optional, Union, cast import httpx -from openai import AsyncAzureOpenAI, AzureOpenAI +from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI import litellm from litellm._logging import verbose_logger @@ -439,12 +439,12 @@ class BaseAzureLLM(BaseOpenAILLM): api_key: Optional[str], api_base: Optional[str], api_version: Optional[str] = None, - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, _is_async: bool = False, model: Optional[str] = None, - ) -> Optional[Union[AzureOpenAI, AsyncAzureOpenAI]]: - openai_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = None + ) -> Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]]: + openai_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None client_initialization_params: dict = locals() client_initialization_params["is_async"] = _is_async if client is None: @@ -453,9 +453,7 @@ class BaseAzureLLM(BaseOpenAILLM): client_type="azure", ) if cached_client: - if isinstance(cached_client, AzureOpenAI) or isinstance( - cached_client, AsyncAzureOpenAI - ): + if isinstance(cached_client, (AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI)): return cached_client azure_client_params = self.initialize_azure_sdk_client( @@ -466,15 +464,40 @@ class BaseAzureLLM(BaseOpenAILLM): api_version=api_version, is_async=_is_async, ) - if _is_async is True: - openai_client = AsyncAzureOpenAI(**azure_client_params) + + # For Azure v1 API, use standard OpenAI client instead of AzureOpenAI + # See: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#api-specs + if self._is_azure_v1_api_version(api_version): + # Extract only params that OpenAI client accepts + # Always use /openai/v1/ regardless of whether user passed "v1", "latest", or "preview" + v1_params = { + "api_key": azure_client_params.get("api_key"), + "base_url": f"{api_base}/openai/v1/", + } + if "timeout" in azure_client_params: + v1_params["timeout"] = azure_client_params["timeout"] + if "max_retries" in azure_client_params: + v1_params["max_retries"] = azure_client_params["max_retries"] + if "http_client" in azure_client_params: + v1_params["http_client"] = azure_client_params["http_client"] + + verbose_logger.debug(f"Using Azure v1 API with base_url: {v1_params['base_url']}") + + if _is_async is True: + openai_client = AsyncOpenAI(**v1_params) # type: ignore + else: + openai_client = OpenAI(**v1_params) # type: ignore else: - openai_client = AzureOpenAI(**azure_client_params) # type: ignore + # Traditional Azure API uses AzureOpenAI client + if _is_async is True: + openai_client = AsyncAzureOpenAI(**azure_client_params) + else: + openai_client = AzureOpenAI(**azure_client_params) # type: ignore else: openai_client = client if api_version is not None and isinstance( - openai_client._custom_query, dict - ): + openai_client, (AzureOpenAI, AsyncAzureOpenAI) + ) and isinstance(openai_client._custom_query, dict): # set api_version to version passed by user openai_client._custom_query.setdefault("api-version", api_version) diff --git a/litellm/llms/azure/files/handler.py b/litellm/llms/azure/files/handler.py index 69b2d71753b..e53ced6b0e2 100644 --- a/litellm/llms/azure/files/handler.py +++ b/litellm/llms/azure/files/handler.py @@ -1,7 +1,7 @@ from typing import Any, Coroutine, Optional, Union, cast import httpx -from openai import AsyncAzureOpenAI, AzureOpenAI +from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI from openai.types.file_deleted import FileDeleted from litellm._logging import verbose_logger @@ -40,7 +40,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): async def acreate_file( self, create_file_data: CreateFileRequest, - openai_client: AsyncAzureOpenAI, + openai_client: Union[AsyncAzureOpenAI, AsyncOpenAI], ) -> OpenAIFileObject: verbose_logger.debug("create_file_data=%s", create_file_data) response = await openai_client.files.create(**self._prepare_create_file_data(create_file_data)) # type: ignore[arg-type] @@ -56,11 +56,11 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): api_version: Optional[str], timeout: Union[float, httpx.Timeout], max_retries: Optional[int], - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, ) -> Union[OpenAIFileObject, Coroutine[Any, Any, OpenAIFileObject]]: openai_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI] + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] ] = self.get_azure_openai_client( litellm_params=litellm_params or {}, api_key=api_key, @@ -75,20 +75,20 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): ) if _is_async is True: - if not isinstance(openai_client, AsyncAzureOpenAI): + if not isinstance(openai_client, (AsyncAzureOpenAI, AsyncOpenAI)): raise ValueError( "AzureOpenAI client is not an instance of AsyncAzureOpenAI. Make sure you passed an AsyncAzureOpenAI client." ) return self.acreate_file( create_file_data=create_file_data, openai_client=openai_client ) - response = cast(AzureOpenAI, openai_client).files.create(**self._prepare_create_file_data(create_file_data)) # type: ignore[arg-type] + response = cast(Union[AzureOpenAI, OpenAI], openai_client).files.create(**self._prepare_create_file_data(create_file_data)) # type: ignore[arg-type] return OpenAIFileObject(**response.model_dump()) async def afile_content( self, file_content_request: FileContentRequest, - openai_client: AsyncAzureOpenAI, + openai_client: Union[AsyncAzureOpenAI, AsyncOpenAI], ) -> HttpxBinaryResponseContent: response = await openai_client.files.content(**file_content_request) return HttpxBinaryResponseContent(response=response.response) @@ -102,13 +102,13 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): timeout: Union[float, httpx.Timeout], max_retries: Optional[int], api_version: Optional[str] = None, - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, ) -> Union[ HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent] ]: openai_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI] + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] ] = self.get_azure_openai_client( litellm_params=litellm_params or {}, api_key=api_key, @@ -123,7 +123,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): ) if _is_async is True: - if not isinstance(openai_client, AsyncAzureOpenAI): + if not isinstance(openai_client, (AsyncAzureOpenAI, AsyncOpenAI)): raise ValueError( "AzureOpenAI client is not an instance of AsyncAzureOpenAI. Make sure you passed an AsyncAzureOpenAI client." ) @@ -131,7 +131,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): file_content_request=file_content_request, openai_client=openai_client, ) - response = cast(AzureOpenAI, openai_client).files.content( + response = cast(Union[AzureOpenAI, OpenAI], openai_client).files.content( **file_content_request ) @@ -140,7 +140,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): async def aretrieve_file( self, file_id: str, - openai_client: AsyncAzureOpenAI, + openai_client: Union[AsyncAzureOpenAI, AsyncOpenAI], ) -> FileObject: response = await openai_client.files.retrieve(file_id=file_id) return response @@ -154,11 +154,11 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): timeout: Union[float, httpx.Timeout], max_retries: Optional[int], api_version: Optional[str] = None, - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, ): openai_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI] + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] ] = self.get_azure_openai_client( litellm_params=litellm_params or {}, api_key=api_key, @@ -173,7 +173,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): ) if _is_async is True: - if not isinstance(openai_client, AsyncAzureOpenAI): + if not isinstance(openai_client, (AsyncAzureOpenAI, AsyncOpenAI)): raise ValueError( "AzureOpenAI client is not an instance of AsyncAzureOpenAI. Make sure you passed an AsyncAzureOpenAI client." ) @@ -188,7 +188,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): async def adelete_file( self, file_id: str, - openai_client: AsyncAzureOpenAI, + openai_client: Union[AsyncAzureOpenAI, AsyncOpenAI], ) -> FileDeleted: response = await openai_client.files.delete(file_id=file_id) @@ -206,11 +206,11 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): max_retries: Optional[int], organization: Optional[str] = None, api_version: Optional[str] = None, - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, ): openai_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI] + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] ] = self.get_azure_openai_client( litellm_params=litellm_params or {}, api_key=api_key, @@ -225,7 +225,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): ) if _is_async is True: - if not isinstance(openai_client, AsyncAzureOpenAI): + if not isinstance(openai_client, (AsyncAzureOpenAI, AsyncOpenAI)): raise ValueError( "AzureOpenAI client is not an instance of AsyncAzureOpenAI. Make sure you passed an AsyncAzureOpenAI client." ) @@ -242,7 +242,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): async def alist_files( self, - openai_client: AsyncAzureOpenAI, + openai_client: Union[AsyncAzureOpenAI, AsyncOpenAI], purpose: Optional[str] = None, ): if isinstance(purpose, str): @@ -260,11 +260,11 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): max_retries: Optional[int], purpose: Optional[str] = None, api_version: Optional[str] = None, - client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI]] = None, litellm_params: Optional[dict] = None, ): openai_client: Optional[ - Union[AzureOpenAI, AsyncAzureOpenAI] + Union[AzureOpenAI, AsyncAzureOpenAI, OpenAI, AsyncOpenAI] ] = self.get_azure_openai_client( litellm_params=litellm_params or {}, api_key=api_key, @@ -279,7 +279,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): ) if _is_async is True: - if not isinstance(openai_client, AsyncAzureOpenAI): + if not isinstance(openai_client, (AsyncAzureOpenAI, AsyncOpenAI)): raise ValueError( "AzureOpenAI client is not an instance of AsyncAzureOpenAI. Make sure you passed an AsyncAzureOpenAI client." ) diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index 654720183a7..5a0680d66bf 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -1542,3 +1542,119 @@ def test_is_azure_v1_api_version(api_version, expected): """ result = BaseAzureLLM._is_azure_v1_api_version(api_version=api_version) assert result == expected + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_api_uses_openai_client(api_version): + """ + Test that Azure v1 API versions use OpenAI client instead of AzureOpenAI. + + When api_version is 'v1', 'latest', or 'preview', the client should be + instantiated as OpenAI/AsyncOpenAI with base_url pointing to /openai/v1/ + instead of the traditional AzureOpenAI client with /deployments/ URL pattern. + + See: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#api-specs + """ + from openai import AsyncOpenAI, OpenAI + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + + # Test sync client + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = { + "api_key": "test-key", + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": None, + } + + client = base_llm.get_azure_openai_client( + api_key="test-key", + api_base=api_base, + api_version=api_version, + _is_async=False, + ) + + # Should be OpenAI client, not AzureOpenAI + assert isinstance(client, OpenAI), f"Expected OpenAI client for api_version={api_version}" + # base_url should be /openai/v1/ (not /deployments/) + assert "/openai/v1/" in str(client.base_url), f"base_url should contain /openai/v1/, got {client.base_url}" + + # Test async client + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = { + "api_key": "test-key", + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": None, + } + + async_client = base_llm.get_azure_openai_client( + api_key="test-key", + api_base=api_base, + api_version=api_version, + _is_async=True, + ) + + # Should be AsyncOpenAI client, not AsyncAzureOpenAI + assert isinstance(async_client, AsyncOpenAI), f"Expected AsyncOpenAI client for api_version={api_version}" + # base_url should be /openai/v1/ + assert "/openai/v1/" in str(async_client.base_url), f"base_url should contain /openai/v1/, got {async_client.base_url}" + + +def test_azure_traditional_api_uses_azure_openai_client(): + """ + Test that traditional Azure API versions still use AzureOpenAI client. + + When api_version is a dated version like '2023-05-15', the client should + be instantiated as AzureOpenAI/AsyncAzureOpenAI with the traditional + /deployments/ URL pattern. + """ + from openai import AsyncAzureOpenAI, AzureOpenAI + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + api_version = "2023-05-15" + + # Test sync client + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = { + "api_key": "test-key", + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": None, + } + + client = base_llm.get_azure_openai_client( + api_key="test-key", + api_base=api_base, + api_version=api_version, + _is_async=False, + ) + + # Should be AzureOpenAI client + assert isinstance(client, AzureOpenAI), f"Expected AzureOpenAI client for api_version={api_version}" + + # Test async client + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = { + "api_key": "test-key", + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": None, + } + + async_client = base_llm.get_azure_openai_client( + api_key="test-key", + api_base=api_base, + api_version=api_version, + _is_async=True, + ) + + # Should be AsyncAzureOpenAI client + assert isinstance(async_client, AsyncAzureOpenAI), f"Expected AsyncAzureOpenAI client for api_version={api_version}" From d30c25af21efa00d75b8ceb0488225a18d318edb Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Mon, 19 Jan 2026 15:45:37 -0300 Subject: [PATCH 17/58] feat(gemini): use responseJsonSchema for Gemini 2.0+ models (#19314) * feat(gemini): add opt-in support for responseJsonSchema Add support for Gemini's native responseJsonSchema parameter which uses standard JSON Schema format instead of OpenAPI-style responseSchema. Benefits of responseJsonSchema (Gemini 2.0+ only): - Standard JSON Schema format (lowercase types) - Supports additionalProperties for stricter validation - Better compatibility with Pydantic's model_json_schema() - No propertyOrdering required Usage: ```python response_format={ "type": "json_schema", "json_schema": {"schema": {...}}, "use_json_schema": True # opt-in } ``` This is backwards compatible - existing code continues to use responseSchema by default. Closes #16340 * docs: add documentation for use_json_schema parameter Document the new use_json_schema option for Gemini 2.0+ models in the JSON Mode documentation. * refactor(gemini): use responseJsonSchema by default for Gemini 2.0+ Remove opt-in flag `use_json_schema` and automatically detect model version: - Gemini 2.0+: uses responseJsonSchema (standard JSON Schema, supports additionalProperties) - Gemini 1.5: uses responseSchema (OpenAPI format, legacy) This follows LiteLLM's philosophy of abstracting provider differences - users write the same code regardless of model version. * test(vertex): update json_schema tests to accept both responseSchema formats Gemini 2.x+ uses responseJsonSchema while Gemini 1.x uses responseSchema. Update tests to accept both formats since litellm now auto-selects based on model version. --- docs/my-website/docs/completion/json_mode.md | 88 ++++++++++++++++- litellm/llms/vertex_ai/common_utils.py | 66 +++++++++++++ .../vertex_and_google_ai_studio_gemini.py | 66 +++++++++---- litellm/types/llms/vertex_ai.py | 1 + .../test_amazing_vertex_completion.py | 22 +++-- ...test_vertex_and_google_ai_studio_gemini.py | 99 ++++++++++++++++++- 6 files changed, 310 insertions(+), 32 deletions(-) diff --git a/docs/my-website/docs/completion/json_mode.md b/docs/my-website/docs/completion/json_mode.md index 0122e202610..14477f99153 100644 --- a/docs/my-website/docs/completion/json_mode.md +++ b/docs/my-website/docs/completion/json_mode.md @@ -341,4 +341,90 @@ curl http://0.0.0.0:4000/v1/chat/completions \ ``` - \ No newline at end of file + + +## Gemini - Native JSON Schema Format (Gemini 2.0+) + +Gemini 2.0+ models automatically use the native `responseJsonSchema` parameter, which provides better compatibility with standard JSON Schema format. + +### Benefits (Gemini 2.0+): +- Standard JSON Schema format (lowercase types like `string`, `object`) +- Supports `additionalProperties: false` for stricter validation +- Better compatibility with Pydantic's `model_json_schema()` +- No `propertyOrdering` required + +### Usage + + + + +```python +from litellm import completion +from pydantic import BaseModel + +class UserInfo(BaseModel): + name: str + age: int + +response = completion( + model="gemini/gemini-2.0-flash", + messages=[{"role": "user", "content": "Extract: John is 25 years old"}], + response_format={ + "type": "json_schema", + "json_schema": { + "name": "user_info", + "schema": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "age": {"type": "integer"} + }, + "required": ["name", "age"], + "additionalProperties": False # Supported on Gemini 2.0+ + } + } + } +) +``` + + + + +```bash +curl http://0.0.0.0:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer $LITELLM_API_KEY" \ + -d '{ + "model": "gemini-2.0-flash", + "messages": [ + {"role": "user", "content": "Extract: John is 25 years old"} + ], + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "user_info", + "schema": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "age": {"type": "integer"} + }, + "required": ["name", "age"], + "additionalProperties": false + } + } + } + }' +``` + + + + +### Model Behavior + +| Model | Format Used | `additionalProperties` Support | +|-------|-------------|-------------------------------| +| Gemini 2.0+ | `responseJsonSchema` (JSON Schema) | ✅ Yes | +| Gemini 1.5 | `responseSchema` (OpenAPI) | ❌ No | + +LiteLLM automatically selects the appropriate format based on the model version. \ No newline at end of file diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 5aa7662f175..6904807741c 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -150,6 +150,34 @@ def get_supports_response_schema( return _supports_response_schema +def supports_response_json_schema(model: str) -> bool: + """ + Check if the model supports responseJsonSchema (JSON Schema format). + + responseJsonSchema is supported by Gemini 2.0+ models and uses standard + JSON Schema format with lowercase types (string, object, etc.) instead of + the OpenAPI-style responseSchema with uppercase types (STRING, OBJECT, etc.). + + Benefits of responseJsonSchema: + - Supports additionalProperties for stricter schema validation + - Uses standard JSON Schema format (no type conversion needed) + - Better compatibility with Pydantic's model_json_schema() + + Args: + model: The model name (e.g., "gemini-2.0-flash", "gemini-2.5-pro") + + Returns: + True if the model supports responseJsonSchema, False otherwise + """ + model_lower = model.lower() + + # Gemini 2.0+ and 2.5+ models support responseJsonSchema + # Pattern matches: gemini-2.0-*, gemini-2.5-*, gemini-3-*, etc. + gemini_2_plus_pattern = re.compile(r"gemini-([2-9]|[1-9]\d+)\.") + + return bool(gemini_2_plus_pattern.search(model_lower)) + + from typing import Literal, Optional all_gemini_url_modes = Literal[ @@ -486,6 +514,44 @@ def _build_vertex_schema(parameters: dict, add_property_ordering: bool = False): return parameters +def _build_json_schema(parameters: dict) -> dict: + """ + Build a JSON Schema for use with Gemini's responseJsonSchema parameter. + + Unlike _build_vertex_schema (used for responseSchema), this function: + - Does NOT convert types to uppercase (keeps standard JSON Schema format) + - Does NOT add propertyOrdering + - Does NOT filter fields (allows additionalProperties) + - Still unpacks $defs/$ref (Gemini doesn't support JSON Schema references) + + Parameters: + parameters: dict - the JSON schema to process + + Returns: + dict - the processed schema in standard JSON Schema format + """ + # Unpack $defs references (Gemini doesn't support $ref) + defs = parameters.pop("$defs", {}) + for name, value in defs.items(): + unpack_defs(value, defs) + unpack_defs(parameters, defs) + + # Convert anyOf with null to nullable + convert_anyof_null_to_nullable(parameters) + + # Handle empty strings in enum values - Gemini doesn't accept empty strings in enums + _fix_enum_empty_strings(parameters) + + # Remove enums for non-string typed fields (Gemini requires enum only on strings) + _fix_enum_types(parameters) + + # Handle empty items objects + process_items(parameters) + add_object_type(parameters) + + return parameters + + def _filter_anyof_fields(schema_dict: Dict[str, Any]) -> Dict[str, Any]: """ When anyof is present, only keep the anyof field and its contents - otherwise VertexAI will throw an error - https://github.com/BerriAI/litellm/issues/11164 diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index f65a19ac46f..dd34fbd772b 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -92,7 +92,12 @@ from litellm.utils import ( ) from ....utils import _remove_additional_properties, _remove_strict_from_schema -from ..common_utils import VertexAIError, _build_vertex_schema +from ..common_utils import ( + VertexAIError, + _build_json_schema, + _build_vertex_schema, + supports_response_json_schema, +) from ..vertex_llm_base import VertexBase from .transformation import ( _gemini_convert_messages_with_history, @@ -624,30 +629,55 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) return old_schema - def apply_response_schema_transformation(self, value: dict, optional_params: dict): + def apply_response_schema_transformation( + self, value: dict, optional_params: dict, model: str + ): new_value = deepcopy(value) - # remove 'additionalProperties' from json schema - new_value = _remove_additional_properties(new_value) - # remove 'strict' from json schema + # remove 'strict' from json schema (not supported by Gemini) new_value = _remove_strict_from_schema(new_value) - if new_value["type"] == "json_object": + + # Automatically use responseJsonSchema for Gemini 2.0+ models + # responseJsonSchema uses standard JSON Schema format and supports additionalProperties + # For older models (Gemini 1.5), fall back to responseSchema (OpenAPI format) + use_json_schema = supports_response_json_schema(model) + + if not use_json_schema: + # For responseSchema, remove 'additionalProperties' (not supported) + new_value = _remove_additional_properties(new_value) + + # Handle response type + if new_value.get("type") == "json_object": optional_params["response_mime_type"] = "application/json" - elif new_value["type"] == "text": + elif new_value.get("type") == "text": optional_params["response_mime_type"] = "text/plain" + + # Extract schema from response_format + schema = None if "response_schema" in new_value: optional_params["response_mime_type"] = "application/json" - optional_params["response_schema"] = new_value["response_schema"] - elif new_value["type"] == "json_schema": # type: ignore - if "json_schema" in new_value and "schema" in new_value["json_schema"]: # type: ignore + schema = new_value["response_schema"] + elif new_value.get("type") == "json_schema": + if "json_schema" in new_value and "schema" in new_value["json_schema"]: optional_params["response_mime_type"] = "application/json" - optional_params["response_schema"] = new_value["json_schema"]["schema"] # type: ignore + schema = new_value["json_schema"]["schema"] - if "response_schema" in optional_params and isinstance( - optional_params["response_schema"], dict - ): - optional_params["response_schema"] = self._map_response_schema( - value=optional_params["response_schema"] - ) + if schema and isinstance(schema, dict): + if use_json_schema: + # Use responseJsonSchema (Gemini 2.0+ only, opt-in) + # - Standard JSON Schema format (lowercase types) + # - Supports additionalProperties + # - No propertyOrdering needed + optional_params["response_json_schema"] = _build_json_schema( + deepcopy(schema) + ) + else: + # Use responseSchema (default, backwards compatible) + # - OpenAPI-style format (uppercase types) + # - No additionalProperties support + # - Requires propertyOrdering + optional_params["response_schema"] = self._map_response_schema( + value=schema + ) @staticmethod def _map_reasoning_effort_to_thinking_budget( @@ -947,7 +977,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): optional_params["max_output_tokens"] = value elif param == "response_format" and isinstance(value, dict): # type: ignore self.apply_response_schema_transformation( - value=value, optional_params=optional_params + value=value, optional_params=optional_params, model=model ) elif param == "frequency_penalty": if self._supports_penalty_parameters(model): diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 74d6339887b..0a4a2d0f14c 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -207,6 +207,7 @@ class GenerationConfig(TypedDict, total=False): frequency_penalty: float response_mime_type: Literal["text/plain", "application/json"] response_schema: dict + response_json_schema: dict seed: int responseLogprobs: bool logprobs: int diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 0373f5f4356..af8942f89eb 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -1386,14 +1386,15 @@ async def test_gemini_pro_json_schema_args_sent_httpx( print(mock_call.call_args.kwargs["json"]["generationConfig"]) if supports_response_schema: + # Gemini 2.x+ uses response_json_schema, Gemini 1.x uses response_schema + gen_config = mock_call.call_args.kwargs["json"]["generationConfig"] assert ( - "response_schema" - in mock_call.call_args.kwargs["json"]["generationConfig"] - ) + "response_schema" in gen_config or "response_json_schema" in gen_config + ), f"Expected response_schema or response_json_schema in {gen_config}" else: + gen_config = mock_call.call_args.kwargs["json"]["generationConfig"] assert ( - "response_schema" - not in mock_call.call_args.kwargs["json"]["generationConfig"] + "response_schema" not in gen_config and "response_json_schema" not in gen_config ) assert ( "Use this JSON schema:" @@ -1566,10 +1567,11 @@ async def test_gemini_pro_json_schema_args_sent_httpx_openai_schema( print(mock_call.call_args.kwargs["json"]["generationConfig"]) if supports_response_schema: + # Gemini 2.x+ uses response_json_schema, Gemini 1.x uses response_schema + gen_config = mock_call.call_args.kwargs["json"]["generationConfig"] assert ( - "response_schema" - in mock_call.call_args.kwargs["json"]["generationConfig"] - ) + "response_schema" in gen_config or "response_json_schema" in gen_config + ), f"Expected response_schema or response_json_schema in {gen_config}" assert ( "response_mime_type" in mock_call.call_args.kwargs["json"]["generationConfig"] @@ -1581,9 +1583,9 @@ async def test_gemini_pro_json_schema_args_sent_httpx_openai_schema( == "application/json" ) else: + gen_config = mock_call.call_args.kwargs["json"]["generationConfig"] assert ( - "response_schema" - not in mock_call.call_args.kwargs["json"]["generationConfig"] + "response_schema" not in gen_config and "response_json_schema" not in gen_config ) assert ( "Use this JSON schema:" diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index f44e8640105..969199f6ad0 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -74,6 +74,10 @@ def test_get_model_name_from_gemini_spec_model(): def test_vertex_ai_response_schema_dict(): + """ + Test that older Gemini models (1.5) use responseSchema (OpenAPI format). + responseSchema requires propertyOrdering and doesn't support additionalProperties. + """ v = VertexGeminiConfig() non_default_params = { "messages": [{"role": "user", "content": "Hello, world!"}], @@ -109,7 +113,7 @@ def test_vertex_ai_response_schema_dict(): transformed_request = v.map_openai_params( non_default_params=non_default_params, optional_params={}, - model="gemini-2.0-flash-lite", + model="gemini-1.5-flash", # Old model uses responseSchema (OpenAPI format) drop_params=False, ) @@ -160,6 +164,9 @@ class Step(BaseModel): def test_vertex_ai_response_schema_defs(): + """ + Test that $defs are unpacked for older Gemini models using responseSchema. + """ v = VertexGeminiConfig() schema = cast(dict, v.get_json_schema_from_pydantic_object(MathReasoning)) @@ -173,7 +180,7 @@ def test_vertex_ai_response_schema_defs(): "response_format": schema, }, optional_params={}, - model="gemini-2.0-flash-lite", + model="gemini-1.5-flash", # Old model uses responseSchema (OpenAPI format) drop_params=False, ) @@ -203,7 +210,93 @@ def test_vertex_ai_response_schema_defs(): } +def test_vertex_ai_response_json_schema_for_gemini_2(): + """ + Test that Gemini 2.0+ models automatically use responseJsonSchema. + + responseJsonSchema uses standard JSON Schema format: + - lowercase types (string, object, etc.) + - no propertyOrdering required + - supports additionalProperties + """ + v = VertexGeminiConfig() + + transformed_request = v.map_openai_params( + non_default_params={ + "messages": [{"role": "user", "content": "Hello, world!"}], + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "test_schema", + "schema": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "age": {"type": "integer"}, + }, + "required": ["name"], + "additionalProperties": False, + }, + }, + }, + }, + optional_params={}, + model="gemini-2.0-flash", # Gemini 2.0+ automatically uses responseJsonSchema + drop_params=False, + ) + + # Should use response_json_schema, not response_schema + assert "response_json_schema" in transformed_request + assert "response_schema" not in transformed_request + + # Types should be lowercase (standard JSON Schema format) + assert transformed_request["response_json_schema"]["type"] == "object" + assert transformed_request["response_json_schema"]["properties"]["name"]["type"] == "string" + assert transformed_request["response_json_schema"]["properties"]["age"]["type"] == "integer" + + # Should NOT have propertyOrdering (not needed for responseJsonSchema) + assert "propertyOrdering" not in transformed_request["response_json_schema"] + + # additionalProperties should be preserved (supported by responseJsonSchema) + assert transformed_request["response_json_schema"].get("additionalProperties") == False + + +def test_vertex_ai_response_schema_for_old_models(): + """ + Test that older models (Gemini 1.5) automatically use responseSchema. + """ + v = VertexGeminiConfig() + + transformed_request = v.map_openai_params( + non_default_params={ + "messages": [{"role": "user", "content": "Hello, world!"}], + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "test_schema", + "schema": { + "type": "object", + "properties": { + "name": {"type": "string"}, + }, + }, + }, + }, + }, + optional_params={}, + model="gemini-1.5-flash", # Old model automatically uses responseSchema + drop_params=False, + ) + + # Should use response_schema for older models + assert "response_schema" in transformed_request + assert "response_json_schema" not in transformed_request + + def test_vertex_ai_retain_property_ordering(): + """ + Test that existing propertyOrdering is preserved for older models using responseSchema. + """ v = VertexGeminiConfig() transformed_request = v.map_openai_params( non_default_params={ @@ -224,7 +317,7 @@ def test_vertex_ai_retain_property_ordering(): }, }, optional_params={}, - model="gemini-2.0-flash-lite", + model="gemini-1.5-flash", # Old model uses responseSchema which needs propertyOrdering drop_params=False, ) From 16b8ed6786280e0c9c70565cb96aeb1590a4e1c2 Mon Sep 17 00:00:00 2001 From: Lucky-Lodhi2004 Date: Tue, 20 Jan 2026 00:22:58 +0530 Subject: [PATCH 18/58] fixed litellm params (#19315) --- litellm/main.py | 54 ++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 51 insertions(+), 3 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 08384960b16..7b3307cf778 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -580,6 +580,10 @@ async def acompletion( model=model, custom_llm_provider=custom_llm_provider, api_base=completion_kwargs.get("base_url", None), + litellm_params=litellm.types.router.LiteLLM_Params( + model=model, + use_litellm_proxy=kwargs.get("use_litellm_proxy", False), + ), ) fallbacks = fallbacks or litellm.model_fallbacks @@ -1351,6 +1355,10 @@ def completion( # type: ignore # noqa: PLR0915 custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key, + litellm_params=litellm.types.router.LiteLLM_Params( + model=model, + use_litellm_proxy=kwargs.get("use_litellm_proxy", False), + ), ) if not _should_allow_input_examples( @@ -4500,6 +4508,10 @@ async def aembedding(*args, **kwargs) -> EmbeddingResponse: model=model, custom_llm_provider=custom_llm_provider, api_base=kwargs.get("api_base", None), + litellm_params=litellm.types.router.LiteLLM_Params( + model=model, + use_litellm_proxy=kwargs.get("use_litellm_proxy", False), + ), ) # Await normally @@ -4685,6 +4697,10 @@ def embedding( # noqa: PLR0915 custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key, + litellm_params=litellm.types.router.LiteLLM_Params( + model=model, + use_litellm_proxy=kwargs.get("use_litellm_proxy", False), + ), ) if dynamic_api_key is not None: @@ -5820,6 +5836,10 @@ def text_completion( # noqa: PLR0915 model=model, # type: ignore custom_llm_provider=custom_llm_provider, api_base=api_base, + litellm_params=litellm.types.router.LiteLLM_Params( + model=model, # type: ignore + use_litellm_proxy=kwargs.get("use_litellm_proxy", False), + ), ) if custom_llm_provider == "huggingface": @@ -6110,6 +6130,10 @@ async def amoderation( custom_llm_provider=custom_llm_provider, api_base=optional_params.api_base, api_key=optional_params.api_key, + litellm_params=litellm.types.router.LiteLLM_Params( + model=model or "", + use_litellm_proxy=kwargs.get("use_litellm_proxy", False), + ), ) except litellm.BadRequestError: # `model` is optional field for moderation - get_llm_provider will throw BadRequestError if model is not set / not recognized @@ -6175,7 +6199,12 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse: func_with_context = partial(ctx.run, func) _, custom_llm_provider, _, _ = get_llm_provider( - model=model, api_base=kwargs.get("api_base", None) + model=model, + api_base=kwargs.get("api_base", None), + litellm_params=litellm.types.router.LiteLLM_Params( + model=model, + use_litellm_proxy=kwargs.get("use_litellm_proxy", False), + ), ) # Await normally @@ -6279,6 +6308,10 @@ def transcription( custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key, + litellm_params=litellm.types.router.LiteLLM_Params( + model=model, + use_litellm_proxy=kwargs.get("use_litellm_proxy", False), + ), ) # type: ignore if dynamic_api_key is not None: @@ -6454,7 +6487,12 @@ async def aspeech(*args, **kwargs) -> HttpxBinaryResponseContent: func_with_context = partial(ctx.run, func) _, custom_llm_provider, _, _ = get_llm_provider( - model=model, api_base=kwargs.get("api_base", None) + model=model, + api_base=kwargs.get("api_base", None), + litellm_params=litellm.types.router.LiteLLM_Params( + model=model, + use_litellm_proxy=kwargs.get("use_litellm_proxy", False), + ), ) # Await normally @@ -6505,7 +6543,13 @@ def speech( # noqa: PLR0915 model_info = kwargs.get("model_info", None) shared_session = kwargs.get("shared_session", None) model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( - model=model, custom_llm_provider=custom_llm_provider, api_base=api_base + model=model, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + litellm_params=litellm.types.router.LiteLLM_Params( + model=model, + use_litellm_proxy=kwargs.get("use_litellm_proxy", False), + ), ) # type: ignore kwargs.pop("tags", []) @@ -7013,6 +7057,10 @@ async def ahealth_check( custom_llm_provider=custom_llm_provider_from_params, api_base=api_base_from_params, api_key=api_key_from_params, + litellm_params=litellm.types.router.LiteLLM_Params( + model=model, + use_litellm_proxy=model_params.get("use_litellm_proxy", False), + ), ) if model in litellm.model_cost and mode is None: mode = litellm.model_cost[model].get("mode") From 1cce718551404198c89600b403a39279289c7d85 Mon Sep 17 00:00:00 2001 From: 0x1f99d Date: Tue, 20 Jan 2026 05:56:49 +1100 Subject: [PATCH 19/58] fix(bedrock): deduplicate tool calls in assistant history (#15178) (#19324) * fix: Avoid attaching tool calls when a call_id already exists * fix: Prevent MCP responses from reviving past tool calls via previous_response_id * test: Parametrize MCP streaming test to cover OpenAI and Anthropic models * test: Fail MCP streaming test when LiteLLM logs errors during follow-up calls * test: Let MCP tool-execution mock accept new kwargs for streaming tests * chore: fix lint error * docs: Add Google Workload Identity Federation (WIF) documentation to Vertex AI (#19320) - Added new section documenting WIF support for Vertex AI authentication - Included SDK and Proxy configuration examples - Added sample WIF credentials file format for AWS federation - Mentioned LLM Credentials UI as an alternative for credential management - Added link to Google Cloud WIF documentation Co-authored-by: Cursor Agent * fix(bedrock): deduplicate tool calls in assistant history (#15178) * fix(types): add missing Set import to factory.py --------- Co-authored-by: Yuta Saito Co-authored-by: Krish Dholakia Co-authored-by: Cursor Agent Co-authored-by: YutaSaito <36355491+uc4w6c@users.noreply.github.com> --- docs/my-website/docs/providers/vertex.md | 71 +++++++ .../prompt_templates/factory.py | 29 ++- .../transformation.py | 96 +++++++++- .../responses/mcp/mcp_streaming_iterator.py | 1 - .../test_anthropic_dedup_factory.py | 83 ++++++++ .../mcp_tests/test_aresponses_api_with_mcp.py | 178 +++++++++++------- 6 files changed, 375 insertions(+), 83 deletions(-) create mode 100644 tests/litellm_core_utils/test_anthropic_dedup_factory.py diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index 33ebf535d29..be2bf86ab10 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -1390,6 +1390,77 @@ model_list: +### **Workload Identity Federation** + +LiteLLM supports [Google Cloud Workload Identity Federation (WIF)](https://cloud.google.com/iam/docs/workload-identity-federation), which allows you to grant on-premises or multi-cloud workloads access to Google Cloud resources without using a service account key. This is the recommended approach for workloads running in other cloud environments (AWS, Azure, etc.) or on-premises. + +To use Workload Identity Federation, pass the path to your WIF credentials configuration file via `vertex_credentials`: + + + + +```python +from litellm import completion + +response = completion( + model="vertex_ai/gemini-1.5-pro", + messages=[{"role": "user", "content": "Hello!"}], + vertex_credentials="/path/to/wif-credentials.json", # 👈 WIF credentials file + vertex_project="your-gcp-project-id", + vertex_location="us-central1" +) +``` + + + + +```yaml +model_list: + - model_name: gemini-model + litellm_params: + model: vertex_ai/gemini-1.5-pro + vertex_project: your-gcp-project-id + vertex_location: us-central1 + vertex_credentials: /path/to/wif-credentials.json # 👈 WIF credentials file +``` + +Alternatively, you can create credentials in **LLM Credentials** in the LiteLLM UI and use those to authenticate your models: + +```yaml +model_list: + - model_name: gemini-model + litellm_params: + model: vertex_ai/gemini-1.5-pro + vertex_project: your-gcp-project-id + vertex_location: us-central1 + litellm_credential_name: my-vertex-wif-credential # 👈 Reference credential stored in UI +``` + + + + +**WIF Credentials File Format** + +Your WIF credentials JSON file typically looks like this (for AWS federation): + +```json +{ + "type": "external_account", + "audience": "//iam.googleapis.com/projects/PROJECT_NUMBER/locations/global/workloadIdentityPools/POOL_ID/providers/PROVIDER_ID", + "subject_token_type": "urn:ietf:params:aws:token-type:aws4_request", + "service_account_impersonation_url": "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/SERVICE_ACCOUNT_EMAIL:generateAccessToken", + "token_url": "https://sts.googleapis.com/v1/token", + "credential_source": { + "environment_id": "aws1", + "region_url": "http://169.254.169.254/latest/meta-data/placement/availability-zone", + "url": "http://169.254.169.254/latest/meta-data/iam/security-credentials", + "regional_cred_verification_url": "https://sts.{region}.amazonaws.com?Action=GetCallerIdentity&Version=2011-06-15" + } +} +``` + +For more details on setting up Workload Identity Federation, see [Google Cloud WIF documentation](https://cloud.google.com/iam/docs/workload-identity-federation). + ### **Environment Variables** You can set: diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 43ed23587d8..a5fb6f1597c 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -6,7 +6,7 @@ import mimetypes import re import xml.etree.ElementTree as ET from enum import Enum -from typing import Any, Dict, List, Optional, Tuple, Union, cast, overload +from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast, overload from jinja2.sandbox import ImmutableSandboxedEnvironment @@ -2143,6 +2143,9 @@ def anthropic_messages_pt( # noqa: PLR0915 if user_content: new_messages.append({"role": "user", "content": user_content}) + # Track unique tool IDs in this merge block to avoid duplication + unique_tool_ids: Set[str] = set() + assistant_content: List[AnthropicMessagesAssistantMessageValues] = [] ## MERGE CONSECUTIVE ASSISTANT CONTENT ## while msg_i < len(messages) and messages[msg_i]["role"] == "assistant": @@ -2236,13 +2239,25 @@ def anthropic_messages_pt( # noqa: PLR0915 assistant_tool_calls, web_search_results=_web_search_results, ) - # AnthropicMessagesAssistantMessageValues includes AnthropicMessagesToolUseParam - assistant_content.extend( - cast( - List[AnthropicMessagesAssistantMessageValues], - tool_invoke_results, + + # Prevent "tool_use ids must be unique" errors by filtering duplicates + # This can happen when merging history that already contains the tool calls + for item in tool_invoke_results: + # tool_use items are typically dicts, but handle objects just in case + item_id = ( + item.get("id") + if isinstance(item, dict) + else getattr(item, "id", None) + ) + + if item_id: + if item_id in unique_tool_ids: + continue + unique_tool_ids.add(item_id) + + assistant_content.append( + cast(AnthropicMessagesAssistantMessageValues, item) ) - ) assistant_function_call = assistant_content_block.get("function_call") diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index ad910c9cd97..eaa80c6cfe4 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2,7 +2,8 @@ Handles transforming from Responses API -> LiteLLM completion (Chat Completion API) """ -from typing import Any, Dict, List, Literal, Optional, Tuple, Union, cast +from collections.abc import Sequence +from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union, cast from openai.types.responses import ResponseFunctionToolCall from openai.types.responses.tool_param import FunctionToolParam @@ -378,6 +379,7 @@ class LiteLLMCompletionResponsesConfig: if isinstance(input, str): messages.append(ChatCompletionUserMessage(role="user", content=input)) elif isinstance(input, list): + existing_tool_call_ids: Set[str] = set() for _input in input: chat_completion_messages = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( input_item=_input @@ -390,12 +392,98 @@ class LiteLLMCompletionResponsesConfig: input_item=_input ): tool_call_output_messages.extend(chat_completion_messages) - else: - messages.extend(chat_completion_messages) + continue - messages.extend(tool_call_output_messages) + if LiteLLMCompletionResponsesConfig._is_input_item_function_call( + input_item=_input + ): + call_id_raw = _input.get("call_id") or _input.get("id") or "" + if call_id_raw: + existing_tool_call_ids.add(str(call_id_raw)) + + messages.extend(chat_completion_messages) + + deduped_tool_call_messages = ( + LiteLLMCompletionResponsesConfig._deduplicate_tool_call_output_messages( + tool_call_output_messages=tool_call_output_messages, + existing_tool_call_ids=existing_tool_call_ids, + ) + ) + messages.extend(deduped_tool_call_messages) return messages + @staticmethod + def _deduplicate_tool_call_output_messages( + tool_call_output_messages: List[ + Union[ + AllMessageValues, + GenericChatCompletionMessage, + ChatCompletionMessageToolCall, + ChatCompletionResponseMessage, + ] + ], + existing_tool_call_ids: Set[str], + ) -> List[ + Union[ + AllMessageValues, + GenericChatCompletionMessage, + ChatCompletionMessageToolCall, + ChatCompletionResponseMessage, + ] + ]: + """Return tool call outputs after dropping assistant entries with duplicate call_ids.""" + if not tool_call_output_messages: + return [] + + filtered_messages: List[ + Union[ + AllMessageValues, + GenericChatCompletionMessage, + ChatCompletionMessageToolCall, + ChatCompletionResponseMessage, + ] + ] = [] + seen_tool_call_ids: Set[str] = set(existing_tool_call_ids) + + for tool_call_message in tool_call_output_messages: + if isinstance(tool_call_message, dict): + role = tool_call_message.get("role", "") + else: + role = getattr(tool_call_message, "role", "") + call_id = "" + + if role == "assistant": + tool_calls: Any = None + if isinstance(tool_call_message, dict): + tool_calls = tool_call_message.get("tool_calls") + else: + tool_calls = getattr(tool_call_message, "tool_calls", None) + + if ( + isinstance(tool_calls, Sequence) + and not isinstance(tool_calls, (str, bytes)) + and len(tool_calls) > 0 + ): + first_call = tool_calls[0] + call_id_raw = None + if isinstance(first_call, dict): + call_id_raw = first_call.get("id") + else: + call_id_raw = getattr(first_call, "id", None) + + if call_id_raw: + call_id = str(call_id_raw) + + if call_id and call_id in seen_tool_call_ids and role == "assistant": + continue + + if call_id and role == "assistant": + seen_tool_call_ids.add(call_id) + + filtered_messages.append(tool_call_message) + + return filtered_messages + @staticmethod def _ensure_tool_call_output_has_corresponding_tool_call( messages: List[Union[AllMessageValues, GenericChatCompletionMessage]], diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index c00c2a2f3b2..ac040d3d6ec 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -655,7 +655,6 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): follow_up_params.update( { "input": follow_up_input, - "previous_response_id": self.collected_response.id, # type: ignore[attr-defined] "stream": True, } ) diff --git a/tests/litellm_core_utils/test_anthropic_dedup_factory.py b/tests/litellm_core_utils/test_anthropic_dedup_factory.py new file mode 100644 index 00000000000..99248db0aec --- /dev/null +++ b/tests/litellm_core_utils/test_anthropic_dedup_factory.py @@ -0,0 +1,83 @@ + +import sys +import os +import pytest +sys.path.insert(0, os.path.abspath(".")) + +from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt + +def test_anthropic_deduplication_logic(): + """ + Verify that anthropic_messages_pt correctly deduplicates tool calls + when merging consecutive assistant messages. + + Scenario: + - User message + - Assistant message with tool call A + - Assistant message (merged) with tool call A (duplicate) and tool call B (new) + + Expected Result: + - Assistant message contains tool call A and tool call B exactly once. + """ + messages = [ + {"role": "user", "content": "Weather in Paris?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "tool_call_unique_1", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"} + } + ] + }, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "tool_call_unique_1", # Duplicate! Should be removed + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"} + }, + { + "id": "tool_call_unique_2", # New! Should be kept + "type": "function", + "function": {"name": "get_time", "arguments": "{}"} + } + ] + } + ] + + # Run transformation + # We pass dummy model/provider args as they are required but valid for this test + result = anthropic_messages_pt( + messages=messages, + model="claude-3-opus-20240229", + llm_provider="anthropic" + ) + + # Inspect results + # We expect 1 user message and 1 merged assistant message + assert len(result) == 2 + assert result[0]["role"] == "user" + assert result[1]["role"] == "assistant" + + assistant_content = result[1]["content"] + + # Filter for tool_use blocks + tool_uses = [ + b for b in assistant_content + if isinstance(b, dict) and b.get("type") == "tool_use" + ] + + # We expect exactly 2 tool uses (one for unique_1, one for unique_2) + # The duplicate unique_1 should be gone. + assert len(tool_uses) == 2 + + ids = [t["id"] for t in tool_uses] + assert "tool_call_unique_1" in ids + assert "tool_call_unique_2" in ids + assert ids.count("tool_call_unique_1") == 1 + assert ids.count("tool_call_unique_2") == 1 diff --git a/tests/mcp_tests/test_aresponses_api_with_mcp.py b/tests/mcp_tests/test_aresponses_api_with_mcp.py index 865a580f0ca..57c79039ee0 100644 --- a/tests/mcp_tests/test_aresponses_api_with_mcp.py +++ b/tests/mcp_tests/test_aresponses_api_with_mcp.py @@ -1,3 +1,4 @@ +import logging import os import sys import pytest @@ -660,16 +661,25 @@ async def test_streaming_mcp_events_validation(): @pytest.mark.asyncio -async def test_streaming_responses_api_with_mcp_tools(): +@pytest.mark.parametrize( + "model", + [ + pytest.param("gpt-4o-mini", id="openai"), + pytest.param("claude-haiku-4-5", id="anthropic"), + ], +) +async def test_streaming_responses_api_with_mcp_tools( + model: str, caplog: pytest.LogCaptureFixture +): """ Test the streaming responses API with MCP tools when using server_url="litellm_proxy" Under the hood the follow occurs - MCP: responses called litellm MCP manager.list_tools (MOCKED) - - Request 1: Made to gpt-4o with fetched tools (REAL LLM CALL) + - Request 1: Made to model under test with fetched tools (REAL LLM CALL) - MCP: Execute tool call from request 1 and returns result (MOCKED) - - Request 2: Made to gpt-4o with fetched tools and tool results (REAL LLM CALL) + - Request 2: Made to model under test with fetched tools and tool results (REAL LLM CALL) Return the user the result of request 2 """ @@ -693,75 +703,101 @@ async def test_streaming_responses_api_with_mcp_tools(): ] # Only mock the MCP-specific operations, let LLM responses be real - with patch.object(LiteLLM_Proxy_MCP_Handler, '_get_mcp_tools_from_manager', new_callable=AsyncMock) as mock_get_tools, \ - patch.object(LiteLLM_Proxy_MCP_Handler, '_execute_tool_calls', new_callable=AsyncMock) as mock_execute_tools: - - # Setup MCP mocks only - mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"]) - - # Create a dynamic mock that will match the actual tool call ID from the LLM response - def mock_execute_tool_calls_side_effect(tool_calls, user_api_key_auth): - """Mock function that returns results matching the actual tool call IDs from the LLM""" - results = [] - for tool_call in tool_calls: - # Extract call_id from the tool call - call_id = None - if isinstance(tool_call, dict): - call_id = tool_call.get("call_id") or tool_call.get("id") - elif hasattr(tool_call, 'call_id'): - call_id = tool_call.call_id - elif hasattr(tool_call, 'id'): - call_id = tool_call.id - - if call_id: - results.append({ - "tool_call_id": call_id, - "result": "LiteLLM is a unified interface for 100+ LLMs that translates inputs to provider-specific completion endpoints and provides consistent OpenAI-format output." - }) - return results - - mock_execute_tools.side_effect = mock_execute_tool_calls_side_effect - - # Make the actual call - LLM responses will be real - mcp_tool_config = cast(Any, { - "type": "mcp", - "server_url": "litellm_proxy", - "require_approval": "never" - }) - response = await litellm.aresponses( - model="gpt-4o-mini", - tools=[mcp_tool_config], - tool_choice="required", - input=[ + with caplog.at_level(logging.ERROR): + with patch.object( + LiteLLM_Proxy_MCP_Handler, + '_get_mcp_tools_from_manager', + new_callable=AsyncMock, + ) as mock_get_tools, patch.object( + LiteLLM_Proxy_MCP_Handler, + '_execute_tool_calls', + new_callable=AsyncMock, + ) as mock_execute_tools: + # Setup MCP mocks only + mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"]) + + # Create a dynamic mock that will match the actual tool call ID from the LLM response + def mock_execute_tool_calls_side_effect( + tool_calls, user_api_key_auth, **kwargs + ): + """Mock function that returns results matching the actual tool call IDs from the LLM""" + results = [] + for tool_call in tool_calls: + # Extract call_id from the tool call + call_id = None + if isinstance(tool_call, dict): + call_id = tool_call.get("call_id") or tool_call.get("id") + elif hasattr(tool_call, 'call_id'): + call_id = tool_call.call_id + elif hasattr(tool_call, 'id'): + call_id = tool_call.id + + if call_id: + results.append( + { + "tool_call_id": call_id, + "result": "LiteLLM is a unified interface for 100+ LLMs that translates inputs to provider-specific completion endpoints and provides consistent OpenAI-format output.", + } + ) + return results + + mock_execute_tools.side_effect = mock_execute_tool_calls_side_effect + + # Make the actual call - LLM responses will be real + mcp_tool_config = cast( + Any, { - "role": "user", - "type": "message", - "content": "give me a TLDR of what BerriAI/litellm is about" - } - ], - stream=True - ) - - print(f"📋 Response type: {type(response)}") - assert hasattr(response, '__aiter__'), "Response should be an async streaming response" - - # Collect streaming chunks - chunks = [] - async for chunk in response: - chunks.append(chunk) - print(f"📦 Chunk type: {getattr(chunk, 'type', 'unknown')}") - - print(f"📊 Total chunks received: {len(chunks)}") - - # Verify MCP mocks were called (may be called multiple times in streaming) - assert mock_get_tools.call_count >= 1, f"Expected MCP tools to be fetched at least once, got {mock_get_tools.call_count}" - print(f"MCP tools fetched: {len(mock_mcp_tools)}") - - # Verify we got a response - assert response is not None - assert len(chunks) > 0, "Should have received streaming chunks" - - print("Basic streaming responses API with MCP tools test passed!") + "type": "mcp", + "server_url": "litellm_proxy", + "require_approval": "never", + }, + ) + response = await litellm.aresponses( + model=model, + tools=[mcp_tool_config], + tool_choice="required", + input=[ + { + "role": "user", + "type": "message", + "content": "give me a TLDR of what BerriAI/litellm is about", + } + ], + stream=True, + ) + + print(f"📋 Response type: {type(response)}") + assert hasattr(response, '__aiter__'), "Response should be an async streaming response" + + # Collect streaming chunks + chunks = [] + async for chunk in response: + chunks.append(chunk) + print(f"📦 Chunk type: {getattr(chunk, 'type', 'unknown')}") + + print(f"📊 Total chunks received: {len(chunks)}") + + # Verify MCP mocks were called (may be called multiple times in streaming) + assert ( + mock_get_tools.call_count >= 1 + ), f"Expected MCP tools to be fetched at least once, got {mock_get_tools.call_count}" + print(f"MCP tools fetched: {len(mock_mcp_tools)}") + + # Verify we got a response + assert response is not None + assert len(chunks) > 0, "Should have received streaming chunks" + + print("Basic streaming responses API with MCP tools test passed!") + + lite_errors = [ + record + for record in caplog.records + if record.levelno >= logging.ERROR + and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage()) + ] + assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join( + record.getMessage() for record in lite_errors + ) @pytest.mark.asyncio From 1645be3a2f86be5a73af31c42d3b4923c78ecdce Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lo=C3=AFc=20Rakotoson?= Date: Tue, 20 Jan 2026 04:24:03 +0100 Subject: [PATCH 20/58] feat: implement SafeAttributeModel for safe attribute access in models (#18321) --- litellm/types/utils.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e1f1c1c4b5b..e71fcac3897 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -63,6 +63,19 @@ def _generate_id(): # private helper function return "chatcmpl-" + str(uuid.uuid4()) + +class SafeAttributeModel: + """ + A base model that provides safe attribute access. + """ + def __delattr__(self, name): + try: + super().__delattr__(name) + except AttributeError: + # noop if attribute does not exist + pass + + class LiteLLMCommonStrings(Enum): redacted_by_litellm = "redacted by litellm. 'litellm.turn_off_message_logging=True'" llm_provider_not_provided = "Unmapped LLM provider for this endpoint. You passed model={model}, custom_llm_provider={custom_llm_provider}. Check supported provider and route: https://docs.litellm.ai/docs/providers" @@ -1021,7 +1034,7 @@ def add_provider_specific_fields( setattr(object, "provider_specific_fields", provider_specific_fields) -class Message(OpenAIObject): +class Message(SafeAttributeModel, OpenAIObject): content: Optional[str] role: Literal["assistant", "user", "system", "tool", "function"] tool_calls: Optional[List[ChatCompletionMessageToolCall]] @@ -1141,7 +1154,7 @@ class Message(OpenAIObject): return self.dict() -class Delta(OpenAIObject): +class Delta(SafeAttributeModel, OpenAIObject): reasoning_content: Optional[str] = None thinking_blocks: Optional[ List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]] @@ -1238,7 +1251,7 @@ class Delta(OpenAIObject): setattr(self, key, value) -class Choices(OpenAIObject): +class Choices(SafeAttributeModel, OpenAIObject): finish_reason: str index: int message: Message @@ -1331,6 +1344,7 @@ class CacheCreationTokenDetails(BaseModel): class PromptTokensDetailsWrapper( + SafeAttributeModel, PromptTokensDetails ): # extends with image generation fields (text_tokens, image_tokens) text_tokens: Optional[int] = None @@ -1378,7 +1392,7 @@ class ServerToolUse(BaseModel): tool_search_requests: Optional[int] = None -class Usage(CompletionUsage): +class Usage(SafeAttributeModel, CompletionUsage): _cache_creation_input_tokens: int = PrivateAttr( 0 ) # hidden param for prompt caching. Might change, once openai introduces their equivalent. From f945fd9a841bfd4f4c87262547467771223b6a39 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 20 Jan 2026 11:15:37 +0530 Subject: [PATCH 21/58] Fix: ID mismatch between text-start and text-delta --- .../streaming_iterator.py | 34 ++- .../test_litellm_completion_responses.py | 217 +++++++++++++++++- 2 files changed, 243 insertions(+), 8 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index def2f72437d..ece94cd4214 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -79,6 +79,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): Union[ModelResponse, TextCompletionResponse] ] = None self.final_text: str = "" + self._cached_item_id: Optional[str] = None + self._cached_response_id: Optional[str] = None def _default_response_created_event_data(self) -> dict: response_created_event_data = { @@ -150,12 +152,15 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) def create_output_item_added_event(self) -> OutputItemAddedEvent: + if self._cached_item_id is None: + self._cached_item_id = f"msg_{str(uuid.uuid4())}" + return OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, output_index=0, item=BaseLiteLLMOpenAIResponseObject( **{ - "id": f"msg_{str(uuid.uuid4())}", + "id": self._cached_item_id, "type": "message", "status": "in_progress", "role": "assistant", @@ -165,9 +170,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) def create_content_part_added_event(self) -> ContentPartAddedEvent: + if self._cached_item_id is None: + self._cached_item_id = f"msg_{str(uuid.uuid4())}" + return ContentPartAddedEvent( type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED, - item_id=f"msg_{str(uuid.uuid4())}", + item_id=self._cached_item_id, output_index=0, content_index=0, part=BaseLiteLLMOpenAIResponseObject( @@ -189,9 +197,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def create_output_text_done_event( self, litellm_complete_object: ModelResponse ) -> OutputTextDoneEvent: + if self._cached_item_id is None: + self._cached_item_id = f"msg_{str(uuid.uuid4())}" + return OutputTextDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, - item_id=f"msg_{str(uuid.uuid4())}", + item_id=self._cached_item_id, output_index=0, content_index=0, text=getattr(litellm_complete_object.choices[0].message, "content", "") # type: ignore @@ -201,6 +212,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def create_output_content_part_done_event( self, litellm_complete_object: ModelResponse ) -> ContentPartDoneEvent: + if self._cached_item_id is None: + self._cached_item_id = f"msg_{str(uuid.uuid4())}" text = getattr(litellm_complete_object.choices[0].message, "content", "") or "" # type: ignore reasoning_content = getattr(litellm_complete_object.choices[0].message, "reasoning_content", "") or "" # type: ignore @@ -226,7 +239,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return ContentPartDoneEvent( type=ResponsesAPIStreamEvents.CONTENT_PART_DONE, - item_id=f"msg_{str(uuid.uuid4())}", + item_id=self._cached_item_id, output_index=0, content_index=0, part=part, @@ -235,6 +248,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): def create_output_item_done_event( self, litellm_complete_object: ModelResponse ) -> OutputItemDoneEvent: + if self._cached_item_id is None: + self._cached_item_id = f"msg_{str(uuid.uuid4())}" + text = self.litellm_model_response.choices[0].message.content or "" # type: ignore annotations = getattr(self.litellm_model_response.choices[0].message, "annotations", None) # type: ignore @@ -247,7 +263,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): sequence_number=1, item=BaseLiteLLMOpenAIResponseObject( **{ - "id": f"msg_{str(uuid.uuid4())}", + "id": self._cached_item_id, "status": "completed", "type": "message", "role": "assistant", @@ -413,6 +429,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): and the ReasoningSummaryTextDeltaEvent, which is used by the responses API to emit reasoning content. It also handles emitting annotation.added events when annotations are detected in the chunk. """ + if self._cached_item_id is None and chunk.id: + self._cached_item_id = chunk.id + item_id = self._cached_item_id or chunk.id + # Check if this chunk has annotations first (before processing text/reasoning) # This ensures we detect and queue annotation events from the annotation chunk if chunk.choices and hasattr(chunk.choices[0].delta, "annotations"): @@ -430,7 +450,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): annotation_dict = annotation.model_dump() if hasattr(annotation, 'model_dump') else dict(annotation) event = OutputTextAnnotationAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, - item_id=chunk.id, + item_id=item_id, output_index=0, content_index=0, annotation_index=idx, @@ -457,7 +477,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if delta_content: return OutputTextDeltaEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, - item_id=chunk.id, + item_id=item_id, output_index=0, content_index=0, delta=delta_content, 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 59c630b6a5b..f7a1984d32f 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 @@ -1351,4 +1351,219 @@ class TestUsageTransformation: assert response_usage.output_tokens == 27 assert response_usage.total_tokens == 36 assert response_usage.input_tokens_details is None - assert response_usage.output_tokens_details is None \ No newline at end of file + assert response_usage.output_tokens_details is None + + +class TestStreamingIDConsistency: + """Test cases for consistent IDs across streaming events (issue #14962)""" + + def test_streaming_iterator_uses_consistent_item_ids(self): + """ + Test that all streaming events use the same item_id throughout the stream. + This fixes the issue where text-start, text-delta, and text-end events + had different IDs, breaking SDK text accumulation. + + Reproduces: https://github.com/BerriAI/litellm/issues/14962 + """ + from unittest.mock import Mock + + import litellm + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + # Create a mock stream wrapper + mock_stream_wrapper = Mock(spec=litellm.CustomStreamWrapper) + mock_logging_obj = Mock() + mock_stream_wrapper.logging_obj = mock_logging_obj + + # Create the streaming iterator + iterator = LiteLLMCompletionStreamingIterator( + model="gemini/gemini-2.5-flash-lite", + litellm_custom_stream_wrapper=mock_stream_wrapper, + request_input="Say Hello World", + responses_api_request={}, + custom_llm_provider="gemini", + ) + + # Simulate streaming chunks with different IDs (as Gemini does) + chunk1 = ModelResponseStream( + id="chatcmpl-first-id", + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="Hello", role="assistant"), + finish_reason=None, + ) + ], + created=1234567890, + model="gemini-2.5-flash-lite", + object="chat.completion.chunk", + ) + + chunk2 = ModelResponseStream( + id="chatcmpl-second-id", # Different ID from chunk1 + choices=[ + StreamingChoices( + index=0, + delta=Delta(content=" World", role=None), + finish_reason=None, + ) + ], + created=1234567890, + model="gemini-2.5-flash-lite", + object="chat.completion.chunk", + ) + + chunk3 = ModelResponseStream( + id="chatcmpl-third-id", # Different ID from chunk1 and chunk2 + choices=[ + StreamingChoices( + index=0, + delta=Delta(content="", role=None), + finish_reason="stop", + ) + ], + created=1234567890, + model="gemini-2.5-flash-lite", + object="chat.completion.chunk", + ) + + # Transform chunks to response API events + event1 = iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk1) + event2 = iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk2) + event3 = iterator._transform_chat_completion_chunk_to_response_api_chunk(chunk3) + + # Assert: All events should use the same item_id (from the first chunk) + assert event1 is not None, "First event should not be None" + assert event2 is not None, "Second event should not be None" + + # Extract item_ids from events + item_id_1 = getattr(event1, "item_id", None) + item_id_2 = getattr(event2, "item_id", None) + + assert item_id_1 is not None, "First event should have an item_id" + assert item_id_2 is not None, "Second event should have an item_id" + + # The critical assertion: IDs should match across all events + assert item_id_1 == item_id_2, ( + f"Item IDs should be consistent across streaming events. " + f"Got {item_id_1} and {item_id_2}. " + f"This breaks SDK text accumulation (issue #14962)." + ) + + # Verify the cached ID is set and matches + assert iterator._cached_item_id is not None, "Iterator should cache the item_id" + assert iterator._cached_item_id == item_id_1, "Cached ID should match event IDs" + assert iterator._cached_item_id == "chatcmpl-first-id", "Should use the first chunk's ID" + + def test_streaming_iterator_initial_events_use_cached_id(self): + """ + Test that initial events (output_item_added, content_part_added) also use the cached ID. + """ + from unittest.mock import Mock + + import litellm + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + + # Create a mock stream wrapper + mock_stream_wrapper = Mock(spec=litellm.CustomStreamWrapper) + mock_logging_obj = Mock() + mock_stream_wrapper.logging_obj = mock_logging_obj + + # Create the streaming iterator + iterator = LiteLLMCompletionStreamingIterator( + model="gemini/gemini-2.5-flash-lite", + litellm_custom_stream_wrapper=mock_stream_wrapper, + request_input="Test", + responses_api_request={}, + ) + + # Create initial events + output_item_event = iterator.create_output_item_added_event() + content_part_event = iterator.create_content_part_added_event() + + # Extract IDs + output_item_id = getattr(output_item_event.item, "id", None) + content_part_id = getattr(content_part_event, "item_id", None) + + # Assert: Both should use the same cached ID + assert output_item_id is not None, "Output item should have an ID" + assert content_part_id is not None, "Content part should have an item_id" + assert output_item_id == content_part_id, ( + f"Initial events should use consistent IDs. " + f"Got output_item_id={output_item_id}, content_part_id={content_part_id}" + ) + + # Verify it matches the cached ID + assert iterator._cached_item_id is not None + assert iterator._cached_item_id == output_item_id + + def test_streaming_iterator_done_events_use_cached_id(self): + """ + Test that done events (output_text_done, content_part_done, output_item_done) use the cached ID. + """ + from unittest.mock import Mock + + import litellm + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + from litellm.types.utils import Choices, Message, ModelResponse + + # Create a mock stream wrapper + mock_stream_wrapper = Mock(spec=litellm.CustomStreamWrapper) + mock_logging_obj = Mock() + mock_stream_wrapper.logging_obj = mock_logging_obj + mock_logging_obj._response_cost_calculator = Mock(return_value=0.001) + + # Create the streaming iterator + iterator = LiteLLMCompletionStreamingIterator( + model="gemini/gemini-2.5-flash-lite", + litellm_custom_stream_wrapper=mock_stream_wrapper, + request_input="Test", + responses_api_request={}, + ) + + # Set up a complete model response + complete_response = ModelResponse( + id="test-response-id", + created=1234567890, + model="gemini-2.5-flash-lite", + object="chat.completion", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="Hello World", role="assistant"), + ) + ], + ) + iterator.litellm_model_response = complete_response + + # Create done events + text_done_event = iterator.create_output_text_done_event(complete_response) + content_done_event = iterator.create_output_content_part_done_event(complete_response) + item_done_event = iterator.create_output_item_done_event(complete_response) + + # Extract IDs + text_done_id = getattr(text_done_event, "item_id", None) + content_done_id = getattr(content_done_event, "item_id", None) + item_done_id = getattr(item_done_event.item, "id", None) + + # Assert: All done events should use the same cached ID + assert text_done_id is not None, "Text done event should have an item_id" + assert content_done_id is not None, "Content done event should have an item_id" + assert item_done_id is not None, "Item done event should have an id" + + assert text_done_id == content_done_id == item_done_id, ( + f"All done events should use consistent IDs. " + f"Got text_done={text_done_id}, content_done={content_done_id}, item_done={item_done_id}" + ) + + # Verify it matches the cached ID + assert iterator._cached_item_id is not None + assert iterator._cached_item_id == text_done_id \ No newline at end of file From 5f80e8d5e8b52a647ec757b630c134c4b097ed2b Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 20 Jan 2026 15:28:09 +0530 Subject: [PATCH 22/58] Fix for Prometheus Metric Cardinality Issue with /responses Endpoint --- litellm/proxy/auth/auth_utils.py | 82 +++++++++ litellm/proxy/auth/user_api_key_auth.py | 5 +- .../integrations/test_prometheus_labels.py | 160 +++++++++++++++++- 3 files changed, 242 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 1a7f05716b3..9b9a988c07a 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -311,6 +311,88 @@ def get_request_route(request: Request) -> str: return request.url.path +def normalize_request_route(route: str) -> str: + """ + Normalize request routes by replacing dynamic path parameters with placeholders. + + This prevents high cardinality in Prometheus metrics by collapsing routes like: + - /v1/responses/1234567890 -> /v1/responses/{response_id} + - /v1/threads/thread_123 -> /v1/threads/{thread_id} + + Args: + route: The request route path + + Returns: + Normalized route with dynamic parameters replaced by placeholders + + Examples: + >>> normalize_request_route("/v1/responses/abc123") + '/v1/responses/{response_id}' + >>> normalize_request_route("/v1/responses/abc123/cancel") + '/v1/responses/{response_id}/cancel' + >>> normalize_request_route("/chat/completions") + '/chat/completions' + """ + # Define patterns for routes with dynamic IDs + # Format: (regex_pattern, replacement_template) + patterns = [ + # Responses API - must come before generic patterns + (r'^(/(?:openai/)?v1/responses)/([^/]+)(/input_items)$', r'\1/{response_id}\3'), + (r'^(/(?:openai/)?v1/responses)/([^/]+)(/cancel)$', r'\1/{response_id}\3'), + (r'^(/(?:openai/)?v1/responses)/([^/]+)$', r'\1/{response_id}'), + (r'^(/responses)/([^/]+)(/input_items)$', r'\1/{response_id}\3'), + (r'^(/responses)/([^/]+)(/cancel)$', r'\1/{response_id}\3'), + (r'^(/responses)/([^/]+)$', r'\1/{response_id}'), + + # Threads API + (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)/([^/]+)$', r'\1/{thread_id}\3/{run_id}\5/{step_id}'), + (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)$', r'\1/{thread_id}\3/{run_id}\5'), + (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/cancel)$', r'\1/{thread_id}\3/{run_id}\5'), + (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/submit_tool_outputs)$', r'\1/{thread_id}\3/{run_id}\5'), + (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)$', r'\1/{thread_id}\3/{run_id}'), + (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)$', r'\1/{thread_id}\3'), + (r'^(/(?:openai/)?v1/threads)/([^/]+)(/messages)/([^/]+)$', r'\1/{thread_id}\3/{message_id}'), + (r'^(/(?:openai/)?v1/threads)/([^/]+)(/messages)$', r'\1/{thread_id}\3'), + (r'^(/(?:openai/)?v1/threads)/([^/]+)$', r'\1/{thread_id}'), + + # Vector Stores API + (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)/([^/]+)$', r'\1/{vector_store_id}\3/{file_id}'), + (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)$', r'\1/{vector_store_id}\3'), + (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)/([^/]+)$', r'\1/{vector_store_id}\3/{batch_id}'), + (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)$', r'\1/{vector_store_id}\3'), + (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)$', r'\1/{vector_store_id}'), + + # Assistants API + (r'^(/(?:openai/)?v1/assistants)/([^/]+)$', r'\1/{assistant_id}'), + + # Files API + (r'^(/(?:openai/)?v1/files)/([^/]+)(/content)$', r'\1/{file_id}\3'), + (r'^(/(?:openai/)?v1/files)/([^/]+)$', r'\1/{file_id}'), + + # Batches API + (r'^(/(?:openai/)?v1/batches)/([^/]+)(/cancel)$', r'\1/{batch_id}\3'), + (r'^(/(?:openai/)?v1/batches)/([^/]+)$', r'\1/{batch_id}'), + + # Fine-tuning API + (r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/events)$', r'\1/{fine_tuning_job_id}\3'), + (r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/cancel)$', r'\1/{fine_tuning_job_id}\3'), + (r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/checkpoints)$', r'\1/{fine_tuning_job_id}\3'), + (r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)$', r'\1/{fine_tuning_job_id}'), + + # Models API + (r'^(/(?:openai/)?v1/models)/([^/]+)$', r'\1/{model}'), + ] + + # Apply patterns in order + for pattern, replacement in patterns: + normalized = re.sub(pattern, replacement, route) + if normalized != route: + return normalized + + # Return original route if no pattern matched + return route + + async def check_if_request_size_is_safe(request: Request) -> bool: """ Enterprise Only: diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index bc0c164a0ad..7e7c7c8c90c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -28,8 +28,8 @@ from litellm.proxy.auth.auth_checks import ( _delete_cache_key_object, _get_user_role, _is_user_proxy_admin, - _virtual_key_max_budget_check, _virtual_key_max_budget_alert_check, + _virtual_key_max_budget_check, _virtual_key_soft_budget_check, can_key_call_model, common_checks, @@ -45,6 +45,7 @@ from litellm.proxy.auth.auth_utils import ( get_end_user_id_from_request_body, get_model_from_request, get_request_route, + normalize_request_route, pre_db_read_auth_checks, route_in_additonal_public_routes, ) @@ -1261,7 +1262,7 @@ async def user_api_key_auth( if end_user_id is not None: user_api_key_auth_obj.end_user_id = end_user_id - user_api_key_auth_obj.request_route = route + user_api_key_auth_obj.request_route = normalize_request_route(route) return user_api_key_auth_obj diff --git a/tests/test_litellm/integrations/test_prometheus_labels.py b/tests/test_litellm/integrations/test_prometheus_labels.py index 8a4295e98fc..c0b863ef6ee 100644 --- a/tests/test_litellm/integrations/test_prometheus_labels.py +++ b/tests/test_litellm/integrations/test_prometheus_labels.py @@ -3,7 +3,7 @@ Unit tests for prometheus metric labels configuration """ from litellm.types.integrations.prometheus import ( PrometheusMetricLabels, - UserAPIKeyLabelNames + UserAPIKeyLabelNames, ) @@ -42,9 +42,10 @@ def test_user_email_label_exists(): def test_prometheus_metric_labels_structure(): """Test that all required prometheus metrics have proper label structure""" - from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS from typing import get_args + from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS + # Test a few key metrics to ensure they have proper label structure test_metrics = [ "litellm_proxy_total_requests_metric", @@ -69,8 +70,161 @@ def test_prometheus_metric_labels_structure(): print(f"✅ {metric_name} has proper label structure with user_email") +def test_route_normalization_for_responses_api(): + """ + Test that route normalization prevents high cardinality in Prometheus metrics + for the /v1/responses/{response_id} endpoint. + + Issue: https://github.com/BerriAI/litellm/issues/XXXX + Each unique response ID was creating a separate metric line, causing the + /metrics endpoint to grow to ~30MB and take ~40 seconds to respond. + + Fix: Routes are normalized to collapse dynamic IDs into placeholders. + """ + from litellm.proxy.auth.auth_utils import normalize_request_route + + # Test responses API routes + responses_routes = [ + ("/v1/responses/1234567890", "/v1/responses/{response_id}"), + ("/v1/responses/9876543210", "/v1/responses/{response_id}"), + ("/v1/responses/abcdefghij", "/v1/responses/{response_id}"), + ("/v1/responses/resp_abc123", "/v1/responses/{response_id}"), + ("/v1/responses/litellm_poll_xyz", "/v1/responses/{response_id}"), + ] + + for original, expected in responses_routes: + normalized = normalize_request_route(original) + assert normalized == expected, \ + f"Failed: {original} -> {normalized} (expected {expected})" + + # Verify cardinality reduction + unique_normalized = set(normalize_request_route(route) for route, _ in responses_routes) + assert len(unique_normalized) == 1, \ + f"Expected 1 unique normalized route, got {len(unique_normalized)}: {unique_normalized}" + + print(f"✅ Responses API routes: {len(responses_routes)} different IDs normalized to 1 metric label") + + +def test_route_normalization_for_sub_routes(): + """Test that sub-routes like /cancel and /input_items are normalized correctly""" + from litellm.proxy.auth.auth_utils import normalize_request_route + + sub_routes = [ + ("/v1/responses/id1/cancel", "/v1/responses/{response_id}/cancel"), + ("/v1/responses/id2/cancel", "/v1/responses/{response_id}/cancel"), + ("/v1/responses/id3/input_items", "/v1/responses/{response_id}/input_items"), + ("/openai/v1/responses/id4/input_items", "/openai/v1/responses/{response_id}/input_items"), + ] + + for original, expected in sub_routes: + normalized = normalize_request_route(original) + assert normalized == expected, \ + f"Failed: {original} -> {normalized} (expected {expected})" + + print("✅ Sub-routes normalized correctly") + + +def test_route_normalization_preserves_static_routes(): + """Test that static routes are not affected by normalization""" + from litellm.proxy.auth.auth_utils import normalize_request_route + + static_routes = [ + "/chat/completions", + "/v1/chat/completions", + "/v1/embeddings", + "/health", + "/metrics", + "/v1/models", + "/v1/responses", # List endpoint without ID + ] + + for route in static_routes: + normalized = normalize_request_route(route) + assert normalized == route, \ + f"Static route should not be modified: {route} -> {normalized}" + + print(f"✅ {len(static_routes)} static routes preserved") + + +def test_route_normalization_other_dynamic_apis(): + """Test normalization for other OpenAI-compatible APIs with dynamic IDs""" + from litellm.proxy.auth.auth_utils import normalize_request_route + + test_cases = [ + # Threads API + ("/v1/threads/thread_123", "/v1/threads/{thread_id}"), + ("/v1/threads/thread_abc/messages", "/v1/threads/{thread_id}/messages"), + ("/v1/threads/thread_abc/runs/run_123", "/v1/threads/{thread_id}/runs/{run_id}"), + + # Vector Stores API + ("/v1/vector_stores/vs_123", "/v1/vector_stores/{vector_store_id}"), + ("/v1/vector_stores/vs_123/files", "/v1/vector_stores/{vector_store_id}/files"), + + # Assistants API + ("/v1/assistants/asst_123", "/v1/assistants/{assistant_id}"), + + # Files API + ("/v1/files/file_123", "/v1/files/{file_id}"), + ("/v1/files/file_123/content", "/v1/files/{file_id}/content"), + + # Batches API + ("/v1/batches/batch_123", "/v1/batches/{batch_id}"), + ("/v1/batches/batch_123/cancel", "/v1/batches/{batch_id}/cancel"), + ] + + for original, expected in test_cases: + normalized = normalize_request_route(original) + assert normalized == expected, \ + f"Failed: {original} -> {normalized} (expected {expected})" + + print(f"✅ {len(test_cases)} other API routes normalized correctly") + + +def test_prometheus_metrics_use_normalized_routes(): + """ + Test that Prometheus metrics use the normalized route in labels + to prevent high cardinality. + """ + from unittest.mock import MagicMock + + from litellm.integrations.prometheus import ( + PrometheusLogger, + UserAPIKeyLabelValues, + prometheus_label_factory, + ) + + # Create a mock PrometheusLogger + prometheus_logger = MagicMock() + prometheus_logger.get_labels_for_metric = PrometheusLogger.get_labels_for_metric.__get__(prometheus_logger) + + # Test with a normalized route + enum_values = UserAPIKeyLabelValues( + route="/v1/responses/{response_id}", # Normalized route + status_code="200", + requested_model="gpt-4", + ) + + labels = prometheus_label_factory( + supported_enum_labels=prometheus_logger.get_labels_for_metric( + metric_name="litellm_proxy_total_requests_metric" + ), + enum_values=enum_values, + ) + + # Verify the route is normalized in labels + assert labels["route"] == "/v1/responses/{response_id}", \ + f"Expected normalized route in labels, got: {labels.get('route')}" + + print("✅ Prometheus metrics use normalized routes in labels") + + if __name__ == "__main__": test_user_email_in_required_metrics() test_user_email_label_exists() test_prometheus_metric_labels_structure() - print("All prometheus label tests passed!") \ No newline at end of file + test_route_normalization_for_responses_api() + test_route_normalization_for_sub_routes() + test_route_normalization_preserves_static_routes() + test_route_normalization_other_dynamic_apis() + test_prometheus_metrics_use_normalized_routes() + print("\n✅ All prometheus label tests passed!") \ No newline at end of file From cebcad48d32188744e8e2c21a5fa3122e23495ce Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 20 Jan 2026 15:53:02 +0530 Subject: [PATCH 23/58] Add gemini-2.5-computer-use-preview-10-2025 model for vertex ai provider --- ...odel_prices_and_context_window_backup.json | 25 +++++++++++++++++++ model_prices_and_context_window.json | 25 +++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4599cafe708..6b20965cd9e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -13465,6 +13465,31 @@ "supports_vision": true, "supports_web_search": true }, + "gemini-2.5-computer-use-preview-10-2025": { + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "litellm_provider": "vertex_ai-language-models", + "max_images_per_prompt": 3000, + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_above_200k_tokens": 1.5e-05, + "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/computer-use", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "gemini-embedding-001": { "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-embedding-models", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4599cafe708..6b20965cd9e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -13465,6 +13465,31 @@ "supports_vision": true, "supports_web_search": true }, + "gemini-2.5-computer-use-preview-10-2025": { + "input_cost_per_token": 1.25e-06, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "litellm_provider": "vertex_ai-language-models", + "max_images_per_prompt": 3000, + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_above_200k_tokens": 1.5e-05, + "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/computer-use", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "gemini-embedding-001": { "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-embedding-models", From be0f61854f08e3c7c683669a3ddf448ac805d790 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 20 Jan 2026 16:23:47 +0530 Subject: [PATCH 24/58] Add input_cost_per_video_per_second in ModelInfoBase --- litellm/litellm_core_utils/llm_cost_calc/utils.py | 10 +++++----- .../vertex_ai/multimodal_embeddings/transformation.py | 2 +- litellm/utils.py | 7 +++++++ 3 files changed, 13 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 65e77f014a3..785976ed319 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -354,7 +354,7 @@ class PromptTokensDetailsResult(TypedDict): image_tokens: int character_count: int image_count: int - video_length_seconds: int + video_length_seconds: float def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: @@ -400,10 +400,10 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: ) video_length_seconds = ( cast( - Optional[int], + Optional[float], getattr(usage.prompt_tokens_details, "video_length_seconds", 0), ) - or 0 + or 0.0 ) return PromptTokensDetailsResult( @@ -415,7 +415,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: image_tokens=image_tokens, character_count=character_count, image_count=image_count, - video_length_seconds=video_length_seconds, + video_length_seconds=float(video_length_seconds), ) @@ -561,7 +561,7 @@ def generic_cost_per_token( # noqa: PLR0915 image_tokens=0, character_count=0, image_count=0, - video_length_seconds=0, + video_length_seconds=0.0, ) if usage.prompt_tokens_details: prompt_tokens_details = _parse_prompt_tokens_details(usage) diff --git a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py index 2cb2ac9ed8f..d82c2bebb7f 100644 --- a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py @@ -265,7 +265,7 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): image_count += 1 ## Calculate video embeddings usage - video_length_seconds = 0 + video_length_seconds = 0.0 for prediction in vertex_predictions["predictions"]: video_embeddings = prediction.get("videoEmbeddings") if video_embeddings: diff --git a/litellm/utils.py b/litellm/utils.py index ac194e4f339..70cf9f02659 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5515,6 +5515,13 @@ def _get_model_info_helper( # noqa: PLR0915 input_cost_per_image_token=_model_info.get( "input_cost_per_image_token", None ), + input_cost_per_image=_model_info.get("input_cost_per_image", None), + input_cost_per_audio_per_second=_model_info.get( + "input_cost_per_audio_per_second", None + ), + input_cost_per_video_per_second=_model_info.get( + "input_cost_per_video_per_second", None + ), input_cost_per_token_batches=_model_info.get( "input_cost_per_token_batches" ), From ae414ed4627cf32a3cdfb15ef0eb42779b219ede Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 20 Jan 2026 17:07:00 +0530 Subject: [PATCH 25/58] =?UTF-8?q?Revert=20"feat:=20add=20retry=5Fdelay,=20?= =?UTF-8?q?exponential=5Fbackoff,=20and=20jitter=20to=20completion(?= =?UTF-8?q?=E2=80=A6"?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 1678f621db8305c7e5f9366f5ee46c7f4edb9a47. --- docs/my-website/docs/completion/input.md | 6 - .../docs/completion/reliable_completions.md | 30 --- .../litellm_core_utils/get_litellm_params.py | 6 - litellm/main.py | 181 +++--------------- litellm/types/utils.py | 8 +- 5 files changed, 27 insertions(+), 204 deletions(-) diff --git a/docs/my-website/docs/completion/input.md b/docs/my-website/docs/completion/input.md index 1272c7eb161..2f6da4bedcd 100644 --- a/docs/my-website/docs/completion/input.md +++ b/docs/my-website/docs/completion/input.md @@ -264,12 +264,6 @@ messages=[{"role": "user", "content": [ - `num_retries`: *int (optional)* - The number of times to retry the API call if an APIError, TimeoutError or ServiceUnavailableError occurs -- `retry_delay`: *float (optional)* - Time in seconds to wait between retries. - -- `exponential_backoff`: *float (optional)* - If true, wait time doubles after each failure (1s, 2s, 4s...) - -- `jitter`: *float (optional)* - If true, adds randomness to the wait time to prevent thundering herd. - - `context_window_fallback_dict`: *dict (optional)* - A mapping of model to use if call fails due to context window error - `fallbacks`: *list (optional)* - A list of model names + params to be used, in case the initial call fails diff --git a/docs/my-website/docs/completion/reliable_completions.md b/docs/my-website/docs/completion/reliable_completions.md index 431df66c9dc..f38917fe53d 100644 --- a/docs/my-website/docs/completion/reliable_completions.md +++ b/docs/my-website/docs/completion/reliable_completions.md @@ -31,36 +31,6 @@ response = completion( ) ``` -### Configurable Retries (Exponential Backoff, Jitter) - -You can also specify the `retry_delay` and `exponential_backoff` or `jitter` to the completion call. - -* `retry_delay`: (float) Time in seconds to wait between retries. -* `exponential_backoff`: (bool) If true, wait time doubles after each failure (1s, 2s, 4s...). -* `jitter`: (bool) If true, adds randomness to the wait time to prevent thundering herd. - -```python -import litellm - -# Exponential Backoff -response = litellm.completion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hi"}], - num_retries=3, - retry_delay=1.0, # Start with 1s wait - exponential_backoff=True # Wait 1s, 2s, 4s... -) - -# Jitter -response = litellm.completion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hi"}], - num_retries=3, - retry_delay=1.0, - jitter=True # Randomize wait times -) -``` - ## Fallbacks (SDK) :::info diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index e9740ed1da6..0d35cfa3140 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -64,9 +64,6 @@ def get_litellm_params( api_version: Optional[str] = None, max_retries: Optional[int] = None, litellm_request_debug: Optional[bool] = None, - retry_delay: Optional[float] = None, - exponential_backoff: Optional[bool] = None, - jitter: Optional[bool] = None, **kwargs, ) -> dict: litellm_params = { @@ -119,9 +116,6 @@ def get_litellm_params( "azure_password": kwargs.get("azure_password"), "azure_scope": kwargs.get("azure_scope"), "max_retries": max_retries, - "retry_delay": retry_delay, - "exponential_backoff": exponential_backoff, - "jitter": jitter, "timeout": kwargs.get("timeout"), "bucket_name": kwargs.get("bucket_name"), "vertex_credentials": kwargs.get("vertex_credentials"), diff --git a/litellm/main.py b/litellm/main.py index e62484e40f6..e199aa3001a 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -415,11 +415,6 @@ async def acompletion( web_search_options: Optional[OpenAIWebSearchOptions] = None, # Session management shared_session: Optional["ClientSession"] = None, - # Retry params - retry_delay: Optional[float] = None, - exponential_backoff: Optional[bool] = None, - jitter: Optional[bool] = None, - num_retries: Optional[int] = None, **kwargs, ) -> Union[ModelResponse, CustomStreamWrapper]: """ @@ -465,16 +460,6 @@ async def acompletion( - The `completion` function is called using `run_in_executor` to execute synchronously in the event loop. - If `stream` is True, the function returns an async generator that yields completion lines. """ - if _update_kwargs_with_retry_params( - retry_delay, - exponential_backoff, - jitter, - num_retries, - locals().get("max_retries"), - kwargs, - ): - return await acompletion_with_retries(model=model, messages=messages, **kwargs) - fallbacks = kwargs.get("fallbacks", None) mock_timeout = kwargs.get("mock_timeout", None) @@ -1007,32 +992,8 @@ def _drop_input_examples_from_tools( return cleaned_tools -def _update_kwargs_with_retry_params( - retry_delay: Optional[float], - exponential_backoff: Optional[bool], - jitter: Optional[bool], - num_retries: Optional[int], - max_retries: Optional[int], - kwargs: dict, -) -> bool: - """ - Updates kwargs with retry parameters if any are provided. - Returns True if retry logic should be triggered, False otherwise. - """ - if retry_delay is not None or exponential_backoff is not None or jitter is not None: - kwargs["retry_delay"] = retry_delay - kwargs["exponential_backoff"] = exponential_backoff - kwargs["jitter"] = jitter - - if num_retries is not None: - kwargs["num_retries"] = num_retries - elif max_retries is not None and "num_retries" not in kwargs: - kwargs["num_retries"] = max_retries - - return True - return False - - +@tracer.wrap() +@client def completion( # type: ignore # noqa: PLR0915 model: str, # Optional OpenAI params: see https://platform.openai.com/docs/api-reference/chat/create @@ -1082,11 +1043,6 @@ def completion( # type: ignore # noqa: PLR0915 thinking: Optional[AnthropicThinkingParam] = None, # Session management shared_session: Optional["ClientSession"] = None, - # Retry params - retry_delay: Optional[float] = None, - exponential_backoff: Optional[bool] = None, - jitter: Optional[bool] = None, - num_retries: Optional[int] = None, **kwargs, ) -> Union[ModelResponse, CustomStreamWrapper]: """ @@ -1134,20 +1090,6 @@ def completion( # type: ignore # noqa: PLR0915 - It supports various optional parameters for customizing the completion behavior. - If 'mock_response' is provided, a mock completion response is returned for testing or debugging. """ - if _update_kwargs_with_retry_params( - retry_delay, - exponential_backoff, - jitter, - num_retries, - locals().get("max_retries"), - kwargs, - ): - # check if this is an async call (acompletion=True in kwargs) - if kwargs.get("acompletion", False) is True: - return acompletion_with_retries(model=model, messages=messages, **kwargs) - - return completion_with_retries(model=model, messages=messages, **kwargs) - ### VALIDATE Request ### if model is None: raise ValueError("model param not passed in.") @@ -1173,9 +1115,7 @@ def completion( # type: ignore # noqa: PLR0915 # Check if MCP tools are present (following responses pattern) # Cast tools to Optional[Iterable[ToolParam]] for type checking tools_for_mcp = cast(Optional[Iterable[ToolParam]], tools) - if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway( - tools=tools_for_mcp - ): + if LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools_for_mcp): # Return coroutine - acompletion will await it # completion() can return a coroutine when MCP tools are present, which acompletion() awaits return acompletion_with_mcp( # type: ignore[return-value] @@ -2405,7 +2345,11 @@ def completion( # type: ignore # noqa: PLR0915 input=messages, api_key=api_key, original_response=response ) elif custom_llm_provider == "minimax": - api_key = api_key or get_secret_str("MINIMAX_API_KEY") or litellm.api_key + api_key = ( + api_key + or get_secret_str("MINIMAX_API_KEY") + or litellm.api_key + ) api_base = ( api_base @@ -2453,9 +2397,7 @@ def completion( # type: ignore # noqa: PLR0915 or custom_llm_provider == "wandb" or custom_llm_provider == "clarifai" or custom_llm_provider in litellm.openai_compatible_providers - or JSONProviderRegistry.exists( - custom_llm_provider - ) # JSON-configured providers + or JSONProviderRegistry.exists(custom_llm_provider) # JSON-configured providers or "ft:gpt-3.5-turbo" in model # finetune gpt-3.5-turbo ): # allow user to make an openai call with a custom base # note: if a user sets a custom base - we should ensure this works @@ -4295,51 +4237,24 @@ def completion_with_retries(*args, **kwargs): retry_strategy: Literal["exponential_backoff_retry", "constant_retry"] = kwargs.pop( "retry_strategy", "constant_retry" ) # type: ignore - retry_delay = kwargs.pop("retry_delay", None) - exponential_backoff = kwargs.pop("exponential_backoff", False) - jitter = kwargs.pop("jitter", False) - original_function = kwargs.pop("original_function", completion) - - # +1 because stop_after_attempt includes the initial attempt - stop_after = tenacity.stop_after_attempt(num_retries + 1) - - if retry_strategy == "exponential_backoff_retry" or exponential_backoff: - # Defaults for exponential backoff - multiplier = 1 - min_wait = 0 - if retry_delay is not None: - multiplier = retry_delay - min_wait = retry_delay - - wait_args = {"multiplier": multiplier, "min": min_wait, "max": 10} - - if jitter: - wait_strategy = tenacity.wait_random_exponential(**wait_args) - else: - wait_strategy = tenacity.wait_exponential(**wait_args) - + if retry_strategy == "exponential_backoff_retry": retryer = tenacity.Retrying( - wait=wait_strategy, - stop=stop_after, + wait=tenacity.wait_exponential(multiplier=1, max=10), + stop=tenacity.stop_after_attempt(num_retries), reraise=True, ) else: - wait_strategy = tenacity.wait_none() - if retry_delay: - wait_strategy = tenacity.wait_fixed(retry_delay) - retryer = tenacity.Retrying( - wait=wait_strategy, - stop=stop_after, - reraise=True, + stop=tenacity.stop_after_attempt(num_retries), reraise=True ) return retryer(original_function, *args, **kwargs) async def acompletion_with_retries(*args, **kwargs): """ - Executes a litellm.completion() with retries. + [DEPRECATED]. Use 'acompletion' or router.acompletion instead! + Executes a litellm.completion() with 3 retries """ try: import tenacity @@ -4352,65 +4267,18 @@ async def acompletion_with_retries(*args, **kwargs): kwargs["max_retries"] = 0 kwargs["num_retries"] = 0 retry_strategy = kwargs.pop("retry_strategy", "constant_retry") - retry_delay = kwargs.pop("retry_delay", None) - exponential_backoff = kwargs.pop("exponential_backoff", False) - jitter = kwargs.pop("jitter", False) original_function = kwargs.pop("original_function", completion) - - # +1 because stop_after_attempt includes the initial attempt - stop_after = tenacity.stop_after_attempt(num_retries + 1) - - # If the original function is completion but we are doing async retries - # we need to ensure it's treated as an async function if it returns a coroutine - # or wraps it. - # However, since we are in acompletion_with_retries, we expect to be waiting. - # If original_function is completion(acompletion=True), it returns a coro. - # Tenacity AsyncRetrying expects the function to be awaitable or return awaitable? - # Actually AsyncRetrying works with async def functions. - # If original_function is sync (but returns coro), we might need a wrapper. - - async def _async_original_function(*args, **kwargs): - # Ensure we await the result if it is a coroutine - result = original_function(*args, **kwargs) - if asyncio.iscoroutine(result): - return await result - return result - - if retry_strategy == "exponential_backoff_retry" or exponential_backoff: - # Defaults for exponential backoff - multiplier = 1 - min_wait = 0 - if retry_delay is not None: - multiplier = retry_delay - min_wait = retry_delay - - wait_args = {"multiplier": multiplier, "min": min_wait, "max": 10} - - if jitter: - # Use wait_random_exponential if available or combine - # Using wait_exponential + wait_random is one way, or wait_random_exponential - # tenacity.wait_random_exponential(multiplier=1, max=10) - wait_strategy = tenacity.wait_random_exponential(**wait_args) - else: - wait_strategy = tenacity.wait_exponential(**wait_args) - + if retry_strategy == "exponential_backoff_retry": retryer = tenacity.AsyncRetrying( - wait=wait_strategy, - stop=stop_after, + wait=tenacity.wait_exponential(multiplier=1, max=10), + stop=tenacity.stop_after_attempt(num_retries), reraise=True, ) else: - wait_strategy = tenacity.wait_none() - if retry_delay: - wait_strategy = tenacity.wait_fixed(retry_delay) - retryer = tenacity.AsyncRetrying( - wait=wait_strategy, - stop=stop_after, - reraise=True, + stop=tenacity.stop_after_attempt(num_retries), reraise=True ) - - return await retryer(_async_original_function, *args, **kwargs) + return await retryer(original_function, *args, **kwargs) def responses_with_retries(*args, **kwargs): @@ -4848,7 +4716,7 @@ def embedding( # noqa: PLR0915 if headers is not None and headers != {}: optional_params["extra_headers"] = headers - + if encoding_format is not None: optional_params["encoding_format"] = encoding_format else: @@ -6911,7 +6779,9 @@ def speech( # noqa: PLR0915 if text_to_speech_provider_config is None: text_to_speech_provider_config = MinimaxTextToSpeechConfig() - minimax_config = cast(MinimaxTextToSpeechConfig, text_to_speech_provider_config) + minimax_config = cast( + MinimaxTextToSpeechConfig, text_to_speech_provider_config + ) if api_base is not None: litellm_params_dict["api_base"] = api_base @@ -7051,7 +6921,7 @@ async def ahealth_check( custom_llm_provider_from_params = model_params.get("custom_llm_provider", None) api_base_from_params = model_params.get("api_base", None) api_key_from_params = model_params.get("api_key", None) - + model, custom_llm_provider, _, _ = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider_from_params, @@ -7429,7 +7299,6 @@ def __getattr__(name: str) -> Any: _encoding = tiktoken.get_encoding("cl100k_base") # Cache it in the module's __dict__ for subsequent accesses import sys - sys.modules[__name__].__dict__["encoding"] = _encoding global _encoding_cache _encoding_cache = _encoding diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e71fcac3897..f5e217d8b46 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -125,14 +125,13 @@ class SearchContextCostPerQuery(TypedDict, total=False): class AgenticLoopParams(TypedDict, total=False): """ Parameters passed to agentic loop hooks (e.g., WebSearch interception). - + Stored in logging_obj.model_call_details["agentic_loop_params"] to provide agentic hooks with the original request context needed for follow-up calls. """ - model: str """The model string with provider prefix (e.g., 'bedrock/invoke/...')""" - + custom_llm_provider: str """The LLM provider name (e.g., 'bedrock', 'anthropic')""" @@ -2925,9 +2924,6 @@ all_litellm_params = ( "shared_session", "search_tool_name", "order", - "retry_delay", - "exponential_backoff", - "jitter", ] + list(StandardCallbackDynamicParams.__annotations__.keys()) + list(CustomPricingLiteLLMParams.model_fields.keys()) From f6fcd0cb8591a1308d5b0ed4dae3702a49e34d35 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 20 Jan 2026 17:50:45 +0530 Subject: [PATCH 26/58] Revert "fixed litellm params (#19315)" This reverts commit 16b8ed6786280e0c9c70565cb96aeb1590a4e1c2. --- litellm/main.py | 54 +++---------------------------------------------- 1 file changed, 3 insertions(+), 51 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index e199aa3001a..ae27b4145b3 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -565,10 +565,6 @@ async def acompletion( model=model, custom_llm_provider=custom_llm_provider, api_base=completion_kwargs.get("base_url", None), - litellm_params=litellm.types.router.LiteLLM_Params( - model=model, - use_litellm_proxy=kwargs.get("use_litellm_proxy", False), - ), ) fallbacks = fallbacks or litellm.model_fallbacks @@ -1295,10 +1291,6 @@ def completion( # type: ignore # noqa: PLR0915 custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key, - litellm_params=litellm.types.router.LiteLLM_Params( - model=model, - use_litellm_proxy=kwargs.get("use_litellm_proxy", False), - ), ) if not _should_allow_input_examples( @@ -4376,10 +4368,6 @@ async def aembedding(*args, **kwargs) -> EmbeddingResponse: model=model, custom_llm_provider=custom_llm_provider, api_base=kwargs.get("api_base", None), - litellm_params=litellm.types.router.LiteLLM_Params( - model=model, - use_litellm_proxy=kwargs.get("use_litellm_proxy", False), - ), ) # Await normally @@ -4565,10 +4553,6 @@ def embedding( # noqa: PLR0915 custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key, - litellm_params=litellm.types.router.LiteLLM_Params( - model=model, - use_litellm_proxy=kwargs.get("use_litellm_proxy", False), - ), ) if dynamic_api_key is not None: @@ -5704,10 +5688,6 @@ def text_completion( # noqa: PLR0915 model=model, # type: ignore custom_llm_provider=custom_llm_provider, api_base=api_base, - litellm_params=litellm.types.router.LiteLLM_Params( - model=model, # type: ignore - use_litellm_proxy=kwargs.get("use_litellm_proxy", False), - ), ) if custom_llm_provider == "huggingface": @@ -5998,10 +5978,6 @@ async def amoderation( custom_llm_provider=custom_llm_provider, api_base=optional_params.api_base, api_key=optional_params.api_key, - litellm_params=litellm.types.router.LiteLLM_Params( - model=model or "", - use_litellm_proxy=kwargs.get("use_litellm_proxy", False), - ), ) except litellm.BadRequestError: # `model` is optional field for moderation - get_llm_provider will throw BadRequestError if model is not set / not recognized @@ -6067,12 +6043,7 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse: func_with_context = partial(ctx.run, func) _, custom_llm_provider, _, _ = get_llm_provider( - model=model, - api_base=kwargs.get("api_base", None), - litellm_params=litellm.types.router.LiteLLM_Params( - model=model, - use_litellm_proxy=kwargs.get("use_litellm_proxy", False), - ), + model=model, api_base=kwargs.get("api_base", None) ) # Await normally @@ -6176,10 +6147,6 @@ def transcription( custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key, - litellm_params=litellm.types.router.LiteLLM_Params( - model=model, - use_litellm_proxy=kwargs.get("use_litellm_proxy", False), - ), ) # type: ignore if dynamic_api_key is not None: @@ -6355,12 +6322,7 @@ async def aspeech(*args, **kwargs) -> HttpxBinaryResponseContent: func_with_context = partial(ctx.run, func) _, custom_llm_provider, _, _ = get_llm_provider( - model=model, - api_base=kwargs.get("api_base", None), - litellm_params=litellm.types.router.LiteLLM_Params( - model=model, - use_litellm_proxy=kwargs.get("use_litellm_proxy", False), - ), + model=model, api_base=kwargs.get("api_base", None) ) # Await normally @@ -6411,13 +6373,7 @@ def speech( # noqa: PLR0915 model_info = kwargs.get("model_info", None) shared_session = kwargs.get("shared_session", None) model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( - model=model, - custom_llm_provider=custom_llm_provider, - api_base=api_base, - litellm_params=litellm.types.router.LiteLLM_Params( - model=model, - use_litellm_proxy=kwargs.get("use_litellm_proxy", False), - ), + model=model, custom_llm_provider=custom_llm_provider, api_base=api_base ) # type: ignore kwargs.pop("tags", []) @@ -6927,10 +6883,6 @@ async def ahealth_check( custom_llm_provider=custom_llm_provider_from_params, api_base=api_base_from_params, api_key=api_key_from_params, - litellm_params=litellm.types.router.LiteLLM_Params( - model=model, - use_litellm_proxy=model_params.get("use_litellm_proxy", False), - ), ) if model in litellm.model_cost and mode is None: mode = litellm.model_cost[model].get("mode") From 2153db5e64f2f77692342dafe9880582081f2ef3 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 20 Jan 2026 18:27:36 +0530 Subject: [PATCH 27/58] fix: test_convert_to_bedrock_format_post_call_streaming_hook --- .../test_bedrock_guardrails.py | 22 ++++++++++--------- .../vertex_ai/test_vertex_ai_common_utils.py | 2 +- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index 1795ff23605..cb594c221c4 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -1139,17 +1139,16 @@ async def test_convert_to_bedrock_format_post_call_streaming_hook(): result_chunks.append(chunk) # Verify bedrock API calls were made + # Note: When event_hook is None (default), the guardrail is considered enabled for all hooks. + # In post_call, INPUT validation is skipped if pre_call/during_call is already enabled + # to avoid redundant validation. Since event_hook=None means all hooks are enabled, + # only OUTPUT validation should be performed in post_call. assert ( - len(bedrock_calls) == 2 - ), f"Expected 2 bedrock calls (INPUT and OUTPUT), got {len(bedrock_calls)}" + len(bedrock_calls) == 1 + ), f"Expected 1 bedrock call (OUTPUT only), got {len(bedrock_calls)}" - # Find the OUTPUT call - output_calls = [call for call in bedrock_calls if call["source"] == "OUTPUT"] - assert ( - len(output_calls) == 1 - ), f"Expected 1 OUTPUT call, got {len(output_calls)}" - - output_call = output_calls[0] + # Verify the OUTPUT call + output_call = bedrock_calls[0] assert output_call["source"] == "OUTPUT" assert output_call["response"] is not None assert output_call["messages"] is None # OUTPUT calls don't need messages @@ -1176,7 +1175,10 @@ async def test_convert_to_bedrock_format_post_call_streaming_hook(): print( "✅ Post-call streaming hook test passed - OUTPUT source used for masking" ) - print(f"✅ Bedrock calls made: {[call['source'] for call in bedrock_calls]}") + print( + f"✅ Bedrock calls made: {[call['source'] for call in bedrock_calls]} " + "(INPUT validation skipped due to event_hook=None implying pre_call/during_call enabled)" + ) print(f"✅ Final masked content: {full_content}") diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index 8fdaf4ea2df..953f54e4a54 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -440,7 +440,7 @@ def test_vertex_ai_complex_response_schema(): optional_params = {} v.apply_response_schema_transformation( - value=non_default_params["response_format"], optional_params=optional_params + value=non_default_params["response_format"], optional_params=optional_params, model="gemini-1.5-pro-preview-0409" ) # Assertions for the transformed schema From dc3ee6335955aa082ef190266dc21176fe3bb863 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 20 Jan 2026 18:37:56 +0530 Subject: [PATCH 28/58] fix: test_env_keys --- docs/my-website/docs/proxy/config_settings.md | 8 +++++++ tests/documentation_tests/test_env_keys.py | 22 ++++++++++++++++++- 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 8b4514f5688..53b0f3eea71 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -449,6 +449,13 @@ router_settings: | BRAINTRUST_API_KEY | API key for Braintrust integration | BRAINTRUST_API_BASE | Base URL for Braintrust API. Default is https://api.braintrustdata.com/v1 | CACHED_STREAMING_CHUNK_DELAY | Delay in seconds for cached streaming chunks. Default is 0.02 +| CHATGPT_API_BASE | Base URL for ChatGPT API. Default is https://chatgpt.com/backend-api/codex +| CHATGPT_AUTH_FILE | Filename for ChatGPT authentication data. Default is "auth.json" +| CHATGPT_DEFAULT_INSTRUCTIONS | Default system instructions for ChatGPT provider +| CHATGPT_ORIGINATOR | Originator identifier for ChatGPT API requests. Default is "codex_cli_rs" +| CHATGPT_TOKEN_DIR | Directory to store ChatGPT authentication tokens. Default is "~/.config/litellm/chatgpt" +| CHATGPT_USER_AGENT | Custom user agent string for ChatGPT API requests +| CHATGPT_USER_AGENT_SUFFIX | Suffix to append to the ChatGPT user agent string | CIRCLE_OIDC_TOKEN | OpenID Connect token for CircleCI | CIRCLE_OIDC_TOKEN_V2 | Version 2 of the OpenID Connect token for CircleCI | CLOUDZERO_API_KEY | CloudZero API key for authentication @@ -794,6 +801,7 @@ router_settings: | OPENAI_BASE_URL | Base URL for OpenAI API | OPENAI_API_BASE | Base URL for OpenAI API. Default is https://api.openai.com/ | OPENAI_API_KEY | API key for OpenAI services +| OPENAI_CHATGPT_API_BASE | Alternative to CHATGPT_API_BASE. Base URL for ChatGPT API | OPENAI_FILE_SEARCH_COST_PER_1K_CALLS | Cost per 1000 calls for OpenAI file search. Default is 0.0025 | OPENAI_ORGANIZATION | Organization identifier for OpenAI | OPENID_BASE_URL | Base URL for OpenID Connect services diff --git a/tests/documentation_tests/test_env_keys.py b/tests/documentation_tests/test_env_keys.py index 6b7c15e2b9d..4fed84d0431 100644 --- a/tests/documentation_tests/test_env_keys.py +++ b/tests/documentation_tests/test_env_keys.py @@ -16,6 +16,25 @@ get_secret_str_pattern = re.compile( # Set to store unique keys from the code env_keys = set() +# Terminal/environment detection variables that should not be documented +# These are internal variables used for terminal detection, not user-configurable settings +EXCLUDED_TERMINAL_VARS = { + "TERM", + "TERM_PROGRAM", + "TERM_PROGRAM_VERSION", + "TERM_SESSION_ID", + "VTE_VERSION", + "KITTY_WINDOW_ID", + "KONSOLE_VERSION", + "ITERM_PROFILE", + "ITERM_PROFILE_NAME", + "ITERM_SESSION_ID", + "WEZTERM_VERSION", + "WT_SESSION", + "GNOME_TERMINAL_SCREEN", + "ALACRITTY_SOCKET", +} + # Walk through all files in the litellm repo to find references of os.getenv() and litellm.get_secret() for root, dirs, files in os.walk(repo_base): for file in files: @@ -28,7 +47,8 @@ for root, dirs, files in os.walk(repo_base): getenv_matches = getenv_pattern.findall(content) env_keys.update( match for match in getenv_matches - ) # Extract only the key part + if match not in EXCLUDED_TERMINAL_VARS + ) # Extract only the key part, excluding terminal vars # Find all keys using litellm.get_secret() get_secret_matches = get_secret_pattern.findall(content) From 3cc19c56bab2fa42b93ff39da42624ba0585a127 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 20 Jan 2026 18:46:26 +0530 Subject: [PATCH 29/58] Revert "feat: Add Redis-based migration lock with bug fixes (#19261)" This reverts commit 98e87c3e67ab36948020f769da6eea87fd3ae2c3. --- .../litellm_proxy_extras/utils.py | 250 ++---------------- litellm/proxy/db/prisma_client.py | 13 +- .../test_litellm_proxy_extras_utils.py | 214 +-------------- 3 files changed, 25 insertions(+), 452 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 4e95267d8ab..7ffbe95be13 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -11,11 +11,6 @@ from typing import Optional from litellm_proxy_extras._logging import logger -try: - from litellm.caching.redis_cache import RedisCache -except ImportError: - RedisCache = None # type: ignore - def str_to_bool(value: Optional[str]) -> bool: if value is None: @@ -23,154 +18,6 @@ def str_to_bool(value: Optional[str]) -> bool: return value.lower() in ("true", "1", "t", "y", "yes") -class MigrationLockManager: - """Redis-based lock manager for database migrations. - - Prevents concurrent Prisma migrations in multi-pod deployments by using - a distributed lock. Only one pod can hold the lock and run migrations at a time. - """ - - MIGRATION_LOCK_KEY = "migration_lock" - LOCK_TTL_SECONDS = 300 # 5 minutes TTL - - def __init__(self, redis_cache: Optional["RedisCache"] = None): - """Initialize the migration lock manager. - - Args: - redis_cache: Optional RedisCache instance for distributed locking. - If None, migrations run without lock protection (single instance mode). - """ - self.redis_cache = redis_cache - self.lock_acquired = False - self.pod_id = f"pod_{os.getpid()}_{int(time.time())}" - - def _get_redis_lock_key(self) -> str: - """Get Redis lock key for migration.""" - return f"migration_lock:{self.MIGRATION_LOCK_KEY}" - - def acquire_lock(self) -> bool: - """Acquire migration lock using Redis SET NX. - - Returns: - bool: True if lock acquired, False otherwise. - """ - if self.redis_cache is None: - logger.warning( - "Redis cache is not available, running migration without lock protection" - ) - self.lock_acquired = True - return True - - try: - lock_key = self._get_redis_lock_key() - - # FIX: Use native Redis SET NX instead of RedisCache.set_cache() - # Original bug: set_cache() doesn't support nx parameter and returns None - # Fixed: Use redis_client.set() directly which returns True/False - acquired = self.redis_cache.redis_client.set( - name=lock_key, - value=self.pod_id, - nx=True, # Only set if key doesn't exist - ex=self.LOCK_TTL_SECONDS, # Set expiration time - ) - - if acquired: - self.lock_acquired = True - logger.info(f"Migration lock acquired by pod {self.pod_id}") - return True - else: - logger.info("Migration lock is already held by another pod") - return False - - except Exception as e: - logger.warning(f"Failed to acquire migration lock: {e}") - return False - - def wait_for_lock_release( - self, check_interval: int = 5, max_wait: int = 300 - ) -> bool: - """Wait for another process to release the lock. - - Args: - check_interval: Seconds to wait between lock acquisition attempts. - max_wait: Maximum seconds to wait for lock release. - - Returns: - bool: True if lock acquired after waiting, False if timeout. - """ - if self.redis_cache is None: - logger.warning("Redis cache is not available, cannot wait for lock") - return False - - logger.info(f"Waiting for migration lock to be released (max {max_wait}s)...") - start_time = time.time() - - while time.time() - start_time < max_wait: - # Try to acquire lock using the public acquire_lock method - if self.acquire_lock(): - logger.info( - f"Migration lock acquired after waiting by pod {self.pod_id}" - ) - return True - - time.sleep(check_interval) - - logger.warning(f"Failed to acquire migration lock within {max_wait} seconds") - return False - - def release_lock(self): - """Release migration lock atomically using Lua script. - - FIX: Use Lua script for atomic compare-and-delete to prevent race conditions. - Original bug: Non-atomic GET then DELETE allows another pod to acquire lock - between the GET and DELETE operations. - """ - if not self.lock_acquired or self.redis_cache is None: - return - - try: - lock_key = self._get_redis_lock_key() - - # FIX: Use Lua script for atomic compare-and-delete - # This prevents race condition where: - # 1. Pod A reads lock value (sees its own pod_id) - # 2. Lock TTL expires - # 3. Pod B acquires lock - # 4. Pod A deletes lock (deletes Pod B's lock!) - lua_script = """ - if redis.call("GET", KEYS[1]) == ARGV[1] then - return redis.call("DEL", KEYS[1]) - else - return 0 - end - """ - - result = self.redis_cache.redis_client.eval( - lua_script, - 1, # Number of keys - lock_key, # KEYS[1] - self.pod_id, # ARGV[1] - ) - - if result == 1: - logger.info(f"Migration lock released by pod {self.pod_id}") - else: - logger.warning(f"Pod {self.pod_id} cannot release lock (not owner)") - - except Exception as e: - logger.warning(f"Failed to release migration lock: {e}") - finally: - self.lock_acquired = False - - def __enter__(self): - """Context manager entry - acquire lock when entering with statement.""" - self.acquire_lock() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - """Context manager exit - release lock when exiting with statement.""" - self.release_lock() - def _get_prisma_env() -> dict: """Get environment variables for Prisma, handling offline mode if configured.""" @@ -178,9 +25,7 @@ def _get_prisma_env() -> dict: if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")): # These env vars prevent Prisma from attempting downloads prisma_env["NPM_CONFIG_PREFER_OFFLINE"] = "true" - prisma_env["NPM_CONFIG_CACHE"] = os.getenv( - "NPM_CONFIG_CACHE", "/app/.cache/npm" - ) + prisma_env["NPM_CONFIG_CACHE"] = os.getenv("NPM_CONFIG_CACHE", "/app/.cache/npm") return prisma_env @@ -189,28 +34,29 @@ def _get_prisma_command() -> str: if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")): # Primary location where Prisma Python package installs the CLI default_cli_path = "/app/.cache/prisma-python/binaries/node_modules/.bin/prisma" - + # Check if custom path is provided (for flexibility) custom_cli_path = os.getenv("PRISMA_CLI_PATH") if custom_cli_path and os.path.exists(custom_cli_path): logger.info(f"Using custom Prisma CLI at {custom_cli_path}") return custom_cli_path - + # Check the default location if os.path.exists(default_cli_path): logger.info(f"Using cached Prisma CLI at {default_cli_path}") return default_cli_path - + # If not found, log warning and fall back logger.warning( f"Prisma CLI not found at {default_cli_path}. " "Falling back to Python wrapper (may attempt downloads)" ) - + # Fall back to the Python wrapper (will work in online mode) return "prisma" + class ProxyExtrasDBManager: @staticmethod def _get_prisma_dir() -> str: @@ -273,7 +119,7 @@ class ProxyExtrasDBManager: stdout=open(migration_file, "w"), check=True, timeout=30, - env=prisma_env, + env=prisma_env ) # 3. Mark the migration as applied since it represents current state @@ -288,7 +134,7 @@ class ProxyExtrasDBManager: ], check=True, timeout=30, - env=prisma_env, + env=prisma_env ) return True @@ -313,20 +159,14 @@ class ProxyExtrasDBManager: @staticmethod def _roll_back_migration(migration_name: str): """Mark a specific migration as rolled back""" - # Set up environment for offline mode if configured + # Set up environment for offline mode if configured prisma_env = _get_prisma_env() subprocess.run( - [ - _get_prisma_command(), - "migrate", - "resolve", - "--rolled-back", - migration_name, - ], + [_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name], timeout=60, check=True, capture_output=True, - env=prisma_env, + env=prisma_env ) @staticmethod @@ -338,7 +178,7 @@ class ProxyExtrasDBManager: timeout=60, check=True, capture_output=True, - env=prisma_env, + env=prisma_env ) @staticmethod @@ -408,7 +248,7 @@ class ProxyExtrasDBManager: if not database_url: logger.error("DATABASE_URL not set") return - + diff_dir = ( Path(migrations_dir) / "migrations" @@ -443,7 +283,7 @@ class ProxyExtrasDBManager: check=True, timeout=60, stdout=f, - env=_get_prisma_env(), + env=_get_prisma_env() ) except subprocess.CalledProcessError as e: logger.warning(f"Failed to generate migration diff: {e.stderr}") @@ -473,7 +313,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env(), + env=_get_prisma_env() ) logger.info(f"prisma db execute stdout: {result.stdout}") logger.info("✅ Migration diff applied successfully") @@ -491,18 +331,12 @@ class ProxyExtrasDBManager: try: logger.info(f"Resolving migration: {migration_name}") subprocess.run( - [ - _get_prisma_command(), - "migrate", - "resolve", - "--applied", - migration_name, - ], + [_get_prisma_command(), "migrate", "resolve", "--applied", migration_name], timeout=60, check=True, capture_output=True, text=True, - env=_get_prisma_env(), + env=_get_prisma_env() ) logger.debug(f"Resolved migration: {migration_name}") except subprocess.CalledProcessError as e: @@ -512,57 +346,19 @@ class ProxyExtrasDBManager: ) @staticmethod - def setup_database( - use_migrate: bool = False, redis_cache: Optional["RedisCache"] = None - ) -> bool: + def setup_database(use_migrate: bool = False) -> bool: """ - Set up the database using either prisma migrate or prisma db push. - Uses migrations from litellm-proxy-extras package. - In multi-instance environment, use redis lock to prevent concurrent execution. + Set up the database using either prisma migrate or prisma db push + Uses migrations from litellm-proxy-extras package Args: + schema_path (str): Path to the Prisma schema file use_migrate (bool): Whether to use prisma migrate instead of db push - redis_cache: Redis cache instance for distributed locking Returns: bool: True if setup was successful, False otherwise """ schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma" - - database_url = os.getenv("DATABASE_URL") - if not database_url: - logger.error("DATABASE_URL environment variable is not set") - return False - - # Use MigrationLockManager to prevent concurrent migration execution - with MigrationLockManager(redis_cache) as lock_manager: - # Lock is already acquired in __enter__, check if it was successful - if not lock_manager.lock_acquired: - # Cannot acquire lock, another process is running migration - logger.info( - "Another pod is running migration, waiting for completion..." - ) - - # Wait for other process to complete migration - if not lock_manager.wait_for_lock_release(): - logger.error("Failed to acquire migration lock after waiting") - return False - - # Successfully acquired lock, proceed with migration - logger.info("Acquired migration lock, proceeding with migration") - return ProxyExtrasDBManager._execute_migration(use_migrate, schema_path) - - @staticmethod - def _execute_migration(use_migrate: bool, schema_path: str) -> bool: - """Execute the actual migration. - - Args: - use_migrate: Whether to use prisma migrate instead of db push - schema_path: Path to the Prisma schema file - - Returns: - bool: True if migration was successful, False otherwise - """ for attempt in range(4): original_dir = os.getcwd() migrations_dir = ProxyExtrasDBManager._get_prisma_dir() @@ -579,7 +375,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env(), + env=_get_prisma_env() ) logger.info(f"prisma migrate deploy stdout: {result.stdout}") @@ -617,7 +413,7 @@ class ProxyExtrasDBManager: check=True, capture_output=True, text=True, - env=_get_prisma_env(), + env=_get_prisma_env() ) logger.info( f"✅ Migration {failed_migration} marked as rolled back... retrying" diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 65fed92340b..c9c0cfe8f68 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -383,18 +383,7 @@ class PrismaManager: prisma_dir = PrismaManager._get_prisma_dir() - # Import redis_usage_cache for distributed locking - try: - from litellm.proxy.proxy_server import redis_usage_cache - except ImportError: - verbose_proxy_logger.warning( - "Could not import redis_usage_cache, migrations will run without distributed locking" - ) - redis_usage_cache = None - - return ProxyExtrasDBManager.setup_database( - use_migrate=use_migrate, redis_cache=redis_usage_cache - ) + return ProxyExtrasDBManager.setup_database(use_migrate=use_migrate) else: # Use prisma db push with increased timeout subprocess.run( diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index cff3a88ab6c..5714cd5c487 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -1,13 +1,11 @@ import os import sys -import time -from unittest.mock import Mock, patch sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path -from litellm_proxy_extras.utils import ProxyExtrasDBManager, MigrationLockManager +from litellm_proxy_extras.utils import ProxyExtrasDBManager def test_custom_prisma_dir(monkeypatch): @@ -127,213 +125,3 @@ class TestErrorClassificationPriority: error_message = "connection timeout" assert ProxyExtrasDBManager._is_permission_error(error_message) is False assert ProxyExtrasDBManager._is_idempotent_error(error_message) is False - - -class TestMigrationLockManager: - """Test cases for MigrationLockManager""" - - def test_acquire_lock_success(self): - """Test successful lock acquisition using Redis SET NX""" - mock_redis = Mock() - # FIX VERIFICATION: redis_client.set() returns True for successful SET NX - mock_redis.redis_client.set.return_value = True - - lock_manager = MigrationLockManager(mock_redis) - result = lock_manager.acquire_lock() - - assert result is True - assert lock_manager.lock_acquired is True - # Verify SET NX was called with correct parameters - mock_redis.redis_client.set.assert_called_once() - call_args = mock_redis.redis_client.set.call_args - assert call_args[1]["nx"] is True # SET NX parameter - assert call_args[1]["ex"] == 300 # TTL - - def test_acquire_lock_failure(self): - """Test lock acquisition failure when lock is already held""" - mock_redis = Mock() - # FIX VERIFICATION: redis_client.set() returns False when key exists - mock_redis.redis_client.set.return_value = False - - lock_manager = MigrationLockManager(mock_redis) - result = lock_manager.acquire_lock() - - assert result is False - assert lock_manager.lock_acquired is False - - def test_acquire_lock_without_redis(self): - """Test lock acquisition without Redis (graceful fallback)""" - lock_manager = MigrationLockManager(redis_cache=None) - result = lock_manager.acquire_lock() - - assert result is True - assert lock_manager.lock_acquired is True - - def test_release_lock_success(self): - """Test successful lock release using Lua script""" - mock_redis = Mock() - # FIX VERIFICATION: Lua script returns 1 for successful delete - mock_redis.redis_client.eval.return_value = 1 - - lock_manager = MigrationLockManager(mock_redis) - lock_manager.pod_id = "pod_123_456" - lock_manager.lock_acquired = True - - lock_manager.release_lock() - - # Verify Lua script was called for atomic compare-and-delete - mock_redis.redis_client.eval.assert_called_once() - call_args = mock_redis.redis_client.eval.call_args - # Verify Lua script contains atomic compare-and-delete logic - lua_script = call_args[0][0] - assert "GET" in lua_script - assert "DEL" in lua_script - assert lock_manager.lock_acquired is False - - def test_release_lock_wrong_owner(self): - """Test releasing lock when not the owner""" - mock_redis = Mock() - # FIX VERIFICATION: Lua script returns 0 when ownership check fails - mock_redis.redis_client.eval.return_value = 0 - - lock_manager = MigrationLockManager(mock_redis) - lock_manager.pod_id = "pod_123_456" - lock_manager.lock_acquired = True - - lock_manager.release_lock() - - mock_redis.redis_client.eval.assert_called_once() - assert lock_manager.lock_acquired is False - - def test_context_manager(self): - """Test MigrationLockManager as context manager""" - mock_redis = Mock() - mock_redis.redis_client.set.return_value = True - mock_redis.redis_client.eval.return_value = 1 - - lock_manager = MigrationLockManager(mock_redis) - lock_manager.pod_id = "pod_123_456" - - with lock_manager: - assert lock_manager.lock_acquired is True - - # Should call release_lock when exiting context - mock_redis.redis_client.eval.assert_called_once() - - def test_wait_for_lock_release_success(self): - """Test waiting for lock release and acquiring it""" - mock_redis = Mock() - # First call fails, second call succeeds - mock_redis.redis_client.set.side_effect = [False, True] - - lock_manager = MigrationLockManager(mock_redis) - result = lock_manager.wait_for_lock_release(check_interval=0.1, max_wait=1) - - assert result is True - assert lock_manager.lock_acquired is True - assert mock_redis.redis_client.set.call_count == 2 - - def test_wait_for_lock_release_timeout(self): - """Test timeout when waiting for lock release""" - mock_redis = Mock() - # Always fails to acquire lock - mock_redis.redis_client.set.return_value = False - - lock_manager = MigrationLockManager(mock_redis) - start_time = time.time() - result = lock_manager.wait_for_lock_release(check_interval=0.1, max_wait=0.5) - end_time = time.time() - - assert result is False - assert lock_manager.lock_acquired is False - assert end_time - start_time >= 0.5 # Should wait at least max_wait time - assert mock_redis.redis_client.set.call_count > 1 # Multiple attempts - - -class TestProxyExtrasDBManagerMigrationLock: - """Test cases for ProxyExtrasDBManager with migration locking""" - - @patch("litellm_proxy_extras.utils.ProxyExtrasDBManager._execute_migration") - def test_setup_database_with_redis_lock_success( - self, mock_execute_migration, monkeypatch - ): - """Test successful database setup with Redis lock""" - monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test") - mock_execute_migration.return_value = True - - # Mock Redis cache - mock_redis = Mock() - mock_redis.redis_client.set.return_value = True # Lock acquired - mock_redis.redis_client.eval.return_value = 1 # Lock released - - result = ProxyExtrasDBManager.setup_database( - use_migrate=True, redis_cache=mock_redis - ) - - assert result is True - mock_execute_migration.assert_called_once() - # Verify lock was acquired - mock_redis.redis_client.set.assert_called_once() - - @patch("litellm_proxy_extras.utils.ProxyExtrasDBManager._execute_migration") - def test_setup_database_with_redis_lock_wait_and_acquire( - self, mock_execute_migration, monkeypatch - ): - """Test database setup when lock is held, then acquired after waiting""" - monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test") - mock_execute_migration.return_value = True - - # Mock Redis cache - first call fails, second call succeeds - mock_redis = Mock() - mock_redis.redis_client.set.side_effect = [False, True] - mock_redis.redis_client.eval.return_value = 1 - - result = ProxyExtrasDBManager.setup_database( - use_migrate=True, redis_cache=mock_redis - ) - - assert result is True - mock_execute_migration.assert_called_once() - # Should have tried to acquire lock twice - assert mock_redis.redis_client.set.call_count == 2 - - @patch("litellm_proxy_extras.utils.ProxyExtrasDBManager._execute_migration") - def test_setup_database_without_redis(self, mock_execute_migration, monkeypatch): - """Test database setup without Redis cache (single instance mode)""" - monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test") - mock_execute_migration.return_value = True - - result = ProxyExtrasDBManager.setup_database(use_migrate=True, redis_cache=None) - - assert result is True - mock_execute_migration.assert_called_once() - - def test_setup_database_no_database_url(self): - """Test database setup without DATABASE_URL""" - with patch.dict(os.environ, {}, clear=True): - result = ProxyExtrasDBManager.setup_database( - use_migrate=True, redis_cache=None - ) - - assert result is False - - @patch("litellm_proxy_extras.utils.ProxyExtrasDBManager._execute_migration") - def test_setup_database_lock_timeout(self, mock_execute_migration, monkeypatch): - """Test database setup when lock acquisition times out""" - monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test") - - # Mock Redis cache - always fails to acquire lock - mock_redis = Mock() - mock_redis.redis_client.set.return_value = False - - # Patch wait_for_lock_release to simulate timeout - with patch.object(MigrationLockManager, "wait_for_lock_release") as mock_wait: - mock_wait.return_value = False # Simulate timeout - - result = ProxyExtrasDBManager.setup_database( - use_migrate=True, redis_cache=mock_redis - ) - - assert result is False - mock_execute_migration.assert_not_called() # Should not run migration - mock_wait.assert_called_once() From f6d6455cbc3cd64324981364d4766f78b7d4329f Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 20 Jan 2026 08:39:17 -0800 Subject: [PATCH 30/58] fix rc --- .../my-website/release_notes/v1.81.0/index.md | 2 +- ...odel_prices_and_context_window_backup.json | 31 +++++++++---------- 2 files changed, 16 insertions(+), 17 deletions(-) diff --git a/docs/my-website/release_notes/v1.81.0/index.md b/docs/my-website/release_notes/v1.81.0/index.md index edd720f6bfd..88ac240c614 100644 --- a/docs/my-website/release_notes/v1.81.0/index.md +++ b/docs/my-website/release_notes/v1.81.0/index.md @@ -27,7 +27,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -docker.litellm.ai/berriai/litellm:v1.81.0 +docker.litellm.ai/berriai/litellm:v1.81.0.rc.1 ``` diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 517e5c0d695..135b0d46ed0 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -15989,6 +15989,21 @@ "supports_response_schema": true, "supports_vision": true }, + "chatgpt/gpt-5.2": { + "litellm_provider": "chatgpt", + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "responses", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_vision": true + }, "chatgpt/gpt-5.1-codex-max": { "litellm_provider": "chatgpt", "max_input_tokens": 128000, @@ -16017,22 +16032,6 @@ "supports_response_schema": true, "supports_vision": true }, - "chatgpt/gpt-5.2": { - "litellm_provider": "chatgpt", - "max_input_tokens": 128000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "mode": "chat", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/messages", - "/v1/responses" - ], - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_vision": true - }, "gigachat/GigaChat-2-Lite": { "input_cost_per_token": 0.0, "litellm_provider": "gigachat", From ce37729da4c901ff416a768cec3daeb68f0a9cd4 Mon Sep 17 00:00:00 2001 From: Otavio Brito <69211663+otaviofbrito@users.noreply.github.com> Date: Tue, 20 Jan 2026 14:03:27 -0300 Subject: [PATCH 31/58] remove count tokens optional param before request is sent to vertex (#19359) --- litellm/llms/anthropic/chat/handler.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 9cc8c24ed13..6a9aafd076b 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -317,6 +317,7 @@ class AnthropicChatCompletion(BaseLLM): stream = optional_params.pop("stream", None) json_mode: bool = optional_params.pop("json_mode", False) is_vertex_request: bool = optional_params.pop("is_vertex_request", False) + optional_params.pop("vertex_count_tokens_location", None) _is_function_call = False messages = copy.deepcopy(messages) headers = AnthropicConfig().validate_environment( From 1377721715f36d3847d862fc84a7c29eb164377f Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Tue, 20 Jan 2026 10:44:31 -0800 Subject: [PATCH 32/58] Fix: Handle PostgreSQL cached plan errors during rolling deployments (#19424) --- litellm/proxy/utils.py | 41 ++++++++++++++++++++++++++++++++++++++++- 1 file changed, 40 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9ea2ea7d5c9..fcb678ef02a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2214,6 +2214,45 @@ class PrismaClient: raise e + async def _query_first_with_cached_plan_fallback( + self, sql_query: str + ) -> Optional[dict]: + """ + Execute a query with automatic fallback for PostgreSQL cached plan errors. + + This handles the "cached plan must not change result type" error that occurs + during rolling deployments when schema changes are applied while old pods + still have cached query plans expecting the old schema. + + Args: + sql_query: SQL query string to execute + + Returns: + Query result or None + + Raises: + Original exception if not a cached plan error + """ + try: + return await self.db.query_first(query=sql_query) + except Exception as e: + error_str = str(e) + if "cached plan must not change result type" in error_str: + # Force PostgreSQL to re-plan by invalidating the cache + # Add a unique comment to make the query different + sql_query_retry = sql_query.replace( + "SELECT", + f"SELECT /* cache_invalidated_{int(time.time() * 1000)} */" + ) + verbose_proxy_logger.warning( + "PostgreSQL cached plan error detected for token lookup, " + "retrying with fresh plan. This may occur during rolling deployments " + "when schema changes are applied." + ) + return await self.db.query_first(query=sql_query_retry) + else: + raise + @backoff.on_exception( backoff.expo, Exception, # base exception to catch for the backoff @@ -2545,7 +2584,7 @@ class PrismaClient: WHERE v.token = '{token}' """ - response = await self.db.query_first(query=sql_query) + response = await self._query_first_with_cached_plan_fallback(sql_query) if response is not None: if response["team_models"] is None: From f95f5563ea5501cad2ffff6ee32f8883ff0f76c5 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 20 Jan 2026 10:44:49 -0800 Subject: [PATCH 33/58] docs: document input/output/total tokens behaviour Closes https://github.com/BerriAI/litellm/issues/17480 --- docs/my-website/docs/proxy/users.md | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/docs/my-website/docs/proxy/users.md b/docs/my-website/docs/proxy/users.md index 3e0e00dfa52..a389f0bd443 100644 --- a/docs/my-website/docs/proxy/users.md +++ b/docs/my-website/docs/proxy/users.md @@ -545,6 +545,26 @@ You can set: - max parallel requests - rpm / tpm limits per model for a given key +### TPM Rate Limit Type (Input/Output/Total) + +By default, TPM (tokens per minute) rate limits count **total tokens** (input + output). You can configure this to count only input tokens or only output tokens instead. + +Set `token_rate_limit_type` in your `config.yaml`: + +```yaml +general_settings: + master_key: sk-1234 + token_rate_limit_type: "output" # Options: "input", "output", "total" (default) +``` + +| Value | Description | +|-------|-------------| +| `total` | Count total tokens (prompt + completion). **Default behavior.** | +| `input` | Count only prompt/input tokens | +| `output` | Count only completion/output tokens | + +This setting applies globally to all TPM rate limit checks (keys, users, teams, etc.). + From 5a068686524295f2171e3f79abda157b9e02095a Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Tue, 20 Jan 2026 12:17:06 -0800 Subject: [PATCH 34/58] Fix in-flight request termination on SIGTERM when health-check runs in a separate process (#19427) --- docker/prod_entrypoint.sh | 1 + docker/supervisord.conf | 2 ++ docs/my-website/docs/proxy/config_settings.md | 1 + docs/my-website/docs/proxy/prod.md | 5 +++++ 4 files changed, 9 insertions(+) diff --git a/docker/prod_entrypoint.sh b/docker/prod_entrypoint.sh index 1fc09d2c864..28d1bdcc294 100644 --- a/docker/prod_entrypoint.sh +++ b/docker/prod_entrypoint.sh @@ -2,6 +2,7 @@ if [ "$SEPARATE_HEALTH_APP" = "1" ]; then export LITELLM_ARGS="$@" + export SUPERVISORD_STOPWAITSECS="${SUPERVISORD_STOPWAITSECS:-3600}" exec supervisord -c /etc/supervisord.conf fi diff --git a/docker/supervisord.conf b/docker/supervisord.conf index 877335804fe..ba9d99d18a5 100644 --- a/docker/supervisord.conf +++ b/docker/supervisord.conf @@ -16,6 +16,7 @@ priority=1 exitcodes=0 stopasgroup=true killasgroup=true +stopwaitsecs=%(ENV_SUPERVISORD_STOPWAITSECS)s stdout_logfile=/dev/stdout stderr_logfile=/dev/stderr stdout_logfile_maxbytes = 0 @@ -31,6 +32,7 @@ priority=2 exitcodes=0 stopasgroup=true killasgroup=true +stopwaitsecs=%(ENV_SUPERVISORD_STOPWAITSECS)s stdout_logfile=/dev/stdout stderr_logfile=/dev/stderr stdout_logfile_maxbytes = 0 diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 53b0f3eea71..f76fd214682 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -875,6 +875,7 @@ router_settings: | SECRET_MANAGER_REFRESH_INTERVAL | Refresh interval in seconds for secret manager. Default is 86400 (24 hours) | SEPARATE_HEALTH_APP | If set to '1', runs health endpoints on a separate ASGI app and port. Default: '0'. | SEPARATE_HEALTH_PORT | Port for the separate health endpoints app. Only used if SEPARATE_HEALTH_APP=1. Default: 4001. +| SUPERVISORD_STOPWAITSECS | Upper bound timeout in seconds for graceful shutdown when SEPARATE_HEALTH_APP=1. Default: 3600 (1 hour). | SERVER_ROOT_PATH | Root path for the server application | SEND_USER_API_KEY_ALIAS | Flag to send user API key alias to Zscaler AI Guard. Default is False | SEND_USER_API_KEY_TEAM_ID | Flag to send user API key team ID to Zscaler AI Guard. Default is False diff --git a/docs/my-website/docs/proxy/prod.md b/docs/my-website/docs/proxy/prod.md index 9216b0fbf30..a42d91a7d5f 100644 --- a/docs/my-website/docs/proxy/prod.md +++ b/docs/my-website/docs/proxy/prod.md @@ -277,8 +277,13 @@ Set the following environment variable(s): ```bash SEPARATE_HEALTH_APP="1" # Default "0" SEPARATE_HEALTH_PORT="8001" # Default "4001", Works only if `SEPARATE_HEALTH_APP` is "1" +SUPERVISORD_STOPWAITSECS="3600" # Optional: Upper bound timeout in seconds for graceful shutdown. Default: 3600 (1 hour). Only used when SEPARATE_HEALTH_APP=1. ``` +**Graceful Shutdown:** + +Previously, `stopwaitsecs` was not set, defaulting to 10 seconds and causing in-flight requests to fail. `SUPERVISORD_STOPWAITSECS` (default: 3600) provides an upper bound for graceful shutdown, allowing uvicorn to wait for all in-flight requests to complete. +