From 4079922c1c946b430ffa00e2298634796148b29d Mon Sep 17 00:00:00 2001 From: balazss Date: Tue, 17 Mar 2026 21:20:07 -0700 Subject: [PATCH] fix(gemini): gate auto-fallback and harden oauth/token paths --- litellm/__init__.py | 16 +- litellm/llms/gemini/authenticator.py | 5 +- litellm/llms/gemini/common_utils.py | 6 +- litellm/llms/gemini/fallback_handler.py | 8 + litellm/llms/vertex_ai/common_utils.py | 2 +- litellm/main.py | 2 + .../llms/gemini/test_fallback_handler.py | 57 +++++ .../vertex_ai/test_vertex_ai_common_utils.py | 220 +++++++++++------- 8 files changed, 220 insertions(+), 96 deletions(-) create mode 100644 tests/test_litellm/llms/gemini/test_fallback_handler.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 51c66838613..52f344fb797 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -433,6 +433,10 @@ default_fallbacks: Optional[List] = None fallbacks: Optional[List] = None context_window_fallbacks: Optional[List] = None content_policy_fallbacks: Optional[List] = None +# Backward-compatible default: do not silently reroute gemini/* calls unless explicitly enabled. +auto_fallback_to_google_code_assist: bool = os.getenv( + "LITELLM_AUTO_FALLBACK_TO_GOOGLE_CODE_ASSIST", "false" +).lower() in ("1", "true", "yes", "on") allowed_fails: int = 3 allow_dynamic_callback_disabling: bool = True num_retries_per_request: Optional[ @@ -1465,9 +1469,15 @@ if TYPE_CHECKING: from .llms.petals.completion.transformation import PetalsConfig as PetalsConfig from .llms.ollama.chat.transformation import OllamaChatConfig as OllamaChatConfig from .llms.ollama.completion.transformation import OllamaConfig as OllamaConfig - from .llms.sagemaker.completion.transformation import SagemakerConfig as SagemakerConfig - from .llms.sagemaker.chat.transformation import SagemakerChatConfig as SagemakerChatConfig - from .llms.sagemaker.nova.transformation import SagemakerNovaConfig as SagemakerNovaConfig + from .llms.sagemaker.completion.transformation import ( + SagemakerConfig as SagemakerConfig, + ) + from .llms.sagemaker.chat.transformation import ( + SagemakerChatConfig as SagemakerChatConfig, + ) + from .llms.sagemaker.nova.transformation import ( + SagemakerNovaConfig as SagemakerNovaConfig, + ) from .llms.cohere.chat.transformation import CohereChatConfig as CohereChatConfig from .llms.anthropic.experimental_pass_through.messages.transformation import ( AnthropicMessagesConfig as AnthropicMessagesConfig, diff --git a/litellm/llms/gemini/authenticator.py b/litellm/llms/gemini/authenticator.py index c5f6aab8a70..bd7d6d0b8dc 100644 --- a/litellm/llms/gemini/authenticator.py +++ b/litellm/llms/gemini/authenticator.py @@ -233,8 +233,9 @@ class GeminiAuthenticator: ) webbrowser.open(auth_url) - # Wait for callback - server.handle_request() + # Wait for callback; browsers may hit non-callback paths first (e.g. /favicon.ico). + while auth_code is None and error is None: + server.handle_request() server.server_close() if error: diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index a1cf0564f68..99fcedc8a44 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -324,7 +324,11 @@ def get_gemini_oauth_token() -> Optional[dict]: # noqa: PLR0915 continue token = creds_data.get("access_token") if not token and "token" in creds_data: - token = creds_data["token"].get("accessToken") + token_field = creds_data["token"] + if isinstance(token_field, dict): + token = token_field.get("accessToken") + elif isinstance(token_field, str): + token = token_field if token: result = {"token": token} diff --git a/litellm/llms/gemini/fallback_handler.py b/litellm/llms/gemini/fallback_handler.py index 60343b884cd..a59497657e6 100644 --- a/litellm/llms/gemini/fallback_handler.py +++ b/litellm/llms/gemini/fallback_handler.py @@ -8,6 +8,7 @@ from litellm.llms.google_code_assist.chat import GoogleCodeAssistChat async def run_gemini_acompletion_with_code_assist_fallback( primary_call: Awaitable[Any], fallback_kwargs: Dict[str, Any], + auto_fallback_to_google_code_assist: bool = False, ) -> Any: """ Execute Gemini async completion and fallback to Google Code Assist when @@ -16,6 +17,9 @@ async def run_gemini_acompletion_with_code_assist_fallback( try: return await primary_call except Exception as e: + if not auto_fallback_to_google_code_assist: + raise e + if not should_fallback_to_google_code_assist(e): raise e @@ -29,6 +33,7 @@ async def run_gemini_acompletion_with_code_assist_fallback( def run_gemini_completion_with_code_assist_fallback( primary_call: Callable[[], Any], fallback_kwargs: Dict[str, Any], + auto_fallback_to_google_code_assist: bool = False, ) -> Any: """ Execute Gemini sync completion and fallback to Google Code Assist when @@ -37,6 +42,9 @@ def run_gemini_completion_with_code_assist_fallback( try: return primary_call() except Exception as e: + if not auto_fallback_to_google_code_assist: + raise e + if not should_fallback_to_google_code_assist(e): raise e diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 9bd2e1f45f7..754227e8873 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -374,7 +374,7 @@ def _get_gemini_url( params = [] if gemini_api_key and not gemini_oauth_token: params.append(f"key={gemini_api_key}") - if stream: + if mode == "chat" and stream: params.append("alt=sse") if params: diff --git a/litellm/main.py b/litellm/main.py index 41140e0cae8..b59ea1bc801 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3476,6 +3476,7 @@ def completion( # type: ignore # noqa: PLR0915 "logging_obj": logging, "logger_fn": logger_fn, }, + auto_fallback_to_google_code_assist=litellm.auto_fallback_to_google_code_assist, ) else: response = run_gemini_completion_with_code_assist_fallback( @@ -3510,6 +3511,7 @@ def completion( # type: ignore # noqa: PLR0915 "logging_obj": logging, "logger_fn": logger_fn, }, + auto_fallback_to_google_code_assist=litellm.auto_fallback_to_google_code_assist, ) elif custom_llm_provider == "vertex_ai": diff --git a/tests/test_litellm/llms/gemini/test_fallback_handler.py b/tests/test_litellm/llms/gemini/test_fallback_handler.py new file mode 100644 index 00000000000..544f241a99b --- /dev/null +++ b/tests/test_litellm/llms/gemini/test_fallback_handler.py @@ -0,0 +1,57 @@ +from unittest.mock import AsyncMock, patch + +import pytest + +from litellm.llms.gemini.fallback_handler import ( + run_gemini_acompletion_with_code_assist_fallback, + run_gemini_completion_with_code_assist_fallback, +) + + +def test_run_gemini_completion_with_code_assist_fallback_disabled(): + def _raise_scope_error(): + raise Exception("ACCESS_TOKEN_SCOPE_INSUFFICIENT") + + with ( + patch( + "litellm.llms.gemini.fallback_handler.should_fallback_to_google_code_assist", + return_value=True, + ), + patch( + "litellm.llms.gemini.fallback_handler.GoogleCodeAssistChat.completion" + ) as mock_completion, + ): + with pytest.raises(Exception, match="ACCESS_TOKEN_SCOPE_INSUFFICIENT"): + run_gemini_completion_with_code_assist_fallback( + primary_call=_raise_scope_error, + fallback_kwargs={}, + auto_fallback_to_google_code_assist=False, + ) + + mock_completion.assert_not_called() + + +@pytest.mark.asyncio +async def test_run_gemini_acompletion_with_code_assist_fallback_enabled(): + async def _raise_scope_error(): + raise Exception("ACCESS_TOKEN_SCOPE_INSUFFICIENT") + + with ( + patch( + "litellm.llms.gemini.fallback_handler.should_fallback_to_google_code_assist", + return_value=True, + ), + patch( + "litellm.llms.gemini.fallback_handler.GoogleCodeAssistChat.acompletion", + new_callable=AsyncMock, + ) as mock_acompletion, + ): + mock_acompletion.return_value = "fallback-ok" + result = await run_gemini_acompletion_with_code_assist_fallback( + primary_call=_raise_scope_error(), + fallback_kwargs={}, + auto_fallback_to_google_code_assist=True, + ) + + assert result == "fallback-ok" + mock_acompletion.assert_awaited_once() 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 d483a81a349..79db733e4aa 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 @@ -11,6 +11,7 @@ sys.path.insert( ) # Adds the parent directory to the system path from litellm.llms.vertex_ai.common_utils import ( + _get_gemini_url, _get_vertex_url, convert_anyof_null_to_nullable, get_vertex_location_from_url, @@ -440,7 +441,9 @@ def test_vertex_ai_complex_response_schema(): optional_params = {} v.apply_response_schema_transformation( - value=non_default_params["response_format"], optional_params=optional_params, model="gemini-1.5-pro-preview-0409" + value=non_default_params["response_format"], + optional_params=optional_params, + model="gemini-1.5-pro-preview-0409", ) # Assertions for the transformed schema @@ -558,6 +561,25 @@ def test_get_vertex_url_global_region(stream, expected_endpoint_suffix): assert url == expected_url +def test_get_gemini_url_stream_query_param_only_for_chat_mode(): + chat_url, _ = _get_gemini_url( + mode="chat", + model="gemini-1.5-flash", + stream=True, + gemini_api_key="test-key", + gemini_oauth_token=None, + ) + embedding_url, _ = _get_gemini_url( + mode="embedding", + model="gemini-1.5-flash", + stream=True, + gemini_api_key="test-key", + gemini_oauth_token=None, + ) + + assert "alt=sse" in chat_url + assert "alt=sse" not in embedding_url + @pytest.mark.parametrize( "model_cost_entry, vertex_region, expected_region", @@ -571,9 +593,17 @@ def test_get_vertex_url_global_region(stream, expected_endpoint_suffix): # Model with supported_regions=["us-west2"], no user region -> use "us-west2" ({"supported_regions": ["us-west2"]}, None, "us-west2"), # Model with supported_regions=["us-west2", "us-central1"], user passes supported region -> respect it - ({"supported_regions": ["us-west2", "us-central1"]}, "us-central1", "us-central1"), + ( + {"supported_regions": ["us-west2", "us-central1"]}, + "us-central1", + "us-central1", + ), # Model with supported_regions=["us-west2", "us-central1"], user passes unsupported region -> override - ({"supported_regions": ["us-west2", "us-central1"]}, "europe-west1", "us-west2"), + ( + {"supported_regions": ["us-west2", "us-central1"]}, + "europe-west1", + "us-west2", + ), # No model_cost entry, no user region -> default us-central1 ({}, None, "us-central1"), # No model_cost entry, user specifies region -> use specified region @@ -656,11 +686,12 @@ def test_vertex_filter_format_uri(): assert "uri" not in json.dumps(new_parameters) + def test_convert_schema_types_type_array_conversion(): """ Test _convert_schema_types function handles type arrays and case conversion. - - This test verifies the fix for the issue where type arrays like ["string", "number"] + + This test verifies the fix for the issue where type arrays like ["string", "number"] would raise an exception in Vertex AI schema validation. Relevant issue: https://github.com/BerriAI/litellm/issues/14091 @@ -673,12 +704,12 @@ def test_convert_schema_types_type_array_conversion(): "properties": { "studio": { "type": ["string", "number"], - "description": "The studio ID or name" + "description": "The studio ID or name", } }, "required": ["studio"], "additionalProperties": False, - "$schema": "http://json-schema.org/draft-07/schema#" + "$schema": "http://json-schema.org/draft-07/schema#", } # Expected output: Vertex AI compatible schema with anyOf and uppercase types @@ -686,16 +717,13 @@ def test_convert_schema_types_type_array_conversion(): "type": "object", "properties": { "studio": { - "anyOf": [ - {"type": "string"}, - {"type": "number"} - ], - "description": "The studio ID or name" + "anyOf": [{"type": "string"}, {"type": "number"}], + "description": "The studio ID or name", } }, "required": ["studio"], "additionalProperties": False, - "$schema": "http://json-schema.org/draft-07/schema#" + "$schema": "http://json-schema.org/draft-07/schema#", } # Apply the transformation @@ -718,15 +746,17 @@ def test_convert_schema_types_type_array_conversion(): assert anyof_types[1]["type"] == "number" # 4. Other properties preserved - assert input_schema["properties"]["studio"]["description"] == "The studio ID or name" + assert ( + input_schema["properties"]["studio"]["description"] == "The studio ID or name" + ) assert input_schema["required"] == ["studio"] def test_fix_enum_empty_strings(): """ Test _fix_enum_empty_strings function replaces empty strings with None in enum arrays. - - This test verifies the fix for the issue where Gemini rejects tool definitions + + This test verifies the fix for the issue where Gemini rejects tool definitions with empty strings in enum values, causing API failures. Relevant issue: Gemini does not accept empty strings in enum values @@ -740,23 +770,23 @@ def test_fix_enum_empty_strings(): "user_agent_type": { "enum": ["", "desktop", "mobile", "tablet"], "type": "string", - "description": "Device type for user agent" + "description": "Device type for user agent", } }, - "required": ["user_agent_type"] + "required": ["user_agent_type"], } # Expected output: Empty strings replaced with None expected_output = { - "type": "object", + "type": "object", "properties": { "user_agent_type": { "enum": [None, "desktop", "mobile", "tablet"], "type": "string", - "description": "Device type for user agent" + "description": "Device type for user agent", } }, - "required": ["user_agent_type"] + "required": ["user_agent_type"], } # Apply the transformation @@ -859,7 +889,7 @@ def test_construct_target_url_with_version_prefix(): def test_fix_enum_types(): """ Test _fix_enum_types function removes enum fields when type is not string. - + This test verifies the fix for the issue where Gemini rejects cached content with function parameter enums on non-string types, causing API failures. @@ -874,38 +904,41 @@ def test_fix_enum_types(): "truncateMode": { "enum": ["auto", "none", "start", "end"], "type": "string", # This should keep the enum - "description": "How to truncate content" + "description": "How to truncate content", }, "maxLength": { "enum": [100, 200, 500], # This should be removed "type": "integer", - "description": "Maximum length" + "description": "Maximum length", }, "enabled": { "enum": [True, False], # This should be removed "type": "boolean", - "description": "Whether feature is enabled" + "description": "Whether feature is enabled", }, "nested": { "type": "object", "properties": { "innerEnum": { "enum": ["a", "b", "c"], # This should be kept - "type": "string" + "type": "string", }, "innerNonStringEnum": { "enum": [1, 2, 3], # This should be removed - "type": "integer" - } - } + "type": "integer", + }, + }, }, "anyOfField": { "anyOf": [ - {"type": "string", "enum": ["option1", "option2"]}, # This should be kept - {"type": "integer", "enum": [1, 2, 3]} # This should be removed + { + "type": "string", + "enum": ["option1", "option2"], + }, # This should be kept + {"type": "integer", "enum": [1, 2, 3]}, # This should be removed ] - } - } + }, + }, } # Expected output: Non-string enums removed, string enums kept @@ -919,31 +952,32 @@ def test_fix_enum_types(): }, "maxLength": { # enum removed "type": "integer", - "description": "Maximum length" + "description": "Maximum length", }, "enabled": { # enum removed "type": "boolean", - "description": "Whether feature is enabled" + "description": "Whether feature is enabled", }, "nested": { "type": "object", "properties": { "innerEnum": { "enum": ["a", "b", "c"], # Kept - string type - "type": "string" + "type": "string", }, - "innerNonStringEnum": { # enum removed - "type": "integer" - } - } + "innerNonStringEnum": {"type": "integer"}, # enum removed + }, }, "anyOfField": { "anyOf": [ - {"type": "string", "enum": ["option1", "option2"]}, # Kept - has string type - {"type": "integer"} # enum removed + { + "type": "string", + "enum": ["option1", "option2"], + }, # Kept - has string type + {"type": "integer"}, # enum removed ] - } - } + }, + }, } # Apply the transformation @@ -955,15 +989,27 @@ def test_fix_enum_types(): # Verify specific transformations: # 1. String enums are preserved assert "enum" in input_schema["properties"]["truncateMode"] - assert input_schema["properties"]["truncateMode"]["enum"] == ["auto", "none", "start", "end"] - + assert input_schema["properties"]["truncateMode"]["enum"] == [ + "auto", + "none", + "start", + "end", + ] + assert "enum" in input_schema["properties"]["nested"]["properties"]["innerEnum"] - assert input_schema["properties"]["nested"]["properties"]["innerEnum"]["enum"] == ["a", "b", "c"] + assert input_schema["properties"]["nested"]["properties"]["innerEnum"]["enum"] == [ + "a", + "b", + "c", + ] # 2. Non-string enums are removed assert "enum" not in input_schema["properties"]["maxLength"] assert "enum" not in input_schema["properties"]["enabled"] - assert "enum" not in input_schema["properties"]["nested"]["properties"]["innerNonStringEnum"] + assert ( + "enum" + not in input_schema["properties"]["nested"]["properties"]["innerNonStringEnum"] + ) # 3. anyOf with string type keeps enum, non-string removes it assert "enum" in input_schema["properties"]["anyOfField"]["anyOf"][0] @@ -1003,8 +1049,6 @@ def test_get_token_url(): print("url=", url) - - should_use_v1beta1_features = vertex_llm.is_using_v1beta1_features( optional_params={"temperature": 0.1} ) @@ -1210,9 +1254,7 @@ def test_vertex_ai_minimax_uses_openai_handler(): VertexAIPartnerModels, ) - assert VertexAIPartnerModels.should_use_openai_handler( - "minimaxai/minimax-m2-maas" - ) + assert VertexAIPartnerModels.should_use_openai_handler("minimaxai/minimax-m2-maas") def test_vertex_ai_moonshot_uses_openai_handler(): @@ -1236,9 +1278,7 @@ def test_vertex_ai_zai_uses_openai_handler(): VertexAIPartnerModels, ) - assert VertexAIPartnerModels.should_use_openai_handler( - "zai-org/glm-4.7-maas" - ) + assert VertexAIPartnerModels.should_use_openai_handler("zai-org/glm-4.7-maas") def test_vertex_ai_zai_is_partner_model(): @@ -1255,14 +1295,14 @@ def test_vertex_ai_zai_is_partner_model(): def test_build_vertex_schema_empty_properties(): """ Test _build_vertex_schema handles empty properties objects correctly. - - This test verifies the fix for the issue where Gemini rejects schemas + + This test verifies the fix for the issue where Gemini rejects schemas with empty properties objects like {"properties": {}, "type": "object"}. - + Error from Gemini: "GenerateContentRequest.generation_config.response_schema - .properties[\"action\"].items.any_of[0].properties[\"go_back\"].properties: + .properties[\"action\"].items.any_of[0].properties[\"go_back\"].properties: should be non-empty for OBJECT type" - + The fix removes empty properties objects and their associated type/required fields. """ from litellm.llms.vertex_ai.common_utils import _build_vertex_schema @@ -1281,20 +1321,20 @@ def test_build_vertex_schema_empty_properties(): "type": "object", "additionalProperties": False, "description": "Go back", - "required": [] + "required": [], } }, "required": ["go_back"], "type": "object", - "additionalProperties": False + "additionalProperties": False, } ] }, - "type": "array" + "type": "array", } }, "type": "object", - "additionalProperties": False + "additionalProperties": False, } # Apply the transformation @@ -1302,24 +1342,36 @@ def test_build_vertex_schema_empty_properties(): # Verify the transformation removed empty properties # Navigate to the go_back schema - go_back_schema = result["properties"]["action"]["items"]["anyOf"][0]["properties"]["go_back"] - + go_back_schema = result["properties"]["action"]["items"]["anyOf"][0]["properties"][ + "go_back" + ] + # Verify empty properties was removed assert "properties" not in go_back_schema, "Empty properties should be removed" - + # Verify type is kept as object (Gemini requires type: object even without properties) - assert go_back_schema.get("type") == "object", "Type should be kept as object when properties is empty" - + assert ( + go_back_schema.get("type") == "object" + ), "Type should be kept as object when properties is empty" + # Verify required was also removed - assert "required" not in go_back_schema, "Required should be removed when properties is empty" - + assert ( + "required" not in go_back_schema + ), "Required should be removed when properties is empty" + # Verify description is preserved - assert go_back_schema.get("description") == "Go back", "Description should be preserved" - + assert ( + go_back_schema.get("description") == "Go back" + ), "Description should be preserved" + # Verify parent schema still has proper structure parent_schema = result["properties"]["action"]["items"]["anyOf"][0] - assert parent_schema["type"] == "object", "Parent schema should still have object type" - assert "go_back" in parent_schema["properties"], "go_back should still be in parent properties" + assert ( + parent_schema["type"] == "object" + ), "Parent schema should still have object type" + assert ( + "go_back" in parent_schema["properties"] + ), "go_back should still be in parent properties" def test_add_object_type_schema_with_no_properties_and_no_type(): @@ -1330,9 +1382,7 @@ def test_add_object_type_schema_with_no_properties_and_no_type(): from litellm.llms.vertex_ai.common_utils import add_object_type # Input: Schema with no properties and no type (the problematic case) - input_schema = { - "$schema": "https://json-schema.org/draft/2020-12/schema" - } + input_schema = {"$schema": "https://json-schema.org/draft/2020-12/schema"} # Apply the transformation add_object_type(input_schema) @@ -1351,10 +1401,7 @@ def test_add_object_type_does_not_override_existing_type(): from litellm.llms.vertex_ai.common_utils import add_object_type # Input: Schema with existing type - input_schema = { - "type": "string", - "description": "A string field" - } + input_schema = {"type": "string", "description": "A string field"} # Apply the transformation add_object_type(input_schema) @@ -1370,12 +1417,7 @@ def test_add_object_type_does_not_add_type_when_anyof_present(): from litellm.llms.vertex_ai.common_utils import add_object_type # Input: Schema with anyOf but no type - input_schema = { - "anyOf": [ - {"type": "string"}, - {"type": "null"} - ] - } + input_schema = {"anyOf": [{"type": "string"}, {"type": "null"}]} # Apply the transformation add_object_type(input_schema)