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 d4d23cd2e37..5cbd05dedfc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -804,6 +804,8 @@ class UnifiedLLMGuardrails(CustomLogger): user_api_key_dict: UserAPIKeyAuth, response: Any, request_data: dict, + guardrail_to_apply: Union[CustomGuardrail, None] = None, + buffer_until_moderated_default: bool = False, ) -> AsyncGenerator[Any, None]: """ Passes the entire stream to the guardrail @@ -824,7 +826,8 @@ class UnifiedLLMGuardrails(CustomLogger): # litellm.integrations.custom_guardrail. from litellm.integrations.custom_guardrail import ModifyResponseException - guardrail_to_apply: CustomGuardrail = request_data.pop("guardrail_to_apply", None) + if guardrail_to_apply is None: + guardrail_to_apply = request_data.pop("guardrail_to_apply", None) # Get streaming configuration. Resolution order (later wins): default # < guardrail attribute < guardrail_config dict < this callback's @@ -852,7 +855,7 @@ class UnifiedLLMGuardrails(CustomLogger): # release the original chunks are replayed as-is, so a # content-rewriting guardrail (e.g. PII masking) would leak # unredacted content. Guarded below via mask_response_content. - buffer_until_moderated = _streaming_flag("streaming_buffer_until_moderated", False) + buffer_until_moderated = _streaming_flag("streaming_buffer_until_moderated", buffer_until_moderated_default) if ( buffer_until_moderated diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1e80f15a761..d95d64c47f4 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -187,6 +187,8 @@ else: unified_guardrail = UnifiedLLMGuardrails() +NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES: "frozenset[CallTypes]" = frozenset({CallTypes.anthropic_messages}) + def print_verbose(print_statement): """ @@ -1762,6 +1764,20 @@ class ProxyLogging: cache[sig] = caps return caps + @staticmethod + def _stream_requires_guardrail_translation(user_api_key_dict: UserAPIKeyAuth) -> bool: + from litellm.litellm_core_utils.api_route_to_call_types import ( + get_call_types_for_route, + ) + + route = user_api_key_dict.request_route + if not route: + return False + call_types = get_call_types_for_route(route) + if not call_types: + return False + return call_types[0] in NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES + @staticmethod def has_post_call_response_headers_callbacks() -> bool: return ProxyLogging._callback_capabilities().has_post_call_response_headers @@ -2723,6 +2739,7 @@ class ProxyLogging: request_data = _check_and_merge_model_level_guardrails(data=request_data, llm_router=llm_router) current_response = response + stream_needs_translation = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict) for resolved_callback, kind in caps.iterator_overrides: if isinstance(resolved_callback, CustomGuardrail): @@ -2731,7 +2748,18 @@ class ProxyLogging: is not True ): continue - if kind == "override": + effective_kind = ( + "apply_guardrail" + if ( + kind == "override" + and stream_needs_translation + and isinstance(resolved_callback, CustomGuardrail) + and resolved_callback.uses_apply_guardrail_interface() + and not resolved_callback.mask_response_content + ) + else kind + ) + if effective_kind == "override": current_response = self._wrap_streaming_iterator_with_enrichment( resolved_callback, resolved_callback.async_post_call_streaming_iterator_hook( @@ -2742,13 +2770,14 @@ class ProxyLogging: ) else: # kind == "apply_guardrail": route through unified_guardrail - request_data["guardrail_to_apply"] = resolved_callback current_response = self._wrap_streaming_iterator_with_enrichment( resolved_callback, unified_guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, request_data=request_data, response=current_response, + guardrail_to_apply=resolved_callback, + buffer_until_moderated_default=(kind == "override"), ), ) diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index f5967030561..015dcd9b5db 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -148,3 +148,334 @@ def test_callback_capabilities_cache_invalidates_on_list_change(monkeypatch): caps = ProxyLogging._callback_capabilities() assert caps.has_pre_call_override is True assert pre in caps.resolved_callbacks + + +def _sse_bytes(event: str, payload: dict) -> bytes: + import json + + return f"event: {event}\ndata: {json.dumps(payload)}\n\n".encode() + + +def _anthropic_stream_chunks(text_parts): + chunks = [ + _sse_bytes( + "message_start", + { + "type": "message_start", + "message": { + "model": "claude-sonnet-5", + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 20, "output_tokens": 1}, + }, + }, + ), + _sse_bytes( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + ] + for part in text_parts: + chunks.append( + _sse_bytes( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": part}}, + ) + ) + chunks.append(_sse_bytes("content_block_stop", {"type": "content_block_stop", "index": 0})) + chunks.append( + _sse_bytes( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"input_tokens": 20, "output_tokens": 8}, + }, + ) + ) + chunks.append(_sse_bytes("message_stop", {"type": "message_stop"})) + return chunks + + +def _content_filter_guardrail(action: str, guardrail_cls=None, **guardrail_kwargs): + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction + + cls = guardrail_cls or ContentFilterGuardrail + return cls( + guardrail_name="output-filter", + blocked_words=[BlockedWord(keyword="zebra", action=ContentFilterAction(action))], + event_hook="post_call", + default_on=True, + **guardrail_kwargs, + ) + + +def _streaming_logging_obj(): + import datetime + import uuid + + from litellm.litellm_core_utils.litellm_logging import Logging + + return Logging( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "Reply with exactly: the zebra runs"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.datetime.now(), + litellm_call_id=str(uuid.uuid4()), + function_id="test", + ) + + +def test_stream_requires_guardrail_translation_route_detection(): + from litellm.proxy._types import UserAPIKeyAuth + + assert ( + ProxyLogging._stream_requires_guardrail_translation( + UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages") + ) + is True + ) + assert ( + ProxyLogging._stream_requires_guardrail_translation( + UserAPIKeyAuth(api_key="sk-1234", request_route="/chat/completions") + ) + is False + ) + assert ProxyLogging._stream_requires_guardrail_translation(UserAPIKeyAuth(api_key="sk-1234")) is False + assert ( + ProxyLogging._stream_requires_guardrail_translation( + UserAPIKeyAuth(api_key="sk-1234", request_route="/route/without/call/types") + ) + is False + ) + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_blocks_anthropic_messages_stream(monkeypatch): + """ + Regression test for https://github.com/BerriAI/litellm/issues/35257. + + /v1/messages streams raw Anthropic SSE bytes. A guardrail whose custom + iterator hook only understands OpenAI ModelResponseStream chunks used to + receive those bytes directly and silently pass every chunk through + unscanned. The dispatch must route apply_guardrail-capable guardrails + through unified_guardrail's anthropic translation so blocked output + raises instead of streaming to the client. Because the guardrail's own + iterator hook withheld content until scanned, the rerouted invocation + defaults to buffer_until_moderated, so nothing may reach the client + before the block fires. + """ + from fastapi import HTTPException + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + + guardrail = _content_filter_guardrail("BLOCK") + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + request_data = { + "model": "claude-sonnet-5", + "litellm_logging_obj": _streaming_logging_obj(), + "metadata": {}, + } + + async def fake_stream(): + for chunk in _anthropic_stream_chunks(["the", " zebra runs"]): + yield chunk + + delivered = [] + with pytest.raises(HTTPException) as exc_info: + async for chunk in proxy_logging.async_post_call_streaming_iterator_hook( + response=fake_stream(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"), + request_data=request_data, + ): + delivered.append(chunk) + + detail = exc_info.value.detail + assert detail["guardrail_name"] == "output-filter" + assert detail["keyword"] == "zebra" + assert delivered == [] + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_keeps_own_iterator_on_chat_completions(monkeypatch): + """ + On /chat/completions the guardrail's own iterator hook must keep running: + it masks incrementally inside ModelResponseStream chunks, which the + unified block_only path never does. Masked output proves the own-hook + path was used. + """ + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + guardrail = _content_filter_guardrail("MASK") + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + async def fake_stream(): + yield ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content="the zebra runs"))] + ) + yield ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")] + ) + + delivered_text = "" + async for chunk in proxy_logging.async_post_call_streaming_iterator_hook( + response=fake_stream(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/chat/completions"), + request_data={"model": "gpt-4o-mini", "metadata": {}}, + ): + for choice in chunk.choices: + delivered_text += choice.delta.content or "" + + assert "zebra" not in delivered_text + assert delivered_text != "" + + +@pytest.mark.asyncio +async def test_unified_guardrail_iterator_accepts_explicit_guardrail(monkeypatch): + """ + The dispatch passes each guardrail explicitly instead of through a shared + request_data key, so chaining two unified-routed guardrails cannot drop + all but the last one. + """ + from fastapi import HTTPException + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import unified_guardrail + + guardrail = _content_filter_guardrail("BLOCK") + request_data = { + "model": "claude-sonnet-5", + "litellm_logging_obj": _streaming_logging_obj(), + "metadata": {}, + } + + async def fake_stream(): + for chunk in _anthropic_stream_chunks(["the", " zebra runs"]): + yield chunk + + with pytest.raises(HTTPException): + async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"), + response=fake_stream(), + request_data=request_data, + guardrail_to_apply=guardrail, + ): + pass + + +@pytest.mark.asyncio +async def test_post_call_stream_guardrail_reroutes_inherited_apply_guardrail(monkeypatch): + """ + The reroute predicate must recognize apply_guardrail implementations + inherited from a parent class, not only ones defined on the registered + leaf class. A vendor base class can carry apply_guardrail while the leaf + only overrides the streaming iterator; a leaf-class ``__dict__`` check + would leave that guardrail on the raw Anthropic SSE path unscanned. + """ + from fastapi import HTTPException + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + + class _InheritsApplyGuardrail(ContentFilterGuardrail): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + async for item in response: + yield item + + guardrail = _content_filter_guardrail("BLOCK", guardrail_cls=_InheritsApplyGuardrail) + assert "apply_guardrail" not in type(guardrail).__dict__ + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + request_data = { + "model": "claude-sonnet-5", + "litellm_logging_obj": _streaming_logging_obj(), + "metadata": {}, + } + + async def fake_stream(): + for chunk in _anthropic_stream_chunks(["the", " zebra runs"]): + yield chunk + + delivered = [] + with pytest.raises(HTTPException) as exc_info: + async for chunk in proxy_logging.async_post_call_streaming_iterator_hook( + response=fake_stream(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"), + request_data=request_data, + ): + delivered.append(chunk) + + assert exc_info.value.detail["keyword"] == "zebra" + assert delivered == [] + + +@pytest.mark.asyncio +async def test_post_call_stream_masking_guardrail_keeps_own_iterator_on_anthropic(monkeypatch): + """ + A guardrail with mask_response_content=True must stay on its own iterator + hook on /v1/messages. The unified streaming path cannot re-emit rewritten + text on raw Anthropic SSE (block_only drops rewrites and buffered replay + releases the unredacted originals), so rerouting such a guardrail would + deliver content it decided to mask. PANW Prisma AIRS is the concrete + case: its own hook parses the raw bytes and blocks instead of masking. + """ + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + + own_hook_streams = [] + + class _MasksViaOwnRawStreamHook(ContentFilterGuardrail): + apply_guardrail = ContentFilterGuardrail.apply_guardrail + + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + own_hook_streams.append(request_data.get("model")) + async for item in response: + yield item + + guardrail = _content_filter_guardrail( + "BLOCK", guardrail_cls=_MasksViaOwnRawStreamHook, mask_response_content=True + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + chunks = _anthropic_stream_chunks(["the", " zebra runs"]) + + async def fake_stream(): + for chunk in chunks: + yield chunk + + delivered = [] + async for chunk in proxy_logging.async_post_call_streaming_iterator_hook( + response=fake_stream(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"), + request_data={ + "model": "claude-sonnet-5", + "litellm_logging_obj": _streaming_logging_obj(), + "metadata": {}, + }, + ): + delivered.append(chunk) + + assert own_hook_streams == ["claude-sonnet-5"] + assert delivered == chunks