diff --git a/AGENTS.md b/AGENTS.md index 58f359081..5e6ebbfcb 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -12,7 +12,8 @@ historical decisions in `docs/changelog.md`. ``` client-samples/ # Platform clients (Android, iOS/visionOS, Web) -agent-sdk/ # Six packages: +agent-sdk/ # Seven packages: + # xr-ai-agent-runtime — agent lifecycles, tools + pub/sub # xr-ai-hub-client — IPC client library (pyzmq + msgpack only) # xr-ai-models — LLM/VLM/STT/TTS service protocols + OpenAI-compat clients # xr-ai-pipecat — optional Pipecat transport bridge (heavier deps) @@ -59,14 +60,22 @@ deps/ # Gitignored downloaded binaries (e.g. LOVR AppImage) `anthropic`, no `litellm`); all in-tree backends speak OpenAI-compatible HTTP. - **Workers never import from `xr_media_hub` or `xr_ai_launcher`.** Use the - public `xr_ai_hub`, `xr_ai_models`, `xr_ai_tools`, `xr_ai_nat`, and - `xr_ai_voice` SDK surfaces plus task-specific libraries (numpy, torch, …). + public `xr_ai_runtime`, `xr_ai_hub`, `xr_ai_models`, `xr_ai_tools`, + `xr_ai_nat`, and `xr_ai_voice` SDK surfaces plus task-specific libraries + (numpy, torch, …). - **Agentic functions are native and in-process.** New and migrated tools live in `xr-ai-tools`; every tool execution passes through NeMo Relay, and all model I/O remains in `xr-ai-models`. `xr-ai-nat` retains NeMo Agent Toolkit compatibility only while existing function groups migrate. Existing MCP servers remain compatibility surfaces while their capabilities migrate. +- **Agents expose the existing tools from `xr-ai-tools`.** Each `Agent` owns + state and a set of `Tool` or `AsyncTool` instances. Agents call one another's + tools directly through `execute()` or `stream()`; model-selected calls use + the same normal tool-calling helpers. `xr-ai-agent-runtime` manages typed + `publish()` fan-out and only the delivery tasks it creates. Agents own their + resources, tasks, lifecycle, and concurrency policy. Model loops, planning, + memory, and raw media transport remain outside the runtime. - **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 eb500df04..bc91723bc 100644 --- a/DEPENDENCIES.md +++ b/DEPENDENCIES.md @@ -37,6 +37,15 @@ CI matrices: ## Internal packages ``` +xr-ai-agent-runtime (agent-sdk/xr-ai-agent-runtime/) + └── pydantic >=2.10 + └── xr-ai-tools [editable: ../xr-ai-tools] + In-process runtime for agent resource lifetimes, runtime-owned background + tasks, and typed ``publish`` fan-out. Agents expose ordinary ``Tool`` and + ``AsyncTool`` instances from ``xr-ai-tools`` and own their synchronization. + Tool execution, model clients, tool loops, planning, memory, and raw media + transport are not runtime responsibilities. + xr-ai-hub-client (agent-sdk/xr-ai-hub-client/) └── pyzmq >=27.0 └── msgpack >=1.0 @@ -333,6 +342,7 @@ vec-mcp-server (agent-mcp-servers/vec-mcp/) remains compatibility-only and is not part of the native function group. xr-ai-tests (tests/) + └── xr-ai-agent-runtime [editable: ../agent-sdk/xr-ai-agent-runtime] └── 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] diff --git a/README.md b/README.md index 7c1b8838f..138abcfcc 100644 --- a/README.md +++ b/README.md @@ -117,6 +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 runtime | `agent-sdk/xr-ai-agent-runtime/` | Agent resource lifetimes, background tasks, existing native tools, and pub/sub | | Agent tools | `agent-sdk/xr-ai-tools/` | Toolkit-independent Relay-managed native tools | | Legacy NAT | `agent-sdk/xr-ai-nat/` | NeMo Agent Toolkit compatibility during migration | | Reusable services | `services/` | Model-serving and typed capability processes | diff --git a/agent-sdk/xr-ai-agent-runtime/README.md b/agent-sdk/xr-ai-agent-runtime/README.md new file mode 100644 index 000000000..3ec413834 --- /dev/null +++ b/agent-sdk/xr-ai-agent-runtime/README.md @@ -0,0 +1,94 @@ + + +# XR AI agent runtime + +`xr-ai-agent-runtime` provides typed pub/sub for composable XR AI agents. An +`Agent` owns private state and exposes ordinary `Tool` or `AsyncTool` instances +from `xr-ai-tools`. + +Tools have one invocation path everywhere: + +- Direct callers and other agents use `Tool.execute()` or `AsyncTool.stream()`. +- Model loops expose unary tools with `ToolSet` and `handle_tool_call()`. +- The agent runtime never wraps or redispatches a tool call. + +Unary tools return one validated Pydantic response or `None` for an +acknowledged side effect. Streaming tools yield validated chunks without +buffering the complete result. + +```python +from pydantic import BaseModel +from xr_ai_runtime import Agent, AgentRuntime +from xr_ai_tools import Tool + + +class TextRequest(BaseModel): + text: str + + +class TextResponse(BaseModel): + text: str + + +class TextAgent(Agent): + def __init__(self) -> None: + self.echo = Tool( + "echo", + "Echo the supplied text.", + TextRequest, + TextResponse, + self._echo, + ) + self.uppercase = Tool( + "uppercase", + "Uppercase the supplied text.", + TextRequest, + TextResponse, + self._uppercase, + ) + super().__init__((self.echo, self.uppercase)) + + async def _echo(self, request: TextRequest) -> TextResponse: + return TextResponse(text=request.text) + + async def _uppercase(self, request: TextRequest) -> TextResponse: + return TextResponse(text=request.text.upper()) + + +runtime = AgentRuntime() +text = runtime.register("text", TextAgent()) + +async with runtime: + result = await text.uppercase.execute(TextRequest(text="hello")) +``` + +Another agent receives the concrete tools it needs and calls them normally. +No runtime address or adapter is involved. + +When tools from several agents are combined for one model, namespace them at +the workflow boundary with +`ToolSet.namespaced({"vision": vision.tools, "planner": planner.tools})`. +This remaps only model-visible catalog names; the agents and underlying tools +remain unchanged. Participant identity needed by a direct tool belongs in that +tool's typed request. Relay supplies nested execution tracing, while runtime +message metadata applies only to pub/sub. + +`publish(topic, event)` is the separate asynchronous fan-out operation for +events. An agent that owns resources or background work is responsible for +controlling them, including creating, cancelling, and awaiting its own tasks. +The runtime neither knows nor controls whether an agent's internal work is +running. `publish()` waits for every fan-out delivery to settle before +propagating any subscriber failures. + +Tools and subscription callbacks may run concurrently. An agent whose mutable +state is shared between them owns the appropriate synchronization, such as an +`asyncio.Lock` or a private queue. This avoids imposing serialization and +head-of-line blocking on unrelated or streaming tools. + +Domain controls such as `start_monitoring`, `stop_monitoring`, and `status` are +ordinary tools. Agent lifetime itself is not a model tool. Model loops, +planning, memory, and model clients remain agent implementations. Raw audio and +video stay on the XR-Media-Hub path. diff --git a/agent-sdk/xr-ai-agent-runtime/pyproject.toml b/agent-sdk/xr-ai-agent-runtime/pyproject.toml new file mode 100644 index 000000000..c38bbb6e4 --- /dev/null +++ b/agent-sdk/xr-ai-agent-runtime/pyproject.toml @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "xr-ai-agent-runtime" +version = "0.1.0" +description = "Typed pub/sub runtime for XR AI agents." +requires-python = ">=3.11,<3.13" +dependencies = [ + "pydantic>=2.10", + "xr-ai-tools", +] + +[tool.uv.sources] +xr-ai-tools = { path = "../xr-ai-tools", editable = true } + +[tool.hatch.build.targets.wheel] +packages = ["xr_ai_runtime"] diff --git a/agent-sdk/xr-ai-agent-runtime/xr_ai_runtime/__init__.py b/agent-sdk/xr-ai-agent-runtime/xr_ai_runtime/__init__.py new file mode 100644 index 000000000..13647177c --- /dev/null +++ b/agent-sdk/xr-ai-agent-runtime/xr_ai_runtime/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed in-process runtime for composable XR AI agents.""" + +from .agent import Agent +from .events import MessageMetadata, Topic, subscribe +from .runtime import AgentRuntime, RuntimeClosedError, RuntimeContext + +__all__ = [ + "Agent", + "AgentRuntime", + "MessageMetadata", + "RuntimeClosedError", + "RuntimeContext", + "Topic", + "subscribe", +] diff --git a/agent-sdk/xr-ai-agent-runtime/xr_ai_runtime/agent.py b/agent-sdk/xr-ai-agent-runtime/xr_ai_runtime/agent.py new file mode 100644 index 000000000..d187b8622 --- /dev/null +++ b/agent-sdk/xr-ai-agent-runtime/xr_ai_runtime/agent.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Agents and their exposed native tools.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +from xr_ai_tools import AsyncTool, Tool + + +class Agent: + """An object that owns state and exposes ordinary native tools.""" + + def __init__( + self, + tools: Iterable[Tool[Any, Any] | AsyncTool[Any, Any]] = (), + ) -> None: + owned = tuple(tools) + names: set[str] = set() + for tool in owned: + if not isinstance(tool, (Tool, AsyncTool)): + raise TypeError("agents may expose only Tool or AsyncTool instances") + if tool.name in names: + raise ValueError(f"duplicate agent tool name: {tool.name}") + names.add(tool.name) + self._tools = owned + + @property + def tools(self) -> tuple[Tool[Any, Any] | AsyncTool[Any, Any], ...]: + """Return the existing native tools exposed by this agent.""" + + return self._tools + +__all__ = ["Agent"] diff --git a/agent-sdk/xr-ai-agent-runtime/xr_ai_runtime/events.py b/agent-sdk/xr-ai-agent-runtime/xr_ai_runtime/events.py new file mode 100644 index 000000000..66fbb2bdc --- /dev/null +++ b/agent-sdk/xr-ai-agent-runtime/xr_ai_runtime/events.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed event contracts for agent communication.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import Any, Generic, TypeVar + +from pydantic import BaseModel + +MessageT = TypeVar("MessageT", bound=BaseModel) + + +@dataclass(frozen=True, slots=True) +class Topic(Generic[MessageT]): + """A stable publish/subscribe name paired with its payload model.""" + + name: str + message_type: type[MessageT] + + def __post_init__(self) -> None: + if not self.name.strip(): + raise ValueError("topic name must not be empty") + if not issubclass(self.message_type, BaseModel): + raise TypeError("topic messages must be Pydantic models") + + def validate(self, message: MessageT | dict[str, Any]) -> MessageT: + """Validate a message before delivery.""" + + return self.message_type.model_validate(message) + + +@dataclass(frozen=True, slots=True) +class MessageMetadata: + """Routing and trace context for one published event.""" + + message_id: str + correlation_id: str + participant_id: str + source: str + parent_message_id: str | None + timestamp_us: int + + +def subscribe( + topic: Topic[MessageT], +) -> Callable[[Callable[..., Awaitable[None]]], Callable[..., Awaitable[None]]]: + """Register an agent method as a typed topic subscriber.""" + + def decorate(method: Callable[..., Awaitable[None]]) -> Callable[..., Awaitable[None]]: + topics = (*getattr(method, "__xr_ai_topics__", ()), topic) + setattr(method, "__xr_ai_topics__", topics) + return method + + return decorate + + +__all__ = ["MessageMetadata", "Topic", "subscribe"] diff --git a/agent-sdk/xr-ai-agent-runtime/xr_ai_runtime/runtime.py b/agent-sdk/xr-ai-agent-runtime/xr_ai_runtime/runtime.py new file mode 100644 index 000000000..a8420c4bf --- /dev/null +++ b/agent-sdk/xr-ai-agent-runtime/xr_ai_runtime/runtime.py @@ -0,0 +1,298 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed event routing between agents.""" + +from __future__ import annotations + +import asyncio +import inspect +import time +import uuid +from builtins import BaseExceptionGroup +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from typing import Any, TypeAlias, TypeVar, cast, get_type_hints + +from pydantic import BaseModel + +from .agent import Agent +from .events import MessageMetadata, Topic + +MessageT = TypeVar("MessageT", bound=BaseModel) +AgentT = TypeVar("AgentT", bound="Agent") + +BoundSubscriber: TypeAlias = Callable[[BaseModel, "RuntimeContext"], Awaitable[None]] + + +class RuntimeClosedError(RuntimeError): + """Raised when work is submitted to a runtime that is not running.""" + + +@dataclass(slots=True) +class _AgentState: + name: str + deliveries: set[asyncio.Task[None]] = field(default_factory=set) + + +class RuntimeContext: + """Runtime operations available during a subscription delivery.""" + + __slots__ = ("_runtime", "_agent_name", "_metadata") + + def __init__( + self, + runtime: AgentRuntime, + agent_name: str, + metadata: MessageMetadata, + ) -> None: + self._runtime = runtime + self._agent_name = agent_name + self._metadata = metadata + + @property + def agent_name(self) -> str: + """Return this agent's runtime-local name.""" + + return self._agent_name + + @property + def metadata(self) -> MessageMetadata: + """Return metadata for the current subscription delivery.""" + + return self._metadata + + async def publish( + self, + topic: Topic[MessageT], + message: MessageT | dict[str, Any], + *, + participant_id: str | None = None, + ) -> None: + """Publish an event while preserving current trace context.""" + + await self._runtime._publish( + topic, + message, + participant_id=self._resolve_participant(participant_id), + source=self._agent_name, + correlation_id=self._metadata.correlation_id, + parent_message_id=self._metadata.message_id, + ) + + def _resolve_participant(self, participant_id: str | None) -> str: + if participant_id is not None: + return participant_id + return self._metadata.participant_id + + +class AgentRuntime: + """Provide typed pub/sub routing between registered agents.""" + + def __init__(self) -> None: + self._agents: dict[str, _AgentState] = {} + self._topics: dict[str, Topic[Any]] = {} + self._subscribers: dict[str, list[tuple[_AgentState, BoundSubscriber]]] = {} + self._running = False + self._closed = False + + @property + def running(self) -> bool: + """Whether the runtime currently accepts event work.""" + + return self._running and not self._closed + + def register(self, name: str, agent: AgentT) -> AgentT: + """Register one agent before startup and return the same typed object.""" + + if self._running or self._closed: + raise RuntimeError("agents must be registered before the runtime starts") + if not name.strip(): + raise ValueError("agent name must not be empty") + if name in self._agents: + raise ValueError(f"agent {name!r} is already registered") + if not isinstance(agent, Agent): + raise TypeError("registered agents must inherit Agent") + try: + agent.tools + except AttributeError as exc: + raise TypeError("Agent subclasses must call super().__init__()") from exc + + state = _AgentState(name=name) + subscriptions = self._discover_subscriptions(agent) + known_topics = dict(self._topics) + for topic, _method in subscriptions: + known = known_topics.setdefault(topic.name, topic) + if known.message_type is not topic.message_type: + raise ValueError(f"topic {topic.name!r} already uses another message type") + + self._agents[name] = state + for topic, method in subscriptions: + self._register_topic(topic) + self._subscribers.setdefault(topic.name, []).append((state, method)) + return agent + + async def start(self) -> None: + """Start accepting event publications.""" + + if self._closed: + raise RuntimeClosedError("agent runtime is closed") + if self._running: + return + self._running = True + + async def publish( + self, + topic: Topic[MessageT], + message: MessageT | dict[str, Any], + *, + participant_id: str, + source: str = "application", + ) -> None: + """Validate and deliver one event to every topic subscriber.""" + + await self._publish( + topic, + message, + participant_id=participant_id, + source=source, + ) + + async def stop(self) -> None: + """Stop accepting events and cancel in-flight deliveries.""" + + if self._closed: + return + self._running = False + tasks = tuple( + task for state in self._agents.values() for task in state.deliveries + ) + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + self._closed = True + + async def __aenter__(self) -> AgentRuntime: + await self.start() + return self + + async def __aexit__(self, *_exc: object) -> None: + await self.stop() + + async def _publish( + self, + topic: Topic[MessageT], + message: MessageT | dict[str, Any], + *, + participant_id: str, + source: str, + correlation_id: str | None = None, + parent_message_id: str | None = None, + ) -> None: + self._ensure_running() + self._register_topic(topic) + value = topic.validate(message) + metadata = self._metadata( + participant_id=participant_id, + source=source, + correlation_id=correlation_id, + parent_message_id=parent_message_id, + ) + deliveries: list[asyncio.Task[None]] = [] + for state, method in tuple(self._subscribers.get(topic.name, ())): + task = asyncio.create_task( + self._deliver( + state, + method, + value.model_copy(deep=True), + metadata, + ), + name=f"agent:{state.name}:subscription:{topic.name}", + ) + state.deliveries.add(task) + task.add_done_callback(state.deliveries.discard) + deliveries.append(task) + if deliveries: + results = await asyncio.gather(*deliveries, return_exceptions=True) + errors = [result for result in results if isinstance(result, BaseException)] + if errors: + raise BaseExceptionGroup("errors during event publication", errors) + + async def _deliver( + self, + state: _AgentState, + method: BoundSubscriber, + message: BaseModel, + metadata: MessageMetadata, + ) -> None: + await method(message, RuntimeContext(self, state.name, metadata)) + + def _discover_subscriptions( + self, + agent: Agent, + ) -> list[tuple[Topic[Any], BoundSubscriber]]: + subscriptions: list[tuple[Topic[Any], BoundSubscriber]] = [] + for _name, method in inspect.getmembers(agent, predicate=inspect.ismethod): + topics = cast(tuple[Topic[Any], ...], getattr(method, "__xr_ai_topics__", ())) + if not topics: + continue + if not inspect.iscoroutinefunction(method): + raise TypeError("topic subscribers must be async") + request_type = self._request_type(method) + for topic in topics: + if topic.message_type is not request_type: + raise TypeError( + f"subscriber for {topic.name!r} must accept " + f"{topic.message_type.__name__}" + ) + subscriptions.append((topic, cast(BoundSubscriber, method))) + return subscriptions + + @staticmethod + def _request_type(method: Callable[..., Any]) -> type[BaseModel]: + parameters = tuple(inspect.signature(method).parameters.values()) + if len(parameters) != 2: + raise TypeError("subscribers must accept exactly (message, context)") + annotation = get_type_hints(method).get(parameters[0].name) + if not isinstance(annotation, type) or not issubclass(annotation, BaseModel): + raise TypeError("subscriber messages must be annotated Pydantic models") + return annotation + + def _register_topic(self, topic: Topic[Any]) -> None: + known = self._topics.setdefault(topic.name, topic) + if known.message_type is not topic.message_type: + raise ValueError(f"topic {topic.name!r} already uses another message type") + + def _ensure_running(self) -> None: + if not self.running: + raise RuntimeClosedError("agent runtime is not running") + + @staticmethod + def _metadata( + *, + participant_id: str, + source: str, + correlation_id: str | None, + parent_message_id: str | None, + ) -> MessageMetadata: + if not participant_id.strip(): + raise ValueError("participant_id must not be empty") + if not source.strip(): + raise ValueError("message source must not be empty") + message_id = uuid.uuid4().hex + return MessageMetadata( + message_id=message_id, + correlation_id=correlation_id or message_id, + participant_id=participant_id, + source=source, + parent_message_id=parent_message_id, + timestamp_us=time.time_ns() // 1_000, + ) + + +__all__ = [ + "AgentRuntime", + "RuntimeClosedError", + "RuntimeContext", +] diff --git a/agent-sdk/xr-ai-tools/README.md b/agent-sdk/xr-ai-tools/README.md index f33ee8483..f31eaf030 100644 --- a/agent-sdk/xr-ai-tools/README.md +++ b/agent-sdk/xr-ai-tools/README.md @@ -64,7 +64,30 @@ for call in response.tool_calls or (): `handle_tool_call(...)` validates and invokes one model-produced `ToolCall`, then returns a tool-role `ChatMessage` plus its `return_direct` hint. The application or agent owns prompts, model calls, conversation state, iteration policy, and -whether calls run sequentially or concurrently. +whether calls run sequentially or concurrently. A unary side-effect tool uses +`result_model=None`, returns `None`, and produces `null` as its model-visible +result. + +Tool catalogs can assign model-visible aliases without wrapping or changing the +underlying tools: + +```python +tools = ToolSet({"camera_status": vision.status}) +``` + +When composing independently named tool groups, namespace them at the workflow +boundary: + +```python +tools = ToolSet.namespaced({ + "vision": vision.tools, + "planner": planner.tools, +}) +``` + +This exposes names such as `vision__status` and `planner__status` to the model. +Only finite `Tool` instances belong in a `ToolSet`; streaming `AsyncTool` +instances are consumed explicitly with `stream()`. ## Finite and streaming live vision tools diff --git a/agent-sdk/xr-ai-tools/xr_ai_tools/tool_calling.py b/agent-sdk/xr-ai-tools/xr_ai_tools/tool_calling.py index 5799e8232..4b5da8278 100644 --- a/agent-sdk/xr-ai-tools/xr_ai_tools/tool_calling.py +++ b/agent-sdk/xr-ai-tools/xr_ai_tools/tool_calling.py @@ -23,16 +23,23 @@ class ToolCallResult: return_direct: bool -def tool_definitions(tools: Iterable[Tool[Any, Any]]) -> tuple[ToolDef, ...]: +def tool_definitions( + tools: Iterable[Tool[Any, Any]] | ToolSet, +) -> tuple[ToolDef, ...]: """Return model-service definitions for native tools.""" + entries = ( + tools.items() + if isinstance(tools, ToolSet) + else tuple((tool.name, tool) for tool in tools) + ) return tuple( ToolDef( - name=tool.name, + name=name, description=tool.description, parameters=tool.request_model.model_json_schema(), ) - for tool in tools + for name, tool in entries ) diff --git a/agent-sdk/xr-ai-tools/xr_ai_tools/tools.py b/agent-sdk/xr-ai-tools/xr_ai_tools/tools.py index 691eb2df5..099b71ffb 100644 --- a/agent-sdk/xr-ai-tools/xr_ai_tools/tools.py +++ b/agent-sdk/xr-ai-tools/xr_ai_tools/tools.py @@ -6,7 +6,7 @@ from __future__ import annotations import json -from collections.abc import Awaitable, Callable, Iterable +from collections.abc import Awaitable, Callable, Iterable, Mapping from dataclasses import dataclass from inspect import isawaitable from typing import Any, Generic, TypeVar, cast @@ -15,7 +15,7 @@ from pydantic import BaseModel, ValidationError RequestT = TypeVar("RequestT", bound=BaseModel) -ResultT = TypeVar("ResultT", bound=BaseModel) +ResultT = TypeVar("ResultT") @dataclass(frozen=True, slots=True) @@ -34,7 +34,7 @@ def __init__( name: str, description: str, request_model: type[RequestT], - result_model: type[ResultT], + result_model: type[BaseModel] | None, handler: Callable[[RequestT], Awaitable[ResultT] | ResultT], *, return_direct: bool = False, @@ -51,7 +51,14 @@ def __init__( self.handler = handler self.return_direct = return_direct self._request_codec = typed.PydanticCodec(request_model) - self._result_codec = typed.PydanticCodec(result_model) + self._result_codec: typed.Codec[ResultT] + if result_model is None: + self._result_codec = cast(typed.Codec[ResultT], _NoneCodec()) + else: + self._result_codec = cast( + typed.Codec[ResultT], + typed.PydanticCodec(result_model), + ) self._render_result = render_result or _json_result async def execute(self, request: RequestT) -> ResultT: @@ -97,24 +104,79 @@ async def _execute_handler(self, request: RequestT) -> ResultT: class ToolSet: - """A non-overlapping native tool catalog.""" + """A native tool catalog with model-visible dispatch names.""" - def __init__(self, tools: Iterable[Tool[Any, Any]]) -> None: + def __init__( + self, + tools: Iterable[Tool[Any, Any]] | Mapping[str, 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 + if isinstance(tools, Mapping): + entries = tuple( + cast(Mapping[str, Tool[Any, Any]], tools).items() + ) + else: + entries = tuple( + (tool.name, tool) + for tool in cast(Iterable[Tool[Any, Any]], tools) + ) + for name, tool in entries: + if not name: + raise ValueError("tool alias must not be empty") + if not isinstance(tool, Tool): + raise TypeError("tool sets may contain only finite Tool instances") + if name in by_name: + raise ValueError(f"duplicate tool name: {name}") + by_name[name] = tool self._by_name = by_name + @classmethod + def namespaced( + cls, + namespaces: Mapping[str, Iterable[Tool[Any, Any]]], + ) -> ToolSet: + """Build a catalog named ``__`` for each group.""" + + aliases: dict[str, Tool[Any, Any]] = {} + for namespace, tools in namespaces.items(): + if not namespace: + raise ValueError("tool namespace must not be empty") + for tool in tools: + if not isinstance(tool, Tool): + raise TypeError("tool sets may contain only finite Tool instances") + name = f"{namespace}__{tool.name}" + if name in aliases: + raise ValueError(f"duplicate tool name: {name}") + aliases[name] = tool + return cls(aliases) + 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 items(self) -> tuple[tuple[str, Tool[Any, Any]], ...]: + """Return model-visible names paired with their underlying tools.""" + + return tuple(self._by_name.items()) + + +class _NoneCodec(typed.Codec[None]): + def to_json(self, value: None) -> None: + if value is not None: + raise TypeError("side-effect tools must return None") + return None + + def from_json(self, data: Any) -> None: + if data is not None: + raise TypeError("side-effect tools must return None") + return None + -def _json_result(result: BaseModel) -> str: - return result.model_dump_json() +def _json_result(result: Any) -> str: + if result is None: + return "null" + return cast(BaseModel, result).model_dump_json() __all__ = ["Tool", "ToolInvocationResult", "ToolSet"] diff --git a/docs/architecture.md b/docs/architecture.md index ff4ce0c1b..4ea4eeee0 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -13,7 +13,7 @@ design decisions see `docs/changelog.md`. ``` client-samples/ # Platform clients (Android, iOS/visionOS, Web) -agent-sdk/ # IPC, model, NAT-function, capability, and voice SDK packages +agent-sdk/ # Agent runtime, IPC, model, NAT-function, capability, and voice SDK packages utils/ # Shared infra: stdlib-only launcher + loguru logging bridge services/ # XR hub, CloudXR, model-serving, and typed XR capability services agent-mcp-servers/ # Optional MCP compatibility adapters for non-NAT consumers @@ -37,9 +37,16 @@ docs/ # Design docs and topic deep-dives - **`agent-sdk/xr-ai-hub-client`** contains only the agent-facing IPC layer. Its sole runtime dependencies are `pyzmq` and `msgpack` — no LiveKit, FastAPI, or uvicorn. -- **Native agents compose typed NAT functions in process.** Runtime-backed - functions call typed capability services, while deterministic functions run - locally. MCP adapters only republish selected functions for MCP consumers. +- **`agent-sdk/xr-ai-agent-runtime`** owns agent resource lifetimes, + runtime-owned background tasks, and typed publish/subscribe. Agents expose + existing `Tool` and `AsyncTool` instances from `xr-ai-tools`; direct callers + and model loops use those tools without a second runtime dispatch API. Each + agent owns its concurrency policy, including any locks or private queues + needed to coordinate tools and subscriptions. Model loops, planning, memory, + and raw media transport remain outside the runtime. +- **Native agents compose typed tools in process.** Model-backed tools call + typed capability services, while deterministic tools run locally. MCP + adapters only republish selected tools for MCP consumers. - **No API keys or tokens in source files** — use env vars or `xr_media_hub.yaml` (see `docs/credentials.md`). diff --git a/docs/changelog.md b/docs/changelog.md index 25d4964ad..2b457ccac 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -29,6 +29,29 @@ chunk is consumed. Producer cleanup remains unbounded: a timeout and detached cleanup policy requires a separate decision because abandoning cleanup could leave Relay state or participant status unfinished. +### 2026-08-12 — Agents own existing tools and their concurrency policy + +`xr-ai-agent-runtime` defines an `Agent` as a plain object containing private +state and ordinary `Tool` or `AsyncTool` instances from +`xr-ai-tools`. There is no separate function contract, agent address, runtime +call API, or tool adapter. Direct callers and other agents use `execute()` or +`stream()`; model-selected calls use the same `ToolSet` and +`handle_tool_call()` path as standalone tools. + +The runtime owns typed, participant-scoped `publish()` fan-out and quiesces only +the in-flight subscription deliveries it creates. Agent resources, background +tasks, and lifecycle remain entirely agent-owned. Tool handlers and subscription +callbacks may run concurrently; each agent owns any lock or private queue needed +for its state instead of receiving a mandatory mailbox. This avoids implicit +head-of-line blocking, especially for streams. + +Unary tools now permit `None` for acknowledged side effects. Streaming tools +remain `AsyncTool` instances and are never silently buffered into model tool +results. Agent lifetime is not a model tool; domain controls such as starting +or stopping monitoring remain ordinary tools. Raw media stays on the hub path, +and Relay remains the execution and telemetry boundary around model calls and +tool execution. + ### 2026-08-12 — Tool-call handling is not an agent runtime `agents.py`, `agent_runner.py`, `Agent`, and `AgentRunner` are removed. diff --git a/docs/nemo-agent-toolkit-migration.md b/docs/nemo-agent-toolkit-migration.md index d1c5a39f4..88fc7be1d 100644 --- a/docs/nemo-agent-toolkit-migration.md +++ b/docs/nemo-agent-toolkit-migration.md @@ -14,9 +14,12 @@ for all model HTTP, and reach clients only through the Hub IPC SDK. NeMo Relay is the local execution boundary. The XR-owned `xr-ai-tools` package is the toolkit-independent tools layer: Pydantic tool schemas, trigger dispatch, and small helpers that adapt those tools to `xr-ai-models` tool-call -types. Applications own their agents, model calls, history, and loop policy. -Relay owns tool lifecycles, middleware, guardrails, and telemetry. Existing -NeMo Agent Toolkit function groups remain compatibility extras until their concrete tools migrate. +types. `xr-ai-agent-runtime` separately owns agent resource lifetimes, +runtime-owned background tasks, and fan-out `publish()`. Agents expose the +existing `Tool` and `AsyncTool` objects directly and own their synchronization. +Applications own their model calls, history, and loop policy. Relay owns 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 @@ -27,6 +30,7 @@ remain optional launch targets after the local runtime has migrated. ```text XR worker + -> xr-ai-agent-runtime: agent lifetimes, background tasks, and publish -> xr-ai-tools: typed tools and trigger dispatch -> NeMo Relay: managed tool execution, guardrails, telemetry -> xr-ai-models: private model boundary used by model-backed tools @@ -57,15 +61,18 @@ acceptance behavior rather than an implementation dependency. 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, +4. **Agent runtime** — add runtime-owned agent resource lifetimes, background + tasks, typed `publish()`, and agents that expose existing unary and streaming + tools directly while owning their concurrency policy. +5. **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 +6. **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 +7. **Tea-making sample** — land the sample in slices: workflow/state core, + foreground and resident agents, observation applications, then activity + viewer and end-to-end evaluations. +8. **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. @@ -73,6 +80,10 @@ 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. +Fabric integration should adapt the runtime lifetime at the hosting boundary +or provide an `Agent` implementation. It must use the existing tool APIs rather +than introduce a second application invocation path alongside `publish()`. + ## Exit criteria The retirement is complete only when repository-wide search finds no runtime diff --git a/docs/source/components/agent-sdk.md b/docs/source/components/agent-sdk.md index 579c5d0bf..0f69ed052 100644 --- a/docs/source/components/agent-sdk.md +++ b/docs/source/components/agent-sdk.md @@ -8,6 +8,8 @@ The `agent-sdk/` workspace holds the libraries an xr-ai agent is built from: +- **`xr-ai-agent-runtime`** — agents exposing existing native tools and typed + publish/subscribe. - **`xr-ai-models`** — unified service protocols (`LLMService`, `VLMService`, `STTService`, `TTSService`, `EmbeddingService`) plus OpenAI-compatible HTTP clients, driven by a structured model deployment profile. Swapping a backend is a configuration @@ -28,6 +30,31 @@ from: --- +## xr-ai-agent-runtime + +`AgentRuntime` provides participant-scoped publish/subscribe. An agent exposes +ordinary `Tool` and `AsyncTool` instances from `xr-ai-tools`. Direct callers +use `execute()` or `stream()`, and model loops use the same `ToolSet` and +`handle_tool_call()` path as standalone tools. There is no runtime call +adapter. + +Agent lifetime is not itself a runtime concern. Domain controls such as starting +or stopping monitoring remain ordinary tools. Agents are registered before +runtime startup and own their resources, tasks, and concurrency policy. Shared +state that is touched by tools and subscriptions must be protected by the +agent's lock or private queue. Model loops, planning, and memory remain agent +implementations. Raw audio and video remain on the hub path. + +`ToolSet.namespaced({"vision": vision.tools, "planner": planner.tools})` +assigns unique model-visible names when tools from multiple agents are combined; +the agents and underlying tools remain unchanged. Participant identity needed +by direct execution belongs in the tool's request schema; participant and +correlation metadata on `RuntimeContext` applies to pub/sub. Agents create, +cancel, and await their own background tasks. `publish()` settles all fan-out +deliveries before propagating subscriber failures. + +--- + ## xr-ai-models Worker code depends on the service protocols and constructs concrete @@ -227,6 +254,11 @@ boundary. It has no voice or output-transport dependency. Both tools own their own frame sources and redact inline camera data from Relay events without changing provider input. +Applications place native tools and model loops inside an agent to group tools +with their state, and register the agent with `xr-ai-agent-runtime` when they +need pub/sub. Relay remains responsible for tool and model execution; the +runtime does not duplicate the tool loop or model boundary. + ## xr-ai-nat model bridge Unmigrated workflows install `xr-ai-nat[agents]` when they use NAT's built-in diff --git a/tests/pyproject.toml b/tests/pyproject.toml index a97fe75cb..9b999958b 100644 --- a/tests/pyproject.toml +++ b/tests/pyproject.toml @@ -12,6 +12,7 @@ description = "Multi-client / multi-agent integration tests for xr-ai." requires-python = ">=3.11,<3.13" dependencies = [ # Pulled in via editable installs of the workspace packages. + "xr-ai-agent-runtime", "xr-ai-hub-client", "xr-ai-models", "xr-ai-nat[agents,services,vision]", @@ -44,6 +45,7 @@ dependencies = [ ] [tool.uv.sources] +xr-ai-agent-runtime = { path = "../agent-sdk/xr-ai-agent-runtime", editable = true } xr-ai-hub-client = { path = "../agent-sdk/xr-ai-hub-client", editable = true } xr-ai-models = { path = "../agent-sdk/xr-ai-models", editable = true } xr-ai-nat = { path = "../agent-sdk/xr-ai-nat", editable = true } diff --git a/tests/test_agent_runtime.py b/tests/test_agent_runtime.py new file mode 100644 index 000000000..d7a44ef65 --- /dev/null +++ b/tests/test_agent_runtime.py @@ -0,0 +1,444 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Contracts for agent-owned tools, state, and typed pub/sub.""" + +from __future__ import annotations + +import asyncio +from builtins import ExceptionGroup +from collections.abc import AsyncIterator +from typing import assert_type + +import pytest +from pydantic import BaseModel +from xr_ai_models import ChatMessage, ToolCall +from xr_ai_runtime import ( + Agent, + AgentRuntime, + RuntimeClosedError, + RuntimeContext, + Topic, + subscribe, +) +from xr_ai_tools import AsyncTool, Tool, ToolSet +from xr_ai_tools.tool_calling import handle_tool_call + + +class _Echo(BaseModel): + text: str + + +class _Count(BaseModel): + amount: int + + +class _Counted(BaseModel): + count: int + + +class _Observation(BaseModel): + labels: list[str] + + +class _Speak(BaseModel): + text: str + + +OBSERVATIONS = Topic("vision.observation", _Observation) +VOICE_OUTPUT = Topic("voice.output", _Speak) +TRANSACTIONAL_ECHO = Topic("transactional", _Echo) +TRANSACTIONAL_COUNT = Topic("transactional", _Count) + + +class _TextAgent(Agent): + def __init__(self) -> None: + self.recorded: list[str] = [] + self.echo = Tool("echo", "Echo text.", _Echo, _Echo, self._echo) + self.uppercase = Tool( + "uppercase", + "Uppercase text.", + _Echo, + _Echo, + self._uppercase, + ) + self.record = Tool("record", "Record text.", _Echo, None, self._record) + self.stream_echo = AsyncTool( + "stream_echo", + "Stream echoed text.", + _Echo, + _Echo, + self._stream_echo, + ) + super().__init__((self.echo, self.uppercase, self.record, self.stream_echo)) + + async def _echo(self, request: _Echo) -> _Echo: + return request + + async def _uppercase(self, request: _Echo) -> _Echo: + return _Echo(text=request.text.upper()) + + async def _record(self, request: _Echo) -> None: + self.recorded.append(request.text) + + async def _stream_echo(self, request: _Echo) -> AsyncIterator[_Echo]: + yield _Echo(text=f"{request.text}-one") + yield _Echo(text=f"{request.text}-two") + + +async def test_agent_exposes_existing_unary_and_streaming_tools() -> None: + agent = _TextAgent() + runtime = AgentRuntime() + + registered = runtime.register("text", agent) + echoed = await agent.echo.execute(_Echo(text="hello")) + chunks = [ + chunk async for chunk in agent.stream_echo.stream({"text": "chunk"}) + ] + + assert registered is agent + assert registered.tools == ( + agent.echo, + agent.uppercase, + agent.record, + agent.stream_echo, + ) + assert_type(echoed, _Echo) + assert echoed == _Echo(text="hello") + assert chunks == [_Echo(text="chunk-one"), _Echo(text="chunk-two")] + + +async def test_agent_tool_uses_normal_model_tool_calling() -> None: + agent = _TextAgent() + tools = ToolSet.namespaced({"text": (agent.uppercase,)}) + + result = await handle_tool_call( + ToolCall( + id="uppercase-call", + name="text__uppercase", + arguments='{"text":"hello"}', + ), + tools, + ) + + assert result.message == ChatMessage( + role="tool", + content='{"text":"HELLO"}', + tool_call_id="uppercase-call", + ) + + +async def test_side_effect_tool_returns_none_through_direct_and_model_calls() -> None: + agent = _TextAgent() + + direct = await agent.record.execute(_Echo(text="direct")) + model = await agent.record.invoke('{"text":"model"}') + + assert direct is None + assert model.content == "null" + assert agent.recorded == ["direct", "model"] + + +class _ForwardingAgent(Agent): + def __init__(self, target: Tool[_Echo, _Echo]) -> None: + self._target = target + self.forward = Tool("forward", "Forward text.", _Echo, _Echo, self._forward) + super().__init__((self.forward,)) + + async def _forward(self, request: _Echo) -> _Echo: + return await self._target.execute(request) + + +async def test_agent_uses_another_agents_tool_without_runtime_dispatch() -> None: + target = _TextAgent() + caller = _ForwardingAgent(target.echo) + runtime = AgentRuntime() + runtime.register("target", target) + runtime.register("caller", caller) + + result = await caller.forward.execute(_Echo(text="hello")) + + assert result == _Echo(text="hello") + + +class _VoiceOutput(Agent): + def __init__(self, *, mutate: bool = False) -> None: + super().__init__() + self.mutate = mutate + self.received: list[tuple[str, str, list[str]]] = [] + + @subscribe(OBSERVATIONS) + async def observe(self, event: _Observation, ctx: RuntimeContext) -> None: + self.received.append( + (ctx.metadata.participant_id, ctx.metadata.source, list(event.labels)) + ) + if self.mutate: + event.labels.append("mutated") + + +async def test_publish_fans_out_isolated_typed_payloads() -> None: + first = _VoiceOutput(mutate=True) + second = _VoiceOutput() + runtime = AgentRuntime() + runtime.register("first", first) + runtime.register("second", second) + original = _Observation(labels=["kettle"]) + + async with runtime: + result = await runtime.publish( + OBSERVATIONS, + original, + participant_id="alice", + source="camera", + ) + + assert result is None + assert original == _Observation(labels=["kettle"]) + assert first.received == [("alice", "camera", ["kettle"])] + assert second.received == [("alice", "camera", ["kettle"])] + + +class _FailingSubscriber(Agent): + def __init__(self, *, failure: Exception | None = None) -> None: + super().__init__() + self.failure = failure + self.completed = False + + @subscribe(OBSERVATIONS) + async def observe(self, _event: _Observation, _ctx: RuntimeContext) -> None: + await asyncio.sleep(0) + self.completed = True + if self.failure is not None: + raise self.failure + + +async def test_publish_settles_every_delivery_then_propagates_failures() -> None: + failed = _FailingSubscriber(failure=RuntimeError("subscriber failed")) + completed = _FailingSubscriber() + runtime = AgentRuntime() + runtime.register("failed", failed) + runtime.register("completed", completed) + await runtime.start() + + with pytest.raises(ExceptionGroup, match="event publication") as raised: + await runtime.publish( + OBSERVATIONS, + _Observation(labels=[]), + participant_id="alice", + ) + + await runtime.stop() + assert failed.completed + assert completed.completed + assert [str(error) for error in raised.value.exceptions] == ["subscriber failed"] + + +class _EventForwarder(Agent): + def __init__(self) -> None: + super().__init__() + self.input_message_id = "" + + @subscribe(OBSERVATIONS) + async def forward(self, _event: _Observation, ctx: RuntimeContext) -> None: + self.input_message_id = ctx.metadata.message_id + await ctx.publish(VOICE_OUTPUT, _Speak(text="Seen.")) + + +class _Speaker(Agent): + def __init__(self) -> None: + super().__init__() + self.metadata = None + + @subscribe(VOICE_OUTPUT) + async def speak(self, _event: _Speak, ctx: RuntimeContext) -> None: + self.metadata = ctx.metadata + + +async def test_nested_publish_preserves_participant_and_trace_context() -> None: + forwarder = _EventForwarder() + speaker = _Speaker() + runtime = AgentRuntime() + runtime.register("forwarder", forwarder) + runtime.register("speaker", speaker) + + async with runtime: + await runtime.publish( + OBSERVATIONS, + _Observation(labels=["cup"]), + participant_id="alice", + source="camera", + ) + + assert speaker.metadata is not None + assert speaker.metadata.participant_id == "alice" + assert speaker.metadata.source == "forwarder" + assert speaker.metadata.parent_message_id == forwarder.input_message_id + assert speaker.metadata.correlation_id == forwarder.input_message_id + + +class _SerializedAgent(Agent): + def __init__(self) -> None: + self._lock = asyncio.Lock() + self._release = asyncio.Event() + self.tool_entered = asyncio.Event() + self.active = 0 + self.max_active = 0 + self.count = 0 + self.increment = Tool( + "increment", + "Increment shared state.", + _Count, + _Counted, + self._increment, + ) + super().__init__((self.increment,)) + + async def _mutate(self, amount: int, *, wait: bool = False) -> None: + async with self._lock: + self.active += 1 + self.max_active = max(self.max_active, self.active) + if wait: + self.tool_entered.set() + await self._release.wait() + self.count += amount + self.active -= 1 + + async def _increment(self, request: _Count) -> _Counted: + await self._mutate(request.amount, wait=True) + return _Counted(count=self.count) + + @subscribe(OBSERVATIONS) + async def observe(self, _event: _Observation, _ctx: RuntimeContext) -> None: + await self._mutate(1) + + +async def test_agent_can_serialize_tools_and_subscriptions_internally() -> None: + agent = _SerializedAgent() + runtime = AgentRuntime() + runtime.register("serialized", agent) + + async with runtime: + tool_call = asyncio.create_task(agent.increment.execute(_Count(amount=2))) + await agent.tool_entered.wait() + publication = asyncio.create_task( + runtime.publish( + OBSERVATIONS, + _Observation(labels=["kettle"]), + participant_id="alice", + ) + ) + await asyncio.sleep(0) + assert agent.max_active == 1 + agent._release.set() + assert await tool_call == _Counted(count=2) + assert await publication is None + + assert agent.count == 3 + assert agent.max_active == 1 + + +class _SlowSubscriber(Agent): + def __init__(self) -> None: + super().__init__() + self.started = asyncio.Event() + self.stopped = asyncio.Event() + + @subscribe(OBSERVATIONS) + async def observe(self, _event: _Observation, _ctx: RuntimeContext) -> None: + self.started.set() + try: + await asyncio.Event().wait() + finally: + self.stopped.set() + + +async def test_stop_cancels_in_flight_deliveries() -> None: + agent = _SlowSubscriber() + runtime = AgentRuntime() + runtime.register("slow", agent) + await runtime.start() + publication = asyncio.create_task( + runtime.publish( + OBSERVATIONS, + _Observation(labels=[]), + participant_id="alice", + ) + ) + await agent.started.wait() + + await runtime.stop() + await asyncio.gather(publication, return_exceptions=True) + + assert agent.stopped.is_set() + + +class _PartiallyInvalidSubscriber(Agent): + def __init__(self) -> None: + super().__init__() + + @subscribe(TRANSACTIONAL_ECHO) + async def a_valid(self, _message: _Echo, _ctx: RuntimeContext) -> None: + return None + + @subscribe(TRANSACTIONAL_ECHO) + async def z_invalid(self, _message: str, _ctx: RuntimeContext) -> None: + return None + + +class _TransactionalSubscriber(Agent): + def __init__(self) -> None: + super().__init__() + + @subscribe(TRANSACTIONAL_COUNT) + async def receive(self, _message: _Count, _ctx: RuntimeContext) -> None: + return None + + +def test_registration_failure_does_not_mutate_topic_state() -> None: + runtime = AgentRuntime() + + with pytest.raises(TypeError, match="Pydantic models"): + runtime.register("invalid", _PartiallyInvalidSubscriber()) + + runtime.register("valid", _TransactionalSubscriber()) + + +def test_agent_rejects_duplicate_or_non_tool_members() -> None: + tool = Tool("echo", "Echo text.", _Echo, _Echo, lambda request: request) + + with pytest.raises(ValueError, match="duplicate"): + Agent((tool, tool)) + with pytest.raises(TypeError, match="Tool or AsyncTool"): + Agent((object(),)) # type: ignore[arg-type] + + +async def test_runtime_rejects_invalid_registration_and_publication() -> None: + runtime = AgentRuntime() + runtime.register("agent", Agent()) + uninitialized_agent = Agent.__new__(Agent) + + with pytest.raises(TypeError, match=r"super\(\)"): + AgentRuntime().register("invalid", uninitialized_agent) + with pytest.raises(ValueError, match="already registered"): + runtime.register("agent", Agent()) + with pytest.raises(RuntimeClosedError): + await runtime.publish( + OBSERVATIONS, + _Observation(labels=[]), + participant_id="alice", + ) + + async with runtime: + with pytest.raises(ValueError, match="must not be empty"): + await runtime.publish( + OBSERVATIONS, + _Observation(labels=[]), + participant_id="", + ) + + with pytest.raises(RuntimeClosedError): + await runtime.publish( + OBSERVATIONS, + _Observation(labels=[]), + participant_id="alice", + ) diff --git a/tests/test_native_tools.py b/tests/test_native_tools.py index 0cde60793..8b1fb8c94 100644 --- a/tests/test_native_tools.py +++ b/tests/test_native_tools.py @@ -40,6 +40,10 @@ async def add_stream(request: AddRequest) -> AsyncIterator[AddResult]: yield AddResult(total=request.left + request.right) +async def acknowledge(_request: AddRequest) -> None: + return None + + async def test_async_tool_validates_and_yields_typed_chunks() -> None: tool = AsyncTool( "stream_add", @@ -224,6 +228,16 @@ async def test_handle_tool_call_returns_a_model_ready_tool_message() -> None: assert result.return_direct is False +async def test_side_effect_tool_returns_none_and_renders_null() -> None: + tool = Tool("acknowledge", "Acknowledge input.", AddRequest, None, acknowledge) + + direct = await tool.execute(AddRequest(left=2, right=3)) + model = await tool.invoke('{"left":2,"right":3}') + + assert direct is None + assert model.content == "null" + + async def test_handled_tool_calls_use_the_relay_tool_lifecycle() -> None: tool = Tool("add", "Add two integers.", AddRequest, AddResult, add) events = [] @@ -294,3 +308,41 @@ def test_tool_sets_reject_duplicate_names() -> None: with pytest.raises(ValueError, match="duplicate tool name"): ToolSet((tool, tool)) + + +async def test_tool_set_aliases_remap_model_definition_and_dispatch_names() -> None: + tool = Tool("add", "Add two integers.", AddRequest, AddResult, add) + tools = ToolSet({"sum": tool}) + + assert tool_definitions(tools) == ( + ToolDef( + name="sum", + description="Add two integers.", + parameters=AddRequest.model_json_schema(), + ), + ) + result = await handle_tool_call( + ToolCall(id="sum-call", name="sum", arguments='{"left":2,"right":3}'), + tools, + ) + + assert result.message.content == '{"total":5}' + assert tool.name == "add" + + +def test_tool_set_namespaces_similarly_named_tool_groups() -> None: + vision_status = Tool("status", "Vision status.", AddRequest, AddResult, add) + planner_status = Tool("status", "Planner status.", AddRequest, AddResult, add) + tools = ToolSet.namespaced( + { + "vision": (vision_status,), + "planner": (planner_status,), + } + ) + + assert [definition.name for definition in tool_definitions(tools)] == [ + "vision__status", + "planner__status", + ] + assert tools.get("vision__status") is vision_status + assert tools.get("planner__status") is planner_status