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 1ffe200169a..a210762d1df 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 @@ -1963,7 +1963,17 @@ class ContentFilterGuardrail(CustomGuardrail): exception_str=exception_str, ) - def _trim_streamed_choice_buffer(self, state: _StreamedChoiceState, masked_text: str) -> _StreamedChoiceState: + def _streamed_scan_context_chars(self) -> int: + """Retained tail length: the default context, widened to the longest configured keyword.""" + longest_keyword: Final = max( + map(len, (*self.blocked_words, *self.category_keywords, *self.always_block_category_keywords)), + default=0, + ) + return max(CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS, longest_keyword) + + def _trim_streamed_choice_buffer( + self, state: _StreamedChoiceState, masked_text: str, scan_context_chars: int + ) -> _StreamedChoiceState: """ Bound the per-choice buffer rescanned on every streamed chunk. @@ -1974,10 +1984,10 @@ class ContentFilterGuardrail(CustomGuardrail): 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: + if len(state.buffered_text) <= 2 * 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: Final = state.buffered_text[:-scan_context_chars] + tail: Final = state.buffered_text[-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) @@ -2016,6 +2026,7 @@ class ContentFilterGuardrail(CustomGuardrail): contract. """ state_by_choice: Final[dict[int, _StreamedChoiceState]] = {} + scan_context_chars: Final = self._streamed_scan_context_chars() start_time: Final = datetime.now() scan_seconds: float = 0.0 # rebind-ok: accumulates per-chunk scan time across the stream @@ -2085,7 +2096,9 @@ class ContentFilterGuardrail(CustomGuardrail): continue trim_started = time.perf_counter() - state_by_choice[choice_index] = self._trim_streamed_choice_buffer(next_state, masked_text) + state_by_choice[choice_index] = self._trim_streamed_choice_buffer( + next_state, masked_text, scan_context_chars + ) scan_seconds += time.perf_counter() - trim_started yield item 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 71f46ca3047..b91abb98dbf 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 @@ -1009,6 +1009,33 @@ class TestContentFilterGuardrail: assert entry["guardrail_status"] == "guardrail_intervened" assert [d["keyword"] for d in entry["guardrail_response"]] == [phrase] + @pytest.mark.asyncio + async def test_streaming_hook_blocks_keyword_longer_than_scan_context(self): + """ + A blocked keyword longer than the default retained context arrives after + enough text that the buffer has already been trimmed at least once. The + retained tail must be wide enough that the keyword's start is still in the + buffer when its end arrives, so the stream is blocked. + """ + phrase = " ".join(f"token{i:03d}" for i in range(80)) + assert len(phrase) > CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS + guardrail = ContentFilterGuardrail( + guardrail_name="test-streaming-keyword-wider-than-context", + blocked_words=[BlockedWord(keyword=phrase, action=ContentFilterAction.BLOCK)], + event_hook=GuardrailEventHooks.post_call, + ) + filler = "plain filler sentence. " * (3 * CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS // 23) + text = filler + phrase + " and that is all." + chunks = [text[i : i + 16] for i in range(0, len(text), 16)] + 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" + @pytest.mark.asyncio async def test_streaming_hook_masks_every_email_in_long_stream_and_logs_once( self,