mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Fix: Preserve full path structure for Gemini custom api_base (#12215)
* Fix: Preserve full path structure for Gemini custom api_base (Fixes #11959) This fix addresses an issue where custom api_base URLs (like Cloudflare AI Gateway) were not working correctly with Google AI Studio (Gemini) models. The problem was that the _check_custom_proxy method was simply appending the endpoint to the custom base URL, resulting in malformed URLs like: https://gateway.ai.cloudflare.com/v1/my-id/my-gateway/google-ai-studio:generateContent Instead of the correct format: https://gateway.ai.cloudflare.com/v1/my-id/my-gateway/google-ai-studio/v1beta/models/gemini-2.5-flash:generateContent Changes: - Modified _check_custom_proxy to preserve the full path structure from the original URL - Extracts the path from the original Google AI URL and appends it to the custom base - Maintains backward compatibility for Vertex AI models (unchanged behavior) - Added comprehensive tests to verify the fix works correctly Fixes #11959 * Fix: Update test to match actual Gemini URL format and fix double colon issue - Fixed test expectation to include the full model path with 'gemini/' prefix - Fixed double colon issue in Vertex AI URL construction when using custom api_base - All tests now pass successfully
This commit is contained in:
Cole McIntosh
GitHub
parent
f6589d81ec
commit
f47254ecab
2 changed files with 151 additions and 8 deletions
|
|
@ -83,7 +83,11 @@ class VertexBase:
|
|||
if "type" in json_obj and json_obj["type"] == "external_account":
|
||||
# If environment_id key contains "aws" value it corresponds to an AWS config file
|
||||
credential_source = json_obj.get("credential_source", {})
|
||||
environment_id = credential_source.get("environment_id", "") if isinstance(credential_source, dict) else ""
|
||||
environment_id = (
|
||||
credential_source.get("environment_id", "")
|
||||
if isinstance(credential_source, dict)
|
||||
else ""
|
||||
)
|
||||
if isinstance(environment_id, str) and "aws" in environment_id:
|
||||
creds = self._credentials_from_identity_pool_with_aws(json_obj)
|
||||
else:
|
||||
|
|
@ -130,7 +134,7 @@ class VertexBase:
|
|||
from google.auth import identity_pool
|
||||
|
||||
return identity_pool.Credentials.from_info(json_obj)
|
||||
|
||||
|
||||
def _credentials_from_identity_pool_with_aws(self, json_obj):
|
||||
from google.auth import aws
|
||||
|
||||
|
|
@ -297,7 +301,21 @@ class VertexBase:
|
|||
"""
|
||||
if api_base:
|
||||
if custom_llm_provider == "gemini":
|
||||
url = "{}:{}".format(api_base, endpoint)
|
||||
# Extract the path from the original URL to preserve the model structure
|
||||
# Original URL format: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent?key=...
|
||||
# We need to extract: /v1beta/models/gemini-2.5-flash:generateContent
|
||||
from urllib.parse import urlparse
|
||||
|
||||
parsed_original = urlparse(url)
|
||||
path_and_query = parsed_original.path
|
||||
|
||||
# Remove any query parameters from the path (they'll be re-added if needed)
|
||||
if "?" in path_and_query:
|
||||
path_and_query = path_and_query.split("?")[0]
|
||||
|
||||
# Construct the new URL with the custom base and the original path
|
||||
url = "{}{}".format(api_base.rstrip("/"), path_and_query)
|
||||
|
||||
if gemini_api_key is None:
|
||||
raise ValueError(
|
||||
"Missing gemini_api_key, please set `GEMINI_API_KEY`"
|
||||
|
|
@ -306,7 +324,7 @@ class VertexBase:
|
|||
gemini_api_key # cloudflare expects api key as bearer token
|
||||
)
|
||||
else:
|
||||
url = "{}:{}".format(api_base, endpoint)
|
||||
url = "{}{}".format(api_base, endpoint)
|
||||
|
||||
if stream is True:
|
||||
url = url + "?alt=sse"
|
||||
|
|
@ -341,7 +359,7 @@ class VertexBase:
|
|||
stream=stream,
|
||||
gemini_api_key=gemini_api_key,
|
||||
)
|
||||
auth_header = None # this field is not used for gemin
|
||||
auth_header = None # this field is not used for gemini
|
||||
else:
|
||||
vertex_location = self.get_vertex_region(
|
||||
vertex_region=vertex_location,
|
||||
|
|
@ -490,7 +508,7 @@ class VertexBase:
|
|||
headers.update(extra_headers)
|
||||
|
||||
return headers
|
||||
|
||||
|
||||
@staticmethod
|
||||
def get_vertex_ai_project(litellm_params: dict) -> Optional[str]:
|
||||
return (
|
||||
|
|
@ -499,7 +517,7 @@ class VertexBase:
|
|||
or litellm.vertex_project
|
||||
or get_secret_str("VERTEXAI_PROJECT")
|
||||
)
|
||||
|
||||
|
||||
@staticmethod
|
||||
def get_vertex_ai_credentials(litellm_params: dict) -> Optional[str]:
|
||||
return (
|
||||
|
|
@ -507,7 +525,7 @@ class VertexBase:
|
|||
or litellm_params.pop("vertex_ai_credentials", None)
|
||||
or get_secret_str("VERTEXAI_CREDENTIALS")
|
||||
)
|
||||
|
||||
|
||||
@staticmethod
|
||||
def get_vertex_ai_location(litellm_params: dict) -> Optional[str]:
|
||||
return (
|
||||
|
|
|
|||
125
tests/test_litellm/llms/vertex_ai/test_custom_api_base_fix.py
Normal file
125
tests/test_litellm/llms/vertex_ai/test_custom_api_base_fix.py
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
import os
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
import litellm
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
|
||||
|
||||
class TestCustomApiBaseFix:
|
||||
"""Test that custom api_base works correctly for both Gemini and Vertex AI models"""
|
||||
|
||||
def test_gemini_custom_api_base_preserves_path(self):
|
||||
"""Test that Gemini custom api_base preserves the full path structure"""
|
||||
vertex_base = VertexBase()
|
||||
|
||||
# Test case from the issue: Cloudflare AI Gateway
|
||||
cloudflare_base = "https://gateway.ai.cloudflare.com/v1/my-id/my-gateway/google-ai-studio"
|
||||
original_url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent?key=test-key"
|
||||
|
||||
token, final_url = vertex_base._check_custom_proxy(
|
||||
api_base=cloudflare_base,
|
||||
auth_header=None,
|
||||
custom_llm_provider="gemini",
|
||||
gemini_api_key="test-api-key",
|
||||
endpoint=":generateContent",
|
||||
stream=False,
|
||||
url=original_url,
|
||||
)
|
||||
|
||||
# Should preserve the full path structure
|
||||
expected_url = "https://gateway.ai.cloudflare.com/v1/my-id/my-gateway/google-ai-studio/v1beta/models/gemini-2.5-flash:generateContent"
|
||||
assert final_url == expected_url
|
||||
assert token == "test-api-key"
|
||||
|
||||
def test_gemini_custom_api_base_with_streaming(self):
|
||||
"""Test that streaming adds the alt=sse parameter correctly"""
|
||||
vertex_base = VertexBase()
|
||||
|
||||
cloudflare_base = "https://gateway.ai.cloudflare.com/v1/my-id/my-gateway/google-ai-studio"
|
||||
original_url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:streamGenerateContent?key=test-key&alt=sse"
|
||||
|
||||
token, final_url = vertex_base._check_custom_proxy(
|
||||
api_base=cloudflare_base,
|
||||
auth_header=None,
|
||||
custom_llm_provider="gemini",
|
||||
gemini_api_key="test-api-key",
|
||||
endpoint=":streamGenerateContent",
|
||||
stream=True,
|
||||
url=original_url,
|
||||
)
|
||||
|
||||
# Should preserve path and add streaming parameter
|
||||
expected_url = "https://gateway.ai.cloudflare.com/v1/my-id/my-gateway/google-ai-studio/v1beta/models/gemini-pro:streamGenerateContent?alt=sse"
|
||||
assert final_url == expected_url
|
||||
assert token == "test-api-key"
|
||||
|
||||
def test_vertex_ai_custom_api_base_unchanged(self):
|
||||
"""Test that Vertex AI behavior remains unchanged"""
|
||||
vertex_base = VertexBase()
|
||||
|
||||
# Vertex AI uses a different URL structure
|
||||
cloudflare_base = "https://gateway.ai.cloudflare.com/v1/my-id/my-gateway/google-vertex-ai/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-2.0-flash"
|
||||
original_url = "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-2.0-flash:generateContent"
|
||||
|
||||
token, final_url = vertex_base._check_custom_proxy(
|
||||
api_base=cloudflare_base,
|
||||
auth_header="Bearer mock-token",
|
||||
custom_llm_provider="vertex_ai",
|
||||
gemini_api_key=None,
|
||||
endpoint=":generateContent",
|
||||
stream=False,
|
||||
url=original_url,
|
||||
)
|
||||
|
||||
# Vertex AI should still use the old behavior (appending endpoint)
|
||||
expected_url = "https://gateway.ai.cloudflare.com/v1/my-id/my-gateway/google-vertex-ai/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-2.0-flash:generateContent"
|
||||
assert final_url == expected_url
|
||||
assert token == "Bearer mock-token" # Auth header unchanged for Vertex AI
|
||||
|
||||
def test_full_flow_with_get_token_and_url(self):
|
||||
"""Test the full flow through _get_token_and_url"""
|
||||
vertex_base = VertexBase()
|
||||
|
||||
cloudflare_base = "https://gateway.ai.cloudflare.com/v1/my-id/my-gateway/google-ai-studio"
|
||||
|
||||
token, url = vertex_base._get_token_and_url(
|
||||
model="gemini/gemini-2.5-flash",
|
||||
auth_header=None,
|
||||
gemini_api_key="test-api-key",
|
||||
vertex_project=None,
|
||||
vertex_location=None,
|
||||
vertex_credentials=None,
|
||||
stream=False,
|
||||
custom_llm_provider="gemini",
|
||||
api_base=cloudflare_base,
|
||||
)
|
||||
|
||||
# Should produce the correct URL with custom base
|
||||
assert cloudflare_base in url
|
||||
assert "/v1beta/models/gemini/gemini-2.5-flash:generateContent" in url
|
||||
assert token == "test-api-key"
|
||||
|
||||
def test_missing_gemini_api_key_raises_error(self):
|
||||
"""Test that missing Gemini API key raises an error"""
|
||||
vertex_base = VertexBase()
|
||||
|
||||
cloudflare_base = "https://gateway.ai.cloudflare.com/v1/my-id/my-gateway/google-ai-studio"
|
||||
original_url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent"
|
||||
|
||||
with pytest.raises(ValueError, match="Missing gemini_api_key"):
|
||||
vertex_base._check_custom_proxy(
|
||||
api_base=cloudflare_base,
|
||||
auth_header=None,
|
||||
custom_llm_provider="gemini",
|
||||
gemini_api_key=None, # Missing API key
|
||||
endpoint=":generateContent",
|
||||
stream=False,
|
||||
url=original_url,
|
||||
)
|
||||
Loading…
Add table
Reference in a new issue