From 55ff3757f88fea48c39dd60630c05ffc9761ea62 Mon Sep 17 00:00:00 2001 From: Mahika Wason Date: Thu, 25 Jun 2026 07:25:01 +0000 Subject: [PATCH 1/5] Surface agentic retrieval progress on stderr instead of running blind The --agentic path wrapped the run in quiet_capture(), which discards stdout/ stderr on success, so a multi-step ReAct query produced no feedback. Mirror the `pipeline run` logging setup (logging.basicConfig at INFO; defaults to stderr) so the operators' existing per-query/step/tool-call logs surface, while stdout stays clean JSON. Dense path unchanged. Signed-off-by: Mahika Wason --- nemo_retriever/src/nemo_retriever/cli/query/app.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/nemo_retriever/src/nemo_retriever/cli/query/app.py b/nemo_retriever/src/nemo_retriever/cli/query/app.py index de9f8e4606..16bea40bb8 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/app.py +++ b/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 @@ -236,8 +237,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 From fe5d1faf86e9478c2378660bcfcede8adfb78417 Mon Sep 17 00:00:00 2001 From: Mahika Wason Date: Thu, 25 Jun 2026 07:25:01 +0000 Subject: [PATCH 2/5] Return full hit metadata from agentic retrieval (re-hydrate by doc_id) Agentic retrieval reduced each hit to {doc_id,text,score} at the agent boundary and returned bare doc_ids, unlike the dense path which returns full RetrievalHit. Cache the full hit (keyed by doc_id) in _retrieve_for_agent, then re-hydrate the selected doc_ids at output time so agentic_query_documents returns RetrievalHit metadata (text/source/page_number/pdf_page/metadata/...) plus the agentic annotations (doc_id/rank/result_source). Addresses review feedback (jperez999). Signed-off-by: Mahika Wason --- .../src/nemo_retriever/query/agentic.py | 20 +++++++++++-- .../src/nemo_retriever/query/workflow.py | 28 ++++++++++--------- nemo_retriever/tests/test_agentic_eval.py | 8 +++++- 3 files changed, 39 insertions(+), 17 deletions(-) diff --git a/nemo_retriever/src/nemo_retriever/query/agentic.py b/nemo_retriever/src/nemo_retriever/query/agentic.py index aa75034470..80c1ca45d6 100644 --- a/nemo_retriever/src/nemo_retriever/query/agentic.py +++ b/nemo_retriever/src/nemo_retriever/query/agentic.py @@ -192,16 +192,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 +260,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``. @@ -280,6 +290,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)] diff --git a/nemo_retriever/src/nemo_retriever/query/workflow.py b/nemo_retriever/src/nemo_retriever/query/workflow.py index 3c93fc6095..b07f934cc8 100644 --- a/nemo_retriever/src/nemo_retriever/query/workflow.py +++ b/nemo_retriever/src/nemo_retriever/query/workflow.py @@ -95,17 +95,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 @@ -144,13 +143,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 diff --git a/nemo_retriever/tests/test_agentic_eval.py b/nemo_retriever/tests/test_agentic_eval.py index 8d8a500d2b..3a689f967f 100644 --- a/nemo_retriever/tests/test_agentic_eval.py +++ b/nemo_retriever/tests/test_agentic_eval.py @@ -90,10 +90,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") From a3ae980dbd5419a2a68c0fece3080108dc8b31a5 Mon Sep 17 00:00:00 2001 From: Mahika Wason Date: Thu, 25 Jun 2026 07:44:35 +0000 Subject: [PATCH 3/5] Default agentic LLM model to llama-3.3-nemotron-super-49b-v1.5 --agentic-llm-model is no longer required: it defaults to nvidia/llama-3.3-nemotron-super-49b-v1.5 (also settable via the NEMO_RETRIEVER_AGENTIC_LLM_MODEL env var), so --agentic works standalone. Signed-off-by: Mahika Wason --- .../src/nemo_retriever/cli/query/app.py | 3 +- .../src/nemo_retriever/cli/query/options.py | 8 +++-- .../src/nemo_retriever/query/options.py | 6 +++- nemo_retriever/tests/test_root_query_cli.py | 29 ++++++++++++++++--- 4 files changed, 38 insertions(+), 8 deletions(-) diff --git a/nemo_retriever/src/nemo_retriever/cli/query/app.py b/nemo_retriever/src/nemo_retriever/cli/query/app.py index 16bea40bb8..fa4fb5f877 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/app.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/app.py @@ -24,6 +24,7 @@ ) from nemo_retriever.common.vdb.records import RetrievalHit from nemo_retriever.query.options import ( + DEFAULT_AGENTIC_LLM_MODEL, QueryAgenticOptions, QueryEmbedOptions, QueryRerankOptions, @@ -178,7 +179,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, diff --git a/nemo_retriever/src/nemo_retriever/cli/query/options.py b/nemo_retriever/src/nemo_retriever/cli/query/options.py index a64bd802d2..065b915c70 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/options.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/options.py @@ -8,7 +8,6 @@ import typer - QueryArgument = Annotated[str, typer.Argument(..., help="Query text.")] TopKOption = Annotated[ int, @@ -152,7 +151,12 @@ 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=( + "Chat model the agent drives. Defaults to " + "nvidia/llama-3.3-nemotron-super-49b-v1.5; override here or via " + "NEMO_RETRIEVER_AGENTIC_LLM_MODEL." + ), ), ] AgenticInvokeUrlOption = Annotated[ diff --git a/nemo_retriever/src/nemo_retriever/query/options.py b/nemo_retriever/src/nemo_retriever/query/options.py index 237001a6e9..1b5ef87025 100644 --- a/nemo_retriever/src/nemo_retriever/query/options.py +++ b/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 diff --git a/nemo_retriever/tests/test_root_query_cli.py b/nemo_retriever/tests/test_root_query_cli.py index 781a49bb21..9a28d7a55d 100644 --- a/nemo_retriever/tests/test_root_query_cli.py +++ b/nemo_retriever/tests/test_root_query_cli.py @@ -399,12 +399,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: From 26d60d68d22a9a8c8a69fbcee01276c5ac5add20 Mon Sep 17 00:00:00 2001 From: Mahika Wason Date: Thu, 25 Jun 2026 07:44:35 +0000 Subject: [PATCH 4/5] 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 --- .../src/nemo_retriever/query/agentic.py | 16 +++++++- .../src/nemo_retriever/query/workflow.py | 3 ++ nemo_retriever/tests/test_agentic_eval.py | 38 ++++++++++++++++++- 3 files changed, 55 insertions(+), 2 deletions(-) diff --git a/nemo_retriever/src/nemo_retriever/query/agentic.py b/nemo_retriever/src/nemo_retriever/query/agentic.py index 80c1ca45d6..2813311b54 100644 --- a/nemo_retriever/src/nemo_retriever/query/agentic.py +++ b/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) diff --git a/nemo_retriever/src/nemo_retriever/query/workflow.py b/nemo_retriever/src/nemo_retriever/query/workflow.py index b07f934cc8..4cd7312512 100644 --- a/nemo_retriever/src/nemo_retriever/query/workflow.py +++ b/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, diff --git a/nemo_retriever/tests/test_agentic_eval.py b/nemo_retriever/tests/test_agentic_eval.py index 3a689f967f..e76ed7492f 100644 --- a/nemo_retriever/tests/test_agentic_eval.py +++ b/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 From fed8ce07432aa7c503e68d77e442182ce10b59fc Mon Sep 17 00:00:00 2001 From: Mahika Wason Date: Tue, 30 Jun 2026 16:34:06 -0700 Subject: [PATCH 5/5] Reuse agentic LLM default in CLI help --- nemo_retriever/src/nemo_retriever/cli/query/options.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/nemo_retriever/src/nemo_retriever/cli/query/options.py b/nemo_retriever/src/nemo_retriever/cli/query/options.py index dce2523244..73fa3aa1e3 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/options.py +++ b/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 @@ -162,9 +163,8 @@ "--agentic-llm-model", envvar="NEMO_RETRIEVER_AGENTIC_LLM_MODEL", help=( - "Chat model the agent drives. Defaults to " - "nvidia/llama-3.3-nemotron-super-49b-v1.5; override here or via " - "NEMO_RETRIEVER_AGENTIC_LLM_MODEL." + f"Chat model the agent drives. Defaults to {DEFAULT_AGENTIC_LLM_MODEL}; " + "override here or via NEMO_RETRIEVER_AGENTIC_LLM_MODEL." ), ), ]