From 94e1a1ecacec835bdbef658cc3d128f89c865fd8 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 7 Nov 2025 09:05:22 +0530 Subject: [PATCH] Use vertex creds passed via arguments (#16266) --- .../llms/azure_ai/rerank/transformation.py | 8 +- .../llms/base_llm/rerank/transformation.py | 8 +- litellm/llms/cohere/rerank/transformation.py | 8 +- .../llms/cohere/rerank_v2/transformation.py | 7 +- litellm/llms/custom_httpx/llm_http_handler.py | 2 + .../llms/deepinfra/rerank/transformation.py | 8 +- .../llms/hosted_vllm/rerank/transformation.py | 8 +- .../llms/huggingface/rerank/transformation.py | 8 +- .../llms/infinity/rerank/transformation.py | 8 +- litellm/llms/jina_ai/rerank/transformation.py | 13 ++- .../llms/nvidia_nim/rerank/transformation.py | 8 +- .../llms/vertex_ai/rerank/transformation.py | 41 +++++++-- .../test_vertex_ai_rerank_transformation.py | 92 +++++++++++++++++++ 13 files changed, 200 insertions(+), 19 deletions(-) diff --git a/litellm/llms/azure_ai/rerank/transformation.py b/litellm/llms/azure_ai/rerank/transformation.py index 4465e0d70a..a47b6082c3 100644 --- a/litellm/llms/azure_ai/rerank/transformation.py +++ b/litellm/llms/azure_ai/rerank/transformation.py @@ -18,7 +18,12 @@ class AzureAIRerankConfig(CohereRerankConfig): Azure AI Rerank - Follows the same Spec as Cohere Rerank """ - def get_complete_url(self, api_base: Optional[str], model: str) -> str: + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: Optional[dict] = None, + ) -> str: if api_base is None: raise ValueError( "Azure AI API Base is required. api_base=None. Set in call or via `AZURE_AI_API_BASE` env var." @@ -32,6 +37,7 @@ class AzureAIRerankConfig(CohereRerankConfig): headers: dict, model: str, api_key: Optional[str] = None, + optional_params: Optional[dict] = None, ) -> dict: if api_key is None: api_key = get_secret_str("AZURE_AI_API_KEY") or litellm.azure_key diff --git a/litellm/llms/base_llm/rerank/transformation.py b/litellm/llms/base_llm/rerank/transformation.py index 6e9c03dee8..b22d85e82b 100644 --- a/litellm/llms/base_llm/rerank/transformation.py +++ b/litellm/llms/base_llm/rerank/transformation.py @@ -23,6 +23,7 @@ class BaseRerankConfig(ABC): headers: dict, model: str, api_key: Optional[str] = None, + optional_params: Optional[dict] = None, ) -> dict: pass @@ -50,7 +51,12 @@ class BaseRerankConfig(ABC): return model_response @abstractmethod - def get_complete_url(self, api_base: Optional[str], model: str) -> str: + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: Optional[dict] = None, + ) -> str: """ OPTIONAL diff --git a/litellm/llms/cohere/rerank/transformation.py b/litellm/llms/cohere/rerank/transformation.py index f9c979712d..d085cb13c4 100644 --- a/litellm/llms/cohere/rerank/transformation.py +++ b/litellm/llms/cohere/rerank/transformation.py @@ -20,7 +20,12 @@ class CohereRerankConfig(BaseRerankConfig): def __init__(self) -> None: pass - def get_complete_url(self, api_base: Optional[str], model: str) -> str: + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: Optional[dict] = None, + ) -> str: if api_base: # Remove trailing slashes and ensure clean base URL api_base = api_base.rstrip("/") @@ -72,6 +77,7 @@ class CohereRerankConfig(BaseRerankConfig): headers: dict, model: str, api_key: Optional[str] = None, + optional_params: Optional[dict] = None, ) -> dict: if api_key is None: api_key = ( diff --git a/litellm/llms/cohere/rerank_v2/transformation.py b/litellm/llms/cohere/rerank_v2/transformation.py index eb551a8a94..01309d937f 100644 --- a/litellm/llms/cohere/rerank_v2/transformation.py +++ b/litellm/llms/cohere/rerank_v2/transformation.py @@ -12,7 +12,12 @@ class CohereRerankV2Config(CohereRerankConfig): def __init__(self) -> None: pass - def get_complete_url(self, api_base: Optional[str], model: str) -> str: + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: Optional[dict] = None, + ) -> str: if api_base: # Remove trailing slashes and ensure clean base URL api_base = api_base.rstrip("/") diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3a290e7ce8..fb3498fb25 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -921,11 +921,13 @@ class BaseLLMHTTPHandler: api_key=api_key, headers=headers or {}, model=model, + optional_params=optional_rerank_params, ) api_base = provider_config.get_complete_url( api_base=api_base, model=model, + optional_params=optional_rerank_params, ) data = provider_config.transform_rerank_request( diff --git a/litellm/llms/deepinfra/rerank/transformation.py b/litellm/llms/deepinfra/rerank/transformation.py index 69c7dabebd..47f47418cb 100644 --- a/litellm/llms/deepinfra/rerank/transformation.py +++ b/litellm/llms/deepinfra/rerank/transformation.py @@ -28,7 +28,12 @@ class DeepinfraRerankConfig(BaseRerankConfig): Deepinfra Rerank - Follows the same Spec as Cohere Rerank """ - def get_complete_url(self, api_base: Optional[str], model: str) -> str: + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: Optional[dict] = None, + ) -> str: """ Constructs the complete DeepInfra inference endpoint URL for rerank. @@ -63,6 +68,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): headers: dict, model: str, api_key: Optional[str] = None, + optional_params: Optional[dict] = None, ) -> dict: if api_key is None: api_key = get_secret_str("DEEPINFRA_API_KEY") diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 2faef2c4c7..9d6a797afd 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -37,7 +37,12 @@ class HostedVLLMRerankConfig(BaseRerankConfig): def __init__(self) -> None: pass - def get_complete_url(self, api_base: Optional[str], model: str) -> str: + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: Optional[dict] = None, + ) -> str: if api_base: # Remove trailing slashes and ensure clean base URL api_base = api_base.rstrip("/") @@ -91,6 +96,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): headers: dict, model: str, api_key: Optional[str] = None, + optional_params: Optional[dict] = None, ) -> dict: if api_key is None: api_key = get_secret_str("HOSTED_VLLM_API_KEY") or "fake-api-key" diff --git a/litellm/llms/huggingface/rerank/transformation.py b/litellm/llms/huggingface/rerank/transformation.py index 1454328cc1..b386daf1c8 100644 --- a/litellm/llms/huggingface/rerank/transformation.py +++ b/litellm/llms/huggingface/rerank/transformation.py @@ -60,7 +60,12 @@ class HuggingFaceRerankConfig(BaseRerankConfig): else: return "https://api-inference.huggingface.co" - def get_complete_url(self, api_base: Optional[str], model: str) -> str: + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: Optional[dict] = None, + ) -> str: """ Get the complete URL for the API call, including the /rerank suffix if necessary. """ @@ -117,6 +122,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig): headers: dict, model: str, api_key: Optional[str] = None, + optional_params: Optional[dict] = None, api_base: Optional[str] = None, ) -> dict: # Get API credentials diff --git a/litellm/llms/infinity/rerank/transformation.py b/litellm/llms/infinity/rerank/transformation.py index 55aac6033d..1c15de714b 100644 --- a/litellm/llms/infinity/rerank/transformation.py +++ b/litellm/llms/infinity/rerank/transformation.py @@ -26,7 +26,12 @@ from ..common_utils import InfinityError class InfinityRerankConfig(CohereRerankConfig): - def get_complete_url(self, api_base: Optional[str], model: str) -> str: + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: Optional[dict] = None, + ) -> str: if api_base is None: raise ValueError("api_base is required for Infinity rerank") # Remove trailing slashes and ensure clean base URL @@ -40,6 +45,7 @@ class InfinityRerankConfig(CohereRerankConfig): headers: dict, model: str, api_key: Optional[str] = None, + optional_params: Optional[dict] = None, ) -> dict: if api_key is None: api_key = ( diff --git a/litellm/llms/jina_ai/rerank/transformation.py b/litellm/llms/jina_ai/rerank/transformation.py index 3ba24680fd..0fddd754a9 100644 --- a/litellm/llms/jina_ai/rerank/transformation.py +++ b/litellm/llms/jina_ai/rerank/transformation.py @@ -55,7 +55,12 @@ class JinaAIRerankConfig(BaseRerankConfig): **optional_params, )) - def get_complete_url(self, api_base: Optional[str], model: str) -> str: + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: Optional[dict] = None, + ) -> str: base_path = "/v1/rerank" if api_base is None: @@ -122,7 +127,11 @@ class JinaAIRerankConfig(BaseRerankConfig): ) # Return response def validate_environment( - self, headers: Dict, model: str, api_key: Optional[str] = None + self, + headers: Dict, + model: str, + api_key: Optional[str] = None, + optional_params: Optional[dict] = None, ) -> Dict: if api_key is None: raise ValueError( diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index cb9fd4beba..5bbe16e538 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -55,7 +55,12 @@ class NvidiaNimRerankConfig(BaseRerankConfig): def __init__(self) -> None: pass - def get_complete_url(self, api_base: Optional[str], model: str) -> str: + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: Optional[dict] = None, + ) -> str: """ Construct the Nvidia NIM rerank URL. @@ -131,6 +136,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig): headers: dict, model: str, api_key: Optional[str] = None, + optional_params: Optional[dict] = None, ) -> dict: """ Validate that the Nvidia NIM API key is present. diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py index 966368bc3c..c3cdd2b0fb 100644 --- a/litellm/llms/vertex_ai/rerank/transformation.py +++ b/litellm/llms/vertex_ai/rerank/transformation.py @@ -27,19 +27,40 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): def __init__(self) -> None: super().__init__() - def get_complete_url(self, api_base: Optional[str], model: str) -> str: + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: Optional[Dict] = None, + ) -> str: """ Get the complete URL for the Vertex AI Discovery Engine ranking API """ - # Get project ID from environment or litellm config + # Try to get project ID from optional_params first (e.g., vertex_project parameter) + params = optional_params or {} + + # Get credentials to extract project ID if needed + vertex_credentials = self.get_vertex_ai_credentials(params.copy()) + vertex_project = self.get_vertex_ai_project(params.copy()) + + # Use _ensure_access_token to extract project_id from credentials + # This is the same method used in vertex embeddings + _, vertex_project = self._ensure_access_token( + credentials=vertex_credentials, + project_id=vertex_project, + custom_llm_provider="vertex_ai", + ) + + # Fallback to environment or litellm config project_id = ( - get_secret_str("VERTEXAI_PROJECT") + vertex_project + or get_secret_str("VERTEXAI_PROJECT") or litellm.vertex_project ) if not project_id: raise ValueError( - "Vertex AI project ID is required. Please set 'VERTEXAI_PROJECT' or 'litellm.vertex_project'" + "Vertex AI project ID is required. Please set 'VERTEXAI_PROJECT', 'litellm.vertex_project', or pass 'vertex_project' parameter" ) return f"https://discoveryengine.googleapis.com/v1/projects/{project_id}/locations/global/rankingConfigs/default_ranking_config:rank" @@ -49,13 +70,15 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): headers: dict, model: str, api_key: Optional[str] = None, + optional_params: Optional[Dict] = None, ) -> dict: """ Validate and set up authentication for Vertex AI Discovery Engine API """ - # Get credentials and project info - vertex_credentials = self.get_vertex_ai_credentials({}) - vertex_project = self.get_vertex_ai_project({}) + # Get credentials and project info from optional_params (which contains vertex_credentials, etc.) + litellm_params = optional_params or {} + vertex_credentials = self.get_vertex_ai_credentials(litellm_params) + vertex_project = self.get_vertex_ai_project(litellm_params) # Get access token using the base class method access_token, project_id = self._ensure_access_token( @@ -218,10 +241,12 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): """ Map Cohere rerank params to Vertex AI format """ - return { + result = { "query": query, "documents": documents, "top_n": top_n, "return_documents": return_documents, } + result.update(non_default_params) + return result diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py index 9a1029b172..66375eef7b 100644 --- a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py @@ -351,3 +351,95 @@ class TestVertexAIRerankTransform: # Verify 0-based indexing for i, record in enumerate(request_data["records"]): assert record["id"] == str(i) + + def test_map_cohere_rerank_params_preserves_vertex_credentials(self): + """Test that map_cohere_rerank_params preserves vertex-specific parameters.""" + # Test with vertex_credentials + non_default_params = { + "documents": ["doc1", "doc2"], + "vertex_credentials": "path/to/credentials.json", + "vertex_project": "my-project-id", + "vertex_location": "us-central1" + } + + params = self.config.map_cohere_rerank_params( + non_default_params=non_default_params, + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + top_n=2, + return_documents=True + ) + + # Verify vertex-specific parameters are preserved + assert params["vertex_credentials"] == "path/to/credentials.json" + assert params["vertex_project"] == "my-project-id" + assert params["vertex_location"] == "us-central1" + + # Verify standard params are still present + assert params["query"] == "test query" + assert params["documents"] == ["doc1", "doc2"] + assert params["top_n"] == 2 + assert params["return_documents"] == True + + def test_map_cohere_rerank_params_without_vertex_credentials(self): + """Test that map_cohere_rerank_params works when vertex credentials are not provided.""" + non_default_params = { + "documents": ["doc1", "doc2"] + } + + params = self.config.map_cohere_rerank_params( + non_default_params=non_default_params, + model=self.model, + drop_params=False, + query="test query", + documents=["doc1", "doc2"], + top_n=2, + return_documents=True + ) + + # Verify no vertex-specific parameters are added when not provided + assert "vertex_credentials" not in params + assert "vertex_project" not in params + assert "vertex_location" not in params + + # Verify standard params are still present + assert params["query"] == "test query" + assert params["documents"] == ["doc1", "doc2"] + assert params["top_n"] == 2 + assert params["return_documents"] == True + + @patch('litellm.llms.vertex_ai.rerank.transformation.VertexAIRerankConfig._ensure_access_token') + def test_validate_environment_with_optional_params(self, mock_ensure_access_token): + """Test that validate_environment accepts and uses optional_params for credentials.""" + # Mock the authentication + mock_ensure_access_token.return_value = ("test-access-token", "test-project-123") + + optional_params = { + "vertex_credentials": "path/to/credentials.json", + "vertex_project": "custom-project-id", + "query": "test query", + "documents": ["doc1"] + } + + headers = self.config.validate_environment( + headers={}, + model=self.model, + api_key=None, + optional_params=optional_params + ) + + # Verify that _ensure_access_token was called with the credentials from optional_params + mock_ensure_access_token.assert_called_once() + call_args = mock_ensure_access_token.call_args + # The first call argument should be credentials (which will be the value from optional_params) + # We can't check the exact value easily due to how get_vertex_ai_credentials pops values, + # but we can verify the headers were set correctly + + expected_headers = { + "Authorization": "Bearer test-access-token", + "Content-Type": "application/json", + "X-Goog-User-Project": "test-project-123" + } + assert headers == expected_headers