fix(vertex_ai): graft default vertex path when api_base has a version-only path (#38986)

* fix(vertex_ai): graft default vertex path when api_base has a version-only path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(vertex_ai): keep query and fragment placement when grafting vertex path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(vertex_ai): merge alt=sse into existing query when streaming Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-09-08 22:21:35 +00:00 · 2026-08-31 12:56:34 -07:00 · 2026-08-31 12:56:34 -07:00 · 1249f84b10
commit 1249f84b10
parent 0c21b30cb7
2 changed files with 127 additions and 3 deletions
--- a/litellm/llms/vertex_ai/vertex_llm_base.py
+++ b/litellm/llms/vertex_ai/vertex_llm_base.py
@ -27,6 +27,15 @@ from .common_utils import (
    get_vertex_base_url,
 )

+
+def _graft_default_vertex_path(api_base: str, default_url: str) -> str:
+    parsed_api_base: Final = urlparse(api_base)
+    default_segments: Final = urlparse(default_url).path.lstrip("/").split("/")
+    graft_segments: Final = default_segments[1:] if default_segments[0] in ("v1", "v1beta1") else default_segments
+    grafted_path: Final = parsed_api_base.path.rstrip("/") + "/" + "/".join(graft_segments)
+    return parsed_api_base._replace(path=grafted_path).geturl()
+
+
 GOOGLE_IMPORT_ERROR_MESSAGE: Final = (
    "Google Cloud SDK not found. Install it with: pip install 'litellm[google]' or pip install google-cloud-aiplatform"
 )
@ -621,8 +630,9 @@ class VertexBase:

        Handles custom api_base for:
        1. Gemini (Google AI Studio) - constructs /models/{model}:{endpoint}
-        2. Vertex AI with standard proxies - constructs {api_base}:{endpoint};
-           if api_base has no path (bare host), grafts the default vertex URL path onto it
+        2. Vertex AI with standard proxies - grafts the default vertex URL path onto the
+           api_base when its path is empty or only an API version (/v1, /v1beta1);
+           otherwise constructs {api_base}:{endpoint}
        3. Vertex AI with PSC endpoints - constructs full path structure
           {api_base}/v1/projects/{project}/locations/{location}/endpoints/{model}:{endpoint}
           (only when use_psc_endpoint_format=True)
@ -669,10 +679,14 @@ class VertexBase:
                    )
                elif urlparse(api_base).path in ("", "/"):
                    url = api_base.rstrip("/") + urlparse(url).path
+                elif urlparse(api_base).path.rstrip("/") in ("/v1", "/v1beta1") and "/projects/" in urlparse(url).path:
+                    url = _graft_default_vertex_path(api_base=api_base, default_url=url)
                else:
                    url = f"{api_base}:{endpoint}"
            if stream is True:
-                url = url + "?alt=sse"
+                parsed_stream_url: Final = urlparse(url)
+                stream_query: Final = f"{parsed_stream_url.query}&alt=sse" if parsed_stream_url.query else "alt=sse"
+                url = parsed_stream_url._replace(query=stream_query).geturl()
        return auth_header, url

    def _get_token_and_url(
--- a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py
+++ b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py
@ -982,6 +982,116 @@ class TestVertexBase:

        assert result_url == f"{gateway_api_base}:embedContent"

+    def test_check_custom_proxy_vertex_api_base_with_version_path_grafts_default_path(self):
+        vertex_base = VertexBase()
+
+        _, result_url = vertex_base._check_custom_proxy(
+            api_base="https://aiplatform.googleapis.com/v1beta1",
+            custom_llm_provider="vertex_ai",
+            gemini_api_key=None,
+            endpoint="generateContent",
+            stream=None,
+            auth_header="Bearer token123",
+            url="https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent",
+            model="gemini-3.5-flash-lite",
+        )
+
+        assert (
+            result_url
+            == "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent"
+        )
+
+    def test_check_custom_proxy_vertex_api_base_with_version_path_trailing_slash_grafts_default_path(self):
+        vertex_base = VertexBase()
+
+        _, result_url = vertex_base._check_custom_proxy(
+            api_base="https://internal-gateway.example.com/v1/",
+            custom_llm_provider="vertex_ai",
+            gemini_api_key=None,
+            endpoint="generateContent",
+            stream=None,
+            auth_header="Bearer token123",
+            url="https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent",
+            model="gemini-3.5-flash-lite",
+        )
+
+        assert (
+            result_url
+            == "https://internal-gateway.example.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent"
+        )
+
+    def test_check_custom_proxy_vertex_api_base_with_version_path_and_query_grafts_before_query(self):
+        vertex_base = VertexBase()
+
+        _, result_url = vertex_base._check_custom_proxy(
+            api_base="https://internal-gateway.example.com/v1beta1?key=abc",
+            custom_llm_provider="vertex_ai",
+            gemini_api_key=None,
+            endpoint="generateContent",
+            stream=None,
+            auth_header="Bearer token123",
+            url="https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent",
+            model="gemini-3.5-flash-lite",
+        )
+
+        assert (
+            result_url
+            == "https://internal-gateway.example.com/v1beta1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent?key=abc"
+        )
+
+    def test_check_custom_proxy_vertex_api_base_with_version_path_and_query_streaming_appends_alt_sse(self):
+        vertex_base = VertexBase()
+
+        _, result_url = vertex_base._check_custom_proxy(
+            api_base="https://internal-gateway.example.com/v1beta1?key=abc",
+            custom_llm_provider="vertex_ai",
+            gemini_api_key=None,
+            endpoint="streamGenerateContent",
+            stream=True,
+            auth_header="Bearer token123",
+            url="https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:streamGenerateContent",
+            model="gemini-3.5-flash-lite",
+        )
+
+        assert (
+            result_url
+            == "https://internal-gateway.example.com/v1beta1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:streamGenerateContent?key=abc&alt=sse"
+        )
+
+    def test_check_custom_proxy_vertex_api_base_with_non_version_path_keeps_endpoint_append(self):
+        vertex_base = VertexBase()
+        gateway_api_base = "https://gateway.example.com/vertex-proxy"
+
+        _, result_url = vertex_base._check_custom_proxy(
+            api_base=gateway_api_base,
+            custom_llm_provider="vertex_ai",
+            gemini_api_key=None,
+            endpoint="generateContent",
+            stream=None,
+            auth_header="Bearer token123",
+            url="https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent",
+            model="gemini-3.5-flash-lite",
+        )
+
+        assert result_url == f"{gateway_api_base}:generateContent"
+
+    def test_check_custom_proxy_vertex_api_base_without_projects_in_default_url_keeps_endpoint_append(self):
+        vertex_base = VertexBase()
+        gemma_api_base = "https://example.com/custom/gemma-deployment"
+
+        _, result_url = vertex_base._check_custom_proxy(
+            api_base=gemma_api_base,
+            custom_llm_provider="vertex_ai",
+            gemini_api_key=None,
+            endpoint="predict",
+            stream=False,
+            auth_header=None,
+            url=gemma_api_base,
+            model="gemma-3-27b-it",
+        )
+
+        assert result_url == f"{gemma_api_base}:predict"
+
    def test_check_custom_proxy_vertex_bare_host_streaming_keeps_single_alt_sse(self):
        vertex_base = VertexBase()