fix: auth header for custom api base in generate Content request

This commit is contained in:
Sameer Kankute
2026-01-05 12:25:13 +05:30
parent fcabc059ca
commit 0730a74dda
2 changed files with 84 additions and 1 deletions
@@ -153,7 +153,9 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM):
gemini_api_key = api_key or self._get_google_ai_studio_api_key(
dict(litellm_params or {})
)
if gemini_api_key is not None:
if isinstance(gemini_api_key, dict):
default_headers.update(gemini_api_key)
elif gemini_api_key is not None:
default_headers[self.XGOOGLE_API_KEY] = gemini_api_key
if headers is not None:
default_headers.update(headers)
@@ -366,3 +366,84 @@ def test_transform_generate_content_request_system_instruction_with_tools():
assert "tools" in result, "tools should be in request body"
assert result["tools"] == tools
assert result["model"] == "gemini-3-flash-preview"
def test_validate_environment_with_dict_api_key():
"""
Test that validate_environment correctly handles api_key as a dict.
This happens when using custom api_base with Gemini - the auth_header
is returned as {"x-goog-api-key": "sk-test"} and should be merged into
headers instead of being set as a header value.
Regression test for: https://github.com/BerriAI/litellm/issues/xxxxx
"""
config = GoogleGenAIConfig()
# Simulate the case where auth_header is a dict (custom api_base scenario)
auth_header_dict = {"x-goog-api-key": "sk-test-key-123"}
result = config.validate_environment(
api_key=auth_header_dict,
headers=None,
model="gemini-2.5-pro",
litellm_params={}
)
# The dict should be merged into headers, not set as a value
assert "x-goog-api-key" in result, "x-goog-api-key should be in headers"
assert result["x-goog-api-key"] == "sk-test-key-123", "API key should be the string value, not a dict"
assert isinstance(result["x-goog-api-key"], str), "Header value should be a string, not a dict"
assert "Content-Type" in result, "Content-Type should be in headers"
assert result["Content-Type"] == "application/json"
def test_validate_environment_with_string_api_key():
"""
Test that validate_environment correctly handles api_key as a string.
This is the normal case when using standard Gemini API.
"""
config = GoogleGenAIConfig()
# Normal case: api_key is a string
api_key_string = "sk-test-key-456"
result = config.validate_environment(
api_key=api_key_string,
headers=None,
model="gemini-2.5-pro",
litellm_params={}
)
# The string should be set as the header value
assert "x-goog-api-key" in result, "x-goog-api-key should be in headers"
assert result["x-goog-api-key"] == "sk-test-key-456", "API key should match input"
assert isinstance(result["x-goog-api-key"], str), "Header value should be a string"
assert "Content-Type" in result, "Content-Type should be in headers"
def test_validate_environment_with_extra_headers():
"""
Test that validate_environment correctly merges extra headers with dict api_key.
"""
config = GoogleGenAIConfig()
# Custom api_base scenario with additional headers
auth_header_dict = {"x-goog-api-key": "sk-test-key-789"}
extra_headers = {"X-Custom-Header": "custom-value"}
result = config.validate_environment(
api_key=auth_header_dict,
headers=extra_headers,
model="gemini-2.5-pro",
litellm_params={}
)
# Both the auth dict and extra headers should be merged
assert "x-goog-api-key" in result, "x-goog-api-key should be in headers"
assert result["x-goog-api-key"] == "sk-test-key-789", "API key should be correctly set"
assert isinstance(result["x-goog-api-key"], str), "Header value should be a string"
assert "X-Custom-Header" in result, "Extra headers should be merged"
assert result["X-Custom-Header"] == "custom-value"
assert "Content-Type" in result