From 60642e875bff8838d452d3db29e017df4d3651b4 Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 16 Sep 2026 11:51:06 +0000 Subject: [PATCH] perf(content_filter): back off refused streamed buffer cuts by one context length Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_content_filter/content_filter.py | 20 ++++++--- .../content_filter/test_content_filter.py | 43 +++++++++++++++++++ 2 files changed, 56 insertions(+), 7 deletions(-) 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 111f1d4d59a..d4edf618d1b 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 @@ -124,6 +124,7 @@ class _StreamedChoiceState: yielded_masked_text_len: int = 0 committed_detections: tuple[ContentFilterDetection, ...] = () latest_detections: tuple[ContentFilterDetection, ...] = () + next_trim_len: int = 0 @dataclass(frozen=True, slots=True) @@ -2013,29 +2014,31 @@ class ContentFilterGuardrail(CustomGuardrail): context-sized tail, provided no exception phrase or unfinished conditional sentence would leave the buffer, the two halves mask to the same output as the whole (so no match or phrase straddles the cut), and the dropped prefix has already been yielded. - Otherwise keep the buffer and retry on the next chunk. + Otherwise keep the buffer and retry once it has grown by another context length. Detections found in the dropped prefix move to the state's committed detections. """ - if len(state.buffered_text) <= 2 * plan.context_chars: + if len(state.buffered_text) <= max(2 * plan.context_chars, state.next_trim_len): return state + deferred: Final = replace(state, next_trim_len=len(state.buffered_text) + plan.context_chars) head: Final = state.buffered_text[: -plan.context_chars] tail: Final = state.buffered_text[-plan.context_chars :] if self._cut_breaks_wider_context(state.buffered_text, head, tail, plan): - return state + return deferred 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 + return deferred if masked_head + masked_tail != masked_text or len(masked_head) > state.yielded_masked_text_len: - return state + return deferred 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), + next_trim_len=0, ) @staticmethod @@ -2123,8 +2126,11 @@ class ContentFilterGuardrail(CustomGuardrail): 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 + 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 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 a35f6f26983..ca0f9c474a8 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 @@ -981,6 +981,49 @@ class TestContentFilterGuardrail: f"{len(chunk)}-char chunk; expected at most {window_bound}" ) + @pytest.mark.asyncio + async def test_streaming_hook_retries_refused_cut_once_per_context_length(self): + """ + A single URL that keeps growing crosses every proposed cut, so no cut is + ever safe. The trim check must then back off instead of adding two extra + scans on every chunk, and the whole URL must still come out masked. + """ + 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-refused-cut-backoff", + patterns=[ + ContentFilterPattern( + pattern_type="prebuilt", + pattern_name="url", + action=ContentFilterAction.MASK, + ) + ], + event_hook=GuardrailEventHooks.post_call, + ) + text = "See https://example.com/" + "a" * (8 * CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS) + " now." + chunks = [text[i : i + 16] for i in range(0, len(text), 16)] + request_data = {"messages": [], "model": "gpt-4o", "metadata": {}} + + streamed = await self._collect_streamed_text(guardrail, chunks, request_data) + streamed_scans = len(scanned_lengths) + + full_scan = await guardrail.apply_guardrail(inputs={"texts": [text]}, request_data={}, input_type="response") + assert streamed == full_scan["texts"][0] == "See [URL_REDACTED] now." + extra_scans = streamed_scans - len(chunks) + assert extra_scans <= 2 * (len(text) // CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS), ( + f"{extra_scans} scans beyond one per chunk for {len(chunks)} chunks; the refused cut must back off" + ) + @pytest.mark.asyncio async def test_streaming_hook_blocks_match_longer_than_holdback_across_chunks( self,