mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-20 04:23:56 +00:00
fix: auth header for custom api base in generate Content request
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user