diff --git a/nemo_retriever/src/nemo_retriever/cli/query/app.py b/nemo_retriever/src/nemo_retriever/cli/query/app.py index 404c8c14d..b1d7c1c70 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 @@ -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 diff --git a/nemo_retriever/src/nemo_retriever/cli/query/options.py b/nemo_retriever/src/nemo_retriever/cli/query/options.py index 0c60b78d2..73fa3aa1e 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 @@ -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[ diff --git a/nemo_retriever/src/nemo_retriever/query/agentic.py b/nemo_retriever/src/nemo_retriever/query/agentic.py index aa7503447..2813311b5 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(): @@ -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)] diff --git a/nemo_retriever/src/nemo_retriever/query/options.py b/nemo_retriever/src/nemo_retriever/query/options.py index 237001a6e..1b5ef8702 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/src/nemo_retriever/query/workflow.py b/nemo_retriever/src/nemo_retriever/query/workflow.py index 5c777d537..d6ffc5c14 100644 --- a/nemo_retriever/src/nemo_retriever/query/workflow.py +++ b/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 diff --git a/nemo_retriever/tests/test_agentic_eval.py b/nemo_retriever/tests/test_agentic_eval.py index 8d8a500d2..e76ed7492 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", @@ -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 diff --git a/nemo_retriever/tests/test_root_query_cli.py b/nemo_retriever/tests/test_root_query_cli.py index 28bfeec9b..437da08fa 100644 --- a/nemo_retriever/tests/test_root_query_cli.py +++ b/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: