feat: improve vertex AI/gemini api_base handling for proxy services (#15039)

2026-07-11 03:04:02 +00:00 · 2025-09-30 03:12:32 +08:00
parent 0d738b2899
commit ff2d19e4ca
2 changed files with 193 additions and 156 deletions
@@ -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",
        }
@@ -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}"