feat(iorails): Add tool-calling rails to RailsManager and ModelRegistry#2030
Merged
Conversation
tgasser-nv
changed the base branch from
develop
to
feat/iorails-tool-calling-streaming
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Contributor
Greptile SummaryThis PR adds two new tool-calling rail validators —
|
| Filename | Overview |
|---|---|
| nemoguardrails/guardrails/tool_schema.py | New file: frozen Tool and ToolResult dataclasses, immutable Toolset backed by a single dict index, and validate_arguments wrapping jsonschema; well-designed with duplicate-key rejection and proper SchemaError handling. |
| nemoguardrails/guardrails/tool_rail_action.py | New base class for model-free tool rails; _guarded wraps any synchronous check in an action span and converts unexpected exceptions to blocking results (fail-closed), consistent with RailAction's error contract. |
| nemoguardrails/guardrails/actions/tool_call_action.py | New ToolCallRailAction: iterates tool calls, checks each against the toolset allowlist, then validates arguments via jsonschema; correctly falls back to call.type for hosted/server tools. |
| nemoguardrails/guardrails/actions/tool_result_action.py | New ToolResultRailAction: validates call_id linkage to prior calls, name consistency, and content shape; blocks on duplicate call IDs and missing/orphaned results; previous review concern about content type check is fully addressed. |
| nemoguardrails/guardrails/rails_manager.py | Adds tool_call_flows/tool_result_flows constructor params, _build_tool_actions with duplicate/direction guards, and are_tool_calls_safe/are_tool_results_safe; correctly reuses _run_rails_sequential with proper coroutine cleanup. |
| nemoguardrails/guardrails/model_engine.py | Adds _parse_tools_openai/NIM parsers and _extract_tool_results_openai/NIM extractors with fallback to OpenAI shape for unknown engines; parse_tools and extract_tool_results methods properly skip malformed entries. |
| nemoguardrails/guardrails/engine_registry.py | Adds parse_tools and extract_tool_results delegation methods that merge body_param_defaults with per-call llm_params before parsing; raises KeyError/TypeError appropriately for unknown or non-model engines. |
| tests/guardrails/test_tool_rails_e2e.py | End-to-end tests exercising the full parse→validate seam for both non-streaming and streaming tool calls, plus tool result linkage; mocks only aiohttp transport so real ModelEngine parsing runs. |
Sequence Diagram
%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant Caller
participant EngineRegistry
participant ModelEngine
participant RailsManager
participant ToolCallRailAction
participant ToolResultRailAction
Caller->>EngineRegistry: parse_tools(model_type, llm_params)
EngineRegistry->>ModelEngine: parse_tools(merged_params)
ModelEngine-->>EngineRegistry: Toolset
EngineRegistry-->>Caller: Toolset
Caller->>EngineRegistry: "model_call(messages, tools=...)"
EngineRegistry->>ModelEngine: HTTP POST /v1/chat/completions
ModelEngine-->>EngineRegistry: "LLMResponse(tool_calls=[...])"
EngineRegistry-->>Caller: LLMResponse
Caller->>RailsManager: are_tool_calls_safe(tool_calls, toolset)
RailsManager->>ToolCallRailAction: run(toolset, tool_calls)
Note over ToolCallRailAction: Allowlist check (name in toolset) + Schema validation (jsonschema)
ToolCallRailAction-->>RailsManager: "RailResult(is_safe=True/False)"
RailsManager-->>Caller: RailResult
Note over Caller: Execute tools if safe, collect results
Caller->>EngineRegistry: extract_tool_results(model_type, messages)
EngineRegistry->>ModelEngine: extract_tool_results(messages)
ModelEngine-->>EngineRegistry: list[ToolResult]
EngineRegistry-->>Caller: list[ToolResult]
Caller->>RailsManager: are_tool_results_safe(tool_results, prior_calls)
RailsManager->>ToolResultRailAction: run(tool_results, prior_calls)
Note over ToolResultRailAction: call_id linkage + name consistency + content shape
ToolResultRailAction-->>RailsManager: "RailResult(is_safe=True/False)"
RailsManager-->>Caller: RailResult
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant Caller
participant EngineRegistry
participant ModelEngine
participant RailsManager
participant ToolCallRailAction
participant ToolResultRailAction
Caller->>EngineRegistry: parse_tools(model_type, llm_params)
EngineRegistry->>ModelEngine: parse_tools(merged_params)
ModelEngine-->>EngineRegistry: Toolset
EngineRegistry-->>Caller: Toolset
Caller->>EngineRegistry: "model_call(messages, tools=...)"
EngineRegistry->>ModelEngine: HTTP POST /v1/chat/completions
ModelEngine-->>EngineRegistry: "LLMResponse(tool_calls=[...])"
EngineRegistry-->>Caller: LLMResponse
Caller->>RailsManager: are_tool_calls_safe(tool_calls, toolset)
RailsManager->>ToolCallRailAction: run(toolset, tool_calls)
Note over ToolCallRailAction: Allowlist check (name in toolset) + Schema validation (jsonschema)
ToolCallRailAction-->>RailsManager: "RailResult(is_safe=True/False)"
RailsManager-->>Caller: RailResult
Note over Caller: Execute tools if safe, collect results
Caller->>EngineRegistry: extract_tool_results(model_type, messages)
EngineRegistry->>ModelEngine: extract_tool_results(messages)
ModelEngine-->>EngineRegistry: list[ToolResult]
EngineRegistry-->>Caller: list[ToolResult]
Caller->>RailsManager: are_tool_results_safe(tool_results, prior_calls)
RailsManager->>ToolResultRailAction: run(tool_results, prior_calls)
Note over ToolResultRailAction: call_id linkage + name consistency + content shape
ToolResultRailAction-->>RailsManager: "RailResult(is_safe=True/False)"
RailsManager-->>Caller: RailResult
Reviews (6): Last reviewed commit: "Merge tools from statis config and infer..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
9 tasks
tgasser-nv
force-pushed
the
feat/iorails-tool-calling-rails
branch
from
22a5d8c to
9f2eb51
Compare
Collaborator
Author
|
@greptile-apps Review this PR |
Collaborator
Author
|
@coderabbitai Review this PR |
Contributor
|
✅ Action performedFull review finished. |
Contributor
📝 WalkthroughWalkthroughThis PR adds a full "tool rails" feature to NeMo Guardrails. It introduces provider-neutral ChangesTool Call and Result Validation Rails
Sequence Diagram(s)sequenceDiagram
participant Caller
participant EngineRegistry
participant ModelEngine
participant RailsManager
participant ToolCallRailAction
participant ToolResultRailAction
rect rgba(70, 130, 180, 0.5)
Note over Caller, ModelEngine: Tool parsing (pre-inference)
Caller->>EngineRegistry: parse_tools(model_type, llm_params)
EngineRegistry->>ModelEngine: parse_tools(llm_params)
ModelEngine-->>EngineRegistry: Toolset
EngineRegistry-->>Caller: Toolset
end
rect rgba(100, 160, 100, 0.5)
Note over Caller, RailsManager: Tool call validation (post-inference)
Caller->>RailsManager: are_tool_calls_safe(tool_calls, toolset)
RailsManager->>ToolCallRailAction: run(toolset, tool_calls)
Note over ToolCallRailAction: allowlist + jsonschema validation
ToolCallRailAction-->>RailsManager: RailResult
RailsManager-->>Caller: RailResult
end
rect rgba(180, 100, 70, 0.5)
Note over Caller, ModelEngine: Tool result extraction
Caller->>EngineRegistry: extract_tool_results(model_type, messages)
EngineRegistry->>ModelEngine: extract_tool_results(messages)
ModelEngine-->>EngineRegistry: list[ToolResult]
EngineRegistry-->>Caller: list[ToolResult]
end
rect rgba(150, 80, 180, 0.5)
Note over Caller, RailsManager: Tool result validation
Caller->>RailsManager: are_tool_results_safe(tool_results, prior_calls)
RailsManager->>ToolResultRailAction: run(tool_results, prior_calls)
Note over ToolResultRailAction: call_id linkage + name + content checks
ToolResultRailAction-->>RailsManager: RailResult
RailsManager-->>Caller: RailResult
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@nemoguardrails/guardrails/actions/tool_result_action.py`:
- Around line 69-73: The isinstance check for result.content on line 69
currently accepts any list without validating the structure of items within it,
which allows lists containing non-parseable item types. Strengthen the
validation by adding an additional check after confirming content is a list:
iterate through the list items and verify that each item is a dict-like object
(isinstance check against dict). If any list item is not dict-like, return a
RailResult with is_safe=False and an appropriate reason message indicating
malformed list item structure. This ensures the rail enforces structural
validity for list content.
- Around line 50-60: Before building the calls_by_id lookup dictionary from
prior_calls, add validation to detect and reject duplicate non-empty IDs. Check
if any call.id values appear more than once in prior_calls and return a
RailResult with is_safe=False and an appropriate reason message when duplicates
are detected. Only proceed to construct the calls_by_id dictionary and process
tool_results after this validation passes.
In `@nemoguardrails/guardrails/model_engine.py`:
- Around line 313-324: The code is parsing function tools without validating
that each function has a valid non-empty name, which causes duplicate key errors
when creating a Toolset later. Add a validation check after the isinstance check
to verify that function.get("name") returns a non-empty value, and skip the
entry by using continue if the name is missing or empty, ensuring only
well-formed tools with valid names are appended to the parsed list.
In `@nemoguardrails/guardrails/rails_manager.py`:
- Around line 218-219: The dictionary comprehension building the `rails`
variable uses `flow` as the key, which means duplicate flow strings in
`self.tool_call_flows` will overwrite earlier entries and cause un-awaited
coroutines. Before the dict comprehension that calls
`self._run_tool_call_rail()`, either deduplicate `self.tool_call_flows` (by
converting to a set then back to a list) or add validation to ensure no
duplicates exist in the flows. Apply the same fix to the similar dictionary
comprehension at lines 231-232.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e839b17e-d34a-4974-a58c-0e53050c1834
📒 Files selected for processing (16)
nemoguardrails/guardrails/actions/tool_call_action.pynemoguardrails/guardrails/actions/tool_result_action.pynemoguardrails/guardrails/engine_registry.pynemoguardrails/guardrails/model_engine.pynemoguardrails/guardrails/rails_manager.pynemoguardrails/guardrails/tool_rail_action.pynemoguardrails/guardrails/tool_schema.pytests/guardrails/test_engine_registry.pytests/guardrails/test_model_engine.pytests/guardrails/test_rails_manager.pytests/guardrails/test_tool_call_action.pytests/guardrails/test_tool_rail_action.pytests/guardrails/test_tool_rails_e2e.pytests/guardrails/test_tool_result_action.pytests/guardrails/test_tool_schema.pytests/guardrails/tool_helpers.py
11 tasks
11 tasks
Pouyanpi
approved these changes
Jun 23, 2026
Pouyanpi
left a comment
Pouyanpi
left a comment
Collaborator
There was a problem hiding this comment.
@tgasser-nv should be good to merge after fixing following potential bugs 👍🏻
Collaborator
|
@tgasser-nv for the future PRs, please change the status from |
tgasser-nv
added
the
status: triaged
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Description
This PR adds Tool Rails to RailsManager and ModelRegistry. It adds two ToolRailActions,
tool call validationandtool result validation.tool call validation: Ensures the tool being called is available for use, and the argument schema matches the tool invocation itself.tool result validation: Ensures the result of a tool call has a call_id matching a previous ToolCall id, that function names, and the response is parseable.There's no local integration of tool-calling in this PR, since that only makes sense in the next PR in the stack where we integrate into the IORails top-level class.
PR Stack
Related Issue(s)
Stacked PR on top of #2024 (which added non-streaming tool-calling).
Test Plan
Pre-commit
Unit-test
Integration test with Chat
AI Assistance
Checklist
Summary by CodeRabbit
Release Notes