diff --git a/litellm/utils.py b/litellm/utils.py index e2820f9c682..4315ae426cf 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9571,7 +9571,7 @@ class ProviderConfigManager: return get_nvidia_nim_rerank_config(model) elif litellm.LlmProviders.VERTEX_AI == provider: - return litellm.VertexAIRerankConfig() + return ProviderConfigManager._get_vertex_ai_rerank_config_cached() elif litellm.LlmProviders.FIREWORKS_AI == provider: return litellm.FireworksAIRerankConfig() elif litellm.LlmProviders.VOYAGE == provider: @@ -9590,6 +9590,11 @@ class ProviderConfigManager: return get_dashscope_family_rerank_config(provider.value) return litellm.CohereRerankConfig() + @staticmethod + @lru_cache(maxsize=1) + def _get_vertex_ai_rerank_config_cached() -> BaseRerankConfig: + return litellm.VertexAIRerankConfig() + @staticmethod def get_provider_decisions_config(model: str, provider: LlmProviders) -> BaseDecisionsConfig | None: if provider == LlmProviders.PERPLEXITY: 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..8abd012f9d8 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,13 +5,19 @@ Based on the test patterns from other rerank providers and the current Vertex AI import json import os +from datetime import datetime, timedelta +from typing import Final from unittest.mock import MagicMock, patch import httpx import pytest +import responses +import respx +import litellm from litellm.llms.vertex_ai.rerank.transformation import VertexAIRerankConfig from litellm.types.rerank import RerankResponse +from litellm.utils import ProviderConfigManager class TestVertexAIRerankTransform: @@ -646,3 +652,83 @@ class TestVertexAIRerankTransform: second_call = mock_ensure_access_token.call_args_list[1] assert first_call.kwargs["credentials"] == "path/to/credentials.json" assert second_call.kwargs["credentials"] == "path/to/credentials.json" + + +@pytest.fixture +def clear_vertex_rerank_cache(request: pytest.FixtureRequest) -> None: + ProviderConfigManager._get_vertex_ai_rerank_config_cached.cache_clear() + request.addfinalizer( + ProviderConfigManager._get_vertex_ai_rerank_config_cached.cache_clear + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("clear_vertex_rerank_cache") +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_rerank_reuses_scoped_credentials_and_refreshes_expired_tokens( + is_async: bool, +) -> None: + now: Final = datetime(2025, 1, 1) + credentials: Final = tuple( + json.dumps( + { + "type": "authorized_user", + "client_id": identity, + "client_secret": "synthetic-client-secret", + "refresh_token": f"synthetic-refresh-{identity}", + } + ) + for identity in ("first-principal", "second-principal") + ) + with ( + patch("litellm.disable_aiohttp_transport", True), + patch("google.auth._helpers.utcnow", return_value=now) as clock, + responses.RequestsMock(assert_all_requests_are_fired=False) as token_http, + respx.mock as rank_http, + ): + for token in ("initial", "other-credential", "other-project", "renewed"): + token_http.post( + "https://oauth2.googleapis.com/token", + json={"access_token": token, "expires_in": 3600, "token_type": "Bearer"}, + ) + rank_http.post( + url__regex=r"https://discoveryengine\.googleapis\.com/v1/projects/[^/]+/" + r"locations/global/rankingConfigs/default_ranking_config:rank" + ).respond(200, json={"records": [{"id": "0", "score": 0.9}]}) + + async def rank( + credential: str, project: str, expected_token: str, refresh_count: int + ) -> None: + params: Final = { + "model": "vertex_ai/semantic-ranker-default@latest", + "query": f"Query for {expected_token}", + "documents": ["A relevant document"], + "vertex_credentials": credential, + "vertex_project": project, + "caching": False, + "num_retries": 0, + } + result: Final = ( + await litellm.arerank(**params) + if is_async + else litellm.rerank(**params) + ) + assert isinstance(result, RerankResponse) + assert result.results[0]["index"] == 0 + assert len(token_http.calls) == refresh_count + request: Final = rank_http.calls.last.request + assert request.headers["Authorization"] == f"Bearer {expected_token}" + assert request.headers["X-Goog-User-Project"] == project + assert str(request.url) == ( + f"https://discoveryengine.googleapis.com/v1/projects/{project}/" + "locations/global/rankingConfigs/default_ranking_config:rank" + ) + assert json.loads(request.content)["query"] == params["query"] + + await rank(credentials[0], "first-project", "initial", 1) + await rank(credentials[0], "first-project", "initial", 1) + await rank(credentials[1], "first-project", "other-credential", 2) + await rank(credentials[0], "second-project", "other-project", 3) + await rank(credentials[0], "first-project", "initial", 3) + clock.return_value = now + timedelta(hours=2) + await rank(credentials[0], "first-project", "renewed", 4)