Dev/mahikaw/agentic enhancements

nemo_retriever/src/nemo_retriever/cli/query/app.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import json
+import logging
 import os
 from typing import cast
 
@@ -23,6 +24,7 @@
 )
 from nemo_retriever.common.vdb.records import RetrievalHit
 from nemo_retriever.query.options import (
+    DEFAULT_AGENTIC_LLM_MODEL,
     QueryAgenticOptions,
     QueryEmbedOptions,
     QueryRerankOptions,
@@ -184,7 +186,7 @@ def _local_command(
     output_format: opts.OutputFormatOption = "hits",
     max_text_chars: opts.MaxTextCharsOption = None,
     agentic: opts.AgenticOption = False,
-    agentic_llm_model: opts.AgenticLlmModelOption = None,
+    agentic_llm_model: opts.AgenticLlmModelOption = DEFAULT_AGENTIC_LLM_MODEL,
     agentic_invoke_url: opts.AgenticInvokeUrlOption = None,
     agentic_reasoning_effort: opts.AgenticReasoningEffortOption = "high",
     agentic_backend_top_k: opts.AgenticBackendTopKOption = 20,
@@ -243,8 +245,15 @@ def _local_command(
                     temperature=agentic_temperature,
                 ),
             )
-            with quiet_capture():
-                ranked = query_agentic_documents(request)
+            # Agentic retrieval is a multi-step ReAct loop, not a single dense pass, so
+            # surface per-query/step progress instead of running blind. Mirrors the
+            # `pipeline run` logging setup: INFO to stderr (stdout stays clean JSON).
+            logging.basicConfig(
+                level=logging.INFO,
+                format="%(asctime)s %(levelname)s %(name)s: %(message)s",
+                force=True,
+            )
+            ranked = query_agentic_documents(request)
             typer.echo(json.dumps(ranked, indent=2, sort_keys=True, default=str))
             return
 

nemo_retriever/src/nemo_retriever/cli/query/options.py

@@ -9,6 +9,7 @@
 import typer
 
 from nemo_retriever.models import VL_EMBED_MODEL, VL_RERANK_MODEL
+from nemo_retriever.query.options import DEFAULT_AGENTIC_LLM_MODEL
 
 DEFAULT_EMBED_MODEL = VL_EMBED_MODEL
 DEFAULT_RERANK_MODEL = VL_RERANK_MODEL
@@ -160,7 +161,11 @@
     str | None,
     typer.Option(
         "--agentic-llm-model",
-        help="Chat model the agent drives. Required when --agentic is set.",
+        envvar="NEMO_RETRIEVER_AGENTIC_LLM_MODEL",
+        help=(
+            f"Chat model the agent drives. Defaults to {DEFAULT_AGENTIC_LLM_MODEL}; "
+            "override here or via NEMO_RETRIEVER_AGENTIC_LLM_MODEL."
+        ),
     ),
 ]
 AgenticInvokeUrlOption = Annotated[

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():
@@ -192,16 +197,22 @@ def __init__(
             },
         )
         self._lock = threading.Lock()
+        # Full hits keyed by doc_id, captured at the retrieval boundary (before the
+        # agent loop reduces them to doc_id/text/score) so the final output can
+        # re-hydrate the metadata the agent drops. Reset at the start of retrieve().
+        self._hit_cache: dict[str, dict[str, Any]] = {}
 
     def retrieve(self, query_ids: Sequence[str], query_texts: Sequence[str]) -> pd.DataFrame:
         """Return selected ranked documents for each query.
 
-        The output schema matches ``SelectionAgentOperator``: ``query_id``,
-        ``doc_id``, ``rank``, and ``message``.
+        Columns: ``query_id``, ``doc_id``, ``rank``, ``message``, ``result_source``,
+        and ``hit`` — the full ``RetrievalHit`` (text/metadata/source/page/…) captured
+        at retrieval time and re-hydrated by ``doc_id`` (``{}`` when unavailable).
         """
 
         if len(query_ids) != len(query_texts):
             raise ValueError("query_ids and query_texts must have the same length.")
+        self._hit_cache.clear()
 
         from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator
         from nemo_retriever.operators.graph_ops.rrf_aggregator_operator import RRFAggregatorOperator
