Honor --candidate-k/--page-dedup/--content-types in agentic retrieval

These retrieval-shaping flags were silently dropped in agentic mode. Thread them
from QueryRetrievalOptions into AgenticRetrievalConfig and forward to the per-hop
Retriever.query call (candidate_k floored at the hop's top_k), matching how the
dense path applies them.

Signed-off-by: Mahika Wason <mwason@nvidia.com>

Author: mahikaw

nemo_retriever/src/nemo_retriever/query/agentic.py

@@ -139,6 +139,11 @@ class AgenticRetrievalConfig:
     # Drives the ReAct target, the RRF/selection cut, and the per-hop fetch depth
     # (which is raised to at least this). Defaults to 10.
     top_k: int = AGENTIC_TARGET_TOP_K
+    # Per-hop retrieval shaping, forwarded to ``Retriever.query`` on every agent
+    # retrieval hop (mirrors the dense path's --candidate-k/--page-dedup/--content-types).
+    candidate_k: Optional[int] = None
+    page_dedup: bool = False
+    content_types: str | Sequence[str] | None = None
 
     def __post_init__(self) -> None:
         if self.llm_model is None or not str(self.llm_model).strip():
@@ -276,8 +281,17 @@ def _retrieve_for_agent(self, query_text: str, top_k: int) -> list[dict[str, Any
         which still run concurrently under ``num_concurrent > 1``.
         """
 
+        # candidate_k must be >= the hop's top_k, which grows as the agent paginates,
+        # so floor it at top_k when the caller requested a wider pool.
+        candidate_k = max(int(self._cfg.candidate_k), int(top_k)) if self._cfg.candidate_k else None
         with self._lock:
-            hits = self._retriever.query(str(query_text), top_k=int(top_k))
+            hits = self._retriever.query(
+                str(query_text),
+                top_k=int(top_k),
+                candidate_k=candidate_k,
+                page_dedup=bool(self._cfg.page_dedup),
+                content_types=self._cfg.content_types,
+            )
 
         docs: list[dict[str, Any]] = []
         doc_id_field = getattr(self, "_doc_id_field", None)

nemo_retriever/src/nemo_retriever/query/workflow.py

@@ -116,6 +116,9 @@ def agentic_query_documents(request: QueryRequest) -> list[dict[str, Any]]:
         "vdb_op": "lancedb",
         "vdb_kwargs": vdb_kwargs,
         "top_k": int(request.retrieval.top_k),
+        "candidate_k": request.retrieval.candidate_k,
+        "page_dedup": bool(request.retrieval.page_dedup),
+        "content_types": request.retrieval.content_types,
         "embedding_endpoint": request.embed.embed_invoke_url,
         "embedding_api_key": api_key or "",
         "llm_model": request.agentic.llm_model,

nemo_retriever/tests/test_agentic_eval.py

@@ -36,11 +36,15 @@ def __init__(self, **kwargs):
         self.kwargs = kwargs
         self.graph = kwargs.get("graph")
         self.top_k = int(kwargs.get("top_k", 10))
+        self.query_calls: list[dict] = []
 
-    def query(self, query: str, *, top_k: int | None = None):
+    def query(self, query: str, *, top_k=None, candidate_k=None, page_dedup=False, content_types=None):
         if self.graph is not None:
             return self.queries([query], top_k=top_k)[0]
         _ = query
+        self.query_calls.append(
+            {"top_k": top_k, "candidate_k": candidate_k, "page_dedup": page_dedup, "content_types": content_types}
+        )
         hits = [
             {
                 "source": "/tmp/doc.pdf",
@@ -125,6 +129,38 @@ def test_agentic_retriever_honors_top_k(mock_react_step, mock_selection_step):
     assert result["rank"].tolist() == list(range(1, 6))  # 5 rows, honoring top_k=5
 
 
+@patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step")
+@patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step")
+@patch("nemo_retriever.query.agentic.Retriever", FakeRetriever)
+def test_agentic_retriever_forwards_retrieval_knobs(mock_react_step, mock_selection_step):
+    """candidate_k/page_dedup/content_types reach the per-hop Retriever.query call."""
+    from nemo_retriever.query.agentic import AgenticRetrievalConfig, AgenticRetriever
+
+    mock_react_step.return_value = _make_tool_call_response(
+        "final_results", {"doc_ids": ["doc_1"], "message": "done", "search_successful": "true"}
+    )
+    mock_selection_step.return_value = _make_tool_call_response(
+        "log_selected_documents", {"doc_ids": ["doc_1"], "message": "best"}
+    )
+
+    cfg = AgenticRetrievalConfig(
+        llm_model="m",
+        invoke_url="http://localhost/v1/chat/completions",
+        top_k=1,
+        candidate_k=40,
+        page_dedup=True,
+        content_types="text",
+    )
+    retriever = AgenticRetriever(cfg, match_mode="pdf_page")
+    retriever.retrieve(["0"], ["find doc"])
+
+    calls = retriever._retriever.query_calls
+    assert calls, "expected at least one per-hop retriever.query call"
+    assert all(c["page_dedup"] is True for c in calls)
+    assert all(c["content_types"] == "text" for c in calls)
+    assert all(c["candidate_k"] >= c["top_k"] for c in calls)  # floored at the hop's top_k
+
+
 def test_agentic_config_requires_llm_model():
     from nemo_retriever.query.agentic import AgenticRetrievalConfig