fix: Preserve tool calls for LLMRails tool rails

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")

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]

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()

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"]