test: add streaming coverage for per-tool rails

tests/guardrails/test_per_tool_regex_rails_iorails.py

@@ -21,20 +21,26 @@
 duplicating them.
 """
 
+import json
+
 import pytest
 import pytest_asyncio
 
 from nemoguardrails.guardrails.iorails import REFUSAL_MESSAGE
 from tests.guardrails.async_helpers import started_iorails
 from tests.guardrails.test_tool_rails_iorails import (
+    StreamedToolCall,
     _inject_forbidden_transport,
     _inject_json_response,
     _inject_sse_stream,
     _stream_violation_chunks,
     _text_payload,
+    _text_sse_lines,
     _tool_call_payload,
     _tool_call_sse_lines,
+    _tool_calls_sse_lines,
 )
+from tests.guardrails.tool_helpers import make_tool_conversation
 
 BASE_CONFIG = {"models": [{"type": "main", "engine": "nim", "model": "meta/llama-3.3-70b-instruct"}]}
 
@@ -166,30 +172,65 @@ async def test_non_matching_pattern_streams_through(self, call_pattern_iorails):
         chunks = await _collect(call_pattern_iorails.stream_async(MESSAGES))
         assert _stream_violation_chunks(chunks) == []
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("matching_index", [0, 1, 2])
+    async def test_matching_call_at_any_position_blocks_stream(self, call_pattern_iorails, matching_index):
+        # Only one of three streamed calls matches, so this fails if any position skips the rails.
+        queries = ["SELECT 1", "SELECT 2", "SELECT 3"]
+        queries[matching_index] = "DROP TABLE users"
+        calls = [StreamedToolCall("run_sql", json.dumps({"query": query})) for query in queries]
+        _inject_sse_stream(call_pattern_iorails, _tool_calls_sse_lines(*calls))
+        chunks = await _collect(call_pattern_iorails.stream_async(MESSAGES))
+        violations = _stream_violation_chunks(chunks)
+        assert len(violations) == 1
+        assert violations[0]["error"]["param"] == "tool_output_rails"
+        assert not any('"tool_calls"' in chunk for chunk in chunks)
+
+    @pytest.mark.asyncio
+    async def test_several_non_matching_calls_stream_through(self, call_pattern_iorails):
+        arguments = [json.dumps({"query": query}) for query in ("SELECT 1", "SELECT 2", "SELECT 3")]
+        _inject_sse_stream(
+            call_pattern_iorails, _tool_calls_sse_lines(*(StreamedToolCall("run_sql", a) for a in arguments))
+        )
+        chunks = await _collect(call_pattern_iorails.stream_async(MESSAGES))
+        assert _stream_violation_chunks(chunks) == []
+        streamed = json.loads(chunks[-1])["tool_calls"]
+        assert [call["function"]["arguments"] for call in streamed] == arguments
+
 
-class TestNonStreamingPerToolResultRegex:
-    def _tool_conversation(self, content: str) -> list:
-        return [
-            {"role": "user", "content": "run a query"},
-            {
-                "role": "assistant",
-                "content": None,
-                "tool_calls": [
-                    {"id": "call_1", "type": "function", "function": {"name": "run_sql", "arguments": "{}"}}
-                ],
-            },
-            {"role": "tool", "tool_call_id": "call_1", "name": "run_sql", "content": content},
-        ]
+def _tool_conversation(content: str) -> list:
+    """A run_sql call answered by a tool result carrying *content*."""
+    return make_tool_conversation(result_name="run_sql", tool_name="run_sql", arguments="{}", content=content)
 
+
+class TestNonStreamingPerToolResultRegex:
     @pytest.mark.asyncio
     async def test_matching_pattern_blocked_before_generation(self, result_pattern_iorails):
         forbidden_post = _inject_forbidden_transport(result_pattern_iorails)
-        result = await result_pattern_iorails.generate_async(messages=self._tool_conversation("ssn: 123-45-6789"))
+        result = await result_pattern_iorails.generate_async(messages=_tool_conversation("ssn: 123-45-6789"))
         assert result == {"role": "assistant", "content": REFUSAL_MESSAGE}
         forbidden_post.assert_not_called()
 
     @pytest.mark.asyncio
     async def test_non_matching_pattern_passes(self, result_pattern_iorails):
         _inject_json_response(result_pattern_iorails, _text_payload("no sensitive data found"))
-        result = await result_pattern_iorails.generate_async(messages=self._tool_conversation("no sensitive data"))
+        result = await result_pattern_iorails.generate_async(messages=_tool_conversation("no sensitive data"))
         assert result == {"role": "assistant", "content": "no sensitive data found"}
+
+
+class TestStreamingPerToolResultRegex:
+    @pytest.mark.asyncio
+    async def test_matching_pattern_blocks_stream_before_generation(self, result_pattern_iorails):
+        forbidden_post = _inject_forbidden_transport(result_pattern_iorails)
+        chunks = await _collect(result_pattern_iorails.stream_async(_tool_conversation("ssn: 123-45-6789")))
+        violations = _stream_violation_chunks(chunks)
+        assert len(violations) == 1
+        assert violations[0]["error"]["param"] == "tool_input_rails"
+        forbidden_post.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_non_matching_pattern_streams_through(self, result_pattern_iorails):
+        _inject_sse_stream(result_pattern_iorails, _text_sse_lines("ok"))
+        chunks = await _collect(result_pattern_iorails.stream_async(_tool_conversation("no sensitive data")))
+        assert _stream_violation_chunks(chunks) == []
+        assert "ok" in chunks

tests/guardrails/test_tool_rails_iorails.py

@@ -23,6 +23,8 @@
 """
 
 import json
