From 4dc23e807995654d85b6a5d924edb4c6afdd72fa Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 29 Sep 2025 13:03:46 -0700 Subject: [PATCH] =?UTF-8?q?Revert=20"feat:=20improve=20vertex=20AI/gemini?= =?UTF-8?q?=20api=5Fbase=20handling=20for=20proxy=20services=20(=E2=80=A6"?= =?UTF-8?q?=20(#15042)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit ff2d19e4cab88c7ce0a417f47e5ff7be85da515a. --- litellm/llms/vertex_ai/vertex_llm_base.py | 165 ++++++++-------- .../llms/vertex_ai/test_vertex_llm_base.py | 184 ++++++++---------- 2 files changed, 156 insertions(+), 193 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 6769bc3fb28..6d194d41add 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -46,10 +46,14 @@ 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, @@ -63,18 +67,26 @@ 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) @@ -111,7 +123,9 @@ 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 @@ -129,12 +143,16 @@ 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 @@ -144,7 +162,9 @@ 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": @@ -194,7 +214,9 @@ 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, @@ -256,45 +278,6 @@ 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], @@ -316,32 +299,19 @@ 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") - - # 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 + 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) if stream is True: url = url + "?alt=sse" @@ -384,7 +354,9 @@ 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, @@ -431,7 +403,8 @@ class VertexBase: The original error if reauthentication fails """ verbose_logger.debug( - f"Handling reauthentication for project_id: {project_id}. Clearing cache and retrying once." + f"Handling reauthentication for project_id: {project_id}. " + f"Clearing cache and retrying once." ) # Clear the cached credentials @@ -478,14 +451,20 @@ 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) @@ -494,7 +473,9 @@ 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, @@ -506,7 +487,9 @@ 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)}" @@ -527,7 +510,11 @@ 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) @@ -543,7 +530,9 @@ 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, @@ -564,7 +553,9 @@ 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: @@ -594,7 +585,9 @@ 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 e65db6614d6..c85f6070084 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 (new behavior - appends full path) + # Test case 1: Gemini with custom API base ( - "https://proxy.example.com", + "https://proxy.example.com/generativelanguage.googleapis.com/v1beta", "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/v1beta/models/gemini-2.5-flash-lite:generateContent" + "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 (new behavior) + # Test case 2: Gemini with custom API base and streaming ( - "https://proxy.example.com", + "https://proxy.example.com/generativelanguage.googleapis.com/v1beta", "gemini", "test-api-key", "generateContent", @@ -732,22 +732,9 @@ 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/v1beta/models/gemini-2.5-flash-lite:generateContent?alt=sse" + "https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent?alt=sse" ), - # 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) + # Test case 3: Non-Gemini provider with custom API base ( "https://custom-vertex-api.com", "vertex_ai", @@ -758,22 +745,9 @@ 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/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-pro:generateContent" + "https://custom-vertex-api.com:generateContent" ), - # 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) + # Test case 4: No API base provided (should return original values) ( None, "gemini", @@ -786,52 +760,80 @@ 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() - - 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}" + + 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}" 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 (new behavior) + + # Test various Gemini models with custom API base test_cases = [ - ("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"), + ("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"), ] - + for model, endpoint, expected_url in test_cases: _, result_url = vertex_base._check_custom_proxy( - api_base="https://proxy.example.com", + api_base="https://proxy.example.com/generativelanguage.googleapis.com/v1beta", custom_llm_provider="gemini", gemini_api_key="test-api-key", endpoint=endpoint, @@ -840,16 +842,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", + api_base="https://proxy.example.com/generativelanguage.googleapis.com/v1beta", custom_llm_provider="gemini", gemini_api_key="test-api-key", endpoint="generateContent", @@ -858,13 +860,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/v1beta/models/gemini-2.5-flash-lite:generateContent?alt=sse" + + expected_streaming_url = "https://proxy.example.com/generativelanguage.googleapis.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", + api_base="https://proxy.example.com/generativelanguage.googleapis.com/v1beta", custom_llm_provider="gemini", gemini_api_key="test-api-key", endpoint="generateContent", @@ -873,38 +875,6 @@ 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/v1beta/models/gemini-2.5-flash-lite:generateContent" + + 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}" - - 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}"