feat(tool-rails): add support for tool input rails and validation #1386

Pouyanpi · 2025-09-10T12:15:02Z

PR Description

Tool Input Rails Implementation

This PR implements comprehensive tool input rails functionality to validate and secure tool messages before they are processed by the LLM.

Key Changes:

Added UserToolMessages event processing with grouped tool message handling
Fixed message-to-event conversion for proper tool message support in conversation flows
Enhanced passthrough mode to preserve full conversation context including tool calls
Implemented tool input rail flows supporting validation, safety checks, and sanitization
Comprehensive test coverage for all tool input rails scenarios

Benefits:

Enables validation and filtering of potentially dangerous tool outputs
Provides sanitization capabilities for sensitive information in tool responses
Maintains conversation context integrity during tool interactions
Supports configurable tool input validation pipelines

Copilot

Pull Request Overview

This PR implements comprehensive tool input rails functionality to validate and secure tool messages before they are processed by the LLM.

Added UserToolMessages event processing with grouped tool message handling
Fixed message-to-event conversion for proper tool message support in conversation flows
Enhanced passthrough mode to preserve full conversation context including tool calls

Reviewed Changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
tests/test_tool_calls_event_extraction.py	Updates test for event name change from BotToolCall to BotToolCalls
tests/test_tool_calling_passthrough_only.py	Updates test assertions and comments for event name change
tests/test_input_tool_rails.py	Comprehensive test suite for tool input rails functionality
tests/test_bot_tool_call_events.py	Updates comments to reflect event name change
tests/runnable_rails/test_tool_calling.py	Updates test comment for event name change
tests/runnable_rails/test_runnable_rails.py	Updates test assertion for improved passthrough mode behavior
tests/input_tool_rails_actions.py	Test utility actions for tool input validation, safety checking, and sanitization
nemoguardrails/rails/llm/llmrails.py	Core implementation for message-to-event conversion with tool message support
nemoguardrails/rails/llm/llm_flows.co	Flow definitions for tool input rails processing
nemoguardrails/integrations/langchain/runnable_rails.py	Enhanced message conversion with tool call support
nemoguardrails/actions/llm/utils.py	Updated tool call extraction functions for new event name
nemoguardrails/actions/llm/generation.py	Enhanced passthrough mode to preserve tool message context

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

tests/test_tool_calls_event_extraction.py

nemoguardrails/rails/llm/llmrails.py

…ion and processing - Add UserToolMessages event handling and tool input rails processing - Fix message-to-event conversion to properly handle tool messages in conversation history - Preserve tool call context in passthrough mode by using full conversation history - Support tool_calls and tool message metadata in LangChain format conversion - Include comprehensive test suite for tool input rails functionality test(runnable_rails): fix prompt format in passthrough mode feat: support ToolMessage in message dicts refactor: rename BotToolCall to BotToolCalls

Copilot

Pull Request Overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 2 comments.

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

Copilot · 2025-09-18T14:42:55Z

nemoguardrails/actions/llm/generation.py

+                    raw_prompt = raw_llm_request.get()
+
+                    if raw_prompt is not None and isinstance(raw_prompt, list):


The raw_llm_request.get() call and the logic for handling tool message context in passthrough mode lacks documentation. Consider adding comments explaining when this context variable is set and how it differs from the regular prompt handling.

Copilot · 2025-09-18T14:42:56Z

nemoguardrails/rails/llm/llm_flows.co

+        $tool_message = $tool_messages[$i].content
+        $tool_name = $tool_messages[$i].name
+        if "tool_call_id" in $tool_messages[$i]
+          $tool_call_id = $tool_messages[$i].tool_call_id
+        else
+          $tool_call_id = ""


[nitpick] The manual iteration and field access pattern is repeated throughout the code. Consider extracting this into a helper function or using a more structured approach to access tool message fields consistently.

tgasser-nv

Looks good, approving. Please address feedback before merging. This is another >1k line PR which should be broken into smaller stacked PRs to review.

Please also run local integration-tests to make a few tool-calls to production LLMs to check this works before merging. There are also some Github CI tests failing, could you take a look before merging?

tgasser-nv · 2025-09-19T20:22:29Z

nemoguardrails/actions/llm/generation.py

