Support local in-process vLLM LLMs for agentic retrieval

[mahikaw]

Description

• Add an in-process vLLM-backed LLM adapter for agentic retrieval.
• Default agentic retrieval to local vLLM with the supported Nemotron 8B model.
• Wire local LLM options through module config and harness overrides while keeping the public query CLI minimal.
• Restrict in-process support to known internal Nemotron model aliases for now.
• Keep custom/self-hosted LLMs on the existing OpenAI-compatible endpoint path.
• Update docs and tests for the new defaults and validation behavior.

Validation

• uv run pytest tests/test_agentic_local_llm.py tests/test_agentic_eval.py tests/test_harness_agentic_eval.py tests/test_root_query_cli.py -q
• uvx pre-commit run --all-files
• Bounded JP20 local vLLM run completed successfully on 15 queries.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR adds an in-process vLLM LLM adapter (VLLMAgentChatLLM) for agentic retrieval and wires it as the new default backend, replacing the previous requirement to pass an external OpenAI-compatible endpoint. The main structural changes are a new models/local/agent_llm.py module, backend-selection logic in AgenticRetrievalConfig, lazy initialization and explicit unload() lifecycle management in AgenticRetriever, and operator-level injection of a chat_completion_fn callable that both ReActAgentOperator and SelectionAgentOperator now accept.

• The in-process path loads nemotron-8b (or super-49b) with pinned HF revisions, an allowlist guard, and a VLLMAgentChatLLM._lock that serializes inference. All three public entry points (agentic_query_documents, agentic_beir_retrieve, run_agentic_audio_recall_evaluation) now wrap retrieve() in try/finally to guarantee unload() and EngineCore teardown.
• Both agent operators gain chat_completion_fn injection and stronger hallucination guards: non-object tool arguments and out-of-candidate doc IDs now return structured error messages to the LLM instead of silently filtering.
• enforce_top_k is removed from ReActAgentOperator (always enforced now), and the prompt always emits the exact-count instruction.

Confidence Score: 5/5

Safe to merge; the new in-process vLLM path is well-isolated behind an allowlist, all public entry points properly call unload(), and the existing agentic test suite is meaningfully extended.

The core model-lifecycle contract (lazy load, serialized inference via VLLMAgentChatLLM._lock, explicit EngineCore teardown in all three public entry points) is correctly implemented and covered by new tests. The two findings are non-blocking style concerns: a removed public parameter that should have a deprecation cycle, and inconsistent int-casting in the harness resolver. Neither affects runtime correctness on the changed paths.

react_agent_operator.py (enforce_top_k removal) and harness/resolution.py (raw-string locals for local_max_model_len / local_max_num_seqs) deserve a second look before merge.

Important Files Changed

Filename	Overview
nemo_retriever/src/nemo_retriever/models/local/agent_llm.py	New file: in-process vLLM chat-completion adapter with tool-call parsing, message normalization for Llama chat templates, and proper GPU teardown via EngineCore shutdown.
nemo_retriever/src/nemo_retriever/query/agentic.py	Adds in-process LLM backend selection, lazy-loads the local LLM on first retrieve() call, and ensures unload() is called in all public entry points.
nemo_retriever/src/nemo_retriever/harness/resolution.py	Adds four new local vLLM override paths; gpu_memory_utilization and tensor_parallel_size are explicitly cast, but local_max_model_len and local_max_num_seqs are left as raw strings.
nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py	Adds chat_completion_fn injection, removes enforce_top_k (previously a public parameter), unconditionally emits the exact-count instruction, and adds doc-ID hallucination validation.
nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py	Adds chat_completion_fn injection; strengthens input validation by rejecting non-list doc_ids and out-of-candidate-set IDs with structured LLM feedback.
nemo_retriever/tests/test_agentic_local_llm.py	New test file with broad unit coverage of VLLMAgentChatLLM: alias resolution, tool-call parsing, message normalization, EngineCore shutdown, and AgenticRetriever unload lifecycle.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant CLI as CLI / Harness
    participant AR as AgenticRetriever
    participant GCCF as _get_chat_completion_fn()
    participant VLLM as VLLMAgentChatLLM
    participant React as ReActAgentOperator
    participant Select as SelectionAgentOperator

    CLI->>AR: retrieve(query_ids, queries)
    AR->>GCCF: lazy-init (self._lock)
    GCCF-->>VLLM: create_local_agent_llm() [first call only]
    GCCF-->>AR: chat_completion_fn
    AR->>React: "run(df, chat_completion_fn=fn)"
    loop ReAct steps
        React->>VLLM: __call__(messages, tools)
        VLLM-->>React: OpenAI-compatible response dict
    end
    AR->>Select: "run(df, chat_completion_fn=fn)"
    loop Selection steps
        Select->>VLLM: __call__(messages, tools)
        VLLM-->>Select: OpenAI-compatible response dict
    end
    AR-->>CLI: ranked DataFrame
    CLI->>AR: unload()
    AR->>VLLM: unload() → EngineCore.shutdown()

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant CLI as CLI / Harness
    participant AR as AgenticRetriever
    participant GCCF as _get_chat_completion_fn()
    participant VLLM as VLLMAgentChatLLM
    participant React as ReActAgentOperator
    participant Select as SelectionAgentOperator

    CLI->>AR: retrieve(query_ids, queries)
    AR->>GCCF: lazy-init (self._lock)
    GCCF-->>VLLM: create_local_agent_llm() [first call only]
    GCCF-->>AR: chat_completion_fn
    AR->>React: "run(df, chat_completion_fn=fn)"
    loop ReAct steps
        React->>VLLM: __call__(messages, tools)
        VLLM-->>React: OpenAI-compatible response dict
    end
    AR->>Select: "run(df, chat_completion_fn=fn)"
    loop Selection steps
        Select->>VLLM: __call__(messages, tools)
        VLLM-->>Select: OpenAI-compatible response dict
    end
    AR-->>CLI: ranked DataFrame
    CLI->>AR: unload()
    AR->>VLLM: unload() → EngineCore.shutdown()

