From c6aab23fcb0f3820759398371bba841a2abcae76 Mon Sep 17 00:00:00 2001
From: Benjamin Stefan <145866889+Benjamin-Stefan@users.noreply.github.com>
Date: Mon, 12 Oct 2026 12:37:04 +0200
Subject: [PATCH] fix(vertex_ai): route Discovery Engine rerank to the eu/us
 endpoint when vertex_location is set

Only the eu and us multi-regions switch the host and the path. Any other value, including a single region like us-central1, keeps the global endpoint so existing configs keep working, and a value outside that list never reaches the URL. The location is read only from the call or deployment params, so a VERTEXAI_LOCATION set for Gemini does not move rerank traffic
---
 .../llms/vertex_ai/rerank/transformation.py   |  7 ++-
 .../test_vertex_ai_rerank_transformation.py   | 53 +++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py
index dce4d2f2a87..525237e332e 100644
--- a/litellm/llms/vertex_ai/rerank/transformation.py
+++ b/litellm/llms/vertex_ai/rerank/transformation.py
@@ -26,6 +26,8 @@ from litellm.types.rerank import (
     RerankResponseResult,
 )
 
+DISCOVERY_ENGINE_MULTI_REGIONS: Final = frozenset(("eu", "us"))
+
 
 class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
     """
@@ -43,7 +45,7 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
         self,
         api_base: str | None,
         model: str,
-        optional_params: dict | None = None,
+        optional_params: dict[str, object] | None = None,
     ) -> str:
         """
         Get the complete URL for the Vertex AI Discovery Engine ranking API
@@ -71,6 +73,9 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
                 "Vertex AI project ID is required. Please set 'VERTEXAI_PROJECT', 'litellm.vertex_project', or pass 'vertex_project' parameter"
             )
 
+        location: Final = self.explicit_vertex_ai_location(params)
+        if location in DISCOVERY_ENGINE_MULTI_REGIONS:
+            return f"https://{location}-discoveryengine.googleapis.com/v1/projects/{project_id}/locations/{location}/rankingConfigs/default_ranking_config:rank"
         return f"https://discoveryengine.googleapis.com/v1/projects/{project_id}/locations/global/rankingConfigs/default_ranking_config:rank"
 
     def validate_environment(
diff --git a/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py
index 630b2e1eb34..6d27c993532 100644
--- a/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py
+++ b/tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py
@@ -5,6 +5,7 @@ Based on the test patterns from other rerank providers and the current Vertex AI
 
 import json
 import os
+from typing import Final
 from unittest.mock import MagicMock, patch
 
 import httpx
@@ -13,6 +14,11 @@ import pytest
 from litellm.llms.vertex_ai.rerank.transformation import VertexAIRerankConfig
 from litellm.types.rerank import RerankResponse
 
+_GLOBAL_RANK_URL: Final = (
+    "https://discoveryengine.googleapis.com/v1/projects/test-project/"
+    "locations/global/rankingConfigs/default_ranking_config:rank"
+)
+
 
 class TestVertexAIRerankTransform:
     def setup_method(self):
@@ -103,6 +109,53 @@ class TestVertexAIRerankTransform:
             finally:
                 litellm.vertex_project = original_project
 
+    # Multi-region hosts per https://docs.cloud.google.com/generative-ai-app-builder/docs/locations (checked 2026-10-12)
+    @pytest.mark.parametrize(
+        ("location_params", "expected_url"),
+        [
+            pytest.param(
+                {"vertex_ai_location": "eu"},
+                "https://eu-discoveryengine.googleapis.com/v1/projects/test-project/"
+                "locations/eu/rankingConfigs/default_ranking_config:rank",
+                id="eu",
+            ),
+            pytest.param(
+                {"vertex_location": "us"},
+                "https://us-discoveryengine.googleapis.com/v1/projects/test-project/"
+                "locations/us/rankingConfigs/default_ranking_config:rank",
+                id="us",
+            ),
+            pytest.param({}, _GLOBAL_RANK_URL, id="unset"),
+            pytest.param({"vertex_ai_location": "global"}, _GLOBAL_RANK_URL, id="global"),
+            pytest.param({"vertex_ai_location": "us-central1"}, _GLOBAL_RANK_URL, id="single-region"),
+            pytest.param({"vertex_ai_location": "evil.example/x?"}, _GLOBAL_RANK_URL, id="host-injection"),
+        ],
+    )
+    def test_get_complete_url_routes_to_the_configured_location(
+        self, location_params: dict[str, str], expected_url: str
+    ):
+        self.config._ensure_access_token = MagicMock(return_value=("mock-token", "test-project"))
+
+        url: Final = self.config.get_complete_url(
+            api_base=None,
+            model=self.model,
+            optional_params={"vertex_project": "test-project", **location_params},
+        )
+
+        assert url == expected_url
+
+    def test_get_complete_url_ignores_the_location_from_the_environment(self, monkeypatch: pytest.MonkeyPatch):
+        monkeypatch.setenv("VERTEXAI_LOCATION", "eu")
+        self.config._ensure_access_token = MagicMock(return_value=("mock-token", "test-project"))
+
+        url: Final = self.config.get_complete_url(
+            api_base=None,
+            model=self.model,
+            optional_params={"vertex_project": "test-project"},
+        )
+
+        assert url == _GLOBAL_RANK_URL
+
     def test_validate_environment(self):
         """
         Test environment validation and header setup.