+from dataclasses import dataclass
+from typing import Optional, Union
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -188,44 +190,38 @@ def _sse(chunk: dict) -> bytes:
     return ("data: " + json.dumps(chunk) + "\n\n").encode("utf-8")
 
 
-def _tool_call_sse_lines(name: str, arg_fragments: list, call_id: str = "call_1") -> list:
-    """SSE lines streaming a single tool call: id/name first, then argument fragments, then finish."""
-    chunks = [
-        {
-            "id": "chatcmpl-1",
-            "choices": [
-                {
-                    "index": 0,
-                    "delta": {
-                        "role": "assistant",
-                        "tool_calls": [
-                            {"index": 0, "id": call_id, "type": "function", "function": {"name": name, "arguments": ""}}
-                        ],
-                    },
-                    "finish_reason": None,
-                }
-            ],
-        }
+@dataclass(frozen=True)
+class StreamedToolCall:
+    """One tool call for a streamed response; ``call_id`` defaults to ``call_<its index>``."""
+
+    name: str
+    # One string streams the arguments in a single chunk; a list streams one chunk per fragment.
+    arguments: Union[str, list]
+    call_id: Optional[str] = None
+
+
+def _tool_calls_sse_lines(*calls: StreamedToolCall) -> list:
+    """SSE lines streaming *calls* at ``tool_calls`` indexes 0, 1, ...: each id and name, then its arguments."""
+    deltas = []
+    for index, call in enumerate(calls):
+        header = {"index": index, "id": call.call_id or f"call_{index}", "type": "function"}
+        deltas.append({"tool_calls": [{**header, "function": {"name": call.name, "arguments": ""}}]})
+        fragments = [call.arguments] if isinstance(call.arguments, str) else call.arguments
+        deltas += [{"tool_calls": [{"index": index, "function": {"arguments": fragment}}]} for fragment in fragments]
+    deltas[0] = {"role": "assistant", **deltas[0]}
+    lines = [
+        _sse({"id": "chatcmpl-1", "choices": [{"index": 0, "delta": delta, "finish_reason": None}]}) for delta in deltas
     ]
-    for fragment in arg_fragments:
-        chunks.append(
-            {
-                "id": "chatcmpl-1",
-                "choices": [
-                    {
-                        "index": 0,
-                        "delta": {"tool_calls": [{"index": 0, "function": {"arguments": fragment}}]},
-                        "finish_reason": None,
-                    }
-                ],
-            }
-        )
-    chunks.append({"id": "chatcmpl-1", "choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]})
-    lines = [_sse(chunk) for chunk in chunks]
+    lines.append(_sse({"id": "chatcmpl-1", "choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]}))
     lines.append(b"data: [DONE]\n\n")
     return lines
 
 
+def _tool_call_sse_lines(name: str, arg_fragments: list, call_id: str = "call_1") -> list:
+    """SSE lines streaming a single tool call: id/name first, then argument fragments, then finish."""
+    return _tool_calls_sse_lines(StreamedToolCall(name, arg_fragments, call_id))
+
+
 def _text_sse_lines(text: str) -> list:
     """SSE lines streaming *text* as a single assistant delta, then finish."""
     return [

tests/guardrails/tool_helpers.py

@@ -93,14 +93,23 @@ def violations_in(outcome) -> list[ToolViolation]:
     return [ToolViolation.model_validate(violation) for violation in outcome.metadata.get("tool_violations", [])]
 
 
-def make_tool_conversation(result_call_id: str = "call_1", result_name: Optional[str] = "get_weather") -> list:
-    """A user turn, an assistant ``get_weather`` tool call (id ``call_1``), then a tool result.
+def make_tool_conversation(
+    result_call_id: str = "call_1",
+    result_name: Optional[str] = "get_weather",
+    *,
+    tool_name: str = "get_weather",
+    arguments: str = '{"city": "Paris"}',
+    content: str = "18C",
+) -> list:
+    """A user turn, an assistant tool call (id ``call_1``), then a tool result.
 
     ``result_call_id`` sets the tool result's ``tool_call_id`` so callers can test
     linked (``call_1``) and unlinked (anything else) results. ``result_name`` sets the
     result's ``name``; ``None`` omits the key, as an OpenAI tool message does.
+    ``tool_name``, ``arguments`` and ``content`` set the called tool, its arguments and
+    the result's content, defaulting to a ``get_weather`` call answered with ``18C``.
     """
-    tool_message = {"role": "tool", "tool_call_id": result_call_id, "content": "18C"}
+    tool_message = {"role": "tool", "tool_call_id": result_call_id, "content": content}
     if result_name is not None:
         tool_message["name"] = result_name
     return [
@@ -112,7 +121,7 @@ def make_tool_conversation(result_call_id: str = "call_1", result_name: Optional
                 {
                     "id": "call_1",
                     "type": "function",
-                    "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
+                    "function": {"name": tool_name, "arguments": arguments},
                 }
             ],
         },