diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index c87357122ce..9227bb25126 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -16,6 +16,8 @@ from litellm.constants import VLLM_CHOICE_FIELDS from litellm.litellm_core_utils.core_helpers import ( PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS, UPSTREAM_USAGE_COST, + as_object_list, + as_str_mapping, chunk_usage, finite_cost, provider_reported_usage_cost, @@ -1800,6 +1802,53 @@ def _assembled_message_fields(processor: ChunkProcessor, chunks: Sequence) -> It yield "provider_specific_fields", combined_provider_fields +_RAW_FRAME_TYPES: Final = (str, bytes) +_FRAME_LINE_STARTS: Final = (":", "data:", "event:", "id:", "retry:") +_FRAME_PAYLOAD: Final = TypeAdapter(dict[str, object]) + + +def _is_chat_chunk(payload: Mapping[str, object]) -> bool: + choices: Final = as_object_list(payload.get("choices")) + return ( + "error" not in payload + and choices is not None + and all( + (choice := as_str_mapping(raw_choice)) is not None + and (delta := as_str_mapping(choice.get("delta"))) is not None + and isinstance(delta.get("content"), (str, type(None))) + for raw_choice in choices + ) + ) + + +def _frame_chunks(frame: str | bytes) -> tuple[Mapping[str, object], ...]: + text: Final = frame if isinstance(frame, str) else frame.decode("utf-8") + lines: Final = tuple(filter(None, (line.strip() for line in text.splitlines()))) + if not all(line.startswith(_FRAME_LINE_STARTS) for line in lines): + raise ValueError("a raw frame beside parsed chunks carries a line that is not an SSE comment or field") + payloads: Final = ( + payload + for line in lines + if line.startswith("data:") and (payload := line.removeprefix("data:").strip()) not in ("", "[DONE]") + ) + chunks: Final = tuple(map(_FRAME_PAYLOAD.validate_json, payloads)) + if not all(map(_is_chat_chunk, chunks)): + raise ValueError("a raw frame beside parsed chunks carries data that is not a chat completion chunk") + return chunks + + +def _with_raw_frames_unpacked(chunks: Sequence[object]) -> Sequence[object]: + frames: Final = sum(isinstance(chunk, _RAW_FRAME_TYPES) for chunk in chunks) + if frames in (0, len(chunks)): + return chunks + verbose_logger.warning("stream_chunk_builder: unpacking %d raw SSE frame(s) found among parsed chunks", frames) + return tuple( + chain.from_iterable( + _frame_chunks(chunk) if isinstance(chunk, _RAW_FRAME_TYPES) else (chunk,) for chunk in chunks + ) + ) + + def stream_chunk_builder( chunks: list, messages: Sequence | None = None, @@ -1832,7 +1881,7 @@ def _assembled_stream_response( if not chunks: return None - processor: Final = ChunkProcessor(list(chunks), messages) + processor: Final = ChunkProcessor(list(_with_raw_frames_unpacked(chunks)), messages) chunks = processor.chunks ### BASE-CASE ### diff --git a/litellm/llms/anthropic/messages/stream_assembly.py b/litellm/llms/anthropic/messages/stream_assembly.py index 843e8a6e495..1975e2e0b9b 100644 --- a/litellm/llms/anthropic/messages/stream_assembly.py +++ b/litellm/llms/anthropic/messages/stream_assembly.py @@ -436,7 +436,7 @@ def _is_anthropic_event_type(event_type: object) -> bool: def is_raw_sse_stream(all_chunks: Sequence[object]) -> bool: - return any(isinstance(chunk, (str, bytes)) for chunk in all_chunks) + return len(all_chunks) > 0 and all(isinstance(chunk, (str, bytes)) for chunk in all_chunks) def _joined_sse_bytes(all_chunks: Sequence[object]) -> bytes: diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py index 2920c6bec9a..8ea28d6c798 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py @@ -477,7 +477,7 @@ class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail): if not all_chunks: return - if not isinstance(all_chunks[0], (ModelResponse, ModelResponseStream)): + if not any(isinstance(chunk, (ModelResponse, ModelResponseStream)) for chunk in all_chunks): verbose_proxy_logger.warning( "Cisco AI Defense guardrail (%s): unsupported streaming chunk shape (%s) — failing closed.", self.guardrail_name, diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index b9871b1e276..69190884731 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1003,7 +1003,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): isinstance(event_type := getattr(chunk, "type", None), str) and event_type.startswith("response.") for chunk in all_chunks ): - return _StreamSurface.RESPONSES + has_raw_frame: Final = any(isinstance(chunk, (str, bytes)) for chunk in all_chunks) + return _StreamSurface.OPAQUE_SSE if has_raw_frame else _StreamSurface.RESPONSES return _StreamSurface.CHAT_COMPLETIONS @staticmethod diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 379c2c4b075..03243b49fb5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -1446,26 +1446,27 @@ class PanwPrismaAirsHandler(CustomGuardrail): verbose_proxy_logger.info("PANW Prisma AIRS: Running post-call streaming scan") - all_chunks: Final = [] content_was_modified = False try: start_time: Final = datetime.now() - # Collect all chunks - async for chunk in response: - all_chunks.append(chunk) + all_chunks: Final = [chunk async for chunk in response] + byte_frames: Final = tuple(chunk for chunk in all_chunks if isinstance(chunk, bytes)) + has_parsed_chunks: Final = any( + isinstance(chunk, (ModelResponse, ModelResponseStream)) for chunk in all_chunks + ) # Handle /v1/messages streaming: chunks are raw bytes (Anthropic SSE) - if all_chunks and isinstance(all_chunks[0], bytes): - text = self._extract_text_from_sse_bytes(all_chunks) + if byte_frames and len(byte_frames) == len(all_chunks): + text = self._extract_text_from_sse_bytes(byte_frames) await self._scan_raw_streaming_text(text, request_data, start_time) for chunk in all_chunks: yield chunk return # Handle /v1/responses streaming: chunks are Pydantic events (not ModelResponse/ModelResponseStream) - if all_chunks and not isinstance(all_chunks[0], (ModelResponse, ModelResponseStream)): + if all_chunks and not has_parsed_chunks and not byte_frames: text = self._extract_text_from_streaming_events(all_chunks) await self._scan_raw_streaming_text(text, request_data, start_time) for chunk in all_chunks: diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index 527c83d5760..be4452d2427 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1,5 +1,6 @@ import base64 import json +import logging from collections.abc import Callable, Mapping, Sequence from typing import Final from unittest.mock import MagicMock @@ -2880,3 +2881,165 @@ def test_stream_chunk_builder_counts_a_custom_named_deployments_missing_usage_wi model="my-gpt-6", messages=messages ) assert built.usage.prompt_tokens == litellm.token_counter(model=counted_as, messages=messages) + + +def _streamed_text(content: str) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-raw-frames", + created=1, + model="gpt-5.4-mini", + choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content=content))], + ) + + +_RAW_FRAMES_CARRYING_NO_CHUNK: Final = ( + pytest.param(b"data: [DONE]\n\n", id="done_bytes"), + pytest.param("data: [DONE]\n\n", id="done_str"), + pytest.param(b": keepalive\n\n", id="comment"), + pytest.param(b"event: ping\n\n", id="event_only"), +) + + +@pytest.mark.parametrize("position", [0, 1, 2]) +@pytest.mark.parametrize("frame", _RAW_FRAMES_CARRYING_NO_CHUNK) +def test_stream_chunk_builder_assembles_the_parsed_chunks_beside_a_raw_frame(frame: str | bytes, position: int) -> None: + parsed: Final = [_streamed_text("Hello"), _streamed_text(" world")] + + built: Final = stream_chunk_builder(chunks=[*parsed[:position], frame, *parsed[position:]]) + + assert built is not None + assert built.choices[0].message.content == "Hello world" + + +def test_stream_chunk_builder_reads_the_chunk_a_raw_frame_carries() -> None: + carried: Final = _openai_sse( + {"model": "gpt-5.4-mini", "choices": [{"index": 0, "delta": {"content": " there"}, "finish_reason": None}]} + ) + + built: Final = stream_chunk_builder(chunks=[_streamed_text("Hello"), carried.encode(), _streamed_text(" world")]) + + assert built is not None + assert built.choices[0].message.content == "Hello there world" + + +_SECRET_CHUNK: Final[dict[str, object]] = { + "id": "chatcmpl-raw-frames", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-5.4-mini", + "choices": [{"index": 0, "delta": {"content": "LEAKED-SECRET"}}], +} + + +def _data_line(payload: Mapping[str, object]) -> str: + return f"data: {json.dumps(payload)}\n\n" + + +@pytest.mark.parametrize("line_break", ["\r", "\r\n", "\x0c", "\u2028"]) +def test_stream_chunk_builder_reads_a_data_line_after_any_line_break_an_sse_reader_honours(line_break: str) -> None: + frame: Final = f"event: message{line_break}{_data_line(_SECRET_CHUNK)}".encode() + + built: Final = stream_chunk_builder(chunks=[_streamed_text("Hello"), frame, _streamed_text(" world")]) + + assert built is not None + assert built.choices[0].message.content == "HelloLEAKED-SECRET world" + + +@pytest.mark.parametrize( + "frame", + [ + pytest.param(b'data: {"error": {"message": "blocked", "code": 400}}\n\n', id="error_payload"), + pytest.param(b'data: {"type": "content_block_delta"}\n\n', id="payload_that_is_not_a_chat_chunk"), + pytest.param(b"data: not json\n\n", id="payload_that_is_not_json"), + pytest.param(_data_line({**_SECRET_CHUNK, "error": {"message": "blocked"}}), id="chunk_and_error"), + pytest.param( + _data_line({**_SECRET_CHUNK, "choices": [{"text": "LEAKED-SECRET"}]}), id="choice_without_a_delta" + ), + pytest.param( + _data_line({**_SECRET_CHUNK, "choices": [{"message": {"content": "LEAKED-SECRET"}}]}), + id="choice_with_a_message", + ), + *( + pytest.param( + _data_line({**_SECRET_CHUNK, "choices": [{"index": 0, "delta": {"content": content}}]}), + id=f"delta_content_{name}", + ) + for name, content in ( + ("list", [{"type": "text", "text": "LEAKED-SECRET"}]), + ("int", 7), + ("dict", {"text": "LEAKED-SECRET"}), + ) + ), + pytest.param(_data_line(_SECRET_CHUNK).replace("LEAKED", "LEAK\u2028ED"), id="line_break_inside_the_json"), + pytest.param(json.dumps(_SECRET_CHUNK), id="chunk_without_a_data_field"), + pytest.param("a plain text frame", id="plain_text_str"), + pytest.param(b"a plain text frame", id="plain_text_bytes"), + pytest.param("[DONE]", id="bare_done"), + pytest.param(b"data\n\n", id="bare_data_field"), + pytest.param(b"\xef\xbb\xbfdata: [DONE]\n\n", id="byte_order_mark"), + pytest.param(bytearray(b"data: [DONE]\n\n"), id="bytearray"), + ], +) +def test_stream_chunk_builder_still_fails_on_a_raw_frame_it_cannot_account_for(frame: bytes | str) -> None: + with pytest.raises(litellm.APIError): + stream_chunk_builder(chunks=[_streamed_text("Hello"), frame]) + + +def _built_content(chunks: list[object]) -> str | None: + try: + built: Final = stream_chunk_builder(chunks=chunks) + except litellm.APIError: + return None + return None if built is None else built.choices[0].message.content + + +@pytest.mark.parametrize("line_break", ["\n", "\r\n", "\r"]) +def test_stream_chunk_builder_never_skips_a_data_line_when_its_frame_is_split_in_two(line_break: str) -> None: + frame: Final = f"event: message{line_break}data: {json.dumps(_SECRET_CHUNK)}{line_break}{line_break}".encode() + + skipped: Final = [ + at + for at in range(1, len(frame)) + if ( + content := _built_content([_streamed_text("Hello"), frame[:at], frame[at:], _streamed_text("!")]) + ) + is not None + and "LEAKED-SECRET" not in content + ] + + assert skipped == [] + + +@pytest.mark.parametrize( + "frames", + [ + pytest.param([b"data: [DONE]\n\n"], id="done_only"), + pytest.param([_openai_sse({"choices": [{"index": 0, "delta": {"content": "Hi"}}]}).encode()], id="chunk_frame"), + ], +) +def test_stream_chunk_builder_still_fails_on_a_stream_of_raw_frames_only(frames: list[bytes]) -> None: + with pytest.raises(litellm.APIError): + stream_chunk_builder(chunks=frames) + + +@pytest.mark.parametrize("frames", [1, 3]) +def test_stream_chunk_builder_warns_with_the_number_of_raw_frames_it_unpacks( + frames: int, caplog: pytest.LogCaptureFixture +) -> None: + stream: Final = [_streamed_text("Hello"), *[b"data: [DONE]\n\n"] * frames, _streamed_text(" world")] + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + built: Final = stream_chunk_builder(chunks=stream) + + assert built is not None + assert [record.getMessage() for record in caplog.records if "raw SSE frame" in record.getMessage()] == [ + f"stream_chunk_builder: unpacking {frames} raw SSE frame(s) found among parsed chunks" + ] + + +def test_stream_chunk_builder_does_not_warn_when_no_raw_frame_is_present(caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + built: Final = stream_chunk_builder(chunks=[_streamed_text("Hello"), _streamed_text(" world")]) + + assert built is not None + assert not [record for record in caplog.records if "raw SSE frame" in record.getMessage()] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 1fb56509b7f..ee1583eff6a 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -8168,3 +8168,76 @@ async def test_guardrail_secret_lookups_leave_the_event_loop_free( assert vault.read.is_set() assert not vault.read_on_loop.is_set() assert upstream.calls.last.request.url.host == "bedrock-runtime.us-west-2.amazonaws.com" + + +def _streamed_text(content: str) -> litellm.ModelResponseStream: + return litellm.ModelResponseStream( + id="tid", + choices=[ + litellm.types.utils.StreamingChoices( + delta=litellm.types.utils.Delta(content=content, role="assistant"), index=0 + ) + ], + created=1, + model="gpt-4o-mini", + object="chat.completion.chunk", + ) + + +def _post_call_guardrail() -> BedrockGuardrail: + return BedrockGuardrail( + guardrail_name="bedrock-raw-frame", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + + +async def _streamed(items: list[object]): + for item in items: + yield item + + +@pytest.mark.asyncio +@pytest.mark.parametrize("position", [0, 1, 2]) +async def test_streaming_post_call_scans_the_parsed_chunks_beside_a_raw_frame(position: int): + guardrail: Final = _post_call_guardrail() + parsed: Final = [_streamed_text("Hello"), _streamed_text(" world")] + stream: Final = [*parsed[:position], b"data: [DONE]\n\n", *parsed[position:]] + + clean_scan: Final = {"action": "NONE", "assessments": [], "outputs": []} + + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=clean_scan)) as mock_make: + delivered = [ + chunk + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), response=_streamed(stream), request_data={"messages": []} + ) + ] + + scanned: Final = [ + call.kwargs["response"].choices[0].message.content + for call in mock_make.call_args_list + if call.kwargs.get("source") == "OUTPUT" + ] + assert scanned == ["Hello world"] + assert "".join(chunk.choices[0].delta.content or "" for chunk in delivered) == "Hello world" + + +@pytest.mark.asyncio +async def test_streaming_post_call_still_blocks_a_violating_response_that_carries_a_raw_frame(): + guardrail: Final = _post_call_guardrail() + violation: Final = HTTPException(status_code=400, detail={"error": "Violated guardrail policy"}) + stream: Final = [_streamed_text("Hello"), b"data: [DONE]\n\n", _streamed_text(" world")] + + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(side_effect=violation)): + with pytest.raises(HTTPException) as blocked: + _ = [ + chunk + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), response=_streamed(stream), request_data={"messages": []} + ) + ] + + assert blocked.value is violation diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py index 4cdabe5c360..ed2701a03d8 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_cisco_ai_defense_chat.py @@ -3067,3 +3067,63 @@ def test_config_model_exposed(): CiscoAIDefenseGuardrail.get_config_model() is CiscoAIDefenseGuardrailConfigModel ) assert CiscoAIDefenseGuardrailConfigModel.ui_friendly_name() == "Cisco AI Defense" + + +class TestCiscoAIDefenseRawFramesInAStream: + @pytest.mark.asyncio + @pytest.mark.parametrize("position", [0, 1, 2]) + async def test_the_parsed_chunks_are_inspected_and_delivered_beside_a_raw_frame(self, position): + g = _make_guardrail(event_hook=["pre_call", "post_call"]) + parsed = _make_streaming_chunks(["Hello", " world"]) + stream = [*parsed[:position], b"data: [DONE]\n\n", *parsed[position:]] + + received, post_mock = await _streaming_setup(g, stream, cisco_response=_safe_response()) + + inspected = [message["content"] for message in post_mock.call_args.kwargs["json"]["messages"]] + assert inspected[-1] == "Hello world" + assert received == stream + + @pytest.mark.asyncio + async def test_a_violation_withholds_the_chunks_when_a_raw_frame_sits_beside_them(self): + g = _make_guardrail(event_hook=["pre_call", "post_call"]) + parsed = _make_streaming_chunks(["Here is your SSN: ", "123-45-6789."]) + + received, post_mock = await _streaming_setup( + g, [parsed[0], b"data: [DONE]\n\n", parsed[1]], cisco_response=_violation_response() + ) + + assert post_mock.called + assert all(isinstance(chunk, str) and chunk.startswith("data: ") for chunk in received) + assert any('"error"' in chunk and "Cisco AI Defense" in chunk for chunk in received) + + @pytest.mark.asyncio + async def test_a_data_line_after_a_carriage_return_is_inspected_and_delivered(self): + g = _make_guardrail(event_hook=["pre_call", "post_call"]) + parsed = _make_streaming_chunks(["Hello", " world"]) + carried = '{"created": 1, "model": "m", "choices": [{"index": 0, "delta": {"content": " LEAKED-SECRET"}}]}' + stream = [parsed[0], f'event: message\rdata: {carried}\n\n'.encode(), parsed[1]] + + received, post_mock = await _streaming_setup(g, stream, cisco_response=_safe_response()) + + inspected = [message["content"] for message in post_mock.call_args.kwargs["json"]["messages"]] + delivered = "".join(chunk.choices[0].delta.content or "" for chunk in received) + assert delivered == "Hello LEAKED-SECRET world" + assert delivered in inspected + + @pytest.mark.asyncio + @pytest.mark.parametrize("frame", [b"LEAKED-SECRET", b'data: {"choices": [{"text": "LEAKED-SECRET"}]}\n\n']) + async def test_a_frame_that_cannot_be_read_is_not_inspected_or_delivered(self, frame): + g = _make_guardrail(event_hook=["pre_call", "post_call"]) + post_mock = AsyncMock(return_value=_safe_response()) + delivered: list = [] + + with _patch_inspection_post(g, post_mock), pytest.raises(litellm.APIError): + async for chunk in g.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), + response=_aiter([*_make_streaming_chunks(["Hello"]), frame]), + request_data={"messages": [{"role": "user", "content": "hi"}]}, + ): + delivered.append(chunk) + + assert delivered == [] + post_mock.assert_not_called() diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py index 34221720205..bb2d3a431ab 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_microsoft_purview.py @@ -2803,3 +2803,54 @@ class TestRegistration: guardrail_class_registry["microsoft_purview"] is MicrosoftPurviewDLPGuardrail ) + + +class TestRawFramesInAStreamedResponse: + @staticmethod + def _streamed_text(content: str): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + return ModelResponseStream( + id="chatcmpl-raw-frame", + created=1, + model="gpt-4o", + choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content=content))], + ) + + async def _scan(self, items: list[object], check: AsyncMock) -> list[object]: + guardrail: Final = _make_guardrail() + + async def upstream(): + for item in items: + yield item + + with patch.object(guardrail, "_check_content", check): + return [ + chunk + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"), + response=upstream(), + request_data={"metadata": {"user_id": "user-123"}}, + ) + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("position", [0, 1, 2]) + async def test_the_parsed_chunks_are_scanned_and_delivered_beside_a_raw_frame(self, position: int): + parsed: Final = [self._streamed_text("Hello"), self._streamed_text(" world")] + check: Final = AsyncMock(return_value={"policyActions": []}) + + delivered: Final = await self._scan([*parsed[:position], b"data: [DONE]\n\n", *parsed[position:]], check) + + assert [call.kwargs["text"] for call in check.call_args_list] == ["Hello world"] + assert "".join(chunk.choices[0].delta.content or "" for chunk in delivered) == "Hello world" + + @pytest.mark.asyncio + async def test_a_violation_is_raised_when_a_raw_frame_sits_beside_the_chunks(self): + parsed: Final = [self._streamed_text("SSN: 123-45"), self._streamed_text("-6789")] + blocked: Final = HTTPException(status_code=400, detail={"error": "Content blocked by policy"}) + + with pytest.raises(HTTPException) as raised: + await self._scan([parsed[0], b"data: [DONE]\n\n", parsed[1]], AsyncMock(side_effect=blocked)) + + assert raised.value is blocked diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py index 5aa3c7dfc17..0650835cb26 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -5318,3 +5318,98 @@ async def test_apply_guardrail_raises_on_flagged_when_not_logging_only(): entries = request_data["metadata"]["standard_logging_guardrail_information"] flagged = [e for e in entries if e["guardrail_status"] == "guardrail_flagged"] assert len(flagged) == 1 + + +def _streamed_text(content: str) -> litellm.ModelResponseStream: + return litellm.ModelResponseStream( + choices=[litellm.types.utils.StreamingChoices(delta=litellm.types.utils.Delta(content=content))] + ) + + +def _clean_scan_response() -> dict[str, object]: + return {"sanitizationResult": {"filterMatchState": "NO_MATCH_FOUND", "filterResults": {}}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("position", [0, 1, 2]) +async def test_streaming_hook_scans_the_parsed_chunks_beside_a_raw_frame(position: int): + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + ) + guardrail.make_model_armor_request = AsyncMock(return_value=_clean_scan_response()) + guardrail.should_run_guardrail = Mock(return_value=True) + parsed = [_streamed_text("Hello"), _streamed_text(" world")] + stream = [*parsed[:position], b"data: [DONE]\n\n", *parsed[position:]] + + delivered = await _drain_surface_hook(guardrail, stream) + + scanned = [call.kwargs["content"] for call in guardrail.make_model_armor_request.call_args_list] + assert scanned == ["Hello world"] + assert delivered == stream + + +@pytest.mark.asyncio +async def test_streaming_hook_blocks_a_flagged_response_that_carries_a_raw_frame(): + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + ) + guardrail.make_model_armor_request = AsyncMock( + return_value={ + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": { + "sdp": { + "sdpFilterResult": { + "inspectResult": { + "matchState": "MATCH_FOUND", + "findings": [{"infoType": "PASSWORD", "likelihood": "VERY_LIKELY"}], + } + } + } + }, + } + } + ) + guardrail.should_run_guardrail = Mock(return_value=True) + parsed = [_streamed_text("My password is "), _streamed_text("hunter2")] + + delivered = await _drain_surface_hook(guardrail, [parsed[0], b"data: [DONE]\n\n", parsed[1]]) + + assert guardrail.make_model_armor_request.call_count == 1 + assert len(delivered) == 1 + assert isinstance(delivered[0], str) + assert int(json.loads(delivered[0].removeprefix("data: "))["error"]["code"]) == 400 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "frame", + [ + pytest.param(b"data: [DONE]\n\n", id="done_bytes"), + pytest.param(b'data: {"type":"response.output_text.delta","delta":"LEAKED-SECRET"}\n\n', id="delta_bytes"), + pytest.param('data: {"type":"response.output_text.delta","delta":"LEAKED-SECRET"}\n\n', id="delta_str"), + ], +) +@pytest.mark.parametrize("frame_first", [False, True], ids=["frame_last", "frame_first"]) +async def test_streaming_hook_fails_closed_on_a_raw_frame_beside_responses_events( + frame: bytes | str, frame_first: bool +): + guardrail = _surface_guardrail() + post = _armor_post_mock(_MODEL_ARMOR_CLEAN) + events = _responses_api_events() + stream = (frame, *events) if frame_first else (*events, frame) + + with patch.object(guardrail.async_handler, "post", post): + delivered = await _drain_surface_hook(guardrail, stream) + + post.assert_not_called() + assert len(delivered) == 1 + assert isinstance(delivered[0], str) + assert "could not be assembled for scanning" in delivered[0] + assert "LEAKED-SECRET" not in delivered[0] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_noma.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_noma.py index 574550e1bdd..2f2a6164bf9 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_noma.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_noma.py @@ -2494,3 +2494,70 @@ class TestNomaAnonymizationFlow: user_api_key_dict=mock_user_api_key_dict, response=llm_response, ) + + +class TestNomaRawFramesInAStream: + @staticmethod + def _guardrail(): + return NomaGuardrail( + api_key="test-api-key", + api_base="https://api.test.noma.security/", + monitor_mode=False, + block_failures=True, + guardrail_name="noma-post-call", + event_hook="post_call", + default_on=True, + ) + + @staticmethod + def _streamed_text(content): + return ModelResponseStream( + id="resp_1", + choices=[StreamingChoices(delta=Delta(content=content), index=0)], + created=1234567890, + model="gpt-4.1-mini", + object="chat.completion.chunk", + ) + + async def _scan(self, items, verdict_for_text): + guardrail = self._guardrail() + + async def upstream(): + for item in items: + yield item + + with patch.object(guardrail.async_handler, "post", side_effect=_noma_reply_per_text(verdict_for_text)): + return [ + chunk + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=upstream(), + request_data={"litellm_call_id": "call-1", "metadata": {}}, + ) + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("position", [0, 1, 2]) + async def test_the_parsed_chunks_are_scanned_and_delivered_beside_a_raw_frame(self, position): + scanned = [] + + def verdict_for_text(text): + scanned.append(text) + return {"aggregatedScanResult": False, "scanResult": []} + + parsed = [self._streamed_text("Hello"), self._streamed_text(" world")] + + delivered = await self._scan([*parsed[:position], b"data: [DONE]\n\n", *parsed[position:]], verdict_for_text) + + assert scanned == ["Hello world"] + assert "".join(chunk.choices[0].delta.content or "" for chunk in delivered) == "Hello world" + + @pytest.mark.asyncio + async def test_a_flagged_response_is_blocked_when_a_raw_frame_sits_beside_it(self): + parsed = [self._streamed_text("The launch code"), self._streamed_text(" is 0000.")] + + with pytest.raises(NomaBlockedMessage): + await self._scan( + [parsed[0], b"data: [DONE]\n\n", parsed[1]], + lambda text: {"aggregatedScanResult": "launch code" in text, "scanResult": []}, + ) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py index 35cea840df3..94c9018fefc 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py @@ -15,6 +15,7 @@ import logging from collections.abc import Mapping, Sequence from contextlib import AbstractContextManager from datetime import datetime +from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -6598,5 +6599,101 @@ class TestPanwAirsServerToolCalls: assert message.tool_calls is None +class TestRawFramesInAStreamedResponse: + @staticmethod + def _streamed_text(content: str) -> ModelResponseStream: + return ModelResponseStream( + id="test_id", + created=1234567890, + model="gpt-4", + choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content=content))], + ) + + async def _scan(self, stream: list[object], scan_result: dict[str, object]) -> tuple[list[object], AsyncMock]: + handler = make_handler() + + async def upstream(): + for item in stream: + yield item + + with patch.object(handler, "_call_panw_api", new_callable=AsyncMock, return_value=scan_result) as scan: + delivered = [ + chunk + async for chunk in handler.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + response=upstream(), + request_data={"model": "gpt-4", "messages": [], "litellm_call_id": "call-1"}, + ) + ] + return delivered, scan + + @pytest.mark.asyncio + @pytest.mark.parametrize("position", [0, 1, 2]) + async def test_the_parsed_chunks_are_scanned_and_delivered_beside_a_raw_frame(self, position: int): + parsed = [self._streamed_text("Hello"), self._streamed_text(" world")] + + delivered, scan = await self._scan( + [*parsed[:position], b"data: [DONE]\n\n", *parsed[position:]], {"action": "allow", "category": "safe"} + ) + + assert [call.kwargs["content"] for call in scan.call_args_list] == ["Hello world"] + assert delivered == [*parsed[:position], b"data: [DONE]\n\n", *parsed[position:]] + + @pytest.mark.asyncio + async def test_a_blocked_response_is_not_delivered_when_a_raw_frame_sits_beside_it(self): + parsed = [self._streamed_text("Hello"), self._streamed_text(" world")] + + delivered, scan = await self._scan( + [parsed[0], b"data: [DONE]\n\n", parsed[1]], {"action": "block", "category": "malicious"} + ) + + assert scan.call_count == 1 + assert all(not isinstance(chunk, (ModelResponseStream, bytes)) for chunk in delivered) + assert any("error" in str(chunk) for chunk in delivered) + + @pytest.mark.asyncio + async def test_a_stream_of_raw_frames_stays_blocked_when_a_str_frame_sits_beside_bytes(self): + frames = [ + b'event: content_block_delta\ndata: {"type": "content_block_delta"}\n\n', + "data: [DONE]\n\n", + ] + + delivered, scan = await self._scan(frames, {"action": "allow", "category": "safe"}) + + assert scan.call_count == 0 + assert all(isinstance(chunk, str) and '"error"' in chunk for chunk in delivered) + + @pytest.mark.asyncio + async def test_a_data_line_split_across_two_frames_is_blocked_not_delivered_unscanned(self): + carried = json.dumps( + { + "id": "c1", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4", + "choices": [{"index": 0, "delta": {"content": "LEAKED-SECRET"}}], + } + ) + stream = [self._streamed_text("Hello"), b"data", f": {carried}\n\n".encode(), self._streamed_text(" world")] + + delivered, scan = await self._scan(stream, {"action": "allow", "category": "safe"}) + + assert scan.call_count == 0 + assert all(isinstance(chunk, str) and '"error"' in chunk for chunk in delivered) + + @pytest.mark.asyncio + @pytest.mark.parametrize("frame_first", [False, True]) + async def test_a_bytes_frame_beside_responses_events_is_blocked_not_scanned_in_part(self, frame_first: bool): + events = [SimpleNamespace(type="response.output_text.delta", delta="test content")] + frame = b"data: [DONE]\n\n" + + delivered, scan = await self._scan( + [frame, *events] if frame_first else [*events, frame], {"action": "allow", "category": "safe"} + ) + + assert scan.call_count == 0 + assert all(isinstance(chunk, str) and '"error"' in chunk for chunk in delivered) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_repelloai.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_repelloai.py index ff61af9ee62..42e964ec315 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_repelloai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_repelloai.py @@ -1087,6 +1087,15 @@ class TestRepelloAIStreaming: return _gen() + @staticmethod + def _stream_items(*contents): + from litellm.types.utils import Delta, StreamingChoices + + return tuple( + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=content))]) + for content in contents + ) + @pytest.mark.asyncio async def test_streaming_passed_reemits_chunks(self, monkeypatch): guardrail = _guardrail(event_hook="post_call") @@ -1106,6 +1115,58 @@ class TestRepelloAIStreaming: ] assert len(out) == 2 + @pytest.mark.asyncio + @pytest.mark.parametrize("position", [0, 1, 2]) + async def test_streaming_scans_the_parsed_chunks_beside_a_raw_frame(self, monkeypatch, position): + guardrail = _guardrail(event_hook="post_call") + captured = {} + + async def capture(url, headers, json, **_kwargs): + captured["json"] = json + return _verdict_response("passed", url) + + monkeypatch.setattr(guardrail.async_handler, "post", capture) + stream = list(self._stream_items("unsafe ", "answer")) + items = [*stream[:position], b"data: [DONE]\n\n", *stream[position:]] + + async def upstream(): + for item in items: + yield item + + out = [ + chunk + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), + response=upstream(), + request_data={"messages": [{"role": "user", "content": "q"}]}, + ) + ] + + assert captured["json"]["scan_data"]["response"] == "unsafe answer" + assert out == items + + @pytest.mark.asyncio + async def test_streaming_blocks_a_flagged_response_that_carries_a_raw_frame(self, monkeypatch): + from litellm.proxy.proxy_server import StreamingCallbackError + + guardrail = _guardrail(event_hook="post_call") + monkeypatch.setattr( + guardrail.async_handler, "post", _async_return(_verdict_response("blocked", ANALYZE_RESPONSE_URL)) + ) + stream = list(self._stream_items("unsafe ", "answer")) + + async def upstream(): + for item in (stream[0], b"data: [DONE]\n\n", stream[1]): + yield item + + with pytest.raises(StreamingCallbackError): + async for _ in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), + response=upstream(), + request_data={"messages": [{"role": "user", "content": "q"}]}, + ): + pass + @pytest.mark.asyncio async def test_streaming_blocked_raises(self, monkeypatch): from litellm.proxy.proxy_server import StreamingCallbackError diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py index a46f1ddeed9..a46ced40127 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -4162,3 +4162,70 @@ class TestToolPermissionGuardrailAnthropicMessages: with patch.object(self.blocking, "should_run_guardrail", return_value=True): with pytest.raises(GuardrailRaisedException): await self._drain(self.blocking, [b"data: not-json\n\n", b"event: weird\n\n"]) + + +class TestToolPermissionGuardrailRawFramesInAChatStream: + def setup_method(self): + self.blocking = ToolPermissionGuardrail( + guardrail_name="chat-block", + rules=[ + {"id": "allow_bash", "tool_name": r"^Bash$", "decision": "allow"}, + {"id": "deny_read", "tool_name": r"^Read$", "decision": "deny"}, + ], + default_action="deny", + on_disallowed_action="block", + ) + + @staticmethod + def _tool_call_stream(tool_name: str) -> list[ModelResponseStream]: + def chunk(delta: Delta, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-raw-frame", + created=1, + model="gpt-5.4-mini", + choices=[StreamingChoices(index=0, delta=delta, finish_reason=finish_reason)], + ) + + call: Final = { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": tool_name, "arguments": '{"file_path": "/etc/passwd"}'}, + } + return [chunk(Delta(role="assistant", tool_calls=[call])), chunk(Delta(), finish_reason="tool_calls")] + + async def _drain(self, items: list[object]) -> list[object]: + async def _stream(): + for item in items: + yield item + + with patch.object(self.blocking, "should_run_guardrail", return_value=True): + return [ + chunk + async for chunk in self.blocking.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), response=_stream(), request_data={} + ) + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("position", [0, 1, 2]) + async def test_denied_tool_call_is_blocked_whichever_way_a_raw_frame_sits_in_the_stream(self, position: int): + parsed: Final = self._tool_call_stream("Read") + + with pytest.raises(GuardrailRaisedException, match="denied by rule 'deny_read'"): + await self._drain([*parsed[:position], b"data: [DONE]\n\n", *parsed[position:]]) + + @pytest.mark.asyncio + async def test_allowed_tool_call_is_delivered_when_a_raw_frame_sits_in_the_stream(self): + parsed: Final = self._tool_call_stream("Bash") + + delivered: Final = await self._drain([parsed[0], b"data: [DONE]\n\n", parsed[1]]) + + names: Final = [ + call.function.name + for chunk in delivered + if isinstance(chunk, ModelResponseStream) + for call in (chunk.choices[0].delta.tool_calls or ()) + if call.function.name + ] + assert names == ["Bash"] diff --git a/tests/unit/proxy/guardrails/test_anthropic_sse.py b/tests/unit/proxy/guardrails/test_anthropic_sse.py index 86edfdab218..e13df30b906 100644 --- a/tests/unit/proxy/guardrails/test_anthropic_sse.py +++ b/tests/unit/proxy/guardrails/test_anthropic_sse.py @@ -8,8 +8,9 @@ import pytest from litellm.proxy.guardrails.anthropic_sse import ( anthropic_sse_chunks_from_response, assemble_anthropic_sse_stream, + is_raw_sse_stream, ) -from litellm.types.utils import Choices, ModelResponse +from litellm.types.utils import Choices, Delta, ModelResponse, ModelResponseStream, StreamingChoices def _frame(event: str, data: dict[str, object]) -> bytes: @@ -429,3 +430,19 @@ def test_healthy_stream_with_a_tool_input_that_is_not_json_still_raises(): with pytest.raises(ValueError, match="Failed to parse tool call arguments"): anthropic_sse_chunks_from_response(assembled, upstream) + + +_PARSED_CHUNK: Final = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="Hi"))]) + + +@pytest.mark.parametrize( + ("chunks", "is_raw"), + [ + pytest.param([b"event: message_start\ndata: {}\n\n", "data: [DONE]\n\n"], True, id="frames_only"), + pytest.param([_PARSED_CHUNK, _PARSED_CHUNK], False, id="parsed_chunks_only"), + pytest.param([_PARSED_CHUNK, b"data: [DONE]\n\n", _PARSED_CHUNK], False, id="parsed_chunks_and_a_frame"), + pytest.param([], False, id="empty"), + ], +) +def test_a_stream_is_raw_sse_only_when_every_chunk_is_a_frame(chunks: list[object], is_raw: bool): + assert is_raw_sse_stream(chunks) is is_raw