diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index f1c7451796d..5c30ff4747a 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -17,7 +17,10 @@ from typing import TYPE_CHECKING, Any, Final, Optional from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, +) from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -313,9 +316,14 @@ class A2AGuardrailHandler(BaseTranslation): return responses_so_far + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + _, valid_parsed = self._parse_streaming_responses(responses_so_far) + combined_text, _ = self._collect_text_from_parsed_chunks(valid_parsed) + return StreamingScanKey(texts=(combined_text,)) + def _parse_streaming_responses( self, - responses_so_far: list[object], + responses_so_far: Sequence[object], ) -> tuple[list[dict[str, object] | None], list[tuple[int, dict[str, object]]]]: """Parse JSON-RPC items, returning aligned parsed list and valid entries.""" parsed: Final[list[dict[str, object] | None]] = [None] * len(responses_so_far) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index c7d12e5cf3a..c23797f72af 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -26,7 +26,10 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im LiteLLMAnthropicMessagesAdapter, is_provider_native_tool_dict, ) -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, +) from litellm.llms.base_llm.guardrail_translation.utils import ( anthropic_tool_name, anthropic_tool_names, @@ -36,6 +39,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( merge_guardrailed_scoped_messages, merge_returned_tools_into_request_tools, scoped_structured_message_indices, + stream_item_fingerprint, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, @@ -1176,6 +1180,25 @@ class AnthropicMessagesHandler(BaseTranslation): inputs["model"] = response_model return inputs + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + stream_ended: Final = self._check_streaming_has_ended(responses_so_far) + return StreamingScanKey( + texts=(self.get_streaming_string_so_far(responses_so_far),), + tool_calls=self._streamed_tool_use_fingerprints(responses_so_far) if stream_ended else (), + stream_ended=stream_ended, + ) + + @classmethod + def _streamed_tool_use_fingerprints(cls, responses_so_far: Sequence[object]) -> tuple[str, ...]: + return tuple( + stream_item_fingerprint(block) + for item in responses_so_far + for event in cls._iter_sse_events(item) + if event.get("type") == "content_block_start" + and isinstance(block := event.get("content_block"), Mapping) + and block.get("type") == "tool_use" + ) + def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str: """ Parse streaming responses and extract accumulated text content. diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 220fcedb0f8..b28daf73bc4 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -35,6 +35,22 @@ class StreamTransformSink: holdback_per_choice: dict[int, int] = field(default_factory=dict) +@dataclass(frozen=True, slots=True) +class StreamingScanKey: + """What a streaming guardrail round would hand to ``apply_guardrail``. Two keys + compare equal when the round would scan the same content again; ``stream_ended`` + stays out of the comparison and only says whether the handler is on its + end-of-stream path, where an empty payload is still scanned today.""" + + texts: tuple[str, ...] + tool_calls: tuple[str, ...] = () + stream_ended: bool = field(default=False, compare=False) + + @property + def has_nothing_to_scan(self) -> bool: + return not self.stream_ended and not any(self.texts) and not self.tool_calls + + class BaseTranslation(ABC): @staticmethod def transform_user_api_key_dict_to_metadata( @@ -151,6 +167,9 @@ class BaseTranslation(ABC): """ return responses_so_far + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + return None + def build_block_sse_chunks( self, exc: "ModifyResponseException", diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 9b6f9c47105..8dee262001d 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -4,6 +4,8 @@ import json from collections.abc import Callable, Iterator, Sequence from typing import Any, Final, TypeVar +from pydantic import BaseModel + from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage @@ -130,6 +132,16 @@ def stream_item_field(item: object, field: str) -> object | None: return getattr(item, field, None) +def stream_item_fingerprint(item: object) -> str: + plain: Final = item.model_dump() if isinstance(item, BaseModel) else item + return json.dumps(plain, sort_keys=True, default=str) + + +def stream_item_items(item: object, field: str) -> tuple[object, ...]: + value: Final = stream_item_field(item, field) + return tuple(value) if isinstance(value, (list, tuple)) else () + + def blocked_chat_stream_usage(original_response: object) -> tuple[int, int]: """ ``(prompt_tokens, completion_tokens)`` for a synthetic guardrail-blocked diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 96a5ed663fc..d41c8557d72 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -26,6 +26,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, + StreamingScanKey, StreamTransformSink, ) from litellm.llms.base_llm.guardrail_translation.utils import ( @@ -39,6 +40,8 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( role_out_of_guardrail_scope, scoped_structured_message_indices, stream_item_field, + stream_item_fingerprint, + stream_item_items, ) from litellm.main import stream_chunk_builder from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -503,12 +506,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): """Block-only streaming path: run the guardrail so an in-flight BLOCK can terminate the stream. Text rewrites are not propagated to the client here (see ``_process_streaming_transform`` for the incremental_diff path).""" - # check if the stream has ended - has_stream_ended = False - for chunk in responses_so_far: - if chunk.choices and chunk.choices[0].finish_reason is not None: - has_stream_ended = True - break + has_stream_ended: Final = self._first_choice_has_finished(responses_so_far) if has_stream_ended: # convert to model response @@ -706,8 +704,33 @@ class OpenAIChatCompletionsHandler(BaseTranslation): indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback) } + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream)) + stream_ended: Final = self._first_choice_has_finished(responses_so_far) + return StreamingScanKey( + texts=tuple(self._combine_streaming_texts(chunks).values()), + tool_calls=self._streamed_tool_call_fingerprints(responses_so_far) if stream_ended else (), + stream_ended=stream_ended, + ) + + @staticmethod + def _streamed_tool_call_fingerprints(responses_so_far: Sequence[object]) -> tuple[str, ...]: + return tuple( + stream_item_fingerprint(tool_call) + for chunk in responses_so_far + for choice in _stream_chunk_choices(chunk) + for tool_call in stream_item_items(stream_item_field(choice, "delta"), "tool_calls") + ) + + @staticmethod + def _first_choice_has_finished(responses_so_far: Sequence[object]) -> bool: + first_choices: Final = tuple( + choices[0] for choices in (_stream_chunk_choices(chunk) for chunk in responses_so_far) if choices + ) + return any(stream_item_field(choice, "finish_reason") is not None for choice in first_choices) + def _combine_streaming_texts( - self, responses_so_far: list["ModelResponseStream"] + self, responses_so_far: Sequence["ModelResponseStream"] ) -> dict[tuple[int, int | None], str]: """ Combine all streaming chunks into complete text per choice. diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 5a5970fb867..a0db7aadb9e 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -44,10 +44,15 @@ from litellm._logging import verbose_proxy_logger from litellm.completion_extras.litellm_responses_transformation.transformation import ( OpenAiResponsesToChatCompletionStreamIterator, ) -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, +) from litellm.llms.base_llm.guardrail_translation.utils import ( blocked_responses_stream_usage, stream_item_field, + stream_item_fingerprint, + stream_item_items, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools from litellm.responses.litellm_completion_transformation.transformation import ( @@ -593,18 +598,55 @@ class OpenAIResponsesHandler(BaseTranslation): ) return responses_so_far - def _check_streaming_has_ended(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> bool: + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: """ Check if the streaming has ended. """ if not responses_so_far: return False - terminal_types: Final = { - ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, - ResponsesAPIStreamEvents.RESPONSE_FAILED.value, - ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value, - } - return responses_so_far[-1].get("type") in terminal_types + terminal_types: Final = frozenset( + ( + ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, + ResponsesAPIStreamEvents.RESPONSE_FAILED.value, + ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value, + ) + ) + return stream_item_field(responses_so_far[-1], "type") in terminal_types + + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + if not responses_so_far or not hasattr(responses_so_far[-1], "get"): + return None + last_event: Final = responses_so_far[-1] + last_event_type: Final = stream_item_field(last_event, "type") + if last_event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE.value: + return None + if last_event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value: + return self._completed_response_scan_key(stream_item_field(last_event, "response")) + return StreamingScanKey( + texts=(self.get_streaming_string_so_far(responses_so_far),), + stream_ended=self._check_streaming_has_ended(responses_so_far), + ) + + @staticmethod + def _completed_response_scan_key(response: object) -> StreamingScanKey: + output_items: Final = stream_item_items(response, "output") + message_items: Final = tuple( + item for item in output_items if stream_item_field(item, "type") != "function_call" + ) + return StreamingScanKey( + texts=tuple( + text + for item in message_items + for part in stream_item_items(item, "content") + if isinstance(text := stream_item_field(part, "text"), str) and text + ), + tool_calls=tuple( + stream_item_fingerprint(item) + for item in output_items + if stream_item_field(item, "type") == "function_call" + ), + stream_ended=True, + ) def build_stream_error_items( self, @@ -629,7 +671,7 @@ class OpenAIResponsesHandler(BaseTranslation): ), ) - def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str: + def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str: """ Get the string so far from the responses so far. @@ -641,12 +683,16 @@ class OpenAIResponsesHandler(BaseTranslation): """ keyed_events: Final = tuple( ( - (event.get("item_id"), event.get("output_index"), event.get("content_index")), - event.get("text"), - event.get("delta"), + ( + stream_item_field(event, "item_id"), + stream_item_field(event, "output_index"), + stream_item_field(event, "content_index"), + ), + stream_item_field(event, "text"), + stream_item_field(event, "delta"), ) for event in responses_so_far - if isinstance(event.get("text"), str) or isinstance(event.get("delta"), str) + if isinstance(stream_item_field(event, "text"), str) or isinstance(stream_item_field(event, "delta"), str) ) def part_text(part_key: tuple[object, object, object]) -> str: diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 46b00829b74..c6b8df1b493 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -36,6 +36,7 @@ if TYPE_CHECKING: from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, + StreamingScanKey, ) # Call types that stream JSON-RPC events (A2A); guardrail HTTPException is emitted as in-stream error @@ -54,6 +55,9 @@ class _EndpointTranslation(Protocol): @property def process_output_streaming_response(self) -> "Callable[..., Awaitable[object]]": ... + @property + def get_streaming_scan_key(self) -> "Callable[[Sequence[object]], StreamingScanKey | None]": ... + @property def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ... @@ -70,6 +74,12 @@ def _chunk_choices(item: object) -> Sequence[object]: return choices +def _is_redundant_scan(scan_key: "StreamingScanKey | None", last_scan_key: "StreamingScanKey | None") -> bool: + if scan_key is None: + return False + return scan_key == last_scan_key or scan_key.has_nothing_to_scan + + class _StreamTerminated(Exception): """Internal signal that the incremental transform stream has already emitted its terminal chunks (block message or in-stream error) and must stop.""" @@ -1011,6 +1021,7 @@ class UnifiedLLMGuardrails(CustomLogger): # Drives how a block terminates the stream: continue the in-progress # message (True) vs emit a standalone block message (False, buffered). chunks_yielded = False + last_scan_key: StreamingScanKey | None = None # rebind-ok: replaced after every scan round async for item in response: chunk_counter += 1 @@ -1052,6 +1063,19 @@ class UnifiedLLMGuardrails(CustomLogger): # Process chunk based on sampling rate if chunk_counter % sampling_rate == 0: + endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far) + if _is_redundant_scan(scan_key, last_scan_key): + verbose_proxy_logger.debug( + "Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round", + chunk_counter, + guardrail_to_apply.guardrail_name, + ) + chunks_yielded = True + responses_yielded.append(item) + yield item + continue + verbose_proxy_logger.debug( "Processing streaming chunk %s (sampling_rate=%s) with guardrail %s", chunk_counter, @@ -1067,8 +1091,6 @@ class UnifiedLLMGuardrails(CustomLogger): # string, permanently losing this chunk's content. original_item = copy.deepcopy(item) - endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() - try: await endpoint_translation.process_output_streaming_response( responses_so_far=responses_so_far, @@ -1110,6 +1132,8 @@ class UnifiedLLMGuardrails(CustomLogger): ): yield error_item return + if scan_key is not None: + last_scan_key = scan_key chunks_yielded = True responses_yielded.append(original_item) yield original_item @@ -1136,6 +1160,18 @@ class UnifiedLLMGuardrails(CustomLogger): # preserve the list, not clone every chunk (deepcopy would double # peak memory for large responses). buffered_items: Final = list(responses_so_far) if buffer_until_moderated else None + end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far) + if _is_redundant_scan(end_scan_key, last_scan_key): + verbose_proxy_logger.debug( + "Skipping end-of-stream scan for guardrail %s: the last sampled round already scanned it all", + guardrail_to_apply.guardrail_name, + ) + for buffered_item in buffered_items or (): + yield buffered_item + for pending_item in pending_end_of_stream_items: + responses_yielded.append(pending_item) + yield pending_item + return try: await endpoint_translation.process_output_streaming_response( diff --git a/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py b/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py new file mode 100644 index 00000000000..dd7e8fadcd8 --- /dev/null +++ b/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_a2a_guardrail_handler.py @@ -0,0 +1,35 @@ +"""Tests for litellm/llms/a2a/chat/guardrail_translation/handler.py.""" + +import json + +from litellm.llms.a2a.chat.guardrail_translation.handler import A2AGuardrailHandler +from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey + + +def _text_event(text: str) -> str: + return json.dumps( + { + "jsonrpc": "2.0", + "id": "req-1", + "result": {"kind": "message", "role": "agent", "parts": [{"kind": "text", "text": text}]}, + } + ) + + +def _status_event() -> str: + return json.dumps({"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "status-update", "status": {}}}) + + +class TestA2AGuardrailHandlerStreamingScanKey: + def test_key_joins_the_text_of_every_message_event(self): + key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hello "), _text_event("world")]) + assert key == StreamingScanKey(texts=("hello world",)) + + def test_events_without_text_leave_the_key_unchanged(self): + handler = A2AGuardrailHandler() + events = [_text_event("hello")] + assert handler.get_streaming_scan_key(events + [_status_event()]) == handler.get_streaming_scan_key(events) + + def test_unparseable_items_are_ignored(self): + key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hi"), "not json", b"bytes"]) + assert key.texts == ("hi",) diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 0fe7730e91e..3044a321aa6 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -13,6 +13,7 @@ import pytest from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey from litellm.llms.anthropic.chat.guardrail_translation.handler import ( AnthropicMessagesHandler, ) @@ -1991,3 +1992,56 @@ class TestStructuredWriteBackKeepsToolResults: } later_blocks = [b for m in messages[tool_use_index + 1 :] for b in self._blocks(m)] assert {"type": "text", "text": "Now fetch the page."} in later_blocks + + +class TestAnthropicMessagesHandlerStreamingScanKey: + """get_streaming_scan_key mirrors what process_output_streaming_response would scan""" + + @staticmethod + def _sse(event_type, data): + return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode() + + def _text_delta(self, text): + return self._sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ) + + def test_key_is_empty_before_any_text_arrives(self): + head = self._sse("message_start", {"type": "message_start", "message": {"stop_reason": None}}) + key = AnthropicMessagesHandler().get_streaming_scan_key([head]) + assert key == StreamingScanKey(texts=("",)) + + def test_key_accumulates_text_deltas(self): + key = AnthropicMessagesHandler().get_streaming_scan_key([self._text_delta("hello "), self._text_delta("world")]) + assert key.texts == ("hello world",) + assert key.stream_ended is False + + def _stop(self, stop_reason): + return self._sse( + "message_delta", + {"type": "message_delta", "delta": {"stop_reason": stop_reason, "stop_sequence": None}, "usage": {}}, + ) + + def test_stop_without_tool_use_scans_the_same_payload(self): + handler = AnthropicMessagesHandler() + open_key = handler.get_streaming_scan_key([self._text_delta("hi")]) + ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), self._stop("end_turn")]) + assert ended_key.stream_ended is True + assert ended_key == open_key + + def test_tool_use_blocks_enter_the_key_once_the_stream_has_ended(self): + handler = AnthropicMessagesHandler() + tool_use = self._sse( + "content_block_start", + { + "type": "content_block_start", + "index": 1, + "content_block": {"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {}}, + }, + ) + open_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use]) + ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use, self._stop("tool_use")]) + assert open_key == StreamingScanKey(texts=("hi",)) + assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key != open_key diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 7dd6065063a..cebab2512d0 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -12,6 +12,7 @@ import pytest from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey from litellm.llms.openai.chat.guardrail_translation.handler import ( OpenAIChatCompletionsHandler, ) @@ -1643,3 +1644,74 @@ class TestCheckStreamingHasEnded: ) ] assert handler._check_streaming_has_ended(chunks) is True + + +class TestStreamingScanKey: + """get_streaming_scan_key identifies what a sampled round would scan so the + unified hook can skip rounds that would re-scan already-cleared text""" + + @staticmethod + def _chunk(content, finish_reason=None, index=0): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + return ModelResponseStream( + choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)] + ) + + def test_key_carries_accumulated_text_and_open_stream(self): + handler = OpenAIChatCompletionsHandler() + key = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo")]) + assert key == StreamingScanKey(texts=("hello",)) + + def test_chunks_without_text_leave_the_key_unchanged(self): + handler = OpenAIChatCompletionsHandler() + before = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo")]) + after = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo"), self._chunk(None)]) + assert after == before + + def test_finish_chunk_without_tool_calls_scans_the_same_payload(self): + handler = OpenAIChatCompletionsHandler() + open_key = handler.get_streaming_scan_key([self._chunk("hi")]) + ended_key = handler.get_streaming_scan_key([self._chunk("hi"), self._chunk(None, finish_reason="stop")]) + assert open_key.stream_ended is False + assert ended_key.stream_ended is True + assert ended_key == open_key + + def test_tool_calls_only_enter_the_key_once_the_stream_has_ended(self): + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + ModelResponseStream, + StreamingChoices, + ) + + handler = OpenAIChatCompletionsHandler() + tool_call = ChatCompletionDeltaToolCall( + id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Paris"}') + ) + tool_chunk = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=None, tool_calls=[tool_call]), finish_reason=None)] + ) + open_key = handler.get_streaming_scan_key([self._chunk("hi"), tool_chunk]) + ended_key = handler.get_streaming_scan_key( + [self._chunk("hi"), tool_chunk, self._chunk(None, finish_reason="stop")] + ) + assert open_key == StreamingScanKey(texts=("hi",)) + assert ended_key.texts == ("hi",) + assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key != open_key + + def test_text_after_the_first_choice_finishes_still_changes_the_key(self): + handler = OpenAIChatCompletionsHandler() + first_done = [self._chunk("a", index=0), self._chunk("b", finish_reason="stop", index=0)] + key_at_first_finish = handler.get_streaming_scan_key(first_done) + key_after_more_text = handler.get_streaming_scan_key(first_done + [self._chunk("y", index=1)]) + assert key_at_first_finish.stream_ended is True + assert key_after_more_text.stream_ended is True + assert key_after_more_text != key_at_first_finish + + def test_non_stream_items_are_ignored(self): + handler = OpenAIChatCompletionsHandler() + key = handler.get_streaming_scan_key([self._chunk("hi"), b"data: [DONE]"]) + assert key.texts == ("hi",) diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index d071ef78c2d..295121167d6 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1731,3 +1731,93 @@ class TestBuildBlockSseChunks: dones = [payload for payload in payloads if payload["type"] == "response.output_item.done"] assert len(dones) == 1 assert dones[0]["item"]["content"][0]["text"] == "Blocked by policy." + + +class TestOpenAIResponsesHandlerStreamingScanKey: + """get_streaming_scan_key mirrors what process_output_streaming_response would scan""" + + @staticmethod + def _delta(sequence_number, text): + return { + "type": "response.output_text.delta", + "sequence_number": sequence_number, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": text, + } + + def test_no_events_yields_no_key(self): + assert OpenAIResponsesHandler().get_streaming_scan_key([]) is None + + def test_key_accumulates_deltas_while_the_stream_is_open(self): + from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey + + key = OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hel"), self._delta(1, "lo")]) + assert key == StreamingScanKey(texts=("hello",)) + + def test_typed_delta_events_accumulate_like_dicts(self): + from litellm.types.llms.openai import OutputTextDeltaEvent + + events = [ + OutputTextDeltaEvent( + type="response.output_text.delta", + item_id="msg_1", + output_index=0, + content_index=0, + delta=text, + sequence_number=i, + ) + for i, text in enumerate(("hel", "lo")) + ] + key = OpenAIResponsesHandler().get_streaming_scan_key(events) + assert key.texts == ("hello",) + assert key.stream_ended is False + + def test_events_without_text_leave_the_key_unchanged(self): + handler = OpenAIResponsesHandler() + events = [self._delta(0, "hi")] + quiet = events + [{"type": "response.in_progress", "sequence_number": 1}] + assert handler.get_streaming_scan_key(quiet) == handler.get_streaming_scan_key(events) + + @staticmethod + def _completed(sequence_number, output): + return {"type": "response.completed", "sequence_number": sequence_number, "response": {"output": output}} + + def test_completed_event_keys_on_the_final_output_text(self): + handler = OpenAIResponsesHandler() + message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]} + open_key = handler.get_streaming_scan_key([self._delta(0, "hi")]) + ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), self._completed(1, [message])]) + assert ended_key.stream_ended is True + assert ended_key == open_key + + def test_completed_event_with_a_function_call_changes_the_key(self): + handler = OpenAIResponsesHandler() + message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]} + function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": "{}"} + open_key = handler.get_streaming_scan_key([self._delta(0, "hi")]) + ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), self._completed(1, [message, function_call])]) + assert ended_key.texts == ("hi",) + assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key != open_key + + def test_completed_event_reads_every_output_text_part(self): + from litellm.types.responses.main import GenericResponseOutputItem, OutputText + + item = GenericResponseOutputItem( + type="message", + id="msg_1", + status="completed", + role="assistant", + content=[ + OutputText(type="output_text", text="one", annotations=[]), + OutputText(type="output_text", text="two", annotations=[]), + ], + ) + key = OpenAIResponsesHandler().get_streaming_scan_key([self._completed(0, [item])]) + assert key.texts == ("one", "two") + + def test_output_item_done_round_is_never_deduped(self): + done = {"type": "response.output_item.done", "sequence_number": 1, "item": {"type": "function_call"}} + assert OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hi"), done]) is None diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py index 476d443d8d8..cb6772977ec 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py @@ -310,8 +310,8 @@ def _make_stream_chunk(content: str, finish_reason=None): @pytest.mark.asyncio async def test_openai_moderation_streaming_default_uses_sampled_cadence(): """Default config samples every 5th streamed chunk and runs a final aggregate - pass after the stream ends. 10 chunks → sampled at chunks 5 and 10 → 2 in-stream - calls, plus 1 final = 3 total. + pass after the stream ends. 10 chunks are sampled at 5 and 10; the end-of-stream + round is skipped because chunk 10 already scanned the full text, for 2 total calls """ import litellm @@ -370,8 +370,9 @@ async def test_openai_moderation_streaming_default_uses_sampled_cadence(): ): pass - assert patched_make_request.await_count == 3, ( - f"Expected 3 moderation calls (2 sampled at chunks 5 / 10 + 1 final), " + assert patched_make_request.await_count == 2, ( + f"Expected 2 moderation calls (2 sampled at chunks 5 / 10; " + f"the end-of-stream round is skipped because chunk 10 already scanned the full text), " f"got {patched_make_request.await_count}" ) @@ -448,7 +449,8 @@ async def test_openai_moderation_streaming_end_of_stream_only_opt_in_calls_moder @pytest.mark.asyncio async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disabled(): """With streaming_end_of_stream_only=False and streaming_sampling_rate=2, - moderation runs every 2nd chunk during the stream, plus once more at end. + moderation runs every 2nd chunk during the stream. The terminal chunk scan covers + the final aggregate, for 3 total calls """ import litellm @@ -509,9 +511,8 @@ async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disab ): pass - # 6 chunks, sampling_rate=2 → in-stream calls at chunks 2, 4, 6 (3 calls), - # plus the final aggregate pass after the stream ends (1 call) = 4 total. - assert patched_make_request.await_count == 4, ( - f"Expected 4 moderation calls (3 sampled + 1 final aggregate), " + assert patched_make_request.await_count == 3, ( + f"Expected 3 moderation calls (3 sampled; the end-of-stream round is skipped " + f"because chunk 6 already scanned the full text), " f"got {patched_make_request.await_count}" ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 523ec1a37b4..83cc9ae8bb9 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -1517,7 +1517,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: @pytest.mark.asyncio async def test_streaming_default_uses_sampled_cadence(self): - """Default samples every 5th chunk + final pass: 10 chunks → calls at 5, 10, and final = 3.""" + """Default samples every 5th chunk. For 10 chunks, sampled scans at 5 and 10 + cover the full text, so the end-of-stream round is skipped and there are 2 calls + """ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) @@ -1566,8 +1568,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: ): pass - assert mock_post.await_count == 3, ( - f"Expected 3 guardrail calls (2 sampled at chunks 5 / 10 + 1 final), " + assert mock_post.await_count == 2, ( + f"Expected 2 guardrail calls (2 sampled at chunks 5 / 10; " + f"the end-of-stream round is skipped because chunk 10 already scanned the full text), " f"got {mock_post.await_count}" ) for call in mock_post.await_args_list: @@ -1631,7 +1634,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: @pytest.mark.asyncio async def test_streaming_sampling_rate_override(self): - """sampling_rate=2 on 6 chunks → in-stream at 2,4,6 plus final = 4 calls.""" + """sampling_rate=2 on 6 chunks. Scans at 2, 4, and 6 cover the full text, so + the end-of-stream round is skipped and there are 3 calls + """ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) @@ -1680,8 +1685,9 @@ class TestGenericGuardrailAPIStreamingViaUnified: ): pass - assert mock_post.await_count == 4, ( - f"Expected 4 guardrail calls (3 sampled + 1 final aggregate), " + assert mock_post.await_count == 3, ( + f"Expected 3 guardrail calls (3 sampled; the end-of-stream round is skipped " + f"because chunk 6 already scanned the full text), " f"got {mock_post.await_count}" ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 8cad1c634a9..a28a2a71613 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -1971,3 +1971,271 @@ class TestStreamingGuardrailInformationBucket: assert recorded[0]["guardrail_name"] == "audit-recorder" assert recorded[0]["guardrail_status"] == "success" assert request_data["metadata"]["user_api_key_user_id"] == "user-1" + + +class _ScanCountingGuardrail(CustomGuardrail): + """Pass-through guardrail that records every response-side scan payload.""" + + def __init__(self, *, sampling_rate=5, end_of_stream_only=False, buffer_until_moderated=False): + super().__init__(guardrail_name="scan-counter") + self.streaming_sampling_rate = sampling_rate + self.streaming_end_of_stream_only = end_of_stream_only + self.streaming_buffer_until_moderated = buffer_until_moderated + self.guardrail_config = {} + self.scans: tuple[dict[str, object], ...] = () + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + return True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.scans = ( + *self.scans, + { + "texts": list(inputs.get("texts") or []), + "tool_calls": list(inputs.get("tool_calls") or []), + "model": inputs.get("model"), + }, + ) + return inputs + + +def _responses_delta(sequence_number, text): + return { + "type": "response.output_text.delta", + "sequence_number": sequence_number, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": text, + } + + +def _responses_tail(sequence_number, text): + return [ + { + "type": "response.output_text.done", + "sequence_number": sequence_number, + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "text": text, + }, + { + "type": "response.completed", + "sequence_number": sequence_number + 1, + "response": { + "model": "gpt-5.6", + "output": [{"type": "message", "content": [{"type": "output_text", "text": text}]}], + }, + }, + ] + + +class TestStreamingScanDedup: + """A sampled round whose scan payload matches the previous round (or carries + no text yet) is skipped, so a stream is never re-scanned for output the + guardrail already cleared. Regression for LIT-6692.""" + + @pytest.fixture(autouse=True) + def _use_real_mappings(self, monkeypatch): + monkeypatch.setattr( + unified_module, + "endpoint_guardrail_translation_mappings", + load_guardrail_translation_mappings(), + ) + + @pytest.mark.asyncio + async def test_chat_terminal_chunk_on_sampled_index_is_scanned_once(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c", finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 3 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"]] + + @pytest.mark.asyncio + async def test_chat_round_with_unchanged_text_is_skipped(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [ + _stream_chunk("a"), + _stream_chunk("b"), + _stream_chunk("c"), + _stream_chunk(None), + _stream_chunk(None), + _stream_chunk(None), + _stream_chunk("d", finish_reason="stop"), + ] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 7 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"], ["abcd"]] + + @pytest.mark.asyncio + async def test_chat_finish_chunk_right_after_a_sampled_round_is_not_rescanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c"), _stream_chunk(None, finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 4 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"]] + + @pytest.mark.asyncio + async def test_chat_finish_chunk_carrying_tool_calls_is_still_scanned(self): + from litellm.types.utils import ChatCompletionDeltaToolCall, Function + + guardrail = _ScanCountingGuardrail(sampling_rate=3) + tool_call = ChatCompletionDeltaToolCall( + id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Paris"}') + ) + finish = ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content=None, tool_calls=[tool_call]), finish_reason="tool_calls") + ] + ) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c"), finish] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(out) == 4 + assert [scan["texts"] for scan in guardrail.scans] == [["abc"], ["abc"]] + assert [call["function"]["name"] for call in guardrail.scans[1]["tool_calls"]] == ["get_weather"] + + @pytest.mark.asyncio + async def test_chat_second_choice_finishing_later_still_gets_the_end_scan(self): + guardrail = _ScanCountingGuardrail(sampling_rate=3) + chunks = [ + _stream_chunk("a", index=0), + _stream_chunk("x", index=1), + _stream_chunk("b", finish_reason="stop", index=0), + _stream_chunk("y", index=1), + _stream_chunk("z", finish_reason="stop", index=1), + ] + + await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert len(guardrail.scans) == 2 + assert any("yz" in text for text in guardrail.scans[-1]["texts"]) + + @pytest.mark.asyncio + async def test_responses_completed_event_on_sampled_index_is_scanned_once(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(8)] + full_text = "".join(f"t{i}" for i in range(8)) + chunks = deltas + _responses_tail(8, full_text) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 10 + assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"], [full_text]] + assert guardrail.scans[-1]["model"] == "gpt-5.6" + + @pytest.mark.asyncio + async def test_responses_completed_right_after_a_sampled_round_is_not_rescanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(5)] + chunks = deltas + _responses_tail(5, "t0t1t2t3t4") + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 7 + assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"]] + + @pytest.mark.asyncio + async def test_responses_completed_carrying_a_function_call_is_still_scanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(5)] + completed = { + "type": "response.completed", + "sequence_number": 5, + "response": { + "model": "gpt-5.6", + "output": [ + {"type": "message", "content": [{"type": "output_text", "text": "t0t1t2t3t4"}]}, + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "status": "completed", + }, + ], + }, + } + chunks = deltas + [completed] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 6 + assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"], ["t0t1t2t3t4"]] + assert [call["function"]["name"] for call in guardrail.scans[1]["tool_calls"]] == ["get_weather"] + + @pytest.mark.asyncio + async def test_responses_round_with_unchanged_text_is_skipped(self): + guardrail = _ScanCountingGuardrail(sampling_rate=5) + deltas = [_responses_delta(i, f"t{i}") for i in range(5)] + quiet = [{"type": "response.in_progress", "sequence_number": i} for i in range(5, 10)] + chunks = deltas + quiet + _responses_tail(10, "t0t1t2t3t4") + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 12 + assert guardrail.scans == ({"texts": ["t0t1t2t3t4"], "tool_calls": [], "model": None},) + + @pytest.mark.asyncio + async def test_responses_tool_call_done_event_is_still_scanned(self): + guardrail = _ScanCountingGuardrail(sampling_rate=2) + tool_call_done = { + "type": "response.output_item.done", + "sequence_number": 1, + "output_index": 1, + "item": { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "status": "completed", + }, + } + chunks = [_responses_delta(0, "hi"), tool_call_done] + _responses_tail(2, "hi") + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") + + assert len(out) == 4 + assert len(guardrail.scans) == 2 + assert [call["function"]["name"] for call in guardrail.scans[0]["tool_calls"]] == ["get_weather"] + assert guardrail.scans[1]["texts"] == ["hi"] + + @pytest.mark.asyncio + async def test_anthropic_skips_empty_round_and_terminal_duplicate(self): + guardrail = _ScanCountingGuardrail(sampling_rate=2) + chunks = _anthropic_message_chunks(["hello ", "world"]) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages") + + assert out == chunks + assert [scan["texts"] for scan in guardrail.scans] == [["hello world"]] + + @pytest.mark.asyncio + async def test_end_of_stream_only_still_scans_exactly_once(self): + guardrail = _ScanCountingGuardrail(sampling_rate=2, end_of_stream_only=True) + chunks = _anthropic_message_chunks(["hello ", "world"]) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages") + + assert out == chunks + assert [scan["texts"] for scan in guardrail.scans] == [["hello world"]] + + @pytest.mark.asyncio + async def test_buffer_until_moderated_still_scans_exactly_once_and_releases_every_chunk(self): + guardrail = _ScanCountingGuardrail(sampling_rate=1, buffer_until_moderated=True) + chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c", finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert out == chunks + assert [scan["texts"] for scan in guardrail.scans] == [["abc"]]