diff --git a/tests/guardrails/test_per_tool_regex_rails_iorails.py b/tests/guardrails/test_per_tool_regex_rails_iorails.py index b256f6bda2..34e3d6b9c3 100644 --- a/tests/guardrails/test_per_tool_regex_rails_iorails.py +++ b/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 diff --git a/tests/guardrails/test_tool_rails_iorails.py b/tests/guardrails/test_tool_rails_iorails.py index 1633156819..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,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_``.""" + + 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 [ 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}, } ], },