From 9bb31caa28843b734a694c1aa343ac0363902a7a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 22:56:49 -0700 Subject: [PATCH] fix(guardrails): release an unfinished choice's held text and omit usage the upstream never sent The content filter's streaming hook held back the end of a choice until its finish_reason, which litellm sends only for the first choice of an n>1 stream, so choice 1 lost its last characters on text and chat streams. A closing chunk now releases every unfinished choice when the stream ends. A blocked text stream's closing chunk reported 0/0 usage when the upstream sent none; it now carries usage only when the upstream reported it. --- .../guardrail_translation/handler.py | 21 +++--- .../litellm_content_filter/content_filter.py | 53 ++++++++++++-- .../test_text_completion_guardrail_handler.py | 13 ++++ .../content_filter/test_content_filter.py | 69 ++++++++++++++++++- 4 files changed, 142 insertions(+), 14 deletions(-) diff --git a/litellm/llms/openai/completion/guardrail_translation/handler.py b/litellm/llms/openai/completion/guardrail_translation/handler.py index 5d86ddec014..54a579401aa 100644 --- a/litellm/llms/openai/completion/guardrail_translation/handler.py +++ b/litellm/llms/openai/completion/guardrail_translation/handler.py @@ -270,26 +270,20 @@ class OpenAITextCompletionHandler(BaseTranslation): ) -> 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)], - ) + ).model_copy(update={"usage": None}) 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, - ), - ) + ).model_copy(update={"usage": _blocked_stream_usage(exc.original_response)}) return tuple( f"data: {chunk.model_dump_json(exclude_none=True)}\n\n".encode() for chunk in (message_chunk, final_chunk) ) @@ -347,6 +341,17 @@ class OpenAITextCompletionHandler(BaseTranslation): choice.text = "" +def _blocked_stream_usage(original_response: object) -> Usage | None: + prompt_tokens, completion_tokens = blocked_chat_stream_usage(original_response) + if isinstance(original_response, list) and prompt_tokens == completion_tokens == 0: + return None + return Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) + + def _text_completion_chunks(responses_so_far: Sequence[object]) -> tuple[TextCompletionResponse, ...]: return tuple(item for item in responses_so_far if isinstance(item, TextCompletionResponse)) 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 76edb2cf59c..9217b8ba9e6 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 @@ -12,11 +12,11 @@ import json import os import re import time -from collections.abc import AsyncGenerator, Coroutine, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, Callable, 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 +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, TypeVar, cast import yaml from fastapi import HTTPException @@ -40,6 +40,7 @@ from litellm.proxy.guardrails.content_filter_data import ( ) from litellm.types.utils import ( CallTypes, + Delta, Function, GenericGuardrailAPIInputs, GuardrailStatus, @@ -81,6 +82,7 @@ from .patterns import PATTERN_EXTRA_CONFIG, get_compiled_pattern MAX_KEYWORD_VALUE_GAP_WORDS: Final = 1 _DECODED_ARGUMENTS: Final = TypeAdapter(object) +_StreamItemT: Final = TypeVar("_StreamItemT") _JSON_ENTRIES: Final = TypeAdapter(list[object]) GAP_WORD_TOKENIZER: Final = re.compile(r"\b\w+\b") SENTENCE_TERMINATORS: Final = re.compile(r"[.!?]+") @@ -165,6 +167,43 @@ class _StreamedChoiceState: committed_detections: tuple[ContentFilterDetection, ...] = () latest_detections: tuple[ContentFilterDetection, ...] = () next_trim_len: int = 0 + finished: bool = False + + +def _end_of_stream_chunk( + last_item: object, state_by_choice: Mapping[int, _StreamedChoiceState] +) -> ModelResponseStream | TextCompletionResponse | None: + unfinished: Final = tuple(index for index, state in state_by_choice.items() if not state.finished) + if not unfinished: + return None + if isinstance(last_item, TextCompletionResponse): + return TextCompletionResponse( + id=last_item.id, + created=last_item.created, + model=last_item.model, + object="text_completion", + choices=[TextChoices(index=index, text="") for index in unfinished], + ) + if isinstance(last_item, ModelResponseStream): + return ModelResponseStream( + id=last_item.id, + created=last_item.created, + model=last_item.model, + choices=[StreamingChoices(index=index, delta=Delta(content="")) for index in unfinished], + ) + return None + + +async def _stream_then_end_of_stream_chunk( + response: AsyncIterable[_StreamItemT], end_of_stream_chunk: Callable[[_StreamItemT], _StreamItemT | None] +) -> AsyncGenerator[tuple[_StreamItemT, bool], None]: + last_item: _StreamItemT | None = None # rebind-ok: the end-of-stream chunk copies the last chunk's id and model + async for item in response: + last_item = item + yield item, False + closing_item: Final = None if last_item is None else end_of_stream_chunk(last_item) + if closing_item is not None: + yield closing_item, True @dataclass(frozen=True, slots=True) @@ -2123,7 +2162,9 @@ class ContentFilterGuardrail(CustomGuardrail): ) try: - async for item in response: + async for item, stream_ended in _stream_then_end_of_stream_chunk( + response, lambda last_item: _end_of_stream_chunk(last_item, state_by_choice) + ): for choice in _streamed_choices(item): new_content = _streamed_choice_content(choice) if new_content is None: @@ -2133,7 +2174,7 @@ class ContentFilterGuardrail(CustomGuardrail): if not isinstance(choice_index, int): choice_index = 0 - is_final = bool(getattr(choice, "finish_reason", None)) + is_final = stream_ended or bool(getattr(choice, "finish_reason", None)) if not new_content and not is_final: continue @@ -2180,13 +2221,15 @@ class ContentFilterGuardrail(CustomGuardrail): latest_detections=latest_detections, ) if is_final: - state_by_choice[choice_index] = next_state + state_by_choice[choice_index] = replace(next_state, finished=True) continue 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 + if stream_ended and not any(_streamed_choice_content(choice) for choice in _streamed_choices(item)): + continue yield item except HTTPException: status = "guardrail_intervened" diff --git a/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py b/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py index cc213cf35f5..dd2bd6f0f13 100644 --- a/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py +++ b/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py @@ -510,3 +510,16 @@ class TestStreamingOutputProcessing: ("", "content_filter"), ] assert frames[-1]["usage"]["total_tokens"] == 8 + + @pytest.mark.asyncio + async def test_flat_stream_passthrough_block_reports_no_usage_the_upstream_never_sent(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")] + + delivered = await _drive_flat_stream(guardrail, [chunk.model_copy(update={"usage": None}) for chunk in chunks]) + + frames = [json.loads(item.decode().removeprefix("data: ")) for item in delivered if isinstance(item, bytes)] + assert frames[-1]["choices"][0]["finish_reason"] == "content_filter" + assert [frame for frame in frames if "usage" in frame] == [] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index 169cd8c0651..43b5b121e18 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -30,7 +30,15 @@ from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter impor ContentFilterCategoryConfig, ContentFilterDetection, ) -from litellm.types.utils import StandardLoggingGuardrailInformation, TextChoices, TextCompletionResponse +from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field +from litellm.types.utils import ( + Delta, + ModelResponseStream, + StandardLoggingGuardrailInformation, + StreamingChoices, + TextChoices, + TextCompletionResponse, +) class TestContentFilterGuardrail: @@ -3563,6 +3571,65 @@ async def test_streaming_hook_masks_a_word_split_across_text_completion_chunks() assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_status"] == "success" +_LONG_TAIL: Final = "the quick brown fox jumps over the lazy dog and then eats a ripe persimmon" + + +async def _text_completion_n2_stream_without_choice_1_finish(): + for index, text, finish_reason in ( + (0, "A short answer", None), + (1, _LONG_TAIL[:30], None), + (1, _LONG_TAIL[30:], None), + (0, None, "stop"), + ): + yield TextCompletionResponse( + id="cmpl-stream", + model="gpt-3.5-turbo-instruct", + object="text_completion", + choices=[TextChoices(index=index, text=text, finish_reason=finish_reason)], + ) + + +async def _chat_n2_stream_without_choice_1_finish(): + for index, content, finish_reason in ( + (0, "A short answer", None), + (1, _LONG_TAIL[:30], None), + (1, _LONG_TAIL[30:], None), + (0, None, "stop"), + ): + yield ModelResponseStream( + id="chatcmpl-stream", + model="gpt-4.1-mini", + choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)], + ) + + +def _streamed_text(chunk: ModelResponseStream | TextCompletionResponse, index: int) -> str: + return "".join(_choice_text(choice) for choice in chunk.choices if choice.index == index) + + +def _choice_text(choice: TextChoices | StreamingChoices) -> str: + content: Final = stream_item_field(choice, "text") if isinstance(choice, TextChoices) else choice.delta.content + return content if isinstance(content, str) else "" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "stream", [_text_completion_n2_stream_without_choice_1_finish, _chat_n2_stream_without_choice_1_finish] +) +async def test_streaming_hook_releases_the_held_tail_of_a_choice_the_stream_never_finished(stream): + delivered: Final = [ + chunk + async for chunk in _fruit_word_filter().async_post_call_streaming_iterator_hook( + user_api_key_dict=MagicMock(), response=stream(), request_data={} + ) + ] + + assert "".join(_streamed_text(chunk, 0) for chunk in delivered) == "A short answer" + assert "".join(_streamed_text(chunk, 1) for chunk in delivered) == _LONG_TAIL.replace( + "persimmon", "[KEYWORD_REDACTED]" + ) + + @pytest.mark.asyncio async def test_streaming_hook_blocks_a_word_split_across_text_completion_chunks(): with pytest.raises(HTTPException) as exc_info: