From 06235b5c84fd31b4d05020d69d87bd58272d0ad7 Mon Sep 17 00:00:00 2001 From: Mahika Wason Date: Fri, 17 Jul 2026 18:57:24 -0700 Subject: [PATCH 01/10] Add in-process vLLM backend for agentic retrieval --- .../src/nemo_retriever/models/__init__.py | 41 ++ .../models/hf_model_registry.py | 2 + .../nemo_retriever/models/local/agent_llm.py | 593 ++++++++++++++++++ .../graph_ops/react_agent_operator.py | 20 +- .../graph_ops/selection_agent_operator.py | 7 +- .../src/nemo_retriever/query/agentic.py | 150 ++++- nemo_retriever/tests/test_agentic_eval.py | 131 +++- .../tests/test_agentic_local_llm.py | 174 +++++ .../tests/test_agentic_operators.py | 55 ++ 9 files changed, 1162 insertions(+), 11 deletions(-) create mode 100644 nemo_retriever/src/nemo_retriever/models/local/agent_llm.py create mode 100644 nemo_retriever/tests/test_agentic_local_llm.py diff --git a/nemo_retriever/src/nemo_retriever/models/__init__.py b/nemo_retriever/src/nemo_retriever/models/__init__.py index df6c72d45d..4864af41f9 100644 --- a/nemo_retriever/src/nemo_retriever/models/__init__.py +++ b/nemo_retriever/src/nemo_retriever/models/__init__.py @@ -258,3 +258,44 @@ def create_local_reranker( device=device, hf_cache_dir=hf_cache_dir, ) + + +_LOCAL_AGENT_LLM_BACKENDS = frozenset({"vllm"}) + + +def create_local_agent_llm( + model_name: str, + *, + backend: str = "vllm", + device: str | None = None, + hf_cache_dir: str | None = None, + gpu_memory_utilization: float = 0.8, + tensor_parallel_size: int = 1, + max_model_len: int | None = None, + max_num_seqs: int | None = None, + tool_call_parser: str | None = None, +) -> Any: + """Create a cached local agent LLM chat-completion callable. + + The callable mirrors ``invoke_chat_completion_step`` and returns an + OpenAI-compatible chat-completions response dict. V1 supports the in-process + vLLM backend and uses model/profile-specific tool-call normalization where + known. + """ + + b = normalize_backend(backend, _LOCAL_AGENT_LLM_BACKENDS, field_name="backend", default="vllm") + if b == "vllm": + from nemo_retriever.models.local.agent_llm import create_cached_vllm_agent_chat_llm + + return create_cached_vllm_agent_chat_llm( + model_name, + device=device, + hf_cache_dir=hf_cache_dir, + gpu_memory_utilization=gpu_memory_utilization, + tensor_parallel_size=tensor_parallel_size, + max_model_len=max_model_len, + max_num_seqs=max_num_seqs, + tool_call_parser=tool_call_parser, + ) + + raise ValueError(f"Unsupported local agent LLM backend {backend!r}") diff --git a/nemo_retriever/src/nemo_retriever/models/hf_model_registry.py b/nemo_retriever/src/nemo_retriever/models/hf_model_registry.py index 0afa266fb0..493195834c 100644 --- a/nemo_retriever/src/nemo_retriever/models/hf_model_registry.py +++ b/nemo_retriever/src/nemo_retriever/models/hf_model_registry.py @@ -39,6 +39,8 @@ "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16": "24e67ea000b7c2837fc8f9488aa2008524fac8ba", "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-FP8": "6647b845a4b786c6e2c7adb1b6a909e1aa71fac2", "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-NVFP4": "dc5f0b0bfddf8b6e0f5891475be9af05b80126fe", + "nvidia/Llama-3.1-Nemotron-Nano-8B-v1": "54641c1611fcff44fa4865626462445e0a153fc7", + "nvidia/Llama-3_3-Nemotron-Super-49B-v1": "387156d8d6868c19f3472fa607aa9bfc4f662333", "nvidia/nemotron-ocr-v1": "8657d08d3279f4864002d5fd3fdcd47ad8c96bcb", "nvidia/nemotron-ocr-v2": "86cacb0467fa4f7ce54342fdb250825e0d928ae7", "nvidia/nemotron-page-elements-v3": "df62dbb631502575ac4d43b44d700b1674ab1d56", diff --git a/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py new file mode 100644 index 0000000000..03ae56cfa6 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py @@ -0,0 +1,593 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import logging +import os +import random +import string +import threading +from copy import deepcopy +from dataclasses import dataclass, field +from typing import Any, Mapping, Optional, Sequence + +from nemo_retriever.models.hf_cache import configure_global_hf_cache_base +from nemo_retriever.models.hf_model_registry import get_hf_revision +from nemo_retriever.models.model import BaseModel, ModelRunMode + +logger = logging.getLogger(__name__) + +_DEFAULT_MAX_TOKENS = 512 +_LOCAL_AGENT_LLM_CACHE: dict[tuple[Any, ...], "VLLMAgentChatLLM"] = {} +_LOCAL_AGENT_LLM_CACHE_LOCK = threading.Lock() + + +@dataclass(frozen=True) +class AgentLLMProfile: + """Metadata for local agent LLMs with known offline tool-call behavior.""" + + model_id: str + aliases: tuple[str, ...] = () + engine_kwargs: Mapping[str, Any] = field(default_factory=dict) + request_extras: Mapping[str, Any] = field(default_factory=dict) + tool_call_parser: str = "json" + tool_prompt_format: str = "none" + pass_tools_to_chat_template: bool = True + + def engine_kwargs_for_local(self) -> dict[str, Any]: + return _mutable_copy(self.engine_kwargs) + + def request_extras_for_local(self) -> dict[str, Any]: + return _mutable_copy(self.request_extras) + + +def _mutable_copy(value: Any) -> Any: + if isinstance(value, Mapping): + return {key: _mutable_copy(item) for key, item in value.items()} + if isinstance(value, tuple): + return [_mutable_copy(item) for item in value] + return deepcopy(value) + + +def _normalize_name(name: str) -> str: + return name.strip().casefold() + + +_NEMOTRON_JSON_TOOL_PROMPT_EXTRAS = {"chat_template_kwargs": {"tools_in_user_message": True}} + +_AGENT_LLM_PROFILES: tuple[AgentLLMProfile, ...] = ( + AgentLLMProfile( + model_id="nvidia/Llama-3.1-Nemotron-Nano-8B-v1", + aliases=( + "llama-3.1-nemotron-nano-8b-v1", + "nemotron-nano-8b", + "nemotron-8b", + "nvidia/llama-3.1-nemotron-nano-8b-v1", + ), + request_extras=_NEMOTRON_JSON_TOOL_PROMPT_EXTRAS, + tool_prompt_format="json", + ), + AgentLLMProfile( + model_id="nvidia/Llama-3_3-Nemotron-Super-49B-v1", + aliases=( + "llama-3.3-nemotron-super-49b-v1", + "nemotron-super-49b", + "super-49b", + "nvidia/llama-3.3-nemotron-super-49b-v1", + "nvidia/llama-3_3-nemotron-super-49b-v1", + ), + request_extras=_NEMOTRON_JSON_TOOL_PROMPT_EXTRAS, + tool_prompt_format="json", + ), +) + +_PROFILE_LOOKUP: dict[str, AgentLLMProfile] = {} +for _profile in _AGENT_LLM_PROFILES: + _PROFILE_LOOKUP[_normalize_name(_profile.model_id)] = _profile + for _alias in _profile.aliases: + _PROFILE_LOOKUP[_normalize_name(_alias)] = _profile + + +def _supported_agent_llm_names() -> str: + names: list[str] = [] + for profile in _AGENT_LLM_PROFILES: + names.append(profile.model_id) + names.extend(profile.aliases) + return ", ".join(sorted(names)) + + +def get_agent_llm_profile(name: str) -> AgentLLMProfile | None: + """Return a supported local agent LLM profile, or ``None`` when unsupported.""" + + return _PROFILE_LOOKUP.get(_normalize_name(name)) + + +def resolve_agent_llm_model_name(name: str) -> str: + """Resolve a local agent LLM alias to its Hugging Face model ID.""" + + profile = get_agent_llm_profile(name) + return profile.model_id if profile is not None else name + + +class VLLMAgentChatLLM(BaseModel): + """In-process vLLM chat-completions adapter for agentic retrieval. + + The public call signature mirrors ``invoke_chat_completion_step`` and returns + an OpenAI-compatible response dict. This lets the existing ReAct and + selection agents consume local GPU inference without changing their loop + semantics. + """ + + def __init__( + self, + model_path: str, + *, + device: Optional[str] = None, + hf_cache_dir: Optional[str] = None, + gpu_memory_utilization: float = 0.8, + tensor_parallel_size: int = 1, + max_model_len: Optional[int] = None, + max_num_seqs: Optional[int] = None, + max_tokens: int = _DEFAULT_MAX_TOKENS, + tool_call_parser: Optional[str] = None, + ) -> None: + super().__init__() + + profile = get_agent_llm_profile(model_path) + if profile is None: + raise ValueError( + f"Unsupported local agent LLM model {model_path!r}. " + f"Supported models and aliases: {_supported_agent_llm_names()}." + ) + + model_path = profile.model_id + cuda_visible_devices = _cuda_visible_devices_from_device(device) + if cuda_visible_devices is not None: + os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices + + try: + from vllm import LLM, SamplingParams # noqa: F401 + except ImportError as e: + raise ImportError( + 'Local agentic LLM inference requires vLLM. Install with: pip install "nemo-retriever[local]"' + ) from e + + _raise_if_cuda_unavailable() + self._model_path = model_path + self._max_tokens = int(max_tokens) + self._tool_call_parser = (tool_call_parser or profile.tool_call_parser or "json").strip().lower() + self._request_extras = profile.request_extras_for_local() + self._tool_prompt_format = profile.tool_prompt_format + self._pass_tools_to_chat_template = profile.pass_tools_to_chat_template + self._lock = threading.Lock() + + configure_global_hf_cache_base(hf_cache_dir) + revision = get_hf_revision(model_path, strict=False) + + engine_kwargs = profile.engine_kwargs_for_local() + if max_model_len is not None: + engine_kwargs["max_model_len"] = int(max_model_len) + if max_num_seqs is not None: + engine_kwargs["max_num_seqs"] = int(max_num_seqs) + + self._sampling_params_cls = SamplingParams + self._llm: Any | None = LLM( + model=model_path, + revision=revision, + trust_remote_code=True, + tensor_parallel_size=int(tensor_parallel_size), + gpu_memory_utilization=float(gpu_memory_utilization), + **engine_kwargs, + ) + + def __call__( + self, + *, + invoke_url: Optional[str] = None, + messages: list[dict[str, Any]], + model: Optional[str] = None, + api_key: Optional[str] = None, + tools: Optional[list[dict[str, Any]]] = None, + tool_choice: str | dict[str, Any] = "auto", + timeout_s: float = 120.0, + temperature: float = 0.0, + max_tokens: Optional[int] = None, + extra_body: Optional[dict[str, Any]] = None, + max_retries: int = 10, + max_429_retries: int = 5, + ) -> dict[str, Any]: + _ = (invoke_url, api_key, timeout_s, max_retries, max_429_retries) + if model and model != self._model_path: + logger.debug("Ignoring per-call model=%r for local agent LLM already loaded as %r", model, self._model_path) + + sampling_params = self._sampling_params_cls( + temperature=float(temperature), + max_tokens=int(max_tokens) if max_tokens is not None else self._max_tokens, + ) + chat_kwargs = self._build_chat_kwargs(extra_body) + active_tools = tools if tools and tool_choice != "none" else None + if active_tools and self._pass_tools_to_chat_template: + chat_kwargs["tools"] = active_tools + + local_messages = self._normalize_messages(messages, tools=active_tools) + with self._lock: + self._require_loaded() + outputs = self._llm.chat(local_messages, sampling_params=sampling_params, **chat_kwargs) + + request_output = outputs[0] + completion = request_output.outputs[0] + text = str(getattr(completion, "text", "") or "").strip() + tool_calls = _tool_calls_from_completion(completion) or parse_tool_calls_from_text( + text, + parser=self._tool_call_parser, + ) + finish_reason = "tool_calls" if tool_calls else str(getattr(completion, "finish_reason", None) or "stop") + + message: dict[str, Any] = {"role": "assistant"} + if tool_calls: + message["content"] = None + message["tool_calls"] = tool_calls + else: + message["content"] = text + + return { + "id": _new_response_id(), + "object": "chat.completion", + "model": self._model_path, + "choices": [ + { + "index": 0, + "message": message, + "finish_reason": finish_reason, + } + ], + "usage": _usage_from_outputs(request_output, completion), + } + + def _build_chat_kwargs(self, extra_body: Optional[dict[str, Any]]) -> dict[str, Any]: + chat_kwargs = _mutable_copy(self._request_extras) + for key, value in (extra_body or {}).items(): + if key in {"parallel_tool_calls", "reasoning_effort"}: + continue + chat_kwargs[key] = value + chat_kwargs.setdefault("use_tqdm", False) + return chat_kwargs + + def _normalize_messages( + self, + messages: Sequence[dict[str, Any]], + *, + tools: Optional[Sequence[dict[str, Any]]] = None, + ) -> list[dict[str, Any]]: + normalized: list[dict[str, Any]] = [] + tool_prompt = _json_tool_prompt(tools) if tools and self._tool_prompt_format == "json" else None + for idx, message in enumerate(messages): + msg = dict(message) + content = msg.get("content") + if isinstance(content, list): + msg["content"] = _flatten_text_content(content) + if tool_prompt and idx == 0 and msg.get("role") == "system": + msg["content"] = f"{tool_prompt}\n\n{msg.get('content') or ''}" + tool_prompt = None + normalized.append(msg) + if tool_prompt: + normalized.insert(0, {"role": "system", "content": tool_prompt}) + return _collapse_consecutive_tool_messages(normalized) + + def unload(self) -> None: + """Release GPU memory held by the local vLLM engine.""" + + with self._lock: + if self._llm is None: + return + del self._llm + self._llm = None + + try: + import torch + except ImportError: + return + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def _require_loaded(self) -> None: + if self._llm is None: + raise RuntimeError( + "VLLMAgentChatLLM has been unloaded; create a new local agent LLM before generating responses." + ) + + @property + def model_name(self) -> str: + return self._model_path + + @property + def model_type(self) -> str: + return "agent-chat-llm" + + @property + def model_runmode(self) -> ModelRunMode: + return "local" + + @property + def input(self) -> Any: + return {"type": "chat_messages", "format": "openai_chat_completions"} + + @property + def output(self) -> Any: + return {"type": "chat_completion", "format": "openai_compatible"} + + @property + def input_batch_size(self) -> int: + return 1 + + +def _cuda_visible_devices_from_device(device: Optional[str]) -> str | None: + if device is None: + return None + normalized = str(device).strip() + if not normalized: + return None + if normalized.casefold() == "cpu": + raise ValueError( + "The local agent LLM vLLM backend requires CUDA. Pass GPU ids such as '0' or '0,1', " + "or use llm_backend='openai_compatible' for a remote endpoint." + ) + return normalized.split(":", 1)[1] if normalized.startswith("cuda:") else normalized + + +def _raise_if_cuda_unavailable() -> None: + try: + import torch + except ImportError: + return + if not torch.cuda.is_available(): + raise RuntimeError( + "The local agent LLM vLLM backend requires CUDA, but torch reports no available CUDA device. " + "Run on a GPU host or use llm_backend='openai_compatible' for a remote endpoint." + ) + + +def create_cached_vllm_agent_chat_llm( + model_name: str, + *, + device: Optional[str] = None, + hf_cache_dir: Optional[str] = None, + gpu_memory_utilization: float = 0.8, + tensor_parallel_size: int = 1, + max_model_len: Optional[int] = None, + max_num_seqs: Optional[int] = None, + tool_call_parser: Optional[str] = None, +) -> VLLMAgentChatLLM: + """Create or reuse a local agent LLM keyed by heavyweight load settings.""" + + key = ( + resolve_agent_llm_model_name(model_name), + device, + hf_cache_dir, + float(gpu_memory_utilization), + int(tensor_parallel_size), + int(max_model_len) if max_model_len is not None else None, + int(max_num_seqs) if max_num_seqs is not None else None, + (tool_call_parser or "").strip().lower() or None, + ) + with _LOCAL_AGENT_LLM_CACHE_LOCK: + cached = _LOCAL_AGENT_LLM_CACHE.get(key) + if cached is None or cached._llm is None: + cached = VLLMAgentChatLLM( + model_path=model_name, + device=device, + hf_cache_dir=hf_cache_dir, + gpu_memory_utilization=gpu_memory_utilization, + tensor_parallel_size=tensor_parallel_size, + max_model_len=max_model_len, + max_num_seqs=max_num_seqs, + tool_call_parser=tool_call_parser, + ) + _LOCAL_AGENT_LLM_CACHE[key] = cached + return cached + + +def unload_cached_vllm_agent_chat_llms() -> None: + """Unload and forget all cached local agent vLLM engines.""" + + with _LOCAL_AGENT_LLM_CACHE_LOCK: + cached_models = list(_LOCAL_AGENT_LLM_CACHE.values()) + _LOCAL_AGENT_LLM_CACHE.clear() + + for model in cached_models: + model.unload() + + +def parse_tool_calls_from_text(text: str, *, parser: str = "json") -> list[dict[str, Any]]: + """Parse common offline tool-call JSON into OpenAI ``tool_calls`` shape.""" + + if parser not in {"auto", "json"}: + raise ValueError(f"Unsupported local tool-call parser {parser!r}; supported parsers: auto, json") + payload = _load_json_payload(text) + if payload is None: + return [] + return _coerce_tool_calls(payload) + + +def _load_json_payload(text: str) -> Any | None: + cleaned = _strip_code_fence(str(text or "").strip()) + if not cleaned: + return None + try: + return json.loads(cleaned) + except json.JSONDecodeError: + pass + + decoder = json.JSONDecoder() + starts = [idx for idx, ch in enumerate(cleaned) if ch in "[{"] + for start in starts: + try: + value, _end = decoder.raw_decode(cleaned[start:]) + except json.JSONDecodeError: + continue + return value + return None + + +def _strip_code_fence(text: str) -> str: + if not text.startswith("```"): + return text + lines = text.splitlines() + if len(lines) >= 2 and lines[-1].strip() == "```": + return "\n".join(lines[1:-1]).strip() + return text + + +def _coerce_tool_calls(payload: Any) -> list[dict[str, Any]]: + if isinstance(payload, Mapping) and "tool_calls" in payload: + return _coerce_tool_calls(payload.get("tool_calls")) + if isinstance(payload, Mapping): + call = _coerce_single_tool_call(payload) + return [call] if call is not None else [] + if isinstance(payload, list): + calls: list[dict[str, Any]] = [] + for item in payload: + call = _coerce_single_tool_call(item) + if call is not None: + calls.append(call) + return calls + return [] + + +def _coerce_single_tool_call(item: Any) -> dict[str, Any] | None: + if not isinstance(item, Mapping): + return None + + if isinstance(item.get("function"), Mapping): + function = dict(item["function"]) + name = function.get("name") + arguments = function.get("arguments", {}) + else: + name = item.get("name") or item.get("tool_name") + arguments = item.get("arguments", item.get("parameters", {})) + + if not name: + return None + + return { + "id": str(item.get("id") or _new_tool_call_id()), + "type": "function", + "function": { + "name": str(name), + "arguments": _arguments_to_json_string(arguments), + }, + } + + +def _arguments_to_json_string(arguments: Any) -> str: + if isinstance(arguments, str): + try: + json.loads(arguments) + except json.JSONDecodeError: + return json.dumps(arguments) + return arguments + return json.dumps(arguments or {}) + + +def _tool_calls_from_completion(completion: Any) -> list[dict[str, Any]]: + for attr in ("tool_calls", "tool_call"): + value = getattr(completion, attr, None) + if value: + return _coerce_tool_calls(value) + return [] + + +def _usage_from_outputs(request_output: Any, completion: Any) -> dict[str, int]: + prompt_tokens = len(getattr(request_output, "prompt_token_ids", None) or []) + completion_tokens = len(getattr(completion, "token_ids", None) or []) + usage: dict[str, int] = {} + if prompt_tokens: + usage["prompt_tokens"] = prompt_tokens + if completion_tokens: + usage["completion_tokens"] = completion_tokens + if usage: + usage["total_tokens"] = prompt_tokens + completion_tokens + return usage + + +def _flatten_text_content(content: Sequence[Any]) -> str: + parts: list[str] = [] + for item in content: + if isinstance(item, Mapping): + if item.get("type") == "text": + parts.append(str(item.get("text", ""))) + else: + parts.append(json.dumps(item)) + else: + parts.append(str(item)) + return "\n\n".join(part for part in parts if part) + + +def _collapse_consecutive_tool_messages(messages: Sequence[dict[str, Any]]) -> list[dict[str, Any]]: + """Merge parallel OpenAI tool results into one turn for local chat templates. + + Llama-family chat templates used by Nemotron require user/tool and + assistant roles to alternate. The agent may execute multiple OpenAI tool + calls from one assistant response, which creates consecutive ``tool`` turns. + Collapsing them keeps the transcript semantically equivalent for local + generation without changing the operator contract. + """ + + collapsed: list[dict[str, Any]] = [] + tool_names_by_id: dict[str, str] = {} + + for message in messages: + msg = dict(message) + if msg.get("role") == "assistant": + for tool_call in msg.get("tool_calls") or []: + if not isinstance(tool_call, Mapping): + continue + tool_call_id = str(tool_call.get("id") or "") + function = tool_call.get("function") or {} + if tool_call_id and isinstance(function, Mapping): + tool_names_by_id[tool_call_id] = str(function.get("name") or "tool") + + if msg.get("role") != "tool": + collapsed.append(msg) + continue + + msg["content"] = _format_tool_message_content(msg, tool_names_by_id) + if collapsed and collapsed[-1].get("role") == "tool": + previous = str(collapsed[-1].get("content") or "") + current = str(msg.get("content") or "") + collapsed[-1]["content"] = "\n\n".join(part for part in (previous, current) if part) + continue + collapsed.append(msg) + + return collapsed + + +def _format_tool_message_content(message: Mapping[str, Any], tool_names_by_id: Mapping[str, str]) -> str: + tool_call_id = str(message.get("tool_call_id") or "") + tool_name = tool_names_by_id.get(tool_call_id, "tool") + content = str(message.get("content") or "") + if tool_call_id: + return f"Tool result for {tool_name} ({tool_call_id}):\n{content}" + return f"Tool result for {tool_name}:\n{content}" + + +def _json_tool_prompt(tools: Sequence[Mapping[str, Any]]) -> str: + return ( + "You have access to the following tools. Your entire response must be a JSON array of one or more tool calls. " + "Each item must be shaped as " + '{"name": "tool_name", "arguments": {"arg_name": "arg_value"}}. ' + "Use only the listed tool names. Do not include prose, markdown, or text outside the JSON array.\n" + f"{json.dumps(list(tools), ensure_ascii=False)}" + ) + + +def _new_tool_call_id() -> str: + alphabet = string.ascii_letters + string.digits + return "".join(random.SystemRandom().choice(alphabet) for _ in range(9)) + + +def _new_response_id() -> str: + return "chatcmpl_" + _new_tool_call_id() diff --git a/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py b/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py index 39306154de..8d9419c3d5 100644 --- a/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py +++ b/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py @@ -414,6 +414,7 @@ def __init__( reasoning_effort: Optional[str] = None, backend_top_k: Optional[int] = None, temperature: float = 0.0, + chat_completion_fn: Optional[Callable[..., Dict[str, Any]]] = None, ) -> None: super().__init__() self._invoke_url = invoke_url or self._NVIDIA_BUILD_ENDPOINT @@ -432,6 +433,7 @@ def __init__( self._reasoning_effort = reasoning_effort self._backend_top_k = backend_top_k self._temperature = temperature + self._chat_completion_fn = chat_completion_fn def _build_extra_body(self) -> Optional[Dict[str, Any]]: """Assemble per-call extra payload fields (parallel_tool_calls, reasoning_effort).""" @@ -574,7 +576,8 @@ def _run_single_query( for _step in range(self._max_steps): logger.info("ReActAgentOperator: query=%s step=%d begin seen_docs=%d", query_id, _step, len(seen_doc_ids)) try: - response = invoke_chat_completion_step( + chat_completion_fn = self._chat_completion_fn or invoke_chat_completion_step + response = chat_completion_fn( invoke_url=self._invoke_url, messages=messages, model=self._llm_model, @@ -740,7 +743,7 @@ def _run_single_query( _step, _preview_text(fn_args.get("message")), ) - validation_error = self._validate_final_results_args(fn_args) + validation_error = self._validate_final_results_args(fn_args, valid_doc_ids=seen_doc_ids) if validation_error is None: final_doc_ids = list(raw_ids) tool_messages.append( @@ -845,7 +848,12 @@ def _call_retriever( return results - def _validate_final_results_args(self, fn_args: Dict[str, Any]) -> Optional[str]: + def _validate_final_results_args( + self, + fn_args: Dict[str, Any], + *, + valid_doc_ids: Optional[set[str]] = None, + ) -> Optional[str]: """Validate final_results tool args outside the prompt/schema.""" message = fn_args.get("message") if not isinstance(message, str): @@ -869,6 +877,12 @@ def _validate_final_results_args(self, fn_args: Dict[str, Any]) -> Optional[str] f"`search_successful` must be one of `true`, `false`, or `partial`. Got `{search_successful}` instead." ) + if valid_doc_ids is not None: + invalid_doc_ids = [doc_id for doc_id in doc_ids if doc_id not in valid_doc_ids] + if invalid_doc_ids: + preview = invalid_doc_ids[:_LOG_DOC_ID_LIMIT] + return f"`doc_ids` contains IDs that were not retrieved: {preview}." + if self._enforce_top_k and len(doc_ids) != self._target_top_k: return ( f"`doc_ids` must contain exactly {self._target_top_k} documents. " diff --git a/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py b/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py index 977e4e7db5..352f628b02 100644 --- a/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py +++ b/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py @@ -11,7 +11,7 @@ import os import requests -from typing import Any, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional import pandas as pd @@ -201,6 +201,7 @@ def __init__( base_url: Optional[str] = None, reasoning_effort: Optional[str] = None, temperature: float = 0.0, + chat_completion_fn: Optional[Callable[..., Dict[str, Any]]] = None, ) -> None: super().__init__() self._reasoning_effort = reasoning_effort @@ -214,6 +215,7 @@ def __init__( self._system_prompt_override = system_prompt_override self._text_truncation = text_truncation self._parallel_tool_calls = parallel_tool_calls + self._chat_completion_fn = chat_completion_fn if invoke_url is not None: self._invoke_url = invoke_url @@ -494,7 +496,8 @@ def _select_documents( feasible_k, ) try: - response = invoke_chat_completion_step( + chat_completion_fn = self._chat_completion_fn or invoke_chat_completion_step + response = chat_completion_fn( invoke_url=self._invoke_url, messages=messages, model=self._llm_model, diff --git a/nemo_retriever/src/nemo_retriever/query/agentic.py b/nemo_retriever/src/nemo_retriever/query/agentic.py index f65d262041..8b121321fd 100644 --- a/nemo_retriever/src/nemo_retriever/query/agentic.py +++ b/nemo_retriever/src/nemo_retriever/query/agentic.py @@ -12,6 +12,7 @@ from __future__ import annotations import logging +import math import threading from dataclasses import dataclass, field from pathlib import Path @@ -55,6 +56,11 @@ AGENTIC_RRF_K = 60 AGENTIC_REACT_MAX_STEPS = 50 AGENTIC_TEMPERATURE = 0.0 # agent LLM sampling temperature (0.0 = greedy) +AGENTIC_MAX_TOKENS: Optional[int] = None +AGENTIC_LLM_BACKEND = "openai_compatible" +AGENTIC_LLM_BACKENDS = frozenset({"openai_compatible", "in_process"}) +AGENTIC_LOCAL_LLM_BACKEND = "vllm" +AGENTIC_LOCAL_LLM_BACKENDS = frozenset({"vllm"}) class AgenticQueryInputOperator(AbstractOperator): @@ -137,8 +143,17 @@ class AgenticRetrievalConfig: reranker_api_key: str = "" local_reranker_backend: str = "vllm" embed_modality: str = "text" + llm_backend: str = AGENTIC_LLM_BACKEND llm_model: str = "" invoke_url: Optional[str] = None + local_llm_backend: str = AGENTIC_LOCAL_LLM_BACKEND + local_hf_cache_dir: Optional[str] = None + local_device: Optional[str] = None + local_gpu_memory_utilization: float = 0.8 + local_tensor_parallel_size: int = 1 + local_max_model_len: Optional[int] = None + local_max_num_seqs: Optional[int] = None + local_tool_call_parser: Optional[str] = None api_key: Optional[str] = None react_max_steps: int = AGENTIC_REACT_MAX_STEPS text_truncation: int = AGENTIC_TEXT_TRUNCATION @@ -150,14 +165,41 @@ class AgenticRetrievalConfig: backend_top_k: int = AGENTIC_BACKEND_TOP_K # Sampling temperature sent on every agent LLM call (0.0 = greedy). temperature: float = AGENTIC_TEMPERATURE + # Optional upper bound on tokens in each agent LLM response. + max_tokens: Optional[int] = AGENTIC_MAX_TOKENS # Final number of documents the agent targets/selects and the pipeline returns. # 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 + # When true, ReAct final_results must contain exactly top_k doc IDs. Keep this + # strict by default for existing eval behavior, but allow local/smaller LLMs to + # return partial final_results when explicitly requested. + enforce_top_k: bool = True def __post_init__(self) -> None: + llm_backend = _normalize_agentic_choice( + self.llm_backend, + AGENTIC_LLM_BACKENDS, + field_name="llm_backend", + default=AGENTIC_LLM_BACKEND, + ) + object.__setattr__(self, "llm_backend", llm_backend) + + local_llm_backend = _normalize_agentic_choice( + self.local_llm_backend, + AGENTIC_LOCAL_LLM_BACKENDS, + field_name="local_llm_backend", + default=AGENTIC_LOCAL_LLM_BACKEND, + ) + object.__setattr__(self, "local_llm_backend", local_llm_backend) + if self.llm_model is None or not str(self.llm_model).strip(): raise ValueError("Agentic retrieval requires a non-empty llm_model.") + object.__setattr__( + self, + "enforce_top_k", + _agentic_bool_value(self.enforce_top_k, field_name="enforce_top_k"), + ) for field_name, value, min_value in ( ("react_max_steps", self.react_max_steps, 1), ("text_truncation", self.text_truncation, 0), @@ -178,9 +220,39 @@ def __post_init__(self) -> None: raise ValueError(backend_error) object.__setattr__(self, "backend_top_k", agentic_int_value(self.backend_top_k, field_name="backend_top_k")) + local_tp_error = agentic_int_min_error( + self.local_tensor_parallel_size, field_name="local_tensor_parallel_size", min_value=1 + ) + if local_tp_error: + raise ValueError(local_tp_error) + object.__setattr__( + self, + "local_tensor_parallel_size", + agentic_int_value(self.local_tensor_parallel_size, field_name="local_tensor_parallel_size"), + ) + + for field_name in ("local_max_model_len", "local_max_num_seqs", "max_tokens"): + value = getattr(self, field_name) + if value is None: + continue + integer_error = agentic_int_min_error(value, field_name=field_name, min_value=1) + if integer_error: + raise ValueError(integer_error) + object.__setattr__(self, field_name, agentic_int_value(value, field_name=field_name)) + + local_gpu_memory_utilization = _agentic_float_range_value( + self.local_gpu_memory_utilization, + field_name="local_gpu_memory_utilization", + min_value=0.0, + max_value=1.0, + min_exclusive=True, + ) + object.__setattr__(self, "local_gpu_memory_utilization", local_gpu_memory_utilization) + + temperature_invoke_url = self.invoke_url if self.llm_backend == "openai_compatible" else "local://in-process" temperature_error = agentic_temperature_error( self.temperature, - invoke_url=self.invoke_url, + invoke_url=temperature_invoke_url, field_name="temperature", ) if temperature_error: @@ -188,6 +260,67 @@ def __post_init__(self) -> None: object.__setattr__(self, "temperature", float(self.temperature)) +def _normalize_agentic_choice(value: object, valid: frozenset[str], *, field_name: str, default: str) -> str: + normalized = str(value or default).strip().lower() + if normalized not in valid: + raise ValueError(f"{field_name} must be one of {sorted(valid)}; got {value!r}") + return normalized + + +def _agentic_bool_value(value: object, *, field_name: str) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "y", "on"}: + return True + if normalized in {"0", "false", "no", "n", "off"}: + return False + raise ValueError(f"{field_name} must be a boolean") + + +def _agentic_float_range_value( + value: object, + *, + field_name: str, + min_value: float, + max_value: float, + min_exclusive: bool = False, +) -> float: + try: + parsed = float(value) + except (TypeError, ValueError): + raise ValueError(f"{field_name} must be a number") from None + if not math.isfinite(parsed): + raise ValueError(f"{field_name} must be finite") + if parsed > max_value or (parsed <= min_value if min_exclusive else parsed < min_value): + operator = ">" if min_exclusive else ">=" + raise ValueError(f"{field_name} must be {operator} {min_value} and <= {max_value}") + return parsed + + +def _build_agent_chat_completion_fn(cfg: AgenticRetrievalConfig) -> Any | None: + if cfg.llm_backend == "openai_compatible": + return None + + if cfg.llm_backend == "in_process": + from nemo_retriever.models import create_local_agent_llm + + return create_local_agent_llm( + str(cfg.llm_model), + backend=str(cfg.local_llm_backend), + device=_none_if_empty(cfg.local_device), + hf_cache_dir=_none_if_empty(cfg.local_hf_cache_dir), + gpu_memory_utilization=float(cfg.local_gpu_memory_utilization), + tensor_parallel_size=int(cfg.local_tensor_parallel_size), + max_model_len=cfg.local_max_model_len, + max_num_seqs=cfg.local_max_num_seqs, + tool_call_parser=cfg.local_tool_call_parser, + ) + + raise ValueError(f"Unsupported agentic llm_backend {cfg.llm_backend!r}") + + class AgenticRetriever: """Run graph-backed agentic retrieval over query IDs and query texts.""" @@ -230,6 +363,15 @@ def __init__( }, ) self._lock = threading.Lock() + self._chat_completion_fn: Any | None = None + + def _get_chat_completion_fn(self) -> Any | None: + if self._cfg.llm_backend == "openai_compatible": + return None + with self._lock: + if self._chat_completion_fn is None: + self._chat_completion_fn = _build_agent_chat_completion_fn(self._cfg) + return self._chat_completion_fn def retrieve(self, query_ids: Sequence[str], query_texts: Sequence[str]) -> pd.DataFrame: """Return selected ranked documents for each query. @@ -250,6 +392,7 @@ def retrieve(self, query_ids: Sequence[str], query_texts: Sequence[str]) -> pd.D # small top_k. target_top_k = int(self._cfg.top_k) per_hop_top_k = max(AGENTIC_RETRIEVER_TOP_K, target_top_k) + chat_completion_fn = self._get_chat_completion_fn() pipeline = ( AgenticQueryInputOperator() @@ -267,7 +410,10 @@ def retrieve(self, query_ids: Sequence[str], query_texts: Sequence[str]) -> pd.D num_concurrent=int(self._cfg.num_concurrent), reasoning_effort=self._cfg.reasoning_effort, backend_top_k=self._cfg.backend_top_k, + enforce_top_k=bool(self._cfg.enforce_top_k), temperature=float(self._cfg.temperature), + max_tokens=self._cfg.max_tokens, + chat_completion_fn=chat_completion_fn, ) >> RRFAggregatorOperator(k=AGENTIC_RRF_K) >> SelectionAgentOperator( @@ -280,6 +426,8 @@ def retrieve(self, query_ids: Sequence[str], query_texts: Sequence[str]) -> pd.D text_truncation=int(self._cfg.text_truncation), reasoning_effort=self._cfg.reasoning_effort, temperature=float(self._cfg.temperature), + max_tokens=self._cfg.max_tokens, + chat_completion_fn=chat_completion_fn, ) >> AgenticSelectionOutputOperator() ) diff --git a/nemo_retriever/tests/test_agentic_eval.py b/nemo_retriever/tests/test_agentic_eval.py index 21c6e2a03d..9b06c19e8e 100644 --- a/nemo_retriever/tests/test_agentic_eval.py +++ b/nemo_retriever/tests/test_agentic_eval.py @@ -5,7 +5,7 @@ from __future__ import annotations import json -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pandas as pd import pytest @@ -60,6 +60,17 @@ def query(self, query: str, *, top_k: int | None = None): "_score": 0.1, }, ] + hits.extend( + { + "source": f"/tmp/extra_{idx}.pdf", + "source_id": f"/tmp/extra_{idx}.pdf", + "page_number": idx + 3, + "pdf_page": f"extra_{idx}", + "text": f"extra document {idx}", + "_score": 0.05, + } + for idx in range(8) + ) return hits[:top_k] def queries(self, queries, *, top_k: int | None = None): @@ -88,13 +99,23 @@ def test_build_beir_run_from_ranked_doc_ids_rejects_length_mismatch(): build_beir_run_from_ranked_doc_ids(["q1", "q2"], [["d1"]]) +def test_agentic_config_validates_max_tokens(): + from nemo_retriever.query.agentic import AgenticRetrievalConfig + + cfg = AgenticRetrievalConfig(llm_model="test-model", max_tokens="128") + + assert cfg.max_tokens == 128 + with pytest.raises(ValueError, match="max_tokens"): + AgenticRetrievalConfig(llm_model="test-model", max_tokens=0) + + @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_runs_graph_with_wrapped_retriever(mock_react_step, mock_selection_step): from nemo_retriever.query.agentic import AgenticRetrievalConfig, AgenticRetriever - final_ids = ["doc_1"] + [f"extra_{i}" for i in range(9)] + final_ids = ["doc_1", "other_2"] + [f"extra_{i}" for i in range(8)] mock_react_step.return_value = _make_tool_call_response( "final_results", {"doc_ids": final_ids, "message": "done", "search_successful": "true"}, @@ -104,9 +125,12 @@ def test_agentic_retriever_runs_graph_with_wrapped_retriever(mock_react_step, mo {"doc_ids": ["doc_1"], "message": "doc_1 is best"}, ) - cfg = AgenticRetrievalConfig(llm_model="test-model", invoke_url="http://localhost/v1/chat/completions") + cfg = AgenticRetrievalConfig( + llm_model="test-model", invoke_url="http://localhost/v1/chat/completions", max_tokens=77 + ) result = AgenticRetriever(cfg, match_mode="pdf_page").retrieve(["0"], ["find doc"]) + assert mock_react_step.call_args.kwargs["max_tokens"] == 77 assert list(result.columns) == ["query_id", "doc_id", "rank", "message", "result_source"] assert result["query_id"].tolist() == ["0"] * 10 assert result["doc_id"].tolist()[0] == "doc_1" @@ -120,7 +144,7 @@ def test_agentic_retriever_honors_top_k(mock_react_step, mock_selection_step): """cfg.top_k drives the pipeline output count, not the hardcoded default of 10.""" from nemo_retriever.query.agentic import AgenticRetrievalConfig, AgenticRetriever - final_ids = ["doc_1"] + [f"extra_{i}" for i in range(4)] # exactly 5 + final_ids = ["doc_1", "other_2"] + [f"extra_{i}" for i in range(3)] # exactly 5 mock_react_step.return_value = _make_tool_call_response( "final_results", {"doc_ids": final_ids, "message": "done", "search_successful": "true"}, @@ -136,6 +160,69 @@ 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.models.create_local_agent_llm") +@patch("nemo_retriever.query.agentic.Retriever", FakeRetriever) +def test_agentic_retriever_builds_in_process_llm_lazily(mock_create_local_agent_llm): + from nemo_retriever.query.agentic import AgenticRetrievalConfig, AgenticRetriever + + local_chat = MagicMock( + return_value=_make_tool_call_response( + "final_results", + {"doc_ids": ["doc_1"], "message": "done", "search_successful": "true"}, + ) + ) + mock_create_local_agent_llm.return_value = local_chat + + cfg = AgenticRetrievalConfig(llm_model="nemotron-8b", llm_backend="in_process", top_k=1) + retriever = AgenticRetriever(cfg, match_mode="pdf_page") + + mock_create_local_agent_llm.assert_not_called() + + result = retriever.retrieve(["0"], ["find doc"]) + + mock_create_local_agent_llm.assert_called_once_with( + "nemotron-8b", + backend="vllm", + device=None, + hf_cache_dir=None, + gpu_memory_utilization=0.8, + tensor_parallel_size=1, + max_model_len=None, + max_num_seqs=None, + tool_call_parser=None, + ) + assert local_chat.call_count == 1 + assert result["doc_id"].tolist() == ["doc_1"] + + +@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_can_accept_partial_react_final_results(mock_react_step, mock_selection_step): + from nemo_retriever.query.agentic import AgenticRetrievalConfig, AgenticRetriever + + mock_react_step.return_value = _make_tool_call_response( + "final_results", + {"doc_ids": ["doc_1"], "message": "partial but valid", "search_successful": "true"}, + ) + mock_selection_step.return_value = _make_tool_call_response( + "log_selected_documents", + {"doc_ids": ["doc_1"], "message": "selection should not run"}, + ) + + cfg = AgenticRetrievalConfig( + llm_model="test-model", + invoke_url="http://localhost/v1/chat/completions", + top_k=5, + enforce_top_k=False, + ) + result = AgenticRetriever(cfg, match_mode="pdf_page").retrieve(["0"], ["find doc"]) + + assert result["doc_id"].tolist() == ["doc_1"] + assert result["result_source"].tolist() == ["final_results"] + mock_selection_step.assert_not_called() + + @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) @@ -184,7 +271,7 @@ def test_run_agentic_beir_evaluation_loads_queries_and_qrels(mock_react_step, mo from nemo_retriever.query.agentic import AgenticRetrievalConfig, run_agentic_beir_evaluation from nemo_retriever.tools.recall.beir import BeirDataset - final_ids = ["doc"] + [f"extra_{i}" for i in range(9)] + final_ids = ["doc", "other"] + [f"extra_{i}" for i in range(8)] mock_react_step.return_value = _make_tool_call_response( "final_results", {"doc_ids": final_ids, "message": "done", "search_successful": "true"}, @@ -278,3 +365,37 @@ def test_agentic_config_rejects_nonfinite_temperature(): with pytest.raises(ValueError, match="temperature must be finite"): AgenticRetrievalConfig(llm_model="m", temperature=float("nan")) + + +def test_agentic_config_accepts_in_process_temperature_above_nvidia_limit(): + from nemo_retriever.query.agentic import AgenticRetrievalConfig + + cfg = AgenticRetrievalConfig(llm_model="m", llm_backend="in_process", temperature=1.5) + + assert cfg.llm_backend == "in_process" + assert cfg.temperature == pytest.approx(1.5) + + +def test_agentic_config_rejects_invalid_local_llm_backend(): + from nemo_retriever.query.agentic import AgenticRetrievalConfig + + with pytest.raises(ValueError, match="local_llm_backend"): + AgenticRetrievalConfig(llm_model="m", llm_backend="in_process", local_llm_backend="hf") + + +def test_agentic_config_validates_local_vllm_knobs(): + from nemo_retriever.query.agentic import AgenticRetrievalConfig + + cfg = AgenticRetrievalConfig( + llm_model="m", + llm_backend="in_process", + local_gpu_memory_utilization="0.6", + local_tensor_parallel_size="2.0", + local_max_model_len="8192", + local_max_num_seqs="4.0", + ) + + assert cfg.local_gpu_memory_utilization == pytest.approx(0.6) + assert cfg.local_tensor_parallel_size == 2 + assert cfg.local_max_model_len == 8192 + assert cfg.local_max_num_seqs == 4 diff --git a/nemo_retriever/tests/test_agentic_local_llm.py b/nemo_retriever/tests/test_agentic_local_llm.py new file mode 100644 index 0000000000..765ab23ea6 --- /dev/null +++ b/nemo_retriever/tests/test_agentic_local_llm.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import threading +from unittest.mock import MagicMock + + +def test_resolve_agent_llm_profile_aliases() -> None: + from nemo_retriever.models.local.agent_llm import resolve_agent_llm_model_name + + assert resolve_agent_llm_model_name("nemotron-8b") == "nvidia/Llama-3.1-Nemotron-Nano-8B-v1" + assert resolve_agent_llm_model_name("nemotron-super-49b") == "nvidia/Llama-3_3-Nemotron-Super-49B-v1" + + +def test_vllm_agent_llm_rejects_unsupported_profile_before_vllm_import() -> None: + import pytest + + from nemo_retriever.models.local.agent_llm import VLLMAgentChatLLM + + with pytest.raises(ValueError, match="Unsupported local agent LLM model"): + VLLMAgentChatLLM("mistral-7b") + + +def test_parse_json_tool_call_output() -> None: + from nemo_retriever.models.local.agent_llm import parse_tool_calls_from_text + + calls = parse_tool_calls_from_text( + json.dumps([{"name": "retrieve", "arguments": {"query": "monetary policy"}}]), + parser="json", + ) + + assert calls == [ + { + "id": calls[0]["id"], + "type": "function", + "function": {"name": "retrieve", "arguments": '{"query": "monetary policy"}'}, + } + ] + + +def test_parse_openai_style_tool_call_output() -> None: + from nemo_retriever.models.local.agent_llm import parse_tool_calls_from_text + + calls = parse_tool_calls_from_text( + json.dumps( + { + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "final_results", "arguments": json.dumps({"doc_ids": ["d1"]})}, + } + ] + } + ) + ) + + assert calls == [ + { + "id": "call_1", + "type": "function", + "function": {"name": "final_results", "arguments": '{"doc_ids": ["d1"]}'}, + } + ] + + +def test_parse_tool_call_output_from_code_fence() -> None: + from nemo_retriever.models.local.agent_llm import parse_tool_calls_from_text + + calls = parse_tool_calls_from_text('```json\n[{"name": "think", "arguments": {"thought": "compare docs"}}]\n```') + + assert calls[0]["function"]["name"] == "think" + assert json.loads(calls[0]["function"]["arguments"]) == {"thought": "compare docs"} + + +def test_parse_plain_text_returns_no_tool_calls() -> None: + from nemo_retriever.models.local.agent_llm import parse_tool_calls_from_text + + assert parse_tool_calls_from_text("I should search again") == [] + + +def test_parse_tool_calls_rejects_unsupported_parser() -> None: + import pytest + + from nemo_retriever.models.local.agent_llm import parse_tool_calls_from_text + + with pytest.raises(ValueError, match="supported parsers: auto, json"): + parse_tool_calls_from_text("[]", parser="mistral") + + +def test_collapse_parallel_tool_results_for_local_chat_template() -> None: + from nemo_retriever.models.local.agent_llm import _collapse_consecutive_tool_messages + + messages = [ + {"role": "system", "content": "system"}, + {"role": "user", "content": "question"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_a", + "type": "function", + "function": {"name": "retrieve", "arguments": "{}"}, + }, + { + "id": "call_b", + "type": "function", + "function": {"name": "final_results", "arguments": "{}"}, + }, + ], + }, + {"role": "tool", "tool_call_id": "call_a", "content": "Retrieved 3 documents."}, + {"role": "tool", "tool_call_id": "call_b", "content": "Error: doc_ids must be a list."}, + {"role": "assistant", "content": "next"}, + ] + + collapsed = _collapse_consecutive_tool_messages(messages) + + assert [message["role"] for message in collapsed] == ["system", "user", "assistant", "tool", "assistant"] + assert "Tool result for retrieve (call_a):" in collapsed[3]["content"] + assert "Retrieved 3 documents." in collapsed[3]["content"] + assert "Tool result for final_results (call_b):" in collapsed[3]["content"] + assert "Error: doc_ids must be a list." in collapsed[3]["content"] + + +def test_local_vllm_device_helper_accepts_cuda_visible_devices() -> None: + from nemo_retriever.models.local.agent_llm import _cuda_visible_devices_from_device + + assert _cuda_visible_devices_from_device(None) is None + assert _cuda_visible_devices_from_device("cuda:5,6") == "5,6" + assert _cuda_visible_devices_from_device("5,6") == "5,6" + + +def test_local_vllm_device_helper_rejects_cpu() -> None: + import pytest + + from nemo_retriever.models.local.agent_llm import _cuda_visible_devices_from_device + + with pytest.raises(ValueError, match="requires CUDA"): + _cuda_visible_devices_from_device("cpu") + + +def test_vllm_agent_llm_unload_releases_engine() -> None: + import pytest + + from nemo_retriever.models.local.agent_llm import VLLMAgentChatLLM + + llm = VLLMAgentChatLLM.__new__(VLLMAgentChatLLM) + llm._llm = object() + llm._lock = threading.Lock() + + llm.unload() + + assert llm._llm is None + with pytest.raises(RuntimeError, match="unloaded"): + llm._require_loaded() + + +def test_unload_cached_vllm_agent_chat_llms_clears_cache() -> None: + from nemo_retriever.models.local import agent_llm + + cached = MagicMock() + with agent_llm._LOCAL_AGENT_LLM_CACHE_LOCK: + agent_llm._LOCAL_AGENT_LLM_CACHE.clear() + agent_llm._LOCAL_AGENT_LLM_CACHE[("model",)] = cached + + agent_llm.unload_cached_vllm_agent_chat_llms() + + cached.unload.assert_called_once_with() + assert agent_llm._LOCAL_AGENT_LLM_CACHE == {} diff --git a/nemo_retriever/tests/test_agentic_operators.py b/nemo_retriever/tests/test_agentic_operators.py index 8803963603..8025b71068 100644 --- a/nemo_retriever/tests/test_agentic_operators.py +++ b/nemo_retriever/tests/test_agentic_operators.py @@ -218,6 +218,31 @@ def test_think_then_select(self, mock_step): assert result["doc_id"].tolist() == ["d3"] assert mock_step.call_count == 2 + @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") + def test_injected_chat_completion_fn_replaces_http_call(self, mock_step): + from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator + + local_chat = MagicMock( + return_value=_make_tool_call_response( + "log_selected_documents", + {"doc_ids": ["d1"], "message": "d1 is best"}, + ) + ) + + op = SelectionAgentOperator( + llm_model="test-model", + invoke_url="http://localhost/v1/chat/completions", + top_k=1, + max_tokens=234, + chat_completion_fn=local_chat, + ) + result = op.run(self._make_input()) + + mock_step.assert_not_called() + assert local_chat.call_count == 1 + assert local_chat.call_args.kwargs["max_tokens"] == 234 + assert result["doc_id"].tolist() == ["d1"] + @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") def test_extended_relevance_in_prompt(self, mock_step): from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator @@ -369,6 +394,35 @@ def test_simple_mode_retrieve_then_final(self, mock_step): assert 0 in result["step_idx"].values assert "d1" in result["doc_id"].values + @patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") + def test_injected_chat_completion_fn_replaces_http_call(self, mock_step): + from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator + + local_chat = MagicMock( + return_value=_make_tool_call_response( + "final_results", + {"doc_ids": ["d1"], "message": "ok", "search_successful": "true"}, + ) + ) + retriever = MagicMock(return_value=[{"doc_id": "d1", "text": "monetary policy"}]) + + op = ReActAgentOperator( + invoke_url="http://localhost/v1/chat/completions", + llm_model="test-model", + retriever_fn=retriever, + user_msg_type="with_results", + target_top_k=1, + max_tokens=123, + chat_completion_fn=local_chat, + ) + + result = op.run(self._make_input()) + + mock_step.assert_not_called() + assert local_chat.call_count == 1 + assert local_chat.call_args.kwargs["max_tokens"] == 123 + assert result[result["doc_id"] == "d1"]["is_final_result"].astype(bool).any() + @patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") def test_with_results_mode_initial_retrieval(self, mock_step): from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator @@ -431,6 +485,7 @@ def retriever_fn(query_text, top_k): ({"doc_ids": [""], "message": "empty-string id", "search_successful": "true"}, 1, False), ({"doc_ids": [" "], "message": "whitespace id", "search_successful": "true"}, 1, False), ({"doc_ids": ["d1"], "message": "wrong count", "search_successful": "true"}, 2, True), + ({"doc_ids": ["missing"], "message": "hallucinated id", "search_successful": "true"}, 1, False), ({"doc_ids": ["d1"], "message": "bad status", "search_successful": "yes"}, 1, False), ], ) From 9c1055a4b39c46ebd88e580e09f1e06c7ac86869 Mon Sep 17 00:00:00 2001 From: Mahika Wason Date: Mon, 20 Jul 2026 13:43:35 -0700 Subject: [PATCH 02/10] cleanup+tighten support+docs --- .../prerequisites-support-matrix.md | 2 +- .../extraction/workflow-agentic-retrieval.md | 24 +++ nemo_retriever/docs/cli/README.md | 44 ++++-- nemo_retriever/docs/cli/benchmarking.md | 49 ++++++- .../src/nemo_retriever/cli/query/app.py | 41 +++++- .../src/nemo_retriever/cli/query/options.py | 45 +++++- .../src/nemo_retriever/harness/resolution.py | 16 ++ .../src/nemo_retriever/models/__init__.py | 2 - .../nemo_retriever/models/local/agent_llm.py | 138 +++++++----------- .../src/nemo_retriever/query/agentic.py | 26 +++- .../src/nemo_retriever/query/options.py | 8 + .../src/nemo_retriever/query/workflow.py | 8 + nemo_retriever/tests/test_agentic_eval.py | 70 ++++++--- .../tests/test_agentic_local_llm.py | 14 +- .../tests/test_harness_agentic_eval.py | 46 +++++- nemo_retriever/tests/test_root_query_cli.py | 71 ++++++++- 16 files changed, 446 insertions(+), 158 deletions(-) diff --git a/docs/docs/extraction/prerequisites-support-matrix.md b/docs/docs/extraction/prerequisites-support-matrix.md index b0fd53906f..db489af75e 100644 --- a/docs/docs/extraction/prerequisites-support-matrix.md +++ b/docs/docs/extraction/prerequisites-support-matrix.md @@ -106,7 +106,7 @@ When you call [NVIDIA-hosted NIMs](deployment-options.md#when-to-use-nvidia-host | llama-nemotron-rerank-vl-1b-v2 | `https://ai.api.nvidia.com/v1/retrieval/nvidia/llama-nemotron-rerank-vl-1b-v2/reranking` | Optional VL reranker | | nemotron-parse | `https://integrate.api.nvidia.com/v1/chat/completions` with model ID `nvidia/nemotron-parse` | Optional `extract_method="nemotron_parse"` | | nemotron-3-nano-omni-30b-a3b-reasoning | `https://integrate.api.nvidia.com/v1/chat/completions` with model ID `nvidia/nemotron-3-nano-omni-30b-a3b-reasoning` | Optional image captioning | -| llama-3.3-nemotron-super-49b-v1.5 | `https://integrate.api.nvidia.com/v1/chat/completions` with model ID `nvidia/llama-3.3-nemotron-super-49b-v1.5` | Optional `/v1/answer` (Helm `answer_llm`) and agentic RAG; not part of the default extraction pipeline. Helm auto-wires to the in-cluster NIM when `nimOperator.answer_llm` is enabled | +| llama-3.3-nemotron-super-49b-v1.5 | `https://integrate.api.nvidia.com/v1/chat/completions` with model ID `nvidia/llama-3.3-nemotron-super-49b-v1.5` | Optional `/v1/answer` (Helm `answer_llm`) and OpenAI-compatible agentic RAG endpoint mode; not part of the default extraction pipeline. Agentic query/harness runs default to local in-process vLLM instead. Helm auto-wires to the in-cluster NIM when `nimOperator.answer_llm` is enabled | | parakeet-1-1b-ctc-en-us | `grpc.nvcf.nvidia.com:443` (function ID from [build.nvidia.com](https://build.nvidia.com/)) | Optional ASR; refer to [Parakeet hosted inference](audio-video.md#parakeet-hosted-inference-build-nvidia) | For local Hugging Face OCR language mode (`multi` vs `english`), Helm OCR image overrides, and local model install, refer to [OCR and scanned documents](multimodal-extraction.md#ocr-and-scanned-documents), [OCR NIM configuration](https://github.com/NVIDIA/NeMo-Retriever/blob/main/nemo_retriever/helm/README.md#ocr-nim-configuration), and [CLI — OCR language mode](https://github.com/NVIDIA/NeMo-Retriever/blob/main/nemo_retriever/docs/cli/README.md#ocr-language-mode). diff --git a/docs/docs/extraction/workflow-agentic-retrieval.md b/docs/docs/extraction/workflow-agentic-retrieval.md index d29b36331b..2299d22b7a 100644 --- a/docs/docs/extraction/workflow-agentic-retrieval.md +++ b/docs/docs/extraction/workflow-agentic-retrieval.md @@ -4,6 +4,30 @@ NeMo Retriever Library provides ingestion, embedding, storage, and retrieval building blocks (jobs, chunking, vector stores, reranking) that you orchestrate in application code or frameworks. + +## Local agentic retrieval + +The `retriever query --agentic` and harness BEIR agentic paths default to an +in-process local vLLM agent LLM. If no agent model is provided, the library loads +`nemotron-8b` (`nvidia/Llama-3.1-Nemotron-Nano-8B-v1`) on the local CUDA host. +The larger `super-49b` profile is also supported. Other custom in-process LLMs +are not supported yet because the agent loop depends on OpenAI-style tool-call +messages; use an OpenAI-compatible endpoint for custom models. + +```bash +retriever query "find documents about parser behavior" --agentic +``` + +For custom or already deployed chat models, opt into the endpoint path: + +```bash +retriever query "find documents about parser behavior" \ + --agentic \ + --agentic-llm-backend openai_compatible \ + --agentic-llm-model custom-remote-model \ + --agentic-invoke-url http://localhost:9000/v1/chat/completions +``` + ## MCP access for agents `retriever service start` mounts a FastMCP HTTP endpoint at `/mcp` by default. Agents can use that endpoint to call the running service for health checks, pipeline introspection, document ingestion, job status, VectorDB query, and answer generation. If service auth is enabled, the MCP endpoint uses the same bearer-token middleware as the REST API. diff --git a/nemo_retriever/docs/cli/README.md b/nemo_retriever/docs/cli/README.md index 9806e0d877..e5ad446ac1 100644 --- a/nemo_retriever/docs/cli/README.md +++ b/nemo_retriever/docs/cli/README.md @@ -178,18 +178,29 @@ output are not used for content-type matching. `--agentic` swaps the single dense pass for an LLM-driven ReAct loop: the agent issues several retrieval sub-queries, fuses the candidates, and selects a final ranking. It searches the same LanceDB table built by `retriever ingest`, so it is -a drop-in alternative to standard retrieval — add `--agentic` and name the chat -model the agent drives with `--agentic-llm-model` (required): +a drop-in alternative to standard retrieval. + +By default, agentic retrieval runs the agent LLM in process with local vLLM and +`nemotron-8b` (`nvidia/Llama-3.1-Nemotron-Nano-8B-v1`). This requires a CUDA GPU +host and the local extras installed. Use `--agentic-llm-backend openai_compatible` +when you want a custom model or a separately hosted OpenAI-compatible endpoint. ```bash +# default local vLLM agent LLM: nemotron-8b +retriever query "how does the ingestion pipeline handle tables?" \ + --agentic + +# larger supported local profile retriever query "how does the ingestion pipeline handle tables?" \ --agentic \ - --agentic-llm-model nvidia/llama-3.3-nemotron-super-49b-v1.5 + --agentic-llm-model super-49b \ + --agentic-local-tensor-parallel-size 2 -# remote agent + embedding endpoints, fewer reasoning rounds +# custom/self-hosted model through an OpenAI-compatible endpoint retriever query "summarize the deployment options" \ --agentic \ - --agentic-llm-model nvidia/llama-3.3-nemotron-super-49b-v1.5 \ + --agentic-llm-backend openai_compatible \ + --agentic-llm-model custom-remote-model \ --agentic-invoke-url http://localhost:9000/v1/chat/completions \ --embed-invoke-url http://localhost:8000/v1 \ --agentic-react-max-steps 5 @@ -201,8 +212,8 @@ produced it (`final_results`, `rrf`, or `selection_agent`). It reuses the same `--top-k`, `--lancedb-uri`, `--table-name`, `--embed-invoke-url`, and `--embed-model-name` options as standard retrieval. -**How it works.** Each agentic query runs `Query → ReActAgentOperator → (RRF -fusion) → SelectionAgentOperator → ranked results`: +**How it works.** Each agentic query runs `Query -> ReActAgentOperator -> (RRF +fusion) -> SelectionAgentOperator -> ranked results`: - `ReActAgentOperator` runs the per-query ReAct loop; every `retrieve` tool call delegates to the standard `Retriever`, so the agent searches the same vector @@ -214,10 +225,25 @@ fusion) → SelectionAgentOperator → ranked results`: Agentic-only knobs (apply only with `--agentic`): +- `--agentic-llm-backend` (default `in_process`) — use local in-process vLLM, or + set `openai_compatible` for a hosted/self-hosted chat-completions endpoint. +- `--agentic-llm-model` — local profile alias/model ID for `in_process` + (`nemotron-8b` by default; `super-49b` also supported), or the remote model ID + when using `openai_compatible`. - `--agentic-invoke-url` — OpenAI-compatible chat-completions endpoint for the - agent LLM; defaults to the operators' built-in endpoint when omitted. + agent LLM; used only with `--agentic-llm-backend openai_compatible`. +- `--agentic-local-llm-backend` (default `vllm`) — local runtime for in-process + agent LLMs. Other local runtimes are not supported yet. +- `--agentic-local-device` — CUDA device ids for local vLLM, for example `0`, + `cuda:0`, or `0,1`. +- `--agentic-local-gpu-memory-utilization` (default `0.8`) — vLLM GPU memory + utilization for the local agent LLM. +- `--agentic-local-tensor-parallel-size` (default `1`) — tensor parallel size for + local vLLM. +- `--agentic-local-max-model-len` and `--agentic-local-max-num-seqs` — optional + vLLM engine overrides for the local agent LLM. - `--agentic-reasoning-effort` (default `high`) — `reasoning_effort` forwarded on - agentic LLM calls. + OpenAI-compatible agentic LLM calls; ignored by the local adapter. - `--agentic-backend-top-k` (default `20`) — candidates pulled from the vector DB per retrieval call. - `--agentic-react-max-steps` (default `50`) — maximum ReAct loop iterations. diff --git a/nemo_retriever/docs/cli/benchmarking.md b/nemo_retriever/docs/cli/benchmarking.md index 72a9f69f3b..399ffa914d 100644 --- a/nemo_retriever/docs/cli/benchmarking.md +++ b/nemo_retriever/docs/cli/benchmarking.md @@ -58,20 +58,53 @@ The agentic harness path runs the same ReAct retrieval graph used by root query, but only after ingest and only for BEIR evaluation (`evaluation.mode: beir`). `retriever pipeline run` does not expose agentic evaluation flags. +By default, agentic harness evaluation uses the in-process local vLLM backend +with `nemotron-8b`. Custom LLMs are not supported in process yet; run them behind +an OpenAI-compatible chat-completions endpoint and set +`query.agentic_llm_backend=openai_compatible`. + Minimal BEIR override example: +```bash +retriever harness run jp20_beir \ + --set query.agentic=true +``` + +Larger supported local profile: + +```bash +retriever harness run jp20_beir \ + --set query.agentic=true \ + --set query.agentic_llm_model=super-49b \ + --set query.agentic_local_tensor_parallel_size=2 +``` + +Custom/self-hosted OpenAI-compatible endpoint: + ```bash retriever harness run jp20_beir \ --set query.agentic=true \ - --set query.agentic_llm_model=nvidia/llama-3.3-nemotron-super-49b-v1.5 + --set query.agentic_llm_backend=openai_compatible \ + --set query.agentic_llm_model=custom-remote-model \ + --set query.agentic_invoke_url=http://localhost:9000/v1/chat/completions ``` Useful agentic query overrides: -- `query.agentic_llm_model` — chat model used by the ReAct and selection agents; - required when `query.agentic=true`. -- `query.agentic_invoke_url` — OpenAI-compatible chat-completions endpoint. Omit - to use the built-in NVIDIA endpoint. +- `query.agentic_llm_backend` — `in_process` by default, or + `openai_compatible` for hosted/self-hosted chat-completions endpoints. +- `query.agentic_llm_model` — local profile alias/model ID for `in_process` + (`nemotron-8b` by default; `super-49b` also supported), or the remote model ID + when using `openai_compatible`. +- `query.agentic_invoke_url` — OpenAI-compatible chat-completions endpoint; used + only with `query.agentic_llm_backend=openai_compatible`. +- `query.agentic_local_llm_backend` — local runtime for in-process agent LLMs; + currently only `vllm` is supported. +- `query.agentic_local_hf_cache_dir`, `query.agentic_local_device`, + `query.agentic_local_gpu_memory_utilization`, + `query.agentic_local_tensor_parallel_size`, `query.agentic_local_max_model_len`, + and `query.agentic_local_max_num_seqs` — local vLLM loading and scheduling + controls. - `query.agentic_backend_top_k` — backend candidate pool per ReAct retrieval call. Must be at least the final requested metric depth (`max(evaluation.ks)`). - `query.agentic_react_max_steps` — maximum ReAct loop iterations per query @@ -80,9 +113,9 @@ Useful agentic query overrides: the agent; `0` disables truncation. - `query.agentic_num_concurrent` — number of queries the agent batch runs concurrently (defaults to `1`). -- `query.agentic_temperature` — defaults to `0.0`; hosted/default NVIDIA - endpoints are validated as `0.0..1.0`, while other OpenAI-compatible endpoints - allow `0.0..2.0`. +- `query.agentic_temperature` — defaults to `0.0`; local and non-NVIDIA + OpenAI-compatible endpoints allow `0.0..2.0`, while hosted/default NVIDIA + endpoints are validated as `0.0..1.0`. - `query.agentic_reasoning_effort` — optional provider-specific field forwarded only when configured. diff --git a/nemo_retriever/src/nemo_retriever/cli/query/app.py b/nemo_retriever/src/nemo_retriever/cli/query/app.py index f4bc1d5a92..9995e700aa 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/app.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/app.py @@ -187,7 +187,15 @@ def _local_command( max_text_chars: opts.MaxTextCharsOption = None, agentic: opts.AgenticOption = False, agentic_llm_model: opts.AgenticLlmModelOption = None, + agentic_llm_backend: opts.AgenticLlmBackendOption = "in_process", agentic_invoke_url: opts.AgenticInvokeUrlOption = None, + agentic_local_llm_backend: opts.AgenticLocalLlmBackendOption = "vllm", + agentic_local_hf_cache_dir: opts.AgenticLocalHfCacheDirOption = None, + agentic_local_device: opts.AgenticLocalDeviceOption = None, + agentic_local_gpu_memory_utilization: opts.AgenticLocalGpuMemoryUtilizationOption = 0.8, + agentic_local_tensor_parallel_size: opts.AgenticLocalTensorParallelSizeOption = 1, + agentic_local_max_model_len: opts.AgenticLocalMaxModelLenOption = None, + agentic_local_max_num_seqs: opts.AgenticLocalMaxNumSeqsOption = None, agentic_reasoning_effort: opts.AgenticReasoningEffortOption = "high", agentic_backend_top_k: opts.AgenticBackendTopKOption = 20, agentic_react_max_steps: opts.AgenticReactMaxStepsOption = 50, @@ -201,16 +209,31 @@ def _local_command( embed_invoke_url = os.environ.get("EMBED_INVOKE_URL") or None rerank = rerank or bool(reranker_invoke_url) or bool(reranker_model_name) or bool(reranker_backend) silence_noisy_libraries() - if agentic and not agentic_llm_model: - typer.echo("Error: --agentic requires --agentic-llm-model.", err=True) - raise typer.Exit(1) - if agentic: + normalized_agentic_llm_backend = str(agentic_llm_backend or "in_process").strip().lower() + if normalized_agentic_llm_backend not in {"in_process", "openai_compatible"}: + typer.echo( + "Error: --agentic-llm-backend must be one of ['in_process', 'openai_compatible'].", + err=True, + ) + raise typer.Exit(1) + if normalized_agentic_llm_backend == "in_process" and not agentic_llm_model: + agentic_llm_model = "nemotron-8b" + if normalized_agentic_llm_backend == "openai_compatible" and not agentic_llm_model: + typer.echo( + "Error: --agentic-llm-backend openai_compatible requires --agentic-llm-model.", + err=True, + ) + raise typer.Exit(1) + backend_error = agentic_backend_top_k_error(agentic_backend_top_k, target_top_k=top_k) if backend_error: typer.echo(f"Error: {backend_error}", err=True) raise typer.Exit(1) - temperature_error = agentic_temperature_error(agentic_temperature, invoke_url=agentic_invoke_url) + temperature_invoke_url = ( + agentic_invoke_url if normalized_agentic_llm_backend == "openai_compatible" else "local://in-process" + ) + temperature_error = agentic_temperature_error(agentic_temperature, invoke_url=temperature_invoke_url) if temperature_error: typer.echo(f"Error: {temperature_error}", err=True) raise typer.Exit(1) @@ -248,7 +271,15 @@ def _local_command( agentic=QueryAgenticOptions( enabled=agentic, llm_model=agentic_llm_model, + llm_backend=normalized_agentic_llm_backend, invoke_url=agentic_invoke_url, + local_llm_backend=agentic_local_llm_backend, + local_hf_cache_dir=agentic_local_hf_cache_dir, + local_device=agentic_local_device, + local_gpu_memory_utilization=agentic_local_gpu_memory_utilization, + local_tensor_parallel_size=agentic_local_tensor_parallel_size, + local_max_model_len=agentic_local_max_model_len, + local_max_num_seqs=agentic_local_max_num_seqs, reasoning_effort=agentic_reasoning_effort, backend_top_k=agentic_backend_top_k, react_max_steps=agentic_react_max_steps, diff --git a/nemo_retriever/src/nemo_retriever/cli/query/options.py b/nemo_retriever/src/nemo_retriever/cli/query/options.py index 688e8ea6f2..fb8b93baf2 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/options.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/options.py @@ -159,7 +159,17 @@ str | None, typer.Option( "--agentic-llm-model", - help="Chat model the agent drives. Required when --agentic is set.", + help=( + "Chat model the agent drives. Defaults to nemotron-8b for --agentic-llm-backend=in_process; " + "required for openai_compatible endpoints." + ), + ), +] +AgenticLlmBackendOption = Annotated[ + str, + typer.Option( + "--agentic-llm-backend", + help="Agent LLM backend: in_process (default) or openai_compatible.", ), ] AgenticInvokeUrlOption = Annotated[ @@ -169,6 +179,39 @@ help="OpenAI-compatible chat-completions endpoint for the agent LLM (agentic mode).", ), ] +AgenticLocalLlmBackendOption = Annotated[ + str, + typer.Option("--agentic-local-llm-backend", help="Local in-process agent LLM runtime. Currently only vllm."), +] +AgenticLocalHfCacheDirOption = Annotated[ + str | None, + typer.Option("--agentic-local-hf-cache-dir", help="Hugging Face cache directory for local agent LLM weights."), +] +AgenticLocalDeviceOption = Annotated[ + str | None, + typer.Option("--agentic-local-device", help="CUDA device ids for local vLLM, e.g. 0, cuda:0, or 0,1."), +] +AgenticLocalGpuMemoryUtilizationOption = Annotated[ + float, + typer.Option( + "--agentic-local-gpu-memory-utilization", + min=0.0, + max=1.0, + help="Fraction of GPU memory reserved by local vLLM for the agent LLM.", + ), +] +AgenticLocalTensorParallelSizeOption = Annotated[ + int, + typer.Option("--agentic-local-tensor-parallel-size", min=1, help="Tensor parallel size for local vLLM."), +] +AgenticLocalMaxModelLenOption = Annotated[ + int | None, + typer.Option("--agentic-local-max-model-len", min=1, help="Optional local vLLM max_model_len override."), +] +AgenticLocalMaxNumSeqsOption = Annotated[ + int | None, + typer.Option("--agentic-local-max-num-seqs", min=1, help="Optional local vLLM max_num_seqs override."), +] AgenticReasoningEffortOption = Annotated[ str | None, typer.Option( diff --git a/nemo_retriever/src/nemo_retriever/harness/resolution.py b/nemo_retriever/src/nemo_retriever/harness/resolution.py index fa9336f24f..d9047d4fc9 100644 --- a/nemo_retriever/src/nemo_retriever/harness/resolution.py +++ b/nemo_retriever/src/nemo_retriever/harness/resolution.py @@ -128,7 +128,15 @@ def _override_child_keys(prefix: str, paths: set[str]) -> set[str]: "query.table_name", "query.agentic", "query.agentic_llm_model", + "query.agentic_llm_backend", "query.agentic_invoke_url", + "query.agentic_local_llm_backend", + "query.agentic_local_hf_cache_dir", + "query.agentic_local_device", + "query.agentic_local_gpu_memory_utilization", + "query.agentic_local_tensor_parallel_size", + "query.agentic_local_max_model_len", + "query.agentic_local_max_num_seqs", "query.agentic_reasoning_effort", "query.agentic_backend_top_k", "query.agentic_react_max_steps", @@ -437,7 +445,15 @@ def build_query_request(resolved: dict[str, Any], query_text: str) -> QueryReque agentic=QueryAgenticOptions( enabled=bool(query.get("agentic", False)), llm_model=query.get("agentic_llm_model"), + llm_backend=str(query.get("agentic_llm_backend") or "in_process"), invoke_url=query.get("agentic_invoke_url"), + local_llm_backend=str(query.get("agentic_local_llm_backend") or "vllm"), + local_hf_cache_dir=query.get("agentic_local_hf_cache_dir"), + local_device=query.get("agentic_local_device"), + local_gpu_memory_utilization=float(query.get("agentic_local_gpu_memory_utilization") or 0.8), + local_tensor_parallel_size=int(query.get("agentic_local_tensor_parallel_size") or 1), + local_max_model_len=query.get("agentic_local_max_model_len"), + local_max_num_seqs=query.get("agentic_local_max_num_seqs"), reasoning_effort=query.get("agentic_reasoning_effort"), backend_top_k=int(query.get("agentic_backend_top_k") or 20), react_max_steps=int(query.get("agentic_react_max_steps") or 50), diff --git a/nemo_retriever/src/nemo_retriever/models/__init__.py b/nemo_retriever/src/nemo_retriever/models/__init__.py index 4864af41f9..7ca2451526 100644 --- a/nemo_retriever/src/nemo_retriever/models/__init__.py +++ b/nemo_retriever/src/nemo_retriever/models/__init__.py @@ -273,7 +273,6 @@ def create_local_agent_llm( tensor_parallel_size: int = 1, max_model_len: int | None = None, max_num_seqs: int | None = None, - tool_call_parser: str | None = None, ) -> Any: """Create a cached local agent LLM chat-completion callable. @@ -295,7 +294,6 @@ def create_local_agent_llm( tensor_parallel_size=tensor_parallel_size, max_model_len=max_model_len, max_num_seqs=max_num_seqs, - tool_call_parser=tool_call_parser, ) raise ValueError(f"Unsupported local agent LLM backend {backend!r}") diff --git a/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py index 03ae56cfa6..c61161cc53 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py +++ b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py @@ -11,7 +11,6 @@ import string import threading from copy import deepcopy -from dataclasses import dataclass, field from typing import Any, Mapping, Optional, Sequence from nemo_retriever.models.hf_cache import configure_global_hf_cache_base @@ -20,28 +19,45 @@ logger = logging.getLogger(__name__) +# Conservative cap for short tool-call JSON responses. Larger evals should +# tune this with completion-token, truncation, and tool-parse telemetry. _DEFAULT_MAX_TOKENS = 512 _LOCAL_AGENT_LLM_CACHE: dict[tuple[Any, ...], "VLLMAgentChatLLM"] = {} _LOCAL_AGENT_LLM_CACHE_LOCK = threading.Lock() -@dataclass(frozen=True) -class AgentLLMProfile: - """Metadata for local agent LLMs with known offline tool-call behavior.""" +def _normalize_name(name: str) -> str: + return name.strip().casefold() - model_id: str - aliases: tuple[str, ...] = () - engine_kwargs: Mapping[str, Any] = field(default_factory=dict) - request_extras: Mapping[str, Any] = field(default_factory=dict) - tool_call_parser: str = "json" - tool_prompt_format: str = "none" - pass_tools_to_chat_template: bool = True - def engine_kwargs_for_local(self) -> dict[str, Any]: - return _mutable_copy(self.engine_kwargs) +_NANO_8B_MODEL_ID = "nvidia/Llama-3.1-Nemotron-Nano-8B-v1" +_SUPER_49B_MODEL_ID = "nvidia/Llama-3_3-Nemotron-Super-49B-v1" +_NANO_8B_ALIASES = ( + "llama-3.1-nemotron-nano-8b-v1", + "nemotron-nano-8b", + "nemotron-8b", + "nvidia/llama-3.1-nemotron-nano-8b-v1", +) +_SUPER_49B_ALIASES = ( + "llama-3.3-nemotron-super-49b-v1", + "nemotron-super-49b", + "super-49b", + "nvidia/llama-3.3-nemotron-super-49b-v1", + "nvidia/llama-3_3-nemotron-super-49b-v1", +) - def request_extras_for_local(self) -> dict[str, Any]: - return _mutable_copy(self.request_extras) +# Supported local agent LLM names. Keep this intentionally small: V1 only +# supports internal/NVIDIA models whose local tool-call behavior we have tested. +_AGENT_LLM_MODEL_ALIASES: dict[str, str] = { + _normalize_name(name): _NANO_8B_MODEL_ID for name in (_NANO_8B_MODEL_ID, *_NANO_8B_ALIASES) +} +_AGENT_LLM_MODEL_ALIASES.update( + {_normalize_name(name): _SUPER_49B_MODEL_ID for name in (_SUPER_49B_MODEL_ID, *_SUPER_49B_ALIASES)} +) +_SUPPORTED_AGENT_LLM_NAMES = tuple( + sorted((_NANO_8B_MODEL_ID, *_NANO_8B_ALIASES, _SUPER_49B_MODEL_ID, *_SUPER_49B_ALIASES)) +) +_NEMOTRON_JSON_TOOL_PROMPT_EXTRAS = {"chat_template_kwargs": {"tools_in_user_message": True}} def _mutable_copy(value: Any) -> Any: @@ -52,64 +68,26 @@ def _mutable_copy(value: Any) -> Any: return deepcopy(value) -def _normalize_name(name: str) -> str: - return name.strip().casefold() - +def supported_agent_llm_names() -> tuple[str, ...]: + """Return supported local agent LLM model IDs and aliases.""" -_NEMOTRON_JSON_TOOL_PROMPT_EXTRAS = {"chat_template_kwargs": {"tools_in_user_message": True}} + return _SUPPORTED_AGENT_LLM_NAMES -_AGENT_LLM_PROFILES: tuple[AgentLLMProfile, ...] = ( - AgentLLMProfile( - model_id="nvidia/Llama-3.1-Nemotron-Nano-8B-v1", - aliases=( - "llama-3.1-nemotron-nano-8b-v1", - "nemotron-nano-8b", - "nemotron-8b", - "nvidia/llama-3.1-nemotron-nano-8b-v1", - ), - request_extras=_NEMOTRON_JSON_TOOL_PROMPT_EXTRAS, - tool_prompt_format="json", - ), - AgentLLMProfile( - model_id="nvidia/Llama-3_3-Nemotron-Super-49B-v1", - aliases=( - "llama-3.3-nemotron-super-49b-v1", - "nemotron-super-49b", - "super-49b", - "nvidia/llama-3.3-nemotron-super-49b-v1", - "nvidia/llama-3_3-nemotron-super-49b-v1", - ), - request_extras=_NEMOTRON_JSON_TOOL_PROMPT_EXTRAS, - tool_prompt_format="json", - ), -) -_PROFILE_LOOKUP: dict[str, AgentLLMProfile] = {} -for _profile in _AGENT_LLM_PROFILES: - _PROFILE_LOOKUP[_normalize_name(_profile.model_id)] = _profile - for _alias in _profile.aliases: - _PROFILE_LOOKUP[_normalize_name(_alias)] = _profile +def _supported_agent_llm_names_for_error() -> str: + return ", ".join(supported_agent_llm_names()) -def _supported_agent_llm_names() -> str: - names: list[str] = [] - for profile in _AGENT_LLM_PROFILES: - names.append(profile.model_id) - names.extend(profile.aliases) - return ", ".join(sorted(names)) +def is_supported_agent_llm_model(name: str) -> bool: + """Return whether *name* is a supported in-process local agent LLM.""" - -def get_agent_llm_profile(name: str) -> AgentLLMProfile | None: - """Return a supported local agent LLM profile, or ``None`` when unsupported.""" - - return _PROFILE_LOOKUP.get(_normalize_name(name)) + return _normalize_name(name) in _AGENT_LLM_MODEL_ALIASES def resolve_agent_llm_model_name(name: str) -> str: """Resolve a local agent LLM alias to its Hugging Face model ID.""" - profile = get_agent_llm_profile(name) - return profile.model_id if profile is not None else name + return _AGENT_LLM_MODEL_ALIASES.get(_normalize_name(name), name) class VLLMAgentChatLLM(BaseModel): @@ -132,18 +110,19 @@ def __init__( max_model_len: Optional[int] = None, max_num_seqs: Optional[int] = None, max_tokens: int = _DEFAULT_MAX_TOKENS, - tool_call_parser: Optional[str] = None, ) -> None: super().__init__() - profile = get_agent_llm_profile(model_path) - if profile is None: + resolved_model_path = resolve_agent_llm_model_name(model_path) + if resolved_model_path == model_path and not is_supported_agent_llm_model(model_path): raise ValueError( f"Unsupported local agent LLM model {model_path!r}. " - f"Supported models and aliases: {_supported_agent_llm_names()}." + "Custom in-process agent LLMs are not supported yet. " + "Use llm_backend='openai_compatible' with invoke_url for a custom/self-hosted endpoint, " + f"or choose one of: {_supported_agent_llm_names_for_error()}." ) - model_path = profile.model_id + model_path = resolved_model_path cuda_visible_devices = _cuda_visible_devices_from_device(device) if cuda_visible_devices is not None: os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices @@ -158,16 +137,13 @@ def __init__( _raise_if_cuda_unavailable() self._model_path = model_path self._max_tokens = int(max_tokens) - self._tool_call_parser = (tool_call_parser or profile.tool_call_parser or "json").strip().lower() - self._request_extras = profile.request_extras_for_local() - self._tool_prompt_format = profile.tool_prompt_format - self._pass_tools_to_chat_template = profile.pass_tools_to_chat_template + self._request_extras = _mutable_copy(_NEMOTRON_JSON_TOOL_PROMPT_EXTRAS) self._lock = threading.Lock() configure_global_hf_cache_base(hf_cache_dir) - revision = get_hf_revision(model_path, strict=False) + revision = get_hf_revision(model_path) - engine_kwargs = profile.engine_kwargs_for_local() + engine_kwargs: dict[str, Any] = {} if max_model_len is not None: engine_kwargs["max_model_len"] = int(max_model_len) if max_num_seqs is not None: @@ -209,7 +185,7 @@ def __call__( ) chat_kwargs = self._build_chat_kwargs(extra_body) active_tools = tools if tools and tool_choice != "none" else None - if active_tools and self._pass_tools_to_chat_template: + if active_tools: chat_kwargs["tools"] = active_tools local_messages = self._normalize_messages(messages, tools=active_tools) @@ -220,10 +196,7 @@ def __call__( request_output = outputs[0] completion = request_output.outputs[0] text = str(getattr(completion, "text", "") or "").strip() - tool_calls = _tool_calls_from_completion(completion) or parse_tool_calls_from_text( - text, - parser=self._tool_call_parser, - ) + tool_calls = _tool_calls_from_completion(completion) or parse_tool_calls_from_text(text) finish_reason = "tool_calls" if tool_calls else str(getattr(completion, "finish_reason", None) or "stop") message: dict[str, Any] = {"role": "assistant"} @@ -263,7 +236,7 @@ def _normalize_messages( tools: Optional[Sequence[dict[str, Any]]] = None, ) -> list[dict[str, Any]]: normalized: list[dict[str, Any]] = [] - tool_prompt = _json_tool_prompt(tools) if tools and self._tool_prompt_format == "json" else None + tool_prompt = _json_tool_prompt(tools) if tools else None for idx, message in enumerate(messages): msg = dict(message) content = msg.get("content") @@ -359,7 +332,6 @@ def create_cached_vllm_agent_chat_llm( tensor_parallel_size: int = 1, max_model_len: Optional[int] = None, max_num_seqs: Optional[int] = None, - tool_call_parser: Optional[str] = None, ) -> VLLMAgentChatLLM: """Create or reuse a local agent LLM keyed by heavyweight load settings.""" @@ -371,7 +343,6 @@ def create_cached_vllm_agent_chat_llm( int(tensor_parallel_size), int(max_model_len) if max_model_len is not None else None, int(max_num_seqs) if max_num_seqs is not None else None, - (tool_call_parser or "").strip().lower() or None, ) with _LOCAL_AGENT_LLM_CACHE_LOCK: cached = _LOCAL_AGENT_LLM_CACHE.get(key) @@ -384,7 +355,6 @@ def create_cached_vllm_agent_chat_llm( tensor_parallel_size=tensor_parallel_size, max_model_len=max_model_len, max_num_seqs=max_num_seqs, - tool_call_parser=tool_call_parser, ) _LOCAL_AGENT_LLM_CACHE[key] = cached return cached @@ -401,11 +371,9 @@ def unload_cached_vllm_agent_chat_llms() -> None: model.unload() -def parse_tool_calls_from_text(text: str, *, parser: str = "json") -> list[dict[str, Any]]: +def parse_tool_calls_from_text(text: str) -> list[dict[str, Any]]: """Parse common offline tool-call JSON into OpenAI ``tool_calls`` shape.""" - if parser not in {"auto", "json"}: - raise ValueError(f"Unsupported local tool-call parser {parser!r}; supported parsers: auto, json") payload = _load_json_payload(text) if payload is None: return [] diff --git a/nemo_retriever/src/nemo_retriever/query/agentic.py b/nemo_retriever/src/nemo_retriever/query/agentic.py index 8b121321fd..31a45e4ec7 100644 --- a/nemo_retriever/src/nemo_retriever/query/agentic.py +++ b/nemo_retriever/src/nemo_retriever/query/agentic.py @@ -57,9 +57,10 @@ AGENTIC_REACT_MAX_STEPS = 50 AGENTIC_TEMPERATURE = 0.0 # agent LLM sampling temperature (0.0 = greedy) AGENTIC_MAX_TOKENS: Optional[int] = None -AGENTIC_LLM_BACKEND = "openai_compatible" +AGENTIC_LLM_BACKEND = "in_process" AGENTIC_LLM_BACKENDS = frozenset({"openai_compatible", "in_process"}) AGENTIC_LOCAL_LLM_BACKEND = "vllm" +AGENTIC_LOCAL_LLM_MODEL = "nemotron-8b" AGENTIC_LOCAL_LLM_BACKENDS = frozenset({"vllm"}) @@ -153,7 +154,6 @@ class AgenticRetrievalConfig: local_tensor_parallel_size: int = 1 local_max_model_len: Optional[int] = None local_max_num_seqs: Optional[int] = None - local_tool_call_parser: Optional[str] = None api_key: Optional[str] = None react_max_steps: int = AGENTIC_REACT_MAX_STEPS text_truncation: int = AGENTIC_TEXT_TRUNCATION @@ -193,8 +193,25 @@ def __post_init__(self) -> None: ) object.__setattr__(self, "local_llm_backend", local_llm_backend) - if self.llm_model is None or not str(self.llm_model).strip(): - raise ValueError("Agentic retrieval requires a non-empty llm_model.") + llm_model = str(self.llm_model or "").strip() + if not llm_model: + if llm_backend == "in_process": + llm_model = AGENTIC_LOCAL_LLM_MODEL + else: + raise ValueError("Agentic retrieval requires a non-empty llm_model.") + + if llm_backend == "in_process": + from nemo_retriever.models.local.agent_llm import is_supported_agent_llm_model, supported_agent_llm_names + + if not is_supported_agent_llm_model(llm_model): + supported = ", ".join(supported_agent_llm_names()) + raise ValueError( + f"Unsupported in-process agentic LLM model {llm_model!r}. " + "Custom in-process agent LLMs are not supported yet. " + "Use llm_backend='openai_compatible' with invoke_url for a custom/self-hosted endpoint, " + f"or choose one of: {supported}." + ) + object.__setattr__(self, "llm_model", llm_model) object.__setattr__( self, "enforce_top_k", @@ -315,7 +332,6 @@ def _build_agent_chat_completion_fn(cfg: AgenticRetrievalConfig) -> Any | None: tensor_parallel_size=int(cfg.local_tensor_parallel_size), max_model_len=cfg.local_max_model_len, max_num_seqs=cfg.local_max_num_seqs, - tool_call_parser=cfg.local_tool_call_parser, ) raise ValueError(f"Unsupported agentic llm_backend {cfg.llm_backend!r}") diff --git a/nemo_retriever/src/nemo_retriever/query/options.py b/nemo_retriever/src/nemo_retriever/query/options.py index ad0e344cf2..c09c2628cc 100644 --- a/nemo_retriever/src/nemo_retriever/query/options.py +++ b/nemo_retriever/src/nemo_retriever/query/options.py @@ -55,7 +55,15 @@ class QueryAgenticOptions: enabled: bool = False llm_model: str | None = None + llm_backend: str = "in_process" invoke_url: str | None = None + local_llm_backend: str = "vllm" + local_hf_cache_dir: str | None = None + local_device: str | None = None + local_gpu_memory_utilization: float = 0.8 + local_tensor_parallel_size: int = 1 + local_max_model_len: int | None = None + local_max_num_seqs: int | None = None reasoning_effort: str | None = None backend_top_k: int = 20 react_max_steps: int = 50 diff --git a/nemo_retriever/src/nemo_retriever/query/workflow.py b/nemo_retriever/src/nemo_retriever/query/workflow.py index 2a193e539b..890d456d35 100644 --- a/nemo_retriever/src/nemo_retriever/query/workflow.py +++ b/nemo_retriever/src/nemo_retriever/query/workflow.py @@ -174,7 +174,15 @@ def build_agentic_config(request: QueryRequest, *, top_k: int | None = None) -> "embedding_endpoint": request.embed.embed_invoke_url, "embedding_api_key": api_key or "", "llm_model": request.agentic.llm_model, + "llm_backend": request.agentic.llm_backend, "invoke_url": request.agentic.invoke_url, + "local_llm_backend": request.agentic.local_llm_backend, + "local_hf_cache_dir": request.agentic.local_hf_cache_dir, + "local_device": request.agentic.local_device, + "local_gpu_memory_utilization": request.agentic.local_gpu_memory_utilization, + "local_tensor_parallel_size": request.agentic.local_tensor_parallel_size, + "local_max_model_len": request.agentic.local_max_model_len, + "local_max_num_seqs": request.agentic.local_max_num_seqs, "api_key": api_key, "reasoning_effort": request.agentic.reasoning_effort, "backend_top_k": int(request.agentic.backend_top_k), diff --git a/nemo_retriever/tests/test_agentic_eval.py b/nemo_retriever/tests/test_agentic_eval.py index 9b06c19e8e..42f2ccf5ab 100644 --- a/nemo_retriever/tests/test_agentic_eval.py +++ b/nemo_retriever/tests/test_agentic_eval.py @@ -102,11 +102,11 @@ def test_build_beir_run_from_ranked_doc_ids_rejects_length_mismatch(): def test_agentic_config_validates_max_tokens(): from nemo_retriever.query.agentic import AgenticRetrievalConfig - cfg = AgenticRetrievalConfig(llm_model="test-model", max_tokens="128") + cfg = AgenticRetrievalConfig(llm_model="test-model", llm_backend="openai_compatible", max_tokens="128") assert cfg.max_tokens == 128 with pytest.raises(ValueError, match="max_tokens"): - AgenticRetrievalConfig(llm_model="test-model", max_tokens=0) + AgenticRetrievalConfig(llm_model="test-model", llm_backend="openai_compatible", max_tokens=0) @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") @@ -126,7 +126,10 @@ 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", max_tokens=77 + llm_model="test-model", + llm_backend="openai_compatible", + invoke_url="http://localhost/v1/chat/completions", + max_tokens=77, ) result = AgenticRetriever(cfg, match_mode="pdf_page").retrieve(["0"], ["find doc"]) @@ -154,7 +157,12 @@ def test_agentic_retriever_honors_top_k(mock_react_step, mock_selection_step): {"doc_ids": ["doc_1"], "message": "doc_1 is best"}, ) - cfg = AgenticRetrievalConfig(llm_model="test-model", invoke_url="http://localhost/v1/chat/completions", top_k=5) + cfg = AgenticRetrievalConfig( + llm_model="test-model", + llm_backend="openai_compatible", + invoke_url="http://localhost/v1/chat/completions", + top_k=5, + ) result = AgenticRetriever(cfg, match_mode="pdf_page").retrieve(["0"], ["find doc"]) assert result["rank"].tolist() == list(range(1, 6)) # 5 rows, honoring top_k=5 @@ -173,7 +181,7 @@ def test_agentic_retriever_builds_in_process_llm_lazily(mock_create_local_agent_ ) mock_create_local_agent_llm.return_value = local_chat - cfg = AgenticRetrievalConfig(llm_model="nemotron-8b", llm_backend="in_process", top_k=1) + cfg = AgenticRetrievalConfig(top_k=1) retriever = AgenticRetriever(cfg, match_mode="pdf_page") mock_create_local_agent_llm.assert_not_called() @@ -189,7 +197,6 @@ def test_agentic_retriever_builds_in_process_llm_lazily(mock_create_local_agent_ tensor_parallel_size=1, max_model_len=None, max_num_seqs=None, - tool_call_parser=None, ) assert local_chat.call_count == 1 assert result["doc_id"].tolist() == ["doc_1"] @@ -212,6 +219,7 @@ def test_agentic_retriever_can_accept_partial_react_final_results(mock_react_ste cfg = AgenticRetrievalConfig( llm_model="test-model", + llm_backend="openai_compatible", invoke_url="http://localhost/v1/chat/completions", top_k=5, enforce_top_k=False, @@ -250,7 +258,9 @@ def test_run_agentic_audio_recall_evaluation_computes_metrics(mock_react_step, m {"doc_ids": [audio_doc_id], "message": "clip is best"}, ) - cfg = AgenticRetrievalConfig(llm_model="test-model", invoke_url="http://localhost/v1/chat/completions") + cfg = AgenticRetrievalConfig( + llm_model="test-model", llm_backend="openai_compatible", invoke_url="http://localhost/v1/chat/completions" + ) df_query, result, gold, retrieved, metrics = run_agentic_audio_recall_evaluation( query_csv=query_csv, cfg=cfg, @@ -287,7 +297,9 @@ def test_run_agentic_beir_evaluation_loads_queries_and_qrels(mock_react_step, mo queries=["find doc"], qrels={"q1": {"doc": 1}}, ) - cfg = AgenticRetrievalConfig(llm_model="test-model", invoke_url="http://localhost/v1/chat/completions") + cfg = AgenticRetrievalConfig( + llm_model="test-model", llm_backend="openai_compatible", invoke_url="http://localhost/v1/chat/completions" + ) with patch("nemo_retriever.query.agentic.load_beir_dataset", return_value=beir_dataset) as mock_loader: df_query, result, qrels, run, metrics = run_agentic_beir_evaluation( @@ -306,28 +318,47 @@ def test_run_agentic_beir_evaluation_loads_queries_and_qrels(mock_react_step, mo assert metrics["recall@1"] == 1.0 -def test_agentic_config_requires_llm_model(): +def test_agentic_config_defaults_empty_in_process_llm_model_to_nemotron_8b(): + from nemo_retriever.query.agentic import AgenticRetrievalConfig + + cfg = AgenticRetrievalConfig(llm_model="") + assert cfg.llm_backend == "in_process" + assert cfg.local_llm_backend == "vllm" + assert cfg.llm_model == "nemotron-8b" + + cfg = AgenticRetrievalConfig(llm_model=None) + assert cfg.llm_model == "nemotron-8b" + + +def test_agentic_config_requires_llm_model_for_openai_compatible(): from nemo_retriever.query.agentic import AgenticRetrievalConfig with pytest.raises(ValueError, match="llm_model"): - AgenticRetrievalConfig(llm_model="") + AgenticRetrievalConfig(llm_model="", llm_backend="openai_compatible") # None must not slip through as the literal string "None". with pytest.raises(ValueError, match="llm_model"): - AgenticRetrievalConfig(llm_model=None) + AgenticRetrievalConfig(llm_model=None, llm_backend="openai_compatible") + + +def test_agentic_config_rejects_custom_in_process_llm_model(): + from nemo_retriever.query.agentic import AgenticRetrievalConfig + + with pytest.raises(ValueError, match="Custom in-process agent LLMs are not supported yet"): + AgenticRetrievalConfig(llm_model="custom/local-model", llm_backend="in_process") def test_agentic_config_rejects_nonpositive_top_k(): from nemo_retriever.query.agentic import AgenticRetrievalConfig with pytest.raises(ValueError, match="top_k"): - AgenticRetrievalConfig(llm_model="m", top_k=0) + AgenticRetrievalConfig(llm_model="nemotron-8b", top_k=0) def test_agentic_config_rejects_noninteger_top_k(): from nemo_retriever.query.agentic import AgenticRetrievalConfig with pytest.raises(ValueError, match="top_k must be an integer"): - AgenticRetrievalConfig(llm_model="m", top_k=1.5) + AgenticRetrievalConfig(llm_model="nemotron-8b", top_k=1.5) def test_agentic_config_normalizes_integer_like_values(): @@ -335,6 +366,7 @@ def test_agentic_config_normalizes_integer_like_values(): cfg = AgenticRetrievalConfig( llm_model="m", + llm_backend="openai_compatible", invoke_url="http://localhost/v1/chat/completions", top_k="5.0", backend_top_k="6.0", @@ -350,27 +382,27 @@ def test_agentic_config_rejects_backend_top_k_below_target(): from nemo_retriever.query.agentic import AgenticRetrievalConfig with pytest.raises(ValueError, match="backend_top_k"): - AgenticRetrievalConfig(llm_model="m", backend_top_k=4, top_k=5) + AgenticRetrievalConfig(llm_model="nemotron-8b", backend_top_k=4, top_k=5) def test_agentic_config_rejects_nvidia_temperature_above_max(): from nemo_retriever.query.agentic import AgenticRetrievalConfig with pytest.raises(ValueError, match="between 0.0 and 1.0"): - AgenticRetrievalConfig(llm_model="m", temperature=1.5) + AgenticRetrievalConfig(llm_model="m", llm_backend="openai_compatible", temperature=1.5) def test_agentic_config_rejects_nonfinite_temperature(): from nemo_retriever.query.agentic import AgenticRetrievalConfig with pytest.raises(ValueError, match="temperature must be finite"): - AgenticRetrievalConfig(llm_model="m", temperature=float("nan")) + AgenticRetrievalConfig(llm_model="nemotron-8b", temperature=float("nan")) def test_agentic_config_accepts_in_process_temperature_above_nvidia_limit(): from nemo_retriever.query.agentic import AgenticRetrievalConfig - cfg = AgenticRetrievalConfig(llm_model="m", llm_backend="in_process", temperature=1.5) + cfg = AgenticRetrievalConfig(llm_model="nemotron-8b", llm_backend="in_process", temperature=1.5) assert cfg.llm_backend == "in_process" assert cfg.temperature == pytest.approx(1.5) @@ -380,14 +412,14 @@ def test_agentic_config_rejects_invalid_local_llm_backend(): from nemo_retriever.query.agentic import AgenticRetrievalConfig with pytest.raises(ValueError, match="local_llm_backend"): - AgenticRetrievalConfig(llm_model="m", llm_backend="in_process", local_llm_backend="hf") + AgenticRetrievalConfig(llm_model="nemotron-8b", llm_backend="in_process", local_llm_backend="hf") def test_agentic_config_validates_local_vllm_knobs(): from nemo_retriever.query.agentic import AgenticRetrievalConfig cfg = AgenticRetrievalConfig( - llm_model="m", + llm_model="nemotron-8b", llm_backend="in_process", local_gpu_memory_utilization="0.6", local_tensor_parallel_size="2.0", diff --git a/nemo_retriever/tests/test_agentic_local_llm.py b/nemo_retriever/tests/test_agentic_local_llm.py index 765ab23ea6..aaf154f72a 100644 --- a/nemo_retriever/tests/test_agentic_local_llm.py +++ b/nemo_retriever/tests/test_agentic_local_llm.py @@ -28,10 +28,7 @@ def test_vllm_agent_llm_rejects_unsupported_profile_before_vllm_import() -> None def test_parse_json_tool_call_output() -> None: from nemo_retriever.models.local.agent_llm import parse_tool_calls_from_text - calls = parse_tool_calls_from_text( - json.dumps([{"name": "retrieve", "arguments": {"query": "monetary policy"}}]), - parser="json", - ) + calls = parse_tool_calls_from_text(json.dumps([{"name": "retrieve", "arguments": {"query": "monetary policy"}}])) assert calls == [ { @@ -83,15 +80,6 @@ def test_parse_plain_text_returns_no_tool_calls() -> None: assert parse_tool_calls_from_text("I should search again") == [] -def test_parse_tool_calls_rejects_unsupported_parser() -> None: - import pytest - - from nemo_retriever.models.local.agent_llm import parse_tool_calls_from_text - - with pytest.raises(ValueError, match="supported parsers: auto, json"): - parse_tool_calls_from_text("[]", parser="mistral") - - def test_collapse_parallel_tool_results_for_local_chat_template() -> None: from nemo_retriever.models.local.agent_llm import _collapse_consecutive_tool_messages diff --git a/nemo_retriever/tests/test_harness_agentic_eval.py b/nemo_retriever/tests/test_harness_agentic_eval.py index 7a4981190d..21abfd4b37 100644 --- a/nemo_retriever/tests/test_harness_agentic_eval.py +++ b/nemo_retriever/tests/test_harness_agentic_eval.py @@ -41,7 +41,15 @@ def test_query_override_paths_include_agentic_fields() -> None: for key in ( "query.agentic", "query.agentic_llm_model", + "query.agentic_llm_backend", "query.agentic_invoke_url", + "query.agentic_local_llm_backend", + "query.agentic_local_hf_cache_dir", + "query.agentic_local_device", + "query.agentic_local_gpu_memory_utilization", + "query.agentic_local_tensor_parallel_size", + "query.agentic_local_max_model_len", + "query.agentic_local_max_num_seqs", "query.agentic_reasoning_effort", "query.agentic_backend_top_k", "query.agentic_react_max_steps", @@ -59,7 +67,15 @@ def test_build_query_request_populates_agentic() -> None: "top_k": 10, "agentic": True, "agentic_llm_model": "test-model", + "agentic_llm_backend": "openai_compatible", "agentic_invoke_url": "https://example.invalid/v1", + "agentic_local_llm_backend": "vllm", + "agentic_local_hf_cache_dir": "/tmp/hf", + "agentic_local_device": "0", + "agentic_local_gpu_memory_utilization": 0.6, + "agentic_local_tensor_parallel_size": 2, + "agentic_local_max_model_len": 8192, + "agentic_local_max_num_seqs": 4, "agentic_reasoning_effort": "high", "agentic_backend_top_k": 25, "agentic_react_max_steps": 12, @@ -73,7 +89,15 @@ def test_build_query_request_populates_agentic() -> None: agentic = request.agentic assert agentic.enabled is True assert agentic.llm_model == "test-model" + assert agentic.llm_backend == "openai_compatible" assert agentic.invoke_url == "https://example.invalid/v1" + assert agentic.local_llm_backend == "vllm" + assert agentic.local_hf_cache_dir == "/tmp/hf" + assert agentic.local_device == "0" + assert agentic.local_gpu_memory_utilization == pytest.approx(0.6) + assert agentic.local_tensor_parallel_size == 2 + assert agentic.local_max_model_len == 8192 + assert agentic.local_max_num_seqs == 4 assert agentic.reasoning_effort == "high" assert agentic.backend_top_k == 25 assert agentic.react_max_steps == 12 @@ -94,6 +118,7 @@ def test_build_agentic_config_maps_request_and_top_k_override() -> None: agentic=QueryAgenticOptions( enabled=True, llm_model="test-model", + llm_backend="openai_compatible", backend_top_k=20, num_concurrent=4, temperature=0.0, @@ -106,6 +131,16 @@ def test_build_agentic_config_maps_request_and_top_k_override() -> None: assert cfg.num_concurrent == 4 +def test_build_agentic_config_defaults_to_local_vllm_nemotron_8b() -> None: + request = QueryRequest(query="q", agentic=QueryAgenticOptions(enabled=True)) + + cfg = build_agentic_config(request, top_k=10) + + assert cfg.llm_backend == "in_process" + assert cfg.local_llm_backend == "vllm" + assert cfg.llm_model == "nemotron-8b" + + def test_run_beir_queries_routes_to_agentic(tmp_path) -> None: # Stage 1 (shared) load and stage 3 (shared) scoring/artifacts run for real; # only stage 2 (the agent retrieve+rank) is stubbed at the agentic_beir_retrieve @@ -123,7 +158,14 @@ def test_run_beir_queries_routes_to_agentic(tmp_path) -> None: "query": {}, } request = build_query_request( - _resolved({"top_k": 10, "agentic": True, "agentic_llm_model": "test-model"}), + _resolved( + { + "top_k": 10, + "agentic": True, + "agentic_llm_model": "test-model", + "agentic_llm_backend": "openai_compatible", + } + ), "", ) @@ -174,7 +216,7 @@ def test_run_beir_queries_invalid_agentic_config_is_structured_failure(tmp_path) "query": {}, } request = build_query_request( - _resolved({"top_k": 10, "agentic": True, "agentic_llm_model": "m", "agentic_backend_top_k": 5}), + _resolved({"top_k": 10, "agentic": True, "agentic_llm_model": "nemotron-8b", "agentic_backend_top_k": 5}), "", ) dataset = BeirDataset(dataset_name="demo", query_ids=["q1"], queries=["t1"], qrels={"q1": {"d1": 1}}) diff --git a/nemo_retriever/tests/test_root_query_cli.py b/nemo_retriever/tests/test_root_query_cli.py index e680ab0d55..053da62cd8 100644 --- a/nemo_retriever/tests/test_root_query_cli.py +++ b/nemo_retriever/tests/test_root_query_cli.py @@ -377,8 +377,6 @@ def retrieve(self, query_ids: Any, query_texts: Any) -> Any: "query", "how does ingest work?", "--agentic", - "--agentic-llm-model", - "nvidia/llama-3.3-nemotron-super-49b-v1.5", "--top-k", "2", "--lancedb-uri", @@ -393,7 +391,9 @@ def retrieve(self, query_ids: Any, query_texts: Any) -> Any: cfg = config_calls[0] assert cfg["vdb_op"] == "lancedb" assert cfg["vdb_kwargs"] == {"uri": "/tmp/lancedb", "table_name": "docs"} - assert cfg["llm_model"] == "nvidia/llama-3.3-nemotron-super-49b-v1.5" + assert cfg["llm_backend"] == "in_process" + assert cfg["local_llm_backend"] == "vllm" + assert cfg["llm_model"] == "nemotron-8b" # --top-k is honored end-to-end: plumbed into the agentic config (drives the # ReAct target / RRF / selection cut), not just applied as a post-filter. assert cfg["top_k"] == 2 @@ -404,12 +404,67 @@ 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.""" - result = RUNNER.invoke(cli_main.app, ["query", "hello", "--agentic"]) +def test_root_query_agentic_rejects_custom_in_process_llm_model() -> None: + result = RUNNER.invoke( + cli_main.app, + ["query", "hello", "--agentic", "--agentic-llm-model", "custom/local-model"], + ) + + assert result.exit_code == 1 + assert "Custom in-process agent LLMs are not supported yet" in result.output + assert "--agentic-llm-backend openai_compatible" in result.output or "openai_compatible" in result.output + + +def test_root_query_agentic_openai_compatible_requires_llm_model() -> None: + result = RUNNER.invoke( + cli_main.app, + ["query", "hello", "--agentic", "--agentic-llm-backend", "openai_compatible"], + ) assert result.exit_code == 1 - assert "requires --agentic-llm-model" in result.output + assert "openai_compatible requires --agentic-llm-model" in result.output + + +def test_root_query_agentic_openai_compatible_allows_custom_model(monkeypatch) -> None: + import pandas as pd + + import nemo_retriever.query.agentic as agentic_retrieval + + 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: + self.cfg = cfg + + 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", + "q", + "--agentic", + "--agentic-llm-backend", + "openai_compatible", + "--agentic-llm-model", + "custom-remote-model", + "--agentic-invoke-url", + "http://localhost:8000/v1/chat/completions", + ], + ) + + assert result.exit_code == 0 + assert config_calls[-1]["llm_backend"] == "openai_compatible" + assert config_calls[-1]["llm_model"] == "custom-remote-model" + assert config_calls[-1]["invoke_url"] == "http://localhost:8000/v1/chat/completions" def test_root_query_agentic_plumbs_rerank_into_config(monkeypatch) -> None: @@ -435,7 +490,7 @@ def retrieve(self, query_ids: Any, query_texts: Any) -> Any: monkeypatch.setattr(agentic_retrieval, "AgenticRetrievalConfig", FakeConfig) monkeypatch.setattr(agentic_retrieval, "AgenticRetriever", FakeAgenticRetriever) - base = ["query", "q", "--agentic", "--agentic-llm-model", "m"] + base = ["query", "q", "--agentic"] # 1. Explicit reranker model + endpoint + backend flow through. result = RUNNER.invoke( From 89efb18e517875a19e1edd2d40c15094a999ee95 Mon Sep 17 00:00:00 2001 From: Mahika Wason Date: Mon, 20 Jul 2026 14:12:43 -0700 Subject: [PATCH 03/10] Reduce public agentic local LLM CLI surface --- nemo_retriever/docs/cli/README.md | 16 --------- nemo_retriever/docs/cli/benchmarking.md | 11 +++---- .../src/nemo_retriever/cli/query/app.py | 14 -------- .../src/nemo_retriever/cli/query/options.py | 33 ------------------- .../src/nemo_retriever/harness/resolution.py | 6 ---- .../tests/test_harness_agentic_eval.py | 8 ----- 6 files changed, 5 insertions(+), 83 deletions(-) diff --git a/nemo_retriever/docs/cli/README.md b/nemo_retriever/docs/cli/README.md index e5ad446ac1..7fc384dabf 100644 --- a/nemo_retriever/docs/cli/README.md +++ b/nemo_retriever/docs/cli/README.md @@ -190,12 +190,6 @@ when you want a custom model or a separately hosted OpenAI-compatible endpoint. retriever query "how does the ingestion pipeline handle tables?" \ --agentic -# larger supported local profile -retriever query "how does the ingestion pipeline handle tables?" \ - --agentic \ - --agentic-llm-model super-49b \ - --agentic-local-tensor-parallel-size 2 - # custom/self-hosted model through an OpenAI-compatible endpoint retriever query "summarize the deployment options" \ --agentic \ @@ -232,16 +226,6 @@ Agentic-only knobs (apply only with `--agentic`): when using `openai_compatible`. - `--agentic-invoke-url` — OpenAI-compatible chat-completions endpoint for the agent LLM; used only with `--agentic-llm-backend openai_compatible`. -- `--agentic-local-llm-backend` (default `vllm`) — local runtime for in-process - agent LLMs. Other local runtimes are not supported yet. -- `--agentic-local-device` — CUDA device ids for local vLLM, for example `0`, - `cuda:0`, or `0,1`. -- `--agentic-local-gpu-memory-utilization` (default `0.8`) — vLLM GPU memory - utilization for the local agent LLM. -- `--agentic-local-tensor-parallel-size` (default `1`) — tensor parallel size for - local vLLM. -- `--agentic-local-max-model-len` and `--agentic-local-max-num-seqs` — optional - vLLM engine overrides for the local agent LLM. - `--agentic-reasoning-effort` (default `high`) — `reasoning_effort` forwarded on OpenAI-compatible agentic LLM calls; ignored by the local adapter. - `--agentic-backend-top-k` (default `20`) — candidates pulled from the vector DB diff --git a/nemo_retriever/docs/cli/benchmarking.md b/nemo_retriever/docs/cli/benchmarking.md index 399ffa914d..be7177224d 100644 --- a/nemo_retriever/docs/cli/benchmarking.md +++ b/nemo_retriever/docs/cli/benchmarking.md @@ -98,13 +98,12 @@ Useful agentic query overrides: when using `openai_compatible`. - `query.agentic_invoke_url` — OpenAI-compatible chat-completions endpoint; used only with `query.agentic_llm_backend=openai_compatible`. -- `query.agentic_local_llm_backend` — local runtime for in-process agent LLMs; - currently only `vllm` is supported. -- `query.agentic_local_hf_cache_dir`, `query.agentic_local_device`, - `query.agentic_local_gpu_memory_utilization`, +- `query.agentic_local_gpu_memory_utilization`, `query.agentic_local_tensor_parallel_size`, `query.agentic_local_max_model_len`, - and `query.agentic_local_max_num_seqs` — local vLLM loading and scheduling - controls. + and `query.agentic_local_max_num_seqs` — harness-only local vLLM resource and + scheduling controls for benchmark runs. Use environment variables such as + `CUDA_VISIBLE_DEVICES` and the standard Hugging Face cache environment for + placement and model cache control. - `query.agentic_backend_top_k` — backend candidate pool per ReAct retrieval call. Must be at least the final requested metric depth (`max(evaluation.ks)`). - `query.agentic_react_max_steps` — maximum ReAct loop iterations per query diff --git a/nemo_retriever/src/nemo_retriever/cli/query/app.py b/nemo_retriever/src/nemo_retriever/cli/query/app.py index 9995e700aa..7733b51d5d 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/app.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/app.py @@ -189,13 +189,6 @@ def _local_command( agentic_llm_model: opts.AgenticLlmModelOption = None, agentic_llm_backend: opts.AgenticLlmBackendOption = "in_process", agentic_invoke_url: opts.AgenticInvokeUrlOption = None, - agentic_local_llm_backend: opts.AgenticLocalLlmBackendOption = "vllm", - agentic_local_hf_cache_dir: opts.AgenticLocalHfCacheDirOption = None, - agentic_local_device: opts.AgenticLocalDeviceOption = None, - agentic_local_gpu_memory_utilization: opts.AgenticLocalGpuMemoryUtilizationOption = 0.8, - agentic_local_tensor_parallel_size: opts.AgenticLocalTensorParallelSizeOption = 1, - agentic_local_max_model_len: opts.AgenticLocalMaxModelLenOption = None, - agentic_local_max_num_seqs: opts.AgenticLocalMaxNumSeqsOption = None, agentic_reasoning_effort: opts.AgenticReasoningEffortOption = "high", agentic_backend_top_k: opts.AgenticBackendTopKOption = 20, agentic_react_max_steps: opts.AgenticReactMaxStepsOption = 50, @@ -273,13 +266,6 @@ def _local_command( llm_model=agentic_llm_model, llm_backend=normalized_agentic_llm_backend, invoke_url=agentic_invoke_url, - local_llm_backend=agentic_local_llm_backend, - local_hf_cache_dir=agentic_local_hf_cache_dir, - local_device=agentic_local_device, - local_gpu_memory_utilization=agentic_local_gpu_memory_utilization, - local_tensor_parallel_size=agentic_local_tensor_parallel_size, - local_max_model_len=agentic_local_max_model_len, - local_max_num_seqs=agentic_local_max_num_seqs, reasoning_effort=agentic_reasoning_effort, backend_top_k=agentic_backend_top_k, react_max_steps=agentic_react_max_steps, diff --git a/nemo_retriever/src/nemo_retriever/cli/query/options.py b/nemo_retriever/src/nemo_retriever/cli/query/options.py index fb8b93baf2..eaa5bbf5d0 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/options.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/options.py @@ -179,39 +179,6 @@ help="OpenAI-compatible chat-completions endpoint for the agent LLM (agentic mode).", ), ] -AgenticLocalLlmBackendOption = Annotated[ - str, - typer.Option("--agentic-local-llm-backend", help="Local in-process agent LLM runtime. Currently only vllm."), -] -AgenticLocalHfCacheDirOption = Annotated[ - str | None, - typer.Option("--agentic-local-hf-cache-dir", help="Hugging Face cache directory for local agent LLM weights."), -] -AgenticLocalDeviceOption = Annotated[ - str | None, - typer.Option("--agentic-local-device", help="CUDA device ids for local vLLM, e.g. 0, cuda:0, or 0,1."), -] -AgenticLocalGpuMemoryUtilizationOption = Annotated[ - float, - typer.Option( - "--agentic-local-gpu-memory-utilization", - min=0.0, - max=1.0, - help="Fraction of GPU memory reserved by local vLLM for the agent LLM.", - ), -] -AgenticLocalTensorParallelSizeOption = Annotated[ - int, - typer.Option("--agentic-local-tensor-parallel-size", min=1, help="Tensor parallel size for local vLLM."), -] -AgenticLocalMaxModelLenOption = Annotated[ - int | None, - typer.Option("--agentic-local-max-model-len", min=1, help="Optional local vLLM max_model_len override."), -] -AgenticLocalMaxNumSeqsOption = Annotated[ - int | None, - typer.Option("--agentic-local-max-num-seqs", min=1, help="Optional local vLLM max_num_seqs override."), -] AgenticReasoningEffortOption = Annotated[ str | None, typer.Option( diff --git a/nemo_retriever/src/nemo_retriever/harness/resolution.py b/nemo_retriever/src/nemo_retriever/harness/resolution.py index d9047d4fc9..ee2f9f8405 100644 --- a/nemo_retriever/src/nemo_retriever/harness/resolution.py +++ b/nemo_retriever/src/nemo_retriever/harness/resolution.py @@ -130,9 +130,6 @@ def _override_child_keys(prefix: str, paths: set[str]) -> set[str]: "query.agentic_llm_model", "query.agentic_llm_backend", "query.agentic_invoke_url", - "query.agentic_local_llm_backend", - "query.agentic_local_hf_cache_dir", - "query.agentic_local_device", "query.agentic_local_gpu_memory_utilization", "query.agentic_local_tensor_parallel_size", "query.agentic_local_max_model_len", @@ -447,9 +444,6 @@ def build_query_request(resolved: dict[str, Any], query_text: str) -> QueryReque llm_model=query.get("agentic_llm_model"), llm_backend=str(query.get("agentic_llm_backend") or "in_process"), invoke_url=query.get("agentic_invoke_url"), - local_llm_backend=str(query.get("agentic_local_llm_backend") or "vllm"), - local_hf_cache_dir=query.get("agentic_local_hf_cache_dir"), - local_device=query.get("agentic_local_device"), local_gpu_memory_utilization=float(query.get("agentic_local_gpu_memory_utilization") or 0.8), local_tensor_parallel_size=int(query.get("agentic_local_tensor_parallel_size") or 1), local_max_model_len=query.get("agentic_local_max_model_len"), diff --git a/nemo_retriever/tests/test_harness_agentic_eval.py b/nemo_retriever/tests/test_harness_agentic_eval.py index 21abfd4b37..c93d357ac7 100644 --- a/nemo_retriever/tests/test_harness_agentic_eval.py +++ b/nemo_retriever/tests/test_harness_agentic_eval.py @@ -43,9 +43,6 @@ def test_query_override_paths_include_agentic_fields() -> None: "query.agentic_llm_model", "query.agentic_llm_backend", "query.agentic_invoke_url", - "query.agentic_local_llm_backend", - "query.agentic_local_hf_cache_dir", - "query.agentic_local_device", "query.agentic_local_gpu_memory_utilization", "query.agentic_local_tensor_parallel_size", "query.agentic_local_max_model_len", @@ -69,9 +66,6 @@ def test_build_query_request_populates_agentic() -> None: "agentic_llm_model": "test-model", "agentic_llm_backend": "openai_compatible", "agentic_invoke_url": "https://example.invalid/v1", - "agentic_local_llm_backend": "vllm", - "agentic_local_hf_cache_dir": "/tmp/hf", - "agentic_local_device": "0", "agentic_local_gpu_memory_utilization": 0.6, "agentic_local_tensor_parallel_size": 2, "agentic_local_max_model_len": 8192, @@ -92,8 +86,6 @@ def test_build_query_request_populates_agentic() -> None: assert agentic.llm_backend == "openai_compatible" assert agentic.invoke_url == "https://example.invalid/v1" assert agentic.local_llm_backend == "vllm" - assert agentic.local_hf_cache_dir == "/tmp/hf" - assert agentic.local_device == "0" assert agentic.local_gpu_memory_utilization == pytest.approx(0.6) assert agentic.local_tensor_parallel_size == 2 assert agentic.local_max_model_len == 8192 From 873798f14222aa00be96a7735cb445617e2ec3e2 Mon Sep 17 00:00:00 2001 From: Mahika Wason Date: Mon, 20 Jul 2026 15:08:44 -0700 Subject: [PATCH 04/10] Trim agentic local LLM adapter surface --- nemo_retriever/docs/cli/README.md | 2 - .../src/nemo_retriever/cli/query/app.py | 14 +----- .../src/nemo_retriever/cli/query/options.py | 8 ---- .../nemo_retriever/models/local/agent_llm.py | 47 ++++++++----------- .../tests/test_agentic_local_llm.py | 19 ++++++++ 5 files changed, 40 insertions(+), 50 deletions(-) diff --git a/nemo_retriever/docs/cli/README.md b/nemo_retriever/docs/cli/README.md index 7fc384dabf..6b50e4d6c3 100644 --- a/nemo_retriever/docs/cli/README.md +++ b/nemo_retriever/docs/cli/README.md @@ -233,8 +233,6 @@ Agentic-only knobs (apply only with `--agentic`): - `--agentic-react-max-steps` (default `50`) — maximum ReAct loop iterations. - `--agentic-text-truncation` (default `0`) — max characters of each candidate shown to the agent; `0` disables truncation. -- `--agentic-temperature` (default `0.0`) — sampling temperature for agentic LLM - calls (`0.0` = greedy). diff --git a/nemo_retriever/src/nemo_retriever/cli/query/app.py b/nemo_retriever/src/nemo_retriever/cli/query/app.py index 7733b51d5d..c983f69ca0 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/app.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/app.py @@ -16,10 +16,7 @@ from nemo_retriever.cli.query import options as opts from nemo_retriever.cli.query_workflow import agentic_query_documents as query_agentic_documents from nemo_retriever.cli.query_workflow import query_documents_with_metadata as query_local_documents_with_metadata -from nemo_retriever.query.agentic_options import ( - agentic_backend_top_k_error, - agentic_temperature_error, -) +from nemo_retriever.query.agentic_options import agentic_backend_top_k_error from nemo_retriever.cli.shared import ( ROOT_CLI_ERRORS, quiet_capture, @@ -193,7 +190,6 @@ def _local_command( agentic_backend_top_k: opts.AgenticBackendTopKOption = 20, agentic_react_max_steps: opts.AgenticReactMaxStepsOption = 50, agentic_text_truncation: opts.AgenticTextTruncationOption = 0, - agentic_temperature: opts.AgenticTemperatureOption = 0.0, ) -> None: _validate_output_options(output_format, max_text_chars) if reranker_invoke_url is None: @@ -223,13 +219,6 @@ def _local_command( if backend_error: typer.echo(f"Error: {backend_error}", err=True) raise typer.Exit(1) - temperature_invoke_url = ( - agentic_invoke_url if normalized_agentic_llm_backend == "openai_compatible" else "local://in-process" - ) - temperature_error = agentic_temperature_error(agentic_temperature, invoke_url=temperature_invoke_url) - if temperature_error: - typer.echo(f"Error: {temperature_error}", err=True) - raise typer.Exit(1) try: reranker_api_key = _api_key_from_env_option(reranker_api_key_env) if reranker_invoke_url else None @@ -270,7 +259,6 @@ def _local_command( backend_top_k=agentic_backend_top_k, react_max_steps=agentic_react_max_steps, text_truncation=agentic_text_truncation, - temperature=agentic_temperature, ), ) with quiet_capture(): diff --git a/nemo_retriever/src/nemo_retriever/cli/query/options.py b/nemo_retriever/src/nemo_retriever/cli/query/options.py index eaa5bbf5d0..cdb39f5cc8 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/options.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/options.py @@ -210,14 +210,6 @@ help="Max characters of each candidate shown to the agent; 0 disables truncation.", ), ] -AgenticTemperatureOption = Annotated[ - float, - typer.Option( - "--agentic-temperature", - min=0.0, - help="Sampling temperature for agentic LLM calls (0.0 = greedy).", - ), -] ServiceUrlOption = Annotated[ str, typer.Option("--service-url", help="Base URL of the retriever service."), diff --git a/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py index c61161cc53..d50d9e96ba 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py +++ b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py @@ -7,9 +7,8 @@ import json import logging import os -import random -import string import threading +import uuid from copy import deepcopy from typing import Any, Mapping, Optional, Sequence @@ -60,24 +59,12 @@ def _normalize_name(name: str) -> str: _NEMOTRON_JSON_TOOL_PROMPT_EXTRAS = {"chat_template_kwargs": {"tools_in_user_message": True}} -def _mutable_copy(value: Any) -> Any: - if isinstance(value, Mapping): - return {key: _mutable_copy(item) for key, item in value.items()} - if isinstance(value, tuple): - return [_mutable_copy(item) for item in value] - return deepcopy(value) - - def supported_agent_llm_names() -> tuple[str, ...]: """Return supported local agent LLM model IDs and aliases.""" return _SUPPORTED_AGENT_LLM_NAMES -def _supported_agent_llm_names_for_error() -> str: - return ", ".join(supported_agent_llm_names()) - - def is_supported_agent_llm_model(name: str) -> bool: """Return whether *name* is a supported in-process local agent LLM.""" @@ -113,16 +100,16 @@ def __init__( ) -> None: super().__init__() - resolved_model_path = resolve_agent_llm_model_name(model_path) - if resolved_model_path == model_path and not is_supported_agent_llm_model(model_path): + requested_model_path = model_path + model_path = resolve_agent_llm_model_name(requested_model_path) + if not is_supported_agent_llm_model(model_path): raise ValueError( - f"Unsupported local agent LLM model {model_path!r}. " + f"Unsupported local agent LLM model {requested_model_path!r}. " "Custom in-process agent LLMs are not supported yet. " "Use llm_backend='openai_compatible' with invoke_url for a custom/self-hosted endpoint, " - f"or choose one of: {_supported_agent_llm_names_for_error()}." + f"or choose one of: {', '.join(supported_agent_llm_names())}." ) - model_path = resolved_model_path cuda_visible_devices = _cuda_visible_devices_from_device(device) if cuda_visible_devices is not None: os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices @@ -137,7 +124,7 @@ def __init__( _raise_if_cuda_unavailable() self._model_path = model_path self._max_tokens = int(max_tokens) - self._request_extras = _mutable_copy(_NEMOTRON_JSON_TOOL_PROMPT_EXTRAS) + self._request_extras = deepcopy(_NEMOTRON_JSON_TOOL_PROMPT_EXTRAS) self._lock = threading.Lock() configure_global_hf_cache_base(hf_cache_dir) @@ -221,7 +208,7 @@ def __call__( } def _build_chat_kwargs(self, extra_body: Optional[dict[str, Any]]) -> dict[str, Any]: - chat_kwargs = _mutable_copy(self._request_extras) + chat_kwargs = deepcopy(self._request_extras) for key, value in (extra_body or {}).items(): if key in {"parallel_tool_calls", "reasoning_effort"}: continue @@ -431,11 +418,18 @@ def _coerce_single_tool_call(item: Any) -> dict[str, Any] | None: if isinstance(item.get("function"), Mapping): function = dict(item["function"]) + if "arguments" not in function: + return None name = function.get("name") - arguments = function.get("arguments", {}) + arguments = function.get("arguments") else: name = item.get("name") or item.get("tool_name") - arguments = item.get("arguments", item.get("parameters", {})) + if "arguments" in item: + arguments = item.get("arguments") + elif "parameters" in item: + arguments = item.get("parameters") + else: + return None if not name: return None @@ -457,7 +451,7 @@ def _arguments_to_json_string(arguments: Any) -> str: except json.JSONDecodeError: return json.dumps(arguments) return arguments - return json.dumps(arguments or {}) + return json.dumps(arguments if arguments is not None else {}) def _tool_calls_from_completion(completion: Any) -> list[dict[str, Any]]: @@ -553,9 +547,8 @@ def _json_tool_prompt(tools: Sequence[Mapping[str, Any]]) -> str: def _new_tool_call_id() -> str: - alphabet = string.ascii_letters + string.digits - return "".join(random.SystemRandom().choice(alphabet) for _ in range(9)) + return f"call_{uuid.uuid4().hex[:12]}" def _new_response_id() -> str: - return "chatcmpl_" + _new_tool_call_id() + return f"chatcmpl_{uuid.uuid4().hex[:12]}" diff --git a/nemo_retriever/tests/test_agentic_local_llm.py b/nemo_retriever/tests/test_agentic_local_llm.py index aaf154f72a..99325e8153 100644 --- a/nemo_retriever/tests/test_agentic_local_llm.py +++ b/nemo_retriever/tests/test_agentic_local_llm.py @@ -74,6 +74,25 @@ def test_parse_tool_call_output_from_code_fence() -> None: assert json.loads(calls[0]["function"]["arguments"]) == {"thought": "compare docs"} +def test_parse_tool_call_output_ignores_echoed_tool_schema() -> None: + from nemo_retriever.models.local.agent_llm import parse_tool_calls_from_text + + echoed_schema = json.dumps( + [ + { + "type": "function", + "function": { + "name": "retrieve", + "description": "Retrieve documents.", + "parameters": {"type": "object"}, + }, + } + ] + ) + + assert parse_tool_calls_from_text(echoed_schema) == [] + + def test_parse_plain_text_returns_no_tool_calls() -> None: from nemo_retriever.models.local.agent_llm import parse_tool_calls_from_text From a29070f29005644ed733d24d2bfc0b2e96dcfc81 Mon Sep 17 00:00:00 2001 From: Mahika Wason Date: Mon, 20 Jul 2026 21:50:17 -0700 Subject: [PATCH 05/10] Infer agentic LLM backend and tighten local tool-call validation --- nemo_retriever/docs/cli/README.md | 15 +++--- nemo_retriever/docs/cli/benchmarking.md | 16 +++--- .../src/nemo_retriever/cli/query/app.py | 15 ++---- .../src/nemo_retriever/cli/query/options.py | 11 +--- .../src/nemo_retriever/harness/resolution.py | 2 - .../nemo_retriever/models/local/agent_llm.py | 6 +-- .../graph_ops/react_agent_operator.py | 29 +++-------- .../graph_ops/selection_agent_operator.py | 43 +++++++++++---- .../src/nemo_retriever/query/agentic.py | 51 +++++++----------- .../src/nemo_retriever/query/options.py | 2 +- .../src/nemo_retriever/query/workflow.py | 3 +- nemo_retriever/tests/test_agentic_eval.py | 41 +++++++-------- .../tests/test_agentic_operators.py | 52 ++++++++++++++----- .../tests/test_harness_agentic_eval.py | 9 ++-- nemo_retriever/tests/test_root_query_cli.py | 14 +++-- 15 files changed, 150 insertions(+), 159 deletions(-) diff --git a/nemo_retriever/docs/cli/README.md b/nemo_retriever/docs/cli/README.md index 6b50e4d6c3..c066341029 100644 --- a/nemo_retriever/docs/cli/README.md +++ b/nemo_retriever/docs/cli/README.md @@ -182,8 +182,8 @@ a drop-in alternative to standard retrieval. By default, agentic retrieval runs the agent LLM in process with local vLLM and `nemotron-8b` (`nvidia/Llama-3.1-Nemotron-Nano-8B-v1`). This requires a CUDA GPU -host and the local extras installed. Use `--agentic-llm-backend openai_compatible` -when you want a custom model or a separately hosted OpenAI-compatible endpoint. +host and the local extras installed. Provide `--agentic-invoke-url` when you want +a custom model or a separately hosted OpenAI-compatible endpoint. ```bash # default local vLLM agent LLM: nemotron-8b @@ -193,7 +193,6 @@ retriever query "how does the ingestion pipeline handle tables?" \ # custom/self-hosted model through an OpenAI-compatible endpoint retriever query "summarize the deployment options" \ --agentic \ - --agentic-llm-backend openai_compatible \ --agentic-llm-model custom-remote-model \ --agentic-invoke-url http://localhost:9000/v1/chat/completions \ --embed-invoke-url http://localhost:8000/v1 \ @@ -219,13 +218,11 @@ fusion) -> SelectionAgentOperator -> ranked results`: Agentic-only knobs (apply only with `--agentic`): -- `--agentic-llm-backend` (default `in_process`) — use local in-process vLLM, or - set `openai_compatible` for a hosted/self-hosted chat-completions endpoint. -- `--agentic-llm-model` — local profile alias/model ID for `in_process` - (`nemotron-8b` by default; `super-49b` also supported), or the remote model ID - when using `openai_compatible`. +- `--agentic-llm-model` — local profile alias/model ID when no invoke URL is + provided (`nemotron-8b` by default; `super-49b` also supported), or the remote + model ID when `--agentic-invoke-url` is provided. - `--agentic-invoke-url` — OpenAI-compatible chat-completions endpoint for the - agent LLM; used only with `--agentic-llm-backend openai_compatible`. + agent LLM. Providing it routes agent LLM calls to that remote endpoint. - `--agentic-reasoning-effort` (default `high`) — `reasoning_effort` forwarded on OpenAI-compatible agentic LLM calls; ignored by the local adapter. - `--agentic-backend-top-k` (default `20`) — candidates pulled from the vector DB diff --git a/nemo_retriever/docs/cli/benchmarking.md b/nemo_retriever/docs/cli/benchmarking.md index be7177224d..b24c120484 100644 --- a/nemo_retriever/docs/cli/benchmarking.md +++ b/nemo_retriever/docs/cli/benchmarking.md @@ -60,8 +60,7 @@ but only after ingest and only for BEIR evaluation (`evaluation.mode: beir`). By default, agentic harness evaluation uses the in-process local vLLM backend with `nemotron-8b`. Custom LLMs are not supported in process yet; run them behind -an OpenAI-compatible chat-completions endpoint and set -`query.agentic_llm_backend=openai_compatible`. +an OpenAI-compatible chat-completions endpoint and set `query.agentic_invoke_url`. Minimal BEIR override example: @@ -84,20 +83,17 @@ Custom/self-hosted OpenAI-compatible endpoint: ```bash retriever harness run jp20_beir \ --set query.agentic=true \ - --set query.agentic_llm_backend=openai_compatible \ --set query.agentic_llm_model=custom-remote-model \ --set query.agentic_invoke_url=http://localhost:9000/v1/chat/completions ``` Useful agentic query overrides: -- `query.agentic_llm_backend` — `in_process` by default, or - `openai_compatible` for hosted/self-hosted chat-completions endpoints. -- `query.agentic_llm_model` — local profile alias/model ID for `in_process` - (`nemotron-8b` by default; `super-49b` also supported), or the remote model ID - when using `openai_compatible`. -- `query.agentic_invoke_url` — OpenAI-compatible chat-completions endpoint; used - only with `query.agentic_llm_backend=openai_compatible`. +- `query.agentic_llm_model` — local profile alias/model ID when no invoke URL is + provided (`nemotron-8b` by default; `super-49b` also supported), or the remote + model ID when `query.agentic_invoke_url` is provided. +- `query.agentic_invoke_url` — OpenAI-compatible chat-completions endpoint. + Providing it routes agent LLM calls to that remote endpoint. - `query.agentic_local_gpu_memory_utilization`, `query.agentic_local_tensor_parallel_size`, `query.agentic_local_max_model_len`, and `query.agentic_local_max_num_seqs` — harness-only local vLLM resource and diff --git a/nemo_retriever/src/nemo_retriever/cli/query/app.py b/nemo_retriever/src/nemo_retriever/cli/query/app.py index c983f69ca0..c2b0292ed4 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/app.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/app.py @@ -184,7 +184,6 @@ def _local_command( max_text_chars: opts.MaxTextCharsOption = None, agentic: opts.AgenticOption = False, agentic_llm_model: opts.AgenticLlmModelOption = None, - agentic_llm_backend: opts.AgenticLlmBackendOption = "in_process", agentic_invoke_url: opts.AgenticInvokeUrlOption = None, agentic_reasoning_effort: opts.AgenticReasoningEffortOption = "high", agentic_backend_top_k: opts.AgenticBackendTopKOption = 20, @@ -199,21 +198,14 @@ def _local_command( rerank = rerank or bool(reranker_invoke_url) or bool(reranker_model_name) or bool(reranker_backend) silence_noisy_libraries() if agentic: - normalized_agentic_llm_backend = str(agentic_llm_backend or "in_process").strip().lower() - if normalized_agentic_llm_backend not in {"in_process", "openai_compatible"}: + if agentic_invoke_url and not agentic_llm_model: typer.echo( - "Error: --agentic-llm-backend must be one of ['in_process', 'openai_compatible'].", + "Error: --agentic-invoke-url requires --agentic-llm-model.", err=True, ) raise typer.Exit(1) - if normalized_agentic_llm_backend == "in_process" and not agentic_llm_model: + if not agentic_invoke_url and not agentic_llm_model: agentic_llm_model = "nemotron-8b" - if normalized_agentic_llm_backend == "openai_compatible" and not agentic_llm_model: - typer.echo( - "Error: --agentic-llm-backend openai_compatible requires --agentic-llm-model.", - err=True, - ) - raise typer.Exit(1) backend_error = agentic_backend_top_k_error(agentic_backend_top_k, target_top_k=top_k) if backend_error: @@ -253,7 +245,6 @@ def _local_command( agentic=QueryAgenticOptions( enabled=agentic, llm_model=agentic_llm_model, - llm_backend=normalized_agentic_llm_backend, invoke_url=agentic_invoke_url, reasoning_effort=agentic_reasoning_effort, backend_top_k=agentic_backend_top_k, diff --git a/nemo_retriever/src/nemo_retriever/cli/query/options.py b/nemo_retriever/src/nemo_retriever/cli/query/options.py index cdb39f5cc8..dae6b72924 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/options.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/options.py @@ -160,18 +160,11 @@ typer.Option( "--agentic-llm-model", help=( - "Chat model the agent drives. Defaults to nemotron-8b for --agentic-llm-backend=in_process; " - "required for openai_compatible endpoints." + "Chat model the agent drives. Defaults to nemotron-8b for local in-process runs; " + "required when --agentic-invoke-url is provided." ), ), ] -AgenticLlmBackendOption = Annotated[ - str, - typer.Option( - "--agentic-llm-backend", - help="Agent LLM backend: in_process (default) or openai_compatible.", - ), -] AgenticInvokeUrlOption = Annotated[ str | None, typer.Option( diff --git a/nemo_retriever/src/nemo_retriever/harness/resolution.py b/nemo_retriever/src/nemo_retriever/harness/resolution.py index ee2f9f8405..006530e42a 100644 --- a/nemo_retriever/src/nemo_retriever/harness/resolution.py +++ b/nemo_retriever/src/nemo_retriever/harness/resolution.py @@ -128,7 +128,6 @@ def _override_child_keys(prefix: str, paths: set[str]) -> set[str]: "query.table_name", "query.agentic", "query.agentic_llm_model", - "query.agentic_llm_backend", "query.agentic_invoke_url", "query.agentic_local_gpu_memory_utilization", "query.agentic_local_tensor_parallel_size", @@ -442,7 +441,6 @@ def build_query_request(resolved: dict[str, Any], query_text: str) -> QueryReque agentic=QueryAgenticOptions( enabled=bool(query.get("agentic", False)), llm_model=query.get("agentic_llm_model"), - llm_backend=str(query.get("agentic_llm_backend") or "in_process"), invoke_url=query.get("agentic_invoke_url"), local_gpu_memory_utilization=float(query.get("agentic_local_gpu_memory_utilization") or 0.8), local_tensor_parallel_size=int(query.get("agentic_local_tensor_parallel_size") or 1), diff --git a/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py index d50d9e96ba..6e24638a9c 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py +++ b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py @@ -106,7 +106,7 @@ def __init__( raise ValueError( f"Unsupported local agent LLM model {requested_model_path!r}. " "Custom in-process agent LLMs are not supported yet. " - "Use llm_backend='openai_compatible' with invoke_url for a custom/self-hosted endpoint, " + "Provide invoke_url for a custom/self-hosted OpenAI-compatible endpoint, " f"or choose one of: {', '.join(supported_agent_llm_names())}." ) @@ -293,7 +293,7 @@ def _cuda_visible_devices_from_device(device: Optional[str]) -> str | None: if normalized.casefold() == "cpu": raise ValueError( "The local agent LLM vLLM backend requires CUDA. Pass GPU ids such as '0' or '0,1', " - "or use llm_backend='openai_compatible' for a remote endpoint." + "or provide invoke_url for a remote OpenAI-compatible endpoint." ) return normalized.split(":", 1)[1] if normalized.startswith("cuda:") else normalized @@ -306,7 +306,7 @@ def _raise_if_cuda_unavailable() -> None: if not torch.cuda.is_available(): raise RuntimeError( "The local agent LLM vLLM backend requires CUDA, but torch reports no available CUDA device. " - "Run on a GPU host or use llm_backend='openai_compatible' for a remote endpoint." + "Run on a GPU host or provide invoke_url for a remote OpenAI-compatible endpoint." ) diff --git a/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py b/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py index 8d9419c3d5..ac56bba677 100644 --- a/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py +++ b/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py @@ -92,7 +92,7 @@ def _preview_doc_ids(docs: List[Dict[str, Any]], *, limit: int = _LOG_DOC_ID_LIM - If needed, revise your search queries based on the documents you find in previous steps. - Once you are confident that you have found all the related and somewhat related documents and there are \ no more related documents in the corpus, call the "final_results" tool to finish the task. -{enforce_top_k_line}\ +{final_results_count_line}\ - When calling the "final_results" tool, the list of documents must be sorted in the decreasing level of \ relevance to the query. I.e., the first document is the most relevant to the query, the second document is \ the second most relevant to the query, and so on. @@ -113,7 +113,6 @@ def _render_react_agent_prompt( top_k: int, *, with_init_docs: bool = True, - enforce_top_k: bool = True, extended_relevance: bool = False, ) -> str: """Render the ReAct agent system prompt (verbatim 02_v1.j2 logic).""" @@ -127,16 +126,14 @@ def _render_react_agent_prompt( if extended_relevance else "" ) - enforce_line = ( + final_results_count_line = ( f'- When calling "final_results", you must select exactly the {top_k} most relevant documents ' "among all documents you have retrieved.\n" - if enforce_top_k - else "" ) parts.append( _WORKFLOW_TEMPLATE.format( extended_relevance_line=ext_line, - enforce_top_k_line=enforce_line, + final_results_count_line=final_results_count_line, ) ) @@ -217,10 +214,8 @@ def _make_retrieve_tool_spec(top_k: int) -> Dict[str, Any]: } -def _make_final_results_tool_spec(top_k: Optional[int]) -> Dict[str, Any]: - tk_ins = "" - if top_k is not None: - tk_ins = f"- You must choose exactly {top_k} document IDs when calling this function.\n" +def _make_final_results_tool_spec(top_k: int) -> Dict[str, Any]: + tk_ins = f"- You must choose exactly {top_k} document IDs when calling this function.\n" description = ( "Signals the completion of the search process for the current query.\n\n" @@ -300,8 +295,7 @@ class ReActAgentOperator(AbstractOperator, CPUOperator): Rank Fusion. The system prompt is a verbatim Python rendering of the retrieval-bench - ``02_v1.j2`` template, including optional ``extended_relevance`` and - ``enforce_top_k`` blocks. + ``02_v1.j2`` template, including the optional ``extended_relevance`` block. Input DataFrame schema ---------------------- @@ -333,10 +327,6 @@ class ReActAgentOperator(AbstractOperator, CPUOperator): target_top_k : int Number of final documents to select, communicated to the LLM via the system prompt and ``final_results`` tool spec. Defaults to ``10``. - enforce_top_k : bool - When ``True``, the system prompt instructs the LLM to select exactly - ``target_top_k`` documents in its ``final_results`` call. - Defaults to ``True``. user_msg_type : {"with_results", "simple"} ``"with_results"`` (default): make one upfront retrieval call with the original query and include those documents in the first user message, @@ -403,7 +393,6 @@ def __init__( retriever_fn: Callable[[str, int], List[Dict[str, Any]]], retriever_top_k: int = 500, target_top_k: int = 10, - enforce_top_k: bool = True, user_msg_type: Literal["with_results", "simple"] = "with_results", extended_relevance: bool = False, max_steps: int = 10, @@ -422,7 +411,6 @@ def __init__( self._retriever_fn = retriever_fn self._retriever_top_k = retriever_top_k self._target_top_k = target_top_k - self._enforce_top_k = enforce_top_k self._user_msg_type = user_msg_type self._extended_relevance = extended_relevance self._max_steps = max_steps @@ -522,13 +510,12 @@ def _run_single_query( system_prompt = _render_react_agent_prompt( self._target_top_k, with_init_docs=with_init_docs, - enforce_top_k=self._enforce_top_k, extended_relevance=self._extended_relevance, ) tools = [ _make_think_tool_spec(self._extended_relevance), _make_retrieve_tool_spec(self._retriever_top_k), - _make_final_results_tool_spec(self._target_top_k if self._enforce_top_k else None), + _make_final_results_tool_spec(self._target_top_k), ] messages: List[Dict[str, Any]] = [{"role": "system", "content": system_prompt}] @@ -883,7 +870,7 @@ def _validate_final_results_args( preview = invalid_doc_ids[:_LOG_DOC_ID_LIMIT] return f"`doc_ids` contains IDs that were not retrieved: {preview}." - if self._enforce_top_k and len(doc_ids) != self._target_top_k: + if len(doc_ids) != self._target_top_k: return ( f"`doc_ids` must contain exactly {self._target_top_k} documents. " f"But got {len(doc_ids)} document IDs instead." diff --git a/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py b/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py index 352f628b02..d6d66631ad 100644 --- a/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py +++ b/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py @@ -617,11 +617,41 @@ def _select_documents( raw_doc_ids = json.loads(raw_doc_ids) except json.JSONDecodeError: raw_doc_ids = [] - doc_ids = [d for d in raw_doc_ids if d in valid_id_set][:feasible_k] + if not isinstance(raw_doc_ids, list): + tool_messages.append( + { + "role": "tool", + "tool_call_id": tc_id, + "content": "Error: `doc_ids` must be a list of candidate document IDs.", + } + ) + continue + + invalid_doc_ids = [doc_id for doc_id in raw_doc_ids if doc_id not in valid_id_set] + if invalid_doc_ids: + logger.warning( + "SelectionAgentOperator: LLM returned doc_id(s) outside the candidate set " + "for query %r: %s", + query_text, + invalid_doc_ids[:_LOG_DOC_ID_LIMIT], + ) + tool_messages.append( + { + "role": "tool", + "tool_call_id": tc_id, + "content": ( + "Error: `doc_ids` contains IDs that are not candidate documents: " + f"{invalid_doc_ids[:_LOG_DOC_ID_LIMIT]}. Use only valid candidate IDs." + ), + } + ) + continue + + doc_ids = raw_doc_ids[:feasible_k] logger.info( "SelectionAgentOperator: step=%d log_selected_documents raw=%s accepted=%s", _step, - raw_doc_ids[:_LOG_DOC_ID_LIMIT] if isinstance(raw_doc_ids, list) else raw_doc_ids, + raw_doc_ids[:_LOG_DOC_ID_LIMIT], doc_ids[:_LOG_DOC_ID_LIMIT], ) # Message can quote document text/PII; keep content at DEBUG. @@ -630,15 +660,6 @@ def _select_documents( _step, _preview_text(fn_args.get("message")), ) - if not doc_ids and raw_doc_ids: - logger.warning( - "SelectionAgentOperator: LLM returned %d doc_id(s) for query %r " - "but none matched the candidate set — possible hallucination. " - "Returned IDs: %s", - len(raw_doc_ids), - query_text, - raw_doc_ids[:10], - ) end_kwargs = {"doc_ids": doc_ids, "message": fn_args.get("message", "")} should_end = True diff --git a/nemo_retriever/src/nemo_retriever/query/agentic.py b/nemo_retriever/src/nemo_retriever/query/agentic.py index 31a45e4ec7..93f4d6b57a 100644 --- a/nemo_retriever/src/nemo_retriever/query/agentic.py +++ b/nemo_retriever/src/nemo_retriever/query/agentic.py @@ -144,7 +144,7 @@ class AgenticRetrievalConfig: reranker_api_key: str = "" local_reranker_backend: str = "vllm" embed_modality: str = "text" - llm_backend: str = AGENTIC_LLM_BACKEND + llm_backend: Optional[str] = None llm_model: str = "" invoke_url: Optional[str] = None local_llm_backend: str = AGENTIC_LOCAL_LLM_BACKEND @@ -171,18 +171,25 @@ 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 - # When true, ReAct final_results must contain exactly top_k doc IDs. Keep this - # strict by default for existing eval behavior, but allow local/smaller LLMs to - # return partial final_results when explicitly requested. - enforce_top_k: bool = True def __post_init__(self) -> None: - llm_backend = _normalize_agentic_choice( - self.llm_backend, - AGENTIC_LLM_BACKENDS, - field_name="llm_backend", - default=AGENTIC_LLM_BACKEND, - ) + invoke_url = _none_if_empty(self.invoke_url) + object.__setattr__(self, "invoke_url", invoke_url) + + explicit_llm_backend = str(self.llm_backend or "").strip().lower() + if explicit_llm_backend and explicit_llm_backend not in AGENTIC_LLM_BACKENDS: + raise ValueError(f"llm_backend must be one of {sorted(AGENTIC_LLM_BACKENDS)}; got {self.llm_backend!r}") + inferred_llm_backend = "openai_compatible" if invoke_url else AGENTIC_LLM_BACKEND + llm_backend = explicit_llm_backend or inferred_llm_backend + if invoke_url and llm_backend != "openai_compatible": + raise ValueError( + "invoke_url selects the openai_compatible agentic LLM backend; " + "omit invoke_url for in-process local LLMs." + ) + if not invoke_url and llm_backend == "openai_compatible": + raise ValueError( + "llm_backend='openai_compatible' requires invoke_url. Omit llm_backend to use in-process local LLMs." + ) object.__setattr__(self, "llm_backend", llm_backend) local_llm_backend = _normalize_agentic_choice( @@ -198,7 +205,7 @@ def __post_init__(self) -> None: if llm_backend == "in_process": llm_model = AGENTIC_LOCAL_LLM_MODEL else: - raise ValueError("Agentic retrieval requires a non-empty llm_model.") + raise ValueError("Agentic retrieval with invoke_url requires a non-empty llm_model.") if llm_backend == "in_process": from nemo_retriever.models.local.agent_llm import is_supported_agent_llm_model, supported_agent_llm_names @@ -208,15 +215,10 @@ def __post_init__(self) -> None: raise ValueError( f"Unsupported in-process agentic LLM model {llm_model!r}. " "Custom in-process agent LLMs are not supported yet. " - "Use llm_backend='openai_compatible' with invoke_url for a custom/self-hosted endpoint, " + "Provide invoke_url for a custom/self-hosted OpenAI-compatible endpoint, " f"or choose one of: {supported}." ) object.__setattr__(self, "llm_model", llm_model) - object.__setattr__( - self, - "enforce_top_k", - _agentic_bool_value(self.enforce_top_k, field_name="enforce_top_k"), - ) for field_name, value, min_value in ( ("react_max_steps", self.react_max_steps, 1), ("text_truncation", self.text_truncation, 0), @@ -284,18 +286,6 @@ def _normalize_agentic_choice(value: object, valid: frozenset[str], *, field_nam return normalized -def _agentic_bool_value(value: object, *, field_name: str) -> bool: - if isinstance(value, bool): - return value - if isinstance(value, str): - normalized = value.strip().lower() - if normalized in {"1", "true", "yes", "y", "on"}: - return True - if normalized in {"0", "false", "no", "n", "off"}: - return False - raise ValueError(f"{field_name} must be a boolean") - - def _agentic_float_range_value( value: object, *, @@ -426,7 +416,6 @@ def retrieve(self, query_ids: Sequence[str], query_texts: Sequence[str]) -> pd.D num_concurrent=int(self._cfg.num_concurrent), reasoning_effort=self._cfg.reasoning_effort, backend_top_k=self._cfg.backend_top_k, - enforce_top_k=bool(self._cfg.enforce_top_k), temperature=float(self._cfg.temperature), max_tokens=self._cfg.max_tokens, chat_completion_fn=chat_completion_fn, diff --git a/nemo_retriever/src/nemo_retriever/query/options.py b/nemo_retriever/src/nemo_retriever/query/options.py index c09c2628cc..f87e9cb8e2 100644 --- a/nemo_retriever/src/nemo_retriever/query/options.py +++ b/nemo_retriever/src/nemo_retriever/query/options.py @@ -55,7 +55,7 @@ class QueryAgenticOptions: enabled: bool = False llm_model: str | None = None - llm_backend: str = "in_process" + llm_backend: str | None = None invoke_url: str | None = None local_llm_backend: str = "vllm" local_hf_cache_dir: str | None = None diff --git a/nemo_retriever/src/nemo_retriever/query/workflow.py b/nemo_retriever/src/nemo_retriever/query/workflow.py index 890d456d35..8d16adfda0 100644 --- a/nemo_retriever/src/nemo_retriever/query/workflow.py +++ b/nemo_retriever/src/nemo_retriever/query/workflow.py @@ -174,7 +174,6 @@ def build_agentic_config(request: QueryRequest, *, top_k: int | None = None) -> "embedding_endpoint": request.embed.embed_invoke_url, "embedding_api_key": api_key or "", "llm_model": request.agentic.llm_model, - "llm_backend": request.agentic.llm_backend, "invoke_url": request.agentic.invoke_url, "local_llm_backend": request.agentic.local_llm_backend, "local_hf_cache_dir": request.agentic.local_hf_cache_dir, @@ -191,6 +190,8 @@ def build_agentic_config(request: QueryRequest, *, top_k: int | None = None) -> "num_concurrent": int(request.agentic.num_concurrent), "temperature": float(request.agentic.temperature), } + if request.agentic.llm_backend: + cfg_kwargs["llm_backend"] = request.agentic.llm_backend if request.embed.embed_model_name: cfg_kwargs["query_embedder"] = request.embed.embed_model_name if request.embed.embed_model_provider_prefix: diff --git a/nemo_retriever/tests/test_agentic_eval.py b/nemo_retriever/tests/test_agentic_eval.py index 42f2ccf5ab..0c9ca6dc8a 100644 --- a/nemo_retriever/tests/test_agentic_eval.py +++ b/nemo_retriever/tests/test_agentic_eval.py @@ -102,11 +102,11 @@ def test_build_beir_run_from_ranked_doc_ids_rejects_length_mismatch(): def test_agentic_config_validates_max_tokens(): from nemo_retriever.query.agentic import AgenticRetrievalConfig - cfg = AgenticRetrievalConfig(llm_model="test-model", llm_backend="openai_compatible", max_tokens="128") + cfg = AgenticRetrievalConfig(llm_model="nemotron-8b", max_tokens="128") assert cfg.max_tokens == 128 with pytest.raises(ValueError, match="max_tokens"): - AgenticRetrievalConfig(llm_model="test-model", llm_backend="openai_compatible", max_tokens=0) + AgenticRetrievalConfig(llm_model="nemotron-8b", max_tokens=0) @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") @@ -127,7 +127,6 @@ def test_agentic_retriever_runs_graph_with_wrapped_retriever(mock_react_step, mo cfg = AgenticRetrievalConfig( llm_model="test-model", - llm_backend="openai_compatible", invoke_url="http://localhost/v1/chat/completions", max_tokens=77, ) @@ -159,7 +158,6 @@ def test_agentic_retriever_honors_top_k(mock_react_step, mock_selection_step): cfg = AgenticRetrievalConfig( llm_model="test-model", - llm_backend="openai_compatible", invoke_url="http://localhost/v1/chat/completions", top_k=5, ) @@ -205,7 +203,7 @@ def test_agentic_retriever_builds_in_process_llm_lazily(mock_create_local_agent_ @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_can_accept_partial_react_final_results(mock_react_step, mock_selection_step): +def test_agentic_retriever_rejects_partial_react_final_results(mock_react_step, mock_selection_step): from nemo_retriever.query.agentic import AgenticRetrievalConfig, AgenticRetriever mock_react_step.return_value = _make_tool_call_response( @@ -219,15 +217,14 @@ def test_agentic_retriever_can_accept_partial_react_final_results(mock_react_ste cfg = AgenticRetrievalConfig( llm_model="test-model", - llm_backend="openai_compatible", invoke_url="http://localhost/v1/chat/completions", top_k=5, - enforce_top_k=False, + react_max_steps=1, ) result = AgenticRetriever(cfg, match_mode="pdf_page").retrieve(["0"], ["find doc"]) - assert result["doc_id"].tolist() == ["doc_1"] - assert result["result_source"].tolist() == ["final_results"] + assert len(result) == 5 + assert "final_results" not in set(result["result_source"]) mock_selection_step.assert_not_called() @@ -258,9 +255,7 @@ def test_run_agentic_audio_recall_evaluation_computes_metrics(mock_react_step, m {"doc_ids": [audio_doc_id], "message": "clip is best"}, ) - cfg = AgenticRetrievalConfig( - llm_model="test-model", llm_backend="openai_compatible", invoke_url="http://localhost/v1/chat/completions" - ) + cfg = AgenticRetrievalConfig(llm_model="test-model", invoke_url="http://localhost/v1/chat/completions") df_query, result, gold, retrieved, metrics = run_agentic_audio_recall_evaluation( query_csv=query_csv, cfg=cfg, @@ -297,9 +292,7 @@ def test_run_agentic_beir_evaluation_loads_queries_and_qrels(mock_react_step, mo queries=["find doc"], qrels={"q1": {"doc": 1}}, ) - cfg = AgenticRetrievalConfig( - llm_model="test-model", llm_backend="openai_compatible", invoke_url="http://localhost/v1/chat/completions" - ) + cfg = AgenticRetrievalConfig(llm_model="test-model", invoke_url="http://localhost/v1/chat/completions") with patch("nemo_retriever.query.agentic.load_beir_dataset", return_value=beir_dataset) as mock_loader: df_query, result, qrels, run, metrics = run_agentic_beir_evaluation( @@ -334,17 +327,17 @@ def test_agentic_config_requires_llm_model_for_openai_compatible(): from nemo_retriever.query.agentic import AgenticRetrievalConfig with pytest.raises(ValueError, match="llm_model"): - AgenticRetrievalConfig(llm_model="", llm_backend="openai_compatible") + AgenticRetrievalConfig(llm_model="", invoke_url="http://localhost/v1/chat/completions") # None must not slip through as the literal string "None". with pytest.raises(ValueError, match="llm_model"): - AgenticRetrievalConfig(llm_model=None, llm_backend="openai_compatible") + AgenticRetrievalConfig(llm_model=None, invoke_url="http://localhost/v1/chat/completions") def test_agentic_config_rejects_custom_in_process_llm_model(): from nemo_retriever.query.agentic import AgenticRetrievalConfig with pytest.raises(ValueError, match="Custom in-process agent LLMs are not supported yet"): - AgenticRetrievalConfig(llm_model="custom/local-model", llm_backend="in_process") + AgenticRetrievalConfig(llm_model="custom/local-model") def test_agentic_config_rejects_nonpositive_top_k(): @@ -366,7 +359,6 @@ def test_agentic_config_normalizes_integer_like_values(): cfg = AgenticRetrievalConfig( llm_model="m", - llm_backend="openai_compatible", invoke_url="http://localhost/v1/chat/completions", top_k="5.0", backend_top_k="6.0", @@ -389,7 +381,11 @@ def test_agentic_config_rejects_nvidia_temperature_above_max(): from nemo_retriever.query.agentic import AgenticRetrievalConfig with pytest.raises(ValueError, match="between 0.0 and 1.0"): - AgenticRetrievalConfig(llm_model="m", llm_backend="openai_compatible", temperature=1.5) + AgenticRetrievalConfig( + llm_model="m", + invoke_url="https://integrate.api.nvidia.com/v1/chat/completions", + temperature=1.5, + ) def test_agentic_config_rejects_nonfinite_temperature(): @@ -402,7 +398,7 @@ def test_agentic_config_rejects_nonfinite_temperature(): def test_agentic_config_accepts_in_process_temperature_above_nvidia_limit(): from nemo_retriever.query.agentic import AgenticRetrievalConfig - cfg = AgenticRetrievalConfig(llm_model="nemotron-8b", llm_backend="in_process", temperature=1.5) + cfg = AgenticRetrievalConfig(llm_model="nemotron-8b", temperature=1.5) assert cfg.llm_backend == "in_process" assert cfg.temperature == pytest.approx(1.5) @@ -412,7 +408,7 @@ def test_agentic_config_rejects_invalid_local_llm_backend(): from nemo_retriever.query.agentic import AgenticRetrievalConfig with pytest.raises(ValueError, match="local_llm_backend"): - AgenticRetrievalConfig(llm_model="nemotron-8b", llm_backend="in_process", local_llm_backend="hf") + AgenticRetrievalConfig(llm_model="nemotron-8b", local_llm_backend="hf") def test_agentic_config_validates_local_vllm_knobs(): @@ -420,7 +416,6 @@ def test_agentic_config_validates_local_vllm_knobs(): cfg = AgenticRetrievalConfig( llm_model="nemotron-8b", - llm_backend="in_process", local_gpu_memory_utilization="0.6", local_tensor_parallel_size="2.0", local_max_model_len="8192", diff --git a/nemo_retriever/tests/test_agentic_operators.py b/nemo_retriever/tests/test_agentic_operators.py index 8025b71068..fda6853608 100644 --- a/nemo_retriever/tests/test_agentic_operators.py +++ b/nemo_retriever/tests/test_agentic_operators.py @@ -104,7 +104,7 @@ class TestPromptRendering: def test_react_prompt_no_extended_relevance(self): from nemo_retriever.operators.graph_ops.react_agent_operator import _render_react_agent_prompt - prompt = _render_react_agent_prompt(10, with_init_docs=True, enforce_top_k=True, extended_relevance=False) + prompt = _render_react_agent_prompt(10, with_init_docs=True, extended_relevance=False) assert "" in prompt assert "" in prompt assert "" in prompt @@ -115,9 +115,9 @@ def test_react_prompt_no_extended_relevance(self): def test_react_prompt_with_extended_relevance(self): from nemo_retriever.operators.graph_ops.react_agent_operator import _render_react_agent_prompt - prompt = _render_react_agent_prompt(5, with_init_docs=False, enforce_top_k=False, extended_relevance=True) + prompt = _render_react_agent_prompt(5, with_init_docs=False, extended_relevance=True) assert "RELEVANCE_DEFINITION" in prompt - assert "exactly the 5" not in prompt + assert "exactly the 5" in prompt assert "TIP" not in prompt def test_selection_prompt_no_extended_relevance(self): @@ -198,6 +198,33 @@ def test_happy_path_selects_docs(self, mock_step): assert result["rank"].tolist() == [1, 2] assert result["result_source"].tolist() == ["selection_agent", "selection_agent"] + @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") + def test_retries_when_selection_returns_invalid_doc_ids(self, mock_step): + from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator + + mock_step.side_effect = [ + _make_tool_call_response( + "log_selected_documents", + {"doc_ids": ["d1", "missing"], "message": "mixed valid and invalid"}, + ), + _make_tool_call_response( + "log_selected_documents", + {"doc_ids": ["d1", "d2"], "message": "corrected"}, + ), + ] + + op = SelectionAgentOperator( + llm_model="test-model", + invoke_url="http://localhost/v1/chat/completions", + top_k=2, + max_steps=2, + ) + result = op.run(self._make_input()) + + assert mock_step.call_count == 2 + assert result["doc_id"].tolist() == ["d1", "d2"] + assert result["message"].tolist() == ["corrected", "corrected"] + @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") def test_think_then_select(self, mock_step): from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator @@ -478,19 +505,19 @@ def retriever_fn(query_text, top_k): assert result[1]["text"] == "new two" @pytest.mark.parametrize( - ("fn_args", "target_top_k", "enforce_top_k"), + ("fn_args", "target_top_k"), [ - ({"doc_ids": [1], "message": "bad id type", "search_successful": "true"}, 1, False), - ({"doc_ids": [], "message": "empty", "search_successful": "false"}, 1, False), - ({"doc_ids": [""], "message": "empty-string id", "search_successful": "true"}, 1, False), - ({"doc_ids": [" "], "message": "whitespace id", "search_successful": "true"}, 1, False), - ({"doc_ids": ["d1"], "message": "wrong count", "search_successful": "true"}, 2, True), - ({"doc_ids": ["missing"], "message": "hallucinated id", "search_successful": "true"}, 1, False), - ({"doc_ids": ["d1"], "message": "bad status", "search_successful": "yes"}, 1, False), + ({"doc_ids": [1], "message": "bad id type", "search_successful": "true"}, 1), + ({"doc_ids": [], "message": "empty", "search_successful": "false"}, 1), + ({"doc_ids": [""], "message": "empty-string id", "search_successful": "true"}, 1), + ({"doc_ids": [" "], "message": "whitespace id", "search_successful": "true"}, 1), + ({"doc_ids": ["d1"], "message": "wrong count", "search_successful": "true"}, 2), + ({"doc_ids": ["missing"], "message": "hallucinated id", "search_successful": "true"}, 1), + ({"doc_ids": ["d1"], "message": "bad status", "search_successful": "yes"}, 1), ], ) @patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") - def test_invalid_final_results_are_rejected(self, mock_step, fn_args, target_top_k, enforce_top_k): + def test_invalid_final_results_are_rejected(self, mock_step, fn_args, target_top_k): from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator mock_step.return_value = _make_tool_call_response("final_results", fn_args) @@ -502,7 +529,6 @@ def test_invalid_final_results_are_rejected(self, mock_step, fn_args, target_top retriever_fn=retriever, user_msg_type="with_results", target_top_k=target_top_k, - enforce_top_k=enforce_top_k, ) result = op.run(self._make_input()) diff --git a/nemo_retriever/tests/test_harness_agentic_eval.py b/nemo_retriever/tests/test_harness_agentic_eval.py index c93d357ac7..b47dd26e61 100644 --- a/nemo_retriever/tests/test_harness_agentic_eval.py +++ b/nemo_retriever/tests/test_harness_agentic_eval.py @@ -41,7 +41,6 @@ def test_query_override_paths_include_agentic_fields() -> None: for key in ( "query.agentic", "query.agentic_llm_model", - "query.agentic_llm_backend", "query.agentic_invoke_url", "query.agentic_local_gpu_memory_utilization", "query.agentic_local_tensor_parallel_size", @@ -64,7 +63,6 @@ def test_build_query_request_populates_agentic() -> None: "top_k": 10, "agentic": True, "agentic_llm_model": "test-model", - "agentic_llm_backend": "openai_compatible", "agentic_invoke_url": "https://example.invalid/v1", "agentic_local_gpu_memory_utilization": 0.6, "agentic_local_tensor_parallel_size": 2, @@ -83,7 +81,7 @@ def test_build_query_request_populates_agentic() -> None: agentic = request.agentic assert agentic.enabled is True assert agentic.llm_model == "test-model" - assert agentic.llm_backend == "openai_compatible" + assert agentic.llm_backend is None assert agentic.invoke_url == "https://example.invalid/v1" assert agentic.local_llm_backend == "vllm" assert agentic.local_gpu_memory_utilization == pytest.approx(0.6) @@ -110,7 +108,7 @@ def test_build_agentic_config_maps_request_and_top_k_override() -> None: agentic=QueryAgenticOptions( enabled=True, llm_model="test-model", - llm_backend="openai_compatible", + invoke_url="http://localhost/v1/chat/completions", backend_top_k=20, num_concurrent=4, temperature=0.0, @@ -118,6 +116,7 @@ def test_build_agentic_config_maps_request_and_top_k_override() -> None: ) cfg = build_agentic_config(request, top_k=10) assert cfg.llm_model == "test-model" + assert cfg.llm_backend == "openai_compatible" assert cfg.top_k == 10 # harness sets this to the deepest BEIR k assert cfg.backend_top_k == 20 assert cfg.num_concurrent == 4 @@ -155,7 +154,7 @@ def test_run_beir_queries_routes_to_agentic(tmp_path) -> None: "top_k": 10, "agentic": True, "agentic_llm_model": "test-model", - "agentic_llm_backend": "openai_compatible", + "agentic_invoke_url": "http://localhost/v1/chat/completions", } ), "", diff --git a/nemo_retriever/tests/test_root_query_cli.py b/nemo_retriever/tests/test_root_query_cli.py index 053da62cd8..f719b85b93 100644 --- a/nemo_retriever/tests/test_root_query_cli.py +++ b/nemo_retriever/tests/test_root_query_cli.py @@ -391,7 +391,7 @@ def retrieve(self, query_ids: Any, query_texts: Any) -> Any: cfg = config_calls[0] assert cfg["vdb_op"] == "lancedb" assert cfg["vdb_kwargs"] == {"uri": "/tmp/lancedb", "table_name": "docs"} - assert cfg["llm_backend"] == "in_process" + assert "llm_backend" not in cfg assert cfg["local_llm_backend"] == "vllm" assert cfg["llm_model"] == "nemotron-8b" # --top-k is honored end-to-end: plumbed into the agentic config (drives the @@ -412,17 +412,17 @@ def test_root_query_agentic_rejects_custom_in_process_llm_model() -> None: assert result.exit_code == 1 assert "Custom in-process agent LLMs are not supported yet" in result.output - assert "--agentic-llm-backend openai_compatible" in result.output or "openai_compatible" in result.output + assert "--agentic-invoke-url" in result.output or "invoke_url" in result.output -def test_root_query_agentic_openai_compatible_requires_llm_model() -> None: +def test_root_query_agentic_invoke_url_requires_llm_model() -> None: result = RUNNER.invoke( cli_main.app, - ["query", "hello", "--agentic", "--agentic-llm-backend", "openai_compatible"], + ["query", "hello", "--agentic", "--agentic-invoke-url", "http://localhost:8000/v1/chat/completions"], ) assert result.exit_code == 1 - assert "openai_compatible requires --agentic-llm-model" in result.output + assert "--agentic-invoke-url requires --agentic-llm-model" in result.output def test_root_query_agentic_openai_compatible_allows_custom_model(monkeypatch) -> None: @@ -452,8 +452,6 @@ def retrieve(self, query_ids: Any, query_texts: Any) -> Any: "query", "q", "--agentic", - "--agentic-llm-backend", - "openai_compatible", "--agentic-llm-model", "custom-remote-model", "--agentic-invoke-url", @@ -462,7 +460,7 @@ def retrieve(self, query_ids: Any, query_texts: Any) -> Any: ) assert result.exit_code == 0 - assert config_calls[-1]["llm_backend"] == "openai_compatible" + assert "llm_backend" not in config_calls[-1] assert config_calls[-1]["llm_model"] == "custom-remote-model" assert config_calls[-1]["invoke_url"] == "http://localhost:8000/v1/chat/completions" From 858cbce27985679b18cae8e0ab9a4f5168bb88c6 Mon Sep 17 00:00:00 2001 From: Mahika Wason Date: Mon, 20 Jul 2026 22:19:15 -0700 Subject: [PATCH 06/10] greptile review --- nemo_retriever/docs/cli/README.md | 3 + .../src/nemo_retriever/cli/query/app.py | 9 +- .../src/nemo_retriever/cli/query/options.py | 8 ++ .../nemo_retriever/models/local/agent_llm.py | 92 ++++++++++++------- .../tests/test_agentic_local_llm.py | 42 +++++++++ nemo_retriever/tests/test_root_query_cli.py | 3 + 6 files changed, 122 insertions(+), 35 deletions(-) diff --git a/nemo_retriever/docs/cli/README.md b/nemo_retriever/docs/cli/README.md index c066341029..83b7c8c3e4 100644 --- a/nemo_retriever/docs/cli/README.md +++ b/nemo_retriever/docs/cli/README.md @@ -225,6 +225,9 @@ Agentic-only knobs (apply only with `--agentic`): agent LLM. Providing it routes agent LLM calls to that remote endpoint. - `--agentic-reasoning-effort` (default `high`) — `reasoning_effort` forwarded on OpenAI-compatible agentic LLM calls; ignored by the local adapter. +- `--agentic-temperature` (default `0.0`) — sampling temperature for agent LLM + calls. Local and non-NVIDIA OpenAI-compatible endpoints allow up to `2.0`; + NVIDIA-hosted endpoints allow up to `1.0`. - `--agentic-backend-top-k` (default `20`) — candidates pulled from the vector DB per retrieval call. - `--agentic-react-max-steps` (default `50`) — maximum ReAct loop iterations. diff --git a/nemo_retriever/src/nemo_retriever/cli/query/app.py b/nemo_retriever/src/nemo_retriever/cli/query/app.py index c2b0292ed4..23af798552 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/app.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/app.py @@ -16,7 +16,7 @@ from nemo_retriever.cli.query import options as opts from nemo_retriever.cli.query_workflow import agentic_query_documents as query_agentic_documents from nemo_retriever.cli.query_workflow import query_documents_with_metadata as query_local_documents_with_metadata -from nemo_retriever.query.agentic_options import agentic_backend_top_k_error +from nemo_retriever.query.agentic_options import agentic_backend_top_k_error, agentic_temperature_error from nemo_retriever.cli.shared import ( ROOT_CLI_ERRORS, quiet_capture, @@ -189,6 +189,7 @@ def _local_command( agentic_backend_top_k: opts.AgenticBackendTopKOption = 20, agentic_react_max_steps: opts.AgenticReactMaxStepsOption = 50, agentic_text_truncation: opts.AgenticTextTruncationOption = 0, + agentic_temperature: opts.AgenticTemperatureOption = 0.0, ) -> None: _validate_output_options(output_format, max_text_chars) if reranker_invoke_url is None: @@ -211,6 +212,11 @@ def _local_command( if backend_error: typer.echo(f"Error: {backend_error}", err=True) raise typer.Exit(1) + temperature_invoke_url = agentic_invoke_url or "local://in-process" + temperature_error = agentic_temperature_error(agentic_temperature, invoke_url=temperature_invoke_url) + if temperature_error: + typer.echo(f"Error: {temperature_error}", err=True) + raise typer.Exit(1) try: reranker_api_key = _api_key_from_env_option(reranker_api_key_env) if reranker_invoke_url else None @@ -250,6 +256,7 @@ def _local_command( backend_top_k=agentic_backend_top_k, react_max_steps=agentic_react_max_steps, text_truncation=agentic_text_truncation, + temperature=agentic_temperature, ), ) with quiet_capture(): diff --git a/nemo_retriever/src/nemo_retriever/cli/query/options.py b/nemo_retriever/src/nemo_retriever/cli/query/options.py index dae6b72924..b47d1f3646 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/options.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/options.py @@ -203,6 +203,14 @@ help="Max characters of each candidate shown to the agent; 0 disables truncation.", ), ] +AgenticTemperatureOption = Annotated[ + float, + typer.Option( + "--agentic-temperature", + min=0.0, + help="Sampling temperature for agentic LLM calls (0.0 = greedy).", + ), +] ServiceUrlOption = Annotated[ str, typer.Option("--service-url", help="Base URL of the retriever service."), diff --git a/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py index 6e24638a9c..ac21d06a8d 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py +++ b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py @@ -110,41 +110,51 @@ def __init__( f"or choose one of: {', '.join(supported_agent_llm_names())}." ) + self._previous_cuda_visible_devices: Optional[str] = None + self._cuda_visible_devices_overridden = False cuda_visible_devices = _cuda_visible_devices_from_device(device) if cuda_visible_devices is not None: + self._previous_cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") + self._cuda_visible_devices_overridden = True os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices try: - from vllm import LLM, SamplingParams # noqa: F401 - except ImportError as e: - raise ImportError( - 'Local agentic LLM inference requires vLLM. Install with: pip install "nemo-retriever[local]"' - ) from e - - _raise_if_cuda_unavailable() - self._model_path = model_path - self._max_tokens = int(max_tokens) - self._request_extras = deepcopy(_NEMOTRON_JSON_TOOL_PROMPT_EXTRAS) - self._lock = threading.Lock() - - configure_global_hf_cache_base(hf_cache_dir) - revision = get_hf_revision(model_path) - - engine_kwargs: dict[str, Any] = {} - if max_model_len is not None: - engine_kwargs["max_model_len"] = int(max_model_len) - if max_num_seqs is not None: - engine_kwargs["max_num_seqs"] = int(max_num_seqs) - - self._sampling_params_cls = SamplingParams - self._llm: Any | None = LLM( - model=model_path, - revision=revision, - trust_remote_code=True, - tensor_parallel_size=int(tensor_parallel_size), - gpu_memory_utilization=float(gpu_memory_utilization), - **engine_kwargs, - ) + try: + from vllm import LLM, SamplingParams + except ImportError as e: + raise ImportError( + 'Local agentic LLM inference requires vLLM. Install with: pip install "nemo-retriever[local]"' + ) from e + + _raise_if_cuda_unavailable() + self._model_path = model_path + self._max_tokens = int(max_tokens) + self._request_extras = deepcopy(_NEMOTRON_JSON_TOOL_PROMPT_EXTRAS) + self._lock = threading.Lock() + + configure_global_hf_cache_base(hf_cache_dir) + revision = get_hf_revision(model_path) + + engine_kwargs: dict[str, Any] = {} + if max_model_len is not None: + engine_kwargs["max_model_len"] = int(max_model_len) + if max_num_seqs is not None: + engine_kwargs["max_num_seqs"] = int(max_num_seqs) + + self._sampling_params_cls = SamplingParams + self._llm: Any | None = LLM( + model=model_path, + revision=revision, + # Safe for this local path: supported model names are hard-allowlisted + # above and revisions are pinned in hf_model_registry.py. + trust_remote_code=True, + tensor_parallel_size=int(tensor_parallel_size), + gpu_memory_utilization=float(gpu_memory_utilization), + **engine_kwargs, + ) + except Exception: + self._restore_cuda_visible_devices() + raise def __call__( self, @@ -241,10 +251,14 @@ def unload(self) -> None: """Release GPU memory held by the local vLLM engine.""" with self._lock: - if self._llm is None: - return - del self._llm - self._llm = None + had_llm = self._llm is not None + if had_llm: + del self._llm + self._llm = None + + self._restore_cuda_visible_devices() + if not had_llm: + return try: import torch @@ -253,6 +267,16 @@ def unload(self) -> None: if torch.cuda.is_available(): torch.cuda.empty_cache() + def _restore_cuda_visible_devices(self) -> None: + if not getattr(self, "_cuda_visible_devices_overridden", False): + return + previous = getattr(self, "_previous_cuda_visible_devices", None) + if previous is None: + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + else: + os.environ["CUDA_VISIBLE_DEVICES"] = previous + self._cuda_visible_devices_overridden = False + def _require_loaded(self) -> None: if self._llm is None: raise RuntimeError( diff --git a/nemo_retriever/tests/test_agentic_local_llm.py b/nemo_retriever/tests/test_agentic_local_llm.py index 99325e8153..ae18c857f3 100644 --- a/nemo_retriever/tests/test_agentic_local_llm.py +++ b/nemo_retriever/tests/test_agentic_local_llm.py @@ -5,6 +5,8 @@ from __future__ import annotations import json +import os +import sys import threading from unittest.mock import MagicMock @@ -167,6 +169,46 @@ def test_vllm_agent_llm_unload_releases_engine() -> None: llm._require_loaded() +def test_vllm_agent_llm_unload_restores_cuda_visible_devices(monkeypatch) -> None: + from nemo_retriever.models.local.agent_llm import VLLMAgentChatLLM + + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") + llm = VLLMAgentChatLLM.__new__(VLLMAgentChatLLM) + llm._llm = object() + llm._lock = threading.Lock() + llm._cuda_visible_devices_overridden = True + llm._previous_cuda_visible_devices = "2,3" + + llm.unload() + + assert os.environ["CUDA_VISIBLE_DEVICES"] == "2,3" + + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") + llm = VLLMAgentChatLLM.__new__(VLLMAgentChatLLM) + llm._llm = object() + llm._lock = threading.Lock() + llm._cuda_visible_devices_overridden = True + llm._previous_cuda_visible_devices = None + + llm.unload() + + assert "CUDA_VISIBLE_DEVICES" not in os.environ + + +def test_vllm_agent_llm_restores_cuda_visible_devices_on_init_failure(monkeypatch) -> None: + import pytest + + from nemo_retriever.models.local.agent_llm import VLLMAgentChatLLM + + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "2,3") + monkeypatch.setitem(sys.modules, "vllm", None) + + with pytest.raises(ImportError, match="requires vLLM"): + VLLMAgentChatLLM("nemotron-8b", device="0,1") + + assert os.environ["CUDA_VISIBLE_DEVICES"] == "2,3" + + def test_unload_cached_vllm_agent_chat_llms_clears_cache() -> None: from nemo_retriever.models.local import agent_llm diff --git a/nemo_retriever/tests/test_root_query_cli.py b/nemo_retriever/tests/test_root_query_cli.py index f719b85b93..16c42c52b1 100644 --- a/nemo_retriever/tests/test_root_query_cli.py +++ b/nemo_retriever/tests/test_root_query_cli.py @@ -383,6 +383,8 @@ def retrieve(self, query_ids: Any, query_texts: Any) -> Any: "/tmp/lancedb", "--table-name", "docs", + "--agentic-temperature", + "1.25", ], ) @@ -394,6 +396,7 @@ def retrieve(self, query_ids: Any, query_texts: Any) -> Any: assert "llm_backend" not in cfg assert cfg["local_llm_backend"] == "vllm" assert cfg["llm_model"] == "nemotron-8b" + assert cfg["temperature"] == 1.25 # --top-k is honored end-to-end: plumbed into the agentic config (drives the # ReAct target / RRF / selection cut), not just applied as a post-filter. assert cfg["top_k"] == 2 From 3284e82d0e4c3dbe369496f53855fffdf3bddc8d Mon Sep 17 00:00:00 2001 From: Mahika Wason Date: Tue, 21 Jul 2026 11:14:00 -0700 Subject: [PATCH 07/10] cleanup and added local llm agentic config --- .../src/nemo_retriever/models/__init__.py | 22 +-- .../nemo_retriever/models/local/agent_llm.py | 186 +++++++----------- .../src/nemo_retriever/query/agentic.py | 26 +-- .../nemo_retriever/query/agentic_options.py | 20 ++ .../src/nemo_retriever/query/options.py | 1 - .../src/nemo_retriever/query/workflow.py | 1 - nemo_retriever/tests/test_agentic_eval.py | 1 - .../tests/test_agentic_local_llm.py | 83 +++----- 8 files changed, 121 insertions(+), 219 deletions(-) diff --git a/nemo_retriever/src/nemo_retriever/models/__init__.py b/nemo_retriever/src/nemo_retriever/models/__init__.py index 7ca2451526..5c75ed8445 100644 --- a/nemo_retriever/src/nemo_retriever/models/__init__.py +++ b/nemo_retriever/src/nemo_retriever/models/__init__.py @@ -267,7 +267,6 @@ def create_local_agent_llm( model_name: str, *, backend: str = "vllm", - device: str | None = None, hf_cache_dir: str | None = None, gpu_memory_utilization: float = 0.8, tensor_parallel_size: int = 1, @@ -278,22 +277,23 @@ def create_local_agent_llm( The callable mirrors ``invoke_chat_completion_step`` and returns an OpenAI-compatible chat-completions response dict. V1 supports the in-process - vLLM backend and uses model/profile-specific tool-call normalization where - known. + vLLM backend and uses process-level vLLM placement (for example, + ``CUDA_VISIBLE_DEVICES`` plus ``tensor_parallel_size``). """ b = normalize_backend(backend, _LOCAL_AGENT_LLM_BACKENDS, field_name="backend", default="vllm") if b == "vllm": - from nemo_retriever.models.local.agent_llm import create_cached_vllm_agent_chat_llm + from nemo_retriever.models.local.agent_llm import LocalAgentLLMConfig, create_cached_vllm_agent_chat_llm return create_cached_vllm_agent_chat_llm( - model_name, - device=device, - hf_cache_dir=hf_cache_dir, - gpu_memory_utilization=gpu_memory_utilization, - tensor_parallel_size=tensor_parallel_size, - max_model_len=max_model_len, - max_num_seqs=max_num_seqs, + LocalAgentLLMConfig( + model_path=model_name, + hf_cache_dir=hf_cache_dir, + gpu_memory_utilization=gpu_memory_utilization, + tensor_parallel_size=tensor_parallel_size, + max_model_len=max_model_len, + max_num_seqs=max_num_seqs, + ) ) raise ValueError(f"Unsupported local agent LLM backend {backend!r}") diff --git a/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py index ac21d06a8d..bae03ba4ff 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py +++ b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py @@ -6,10 +6,10 @@ import json import logging -import os import threading import uuid from copy import deepcopy +from dataclasses import dataclass from typing import Any, Mapping, Optional, Sequence from nemo_retriever.models.hf_cache import configure_global_hf_cache_base @@ -77,6 +77,19 @@ def resolve_agent_llm_model_name(name: str) -> str: return _AGENT_LLM_MODEL_ALIASES.get(_normalize_name(name), name) +@dataclass(frozen=True) +class LocalAgentLLMConfig: + """Local vLLM configuration for the agent chat LLM.""" + + model_path: str + hf_cache_dir: Optional[str] = None + gpu_memory_utilization: float = 0.8 + tensor_parallel_size: int = 1 + max_model_len: Optional[int] = None + max_num_seqs: Optional[int] = None + max_tokens: int = _DEFAULT_MAX_TOKENS + + class VLLMAgentChatLLM(BaseModel): """In-process vLLM chat-completions adapter for agentic retrieval. @@ -86,21 +99,10 @@ class VLLMAgentChatLLM(BaseModel): semantics. """ - def __init__( - self, - model_path: str, - *, - device: Optional[str] = None, - hf_cache_dir: Optional[str] = None, - gpu_memory_utilization: float = 0.8, - tensor_parallel_size: int = 1, - max_model_len: Optional[int] = None, - max_num_seqs: Optional[int] = None, - max_tokens: int = _DEFAULT_MAX_TOKENS, - ) -> None: + def __init__(self, config: LocalAgentLLMConfig) -> None: super().__init__() - requested_model_path = model_path + requested_model_path = config.model_path model_path = resolve_agent_llm_model_name(requested_model_path) if not is_supported_agent_llm_model(model_path): raise ValueError( @@ -110,51 +112,39 @@ def __init__( f"or choose one of: {', '.join(supported_agent_llm_names())}." ) - self._previous_cuda_visible_devices: Optional[str] = None - self._cuda_visible_devices_overridden = False - cuda_visible_devices = _cuda_visible_devices_from_device(device) - if cuda_visible_devices is not None: - self._previous_cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") - self._cuda_visible_devices_overridden = True - os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices - try: - try: - from vllm import LLM, SamplingParams - except ImportError as e: - raise ImportError( - 'Local agentic LLM inference requires vLLM. Install with: pip install "nemo-retriever[local]"' - ) from e - - _raise_if_cuda_unavailable() - self._model_path = model_path - self._max_tokens = int(max_tokens) - self._request_extras = deepcopy(_NEMOTRON_JSON_TOOL_PROMPT_EXTRAS) - self._lock = threading.Lock() - - configure_global_hf_cache_base(hf_cache_dir) - revision = get_hf_revision(model_path) - - engine_kwargs: dict[str, Any] = {} - if max_model_len is not None: - engine_kwargs["max_model_len"] = int(max_model_len) - if max_num_seqs is not None: - engine_kwargs["max_num_seqs"] = int(max_num_seqs) - - self._sampling_params_cls = SamplingParams - self._llm: Any | None = LLM( - model=model_path, - revision=revision, - # Safe for this local path: supported model names are hard-allowlisted - # above and revisions are pinned in hf_model_registry.py. - trust_remote_code=True, - tensor_parallel_size=int(tensor_parallel_size), - gpu_memory_utilization=float(gpu_memory_utilization), - **engine_kwargs, - ) - except Exception: - self._restore_cuda_visible_devices() - raise + from vllm import LLM, SamplingParams + except ImportError as e: + raise ImportError( + 'Local agentic LLM inference requires vLLM. Install with: pip install "nemo-retriever[local]"' + ) from e + + _raise_if_cuda_unavailable() + self._model_path = model_path + self._max_tokens = int(config.max_tokens) + self._request_extras = deepcopy(_NEMOTRON_JSON_TOOL_PROMPT_EXTRAS) + self._lock = threading.Lock() + + configure_global_hf_cache_base(config.hf_cache_dir) + revision = get_hf_revision(model_path) + + engine_kwargs: dict[str, Any] = {} + if config.max_model_len is not None: + engine_kwargs["max_model_len"] = int(config.max_model_len) + if config.max_num_seqs is not None: + engine_kwargs["max_num_seqs"] = int(config.max_num_seqs) + + self._sampling_params_cls = SamplingParams + self._llm: Any | None = LLM( + model=model_path, + revision=revision, + # Safe for this local path: supported model names are hard-allowlisted + # above and revisions are pinned in hf_model_registry.py. + trust_remote_code=True, + tensor_parallel_size=int(config.tensor_parallel_size), + gpu_memory_utilization=float(config.gpu_memory_utilization), + **engine_kwargs, + ) def __call__( self, @@ -233,7 +223,15 @@ def _normalize_messages( tools: Optional[Sequence[dict[str, Any]]] = None, ) -> list[dict[str, Any]]: normalized: list[dict[str, Any]] = [] - tool_prompt = _json_tool_prompt(tools) if tools else None + tool_prompt = ( + "You have access to the following tools. Your entire response must be a JSON array of one or more " + "tool calls. Each item must be shaped as " + '{"name": "tool_name", "arguments": {"arg_name": "arg_value"}}. ' + "Use only the listed tool names. Do not include prose, markdown, or text outside the JSON array.\n" + f"{json.dumps(list(tools), ensure_ascii=False)}" + if tools + else None + ) for idx, message in enumerate(messages): msg = dict(message) content = msg.get("content") @@ -256,7 +254,6 @@ def unload(self) -> None: del self._llm self._llm = None - self._restore_cuda_visible_devices() if not had_llm: return @@ -267,16 +264,6 @@ def unload(self) -> None: if torch.cuda.is_available(): torch.cuda.empty_cache() - def _restore_cuda_visible_devices(self) -> None: - if not getattr(self, "_cuda_visible_devices_overridden", False): - return - previous = getattr(self, "_previous_cuda_visible_devices", None) - if previous is None: - os.environ.pop("CUDA_VISIBLE_DEVICES", None) - else: - os.environ["CUDA_VISIBLE_DEVICES"] = previous - self._cuda_visible_devices_overridden = False - def _require_loaded(self) -> None: if self._llm is None: raise RuntimeError( @@ -308,20 +295,6 @@ def input_batch_size(self) -> int: return 1 -def _cuda_visible_devices_from_device(device: Optional[str]) -> str | None: - if device is None: - return None - normalized = str(device).strip() - if not normalized: - return None - if normalized.casefold() == "cpu": - raise ValueError( - "The local agent LLM vLLM backend requires CUDA. Pass GPU ids such as '0' or '0,1', " - "or provide invoke_url for a remote OpenAI-compatible endpoint." - ) - return normalized.split(":", 1)[1] if normalized.startswith("cuda:") else normalized - - def _raise_if_cuda_unavailable() -> None: try: import torch @@ -334,39 +307,22 @@ def _raise_if_cuda_unavailable() -> None: ) -def create_cached_vllm_agent_chat_llm( - model_name: str, - *, - device: Optional[str] = None, - hf_cache_dir: Optional[str] = None, - gpu_memory_utilization: float = 0.8, - tensor_parallel_size: int = 1, - max_model_len: Optional[int] = None, - max_num_seqs: Optional[int] = None, -) -> VLLMAgentChatLLM: +def create_cached_vllm_agent_chat_llm(config: LocalAgentLLMConfig) -> VLLMAgentChatLLM: """Create or reuse a local agent LLM keyed by heavyweight load settings.""" key = ( - resolve_agent_llm_model_name(model_name), - device, - hf_cache_dir, - float(gpu_memory_utilization), - int(tensor_parallel_size), - int(max_model_len) if max_model_len is not None else None, - int(max_num_seqs) if max_num_seqs is not None else None, + resolve_agent_llm_model_name(config.model_path), + config.hf_cache_dir, + float(config.gpu_memory_utilization), + int(config.tensor_parallel_size), + int(config.max_model_len) if config.max_model_len is not None else None, + int(config.max_num_seqs) if config.max_num_seqs is not None else None, + int(config.max_tokens), ) with _LOCAL_AGENT_LLM_CACHE_LOCK: cached = _LOCAL_AGENT_LLM_CACHE.get(key) if cached is None or cached._llm is None: - cached = VLLMAgentChatLLM( - model_path=model_name, - device=device, - hf_cache_dir=hf_cache_dir, - gpu_memory_utilization=gpu_memory_utilization, - tensor_parallel_size=tensor_parallel_size, - max_model_len=max_model_len, - max_num_seqs=max_num_seqs, - ) + cached = VLLMAgentChatLLM(config) _LOCAL_AGENT_LLM_CACHE[key] = cached return cached @@ -560,16 +516,6 @@ def _format_tool_message_content(message: Mapping[str, Any], tool_names_by_id: M return f"Tool result for {tool_name}:\n{content}" -def _json_tool_prompt(tools: Sequence[Mapping[str, Any]]) -> str: - return ( - "You have access to the following tools. Your entire response must be a JSON array of one or more tool calls. " - "Each item must be shaped as " - '{"name": "tool_name", "arguments": {"arg_name": "arg_value"}}. ' - "Use only the listed tool names. Do not include prose, markdown, or text outside the JSON array.\n" - f"{json.dumps(list(tools), ensure_ascii=False)}" - ) - - def _new_tool_call_id() -> str: return f"call_{uuid.uuid4().hex[:12]}" diff --git a/nemo_retriever/src/nemo_retriever/query/agentic.py b/nemo_retriever/src/nemo_retriever/query/agentic.py index 93f4d6b57a..fc7c0c2f2e 100644 --- a/nemo_retriever/src/nemo_retriever/query/agentic.py +++ b/nemo_retriever/src/nemo_retriever/query/agentic.py @@ -12,7 +12,6 @@ from __future__ import annotations import logging -import math import threading from dataclasses import dataclass, field from pathlib import Path @@ -24,6 +23,7 @@ from nemo_retriever.models import VL_EMBED_MODEL, VL_RERANK_MODEL from nemo_retriever.query.agentic_options import ( agentic_backend_top_k_error, + agentic_float_range_value, agentic_int_min_error, agentic_int_value, agentic_temperature_error, @@ -149,7 +149,6 @@ class AgenticRetrievalConfig: invoke_url: Optional[str] = None local_llm_backend: str = AGENTIC_LOCAL_LLM_BACKEND local_hf_cache_dir: Optional[str] = None - local_device: Optional[str] = None local_gpu_memory_utilization: float = 0.8 local_tensor_parallel_size: int = 1 local_max_model_len: Optional[int] = None @@ -259,7 +258,7 @@ def __post_init__(self) -> None: raise ValueError(integer_error) object.__setattr__(self, field_name, agentic_int_value(value, field_name=field_name)) - local_gpu_memory_utilization = _agentic_float_range_value( + local_gpu_memory_utilization = agentic_float_range_value( self.local_gpu_memory_utilization, field_name="local_gpu_memory_utilization", min_value=0.0, @@ -286,26 +285,6 @@ def _normalize_agentic_choice(value: object, valid: frozenset[str], *, field_nam return normalized -def _agentic_float_range_value( - value: object, - *, - field_name: str, - min_value: float, - max_value: float, - min_exclusive: bool = False, -) -> float: - try: - parsed = float(value) - except (TypeError, ValueError): - raise ValueError(f"{field_name} must be a number") from None - if not math.isfinite(parsed): - raise ValueError(f"{field_name} must be finite") - if parsed > max_value or (parsed <= min_value if min_exclusive else parsed < min_value): - operator = ">" if min_exclusive else ">=" - raise ValueError(f"{field_name} must be {operator} {min_value} and <= {max_value}") - return parsed - - def _build_agent_chat_completion_fn(cfg: AgenticRetrievalConfig) -> Any | None: if cfg.llm_backend == "openai_compatible": return None @@ -316,7 +295,6 @@ def _build_agent_chat_completion_fn(cfg: AgenticRetrievalConfig) -> Any | None: return create_local_agent_llm( str(cfg.llm_model), backend=str(cfg.local_llm_backend), - device=_none_if_empty(cfg.local_device), hf_cache_dir=_none_if_empty(cfg.local_hf_cache_dir), gpu_memory_utilization=float(cfg.local_gpu_memory_utilization), tensor_parallel_size=int(cfg.local_tensor_parallel_size), diff --git a/nemo_retriever/src/nemo_retriever/query/agentic_options.py b/nemo_retriever/src/nemo_retriever/query/agentic_options.py index e6f0e4727f..5bcf6823af 100644 --- a/nemo_retriever/src/nemo_retriever/query/agentic_options.py +++ b/nemo_retriever/src/nemo_retriever/query/agentic_options.py @@ -57,6 +57,26 @@ def agentic_int_min_error(value: object, *, field_name: str, min_value: int) -> return None +def agentic_float_range_value( + value: object, + *, + field_name: str, + min_value: float, + max_value: float, + min_exclusive: bool = False, +) -> float: + try: + parsed = float(value) + except (TypeError, ValueError): + raise ValueError(f"{field_name} must be a number") from None + if not math.isfinite(parsed): + raise ValueError(f"{field_name} must be finite") + if parsed > max_value or (parsed <= min_value if min_exclusive else parsed < min_value): + operator = ">" if min_exclusive else ">=" + raise ValueError(f"{field_name} must be {operator} {min_value} and <= {max_value}") + return parsed + + def is_nvidia_agentic_endpoint(invoke_url: str | None) -> bool: """Return true for the default/hosted NVIDIA chat-completions endpoint.""" diff --git a/nemo_retriever/src/nemo_retriever/query/options.py b/nemo_retriever/src/nemo_retriever/query/options.py index f87e9cb8e2..f9a0c4696a 100644 --- a/nemo_retriever/src/nemo_retriever/query/options.py +++ b/nemo_retriever/src/nemo_retriever/query/options.py @@ -59,7 +59,6 @@ class QueryAgenticOptions: invoke_url: str | None = None local_llm_backend: str = "vllm" local_hf_cache_dir: str | None = None - local_device: str | None = None local_gpu_memory_utilization: float = 0.8 local_tensor_parallel_size: int = 1 local_max_model_len: int | None = None diff --git a/nemo_retriever/src/nemo_retriever/query/workflow.py b/nemo_retriever/src/nemo_retriever/query/workflow.py index 8d16adfda0..58176c7d02 100644 --- a/nemo_retriever/src/nemo_retriever/query/workflow.py +++ b/nemo_retriever/src/nemo_retriever/query/workflow.py @@ -177,7 +177,6 @@ def build_agentic_config(request: QueryRequest, *, top_k: int | None = None) -> "invoke_url": request.agentic.invoke_url, "local_llm_backend": request.agentic.local_llm_backend, "local_hf_cache_dir": request.agentic.local_hf_cache_dir, - "local_device": request.agentic.local_device, "local_gpu_memory_utilization": request.agentic.local_gpu_memory_utilization, "local_tensor_parallel_size": request.agentic.local_tensor_parallel_size, "local_max_model_len": request.agentic.local_max_model_len, diff --git a/nemo_retriever/tests/test_agentic_eval.py b/nemo_retriever/tests/test_agentic_eval.py index 0c9ca6dc8a..e714904fb6 100644 --- a/nemo_retriever/tests/test_agentic_eval.py +++ b/nemo_retriever/tests/test_agentic_eval.py @@ -189,7 +189,6 @@ def test_agentic_retriever_builds_in_process_llm_lazily(mock_create_local_agent_ mock_create_local_agent_llm.assert_called_once_with( "nemotron-8b", backend="vllm", - device=None, hf_cache_dir=None, gpu_memory_utilization=0.8, tensor_parallel_size=1, diff --git a/nemo_retriever/tests/test_agentic_local_llm.py b/nemo_retriever/tests/test_agentic_local_llm.py index ae18c857f3..10f2456255 100644 --- a/nemo_retriever/tests/test_agentic_local_llm.py +++ b/nemo_retriever/tests/test_agentic_local_llm.py @@ -5,8 +5,6 @@ from __future__ import annotations import json -import os -import sys import threading from unittest.mock import MagicMock @@ -18,13 +16,33 @@ def test_resolve_agent_llm_profile_aliases() -> None: assert resolve_agent_llm_model_name("nemotron-super-49b") == "nvidia/Llama-3_3-Nemotron-Super-49B-v1" +def test_local_agent_llm_config_carries_vllm_resource_options() -> None: + from nemo_retriever.models.local.agent_llm import LocalAgentLLMConfig + + cfg = LocalAgentLLMConfig( + model_path="nemotron-8b", + hf_cache_dir="/tmp/hf", + gpu_memory_utilization=0.7, + tensor_parallel_size=2, + max_model_len=8192, + max_num_seqs=4, + ) + + assert cfg.model_path == "nemotron-8b" + assert cfg.hf_cache_dir == "/tmp/hf" + assert cfg.gpu_memory_utilization == 0.7 + assert cfg.tensor_parallel_size == 2 + assert cfg.max_model_len == 8192 + assert cfg.max_num_seqs == 4 + + def test_vllm_agent_llm_rejects_unsupported_profile_before_vllm_import() -> None: import pytest - from nemo_retriever.models.local.agent_llm import VLLMAgentChatLLM + from nemo_retriever.models.local.agent_llm import LocalAgentLLMConfig, VLLMAgentChatLLM with pytest.raises(ValueError, match="Unsupported local agent LLM model"): - VLLMAgentChatLLM("mistral-7b") + VLLMAgentChatLLM(LocalAgentLLMConfig(model_path="mistral-7b")) def test_parse_json_tool_call_output() -> None: @@ -136,23 +154,6 @@ def test_collapse_parallel_tool_results_for_local_chat_template() -> None: assert "Error: doc_ids must be a list." in collapsed[3]["content"] -def test_local_vllm_device_helper_accepts_cuda_visible_devices() -> None: - from nemo_retriever.models.local.agent_llm import _cuda_visible_devices_from_device - - assert _cuda_visible_devices_from_device(None) is None - assert _cuda_visible_devices_from_device("cuda:5,6") == "5,6" - assert _cuda_visible_devices_from_device("5,6") == "5,6" - - -def test_local_vllm_device_helper_rejects_cpu() -> None: - import pytest - - from nemo_retriever.models.local.agent_llm import _cuda_visible_devices_from_device - - with pytest.raises(ValueError, match="requires CUDA"): - _cuda_visible_devices_from_device("cpu") - - def test_vllm_agent_llm_unload_releases_engine() -> None: import pytest @@ -169,46 +170,6 @@ def test_vllm_agent_llm_unload_releases_engine() -> None: llm._require_loaded() -def test_vllm_agent_llm_unload_restores_cuda_visible_devices(monkeypatch) -> None: - from nemo_retriever.models.local.agent_llm import VLLMAgentChatLLM - - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") - llm = VLLMAgentChatLLM.__new__(VLLMAgentChatLLM) - llm._llm = object() - llm._lock = threading.Lock() - llm._cuda_visible_devices_overridden = True - llm._previous_cuda_visible_devices = "2,3" - - llm.unload() - - assert os.environ["CUDA_VISIBLE_DEVICES"] == "2,3" - - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") - llm = VLLMAgentChatLLM.__new__(VLLMAgentChatLLM) - llm._llm = object() - llm._lock = threading.Lock() - llm._cuda_visible_devices_overridden = True - llm._previous_cuda_visible_devices = None - - llm.unload() - - assert "CUDA_VISIBLE_DEVICES" not in os.environ - - -def test_vllm_agent_llm_restores_cuda_visible_devices_on_init_failure(monkeypatch) -> None: - import pytest - - from nemo_retriever.models.local.agent_llm import VLLMAgentChatLLM - - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "2,3") - monkeypatch.setitem(sys.modules, "vllm", None) - - with pytest.raises(ImportError, match="requires vLLM"): - VLLMAgentChatLLM("nemotron-8b", device="0,1") - - assert os.environ["CUDA_VISIBLE_DEVICES"] == "2,3" - - def test_unload_cached_vllm_agent_chat_llms_clears_cache() -> None: from nemo_retriever.models.local import agent_llm From 39e53966fb937afc2214f966bdd443c1cd58f3c8 Mon Sep 17 00:00:00 2001 From: Mahika Wason Date: Tue, 21 Jul 2026 13:06:23 -0700 Subject: [PATCH 08/10] Harden local agentic tool-call handling --- .../nemo_retriever/models/local/agent_llm.py | 26 ++++++-- .../graph_ops/react_agent_operator.py | 9 +++ .../graph_ops/selection_agent_operator.py | 9 +++ .../tests/test_agentic_local_llm.py | 33 ++++++++++ .../tests/test_agentic_operators.py | 60 +++++++++++++++++++ 5 files changed, 133 insertions(+), 4 deletions(-) diff --git a/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py index bae03ba4ff..03a3f219c2 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py +++ b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py @@ -237,6 +237,8 @@ def _normalize_messages( content = msg.get("content") if isinstance(content, list): msg["content"] = _flatten_text_content(content) + if msg.get("role") == "assistant" and msg.get("tool_calls") and not msg.get("content"): + msg["content"] = _format_assistant_tool_calls_content(msg.get("tool_calls") or []) if tool_prompt and idx == 0 and msg.get("role") == "system": msg["content"] = f"{tool_prompt}\n\n{msg.get('content') or ''}" tool_prompt = None @@ -426,10 +428,6 @@ def _coerce_single_tool_call(item: Any) -> dict[str, Any] | None: def _arguments_to_json_string(arguments: Any) -> str: if isinstance(arguments, str): - try: - json.loads(arguments) - except json.JSONDecodeError: - return json.dumps(arguments) return arguments return json.dumps(arguments if arguments is not None else {}) @@ -516,6 +514,26 @@ def _format_tool_message_content(message: Mapping[str, Any], tool_names_by_id: M return f"Tool result for {tool_name}:\n{content}" +def _format_assistant_tool_calls_content(tool_calls: Sequence[Any]) -> str: + """Render OpenAI assistant tool calls for local templates that read only content.""" + + serializable_calls: list[dict[str, Any]] = [] + for tool_call in tool_calls: + if not isinstance(tool_call, Mapping): + continue + function = tool_call.get("function") or {} + if not isinstance(function, Mapping): + continue + serializable_calls.append( + { + "id": tool_call.get("id"), + "name": function.get("name"), + "arguments": function.get("arguments"), + } + ) + return "Assistant tool calls:\n" + json.dumps(serializable_calls, ensure_ascii=False) + + def _new_tool_call_id() -> str: return f"call_{uuid.uuid4().hex[:12]}" diff --git a/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py b/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py index ac56bba677..4750706432 100644 --- a/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py +++ b/nemo_retriever/src/nemo_retriever/operators/graph_ops/react_agent_operator.py @@ -673,6 +673,15 @@ def _run_single_query( {"role": "tool", "tool_call_id": tc_id, "content": "Error: could not parse tool arguments."} ) continue + if not isinstance(fn_args, dict): + tool_messages.append( + { + "role": "tool", + "tool_call_id": tc_id, + "content": "Error: tool arguments must be a JSON object.", + } + ) + continue if fn_name == "think": # Agent thoughts can quote document text/PII; keep content at DEBUG. diff --git a/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py b/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py index d6d66631ad..eb8ee037f2 100644 --- a/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py +++ b/nemo_retriever/src/nemo_retriever/operators/graph_ops/selection_agent_operator.py @@ -597,6 +597,15 @@ def _select_documents( {"role": "tool", "tool_call_id": tc_id, "content": "Error: could not parse tool arguments."} ) continue + if not isinstance(fn_args, dict): + tool_messages.append( + { + "role": "tool", + "tool_call_id": tc_id, + "content": "Error: tool arguments must be a JSON object.", + } + ) + continue if fn.get("name") == "think": # Agent thoughts can quote document text/PII; keep content at DEBUG diff --git a/nemo_retriever/tests/test_agentic_local_llm.py b/nemo_retriever/tests/test_agentic_local_llm.py index 10f2456255..58bbf8010e 100644 --- a/nemo_retriever/tests/test_agentic_local_llm.py +++ b/nemo_retriever/tests/test_agentic_local_llm.py @@ -154,6 +154,39 @@ def test_collapse_parallel_tool_results_for_local_chat_template() -> None: assert "Error: doc_ids must be a list." in collapsed[3]["content"] +def test_normalize_messages_serializes_assistant_tool_calls_for_local_templates() -> None: + from nemo_retriever.models.local.agent_llm import VLLMAgentChatLLM + + llm = VLLMAgentChatLLM.__new__(VLLMAgentChatLLM) + + normalized = llm._normalize_messages( + [ + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "retrieve", "arguments": '{"query": "inflation"}'}, + } + ], + } + ] + ) + + assert normalized[0]["tool_calls"][0]["function"]["name"] == "retrieve" + assert normalized[0]["content"].startswith("Assistant tool calls:") + assert "inflation" in normalized[0]["content"] + + +def test_malformed_string_tool_arguments_remain_malformed() -> None: + from nemo_retriever.models.local.agent_llm import parse_tool_calls_from_text + + calls = parse_tool_calls_from_text(json.dumps([{"name": "retrieve", "arguments": "query=inflation"}])) + + assert calls[0]["function"]["arguments"] == "query=inflation" + + def test_vllm_agent_llm_unload_releases_engine() -> None: import pytest diff --git a/nemo_retriever/tests/test_agentic_operators.py b/nemo_retriever/tests/test_agentic_operators.py index fda6853608..c30e7986cf 100644 --- a/nemo_retriever/tests/test_agentic_operators.py +++ b/nemo_retriever/tests/test_agentic_operators.py @@ -164,6 +164,27 @@ def _make_tool_call_response(fn_name: str, fn_args: dict, tc_id: str = "call_1") } +def _make_raw_arguments_tool_call_response(fn_name: str, arguments: str, tc_id: str = "call_1") -> dict: + """Build a canned chat-completions response with raw function arguments.""" + return { + "choices": [ + { + "message": { + "content": None, + "tool_calls": [ + { + "id": tc_id, + "type": "function", + "function": {"name": fn_name, "arguments": arguments}, + } + ], + }, + "finish_reason": "tool_calls", + } + ] + } + + class TestSelectionAgentOperator: def _make_input(self): return pd.DataFrame( @@ -225,6 +246,25 @@ def test_retries_when_selection_returns_invalid_doc_ids(self, mock_step): assert result["doc_id"].tolist() == ["d1", "d2"] assert result["message"].tolist() == ["corrected", "corrected"] + @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") + def test_non_object_tool_arguments_are_reported_and_fall_back(self, mock_step): + from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator + + mock_step.return_value = _make_raw_arguments_tool_call_response( + "log_selected_documents", json.dumps("doc_ids=d1") + ) + + op = SelectionAgentOperator( + llm_model="test-model", + invoke_url="http://localhost/v1/chat/completions", + top_k=2, + max_steps=1, + ) + result = op.run(self._make_input()) + + assert result["doc_id"].tolist() == ["d1", "d2"] + assert result["result_source"].tolist() == ["candidate_ranking", "candidate_ranking"] + @patch("nemo_retriever.operators.graph_ops.selection_agent_operator.invoke_chat_completion_step") def test_think_then_select(self, mock_step): from nemo_retriever.operators.graph_ops.selection_agent_operator import SelectionAgentOperator @@ -450,6 +490,26 @@ def test_injected_chat_completion_fn_replaces_http_call(self, mock_step): assert local_chat.call_args.kwargs["max_tokens"] == 123 assert result[result["doc_id"] == "d1"]["is_final_result"].astype(bool).any() + @patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") + def test_non_object_tool_arguments_are_reported_without_crashing(self, mock_step): + from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator + + mock_step.return_value = _make_raw_arguments_tool_call_response("retrieve", json.dumps("query=inflation")) + retriever = MagicMock(return_value=[{"doc_id": "d1", "text": "monetary policy"}]) + + op = ReActAgentOperator( + invoke_url="http://localhost/v1/chat/completions", + llm_model="test-model", + retriever_fn=retriever, + user_msg_type="with_results", + target_top_k=1, + max_steps=1, + ) + result = op.run(self._make_input()) + + assert result["doc_id"].tolist() == ["d1"] + assert not result["is_final_result"].astype(bool).any() + @patch("nemo_retriever.operators.graph_ops.react_agent_operator.invoke_chat_completion_step") def test_with_results_mode_initial_retrieval(self, mock_step): from nemo_retriever.operators.graph_ops.react_agent_operator import ReActAgentOperator From ac95fbba431b891e3bc4517e290d5c5e4898987f Mon Sep 17 00:00:00 2001 From: Mahika Wason Date: Tue, 21 Jul 2026 13:50:24 -0700 Subject: [PATCH 09/10] document local agentic vllm usage example in readme --- nemo_retriever/README.md | 53 +++++++++++++++++++--------------------- 1 file changed, 25 insertions(+), 28 deletions(-) diff --git a/nemo_retriever/README.md b/nemo_retriever/README.md index 857581ecb4..9bd1c6f825 100644 --- a/nemo_retriever/README.md +++ b/nemo_retriever/README.md @@ -315,31 +315,29 @@ Agentic retrieval runs an LLM-driven ReAct loop over an existing LanceDB index. It does not ingest documents. Build the index with one of the ingestion flows above, then query the same `lancedb_uri`, `table_name`, and embedding model. -The agent LLM supports hosted NVIDIA inference or a local OpenAI-compatible -chat-completions endpoint. Pass `--agentic-invoke-url` for a local vLLM or -self-hosted NIM server; when omitted, the agent uses the built-in NVIDIA hosted -endpoint. Embedding follows the same local/remote split as dense retrieval: -CPU-only hosts default to hosted endpoints; GPU-capable hosts use local GPU -embedding unless you pass `--embed-invoke-url`, for example -`--embed-invoke-url https://integrate.api.nvidia.com/v1/embeddings`. - -**Hosted inference.** For [build.nvidia.com](https://build.nvidia.com/) hosted -inference, set `NVIDIA_API_KEY`. On CPU-only machines, the CPU embedding actor -and agent LLM use the hosted NVIDIA endpoints by default: +By default, the agent LLM runs in process with local vLLM and `nemotron-8b` +(`nvidia/Llama-3.1-Nemotron-Nano-8B-v1`). This requires a CUDA GPU host and the +`[local]` extra. GPU placement follows process-level vLLM behavior, so set +`CUDA_VISIBLE_DEVICES` before starting the command. Embedding follows the same +local/remote split as dense retrieval; pass `--embed-invoke-url` to use a remote +embedding endpoint. -```bash -export NVIDIA_API_KEY=nvapi-... +**Local in-process vLLM agent LLM.** Omit `--agentic-invoke-url` to load the +supported local agent LLM directly in the Python process. `nemotron-8b` is the +default; `super-49b` is also supported when the process has enough visible GPUs. -retriever query "Given their activities, which animal is responsible for the typos in my documents?" \ +```bash +CUDA_VISIBLE_DEVICES=0 retriever query "Given their activities, which animal is responsible for the typos in my documents?" \ --agentic \ - --agentic-llm-model nvidia/llama-3.3-nemotron-super-49b-v1.5 \ --lancedb-uri lancedb \ --table-name nemo-retriever \ --embed-model-name nvidia/llama-nemotron-embed-1b-v2 ``` -**Local agent LLM.** Point `--agentic-invoke-url` at your OpenAI-compatible -chat-completions server (for example vLLM or a self-hosted NIM): +**OpenAI-compatible agent endpoint.** Pass `--agentic-invoke-url` when you want a +custom model or a separately hosted chat-completions server, such as vLLM server +mode or a self-hosted NIM. When an invoke URL is provided, `--agentic-llm-model` +is required and is sent as the remote model ID. ```bash retriever query "What is RAG?" \ @@ -351,18 +349,14 @@ retriever query "What is RAG?" \ --embed-model-name nvidia/llama-nemotron-embed-1b-v2 ``` -Native local Hugging Face LLM inference for the agent is in active development; -this README will be updated when that path ships. - Unlike dense retrieval, agentic mode returns ranked document IDs as JSON, not text-enriched hits. For a quick smoke test, reduce agent work: ```bash -retriever query "What is RAG?" \ +CUDA_VISIBLE_DEVICES=0 retriever query "What is RAG?" \ --agentic \ - --agentic-llm-model nvidia/llama-3.3-nemotron-super-49b-v1.5 \ --lancedb-uri lancedb \ --table-name nemo-retriever \ --embed-model-name nvidia/llama-nemotron-embed-1b-v2 \ @@ -371,9 +365,9 @@ retriever query "What is RAG?" \ --agentic-backend-top-k 1 ``` -You can run the same flow from Python. Set `NVIDIA_API_KEY` for hosted inference, -or pass `invoke_url` on `QueryAgenticOptions` for a local chat-completions -endpoint. +You can run the same flow from Python. Omit `invoke_url` for the default local +in-process vLLM backend, or pass `invoke_url` on `QueryAgenticOptions` for a +separate OpenAI-compatible chat-completions endpoint. ```python from nemo_retriever.cli.query_workflow import agentic_query_documents @@ -385,8 +379,8 @@ from nemo_retriever.query.options import ( QueryStorageOptions, ) -# Hosted: set NVIDIA_API_KEY=nvapi-... in the environment. -# Local agent LLM: pass invoke_url="http://localhost:9000/v1/chat/completions". +# Local in-process vLLM: set CUDA_VISIBLE_DEVICES before starting Python. +# Remote agent LLM: pass invoke_url="http://localhost:9000/v1/chat/completions". results = agentic_query_documents( QueryRequest( query="What is RAG?", @@ -400,7 +394,10 @@ results = agentic_query_documents( ), agentic=QueryAgenticOptions( enabled=True, - llm_model="nvidia/llama-3.3-nemotron-super-49b-v1.5", + llm_model="nemotron-8b", + local_llm_backend="vllm", + local_gpu_memory_utilization=0.8, + local_tensor_parallel_size=1, ), ) ) From e6f8d521fb1ea6385a964d3c6d9cfc65e6b536e5 Mon Sep 17 00:00:00 2001 From: Mahika Wason Date: Tue, 21 Jul 2026 16:14:28 -0700 Subject: [PATCH 10/10] Fix in-process agent LLM lifecycle and EngineCore teardown --- .../src/nemo_retriever/models/__init__.py | 10 +- .../nemo_retriever/models/local/agent_llm.py | 84 ++++++++++------ .../src/nemo_retriever/query/agentic.py | 38 +++++++- .../src/nemo_retriever/query/workflow.py | 36 ++++--- .../tests/test_agentic_local_llm.py | 97 +++++++++++++++++-- nemo_retriever/tests/test_root_query_cli.py | 9 ++ 6 files changed, 208 insertions(+), 66 deletions(-) diff --git a/nemo_retriever/src/nemo_retriever/models/__init__.py b/nemo_retriever/src/nemo_retriever/models/__init__.py index 5c75ed8445..c83d3ac603 100644 --- a/nemo_retriever/src/nemo_retriever/models/__init__.py +++ b/nemo_retriever/src/nemo_retriever/models/__init__.py @@ -273,19 +273,21 @@ def create_local_agent_llm( max_model_len: int | None = None, max_num_seqs: int | None = None, ) -> Any: - """Create a cached local agent LLM chat-completion callable. + """Create a local agent LLM chat-completion callable owned by the caller. The callable mirrors ``invoke_chat_completion_step`` and returns an OpenAI-compatible chat-completions response dict. V1 supports the in-process vLLM backend and uses process-level vLLM placement (for example, - ``CUDA_VISIBLE_DEVICES`` plus ``tensor_parallel_size``). + ``CUDA_VISIBLE_DEVICES`` plus ``tensor_parallel_size``). Callers (typically + ``AgenticRetriever``) should reuse one instance for a harness/CLI job and + call ``unload()`` when the job finishes. """ b = normalize_backend(backend, _LOCAL_AGENT_LLM_BACKENDS, field_name="backend", default="vllm") if b == "vllm": - from nemo_retriever.models.local.agent_llm import LocalAgentLLMConfig, create_cached_vllm_agent_chat_llm + from nemo_retriever.models.local.agent_llm import LocalAgentLLMConfig, create_vllm_agent_chat_llm - return create_cached_vllm_agent_chat_llm( + return create_vllm_agent_chat_llm( LocalAgentLLMConfig( model_path=model_name, hf_cache_dir=hf_cache_dir, diff --git a/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py index 03a3f219c2..b539751adb 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py +++ b/nemo_retriever/src/nemo_retriever/models/local/agent_llm.py @@ -21,8 +21,8 @@ # Conservative cap for short tool-call JSON responses. Larger evals should # tune this with completion-token, truncation, and tool-parse telemetry. _DEFAULT_MAX_TOKENS = 512 -_LOCAL_AGENT_LLM_CACHE: dict[tuple[Any, ...], "VLLMAgentChatLLM"] = {} -_LOCAL_AGENT_LLM_CACHE_LOCK = threading.Lock() +# Bound EngineCore join so a wedged child cannot hang process exit forever. +_VLLM_ENGINE_SHUTDOWN_TIMEOUT_S = 30.0 def _normalize_name(name: str) -> str: @@ -113,12 +113,14 @@ def __init__(self, config: LocalAgentLLMConfig) -> None: ) try: + from nemo_retriever.models.inference.vllm import apply_vllm_startup_defaults from vllm import LLM, SamplingParams except ImportError as e: raise ImportError( 'Local agentic LLM inference requires vLLM. Install with: pip install "nemo-retriever[local]"' ) from e + apply_vllm_startup_defaults() _raise_if_cuda_unavailable() self._model_path = model_path self._max_tokens = int(config.max_tokens) @@ -248,17 +250,27 @@ def _normalize_messages( return _collapse_consecutive_tool_messages(normalized) def unload(self) -> None: - """Release GPU memory held by the local vLLM engine.""" + """Shut down the EngineCore child and release GPU memory. + + Unlike embed/rerank holders that only ``del`` the ``LLM`` and rely on + process exit, the agent chat path must call ``engine_core.shutdown()``: + vLLM V1 generation spawns a multiprocess ``VLLM::EngineCore`` that can + keep the parent process alive after results are returned. + """ with self._lock: - had_llm = self._llm is not None - if had_llm: - del self._llm - self._llm = None + llm = self._llm + self._llm = None - if not had_llm: + if llm is None: return + _shutdown_vllm_engine(llm) + try: + del llm + except Exception: + logger.debug("Ignoring error while deleting local agent vLLM LLM", exc_info=True) + try: import torch except ImportError: @@ -309,35 +321,43 @@ def _raise_if_cuda_unavailable() -> None: ) -def create_cached_vllm_agent_chat_llm(config: LocalAgentLLMConfig) -> VLLMAgentChatLLM: - """Create or reuse a local agent LLM keyed by heavyweight load settings.""" +def _call_shutdown(shutdown: Any, *, timeout_s: float) -> None: + try: + shutdown(timeout=timeout_s) + except TypeError: + shutdown() - key = ( - resolve_agent_llm_model_name(config.model_path), - config.hf_cache_dir, - float(config.gpu_memory_utilization), - int(config.tensor_parallel_size), - int(config.max_model_len) if config.max_model_len is not None else None, - int(config.max_num_seqs) if config.max_num_seqs is not None else None, - int(config.max_tokens), - ) - with _LOCAL_AGENT_LLM_CACHE_LOCK: - cached = _LOCAL_AGENT_LLM_CACHE.get(key) - if cached is None or cached._llm is None: - cached = VLLMAgentChatLLM(config) - _LOCAL_AGENT_LLM_CACHE[key] = cached - return cached +def _shutdown_vllm_engine(llm: Any) -> None: + """Best-effort vLLM V1 EngineCore / executor shutdown before dropping the LLM.""" + + try: + engine = getattr(llm, "llm_engine", None) + if engine is None: + return + engine_core = getattr(engine, "engine_core", None) + shutdown = getattr(engine_core, "shutdown", None) if engine_core is not None else None + if callable(shutdown): + _call_shutdown(shutdown, timeout_s=_VLLM_ENGINE_SHUTDOWN_TIMEOUT_S) + return + shutdown = getattr(engine, "shutdown", None) + if callable(shutdown): + _call_shutdown(shutdown, timeout_s=_VLLM_ENGINE_SHUTDOWN_TIMEOUT_S) + except Exception: + logger.warning( + "Local agent vLLM engine shutdown failed; dropping the LLM reference anyway", + exc_info=True, + ) -def unload_cached_vllm_agent_chat_llms() -> None: - """Unload and forget all cached local agent vLLM engines.""" - with _LOCAL_AGENT_LLM_CACHE_LOCK: - cached_models = list(_LOCAL_AGENT_LLM_CACHE.values()) - _LOCAL_AGENT_LLM_CACHE.clear() +def create_vllm_agent_chat_llm(config: LocalAgentLLMConfig) -> VLLMAgentChatLLM: + """Create a local agent LLM owned by the caller (no process-global cache). + + Mirrors embed/rerank: the harness/CLI job holds one instance for the whole + run and calls ``unload()`` when the job finishes. + """ - for model in cached_models: - model.unload() + return VLLMAgentChatLLM(config) def parse_tool_calls_from_text(text: str) -> list[dict[str, Any]]: diff --git a/nemo_retriever/src/nemo_retriever/query/agentic.py b/nemo_retriever/src/nemo_retriever/query/agentic.py index fc7c0c2f2e..dfd29ee635 100644 --- a/nemo_retriever/src/nemo_retriever/query/agentic.py +++ b/nemo_retriever/src/nemo_retriever/query/agentic.py @@ -354,9 +354,29 @@ def _get_chat_completion_fn(self) -> Any | None: return None with self._lock: if self._chat_completion_fn is None: + # Instance-owned, same pattern as Retriever's cached embed/rerank + # graph: load once per AgenticRetriever and reuse across queries. self._chat_completion_fn = _build_agent_chat_completion_fn(self._cfg) return self._chat_completion_fn + def unload(self) -> None: + """Release the in-process agent LLM owned by this retriever. + + OpenAI-compatible endpoint mode is a no-op. Local vLLM mode shuts down + this instance's EngineCore so CLI/harness jobs can exit cleanly. Embed + and rerank models stay on ``self._retriever`` and are released with the + process, matching dense harness BEIR behavior. + """ + + with self._lock: + chat_fn = self._chat_completion_fn + self._chat_completion_fn = None + if chat_fn is None: + return + unload = getattr(chat_fn, "unload", None) + if callable(unload): + unload() + def retrieve(self, query_ids: Sequence[str], query_texts: Sequence[str]) -> pd.DataFrame: """Return selected ranked documents for each query. @@ -596,7 +616,11 @@ def run_agentic_audio_recall_evaluation( queries = df_query["query"].astype(str).tolist() gold_doc_ids = df_query["golden_answer"].astype(str).tolist() - result = AgenticRetriever(cfg, match_mode="audio_segment").retrieve(query_ids, queries) + retriever = AgenticRetriever(cfg, match_mode="audio_segment") + try: + result = retriever.retrieve(query_ids, queries) + finally: + retriever.unload() retrieved_doc_ids = _agentic_result_to_ranked_doc_ids(query_ids, result) ks_sorted = sorted({int(k) for k in ks if int(k) > 0}) if not ks_sorted: @@ -656,9 +680,13 @@ def agentic_beir_retrieve( that already hold a loaded dataset (e.g. the harness) reuse the agent's retrieve+rank step without re-loading or re-implementing it. """ - result = AgenticRetriever(cfg, match_mode="pdf_page", doc_id_field=doc_id_field).retrieve( - list(dataset.query_ids), - list(dataset.queries), - ) + retriever = AgenticRetriever(cfg, match_mode="pdf_page", doc_id_field=doc_id_field) + try: + result = retriever.retrieve( + list(dataset.query_ids), + list(dataset.queries), + ) + finally: + retriever.unload() ranked_doc_ids = _agentic_result_to_ranked_doc_ids(list(dataset.query_ids), result) return result, ranked_doc_ids diff --git a/nemo_retriever/src/nemo_retriever/query/workflow.py b/nemo_retriever/src/nemo_retriever/query/workflow.py index 58176c7d02..e7d9bb21db 100644 --- a/nemo_retriever/src/nemo_retriever/query/workflow.py +++ b/nemo_retriever/src/nemo_retriever/query/workflow.py @@ -227,18 +227,24 @@ def agentic_query_documents(request: QueryRequest) -> list[dict[str, Any]]: through to the wrapped ``Retriever`` that backs the agent's ``retrieve`` tool. Reranking therefore applies per agent retrieval hop. """ - result = build_agentic_retriever(request).retrieve(["0"], [str(request.query)]) - if "rank" in result.columns: - result = result.sort_values("rank") - ranked: list[dict[str, Any]] = [] - for _, row in result.iterrows(): - ranked.append( - { - "rank": int(row.get("rank", len(ranked) + 1)), - "doc_id": str(row.get("doc_id", "")), - "result_source": str(row.get("result_source", "")), - } - ) - if len(ranked) >= request.retrieval.top_k: - break - return ranked + retriever = build_agentic_retriever(request) + try: + result = retriever.retrieve(["0"], [str(request.query)]) + if "rank" in result.columns: + result = result.sort_values("rank") + ranked: list[dict[str, Any]] = [] + for _, row in result.iterrows(): + ranked.append( + { + "rank": int(row.get("rank", len(ranked) + 1)), + "doc_id": str(row.get("doc_id", "")), + "result_source": str(row.get("result_source", "")), + } + ) + if len(ranked) >= request.retrieval.top_k: + break + return ranked + finally: + # One-shot CLI/Python entry: tear down cached local vLLM EngineCore so the + # process can exit instead of hanging on a live child worker. + retriever.unload() diff --git a/nemo_retriever/tests/test_agentic_local_llm.py b/nemo_retriever/tests/test_agentic_local_llm.py index 58bbf8010e..308f8b5abf 100644 --- a/nemo_retriever/tests/test_agentic_local_llm.py +++ b/nemo_retriever/tests/test_agentic_local_llm.py @@ -6,6 +6,7 @@ import json import threading +from typing import Any from unittest.mock import MagicMock @@ -192,26 +193,102 @@ def test_vllm_agent_llm_unload_releases_engine() -> None: from nemo_retriever.models.local.agent_llm import VLLMAgentChatLLM + engine_core = MagicMock() + llm_engine = MagicMock() + llm_engine.engine_core = engine_core + vllm_llm = MagicMock() + vllm_llm.llm_engine = llm_engine + llm = VLLMAgentChatLLM.__new__(VLLMAgentChatLLM) - llm._llm = object() + llm._llm = vllm_llm llm._lock = threading.Lock() llm.unload() + engine_core.shutdown.assert_called_once_with(timeout=30.0) assert llm._llm is None with pytest.raises(RuntimeError, match="unloaded"): llm._require_loaded() -def test_unload_cached_vllm_agent_chat_llms_clears_cache() -> None: - from nemo_retriever.models.local import agent_llm +def test_vllm_agent_llm_unload_falls_back_when_shutdown_rejects_timeout() -> None: + from nemo_retriever.models.local.agent_llm import VLLMAgentChatLLM + + calls: list[dict[str, Any]] = [] + + def shutdown(**kwargs: Any) -> None: + calls.append(kwargs) + if "timeout" in kwargs: + raise TypeError("shutdown() got an unexpected keyword argument 'timeout'") + + engine_core = MagicMock() + engine_core.shutdown.side_effect = shutdown + llm_engine = MagicMock() + llm_engine.engine_core = engine_core + vllm_llm = MagicMock() + vllm_llm.llm_engine = llm_engine + + llm = VLLMAgentChatLLM.__new__(VLLMAgentChatLLM) + llm._llm = vllm_llm + llm._lock = threading.Lock() + + llm.unload() + + assert calls == [{"timeout": 30.0}, {}] + assert llm._llm is None + + +def test_create_vllm_agent_chat_llm_returns_fresh_instance() -> None: + from nemo_retriever.models.local import agent_llm as agent_llm_mod + from nemo_retriever.models.local.agent_llm import LocalAgentLLMConfig, create_vllm_agent_chat_llm + + created: list[Any] = [] + + def fake_ctor(config: LocalAgentLLMConfig) -> Any: + obj = MagicMock(name="VLLMAgentChatLLM") + obj.config = config + created.append(obj) + return obj + + original = agent_llm_mod.VLLMAgentChatLLM + agent_llm_mod.VLLMAgentChatLLM = fake_ctor # type: ignore[misc,assignment] + try: + first = create_vllm_agent_chat_llm(LocalAgentLLMConfig(model_path="nemotron-8b")) + second = create_vllm_agent_chat_llm(LocalAgentLLMConfig(model_path="nemotron-8b")) + finally: + agent_llm_mod.VLLMAgentChatLLM = original + + assert first is not second + assert len(created) == 2 + + +def test_agentic_retriever_unload_releases_owned_llm() -> None: + from unittest.mock import patch - cached = MagicMock() - with agent_llm._LOCAL_AGENT_LLM_CACHE_LOCK: - agent_llm._LOCAL_AGENT_LLM_CACHE.clear() - agent_llm._LOCAL_AGENT_LLM_CACHE[("model",)] = cached + from nemo_retriever.query.agentic import AgenticRetrievalConfig, AgenticRetriever - agent_llm.unload_cached_vllm_agent_chat_llms() + owned = MagicMock() + with patch("nemo_retriever.query.agentic.Retriever"): + retriever = AgenticRetriever(AgenticRetrievalConfig(llm_model="nemotron-8b")) + retriever._chat_completion_fn = owned + retriever.unload() + + assert retriever._chat_completion_fn is None + owned.unload.assert_called_once_with() + + +def test_agentic_retriever_unload_noop_when_no_local_llm() -> None: + from unittest.mock import patch + + from nemo_retriever.query.agentic import AgenticRetrievalConfig, AgenticRetriever + + with patch("nemo_retriever.query.agentic.Retriever"): + retriever = AgenticRetriever( + AgenticRetrievalConfig( + llm_model="remote-model", + invoke_url="http://localhost/v1/chat/completions", + ) + ) + retriever.unload() - cached.unload.assert_called_once_with() - assert agent_llm._LOCAL_AGENT_LLM_CACHE == {} + assert retriever._chat_completion_fn is None diff --git a/nemo_retriever/tests/test_root_query_cli.py b/nemo_retriever/tests/test_root_query_cli.py index 16c42c52b1..f03c9e824b 100644 --- a/nemo_retriever/tests/test_root_query_cli.py +++ b/nemo_retriever/tests/test_root_query_cli.py @@ -368,6 +368,9 @@ def retrieve(self, query_ids: Any, query_texts: Any) -> Any: ] ) + def unload(self) -> None: + return None + monkeypatch.setattr(agentic_retrieval, "AgenticRetrievalConfig", FakeConfig) monkeypatch.setattr(agentic_retrieval, "AgenticRetriever", FakeAgenticRetriever) @@ -446,6 +449,9 @@ def __init__(self, cfg: Any) -> None: 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"}]) + def unload(self) -> None: + return None + monkeypatch.setattr(agentic_retrieval, "AgenticRetrievalConfig", FakeConfig) monkeypatch.setattr(agentic_retrieval, "AgenticRetriever", FakeAgenticRetriever) @@ -488,6 +494,9 @@ def __init__(self, cfg: Any) -> None: 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"}]) + def unload(self) -> None: + return None + monkeypatch.setattr(agentic_retrieval, "AgenticRetrievalConfig", FakeConfig) monkeypatch.setattr(agentic_retrieval, "AgenticRetriever", FakeAgenticRetriever)