diff --git a/litellm/llms/gemini/google_genai/transformation.py b/litellm/llms/gemini/google_genai/transformation.py index 3474c8abe34..48046dd9dfa 100644 --- a/litellm/llms/gemini/google_genai/transformation.py +++ b/litellm/llms/gemini/google_genai/transformation.py @@ -89,6 +89,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): "audio_timestamp", "automatic_function_calling", "thinking_config", + "image_config", ] def map_generate_content_optional_params( diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 569634ee140..bba5b87024b 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -37,6 +37,11 @@ async def google_generate_content( if "model" not in data: data["model"] = model_name + # Extract generationConfig and pass it as config parameter + generation_config = data.pop("generationConfig", None) + if generation_config: + data["config"] = generation_config + # Add user authentication metadata for cost tracking data = await add_litellm_data_to_request( data=data, @@ -83,6 +88,11 @@ async def google_stream_generate_content( data["stream"] = True # enforce streaming for this endpoint + # Extract generationConfig and pass it as config parameter + generation_config = data.pop("generationConfig", None) + if generation_config: + data["config"] = generation_config + # Add user authentication metadata for cost tracking data = await add_litellm_data_to_request( data=data, 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 11e34cbbea4..26c60b2c303 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 @@ -306,6 +306,88 @@ def test_google_generate_content_with_system_instruction(): assert called_data["systemInstruction"] == system_instruction 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 + assert called_data["contents"][0]["role"] == "user" + + +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. + """ + try: + from fastapi import FastAPI + from fastapi.testclient import TestClient + + 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: + + 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): + 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"}] + }], + "generationConfig": { + "responseModalities": ["TEXT", "IMAGE"], + "imageConfig": { + "aspectRatio": "9:16", + "imageSize": "4K" + } + } + }, + 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