Address reviews

Author: ChrisJar

nemo_retriever/src/nemo_retriever/rerank/__init__.py

@@ -13,7 +13,7 @@
 rerank_hits
     Convenience function to rerank a list of LanceDB hit dicts for a single
     query string, using either a local ``NemotronRerankV2`` model or a remote
-    vLLM / NIM ranking endpoint.
+    vLLM / NIM ``/v1/ranking`` endpoint.
 """
 
 from .rerank import NemotronRerankActor, NemotronRerankCPUActor, NemotronRerankGPUActor, rerank_hits

nemo_retriever/src/nemo_retriever/rerank/rerank.py

@@ -12,7 +12,7 @@
 Remote endpoint
 ---------------
 When ``invoke_url`` is set the actor/function calls a vLLM (>=0.14) or NIM
-server that exposes the OpenAI-compatible ranking REST API. The helper accepts
+server that exposes the NIM ranking REST API. The helper accepts
 either a fully qualified ``.../reranking`` URL or a base URL and appends
 ``/v1/ranking`` automatically::
 

nemo_retriever/tests/test_nemotron_rerank_v2.py

@@ -402,7 +402,7 @@ def test_original_hit_keys_preserved(self):
 
 
 class TestRerankViaEndpoint:
-    def test_posts_to_rerank_url(self):
+    def test_posts_to_ranking_url(self):
         from nemo_retriever.rerank.rerank import _rerank_via_endpoint
 
         mock_resp = MagicMock()
@@ -426,6 +426,8 @@ def test_posts_to_rerank_url(self):
         call_kwargs = mock_post.call_args
         assert call_kwargs[0][0] == "http://localhost:8000/v1/ranking"
         assert call_kwargs[1]["json"]["query"] == {"text": "What is ML?"}
+        assert call_kwargs[1]["json"]["truncate"] == "END"
+        assert call_kwargs[1]["json"]["passages"][0] == {"text": "Machine learning is…"}
         assert len(call_kwargs[1]["json"]["passages"]) == 2
 
         assert scores == [0.9, 0.3]
@@ -459,7 +461,7 @@ def test_authorization_header_sent_when_api_key_provided(self):
         from nemo_retriever.rerank.rerank import _rerank_via_endpoint
 
         mock_resp = MagicMock()
-        mock_resp.json.return_value = {"results": [{"index": 0, "relevance_score": 1.0}]}
+        mock_resp.json.return_value = {"rankings": [{"index": 0, "logit": 1.0}]}
         mock_resp.raise_for_status = MagicMock()
 
         with patch("requests.post", return_value=mock_resp) as mock_post:
@@ -477,7 +479,7 @@ def test_trailing_slash_on_endpoint_normalized(self):
         from nemo_retriever.rerank.rerank import _rerank_via_endpoint
 
         mock_resp = MagicMock()
-        mock_resp.json.return_value = {"results": [{"index": 0, "relevance_score": 0.5}]}
+        mock_resp.json.return_value = {"rankings": [{"index": 0, "logit": 0.5}]}
         mock_resp.raise_for_status = MagicMock()
 
         with patch("requests.post", return_value=mock_resp) as mock_post:
@@ -490,7 +492,7 @@ def test_top_n_not_in_payload_when_not_specified(self):
         from nemo_retriever.rerank.rerank import _rerank_via_endpoint
 
         mock_resp = MagicMock()
-        mock_resp.json.return_value = {"results": [{"index": 0, "relevance_score": 0.5}]}
+        mock_resp.json.return_value = {"rankings": [{"index": 0, "logit": 0.5}]}
         mock_resp.raise_for_status = MagicMock()
 
         with patch("requests.post", return_value=mock_resp) as mock_post:
@@ -532,8 +534,8 @@ def test_actor_call_scores_dataframe(self):
         mock_resp = MagicMock()
         mock_resp.raise_for_status = MagicMock()
         mock_resp.json.side_effect = [
-            {"results": [{"index": 0, "relevance_score": 0.9}]},
-            {"results": [{"index": 0, "relevance_score": 0.4}]},
+            {"rankings": [{"index": 0, "logit": 0.9}]},
+            {"rankings": [{"index": 0, "logit": 0.4}]},
         ]
 
         with patch("requests.post", return_value=mock_resp):
@@ -552,8 +554,8 @@ def test_actor_call_sorts_descending_by_default(self):
         mock_resp = MagicMock()
         mock_resp.raise_for_status = MagicMock()
         mock_resp.json.side_effect = [
-            {"results": [{"index": 0, "relevance_score": 0.1}]},
-            {"results": [{"index": 0, "relevance_score": 0.9}]},
+            {"rankings": [{"index": 0, "logit": 0.1}]},
+            {"rankings": [{"index": 0, "logit": 0.9}]},
         ]
 
         with patch("requests.post", return_value=mock_resp):
@@ -587,7 +589,7 @@ def test_actor_custom_score_column_name(self):
 
         mock_resp = MagicMock()
         mock_resp.raise_for_status = MagicMock()
-        mock_resp.json.return_value = {"results": [{"index": 0, "relevance_score": 0.7}]}
+        mock_resp.json.return_value = {"rankings": [{"index": 0, "logit": 0.7}]}
 
         with patch("requests.post", return_value=mock_resp):
             out = actor(df)