scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-04 02:31:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 70

fix(gemini): gate auto-fallback and harden oauth/token paths

Browse source
This commit is contained in:
balazss 2026-03-17 21:20:07 -07:00
parent 3c143336f7
commit 4079922c1c
8 changed files with 220 additions and 96 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
litellm/__init__.py
View file

@ -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,

5
litellm/llms/gemini/authenticator.py
View file

@ -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:

6
litellm/llms/gemini/common_utils.py
View file

@ -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}

8
litellm/llms/gemini/fallback_handler.py
View file

@ -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

2
litellm/llms/vertex_ai/common_utils.py
View file

@ -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:

2
litellm/main.py
View file

@ -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":

57
tests/test_litellm/llms/gemini/test_fallback_handler.py Normal file
View file

@ -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()

220
tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py
View file

@ -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)