From cb1b6d5684cc6f426133db142cc88ead6f480a06 Mon Sep 17 00:00:00 2001 From: Tanushriya Singh Date: Thu, 8 Oct 2026 12:10:50 -0700 Subject: [PATCH 1/2] test: add streaming coverage for per-tool regex rails with multiple calls and tool results --- .../test_per_tool_regex_rails_iorails.py | 74 +++++++++++++++---- tests/guardrails/test_tool_rails_iorails.py | 26 +++++++ 2 files changed, 86 insertions(+), 14 deletions(-) diff --git a/tests/guardrails/test_per_tool_regex_rails_iorails.py b/tests/guardrails/test_per_tool_regex_rails_iorails.py index b256f6bda2..18ef4963a3 100644 --- a/tests/guardrails/test_per_tool_regex_rails_iorails.py +++ b/tests/guardrails/test_per_tool_regex_rails_iorails.py @@ -21,6 +21,8 @@ duplicating them. """ +import json + import pytest import pytest_asyncio @@ -32,8 +34,10 @@ _inject_sse_stream, _stream_violation_chunks, _text_payload, + _text_sse_lines, _tool_call_payload, _tool_call_sse_lines, + _tool_calls_sse_lines, ) BASE_CONFIG = {"models": [{"type": "main", "engine": "nim", "model": "meta/llama-3.3-70b-instruct"}]} @@ -166,30 +170,72 @@ 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 + async def test_matching_second_call_blocks_stream(self, call_pattern_iorails): + # Only the second call matches, so this fails if a later streamed call skips the rails. + lines = _tool_calls_sse_lines( + ("run_sql", '{"query": "SELECT 1"}'), ("run_sql", '{"query": "DROP TABLE users"}') + ) + _inject_sse_stream(call_pattern_iorails, lines) + 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) -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}, + @pytest.mark.asyncio + async def test_several_non_matching_calls_stream_through(self, call_pattern_iorails): + lines = _tool_calls_sse_lines(("run_sql", '{"query": "SELECT 1"}'), ("run_sql", '{"query": "SELECT 2"}')) + _inject_sse_stream(call_pattern_iorails, lines) + 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] == [ + '{"query": "SELECT 1"}', + '{"query": "SELECT 2"}', ] + +def _tool_conversation(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}, + ] + + +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 diff --git a/tests/guardrails/test_tool_rails_iorails.py b/tests/guardrails/test_tool_rails_iorails.py index 1633156819..22603dfd52 100644 --- a/tests/guardrails/test_tool_rails_iorails.py +++ b/tests/guardrails/test_tool_rails_iorails.py @@ -226,6 +226,32 @@ def _tool_call_sse_lines(name: str, arg_fragments: list, call_id: str = "call_1" return lines +def _tool_calls_sse_lines(*calls: tuple[str, str]) -> list: + """SSE lines streaming one tool call per ``(name, arguments)`` pair, at ``tool_calls`` indexes 0, 1, ...""" + deltas = [] + for index, (name, arguments) in enumerate(calls): + deltas.append( + { + "tool_calls": [ + { + "index": index, + "id": f"call_{index}", + "type": "function", + "function": {"name": name, "arguments": ""}, + } + ] + } + ) + deltas.append({"tool_calls": [{"index": index, "function": {"arguments": arguments}}]}) + deltas[0] = {"role": "assistant", **deltas[0]} + lines = [ + _sse({"id": "chatcmpl-1", "choices": [{"index": 0, "delta": delta, "finish_reason": None}]}) for delta in deltas + ] + lines.append(_sse({"id": "chatcmpl-1", "choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]})) + lines.append(b"data: [DONE]\n\n") + return lines + + def _text_sse_lines(text: str) -> list: """SSE lines streaming *text* as a single assistant delta, then finish.""" return [ From 421e40fd77cdddf3657231e0bc9896dbfa40da95 Mon Sep 17 00:00:00 2001 From: Tanushriya Singh Date: Sat, 10 Oct 2026 13:39:49 -0700 Subject: [PATCH 2/2] test: parametrize multi-call streaming test and consolidate tool SSE helpers --- .../test_per_tool_regex_rails_iorails.py | 37 ++++------ tests/guardrails/test_tool_rails_iorails.py | 74 ++++++------------- tests/guardrails/tool_helpers.py | 17 ++++- 3 files changed, 51 insertions(+), 77 deletions(-) diff --git a/tests/guardrails/test_per_tool_regex_rails_iorails.py b/tests/guardrails/test_per_tool_regex_rails_iorails.py index 18ef4963a3..34e3d6b9c3 100644 --- a/tests/guardrails/test_per_tool_regex_rails_iorails.py +++ b/tests/guardrails/test_per_tool_regex_rails_iorails.py @@ -29,6 +29,7 @@ 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, @@ -39,6 +40,7 @@ _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"}]} @@ -171,12 +173,13 @@ async def test_non_matching_pattern_streams_through(self, call_pattern_iorails): assert _stream_violation_chunks(chunks) == [] @pytest.mark.asyncio - async def test_matching_second_call_blocks_stream(self, call_pattern_iorails): - # Only the second call matches, so this fails if a later streamed call skips the rails. - lines = _tool_calls_sse_lines( - ("run_sql", '{"query": "SELECT 1"}'), ("run_sql", '{"query": "DROP TABLE users"}') - ) - _inject_sse_stream(call_pattern_iorails, lines) + @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 @@ -185,27 +188,19 @@ async def test_matching_second_call_blocks_stream(self, call_pattern_iorails): @pytest.mark.asyncio async def test_several_non_matching_calls_stream_through(self, call_pattern_iorails): - lines = _tool_calls_sse_lines(("run_sql", '{"query": "SELECT 1"}'), ("run_sql", '{"query": "SELECT 2"}')) - _inject_sse_stream(call_pattern_iorails, lines) + 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] == [ - '{"query": "SELECT 1"}', - '{"query": "SELECT 2"}', - ] + assert [call["function"]["arguments"] for call in streamed] == arguments def _tool_conversation(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}, - ] + """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: diff --git a/tests/guardrails/test_tool_rails_iorails.py b/tests/guardrails/test_tool_rails_iorails.py index 22603dfd52..7d0ad43671 100644 --- a/tests/guardrails/test_tool_rails_iorails.py +++ b/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,61 +190,24 @@ 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, - } - ], - } - ] - 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(b"data: [DONE]\n\n") - return lines +@dataclass(frozen=True) +class StreamedToolCall: + """One tool call for a streamed response; ``call_id`` defaults to ``call_``.""" + + 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: tuple[str, str]) -> list: - """SSE lines streaming one tool call per ``(name, arguments)`` pair, at ``tool_calls`` indexes 0, 1, ...""" +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, (name, arguments) in enumerate(calls): - deltas.append( - { - "tool_calls": [ - { - "index": index, - "id": f"call_{index}", - "type": "function", - "function": {"name": name, "arguments": ""}, - } - ] - } - ) - deltas.append({"tool_calls": [{"index": index, "function": {"arguments": arguments}}]}) + 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 @@ -252,6 +217,11 @@ def _tool_calls_sse_lines(*calls: tuple[str, str]) -> list: 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 [ diff --git a/tests/guardrails/tool_helpers.py b/tests/guardrails/tool_helpers.py index efb2671aad..9044580b52 100644 --- a/tests/guardrails/tool_helpers.py +++ b/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}, } ], },