diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 8dc6881d23e..a6e3d000120 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -762,6 +762,9 @@ class CustomGuardrail(CustomLogger): def uses_apply_guardrail_interface(self) -> bool: return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail + def rewrites_streamed_output(self) -> bool: + return self.mask_response_content + def _deployment_pre_call_target(self) -> "CustomLogger": if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks: return self diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index 722f96ef814..bd31882841e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -1947,6 +1947,13 @@ class ContentFilterGuardrail(CustomGuardrail): exception_str=exception_str, ) + def rewrites_streamed_output(self) -> bool: + return ( + super().rewrites_streamed_output() + or any(entry["action"] == ContentFilterAction.MASK for entry in self.compiled_patterns) + or any(action == ContentFilterAction.MASK for action, _ in self.blocked_words.values()) + ) + async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a64b57c1388..138f272af42 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -497,7 +497,7 @@ def _pipeline_step_rewrites_streamed_content(guardrail_name: str) -> bool: if callback is None: return False transform_mode: Final = unified_guardrail.resolve_streaming_flag(callback, "streaming_transform_mode", "block_only") - return callback.mask_response_content or transform_mode == "incremental_diff" + return callback.rewrites_streamed_output() or transform_mode == "incremental_diff" class _PipelineErrorBody(TypedDict): @@ -518,10 +518,10 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_ap Background responses skip the post_call hooks entirely, so a pipeline governing one would silently never execute. Streaming responses execute pipelines against the buffered stream through the endpoint guardrail - translation of the request route, releasing the buffered chunks verbatim - on allow. That needs every step's guardrail to support the unified - apply_guardrail interface and to only allow or block (a step that rewrites - streamed content, via mask_response_content or + translation of the request route, releasing the buffered chunks on allow. + That needs every step's guardrail to support the unified apply_guardrail + interface and to only allow or block (a step that rewrites streamed + content, via mask_response_content, a MASK action, or streaming_transform_mode=incremental_diff, would have its rewrite silently dropped), and needs the route to have a translation at all; anything else keeps the 400 rather than letting ungoverned output stream through. @@ -577,8 +577,8 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_ap "error": { "message": ( "Policies with post_call guardrail pipelines cannot govern streaming responses " - "because these pipeline guardrails rewrite streamed content (mask_response_content " - "or streaming_transform_mode=incremental_diff), which pipeline steps would release " + "because these pipeline guardrails rewrite streamed content (mask_response_content, " + "a MASK action, or streaming_transform_mode=incremental_diff), which pipeline steps would release " f"unmodified: {', '.join(rewriting_guardrails)}. Retry with stream=false, or drop " "them from the pipeline steps so guardrails.add applies them to streamed output." ), diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index be55ac47bde..ffbedfa43ff 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -3068,3 +3068,39 @@ class TestContentFilterToolCallArguments: request_data={}, input_type="response", ) + + +class TestRewritesStreamedOutput: + def test_block_only_rules_do_not_rewrite(self): + guardrail = ContentFilterGuardrail( + guardrail_name="cf", + patterns=[ContentFilterPattern(pattern_type="prebuilt", pattern_name="us_ssn", action=ContentFilterAction.BLOCK)], + blocked_words=[BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK)], + ) + + assert guardrail.rewrites_streamed_output() is False + + def test_mask_blocked_word_rewrites(self): + guardrail = ContentFilterGuardrail( + guardrail_name="cf", + blocked_words=[BlockedWord(keyword="persimmon", action=ContentFilterAction.MASK)], + ) + + assert guardrail.rewrites_streamed_output() is True + + def test_mask_pattern_rewrites(self): + guardrail = ContentFilterGuardrail( + guardrail_name="cf", + patterns=[ContentFilterPattern(pattern_type="prebuilt", pattern_name="us_ssn", action=ContentFilterAction.MASK)], + ) + + assert guardrail.rewrites_streamed_output() is True + + def test_mask_response_content_rewrites(self): + guardrail = ContentFilterGuardrail( + guardrail_name="cf", + blocked_words=[BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK)], + mask_response_content=True, + ) + + assert guardrail.rewrites_streamed_output() is True diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index a258442c79a..ee0a9a10172 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -27,7 +27,8 @@ from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy._types import ProxyException from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.utils import ProxyLogging, _raise_for_streaming_post_call_pipelines -from litellm.types.guardrails import GuardrailEventHooks +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail +from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, @@ -1461,6 +1462,37 @@ async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_rewrites_ assert seen.get("count") is None +@pytest.mark.asyncio +@pytest.mark.parametrize("action, rejected", [(ContentFilterAction.MASK, True), (ContentFilterAction.BLOCK, False)]) +async def test_pre_call_hook_rejects_streaming_only_when_content_filter_step_masks( + proxy_logging, make_user_api_key_auth, monkeypatch, action, rejected +): + guardrail = ContentFilterGuardrail( + guardrail_name="gr-post", + event_hook=GuardrailEventHooks.post_call, + blocked_words=[BlockedWord(keyword="persimmon", action=action)], + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + data = _post_call_pipeline_data(stream=True) + user_api_key_dict = make_user_api_key_auth(request_route="/v1/chat/completions") + + if not rejected: + out = await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True + ) + assert out is not None and out.get("stream") is True + return + + with pytest.raises(HTTPException) as info: + await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True + ) + + assert info.value.status_code == 400 + assert info.value.detail["error"]["guardrails"] == ("gr-post",) + assert "a MASK action" in info.value.detail["error"]["message"] + + @pytest.mark.asyncio async def test_pre_call_hook_rejects_streaming_when_route_has_no_guardrail_translation( proxy_logging, make_user_api_key_auth, monkeypatch