tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-18 18:18:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Feat] Add Nvidia NIM Rerank Support (#15152)

Browse Source 
* feat: add NvidiaNimRerankConfig

* fix: NvidiaNimRerankConfig

* fix: NvidiaNimRerankConfig

* fix routing to nvidia nim

* docs nvidia nim rerank

* TestNvidiaNim

* nvidia nim rerank fixes

* fix rerank

* transform_rerank_response

* Usage with LiteLLM Proxy

* fixes linting

* NvidiaNimRerankConfig.DEFAULT_NIM_RERANK_API_BASE

* fix Custom API Base URL

* fix rerank base

* fix main.py

* fix transform

* fix linting

* map_cohere_rerank_params

* ruff fix

* linting fixes

* ruff fix
This commit is contained in:
Ishaan Jaff
2025-10-02 18:58:52 -07:00
committed by GitHub
parent 79c24be48b
commit efa782d6d2
21 changed files with 739 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/adding_provider/new_rerank_provider.md
+1 -1
View File
@@ -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)
docs/my-website/docs/providers/nvidia_nim.md
+2 -2
View File
@@ -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
docs/my-website/docs/providers/nvidia_nim_rerank.md
+261
View File
@@ -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
<Tabs>
<TabItem value="llama-1b" label="LLaMa 1B Model">
```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)
```
</TabItem>
<TabItem value="mistral-4b" label="Mistral 4B Model">
```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)
```
</TabItem>
</Tabs>
**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:
<Tabs>
<TabItem value="env" label="Environment Variable">
```bash
export NVIDIA_NIM_API_KEY="nvapi-..."
```
</TabItem>
<TabItem value="python" label="Python">
```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-...",
)
```
</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
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/)
docs/my-website/sidebars.js
+8 -1
View File
@@ -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",
litellm/__init__.py
+1
View File
@@ -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
litellm/llms/base_llm/rerank/transformation.py
+3 -3
View File
@@ -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(
litellm/llms/cohere/rerank/transformation.py
+5 -5
View File
@@ -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:
litellm/llms/cohere/rerank_v2/transformation.py
+4 -4
View File
@@ -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:
litellm/llms/custom_httpx/llm_http_handler.py
+2 -2
View File
@@ -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,
litellm/llms/deepinfra/rerank/transformation.py
+3 -3
View File
@@ -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
litellm/llms/hosted_vllm/rerank/transformation.py
+13 -14
View File
@@ -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:
litellm/llms/huggingface/rerank/transformation.py
+2 -2
View File
@@ -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():
litellm/llms/jina_ai/rerank/transformation.py
+5 -5
View File
@@ -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}
litellm/llms/nvidia_nim/rerank/transformation.py
+325
View File
@@ -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,
)
litellm/model_prices_and_context_window_backup.json
+14
View File
@@ -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",
litellm/rerank_api/main.py
+38 -10
View File
@@ -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
litellm/rerank_api/rerank_utils.py
+1 -2
View File
@@ -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
litellm/utils.py
+2
View File
@@ -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
model_prices_and_context_window.json
+14
View File
@@ -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",
tests/llm_translation/base_rerank_unit_tests.py
+20 -1
View File
@@ -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
tests/llm_translation/test_nvidia_nim.py
+15
View File
@@ -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