@@ -254,7 +265,11 @@ def retrieve(self, query_ids: Sequence[str], query_texts: Sequence[str]) -> pd.D
             [str(query_text) for query_text in query_texts],
             top_k=target_top_k,
         )
-        return _raw_hits_to_agentic_result([str(query_id) for query_id in query_ids], raw_hits)
+        result = _raw_hits_to_agentic_result([str(query_id) for query_id in query_ids], raw_hits)
+        # Re-hydrate the metadata dropped at the agent boundary by joining the
+        # selected doc_ids back to the full hits captured in _retrieve_for_agent.
+        result["hit"] = [self._hit_cache.get(str(doc_id), {}) for doc_id in result["doc_id"]]
+        return result
 
     def _retrieve_for_agent(self, query_text: str, top_k: int) -> list[dict[str, Any]]:
         """Retriever callback used by ``ReActAgentOperator``.
@@ -266,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)
@@ -280,6 +304,10 @@ def _retrieve_for_agent(self, query_text: str, top_k: int) -> list[dict[str, Any
             )
             if not doc_id:
                 continue
+            # Capture the full hit (minus the embedding) for output re-hydration;
+            # first occurrence per doc_id wins. Locked: shared across ReAct workers.
+            with self._lock:
+                self._hit_cache.setdefault(doc_id, {k: v for k, v in hit_dict.items() if k != "vector"})
             text = str(hit_dict.get("text", ""))
             if int(self._cfg.text_truncation) > 0:
                 text = text[: int(self._cfg.text_truncation)]

nemo_retriever/src/nemo_retriever/query/options.py

@@ -10,6 +10,10 @@
 
 QueryRetrievalMode = Literal["auto", "dense", "hybrid", "sparse"]
 
+#: Default chat model the agentic ReAct loop drives. Override per-call via
+#: ``--agentic-llm-model`` or the ``NEMO_RETRIEVER_AGENTIC_LLM_MODEL`` env var.
+DEFAULT_AGENTIC_LLM_MODEL = "nvidia/llama-3.3-nemotron-super-49b-v1.5"
+
 
 @dataclass(frozen=True)
 class QueryRetrievalOptions:
@@ -54,7 +58,7 @@ class QueryAgenticOptions:
     """Options for the agentic (ReAct) retrieval strategy."""
 
     enabled: bool = False
-    llm_model: str | None = None
+    llm_model: str | None = DEFAULT_AGENTIC_LLM_MODEL
     invoke_url: str | None = None
     reasoning_effort: str | None = "high"
     backend_top_k: int = 20

nemo_retriever/src/nemo_retriever/query/workflow.py

@@ -144,17 +144,16 @@ def query_documents(request: QueryRequest) -> list[RetrievalHit]:
 
 
 def agentic_query_documents(request: QueryRequest) -> list[dict[str, Any]]:
-    """Run agentic (ReAct) retrieval for a single query and return the agent's
-    ranked document IDs.
-
-    Unlike the dense ``query_documents`` path (which returns enriched hits with
-    text), the agent operates at the document-ID granularity of the configured
-    index, so the result is the ranked ``doc_id`` list the agent selected,
-    annotated with the source that produced it (``final_results`` / ``rrf`` /
-    ``selection_agent``). The LanceDB ``uri``/``table_name``, embedding config,
-    and (when ``--rerank`` is enabled) reranker config are passed straight
-    through to the wrapped ``Retriever`` that backs the agent's ``retrieve``
-    tool. Reranking therefore applies per agent retrieval hop.
+    """Run agentic (ReAct) retrieval for a single query and return ranked hits.
+
+    Each result carries the full ``RetrievalHit`` metadata (``text``, ``source``,
+    ``page_number``, ``pdf_page``, ``metadata``, …) — re-hydrated by ``doc_id`` from
+    the hits captured at retrieval time — plus the agentic annotations ``doc_id``,
+    ``rank``, and ``result_source`` (``final_results`` / ``rrf`` / ``selection_agent``).
+    This matches the metadata the dense ``query_documents`` path returns. The LanceDB
+    ``uri``/``table_name``, embedding config, and (when ``--rerank`` is enabled)
+    reranker config are passed through to the wrapped ``Retriever`` that backs the
+    agent's ``retrieve`` tool; reranking applies per agent retrieval hop.
     """
     from nemo_retriever.query.agentic import AgenticRetrievalConfig, AgenticRetriever
 
