From 601da4a3d1324bfcc9cc0eed94d073493f0f26cf Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 8 Dec 2025 15:25:23 -0800 Subject: [PATCH] [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/providers/nvidia_nim_rerank.md | 117 ++++++++++++++++-- litellm/__init__.py | 1 + .../llms/nvidia_nim/rerank/common_utils.py | 28 +++++ .../rerank/ranking_transformation.py | 75 +++++++++++ ...odel_prices_and_context_window_backup.json | 7 ++ litellm/utils.py | 6 +- model_prices_and_context_window.json | 7 ++ tests/llm_translation/test_nvidia_nim.py | 65 +++++++++- 8 files changed, 293 insertions(+), 13 deletions(-) create mode 100644 litellm/llms/nvidia_nim/rerank/common_utils.py create mode 100644 litellm/llms/nvidia_nim/rerank/ranking_transformation.py diff --git a/docs/my-website/docs/providers/nvidia_nim_rerank.md b/docs/my-website/docs/providers/nvidia_nim_rerank.md index 7373014a96..d28f056c24 100644 --- a/docs/my-website/docs/providers/nvidia_nim_rerank.md +++ b/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/` | Default for most rerank models | +| `/v1/ranking` | `nvidia_nim/ranking/` | 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( -## 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/) - diff --git a/litellm/__init__.py b/litellm/__init__.py index d2766be03c..34bfc77898 100644 --- a/litellm/__init__.py +++ b/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 diff --git a/litellm/llms/nvidia_nim/rerank/common_utils.py b/litellm/llms/nvidia_nim/rerank/common_utils.py new file mode 100644 index 0000000000..2bd8c123c9 --- /dev/null +++ b/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() + diff --git a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py new file mode 100644 index 0000000000..72e3c039d4 --- /dev/null +++ b/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/" 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, + ) + diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index fde60a9237..f81bd214c5 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/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", diff --git a/litellm/utils.py b/litellm/utils.py index d77607fd3e..00fc61b228 100644 --- a/litellm/utils.py +++ b/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: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index fde60a9237..f81bd214c5 100644 --- a/model_prices_and_context_window.json +++ b/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", diff --git a/tests/llm_translation/test_nvidia_nim.py b/tests/llm_translation/test_nvidia_nim.py index 1705871258..d0462efa6d 100644 --- a/tests/llm_translation/test_nvidia_nim.py +++ b/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/" 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: