feat(iorails): Add tool-calling rails to RailsManager and ModelRegistry (#2030)

Author: tgasser-nv

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)

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

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.
 

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)