@@ -166,6 +165,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,
@@ -193,13 +195,16 @@ def agentic_query_documents(request: QueryRequest) -> list[dict[str, Any]]:
         result = result.sort_values("rank")
     ranked: list[dict[str, Any]] = []
     for _, row in result.iterrows():
-        ranked.append(
+        hit = row.get("hit") if "hit" in result.columns else None
+        enriched: dict[str, Any] = dict(hit) if isinstance(hit, dict) else {}
+        enriched.update(
             {
-                "rank": int(row.get("rank", len(ranked) + 1)),
                 "doc_id": str(row.get("doc_id", "")),
+                "rank": int(row.get("rank", len(ranked) + 1)),
                 "result_source": str(row.get("result_source", "")),
             }
         )
+        ranked.append(enriched)
         if len(ranked) >= request.retrieval.top_k:
             break
     return ranked

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",
@@ -90,10 +94,16 @@ def test_agentic_retriever_runs_graph_with_wrapped_retriever(mock_react_step, mo
     cfg = AgenticRetrievalConfig(llm_model="test-model", invoke_url="http://localhost/v1/chat/completions")
     result = AgenticRetriever(cfg, match_mode="pdf_page").retrieve(["0"], ["find doc"])
 
-    assert list(result.columns) == ["query_id", "doc_id", "rank", "message", "result_source"]
+    assert list(result.columns) == ["query_id", "doc_id", "rank", "message", "result_source", "hit"]
     assert result["query_id"].tolist() == ["0"] * 10
     assert result["doc_id"].tolist()[0] == "doc_1"
     assert result["rank"].tolist() == list(range(1, 11))
+    # doc_1 was actually retrieved, so its row re-hydrates the full hit metadata
+    # (dropped at the agent boundary) rather than just a bare doc_id.
+    doc1_hit = result.loc[result["doc_id"] == "doc_1", "hit"].iloc[0]
+    assert doc1_hit.get("source") == "/tmp/doc.pdf"
+    assert doc1_hit.get("page_number") == 1
+    assert "text" in doc1_hit
 
 
 @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step")
@@ -119,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
 

nemo_retriever/tests/test_root_query_cli.py

@@ -400,12 +400,33 @@ def retrieve(self, query_ids: Any, query_texts: Any) -> Any:
     ]
 
 
-def test_root_query_agentic_requires_llm_model() -> None:
-    """Agentic mode is inert without a chat model to drive the loop."""
+def test_root_query_agentic_defaults_llm_model(monkeypatch) -> None:
+    """--agentic without --agentic-llm-model falls back to the default 49B model."""
+    import pandas as pd
+
+    import nemo_retriever.query.agentic as agentic_retrieval
+    from nemo_retriever.query.options import DEFAULT_AGENTIC_LLM_MODEL
+
+    config_calls: list[dict[str, Any]] = []
+
+    class FakeConfig:
+        def __init__(self, **kwargs: Any) -> None:
+            config_calls.append(kwargs)
+
+    class FakeAgenticRetriever:
+        def __init__(self, cfg: Any) -> None:
+            pass
+
+        def retrieve(self, query_ids: Any, query_texts: Any) -> Any:
+            return pd.DataFrame([{"query_id": "0", "doc_id": "a.pdf", "rank": 1, "result_source": "rrf"}])
+
+    monkeypatch.setattr(agentic_retrieval, "AgenticRetrievalConfig", FakeConfig)
+    monkeypatch.setattr(agentic_retrieval, "AgenticRetriever", FakeAgenticRetriever)
+
     result = RUNNER.invoke(cli_main.app, ["query", "hello", "--agentic"])
 
-    assert result.exit_code == 1
-    assert "requires --agentic-llm-model" in result.output
+    assert result.exit_code == 0
+    assert config_calls[0]["llm_model"] == DEFAULT_AGENTIC_LLM_MODEL
 
 
 def test_root_query_agentic_plumbs_rerank_into_config(monkeypatch) -> None: