From 0d4084eda60d22464c6d8eb7520015c2dd8ca886 Mon Sep 17 00:00:00 2001 From: Jash Gulabrai Date: Thu, 25 Jun 2026 15:11:59 -0400 Subject: [PATCH 1/2] fix: Preserve tool calls for LLMRails tool rails --- nemoguardrails/logging/processing_log.py | 31 ++++++++++++++++++++++-- nemoguardrails/rails/llm/llmrails.py | 12 +++++++-- 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/nemoguardrails/logging/processing_log.py b/nemoguardrails/logging/processing_log.py index 00ddc554a9..056a6ef78f 100644 --- a/nemoguardrails/logging/processing_log.py +++ b/nemoguardrails/logging/processing_log.py @@ -47,6 +47,10 @@ def compute_generation_log(processing_log: List[dict]) -> GenerationLog: "run dialog rails", "process bot message", "run output rails", + "process bot tool call", + "process user tool messages", + "run tool output rails", + "run tool input rails", ] generation_flows = [ "generate bot message", @@ -129,6 +133,22 @@ def compute_generation_log(processing_log: List[dict]) -> GenerationLog: ) generation_log.activated_rails.append(activated_rail) + elif event_type == "StartToolOutputRail": + activated_rail = ActivatedRail( + type="tool_output", + name=event_data["flow_id"], + started_at=event["timestamp"], + ) + generation_log.activated_rails.append(activated_rail) + + elif event_type == "StartToolInputRail": + activated_rail = ActivatedRail( + type="tool_input", + name=event_data["flow_id"], + started_at=event["timestamp"], + ) + generation_log.activated_rails.append(activated_rail) + elif event_type == "StartInternalSystemAction": action_name = event_data["action_name"] if action_name in ignored_actions: @@ -154,7 +174,12 @@ def compute_generation_log(processing_log: List[dict]) -> GenerationLog: executed_action.return_value = event_data["return_value"] executed_action = None - elif event_type in ["InputRailFinished", "OutputRailFinished"]: + elif event_type in [ + "InputRailFinished", + "OutputRailFinished", + "ToolOutputRailFinished", + "ToolInputRailFinished", + ]: if activated_rail is not None: activated_rail.finished_at = event["timestamp"] if activated_rail.finished_at is not None and activated_rail.started_at is not None: @@ -171,6 +196,8 @@ def compute_generation_log(processing_log: List[dict]) -> GenerationLog: if activated_rail is not None and activated_rail.type in [ "input", "output", + "tool_output", + "tool_input", ]: activated_rail.stop = True if "stop" not in activated_rail.decisions: @@ -188,7 +215,7 @@ def compute_generation_log(processing_log: List[dict]) -> GenerationLog: if activated_rail.finished_at is not None and activated_rail.started_at is not None: activated_rail.duration = activated_rail.finished_at - activated_rail.started_at - if activated_rail.type in ["input", "output"]: + if activated_rail.type in ["input", "output", "tool_output", "tool_input"]: activated_rail.stop = True if "stop" not in activated_rail.decisions: activated_rail.decisions.append("stop") diff --git a/nemoguardrails/rails/llm/llmrails.py b/nemoguardrails/rails/llm/llmrails.py index 63c0fa27b7..35eb99a0ff 100644 --- a/nemoguardrails/rails/llm/llmrails.py +++ b/nemoguardrails/rails/llm/llmrails.py @@ -1017,8 +1017,16 @@ async def generate_async( # If the last message is from the assistant, rather than the user, then # we move that to the `$bot_message` variable. This is to enable a more - # convenient interface. (only when dialog rails are disabled) - if messages and messages[-1]["role"] == "assistant" and gen_options and gen_options.rails.dialog is False: + # convenient interface for text output rails. Tool-call assistant messages + # must remain in the history so they can be converted into BotToolCalls + # events and evaluated by tool output rails. + if ( + messages + and messages[-1]["role"] == "assistant" + and not messages[-1].get("tool_calls") + and gen_options + and gen_options.rails.dialog is False + ): # We already have the first message with a context update, so we use that messages[0]["content"]["bot_message"] = messages[-1]["content"] messages = messages[0:-1] From 5b46fe9ece58203d1b51e18adcb7a7b6a9925c2e Mon Sep 17 00:00:00 2001 From: Jash Gulabrai Date: Thu, 25 Jun 2026 16:34:08 -0400 Subject: [PATCH 2/2] Add unit tests --- tests/test_logging.py | 21 ++++++++++++- tests/test_tool_output_rails.py | 52 +++++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 1 deletion(-) diff --git a/tests/test_logging.py b/tests/test_logging.py index 8fa230c683..8f3e64c05e 100644 --- a/tests/test_logging.py +++ b/tests/test_logging.py @@ -21,11 +21,30 @@ from nemoguardrails.context import explain_info_var, llm_call_info_var, llm_stats_var from nemoguardrails.logging.explain import ExplainInfo, LLMCallInfo from nemoguardrails.logging.llm_tracker import track_llm_call -from nemoguardrails.logging.processing_log import processing_log_var +from nemoguardrails.logging.processing_log import compute_generation_log, processing_log_var from nemoguardrails.logging.stats import LLMStats from nemoguardrails.types import LLMResponse, UsageInfo +def test_compute_generation_log_includes_tool_rails(): + generation_log = compute_generation_log( + [ + {"type": "step", "flow_id": "process bot tool call", "timestamp": 0.0, "next_steps": []}, + {"type": "event", "timestamp": 1.0, "data": {"type": "StartToolOutputRail", "flow_id": "check tool call"}}, + {"type": "event", "timestamp": 1.25, "data": {"type": "ToolOutputRailFinished"}}, + {"type": "step", "flow_id": "process user tool messages", "timestamp": 2.0, "next_steps": []}, + {"type": "event", "timestamp": 3.0, "data": {"type": "StartToolInputRail", "flow_id": "check tool result"}}, + {"type": "event", "timestamp": 3.5, "data": {"type": "ToolInputRailFinished"}}, + ] + ) + + activated_rails = generation_log.activated_rails + + assert [rail.type for rail in activated_rails] == ["tool_output", "tool_input"] + assert [rail.name for rail in activated_rails] == ["check tool call", "check tool result"] + assert [rail.duration for rail in activated_rails] == [0.25, 0.5] + + @pytest.mark.asyncio async def test_token_usage_tracking_with_usage(): llm_call_info = LLMCallInfo() diff --git a/tests/test_tool_output_rails.py b/tests/test_tool_output_rails.py index f587be5365..6bd29de140 100644 --- a/tests/test_tool_output_rails.py +++ b/tests/test_tool_output_rails.py @@ -19,6 +19,7 @@ from nemoguardrails import LLMRails, RailsConfig from nemoguardrails.actions import action +from nemoguardrails.rails.llm.options import GenerationResponse from nemoguardrails.types import LLMResponse, ToolCall, ToolCallFunction from tests.utils import FakeLLMModel, TestChat @@ -190,3 +191,54 @@ async def test_multiple_tool_output_rails(): assert result["tool_calls"] is not None assert result["tool_calls"][0]["name"] == "test_tool" + + +@pytest.mark.asyncio +async def test_assistant_tool_calls_run_tool_output_rails_when_dialog_disabled(): + config = RailsConfig.from_content( + """ + define subflow validate tool parameters + $valid = execute validate_tool_parameters(tool_calls=$tool_calls) + + if not $valid + bot refuse dangerous tool parameters + abort + + define bot refuse dangerous tool parameters + "I cannot execute this tool request because the parameters may be unsafe." + """, + """ + models: [] + passthrough: true + rails: + tool_output: + flows: + - validate tool parameters + """, + ) + rails = LLMRails(config) + rails.runtime.register_action(validate_tool_parameters, name="validate_tool_parameters") + + messages = [ + {"role": "user", "content": "Use the requested tool"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_bad", + "type": "function", + "function": { + "name": "dangerous_tool", + "arguments": {"param": "eval('malicious code')"}, + }, + } + ], + }, + ] + + result = await rails.generate_async(messages=messages, options={"rails": {"dialog": False}}) + + assert isinstance(result, GenerationResponse) + assert isinstance(result.response, list) + assert "parameters may be unsafe" in result.response[0]["content"]