Fix tests; address PR comments

Author: JashG

nemoguardrails/actions/llm/utils.py

@@ -931,18 +931,23 @@ def get_and_clear_tool_calls_contextvar() -> Optional[list]:
 
 
 def extract_tool_calls_from_events(events: list) -> Optional[list]:
-    """Extract tool_calls from BotToolCalls events.
+    """Extract tool_calls from runtime events.
 
-    Args:
-        events: List of events to search through
-
-    Returns:
-        tool_calls if found in BotToolCalls event, None otherwise
+    ``StartToolCallBotAction`` carries the tool calls that passed tool-output
+    rails and should be returned to the caller. ``BotToolCalls`` is used as a
+    fallback for paths that do not emit the post-rail action event.
     """
+    bot_tool_calls = None
+
     for event in events:
-        if event.get("type") == "BotToolCalls":
-            return event.get("tool_calls")
-    return None
+        if event.get("type") == "StartToolCallBotAction":
+            tool_calls = event.get("tool_calls")
+            if tool_calls is not None:
+                return tool_calls
+        elif event.get("type") == "BotToolCalls":
+            bot_tool_calls = event.get("tool_calls")
+
+    return bot_tool_calls
 
 
 def extract_bot_thinking_from_events(events: list):

tests/test_tool_calls_event_extraction.py

@@ -126,6 +126,53 @@ def mock_get_and_clear():
         assert result["tool_calls"][0]["name"] == "test_tool"
 
 
+@pytest.mark.asyncio
+async def test_extract_tool_calls_from_start_tool_call_action():
+    from nemoguardrails.actions.llm.utils import extract_tool_calls_from_events
+
+    test_tool_calls = [
+        {
+            "id": "call_approved",
+            "type": "function",
+            "function": {
+                "name": "approved_tool",
+                "arguments": {"data": "safe"},
+            },
+        }
+    ]
+
+    events = [{"type": "StartToolCallBotAction", "tool_calls": test_tool_calls}]
+
+    assert extract_tool_calls_from_events(events) == test_tool_calls
+
+
+@pytest.mark.asyncio
+async def test_extract_tool_calls_prefers_post_rail_action_event():
+    from nemoguardrails.actions.llm.utils import extract_tool_calls_from_events
+
+    pre_rail = [
+        {
+            "id": "call_modified",
+            "type": "function",
+            "function": {"name": "lookup", "arguments": {"query": "unfiltered"}},
+        }
+    ]
+    post_rail = [
+        {
+            "id": "call_modified",
+            "type": "function",
+            "function": {"name": "lookup", "arguments": {"query": "filtered"}},
+        }
+    ]
+
+    events = [
+        {"type": "BotToolCalls", "tool_calls": pre_rail},
+        {"type": "StartToolCallBotAction", "tool_calls": post_rail},
+    ]
+
+    assert extract_tool_calls_from_events(events) == post_rail
+
+
 @pytest.mark.asyncio
 async def test_llmrails_extracts_tool_calls_from_events():
     config = RailsConfig.from_content(config={"models": [], "passthrough": True})

tests/test_tool_output_rails.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 from unittest.mock import patch
 
 import pytest
@@ -24,6 +25,17 @@
 from tests.utils import FakeLLMModel, TestChat
 
 
+def _tool_arguments(func: dict) -> dict:
+    """Parse OpenAI-style tool call arguments (JSON string or dict)."""
+    arguments = func.get("arguments", {})
+    if isinstance(arguments, str):
+        try:
+            arguments = json.loads(arguments)
+        except json.JSONDecodeError:
+            return {}
+    return arguments if isinstance(arguments, dict) else {}
+
+
 @action(is_system_action=True)
 async def validate_tool_parameters(tool_calls, context=None, **kwargs):
     tool_calls = tool_calls or (context.get("tool_calls", []) if context else [])
@@ -32,7 +44,7 @@ async def validate_tool_parameters(tool_calls, context=None, **kwargs):
 
     for tool_call in tool_calls:
         func = tool_call.get("function", {})
-        args = func.get("arguments", {})
+        args = _tool_arguments(func)
         for param_value in args.values():
             if isinstance(param_value, str):
                 if any(pattern.lower() in param_value.lower() for pattern in dangerous_patterns):
@@ -230,7 +242,7 @@ async def test_assistant_tool_calls_run_tool_output_rails_when_dialog_disabled()
                     "type": "function",
                     "function": {
                         "name": "dangerous_tool",
-                        "arguments": {"param": "eval('malicious code')"},
+                        "arguments": '{"param": "eval(\'malicious code\')"}',
                     },
                 }
             ],
@@ -242,3 +254,62 @@ async def test_assistant_tool_calls_run_tool_output_rails_when_dialog_disabled()
     assert isinstance(result, GenerationResponse)
     assert isinstance(result.response, list)
     assert "parameters may be unsafe" in result.response[0]["content"]
+
+
+@pytest.mark.asyncio
+async def test_approved_assistant_tool_calls_are_returned_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_safe",
+                    "type": "function",
+                    "function": {
+                        "name": "safe_tool",
+                        "arguments": '{"param": "safe value"}',
+                    },
+                }
+            ],
+        },
+    ]
+
+    result = await rails.generate_async(messages=messages, options={"rails": {"dialog": False}})
+
+    assert isinstance(result, GenerationResponse)
+    assert result.tool_calls == [
+        {
+            "id": "call_safe",
+            "type": "function",
+            "function": {
+                "name": "safe_tool",
+                "arguments": '{"param": "safe value"}',
+            },
+        }
+    ]