diff --git a/AGENTS.md b/AGENTS.md index 5feebb6cc..df901d05c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -17,7 +17,7 @@ agent-sdk/ # Five packages: # xr-ai-models — LLM/VLM/STT/TTS service protocols + OpenAI-compat clients # xr-ai-pipecat — optional Pipecat transport bridge (heavier deps) # xr-ai-voice — voice runtime (VoiceSession); introduced alongside xr-ai-pipecat - # xr-ai-nat — typed, in-process NAT functions for XR capabilities + # xr-ai-nat — native Relay-managed tools; legacy NAT compatibility during migration utils/ # Shared infra: launcher, logging, vad, vllm, voicegate services/ # XR hub, CloudXR, model-serving, and typed capability services agent-mcp-servers/ # MCP adapters: oxr, render, transcript, vec, video, vlm @@ -60,9 +60,12 @@ deps/ # Gitignored downloaded binaries (e.g. LOVR AppImage) - **Workers never import from `xr_media_hub` or `xr_ai_launcher`.** Use the public `xr_ai_hub`, `xr_ai_models`, `xr_ai_nat`, and `xr_ai_voice` SDK surfaces plus task-specific libraries (numpy, torch, …). -- **Agentic functions are NAT-first and in-process.** Reusable deterministic - functions live in `xr-ai-nat` as typed NAT function groups. Existing MCP - servers remain compatibility surfaces while their capabilities migrate. +- **Agentic functions are native and in-process.** New and migrated tools live + in `xr-ai-nat`; every tool and tool-driven agent lifecycle passes through + NeMo Relay, and all model I/O remains in `xr-ai-models`. Its legacy extras + retain NeMo Agent Toolkit compatibility only while existing function groups + migrate. Existing MCP servers remain compatibility surfaces while their + capabilities migrate. - **RAG is a native typed capability.** `rag-service` owns document chunking, embedding caches, and dense retrieval behind private msgpack/ZMQ; `RAGFunctionsConfig` exposes it as the `xr_rag` NAT function group. diff --git a/DEPENDENCIES.md b/DEPENDENCIES.md index fd7960d3f..adc263500 100644 --- a/DEPENDENCIES.md +++ b/DEPENDENCIES.md @@ -107,7 +107,9 @@ xr-ai-models (agent-sdk/xr-ai-models/) and OpenAI-compatible HTTP clients that cover every in-tree model backend (vLLM-served VLM/LLMs, NeMo Parakeet STT, Piper/Magpie TTS). Per-model profiles separate adapter behavior, endpoint connectivity/readiness, and - launcher-facing deployment ownership. Per-model quirks remain behind one + launcher-facing deployment ownership. Relay may pass controlled per-call + context headers; configured model credentials remain non-overridable. + Per-model quirks remain behind one seam: reasoning-field aliasing (nano_v3 → `reasoning`, nemotron_v3 → `reasoning_content`), `chat_template_kwargs` plumbing for `enable_thinking` / `thinking_budget`, and built-in presets @@ -118,15 +120,23 @@ xr-ai-models (agent-sdk/xr-ai-models/) metadata while the existing flat YAML schema remains valid. xr-ai-nat (agent-sdk/xr-ai-nat/) - └── nvidia-nat-core ==1.8.0 + └── nemo-relay >=0.7.2,<0.8 └── pydantic >=2.10 - └── [agents] nvidia-nat-langchain ==1.8.0, xr-ai-models [editable: ../xr-ai-models] - └── [mcp] fastmcp >=3.4,<4 - └── [services] msgpack >=1.0, pyzmq >=27.0 - └── [vision] httpx >=0.27, numpy >=1.24, Pillow >=10.0, xr-ai-hub-client [editable: ../xr-ai-hub-client], xr-ai-models [editable: ../xr-ai-models] - └── [voice] xr-ai-voice [editable: ../xr-ai-voice] - Typed, in-process NeMo Agent Toolkit functions for XR capabilities. The - ``xr_spatial_math`` function group accepts explicit coordinate frames and + └── [relay] xr-ai-models [editable: ../xr-ai-models] + └── [live-vision] numpy >=1.24, Pillow >=10.0, xr-ai-hub-client [editable: ../xr-ai-hub-client], xr-ai-models [editable: ../xr-ai-models] + └── [agents] nvidia-nat-core ==1.8.0, nvidia-nat-langchain ==1.8.0, xr-ai-models [editable: ../xr-ai-models] + └── [mcp] nvidia-nat-core ==1.8.0, fastmcp >=3.4,<4 + └── [services] nvidia-nat-core ==1.8.0, msgpack >=1.0, pyzmq >=27.0 + └── [vision] nvidia-nat-core ==1.8.0, httpx >=0.27, numpy >=1.24, Pillow >=10.0, xr-ai-hub-client [editable: ../xr-ai-hub-client], xr-ai-models [editable: ../xr-ai-models] + └── [voice] nvidia-nat-core ==1.8.0, xr-ai-voice [editable: ../xr-ai-voice] + The base package is the toolkit-independent native tools layer: Pydantic + request and response models, Relay-managed execution, the generic + ``AgentRunner`` protocol, and a bounded default tool loop over + `xr-ai-models`. The ``[relay]`` and + ``[live-vision]`` extras add model-backed tools without selecting NeMo Agent + Toolkit. The existing function groups remain behind legacy compatibility + extras while they migrate. The ``xr_spatial_math`` function group accepts + explicit coordinate frames and performs deterministic spatial calculations without OpenXR, model, or MCP dependencies. ``xr_text_memory`` owns persistent per-source JSONL text history, and ``xr_conversation_memory`` composes it into a participant- @@ -324,7 +334,7 @@ vec-mcp-server (agent-mcp-servers/vec-mcp/) xr-ai-tests (tests/) └── xr-ai-hub-client [editable: ../agent-sdk/xr-ai-hub-client] └── xr-ai-models [editable: ../agent-sdk/xr-ai-models] - └── xr-ai-nat[agents,services,vision] [editable: ../agent-sdk/xr-ai-nat] + └── xr-ai-nat[agents,relay,services,vision] [editable: ../agent-sdk/xr-ai-nat] └── xr-rag-service [editable: ../services/rag-service] └── xr-ai-pipecat [editable: ../agent-sdk/xr-ai-pipecat] └── xr-ai-voice [editable: ../agent-sdk/xr-ai-voice] diff --git a/README.md b/README.md index 691ebfc11..288fae02e 100644 --- a/README.md +++ b/README.md @@ -117,7 +117,7 @@ frames are dropped if it is closed. | Hub service | `services/xr-media-hub/` | XR-Media-Hub + LiveKit internal transport | | Launcher | `utils/xr-ai-launcher/` | stdlib-only process manager used by samples | | Logging | `utils/xr-ai-logging/` | shared loguru sink + stdlib bridge for every process | -| Agent functions | `agent-sdk/xr-ai-nat/` | Typed, in-process NAT functions for XR capabilities | +| Agent tools | `agent-sdk/xr-ai-nat/` | Relay-managed native tools and legacy NAT compatibility during migration | | Reusable services | `services/` | Model-serving and typed capability processes | | Agent interfaces | `agent-mcp-servers/` | MCP compatibility processes for XR data & rendering | | Agent demos | `agent-samples/` | End-to-end agent pipelines | diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 4bf700dc5..6e7af716f 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -30,6 +30,7 @@ For the per-package dependency mapping, see [`DEPENDENCIES.md`](DEPENDENCIES.md) | `livekit-api` | 0.7.0 | Apache-2.0 | https://github.com/livekit/python-sdks | | `numpy` | 1.24.0 | BSD-3-Clause | https://github.com/numpy/numpy | | `nvidia-nat-core` | 1.8.0 | Apache-2.0 | https://github.com/NVIDIA/NeMo-Agent-Toolkit | +| `nemo-relay` | >=0.7.2,<0.8 | Apache-2.0 | https://github.com/NVIDIA/NeMo-Relay | | `Pillow` | 10.0.0 | HPND | https://github.com/python-pillow/Pillow | | `pydantic` | >=2.10 | MIT | https://github.com/pydantic/pydantic | | `websockets` | 12.0 | BSD-3-Clause | https://github.com/python-websockets/websockets | diff --git a/agent-sdk/xr-ai-models/README.md b/agent-sdk/xr-ai-models/README.md index faf1f47d4..9b634f289 100644 --- a/agent-sdk/xr-ai-models/README.md +++ b/agent-sdk/xr-ai-models/README.md @@ -140,7 +140,8 @@ class LLMService(Protocol): capabilities: Capabilities async def chat(self, messages, *, tools=None, max_tokens=None, temperature=None, enable_thinking=False, - thinking_budget=None, timeout=None) -> ChatResponse: ... + thinking_budget=None, timeout=None, + headers=None) -> ChatResponse: ... def stream(self, messages, *, ...) -> AsyncIterator[str]: ... async def health(self) -> bool: ... async def close(self) -> None: ... @@ -174,6 +175,11 @@ class EmbeddingService(Protocol): `reasoning_field` knob normalizes `reasoning_content` (nemotron_v3 parser) into the same surface. +`LLMService.chat` and `LLMService.stream` accept optional string-valued +per-call headers for execution context such as Relay session lineage. The +model profile remains the authority for credentials: callers cannot supply an +`Authorization` header. + ## Remote / hosted-NIM endpoints Cloud / remote endpoints (e.g. hosted [NVIDIA NIM](https://build.nvidia.com)) diff --git a/agent-sdk/xr-ai-models/xr_ai_models/_openai_compat.py b/agent-sdk/xr-ai-models/xr_ai_models/_openai_compat.py index be981cc1e..d066308ae 100644 --- a/agent-sdk/xr-ai-models/xr_ai_models/_openai_compat.py +++ b/agent-sdk/xr-ai-models/xr_ai_models/_openai_compat.py @@ -16,7 +16,7 @@ import os import wave from pathlib import Path -from typing import Any, AsyncIterator, Sequence +from typing import Any, AsyncIterator, Mapping, Sequence from urllib.parse import urlparse import httpx @@ -74,6 +74,22 @@ def _auth_headers(api_key: str | None) -> dict[str, str]: return {"Authorization": f"Bearer {api_key}"} if api_key else {} +def _request_headers( + api_key: str | None, + headers: Mapping[str, str] | None, +) -> dict[str, str]: + """Merge per-call context while keeping model credentials configuration-owned.""" + result: dict[str, str] = {} + for name, value in (headers or {}).items(): + if not isinstance(name, str) or not isinstance(value, str): + raise TypeError("LLM request headers must be strings") + if name.lower() == "authorization": + raise ValueError("LLM request headers cannot override Authorization") + result[name] = value + result.update(_auth_headers(api_key)) + return result + + async def _http_health(client: httpx.AsyncClient, url: str, enabled: bool) -> bool: # Remote endpoints (hosted NIM) expose no local /health route; the spec # sets health_check=false, in which case readiness is assumed. @@ -322,6 +338,7 @@ async def chat( enable_thinking: bool = False, thinking_budget: int | None = None, timeout: float | None = None, + headers: Mapping[str, str] | None = None, ) -> ChatResponse: payload = self._build_payload( messages, @@ -329,7 +346,7 @@ async def chat( enable_thinking=enable_thinking, thinking_budget=thinking_budget, stream=False, ) - kwargs: dict[str, Any] = {"json": payload, "headers": _auth_headers(self._api_key)} + kwargs: dict[str, Any] = {"json": payload, "headers": _request_headers(self._api_key, headers)} if timeout is not None: kwargs["timeout"] = timeout resp = await self._client.post(self._chat_url, **kwargs) @@ -348,6 +365,7 @@ async def stream( enable_thinking: bool = False, thinking_budget: int | None = None, timeout: float | None = None, + headers: Mapping[str, str] | None = None, ) -> AsyncIterator[str]: payload = self._build_payload( messages, @@ -355,7 +373,7 @@ async def stream( enable_thinking=enable_thinking, thinking_budget=thinking_budget, stream=True, ) - kwargs: dict[str, Any] = {"json": payload, "headers": _auth_headers(self._api_key)} + kwargs: dict[str, Any] = {"json": payload, "headers": _request_headers(self._api_key, headers)} if timeout is not None: kwargs["timeout"] = timeout async with self._client.stream("POST", self._chat_url, **kwargs) as resp: diff --git a/agent-sdk/xr-ai-models/xr_ai_models/_protocols.py b/agent-sdk/xr-ai-models/xr_ai_models/_protocols.py index 37e649d63..05ac730ce 100644 --- a/agent-sdk/xr-ai-models/xr_ai_models/_protocols.py +++ b/agent-sdk/xr-ai-models/xr_ai_models/_protocols.py @@ -13,7 +13,7 @@ from dataclasses import dataclass from pathlib import Path -from typing import Any, AsyncIterator, Literal, Protocol, Sequence, runtime_checkable +from typing import Any, AsyncIterator, Literal, Mapping, Protocol, Sequence, runtime_checkable ImageInput = bytes | Path | str @@ -109,6 +109,7 @@ async def chat( enable_thinking: bool = False, thinking_budget: int | None = None, timeout: float | None = None, + headers: Mapping[str, str] | None = None, ) -> ChatResponse: pass def stream( @@ -121,6 +122,7 @@ def stream( enable_thinking: bool = False, thinking_budget: int | None = None, timeout: float | None = None, + headers: Mapping[str, str] | None = None, ) -> AsyncIterator[str]: pass async def health(self) -> bool: pass diff --git a/agent-sdk/xr-ai-nat/README.md b/agent-sdk/xr-ai-nat/README.md index b5d57594b..cf8822796 100644 --- a/agent-sdk/xr-ai-nat/README.md +++ b/agent-sdk/xr-ai-nat/README.md @@ -3,11 +3,64 @@ SPDX-License-Identifier: Apache-2.0 --> -# XR AI functions for NeMo Agent Toolkit +# XR AI native tools -`xr-ai-nat` provides typed, in-process XR functions for NVIDIA NeMo Agent -Toolkit (NAT). Applications compose these functions directly; process-backed -or MCP compatibility adapters remain separate boundaries. +`xr-ai-nat` is the toolkit-independent native tools layer for XR AI. +`Tool` gives voice, background triggers, and model-driven agents one typed +Pydantic invocation interface. NeMo Relay manages every new tool execution; +model-backed tools use injected `xr-ai-models` services rather than exposing a +model client to an application trigger. + +The existing NeMo Agent Toolkit function groups remain available through legacy +extras while their concrete capabilities migrate. They are compatibility code, +not the destination for new tools. + +## Native tools and tool-driven agents + +The base install supplies `Tool`, `AgentRunner`, and `as_agent_tool`. Install +`xr-ai-nat[relay]` for the bundled bounded tool-driven `Agent`: + +```python +from pydantic import BaseModel +from xr_ai_nat import Tool +from xr_ai_nat.agents import Agent + + +class LookupRequest(BaseModel): + query: str + + +class LookupResult(BaseModel): + answer: str + + +async def lookup(request: LookupRequest) -> LookupResult: + return LookupResult(answer=request.query) + + +lookup_tool = Tool( + "lookup", + "Look up one answer.", + LookupRequest, + LookupResult, + lookup, +) +agent = Agent( + name="assistant", + llm=llm, + system_prompt="Use the available tools.", + tools=(lookup_tool,), +) +``` + +`AgentRunner` is the small async turn protocol behind `as_agent_tool(...)`. +The bundled `Agent` is the basic stateless tool loop; applications can expose a +custom, Fabric-backed, or framework-backed runner through the same registered +`Tool`. That keeps voice, text, and autonomous background work on one +invocation path. Relay observes model calls inside a tool-backed runner; the +application never calls an LLM client as a separate control path. + +## Legacy NAT compatibility ## Shared value models and the service boundary @@ -22,7 +75,7 @@ Capabilities that talk to an out-of-process service share one private transport, `RPCServer`). `_service` owns only the transport; the value models above live in `functions.types`, not in `_service`. -## Model-backed agents +## Legacy NAT model bridge Install `xr-ai-nat[agents]` to make an `xr-ai-models` `LLMService` available to NAT's built-in LangChain-backed agent types without bypassing the repository diff --git a/agent-sdk/xr-ai-nat/pyproject.toml b/agent-sdk/xr-ai-nat/pyproject.toml index ea9575182..8c3545982 100644 --- a/agent-sdk/xr-ai-nat/pyproject.toml +++ b/agent-sdk/xr-ai-nat/pyproject.toml @@ -8,19 +8,21 @@ build-backend = "hatchling.build" [project] name = "xr-ai-nat" version = "0.1.0" -description = "NVIDIA NeMo Agent Toolkit functions for XR AI." +description = "Native Relay-managed tools and legacy NeMo Agent Toolkit compatibility for XR AI." requires-python = ">=3.11,<3.13" dependencies = [ - "nvidia-nat-core==1.8.0", + "nemo-relay>=0.7.2,<0.8", "pydantic>=2.10", ] [project.optional-dependencies] -agents = ["nvidia-nat-langchain==1.8.0", "xr-ai-models"] -mcp = ["fastmcp>=3.4,<4"] -services = ["msgpack>=1.0", "pyzmq>=27.0"] -vision = ["httpx>=0.27", "numpy>=1.24", "Pillow>=10.0", "xr-ai-hub-client", "xr-ai-models"] -voice = ["xr-ai-voice"] +relay = ["xr-ai-models"] +live-vision = ["numpy>=1.24", "Pillow>=10.0", "xr-ai-hub-client", "xr-ai-models"] +agents = ["nvidia-nat-core==1.8.0", "nvidia-nat-langchain==1.8.0", "xr-ai-models"] +mcp = ["nvidia-nat-core==1.8.0", "fastmcp>=3.4,<4"] +services = ["nvidia-nat-core==1.8.0", "msgpack>=1.0", "pyzmq>=27.0"] +vision = ["nvidia-nat-core==1.8.0", "httpx>=0.27", "numpy>=1.24", "Pillow>=10.0", "xr-ai-hub-client", "xr-ai-models"] +voice = ["nvidia-nat-core==1.8.0", "xr-ai-voice"] [project.entry-points."nat.plugins"] xr_ai_nat_llm = "xr_ai_nat.llm.config" diff --git a/agent-sdk/xr-ai-nat/xr_ai_nat/__init__.py b/agent-sdk/xr-ai-nat/xr_ai_nat/__init__.py index 282ee8956..dfcdb78c0 100644 --- a/agent-sdk/xr-ai-nat/xr_ai_nat/__init__.py +++ b/agent-sdk/xr-ai-nat/xr_ai_nat/__init__.py @@ -1,6 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""NAT-native XR functions.""" +"""Toolkit-independent native XR tools with legacy NAT compatibility.""" -__all__: list[str] = [] +from .agent_runner import AgentRunner, as_agent_tool +from .tools import Tool, ToolInvocationResult, ToolSet + +__all__ = ["AgentRunner", "Tool", "ToolInvocationResult", "ToolSet", "as_agent_tool"] diff --git a/agent-sdk/xr-ai-nat/xr_ai_nat/agent_runner.py b/agent-sdk/xr-ai-nat/xr_ai_nat/agent_runner.py new file mode 100644 index 000000000..6c65dc0d0 --- /dev/null +++ b/agent-sdk/xr-ai-nat/xr_ai_nat/agent_runner.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Framework-neutral agent runners exposed through native tools.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Protocol, TypeVar + +from pydantic import BaseModel + +from .tools import Tool + +RunnerRequestT = TypeVar("RunnerRequestT", contravariant=True) +RunnerResultT = TypeVar("RunnerResultT", covariant=True) +ToolRequestT = TypeVar("ToolRequestT", bound=BaseModel) +ToolResultT = TypeVar("ToolResultT", bound=BaseModel) + + +class AgentRunner(Protocol[RunnerRequestT, RunnerResultT]): + """An application-owned agent implementation that completes one asynchronous turn.""" + + async def run(self, request: RunnerRequestT) -> RunnerResultT: + """Run one turn and return the implementation-specific result.""" + raise NotImplementedError + + +def as_agent_tool( + *, + name: str, + description: str, + agent: AgentRunner[RunnerRequestT, RunnerResultT], + request_model: type[ToolRequestT], + result_model: type[ToolResultT], + request: Callable[[ToolRequestT], RunnerRequestT], + response: Callable[[RunnerResultT], ToolResultT], + return_direct: bool = False, +) -> Tool[ToolRequestT, ToolResultT]: + """Expose any ``AgentRunner`` through the same ``Tool`` interface as capabilities.""" + + async def invoke(value: ToolRequestT) -> ToolResultT: + return response(await agent.run(request(value))) + + return Tool( + name, + description, + request_model, + result_model, + invoke, + return_direct=return_direct, + ) + + +__all__ = ["AgentRunner", "as_agent_tool"] diff --git a/agent-sdk/xr-ai-nat/xr_ai_nat/agents.py b/agent-sdk/xr-ai-nat/xr_ai_nat/agents.py new file mode 100644 index 000000000..c98b8cf65 --- /dev/null +++ b/agent-sdk/xr-ai-nat/xr_ai_nat/agents.py @@ -0,0 +1,316 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tool-driven agents whose model calls remain internal to the tools layer.""" + +from __future__ import annotations + +import json +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import nemo_relay +from nemo_relay.codecs import OpenAIChatCodec +from xr_ai_models import ChatMessage, ChatResponse, LLMService, ToolCall, ToolDef + +from .agent_runner import AgentRunner, as_agent_tool +from .tools import Tool, ToolSet + + +class ToolLoopLimitError(RuntimeError): + """Raised when a model has not produced a final answer within the configured budget.""" + + +@dataclass(frozen=True, slots=True) +class AgentResult: + """The final text and messages generated during one tool-driven agent turn.""" + + text: str + messages: tuple[ChatMessage, ...] + + +class Agent(AgentRunner[str, AgentResult]): + """Run one small stateless tool-calling turn over a catalog of native tools.""" + + def __init__( + self, + *, + name: str, + llm: LLMService, + system_prompt: str, + tools: Sequence[Tool[Any, Any]], + model_name: str = "xr-ai-model", + max_iterations: int = 4, + max_tokens: int = 1024, + temperature: float = 0.0, + enable_thinking: bool = False, + thinking_budget: int | None = None, + ) -> None: + if not name: + raise ValueError("agent name must not be empty") + if not system_prompt: + raise ValueError(f"agent {name!r} needs a system prompt") + if max_iterations < 1: + raise ValueError("max_iterations must be at least one") + self.name = name + self.llm = llm + self.system_prompt = system_prompt + self.tools = ToolSet(tools) + self.model_name = model_name + self.max_iterations = max_iterations + self.max_tokens = max_tokens + self.temperature = temperature + self.enable_thinking = enable_thinking + self.thinking_budget = thinking_budget + + async def run(self, request: str) -> AgentResult: + """Execute one user request without carrying hidden conversation history.""" + + if not request.strip(): + raise ValueError("agent request must not be blank") + messages = [ + ChatMessage(role="system", content=self.system_prompt), + ChatMessage(role="user", content=request), + ] + with nemo_relay.scope.scope( + self.name, + nemo_relay.ScopeType.Agent, + input={"request": request}, + ): + for _ in range(self.max_iterations): + response = await self._chat(messages) + messages.append( + ChatMessage( + role="assistant", + content=response.content, + tool_calls=response.tool_calls, + ) + ) + if not response.tool_calls: + return AgentResult(response.content.strip(), tuple(messages)) + for call in response.tool_calls: + tool = self.tools.get(call.name) + if tool is None: + outcome = json.dumps({"error": "unknown_tool", "tool": call.name}) + return_direct = False + else: + invocation = await tool.invoke(call.arguments) + outcome = invocation.content + return_direct = invocation.return_direct + messages.append( + ChatMessage( + role="tool", + content=outcome, + tool_call_id=call.id, + ) + ) + if return_direct: + return AgentResult(outcome, tuple(messages)) + raise ToolLoopLimitError( + f"agent {self.name!r} exhausted {self.max_iterations} model iterations", + ) + + async def _chat(self, messages: Sequence[ChatMessage]) -> ChatResponse: + tool_definitions = list(self.tools.definitions) + content: dict[str, Any] = { + "model": self.model_name, + "messages": [_message_to_openai(message) for message in messages], + "max_tokens": self.max_tokens, + "temperature": self.temperature, + "enable_thinking": self.enable_thinking, + "thinking_budget": self.thinking_budget, + } + if tool_definitions: + content["tools"] = tool_definitions + relay_request = nemo_relay.LLMRequest( + {}, + content, + ) + + async def invoke(request: nemo_relay.LLMRequest) -> dict[str, Any]: + content = request.content + response = await self.llm.chat( + _messages_from_openai(content.get("messages")), + tools=_tools_from_openai(content.get("tools")), + max_tokens=_optional_int(content.get("max_tokens")), + temperature=_optional_float(content.get("temperature")), + enable_thinking=bool(content.get("enable_thinking", False)), + thinking_budget=_optional_int(content.get("thinking_budget")), + headers=_headers_from_relay(request.headers), + ) + return _response_to_openai(response) + + raw_response = await nemo_relay.llm.execute( + self.name, + relay_request, + invoke, + model_name=self.model_name, + codec=OpenAIChatCodec(), + response_codec=OpenAIChatCodec(), + ) + return _response_from_openai(raw_response) + + +def _message_to_openai(message: ChatMessage) -> dict[str, Any]: + if not isinstance(message.content, str): + raise TypeError("the native agent accepts text-only LLM messages") + content: str | None = message.content + if message.role == "assistant" and not content and message.tool_calls: + content = None + result: dict[str, Any] = {"role": message.role, "content": content} + if message.tool_calls: + result["tool_calls"] = [ + { + "id": call.id, + "type": "function", + "function": {"name": call.name, "arguments": call.arguments}, + } + for call in message.tool_calls + ] + if message.tool_call_id: + result["tool_call_id"] = message.tool_call_id + return result + + +def _messages_from_openai(raw: object) -> list[ChatMessage]: + if not isinstance(raw, list): + raise TypeError("Relay LLM request must contain a message array") + messages: list[ChatMessage] = [] + for item in raw: + if not isinstance(item, dict): + raise TypeError("Relay LLM messages must be objects") + role = item.get("role") + if role not in {"system", "user", "assistant", "tool"}: + raise ValueError(f"unsupported Relay LLM role: {role!r}") + content = item.get("content", "") + tool_calls = _tool_calls_from_openai(item.get("tool_calls")) or None + if content is None and tool_calls: + content = "" + if not isinstance(content, str): + raise TypeError("the native agent accepts text-only LLM messages") + tool_call_id = item.get("tool_call_id") + if tool_call_id is not None and not isinstance(tool_call_id, str): + raise TypeError("tool_call_id must be a string") + messages.append( + ChatMessage( + role=role, + content=content, + tool_calls=tool_calls, + tool_call_id=tool_call_id, + ) + ) + return messages + + +def _tool_calls_from_openai(raw: object) -> list[ToolCall]: + if raw is None: + return [] + if not isinstance(raw, list): + raise TypeError("tool_calls must be an array") + calls: list[ToolCall] = [] + for item in raw: + if not isinstance(item, dict) or not isinstance(item.get("function"), dict): + raise TypeError("tool call must contain a function object") + function = item["function"] + name = function.get("name") + arguments = function.get("arguments") + identifier = item.get("id") + if not isinstance(identifier, str) or not isinstance(name, str) or not isinstance(arguments, str): + raise TypeError("tool call id, name, and arguments must be strings") + calls.append(ToolCall(id=identifier, name=name, arguments=arguments)) + return calls + + +def _tools_from_openai(raw: object) -> list[ToolDef] | None: + if raw is None: + return None + if not isinstance(raw, list): + raise TypeError("tools must be an array") + definitions: list[ToolDef] = [] + for item in raw: + if not isinstance(item, dict): + raise TypeError("tool definition must be an object") + function = item.get("function", item) + if not isinstance(function, dict): + raise TypeError("tool definition function must be an object") + name = function.get("name") + description = function.get("description", "") + parameters = function.get("parameters", {"type": "object"}) + if not isinstance(name, str) or not isinstance(description, str) or not isinstance(parameters, dict): + raise TypeError("tool definition has invalid fields") + definitions.append(ToolDef(name=name, description=description, parameters=parameters)) + return definitions + + +def _headers_from_relay(raw: object) -> dict[str, str]: + if not isinstance(raw, dict): + raise TypeError("Relay LLM request headers must be an object") + headers: dict[str, str] = {} + for name, value in raw.items(): + if not isinstance(name, str) or not isinstance(value, str): + raise TypeError("Relay LLM request headers must be strings") + headers[name] = value + return headers + + +def _response_to_openai(response: ChatResponse) -> dict[str, Any]: + message = _message_to_openai( + ChatMessage(role="assistant", content=response.content, tool_calls=response.tool_calls), + ) + return { + "model": response.raw.get("model", "xr-ai-model") if isinstance(response.raw, dict) else "xr-ai-model", + "choices": [{"message": message, "finish_reason": response.finish_reason}], + "usage": response.raw.get("usage", {}) if isinstance(response.raw, dict) else {}, + "xr_ai_reasoning": response.reasoning, + } + + +def _response_from_openai(raw: object) -> ChatResponse: + if not isinstance(raw, dict): + raise TypeError("Relay LLM response must be an object") + choices = raw.get("choices") + if not isinstance(choices, list) or not choices or not isinstance(choices[0], dict): + raise ValueError("Relay LLM response must contain one choice") + choice = choices[0] + message = choice.get("message") + if not isinstance(message, dict): + raise TypeError("Relay LLM response choice must contain a message") + content = message.get("content", "") + tool_calls = _tool_calls_from_openai(message.get("tool_calls")) or None + if content is None and tool_calls: + content = "" + if not isinstance(content, str): + raise TypeError("Relay LLM response content must be a string") + finish_reason = choice.get("finish_reason") + if finish_reason is not None and not isinstance(finish_reason, str): + raise TypeError("Relay LLM finish_reason must be a string") + reasoning = raw.get("xr_ai_reasoning") + if reasoning is not None and not isinstance(reasoning, str): + raise TypeError("Relay LLM reasoning must be a string") + return ChatResponse( + content=content, + reasoning=reasoning, + tool_calls=tool_calls, + finish_reason=finish_reason, + raw=raw, + ) + + +def _optional_int(value: object) -> int | None: + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError("expected an integer or null") + return value + + +def _optional_float(value: object) -> float | None: + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError("expected a number or null") + return float(value) + + +__all__ = ["Agent", "AgentResult", "AgentRunner", "ToolLoopLimitError", "as_agent_tool"] diff --git a/agent-sdk/xr-ai-nat/xr_ai_nat/tools.py b/agent-sdk/xr-ai-nat/xr_ai_nat/tools.py new file mode 100644 index 000000000..37ae59cee --- /dev/null +++ b/agent-sdk/xr-ai-nat/xr_ai_nat/tools.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed native tools with a Relay-managed execution boundary.""" + +from __future__ import annotations + +import json +from collections.abc import Awaitable, Callable, Iterable +from dataclasses import dataclass +from inspect import isawaitable +from typing import Any, Generic, TypeVar, cast + +from nemo_relay import typed +from pydantic import BaseModel, ValidationError + +RequestT = TypeVar("RequestT", bound=BaseModel) +ResultT = TypeVar("ResultT", bound=BaseModel) + + +@dataclass(frozen=True, slots=True) +class ToolInvocationResult: + """A model-visible result from one local tool invocation.""" + + content: str + return_direct: bool + + +class Tool(Generic[RequestT, ResultT]): + """A Pydantic-validated tool shared by agents, voice, and background triggers.""" + + def __init__( + self, + name: str, + description: str, + request_model: type[RequestT], + result_model: type[ResultT], + handler: Callable[[RequestT], Awaitable[ResultT] | ResultT], + *, + return_direct: bool = False, + render_result: Callable[[ResultT], str] | None = None, + ) -> None: + if not name: + raise ValueError("tool name must not be empty") + if not description: + raise ValueError(f"tool {name!r} needs a description") + self.name = name + self.description = description + self.request_model = request_model + self.result_model = result_model + self.handler = handler + self.return_direct = return_direct + self._request_codec = typed.PydanticCodec(request_model) + self._result_codec = typed.PydanticCodec(result_model) + self._render_result = render_result or _json_result + + def to_openai(self) -> dict[str, Any]: + """Return the OpenAI-compatible definition supplied to an agent model.""" + + return { + "type": "function", + "function": { + "name": self.name, + "description": self.description, + "parameters": self.request_model.model_json_schema(), + }, + } + + async def execute(self, request: RequestT) -> ResultT: + """Run one validated request through the shared Relay tool lifecycle.""" + + return await typed.tool_execute( + self.name, + request, + self._execute_handler, + self._request_codec, + self._result_codec, + ) + + async def invoke(self, arguments: str) -> ToolInvocationResult: + """Validate and run one model-supplied JSON argument payload.""" + + try: + raw_arguments = json.loads(arguments) + except json.JSONDecodeError as exc: + return self._validation_error(f"arguments must be valid JSON: {exc.msg}") + if not isinstance(raw_arguments, dict): + return self._validation_error("arguments must be a JSON object") + try: + request = self.request_model.model_validate(raw_arguments) + except ValidationError as exc: + return self._validation_error(exc.json(include_url=False)) + return ToolInvocationResult( + self._render_result(await self.execute(request)), + self.return_direct, + ) + + def _validation_error(self, detail: str) -> ToolInvocationResult: + return ToolInvocationResult( + content=json.dumps({"error": "invalid_tool_arguments", "detail": detail}), + return_direct=False, + ) + + async def _execute_handler(self, request: RequestT) -> ResultT: + result = self.handler(request) + if isawaitable(result): + return await cast(Awaitable[ResultT], result) + return cast(ResultT, result) + + +class ToolSet: + """A non-overlapping tool catalog used by one native agent.""" + + def __init__(self, tools: Iterable[Tool[Any, Any]]) -> None: + by_name: dict[str, Tool[Any, Any]] = {} + for tool in tools: + if tool.name in by_name: + raise ValueError(f"duplicate tool name: {tool.name}") + by_name[tool.name] = tool + self._by_name = by_name + + @property + def definitions(self) -> tuple[dict[str, Any], ...]: + """Return model-visible schemas in registration order.""" + + return tuple(tool.to_openai() for tool in self._by_name.values()) + + def get(self, name: str) -> Tool[Any, Any] | None: + """Return the named tool when this catalog owns it.""" + + return self._by_name.get(name) + + +def _json_result(result: BaseModel) -> str: + return result.model_dump_json() + + +__all__ = ["Tool", "ToolInvocationResult", "ToolSet"] diff --git a/docs/changelog.md b/docs/changelog.md index b7aebe15d..b7230b29d 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -9,6 +9,26 @@ Significant decisions, in reverse-chronological order. Update this whenever a non-trivial architectural or design decision is made so the rationale is preserved and not re-litigated. +### 2026-08-11 — Native tools own agent composition; Relay owns their execution + +NeMo Agent Toolkit is being retired from XR AI in focused migrations rather +than a framework-wide swap. `xr-ai-nat` is the public, toolkit-independent +native tools layer: typed tools, an `AgentRunner` async-turn protocol, and a +bounded default tool loop. A model is a private implementation dependency of a +model-backed tool or agent, reached only through `xr-ai-models`; voice, text, +and background triggers invoke registered tools rather than model clients. NeMo +Relay runs the tool and model lifecycles, supplying middleware, guardrails, and +telemetry. `as_agent_tool` lets a custom or future Fabric-backed runner use the +same registered-tool path without making a framework part of the public trigger +boundary. + +The existing NeMo Agent Toolkit function groups remain behind legacy extras +while they migrate. Relay does not own XR application routing, participant +state, media IPC, or deployment. The tea-making sample defines the required +foreground/background application behavior for the migration, while NeMo +Platform and NeMo Fabric remain optional evaluation and harness deployment +targets outside the local worker dependency graph. + ### 2026-08-10 — Client readiness is hub-owned and routability-gated Process readiness and client readiness were the same signal, and both were diff --git a/docs/nemo-agent-toolkit-migration.md b/docs/nemo-agent-toolkit-migration.md new file mode 100644 index 000000000..409de963c --- /dev/null +++ b/docs/nemo-agent-toolkit-migration.md @@ -0,0 +1,83 @@ + + +# NeMo Agent Toolkit migration + +XR AI is retiring its NeMo Agent Toolkit dependency without changing the +application boundary: agents remain in-process XR workers, use `xr-ai-models` +for all model HTTP, and reach clients only through the Hub IPC SDK. + +## Target architecture + +NeMo Relay is the local execution boundary. The XR-owned `xr-ai-nat` package +is becoming a toolkit-independent tools layer: Pydantic tool schemas, trigger +dispatch, and the generic `AgentRunner` protocol. Its bundled agent is a +bounded tool loop, while `as_agent_tool(...)` lets custom or future +Fabric-backed runners use the same trigger path. Relay owns LLM and tool +lifecycles, middleware, guardrails, and telemetry. Existing NeMo Agent Toolkit +function groups remain compatibility extras until their concrete tools migrate. + +NeMo Platform and NeMo Fabric are deployment and evaluation integrations, not +worker dependencies. Platform currently requires Python 3.12 or 3.13 and owns +local services, evaluation, tuning, and deployments; Fabric runs a selected +agent harness in a configured environment. XR AI supports Python 3.11 and 3.12 +and needs direct local Hub, camera, and voice ownership, so both integrations +remain optional launch targets after the local runtime has migrated. + +```text +XR worker + -> xr-ai-nat: typed tools and trigger dispatch + -> NeMo Relay: managed tool and agent execution, guardrails, telemetry + -> xr-ai-models: private model boundary used by model-backed tools + -> Hub IPC: media and client data + +Optional operations path + -> NeMo Platform: evaluation, tuning, deployment, monitoring + -> NeMo Fabric: configured harness execution +``` + +## Tea-making acceptance reference + +`origin/devdeepr/tea-making-sample` is the behavioral reference. Its eventual +port must retain deterministic foreground ownership, independently lifecycled +background applications, typed state commits, evidence-gated workflow progress, +RAG only for missing known-product guidance, participant isolation, and the +read-only activity viewer. The branch currently uses NAT, so it defines +acceptance behavior rather than an implementation dependency. + +## Focused PR sequence + +1. **Native tools foundation** — make `xr-ai-nat` independent of NeMo Agent + Toolkit by default, add Relay-managed tools, the generic `AgentRunner` + seam, and a bounded default tool loop, and retain existing function groups + behind legacy extras. +2. **Simple VLM tool** — move the single-turn streaming-vision path to a normal + native tool and prove the lightweight voice sample selects no legacy extra. +3. **Native event dispatcher** — port tea-making's typed participant-scoped + subscriptions and periodic background sources so voice and autonomous work + invoke the same registered tools. +4. **Deterministic and service capabilities** — port spatial math, text memory, + RAG, vision, XR tracking, and video memory to the native tool surface; keep + MCP adapters as explicit compatibility publishers. +5. **Existing agent workflows** — port render-demo's tool catalog and evaluation + harness to the tools layer, then replace its NAT builder and LangChain bridge. +6. **Tea-making sample** — land the sample in slices: workflow/state core, + foreground/background application manager, observation applications, then + activity viewer and end-to-end evaluations. +7. **Retirement and operations** — remove all `nat.*`, `nvidia-nat-*`, and + `nemo_toolkit` legacy references; remove lockfile and notice entries; + add opt-in Platform evaluation/deployment and Fabric harness profiles. + +Every migration PR must preserve participant scoping, add or update the direct +tests for the surface it changes, update `DEPENDENCIES.md`, and leave no +hand-rolled model HTTP client behind. + +## Exit criteria + +The retirement is complete only when repository-wide search finds no runtime +imports from `nat` and no `nvidia-nat-*` dependency, worker requirements select +only toolkit-independent `xr-ai-nat` extras, all migrated samples pass their +unit and evaluation suites, and `THIRD_PARTY_NOTICES.md` no longer lists NeMo +Agent Toolkit. diff --git a/docs/source/components/agent-sdk.md b/docs/source/components/agent-sdk.md index 659ad84e9..bd480079a 100644 --- a/docs/source/components/agent-sdk.md +++ b/docs/source/components/agent-sdk.md @@ -21,9 +21,9 @@ from: - **`xr-ai-hub-client`** — the minimal pyzmq + msgpack IPC library every agent uses to talk to the XR-Media-Hub (refer to {doc}`server-runtime`). No LiveKit or FastAPI dependency. -- **`xr-ai-nat`** — typed in-process XR functions and the model bridge used by - NAT's built-in agents. NAT composition stays in process while all model I/O - continues through `xr-ai-models`. +- **`xr-ai-nat`** — Relay-managed native tools, the generic `AgentRunner` + protocol, and a basic tool-driven agent. Its legacy extras retain existing + NeMo Agent Toolkit function groups while their concrete capabilities migrate. --- @@ -139,7 +139,8 @@ class LLMService(Protocol): capabilities: Capabilities async def chat(self, messages, *, tools=None, max_tokens=None, temperature=None, enable_thinking=False, - thinking_budget=None, timeout=None) -> ChatResponse: ... + thinking_budget=None, timeout=None, + headers=None) -> ChatResponse: ... def stream(self, messages, *, ...) -> AsyncIterator[str]: ... async def health(self) -> bool: ... async def close(self) -> None: ... @@ -202,9 +203,25 @@ The clients can be exercised without a GPU. --- +## Native tools and agents + +`xr-ai-nat` is the native migration target for model-driven XR composition. +`Tool` declares Pydantic request and response boundaries and executes its +handler through NeMo Relay. `AgentRunner` is the generic async-turn protocol; +`xr_ai_nat.agents.Agent` is its bounded, stateless tool-loop implementation. +It builds OpenAI-compatible tool definitions from those schemas and sends each +model request through an injected `LLMService`. + +Relay scopes each turn and manages the LLM and tool lifecycles. The base package +does not add an HTTP client, a Hub transport, or an implicit conversation store. +Applications use `as_agent_tool(...)` to expose any `AgentRunner` as a +registered tool; foreground selection, workflow state, and background work stay +explicit in application code. + ## xr-ai-nat model bridge -Install `xr-ai-nat[agents]` when a workflow uses NAT's built-in agent graphs. +Unmigrated workflows install `xr-ai-nat[agents]` when they use NAT's built-in +agent graphs. `ModelsLLMConfig` adapts an `xr-ai-models` `LLMService` to NAT's LangChain client contract: diff --git a/tests/pyproject.toml b/tests/pyproject.toml index 5d6ae21c3..e1150b59f 100644 --- a/tests/pyproject.toml +++ b/tests/pyproject.toml @@ -14,7 +14,7 @@ dependencies = [ # Pulled in via editable installs of the workspace packages. "xr-ai-hub-client", "xr-ai-models", - "xr-ai-nat[agents,services,vision]", + "xr-ai-nat[agents,relay,services,vision]", "xr-ai-pipecat", "xr-ai-voice", "xr-media-hub", diff --git a/tests/test_models_openai_compat.py b/tests/test_models_openai_compat.py index 8353e80d8..7cc50424e 100644 --- a/tests/test_models_openai_compat.py +++ b/tests/test_models_openai_compat.py @@ -262,6 +262,30 @@ async def test_llm_chat_no_auth_header_without_api_key_env() -> None: assert "Authorization" not in stub.last_request().headers +async def test_llm_chat_forwards_controlled_per_call_headers() -> None: + stub = StubOpenAI() + async with OpenAICompatLLM( + "http://stub", "llm", client=stub.client(), + ) as llm: + await llm.chat( + [ChatMessage(role="user", content="x")], + headers={"X-Relay-Session": "turn-7"}, + ) + assert stub.last_request().headers["X-Relay-Session"] == "turn-7" + + +async def test_llm_chat_rejects_per_call_authorization_header() -> None: + stub = StubOpenAI() + async with OpenAICompatLLM( + "http://stub", "llm", client=stub.client(), + ) as llm: + with pytest.raises(ValueError, match="cannot override Authorization"): + await llm.chat( + [ChatMessage(role="user", content="x")], + headers={"Authorization": "Bearer untrusted"}, + ) + + async def test_llm_chat_raises_on_http_error() -> None: import httpx stub = StubOpenAI() diff --git a/tests/test_native_tools.py b/tests/test_native_tools.py new file mode 100644 index 000000000..52335e291 --- /dev/null +++ b/tests/test_native_tools.py @@ -0,0 +1,366 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Contracts for Relay-managed native tools and tool-driven agents.""" + +from __future__ import annotations + +import json +from collections.abc import AsyncIterator + +import nemo_relay +import pytest +from nemo_relay.codecs import OpenAIChatCodec +from pydantic import BaseModel +from xr_ai_models import Capabilities, ChatMessage, ChatResponse, ToolCall, ToolDef +from xr_ai_nat import AgentRunner, Tool, ToolSet, as_agent_tool +from xr_ai_nat.agents import Agent, ToolLoopLimitError, _response_from_openai + + +class AddRequest(BaseModel): + """Two integers to add.""" + + left: int + right: int + + +class AddResult(BaseModel): + """The computed total.""" + + total: int + + +class AskRequest(BaseModel): + """One text request delegated to an agent tool.""" + + text: str + + +class AskResult(BaseModel): + """The agent tool's text result.""" + + text: str + + +async def add(request: AddRequest) -> AddResult: + return AddResult(total=request.left + request.right) + + +class _ToolCallingLLM: + capabilities = Capabilities(tool_calls=True) + + def __init__(self) -> None: + self.calls: list[tuple[list[ChatMessage], list[ToolDef] | None]] = [] + + async def chat(self, messages, *, tools=None, **_kwargs) -> ChatResponse: + self.calls.append((list(messages), list(tools) if tools else None)) + if len(self.calls) == 1: + return ChatResponse( + content="", + reasoning=None, + tool_calls=[ + ToolCall( + id="add-call", + name="add", + arguments=json.dumps({"left": 2, "right": 3}), + ) + ], + finish_reason="tool_calls", + raw={"usage": {"prompt_tokens": 10, "completion_tokens": 3}}, + ) + return ChatResponse( + content="The answer is five.", + reasoning=None, + tool_calls=None, + finish_reason="stop", + raw={"usage": {"prompt_tokens": 20, "completion_tokens": 5}}, + ) + + async def health(self) -> bool: + return True + + async def close(self) -> None: + return None + + async def stream(self, *_args, **_kwargs) -> AsyncIterator[str]: + if False: + yield "" + + +async def test_agent_executes_a_native_tool_and_returns_the_final_model_answer() -> None: + llm = _ToolCallingLLM() + agent = Agent( + name="calculator", + llm=llm, + system_prompt="Use the add tool.", + tools=(Tool("add", "Add two integers.", AddRequest, AddResult, add),), + max_iterations=2, + ) + + result = await agent.run("What is two plus three?") + + assert result.text == "The answer is five." + assert len(llm.calls) == 2 + first_messages, first_tools = llm.calls[0] + assert [message.role for message in first_messages] == ["system", "user"] + assert first_tools is not None + assert first_tools[0].name == "add" + second_messages, _ = llm.calls[1] + assert [(message.role, message.content) for message in second_messages[-2:]] == [ + ("assistant", ""), + ("tool", '{"total":5}'), + ] + + +async def test_tools_are_relay_managed_for_agent_and_direct_invocation() -> None: + tool = Tool("add", "Add two integers.", AddRequest, AddResult, add) + llm = _ToolCallingLLM() + agent = Agent( + name="calculator-lifecycle", + llm=llm, + system_prompt="Use the add tool.", + tools=(tool,), + max_iterations=2, + ) + events = [] + subscriber = "xr-ai-native-tools-lifecycle" + nemo_relay.subscribers.register(subscriber, events.append) + try: + assert await tool.execute(AddRequest(left=1, right=2)) == AddResult(total=3) + await agent.run("What is two plus three?") + await nemo_relay.subscribers.flush_async() + finally: + nemo_relay.subscribers.deregister(subscriber) + + categories = {getattr(event, "category", None) for event in events} + assert {"llm", "tool"} <= categories + + +async def test_agent_can_be_exposed_as_a_normal_native_tool() -> None: + class _FinalAnswerLLM(_ToolCallingLLM): + async def chat(self, messages, *, tools=None, **_kwargs) -> ChatResponse: + self.calls.append((list(messages), list(tools) if tools else None)) + return ChatResponse( + content="Handled by the agent tool.", + reasoning=None, + tool_calls=None, + finish_reason="stop", + raw={}, + ) + + agent = Agent( + name="delegate", + llm=_FinalAnswerLLM(), + system_prompt="Answer directly.", + tools=(), + ) + tool = as_agent_tool( + name="delegate", + description="Delegate one request to the agent.", + agent=agent, + request_model=AskRequest, + result_model=AskResult, + request=lambda value: value.text, + response=lambda result: AskResult(text=result.text), + ) + + assert await tool.execute(AskRequest(text="Hello")) == AskResult( + text="Handled by the agent tool.", + ) + + +async def test_custom_agent_runner_can_be_exposed_as_a_normal_native_tool() -> None: + class _CustomRunner: + async def run(self, request: str) -> AskResult: + return AskResult(text=f"Custom runner: {request}") + + runner: AgentRunner[str, AskResult] = _CustomRunner() + tool = as_agent_tool( + name="custom_delegate", + description="Delegate one request to a custom agent runner.", + agent=runner, + request_model=AskRequest, + result_model=AskResult, + request=lambda value: value.text, + response=lambda result: result, + ) + + assert await tool.execute(AskRequest(text="Hello")) == AskResult( + text="Custom runner: Hello", + ) + + +async def test_agent_forwards_relay_rewritten_request_to_the_model(monkeypatch) -> None: + class _RecordingLLM: + capabilities = Capabilities() + + def __init__(self) -> None: + self.messages: list[ChatMessage] = [] + self.headers: dict[str, str] = {} + + async def chat(self, messages, *, headers=None, **_kwargs) -> ChatResponse: + self.messages = list(messages) + self.headers = dict(headers or {}) + return ChatResponse("Rewritten.", None, None, "stop", {}) + + async def health(self) -> bool: + return True + + async def close(self) -> None: + return None + + async def stream(self, *_args, **_kwargs) -> AsyncIterator[str]: + if False: + yield "" + + llm = _RecordingLLM() + observed: dict[str, object] = {} + + async def execute(_name, request, invoke, **kwargs): + observed["content"] = request.content + observed["codec"] = kwargs["codec"] + observed["response_codec"] = kwargs["response_codec"] + rewritten = dict(request.content) + rewritten["messages"] = [ + {"role": "system", "content": "Answer directly."}, + {"role": "user", "content": "rewritten request"}, + ] + return await invoke(nemo_relay.LLMRequest({"X-Relay-Session": "turn-7"}, rewritten)) + + monkeypatch.setattr(nemo_relay.llm, "execute", execute) + result = await Agent( + name="relay-boundary", + llm=llm, + system_prompt="Answer directly.", + tools=(), + ).run("original request") + + assert result.text == "Rewritten." + initial_content = observed["content"] + assert isinstance(initial_content, dict) + assert "tools" not in initial_content + assert isinstance(observed["codec"], OpenAIChatCodec) + assert isinstance(observed["response_codec"], OpenAIChatCodec) + assert llm.headers == {"X-Relay-Session": "turn-7"} + assert llm.messages[-1].content == "rewritten request" + + +def test_agent_accepts_null_content_for_tool_call_only_response() -> None: + response = _response_from_openai( + { + "choices": [ + { + "message": { + "content": None, + "tool_calls": [ + { + "id": "lookup-1", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + }, + "finish_reason": "tool_calls", + } + ] + } + ) + + assert response.content == "" + assert response.tool_calls == [ToolCall(id="lookup-1", name="lookup", arguments="{}")] + + +async def test_invalid_tool_arguments_are_returned_to_the_model_for_repair() -> None: + tool = Tool("add", "Add two integers.", AddRequest, AddResult, add) + + result = await tool.invoke('{"left":"not-an-int"}') + + assert result.return_direct is False + payload = json.loads(result.content) + assert payload["error"] == "invalid_tool_arguments" + assert "right" in payload["detail"] + + +async def test_return_direct_tool_finishes_without_an_extra_model_call() -> None: + llm = _ToolCallingLLM() + agent = Agent( + name="calculator", + llm=llm, + system_prompt="Use the add tool.", + tools=( + Tool( + "add", + "Add two integers.", + AddRequest, + AddResult, + add, + return_direct=True, + ), + ), + ) + + result = await agent.run("What is two plus three?") + + assert result.text == '{"total":5}' + assert len(llm.calls) == 1 + + +async def test_unknown_tool_is_returned_to_the_model_for_repair() -> None: + class _UnknownToolLLM(_ToolCallingLLM): + async def chat(self, messages, *, tools=None, **kwargs) -> ChatResponse: + response = await super().chat(messages, tools=tools, **kwargs) + if len(self.calls) == 1: + return ChatResponse( + content="", + reasoning=None, + tool_calls=[ToolCall(id="missing", name="missing", arguments="{}")], + finish_reason="tool_calls", + raw={}, + ) + return response + + llm = _UnknownToolLLM() + agent = Agent( + name="calculator", + llm=llm, + system_prompt="Use the available tools.", + tools=(Tool("add", "Add two integers.", AddRequest, AddResult, add),), + max_iterations=2, + ) + + result = await agent.run("Calculate.") + + assert result.text == "The answer is five." + content = llm.calls[1][0][-1].content + assert isinstance(content, str) + assert json.loads(content) == {"error": "unknown_tool", "tool": "missing"} + + +async def test_agent_enforces_its_model_iteration_budget() -> None: + class _LoopingLLM(_ToolCallingLLM): + async def chat(self, messages, *, tools=None, **kwargs) -> ChatResponse: + self.calls.append((list(messages), list(tools) if tools else None)) + return ChatResponse( + content="", + reasoning=None, + tool_calls=[ToolCall(id=str(len(self.calls)), name="add", arguments='{"left":1,"right":1}')], + finish_reason="tool_calls", + raw={}, + ) + + agent = Agent( + name="calculator", + llm=_LoopingLLM(), + system_prompt="Use the add tool.", + tools=(Tool("add", "Add two integers.", AddRequest, AddResult, add),), + max_iterations=2, + ) + + with pytest.raises(ToolLoopLimitError, match="exhausted 2"): + await agent.run("Loop forever.") + + +def test_tool_sets_reject_duplicate_names() -> None: + tool = Tool("add", "Add two integers.", AddRequest, AddResult, add) + + with pytest.raises(ValueError, match="duplicate tool name"): + ToolSet((tool, tool))