[Feat] New model - add nvidia nim llama-3.2-nv-rerankqa-1b-v2 (#17670)

* fix get_nvidia_nim_rerank_config

* add NvidiaNimRankingConfig

* add get_nvidia_nim_rerank_config

* add test_nvidia_nim_rerank_ranking_endpoint

* add /ranking model provider support

* feat: add nvidia/llama-3.2-nv-rerankqa-1b-v2

docs/my-website/docs/providers/nvidia_nim_rerank.md

@@ -141,6 +141,111 @@ curl -X POST http://0.0.0.0:4000/rerank \
   }'
 ```
 
+## `/v1/ranking` Models (llama-3.2-nv-rerankqa-1b-v2)
+
+Some Nvidia NIM rerank models use the `/v1/ranking` endpoint instead of the default `/v1/retrieval/{model}/reranking` endpoint.
+
+Use the `ranking/` prefix to force requests to the `/v1/ranking` endpoint:
+
+### LiteLLM Python SDK
+
+```python showLineNumbers title="Force /v1/ranking endpoint with ranking/ prefix"
+import litellm
+import os
+
+os.environ['NVIDIA_NIM_API_KEY'] = "nvapi-..."
+
+# Use "ranking/" prefix to force /v1/ranking endpoint
+response = litellm.rerank(
+    model="nvidia_nim/ranking/nvidia/llama-3.2-nv-rerankqa-1b-v2",
+    query="which way did the traveler go?",
+    documents=[
+        "two roads diverged in a yellow wood...",
+        "then took the other, as just as fair...",
+        "i shall be telling this with a sigh somewhere ages and ages hence..."
+    ],
+    top_n=3,
+    truncate="END",  # Optional: truncate long text from the end
+)
+
+print(response)
+```
+
+### LiteLLM Proxy
+
+```yaml showLineNumbers title="config.yaml"
+model_list:
+  - model_name: nvidia-ranking
+    litellm_params:
+      model: nvidia_nim/ranking/nvidia/llama-3.2-nv-rerankqa-1b-v2
+      api_key: os.environ/NVIDIA_NIM_API_KEY
+```
+
+```bash title="Request to LiteLLM Proxy"
+curl -X POST http://0.0.0.0:4000/rerank \
+  -H "Authorization: Bearer sk-1234" \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "nvidia-ranking",
+    "query": "which way did the traveler go?",
+    "documents": [
+      "two roads diverged in a yellow wood...",
+      "then took the other, as just as fair..."
+    ],
+    "top_n": 2
+  }'
+```
+
+### Understanding Model Resolution
+
+**Ranking Endpoint (`/v1/ranking`):**
+
+```
+model: nvidia_nim/ranking/nvidia/llama-3.2-nv-rerankqa-1b-v2
+       └────┬────┘ └──┬──┘ └─────────────┬──────────────────┘
+            │        │                   │
+            │        │                   └────▶ Model name sent to provider
+            │        │
+            │        └────────────────────────▶ Tells LiteLLM the request/response and url should be sent to Nvidia NIM /v1/ranking endpoint
+            │
+            └─────────────────────────────────▶ Provider prefix
+
+API URL: https://ai.api.nvidia.com/v1/ranking
+```
+
+**Visual Flow:**
+
+```
+Client Request                LiteLLM                              Provider API
+──────────────              ────────────                         ─────────────
+
+# Default reranking endpoint
+model: "nvidia_nim/nvidia/model-name"
+                            1. Extracts model: nvidia/model-name
+                            2. Routes to default endpoint ──────▶ POST /v1/retrieval/nvidia/model-name/reranking
+
+
+# Forced ranking endpoint  
+model: "nvidia_nim/ranking/nvidia/model-name"
+                            1. Detects "ranking/" prefix
+                            2. Extracts model: nvidia/model-name
+                            3. Routes to ranking endpoint ──────▶ POST /v1/ranking
+                                                                  Body: {"model": "nvidia/model-name", ...}
+```
+
+**When to use each endpoint:**
+
+| Endpoint | Model Prefix | Use Case |
+|----------|--------------|----------|
+| `/v1/retrieval/{model}/reranking` | `nvidia_nim/<model>` | Default for most rerank models |
+| `/v1/ranking` | `nvidia_nim/ranking/<model>` | For models like `nvidia/llama-3.2-nv-rerankqa-1b-v2` that require this endpoint |
+
+:::tip
+
+Check the [Nvidia NIM model deployment page](https://build.nvidia.com/nvidia/llama-3_2-nv-rerankqa-1b-v2/deploy) to see which endpoint your model requires.
+
+:::
+
 ## API Parameters
 
 ### Required Parameters
@@ -203,16 +308,7 @@ response = litellm.rerank(
 </TabItem>
 </Tabs>
 
-## API Endpoint
-
-The rerank endpoint uses a different base URL than chat/embeddings:
-
-- **Chat/Embeddings:** `https://integrate.api.nvidia.com/v1/`
-- **Rerank:** `https://ai.api.nvidia.com/v1/`
-
-LiteLLM automatically uses the correct endpoint for rerank requests.
-
-### Custom API Base URL
+## Custom API Base URL
 
 You can override the default base URL in several ways:
 
