From c45fad3855847715afaebba71926f8e84eb7b355 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 5 Nov 2025 20:42:13 +0530 Subject: [PATCH] Fix: Send Gemini API key via x-goog-api-key header with custom api_base (#16085) * Add gemini api key in the custom api url * Update tests * Use api key n the header * Use api key n the header * fix mypy error * fix mypy error * fix test gemini auth --- .../vertex_and_google_ai_studio_gemini.py | 6 +- litellm/llms/vertex_ai/vertex_llm_base.py | 5 +- .../test_google_gemini_proxy_request.py | 10 ++- .../llms/vertex_ai/test_vertex_llm_base.py | 63 ++++++++++++++++++- 4 files changed, 74 insertions(+), 10 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 849d27d1ea0..b8370d5fef2 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1743,13 +1743,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): messages: List[AllMessageValues], optional_params: Dict, litellm_params: Dict, - api_key: Optional[str] = None, + api_key: Optional[Union[str, Dict]] = None, api_base: Optional[str] = None, ) -> Dict: default_headers = { "Content-Type": "application/json", } - if api_key is not None: + if isinstance(api_key, dict): + default_headers.update(api_key) + elif api_key is not None: default_headers["Authorization"] = f"Bearer {api_key}" if headers is not None: default_headers.update(headers) diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 54608991595..9ddbc461a70 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -308,9 +308,8 @@ class VertexBase: raise ValueError( "Missing gemini_api_key, please set `GEMINI_API_KEY`" ) - auth_header = ( - gemini_api_key # cloudflare expects api key as bearer token - ) + if gemini_api_key is not None: + auth_header = {"x-goog-api-key": gemini_api_key} # type: ignore[assignment] else: url = "{}:{}".format(api_base, endpoint) diff --git a/tests/proxy_unit_tests/test_google_gemini_proxy_request.py b/tests/proxy_unit_tests/test_google_gemini_proxy_request.py index 9d16246aa8f..90c2cac18d0 100644 --- a/tests/proxy_unit_tests/test_google_gemini_proxy_request.py +++ b/tests/proxy_unit_tests/test_google_gemini_proxy_request.py @@ -378,8 +378,8 @@ async def test_gemini_custom_api_base_proxy_integration(): expected_url = f"{custom_api_base}/models/{model}:{endpoint}" assert result_url == expected_url, f"Expected {expected_url}, got {result_url}" - # Verify the auth header is set to the API key - assert auth_header == "test-api-key", f"Expected 'test-api-key', got {auth_header}" + # Verify the auth header is set to the API key as a dictionary + assert auth_header == {"x-goog-api-key": "test-api-key"}, f"Expected {{'x-goog-api-key': 'test-api-key'}}, got {auth_header}" print(f"✅ Custom API base URL construction test passed: {result_url}") @@ -399,6 +399,9 @@ async def test_gemini_custom_api_base_proxy_integration(): expected_streaming_url = f"{custom_api_base}/models/{model}:{endpoint}?alt=sse" assert result_url_streaming == expected_streaming_url, f"Expected {expected_streaming_url}, got {result_url_streaming}" + # Verify the auth header is also set correctly for streaming + assert auth_header_streaming == {"x-goog-api-key": "test-api-key"}, f"Expected {{'x-goog-api-key': 'test-api-key'}}, got {auth_header_streaming}" + print(f"✅ Custom API base streaming URL test passed: {result_url_streaming}") # Test case 3: Error handling - missing API key @@ -462,7 +465,8 @@ async def test_gemini_proxy_config_with_custom_api_base(): expected_url = f"{model_config['litellm_params']['api_base']}/models/{model}:generateContent" assert result_url == expected_url, f"Expected {expected_url}, got {result_url} for model {model}" - assert auth_header == model_config["litellm_params"]["api_key"], f"Expected API key, got {auth_header} for model {model}" + expected_auth_header = {"x-goog-api-key": model_config["litellm_params"]["api_key"]} + assert auth_header == expected_auth_header, f"Expected {expected_auth_header}, got {auth_header} for model {model}" print(f"✅ Model {model} configuration test passed: {result_url}") diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py index c85f6070084..129534ffc7a 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py @@ -718,7 +718,7 @@ class TestVertexBase: None, "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", "gemini-2.5-flash-lite", - "test-api-key", + {"x-goog-api-key": "test-api-key"}, "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent" ), # Test case 2: Gemini with custom API base and streaming @@ -731,7 +731,7 @@ class TestVertexBase: None, "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", "gemini-2.5-flash-lite", - "test-api-key", + {"x-goog-api-key": "test-api-key"}, "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent?alt=sse" ), # Test case 3: Non-Gemini provider with custom API base @@ -878,3 +878,62 @@ class TestVertexBase: expected_no_streaming_url = "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent" assert result_url_no_streaming == expected_no_streaming_url, f"Expected {expected_no_streaming_url}, got {result_url_no_streaming}" + + @pytest.mark.parametrize( + "api_base, custom_llm_provider, gemini_api_key, endpoint, stream, auth_header, url, model, expected_auth_header, expected_url", + [ + # Test case 1: Gemini with custom API base + ( + "https://proxy.example.com/generativelanguage.googleapis.com/v1beta", + "gemini", + "test-api-key", + "generateContent", + False, + None, + "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", + "gemini-2.5-flash-lite", + {"x-goog-api-key": "test-api-key"}, + "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent" + ), + # Test case 2: Gemini with custom API base and streaming + ( + "https://proxy.example.com/generativelanguage.googleapis.com/v1beta", + "gemini", + "test-api-key", + "generateContent", + True, + None, + "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", + "gemini-2.5-flash-lite", + {"x-goog-api-key": "test-api-key"}, + "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent?alt=sse" + ), + ], + ) + def test_check_custom_proxy_minimal_gemini_key_param( + self, + api_base, + custom_llm_provider, + gemini_api_key, + endpoint, + stream, + auth_header, + url, + model, + expected_auth_header, + expected_url, + ): + """Single focused test to ensure ?key is appended (and &alt=sse for streaming).""" + vertex_base = VertexBase() + result_auth_header, result_url = vertex_base._check_custom_proxy( + api_base=api_base, + custom_llm_provider=custom_llm_provider, + gemini_api_key=gemini_api_key, + endpoint=endpoint, + stream=stream, + auth_header=auth_header, + url=url, + model=model, + ) + assert result_auth_header == expected_auth_header + assert result_url == expected_url