-
Notifications
You must be signed in to change notification settings - Fork 754
feat(iorails): Add tool-calling rails to RailsManager and ModelRegistry #2030
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from all commits
Commits
Show all changes
16 commits
Select commit
Hold shift + click to select a range
acb9f7b
Add internal canonical tool types and unit-tests
tgasser-nv 06bcff7
Add tool parser for OpenAI tools/functions
tgasser-nv 308858f
Add tool result parser
tgasser-nv 0d80d9f
Call parsing argument and results in EngineRegistry
tgasser-nv 289d7a1
Add tool rail action to check allowed functions and schema
tgasser-nv b238a50
Connect up Tool Call Rails in RailsManager
tgasser-nv 082f4c4
Change error message back
tgasser-nv 9f2eb51
Compact tests with helpers
tgasser-nv c11e829
Address greptile feedback: Toolset tools is a dict for O(1) tool acce…
tgasser-nv b64f337
Capture full tool-calling context
tgasser-nv ad0dd3b
Add validation for content
tgasser-nv 3704c79
Address code rabbit feedback
tgasser-nv 47a0cb7
Refactor tool result action into helpers, add check for duplicate cal…
tgasser-nv 4397ebb
Add test coverage to get to 100%
tgasser-nv e0a5140
Use ToolCall type for hosted functions if function.name is missing
tgasser-nv 98fed69
Merge tools from statis config and inference-time llm_params
tgasser-nv File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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
147
nemoguardrails/guardrails/actions/tool_result_action.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.