feat(iorails): Connect tool-calling up in IORails#2058
Conversation
github-actions
Bot
added
status: needs triage
tgasser-nv
changed the base branch from
develop
to
feat/iorails-tool-calling-rails
|
@greptile-apps Review this PR |
|
@coderabbitai Review this PR |
|
✅ Action performedReview finished.
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Greptile SummaryThis PR completes the IORails tool-calling integration by wiring tool-call and tool-result rails into both the non-streaming (
|
| Filename | Overview |
|---|---|
| nemoguardrails/guardrails/iorails.py | Core pipeline extended with tool-result and tool-call rail stages; _frame_for_stream helper fixes include_metadata wrapping for output-rail violations; streaming input block changed from REFUSAL_MESSAGE to guardrails_violation JSON (intentional but undocumented breaking change). |
| nemoguardrails/guardrails/rails_manager.py | are_tool_calls_safe and are_tool_results_safe refactored to accept raw messages/llm_params and parse internally; per-request enabled toggle and _enabled_flows helper added cleanly. |
| nemoguardrails/guardrails/model_engine.py | _finalize_tool_calls now fails closed on malformed JSON (raises ValueError instead of degrading to {}); ToolExchange extraction with per-turn grouping added correctly. |
| nemoguardrails/guardrails/tool_schema.py | ToolExchange NamedTuple added; validate_arguments tightened to reject arguments for no-parameter function tools via _schema_accepts_no_arguments logic. |
| tests/guardrails/test_tool_rails_iorails.py | 577-line new test file covering routing fallback, non-streaming/streaming tool call and result scenarios, per-request toggles, metrics capture, and fail-closed behavior. |
| nemoguardrails/guardrails/actions/tool_result_action.py | _validate_result_name tightened to also reject missing result names when prior call name is known, closing the slip-through on call_id linkage alone. |
| nemoguardrails/guardrails/engine_registry.py | extract_tool_exchanges delegation added cleanly; follows existing extract_tool_results pattern. |
Sequence Diagram
%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant Client
participant IORails
participant RailsManager
participant LLM
Client->>IORails: generate_async(messages, options)
IORails->>RailsManager: are_tool_results_safe(messages)
Note over RailsManager: Extract ToolExchanges (per-turn grouping)
alt tool result unsafe
RailsManager-->>IORails: "RailResult(is_safe=False)"
IORails-->>Client: "{role:assistant, content:REFUSAL_MESSAGE}"
else safe
RailsManager-->>IORails: "RailResult(is_safe=True)"
IORails->>RailsManager: is_input_safe(messages)
alt input unsafe
RailsManager-->>IORails: "RailResult(is_safe=False)"
IORails-->>Client: "{role:assistant, content:REFUSAL_MESSAGE}"
else safe
IORails->>LLM: model_call(messages, llm_params)
LLM-->>IORails: LLMResponse(tool_calls / content)
alt response has tool_calls
IORails->>RailsManager: are_tool_calls_safe(tool_calls, llm_params)
Note over RailsManager: parse_tools then validate against toolset
alt tool call unsafe
RailsManager-->>IORails: "RailResult(is_safe=False)"
IORails-->>Client: "{role:assistant, content:REFUSAL_MESSAGE}"
else safe
IORails-->>Client: "{role:assistant, tool_calls:[...]}"
end
else text response
IORails->>RailsManager: is_output_safe(messages, response_text)
alt output unsafe
RailsManager-->>IORails: "RailResult(is_safe=False)"
IORails-->>Client: "{role:assistant, content:REFUSAL_MESSAGE}"
else safe
IORails-->>Client: "{role:assistant, content:response_text}"
end
end
end
end
%%{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 Client
participant IORails
participant RailsManager
participant LLM
Client->>IORails: generate_async(messages, options)
IORails->>RailsManager: are_tool_results_safe(messages)
Note over RailsManager: Extract ToolExchanges (per-turn grouping)
alt tool result unsafe
RailsManager-->>IORails: "RailResult(is_safe=False)"
IORails-->>Client: "{role:assistant, content:REFUSAL_MESSAGE}"
else safe
RailsManager-->>IORails: "RailResult(is_safe=True)"
IORails->>RailsManager: is_input_safe(messages)
alt input unsafe
RailsManager-->>IORails: "RailResult(is_safe=False)"
IORails-->>Client: "{role:assistant, content:REFUSAL_MESSAGE}"
else safe
IORails->>LLM: model_call(messages, llm_params)
LLM-->>IORails: LLMResponse(tool_calls / content)
alt response has tool_calls
IORails->>RailsManager: are_tool_calls_safe(tool_calls, llm_params)
Note over RailsManager: parse_tools then validate against toolset
alt tool call unsafe
RailsManager-->>IORails: "RailResult(is_safe=False)"
IORails-->>Client: "{role:assistant, content:REFUSAL_MESSAGE}"
else safe
IORails-->>Client: "{role:assistant, tool_calls:[...]}"
end
else text response
IORails->>RailsManager: is_output_safe(messages, response_text)
alt output unsafe
RailsManager-->>IORails: "RailResult(is_safe=False)"
IORails-->>Client: "{role:assistant, content:REFUSAL_MESSAGE}"
else safe
IORails-->>Client: "{role:assistant, content:response_text}"
end
end
end
end
Reviews (4): Last reviewed commit: "Add tests to make sure static Model para..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
📝 WalkthroughWalkthrough
ChangesTool Rails: per-request parsing, extraction, and IORails integration
Sequence Diagram(s)sequenceDiagram
participant Client
participant IORails
participant RailsManager
participant EngineRegistry
participant ModelEngine
participant LLM
rect rgba(135, 206, 250, 0.5)
note over Client,IORails: Tool-result pre-check (tool_input rail)
Client->>IORails: generate_async(messages, options)
IORails->>RailsManager: are_tool_results_safe(messages, enabled=tool_input_enabled)
RailsManager->>EngineRegistry: extract_tool_results(model_type, messages)
RailsManager->>EngineRegistry: extract_tool_calls(model_type, messages)
EngineRegistry->>ModelEngine: extract_tool_calls(messages)
ModelEngine-->>EngineRegistry: list[ToolCall]
RailsManager-->>IORails: RailResult (safe/unsafe)
end
alt tool result unsafe
IORails-->>Client: REFUSAL_MESSAGE
end
IORails->>IORails: run input rails
IORails->>LLM: prompt
LLM-->>IORails: response with tool_calls
rect rgba(255, 165, 0, 0.5)
note over IORails,RailsManager: Tool-call post-check (tool_output rail)
IORails->>RailsManager: are_tool_calls_safe(tool_calls, llm_params, enabled=tool_output_enabled)
RailsManager->>EngineRegistry: parse_tools(model_type, llm_params)
RailsManager-->>IORails: RailResult (safe/unsafe)
end
alt tool call unsafe
IORails-->>Client: REFUSAL_MESSAGE / guardrails_violation
end
IORails->>IORails: run output rails
IORails-->>Client: final response
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}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tests/guardrails/test_tool_rails_iorails.py`:
- Around line 299-303: The test
test_unlinked_tool_result_blocked_before_generation relies on guard behavior to
prevent provider calls but does not explicitly mock or assert against network
transport usage, which could allow accidental regressions to hit live provider
services. Inject a mock HTTP client (typically the post method used for provider
communication) into the iorails fixture or test setup, then add an assertion
after the generate_async call to verify that the mocked post method was never
invoked. Apply this same hardening pattern to the other related test mentioned
at line 330-332 to ensure both tests are network-impossible rather than
network-accidental.
🪄 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: 5a59e343-73c9-4aa7-8c9a-6c2248e662cc
📒 Files selected for processing (10)
nemoguardrails/guardrails/engine_registry.pynemoguardrails/guardrails/iorails.pynemoguardrails/guardrails/model_engine.pynemoguardrails/guardrails/rails_manager.pytests/guardrails/test_engine_registry.pytests/guardrails/test_guardrails.pytests/guardrails/test_model_engine.pytests/guardrails/test_rails_manager.pytests/guardrails/test_tool_rails_e2e.pytests/guardrails/test_tool_rails_iorails.py
tgasser-nv
changed the title
|
@tgasser-nv this is what claude code review returns, might be worth looking at Code Review — PR #2058: IORails tool-calling railsScope: Correctness1. Truncated/malformed streamed tool-call arguments silently degrade to
|
Pouyanpi
left a comment
Pouyanpi
left a comment
There was a problem hiding this comment.
Thanks @tgasser-nv . LGTM!
You can have a look at claude code's review above. For me the only thing we might need to address in future is commented below. Not blocking this PR
tgasser-nv
force-pushed
the
feat/iorails-tool-calling-api
branch
from
430f067 to
65d3feb
Compare
… if no inference-time values are passed
tgasser-nv
added
status: triaged
Addressed all but the refactoring one at the end (out-of-scope for this PR IMO). |
2 participants
Description
Stacked PR connecting up the tool-calling rails from the previous PR #2030 at IORails-level. Adds validation of tool-calling arguments and schema, and results of tool execution by the agent harness.
PR Stack
Related Issue(s)
Stacked PR on top of #2030 (which added tool-calling argument and results rails).
Verification
Pre-commit
Unit-test
Integration test with Chat
Integration test with Prod endpoint (using e2e_tool_rail_validation.py
Non-streaming
Streaming
AI Assistance
Checklist
Summary by CodeRabbit
New Features
Tests