diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index 12377d8f854..508cde83ce1 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -34,7 +34,6 @@ class NvidiaNimRerankRequest(TypedDict, total=False): query: Required[NvidiaNimQueryObject] passages: Required[list[NvidiaNimPassageObject]] truncate: Literal["NONE", "END"] - top_k: int class NvidiaNimRankingResult(TypedDict): @@ -70,7 +69,9 @@ class NvidiaNimRerankConfig(BaseRerankConfig): SUPPORTED_PASSAGE_FIELDS: tuple[str, ...] = ("text",) def __init__(self) -> None: - pass + # Provider configs are created per request; retain the Cohere result + # limit here rather than forwarding unsupported top_k to NVIDIA. + self._client_side_top_n: int | None = None def _get_clean_model_name(self, model: str) -> str: """Strip 'nvidia_nim/' prefix from model name if present.""" @@ -137,8 +138,8 @@ class NvidiaNimRerankConfig(BaseRerankConfig): """ Map Cohere/OpenAI rerank params to Nvidia NIM format. - Parameter mapping: - - top_n (Cohere) -> top_k (Nvidia) + Cohere's top_n is applied client-side: the NVIDIA retrieval endpoint + rejects the top_k field. Nvidia NIM specific params (passed through as-is from non_default_params): - truncate: How to truncate input if too long (NONE, END) @@ -148,9 +149,9 @@ class NvidiaNimRerankConfig(BaseRerankConfig): "documents": documents, } - # Map Cohere's top_n to Nvidia's top_k + # Keep top_n for client-side truncation after NVIDIA returns rankings. if top_n is not None: - optional_nvidia_nim_rerank_params["top_k"] = top_n + optional_nvidia_nim_rerank_params["top_n"] = top_n # Pass through Nvidia-specific params from non_default_params if non_default_params: @@ -200,17 +201,26 @@ class NvidiaNimRerankConfig(BaseRerankConfig): Nvidia NIM expects: - query as {text: "..."} - documents as passages: [{text: "..."}, ...] - - Optional: truncate (NONE or END), top_k + - Optional: truncate (NONE or END) - 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. + top_n is applied locally after receiving the rankings. Both top_n and + top_k are excluded from the NVIDIA request body. """ 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") + # Validate the boundary as Mapping[str, object] to avoid leaking an + # Unknown type from the provider's untyped optional params. + top_n: Final[object] = _JSON_OBJECT.validate_python( + {"top_n": optional_rerank_params.get("top_n")} + )["top_n"] + if top_n is not None: + if isinstance(top_n, bool) or not isinstance(top_n, int) or top_n < 1: + raise ValueError(f"top_n must be a positive integer, got: {top_n!r}") + self._client_side_top_n = top_n + query: Final = optional_rerank_params["query"] documents: Final = optional_rerank_params["documents"] @@ -254,10 +264,6 @@ class NvidiaNimRerankConfig(BaseRerankConfig): "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: - request_data["top_k"] = optional_rerank_params.get("top_k") - # 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: @@ -347,6 +353,14 @@ class NvidiaNimRerankConfig(BaseRerankConfig): meta: Final[RerankResponseMeta] = {"billed_units": billed_units} + # The hosted /v1/retrieval/{model}/reranking API rejects top_k. + # Sort before limiting in case the provider returns unsorted rankings. + # Work with the typed local results, not the Pydantic response attribute. + top_n: Final[int | None] = self._client_side_top_n + if top_n is not None and len(results) > top_n: + results.sort(key=lambda result: result["relevance_score"], reverse=True) + del results[top_n:] + return RerankResponse( id=_STR.validate_python(payload.get("id") or str(uuid.uuid4())), results=results, diff --git a/tests/unit/llms/base_llm/rerank/test_provider_rerank_translation.py b/tests/unit/llms/base_llm/rerank/test_provider_rerank_translation.py index 8b630f6b946..2f23789b44c 100644 --- a/tests/unit/llms/base_llm/rerank/test_provider_rerank_translation.py +++ b/tests/unit/llms/base_llm/rerank/test_provider_rerank_translation.py @@ -138,7 +138,8 @@ def _assert_translated_request(case: _Case, body: Mapping[str, JsonValue]) -> No assert body["model"] == "nvidia/llama-3.2-nv-rerankqa-1b-v2" assert body["query"] == {"text": "hello"} assert body["passages"] == [{"text": "hello"}, {"text": "world"}] - assert body["top_k"] == 2 + assert "top_k" not in body + assert "top_n" not in body else: assert body["model"] == "jina-reranker-v2-base-multilingual" assert body["query"] == "hello" diff --git a/tests/unit/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py b/tests/unit/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py index 39ba396d8d2..1f805342b94 100644 --- a/tests/unit/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py +++ b/tests/unit/llms/nvidia_nim/rerank/test_nvidia_nim_rerank_transformation.py @@ -189,15 +189,14 @@ async def test_nvidia_nim_ranking_endpoint_image_documents_and_top_n(): class TestNvidiaNimRetrievalRerankRequestTransform: """ - The default /v1/retrieval/{model}/reranking route keeps its existing - contract: top_n still maps to top_k, and structured documents keep the - prior text-only passage behavior. + The /v1/retrieval/{model}/reranking route omits unsupported top_k and + applies top_n client-side while preserving text-only passage behavior. """ - def _build_request(self, documents, top_n=None): + def _build_request(self, documents, top_n=None, non_default_params=None): config = NvidiaNimRerankConfig() optional_params = config.map_cohere_rerank_params( - non_default_params=None, + non_default_params=non_default_params, model="nvidia/llama-3_2-nv-rerankqa-1b-v2", drop_params=False, query="which passage shows a cat?", @@ -210,11 +209,20 @@ class TestNvidiaNimRetrievalRerankRequestTransform: headers={}, ) - def test_top_n_still_maps_to_top_k(self): + def test_top_n_is_not_sent_as_top_k(self): request_data = self._build_request(["a", "b"], top_n=1) - assert request_data["top_k"] == 1 + assert "top_k" not in request_data assert "top_n" not in request_data + def test_provider_specific_top_k_is_not_forwarded(self): + request_data = self._build_request(["a", "b"], non_default_params={"top_k": 300}) + assert "top_k" not in request_data + + @pytest.mark.parametrize("invalid_top_n", [0, -1, 1.5, "2", True]) + def test_invalid_top_n_raises_value_error(self, invalid_top_n): + with pytest.raises(ValueError, match="top_n"): + self._build_request(["a", "b"], top_n=invalid_top_n) + def test_string_documents_unchanged(self): request_data = self._build_request(["passage one", "passage two"]) assert request_data["passages"] == [ @@ -243,6 +251,85 @@ class TestNvidiaNimRetrievalRerankRequestTransform: assert request_data["passages"] == [{"text": json.dumps(doc)}] +def _rerank_retrieval_with_top_n(top_n, rankings): + config = NvidiaNimRerankConfig() + params = config.map_cohere_rerank_params( + non_default_params=None, + model="nvidia/llama-nemotron-rerank-vl-1b-v2", + drop_params=False, + query="test", + documents=["first", "second", "third"], + top_n=top_n, + ) + request = config.transform_rerank_request( + model="nvidia/llama-nemotron-rerank-vl-1b-v2", + optional_rerank_params=params, + headers={}, + ) + response = config.transform_rerank_response( + model="nvidia/llama-nemotron-rerank-vl-1b-v2", + raw_response=httpx.Response(200, json={"rankings": rankings}), + model_response=RerankResponse(), + logging_obj=MagicMock(), + request_data=request, + ) + return request, response + + +class TestNvidiaNimRetrievalRerankResponseTransform: + RANKINGS = [ + {"index": 0, "logit": 0.10}, + {"index": 1, "logit": 0.90}, + {"index": 2, "logit": 0.50}, + ] + + def test_top_n_truncates_by_relevance(self): + request, response = _rerank_retrieval_with_top_n(2, self.RANKINGS) + assert "top_k" not in request + assert [result["index"] for result in response.results] == [1, 2] + assert response.meta == {"billed_units": {"total_tokens": 3}} + + def test_top_n_greater_than_document_count_keeps_all(self): + _, response = _rerank_retrieval_with_top_n(300, self.RANKINGS) + assert len(response.results) == 3 + + def test_no_top_n_keeps_all_results(self): + _, response = _rerank_retrieval_with_top_n(None, self.RANKINGS) + assert len(response.results) == 3 + + +@pytest.mark.asyncio() +async def test_nvidia_nim_retrieval_endpoint_does_not_send_top_k(): + """Mocked HTTP transport regression: NVIDIA rejects top_k in the body.""" + mock_response = AsyncMock() + mock_response.json = lambda: { + "rankings": [{"index": 0, "logit": 0.2}, {"index": 1, "logit": 0.9}] + } + 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: + response = await litellm.arerank( + model="nvidia_nim/nvidia/llama-nemotron-rerank-vl-1b-v2", + query="which document is more relevant?", + documents=["unrelated text", "relevant text"], + top_n=1, + api_key="fake-api-key", + ) + mock_post.assert_called_once() + assert mock_post.call_args.kwargs["url"] == ( + "https://ai.api.nvidia.com/v1/retrieval/nvidia/llama-nemotron-rerank-vl-1b-v2/reranking" + ) + body = json.loads(mock_post.call_args.kwargs["data"]) + assert "top_k" not in body + assert "top_n" not in body + assert body["passages"] == [{"text": "unrelated text"}, {"text": "relevant text"}] + assert [result["index"] for result in response.results] == [1] + + def _transform_retrieval_response(payload: object) -> RerankResponse: return NvidiaNimRerankConfig().transform_rerank_response( model="nvidia/llama-3_2-nv-rerankqa-1b-v2", @@ -309,7 +396,7 @@ def test_transform_rerank_response_usage_shape_errors_do_not_echo_the_payload(): assert "leaked usage text" not in str(exc_info.value) -def test_map_cohere_rerank_params_passes_provider_params_through_and_maps_top_n(): +def test_map_cohere_rerank_params_preserves_top_n_for_local_truncation(): params = NvidiaNimRerankConfig().map_cohere_rerank_params( non_default_params={"truncate": "END"}, model="nvidia/llama-3_2-nv-rerankqa-1b-v2", @@ -322,6 +409,6 @@ def test_map_cohere_rerank_params_passes_provider_params_through_and_maps_top_n( assert params == { "query": "which passage shows a cat?", "documents": ["a", {"text": "b"}], - "top_k": 1, + "top_n": 1, "truncate": "END", }