Merge d13c9be510 into e16d79d635

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,

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"

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",
     }