diff --git a/docs/my-website/docs/adding_provider/new_rerank_provider.md b/docs/my-website/docs/adding_provider/new_rerank_provider.md
index 84c363261c..628c099443 100644
--- a/docs/my-website/docs/adding_provider/new_rerank_provider.md
+++ b/docs/my-website/docs/adding_provider/new_rerank_provider.md
@@ -17,7 +17,7 @@ class YourProviderRerankConfig(BaseRerankConfig):
# ... other supported params
]
- def transform_rerank_request(self, model: str, optional_rerank_params: OptionalRerankParams, headers: dict) -> dict:
+ def transform_rerank_request(self, model: str, optional_rerank_params: Dict, headers: dict) -> dict:
# Transform request to RerankRequest spec
return rerank_request.model_dump(exclude_none=True)
diff --git a/docs/my-website/docs/providers/nvidia_nim.md b/docs/my-website/docs/providers/nvidia_nim.md
index 270b356c91..9dbfc80f4e 100644
--- a/docs/my-website/docs/providers/nvidia_nim.md
+++ b/docs/my-website/docs/providers/nvidia_nim.md
@@ -15,8 +15,8 @@ https://docs.api.nvidia.com/nim/reference/
| Description | Nvidia NIM is a platform that provides a simple API for deploying and using AI models. LiteLLM supports all models from [Nvidia NIM](https://developer.nvidia.com/nim/) |
| Provider Route on LiteLLM | `nvidia_nim/` |
| Provider Doc | [Nvidia NIM Docs ↗](https://developer.nvidia.com/nim/) |
-| API Endpoint for Provider | https://integrate.api.nvidia.com/v1/ |
-| Supported OpenAI Endpoints | `/chat/completions`, `/completions`, `/responses`, `/embeddings` |
+| API Endpoint for Provider | https://integrate.api.nvidia.com/v1/ (chat/embeddings), https://ai.api.nvidia.com/v1/ (rerank) |
+| Supported OpenAI Endpoints | `/chat/completions`, `/completions`, `/responses`, `/embeddings`, `/rerank` |
## API Key
```python
diff --git a/docs/my-website/docs/providers/nvidia_nim_rerank.md b/docs/my-website/docs/providers/nvidia_nim_rerank.md
new file mode 100644
index 0000000000..7373014a96
--- /dev/null
+++ b/docs/my-website/docs/providers/nvidia_nim_rerank.md
@@ -0,0 +1,261 @@
+import Tabs from '@theme/Tabs';
+import TabItem from '@theme/TabItem';
+
+# Nvidia NIM - Rerank
+
+Use Nvidia NIM Rerank models through LiteLLM.
+
+| Property | Details |
+|----------|---------|
+| Description | Nvidia NIM provides high-performance reranking models for semantic search and retrieval-augmented generation (RAG) |
+| Provider Doc | [Nvidia NIM Rerank API ↗](https://docs.api.nvidia.com/nim/reference/nvidia-llama-3_2-nv-rerankqa-1b-v2-infer) |
+| Supported Endpoint | `/rerank` |
+
+## Overview
+
+Nvidia NIM rerank models help you:
+- Reorder search results by relevance to a query
+- Improve RAG (Retrieval-Augmented Generation) accuracy
+- Filter and rank large document sets efficiently
+
+**Supported Models:**
+- All Nvidia NIM rerank models on their platform
+
+:::tip
+
+See the full list of LiteLLM supported Nvidia NIM rerank models on [Nvidia NIM](https://models.litellm.ai)
+
+:::
+
+## Usage
+
+### LiteLLM Python SDK
+
+
+
+
+```python
+import litellm
+import os
+
+os.environ['NVIDIA_NIM_API_KEY'] = "nvapi-..."
+
+response = litellm.rerank(
+ model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2",
+ query="What is the GPU memory bandwidth of H100 SXM?",
+ documents=[
+ "The Hopper GPU is paired with the Grace CPU using NVIDIA's ultra-fast chip-to-chip interconnect, delivering 900GB/s of bandwidth.",
+ "A100 provides up to 20X higher performance over the prior generation.",
+ "Accelerated servers with H100 deliver 3 terabytes per second (TB/s) of memory bandwidth per GPU."
+ ],
+ top_n=3,
+)
+
+print(response)
+```
+
+
+
+
+```python
+import litellm
+import os
+
+os.environ['NVIDIA_NIM_API_KEY'] = "nvapi-..."
+
+response = litellm.rerank(
+ model="nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3",
+ query="What is the GPU memory bandwidth of H100 SXM?",
+ documents=[
+ "The Hopper GPU is paired with the Grace CPU using NVIDIA's ultra-fast chip-to-chip interconnect, delivering 900GB/s of bandwidth.",
+ "A100 provides up to 20X higher performance over the prior generation.",
+ "Accelerated servers with H100 deliver 3 terabytes per second (TB/s) of memory bandwidth per GPU."
+ ],
+ top_n=3,
+)
+
+print(response)
+```
+
+
+
+
+**Response:**
+```json
+{
+ "results": [
+ {
+ "index": 2,
+ "relevance_score": 6.828125,
+ "document": {
+ "text": "Accelerated servers with H100 deliver 3 terabytes per second (TB/s) of memory bandwidth per GPU."
+ }
+ },
+ {
+ "index": 0,
+ "relevance_score": -1.564453125,
+ "document": {
+ "text": "The Hopper GPU is paired with the Grace CPU using NVIDIA's ultra-fast chip-to-chip interconnect, delivering 900GB/s of bandwidth."
+ }
+ }
+ ]
+}
+```
+
+
+## Usage with LiteLLM Proxy
+
+### 1. Setup Config
+
+Add Nvidia NIM rerank models to your proxy configuration:
+
+```yaml
+model_list:
+ - model_name: nvidia-rerank
+ litellm_params:
+ model: nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2
+ api_key: os.environ/NVIDIA_NIM_API_KEY
+```
+
+### 2. Start Proxy
+
+```bash
+litellm --config /path/to/config.yaml
+```
+
+### 3. Make Rerank Requests
+
+```bash
+curl -X POST http://0.0.0.0:4000/rerank \
+ -H "Authorization: Bearer sk-1234" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "model": "nvidia-rerank",
+ "query": "What is the GPU memory bandwidth of H100?",
+ "documents": [
+ "H100 delivers 3TB/s memory bandwidth",
+ "A100 has 2TB/s memory bandwidth",
+ "V100 offers 900GB/s memory bandwidth"
+ ],
+ "top_n": 2
+ }'
+```
+
+## API Parameters
+
+### Required Parameters
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `model` | string | The Nvidia NIM rerank model name with `nvidia_nim/` prefix |
+| `query` | string | The search query to rank documents against |
+| `documents` | array | List of documents to rank (1-1000 documents) |
+
+### Optional Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `top_n` | integer | All documents | Number of top-ranked documents to return |
+
+### Nvidia-Specific Parameters
+
+**`truncate`**: Controls how text is truncated if it exceeds the model's context window
+- `"NONE"`: No truncation (request may fail if too long)
+- `"END"`: Truncate from the end of the text
+
+```python
+response = litellm.rerank(
+ model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2",
+ query="GPU performance",
+ documents=["High performance computing", "Fast GPU processing"],
+ top_n=2,
+ truncate="END", # Nvidia-specific parameter
+)
+```
+
+## Authentication
+
+Set your Nvidia NIM API key:
+
+
+
+
+```bash
+export NVIDIA_NIM_API_KEY="nvapi-..."
+```
+
+
+
+
+```python
+import os
+os.environ['NVIDIA_NIM_API_KEY'] = "nvapi-..."
+
+# Or pass directly
+response = litellm.rerank(
+ model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2",
+ query="test",
+ documents=["doc1"],
+ api_key="nvapi-...",
+)
+```
+
+
+
+
+## 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
+
+You can override the default base URL in several ways:
+
+**Option 1: Environment Variable**
+
+```bash
+export NVIDIA_NIM_API_BASE="https://your-custom-endpoint.com"
+```
+
+**Option 2: Pass as parameter**
+
+```python
+response = litellm.rerank(
+ model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2",
+ query="test",
+ documents=["doc1"],
+ api_base="https://your-custom-endpoint.com",
+)
+```
+
+**Option 3: Full URL (including model path)**
+
+If you have the complete endpoint URL, you can pass it directly:
+
+```python
+response = litellm.rerank(
+ model="nvidia_nim/nvidia/llama-3_2-nv-rerankqa-1b-v2",
+ query="test",
+ documents=["doc1"],
+ api_base="https://your-custom-endpoint.com/v1/retrieval/nvidia/llama-3_2-nv-rerankqa-1b-v2/reranking",
+)
+```
+
+LiteLLM will detect the full URL (by checking for `/retrieval/` in the path) and use it as-is.
+
+### How do I get an API key?
+
+Get your Nvidia NIM API key from [Nvidia's website](https://developer.nvidia.com/nim/).
+
+## Related Documentation
+
+- [Nvidia NIM - Main Documentation](./nvidia_nim)
+- [Nvidia NIM Chat Completions](./nvidia_nim#sample-usage)
+- [LiteLLM Rerank Endpoint](../rerank)
+- [Nvidia NIM Official Docs ↗](https://docs.api.nvidia.com/nim/reference/)
+
diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js
index 56baf1a702..7acd92240c 100644
--- a/docs/my-website/sidebars.js
+++ b/docs/my-website/sidebars.js
@@ -458,7 +458,14 @@ const sidebars = {
"providers/deepgram",
"providers/watsonx",
"providers/predibase",
- "providers/nvidia_nim",
+ {
+ type: "category",
+ label: "Nvidia NIM",
+ items: [
+ "providers/nvidia_nim",
+ "providers/nvidia_nim_rerank",
+ ]
+ },
{ type: "doc", id: "providers/nscale", label: "Nscale (EU Sovereign)" },
"providers/xai",
"providers/moonshot",
diff --git a/litellm/__init__.py b/litellm/__init__.py
index d1f00d3f0d..43ae1750cd 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -1061,6 +1061,7 @@ from .llms.azure_ai.rerank.transformation import AzureAIRerankConfig
from .llms.infinity.rerank.transformation import InfinityRerankConfig
from .llms.jina_ai.rerank.transformation import JinaAIRerankConfig
from .llms.deepinfra.rerank.transformation import DeepinfraRerankConfig
+from .llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig
from .llms.clarifai.chat.transformation import ClarifaiConfig
from .llms.ai21.chat.transformation import AI21ChatConfig, AI21ChatConfig as AI21Config
from .llms.meta_llama.chat.transformation import LlamaAPIConfig
diff --git a/litellm/llms/base_llm/rerank/transformation.py b/litellm/llms/base_llm/rerank/transformation.py
index 8701fe57bf..6e9c03dee8 100644
--- a/litellm/llms/base_llm/rerank/transformation.py
+++ b/litellm/llms/base_llm/rerank/transformation.py
@@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
import httpx
-from litellm.types.rerank import OptionalRerankParams, RerankBilledUnits, RerankResponse
+from litellm.types.rerank import RerankBilledUnits, RerankResponse
from litellm.types.utils import ModelInfo
from ..chat.transformation import BaseLLMException
@@ -30,7 +30,7 @@ class BaseRerankConfig(ABC):
def transform_rerank_request(
self,
model: str,
- optional_rerank_params: OptionalRerankParams,
+ optional_rerank_params: Dict,
headers: dict,
) -> dict:
return {}
@@ -78,7 +78,7 @@ class BaseRerankConfig(ABC):
return_documents: Optional[bool] = True,
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
- ) -> OptionalRerankParams:
+ ) -> Dict:
pass
def get_error_class(
diff --git a/litellm/llms/cohere/rerank/transformation.py b/litellm/llms/cohere/rerank/transformation.py
index 5371b9a4b6..6586a83f06 100644
--- a/litellm/llms/cohere/rerank/transformation.py
+++ b/litellm/llms/cohere/rerank/transformation.py
@@ -1,8 +1,8 @@
from typing import Any, Dict, List, Optional, Union
import httpx
-import litellm
+import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
@@ -52,20 +52,20 @@ class CohereRerankConfig(BaseRerankConfig):
return_documents: Optional[bool] = True,
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
- ) -> OptionalRerankParams:
+ ) -> Dict:
"""
Map Cohere rerank params
No mapping required - returns all supported params
"""
- return OptionalRerankParams(
+ return dict(OptionalRerankParams(
query=query,
documents=documents,
top_n=top_n,
rank_fields=rank_fields,
return_documents=return_documents,
max_chunks_per_doc=max_chunks_per_doc,
- )
+ ))
def validate_environment(
self,
@@ -101,7 +101,7 @@ class CohereRerankConfig(BaseRerankConfig):
def transform_rerank_request(
self,
model: str,
- optional_rerank_params: OptionalRerankParams,
+ optional_rerank_params: Dict,
headers: dict,
) -> dict:
if "query" not in optional_rerank_params:
diff --git a/litellm/llms/cohere/rerank_v2/transformation.py b/litellm/llms/cohere/rerank_v2/transformation.py
index 74e760460d..eb551a8a94 100644
--- a/litellm/llms/cohere/rerank_v2/transformation.py
+++ b/litellm/llms/cohere/rerank_v2/transformation.py
@@ -44,25 +44,25 @@ class CohereRerankV2Config(CohereRerankConfig):
return_documents: Optional[bool] = True,
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
- ) -> OptionalRerankParams:
+ ) -> Dict:
"""
Map Cohere rerank params
No mapping required - returns all supported params
"""
- return OptionalRerankParams(
+ return dict(OptionalRerankParams(
query=query,
documents=documents,
top_n=top_n,
rank_fields=rank_fields,
return_documents=return_documents,
max_tokens_per_doc=max_tokens_per_doc,
- )
+ ))
def transform_rerank_request(
self,
model: str,
- optional_rerank_params: OptionalRerankParams,
+ optional_rerank_params: Dict,
headers: dict,
) -> dict:
if "query" not in optional_rerank_params:
diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py
index 173bb5a2cc..ae449bbb15 100644
--- a/litellm/llms/custom_httpx/llm_http_handler.py
+++ b/litellm/llms/custom_httpx/llm_http_handler.py
@@ -65,7 +65,7 @@ from litellm.types.llms.openai import (
ResponseInputParam,
ResponsesAPIResponse,
)
-from litellm.types.rerank import OptionalRerankParams, RerankResponse
+from litellm.types.rerank import RerankResponse
from litellm.types.responses.main import DeleteResponseResult
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import (
@@ -893,7 +893,7 @@ class BaseLLMHTTPHandler:
custom_llm_provider: str,
logging_obj: LiteLLMLoggingObj,
provider_config: BaseRerankConfig,
- optional_rerank_params: OptionalRerankParams,
+ optional_rerank_params: Dict,
timeout: Optional[Union[float, httpx.Timeout]],
model_response: RerankResponse,
_is_async: bool = False,
diff --git a/litellm/llms/deepinfra/rerank/transformation.py b/litellm/llms/deepinfra/rerank/transformation.py
index 6a3244a3c8..69c7dabebd 100644
--- a/litellm/llms/deepinfra/rerank/transformation.py
+++ b/litellm/llms/deepinfra/rerank/transformation.py
@@ -2,11 +2,11 @@
Translate between Cohere's `/rerank` format and Deepinfra's `/rerank` format.
"""
-from litellm._uuid import uuid
from typing import Any, Dict, List, Optional, Union
import httpx
+from litellm._uuid import uuid
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.rerank.transformation import (
BaseLLMException,
@@ -98,7 +98,7 @@ class DeepinfraRerankConfig(BaseRerankConfig):
return_documents: Optional[bool] = True,
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
- ) -> OptionalRerankParams:
+ ) -> Dict:
# Start with the basic parameters
optional_rerank_params = {}
if query:
@@ -124,7 +124,7 @@ class DeepinfraRerankConfig(BaseRerankConfig):
def transform_rerank_request(
self,
model: str,
- optional_rerank_params: OptionalRerankParams,
+ optional_rerank_params: Dict,
headers: dict,
) -> dict:
# Convert OptionalRerankParams to dict as expected by parent class
diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py
index 4ed604e2c8..2faef2c4c7 100644
--- a/litellm/llms/hosted_vllm/rerank/transformation.py
+++ b/litellm/llms/hosted_vllm/rerank/transformation.py
@@ -2,27 +2,26 @@
Transformation logic for Hosted VLLM rerank
"""
-from litellm._uuid import uuid
from typing import Any, Dict, List, Optional, Union
+import httpx
+
+from litellm._uuid import uuid
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
+from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
+from litellm.secret_managers.main import get_secret_str
from litellm.types.rerank import (
+ OptionalRerankParams,
RerankBilledUnits,
+ RerankRequest,
RerankResponse,
RerankResponseDocument,
RerankResponseMeta,
RerankResponseResult,
RerankTokens,
- OptionalRerankParams,
- RerankRequest,
)
-import httpx
-
-from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
-from litellm.llms.base_llm.chat.transformation import BaseLLMException
-from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
-from litellm.secret_managers.main import get_secret_str
-
class HostedVLLMRerankError(BaseLLMException):
def __init__(
@@ -72,20 +71,20 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
return_documents: Optional[bool] = True,
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
- ) -> OptionalRerankParams:
+ ) -> Dict:
"""
Map parameters for Hosted VLLM rerank
"""
if max_chunks_per_doc is not None:
raise ValueError("Hosted VLLM does not support max_chunks_per_doc")
- return OptionalRerankParams(
+ return dict(OptionalRerankParams(
query=query,
documents=documents,
top_n=top_n,
rank_fields=rank_fields,
return_documents=return_documents,
- )
+ ))
def validate_environment(
self,
@@ -112,7 +111,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
def transform_rerank_request(
self,
model: str,
- optional_rerank_params: OptionalRerankParams,
+ optional_rerank_params: Dict,
headers: dict,
) -> dict:
if "query" not in optional_rerank_params:
diff --git a/litellm/llms/huggingface/rerank/transformation.py b/litellm/llms/huggingface/rerank/transformation.py
index aa0f37bc6b..1454328cc1 100644
--- a/litellm/llms/huggingface/rerank/transformation.py
+++ b/litellm/llms/huggingface/rerank/transformation.py
@@ -1,11 +1,11 @@
import os
-from litellm._uuid import uuid
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
import httpx
from typing_extensions import TypedDict
import litellm
+from litellm._uuid import uuid
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.secret_managers.main import get_secret_str
@@ -95,7 +95,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig):
return_documents: Optional[bool] = True,
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
- ) -> OptionalRerankParams:
+ ) -> Dict:
optional_rerank_params = {}
if non_default_params is not None:
for k, v in non_default_params.items():
diff --git a/litellm/llms/jina_ai/rerank/transformation.py b/litellm/llms/jina_ai/rerank/transformation.py
index d8569b01b8..3bd5a4847f 100644
--- a/litellm/llms/jina_ai/rerank/transformation.py
+++ b/litellm/llms/jina_ai/rerank/transformation.py
@@ -6,11 +6,11 @@ Why separate file? Make it easy to see how transformation works
Docs - https://jina.ai/reranker
"""
-from litellm._uuid import uuid
from typing import Any, Dict, List, Optional, Tuple, Union
from httpx import URL, Response
+from litellm._uuid import uuid
from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.types.rerank import (
@@ -45,15 +45,15 @@ class JinaAIRerankConfig(BaseRerankConfig):
return_documents: Optional[bool] = True,
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
- ) -> OptionalRerankParams:
+ ) -> Dict:
optional_params = {}
supported_params = self.get_supported_cohere_rerank_params(model)
for k, v in non_default_params.items():
if k in supported_params:
optional_params[k] = v
- return OptionalRerankParams(
+ return dict(OptionalRerankParams(
**optional_params,
- )
+ ))
def get_complete_url(self, api_base: Optional[str], model: str) -> str:
base_path = "/v1/rerank"
@@ -67,7 +67,7 @@ class JinaAIRerankConfig(BaseRerankConfig):
return cleaned_base
def transform_rerank_request(
- self, model: str, optional_rerank_params: OptionalRerankParams, headers: Dict
+ self, model: str, optional_rerank_params: Dict, headers: Dict
) -> Dict:
return {"model": model, **optional_rerank_params}
diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py
new file mode 100644
index 0000000000..cb9fd4beba
--- /dev/null
+++ b/litellm/llms/nvidia_nim/rerank/transformation.py
@@ -0,0 +1,325 @@
+from typing import Any, Dict, List, Literal, Optional, Union
+
+import httpx
+from typing_extensions import Required, TypedDict
+
+import litellm
+from litellm._uuid import uuid
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
+from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
+from litellm.secret_managers.main import get_secret_str
+from litellm.types.rerank import (
+ RerankBilledUnits,
+ RerankResponse,
+ RerankResponseMeta,
+ RerankResponseResult,
+)
+
+
+class NvidiaNimQueryObject(TypedDict):
+ text: Required[str]
+
+
+class NvidiaNimPassageObject(TypedDict):
+ text: Required[str]
+
+
+class NvidiaNimRerankRequest(TypedDict, total=False):
+ model: Required[str]
+ query: Required[NvidiaNimQueryObject]
+ passages: Required[List[NvidiaNimPassageObject]]
+ truncate: Literal["NONE", "END"]
+ top_k: int
+
+
+class NvidiaNimRankingResult(TypedDict):
+ index: Required[int]
+ logit: Required[float]
+
+
+class NvidiaNimRerankResponse(TypedDict):
+ rankings: Required[List[NvidiaNimRankingResult]]
+
+
+class NvidiaNimRerankConfig(BaseRerankConfig):
+ """
+ Reference: https://docs.api.nvidia.com/nim/reference/nvidia-llama-3_2-nv-rerankqa-1b-v2-infer
+
+ Nvidia NIM rerank API uses a different format:
+ - query is an object with 'text' field
+ - documents are called 'passages' and have 'text' field
+ """
+ DEFAULT_NIM_RERANK_API_BASE = "https://ai.api.nvidia.com"
+
+ def __init__(self) -> None:
+ pass
+
+ def get_complete_url(self, api_base: Optional[str], model: str) -> str:
+ """
+ Construct the Nvidia NIM rerank URL.
+
+ Format: {api_base}/v1/retrieval/{model}/reranking
+
+ If the user provides a full URL (e.g., {api_base}/v1/retrieval/{model}/reranking),
+ it will be used as-is.
+ """
+ if not api_base:
+ api_base = self.DEFAULT_NIM_RERANK_API_BASE
+
+ api_base = api_base.rstrip("/")
+
+ # Check if user already provided the full URL with /retrieval/ path
+ if "/retrieval/" in api_base:
+ return api_base
+
+ # Ensure we don't have duplicate /v1
+ if api_base.endswith("/v1"):
+ api_base = api_base[:-3]
+
+ return f"{api_base}/v1/retrieval/{model}/reranking"
+
+ def get_supported_cohere_rerank_params(self, model: str) -> list:
+ """
+ Nvidia NIM supports these rerank parameters.
+ """
+ return [
+ "query",
+ "documents",
+ "top_n",
+ ]
+
+ def map_cohere_rerank_params(
+ self,
+ non_default_params: Optional[dict],
+ model: str,
+ drop_params: bool,
+ query: str,
+ documents: List[Union[str, Dict[str, Any]]],
+ custom_llm_provider: Optional[str] = None,
+ top_n: Optional[int] = None,
+ rank_fields: Optional[List[str]] = None,
+ return_documents: Optional[bool] = True,
+ max_chunks_per_doc: Optional[int] = None,
+ max_tokens_per_doc: Optional[int] = None,
+ ) -> Dict:
+ """
+ Map Cohere/OpenAI rerank params to Nvidia NIM format.
+
+ Parameter mapping:
+ - top_n (Cohere) -> top_k (Nvidia)
+
+ Nvidia NIM specific params (passed through as-is from non_default_params):
+ - truncate: How to truncate input if too long (NONE, END)
+ """
+ optional_nvidia_nim_rerank_params: Dict[str, Any] = {
+ "query": query,
+ "documents": documents,
+ }
+
+ # Map Cohere's top_n to Nvidia's top_k
+ if top_n is not None:
+ optional_nvidia_nim_rerank_params["top_k"] = top_n
+
+ # Pass through Nvidia-specific params from non_default_params
+ if non_default_params:
+ optional_nvidia_nim_rerank_params.update(non_default_params)
+ return dict(optional_nvidia_nim_rerank_params)
+
+ def validate_environment(
+ self,
+ headers: dict,
+ model: str,
+ api_key: Optional[str] = None,
+ ) -> dict:
+ """
+ Validate that the Nvidia NIM API key is present.
+ """
+ if api_key is None:
+ api_key = (
+ get_secret_str("NVIDIA_NIM_API_KEY")
+ or litellm.api_key
+ )
+
+ if api_key is None:
+ raise ValueError(
+ "Nvidia NIM API key is required. Please set 'NVIDIA_NIM_API_KEY' in your environment"
+ )
+
+ default_headers = {
+ "Authorization": f"Bearer {api_key}",
+ "accept": "application/json",
+ "content-type": "application/json",
+ }
+
+ # If 'Authorization' is provided in headers, it overrides the default
+ if "Authorization" in headers:
+ default_headers["Authorization"] = headers["Authorization"]
+
+ # Merge other headers, overriding any default ones except Authorization
+ return {**default_headers, **headers}
+
+ def transform_rerank_request(
+ self,
+ model: str,
+ optional_rerank_params: Dict,
+ headers: dict,
+ ) -> dict:
+ """
+ Transform request to Nvidia NIM format.
+
+ Nvidia NIM expects:
+ - query as {text: "..."}
+ - documents as passages: [{text: "..."}, ...]
+ - Optional: truncate (NONE or END), top_k
+
+ Note: optional_rerank_params may contain provider-specific params like 'top_k' and 'truncate'
+ that aren't in the OptionalRerankParams TypedDict but are passed through at runtime.
+ The mapping from Cohere's 'top_n' to Nvidia's 'top_k' already happened in map_cohere_rerank_params.
+ """
+ if "query" not in optional_rerank_params:
+ raise ValueError("query is required for Nvidia NIM rerank")
+ if "documents" not in optional_rerank_params:
+ raise ValueError("documents is required for Nvidia NIM rerank")
+
+ query = optional_rerank_params["query"]
+ documents = optional_rerank_params["documents"]
+
+ # Transform query to object format
+ query_obj: NvidiaNimQueryObject = {"text": query}
+
+ # Transform documents to passages format
+ passages: List[NvidiaNimPassageObject] = []
+ for doc in documents:
+ if isinstance(doc, str):
+ passages.append({"text": doc})
+ elif isinstance(doc, dict):
+ # If document is already a dict, check if it has 'text' field
+ if "text" in doc:
+ passages.append({"text": doc["text"]})
+ else:
+ # Otherwise, stringify the dict
+ import json
+ passages.append({"text": json.dumps(doc)})
+ else:
+ passages.append({"text": str(doc)})
+
+ # Note: URL path uses underscores (llama-3_2) but JSON body uses periods (llama-3.2)
+ # Convert underscores back to periods for the model field in request body
+ model_for_body = model.replace("_", ".")
+
+ # Build request using TypedDict
+ request_data: NvidiaNimRerankRequest = {
+ "model": model_for_body,
+ "query": query_obj,
+ "passages": passages,
+ }
+
+ # Add optional top_k parameter if provided (already mapped from top_n in map_cohere_rerank_params)
+ if "top_k" in optional_rerank_params and optional_rerank_params.get("top_k") is not None: # type: ignore
+ request_data["top_k"] = optional_rerank_params.get("top_k") # type: ignore
+
+ # Add Nvidia-specific truncate parameter if provided
+ # This is passed through from non_default_params, not in base OptionalRerankParams
+ if "truncate" in optional_rerank_params and optional_rerank_params.get("truncate") is not None: # type: ignore
+ truncate_value = optional_rerank_params.get("truncate") # type: ignore
+ if truncate_value in ["NONE", "END"]:
+ request_data["truncate"] = truncate_value # type: ignore
+
+ return dict(request_data)
+
+ def transform_rerank_response(
+ self,
+ model: str,
+ raw_response: httpx.Response,
+ model_response: RerankResponse,
+ logging_obj: LiteLLMLoggingObj,
+ api_key: Optional[str] = None,
+ request_data: dict = {},
+ optional_params: dict = {},
+ litellm_params: dict = {},
+ ) -> RerankResponse:
+ """
+ Transform Nvidia NIM rerank response to LiteLLM format.
+
+ Nvidia NIM returns (NvidiaNimRerankResponse):
+ {
+ "rankings": [
+ {
+ "index": 0,
+ "logit": 0.123
+ }
+ ]
+ }
+
+ LiteLLM expects (RerankResponse):
+ {
+ "results": [
+ {
+ "index": 0,
+ "relevance_score": 0.123,
+ "document": {"text": "..."} # optional
+ }
+ ]
+ }
+ """
+ try:
+ raw_response_json = raw_response.json()
+ except Exception:
+ raise BaseLLMException(
+ status_code=raw_response.status_code,
+ message=raw_response.text,
+ headers=raw_response.headers,
+ )
+
+ # Parse as NvidiaNimRerankResponse
+ nvidia_response: NvidiaNimRerankResponse = raw_response_json
+
+ # Transform Nvidia NIM response to LiteLLM format
+ results: List[RerankResponseResult] = []
+ rankings = nvidia_response.get("rankings", [])
+
+ # Get original documents from request if we need to include them
+ original_passages: List[NvidiaNimPassageObject] = request_data.get("passages", [])
+
+ for ranking in rankings:
+ result_item: RerankResponseResult = {
+ "index": ranking["index"],
+ "relevance_score": ranking["logit"],
+ }
+
+ # Include document if it was in the original request
+ index: int = ranking["index"]
+ if index < len(original_passages):
+ result_item["document"] = {"text": original_passages[index]["text"]} # type: ignore
+
+ results.append(result_item)
+
+ # Construct metadata with billed_units
+ # Nvidia NIM uses "usage" field with "total_tokens"
+ usage = raw_response_json.get("usage", {})
+ total_tokens = usage.get("total_tokens", 0)
+
+ billed_units: RerankBilledUnits = {
+ "total_tokens": total_tokens if total_tokens > 0 else len(results)
+ }
+
+ meta: RerankResponseMeta = {
+ "billed_units": billed_units
+ }
+
+ return RerankResponse(
+ id=raw_response_json.get("id") or str(uuid.uuid4()),
+ results=results,
+ meta=meta,
+ )
+
+ def get_error_class(
+ self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
+ ) -> BaseLLMException:
+ return BaseLLMException(
+ status_code=status_code,
+ message=error_message,
+ headers=headers,
+ )
+
diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index 1b8f7f7c08..521251f50c 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -18413,6 +18413,20 @@
"mode": "rerank",
"output_cost_per_token": 0.0
},
+ "nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3": {
+ "input_cost_per_query": 0.0,
+ "input_cost_per_token": 0.0,
+ "litellm_provider": "nvidia_nim",
+ "mode": "rerank",
+ "output_cost_per_token": 0.0
+ },
+ "nvidia_nim/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",
diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py
index 5b9337e852..fc45266536 100644
--- a/litellm/rerank_api/main.py
+++ b/litellm/rerank_api/main.py
@@ -12,7 +12,7 @@ from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.llms.together_ai.rerank.handler import TogetherAIRerank
from litellm.rerank_api.rerank_utils import get_optional_rerank_params
from litellm.secret_managers.main import get_secret, get_secret_str
-from litellm.types.rerank import OptionalRerankParams, RerankResponse
+from litellm.types.rerank import RerankResponse
from litellm.types.router import *
from litellm.utils import ProviderConfigManager, client, exception_type
@@ -136,7 +136,7 @@ def rerank( # noqa: PLR0915
)
)
- optional_rerank_params: OptionalRerankParams = get_optional_rerank_params(
+ optional_rerank_params: Dict = get_optional_rerank_params(
rerank_provider_config=rerank_provider_config,
model=model,
drop_params=kwargs.get("drop_params") or litellm.drop_params or False,
@@ -173,7 +173,7 @@ def rerank( # noqa: PLR0915
)
# Implement rerank logic here based on the custom_llm_provider
- if _custom_llm_provider == "cohere" or _custom_llm_provider == "litellm_proxy":
+ if _custom_llm_provider == litellm.LlmProviders.COHERE or _custom_llm_provider == litellm.LlmProviders.LITELLM_PROXY:
# Implement Cohere rerank logic
api_key: Optional[str] = (
dynamic_api_key or optional_params.api_key or litellm.api_key
@@ -205,7 +205,7 @@ def rerank( # noqa: PLR0915
client=client,
model_response=model_response,
)
- elif _custom_llm_provider == "azure_ai":
+ elif _custom_llm_provider == litellm.LlmProviders.AZURE_AI:
api_base = (
dynamic_api_base # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there
or optional_params.api_base
@@ -226,7 +226,7 @@ def rerank( # noqa: PLR0915
client=client,
model_response=model_response,
)
- elif _custom_llm_provider == "infinity":
+ elif _custom_llm_provider == litellm.LlmProviders.INFINITY:
# Implement Infinity rerank logic
api_key = dynamic_api_key or optional_params.api_key or litellm.api_key
@@ -256,7 +256,7 @@ def rerank( # noqa: PLR0915
client=client,
model_response=model_response,
)
- elif _custom_llm_provider == "together_ai":
+ elif _custom_llm_provider == litellm.LlmProviders.TOGETHER_AI:
# Implement Together AI rerank logic
api_key = (
dynamic_api_key
@@ -282,7 +282,7 @@ def rerank( # noqa: PLR0915
api_key=api_key,
_is_async=_is_async,
)
- elif _custom_llm_provider == "jina_ai":
+ elif _custom_llm_provider == litellm.LlmProviders.JINA_AI:
if dynamic_api_key is None:
raise ValueError(
"Jina AI API key is required, please set 'JINA_AI_API_KEY' in your environment"
@@ -309,7 +309,35 @@ def rerank( # noqa: PLR0915
client=client,
model_response=model_response,
)
- elif _custom_llm_provider == "bedrock":
+ elif _custom_llm_provider == litellm.LlmProviders.NVIDIA_NIM:
+ if dynamic_api_key is None:
+ raise ValueError(
+ "Nvidia NIM API key is required, please set 'NVIDIA_NIM_API_KEY' in your environment"
+ )
+
+ # Note: For rerank, the base URL is different from chat/embeddings
+ # Rerank uses ai.api.nvidia.com instead of integrate.api.nvidia.com
+ api_base = (
+ optional_params.api_base
+ or get_secret("NVIDIA_NIM_API_BASE") # type: ignore
+ or "https://ai.api.nvidia.com" # Default for rerank
+ )
+
+ response = base_llm_http_handler.rerank(
+ model=model,
+ custom_llm_provider=_custom_llm_provider,
+ optional_rerank_params=optional_rerank_params,
+ logging_obj=litellm_logging_obj,
+ provider_config=rerank_provider_config,
+ timeout=optional_params.timeout,
+ api_key=dynamic_api_key or optional_params.api_key,
+ api_base=api_base,
+ _is_async=_is_async,
+ headers=headers or litellm.headers or {},
+ client=client,
+ model_response=model_response,
+ )
+ elif _custom_llm_provider == litellm.LlmProviders.BEDROCK:
api_base = (
dynamic_api_base
or optional_params.api_base
@@ -331,7 +359,7 @@ def rerank( # noqa: PLR0915
logging_obj=litellm_logging_obj,
client=client,
)
- elif _custom_llm_provider == "hosted_vllm":
+ elif _custom_llm_provider == litellm.LlmProviders.HOSTED_VLLM:
# Implement Hosted VLLM rerank logic
api_key = (
dynamic_api_key
@@ -365,7 +393,7 @@ def rerank( # noqa: PLR0915
model_response=model_response,
)
- elif _custom_llm_provider == "deepinfra":
+ elif _custom_llm_provider == litellm.LlmProviders.DEEPINFRA:
api_key = (
dynamic_api_key
or optional_params.api_key
diff --git a/litellm/rerank_api/rerank_utils.py b/litellm/rerank_api/rerank_utils.py
index f70ec015b6..38e599ef82 100644
--- a/litellm/rerank_api/rerank_utils.py
+++ b/litellm/rerank_api/rerank_utils.py
@@ -1,7 +1,6 @@
from typing import Any, Dict, List, Optional, Union
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
-from litellm.types.rerank import OptionalRerankParams
def get_optional_rerank_params(
@@ -17,7 +16,7 @@ def get_optional_rerank_params(
max_chunks_per_doc: Optional[int] = None,
max_tokens_per_doc: Optional[int] = None,
non_default_params: Optional[dict] = None,
-) -> OptionalRerankParams:
+) -> Dict:
all_non_default_params = non_default_params or {}
if query is not None:
all_non_default_params["query"] = query
diff --git a/litellm/utils.py b/litellm/utils.py
index ad4d2ff11b..f963e8c944 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -7208,6 +7208,8 @@ class ProviderConfigManager:
return litellm.HuggingFaceRerankConfig()
elif litellm.LlmProviders.DEEPINFRA == provider:
return litellm.DeepinfraRerankConfig()
+ elif litellm.LlmProviders.NVIDIA_NIM == provider:
+ return litellm.NvidiaNimRerankConfig()
return litellm.CohereRerankConfig()
@staticmethod
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index 1b8f7f7c08..521251f50c 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -18413,6 +18413,20 @@
"mode": "rerank",
"output_cost_per_token": 0.0
},
+ "nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3": {
+ "input_cost_per_query": 0.0,
+ "input_cost_per_token": 0.0,
+ "litellm_provider": "nvidia_nim",
+ "mode": "rerank",
+ "output_cost_per_token": 0.0
+ },
+ "nvidia_nim/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",
diff --git a/tests/llm_translation/base_rerank_unit_tests.py b/tests/llm_translation/base_rerank_unit_tests.py
index cff4a02753..ac62dbbd8b 100644
--- a/tests/llm_translation/base_rerank_unit_tests.py
+++ b/tests/llm_translation/base_rerank_unit_tests.py
@@ -83,6 +83,14 @@ class BaseLLMRerankTest(ABC):
"""Must return the custom llm provider"""
pass
+ def get_expected_cost(self) -> float:
+ """
+ Override this method to set the expected cost for the rerank call.
+ Default is None, which means the test will check cost > 0.
+ Return 0.0 for free models.
+ """
+ return None
+
@pytest.mark.asyncio()
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_basic_rerank(self, sync_mode):
@@ -105,7 +113,18 @@ class BaseLLMRerankTest(ABC):
assert response.results is not None
assert response._hidden_params["response_cost"] is not None
- assert response._hidden_params["response_cost"] > 0
+
+ # Check expected cost
+ expected_cost = self.get_expected_cost()
+ if expected_cost is not None:
+ # If expected cost is specified, check exact match or >= for 0
+ if expected_cost == 0.0:
+ assert response._hidden_params["response_cost"] >= 0
+ else:
+ assert response._hidden_params["response_cost"] == expected_cost
+ else:
+ # Default behavior: cost should be greater than 0
+ assert response._hidden_params["response_cost"] > 0
assert_response_shape(
response=response, custom_llm_provider=custom_llm_provider.value
diff --git a/tests/llm_translation/test_nvidia_nim.py b/tests/llm_translation/test_nvidia_nim.py
index 89513c58b0..1705871258 100644
--- a/tests/llm_translation/test_nvidia_nim.py
+++ b/tests/llm_translation/test_nvidia_nim.py
@@ -17,6 +17,8 @@ from unittest.mock import patch, MagicMock, AsyncMock
import litellm
from litellm import Choices, Message, ModelResponse, EmbeddingResponse, Usage
from litellm import completion
+from base_rerank_unit_tests import BaseLLMRerankTest
+import litellm
def test_completion_nvidia_nim():
@@ -181,3 +183,16 @@ def test_chat_completion_nvidia_nim_with_tools():
assert request_body["tools"] == tools
assert request_body["tool_choice"] == "auto"
assert request_body["parallel_tool_calls"] == True
+
+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",
+ }
+
+ def get_expected_cost(self) -> float:
+ """Nvidia NIM rerank models are free (cost = 0.0)"""
+ return 0.0
\ No newline at end of file