perf(content_filter): scan a bounded window per streamed chunk

The streaming post-call hook rescanned the whole accumulated choice buffer on every chunk, so scan cost grew quadratically with output length. Keep a bounded per-choice buffer instead: once it exceeds twice the scan context, drop the head when masking the head and tail separately yields the same output as masking the whole buffer, so no pattern, phrase or exception straddles the cut. Detections from the dropped head are kept and merged, deduplicated, into the final log row

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/constants.py

@@ -364,6 +364,8 @@ GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS: Final = int(
     os.getenv("GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60)
 )
 BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS: Final = 25_000
+CONTENT_FILTER_STREAMING_HOLDBACK_CHARS: Final = 50
+CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS: Final = 512
 DEFAULT_PRESIDIO_ANALYZE_CHUNK_SIZE_BYTES: Final = 500_000
 PRESIDIO_ANALYZE_CHUNK_OVERLAP_CHARS: Final = 4096
 PRESIDIO_ANALYZE_CHUNK_CONCURRENCY: Final = 8

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

@@ -11,6 +11,7 @@ import os
 import re
 import time
 from collections.abc import AsyncGenerator, Coroutine, Mapping, Sequence
+from dataclasses import dataclass, replace
 from datetime import datetime
 from re import Pattern
 from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, cast
@@ -20,7 +21,11 @@ from fastapi import HTTPException
 
 from litellm import Router
 from litellm._logging import verbose_proxy_logger
-from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
+from litellm.constants import (
+    CONTENT_FILTER_STREAMING_HOLDBACK_CHARS,
+    CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS,
+    DEFAULT_MAX_RECURSE_DEPTH,
+)
 from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.types.utils import (
@@ -112,6 +117,14 @@ class _CategoryConfigView(TypedDict):
     category_file: str | None
 
 
+@dataclass(frozen=True, slots=True)
+class _StreamedChoiceState:
+    buffered_text: str = ""
+    yielded_masked_text_len: int = 0
+    committed_detections: tuple[ContentFilterDetection, ...] = ()
+    latest_detections: tuple[ContentFilterDetection, ...] = ()
+
+
 class CategoryFileData(TypedDict, total=False):
     category_name: str
     description: str
@@ -1950,6 +1963,40 @@ class ContentFilterGuardrail(CustomGuardrail):
                 exception_str=exception_str,
             )
 
+    def _trim_streamed_choice_buffer(self, state: _StreamedChoiceState, masked_text: str) -> _StreamedChoiceState:
+        """
+        Bound the per-choice buffer rescanned on every streamed chunk.
+
+        Once the buffer exceeds twice the scan context, drop everything but the last
+        context-sized tail, provided the two halves mask to the same output as the whole
+        (so no match, phrase or exception straddles the cut) and the dropped prefix has
+        already been yielded. Otherwise keep the buffer and retry on the next chunk.
+
+        Detections found in the dropped prefix move to the state's committed detections.
+        """
+        if len(state.buffered_text) <= 2 * CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS:
+            return state
+        head: Final = state.buffered_text[:-CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS]
+        tail: Final = state.buffered_text[-CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS:]
+        head_detections: Final[list[ContentFilterDetection]] = []  # mutable-ok: filled by _filter_single_text
+        try:
+            masked_head: Final = self._filter_single_text(head, detections=head_detections)
+            masked_tail: Final = self._filter_single_text(tail)
+        except Exception:
+            return state
+        if masked_head + masked_tail != masked_text or len(masked_head) > state.yielded_masked_text_len:
+            return state
+        return replace(
+            state,
+            buffered_text=tail,
+            yielded_masked_text_len=state.yielded_masked_text_len - len(masked_head),
+            committed_detections=state.committed_detections + tuple(head_detections),
+        )
+
+    @staticmethod
+    def _merge_detections(detections: Sequence[ContentFilterDetection]) -> tuple[ContentFilterDetection, ...]:
+        return tuple(detection for index, detection in enumerate(detections) if detection not in detections[:index])
+
     async def async_post_call_streaming_iterator_hook(
         self,
         user_api_key_dict: UserAPIKeyAuth,
@@ -1968,10 +2015,7 @@ class ContentFilterGuardrail(CustomGuardrail):
         and the UI Request Lifecycle panel. Mirrors apply_guardrail's finally-block
         contract.
         """
-        accumulated_text_by_choice: Final[dict[int, str]] = {}
-        yielded_masked_text_len_by_choice: Final[dict[int, int]] = {}
-        latest_detections_by_choice: Final[dict[int, list[ContentFilterDetection]]] = {}
-        buffer_size: Final = 50  # Increased buffer to catch patterns split across many chunks
+        state_by_choice: Final[dict[int, _StreamedChoiceState]] = {}
 
         start_time: Final = datetime.now()
         scan_seconds: float = 0.0  # rebind-ok: accumulates per-chunk scan time across the stream
@@ -1997,69 +2041,57 @@ class ContentFilterGuardrail(CustomGuardrail):
 
                         content = getattr(choice.delta, "content", None)
                         is_final = bool(getattr(choice, "finish_reason", None))
-                        if isinstance(content, str) and content:
-                            accumulated_text_by_choice[choice_index] = (
-                                accumulated_text_by_choice.get(choice_index, "") + content
-                            )
-                        elif not is_final:
+                        new_content = content if isinstance(content, str) else ""
+                        if not new_content and not is_final:
                             continue
 
-                        text_to_check = accumulated_text_by_choice.get(choice_index, "")
-                        if not text_to_check:
+                        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 = text_to_check + (" " if is_final else "")
+                        text_to_scan = buffered_text + (" " if is_final else "")
                         choice_detections: list[ContentFilterDetection] = []
                         scan_started = time.perf_counter()
 
                         try:
-                            # _filter_single_text scans the whole accumulated
-                            # choice buffer every chunk, so previous-chunk
-                            # matches are guaranteed to be re-found. Keeping
-                            # only each choice's latest scan avoids duplicate
-                            # detections in the final log row.
                             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_by_choice[choice_index] = choice_detections
+                            latest_detections = tuple(choice_detections)
                         except HTTPException:
-                            latest_detections_by_choice[choice_index] = choice_detections
+                            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
 
-                        # Determine how much can be safely yielded
+                        safe_to_yield_len = max(
+                            previous_state.yielded_masked_text_len,
+                            len(masked_text) - (0 if is_final else CONTENT_FILTER_STREAMING_HOLDBACK_CHARS),
+                        )
+                        choice.delta.content = masked_text[previous_state.yielded_masked_text_len : safe_to_yield_len]
+                        next_state = _StreamedChoiceState(
+                            buffered_text, safe_to_yield_len, previous_state.committed_detections, latest_detections
+                        )
                         if is_final:
-                            safe_to_yield_len = len(masked_text)
-                        else:
-                            safe_to_yield_len = max(0, len(masked_text) - buffer_size)
+                            state_by_choice[choice_index] = next_state
+                            continue
 
-                        yielded_masked_text_len = yielded_masked_text_len_by_choice.get(choice_index, 0)
-                        if safe_to_yield_len > yielded_masked_text_len:
-                            new_masked_content = masked_text[yielded_masked_text_len:safe_to_yield_len]
-                            choice.delta.content = new_masked_content
-                            yielded_masked_text_len_by_choice[choice_index] = safe_to_yield_len
-                        else:
-                            # Hold content by yielding empty content on this choice
-                            # while preserving chunk metadata and other choices.
-                            choice.delta.content = ""
+                        trim_started = time.perf_counter()
+                        state_by_choice[choice_index] = self._trim_streamed_choice_buffer(next_state, masked_text)
+                        scan_seconds += time.perf_counter() - trim_started
 
                     yield item
                 else:
                     # Not a ModelResponseStream or no choices - yield as is
                     yield item
-
-            # Any remaining content (should have been handled by is_final, but just in case)
-            if any(
-                yielded_masked_text_len_by_choice.get(choice_index, 0) < len(accumulated_text)
-                for choice_index, accumulated_text in accumulated_text_by_choice.items()
-            ):
-                # We already reached the end of the generator
-                pass
         except HTTPException:
             status = "guardrail_intervened"
             raise
@@ -2070,8 +2102,8 @@ class ContentFilterGuardrail(CustomGuardrail):
         finally:
             detections = [
                 detection
-                for choice_detections in latest_detections_by_choice.values()
-                for detection in choice_detections
+                for state in state_by_choice.values()
+                for detection in self._merge_detections((*state.committed_detections, *state.latest_detections))
             ]
             self._count_masked_entities(detections, masked_entity_count)
             self._log_guardrail_information(

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

@@ -11,6 +11,10 @@ import pytest
 
 from fastapi import HTTPException
 
+from litellm.constants import (
+    CONTENT_FILTER_STREAMING_HOLDBACK_CHARS,
+    CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS,
+)
 from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
     ContentFilterGuardrail,
 )
@@ -22,6 +26,7 @@ from litellm.types.guardrails import (
 )
 from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import (
     ContentFilterCategoryConfig,
+    ContentFilterDetection,
 )
 
 
@@ -900,6 +905,178 @@ class TestContentFilterGuardrail:
         # masked_entity_count for email is the real count, not N×.
         assert entry["masked_entity_count"].get("email") == 1
 
+    @staticmethod
+    async def _collect_streamed_text(
+        guardrail: ContentFilterGuardrail, chunks: list[str], request_data: dict
+    ) -> str:
+        from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
+
+        async def mock_stream():
+            for i, content in enumerate(chunks):
+                yield ModelResponseStream(
+                    id=f"c{i}",
+                    choices=[StreamingChoices(delta=Delta(content=content), index=0)],
+                    model="gpt-4",
+                )
+            yield ModelResponseStream(
+                id="final",
+                choices=[
+                    StreamingChoices(
+                        delta=Delta(content=""), index=0, finish_reason="stop"
+                    )
+                ],
+                model="gpt-4",
+            )
+
+        yielded = []
+        async for chunk in guardrail.async_post_call_streaming_iterator_hook(
+            user_api_key_dict=MagicMock(),
+            response=mock_stream(),
+            request_data=request_data,
+        ):
+            yielded.append(chunk.choices[0].delta.content or "")
+        return "".join(yielded)
+
+    @pytest.mark.asyncio
+    async def test_streaming_hook_scans_bounded_window_per_chunk(self):
+        """
+        Regression: the streaming hook used to re-scan the whole accumulated
+        buffer on every chunk, so scan work grew quadratically with the length
+        of the response. Each scan must now cover only the new chunk plus a
+        bounded tail of what came before, without dropping any output.
+        """
+        scanned_lengths: list[int] = []
+
+        class RecordingGuardrail(ContentFilterGuardrail):
+            def _filter_single_text(
+                self,
+                text: str,
+                detections: list[ContentFilterDetection] | None = None,
+            ) -> str:
+                scanned_lengths.append(len(text))
+                return super()._filter_single_text(text, detections=detections)
+
+        guardrail = RecordingGuardrail(
+            guardrail_name="test-streaming-bounded-scan",
+            patterns=[
+                ContentFilterPattern(
+                    pattern_type="prebuilt",
+                    pattern_name="email",
+                    action=ContentFilterAction.MASK,
+                )
+            ],
+            event_hook=GuardrailEventHooks.post_call,
+        )
+        chunk = "Item: a plain household object description. "
+        chunks = [chunk] * 200
+        request_data = {"messages": [], "model": "gpt-4o", "metadata": {}}
+
+        streamed = await self._collect_streamed_text(guardrail, chunks, request_data)
+
+        assert streamed == chunk * 200
+        assert len(chunk) * 200 > 4 * CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS
+        window_bound = 2 * CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS + len(chunk) + 1
+        assert max(scanned_lengths) <= window_bound, (
+            f"scan input grew to {max(scanned_lengths)} chars for a "
+            f"{len(chunk)}-char chunk; expected at most {window_bound}"
+        )
+
+    @pytest.mark.asyncio
+    async def test_streaming_hook_blocks_match_longer_than_holdback_across_chunks(
+        self,
+    ):
+        """
+        A blocked phrase longer than the holdback window arrives in small chunks,
+        so its start has already been yielded before its end shows up. The scan
+        still has to see the whole phrase and block.
+        """
+        phrase = "alpha bravo charlie delta echo foxtrot golf hotel india juliet kilo lima"
+        assert len(phrase) > CONTENT_FILTER_STREAMING_HOLDBACK_CHARS
+        guardrail = ContentFilterGuardrail(
+            guardrail_name="test-streaming-long-block",
+            blocked_words=[BlockedWord(keyword=phrase, action=ContentFilterAction.BLOCK)],
+            event_hook=GuardrailEventHooks.post_call,
+        )
+        text = "Here is the codeword list: " + phrase + " and that is all."
+        chunks = [text[i : i + 4] for i in range(0, len(text), 4)]
+        request_data = {"messages": [], "model": "gpt-4o", "metadata": {}}
+
+        with pytest.raises(HTTPException) as exc_info:
+            await self._collect_streamed_text(guardrail, chunks, request_data)
+
+        assert exc_info.value.detail["keyword"] == phrase
+        entry = request_data["metadata"]["standard_logging_guardrail_information"][0]
+        assert entry["guardrail_status"] == "guardrail_intervened"
+        assert [d["keyword"] for d in entry["guardrail_response"]] == [phrase]
+
+    @pytest.mark.asyncio
+    async def test_streaming_hook_masks_every_email_in_long_stream_and_logs_once(
+        self,
+    ):
+        """
+        A response made of nothing but emails, several times longer than the
+        rescanned buffer, must come out as nothing but redaction tags, and the log
+        must carry one email detection, matching what a single scan of the full
+        text reports. Wherever the buffer is cut, an email sits on the cut, so
+        dropping text without checking that the cut leaves the masked output
+        unchanged corrupts the stream.
+        """
+        guardrail = ContentFilterGuardrail(
+            guardrail_name="test-streaming-many-emails",
+            patterns=[
+                ContentFilterPattern(
+                    pattern_type="prebuilt",
+                    pattern_name="email",
+                    action=ContentFilterAction.MASK,
+                )
+            ],
+            event_hook=GuardrailEventHooks.post_call,
+        )
+        emails = [f"user{i:03d}@example.com" for i in range(200)]
+        text = " ".join(emails)
+        assert len(text) > 4 * CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS
+        chunks = [text[i : i + 3] for i in range(0, len(text), 3)]
+        request_data = {"messages": [], "model": "gpt-4o", "metadata": {}}
+
+        streamed = await self._collect_streamed_text(guardrail, chunks, request_data)
+
+        assert streamed == " ".join(["[EMAIL_REDACTED]"] * len(emails))
+        entry = request_data["metadata"]["standard_logging_guardrail_information"][0]
+        assert entry["guardrail_status"] == "success"
+        assert [d["pattern_name"] for d in entry["guardrail_response"]] == ["email"]
+        assert entry["masked_entity_count"] == {"email": 1}
+
+    @pytest.mark.asyncio
+    async def test_streaming_hook_logs_detection_masked_long_before_stream_end(self):
+        """
+        An email at the start of a long response is masked and then falls out of
+        the rescanned buffer well before the stream ends. The final log entry must
+        still report it, as a scan of the full text would.
+        """
+        guardrail = ContentFilterGuardrail(
+            guardrail_name="test-streaming-early-detection",
+            patterns=[
+                ContentFilterPattern(
+                    pattern_type="prebuilt",
+                    pattern_name="email",
+                    action=ContentFilterAction.MASK,
+                )
+            ],
+            event_hook=GuardrailEventHooks.post_call,
+        )
+        filler = "filler text " * CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS
+        text = f"Contact one@example.com for details. {filler}"
+        chunks = [text[i : i + 40] for i in range(0, len(text), 40)]
+        request_data = {"messages": [], "model": "gpt-4o", "metadata": {}}
+
+        streamed = await self._collect_streamed_text(guardrail, chunks, request_data)
+
+        assert streamed == text.replace("one@example.com", "[EMAIL_REDACTED]")
+        entry = request_data["metadata"]["standard_logging_guardrail_information"][0]
+        assert entry["guardrail_status"] == "success"
+        assert [d["pattern_name"] for d in entry["guardrail_response"]] == ["email"]
+        assert entry["masked_entity_count"] == {"email": 1}
+
     def test_init_with_plain_dicts(self):
         """
         Test initialization with plain dicts (DB format).