@@ -258,4 +354,3 @@ Get your Nvidia NIM API key from [Nvidia's website](https://developer.nvidia.com
 - [Nvidia NIM Chat Completions](./nvidia_nim#sample-usage)
 - [LiteLLM Rerank Endpoint](../rerank)
 - [Nvidia NIM Official Docs ↗](https://docs.api.nvidia.com/nim/reference/)
-

litellm/__init__.py

@@ -1111,6 +1111,7 @@ from .llms.jina_ai.rerank.transformation import JinaAIRerankConfig
 from .llms.deepinfra.rerank.transformation import DeepinfraRerankConfig
 from .llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig
 from .llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig
+from .llms.nvidia_nim.rerank.ranking_transformation import NvidiaNimRankingConfig
 from .llms.vertex_ai.rerank.transformation import VertexAIRerankConfig
 from .llms.fireworks_ai.rerank.transformation import FireworksAIRerankConfig
 from .llms.clarifai.chat.transformation import ClarifaiConfig

litellm/llms/nvidia_nim/rerank/common_utils.py

@@ -0,0 +1,28 @@
+"""
+Common utilities for NVIDIA NIM rerank provider.
+"""
+
+
+def get_nvidia_nim_rerank_config(model: str):
+    """
+    Get the appropriate NVIDIA NIM rerank config based on the model.
+    
+    Args:
+        model: The model string (e.g., "nvidia/llama-3.2-nv-rerankqa-1b-v2" or "ranking/nvidia/llama-3.2-nv-rerankqa-1b-v2")
+    
+    Returns:
+        NvidiaNimRankingConfig if model starts with "ranking/", else NvidiaNimRerankConfig
+    
+    Example:
+        - "ranking/nvidia/llama-3.2-nv-rerankqa-1b-v2" -> NvidiaNimRankingConfig
+        - "nvidia/llama-3.2-nv-rerankqa-1b-v2" -> NvidiaNimRerankConfig
+    """
+    from litellm.llms.nvidia_nim.rerank.ranking_transformation import (
+        NvidiaNimRankingConfig,
+    )
+    from litellm.llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig
+
+    if model.startswith("ranking/"):
+        return NvidiaNimRankingConfig()
+    return NvidiaNimRerankConfig()
+

litellm/llms/nvidia_nim/rerank/ranking_transformation.py

@@ -0,0 +1,75 @@
+"""
+Transformation for NVIDIA NIM Ranking models that use /v1/ranking endpoint.
+
+Use this by passing "nvidia_nim/ranking/<model>" to force the /v1/ranking endpoint.
+
+Reference: https://build.nvidia.com/nvidia/llama-3_2-nv-rerankqa-1b-v2/deploy
+"""
+
+from typing import Dict, Optional
+
+from litellm.llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig
+
+
+class NvidiaNimRankingConfig(NvidiaNimRerankConfig):
+    """
+    Configuration for NVIDIA NIM models that use the /v1/ranking endpoint.
+    
+    Example:
+        curl -X "POST" 'https://ai.api.nvidia.com/v1/ranking' \
+            -H 'Accept: application/json' \
+            -H 'Content-Type: application/json' \
+            -d '{
+                "model": "nvidia/llama-3.2-nv-rerankqa-1b-v2",
+                "query": {"text": "which way did the traveler go?"},
+                "passages": [{"text": "..."}, {"text": "..."}],
+                "truncate": "END"
+            }'
+    """
+
+    def _get_clean_model_name(self, model: str) -> str:
+        """Strip 'ranking/' prefix from model name."""
+        if model.startswith("ranking/"):
+            return model[len("ranking/"):]
+        return model
+
+    def get_complete_url(
+        self,
+        api_base: Optional[str],
+        model: str,
+        optional_params: Optional[dict] = None,
+    ) -> str:
+        """
+        Construct the Nvidia NIM ranking URL.
+        
+        Format: {api_base}/v1/ranking
+        """
+        if not api_base:
+            api_base = self.DEFAULT_NIM_RERANK_API_BASE
+
+        api_base = api_base.rstrip("/")
+
+        if api_base.endswith("/ranking"):
+            return api_base
+
+        if api_base.endswith("/v1"):
+            api_base = api_base[:-3]
+
+        return f"{api_base}/v1/ranking"
+
+    def transform_rerank_request(
+        self,
+        model: str,
+        optional_rerank_params: Dict,
+        headers: dict,
+    ) -> dict:
+        """
+        Transform request, using clean model name without 'ranking/' prefix.
+        """
+        clean_model = self._get_clean_model_name(model)
+        return super().transform_rerank_request(
+            model=clean_model,
+            optional_rerank_params=optional_rerank_params,
+            headers=headers,
+        )
+

litellm/model_prices_and_context_window_backup.json

@@ -22865,6 +22865,13 @@
         "mode": "rerank",
         "output_cost_per_token": 0.0
     },
+    "nvidia_nim/ranking/nvidia/llama-3.2-nv-rerankqa-1b-v2": {
+        "input_cost_per_query": 0.0,
+        "input_cost_per_token": 0.0,
+        "litellm_provider": "nvidia_nim",
+        "mode": "rerank",
+        "output_cost_per_token": 0.0
+    },
     "sagemaker/meta-textgeneration-llama-2-13b": {
         "input_cost_per_token": 0.0,
         "litellm_provider": "sagemaker",

litellm/utils.py

@@ -7366,7 +7366,11 @@ class ProviderConfigManager:
         elif litellm.LlmProviders.DEEPINFRA == provider:
             return litellm.DeepinfraRerankConfig()
         elif litellm.LlmProviders.NVIDIA_NIM == provider:
-            return litellm.NvidiaNimRerankConfig()
+            from litellm.llms.nvidia_nim.rerank.common_utils import (
+                get_nvidia_nim_rerank_config,
+            )
+
+            return get_nvidia_nim_rerank_config(model)
         elif litellm.LlmProviders.VERTEX_AI == provider:
             return litellm.VertexAIRerankConfig()
         elif litellm.LlmProviders.FIREWORKS_AI == provider:

model_prices_and_context_window.json

@@ -22865,6 +22865,13 @@
         "mode": "rerank",
         "output_cost_per_token": 0.0
     },
+    "nvidia_nim/ranking/nvidia/llama-3.2-nv-rerankqa-1b-v2": {
+        "input_cost_per_query": 0.0,
+        "input_cost_per_token": 0.0,
+        "litellm_provider": "nvidia_nim",
+        "mode": "rerank",
+        "output_cost_per_token": 0.0
+    },
     "sagemaker/meta-textgeneration-llama-2-13b": {
         "input_cost_per_token": 0.0,
         "litellm_provider": "sagemaker",

tests/llm_translation/test_nvidia_nim.py

@@ -184,13 +184,76 @@ def test_chat_completion_nvidia_nim_with_tools():
         assert request_body["tool_choice"] == "auto"
         assert request_body["parallel_tool_calls"] == True
 
+@pytest.mark.asyncio()
+async def test_nvidia_nim_rerank_ranking_endpoint():
+    """
+    Test that using "nvidia_nim/ranking/<model>" forces the /v1/ranking endpoint.
+    
+    This allows users to explicitly use the /v1/ranking endpoint for models like
+    nvidia/llama-3.2-nv-rerankqa-1b-v2.
+    
+    Reference: https://build.nvidia.com/nvidia/llama-3_2-nv-rerankqa-1b-v2/deploy
+    """
+    mock_response = AsyncMock()
+
+    def return_val():
+        return {
+            "rankings": [
+                {"index": 0, "logit": 0.95},
+                {"index": 1, "logit": 0.75},
+            ],
+        }
+
+    mock_response.json = return_val
+    mock_response.headers = {"key": "value"}
+    mock_response.status_code = 200
+
+    with patch(
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        return_value=mock_response,
+    ) as mock_post:
+        # Use "ranking/" prefix to force /v1/ranking endpoint
+        response = await litellm.arerank(
+            model="nvidia_nim/ranking/nvidia/llama-3.2-nv-rerankqa-1b-v2",
+            query="What is the GPU memory bandwidth?",
+            documents=["H100 delivers 3TB/s memory bandwidth", "A100 has 2TB/s memory bandwidth"],
+            top_n=2,
+            api_key="fake-api-key",
+        )
+
+        mock_post.assert_called_once()
+        
+        args_to_api = mock_post.call_args.kwargs["data"]
+        _url = mock_post.call_args.kwargs["url"]
+        print("url = ", _url)
+
+        # Verify URL is /v1/ranking
+        assert _url == "https://ai.api.nvidia.com/v1/ranking"
+
+        # Verify request body structure
+        request_data = json.loads(args_to_api)
+        print("request_data=", request_data)
+
+        # Query should be an object with 'text' field
+        assert request_data["query"] == {"text": "What is the GPU memory bandwidth?"}
+
+        # Documents should be 'passages'
+        assert request_data["passages"] == [
+            {"text": "H100 delivers 3TB/s memory bandwidth"},
+            {"text": "A100 has 2TB/s memory bandwidth"},
+        ]
+
+        # Model name in body should NOT have "ranking/" prefix
+        assert request_data["model"] == "nvidia/llama-3.2-nv-rerankqa-1b-v2"
+
+
 class TestNvidiaNim(BaseLLMRerankTest):
     def get_custom_llm_provider(self) -> litellm.LlmProviders:
         return litellm.LlmProviders.NVIDIA_NIM
 
     def get_base_rerank_call_args(self) -> dict:
         return {
-            "model": "nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2",
+            "model": "nvidia_nim/nvidia/llama-3.2-nv-rerankqa-1b-v2",
         }
     
     def get_expected_cost(self) -> float: