From 1a123b2cd5109977f834d6fdac9d342ccc477a0c Mon Sep 17 00:00:00 2001 From: Sameer Kankute <135028480+kankute-sameer@users.noreply.github.com> Date: Sat, 13 Sep 2025 00:25:26 +0530 Subject: [PATCH] Litellm gemini cli bug fix (#14451) * Fix gemini cli error * Add reasoning request support * Added better handling * remove other PR code * refactored code for better structure following --------- Co-authored-by: sameer@berri.ai --- litellm/google_genai/main.py | 12 + litellm/llms/gemini/count_tokens/handler.py | 1 - .../vertex_ai/google_genai/transformation.py | 73 +++- litellm/proxy/route_llm_request.py | 6 +- .../test_google_gemini_proxy_request.py | 350 ++++++++++++++++++ .../llms/gemini/test_gemini_common_utils.py | 2 +- 6 files changed, 431 insertions(+), 13 deletions(-) create mode 100644 tests/proxy_unit_tests/test_google_gemini_proxy_request.py diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index 87970885355..b480a85c85e 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -224,6 +224,9 @@ async def agenerate_content( loop = asyncio.get_event_loop() kwargs["agenerate_content"] = True + # Handle generationConfig parameter from kwargs for backward compatibility + if "generationConfig" in kwargs and config is None: + config = kwargs.pop("generationConfig") # get custom llm provider so we can use this for mapping exceptions if custom_llm_provider is None: _, custom_llm_provider, _, _ = litellm.get_llm_provider( @@ -288,6 +291,9 @@ def generate_content( try: _is_async = kwargs.pop("agenerate_content", False) is True + # Handle generationConfig parameter from kwargs for backward compatibility + if "generationConfig" in kwargs and config is None: + config = kwargs.pop("generationConfig") # Check for mock response first litellm_params = GenericLiteLLMParams(**kwargs) if litellm_params.mock_response and isinstance( @@ -374,6 +380,9 @@ async def agenerate_content_stream( try: kwargs["agenerate_content_stream"] = True + # Handle generationConfig parameter from kwargs for backward compatibility + if "generationConfig" in kwargs and config is None: + config = kwargs.pop("generationConfig") # get custom llm provider so we can use this for mapping exceptions if custom_llm_provider is None: _, custom_llm_provider, _, _ = litellm.get_llm_provider( @@ -461,6 +470,9 @@ def generate_content_stream( # Remove any async-related flags since this is the sync function _is_async = kwargs.pop("agenerate_content_stream", False) + # Handle generationConfig parameter from kwargs for backward compatibility + if "generationConfig" in kwargs and config is None: + config = kwargs.pop("generationConfig") # Setup the call setup_result = GenerateContentHelper.setup_generate_content_call( model=model, diff --git a/litellm/llms/gemini/count_tokens/handler.py b/litellm/llms/gemini/count_tokens/handler.py index bcc8ab9553d..a0782a838bd 100644 --- a/litellm/llms/gemini/count_tokens/handler.py +++ b/litellm/llms/gemini/count_tokens/handler.py @@ -136,4 +136,3 @@ class GoogleAIStudioTokenCounter: except Exception as e: error_msg = f"Unexpected error during token counting: {str(e)}" raise Exception(error_msg) from e - diff --git a/litellm/llms/vertex_ai/google_genai/transformation.py b/litellm/llms/vertex_ai/google_genai/transformation.py index 47933811196..574000e6bca 100644 --- a/litellm/llms/vertex_ai/google_genai/transformation.py +++ b/litellm/llms/vertex_ai/google_genai/transformation.py @@ -1,7 +1,7 @@ """ Transformation for Calling Google models in their native format. """ -from typing import Literal, Optional, Union +from typing import Dict, Literal, Optional, Union from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig from litellm.types.router import GenericLiteLLMParams @@ -11,20 +11,20 @@ class VertexAIGoogleGenAIConfig(GoogleGenAIConfig): """ Configuration for calling Google models in their native format. """ + HEADER_NAME = "Authorization" BEARER_PREFIX = "Bearer" - + @property def custom_llm_provider(self) -> Literal["gemini", "vertex_ai"]: return "vertex_ai" - def validate_environment( - self, + self, api_key: Optional[str], headers: Optional[dict], model: str, - litellm_params: Optional[Union[GenericLiteLLMParams, dict]] + litellm_params: Optional[Union[GenericLiteLLMParams, dict]], ) -> dict: default_headers = { "Content-Type": "application/json", @@ -36,4 +36,65 @@ class VertexAIGoogleGenAIConfig(GoogleGenAIConfig): default_headers.update(headers) return default_headers - \ No newline at end of file + + def _camel_to_snake(self, camel_str: str) -> str: + """Convert camelCase to snake_case""" + import re + + return re.sub(r"(? dict: + """ + Transform the generate content request for Vertex AI. + Since Vertex AI natively supports Google GenAI format, we can pass most fields directly. + """ + # Build the request in Google GenAI format that Vertex AI expects + result = { + "model": model, + "contents": contents, + } + + # Add tools if provided + if tools: + result["tools"] = tools + + # Add systemInstruction if provided + if system_instruction: + result["systemInstruction"] = system_instruction + + # Handle generationConfig - Vertex AI expects it in the same format + if generate_content_config_dict: + result["generationConfig"] = generate_content_config_dict + + return result diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index eb48f7de2dc..cdeea0094a6 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -85,6 +85,7 @@ async def route_request( """ team_id = get_team_id_from_data(data) router_model_names = llm_router.model_names if llm_router is not None else [] + if "api_key" in data or "api_base" in data: if llm_router is not None: return getattr(llm_router, f"{route_type}")(**data) @@ -123,24 +124,20 @@ async def route_request( data["model"] in router_model_names or data["model"] in llm_router.get_model_ids() ): - return getattr(llm_router, f"{route_type}")(**data) elif ( llm_router.model_group_alias is not None and data["model"] in llm_router.model_group_alias ): - return getattr(llm_router, f"{route_type}")(**data) elif data["model"] in llm_router.deployment_names: - return getattr(llm_router, f"{route_type}")( **data, specific_deployment=True ) elif data["model"] not in router_model_names: - if llm_router.router_general_settings.pass_through_all_models: return getattr(litellm, f"{route_type}")(**data) elif ( @@ -162,7 +159,6 @@ async def route_request( elif user_model is not None: return getattr(litellm, f"{route_type}")(**data) elif route_type == "allm_passthrough_route": - return getattr(litellm, f"{route_type}")(**data) # if no route found then it's a bad request diff --git a/tests/proxy_unit_tests/test_google_gemini_proxy_request.py b/tests/proxy_unit_tests/test_google_gemini_proxy_request.py new file mode 100644 index 00000000000..a8c67419402 --- /dev/null +++ b/tests/proxy_unit_tests/test_google_gemini_proxy_request.py @@ -0,0 +1,350 @@ +""" +Test case for Google Gemini API proxy request handling. + +This test verifies that when a request comes to the proxy endpoint: +http://localhost:4000/v1beta/models/gemini-2.5-flash:generateContent + +The request payload is correctly processed and forwarded to the httpx client. +""" + +import json +import os +import sys +import unittest.mock +from typing import Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +# Add the parent directory to the system path +sys.path.insert(0, os.path.abspath("../..")) + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.google_endpoints.endpoints import google_generate_content +from litellm.proxy.proxy_server import ProxyConfig +from litellm.proxy.utils import ProxyLogging +from fastapi import Request, Response +from fastapi.datastructures import Headers + + +@pytest.fixture +def sample_request_payload(): + """Sample request payload as provided in the user query.""" + return { + "contents": [ + { + "parts": [ + { + "text": "You are an interactive CLI agent specializing in software engineering tasks. Your primary goal is to help users safely and efficiently, adhering strictly to the following instructions and utilizing your available tools" + } + ], + "role": "user" + }, + { + "parts": [{"text": "Got it. Thanks for the context!"}], + "role": "model" + }, + { + "parts": [{"text": "Hello how are you"}], + "role": "user" + }, + { + "parts": [{"text": "I'm doing well, thank you! How can I help you today?\n"}], + "role": "model" + }, + { + "parts": [ + { + "text": "Analyze *only* the content and structure of your immediately preceding response (your last turn in the conversation history)." + } + ], + "role": "user" + } + ], + "systemInstruction": { + "parts": [ + { + "text": "You are an interactive CLI agent specializing in software engineering tasks. Your primary goal is to help users safely and efficiently, adhering strictly to the following instructions and utilizing your available tools" + } + ], + "role": "user" + }, + "generationConfig": { + "temperature": 0, + "topP": 1, + "responseMimeType": "application/json", + "responseJsonSchema": { + "type": "object", + "properties": { + "reasoning": { + "type": "string", + "description": "Brief explanation justifying the 'next_speaker' choice based *strictly* on the applicable rule and the content/structure of the preceding turn." + }, + "next_speaker": { + "type": "string", + "enum": ["user", "model"], + "description": "Who should speak next based *only* on the preceding turn and the decision rules" + } + }, + "required": ["reasoning", "next_speaker"] + } + } + } + + +@pytest.fixture +def mock_user_api_key_dict(): + """Mock user API key dictionary.""" + return UserAPIKeyAuth( + api_key="test_api_key", + user_id="test_user_id", + user_email="test@example.com", + team_id="test_team_id", + max_budget=100.0, + spend=0.0, + user_role="internal_user", + allowed_cache_controls=[], + metadata={}, + tpm_limit=None, + rpm_limit=None, + ) + + +@pytest.fixture +def mock_request(sample_request_payload): + """Create a mock FastAPI request with the sample payload.""" + mock_request = MagicMock(spec=Request) + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.method = "POST" + mock_request.url.path = "/v1beta/models/gemini-2.5-flash:generateContent" + + # Mock the request body reading + async def mock_body(): + return json.dumps(sample_request_payload).encode('utf-8') + + mock_request.body = mock_body + return mock_request + + +@pytest.fixture +def mock_response(): + """Create a mock FastAPI response.""" + return MagicMock(spec=Response) + + +@pytest.mark.asyncio +async def test_google_gemini_httpx_request_direct(): + """ + Test that the Google Gemini generate_content_handler correctly processes the request + and forwards it to the httpx client with the correct parameters. + + This test directly calls the HTTP handler to verify the httpx integration. + """ + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + # Sample request payload + sample_payload = { + "contents": [ + { + "parts": [ + { + "text": "You are an interactive CLI agent specializing in software engineering tasks." + } + ], + "role": "user" + }, + { + "parts": [{"text": "Got it. Thanks for the context!"}], + "role": "model" + }, + { + "parts": [{"text": "Hello how are you"}], + "role": "user" + } + ], + "systemInstruction": { + "parts": [ + { + "text": "You are an interactive CLI agent specializing in software engineering tasks." + } + ], + "role": "user" + }, + "config": { # Note: already transformed from generationConfig + "temperature": 0, + "topP": 1, + "responseMimeType": "application/json", + "responseJsonSchema": { + "type": "object", + "properties": { + "reasoning": {"type": "string"}, + "next_speaker": {"type": "string", "enum": ["user", "model"]} + }, + "required": ["reasoning", "next_speaker"] + } + } + } + + # Mock the HTTP handler to capture the request + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") as mock_post: + # Create mock response + mock_http_response = MagicMock() + mock_http_response.status_code = 200 + mock_http_response.json.return_value = { + "candidates": [ + { + "content": { + "parts": [ + { + "text": '{"reasoning": "The preceding response was a helpful greeting asking how to assist.", "next_speaker": "user"}' + } + ], + "role": "model" + } + } + ] + } + mock_post.return_value = mock_http_response + + # Create the HTTP handler and provider config + from litellm.types.router import GenericLiteLLMParams + + http_handler = BaseLLMHTTPHandler() + provider_config = GoogleGenAIConfig() + + # Create proper litellm params + litellm_params = GenericLiteLLMParams( + api_base="https://generativelanguage.googleapis.com", + api_key="test_api_key" + ) + + logging_obj = LiteLLMLoggingObj( + model="gemini/gemini-2.5-flash", + messages=[], + stream=False, + call_type="agenerate_content", + start_time=None, + litellm_call_id="test_call_id", + function_id="test_function_id" + ) + + try: + # Call the generate_content_handler directly + response = http_handler.generate_content_handler( + model="gemini/gemini-2.5-flash", + contents=sample_payload["contents"], + generate_content_provider_config=provider_config, + generate_content_config_dict=sample_payload["config"], + tools=None, + custom_llm_provider="gemini", + litellm_params=litellm_params, + logging_obj=logging_obj, + extra_headers=None, + extra_body=None, + timeout=30.0, + _is_async=False, + client=None, + stream=False, + litellm_metadata={} + ) + + # Verify that the HTTP post was called + assert mock_post.called, "Expected HTTP POST to be called" + + # Get the call arguments + call_args, call_kwargs = mock_post.call_args + + print(f"POST call args: {call_args}") + print(f"POST call kwargs: {call_kwargs}") + + # Validate that the request data includes the expected fields + request_data = call_kwargs.get('json') + if request_data: + assert 'contents' in request_data, "Expected 'contents' in request data" + + # The config should be included in the request as generationConfig + if 'generationConfig' in request_data: + config = request_data['generationConfig'] + assert config['temperature'] == 0, "Expected temperature to be 0" + assert config['topP'] == 1, "Expected topP to be 1" + assert config['responseMimeType'] == "application/json", "Expected responseMimeType to be application/json" + assert 'responseJsonSchema' in config, "Expected responseJsonSchema in config" + + # Validate the responseJsonSchema structure + schema = config['responseJsonSchema'] + assert schema['type'] == 'object', "Expected schema type to be object" + assert 'properties' in schema, "Expected properties in schema" + assert 'reasoning' in schema['properties'], "Expected reasoning property in schema" + assert 'next_speaker' in schema['properties'], "Expected next_speaker property in schema" + + print("✅ Request data validation passed") + print(f"Request data: {json.dumps(request_data, indent=2)}") + + # Validate URL contains the correct endpoint + if call_args: + url = call_args[0] if len(call_args) > 0 else call_kwargs.get('url') + assert url is not None, "Expected URL to be provided" + print(f"✅ URL validation passed: {url}") + + except Exception as e: + print(f"Exception occurred: {e}") + + # Check if the HTTP handler was called despite the exception + if mock_post.called: + call_args, call_kwargs = mock_post.call_args + print(f"HTTP POST was called with args: {call_args}") + print(f"HTTP POST was called with kwargs: {call_kwargs}") + + # Even with an exception, we can validate the request structure + request_data = call_kwargs.get('json') + if request_data: + assert 'contents' in request_data, "Expected 'contents' in request data" + if 'generationConfig' in request_data: + config = request_data['generationConfig'] + assert config['temperature'] == 0, "Expected temperature to be 0" + assert config['responseMimeType'] == "application/json", "Expected responseMimeType to be application/json" + print("✅ Request structure validation passed despite exception") + else: + # If no HTTP call was made, re-raise the exception for debugging + raise + + +@pytest.mark.asyncio +async def test_generationconfig_to_config_mapping(sample_request_payload): + """ + Test that generationConfig is correctly mapped to config parameter + for Google GenAI compatibility in the main functions. + """ + from litellm.google_genai.main import agenerate_content + + # Create a copy of the payload to avoid modifying the fixture + test_data = sample_request_payload.copy() + + # Test that agenerate_content can handle generationConfig parameter + # This should not raise an error about parameter handling + try: + # This will fail due to missing API key, but should not fail due to parameter handling + await agenerate_content( + model="gemini/gemini-2.5-flash", + contents=test_data["contents"], + generationConfig=test_data["generationConfig"], # Pass as generationConfig + custom_llm_provider="gemini" + ) + except Exception as e: + # Should not fail due to parameter handling issues + error_msg = str(e).lower() + if "generationconfig" in error_msg or "config" in error_msg or "parameter" in error_msg: + pytest.fail(f"Parameter handling failed: {e}") + # Other errors (like API key missing) are expected + print(f"✅ Parameter handling worked (API error expected): {type(e).__name__}") + + print("✅ generationConfig to config mapping test passed") + + +if __name__ == "__main__": + # Run the tests + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/llms/gemini/test_gemini_common_utils.py b/tests/test_litellm/llms/gemini/test_gemini_common_utils.py index 34472b3856d..3b5b683d2d0 100644 --- a/tests/test_litellm/llms/gemini/test_gemini_common_utils.py +++ b/tests/test_litellm/llms/gemini/test_gemini_common_utils.py @@ -157,4 +157,4 @@ class TestGoogleAIStudioTokenCounter: mock_acount_tokens.assert_called_once_with( model=model_to_use, contents=contents - ) + ) \ No newline at end of file