Use vertex creds passed via arguments (#16266)

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

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
 

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

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("/")

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(

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")

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"

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

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

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(

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.

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
 

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