From 4878bc6275b917a094dd7618b8a39f08f10c5cc0 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 29 Jul 2025 13:46:43 -0700 Subject: [PATCH] [Bug Fix] Gemini-CLI - The Gemini Custom API request has an incorrect authorization format (#13098) * fix GoogleGenAIConfig * fix validate_environment * test_agenerate_content_x_goog_api_key_header --- .../gemini/google_genai/transformation.py | 11 ++- litellm/proxy/proxy_config.yaml | 18 +---- .../google_genai/test_google_genai_adapter.py | 77 ++++++++++++++++++- 3 files changed, 87 insertions(+), 19 deletions(-) diff --git a/litellm/llms/gemini/google_genai/transformation.py b/litellm/llms/gemini/google_genai/transformation.py index 9910e478063..0de6ef9043d 100644 --- a/litellm/llms/gemini/google_genai/transformation.py +++ b/litellm/llms/gemini/google_genai/transformation.py @@ -30,6 +30,12 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): """ Configuration for calling Google models in their native format. """ + ############################## + # Constants + ############################## + XGOOGLE_API_KEY = "x-goog-api-key" + ############################## + @property def custom_llm_provider(self) -> Literal["gemini", "vertex_ai"]: return "gemini" @@ -113,8 +119,9 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): default_headers = { "Content-Type": "application/json", } - if api_key is not None: - default_headers["Authorization"] = f"Bearer {api_key}" + gemini_api_key = self._get_google_ai_studio_api_key(dict(litellm_params or {})) + if gemini_api_key is not None: + default_headers[self.XGOOGLE_API_KEY] = gemini_api_key if headers is not None: default_headers.update(headers) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index f1cfe75c7ac..9ee8e675c38 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,19 +1,5 @@ model_list: - - model_name: anthropic/* + - model_name: vertex_ai/* litellm_params: - model: anthropic/* - - model_name: openai/* - litellm_params: - model: openai/* - - - model_name: vertex_ai/gemini-2.5-pro # Vertex AI Gemini - litellm_params: - model: vertex_ai/gemini-2.5-pro - thinking: {"type": "enabled", "budget_tokens": 1024} - merge_reasoning_content_in_choices: true - - -litellm_settings: - callbacks: ["datadog_llm_observability"] - cache: true + model: gemini/* diff --git a/tests/test_litellm/google_genai/test_google_genai_adapter.py b/tests/test_litellm/google_genai/test_google_genai_adapter.py index bf297758888..69ab677e86a 100644 --- a/tests/test_litellm/google_genai/test_google_genai_adapter.py +++ b/tests/test_litellm/google_genai/test_google_genai_adapter.py @@ -1125,4 +1125,79 @@ async def test_google_generate_content_with_openai(): passed_fields = set(call_kwargs.keys()) # remove any GenericLiteLLMParams fields passed_fields = passed_fields - set(GenericLiteLLMParams.model_fields.keys()) - assert passed_fields == set(["model", "messages"]), f"Expected only model, contents, systemInstruction, and safetySettings to be passed through, got {passed_fields}" \ No newline at end of file + assert passed_fields == set(["model", "messages"]), f"Expected only model, contents, systemInstruction, and safetySettings to be passed through, got {passed_fields}" + +@pytest.mark.asyncio +async def test_agenerate_content_x_goog_api_key_header(): + """ + Test that agenerate_content passes x-goog-api-key header correctly. + + This test verifies that when calling agenerate_content with a Google GenAI model, + the HTTP request includes the x-goog-api-key header with the correct API key value. + """ + import os + import unittest.mock + + import httpx + + test_api_key = "test-gemini-api-key-123" + + # Mock environment to ensure we use our test API key + with unittest.mock.patch.dict(os.environ, {"GEMINI_API_KEY": test_api_key}, clear=False): + # Mock the AsyncHTTPHandler's post method to capture headers + with unittest.mock.patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", new_callable=unittest.mock.AsyncMock) as mock_post: + # Mock a successful response + mock_response = unittest.mock.MagicMock() + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "parts": [{"text": "Hello! How can I help you today?"}], + "role": "model" + }, + "finishReason": "STOP", + "index": 0 + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 10, + "totalTokenCount": 15 + } + } + mock_response.status_code = 200 + mock_response.headers = {} + mock_post.return_value = mock_response + + # Call agenerate_content with Google AI Studio model + try: + response = await agenerate_content( + model="gemini/gemini-1.5-flash", + contents=[ + {"role": "user", "parts": [{"text": "Hello, world!"}]} + ], + api_key=test_api_key + ) + except Exception: + # Ignore any response processing errors, we just want to check the headers + pass + + # Verify that AsyncHTTPHandler.post was called + mock_post.assert_called_once() + + # Get the arguments passed to the post call + call_args, call_kwargs = mock_post.call_args + + # Verify that headers contain x-goog-api-key + headers = call_kwargs.get("headers", {}) + assert "x-goog-api-key" in headers, f"x-goog-api-key header not found in headers: {list(headers.keys())}" + + # Verify the API key is set (could be our test key or from api_key parameter) + api_key_value = headers["x-goog-api-key"] + assert api_key_value == test_api_key, f"Expected x-goog-api-key to be {test_api_key}, got {api_key_value}" + + # Verify other expected headers + assert headers.get("Content-Type") == "application/json", f"Expected Content-Type application/json, got {headers.get('Content-Type')}" + + print(f"✓ Test passed: x-goog-api-key header correctly set to {api_key_value}") + print(f"✓ All headers: {list(headers.keys())}") \ No newline at end of file