diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index ef987e46b7..5f1b309ef8 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -242,6 +242,10 @@ def _parse_tool_call_arguments(arguments: Any) -> Any: # Providers that require file_id instead of inline file_data _FILE_ID_REQUIRED_PROVIDERS = frozenset({"openai", "azure"}) +# Routing-only prefix: requests go through a LiteLLM Proxy deployment, but the +# payload must still be shaped for the provider named in the next segment. +_PROXY_PROVIDER = "litellm_proxy" + _MISSING_TOOL_RESULT_MESSAGE = ( "Error: Missing tool result (tool execution may have been interrupted " "before a response was recorded)." @@ -309,6 +313,34 @@ def _map_finish_reason( return _FINISH_REASON_MAPPING.get(finish_reason_str, types.FinishReason.OTHER) +def _strip_proxy_prefix(model: str) -> str: + """Removes a leading ``litellm_proxy/`` routing prefix from a model string. + + ``litellm_proxy`` selects the transport (a LiteLLM Proxy deployment), not the + model family, so the segment after it identifies the provider that actually + serves the request (e.g. ``litellm_proxy/azure/my-deployment`` is served by + Azure). Provider-specific request shaping must follow that underlying + provider, otherwise proxied requests get generic payloads the backend + rejects. + + A bare ``litellm_proxy/`` has no nested provider, so the + remainder is returned as-is and provider detection falls back to the + model-name heuristics. + + Args: + model: The model string (e.g., "litellm_proxy/azure/gpt-4"). + + Returns: + The model string without the ``litellm_proxy/`` prefix. + """ + if not model: + return model + prefix = _PROXY_PROVIDER + "/" + if model.lower().startswith(prefix): + return model[len(prefix) :] + return model + + def _get_provider_from_model(model: str) -> str: """Extracts the provider name from a LiteLLM model string. @@ -320,6 +352,9 @@ def _get_provider_from_model(model: str) -> str: """ if not model: return "" + # `litellm_proxy` is a transport prefix; the provider that actually serves + # the request is the next segment. + model = _strip_proxy_prefix(model) # LiteLLM uses "provider/model" format if "/" in model: provider, _ = model.split("/", 1) @@ -2685,7 +2720,7 @@ def _is_anthropic_model(model_string: str) -> bool: Returns: True if it's an Anthropic Claude model, False otherwise. """ - lower = model_string.lower() + lower = _strip_proxy_prefix(model_string.lower()) if lower.startswith("anthropic/"): return True if lower.startswith("bedrock/"): @@ -2706,7 +2741,7 @@ def _is_litellm_vertex_model(model_string: str) -> bool: Returns: True if it's a Vertex AI model accessed via LiteLLM, False otherwise """ - return model_string.startswith("vertex_ai/") + return _strip_proxy_prefix(model_string).startswith("vertex_ai/") def _is_litellm_gemini_model(model_string: str) -> bool: @@ -2719,7 +2754,9 @@ def _is_litellm_gemini_model(model_string: str) -> bool: Returns: True if it's a Gemini model accessed via LiteLLM, False otherwise """ - return model_string.startswith(("gemini/gemini-", "vertex_ai/gemini-")) + return _strip_proxy_prefix(model_string).startswith( + ("gemini/gemini-", "vertex_ai/gemini-") + ) def _extract_gemini_model_from_litellm(litellm_model: str) -> str: @@ -2731,6 +2768,9 @@ def _extract_gemini_model_from_litellm(litellm_model: str) -> str: Returns: Pure Gemini model name like "gemini-2.5-pro" """ + # Remove the proxy routing prefix first so the provider prefix below is the + # one that actually names the model family. + litellm_model = _strip_proxy_prefix(litellm_model) # Remove LiteLLM provider prefix if "/" in litellm_model: return litellm_model.split("/", 1)[1] diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index bd1729c639..ab4b0fb2da 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -33,6 +33,7 @@ from google.adk.models.lite_llm import _content_to_message_param from google.adk.models.lite_llm import _convert_reasoning_value_to_parts from google.adk.models.lite_llm import _enforce_strict_openai_schema +from google.adk.models.lite_llm import _extract_gemini_model_from_litellm from google.adk.models.lite_llm import _extract_json_from_deepseek_args from google.adk.models.lite_llm import _extract_reasoning_value from google.adk.models.lite_llm import _extract_thought_signature_from_tool_call @@ -45,6 +46,8 @@ from google.adk.models.lite_llm import _is_anthropic_model from google.adk.models.lite_llm import _is_anthropic_provider from google.adk.models.lite_llm import _is_anthropic_route +from google.adk.models.lite_llm import _is_litellm_gemini_model +from google.adk.models.lite_llm import _is_litellm_vertex_model from google.adk.models.lite_llm import _looks_like_openai_file_id from google.adk.models.lite_llm import _message_to_generate_content_response from google.adk.models.lite_llm import _MISSING_TOOL_RESULT_MESSAGE @@ -5172,6 +5175,17 @@ async def test_finish_reason_unknown_maps_to_other( ("groq/llama3-70b", "groq"), ("anthropic/claude-3", "anthropic"), ("vertex_ai/gemini-pro", "vertex_ai"), + # litellm_proxy is a routing prefix: the provider that actually serves + # the request is the segment after it. + ("litellm_proxy/azure/my-deployment", "azure"), + ("litellm_proxy/openai/gpt-4o", "openai"), + ("litellm_proxy/anthropic/claude-3", "anthropic"), + ("litellm_proxy/vertex_ai/gemini-pro", "vertex_ai"), + ("LiteLLM_Proxy/azure/gpt-4", "azure"), + # A bare proxy deployment name has no nested provider, so detection + # falls back to the model-name heuristics. + ("litellm_proxy/azure-gpt-4", "azure"), + ("litellm_proxy/my-deployment", ""), # Fallback heuristics ("gpt-4o", "openai"), ("o1-preview", "openai"), @@ -5187,6 +5201,55 @@ def test_get_provider_from_model(model_string, expected_provider): assert _get_provider_from_model(model_string) == expected_provider +@pytest.mark.parametrize( + "model_string, is_anthropic, is_gemini, is_vertex, gemini_name", + [ + # Proxied models keep the behavior of the provider that serves them. + ( + "litellm_proxy/anthropic/claude-4-sonnet", + True, + False, + False, + "claude-4-sonnet", + ), + ( + "litellm_proxy/vertex_ai/gemini-2.5-flash", + False, + True, + True, + "gemini-2.5-flash", + ), + ( + "litellm_proxy/bedrock/anthropic.claude-3-5-sonnet", + True, + False, + False, + "anthropic.claude-3-5-sonnet", + ), + ( + "litellm_proxy/azure/my-deployment", + False, + False, + False, + "my-deployment", + ), + # Direct (non-proxied) strings are unaffected. + ("anthropic/claude-4-sonnet", True, False, False, "claude-4-sonnet"), + ("vertex_ai/gemini-2.5-flash", False, True, True, "gemini-2.5-flash"), + ("gemini/gemini-2.5-pro", False, True, False, "gemini-2.5-pro"), + ("azure/gpt-4", False, False, False, "gpt-4"), + ], +) +def test_model_family_detection_through_litellm_proxy( + model_string, is_anthropic, is_gemini, is_vertex, gemini_name +): + """Model-family detection must see through the litellm_proxy prefix.""" + assert _is_anthropic_model(model_string) is is_anthropic + assert _is_litellm_gemini_model(model_string) is is_gemini + assert _is_litellm_vertex_model(model_string) is is_vertex + assert _extract_gemini_model_from_litellm(model_string) == gemini_name + + @pytest.mark.parametrize( "provider, expected_in_list", [ @@ -5226,6 +5289,38 @@ async def test_get_content_pdf_openai_uses_file_id(mocker): ) +@pytest.mark.asyncio +async def test_get_content_pdf_proxied_azure_uses_file_id(mocker): + """PDFs sent to a proxied Azure model must upload and send a file_id. + + Regression test: a nested ``litellm_proxy/azure/...`` identifier used to be + classified as the ``litellm_proxy`` provider, which skipped the Azure upload + path and emitted a bare ``file_data`` block that Azure rejects. + """ + mock_file_response = mocker.create_autospec(litellm.FileObject) + mock_file_response.id = "file-abc123" + mock_acreate_file = AsyncMock(return_value=mock_file_response) + mocker.patch.object(litellm, "acreate_file", new=mock_acreate_file) + + model = "litellm_proxy/azure/my-deployment" + parts = [ + types.Part.from_bytes(data=b"test_pdf_data", mime_type="application/pdf") + ] + content = await _get_content( + parts, provider=_get_provider_from_model(model), model=model + ) + + assert content[0]["type"] == "file" + assert content[0]["file"]["file_id"] == "file-abc123" + assert "file_data" not in content[0]["file"] + + mock_acreate_file.assert_called_once_with( + file=b"test_pdf_data", + purpose="assistants", + custom_llm_provider="azure", + ) + + @pytest.mark.asyncio async def test_get_content_pdf_non_openai_uses_file_data(): """Test that PDF files use file_data for non-OpenAI providers."""