refactor(agentic): make agentic retrieval a self-contained module - #2407
refactor(agentic): make agentic retrieval a self-contained module#2407Reza-esfandiarpoor wants to merge 4 commits into
Conversation
Signed-off-by: Reza Esfandiarpoor <resfandiarpo@nvidia.com>
Greptile SummaryThis PR replaces the two inline ReAct/selection loops in
|
| Filename | Overview |
|---|---|
| nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/agent.py | New Agent class implementing the ReAct loop with retrieval-specific bookkeeping; well-structured and documented. |
| nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/loop.py | Shared agent loop engine; bare except Exception: in tool-argument parsing silently swallows non-JSON exceptions without logging. |
| nemo_retriever/src/nemo_retriever/query/agentic.py | Removes backend_top_k field from the public AgenticRetrievalConfig dataclass without a deprecation cycle; changes AGENTIC_PARALLEL_TOOL_CALLS from False to None, silently altering parallel-tool-call behavior. |
| nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py | Greatly simplified to a thin adapter over the new Agent; lazy init, proper close(), and batch ordering are all correct. |
| nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py | Replaced inline LLM loop with SelectionAgent; three-tier gate (final_results → selection_agent → rrf) is clearly structured with proper resource cleanup. |
| nemo_retriever/src/nemo_retriever/cli/query/options.py | Removes AgenticBackendTopKOption (CLI argument removal flagged in prior review) and adds AgenticLlmClientOption. |
| nemo_retriever/src/nemo_retriever/_agentic/nemo_agent/llm/init.py | Clean backend registry with create_llm, create_llm_config, and get_available_backends; unknown-backend errors are informative. |
Sequence Diagram
sequenceDiagram
participant Caller
participant AgenticRetriever
participant ReActAgentOperator
participant Agent
participant BaseLLMBackend
participant RetrieveTool
participant SelectionAgentOperator
participant SelectionAgent
Caller->>AgenticRetriever: retrieve(query_texts, query_ids)
AgenticRetriever->>ReActAgentOperator: process(DataFrame)
loop Per query (ThreadPoolExecutor)
ReActAgentOperator->>Agent: run_sync(query, query_id)
loop ReAct steps (max_steps)
Agent->>BaseLLMBackend: acompletion(messages, tools)
BaseLLMBackend-->>Agent: CompletionResult
alt "tool_calls = retrieve"
Agent->>RetrieveTool: acall(query, top_k)
RetrieveTool-->>Agent: "[{id, score, text}]"
else "tool_calls = final_results"
Agent-->>Agent: record end_payload, break
end
end
Agent-->>ReActAgentOperator: AgentRunResult
end
ReActAgentOperator-->>AgenticRetriever: exploded DataFrame
AgenticRetriever->>SelectionAgentOperator: process(rrf_ranked DataFrame)
loop Per query_id group
alt has react final_results
SelectionAgentOperator-->>SelectionAgentOperator: pass through final_doc_ids
else run selection agent
SelectionAgentOperator->>SelectionAgent: select_sync(query, docs)
SelectionAgent->>BaseLLMBackend: acompletion(messages, tools)
BaseLLMBackend-->>SelectionAgent: CompletionResult
SelectionAgent-->>SelectionAgentOperator: AgentRunResult
else fallback
SelectionAgentOperator-->>SelectionAgentOperator: top RRF-ranked docs
end
end
SelectionAgentOperator-->>AgenticRetriever: ranked DataFrame
AgenticRetriever-->>Caller: AgenticResult
Reviews (2): Last reviewed commit: "remove nim stub" | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| extended_relevance: bool = False, | ||
| max_steps: int = 10, | ||
| max_steps: int = 200, | ||
| num_concurrent: int = 8, |
There was a problem hiding this comment.
max_steps default inconsistent with config-layer default
The operator defaults max_steps=200, but AgenticRetrievalConfig defaults to react_max_steps=50 and always passes its value explicitly. A caller constructing ReActAgentOperator directly (e.g., in tests or custom pipelines) will get 200 steps when the documentation elsewhere suggests 50. These two defaults should be reconciled — either expose a shared constant or document the discrepancy explicitly in the docstring so integrators are not surprised.
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py
Line: 148
Comment:
**`max_steps` default inconsistent with config-layer default**
The operator defaults `max_steps=200`, but `AgenticRetrievalConfig` defaults to `react_max_steps=50` and always passes its value explicitly. A caller constructing `ReActAgentOperator` directly (e.g., in tests or custom pipelines) will get 200 steps when the documentation elsewhere suggests 50. These two defaults should be reconciled — either expose a shared constant or document the discrepancy explicitly in the docstring so integrators are not surprised.
How can I resolve this? If you propose a fix, please make it concise.Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| kwargs = {"completion_fn": self._chat_completion_fn} if self._chat_completion_fn is not None else {} | ||
| return create_llm(config, **kwargs) | ||
|
|
||
| def _retrieve_adapter(self, query: str, top_k: int) -> List[Dict[str, Any]]: |
There was a problem hiding this comment.
nit: can we avoid get away without these dual keys for id across operator/agent core boundary. this way we can avoid adding this rename shim in the operators permanently.
There was a problem hiding this comment.
This is the same check from before this PR. I prefer to keep NRL-specific details the same as before and just focus on the agent implementation.
|
Evaluation results with changes in this PR. The reported numbers are average across three runs.
Command to reproduce these results: retriever harness run "${DS}_beir" \
--set query.agentic=true \
--set query.agentic_llm_model=openai/nvidia/nvidia/nemotron-3-ultra \
--set query.agentic_invoke_url=${AGENTIC_INVOKE_URL} \
--set query.agentic_num_concurrent=8 \
--set query.agentic_reasoning_effort=high \
--set query.agentic_react_max_steps=200 \
--set query.agentic_text_truncation=0 \
--set ingest.embed.embed_model_name=nvidia/nvidia/nemotron-3-embed-1b \
--set ingest.embed.embed_invoke_url=${EMBED_INVOKE_URL} \
--set ingest.embed.embed_modality=text \
--set ingest.embed.embed_granularity=page \
--set query.embed_model_name=nvidia/nvidia/nemotron-3-embed-1b \
--set query.embed_invoke_url=${EMBED_INVOKE_URL} \
--set dataset.path=${CORPUS_PATH} \
--set evaluation.query_language=english \
--set ingest.extract.method=pdfium_hybrid \
--set ingest.extract.ocr_lang=multi \
--set ingest.extract.use_table_structure=true |
Thank you @Reza-esfandiarpoor this result summary is really helpful to confirm consistency with what we currently have in https://docs.google.com/spreadsheets/d/1TPcVZa_1dxnwrNltkMq63lmgGtOE_cij7TTjkX2te8w/edit?gid=0#gid=0 and reported in the release, just a couple of questions: |
Thanks for pointing it out. Yes, this is not the nvfp4 version. But it should be okay in this case. This run is to just show that the refactor in this PR is correct.
The scope of this PR is limited to the implementation of the agent itself. Ingestion and the implementation of retrieval methods are provided by NRL (this code just wraps whatever retriever function it gets from NRL in an agent tool). How and where different reranking and embedding models are implemented and using what backend (e.g., vLLM) is outside of the scope of this PR. We can discuss that for future changes. |
Signed-off-by: Reza Esfandiarpoor <resfandiarpo@nvidia.com>
Signed-off-by: Reza Esfandiarpoor <resfandiarpo@nvidia.com>
Signed-off-by: Reza Esfandiarpoor <resfandiarpo@nvidia.com>
3 participants
Description
Replaces the inline ReAct/selection loops in the graph operators with a
self-contained
nemo_agentlibrary (src/nemo_retriever/_agentic/nemo_agent/),and reconciles it with main's CLI/config layer and in-process local-vLLM support.
What's new
nemo_agentlibrary:Agent(ReAct) +SelectionAgent, retrieve/end/thinktools, Jinja prompts, and an LLM backend layer behind one
_completion_implseam.
litellm(remote) andcallable— the latter adapts anyOpenAI-compatible completion callable, used to wrap main's in-process vLLM
(
VLLMAgentChatLLM).llm_backend(transport:openai_compatible/in_process,inferred from
--agentic-invoke-url, similar to main, not changed) andllm_client(implementation), resolvedin
AgenticRetrievalConfig.__post_init__(detected automatically if not set).CLI
--agentic-backend-top-k(query.agentic_backend_top_k) — the agent nowsets retrieval breadth per tool-call.
--agentic-llm-client(callablelocal default /litellmremote default) - Optional.--agentic-temperaturenow defaults to unset (endpoint/model default) instead of0.0.Checklist