Loading

_{Reviews (8): Last reviewed commit: "Merge branch 'main' into dev/wasonmahika..." | Re-trigger Greptile}

[jperez999]

Also found these two possible snags, would like to know if these are by design or actual problems (found by codex):

• [P2] Preserve assistant tool calls for Super-49B follow-up turns — nemo_retriever/src/nemo_retriever/models/local/agent_llm.py:228

  _normalize_messages() leaves assistant messages containing only tool_calls unchanged. The adapter itself creates exactly that shape whenever a local model calls a tool. However, the pinned Super-49B chat template
  (https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/blob/387156d8d6868c19f3472fa607aa9bfc4f662333/tokenizer_config.json) reads only message["content"] and has no tool_calls handling. Consequently, subsequent ReAct iterations silently omit the previous tool name and arguments, so the
  advertised Super-49B profile receives incomplete conversation history. Serialize assistant tool calls into textual content for templates without native tool-call support.

• [P2] Keep malformed argument strings malformed — nemo_retriever/src/nemo_retriever/models/local/agent_llm.py:452

  When a model emits invalid string arguments such as "query=foo", _arguments_to_json_string() wraps them with json.dumps(). Downstream json.loads() therefore succeeds but returns a string; both agent operators then call fn_args.get(...), raising AttributeError and aborting the query. Preserve
  the invalid string so existing JSONDecodeError handling can return a tool error, or validate that decoded arguments are mappings.

[mahikaw]

Also found these two possible snags, would like to know if these are by design or actual problems (found by codex):

• [P2] Preserve assistant tool calls for Super-49B follow-up turns — nemo_retriever/src/nemo_retriever/models/local/agent_llm.py:228
  _normalize_messages() leaves assistant messages containing only tool_calls unchanged. The adapter itself creates exactly that shape whenever a local model calls a tool. However, the pinned Super-49B chat template
  (https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/blob/387156d8d6868c19f3472fa607aa9bfc4f662333/tokenizer_config.json) reads only message["content"] and has no tool_calls handling. Consequently, subsequent ReAct iterations silently omit the previous tool name and arguments, so the
  advertised Super-49B profile receives incomplete conversation history. Serialize assistant tool calls into textual content for templates without native tool-call support.
• [P2] Keep malformed argument strings malformed — nemo_retriever/src/nemo_retriever/models/local/agent_llm.py:452
  When a model emits invalid string arguments such as "query=foo", _arguments_to_json_string() wraps them with json.dumps(). Downstream json.loads() therefore succeeds but returns a string; both agent operators then call fn_args.get(...), raising AttributeError and aborting the query. Preserve
  the invalid string so existing JSONDecodeError handling can return a tool error, or validate that decoded arguments are mappings.

Handling these in-process boundary issues as described below:

1. not preserving tool_calls when content is empty varies model to model due to different chat templates. We can add a defensive workaround here that serializes assistant tool-call metadata into textual content before local rendering. there will be a tradeoff here of possible redundant context for templates that handles tool_calls as expected though. but we do gain from not silently dropping prior tool history in templates that don't support content reads.
2. incase of malformed tool args in llm response, the right behavior can be to fail at adapter/operator boundary rather than letting those errors trickle down into .get() calls. The adapter can avoid converting malformed args into valid json strings and operators should validate that decoded args are true json objects. if any of these boundaries detect failures, we send a clear tool error back to the agent so it can retry with valid args.