Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 59 additions & 0 deletions nemoguardrails/guardrails/actions/tool_call_action.py
Original file line number Diff line number Diff line change
@@ -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)
147 changes: 147 additions & 0 deletions nemoguardrails/guardrails/actions/tool_result_action.py
Original file line number Diff line number Diff line change
@@ -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
28 changes: 28 additions & 0 deletions nemoguardrails/guardrails/engine_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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.

Expand Down
103 changes: 103 additions & 0 deletions nemoguardrails/guardrails/model_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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"),
)
)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
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."""

Expand Down Expand Up @@ -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()
Comment thread
tgasser-nv marked this conversation as resolved.
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)
Loading
Loading