diff --git a/nemoguardrails/guardrails/actions/tool_call_action.py b/nemoguardrails/guardrails/actions/tool_call_action.py new file mode 100644 index 0000000000..ce04788060 --- /dev/null +++ b/nemoguardrails/guardrails/actions/tool_call_action.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tool-call safety rail for IORails. + +Validates the tool calls a model emitted against the request's declared +``Toolset``: every call must name an allowed tool, and its arguments must +satisfy that tool's JSON Schema. The rail is local and model-free -- it runs +through :meth:`ToolRailAction._guarded`, so a malformed call or an unexpected +error fails closed (blocks) rather than propagating. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, List + +from nemoguardrails.guardrails.guardrails_types import RailResult +from nemoguardrails.guardrails.tool_rail_action import ToolRailAction +from nemoguardrails.guardrails.tool_schema import validate_arguments + +if TYPE_CHECKING: + from nemoguardrails.guardrails.tool_schema import Toolset + from nemoguardrails.types import ToolCall + + +class ToolCallRailAction(ToolRailAction): + """Check the model's tool calls against the declared toolset (allowlist + schema).""" + + action_name = "tool call validation" + + async def run(self, toolset: "Toolset", tool_calls: List["ToolCall"]) -> RailResult: + """Block unless every tool call names an allowed tool with schema-valid arguments.""" + return self._guarded(lambda: self._validate(toolset, tool_calls)) + + def _validate(self, toolset: "Toolset", tool_calls: List["ToolCall"]) -> RailResult: + """Allowlist each call by name, then validate its arguments against the tool schema.""" + for call in tool_calls: + # Hosted/server tools (e.g. web_search) have no function name; fall back to + # call.type, mirroring Tool.key = name or type used when indexing the toolset. + name = call.function.name or call.type + tool = toolset.get(name) + if tool is None: + return RailResult(is_safe=False, reason=f"tool call '{name}' is not an allowed tool") + block_reason = validate_arguments(tool, call.function.arguments) + if block_reason is not None: + return RailResult(is_safe=False, reason=block_reason) + return RailResult(is_safe=True) diff --git a/nemoguardrails/guardrails/actions/tool_result_action.py b/nemoguardrails/guardrails/actions/tool_result_action.py new file mode 100644 index 0000000000..395fd1a15c --- /dev/null +++ b/nemoguardrails/guardrails/actions/tool_result_action.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tool-result validation rail for IORails. + +Structurally validates the tool results carried on an incoming request against +the tool calls the model previously made: every result must link to a prior +call by ``call_id``, name a tool consistent with that call, and carry +well-formed content. This PR validates structure only -- there are no declared +response schemas yet. The rail is local and model-free; it runs through +:meth:`ToolRailAction._guarded`, so a malformed result or an unexpected error +fails closed (blocks) rather than propagating. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, List + +from nemoguardrails.guardrails.guardrails_types import RailResult +from nemoguardrails.guardrails.tool_rail_action import ToolRailAction + +if TYPE_CHECKING: + from nemoguardrails.guardrails.tool_schema import ToolResult + from nemoguardrails.types import ToolCall + + +def _is_well_formed_content(content: object) -> bool: + """Tool-result content is a string, or a list of content-block dicts. + + Matches the declared ``ToolResult.content`` type (``str | list[dict] | None``); + a list of non-dict values (e.g. ``[1, 2, 3]``) is not well-formed. + """ + if isinstance(content, str): + return True + return isinstance(content, list) and all(isinstance(block, dict) for block in content) + + +class ToolResultRailAction(ToolRailAction): + """Check incoming tool results link to a prior call and are structurally well-formed.""" + + action_name = "tool result validation" + + async def run(self, tool_results: List["ToolResult"], prior_calls: List["ToolCall"]) -> RailResult: + """Block unless every tool result links to a prior call with a consistent name and valid content.""" + return self._guarded(lambda: self._validate(tool_results, prior_calls)) + + def _validate(self, tool_results: List["ToolResult"], prior_calls: List["ToolCall"]) -> RailResult: + """Check call_id linkage, name consistency, and content shape for each result.""" + calls_by_id = self._validate_prior_calls(prior_calls) + if isinstance(calls_by_id, RailResult): + return calls_by_id + return self._validate_results(tool_results, calls_by_id) + + def _validate_prior_calls(self, prior_calls: List["ToolCall"]) -> "RailResult | dict[str, ToolCall]": + """Build a call_id index from prior_calls; return a blocking RailResult on duplicate IDs.""" + calls_by_id: dict[str, "ToolCall"] = {} + for call in prior_calls: + if not call.id: + continue + if call.id in calls_by_id: + return RailResult( + is_safe=False, + reason=f"duplicate prior tool call id '{call.id}' makes tool-result linkage ambiguous", + ) + calls_by_id[call.id] = call + return calls_by_id + + def _validate_results(self, tool_results: List["ToolResult"], calls_by_id: "dict[str, ToolCall]") -> RailResult: + """Check each result links to a prior call with a consistent name and well-formed content.""" + rail_result = self._validate_tool_result_ids(tool_results) + if rail_result: + return rail_result + + for result in tool_results: + rail_result = self._validate_result_call_id(result, calls_by_id) + if rail_result: + return rail_result + + prior = calls_by_id[result.call_id] # type: ignore[index] + rail_result = self._validate_result_name(result, prior) + if rail_result: + return rail_result + + rail_result = self._validate_result_content(result) + if rail_result: + return rail_result + + return RailResult(is_safe=True) + + def _validate_tool_result_ids(self, tool_results: List["ToolResult"]) -> "RailResult | None": + """Return a blocking RailResult if any call_id appears more than once in the result list.""" + seen: set[str] = set() + for result in tool_results: + if not result.call_id: + continue + if result.call_id in seen: + return RailResult( + is_safe=False, + reason=f"duplicate tool result for call_id '{result.call_id}': each tool call must have exactly one result", + ) + seen.add(result.call_id) + return None + + def _validate_result_call_id(self, result: "ToolResult", calls_by_id: "dict[str, ToolCall]") -> "RailResult | None": + """Return a blocking RailResult if the result is missing a call_id or it has no prior call.""" + call_id = result.call_id + if not call_id: + return RailResult(is_safe=False, reason="tool result is missing a call_id") + if calls_by_id.get(call_id) is None: + return RailResult( + is_safe=False, + reason=f"tool result for call_id '{call_id}' does not correspond to a prior tool call", + ) + return None + + def _validate_result_name(self, result: "ToolResult", prior: "ToolCall") -> "RailResult | None": + """Return a blocking RailResult if the result name conflicts with the prior call's function name.""" + if result.name and prior.function.name and result.name != prior.function.name: + return RailResult( + is_safe=False, + reason=( + f"tool result name '{result.name}' does not match the called tool " + f"'{prior.function.name}' for call_id '{result.call_id}'" + ), + ) + return None + + def _validate_result_content(self, result: "ToolResult") -> "RailResult | None": + """Return a blocking RailResult if the result content is not a string or list of dicts.""" + if result.content is not None and not _is_well_formed_content(result.content): + return RailResult( + is_safe=False, + reason=f"tool result for call_id '{result.call_id}' has malformed content", + ) + return None diff --git a/nemoguardrails/guardrails/engine_registry.py b/nemoguardrails/guardrails/engine_registry.py index 3fd6b2c249..e92e7dd86f 100644 --- a/nemoguardrails/guardrails/engine_registry.py +++ b/nemoguardrails/guardrails/engine_registry.py @@ -36,6 +36,7 @@ set_llm_request_attributes, set_llm_response_attributes, ) +from nemoguardrails.guardrails.tool_schema import ToolResult, Toolset from nemoguardrails.rails.llm.config import Model, RailsConfigData from nemoguardrails.tracing.constants import ( llm_operation_duration, @@ -370,6 +371,33 @@ async def stream_model_call( if self._metrics_enabled: record_token_usage(engine.model_name, provider_name, operation_name, captured_usage) + def parse_tools(self, model_type: str, llm_params: Optional[dict]) -> Toolset: + """Parse the tool block in ``llm_params`` for the named model engine. + + Delegates to the engine's ``parse_tools`` so the provider-specific shape + (keyed on the engine) is normalized into a ``Toolset`` for the tool rails. + + Raises: + KeyError: If no engine is registered with the given name. + TypeError: If the named engine is not a ModelEngine. + """ + engine = self._get_engine(model_type, ModelEngine) + return engine.parse_tools({**engine.body_param_defaults, **(llm_params or {})}) + + def extract_tool_results(self, model_type: str, messages: list[dict]) -> list[ToolResult]: + """Extract incoming tool results from ``messages`` for the named model engine. + + Delegates to the engine's ``extract_tool_results`` so the provider's + tool-result messages are normalized into the ``ToolResult`` list the + ToolResultRail consumes. + + Raises: + KeyError: If no engine is registered with the given name. + TypeError: If the named engine is not a ModelEngine. + """ + engine = self._get_engine(model_type, ModelEngine) + return engine.extract_tool_results(messages) + async def api_call(self, api_name: str, message: dict[str, Any], **kwargs: Any) -> dict[str, Any]: """Route an API request to the named API engine. diff --git a/nemoguardrails/guardrails/model_engine.py b/nemoguardrails/guardrails/model_engine.py index e00fb31f0c..c662024db6 100644 --- a/nemoguardrails/guardrails/model_engine.py +++ b/nemoguardrails/guardrails/model_engine.py @@ -39,6 +39,7 @@ ) from nemoguardrails.guardrails.base_engine import BaseEngine from nemoguardrails.guardrails.guardrails_types import LLMMessages, get_request_id, truncate +from nemoguardrails.guardrails.tool_schema import Tool, ToolResult, Toolset from nemoguardrails.rails.llm.config import Model from nemoguardrails.types import ChatMessage, LLMResponse, LLMResponseChunk, ToolCall, ToolCallFunction, UsageInfo @@ -296,6 +297,79 @@ def _finalize_tool_calls(tool_calls: dict[int, dict]) -> list[ToolCall]: return result +def _parse_tools_openai(tools: list) -> list[Tool]: + """Parse OpenAI Chat Completions tool definitions into ``Tool`` objects. + + Each entry has the nested shape ``{"type": "function", "function": {"name", + "description", "parameters", "strict"}}``; ``function.parameters`` (the JSON + Schema) maps to ``Tool.arguments_schema``. Entries that are not a dict, lack a + ``function`` block, or whose function has no non-empty ``name`` are skipped. + """ + parsed: list[Tool] = [] + for entry in tools: + if not isinstance(entry, dict): + continue + function = entry.get("function") + if not isinstance(function, dict): + continue + name = function.get("name") + if not isinstance(name, str) or not name: + continue + parsed.append( + Tool( + name=name, + type=entry.get("type", "function"), + description=function.get("description"), + arguments_schema=function.get("parameters"), + strict=function.get("strict"), + ) + ) + return parsed + + +def _parse_tools_nim(tools: list) -> list[Tool]: + """Parse NIM tool definitions. NIM uses the OpenAI Chat Completions tool shape.""" + return _parse_tools_openai(tools) + + +_TOOL_PARSERS = { + "openai": _parse_tools_openai, + "nim": _parse_tools_nim, +} + + +def _extract_tool_results_openai(messages: LLMMessages) -> list[ToolResult]: + """Extract OpenAI Chat Completions tool results into ``ToolResult`` objects. + + Chat Completions carries each tool result as a top-level ``{"role": "tool", + "tool_call_id", "content"}`` message (optionally ``name``). This shape has no + error flag, so ``is_error`` is always ``False``. + """ + results: list[ToolResult] = [] + for message in messages: + if not isinstance(message, dict) or message.get("role") != "tool": + continue + results.append( + ToolResult( + call_id=message.get("tool_call_id"), + name=message.get("name"), + content=message.get("content"), + ) + ) + return results + + +def _extract_tool_results_nim(messages: LLMMessages) -> list[ToolResult]: + """Extract NIM tool results. NIM uses the OpenAI Chat Completions shape.""" + return _extract_tool_results_openai(messages) + + +_RESULT_EXTRACTORS = { + "openai": _extract_tool_results_openai, + "nim": _extract_tool_results_nim, +} + + class ModelEngineError(Exception): """Raised when a model engine call fails.""" @@ -658,3 +732,32 @@ async def stream_chat_completion( """ async for chunk in self.stream_call(messages, **kwargs): yield chunk + + def parse_tools(self, llm_params: Optional[dict]) -> Toolset: + """Parse the provider tool block in ``llm_params`` into a ``Toolset``. + + Reads the opaque ``tools`` block forwarded via + ``GenerationOptions.llm_params`` and normalizes it into the internal + ``Toolset`` the tool rails validate against, keyed on the model's engine + (``_TOOL_PARSERS``). OpenAI and NIM share the Chat Completions shape; an + engine with no registered parser falls back to it. Returns an empty + ``Toolset`` when no tools are declared. + """ + tools = (llm_params or {}).get("tools") + if not tools: + return Toolset() + parser = _TOOL_PARSERS.get(self.model_config.engine, _parse_tools_openai) + return Toolset(tools=parser(tools)) + + def extract_tool_results(self, messages: LLMMessages) -> list[ToolResult]: + """Extract incoming tool results from ``messages`` into ``ToolResult`` objects. + + Pulls the provider's tool-result messages out of the conversation and + normalizes them into the internal ``ToolResult`` shape the ToolResultRail + consumes, keyed on the model's engine (``_RESULT_EXTRACTORS``). OpenAI and + NIM share the Chat Completions shape (``role:"tool"`` messages); an engine + with no registered extractor falls back to it. Returns an empty list when + there are no tool results. + """ + extractor = _RESULT_EXTRACTORS.get(self.model_config.engine, _extract_tool_results_openai) + return extractor(messages) diff --git a/nemoguardrails/guardrails/rails_manager.py b/nemoguardrails/guardrails/rails_manager.py index 8c1a769983..41331ef782 100644 --- a/nemoguardrails/guardrails/rails_manager.py +++ b/nemoguardrails/guardrails/rails_manager.py @@ -24,13 +24,15 @@ import asyncio import logging from collections.abc import Coroutine, Mapping -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any, Optional, TypeVar from nemoguardrails.guardrails.actions.content_safety_action import ( ContentSafetyInputAction, ContentSafetyOutputAction, ) from nemoguardrails.guardrails.actions.jailbreak_detection_action import JailbreakDetectionAction +from nemoguardrails.guardrails.actions.tool_call_action import ToolCallRailAction +from nemoguardrails.guardrails.actions.tool_result_action import ToolResultRailAction from nemoguardrails.guardrails.actions.topic_safety_action import TopicSafetyInputAction from nemoguardrails.guardrails.engine_registry import EngineRegistry from nemoguardrails.guardrails.guardrails_types import ( @@ -40,8 +42,11 @@ ) from nemoguardrails.guardrails.rail_action import RailAction from nemoguardrails.guardrails.telemetry import mark_rail_stop, rail_span, set_rail_content +from nemoguardrails.guardrails.tool_rail_action import ToolRailAction +from nemoguardrails.guardrails.tool_schema import ToolResult, Toolset from nemoguardrails.llm.taskmanager import LLMTaskManager from nemoguardrails.rails.llm.config import _get_flow_name +from nemoguardrails.types import ToolCall if TYPE_CHECKING: from opentelemetry.trace import Tracer @@ -59,6 +64,19 @@ ] } +# All known ToolRailAction subclasses, keyed by their action_name. Tool rails are +# local structural/schema validators (model-free) and so are registered separately +# from the LLM/API-call-shaped RailAction classes above. +_TOOL_ACTION_CLASSES: dict[str, type[ToolRailAction]] = { + cls.action_name: cls + for cls in [ + ToolCallRailAction, + ToolResultRailAction, + ] +} + +_ToolActionT = TypeVar("_ToolActionT", bound=ToolRailAction) + class RailsManager: """Orchestrates input and output safety checks for IORails. @@ -77,6 +95,8 @@ def __init__( output_flows: list[str], input_parallel: bool = False, output_parallel: bool = False, + tool_call_flows: Optional[list[str]] = None, + tool_result_flows: Optional[list[str]] = None, tracer: Optional["Tracer"] = None, content_capture_enabled: bool = False, ) -> None: @@ -102,16 +122,27 @@ def __init__( self.input_parallel: bool = input_parallel self.output_parallel: bool = output_parallel + self.tool_call_flows: list[str] = list(tool_call_flows or []) + self.tool_result_flows: list[str] = list(tool_result_flows or []) + # Build action instances for each configured flow self._actions: dict[str, RailAction] = {} for flow in self.input_flows + self.output_flows: base_name = _get_flow_name(flow) or flow self._actions[flow] = self._create_action(base_name) + # Tool Call Actions run on tool invocations from the main LLM response + # Tool Result Actions run on the results of executing Tool Calls in the harness + self._tool_call_actions = self._build_tool_actions(self.tool_call_flows, ToolCallRailAction) + self._tool_result_actions = self._build_tool_actions(self.tool_result_flows, ToolResultRailAction) + log.info( - "RailsManager initialized: input_flows=%s, output_flows=%s, input_parallel=%s, output_parallel=%s", + "RailsManager initialized: input_flows=%s, output_flows=%s, tool_call_flows=%s, " + "tool_result_flows=%s, input_parallel=%s, output_parallel=%s", self.input_flows, self.output_flows, + self.tool_call_flows, + self.tool_result_flows, self.input_parallel, self.output_parallel, ) @@ -124,6 +155,30 @@ def _create_action(self, base_name: str) -> RailAction: raise RuntimeError(f"Rail flow '{base_name}' not supported. Available: {available}") return action_cls(self.engine_registry, self.task_manager, tracer=self._tracer) + def _build_tool_actions(self, flows: list[str], expected_cls: type[_ToolActionT]) -> dict[str, _ToolActionT]: + """Instantiate the tool rails for *flows*, checking each resolves to *expected_cls*. + + Raises ``RuntimeError`` on a duplicate flow, an unknown flow, or a flow that + resolves to the wrong direction. Duplicates are rejected because the dispatch + keys its coroutine map by flow, so a repeated flow would silently drop a run. + """ + actions: dict[str, _ToolActionT] = {} + for flow in flows: + if flow in actions: + raise RuntimeError(f"Duplicate tool rail flow '{flow}' is not supported") + base_name = _get_flow_name(flow) or flow + action_cls = _TOOL_ACTION_CLASSES.get(base_name) + if action_cls is None: + available = sorted(_TOOL_ACTION_CLASSES.keys()) + raise RuntimeError(f"Tool rail flow '{base_name}' not supported. Available: {available}") + action = action_cls(tracer=self._tracer) + if not isinstance(action, expected_cls): + raise RuntimeError( + f"Tool rail flow '{flow}' resolved to {type(action).__name__}, expected {expected_cls.__name__}" + ) + actions[flow] = action + return actions + async def is_input_safe(self, messages: list[dict]) -> RailResult: """Run all enabled input rails, short-circuiting on the first failure. @@ -155,6 +210,31 @@ async def is_output_safe(self, messages: list[dict], response: str) -> RailResul return await self._run_rails_parallel(rails, RailDirection.OUTPUT) return await self._run_rails_sequential(rails, RailDirection.OUTPUT) + async def are_tool_calls_safe(self, tool_calls: list[ToolCall], toolset: Toolset) -> RailResult: + """Run all enabled tool-call rails against the model's tool calls. + + Rails run sequentially, short-circuiting on the first failure. + Returns safe when no tool-call rails are configured. + """ + if not self.tool_call_flows: + return RailResult(is_safe=True) + + rails = {flow: self._run_tool_call_rail(flow, tool_calls, toolset) for flow in self.tool_call_flows} + return await self._run_rails_sequential(rails, RailDirection.OUTPUT) + + async def are_tool_results_safe(self, tool_results: list[ToolResult], prior_calls: list[ToolCall]) -> RailResult: + """Run all enabled tool-result rails against the results of tool calls. + + Validates the tool results (already normalized to ``ToolResult``) against + the prior tool calls the model made. Rails run sequentially, short-circuiting + on the first failure. Returns safe when no tool-result rails are configured. + """ + if not self.tool_result_flows: + return RailResult(is_safe=True) + + rails = {flow: self._run_tool_result_rail(flow, tool_results, prior_calls) for flow in self.tool_result_flows} + return await self._run_rails_sequential(rails, RailDirection.INPUT) + async def _run_rail( self, flow: str, @@ -180,6 +260,38 @@ async def _run_rail( ) return result + async def _run_tool_call_rail(self, flow: str, tool_calls: list[ToolCall], toolset: Toolset) -> RailResult: + """Dispatch a single tool-call rail to its action, wrapped in an OUTPUT rail span.""" + with rail_span(self._tracer, flow, RailDirection.OUTPUT) as span: + result = await self._tool_call_actions[flow].run(toolset, tool_calls) + mark_rail_stop(span, result.is_safe) + if self._content_capture_enabled: + set_rail_content( + span, + {"tool_calls": [tc.to_dict() for tc in tool_calls]}, + reason=result.reason if not result.is_safe else None, + ) + return result + + async def _run_tool_result_rail( + self, flow: str, tool_results: list[ToolResult], prior_calls: list[ToolCall] + ) -> RailResult: + """Dispatch a single tool-result rail to its action, wrapped in an INPUT rail span.""" + with rail_span(self._tracer, flow, RailDirection.INPUT) as span: + result = await self._tool_result_actions[flow].run(tool_results, prior_calls) + mark_rail_stop(span, result.is_safe) + if self._content_capture_enabled: + set_rail_content( + span, + { + "tool_results": [ + {"call_id": r.call_id, "name": r.name, "is_error": r.is_error} for r in tool_results + ] + }, + reason=result.reason if not result.is_safe else None, + ) + return result + async def _run_rails_sequential( self, rails: Mapping[str, Coroutine[Any, Any, RailResult]], diff --git a/nemoguardrails/guardrails/tool_rail_action.py b/nemoguardrails/guardrails/tool_rail_action.py new file mode 100644 index 0000000000..fe800a411f --- /dev/null +++ b/nemoguardrails/guardrails/tool_rail_action.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Base class for IORails tool-calling rails. + +Tool rails are local structural/schema validators. Unlike ``RailAction`` they make +no LLM or API call, render no prompt, and need no model: they take already-normalized +tool data (a ``Toolset``, the model's ``ToolCall`` list, or incoming ``ToolResult`` +objects — all produced by the engine adapter) and return a ``RailResult``. +``requires_model`` is therefore ``False`` and the only collaborator is an optional +tracer for spans. + +Subclasses set :attr:`ToolRailAction.action_name` and implement an async ``run`` with +their own typed inputs (the inputs differ per rail), performing the check through +:meth:`ToolRailAction._guarded` so every rail gets a consistent action span and turns +an unexpected error into a blocking result rather than letting it propagate. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Callable, Optional + +from nemoguardrails.guardrails.guardrails_types import RailResult, get_request_id +from nemoguardrails.guardrails.telemetry import action_span, record_span_error + +if TYPE_CHECKING: + from opentelemetry.trace import Tracer + +log = logging.getLogger(__name__) + + +class ToolRailAction: + """Base for the local, model-free tool-calling rails (tool-call and tool-result).""" + + action_name: str + requires_model: bool = False + + def __init__(self, tracer: Optional["Tracer"] = None) -> None: + """Store the optional tracer used to emit the action span.""" + self._tracer = tracer + + def _guarded(self, check: Callable[[], RailResult]) -> RailResult: + """Run *check* inside an action span, converting any error into a block. + + Mirrors ``RailAction``'s error contract: an unexpected exception is recorded + on the span and returned as ``RailResult(is_safe=False, ...)`` so a malformed + input or a rail bug fails closed rather than crashing the request. + """ + with action_span(self._tracer, self.action_name) as span: + try: + return check() + except Exception as e: + record_span_error(span, e) + log.error("[%s] %s failed: %s", get_request_id(), self.action_name, e) + return RailResult(is_safe=False, reason=f"{self.action_name} error: {e}") diff --git a/nemoguardrails/guardrails/tool_schema.py b/nemoguardrails/guardrails/tool_schema.py new file mode 100644 index 0000000000..53e3dfe660 --- /dev/null +++ b/nemoguardrails/guardrails/tool_schema.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Internal tool types for IORails tool-calling rails. + +These engine-internal dataclasses are the normalized, provider-neutral shape the +tool rails validate against. They are NOT part of the public API and NOT carried +on ``GenerationOptions``: the request surface stays the provider-native +``llm_params`` block, which ``ModelEngine`` parses into a ``Toolset`` (and incoming +tool results into ``ToolResult`` objects) per inference call. + +``Tool`` is a declared tool definition (what the caller offers); ``ToolCall`` (in +``nemoguardrails.types``) is an invocation the model emitted. Field names are +provider-neutral so the per-provider adapters all produce the same shape: +``arguments_schema`` is OpenAI ``parameters`` / Anthropic ``input_schema`` / Gemini +``parameters``; ``ToolResult.call_id`` is the OpenAI ``tool_call_id`` / Responses +``call_id`` / Anthropic ``tool_use_id`` / Gemini function-call ``id``. OpenAI Chat +Completions is the engine implemented today. +""" + +from collections.abc import Iterable +from dataclasses import dataclass + +import jsonschema + + +@dataclass(frozen=True, slots=True) +class Tool: + """A declared tool definition the caller offered to the model. + + ``name`` is set for function tools and ``None`` for hosted/server tools that + are identified only by ``type`` (e.g. web_search). ``arguments_schema`` is the + JSON Schema for the call arguments, or ``None`` for hosted tools (and any + function tool that declares no parameters), in which case argument validation + is skipped and only the allowlist applies. + """ + + name: str | None = None + type: str = "function" + description: str | None = None + arguments_schema: dict | None = None + strict: bool | None = None + + @property + def key(self) -> str: + """Allowlist / lookup identifier: the function ``name``, else the ``type``.""" + return self.name or self.type + + +class Toolset: + """The set of tools declared on a request, indexed by tool key. + + Backed by a single ``key -> Tool`` mapping built at construction, so there is + no separate list that can drift out of sync with the index. Look tools up with + :meth:`get` (``toolset.get(name)``); iterate or count via the read-only + :attr:`tools`. + """ + + __slots__ = ("_by_key",) + + def __init__(self, tools: Iterable[Tool] | None = None) -> None: + """Index *tools* by their ``key``, rejecting duplicates. + + A toolset must not declare the same function name (or hosted-tool type) + twice, so a repeated ``key`` raises ``ValueError``. A tool with an empty + ``key`` (no name and no type) has no lookup identifier and is dropped. + """ + self._by_key: dict[str, Tool] = {} + for tool in tools or []: + if not tool.key: + continue + if tool.key in self._by_key: + raise ValueError(f"duplicate tool '{tool.key}' in toolset") + self._by_key[tool.key] = tool + + def get(self, key: str) -> Tool | None: + """Return the declared tool registered under *key* (function name or hosted-tool type), or None.""" + return self._by_key.get(key) + + @property + def tools(self) -> tuple[Tool, ...]: + """The declared tools in declaration order (read-only view of the index).""" + return tuple(self._by_key.values()) + + +@dataclass(frozen=True, slots=True) +class ToolResult: + """A normalized tool result extracted from incoming messages by the engine. + + The ToolResultRail consumes a list of these; the per-provider extraction (e.g. + OpenAI ``role:"tool"`` messages) lives in the engine adapter, so the rail never + sees provider wire shapes. ``content`` is a string or a list of content blocks + (the latter covers multimodal results). ``is_error`` flags a failed result + where the provider exposes one (e.g. Anthropic ``is_error`` / Bedrock + ``status:"error"``). + """ + + call_id: str | None = None + name: str | None = None + content: str | list[dict] | None = None + is_error: bool = False + + +def validate_arguments(tool: Tool, arguments: dict) -> str | None: + """Validate model-supplied tool-call arguments against the tool's schema. + + Returns ``None`` when the arguments are valid or the tool declares no schema. + Returns a human-readable reason when the arguments violate the schema, or when + the declared schema itself is not valid JSON Schema (e.g. a non-JSON-Schema + dialect reaching this validator before its engine adapter normalizes it). + """ + if tool.arguments_schema is None: + return None + try: + jsonschema.validate(instance=arguments, schema=tool.arguments_schema) + except jsonschema.ValidationError as exc: + return f"arguments for tool '{tool.key}' do not match its schema: {exc.message}" + except jsonschema.SchemaError as exc: + return f"declared schema for tool '{tool.key}' is not valid JSON Schema: {exc.message}" + return None diff --git a/tests/guardrails/test_engine_registry.py b/tests/guardrails/test_engine_registry.py index 668c3a83f4..de79e6a435 100644 --- a/tests/guardrails/test_engine_registry.py +++ b/tests/guardrails/test_engine_registry.py @@ -30,6 +30,7 @@ from nemoguardrails.guardrails.api_engine import APIEngine from nemoguardrails.guardrails.engine_registry import EngineRegistry from nemoguardrails.guardrails.model_engine import ModelEngine +from nemoguardrails.guardrails.tool_schema import Toolset from nemoguardrails.rails.llm.config import RailsConfig from nemoguardrails.tracing import constants as tracing_constants from nemoguardrails.tracing.constants import SystemConstants @@ -1503,3 +1504,60 @@ async def test_request_attributes_present_on_provider_error(self, manager_with_t assert "gen_ai.response.model" not in attrs assert "gen_ai.usage.input_tokens" not in attrs assert attrs["error.type"] == "RuntimeError" + + +class TestEngineRegistryToolDelegation: + _TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, + } + ] + + def test_parse_tools_delegates_to_model_engine(self, manager): + toolset = manager.parse_tools("main", {"tools": self._TOOLS}) + assert isinstance(toolset, Toolset) + assert [t.key for t in toolset.tools] == ["get_weather"] + + def test_parse_tools_no_tools_returns_empty(self, manager): + assert manager.parse_tools("main", None).tools == () + + def test_extract_tool_results_delegates_to_model_engine(self, manager): + messages = [{"role": "tool", "tool_call_id": "c1", "name": "get_weather", "content": "18C"}] + results = manager.extract_tool_results("main", messages) + assert [(r.call_id, r.name, r.content) for r in results] == [("c1", "get_weather", "18C")] + + def test_parse_tools_unknown_engine_raises_keyerror(self, manager): + with pytest.raises(KeyError): + manager.parse_tools("nonexistent", {"tools": self._TOOLS}) + + def test_extract_tool_results_unknown_engine_raises_keyerror(self, manager): + with pytest.raises(KeyError): + manager.extract_tool_results("nonexistent", []) + + def test_parse_tools_non_model_engine_raises_typeerror(self, manager): + """jailbreak_detection is an APIEngine, not a ModelEngine.""" + with pytest.raises(TypeError): + manager.parse_tools("jailbreak_detection", {"tools": self._TOOLS}) + + def test_parse_tools_includes_tools_from_model_parameters(self): + """Tools declared in model parameters (body_param_defaults) are included even with no per-call llm_params.""" + config = RailsConfig.from_content( + config={ + "models": [ + { + "type": "main", + "engine": "nim", + "model": "meta/llama-3.3-70b-instruct", + "parameters": {"tools": self._TOOLS}, + } + ] + } + ) + with patch.dict("os.environ", {"NVIDIA_API_KEY": "test-key"}): + mgr = EngineRegistry(config.models, config.rails.config) + toolset = mgr.parse_tools("main", None) + assert [t.key for t in toolset.tools] == ["get_weather"] diff --git a/tests/guardrails/test_model_engine.py b/tests/guardrails/test_model_engine.py index 279b71c36c..2fff50c956 100644 --- a/tests/guardrails/test_model_engine.py +++ b/tests/guardrails/test_model_engine.py @@ -35,8 +35,10 @@ _parse_chat_completion, _parse_chat_completion_chunk, ) +from nemoguardrails.guardrails.tool_schema import Toolset from nemoguardrails.rails.llm.config import Model from nemoguardrails.types import LLMResponse, LLMResponseChunk, UsageInfo +from tests.guardrails.tool_helpers import make_tool_conversation def _make_model( @@ -1917,3 +1919,142 @@ def test_passes_through_metadata(self): assert result.model == "gpt-5" assert result.request_id == "chunk-1" assert result.finish_reason == "stop" + + +_OPENAI_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + "strict": True, + }, + }, + {"type": "function", "function": {"name": "noargs"}}, +] + + +class TestParseTools: + def test_parses_openai_chat_completions_tools(self): + engine = ModelEngine(_make_model(engine="openai")) + toolset = engine.parse_tools({"tools": _OPENAI_TOOLS}) + + assert isinstance(toolset, Toolset) + assert sorted(t.key for t in toolset.tools) == ["get_weather", "noargs"] + + weather = toolset.get("get_weather") + assert weather is not None + assert weather.name == "get_weather" + assert weather.type == "function" + assert weather.description == "Get the weather for a city." + assert weather.arguments_schema == _OPENAI_TOOLS[0]["function"]["parameters"] + assert weather.strict is True + + def test_nim_uses_the_same_shape(self): + engine = ModelEngine(_make_model(engine="nim")) + toolset = engine.parse_tools({"tools": _OPENAI_TOOLS}) + assert sorted(t.key for t in toolset.tools) == ["get_weather", "noargs"] + + def test_tool_without_parameters_has_no_arguments_schema(self): + engine = ModelEngine(_make_model(engine="openai")) + toolset = engine.parse_tools({"tools": _OPENAI_TOOLS}) + noargs = toolset.get("noargs") + assert noargs is not None + assert noargs.arguments_schema is None + + def test_no_tools_returns_empty_toolset(self): + engine = ModelEngine(_make_model(engine="openai")) + assert engine.parse_tools({}).tools == () + assert engine.parse_tools(None).tools == () + assert engine.parse_tools({"tools": []}).tools == () + + def test_malformed_entries_are_skipped(self): + """Non-dict / function-less entries are dropped so a malformed tool fails closed.""" + engine = ModelEngine(_make_model(engine="openai")) + tools = [ + "garbage", + {"type": "function"}, + {"type": "function", "function": {"name": "ok"}}, + ] + toolset = engine.parse_tools({"tools": tools}) + assert [t.key for t in toolset.tools] == ["ok"] + + def test_unknown_engine_falls_back_to_openai_parser(self): + engine = ModelEngine(_make_model(engine="vllm", parameters={"base_url": "http://localhost:8000"})) + toolset = engine.parse_tools({"tools": _OPENAI_TOOLS}) + assert sorted(t.key for t in toolset.tools) == ["get_weather", "noargs"] + + def test_function_entries_without_name_are_skipped(self): + """Name-less function entries are dropped, so duplicate empty keys never crash parse_tools.""" + engine = ModelEngine(_make_model(engine="openai")) + tools = [ + {"type": "function", "function": {"description": "no name"}}, + {"type": "function", "function": {"name": ""}}, + {"type": "function", "function": {"name": "ok"}}, + ] + toolset = engine.parse_tools({"tools": tools}) + assert [t.key for t in toolset.tools] == ["ok"] + + +_TOOL_MESSAGES = make_tool_conversation() + + +class TestExtractToolResults: + def test_extracts_openai_tool_result(self): + engine = ModelEngine(_make_model(engine="openai")) + results = engine.extract_tool_results(_TOOL_MESSAGES) + + assert len(results) == 1 + result = results[0] + assert result.call_id == "call_1" + assert result.name == "get_weather" + assert result.content == "18C" + assert result.is_error is False + + def test_ignores_non_tool_messages(self): + engine = ModelEngine(_make_model(engine="openai")) + messages = [ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + assert engine.extract_tool_results(messages) == [] + + def test_extracts_multiple_results_in_order(self): + engine = ModelEngine(_make_model(engine="openai")) + messages = [ + {"role": "tool", "tool_call_id": "call_1", "name": "a", "content": "r1"}, + {"role": "tool", "tool_call_id": "call_2", "name": "b", "content": "r2"}, + ] + results = engine.extract_tool_results(messages) + assert [r.call_id for r in results] == ["call_1", "call_2"] + + def test_skips_non_dict_messages(self): + engine = ModelEngine(_make_model(engine="openai")) + messages = ["garbage", {"role": "tool", "tool_call_id": "call_1", "content": "r1"}] + results = engine.extract_tool_results(messages) + assert len(results) == 1 + assert results[0].call_id == "call_1" + + def test_missing_fields_become_none(self): + """Missing tool_call_id/name still extracts; the rail (not the extractor) judges linkage.""" + engine = ModelEngine(_make_model(engine="openai")) + results = engine.extract_tool_results([{"role": "tool", "content": "r1"}]) + assert len(results) == 1 + assert results[0].call_id is None + assert results[0].name is None + + def test_nim_uses_the_same_shape(self): + engine = ModelEngine(_make_model(engine="nim")) + results = engine.extract_tool_results(_TOOL_MESSAGES) + assert [r.call_id for r in results] == ["call_1"] + + def test_unknown_engine_falls_back_to_openai_extractor(self): + engine = ModelEngine(_make_model(engine="vllm", parameters={"base_url": "http://localhost:8000"})) + results = engine.extract_tool_results(_TOOL_MESSAGES) + assert [r.call_id for r in results] == ["call_1"] diff --git a/tests/guardrails/test_rails_manager.py b/tests/guardrails/test_rails_manager.py index f558f6c6a9..199212791f 100644 --- a/tests/guardrails/test_rails_manager.py +++ b/tests/guardrails/test_rails_manager.py @@ -25,12 +25,17 @@ from unittest.mock import AsyncMock, patch import pytest +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from nemoguardrails.guardrails.engine_registry import EngineRegistry from nemoguardrails.guardrails.rails_manager import RailsManager +from nemoguardrails.guardrails.tool_schema import Tool, ToolResult, Toolset from nemoguardrails.llm.taskmanager import LLMTaskManager from nemoguardrails.rails.llm.config import RailsConfig -from nemoguardrails.types import LLMResponse +from nemoguardrails.tracing.constants import GuardrailsAttributes +from nemoguardrails.types import LLMResponse, ToolCall, ToolCallFunction from tests.guardrails.test_data import ( CONTENT_SAFETY_CONFIG, NEMOGUARDS_CONFIG, @@ -39,6 +44,7 @@ NEMOGUARDS_PARALLEL_OUTPUT_CONFIG, TOPIC_SAFETY_CONFIG, ) +from tests.guardrails.tool_helpers import WEATHER_SCHEMA, assert_blocked SAFE_INPUT_JSON = json.dumps({"User Safety": "safe"}) UNSAFE_INPUT_JSON = json.dumps({"User Safety": "unsafe", "Safety Categories": "S1: Violence"}) @@ -469,3 +475,155 @@ async def test_output_unsafe(self, parallel_rails_manager): ) result = await parallel_rails_manager.is_output_safe(MESSAGES, "response") assert not result.is_safe + + +def _tool_rails_manager(*, tool_call_flows=None, tool_result_flows=None) -> RailsManager: + """Build a RailsManager with only tool rails wired (no LLM input/output flows).""" + config = RailsConfig.from_content(config={"models": []}) + return RailsManager( + engine_registry=EngineRegistry(config.models, config.rails.config), + task_manager=LLMTaskManager(config), + input_flows=[], + output_flows=[], + tool_call_flows=tool_call_flows or [], + tool_result_flows=tool_result_flows or [], + ) + + +def _toolset() -> Toolset: + return Toolset(tools=[Tool(name="get_weather", arguments_schema=WEATHER_SCHEMA)]) + + +def _call(name: str, arguments: dict) -> ToolCall: + return ToolCall(id="c1", function=ToolCallFunction(name=name, arguments=arguments)) + + +class TestRailsManagerToolInit: + def test_tool_flows_populated(self): + mgr = _tool_rails_manager( + tool_call_flows=["tool call validation"], tool_result_flows=["tool result validation"] + ) + assert mgr.tool_call_flows == ["tool call validation"] + assert mgr.tool_result_flows == ["tool result validation"] + + def test_no_tool_flows_by_default(self): + mgr = _tool_rails_manager() + assert mgr.tool_call_flows == [] + assert mgr.tool_result_flows == [] + + def test_unknown_tool_flow_raises(self): + with pytest.raises(RuntimeError, match="not supported"): + _tool_rails_manager(tool_call_flows=["bogus tool rail"]) + + def test_tool_call_flow_with_result_rail_raises(self): + with pytest.raises(RuntimeError, match="expected ToolCallRailAction"): + _tool_rails_manager(tool_call_flows=["tool result validation"]) + + def test_tool_result_flow_with_call_rail_raises(self): + with pytest.raises(RuntimeError, match="expected ToolResultRailAction"): + _tool_rails_manager(tool_result_flows=["tool call validation"]) + + def test_duplicate_tool_call_flow_raises(self): + with pytest.raises(RuntimeError, match="Duplicate tool rail flow"): + _tool_rails_manager(tool_call_flows=["tool call validation", "tool call validation"]) + + def test_duplicate_tool_result_flow_raises(self): + with pytest.raises(RuntimeError, match="Duplicate tool rail flow"): + _tool_rails_manager(tool_result_flows=["tool result validation", "tool result validation"]) + + +class TestRailsManagerToolCalls: + @pytest.mark.asyncio + async def test_allows_valid_call(self): + mgr = _tool_rails_manager(tool_call_flows=["tool call validation"]) + result = await mgr.are_tool_calls_safe([_call("get_weather", {"city": "Paris"})], _toolset()) + assert result.is_safe is True + + @pytest.mark.asyncio + async def test_blocks_undeclared_call(self): + mgr = _tool_rails_manager(tool_call_flows=["tool call validation"]) + result = await mgr.are_tool_calls_safe([_call("rm_rf", {})], _toolset()) + assert_blocked(result, "rm_rf") + + @pytest.mark.asyncio + async def test_no_flows_returns_safe(self): + mgr = _tool_rails_manager() + result = await mgr.are_tool_calls_safe([_call("rm_rf", {})], _toolset()) + assert result.is_safe is True + + +class TestRailsManagerToolResults: + @pytest.mark.asyncio + async def test_allows_linked_result(self): + mgr = _tool_rails_manager(tool_result_flows=["tool result validation"]) + prior = [_call("get_weather", {"city": "Paris"})] + result = await mgr.are_tool_results_safe([ToolResult(call_id="c1", content="18C")], prior) + assert result.is_safe is True + + @pytest.mark.asyncio + async def test_blocks_unlinked_result(self): + mgr = _tool_rails_manager(tool_result_flows=["tool result validation"]) + prior = [_call("get_weather", {"city": "Paris"})] + result = await mgr.are_tool_results_safe([ToolResult(call_id="c9", content="x")], prior) + assert_blocked(result, "c9") + + @pytest.mark.asyncio + async def test_no_flows_returns_safe(self): + mgr = _tool_rails_manager() + result = await mgr.are_tool_results_safe([ToolResult(call_id="c9", content="x")], []) + assert result.is_safe is True + + +def _capture_tool_rails_manager(): + """Build (manager, exporter) with a real tracer + content capture on, both tool rails wired.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + config = RailsConfig.from_content(config={"models": []}) + manager = RailsManager( + engine_registry=EngineRegistry(config.models, config.rails.config), + task_manager=LLMTaskManager(config), + input_flows=[], + output_flows=[], + tool_call_flows=["tool call validation"], + tool_result_flows=["tool result validation"], + tracer=provider.get_tracer("test"), + content_capture_enabled=True, + ) + return manager, exporter + + +def _rail_span(exporter): + """The single finished span that carries rail.input (the rail span, not the action span).""" + spans = [s for s in exporter.get_finished_spans() if GuardrailsAttributes.RAIL_INPUT in s.attributes] + assert len(spans) == 1 + return spans[0] + + +class TestRailsManagerToolContentCapture: + @pytest.mark.asyncio + async def test_tool_call_span_captures_calls_and_reason_on_block(self): + manager, exporter = _capture_tool_rails_manager() + await manager.are_tool_calls_safe([_call("rm_rf", {})], _toolset()) + attrs = _rail_span(exporter).attributes + payload = json.loads(attrs[GuardrailsAttributes.RAIL_INPUT]) + assert payload["tool_calls"][0]["function"]["name"] == "rm_rf" + assert "rm_rf" in attrs[GuardrailsAttributes.RAIL_REASON] + + @pytest.mark.asyncio + async def test_tool_call_span_omits_reason_when_safe(self): + manager, exporter = _capture_tool_rails_manager() + await manager.are_tool_calls_safe([_call("get_weather", {"city": "Paris"})], _toolset()) + attrs = _rail_span(exporter).attributes + assert "tool_calls" in json.loads(attrs[GuardrailsAttributes.RAIL_INPUT]) + assert GuardrailsAttributes.RAIL_REASON not in attrs + + @pytest.mark.asyncio + async def test_tool_result_span_captures_linkage_and_reason_on_block(self): + manager, exporter = _capture_tool_rails_manager() + prior = [_call("get_weather", {"city": "Paris"})] + await manager.are_tool_results_safe([ToolResult(call_id="c9", name="get_weather", content="x")], prior) + attrs = _rail_span(exporter).attributes + payload = json.loads(attrs[GuardrailsAttributes.RAIL_INPUT]) + assert payload["tool_results"][0] == {"call_id": "c9", "name": "get_weather", "is_error": False} + assert "c9" in attrs[GuardrailsAttributes.RAIL_REASON] diff --git a/tests/guardrails/test_tool_call_action.py b/tests/guardrails/test_tool_call_action.py new file mode 100644 index 0000000000..c5e8126ac8 --- /dev/null +++ b/tests/guardrails/test_tool_call_action.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for ToolCallRailAction (allowlist + argument-schema validation).""" + +import pytest + +from nemoguardrails.guardrails.actions.tool_call_action import ToolCallRailAction +from nemoguardrails.guardrails.tool_schema import Tool, Toolset +from nemoguardrails.types import ToolCall, ToolCallFunction +from tests.guardrails.tool_helpers import WEATHER_SCHEMA, assert_blocked + + +def _toolset() -> Toolset: + return Toolset(tools=[Tool(name="get_weather", arguments_schema=WEATHER_SCHEMA), Tool(name="ping")]) + + +def _call(name: str, arguments: dict) -> ToolCall: + return ToolCall(id="c1", function=ToolCallFunction(name=name, arguments=arguments)) + + +class TestToolCallRailAction: + @pytest.mark.asyncio + async def test_allowed_call_with_valid_arguments_is_safe(self): + result = await ToolCallRailAction().run(_toolset(), [_call("get_weather", {"city": "Paris"})]) + assert result.is_safe is True + + @pytest.mark.asyncio + async def test_tool_without_schema_passes_allowlist_only(self): + result = await ToolCallRailAction().run(_toolset(), [_call("ping", {"anything": 1})]) + assert result.is_safe is True + + @pytest.mark.asyncio + async def test_undeclared_tool_is_blocked(self): + result = await ToolCallRailAction().run(_toolset(), [_call("rm_rf", {})]) + assert_blocked(result, "rm_rf", "not an allowed tool") + + @pytest.mark.asyncio + async def test_invalid_arguments_are_blocked(self): + result = await ToolCallRailAction().run(_toolset(), [_call("get_weather", {})]) + assert_blocked(result, "get_weather") + + @pytest.mark.asyncio + async def test_empty_tool_calls_is_safe(self): + result = await ToolCallRailAction().run(_toolset(), []) + assert result.is_safe is True + + @pytest.mark.asyncio + async def test_one_bad_call_blocks_the_batch(self): + calls = [_call("get_weather", {"city": "Paris"}), _call("rm_rf", {})] + result = await ToolCallRailAction().run(_toolset(), calls) + assert_blocked(result, "rm_rf") + + @pytest.mark.asyncio + async def test_hosted_tool_matched_by_type_when_name_is_empty(self): + toolset = Toolset(tools=[Tool(name=None, type="web_search")]) + call = ToolCall(id="c1", type="web_search", function=ToolCallFunction(name="", arguments={})) + result = await ToolCallRailAction().run(toolset, [call]) + assert result.is_safe is True + + @pytest.mark.asyncio + async def test_undeclared_hosted_tool_type_is_blocked(self): + toolset = Toolset(tools=[Tool(name=None, type="web_search")]) + call = ToolCall(id="c1", type="code_interpreter", function=ToolCallFunction(name="", arguments={})) + result = await ToolCallRailAction().run(toolset, [call]) + assert_blocked(result, "code_interpreter", "not an allowed tool") diff --git a/tests/guardrails/test_tool_rail_action.py b/tests/guardrails/test_tool_rail_action.py new file mode 100644 index 0000000000..5df7430688 --- /dev/null +++ b/tests/guardrails/test_tool_rail_action.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the ToolRailAction base (span wrapper + fail-closed error handling).""" + +from nemoguardrails.guardrails.guardrails_types import RailResult +from nemoguardrails.guardrails.tool_rail_action import ToolRailAction + + +class _ProbeRail(ToolRailAction): + action_name = "probe tool rail" + + +class TestToolRailActionGuarded: + def test_requires_model_is_false(self): + assert _ProbeRail().requires_model is False + + def test_returns_safe_check_result(self): + assert _ProbeRail()._guarded(lambda: RailResult(is_safe=True)) == RailResult(is_safe=True) + + def test_returns_unsafe_check_result(self): + result = _ProbeRail()._guarded(lambda: RailResult(is_safe=False, reason="bad call")) + assert result.is_safe is False + assert result.reason == "bad call" + + def test_exception_fails_closed(self): + def boom() -> RailResult: + raise ValueError("kaboom") + + result = _ProbeRail()._guarded(boom) + assert result.is_safe is False + assert result.reason is not None + assert "probe tool rail" in result.reason + assert "kaboom" in result.reason diff --git a/tests/guardrails/test_tool_rails_e2e.py b/tests/guardrails/test_tool_rails_e2e.py new file mode 100644 index 0000000000..eaa4ddaa8b --- /dev/null +++ b/tests/guardrails/test_tool_rails_e2e.py @@ -0,0 +1,268 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""End-to-end tests for the tool rails at the RailsManager level. + +These exercise the full seam this PR builds: the real ModelEngine response +parsing (non-streaming ``chat_completion`` and streaming SSE assembly) feeds +the EngineRegistry tool helpers (``parse_tools`` / ``extract_tool_results``), +whose normalized output is validated by ``RailsManager.are_tool_calls_safe`` / +``are_tool_results_safe``. + +The only mock is the aiohttp transport (the model's HTTP/SSE response body), so +ModelEngine's parsing runs for real. The "model response -> parse -> validate" +orchestration here stands in for the future IORails glue; this PR keeps the +RailsManager surface limited to the two validation methods. +""" + +import json +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from nemoguardrails.guardrails.engine_registry import EngineRegistry +from nemoguardrails.guardrails.guardrails_types import RailResult +from nemoguardrails.guardrails.model_engine import ModelEngine +from nemoguardrails.guardrails.rails_manager import RailsManager +from nemoguardrails.llm.taskmanager import LLMTaskManager +from nemoguardrails.rails.llm.config import RailsConfig +from nemoguardrails.types import ChatMessage +from tests.guardrails.tool_helpers import WEATHER_SCHEMA, assert_blocked, make_tool_conversation + +STACK_CONFIG = {"models": [{"type": "main", "engine": "nim", "model": "meta/llama-3.3-70b-instruct"}]} + +WEATHER_TOOL = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the weather for a city.", + "parameters": WEATHER_SCHEMA, + }, +} +LLM_PARAMS = {"tools": [WEATHER_TOOL], "tool_choice": "auto"} +MESSAGES = [{"role": "user", "content": "What's the weather in Paris?"}] + + +def _build_stack() -> tuple[EngineRegistry, RailsManager]: + """Build an EngineRegistry + RailsManager with both tool rails enabled.""" + config = RailsConfig.from_content(config=STACK_CONFIG) + engine_registry = EngineRegistry(config.models, config.rails.config) + rails_manager = RailsManager( + engine_registry=engine_registry, + task_manager=LLMTaskManager(config), + input_flows=[], + output_flows=[], + tool_call_flows=["tool call validation"], + tool_result_flows=["tool result validation"], + ) + return engine_registry, rails_manager + + +def _inject_json_response(registry: EngineRegistry, payload: dict) -> None: + """Wire the 'main' engine's aiohttp client to return *payload* as JSON.""" + mock_response = AsyncMock() + mock_response.__aenter__ = AsyncMock(return_value=mock_response) + mock_response.status = 200 + mock_response.json = AsyncMock(return_value=payload) + mock_client = AsyncMock() + mock_client.post = MagicMock(return_value=mock_response) + mock_client.closed = False + engine = registry._get_engine("main", ModelEngine) + engine._client = mock_client + engine._running = True + + +def _inject_sse_stream(registry: EngineRegistry, raw_lines: list) -> None: + """Wire the 'main' engine's aiohttp client to stream *raw_lines* via readline().""" + all_lines = [] + for raw in raw_lines: + for part in raw.split(b"\n"): + if part: + all_lines.append(part + b"\n") + line_iter = iter(all_lines) + + async def _readline(): + return next(line_iter, b"") + + mock_content = MagicMock() + mock_content.readline = _readline + + mock_response = AsyncMock() + mock_response.__aenter__ = AsyncMock(return_value=mock_response) + mock_response.status = 200 + mock_response.content = mock_content + + mock_client = AsyncMock() + mock_client.post = MagicMock(return_value=mock_response) + mock_client.closed = False + engine = registry._get_engine("main", ModelEngine) + engine._client = mock_client + engine._running = True + + +def _tool_call_payload(name: str, arguments_json: str, call_id: str = "call_1") -> dict: + """A non-streaming /v1/chat/completions response carrying a single tool call.""" + return { + "id": "chatcmpl-1", + "model": "meta/llama-3.3-70b-instruct", + "choices": [ + { + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": call_id, + "type": "function", + "function": {"name": name, "arguments": arguments_json}, + } + ], + }, + "finish_reason": "tool_calls", + } + ], + } + + +def _sse(chunk: dict) -> bytes: + return ("data: " + json.dumps(chunk) + "\n\n").encode("utf-8") + + +def _tool_call_sse_lines(name: str, arg_fragments: list, call_id: str = "call_1") -> list: + """SSE lines streaming a single tool call: id/name first, then argument fragments, then finish.""" + chunks = [ + { + "id": "chatcmpl-1", + "choices": [ + { + "index": 0, + "delta": { + "role": "assistant", + "tool_calls": [ + { + "index": 0, + "id": call_id, + "type": "function", + "function": {"name": name, "arguments": ""}, + } + ], + }, + "finish_reason": None, + } + ], + } + ] + for fragment in arg_fragments: + chunks.append( + { + "id": "chatcmpl-1", + "choices": [ + { + "index": 0, + "delta": {"tool_calls": [{"index": 0, "function": {"arguments": fragment}}]}, + "finish_reason": None, + } + ], + } + ) + chunks.append({"id": "chatcmpl-1", "choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]}) + lines = [_sse(chunk) for chunk in chunks] + lines.append(b"data: [DONE]\n\n") + return lines + + +async def _drive_nonstream_call(payload: dict) -> RailResult: + """Mock the model's JSON response, then parse tools + model_call + validate the calls.""" + registry, manager = _build_stack() + _inject_json_response(registry, payload) + toolset = registry.parse_tools("main", LLM_PARAMS) + response = await registry.model_call("main", MESSAGES, **LLM_PARAMS) + assert response.tool_calls is not None + return await manager.are_tool_calls_safe(response.tool_calls, toolset) + + +async def _drive_stream_call(sse_lines: list) -> RailResult: + """Mock the model's SSE stream, assemble the tool calls, then validate them.""" + registry, manager = _build_stack() + _inject_sse_stream(registry, sse_lines) + toolset = registry.parse_tools("main", LLM_PARAMS) + collected = [] + async for chunk in registry.stream_model_call("main", MESSAGES, **LLM_PARAMS): + if chunk.delta_tool_calls: + collected.extend(chunk.delta_tool_calls) + assert collected, "expected assembled tool calls from the stream" + return await manager.are_tool_calls_safe(collected, toolset) + + +async def _drive_result(result_call_id: str) -> RailResult: + """Extract tool results + prior calls from a conversation, then validate the results.""" + registry, manager = _build_stack() + messages = make_tool_conversation(result_call_id=result_call_id) + tool_results = registry.extract_tool_results("main", messages) + prior_calls = ChatMessage.from_dict(messages[1]).tool_calls + assert prior_calls is not None + return await manager.are_tool_results_safe(tool_results, prior_calls) + + +class TestToolCallRailEndToEndNonStreaming: + @pytest.mark.parametrize( + "payload, blocked", + [ + (_tool_call_payload("get_weather", '{"city": "Paris"}'), None), + (_tool_call_payload("rm_rf", "{}"), "rm_rf"), + (_tool_call_payload("get_weather", "{}"), "get_weather"), + ], + ids=["allowed", "undeclared", "invalid-args"], + ) + @pytest.mark.asyncio + async def test_tool_call(self, payload, blocked): + result = await _drive_nonstream_call(payload) + if blocked is None: + assert result.is_safe is True + else: + assert_blocked(result, blocked) + + +class TestToolResultRailEndToEnd: + @pytest.mark.parametrize( + "result_call_id, blocked", + [("call_1", None), ("call_999", "call_999")], + ids=["linked", "unlinked"], + ) + @pytest.mark.asyncio + async def test_tool_result(self, result_call_id, blocked): + result = await _drive_result(result_call_id) + if blocked is None: + assert result.is_safe is True + else: + assert_blocked(result, blocked) + + +class TestToolCallRailEndToEndStreaming: + @pytest.mark.parametrize( + "sse_lines, blocked", + [ + (_tool_call_sse_lines("get_weather", ['{"city": ', '"Paris"}']), None), + (_tool_call_sse_lines("rm_rf", ["{}"]), "rm_rf"), + ], + ids=["allowed", "undeclared"], + ) + @pytest.mark.asyncio + async def test_streamed_tool_call(self, sse_lines, blocked): + result = await _drive_stream_call(sse_lines) + if blocked is None: + assert result.is_safe is True + else: + assert_blocked(result, blocked) diff --git a/tests/guardrails/test_tool_result_action.py b/tests/guardrails/test_tool_result_action.py new file mode 100644 index 0000000000..3a302dd295 --- /dev/null +++ b/tests/guardrails/test_tool_result_action.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for ToolResultRailAction (call_id linkage + structural validation).""" + +import pytest + +from nemoguardrails.guardrails.actions.tool_result_action import ToolResultRailAction +from nemoguardrails.guardrails.tool_schema import ToolResult +from nemoguardrails.types import ToolCall, ToolCallFunction +from tests.guardrails.tool_helpers import assert_blocked + + +def _prior_calls() -> list: + return [ + ToolCall(id="c1", function=ToolCallFunction(name="get_weather", arguments={"city": "Paris"})), + ToolCall(id="c2", function=ToolCallFunction(name="search", arguments={"q": "x"})), + ] + + +def _result(call_id, name=None, content: "str | list[dict] | None" = "18C") -> ToolResult: + return ToolResult(call_id=call_id, name=name, content=content) + + +class TestToolResultRailAction: + @pytest.mark.asyncio + async def test_linked_result_with_matching_name_is_safe(self): + result = await ToolResultRailAction().run([_result("c1", name="get_weather")], _prior_calls()) + assert result.is_safe is True + + @pytest.mark.asyncio + async def test_result_without_name_links_on_call_id_only(self): + result = await ToolResultRailAction().run([_result("c2")], _prior_calls()) + assert result.is_safe is True + + @pytest.mark.asyncio + async def test_list_content_is_well_formed(self): + result = await ToolResultRailAction().run( + [_result("c1", content=[{"type": "text", "text": "18C"}])], _prior_calls() + ) + assert result.is_safe is True + + @pytest.mark.asyncio + async def test_empty_results_is_safe(self): + result = await ToolResultRailAction().run([], _prior_calls()) + assert result.is_safe is True + + @pytest.mark.asyncio + async def test_missing_call_id_is_blocked(self): + result = await ToolResultRailAction().run([_result("")], _prior_calls()) + assert_blocked(result, "missing a call_id") + + @pytest.mark.asyncio + async def test_unlinked_call_id_is_blocked(self): + result = await ToolResultRailAction().run([_result("c9")], _prior_calls()) + assert_blocked(result, "c9", "does not correspond to a prior tool call") + + @pytest.mark.asyncio + async def test_name_mismatch_is_blocked(self): + result = await ToolResultRailAction().run([_result("c1", name="search")], _prior_calls()) + assert_blocked(result, "does not match the called tool", "get_weather") + + @pytest.mark.asyncio + async def test_malformed_content_is_blocked(self): + result = await ToolResultRailAction().run( + [_result("c1", content={"unexpected": "shape"})], # type: ignore[arg-type] + _prior_calls(), + ) + assert_blocked(result, "malformed content") + + @pytest.mark.asyncio + async def test_list_of_non_dicts_is_blocked(self): + result = await ToolResultRailAction().run( + [_result("c1", content=[1, 2, 3])], # type: ignore[arg-type] + _prior_calls(), + ) + assert_blocked(result, "malformed content") + + @pytest.mark.asyncio + async def test_empty_content_list_is_well_formed(self): + result = await ToolResultRailAction().run([_result("c1", content=[])], _prior_calls()) + assert result.is_safe is True + + @pytest.mark.asyncio + async def test_one_bad_result_blocks_the_batch(self): + results = [_result("c1", name="get_weather"), _result("c9")] + result = await ToolResultRailAction().run(results, _prior_calls()) + assert_blocked(result, "c9") + + @pytest.mark.asyncio + async def test_duplicate_prior_call_id_is_blocked(self): + prior = [ + ToolCall(id="c1", function=ToolCallFunction(name="get_weather", arguments={})), + ToolCall(id="c1", function=ToolCallFunction(name="search", arguments={})), + ] + result = await ToolResultRailAction().run([_result("c1")], prior) + assert_blocked(result, "duplicate prior tool call id", "c1") + + @pytest.mark.asyncio + async def test_prior_call_with_empty_id_is_skipped(self): + prior = [ToolCall(id="", function=ToolCallFunction(name="get_weather", arguments={}))] + result = await ToolResultRailAction().run([], prior) + assert result.is_safe is True + + @pytest.mark.asyncio + async def test_duplicate_result_ids_are_blocked(self): + results = [_result("c1", name="get_weather"), _result("c1", name="get_weather")] + result = await ToolResultRailAction().run(results, _prior_calls()) + assert_blocked(result, "duplicate tool result", "c1") + + +class TestValidateToolResultIds: + def setup_method(self): + self.action = ToolResultRailAction() + + def test_unique_ids_returns_none(self): + assert self.action._validate_tool_result_ids([_result("c1"), _result("c2")]) is None + + def test_empty_list_returns_none(self): + assert self.action._validate_tool_result_ids([]) is None + + def test_result_without_call_id_is_skipped(self): + assert self.action._validate_tool_result_ids([ToolResult(call_id=None)]) is None + + def test_duplicate_call_id_is_blocked(self): + results = [ + ToolResult(call_id="call_123", name="search", content="real result"), + ToolResult(call_id="call_123", name="search", content="duplicate/injected result"), + ] + assert_blocked(self.action._validate_tool_result_ids(results), "duplicate tool result", "call_123") + + +class TestValidateResultCallId: + def setup_method(self): + self.action = ToolResultRailAction() + self.calls_by_id = { + "c1": ToolCall(id="c1", function=ToolCallFunction(name="get_weather", arguments={})), + } + + def test_linked_result_returns_none(self): + assert self.action._validate_result_call_id(_result("c1"), self.calls_by_id) is None + + def test_missing_call_id_is_blocked(self): + assert_blocked( + self.action._validate_result_call_id(_result(""), self.calls_by_id), + "missing a call_id", + ) + + def test_orphaned_call_id_is_blocked(self): + assert_blocked( + self.action._validate_result_call_id(_result("c9"), self.calls_by_id), + "c9", + "does not correspond to a prior tool call", + ) + + +class TestValidateResultName: + def setup_method(self): + self.action = ToolResultRailAction() + self.prior = ToolCall(id="c1", function=ToolCallFunction(name="get_weather", arguments={})) + + def test_matching_names_returns_none(self): + assert self.action._validate_result_name(_result("c1", name="get_weather"), self.prior) is None + + def test_result_without_name_returns_none(self): + assert self.action._validate_result_name(_result("c1"), self.prior) is None + + def test_prior_without_function_name_returns_none(self): + prior = ToolCall(id="c1", function=ToolCallFunction(name="", arguments={})) + assert self.action._validate_result_name(_result("c1", name="get_weather"), prior) is None + + def test_name_mismatch_is_blocked(self): + assert_blocked( + self.action._validate_result_name(_result("c1", name="search"), self.prior), + "search", + "get_weather", + "does not match", + ) + + +class TestValidateResultContent: + def setup_method(self): + self.action = ToolResultRailAction() + + def test_string_content_returns_none(self): + assert self.action._validate_result_content(_result("c1", content="18C")) is None + + def test_list_of_dicts_returns_none(self): + assert self.action._validate_result_content(_result("c1", content=[{"type": "text", "text": "18C"}])) is None + + def test_none_content_returns_none(self): + assert self.action._validate_result_content(_result("c1", content=None)) is None + + def test_empty_list_returns_none(self): + assert self.action._validate_result_content(_result("c1", content=[])) is None + + def test_dict_content_is_blocked(self): + assert_blocked( + self.action._validate_result_content(_result("c1", content={"key": "val"})), # type: ignore[arg-type] + "malformed content", + ) + + def test_list_of_non_dicts_is_blocked(self): + assert_blocked( + self.action._validate_result_content(_result("c1", content=[1, 2, 3])), # type: ignore[arg-type] + "malformed content", + ) diff --git a/tests/guardrails/test_tool_schema.py b/tests/guardrails/test_tool_schema.py new file mode 100644 index 0000000000..bfbb5f9758 --- /dev/null +++ b/tests/guardrails/test_tool_schema.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the tool_schema module (canonical tool types + arg validation).""" + +from dataclasses import FrozenInstanceError + +import pytest + +from nemoguardrails.guardrails.tool_schema import ( + Tool, + ToolResult, + Toolset, + validate_arguments, +) + +_WEATHER_SCHEMA = { + "type": "object", + "properties": { + "city": {"type": "string"}, + "days": {"type": "integer"}, + }, + "required": ["city"], +} + + +def _weather_tool() -> Tool: + return Tool( + name="get_weather", + description="Get the weather for a city.", + arguments_schema=_WEATHER_SCHEMA, + ) + + +class TestTool: + def test_defaults(self): + tool = Tool() + assert tool.name is None + assert tool.type == "function" + assert tool.description is None + assert tool.arguments_schema is None + assert tool.strict is None + + def test_key_uses_name_when_present(self): + assert Tool(name="get_weather").key == "get_weather" + + def test_key_falls_back_to_type_for_hosted_tool(self): + # Hosted/server tools carry no name; the key is the type. + assert Tool(name=None, type="web_search").key == "web_search" + + def test_is_frozen(self): + tool = Tool(name="get_weather") + with pytest.raises(FrozenInstanceError): + setattr(tool, "name", "other") + + +class TestToolset: + def test_empty(self): + ts = Toolset() + assert ts.tools == () + assert ts.get("anything") is None + + def test_get_returns_function_and_hosted_tools_by_key(self): + fn = Tool(name="get_weather", arguments_schema=_WEATHER_SCHEMA) + hosted = Tool(name=None, type="web_search") + ts = Toolset(tools=[fn, hosted]) + + assert ts.get("get_weather") is fn + assert ts.get("web_search") is hosted + assert ts.get("missing") is None + + def test_duplicate_function_name_raises(self): + with pytest.raises(ValueError, match="duplicate tool 'dup'"): + Toolset(tools=[Tool(name="dup", description="first"), Tool(name="dup", description="second")]) + + def test_duplicate_hosted_tool_type_raises(self): + with pytest.raises(ValueError, match="duplicate tool 'web_search'"): + Toolset(tools=[Tool(name=None, type="web_search"), Tool(name=None, type="web_search")]) + + def test_tool_with_empty_key_is_skipped(self): + # name=None and an empty type => empty key => not indexed. + ts = Toolset(tools=[Tool(name=None, type="")]) + assert ts.get("") is None + + def test_multiple_empty_key_tools_do_not_collide(self): + # Empty-key tools are skipped, so duplicates among them never raise. + ts = Toolset(tools=[Tool(name=None, type=""), Tool(name=None, type="")]) + assert ts.get("") is None + + +class TestToolResult: + def test_defaults(self): + result = ToolResult() + assert result.call_id is None + assert result.name is None + assert result.content is None + assert result.is_error is False + + def test_string_content(self): + result = ToolResult(call_id="call_1", name="get_weather", content="18C") + assert result.call_id == "call_1" + assert result.name == "get_weather" + assert result.content == "18C" + assert result.is_error is False + + def test_block_content_and_error(self): + blocks = [{"type": "text", "text": "boom"}] + result = ToolResult(call_id="call_2", content=blocks, is_error=True) + assert result.content == blocks + assert result.is_error is True + + def test_is_frozen(self): + result = ToolResult(call_id="call_1") + with pytest.raises(FrozenInstanceError): + setattr(result, "call_id", "other") + + +class TestValidateArguments: + def test_valid_arguments_pass(self): + assert validate_arguments(_weather_tool(), {"city": "Paris", "days": 3}) is None + + def test_valid_with_only_required(self): + assert validate_arguments(_weather_tool(), {"city": "Paris"}) is None + + def test_type_mismatch_is_rejected(self): + reason = validate_arguments(_weather_tool(), {"city": 123}) + assert reason is not None + assert "get_weather" in reason + + def test_missing_required_is_rejected(self): + reason = validate_arguments(_weather_tool(), {}) + assert reason is not None + assert "get_weather" in reason + + def test_no_schema_skips_validation(self): + # A tool with no declared schema (hosted tool, or a function with no + # parameters) is allowlist-only: any arguments are accepted here. + hosted = Tool(name=None, type="web_search") + assert validate_arguments(hosted, {"anything": [1, 2, 3]}) is None + + def test_malformed_schema_is_reported_not_raised(self): + # A declared schema that is not valid JSON Schema degrades to a reason + # rather than raising (e.g. a non-JSON-Schema dialect reaching the + # validator before its engine adapter normalizes it). + bad = Tool(name="bad", arguments_schema={"type": "not-a-real-type"}) + reason = validate_arguments(bad, {}) + assert reason is not None + assert "bad" in reason diff --git a/tests/guardrails/tool_helpers.py b/tests/guardrails/tool_helpers.py new file mode 100644 index 0000000000..100728664b --- /dev/null +++ b/tests/guardrails/tool_helpers.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared fixtures and assertions for the IORails tool-rail tests. + +Lives alongside ``async_helpers.py`` / ``metric_helpers.py`` and is not collected +by pytest (no ``test_`` prefix). Holds the small tool shapes and the +blocked-result assertion reused across the tool-rail test modules. Assertions +carry explicit messages since this module is not assertion-rewritten by pytest. +""" + +WEATHER_SCHEMA = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], +} + + +def assert_blocked(result, *substrings: str) -> None: + """Assert *result* blocked the request and its reason contains each substring.""" + assert result.is_safe is False, f"expected blocked, got {result!r}" + assert result.reason is not None, f"expected a block reason, got {result!r}" + for substring in substrings: + assert substring in result.reason, f"{substring!r} not in reason {result.reason!r}" + + +def make_tool_conversation(result_call_id: str = "call_1") -> list: + """A user turn, an assistant ``get_weather`` tool call (id ``call_1``), then a tool result. + + ``result_call_id`` sets the tool result's ``tool_call_id`` so callers can test + linked (``call_1``) and unlinked (anything else) results. + """ + return [ + {"role": "user", "content": "What's the weather in Paris?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": result_call_id, "name": "get_weather", "content": "18C"}, + ]