diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 4506c114208..2814ed4eafd 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -48,7 +48,10 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: - from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, + ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, @@ -70,6 +73,146 @@ class AnthropicMessagesHandler(BaseTranslation): super().__init__() self.adapter = LiteLLMAnthropicMessagesAdapter() + def build_block_sse_chunks( + self, + exc: "ModifyResponseException", + stream_started: bool = False, + responses_so_far: Optional[list[Any]] = None, + ) -> list[bytes]: + """ + Build an Anthropic SSE sequence delivering the guardrail block message + and terminating the stream cleanly. + + - ``stream_started`` False (buffered / pre-stream): nothing has been + sent, so emit a complete standalone message (message_start -> + content_block_* -> message_delta -> message_stop) via + FakeAnthropicMessagesStreamIterator, the same converter the + /v1/messages pre-stream block handler uses. + - ``stream_started`` True (sampling / detect-only end-of-stream): real + chunks were already sent, so *continue* the in-progress message -- + close the open content block, append the block message as a new text + block, then end the message. Emitting a second ``message_start`` here + would make Anthropic clients reject the stream. + """ + if stream_started: + return self._block_continuation_chunks(exc, responses_so_far or []) + return self._standalone_block_chunks(exc) + + def _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]: + import uuid + + from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + from litellm.types.utils import AnthropicMessagesResponse + + block_response = AnthropicMessagesResponse( + id=f"msg_{uuid.uuid4()}", + type="message", + role="assistant", + content=[{"type": "text", "text": exc.message}], + model=exc.model, + stop_reason="end_turn", + usage={"input_tokens": 0, "output_tokens": 0}, + ) + return list(FakeAnthropicMessagesStreamIterator(response=block_response)) + + def _block_continuation_chunks(self, exc: "ModifyResponseException", responses_so_far: list[Any]) -> list[bytes]: + """Continue an already-started message: close the open content block, + append the block message as a new text block, then end the message -- + without a second message_start.""" + + def _sse(event_type: str, payload: dict) -> bytes: + return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode() + + open_index, max_index = self._content_block_state(responses_so_far) + new_index = (max_index + 1) if max_index is not None else 0 + chunks: list[bytes] = [] + if open_index is not None: + chunks.append(_sse("content_block_stop", {"type": "content_block_stop", "index": open_index})) + chunks += [ + _sse( + "content_block_start", + { + "type": "content_block_start", + "index": new_index, + "content_block": {"type": "text", "text": ""}, + }, + ), + _sse( + "content_block_delta", + { + "type": "content_block_delta", + "index": new_index, + "delta": {"type": "text_delta", "text": exc.message}, + }, + ), + _sse("content_block_stop", {"type": "content_block_stop", "index": new_index}), + _sse( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 0}, + }, + ), + _sse("message_stop", {"type": "message_stop"}), + ] + return chunks + + @staticmethod + def _content_block_state( + responses_so_far: list[Any], + ) -> tuple[Optional[int], Optional[int]]: + """From the SSE chunks already sent to the client, return (open + content-block index or None, highest content-block index seen or None). + + A single streamed item may bundle multiple SSE events (raw bytes) or be + an already-parsed event dict, so every event across every item is + considered -- matching how ``get_streaming_string_so_far`` reads the + same stream.""" + open_indices: set[int] = set() + max_index: Optional[int] = None + for item in responses_so_far: + for data in AnthropicMessagesHandler._iter_sse_events(item): + event_type = data.get("type") + index = data.get("index") + if not isinstance(index, int): + continue + if event_type == "content_block_start": + open_indices.add(index) + max_index = index if max_index is None else max(max_index, index) + elif event_type == "content_block_stop": + open_indices.discard(index) + open_index = max(open_indices) if open_indices else None + return open_index, max_index + + @staticmethod + def _iter_sse_events(item: Any) -> list[dict]: + """Yield the event-data dicts in one stream chunk. + + Handles both formats this stream can carry (see + ``get_streaming_string_so_far``): raw SSE ``bytes`` -- which may bundle + several events separated by a blank line -- and an already-parsed event + ``dict``.""" + if isinstance(item, dict): + return [item] + if not isinstance(item, (bytes, bytearray)): + return [] + events: list[dict] = [] + for block in item.decode("utf-8", errors="replace").split("\n\n"): + for line in block.split("\n"): + line = line.strip() + if not line.startswith("data:"): + continue + try: + parsed = json.loads(line[len("data:") :].strip()) + except json.JSONDecodeError: + continue + if isinstance(parsed, dict): + events.append(parsed) + return events + def _translate_to_openai(self, data: dict) -> ChatCompletionRequest: """Translate Anthropic request to OpenAI chat completion format.""" ( diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 68db36b529e..6c41b46cfa0 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -2,7 +2,10 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any, Dict, List, Optional if TYPE_CHECKING: - from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, + ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.openai import AllMessageValues @@ -98,6 +101,30 @@ class BaseTranslation(ABC): """ return responses_so_far + def build_block_sse_chunks( + self, + exc: "ModifyResponseException", + stream_started: bool = False, + responses_so_far: Optional[list[Any]] = None, + ) -> Optional[list[bytes]]: + """ + Build the streaming chunks that deliver a guardrail block message and + cleanly terminate the stream in this provider's wire format. + + ``stream_started`` is True when real chunks were already sent to the + client: the result must *continue* the in-progress message (e.g. close + the open content block and append the block message) rather than start + a new one, which clients reject. ``responses_so_far`` provides the prior + chunks needed to do so. When False, nothing has been sent and a + standalone block message is emitted. + + Returns None when the format has no safe terminator; the caller then + re-raises ``exc`` so the proxy can surface a clean error instead. + Override in provider subclasses that support synthesizing a block + stream. + """ + return None + def get_structured_messages(self, data: dict) -> Optional[List["AllMessageValues"]]: """ Convert request data to OpenAI-spec structured messages. diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 83b787aaaf4..c325c8ba186 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -8,7 +8,7 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint import copy import json -from typing import Any, AsyncGenerator, List, Optional, Union +from typing import TYPE_CHECKING, Any, AsyncGenerator, List, Optional, Union from fastapi import HTTPException @@ -23,6 +23,11 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypes, CallTypesLiteral +if TYPE_CHECKING: + # Imported lazily at runtime (inside the streaming hook) to avoid a + # module-level cyclic import with litellm.integrations.custom_guardrail. + from litellm.integrations.custom_guardrail import ModifyResponseException + # Call types that use NDJSON streaming (A2A); guardrail HTTPException is emitted as in-stream error A2A_CALL_TYPES = (CallTypes.asend_message, CallTypes.send_message) @@ -263,6 +268,30 @@ class UnifiedLLMGuardrails(CustomLogger): return response + async def _handle_streaming_block( + self, + exc: "ModifyResponseException", + endpoint_translation: Any, + stream_started: bool, + responses_so_far: list[Any], + ) -> AsyncGenerator[Any, None]: + """ + Terminate a streamed response cleanly when a guardrail blocks it. + + Format-agnostic routing: delegates to the provider translation handler's + ``build_block_sse_chunks`` (see ``BaseTranslation.build_block_sse_chunks`` + for the ``stream_started`` / ``responses_so_far`` contract). When the + format has no safe terminator the handler returns None and we re-raise + ``exc`` so the proxy can surface a clean error. + """ + block_chunks = endpoint_translation.build_block_sse_chunks( + exc, stream_started=stream_started, responses_so_far=responses_so_far + ) + if block_chunks is None: + raise exc + for chunk in block_chunks: + yield chunk + async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -284,26 +313,38 @@ class UnifiedLLMGuardrails(CustomLogger): global endpoint_guardrail_translation_mappings + # Local import avoids a module-level cyclic import with + # litellm.integrations.custom_guardrail. + from litellm.integrations.custom_guardrail import ModifyResponseException + guardrail_to_apply: CustomGuardrail = request_data.pop("guardrail_to_apply", None) - # Get streaming configuration from guardrail or optional_params - sampling_rate = 5 - end_of_stream_only = False # If True, only apply guardrail at end of stream + # Get streaming configuration, with precedence: guardrail attribute -> + # guardrail_config dict -> this callback's optional_params. + def _streaming_flag(name: str, default: Any) -> Any: + value = default + if guardrail_to_apply is not None: + value = getattr(guardrail_to_apply, name, value) + config = getattr(guardrail_to_apply, "guardrail_config", {}) + if isinstance(config, dict): + value = config.get(name, value) + return self.optional_params.get(name, value) - if guardrail_to_apply is not None: - # Check direct attributes on guardrail first - sampling_rate = getattr(guardrail_to_apply, "streaming_sampling_rate", sampling_rate) - end_of_stream_only = getattr(guardrail_to_apply, "streaming_end_of_stream_only", end_of_stream_only) + sampling_rate = _streaming_flag("streaming_sampling_rate", 5) + # Only apply the guardrail at end of stream (not per chunk). + end_of_stream_only = _streaming_flag("streaming_end_of_stream_only", False) + # Withhold every chunk until end-of-stream moderation passes, then + # release the original chunks (clean) or only the block message + # (blocked) -- moderating the whole response *before* any content + # reaches the client. Intended for allow/block guardrails: on release + # the original chunks are replayed as-is, so content-rewriting + # guardrails (e.g. PII masking) are not applied to a buffered stream. + buffer_until_moderated = _streaming_flag("streaming_buffer_until_moderated", False) - # Also check guardrail_config dict if present - guardrail_config = getattr(guardrail_to_apply, "guardrail_config", {}) - if isinstance(guardrail_config, dict): - sampling_rate = guardrail_config.get("streaming_sampling_rate", sampling_rate) - end_of_stream_only = guardrail_config.get("streaming_end_of_stream_only", end_of_stream_only) - - # Also check optional_params as fallback - sampling_rate = self.optional_params.get("streaming_sampling_rate", sampling_rate) - end_of_stream_only = self.optional_params.get("streaming_end_of_stream_only", end_of_stream_only) + # Buffering can only moderate the assembled response, so it always + # defers to end-of-stream. + if buffer_until_moderated: + end_of_stream_only = True if guardrail_to_apply is None: async for item in response: @@ -328,6 +369,10 @@ class UnifiedLLMGuardrails(CustomLogger): call_type = None chunk_counter = 0 responses_so_far: List[Any] = [] + # Whether any real response chunk has been forwarded to the client. + # Drives how a block terminates the stream: continue the in-progress + # message (True) vs emit a standalone block message (False, buffered). + chunks_yielded = False async for item in response: chunk_counter += 1 @@ -349,9 +394,14 @@ class UnifiedLLMGuardrails(CustomLogger): yield remaining_item return - # If end_of_stream_only mode, yield chunks without processing + # If end_of_stream_only mode, yield chunks without processing. + # When buffering, withhold them instead -- they are released (or + # replaced by the block message) only after end-of-stream + # moderation runs below. if end_of_stream_only: - yield item + if not buffer_until_moderated: + chunks_yielded = True + yield item continue # Process chunk based on sampling rate @@ -381,6 +431,24 @@ class UnifiedLLMGuardrails(CustomLogger): user_api_key_dict=user_api_key_dict, request_data=request_data, ) + except ModifyResponseException as e: + # Guardrail blocked the response mid-stream. Emit a clean + # terminating SSE sequence delivering the block message + # instead of letting the exception propagate into a bare + # `data: {"error": ...}` blob (which truncates the stream). + # Chunks have already been forwarded here, so the block + # continues the in-progress message (stream_started=True). + # The current chunk was appended to responses_so_far but not + # yet yielded, so exclude it: the continuation must reflect + # only what the client has actually received. + async for block_chunk in self._handle_streaming_block( + e, + endpoint_translation, + stream_started=chunks_yielded, + responses_so_far=responses_so_far[:-1], + ): + yield block_chunk + return except HTTPException as e: # Response already started (we already yielded chunks); cannot send 400. # For A2A (NDJSON), yield an in-stream JSON-RPC error so the client sees it. @@ -407,8 +475,10 @@ class UnifiedLLMGuardrails(CustomLogger): yield error_chunk return raise + chunks_yielded = True yield original_item else: + chunks_yielded = True yield item # Stream has ended - do final processing with all collected chunks @@ -421,6 +491,15 @@ class UnifiedLLMGuardrails(CustomLogger): endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + # When buffering, snapshot the original chunks before moderation. + # A shallow copy suffices: end-of-stream + # process_output_streaming_response builds a separate assembled + # response (it does not mutate the individual chunks in place), and + # the chunks themselves are replayed verbatim -- so we only need to + # preserve the list, not clone every chunk (deepcopy would double + # peak memory for large responses). + buffered_items = list(responses_so_far) if buffer_until_moderated else None + try: await endpoint_translation.process_output_streaming_response( responses_so_far=responses_so_far, @@ -429,6 +508,23 @@ class UnifiedLLMGuardrails(CustomLogger): user_api_key_dict=user_api_key_dict, request_data=request_data, ) + # Moderation passed: release the withheld original chunks. + if buffered_items is not None: + for buffered_item in buffered_items: + yield buffered_item + except ModifyResponseException as e: + # Block detected during end-of-stream processing. Emit a clean + # terminating SSE sequence with the block message rather than + # propagating into a bare error blob that truncates the stream. + # The withheld original chunks are never released. + async for block_chunk in self._handle_streaming_block( + e, + endpoint_translation, + stream_started=chunks_yielded, + responses_so_far=responses_so_far, + ): + yield block_chunk + return except HTTPException as e: if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: request_id = _get_a2a_request_id(responses_so_far, request_data) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_anthropic_streaming_block.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_anthropic_streaming_block.py new file mode 100644 index 00000000000..69792fb2449 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_anthropic_streaming_block.py @@ -0,0 +1,275 @@ +""" +Regression tests for blocking an Anthropic streaming response from the +unified guardrail post-call streaming iterator hook. + +When a guardrail's ``apply_guardrail`` raises ``ModifyResponseException`` while +(or at the end of) an Anthropic ``/v1/messages`` stream is being relayed, the +hook must emit a well-formed Anthropic SSE termination sequence carrying the +block message - NOT a bare ``data: {"error": ...}`` blob that truncates the +stream and causes the Anthropic SDK parser to discard the response. +""" + +import json +from typing import Any, List, Literal, Optional + +import pytest + +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, +) +from litellm.types.utils import GenericGuardrailAPIInputs + +BLOCK_MESSAGE = "Blocked by policy: this response was withheld." + + +class _BlockingGuardrail(CustomGuardrail): + """Mock guardrail that always blocks by raising ModifyResponseException.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + raise ModifyResponseException( + message=BLOCK_MESSAGE, + model="claude-3-5-sonnet", + request_data=request_data, + guardrail_name=self.guardrail_name, + ) + + +def _sse_event(event_type: str, data: dict) -> bytes: + return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode() + + +async def _anthropic_stream(end: bool): + """Yield Anthropic SSE byte chunks. If end=True, include a terminating + message_delta (stop_reason set) so the hook's end-of-stream path runs.""" + yield _sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_orig", + "type": "message", + "role": "assistant", + "model": "claude-3-5-sonnet", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + ) + yield _sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ) + for text in ["This ", "is ", "the ", "original ", "answer."]: + yield _sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": text}, + }, + ) + if end: + yield _sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}) + yield _sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, + ) + yield _sse_event("message_stop", {"type": "message_stop"}) + + +def _decode(chunks: List[Any]) -> str: + parts = [] + for chunk in chunks: + parts.append(chunk.decode() if isinstance(chunk, bytes) else str(chunk)) + return "".join(parts) + + +def _parse_sse_event_types(raw: str) -> List[str]: + event_types = [] + for block in raw.split("\n\n"): + for line in block.strip().split("\n"): + if line.startswith("data:"): + payload = line[len("data:") :].strip() + try: + event_types.append(json.loads(payload).get("type")) + except json.JSONDecodeError: + pass + return event_types + + +async def _run_hook(end: bool, sampling_rate: int = 1) -> str: + guardrail = _BlockingGuardrail(guardrail_name="test-blocking-guardrail", event_hook="post_call") + # sampling_rate controls how many chunks are forwarded before the block + # fires: 1 blocks on the first chunk (nothing sent yet); >1 forwards earlier + # chunks first, exercising the mid-stream "continue the message" path. + guardrail.streaming_sampling_rate = sampling_rate + + unified_guardrail = UnifiedLLMGuardrails() + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/messages") + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": ["test-blocking-guardrail"]}, + } + + collected: List[Any] = [] + async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=_anthropic_stream(end=end), + request_data=request_data, + ): + collected.append(chunk) + return _decode(collected) + + +def _assert_clean_block_termination(raw: str) -> None: + # No bare error blob that would truncate the stream. + assert '"error"' not in raw, f"unexpected error blob in stream: {raw!r}" + # The block message is delivered as assistant text. + assert BLOCK_MESSAGE in raw, f"block message missing from stream: {raw!r}" + # A complete, parseable Anthropic SSE termination sequence is present. + event_types = _parse_sse_event_types(raw) + assert "message_start" in event_types + assert "content_block_delta" in event_types + # Exactly one message_start: a block must never inject a second + # message envelope into an already-started stream (clients reject it). + assert event_types.count("message_start") == 1, f"expected a single message_start, got: {event_types}" + assert event_types[-1] == "message_stop", f"stream did not end cleanly: {event_types}" + # message_delta carries a stop_reason. + assert any('"stop_reason"' in block and "message_delta" in block for block in raw.split("\n\n")) + + +@pytest.mark.asyncio +async def test_mid_stream_block_emits_clean_anthropic_sse(): + """Per-chunk block: a clean SSE termination with the block message, no error blob.""" + raw = await _run_hook(end=False) + _assert_clean_block_termination(raw) + + +@pytest.mark.asyncio +async def test_end_of_stream_block_emits_clean_anthropic_sse(): + """End-of-stream block: same clean SSE termination guarantees.""" + raw = await _run_hook(end=True) + _assert_clean_block_termination(raw) + + +@pytest.mark.asyncio +async def test_mid_stream_block_after_prior_chunks_continues_message(): + """Regression: when real chunks were already forwarded (sampling_rate>1), + the block must continue the in-progress message, not start a second one.""" + raw = await _run_hook(end=False, sampling_rate=5) + # Some original content was forwarded before the block... + assert "message_start" in raw + # ...and the block continues that same message (single message_start) with + # the block message appended, ending cleanly. + _assert_clean_block_termination(raw) + + +class TestContentBlockState: + """`_content_block_state` must reflect the true open/last block index across + the two chunk formats the stream can carry (multi-event bytes, parsed dict), + so a mid-stream block closes/opens the right indices.""" + + def _handler(self): + from litellm.llms.anthropic.chat.guardrail_translation.handler import ( + AnthropicMessagesHandler, + ) + + return AnthropicMessagesHandler() + + def test_multi_event_bytes_chunk_is_fully_parsed(self): + # One item bundles start(0) + delta + stop(0): the block is already + # closed, so open_index is None (not 0) and max_index is 0. + bundled = ( + _sse_event( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ) + + _sse_event( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}}, + ) + + _sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}) + ) + open_index, max_index = self._handler()._content_block_state([bundled]) + assert open_index is None + assert max_index == 0 + + def test_open_block_across_separate_chunks(self): + chunks = [ + _sse_event( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}), + _sse_event( + "content_block_start", + {"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}, + ), + ] + open_index, max_index = self._handler()._content_block_state(chunks) + assert open_index == 1 + assert max_index == 1 + + def test_dict_format_chunks_are_parsed(self): + # The backwards-compat parsed-dict format must be understood too. + chunks = [ + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}}, + ] + open_index, max_index = self._handler()._content_block_state(chunks) + assert open_index == 0 + assert max_index == 0 + + def test_continuation_closes_open_block_and_appends_after_it(self): + from litellm.integrations.custom_guardrail import ModifyResponseException + + handler = self._handler() + # Client has seen an open text block at index 0. + seen = [ + _sse_event( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _sse_event( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "partial"}}, + ), + ] + exc = ModifyResponseException( + message=BLOCK_MESSAGE, model="claude-3-5-sonnet", request_data={}, guardrail_name="g" + ) + raw = b"".join(handler.build_block_sse_chunks(exc, stream_started=True, responses_so_far=seen)).decode() + events = _parse_sse_event_types(raw) + # No new message envelope, closes block 0, appends block text at index 1. + assert "message_start" not in events + assert events == [ + "content_block_stop", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ] + assert BLOCK_MESSAGE in raw + assert '"index": 1' in raw diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py new file mode 100644 index 00000000000..857389c6af2 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py @@ -0,0 +1,159 @@ +""" +Tests for ``streaming_buffer_until_moderated`` on the unified guardrail +post-call streaming iterator hook. + +With this flag set, the hook must withhold every upstream chunk until +end-of-stream moderation has run. The decisive guarantee versus the +detect-only ``streaming_end_of_stream_only`` behavior: when the guardrail +blocks, the original (objectionable) content is NEVER yielded to the client -- +only the block message is. On a clean response, all original chunks are +released unchanged after moderation passes. +""" + +import json +from typing import Any, List, Literal, Optional + +import pytest + +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, +) +from litellm.types.utils import GenericGuardrailAPIInputs + +BLOCK_MESSAGE = "Blocked by policy: this response was withheld." +ORIGINAL_MARKER = "ORIGINAL-SECRET-ANSWER" + + +class _BlockingGuardrail(CustomGuardrail): + """Always blocks at moderation time.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + raise ModifyResponseException( + message=BLOCK_MESSAGE, + model="claude-3-5-sonnet", + request_data=request_data, + guardrail_name=self.guardrail_name, + ) + + +class _PassingGuardrail(CustomGuardrail): + """Never blocks; returns inputs unchanged.""" + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + return inputs + + +def _sse_event(event_type: str, data: dict) -> bytes: + return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode() + + +async def _anthropic_stream(): + """A complete Anthropic /v1/messages SSE stream whose assistant text + contains ORIGINAL_MARKER so leakage is unambiguous to assert.""" + yield _sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_orig", + "type": "message", + "role": "assistant", + "model": "claude-3-5-sonnet", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + ) + yield _sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ) + for text in ["Here is ", "the ", ORIGINAL_MARKER, " for you."]: + yield _sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": text}, + }, + ) + yield _sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}) + yield _sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, + ) + yield _sse_event("message_stop", {"type": "message_stop"}) + + +def _decode(chunks: List[Any]) -> str: + return "".join(c.decode() if isinstance(c, bytes) else str(c) for c in chunks) + + +async def _run(guardrail: CustomGuardrail) -> str: + # Rubrik's real config: end-of-stream-only moderation. Without buffering + # this releases every chunk before moderation runs (content leaks on + # block); the buffer flag must change that to moderate-then-release. + guardrail.streaming_end_of_stream_only = True + guardrail.streaming_buffer_until_moderated = True + unified = UnifiedLLMGuardrails() + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/messages") + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": [guardrail.guardrail_name]}, + } + collected: List[Any] = [] + async for chunk in unified.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=_anthropic_stream(), + request_data=request_data, + ): + collected.append(chunk) + return _decode(collected) + + +@pytest.mark.asyncio +async def test_buffered_block_withholds_original_content(): + raw = await _run(_BlockingGuardrail(guardrail_name="blk", event_hook="post_call")) + # The original content must never reach the client... + assert ORIGINAL_MARKER not in raw, f"original content leaked: {raw!r}" + # ...only the block message, in a clean terminating stream. + assert BLOCK_MESSAGE in raw + assert '"error"' not in raw + + +@pytest.mark.asyncio +async def test_buffered_clean_releases_all_content(): + raw = await _run(_PassingGuardrail(guardrail_name="pass", event_hook="post_call")) + # A clean response is released in full after moderation passes. + assert ORIGINAL_MARKER in raw + assert ( + raw.rstrip().endswith('event: message_stop\ndata: {"type": "message_stop"}'.rstrip()) or "message_stop" in raw + ) + assert BLOCK_MESSAGE not in raw