scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-24 05:45:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 72

Merge a376411676 into 2dfeb80e04

Browse source
This commit is contained in:
Benjamin Stefan 2026-10-19 00:59:21 +00:00 • committed by GitHub
commit 73f708ffc8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 59 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/llms/vertex_ai/rerank/transformation.py
View file

@ -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(

53
tests/unit/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py
View file

@ -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.