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 introduces an in-process vLLM backend for agentic retrieval, defaulting retriever query --agentic to the nemotron-8b model loaded locally on GPU instead of requiring a remote OpenAI-compatible endpoint. It wires the new backend through AgenticRetrievalConfig, the harness --set overrides, and the CLI, while keeping the openai_compatible path for custom models.

• New VLLMAgentChatLLM adapter (models/local/agent_llm.py): process-level model caching, thread-safe single-call lock, tool-call JSON fallback parser, and consecutive-tool-message collapsing for Llama chat templates.
• AgenticRetrievalConfig now defaults to in_process: custom model names without invoke_url now raise ValueError at construction time (previously they silently forwarded to the NVIDIA endpoint), which is a breaking change for existing callers.
• ReActAgentOperator removes enforce_top_k: top-k enforcement is now unconditional; both operators accept an optional chat_completion_fn to inject the local adapter without changing the agent loop.

Confidence Score: 4/5

Safe to merge after confirming whether existing callers rely on the previously-implicit NVIDIA-endpoint fallback for non-allowlisted models; the config change now raises at construction time rather than failing at inference time.

AgenticRetrievalConfig now rejects custom model names without invoke_url, breaking any integration that previously relied on the implicit fallback to the NVIDIA endpoint. The enforce_top_k removal from ReActAgentOperator is a secondary concern for downstream operator users.

nemo_retriever/src/nemo_retriever/query/agentic.py (AgenticRetrievalConfig breaking change) and nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py (enforce_top_k removal)

Important Files Changed

Filename	Overview
nemo_retriever/src/nemo_retriever/models/local/agent_llm.py	New file: in-process vLLM adapter for agentic LLM inference. Well-structured with process-level model caching, thread-safe call path, and tool-call JSON parsing fallback. Key concern: VLLMAgentChatLLM is constructed while holding _LOCAL_AGENT_LLM_CACHE_LOCK, blocking concurrent unload operations for the full model-load duration.
nemo_retriever/src/nemo_retriever/query/agentic.py	Adds in-process backend routing, lazy chat_completion_fn initialisation, and new local vLLM config knobs to AgenticRetrievalConfig. Breaking change: custom model names without invoke_url previously reached the NVIDIA endpoint; now they raise ValueError at construction time.
nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py	Wires optional chat_completion_fn injection, removes enforce_top_k parameter (now always-on), and adds valid-doc-id validation to final_results. enforce_top_k removal is a breaking API change without a deprecation shim.
nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py	Adds chat_completion_fn injection and replaces the silent invalid-doc-id filter with an explicit LLM feedback loop that retries on hallucinated IDs.
nemo_retriever/src/nemo_retriever/harness/resolution.py	Adds four new local vLLM override paths and wires them into build_query_request. Uses the or-based fallback for agentic_local_gpu_memory_utilization which silently substitutes 0.8 for any falsy value, though downstream validation catches 0.0.
nemo_retriever/src/nemo_retriever/models/init.py	Adds create_local_agent_llm factory with a backend-normalisation guard. Lazy-imports the vLLM adapter and delegates to create_cached_vllm_agent_chat_llm. Clean addition.
nemo_retriever/src/nemo_retriever/models/hf_model_registry.py	Adds pinned HF revisions for nvidia/Llama-3.1-Nemotron-Nano-8B-v1 and nvidia/Llama-3_3-Nemotron-Super-49B-v1. Straightforward, safe addition.
nemo_retriever/tests/test_agentic_local_llm.py	New test file covering alias resolution, config propagation, tool-call JSON parsing, consecutive-tool-message collapsing, unload lifecycle, and cache clearing, all without real GPU requirements. Good unit coverage.

Comments Outside Diff (1)

1. nemo_retriever/src/nemo_retriever/query/agentic.py, line 223-234 (link)

   Breaking config contract for existing in-process callers

   AgenticRetrievalConfig(llm_model="custom-model") (without invoke_url) previously resolved to the NVIDIA endpoint silently; now it raises ValueError("Unsupported in-process agentic LLM model"). Any integration that constructs AgenticRetrievalConfig with a non-allowlisted model name and no invoke_url will break at config construction time. The error message is clear and the new behaviour is arguably better, but under the api-backward-compatibility rule this is a breaking change to a public configuration class that requires a migration note or deprecation period.

   Rule Used: Changes to public API surfaces (FastAPI endpoints,... (source)

   Prompt To Fix With AI

   This is a comment left during a code review.
   Path: nemo_retriever/src/nemo_retriever/query/agentic.py
   Line: 223-234
   
   Comment:
   **Breaking config contract for existing in-process callers**
   
   `AgenticRetrievalConfig(llm_model="custom-model")` (without `invoke_url`) previously resolved to the NVIDIA endpoint silently; now it raises `ValueError("Unsupported in-process agentic LLM model")`. Any integration that constructs `AgenticRetrievalConfig` with a non-allowlisted model name and no `invoke_url` will break at config construction time. The error message is clear and the new behaviour is arguably better, but under the `api-backward-compatibility` rule this is a breaking change to a public configuration class that requires a migration note or deprecation period.
   
   **Rule Used:** Changes to public API surfaces (FastAPI endpoints,... ([source](.greptile))
   
   How can I resolve this? If you propose a fix, please make it concise.

Prompt To Fix All With AI

Fix the following 1 code review issue. Work through them one at a time, proposing concise fixes.

---

### Issue 1 of 1
nemo_retriever/src/nemo_retriever/query/agentic.py:223-234
**Breaking config contract for existing in-process callers**

`AgenticRetrievalConfig(llm_model="custom-model")` (without `invoke_url`) previously resolved to the NVIDIA endpoint silently; now it raises `ValueError("Unsupported in-process agentic LLM model")`. Any integration that constructs `AgenticRetrievalConfig` with a non-allowlisted model name and no `invoke_url` will break at config construction time. The error message is clear and the new behaviour is arguably better, but under the `api-backward-compatibility` rule this is a breaking change to a public configuration class that requires a migration note or deprecation period.

_{Reviews (4): Last reviewed commit: "cleanup and added local llm agentic conf..." | 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.