@@ -591,7 +591,7 @@ async def generate_user_intent(

                if tool_calls:
                    output_events.append(
-                        new_event_dict("BotToolCall", tool_calls=tool_calls)
+                        new_event_dict("BotToolCalls", tool_calls=tool_calls)


Could you use an enum for these events? I'm guessing there's only a static limited number that are known ahead of time

tgasser-nv · 2025-09-19T20:23:50Z

nemoguardrails/actions/llm/utils.py

+                    AIMessage(content=msg["content"], tool_calls=tool_calls)
+                )
+            else:
+                messages.append(AIMessage(content=msg["content"]))
        elif msg_type == "system":


Also need to add developer which OpenAI uses rather than system for models aftero1

tgasser-nv · 2025-09-19T20:26:13Z

nemoguardrails/rails/llm/llm_flows.co

@@ -130,6 +130,40 @@ define parallel extension flow process bot tool call
  create event StartToolCallBotAction(tool_calls=$tool_calls)


+define parallel flow process user tool messages
+  """Run all the tool input rails on the tool messages."""
+  priority 200


Why is this set to 200? Is the aim to run all tool calls before other LLM calls?

tgasser-nv · 2025-09-19T20:28:51Z

tests/input_tool_rails_actions.py

+
+
+@action(is_system_action=True)
+async def validate_tool_input_safety(


nit: Rename to _is_tool_input_safe()? Then the "polarity" of the bool return is clear

tgasser-nv · 2025-09-19T20:29:32Z

tests/input_tool_rails_actions.py

+
+
+@action(is_system_action=True)
+async def self_check_tool_input(


nit: Rename to is_tool_input_safe_self_check() or similar? The polarity of the bool return type is clearer

tgasser-nv · 2025-09-19T20:30:23Z

tests/input_tool_rails_actions.py

+
+    log.debug(f"Validating safety of tool input from {tool_name}")
+
+    dangerous_patterns = [


Are we recommending using these in production? Seems like there might be valid uses for an api_key argument in a tool call for example

Pouyanpi added this to the v0.17.0 milestone Sep 10, 2025

Pouyanpi self-assigned this Sep 10, 2025

Pouyanpi added the enhancement New feature or request label Sep 10, 2025

Pouyanpi requested a review from Copilot September 10, 2025 12:21

Copilot AI reviewed Sep 10, 2025

View reviewed changes

tests/test_tool_calls_event_extraction.py Show resolved Hide resolved

nemoguardrails/rails/llm/llmrails.py Show resolved Hide resolved

Pouyanpi changed the title ~~Feat/tool calling input~~ eat(tool-rails): add support for tool input rails and validation Sep 10, 2025

Pouyanpi changed the title ~~eat(tool-rails): add support for tool input rails and validation~~ feat(tool-rails): add support for tool input rails and validation Sep 10, 2025

Pouyanpi force-pushed the feat/tool-calling-input branch from 0358cd7 to 3240dc9 Compare September 15, 2025 09:35

Pouyanpi force-pushed the feat/tool-output-rails branch from df75862 to 56cb9dd Compare September 15, 2025 09:38

Pouyanpi force-pushed the feat/tool-calling-input branch 2 times, most recently from 21e33e2 to 2f57ec4 Compare September 15, 2025 09:46

Pouyanpi force-pushed the feat/tool-output-rails branch 2 times, most recently from 44a25f3 to c8a4959 Compare September 15, 2025 09:54

Pouyanpi force-pushed the feat/tool-calling-input branch 2 times, most recently from ed234d6 to 4c34032 Compare September 15, 2025 10:05

Pouyanpi force-pushed the feat/tool-calling-input branch from 4c34032 to c8ff064 Compare September 15, 2025 11:01

Pouyanpi requested review from tgasser-nv and Copilot September 15, 2025 11:01

Copilot AI reviewed Sep 18, 2025

View reviewed changes

tgasser-nv approved these changes Sep 19, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat(tool-rails): add support for tool input rails and validation #1386

feat(tool-rails): add support for tool input rails and validation #1386

Uh oh!

Pouyanpi commented Sep 10, 2025

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Sep 18, 2025

Uh oh!

Copilot AI Sep 18, 2025

Uh oh!

tgasser-nv left a comment •

edited

Loading

Uh oh!

tgasser-nv Sep 19, 2025

Uh oh!

tgasser-nv Sep 19, 2025

Uh oh!

tgasser-nv Sep 19, 2025

Uh oh!

tgasser-nv Sep 19, 2025

Uh oh!

tgasser-nv Sep 19, 2025

Uh oh!

tgasser-nv Sep 19, 2025

Uh oh!

Uh oh!

		raw_prompt = raw_llm_request.get()

		if raw_prompt is not None and isinstance(raw_prompt, list):



		@action(is_system_action=True)
		async def validate_tool_input_safety(



		@action(is_system_action=True)
		async def self_check_tool_input(


		log.debug(f"Validating safety of tool input from {tool_name}")

		dangerous_patterns = [

feat(tool-rails): add support for tool input rails and validation #1386

Are you sure you want to change the base?

feat(tool-rails): add support for tool input rails and validation #1386

Uh oh!

Conversation

Pouyanpi commented Sep 10, 2025

PR Description

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Uh oh!

Copilot AI Sep 18, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Sep 18, 2025

Choose a reason for hiding this comment

Uh oh!

tgasser-nv left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tgasser-nv Sep 19, 2025

Choose a reason for hiding this comment

Uh oh!

tgasser-nv Sep 19, 2025

Choose a reason for hiding this comment

Uh oh!

tgasser-nv Sep 19, 2025

Choose a reason for hiding this comment

Uh oh!

tgasser-nv Sep 19, 2025

Choose a reason for hiding this comment

Uh oh!

tgasser-nv Sep 19, 2025

Choose a reason for hiding this comment

Uh oh!

tgasser-nv Sep 19, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tgasser-nv left a comment •

edited

Loading