From ff2d19e4cab88c7ce0a417f47e5ff7be85da515a Mon Sep 17 00:00:00 2001 From: Zero Clover <13190004+ZeroClover@users.noreply.github.com> Date: Tue, 30 Sep 2025 03:12:32 +0800 Subject: [PATCH] feat: improve vertex AI/gemini api_base handling for proxy services (#15039) --- litellm/llms/vertex_ai/vertex_llm_base.py | 165 ++++++++-------- .../llms/vertex_ai/test_vertex_llm_base.py | 184 ++++++++++-------- 2 files changed, 193 insertions(+), 156 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 6d194d41ad..6769bc3fb2 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -46,14 +46,10 @@ class VertexBase: return "global" return vertex_region or "us-central1" - def load_auth( - self, credentials: Optional[VERTEX_CREDENTIALS_TYPES], project_id: Optional[str] - ) -> Tuple[Any, str]: + def load_auth(self, credentials: Optional[VERTEX_CREDENTIALS_TYPES], project_id: Optional[str]) -> Tuple[Any, str]: if credentials is not None: if isinstance(credentials, str): - verbose_logger.debug( - "Vertex: Loading vertex credentials from %s", credentials - ) + verbose_logger.debug("Vertex: Loading vertex credentials from %s", credentials) verbose_logger.debug( "Vertex: checking if credentials is a valid path, os.path.exists(%s)=%s, current dir %s", credentials, @@ -67,26 +63,18 @@ class VertexBase: else: json_obj = json.loads(credentials) except Exception: - raise Exception( - "Unable to load vertex credentials from environment. Got={}".format( - credentials - ) - ) + raise Exception("Unable to load vertex credentials from environment. Got={}".format(credentials)) elif isinstance(credentials, dict): json_obj = credentials else: - raise ValueError( - "Invalid credentials type: {}".format(type(credentials)) - ) + raise ValueError("Invalid credentials type: {}".format(type(credentials))) # Check if the JSON object contains Workload Identity Federation configuration 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 "" + 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) @@ -123,9 +111,7 @@ class VertexBase: raise ValueError("Could not resolve project_id") if not isinstance(project_id, str): - raise TypeError( - f"Expected project_id to be a str but got {type(project_id)}" - ) + raise TypeError(f"Expected project_id to be a str but got {type(project_id)}") return creds, project_id @@ -143,16 +129,12 @@ class VertexBase: def _credentials_from_authorized_user(self, json_obj, scopes): import google.oauth2.credentials - return google.oauth2.credentials.Credentials.from_authorized_user_info( - json_obj, scopes=scopes - ) + return google.oauth2.credentials.Credentials.from_authorized_user_info(json_obj, scopes=scopes) def _credentials_from_service_account(self, json_obj, scopes): import google.oauth2.service_account - return google.oauth2.service_account.Credentials.from_service_account_info( - json_obj, scopes=scopes - ) + return google.oauth2.service_account.Credentials.from_service_account_info(json_obj, scopes=scopes) def _credentials_from_default_auth(self, scopes): import google.auth as google_auth @@ -162,9 +144,7 @@ class VertexBase: def get_default_vertex_location(self) -> str: return "us-central1" - def get_api_base( - self, api_base: Optional[str], vertex_location: Optional[str] - ) -> str: + def get_api_base(self, api_base: Optional[str], vertex_location: Optional[str]) -> str: if api_base: return api_base elif vertex_location == "global": @@ -214,9 +194,7 @@ class VertexBase: stream: Optional[bool], model: str, ) -> str: - api_base = self.get_api_base( - api_base=custom_api_base, vertex_location=vertex_location - ) + api_base = self.get_api_base(api_base=custom_api_base, vertex_location=vertex_location) default_api_base = VertexBase.create_vertex_url( vertex_location=vertex_location or "us-central1", vertex_project=vertex_project or project_id, @@ -278,6 +256,45 @@ class VertexBase: """ return False + def _is_complete_gemini_url(self, api_base: str, model: str) -> bool: + import re + + # If URL already contains /models/{model_name}, consider it complete + if re.search(r"/models/" + re.escape(model), api_base): + return True + # Or check if it contains /models/ path segment (generic detection) + if "/models/" in api_base: + return True + return False + + def _is_complete_vertex_url(self, api_base: str) -> bool: + import re + + # If contains Vertex AI full path pattern, consider it complete + complete_url_patterns = [ + r"/projects/[^/]+/locations/[^/]+/publishers/", # Partner models + r"/endpoints/\d+", # Model Garden endpoints + ] + + for pattern in complete_url_patterns: + if re.search(pattern, api_base): + return True + + return False + + def _extract_gemini_path(self, url: str) -> str: + from urllib.parse import urlparse + + parsed = urlparse(url) + # Return path without query parameters (Cloudflare may need different auth) + return parsed.path + + def _extract_vertex_path(self, url: str) -> str: + from urllib.parse import urlparse + + parsed = urlparse(url) + return parsed.path + def _check_custom_proxy( self, api_base: Optional[str], @@ -299,19 +316,32 @@ class VertexBase: if custom_llm_provider == "gemini": # For Gemini (Google AI Studio), construct the full path like other providers if model is None: - raise ValueError( - "Model parameter is required for Gemini custom API base URLs" - ) - url = "{}/models/{}:{}".format(api_base, model, endpoint) - if gemini_api_key is None: - raise ValueError( - "Missing gemini_api_key, please set `GEMINI_API_KEY`" - ) - auth_header = ( - gemini_api_key # cloudflare expects api key as bearer token - ) - else: - url = "{}:{}".format(api_base, endpoint) + raise ValueError("Model parameter is required for Gemini custom API base URLs") + + # Smart detection: is api_base a complete path or base URL? + if self._is_complete_gemini_url(api_base, model): + # Old behavior: user provided complete path, only append endpoint + url = "{}:{}".format(api_base, endpoint) + else: + # New behavior: user provided base URL, need to append full path + # Extract path from default url (/v1beta/models/{model}:{endpoint}) + path_with_endpoint = self._extract_gemini_path(url) + url = api_base.rstrip("/") + path_with_endpoint + + # Set auth_header only if gemini_api_key is provided + # Cloudflare AI Gateway can store API keys server-side + if gemini_api_key is not None: + auth_header = gemini_api_key + else: # vertex_ai + # Smart detection: is api_base a complete path or base URL? + if self._is_complete_vertex_url(api_base): + # Old behavior: user provided complete path, only append endpoint + url = "{}:{}".format(api_base, endpoint) + else: + # New behavior: user provided base URL, need to append full path + # Extract path from create_vertex_url() generated url + path_with_endpoint = self._extract_vertex_path(url) + url = api_base.rstrip("/") + path_with_endpoint if stream is True: url = url + "?alt=sse" @@ -354,9 +384,7 @@ class VertexBase: ) ### SET RUNTIME ENDPOINT ### - version: Literal["v1beta1", "v1"] = ( - "v1beta1" if should_use_v1beta1_features is True else "v1" - ) + version: Literal["v1beta1", "v1"] = "v1beta1" if should_use_v1beta1_features is True else "v1" url, endpoint = _get_vertex_url( mode=mode, model=model, @@ -403,8 +431,7 @@ class VertexBase: The original error if reauthentication fails """ verbose_logger.debug( - f"Handling reauthentication for project_id: {project_id}. " - f"Clearing cache and retrying once." + f"Handling reauthentication for project_id: {project_id}. Clearing cache and retrying once." ) # Clear the cached credentials @@ -451,20 +478,14 @@ class VertexBase: """ # Convert dict credentials to string for caching - cache_credentials = ( - json.dumps(credentials) if isinstance(credentials, dict) else credentials - ) + cache_credentials = json.dumps(credentials) if isinstance(credentials, dict) else credentials credential_cache_key = (cache_credentials, project_id) _credentials: Optional[GoogleCredentialsObject] = None - verbose_logger.debug( - f"Checking cached credentials for project_id: {project_id}" - ) + verbose_logger.debug(f"Checking cached credentials for project_id: {project_id}") if credential_cache_key in self._credentials_project_mapping: - verbose_logger.debug( - f"Cached credentials found for project_id: {project_id}." - ) + verbose_logger.debug(f"Cached credentials found for project_id: {project_id}.") # Retrieve both credentials and cached project_id cached_entry = self._credentials_project_mapping[credential_cache_key] verbose_logger.debug("cached_entry: %s", cached_entry) @@ -473,9 +494,7 @@ class VertexBase: else: # Backward compatibility with old cache format _credentials = cached_entry - credential_project_id = _credentials.quota_project_id or getattr( - _credentials, "project_id", None - ) + credential_project_id = _credentials.quota_project_id or getattr(_credentials, "project_id", None) verbose_logger.debug( "Using cached credentials for project_id: %s", credential_project_id, @@ -487,9 +506,7 @@ class VertexBase: ) try: - _credentials, credential_project_id = self.load_auth( - credentials=credentials, project_id=project_id - ) + _credentials, credential_project_id = self.load_auth(credentials=credentials, project_id=project_id) except Exception as e: verbose_logger.exception( f"Failed to load vertex credentials. Check to see if credentials containing partial/invalid information. Error: {str(e)}" @@ -510,11 +527,7 @@ class VertexBase: ## VALIDATE CREDENTIALS verbose_logger.debug(f"Validating credentials for project_id: {project_id}") - if ( - project_id is None - and credential_project_id is not None - and isinstance(credential_project_id, str) - ): + if project_id is None and credential_project_id is not None and isinstance(credential_project_id, str): project_id = credential_project_id # Update cache with resolved project_id for future lookups resolved_cache_key = (cache_credentials, project_id) @@ -530,9 +543,7 @@ class VertexBase: if _credentials.expired: try: - verbose_logger.debug( - f"Credentials expired, refreshing for project_id: {project_id}" - ) + verbose_logger.debug(f"Credentials expired, refreshing for project_id: {project_id}") self.refresh_auth(_credentials) self._credentials_project_mapping[credential_cache_key] = ( _credentials, @@ -553,9 +564,7 @@ class VertexBase: ## VALIDATION STEP if _credentials.token is None or not isinstance(_credentials.token, str): raise ValueError( - "Could not resolve credentials token. Got None or non-string token - {}".format( - _credentials.token - ) + "Could not resolve credentials token. Got None or non-string token - {}".format(_credentials.token) ) if project_id is None: @@ -585,9 +594,7 @@ class VertexBase: except Exception as e: raise e - def set_headers( - self, auth_header: Optional[str], extra_headers: Optional[dict] - ) -> dict: + def set_headers(self, auth_header: Optional[str], extra_headers: Optional[dict]) -> dict: headers = { "Content-Type": "application/json", } 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 c85f607008..e65db6614d 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 @@ -708,9 +708,9 @@ class TestVertexBase: @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 + # Test case 1: Gemini with custom API base (new behavior - appends full path) ( - "https://proxy.example.com/generativelanguage.googleapis.com/v1beta", + "https://proxy.example.com", "gemini", "test-api-key", "generateContent", @@ -719,11 +719,11 @@ class TestVertexBase: "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", "gemini-2.5-flash-lite", "test-api-key", - "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent" + "https://proxy.example.com/v1beta/models/gemini-2.5-flash-lite:generateContent" ), - # Test case 2: Gemini with custom API base and streaming + # Test case 2: Gemini with custom API base and streaming (new behavior) ( - "https://proxy.example.com/generativelanguage.googleapis.com/v1beta", + "https://proxy.example.com", "gemini", "test-api-key", "generateContent", @@ -732,9 +732,22 @@ class TestVertexBase: "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", "gemini-2.5-flash-lite", "test-api-key", - "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent?alt=sse" + "https://proxy.example.com/v1beta/models/gemini-2.5-flash-lite:generateContent?alt=sse" ), - # Test case 3: Non-Gemini provider with custom API base + # Test case 3: Gemini with complete URL (old behavior - backward compatibility) + ( + "https://proxy.example.com/v1beta/models/gemini-2.5-flash-lite", + "gemini", + "test-api-key", + "generateContent", + False, + None, + "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", + "gemini-2.5-flash-lite", + "test-api-key", + "https://proxy.example.com/v1beta/models/gemini-2.5-flash-lite:generateContent" + ), + # Test case 4: Vertex AI with base URL (new behavior - appends full path) ( "https://custom-vertex-api.com", "vertex_ai", @@ -745,9 +758,22 @@ class TestVertexBase: "https://aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-pro:generateContent", "gemini-pro", "Bearer token123", - "https://custom-vertex-api.com:generateContent" + "https://custom-vertex-api.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-pro:generateContent" ), - # Test case 4: No API base provided (should return original values) + # Test case 5: Vertex AI with complete URL (old behavior - backward compatibility) + ( + "https://custom-vertex-api.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-pro", + "vertex_ai", + None, + "generateContent", + False, + "Bearer token123", + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-pro:generateContent", + "gemini-pro", + "Bearer token123", + "https://custom-vertex-api.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-pro:generateContent" + ), + # Test case 6: No API base provided (should return original values) ( None, "gemini", @@ -760,80 +786,52 @@ class TestVertexBase: "Bearer token123", "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent" ), - # Test case 5: Gemini without API key (should raise ValueError) - ( - "https://proxy.example.com/generativelanguage.googleapis.com/v1beta", - "gemini", - None, - "generateContent", - False, - None, - "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", - "gemini-2.5-flash-lite", - None, # This should raise an exception - None - ), ], ) def test_check_custom_proxy( - self, - api_base, - custom_llm_provider, - gemini_api_key, - endpoint, - stream, - auth_header, - url, - model, - expected_auth_header, + self, + api_base, + custom_llm_provider, + gemini_api_key, + endpoint, + stream, + auth_header, + url, + model, + expected_auth_header, expected_url ): """Test the _check_custom_proxy method for handling custom API base URLs""" vertex_base = VertexBase() - - if custom_llm_provider == "gemini" and api_base and gemini_api_key is None: - # Test case 5: Should raise ValueError for Gemini without API key - with pytest.raises(ValueError, match="Missing gemini_api_key"): - 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, - ) - else: - # Test cases 1-4: Should work correctly - 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, f"Expected auth_header {expected_auth_header}, got {result_auth_header}" - assert result_url == expected_url, f"Expected URL {expected_url}, got {result_url}" + + 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, f"Expected auth_header {expected_auth_header}, got {result_auth_header}" + assert result_url == expected_url, f"Expected URL {expected_url}, got {result_url}" def test_check_custom_proxy_gemini_url_construction(self): """Test that Gemini URLs are constructed correctly with custom API base""" vertex_base = VertexBase() - - # Test various Gemini models with custom API base + + # Test various Gemini models with custom API base (new behavior) test_cases = [ - ("gemini-2.5-flash-lite", "generateContent", "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent"), - ("gemini-2.5-pro", "generateContent", "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:generateContent"), - ("gemini-1.5-flash", "streamGenerateContent", "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:streamGenerateContent"), + ("gemini-2.5-flash-lite", "generateContent", "https://proxy.example.com/v1beta/models/gemini-2.5-flash-lite:generateContent"), + ("gemini-2.5-pro", "generateContent", "https://proxy.example.com/v1beta/models/gemini-2.5-pro:generateContent"), + ("gemini-1.5-flash", "streamGenerateContent", "https://proxy.example.com/v1beta/models/gemini-1.5-flash:streamGenerateContent"), ] - + for model, endpoint, expected_url in test_cases: _, result_url = vertex_base._check_custom_proxy( - api_base="https://proxy.example.com/generativelanguage.googleapis.com/v1beta", + api_base="https://proxy.example.com", custom_llm_provider="gemini", gemini_api_key="test-api-key", endpoint=endpoint, @@ -842,16 +840,16 @@ class TestVertexBase: url=f"https://generativelanguage.googleapis.com/v1beta/models/{model}:{endpoint}", model=model, ) - + assert result_url == expected_url, f"Expected {expected_url}, got {result_url} for model {model}" def test_check_custom_proxy_streaming_parameter(self): """Test that streaming parameter correctly adds ?alt=sse to URLs""" vertex_base = VertexBase() - + # Test with streaming enabled _, result_url_streaming = vertex_base._check_custom_proxy( - api_base="https://proxy.example.com/generativelanguage.googleapis.com/v1beta", + api_base="https://proxy.example.com", custom_llm_provider="gemini", gemini_api_key="test-api-key", endpoint="generateContent", @@ -860,13 +858,13 @@ class TestVertexBase: url="https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", model="gemini-2.5-flash-lite", ) - - expected_streaming_url = "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent?alt=sse" + + expected_streaming_url = "https://proxy.example.com/v1beta/models/gemini-2.5-flash-lite:generateContent?alt=sse" assert result_url_streaming == expected_streaming_url, f"Expected {expected_streaming_url}, got {result_url_streaming}" - + # Test with streaming disabled _, result_url_no_streaming = vertex_base._check_custom_proxy( - api_base="https://proxy.example.com/generativelanguage.googleapis.com/v1beta", + api_base="https://proxy.example.com", custom_llm_provider="gemini", gemini_api_key="test-api-key", endpoint="generateContent", @@ -875,6 +873,38 @@ class TestVertexBase: url="https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent", model="gemini-2.5-flash-lite", ) - - expected_no_streaming_url = "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent" + + expected_no_streaming_url = "https://proxy.example.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}" + + def test_check_custom_proxy_cloudflare_ai_gateway(self): + """Test Cloudflare AI Gateway URL construction for both Gemini and Vertex AI""" + vertex_base = VertexBase() + + # Test Gemini with Cloudflare AI Gateway + _, gemini_url = vertex_base._check_custom_proxy( + api_base="https://gateway.ai.cloudflare.com/v1/account123/gateway456/google-ai-studio", + custom_llm_provider="gemini", + gemini_api_key="test-api-key", + endpoint="generateContent", + stream=False, + auth_header=None, + url="https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent", + model="gemini-pro", + ) + expected_gemini_url = "https://gateway.ai.cloudflare.com/v1/account123/gateway456/google-ai-studio/v1beta/models/gemini-pro:generateContent" + assert gemini_url == expected_gemini_url, f"Expected {expected_gemini_url}, got {gemini_url}" + + # Test Vertex AI with Cloudflare AI Gateway + _, vertex_url = vertex_base._check_custom_proxy( + api_base="https://gateway.ai.cloudflare.com/v1/account123/gateway456/google-vertex-ai", + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="streamRawPredict", + stream=False, + auth_header="Bearer token123", + url="https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/anthropic/models/claude-3-sonnet@20240229:streamRawPredict", + model="claude-3-sonnet@20240229", + ) + expected_vertex_url = "https://gateway.ai.cloudflare.com/v1/account123/gateway456/google-vertex-ai/v1/projects/my-project/locations/us-central1/publishers/anthropic/models/claude-3-sonnet@20240229:streamRawPredict" + assert vertex_url == expected_vertex_url, f"Expected {expected_vertex_url}, got {vertex_url}"