fix(guardrails): replace blocked text completion streams like chat, keep Gray Swan's vendor payload, scan text streams in the content filter

litellm/llms/openai/completion/guardrail_translation/handler.py

@@ -16,13 +16,17 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import (
     StreamingScanKey,
     StreamTransformSink,
 )
-from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items
-from litellm.types.utils import GenericGuardrailAPIInputs, TextChoices, TextCompletionResponse
+from litellm.llms.base_llm.guardrail_translation.utils import (
+    blocked_chat_stream_usage,
+    stream_item_field,
+    stream_item_items,
+)
+from litellm.types.utils import GenericGuardrailAPIInputs, TextChoices, TextCompletionResponse, Usage
 
 if TYPE_CHECKING:
     from fastapi import HTTPException
 
-    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.proxy.auth.user_api_key_auth import UserAPIKeyAuth
 
@@ -258,6 +262,38 @@ class OpenAITextCompletionHandler(BaseTranslation):
         )
         return (*responses_so_far, terminator)
 
+    def build_block_sse_chunks(
+        self,
+        exc: "ModifyResponseException",
+        stream_started: bool = False,
+        responses_so_far: Sequence[object] | None = None,
+    ) -> Sequence[bytes]:
+        streamed: Final = _text_completion_chunks(responses_so_far or ())
+        source: Final = streamed[0] if streamed else TextCompletionResponse(model=exc.model)
+        prompt_tokens, completion_tokens = blocked_chat_stream_usage(exc.original_response)
+        message_chunk: Final = TextCompletionResponse(
+            id=source.id,
+            created=source.created,
+            model=source.model or exc.model,
+            object="text_completion",
+            choices=[TextChoices(index=0, text=exc.message)],
+        )
+        final_chunk: Final = TextCompletionResponse(
+            id=source.id,
+            created=source.created,
+            model=source.model or exc.model,
+            object="text_completion",
+            choices=[TextChoices(index=0, text="", finish_reason="content_filter")],
+            usage=Usage(
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                total_tokens=prompt_tokens + completion_tokens,
+            ),
+        )
+        return tuple(
+            f"data: {chunk.model_dump_json(exclude_none=True)}\n\n".encode() for chunk in (message_chunk, final_chunk)
+        )
+
     def build_stream_error_items(
         self,
         exc: "HTTPException",

litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py

@@ -89,12 +89,13 @@ def _str_keyed(value: object) -> Mapping[str, object]:
     return _STR_KEYED_MAPPING.validate_python(value) if isinstance(value, Mapping) else {}
 
 
-def _proxy_metadata(request_data: Mapping[str, object]) -> Mapping[str, object]:
+def _request_route(request_data: Mapping[str, object]) -> str | None:
     litellm_metadata: Final = request_data.get("litellm_metadata")
-    if litellm_metadata is not None:
-        return _str_keyed(litellm_metadata)
-    metadata: Final = _str_keyed(request_data.get("metadata"))
-    return {key: value for key, value in metadata.items() if key.startswith("user_api_key_")}
+    metadata_bucket: Final = _str_keyed(
+        litellm_metadata if litellm_metadata is not None else request_data.get("metadata")
+    )
+    route: Final = metadata_bucket.get("user_api_key_request_route")
+    return route if isinstance(route, str) else None
 
 
 class GraySwanGuardrailMissingSecrets(Exception):
@@ -640,8 +641,8 @@ class GraySwanGuardrail(CustomGuardrail):
         from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
         from litellm.llms import load_guardrail_translation_mappings
 
-        request_route: Final = _proxy_metadata(request_data).get("user_api_key_request_route")
-        route_call_types: Final = get_call_types_for_route(request_route) if isinstance(request_route, str) else None
+        request_route: Final = _request_route(request_data)
+        route_call_types: Final = get_call_types_for_route(request_route) if request_route is not None else None
         call_type: Final = (
             (route_call_types[0].value if route_call_types else None)
             or (logging_obj.call_type if logging_obj is not None else None)
@@ -710,11 +711,11 @@ class GraySwanGuardrail(CustomGuardrail):
 
         inbound_headers: Final = self._extract_inbound_headers(request_data, logging_obj)
 
-        proxy_metadata: Final = _proxy_metadata(request_data)
+        litellm_metadata: Final = _str_keyed(request_data.get("litellm_metadata"))
         cleaned_litellm_metadata: Final = (
-            {**proxy_metadata, "headers": {**_str_keyed(proxy_metadata.get("headers")), **inbound_headers}}
+            {**litellm_metadata, "headers": {**_str_keyed(litellm_metadata.get("headers")), **inbound_headers}}
             if inbound_headers
-            else proxy_metadata
+            else litellm_metadata
         )
         if cleaned_litellm_metadata:
             sanitized: Final = _DECODED_JSON.validate_python(

litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py

@@ -47,6 +47,9 @@ from litellm.types.utils import (
     Message,
     ModelResponse,
     ModelResponseStream,
+    StreamingChoices,
+    TextChoices,
+    TextCompletionResponse,
 )
 
 if TYPE_CHECKING:
@@ -139,6 +142,22 @@ class _CategoryConfigView(TypedDict):
     category_file: str | None
 
 
+def _streamed_choice_content(choice: StreamingChoices | TextChoices) -> str | None:
+    if isinstance(choice, TextChoices):
+        return choice.text if isinstance(choice.text, str) else ""
+    delta: Final = getattr(choice, "delta", None)
+    if not delta:
+        return None
+    content: Final = getattr(delta, "content", None)
+    return content if isinstance(content, str) else ""
+
+
+def _streamed_choices(item: object) -> Sequence[StreamingChoices | TextChoices]:
+    if isinstance(item, (ModelResponseStream, TextCompletionResponse)):
+        return item.choices
+    return ()
+
+
 @dataclass(frozen=True, slots=True)
 class _StreamedChoiceState:
     buffered_text: str = ""
@@ -2105,71 +2124,70 @@ class ContentFilterGuardrail(CustomGuardrail):
 
         try:
             async for item in response:
-                if isinstance(item, ModelResponseStream) and item.choices:
-                    for choice in item.choices:
-                        if not (hasattr(choice, "delta") and choice.delta):
-                            continue
+                for choice in _streamed_choices(item):
+                    new_content = _streamed_choice_content(choice)
+                    if new_content is None:
+                        continue
 
-                        choice_index = getattr(choice, "index", 0)
-                        if not isinstance(choice_index, int):
-                            choice_index = 0
+                    choice_index = getattr(choice, "index", 0)
+                    if not isinstance(choice_index, int):
+                        choice_index = 0
 
-                        content = getattr(choice.delta, "content", None)
-                        is_final = bool(getattr(choice, "finish_reason", None))
-                        new_content = content if isinstance(content, str) else ""
-                        if not new_content and not is_final:
-                            continue
+                    is_final = bool(getattr(choice, "finish_reason", None))
+                    if not new_content and not is_final:
+                        continue
 
-                        previous_state = state_by_choice.get(choice_index, _StreamedChoiceState())
-                        buffered_text = previous_state.buffered_text + new_content
-                        if not buffered_text:
-                            continue
+                    previous_state = state_by_choice.get(choice_index, _StreamedChoiceState())
+                    buffered_text = previous_state.buffered_text + new_content
+                    if not buffered_text:
+                        continue
 
-                        # Add a space at the end if it's the final chunk to trigger word boundaries (\b)
-                        text_to_scan = buffered_text + (" " if is_final else "")
-                        choice_detections: list[ContentFilterDetection] = []
-                        scan_started = time.perf_counter()
+                    # Add a space at the end if it's the final chunk to trigger word boundaries (\b)
+                    text_to_scan = buffered_text + (" " if is_final else "")
+                    choice_detections: list[ContentFilterDetection] = []
+                    scan_started = time.perf_counter()
 
-                        try:
-                            masked_text = self._filter_single_text(text_to_scan, detections=choice_detections)
-                            if is_final and masked_text.endswith(" "):
-                                masked_text = masked_text[:-1]
-                            latest_detections = tuple(choice_detections)
-                        except HTTPException:
-                            state_by_choice[choice_index] = replace(
-                                previous_state, latest_detections=tuple(choice_detections)
-                            )
-                            raise
-                        except Exception as e:
-                            verbose_proxy_logger.error("ContentFilterGuardrail: Error in masking: %s", e)
-                            masked_text = text_to_scan  # Fallback to current text
-                            latest_detections = previous_state.latest_detections
-                        finally:
-                            scan_seconds += time.perf_counter() - scan_started
-
-                        safe_to_yield_len = max(
-                            previous_state.yielded_masked_text_len,
-                            len(masked_text) - (0 if is_final else CONTENT_FILTER_STREAMING_HOLDBACK_CHARS),
+                    try:
+                        masked_text = self._filter_single_text(text_to_scan, detections=choice_detections)
+                        if is_final and masked_text.endswith(" "):
+                            masked_text = masked_text[:-1]
+                        latest_detections = tuple(choice_detections)
+                    except HTTPException:
+                        state_by_choice[choice_index] = replace(
+                            previous_state, latest_detections=tuple(choice_detections)
                         )
-                        choice.delta.content = masked_text[previous_state.yielded_masked_text_len : safe_to_yield_len]
-                        next_state = replace(
-                            previous_state,
-                            buffered_text=buffered_text,
-                            yielded_masked_text_len=safe_to_yield_len,
-                            latest_detections=latest_detections,
-                        )
-                        if is_final:
-                            state_by_choice[choice_index] = next_state
-                            continue
+                        raise
+                    except Exception as e:
+                        verbose_proxy_logger.error("ContentFilterGuardrail: Error in masking: %s", e)
+                        masked_text = text_to_scan  # Fallback to current text
+                        latest_detections = previous_state.latest_detections
+                    finally:
+                        scan_seconds += time.perf_counter() - scan_started
 
-                        trim_started = time.perf_counter()
-                        state_by_choice[choice_index] = self._trim_streamed_choice_buffer(next_state, masked_text, plan)
-                        scan_seconds += time.perf_counter() - trim_started
+                    safe_to_yield_len = max(
+                        previous_state.yielded_masked_text_len,
+                        len(masked_text) - (0 if is_final else CONTENT_FILTER_STREAMING_HOLDBACK_CHARS),
+                    )
+                    released_text = masked_text[previous_state.yielded_masked_text_len : safe_to_yield_len]
+                    if isinstance(choice, TextChoices):
+                        choice.text = released_text
+                    else:
+                        choice.delta.content = released_text
+                    next_state = replace(
+                        previous_state,
+                        buffered_text=buffered_text,
+                        yielded_masked_text_len=safe_to_yield_len,
+                        latest_detections=latest_detections,
+                    )
+                    if is_final:
+                        state_by_choice[choice_index] = next_state
+                        continue
 
-                    yield item
-                else:
-                    # Not a ModelResponseStream or no choices - yield as is
-                    yield item
+                    trim_started = time.perf_counter()
+                    state_by_choice[choice_index] = self._trim_streamed_choice_buffer(next_state, masked_text, plan)
+                    scan_seconds += time.perf_counter() - trim_started
+
+                yield item
         except HTTPException:
             status = "guardrail_intervened"
             raise

tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py

@@ -2,6 +2,7 @@
 Unit tests for OpenAI Text Completion Guardrail Translation Handler
 """
 
+import json
 from typing import List, Optional, Tuple
 from unittest.mock import MagicMock
 
@@ -367,6 +368,13 @@ class _RecordingKeywordFilter(ContentFilterGuardrail):
         return await super().apply_guardrail(inputs, request_data, input_type, logging_obj)
 
 
+class _PassthroughReplaceGuardrail(CustomGuardrail):
+    async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
+        if any("kumquat" in text for text in inputs.get("texts") or []):
+            self.raise_passthrough_exception(violation_message="replaced kumquat", request_data=request_data)
+        return inputs
+
+
 def _text_chunk(index: int, text: str | None, finish_reason: str | None = None) -> TextCompletionResponse:
     return TextCompletionResponse(
         id="cmpl-stream",
@@ -486,3 +494,19 @@ class TestStreamingOutputProcessing:
         assert isinstance(error_frame, bytes)
         assert error_frame.startswith(b'data: {"error": ')
         assert b"kumquat" in error_frame
+
+    @pytest.mark.asyncio
+    async def test_flat_stream_passthrough_block_appends_the_replacement_and_ends_with_content_filter(self):
+        guardrail = _PassthroughReplaceGuardrail(
+            guardrail_name="passthrough", event_hook=GuardrailEventHooks.post_call, default_on=True
+        )
+        chunks = [_text_chunk(0, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop"), _usage_chunk()]
+
+        delivered = await _drive_flat_stream(guardrail, chunks)
+
+        frames = [json.loads(item.decode().removeprefix("data: ")) for item in delivered if isinstance(item, bytes)]
+        assert [(frame["choices"][0]["text"], frame["choices"][0].get("finish_reason")) for frame in frames] == [
+            ("replaced kumquat", None),
+            ("", "content_filter"),
+        ]
+        assert frames[-1]["usage"]["total_tokens"] == 8

tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py

@@ -30,7 +30,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter impor
     ContentFilterCategoryConfig,
     ContentFilterDetection,
 )
-from litellm.types.utils import StandardLoggingGuardrailInformation
+from litellm.types.utils import StandardLoggingGuardrailInformation, TextChoices, TextCompletionResponse
 
 
 class TestContentFilterGuardrail:
@@ -3523,3 +3523,55 @@ async def test_image_whose_description_is_clean_passes_through_unchanged():
     )
 
     assert result == {"images": ["https://example.com/desk.png"], "texts": []}
+
+
+def _fruit_word_filter() -> ContentFilterGuardrail:
+    return ContentFilterGuardrail(
+        guardrail_name="fruit-filter",
+        blocked_words=[
+            BlockedWord(keyword="persimmon", action=ContentFilterAction.MASK),
+            BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK),
+        ],
+        event_hook=GuardrailEventHooks.post_call,
+    )
+
+
+async def _text_completion_stream(*parts: tuple[str | None, str | None]):
+    for text, finish_reason in parts:
+        yield TextCompletionResponse(
+            id="cmpl-stream",
+            model="gpt-3.5-turbo-instruct",
+            object="text_completion",
+            choices=[TextChoices(index=0, text=text, finish_reason=finish_reason)],
+        )
+
+
+@pytest.mark.asyncio
+async def test_streaming_hook_masks_a_word_split_across_text_completion_chunks():
+    request_data: Final[dict] = {}
+
+    delivered: Final = [
+        chunk
+        async for chunk in _fruit_word_filter().async_post_call_streaming_iterator_hook(
+            user_api_key_dict=MagicMock(),
+            response=_text_completion_stream(("A pers", None), ("immon tart", None), (None, "stop")),
+            request_data=request_data,
+        )
+    ]
+
+    assert "".join(chunk.choices[0].text or "" for chunk in delivered) == "A [KEYWORD_REDACTED] tart"
+    assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_status"] == "success"
+
+
+@pytest.mark.asyncio
+async def test_streaming_hook_blocks_a_word_split_across_text_completion_chunks():
+    with pytest.raises(HTTPException) as exc_info:
+        async for _ in _fruit_word_filter().async_post_call_streaming_iterator_hook(
+            user_api_key_dict=MagicMock(),
+            response=_text_completion_stream(("A kum", None), ("quat tart", None), (None, "stop")),
+            request_data={},
+        ):
+            pass
+
+    assert exc_info.value.status_code == 400
+    assert "kumquat" in str(exc_info.value.detail)

tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py

@@ -889,7 +889,7 @@ async def test_pre_call_payload_unchanged() -> None:
 
 
 @pytest.mark.asyncio
-async def test_post_call_on_a_text_completion_forwards_only_the_key_and_team_metadata() -> None:
+async def test_post_call_never_forwards_the_proxy_key_metadata_of_the_request_metadata() -> None:
     guardrail = _post_call_guardrail()
     client = _CapturingClient()
     guardrail.async_handler = client
@@ -899,11 +899,16 @@ async def test_post_call_on_a_text_completion_forwards_only_the_key_and_team_met
             id="cmpl-1", model="gpt-4.1-mini", choices=[TextChoices(index=0, text="hello", finish_reason="stop")]
         ),
         guardrail_to_apply=guardrail,
-        user_api_key_dict=UserAPIKeyAuth(team_id="team-456", key_alias="ci-key"),
-        request_data={"model": "gpt-4.1-mini", "prompt": "say hello", "metadata": {"requester_note": "private"}},
+        user_api_key_dict=UserAPIKeyAuth(team_id="team-456", key_alias="ci-key", user_email="dev@example.com"),
+        request_data={
+            "model": "gpt-4.1-mini",
+            "prompt": "say hello",
+            "metadata": {
+                "user_api_key_hash": "hashed-key",
+                "user_api_key_user_email": "dev@example.com",
+                "user_api_key_request_route": "/v1/completions",
+            },
+        },
     )
 
-    forwarded = client.calls[0]["json"]["litellm_metadata"]
-    assert forwarded["user_api_key_team_id"] == "team-456"
-    assert forwarded["user_api_key_key_alias"] == "ci-key"
-    assert "requester_note" not in forwarded
+    assert "litellm_metadata" not in client.calls[0]["json"]

tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py

@@ -2730,6 +2730,45 @@ async def test_streaming_iterator_hook_pipeline_modify_response_emits_translated
     assert not any(item is chunk for item in delivered for chunk in chunks)
 
 
+@pytest.mark.asyncio
+async def test_streaming_iterator_hook_pipeline_modify_response_replaces_a_text_completion_stream(
+    proxy_logging, make_user_api_key_auth, monkeypatch
+):
+    seen: Final[Dict[str, Any]] = {}
+    monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen, block=True)])
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
+    pipeline: Final = GuardrailPipeline(
+        mode="post_call",
+        steps=[
+            PipelineStep(
+                guardrail="gr-post",
+                on_pass="allow",
+                on_fail="modify_response",
+                modify_response_message="content policy block",
+            )
+        ],
+    )
+    data: Final = {**_post_call_pipeline_data(stream=True), "prompt": "hi"}
+    data["metadata"]["_guardrail_pipelines"] = [("response-governance", pipeline)]
+    chunks: Final = [_text_completion_chunk("hello "), _text_completion_chunk("world"), _text_completion_chunk("", "stop")]
+    auth: Final = make_user_api_key_auth(request_route="/v1/completions")
+
+    await proxy_logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="text_completion", guardrails_only=True)
+    delivered: Final = [
+        item
+        async for item in proxy_logging.async_post_call_streaming_iterator_hook(
+            user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data
+        )
+    ]
+
+    frames: Final = [json.loads(item.decode().removeprefix("data: ")) for item in delivered]
+    assert seen["count"] == 1
+    assert [(frame["choices"][0]["text"], frame["choices"][0].get("finish_reason")) for frame in frames] == [
+        ("content policy block", None),
+        ("", "content_filter"),
+    ]
+
+
 @pytest.mark.asyncio
 async def test_streaming_iterator_hook_pipeline_delivers_text_rewrite_on_anthropic_sse(
     proxy_logging, make_user_api_key_auth, monkeypatch