From 54e2a582acd78253f98f77d1f405364b55b98f3a Mon Sep 17 00:00:00 2001 From: albertbausili Date: Fri, 16 Oct 2026 09:27:18 +0200 Subject: [PATCH 01/36] fix(guardrails): inspect streamed tool calls before delivery --- .../guardrail_translation/base_translation.py | 10 +- .../chat/guardrail_translation/handler.py | 220 +++++++++++------- .../unified_guardrail/unified_guardrail.py | 191 ++++++--------- .../test_openai_guardrail_handler.py | 70 ++++-- .../test_unified_guardrail.py | 180 +++++++++++++- 5 files changed, 433 insertions(+), 238 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 88c4dbfdce6..6280e636700 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional @@ -31,8 +31,12 @@ class StreamTransformSink: single call, not shared state. """ - mutated_text_per_choice: dict[int, str] = field(default_factory=dict) - holdback_per_choice: dict[int, int] = field(default_factory=dict) + mutated_text_per_choice: Mapping[int, str] = field(default_factory=dict) + holdback_per_choice: Mapping[int, int] = field(default_factory=dict) + + def record(self, texts: Mapping[int, str], holdback: Mapping[int, int]) -> None: + self.mutated_text_per_choice = texts + self.holdback_per_choice = holdback @dataclass(frozen=True, slots=True) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index a12f5f164e4..d18fe36e3df 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -18,9 +18,11 @@ import json import time import uuid from collections.abc import Iterator, Mapping, Sequence +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast +from pydantic import TypeAdapter from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -46,7 +48,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( stream_item_items, unappliable_request_rewrite, ) -from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolParam from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( coerce_stream_holdback_value, ) @@ -168,8 +170,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) - guardrailed_texts: Final = guardrailed_inputs.get("texts", []) - guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", []) + guardrailed_texts: Final = guardrailed_inputs.get("texts", ()) + guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", ()) guardrailed_tools: Final = guardrailed_inputs.get("tools") if guardrailed_tools is not None: data["tools"] = ( @@ -258,12 +260,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation): def extract_request_tool_names(self, data: dict) -> list[str]: """Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name).""" names: Final[list[str]] = [] - for tool in data.get("tools") or []: + for tool in data.get("tools") or (): if isinstance(tool, dict) and tool.get("type") == "function": fn = tool.get("function") if isinstance(fn, dict) and fn.get("name"): names.append(str(fn["name"])) - for fn in data.get("functions") or []: + for fn in data.get("functions") or (): if isinstance(fn, dict) and fn.get("name"): names.append(str(fn["name"])) return names @@ -387,7 +389,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, - request_data: dict | None = None, + request_data: dict[str, object] | None = None, ) -> ModelResponse: """ Process output response by applying guardrails to text content. @@ -438,9 +440,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # create a standalone dict (SDK / direct-call path). if request_data is None: request_data = {"response": response} - else: - if "response" not in request_data: - request_data["response"] = response + request_data.setdefault("response", response) self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) @@ -460,7 +460,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) - guardrailed_texts: Final = guardrailed_inputs.get("texts", []) + guardrailed_texts: Final = guardrailed_inputs.get("texts", ()) returned_tool_calls: Final = guardrailed_inputs.get("tool_calls") guardrailed_tool_calls: Final[list[dict[str, object]]] = ( cast(list[dict[str, object]], returned_tool_calls) @@ -494,7 +494,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, - request_data: dict | None = None, + request_data: dict[str, object] | None = None, stream_transform_sink: StreamTransformSink | None = None, deliver_ended_stream_rewrites: bool = False, ) -> list["ModelResponseStream"]: @@ -550,14 +550,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None", user_api_key_dict: "UserAPIKeyAuth | None", - request_data: dict | None, + request_data: dict[str, object] | None, deliver_ended_stream_rewrites: bool = False, ) -> list["ModelResponseStream"]: """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) unless ``deliver_ended_stream_rewrites`` opts the ended-stream branch in.""" - has_stream_ended: Final = self._first_choice_has_finished(responses_so_far) + has_stream_ended: Final = deliver_ended_stream_rewrites or self._first_choice_has_finished(responses_so_far) if has_stream_ended: await self._process_ended_stream( @@ -623,7 +623,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) - guardrailed_texts: Final = guardrailed_inputs.get("texts", []) + guardrailed_texts: Final = guardrailed_inputs.get("texts", ()) # Step 4: Apply guardrailed text back to all streaming chunks # For each choice, replace the combined text across all chunks @@ -656,13 +656,36 @@ class OpenAIChatCompletionsHandler(BaseTranslation): model_response: Final = self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj) pre_guardrail_texts: Final = self._string_choice_contents(model_response) pre_guardrail_tool_calls: Final = self._function_tool_call_shapes(model_response) - await self.process_output_response( - response=model_response, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=litellm_logging_obj, - user_api_key_dict=user_api_key_dict, - request_data=request_data, + inspection_responses: Final = ( + tuple( + model_response.model_copy(update=MappingProxyType({"choices": [choice]})) + for choice in model_response.choices + ) + if pre_guardrail_tool_calls and len(model_response.choices) > 1 + else (model_response,) ) + inspection_request_data: Final = request_data if request_data is not None else {} + try: + for inspection_response in inspection_responses: + inspection_request_data["response"] = inspection_response + inspection_request_data["responses"] = ( + [ + self._narrowed_to_choice(chunk, inspection_response.choices[0].index) + for chunk in responses_so_far + ] + if len(inspection_response.choices) == 1 + else responses_so_far + ) + await self.process_output_response( + response=inspection_response, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=inspection_request_data, + ) + finally: + inspection_request_data["response"] = model_response + inspection_request_data["responses"] = responses_so_far if not deliver_ended_stream_rewrites: return await self._write_ended_stream_text_rewrites( @@ -763,7 +786,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None", user_api_key_dict: "UserAPIKeyAuth | None", - request_data: dict | None, + request_data: dict[str, object] | None, sink: StreamTransformSink, ) -> None: """Run the guardrail over the raw accumulated text and report the @@ -775,8 +798,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): """ raw_by_index: Final = self._accumulate_string_content_by_choice_index(responses_so_far) if not raw_by_index: - sink.mutated_text_per_choice = {} - sink.holdback_per_choice = {} + sink.record(MappingProxyType({}), MappingProxyType({})) return # Fix #2 — sort by StreamingChoices.index so an n>1 stream that emits @@ -788,18 +810,72 @@ class OpenAIChatCompletionsHandler(BaseTranslation): indices: Final = sorted(raw_by_index.keys()) texts_to_check: Final = [raw_by_index[i] for i in indices] - if request_data is None: - request_data = {"responses": responses_so_far} - elif "responses" not in request_data: - request_data["responses"] = responses_so_far - self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) + resolved_request_data: Final[dict[str, object]] = ( + request_data if request_data is not None else {"responses": responses_so_far} + ) + if "responses" not in resolved_request_data: + resolved_request_data["responses"] = responses_so_far + self.merge_user_api_key_metadata_into_request(resolved_request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) + if self._streamed_tool_call_fingerprints(responses_so_far): + assembled: Final = self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj) + resolved_request_data["response"] = assembled + if len(assembled.choices) > 1: + choice_rounds: Final = tuple( + ( + choice, + StreamTransformSink(), + [self._narrowed_to_choice(chunk, choice.index) for chunk in responses_so_far], + ) + for choice in assembled.choices + ) + try: + for choice, choice_sink, choice_chunks in choice_rounds: + resolved_request_data["response"] = assembled.model_copy( + update=MappingProxyType({"choices": [choice]}) + ) + resolved_request_data["responses"] = choice_chunks + await self._process_streaming_transform( + responses_so_far=choice_chunks, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=resolved_request_data, + sink=choice_sink, + ) + finally: + resolved_request_data["response"] = assembled + resolved_request_data["responses"] = responses_so_far + merged_texts: Final = MappingProxyType( + dict( + chain.from_iterable( + choice_sink.mutated_text_per_choice.items() for _, choice_sink, _ in choice_rounds + ) + ) + ) + merged_holdback: Final = MappingProxyType( + dict( + chain.from_iterable( + choice_sink.holdback_per_choice.items() for _, choice_sink, _ in choice_rounds + ) + ) + ) + sink.record(merged_texts, merged_holdback) + return + tool_calls: Final = chain.from_iterable(choice.message.tool_calls or () for choice in assembled.choices) + inputs["tool_calls"] = TypeAdapter(list[ChatCompletionToolCallChunk]).validate_python( + tuple( + MappingProxyType({"index": index, **converted}) + for index, tool_call in enumerate(tool_calls) + if (converted := self._convert_tool_call_to_dict(tool_call)) is not None + ) + ) if responses_so_far and getattr(responses_so_far[0], "model", None): inputs["model"] = responses_so_far[0].model guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, - request_data=request_data, + request_data=resolved_request_data, input_type="response", logging_obj=litellm_logging_obj, ) @@ -819,13 +895,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation): len(texts_to_check), ) - holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or [] - sink.mutated_text_per_choice = { - idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts) - } - sink.holdback_per_choice = { - indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback) - } + holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or () + transformed_texts: Final = MappingProxyType( + {idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts)} + ) + transformed_holdback: Final = MappingProxyType( + {indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback)} + ) + sink.record(transformed_texts, transformed_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)) @@ -1145,20 +1222,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation): def _function_tool_call_fragments( responses_so_far: Sequence["ModelResponseStream"], ) -> tuple[tuple[ChatCompletionDeltaToolCall, ...], ...]: - """Group the stream's function tool-call fragments by their tool-call index, in - the index order ``stream_chunk_builder`` lists the rebuilt tool calls, keeping - only the indices the builder keeps (an id and a name somewhere in the stream).""" fragments: Final = tuple( - tool_call + (choice.index, tool_call) for response in responses_so_far for choice in response.choices for tool_call in choice.delta.tool_calls or () if isinstance(tool_call, ChatCompletionDeltaToolCall) ) - identified: Final = frozenset(fragment.index for fragment in fragments if fragment.id) - named: Final = frozenset(fragment.index for fragment in fragments if fragment.function.name) + identified: Final = frozenset((choice, fragment.index) for choice, fragment in fragments if fragment.id) + named: Final = frozenset((choice, fragment.index) for choice, fragment in fragments if fragment.function.name) return tuple( - tuple(fragment for fragment in fragments if fragment.index == index) for index in sorted(identified & named) + tuple(fragment for choice, fragment in fragments if (choice, fragment.index) == key) + for key in sorted(identified & named) ) def _write_ended_stream_tool_call_rewrites( @@ -1168,28 +1243,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): pre_guardrail_tool_calls: tuple[tuple[str | None, str], ...], guardrail_name: str, ) -> None: - """Write ended-stream guardrail tool-call rewrites back across the buffered - chunks: the rewritten name and full arguments land in the tool call's first - fragment and the arguments of its later fragments are blanked, mirroring the - text write-back. A rewrite on a stream carrying more than one distinct choice - index, or whose fragments do not line up with the rebuilt tool calls, is - reported as undeliverable, so the pipeline executor discards it and releases - the original chunks.""" post_guardrail_tool_calls: Final = self._function_tool_call_shapes(guardrailed_response) if post_guardrail_tool_calls == pre_guardrail_tool_calls: return - stream_choice_indices: Final = frozenset( - choice.index for response in responses_so_far for choice in response.choices - ) fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far) - if len(stream_choice_indices) != 1: - from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - - raise UndeliverableStreamRewrite( - guardrail_name, - f"the stream carries {len(stream_choice_indices)} choices and tool-call rewrites are only written " - "back on single-choice streams", - ) if len(fragments_by_tool_call) != len(post_guardrail_tool_calls): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite @@ -1212,7 +1269,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): async def _apply_guardrail_responses_to_output_streaming( self, responses: list["ModelResponseStream"], - guardrailed_texts: list[str], + guardrailed_texts: Sequence[str], task_mappings: list[tuple[int, int | None]], ) -> None: """ @@ -1230,12 +1287,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Override this method to customize how responses are applied to streaming responses. """ # Build a mapping of what guardrailed text to use for each (choice_idx, content_idx) - guardrail_map: Final[dict[tuple[int, int | None], str]] = {} - for task_idx, guardrail_response in enumerate(guardrailed_texts): - mapping = task_mappings[task_idx] - choice_idx = cast(int, mapping[0]) - content_idx_optional = cast(int | None, mapping[1]) - guardrail_map[(choice_idx, content_idx_optional)] = guardrail_response + guardrail_map: Final = MappingProxyType( + {task_mappings[index]: text for index, text in enumerate(guardrailed_texts)} + ) # Track which choices we've already set the guardrailed text for # Key: (choice_idx, content_idx), Value: boolean (True if already set) @@ -1245,34 +1299,20 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # choice by its index field: on n>1 streams a chunk usually carries one # choice at list position 0 whose index names the logical choice. for response in responses: - for choice in response.choices: - if isinstance(choice, litellm.StreamingChoices): - content = choice.delta.content - elif isinstance(choice, litellm.Choices): - content = choice.message.content - else: - continue - + for choice, message in ( + (choice, choice.delta if isinstance(choice, StreamingChoices) else choice.message) + for choice in response.choices + if isinstance(choice, (StreamingChoices, Choices)) + ): + content = message.content if content is None: continue if isinstance(content, str): - # String content str_key: tuple[int, int | None] = (choice.index, None) if str_key in guardrail_map: - if str_key not in already_set: - # First chunk - set the complete guardrailed text - if isinstance(choice, litellm.StreamingChoices): - choice.delta.content = guardrail_map[str_key] - elif isinstance(choice, litellm.Choices): - choice.message.content = guardrail_map[str_key] - already_set[str_key] = True - else: - # Subsequent chunks - clear the content - if isinstance(choice, litellm.StreamingChoices): - choice.delta.content = "" - elif isinstance(choice, litellm.Choices): - choice.message.content = "" + message.content = guardrail_map[str_key] if str_key not in already_set else "" + already_set[str_key] = True elif isinstance(content, list): # List content - handle each content item 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 a7f813d953e..9134a1cd0f3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, cast, runtime import anyio from fastapi import HTTPException +from pydantic import InstanceOf, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache @@ -49,6 +50,7 @@ if TYPE_CHECKING: A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message) GUARDRAIL_NAME: Final = "unified_llm_guardrails" +GUARDRAIL_ADAPTER: Final[TypeAdapter[CustomGuardrail | None]] = TypeAdapter(InstanceOf[CustomGuardrail] | None) _RequestData: TypeAlias = dict[str, object] @@ -513,9 +515,9 @@ class UnifiedLLMGuardrails(CustomLogger): self, *, reference_chunk: object, - mutated_text_per_choice: dict[int, str], + mutated_text_per_choice: Mapping[int, str], emitted_text_per_choice: dict[int, str], - holdback_per_choice: dict[int, int], + holdback_per_choice: Mapping[int, int], finish_reason_per_choice: dict[int, str | None], held_chars_per_choice: dict[int, int], is_final: bool, @@ -751,46 +753,13 @@ class UnifiedLLMGuardrails(CustomLogger): saw_tool_calls = False saw_text_content = False - tool_calls_released = False # rebind-ok: set once a raw tool call reaches the client unscanned - end_of_stream_inspection_started = False # rebind-ok: set once the end-of-stream inspection owns the verdict try: async for item in response: - # v1 transforms only text. A chunk carrying tool_calls is passed - # through raw so function-calling turns are not dropped, but ONLY - # its tool-call fields are forwarded: content is stripped so any - # response text (in the same delta, or in another choice of an n>1 - # chunk) can never bypass the transform. The original chunk is kept - # in responses_so_far so its text is still accumulated + redacted + - # emitted as synthetic deltas, and so the guardrail inspects the - # assembled tool calls at end of stream (see the block inspection - # below), matching block_only. finish_reason rides on the raw - # tool-only chunk, so it is not recorded for the text flush. if self._chunk_has_tool_calls(item): saw_tool_calls = True responses_so_far.append(item) last_chunk = item - # Fix #3 — flush accumulated text BEFORE the tool-call - # passthrough. Without this, a stream of text chunks that - # hasn't yet hit a sampled round can be trailed by a - # tool-call chunk carrying finish_reason="tool_calls"; an - # SSE-compliant client stops reading at that finish_reason - # and drops the end-of-stream text flush that would follow. - if saw_text_content: - async for out in _round(item, is_final=False): - yield out - # Fix #1 — pass finish_reason_per_choice into the - # passthrough so a mixed content+tool_call chunk defers its - # finish_reason to the final text terminator (see the - # _tool_call_passthrough_chunk docstring). - tool_only = self._tool_call_passthrough_chunk( - item, - finish_reason_per_choice=finish_reason_per_choice, - held_choices=_held_choices(held_chars_per_choice), - ) - responses_yielded.append(tool_only) - tool_calls_released = True - yield tool_only continue if self._is_trailing_metadata_chunk(item): @@ -808,39 +777,62 @@ class UnifiedLLMGuardrails(CustomLogger): # sampled round here would guardrail the same content twice. if ( not end_of_stream_only + and not saw_tool_calls and not self._chunk_has_finish_reason(item) and chunk_counter % sampling_rate == 0 ): async for out in _round(item, is_final=False): yield out - # v1 does not transform streamed tool calls, but they must still go - # through the guardrail's block decision. Run the block_only inspection - # over the full assembled response so tool calls cannot bypass it. - # - # Pass a deep copy of responses_so_far — the block path routes through - # ``_process_streaming_block_only`` which mutates ``delta.content`` - # in-place on the chunk objects it receives. For an n>1 chunk carrying - # text on one choice and tool_calls (with finish_reason) on another, - # ``has_stream_ended`` reads ``choices[0]`` alone and can miss the - # terminal signal, letting the block path rewrite the raw accumulator. - # The subsequent final ``_round`` would then re-read the already-mutated - # text, producing double-application for a non-idempotent guardrail or a - # ``stream_transform_underflow`` 400 from mismatched prefixes. A shallow - # list copy wouldn't help — the mutation is on the chunk objects - # themselves — so we deepcopy. - end_of_stream_inspection_started = True if saw_tool_calls: - async for out in self._inspect_full_response_for_block( + inspected_responses: Final = copy.deepcopy(responses_so_far) + async for out in self._inspect_full_response( endpoint_translation=endpoint_translation, guardrail_to_apply=guardrail_to_apply, request_data=request_data, user_api_key_dict=user_api_key_dict, - responses_so_far=copy.deepcopy(responses_so_far), + responses_so_far=inspected_responses, responses_yielded=responses_yielded, ): yield out + if saw_text_content: + async for out in _round(last_chunk, is_final=False): + yield out + tool_chunks: Final = tuple( + self._tool_call_passthrough_chunk( + buffered_item, + finish_reason_per_choice=finish_reason_per_choice, + held_choices=_held_choices(held_chars_per_choice), + ) + for buffered_item in inspected_responses + if self._chunk_has_tool_calls(buffered_item) + ) + + async def checked_tail() -> AsyncGenerator[object, None]: + try: + async for tail_chunk in self._emit_stream_tail( + last_chunk=last_chunk, + final_round=_round, + responses_so_far=responses_so_far, + responses_yielded=responses_yielded, + ): + yield tail_chunk + except _StreamTerminated as exc: + yield exc + + tail_chunks: Final = tuple([chunk async for chunk in checked_tail()]) + if tail_chunks and isinstance(tail_chunks[-1], _StreamTerminated): + for error_chunk in tail_chunks[:-1]: + yield error_chunk + return + for tool_only in tool_chunks: + responses_yielded.append(tool_only) + yield tool_only + for tail_chunk in tail_chunks: + yield tail_chunk + return + async for out in self._emit_stream_tail( last_chunk=last_chunk, final_round=_round, @@ -850,37 +842,6 @@ class UnifiedLLMGuardrails(CustomLogger): yield out except _StreamTerminated: return - except (GeneratorExit, asyncio.CancelledError): - await self._scan_uninspected_tool_calls_after_disconnect( - uninspected=tool_calls_released and not end_of_stream_inspection_started and not terminated.is_set(), - endpoint_translation=endpoint_translation, - responses_released=responses_yielded, - guardrail_to_apply=guardrail_to_apply, - user_api_key_dict=user_api_key_dict, - request_data=request_data, - ) - raise - - @staticmethod - async def _scan_uninspected_tool_calls_after_disconnect( - *, - uninspected: bool, - endpoint_translation: _EndpointTranslation, - responses_released: Sequence[object], - guardrail_to_apply: CustomGuardrail, - user_api_key_dict: UserAPIKeyAuth, - request_data: _RequestData, - ) -> None: - if not uninspected: - return - await UnifiedLLMGuardrails._scan_released_stream_after_disconnect( - endpoint_translation=endpoint_translation, - responses_released=responses_released, - last_scan_key=None, - guardrail_to_apply=guardrail_to_apply, - user_api_key_dict=user_api_key_dict, - request_data=request_data, - ) async def _emit_stream_tail( self, @@ -899,29 +860,21 @@ class UnifiedLLMGuardrails(CustomLogger): responses_yielded.append(trailing) yield trailing - async def _inspect_full_response_for_block( + async def _inspect_full_response( self, *, endpoint_translation: _EndpointTranslation, guardrail_to_apply: CustomGuardrail, - request_data: dict, + request_data: _RequestData, user_api_key_dict: UserAPIKeyAuth, responses_so_far: Sequence[object], responses_yielded: Sequence[object], ) -> AsyncGenerator[object, None]: - """Run the block-only guardrail inspection over the full assembled - response (text + tool calls) so nothing bypasses the block decision. - - The guardrail's returned transforms are discarded here (v1 does not - transform tool calls); only its block decision matters. A block is - surfaced the same way as elsewhere: ModifyResponseException terminates the - stream via the shared block handler; a GenericGuardrailAPI block raises and - propagates, matching block_only. - """ from litellm.integrations.custom_guardrail import ModifyResponseException + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite try: - with anyio.CancelScope(shield=bool(responses_yielded)): + with anyio.CancelScope(shield=True): await endpoint_translation.process_output_streaming_response( responses_so_far=responses_so_far, guardrail_to_apply=guardrail_to_apply, @@ -929,7 +882,10 @@ class UnifiedLLMGuardrails(CustomLogger): user_api_key_dict=user_api_key_dict, request_data=request_data, stream_transform_sink=None, + deliver_ended_stream_rewrites=True, ) + except UndeliverableStreamRewrite as exc: + raise HTTPException(status_code=400, detail="Guardrail stream rewrite could not be applied") from exc except ModifyResponseException as e: if e.original_response is None: e.original_response = responses_so_far @@ -1113,12 +1069,15 @@ class UnifiedLLMGuardrails(CustomLogger): # litellm.integrations.custom_guardrail. from litellm.integrations.custom_guardrail import ModifyResponseException - if guardrail_to_apply is None: - guardrail_to_apply = request_data.pop("guardrail_to_apply", None) + resolved_guardrail: Final[CustomGuardrail | None] = ( + guardrail_to_apply + if guardrail_to_apply is not None + else GUARDRAIL_ADAPTER.validate_python(request_data.pop("guardrail_to_apply", None)) + ) typed_request_data: Final[_RequestData] = request_data def _streaming_flag(name: str, default: object) -> Any: - return self.resolve_streaming_flag(guardrail_to_apply, name, default) + return self.resolve_streaming_flag(resolved_guardrail, name, default) sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5) # Only apply the guardrail at end of stream (not per chunk). @@ -1141,30 +1100,30 @@ class UnifiedLLMGuardrails(CustomLogger): if ( buffer_until_moderated - and guardrail_to_apply is not None - and getattr(guardrail_to_apply, "mask_response_content", False) + and resolved_guardrail is not None + and getattr(resolved_guardrail, "mask_response_content", False) ): verbose_proxy_logger.warning( "UnifiedLLMGuardrails: streaming_buffer_until_moderated is disabled for %s " "because mask_response_content=True -- buffered replay would release " "unredacted original chunks instead of the moderated output.", - guardrail_to_apply.guardrail_name, + resolved_guardrail.guardrail_name, ) buffer_until_moderated = False if buffer_until_moderated and not release_on_scan: end_of_stream_only = True - if guardrail_to_apply is None: + if resolved_guardrail is None: async for item in response: yield item return event_type: Final[GuardrailEventHooks] = GuardrailEventHooks.post_call - if guardrail_to_apply.should_run_guardrail(data=request_data, event_type=event_type) is not True: + if resolved_guardrail.should_run_guardrail(data=request_data, event_type=event_type) is not True: verbose_proxy_logger.debug( "UnifiedLLMGuardrails: Post-call streaming scanning disabled for %s", - guardrail_to_apply.guardrail_name, + resolved_guardrail.guardrail_name, ) async for item in response: yield item @@ -1184,7 +1143,7 @@ class UnifiedLLMGuardrails(CustomLogger): if transform_call_type is not None: async with contextlib.aclosing( self._run_incremental_transform_stream( - guardrail_to_apply=guardrail_to_apply, + guardrail_to_apply=resolved_guardrail, response=response, request_data=typed_request_data, user_api_key_dict=user_api_key_dict, @@ -1201,7 +1160,7 @@ class UnifiedLLMGuardrails(CustomLogger): "UnifiedLLMGuardrails: streaming_transform_mode=incremental_diff is only supported " "for the OpenAI chat completions streaming path with a resolvable request route; " "falling back to block_only for %s", - getattr(guardrail_to_apply, "guardrail_name", None), + getattr(resolved_guardrail, "guardrail_name", None), ) # Infer call type from first chunk @@ -1273,7 +1232,7 @@ class UnifiedLLMGuardrails(CustomLogger): 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, + resolved_guardrail.guardrail_name, ) if buffer_until_moderated: if hold_window: @@ -1293,7 +1252,7 @@ class UnifiedLLMGuardrails(CustomLogger): "Processing streaming chunk %s (sampling_rate=%s) with guardrail %s", chunk_counter, sampling_rate, - guardrail_to_apply.guardrail_name, + resolved_guardrail.guardrail_name, ) original_items = ( @@ -1303,7 +1262,7 @@ class UnifiedLLMGuardrails(CustomLogger): try: await endpoint_translation.process_output_streaming_response( responses_so_far=responses_so_far, - guardrail_to_apply=guardrail_to_apply, + guardrail_to_apply=resolved_guardrail, litellm_logging_obj=request_data.get("litellm_logging_obj"), user_api_key_dict=user_api_key_dict, request_data=request_data, @@ -1349,7 +1308,7 @@ class UnifiedLLMGuardrails(CustomLogger): verbose_proxy_logger.debug( "Holding %s buffered chunks for guardrail %s: this round could not scan the whole window", len(withheld_items), - guardrail_to_apply.guardrail_name, + resolved_guardrail.guardrail_name, ) withheld_items[:] = original_items continue @@ -1369,7 +1328,7 @@ class UnifiedLLMGuardrails(CustomLogger): verbose_proxy_logger.debug( "Processing final streaming response with all %s chunks for guardrail %s", len(responses_so_far), - guardrail_to_apply.guardrail_name, + resolved_guardrail.guardrail_name, ) endpoint_translation = mappings[CallTypes(call_type)]() @@ -1386,7 +1345,7 @@ class UnifiedLLMGuardrails(CustomLogger): 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, + resolved_guardrail.guardrail_name, ) for buffered_item in buffered_items or (): yield buffered_item @@ -1399,7 +1358,7 @@ class UnifiedLLMGuardrails(CustomLogger): with anyio.CancelScope(shield=chunks_yielded): await endpoint_translation.process_output_streaming_response( responses_so_far=responses_so_far, - guardrail_to_apply=guardrail_to_apply, + guardrail_to_apply=resolved_guardrail, litellm_logging_obj=request_data.get("litellm_logging_obj"), user_api_key_dict=user_api_key_dict, request_data=request_data, @@ -1443,13 +1402,13 @@ class UnifiedLLMGuardrails(CustomLogger): chunks_yielded and not verdict_settled and translation_class is not None - and isinstance(guardrail_to_apply, CustomGuardrail) + and isinstance(resolved_guardrail, CustomGuardrail) ): await self._scan_released_stream_after_disconnect( endpoint_translation=translation_class(), responses_released=responses_yielded, last_scan_key=last_scan_key, - guardrail_to_apply=guardrail_to_apply, + guardrail_to_apply=resolved_guardrail, user_api_key_dict=user_api_key_dict, request_data=typed_request_data, ) diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 5d71cbe9afc..bffe6e4c123 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -7,7 +7,7 @@ with guardrail transformations, including tool calls. import json from collections.abc import Mapping -from typing import Any, Literal, Optional +from typing import Any, Final, Literal, Optional import pytest @@ -1395,29 +1395,65 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: ) assert [ - (tool_call["id"], tool_call["function"]["arguments"]) - for tool_call in guardrail.seen_inputs[-1]["tool_calls"] - ] == [("call_1", '{"fruit": "persimmon"}'), ("call_2", '{"fruit": "durian"}')] + [(tool_call["id"], tool_call["function"]["arguments"]) for tool_call in inputs["tool_calls"]] + for inputs in guardrail.seen_inputs + ] == [[("call_1", '{"fruit": "persimmon"}')], [("call_2", '{"fruit": "durian"}')]] @pytest.mark.asyncio - async def test_deliver_ended_stream_tool_call_rewrite_on_multi_choice_stream_fails_closed(self): - from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + @pytest.mark.parametrize("transform", [False, True]) + async def test_each_choice_refreshes_shared_guardrail_response_context(self, transform: bool) -> None: + from fastapi import HTTPException + from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices - handler = OpenAIChatCompletionsHandler() - chunks = self._two_choice_tool_call_stream_chunks() + class ContextGuardrail(CustomGuardrail): + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object], + input_type: Literal["request", "response"], logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + response: Final = request_data["response"] + assert isinstance(response, ModelResponse) + request_data["visited_choices"] = (*request_data.get("visited_choices", ()), response.choices[0].index) + if response.choices[0].message.content == "forbidden": + raise HTTPException(status_code=400, detail="later choice blocked") + return inputs - with pytest.raises(UndeliverableStreamRewrite, match="the stream carries 2 choices") as raised: - await handler.process_output_streaming_response( - responses_so_far=chunks, - guardrail_to_apply=MockGuardrail(guardrail_name="test"), - litellm_logging_obj=None, - deliver_ended_stream_rewrites=True, + chunks: Final = [ModelResponseStream(choices=[StreamingChoices( + index=index, delta=Delta(content=text, tool_calls=[{ + "index": 0, "id": f"call_{index}", "type": "function", + "function": {"name": "contact", "arguments": "{}"}, + }]), finish_reason="tool_calls", + ) for index, text in enumerate(("allowed", "forbidden"))])] + request_data: Final = {"metadata": {"trace": "retained"}} + with pytest.raises(HTTPException, match="later choice blocked"): + await OpenAIChatCompletionsHandler().process_output_streaming_response( + responses_so_far=chunks, guardrail_to_apply=ContextGuardrail(guardrail_name="context"), + request_data=request_data, deliver_ended_stream_rewrites=True, + stream_transform_sink=StreamTransformSink() if transform else None, ) + assert request_data["visited_choices"] == (0, 1) + assert request_data["metadata"]["trace"] == "retained" + assert request_data["responses"] is chunks + restored: Final = request_data["response"] + assert isinstance(restored, ModelResponse) + assert tuple(choice.index for choice in restored.choices) == (0, 1) - assert raised.value.guardrail_name == "test" - assert raised.value.reason == ( - "the stream carries 2 choices and tool-call rewrites are only written back on single-choice streams" + @pytest.mark.asyncio + async def test_deliver_ended_stream_tool_rewrites_keep_choice_indices(self) -> None: + handler: Final = OpenAIChatCompletionsHandler() + chunks: Final = self._two_choice_tool_call_stream_chunks() + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MockGuardrail(guardrail_name="test"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, ) + arguments: Final = tuple( + "".join(call.function.arguments for chunk in chunks for choice in chunk.choices + if choice.index == index for call in choice.delta.tool_calls or ()) + for index in range(2) + ) + assert arguments == ('{"fruit": "PERSIMMON"}', '{"fruit": "DURIAN"}') @pytest.mark.asyncio async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self): diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 3efbe6475b2..4dd2f2cb6e3 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -8,6 +8,8 @@ from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator from types import SimpleNamespace from typing import TYPE_CHECKING, Final, Literal +from pydantic import TypeAdapter + import anyio import pytest @@ -50,6 +52,7 @@ from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import ( CallTypes, CallTypesLiteral, + ChatCompletionMessageToolCall, Delta, GenericGuardrailAPIInputs, ModelResponse, @@ -1008,6 +1011,42 @@ def _delta_text(item): return item.choices[0].delta.content or "" +class _ToolRedactingGuardrail(CustomGuardrail): + def __init__(self, require_tool_context: bool = False) -> None: + super().__init__(guardrail_name="tool-redactor", event_hook=GuardrailEventHooks.post_call, default_on=True) + self.streaming_transform_mode = "incremental_diff" + self.streaming_sampling_rate = 1 + self.require_tool_context = require_tool_context + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: object | None = None, + ) -> GenericGuardrailAPIInputs: + calls: Final = TypeAdapter(tuple[ChatCompletionMessageToolCall, ...]).validate_python( + inputs.get("tool_calls", ()) + ) + input_texts: Final = inputs.get("texts", ()) + texts: Final = tuple( + "checked:" + text.replace("SECRET", "MASKED") + if not self.require_tool_context or (calls and index == len(input_texts) - 1) else text + for index, text in enumerate(input_texts) + ) + return { + **inputs, + "texts": list(texts), + "stream_holdback_chars": [len(text) for text in texts], + "tool_calls": [ + call.model_copy(update={"function": call.function.model_copy(update={ + "arguments": call.function.arguments.replace("SECRET", "MASKED"), + })}) + for call in calls + ], + } + + class TestStreamingTransform: """Streaming text-transformation (incremental_diff) path on the OpenAI chat completions streaming surface.""" @@ -1016,6 +1055,103 @@ class TestStreamingTransform: def _use_openai_handler_mapping(self, monkeypatch): _patch_translation_mappings(monkeypatch, {CallTypes.acompletion: OpenAIChatCompletionsHandler}) + @pytest.mark.asyncio + @pytest.mark.parametrize("require_tool_context", [False, True]) + @pytest.mark.parametrize("include_text", [False, True]) + @pytest.mark.parametrize("tool_count", [1, 2]) + async def test_buffered_tool_arguments_are_rewritten_before_delivery( + self, include_text: bool, tool_count: int, require_tool_context: bool + ) -> None: + chunks: Final = ( + *([_stream_chunk("hello SECRET")] if include_text else []), + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(tool_calls=[ + {"index": index, "id": f"call_{index}", "type": "function", + "function": {"name": "contact", "arguments": '{"contact":"SEC'}} + for index in range(tool_count) + ]))]), + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(tool_calls=[ + {"index": index, "function": {"arguments": 'RET"}'}} for index in range(tool_count) + ]), finish_reason="tool_calls")]), + ModelResponseStream(choices=[], usage={"prompt_tokens": 7, "completion_tokens": 11, "total_tokens": 18}), + ) + out: Final = await _drive_stream( + UnifiedLLMGuardrails(), _ToolRedactingGuardrail(require_tool_context=require_tool_context), chunks + ) + calls: Final = tuple( + call for chunk in out for choice in chunk.choices for call in choice.delta.tool_calls or () + ) + for index in range(tool_count): + arguments: Final = "".join(call.function.arguments or "" for call in calls if call.index == index) + assert arguments == '{"contact":"MASKED"}' + assert next(call.id for call in calls if call.index == index and call.id) == f"call_{index}" + assert "".join(_delta_text(chunk) for chunk in out) == ("checked:hello MASKED" if include_text else "") + assert any(choice.finish_reason == "tool_calls" for chunk in out for choice in chunk.choices) + assert out[-1].usage.total_tokens == 18 + assert all("SECRET" not in chunk.model_dump_json() for chunk in out) + + @pytest.mark.asyncio + @pytest.mark.parametrize("tool_choice_index", [0, 1]) + async def test_tool_dependent_text_rewrites_keep_choice_context(self, tool_choice_index: int) -> None: + chunks: Final = ( + ModelResponseStream(choices=[StreamingChoices( + index=index, delta=Delta(content="SECRET" if index == tool_choice_index else "plain"), + ) for index in (1, 0)]), + ModelResponseStream(choices=[StreamingChoices( + index=tool_choice_index, + delta=Delta(tool_calls=[{ + "index": 0, "id": "call_context", "type": "function", + "function": {"name": "contact", "arguments": '{"contact":"SECRET"}'}, + }]), finish_reason="tool_calls", + )]), + ) + out: Final = await _drive_stream(UnifiedLLMGuardrails(), _ToolRedactingGuardrail(True), chunks) + for index in (0, 1): + text: Final = "".join( + choice.delta.content or "" for chunk in out for choice in chunk.choices if choice.index == index + ) + assert text == ("checked:MASKED" if index == tool_choice_index else "plain") + assert all("SECRET" not in chunk.model_dump_json() for chunk in out) + + @pytest.mark.asyncio + async def test_tool_rewrites_keep_completion_choices_separate(self) -> None: + async def response() -> AsyncIterator[ModelResponseStream]: + yield ModelResponseStream(choices=[StreamingChoices( + index=index, + delta=Delta(tool_calls=[{"index": 0, "id": f"call_{index}", "type": "function", + "function": {"name": "contact", "arguments": '{"contact":"SECRET"}'}}]), + finish_reason="tool_calls", + ) for index in range(2)]) + + iterator: Final = UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"), + response=response(), request_data={"guardrail_to_apply": _ToolRedactingGuardrail()}, + ) + chunks: Final = tuple([chunk async for chunk in iterator]) + calls: Final = tuple( + (choice.index, call.id, call.function.arguments) + for chunk in chunks for choice in chunk.choices for call in choice.delta.tool_calls or () + ) + assert calls == ((0, "call_0", '{"contact":"MASKED"}'), (1, "call_1", '{"contact":"MASKED"}')) + + @pytest.mark.asyncio + async def test_undeliverable_rewrite_is_a_closed_failure(self) -> None: + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + class UndeliverableTranslation(OpenAIChatCompletionsHandler): + async def process_output_streaming_response( + self, *args: object, **kwargs: object + ) -> list[ModelResponseStream]: + raise UndeliverableStreamRewrite("tool-redactor", "unmappable tool fragments") + + iterator: Final = UnifiedLLMGuardrails()._inspect_full_response( + endpoint_translation=UndeliverableTranslation(), guardrail_to_apply=_ToolRedactingGuardrail(), + request_data={}, user_api_key_dict=UserAPIKeyAuth(), responses_so_far=(), responses_yielded=(), + ) + with pytest.raises(unified_module.HTTPException) as error: + await anext(iterator) + assert error.value.status_code == 400 + assert error.value.detail == "Guardrail stream rewrite could not be applied" + @pytest.mark.asyncio async def test_block_only_drops_text_rewrites(self): """Default block_only: the guardrail's uppercasing never reaches the @@ -1436,14 +1572,15 @@ class TestStreamingTransform: assert out[-2].choices[0].finish_reason == "stop" @pytest.mark.asyncio - async def test_tool_call_blocking_guardrail_is_enforced(self): + @pytest.mark.parametrize(("content", "allowed_scans"), [(None, 0), ("proposal", 0), ("proposal", 1)]) + async def test_tool_call_blocking_guardrail_is_enforced(self, content: str | None, allowed_scans: int): """A guardrail that blocks on tool calls must terminate the incremental_diff stream: tool calls go through the block decision, not bypass it.""" from litellm.exceptions import GuardrailRaisedException class _ToolCallBlocker(_StreamingTextGuardrail): async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): - if input_type == "response" and inputs.get("tool_calls"): + if input_type == "response" and self.response_calls >= allowed_scans: raise GuardrailRaisedException( guardrail_name="tc-block", message="blocked tool call", @@ -1456,7 +1593,7 @@ class TestStreamingTransform: StreamingChoices( index=0, delta=Delta( - content=None, + content=content, tool_calls=[ { "index": 0, @@ -1471,8 +1608,21 @@ class TestStreamingTransform: ], ) + async def upstream(): + yield tool_chunk + + stream: Final = UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/chat/completions"), + response=upstream(), + request_data={"guardrail_to_apply": _ToolCallBlocker(), "model": "gpt-4"}, + ) + async def consume_checked_stream() -> None: + async for chunk in stream: + assert isinstance(chunk, ModelResponseStream) + assert all(not choice.delta.tool_calls and choice.finish_reason is None for choice in chunk.choices) + with pytest.raises(GuardrailRaisedException): - await _drive_stream(UnifiedLLMGuardrails(), _ToolCallBlocker(), [tool_chunk]) + await consume_checked_stream() @pytest.mark.asyncio async def test_mixed_content_and_tool_call_chunk_does_not_leak_text(self): @@ -2690,7 +2840,7 @@ class TestStreamingClientDisconnectScan: yield _stream_chunk(" tail", finish_reason="stop") @pytest.mark.asyncio - async def test_closing_while_an_incremental_diff_block_is_delivered_does_not_inspect_the_tool_call_again(self): + async def test_incremental_diff_block_withholds_tool_calls_and_does_not_inspect_again(self): guardrail = _MarkerBlockingStreamingTextGuardrail() async def upstream() -> AsyncIterator[ModelResponseStream]: @@ -2704,12 +2854,16 @@ class TestStreamingClientDisconnectScan: released = [await stream.__anext__(), await stream.__anext__()] await stream.aclose() - assert [call.function.name for call in released[0].choices[0].delta.tool_calls] == ["get_weather"] - assert guardrail.received_texts == [["BLOCKME"]], guardrail.received_texts - assert guardrail.received_tool_calls == [], guardrail.received_tool_calls + assert all(isinstance(item, bytes) for item in released) + assert b"get_weather" not in b"".join(released) + assert b"Paris" not in b"".join(released) + assert guardrail.received_texts == [["BLOCKME tail"]], guardrail.received_texts + assert [[call["function"]["name"] for call in calls] for calls in guardrail.received_tool_calls] == [ + ["get_weather"] + ] @pytest.mark.asyncio - async def test_closing_after_a_released_tool_call_under_incremental_diff_scans_no_held_back_text(self): + async def test_closing_after_a_buffered_tool_call_inspects_held_text_before_delivery(self): guardrail = _StreamingTextGuardrail(holdback_schedule=[len("held secret")] * 2) async def upstream() -> AsyncIterator[ModelResponseStream]: @@ -2718,12 +2872,14 @@ class TestStreamingClientDisconnectScan: yield tool_chunk stream = self._guarded_stream(guardrail, upstream()) - received = await stream.__anext__() + received: Final = tuple([await stream.__anext__() for _ in range(2)]) + scans_before_close: Final = len(guardrail.received_texts) await stream.aclose() - assert [call.function.name for call in received.choices[0].delta.tool_calls] == ["get_weather"] + assert [call.function.name for call in received[-1].choices[0].delta.tool_calls] == ["get_weather"] assert guardrail.received_tool_calls, guardrail.received_texts - assert all("held secret" not in text for text in guardrail.received_texts[-1]), guardrail.received_texts + assert guardrail.received_texts[0] == ["held secret"] + assert len(guardrail.received_texts) == scans_before_close def _responses_delta(sequence_number, text): From 95348ba610fcd21b742c64786d1fa0eed7d4b9f4 Mon Sep 17 00:00:00 2001 From: albertbausili Date: Fri, 16 Oct 2026 09:36:57 +0200 Subject: [PATCH 02/36] fix(guardrails): isolate streamed completion choices during inspection --- .../guardrail_translation/base_translation.py | 10 +- .../chat/guardrail_translation/handler.py | 220 +++++++++++------- .../test_openai_guardrail_handler.py | 70 ++++-- 3 files changed, 190 insertions(+), 110 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 88c4dbfdce6..6280e636700 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional @@ -31,8 +31,12 @@ class StreamTransformSink: single call, not shared state. """ - mutated_text_per_choice: dict[int, str] = field(default_factory=dict) - holdback_per_choice: dict[int, int] = field(default_factory=dict) + mutated_text_per_choice: Mapping[int, str] = field(default_factory=dict) + holdback_per_choice: Mapping[int, int] = field(default_factory=dict) + + def record(self, texts: Mapping[int, str], holdback: Mapping[int, int]) -> None: + self.mutated_text_per_choice = texts + self.holdback_per_choice = holdback @dataclass(frozen=True, slots=True) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index a12f5f164e4..d18fe36e3df 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -18,9 +18,11 @@ import json import time import uuid from collections.abc import Iterator, Mapping, Sequence +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast +from pydantic import TypeAdapter from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -46,7 +48,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( stream_item_items, unappliable_request_rewrite, ) -from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolParam from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( coerce_stream_holdback_value, ) @@ -168,8 +170,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) - guardrailed_texts: Final = guardrailed_inputs.get("texts", []) - guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", []) + guardrailed_texts: Final = guardrailed_inputs.get("texts", ()) + guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", ()) guardrailed_tools: Final = guardrailed_inputs.get("tools") if guardrailed_tools is not None: data["tools"] = ( @@ -258,12 +260,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation): def extract_request_tool_names(self, data: dict) -> list[str]: """Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name).""" names: Final[list[str]] = [] - for tool in data.get("tools") or []: + for tool in data.get("tools") or (): if isinstance(tool, dict) and tool.get("type") == "function": fn = tool.get("function") if isinstance(fn, dict) and fn.get("name"): names.append(str(fn["name"])) - for fn in data.get("functions") or []: + for fn in data.get("functions") or (): if isinstance(fn, dict) and fn.get("name"): names.append(str(fn["name"])) return names @@ -387,7 +389,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, - request_data: dict | None = None, + request_data: dict[str, object] | None = None, ) -> ModelResponse: """ Process output response by applying guardrails to text content. @@ -438,9 +440,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # create a standalone dict (SDK / direct-call path). if request_data is None: request_data = {"response": response} - else: - if "response" not in request_data: - request_data["response"] = response + request_data.setdefault("response", response) self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) @@ -460,7 +460,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) - guardrailed_texts: Final = guardrailed_inputs.get("texts", []) + guardrailed_texts: Final = guardrailed_inputs.get("texts", ()) returned_tool_calls: Final = guardrailed_inputs.get("tool_calls") guardrailed_tool_calls: Final[list[dict[str, object]]] = ( cast(list[dict[str, object]], returned_tool_calls) @@ -494,7 +494,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, - request_data: dict | None = None, + request_data: dict[str, object] | None = None, stream_transform_sink: StreamTransformSink | None = None, deliver_ended_stream_rewrites: bool = False, ) -> list["ModelResponseStream"]: @@ -550,14 +550,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None", user_api_key_dict: "UserAPIKeyAuth | None", - request_data: dict | None, + request_data: dict[str, object] | None, deliver_ended_stream_rewrites: bool = False, ) -> list["ModelResponseStream"]: """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) unless ``deliver_ended_stream_rewrites`` opts the ended-stream branch in.""" - has_stream_ended: Final = self._first_choice_has_finished(responses_so_far) + has_stream_ended: Final = deliver_ended_stream_rewrites or self._first_choice_has_finished(responses_so_far) if has_stream_ended: await self._process_ended_stream( @@ -623,7 +623,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) - guardrailed_texts: Final = guardrailed_inputs.get("texts", []) + guardrailed_texts: Final = guardrailed_inputs.get("texts", ()) # Step 4: Apply guardrailed text back to all streaming chunks # For each choice, replace the combined text across all chunks @@ -656,13 +656,36 @@ class OpenAIChatCompletionsHandler(BaseTranslation): model_response: Final = self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj) pre_guardrail_texts: Final = self._string_choice_contents(model_response) pre_guardrail_tool_calls: Final = self._function_tool_call_shapes(model_response) - await self.process_output_response( - response=model_response, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=litellm_logging_obj, - user_api_key_dict=user_api_key_dict, - request_data=request_data, + inspection_responses: Final = ( + tuple( + model_response.model_copy(update=MappingProxyType({"choices": [choice]})) + for choice in model_response.choices + ) + if pre_guardrail_tool_calls and len(model_response.choices) > 1 + else (model_response,) ) + inspection_request_data: Final = request_data if request_data is not None else {} + try: + for inspection_response in inspection_responses: + inspection_request_data["response"] = inspection_response + inspection_request_data["responses"] = ( + [ + self._narrowed_to_choice(chunk, inspection_response.choices[0].index) + for chunk in responses_so_far + ] + if len(inspection_response.choices) == 1 + else responses_so_far + ) + await self.process_output_response( + response=inspection_response, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=inspection_request_data, + ) + finally: + inspection_request_data["response"] = model_response + inspection_request_data["responses"] = responses_so_far if not deliver_ended_stream_rewrites: return await self._write_ended_stream_text_rewrites( @@ -763,7 +786,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None", user_api_key_dict: "UserAPIKeyAuth | None", - request_data: dict | None, + request_data: dict[str, object] | None, sink: StreamTransformSink, ) -> None: """Run the guardrail over the raw accumulated text and report the @@ -775,8 +798,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): """ raw_by_index: Final = self._accumulate_string_content_by_choice_index(responses_so_far) if not raw_by_index: - sink.mutated_text_per_choice = {} - sink.holdback_per_choice = {} + sink.record(MappingProxyType({}), MappingProxyType({})) return # Fix #2 — sort by StreamingChoices.index so an n>1 stream that emits @@ -788,18 +810,72 @@ class OpenAIChatCompletionsHandler(BaseTranslation): indices: Final = sorted(raw_by_index.keys()) texts_to_check: Final = [raw_by_index[i] for i in indices] - if request_data is None: - request_data = {"responses": responses_so_far} - elif "responses" not in request_data: - request_data["responses"] = responses_so_far - self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) + resolved_request_data: Final[dict[str, object]] = ( + request_data if request_data is not None else {"responses": responses_so_far} + ) + if "responses" not in resolved_request_data: + resolved_request_data["responses"] = responses_so_far + self.merge_user_api_key_metadata_into_request(resolved_request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) + if self._streamed_tool_call_fingerprints(responses_so_far): + assembled: Final = self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj) + resolved_request_data["response"] = assembled + if len(assembled.choices) > 1: + choice_rounds: Final = tuple( + ( + choice, + StreamTransformSink(), + [self._narrowed_to_choice(chunk, choice.index) for chunk in responses_so_far], + ) + for choice in assembled.choices + ) + try: + for choice, choice_sink, choice_chunks in choice_rounds: + resolved_request_data["response"] = assembled.model_copy( + update=MappingProxyType({"choices": [choice]}) + ) + resolved_request_data["responses"] = choice_chunks + await self._process_streaming_transform( + responses_so_far=choice_chunks, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=resolved_request_data, + sink=choice_sink, + ) + finally: + resolved_request_data["response"] = assembled + resolved_request_data["responses"] = responses_so_far + merged_texts: Final = MappingProxyType( + dict( + chain.from_iterable( + choice_sink.mutated_text_per_choice.items() for _, choice_sink, _ in choice_rounds + ) + ) + ) + merged_holdback: Final = MappingProxyType( + dict( + chain.from_iterable( + choice_sink.holdback_per_choice.items() for _, choice_sink, _ in choice_rounds + ) + ) + ) + sink.record(merged_texts, merged_holdback) + return + tool_calls: Final = chain.from_iterable(choice.message.tool_calls or () for choice in assembled.choices) + inputs["tool_calls"] = TypeAdapter(list[ChatCompletionToolCallChunk]).validate_python( + tuple( + MappingProxyType({"index": index, **converted}) + for index, tool_call in enumerate(tool_calls) + if (converted := self._convert_tool_call_to_dict(tool_call)) is not None + ) + ) if responses_so_far and getattr(responses_so_far[0], "model", None): inputs["model"] = responses_so_far[0].model guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, - request_data=request_data, + request_data=resolved_request_data, input_type="response", logging_obj=litellm_logging_obj, ) @@ -819,13 +895,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation): len(texts_to_check), ) - holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or [] - sink.mutated_text_per_choice = { - idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts) - } - sink.holdback_per_choice = { - indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback) - } + holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or () + transformed_texts: Final = MappingProxyType( + {idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts)} + ) + transformed_holdback: Final = MappingProxyType( + {indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback)} + ) + sink.record(transformed_texts, transformed_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)) @@ -1145,20 +1222,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation): def _function_tool_call_fragments( responses_so_far: Sequence["ModelResponseStream"], ) -> tuple[tuple[ChatCompletionDeltaToolCall, ...], ...]: - """Group the stream's function tool-call fragments by their tool-call index, in - the index order ``stream_chunk_builder`` lists the rebuilt tool calls, keeping - only the indices the builder keeps (an id and a name somewhere in the stream).""" fragments: Final = tuple( - tool_call + (choice.index, tool_call) for response in responses_so_far for choice in response.choices for tool_call in choice.delta.tool_calls or () if isinstance(tool_call, ChatCompletionDeltaToolCall) ) - identified: Final = frozenset(fragment.index for fragment in fragments if fragment.id) - named: Final = frozenset(fragment.index for fragment in fragments if fragment.function.name) + identified: Final = frozenset((choice, fragment.index) for choice, fragment in fragments if fragment.id) + named: Final = frozenset((choice, fragment.index) for choice, fragment in fragments if fragment.function.name) return tuple( - tuple(fragment for fragment in fragments if fragment.index == index) for index in sorted(identified & named) + tuple(fragment for choice, fragment in fragments if (choice, fragment.index) == key) + for key in sorted(identified & named) ) def _write_ended_stream_tool_call_rewrites( @@ -1168,28 +1243,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): pre_guardrail_tool_calls: tuple[tuple[str | None, str], ...], guardrail_name: str, ) -> None: - """Write ended-stream guardrail tool-call rewrites back across the buffered - chunks: the rewritten name and full arguments land in the tool call's first - fragment and the arguments of its later fragments are blanked, mirroring the - text write-back. A rewrite on a stream carrying more than one distinct choice - index, or whose fragments do not line up with the rebuilt tool calls, is - reported as undeliverable, so the pipeline executor discards it and releases - the original chunks.""" post_guardrail_tool_calls: Final = self._function_tool_call_shapes(guardrailed_response) if post_guardrail_tool_calls == pre_guardrail_tool_calls: return - stream_choice_indices: Final = frozenset( - choice.index for response in responses_so_far for choice in response.choices - ) fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far) - if len(stream_choice_indices) != 1: - from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - - raise UndeliverableStreamRewrite( - guardrail_name, - f"the stream carries {len(stream_choice_indices)} choices and tool-call rewrites are only written " - "back on single-choice streams", - ) if len(fragments_by_tool_call) != len(post_guardrail_tool_calls): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite @@ -1212,7 +1269,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): async def _apply_guardrail_responses_to_output_streaming( self, responses: list["ModelResponseStream"], - guardrailed_texts: list[str], + guardrailed_texts: Sequence[str], task_mappings: list[tuple[int, int | None]], ) -> None: """ @@ -1230,12 +1287,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Override this method to customize how responses are applied to streaming responses. """ # Build a mapping of what guardrailed text to use for each (choice_idx, content_idx) - guardrail_map: Final[dict[tuple[int, int | None], str]] = {} - for task_idx, guardrail_response in enumerate(guardrailed_texts): - mapping = task_mappings[task_idx] - choice_idx = cast(int, mapping[0]) - content_idx_optional = cast(int | None, mapping[1]) - guardrail_map[(choice_idx, content_idx_optional)] = guardrail_response + guardrail_map: Final = MappingProxyType( + {task_mappings[index]: text for index, text in enumerate(guardrailed_texts)} + ) # Track which choices we've already set the guardrailed text for # Key: (choice_idx, content_idx), Value: boolean (True if already set) @@ -1245,34 +1299,20 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # choice by its index field: on n>1 streams a chunk usually carries one # choice at list position 0 whose index names the logical choice. for response in responses: - for choice in response.choices: - if isinstance(choice, litellm.StreamingChoices): - content = choice.delta.content - elif isinstance(choice, litellm.Choices): - content = choice.message.content - else: - continue - + for choice, message in ( + (choice, choice.delta if isinstance(choice, StreamingChoices) else choice.message) + for choice in response.choices + if isinstance(choice, (StreamingChoices, Choices)) + ): + content = message.content if content is None: continue if isinstance(content, str): - # String content str_key: tuple[int, int | None] = (choice.index, None) if str_key in guardrail_map: - if str_key not in already_set: - # First chunk - set the complete guardrailed text - if isinstance(choice, litellm.StreamingChoices): - choice.delta.content = guardrail_map[str_key] - elif isinstance(choice, litellm.Choices): - choice.message.content = guardrail_map[str_key] - already_set[str_key] = True - else: - # Subsequent chunks - clear the content - if isinstance(choice, litellm.StreamingChoices): - choice.delta.content = "" - elif isinstance(choice, litellm.Choices): - choice.message.content = "" + message.content = guardrail_map[str_key] if str_key not in already_set else "" + already_set[str_key] = True elif isinstance(content, list): # List content - handle each content item diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 5d71cbe9afc..bffe6e4c123 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -7,7 +7,7 @@ with guardrail transformations, including tool calls. import json from collections.abc import Mapping -from typing import Any, Literal, Optional +from typing import Any, Final, Literal, Optional import pytest @@ -1395,29 +1395,65 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: ) assert [ - (tool_call["id"], tool_call["function"]["arguments"]) - for tool_call in guardrail.seen_inputs[-1]["tool_calls"] - ] == [("call_1", '{"fruit": "persimmon"}'), ("call_2", '{"fruit": "durian"}')] + [(tool_call["id"], tool_call["function"]["arguments"]) for tool_call in inputs["tool_calls"]] + for inputs in guardrail.seen_inputs + ] == [[("call_1", '{"fruit": "persimmon"}')], [("call_2", '{"fruit": "durian"}')]] @pytest.mark.asyncio - async def test_deliver_ended_stream_tool_call_rewrite_on_multi_choice_stream_fails_closed(self): - from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + @pytest.mark.parametrize("transform", [False, True]) + async def test_each_choice_refreshes_shared_guardrail_response_context(self, transform: bool) -> None: + from fastapi import HTTPException + from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices - handler = OpenAIChatCompletionsHandler() - chunks = self._two_choice_tool_call_stream_chunks() + class ContextGuardrail(CustomGuardrail): + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object], + input_type: Literal["request", "response"], logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + response: Final = request_data["response"] + assert isinstance(response, ModelResponse) + request_data["visited_choices"] = (*request_data.get("visited_choices", ()), response.choices[0].index) + if response.choices[0].message.content == "forbidden": + raise HTTPException(status_code=400, detail="later choice blocked") + return inputs - with pytest.raises(UndeliverableStreamRewrite, match="the stream carries 2 choices") as raised: - await handler.process_output_streaming_response( - responses_so_far=chunks, - guardrail_to_apply=MockGuardrail(guardrail_name="test"), - litellm_logging_obj=None, - deliver_ended_stream_rewrites=True, + chunks: Final = [ModelResponseStream(choices=[StreamingChoices( + index=index, delta=Delta(content=text, tool_calls=[{ + "index": 0, "id": f"call_{index}", "type": "function", + "function": {"name": "contact", "arguments": "{}"}, + }]), finish_reason="tool_calls", + ) for index, text in enumerate(("allowed", "forbidden"))])] + request_data: Final = {"metadata": {"trace": "retained"}} + with pytest.raises(HTTPException, match="later choice blocked"): + await OpenAIChatCompletionsHandler().process_output_streaming_response( + responses_so_far=chunks, guardrail_to_apply=ContextGuardrail(guardrail_name="context"), + request_data=request_data, deliver_ended_stream_rewrites=True, + stream_transform_sink=StreamTransformSink() if transform else None, ) + assert request_data["visited_choices"] == (0, 1) + assert request_data["metadata"]["trace"] == "retained" + assert request_data["responses"] is chunks + restored: Final = request_data["response"] + assert isinstance(restored, ModelResponse) + assert tuple(choice.index for choice in restored.choices) == (0, 1) - assert raised.value.guardrail_name == "test" - assert raised.value.reason == ( - "the stream carries 2 choices and tool-call rewrites are only written back on single-choice streams" + @pytest.mark.asyncio + async def test_deliver_ended_stream_tool_rewrites_keep_choice_indices(self) -> None: + handler: Final = OpenAIChatCompletionsHandler() + chunks: Final = self._two_choice_tool_call_stream_chunks() + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MockGuardrail(guardrail_name="test"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, ) + arguments: Final = tuple( + "".join(call.function.arguments for chunk in chunks for choice in chunk.choices + if choice.index == index for call in choice.delta.tool_calls or ()) + for index in range(2) + ) + assert arguments == ('{"fruit": "PERSIMMON"}', '{"fruit": "DURIAN"}') @pytest.mark.asyncio async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self): From 12252219f3f21f0e45e92b9197d99ffd5d35d221 Mon Sep 17 00:00:00 2001 From: albertbausili Date: Fri, 16 Oct 2026 09:37:09 +0200 Subject: [PATCH 03/36] refactor(guardrails): separate completion choice inspection prerequisite --- .../guardrail_translation/base_translation.py | 10 +- .../chat/guardrail_translation/handler.py | 220 +++++++----------- .../test_openai_guardrail_handler.py | 70 ++---- 3 files changed, 110 insertions(+), 190 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 6280e636700..88c4dbfdce6 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional @@ -31,12 +31,8 @@ class StreamTransformSink: single call, not shared state. """ - mutated_text_per_choice: Mapping[int, str] = field(default_factory=dict) - holdback_per_choice: Mapping[int, int] = field(default_factory=dict) - - def record(self, texts: Mapping[int, str], holdback: Mapping[int, int]) -> None: - self.mutated_text_per_choice = texts - self.holdback_per_choice = holdback + mutated_text_per_choice: dict[int, str] = field(default_factory=dict) + holdback_per_choice: dict[int, int] = field(default_factory=dict) @dataclass(frozen=True, slots=True) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index d18fe36e3df..a12f5f164e4 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -18,11 +18,9 @@ import json import time import uuid from collections.abc import Iterator, Mapping, Sequence -from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast -from pydantic import TypeAdapter from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -48,7 +46,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( stream_item_items, unappliable_request_rewrite, ) -from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolParam +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( coerce_stream_holdback_value, ) @@ -170,8 +168,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) - guardrailed_texts: Final = guardrailed_inputs.get("texts", ()) - guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", ()) + guardrailed_texts: Final = guardrailed_inputs.get("texts", []) + guardrailed_tool_calls: Final = guardrailed_inputs.get("tool_calls", []) guardrailed_tools: Final = guardrailed_inputs.get("tools") if guardrailed_tools is not None: data["tools"] = ( @@ -260,12 +258,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation): def extract_request_tool_names(self, data: dict) -> list[str]: """Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name).""" names: Final[list[str]] = [] - for tool in data.get("tools") or (): + for tool in data.get("tools") or []: if isinstance(tool, dict) and tool.get("type") == "function": fn = tool.get("function") if isinstance(fn, dict) and fn.get("name"): names.append(str(fn["name"])) - for fn in data.get("functions") or (): + for fn in data.get("functions") or []: if isinstance(fn, dict) and fn.get("name"): names.append(str(fn["name"])) return names @@ -389,7 +387,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, - request_data: dict[str, object] | None = None, + request_data: dict | None = None, ) -> ModelResponse: """ Process output response by applying guardrails to text content. @@ -440,7 +438,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # create a standalone dict (SDK / direct-call path). if request_data is None: request_data = {"response": response} - request_data.setdefault("response", response) + else: + if "response" not in request_data: + request_data["response"] = response self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) @@ -460,7 +460,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) - guardrailed_texts: Final = guardrailed_inputs.get("texts", ()) + guardrailed_texts: Final = guardrailed_inputs.get("texts", []) returned_tool_calls: Final = guardrailed_inputs.get("tool_calls") guardrailed_tool_calls: Final[list[dict[str, object]]] = ( cast(list[dict[str, object]], returned_tool_calls) @@ -494,7 +494,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, - request_data: dict[str, object] | None = None, + request_data: dict | None = None, stream_transform_sink: StreamTransformSink | None = None, deliver_ended_stream_rewrites: bool = False, ) -> list["ModelResponseStream"]: @@ -550,14 +550,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None", user_api_key_dict: "UserAPIKeyAuth | None", - request_data: dict[str, object] | None, + request_data: dict | None, deliver_ended_stream_rewrites: bool = False, ) -> list["ModelResponseStream"]: """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) unless ``deliver_ended_stream_rewrites`` opts the ended-stream branch in.""" - has_stream_ended: Final = deliver_ended_stream_rewrites or self._first_choice_has_finished(responses_so_far) + has_stream_ended: Final = self._first_choice_has_finished(responses_so_far) if has_stream_ended: await self._process_ended_stream( @@ -623,7 +623,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) - guardrailed_texts: Final = guardrailed_inputs.get("texts", ()) + guardrailed_texts: Final = guardrailed_inputs.get("texts", []) # Step 4: Apply guardrailed text back to all streaming chunks # For each choice, replace the combined text across all chunks @@ -656,36 +656,13 @@ class OpenAIChatCompletionsHandler(BaseTranslation): model_response: Final = self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj) pre_guardrail_texts: Final = self._string_choice_contents(model_response) pre_guardrail_tool_calls: Final = self._function_tool_call_shapes(model_response) - inspection_responses: Final = ( - tuple( - model_response.model_copy(update=MappingProxyType({"choices": [choice]})) - for choice in model_response.choices - ) - if pre_guardrail_tool_calls and len(model_response.choices) > 1 - else (model_response,) + await self.process_output_response( + response=model_response, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=request_data, ) - inspection_request_data: Final = request_data if request_data is not None else {} - try: - for inspection_response in inspection_responses: - inspection_request_data["response"] = inspection_response - inspection_request_data["responses"] = ( - [ - self._narrowed_to_choice(chunk, inspection_response.choices[0].index) - for chunk in responses_so_far - ] - if len(inspection_response.choices) == 1 - else responses_so_far - ) - await self.process_output_response( - response=inspection_response, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=litellm_logging_obj, - user_api_key_dict=user_api_key_dict, - request_data=inspection_request_data, - ) - finally: - inspection_request_data["response"] = model_response - inspection_request_data["responses"] = responses_so_far if not deliver_ended_stream_rewrites: return await self._write_ended_stream_text_rewrites( @@ -786,7 +763,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None", user_api_key_dict: "UserAPIKeyAuth | None", - request_data: dict[str, object] | None, + request_data: dict | None, sink: StreamTransformSink, ) -> None: """Run the guardrail over the raw accumulated text and report the @@ -798,7 +775,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): """ raw_by_index: Final = self._accumulate_string_content_by_choice_index(responses_so_far) if not raw_by_index: - sink.record(MappingProxyType({}), MappingProxyType({})) + sink.mutated_text_per_choice = {} + sink.holdback_per_choice = {} return # Fix #2 — sort by StreamingChoices.index so an n>1 stream that emits @@ -810,72 +788,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation): indices: Final = sorted(raw_by_index.keys()) texts_to_check: Final = [raw_by_index[i] for i in indices] - resolved_request_data: Final[dict[str, object]] = ( - request_data if request_data is not None else {"responses": responses_so_far} - ) - if "responses" not in resolved_request_data: - resolved_request_data["responses"] = responses_so_far - self.merge_user_api_key_metadata_into_request(resolved_request_data, user_api_key_dict) + if request_data is None: + request_data = {"responses": responses_so_far} + elif "responses" not in request_data: + request_data["responses"] = responses_so_far + self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) - if self._streamed_tool_call_fingerprints(responses_so_far): - assembled: Final = self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj) - resolved_request_data["response"] = assembled - if len(assembled.choices) > 1: - choice_rounds: Final = tuple( - ( - choice, - StreamTransformSink(), - [self._narrowed_to_choice(chunk, choice.index) for chunk in responses_so_far], - ) - for choice in assembled.choices - ) - try: - for choice, choice_sink, choice_chunks in choice_rounds: - resolved_request_data["response"] = assembled.model_copy( - update=MappingProxyType({"choices": [choice]}) - ) - resolved_request_data["responses"] = choice_chunks - await self._process_streaming_transform( - responses_so_far=choice_chunks, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=litellm_logging_obj, - user_api_key_dict=user_api_key_dict, - request_data=resolved_request_data, - sink=choice_sink, - ) - finally: - resolved_request_data["response"] = assembled - resolved_request_data["responses"] = responses_so_far - merged_texts: Final = MappingProxyType( - dict( - chain.from_iterable( - choice_sink.mutated_text_per_choice.items() for _, choice_sink, _ in choice_rounds - ) - ) - ) - merged_holdback: Final = MappingProxyType( - dict( - chain.from_iterable( - choice_sink.holdback_per_choice.items() for _, choice_sink, _ in choice_rounds - ) - ) - ) - sink.record(merged_texts, merged_holdback) - return - tool_calls: Final = chain.from_iterable(choice.message.tool_calls or () for choice in assembled.choices) - inputs["tool_calls"] = TypeAdapter(list[ChatCompletionToolCallChunk]).validate_python( - tuple( - MappingProxyType({"index": index, **converted}) - for index, tool_call in enumerate(tool_calls) - if (converted := self._convert_tool_call_to_dict(tool_call)) is not None - ) - ) if responses_so_far and getattr(responses_so_far[0], "model", None): inputs["model"] = responses_so_far[0].model guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, - request_data=resolved_request_data, + request_data=request_data, input_type="response", logging_obj=litellm_logging_obj, ) @@ -895,14 +819,13 @@ class OpenAIChatCompletionsHandler(BaseTranslation): len(texts_to_check), ) - holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or () - transformed_texts: Final = MappingProxyType( - {idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts)} - ) - transformed_holdback: Final = MappingProxyType( - {indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback)} - ) - sink.record(transformed_texts, transformed_holdback) + holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or [] + sink.mutated_text_per_choice = { + idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts) + } + sink.holdback_per_choice = { + 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)) @@ -1222,18 +1145,20 @@ class OpenAIChatCompletionsHandler(BaseTranslation): def _function_tool_call_fragments( responses_so_far: Sequence["ModelResponseStream"], ) -> tuple[tuple[ChatCompletionDeltaToolCall, ...], ...]: + """Group the stream's function tool-call fragments by their tool-call index, in + the index order ``stream_chunk_builder`` lists the rebuilt tool calls, keeping + only the indices the builder keeps (an id and a name somewhere in the stream).""" fragments: Final = tuple( - (choice.index, tool_call) + tool_call for response in responses_so_far for choice in response.choices for tool_call in choice.delta.tool_calls or () if isinstance(tool_call, ChatCompletionDeltaToolCall) ) - identified: Final = frozenset((choice, fragment.index) for choice, fragment in fragments if fragment.id) - named: Final = frozenset((choice, fragment.index) for choice, fragment in fragments if fragment.function.name) + identified: Final = frozenset(fragment.index for fragment in fragments if fragment.id) + named: Final = frozenset(fragment.index for fragment in fragments if fragment.function.name) return tuple( - tuple(fragment for choice, fragment in fragments if (choice, fragment.index) == key) - for key in sorted(identified & named) + tuple(fragment for fragment in fragments if fragment.index == index) for index in sorted(identified & named) ) def _write_ended_stream_tool_call_rewrites( @@ -1243,10 +1168,28 @@ class OpenAIChatCompletionsHandler(BaseTranslation): pre_guardrail_tool_calls: tuple[tuple[str | None, str], ...], guardrail_name: str, ) -> None: + """Write ended-stream guardrail tool-call rewrites back across the buffered + chunks: the rewritten name and full arguments land in the tool call's first + fragment and the arguments of its later fragments are blanked, mirroring the + text write-back. A rewrite on a stream carrying more than one distinct choice + index, or whose fragments do not line up with the rebuilt tool calls, is + reported as undeliverable, so the pipeline executor discards it and releases + the original chunks.""" post_guardrail_tool_calls: Final = self._function_tool_call_shapes(guardrailed_response) if post_guardrail_tool_calls == pre_guardrail_tool_calls: return + stream_choice_indices: Final = frozenset( + choice.index for response in responses_so_far for choice in response.choices + ) fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far) + if len(stream_choice_indices) != 1: + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite( + guardrail_name, + f"the stream carries {len(stream_choice_indices)} choices and tool-call rewrites are only written " + "back on single-choice streams", + ) if len(fragments_by_tool_call) != len(post_guardrail_tool_calls): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite @@ -1269,7 +1212,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): async def _apply_guardrail_responses_to_output_streaming( self, responses: list["ModelResponseStream"], - guardrailed_texts: Sequence[str], + guardrailed_texts: list[str], task_mappings: list[tuple[int, int | None]], ) -> None: """ @@ -1287,9 +1230,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Override this method to customize how responses are applied to streaming responses. """ # Build a mapping of what guardrailed text to use for each (choice_idx, content_idx) - guardrail_map: Final = MappingProxyType( - {task_mappings[index]: text for index, text in enumerate(guardrailed_texts)} - ) + guardrail_map: Final[dict[tuple[int, int | None], str]] = {} + for task_idx, guardrail_response in enumerate(guardrailed_texts): + mapping = task_mappings[task_idx] + choice_idx = cast(int, mapping[0]) + content_idx_optional = cast(int | None, mapping[1]) + guardrail_map[(choice_idx, content_idx_optional)] = guardrail_response # Track which choices we've already set the guardrailed text for # Key: (choice_idx, content_idx), Value: boolean (True if already set) @@ -1299,20 +1245,34 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # choice by its index field: on n>1 streams a chunk usually carries one # choice at list position 0 whose index names the logical choice. for response in responses: - for choice, message in ( - (choice, choice.delta if isinstance(choice, StreamingChoices) else choice.message) - for choice in response.choices - if isinstance(choice, (StreamingChoices, Choices)) - ): - content = message.content + for choice in response.choices: + if isinstance(choice, litellm.StreamingChoices): + content = choice.delta.content + elif isinstance(choice, litellm.Choices): + content = choice.message.content + else: + continue + if content is None: continue if isinstance(content, str): + # String content str_key: tuple[int, int | None] = (choice.index, None) if str_key in guardrail_map: - message.content = guardrail_map[str_key] if str_key not in already_set else "" - already_set[str_key] = True + if str_key not in already_set: + # First chunk - set the complete guardrailed text + if isinstance(choice, litellm.StreamingChoices): + choice.delta.content = guardrail_map[str_key] + elif isinstance(choice, litellm.Choices): + choice.message.content = guardrail_map[str_key] + already_set[str_key] = True + else: + # Subsequent chunks - clear the content + if isinstance(choice, litellm.StreamingChoices): + choice.delta.content = "" + elif isinstance(choice, litellm.Choices): + choice.message.content = "" elif isinstance(content, list): # List content - handle each content item diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index bffe6e4c123..5d71cbe9afc 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -7,7 +7,7 @@ with guardrail transformations, including tool calls. import json from collections.abc import Mapping -from typing import Any, Final, Literal, Optional +from typing import Any, Literal, Optional import pytest @@ -1395,65 +1395,29 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: ) assert [ - [(tool_call["id"], tool_call["function"]["arguments"]) for tool_call in inputs["tool_calls"]] - for inputs in guardrail.seen_inputs - ] == [[("call_1", '{"fruit": "persimmon"}')], [("call_2", '{"fruit": "durian"}')]] + (tool_call["id"], tool_call["function"]["arguments"]) + for tool_call in guardrail.seen_inputs[-1]["tool_calls"] + ] == [("call_1", '{"fruit": "persimmon"}'), ("call_2", '{"fruit": "durian"}')] @pytest.mark.asyncio - @pytest.mark.parametrize("transform", [False, True]) - async def test_each_choice_refreshes_shared_guardrail_response_context(self, transform: bool) -> None: - from fastapi import HTTPException - from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink - from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + async def test_deliver_ended_stream_tool_call_rewrite_on_multi_choice_stream_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite - class ContextGuardrail(CustomGuardrail): - async def apply_guardrail( - self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object], - input_type: Literal["request", "response"], logging_obj: object = None, - ) -> GenericGuardrailAPIInputs: - response: Final = request_data["response"] - assert isinstance(response, ModelResponse) - request_data["visited_choices"] = (*request_data.get("visited_choices", ()), response.choices[0].index) - if response.choices[0].message.content == "forbidden": - raise HTTPException(status_code=400, detail="later choice blocked") - return inputs + handler = OpenAIChatCompletionsHandler() + chunks = self._two_choice_tool_call_stream_chunks() - chunks: Final = [ModelResponseStream(choices=[StreamingChoices( - index=index, delta=Delta(content=text, tool_calls=[{ - "index": 0, "id": f"call_{index}", "type": "function", - "function": {"name": "contact", "arguments": "{}"}, - }]), finish_reason="tool_calls", - ) for index, text in enumerate(("allowed", "forbidden"))])] - request_data: Final = {"metadata": {"trace": "retained"}} - with pytest.raises(HTTPException, match="later choice blocked"): - await OpenAIChatCompletionsHandler().process_output_streaming_response( - responses_so_far=chunks, guardrail_to_apply=ContextGuardrail(guardrail_name="context"), - request_data=request_data, deliver_ended_stream_rewrites=True, - stream_transform_sink=StreamTransformSink() if transform else None, + with pytest.raises(UndeliverableStreamRewrite, match="the stream carries 2 choices") as raised: + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MockGuardrail(guardrail_name="test"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, ) - assert request_data["visited_choices"] == (0, 1) - assert request_data["metadata"]["trace"] == "retained" - assert request_data["responses"] is chunks - restored: Final = request_data["response"] - assert isinstance(restored, ModelResponse) - assert tuple(choice.index for choice in restored.choices) == (0, 1) - @pytest.mark.asyncio - async def test_deliver_ended_stream_tool_rewrites_keep_choice_indices(self) -> None: - handler: Final = OpenAIChatCompletionsHandler() - chunks: Final = self._two_choice_tool_call_stream_chunks() - await handler.process_output_streaming_response( - responses_so_far=chunks, - guardrail_to_apply=MockGuardrail(guardrail_name="test"), - litellm_logging_obj=None, - deliver_ended_stream_rewrites=True, + assert raised.value.guardrail_name == "test" + assert raised.value.reason == ( + "the stream carries 2 choices and tool-call rewrites are only written back on single-choice streams" ) - arguments: Final = tuple( - "".join(call.function.arguments for chunk in chunks for choice in chunk.choices - if choice.index == index for call in choice.delta.tool_calls or ()) - for index in range(2) - ) - assert arguments == ('{"fruit": "PERSIMMON"}', '{"fruit": "DURIAN"}') @pytest.mark.asyncio async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self): From 68055731029c617292923b22d2764c1285d1a99d Mon Sep 17 00:00:00 2001 From: albertbausili Date: Fri, 16 Oct 2026 09:41:45 +0200 Subject: [PATCH 04/36] refactor(guardrails): remove redundant validated guardrail check --- .../guardrail_hooks/unified_guardrail/unified_guardrail.py | 1 - 1 file changed, 1 deletion(-) 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 9134a1cd0f3..a33b7fd3826 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -1402,7 +1402,6 @@ class UnifiedLLMGuardrails(CustomLogger): chunks_yielded and not verdict_settled and translation_class is not None - and isinstance(resolved_guardrail, CustomGuardrail) ): await self._scan_released_stream_after_disconnect( endpoint_translation=translation_class(), From 9901032e4cbbfa99262dca7ff622365c48ad576d Mon Sep 17 00:00:00 2001 From: albertbausili Date: Fri, 16 Oct 2026 09:41:49 +0200 Subject: [PATCH 05/36] style(guardrails): format disconnect guard condition --- .../guardrail_hooks/unified_guardrail/unified_guardrail.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) 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 a33b7fd3826..786e90c6b28 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -1398,11 +1398,7 @@ class UnifiedLLMGuardrails(CustomLogger): yield error_item except (GeneratorExit, asyncio.CancelledError): translation_class: Final = None if call_type is None else mappings.get(CallTypes(call_type)) - if ( - chunks_yielded - and not verdict_settled - and translation_class is not None - ): + if chunks_yielded and not verdict_settled and translation_class is not None: await self._scan_released_stream_after_disconnect( endpoint_translation=translation_class(), responses_released=responses_yielded, From c99cc59ddaeb29f0f30ec0693b9314e67b315e6c Mon Sep 17 00:00:00 2001 From: albertbausili Date: Fri, 16 Oct 2026 09:43:10 +0200 Subject: [PATCH 06/36] fix(guardrails): preserve the existing stream sink value types --- .../base_llm/guardrail_translation/base_translation.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 6280e636700..f3dd67a2bb9 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -31,12 +31,12 @@ class StreamTransformSink: single call, not shared state. """ - mutated_text_per_choice: Mapping[int, str] = field(default_factory=dict) - holdback_per_choice: Mapping[int, int] = field(default_factory=dict) + mutated_text_per_choice: dict[int, str] = field(default_factory=dict) + holdback_per_choice: dict[int, int] = field(default_factory=dict) def record(self, texts: Mapping[int, str], holdback: Mapping[int, int]) -> None: - self.mutated_text_per_choice = texts - self.holdback_per_choice = holdback + self.mutated_text_per_choice = dict(texts) + self.holdback_per_choice = dict(holdback) @dataclass(frozen=True, slots=True) From b2b61d21999dca4fa086d7b19f5fd41f5895b990 Mon Sep 17 00:00:00 2001 From: albertbausili Date: Fri, 16 Oct 2026 10:00:45 +0200 Subject: [PATCH 07/36] test(guardrails): verify disconnect inspection retains tool context --- .../unified_guardrails/test_unified_guardrail.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 3efbe6475b2..9a4fa3ba91b 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -2706,7 +2706,12 @@ class TestStreamingClientDisconnectScan: assert [call.function.name for call in released[0].choices[0].delta.tool_calls] == ["get_weather"] assert guardrail.received_texts == [["BLOCKME"]], guardrail.received_texts - assert guardrail.received_tool_calls == [], guardrail.received_tool_calls + assert guardrail.received_tool_calls == [[{ + "id": "call_1", + "type": "function", + "index": 0, + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + }]] @pytest.mark.asyncio async def test_closing_after_a_released_tool_call_under_incremental_diff_scans_no_held_back_text(self): From 37d6a1cacd3493ffdf50430a2896672776fa105b Mon Sep 17 00:00:00 2001 From: albertbausili Date: Fri, 16 Oct 2026 10:03:34 +0200 Subject: [PATCH 08/36] refactor(guardrails): type the shared constructor extension fields --- litellm/integrations/custom_guardrail.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 45b20211ac5..583e9108af7 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -256,8 +256,8 @@ class CustomGuardrail(CustomLogger): scan_raw_request: bool = False, only_scan_new_messages: bool = False, timeout: float | None = None, - **kwargs, - ): + **kwargs: object, + ) -> None: """ Initialize the CustomGuardrail class @@ -303,7 +303,7 @@ class CustomGuardrail(CustomLogger): self.run_in_parallel: bool = run_in_parallel self.scan_raw_request: bool = scan_raw_request self.only_scan_new_messages: bool = only_scan_new_messages - stream_scope_arg: Final[object] = cast(object, kwargs.pop("stream_scope", None)) # cast-ok: config + stream_scope_arg: Final[object] = kwargs.pop("stream_scope", None) self.apply_stream_scope(stream_scope_arg) self.logging_only_scope = None self.logging_only_continue_on_input_failure = False From 3b3fc05d5cc04cba068dc3122659881e7272afdc Mon Sep 17 00:00:00 2001 From: albertbausili Date: Fri, 16 Oct 2026 10:03:36 +0200 Subject: [PATCH 09/36] refactor(guardrails): collect checked stream tails outside the replay loop --- .../unified_guardrail/unified_guardrail.py | 35 ++++++++++++++----- 1 file changed, 26 insertions(+), 9 deletions(-) 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 786e90c6b28..7b1d646bdec 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -809,19 +809,17 @@ class UnifiedLLMGuardrails(CustomLogger): if self._chunk_has_tool_calls(buffered_item) ) - async def checked_tail() -> AsyncGenerator[object, None]: - try: - async for tail_chunk in self._emit_stream_tail( + tail_chunks: Final = tuple( + [ + chunk + async for chunk in self._checked_stream_tail( last_chunk=last_chunk, final_round=_round, responses_so_far=responses_so_far, responses_yielded=responses_yielded, - ): - yield tail_chunk - except _StreamTerminated as exc: - yield exc - - tail_chunks: Final = tuple([chunk async for chunk in checked_tail()]) + ) + ] + ) if tail_chunks and isinstance(tail_chunks[-1], _StreamTerminated): for error_chunk in tail_chunks[:-1]: yield error_chunk @@ -843,6 +841,25 @@ class UnifiedLLMGuardrails(CustomLogger): except _StreamTerminated: return + async def _checked_stream_tail( + self, + *, + last_chunk: object | None, + final_round: Callable[[object, bool], AsyncGenerator[object, None]], + responses_so_far: Sequence[object], + responses_yielded: list[object], + ) -> AsyncGenerator[object, None]: + try: + async for chunk in self._emit_stream_tail( + last_chunk=last_chunk, + final_round=final_round, + responses_so_far=responses_so_far, + responses_yielded=responses_yielded, + ): + yield chunk + except _StreamTerminated as exc: + yield exc + async def _emit_stream_tail( self, *, From a07ec080365c633968a15c0198f6afac6c4683d1 Mon Sep 17 00:00:00 2001 From: albertbausili Date: Fri, 16 Oct 2026 10:09:17 +0200 Subject: [PATCH 10/36] fix(guardrails): forward typed message logging flags to the logger --- litellm/integrations/custom_guardrail.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 583e9108af7..f7457dd4ff9 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -256,6 +256,8 @@ class CustomGuardrail(CustomLogger): scan_raw_request: bool = False, only_scan_new_messages: bool = False, timeout: float | None = None, + turn_off_message_logging: bool = False, + message_logging: bool = True, **kwargs: object, ) -> None: """ @@ -323,7 +325,11 @@ class CustomGuardrail(CustomLogger): "(default) to enforce validation and fail at startup.", validation_error, ) - super().__init__(**kwargs) + super().__init__( + turn_off_message_logging=turn_off_message_logging, + message_logging=message_logging, + **kwargs, + ) def render_violation_message(self, default: str, context: Mapping[str, object] | None = None) -> str: """Return a custom violation message if template is configured.""" From fa10963e3d3fd03020c07a35d9c5d09bae3a7b80 Mon Sep 17 00:00:00 2001 From: albertbausili Date: Fri, 16 Oct 2026 10:11:45 +0200 Subject: [PATCH 11/36] refactor(guardrails): keep legacy constructor typing outside choice inspection --- litellm/integrations/custom_guardrail.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index f7457dd4ff9..45b20211ac5 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -256,10 +256,8 @@ class CustomGuardrail(CustomLogger): scan_raw_request: bool = False, only_scan_new_messages: bool = False, timeout: float | None = None, - turn_off_message_logging: bool = False, - message_logging: bool = True, - **kwargs: object, - ) -> None: + **kwargs, + ): """ Initialize the CustomGuardrail class @@ -305,7 +303,7 @@ class CustomGuardrail(CustomLogger): self.run_in_parallel: bool = run_in_parallel self.scan_raw_request: bool = scan_raw_request self.only_scan_new_messages: bool = only_scan_new_messages - stream_scope_arg: Final[object] = kwargs.pop("stream_scope", None) + stream_scope_arg: Final[object] = cast(object, kwargs.pop("stream_scope", None)) # cast-ok: config self.apply_stream_scope(stream_scope_arg) self.logging_only_scope = None self.logging_only_continue_on_input_failure = False @@ -325,11 +323,7 @@ class CustomGuardrail(CustomLogger): "(default) to enforce validation and fail at startup.", validation_error, ) - super().__init__( - turn_off_message_logging=turn_off_message_logging, - message_logging=message_logging, - **kwargs, - ) + super().__init__(**kwargs) def render_violation_message(self, default: str, context: Mapping[str, object] | None = None) -> str: """Return a custom violation message if template is configured.""" From 296681bdc330a7d0656d75a48400d2aa02aab416 Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Wed, 21 Oct 2026 23:26:39 +0200 Subject: [PATCH 12/36] fix(guardrails): pass the resolved guardrail to the end-of-stream scan check --- .../guardrail_hooks/unified_guardrail/unified_guardrail.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 4a985234a6c..f86d718bd86 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -1385,7 +1385,7 @@ class UnifiedLLMGuardrails(CustomLogger): end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far) verdict_settled = True if _is_redundant_scan(end_scan_key, last_scan_key) and not _owes_ended_scan( - endpoint_translation, guardrail_to_apply, responses_so_far, last_scan_key + endpoint_translation, resolved_guardrail, responses_so_far, last_scan_key ): verbose_proxy_logger.debug( "Skipping end-of-stream scan for guardrail %s: the last sampled round already scanned it all", From 7254b3ad5ee17b6270a474ee2f8b84833be4cbfd Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Wed, 28 Oct 2026 12:02:25 +0100 Subject: [PATCH 13/36] fix(guardrails): inspect every choice before raising an undeliverable rewrite --- .../chat/guardrail_translation/handler.py | 75 ++++++++++++---- .../test_openai_guardrail_handler.py | 89 ++++++++++++++++++- 2 files changed, 143 insertions(+), 21 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index a4a3956df94..9ad42b18c70 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -669,7 +669,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if pre_guardrail_tool_calls and len(model_response.choices) > 1 else (model_response,) ) + # Local import avoids a module-level cyclic import with the policy engine. + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + inspection_request_data: Final = request_data if request_data is not None else {} + missing: Final = object() + prior_response: Final = inspection_request_data.get("response", missing) + prior_responses: Final = inspection_request_data.get("responses", missing) + # An undeliverable rewrite on one choice must not end the loop: the caller releases the + # original chunks when it catches this, so leaving later choices uninspected would let a + # request with n > 1 pair a redactable choice with a prohibited one and have the whole + # response delivered unchecked. Inspect every choice, then raise the first failure. + undeliverable: UndeliverableStreamRewrite | None = None # rebind-ok: keeps the first failure try: for inspection_response in inspection_responses: inspection_request_data["response"] = inspection_response @@ -681,16 +692,31 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if len(inspection_response.choices) == 1 else responses_so_far ) - await self.process_output_response( - response=inspection_response, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=litellm_logging_obj, - user_api_key_dict=user_api_key_dict, - request_data=inspection_request_data, - ) + try: + await self.process_output_response( + response=inspection_response, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=inspection_request_data, + ) + except UndeliverableStreamRewrite as exc: + if undeliverable is None: + undeliverable = exc finally: - inspection_request_data["response"] = model_response - inspection_request_data["responses"] = responses_so_far + # Restore the caller's keys rather than overwriting them: the inspection loop owns + # these entries only for its own duration, and "responses" is not a key the + # non-streaming path sets at all. + if prior_response is missing: + inspection_request_data.pop("response", None) + else: + inspection_request_data["response"] = prior_response + if prior_responses is missing: + inspection_request_data.pop("responses", None) + else: + inspection_request_data["responses"] = prior_responses + if undeliverable is not None: + raise undeliverable if not deliver_ended_stream_rewrites: return await self._write_ended_stream_text_rewrites( @@ -835,23 +861,38 @@ class OpenAIChatCompletionsHandler(BaseTranslation): ) for choice in assembled.choices ) + # Local import avoids a module-level cyclic import with the policy engine. + from litellm.proxy.policy_engine.pipeline_executor import ( + UndeliverableStreamRewrite, + ) + + # Same rule as the ended-stream inspection loop: an undeliverable rewrite on one + # choice must not leave the rest uninspected, because the caller releases the + # original chunks when it catches this. + transform_undeliverable: UndeliverableStreamRewrite | None = None # rebind-ok: keeps the first failure try: for choice, choice_sink, choice_chunks in choice_rounds: resolved_request_data["response"] = assembled.model_copy( update=MappingProxyType({"choices": [choice]}) ) resolved_request_data["responses"] = choice_chunks - await self._process_streaming_transform( - responses_so_far=choice_chunks, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=litellm_logging_obj, - user_api_key_dict=user_api_key_dict, - request_data=resolved_request_data, - sink=choice_sink, - ) + try: + await self._process_streaming_transform( + responses_so_far=choice_chunks, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=resolved_request_data, + sink=choice_sink, + ) + except UndeliverableStreamRewrite as exc: + if transform_undeliverable is None: + transform_undeliverable = exc finally: resolved_request_data["response"] = assembled resolved_request_data["responses"] = responses_so_far + if transform_undeliverable is not None: + raise transform_undeliverable merged_texts: Final = MappingProxyType( dict( chain.from_iterable( diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index bffe6e4c123..f21e22545f5 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1433,10 +1433,91 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: ) assert request_data["visited_choices"] == (0, 1) assert request_data["metadata"]["trace"] == "retained" - assert request_data["responses"] is chunks - restored: Final = request_data["response"] - assert isinstance(restored, ModelResponse) - assert tuple(choice.index for choice in restored.choices) == (0, 1) + if transform: + # The transform path seeds "response" and "responses" up front as the context the + # guardrail reads, so they are its own values rather than leftovers. + assert request_data["responses"] is chunks + restored: Final = request_data["response"] + assert isinstance(restored, ModelResponse) + assert tuple(choice.index for choice in restored.choices) == (0, 1) + else: + # The ended-stream inspection loop owns these keys only while it runs; neither was in + # request_data beforehand, and the non-streaming path never sets "responses" at all. + assert "response" not in request_data + assert "responses" not in request_data + + @pytest.mark.asyncio + async def test_undeliverable_rewrite_still_inspects_the_remaining_choices(self) -> None: + """An undeliverable rewrite on one choice must not leave later choices uninspected. + + The caller releases the original chunks when it catches UndeliverableStreamRewrite, so + bailing out of the loop would let a request with n > 1 pair a redactable choice with a + prohibited one and have the whole response delivered unchecked. + """ + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + class UndeliverableOnFirstChoice(CustomGuardrail): + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object], + input_type: Literal["request", "response"], logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + response: Final = request_data["response"] + assert isinstance(response, ModelResponse) + index: Final = response.choices[0].index + request_data["visited_choices"] = (*request_data.get("visited_choices", ()), index) + if index == 0: + raise UndeliverableStreamRewrite("undeliverable", "tool_call_mismatch") + return inputs + + chunks: Final = [ModelResponseStream(choices=[StreamingChoices( + index=index, delta=Delta(content=text, tool_calls=[{ + "index": 0, "id": f"call_{index}", "type": "function", + "function": {"name": "contact", "arguments": "{}"}, + }]), finish_reason="tool_calls", + ) for index, text in enumerate(("redacted", "prohibited"))])] + request_data: Final = {} + with pytest.raises(UndeliverableStreamRewrite): + await OpenAIChatCompletionsHandler().process_output_streaming_response( + responses_so_far=chunks, guardrail_to_apply=UndeliverableOnFirstChoice(guardrail_name="undeliverable"), + request_data=request_data, deliver_ended_stream_rewrites=True, + ) + assert request_data["visited_choices"] == (0, 1) + + @pytest.mark.asyncio + async def test_undeliverable_rewrite_in_transform_mode_still_inspects_remaining_choices(self) -> None: + """The transform path runs its own per-choice loop and needs the same rule.""" + from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + class UndeliverableOnFirstChoice(CustomGuardrail): + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object], + input_type: Literal["request", "response"], logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + response: Final = request_data["response"] + assert isinstance(response, ModelResponse) + index: Final = response.choices[0].index + request_data["visited_choices"] = (*request_data.get("visited_choices", ()), index) + if index == 0: + raise UndeliverableStreamRewrite("undeliverable", "tool_call_mismatch") + return inputs + + chunks: Final = [ModelResponseStream(choices=[StreamingChoices( + index=index, delta=Delta(content=text, tool_calls=[{ + "index": 0, "id": f"call_{index}", "type": "function", + "function": {"name": "contact", "arguments": "{}"}, + }]), finish_reason="tool_calls", + ) for index, text in enumerate(("redacted", "prohibited"))])] + request_data: Final = {} + with pytest.raises(UndeliverableStreamRewrite): + await OpenAIChatCompletionsHandler().process_output_streaming_response( + responses_so_far=chunks, guardrail_to_apply=UndeliverableOnFirstChoice(guardrail_name="undeliverable"), + request_data=request_data, deliver_ended_stream_rewrites=True, + stream_transform_sink=StreamTransformSink(), + ) + assert request_data["visited_choices"] == (0, 1) @pytest.mark.asyncio async def test_deliver_ended_stream_tool_rewrites_keep_choice_indices(self) -> None: From c6e386cefd25122d2eb134c3fe8a0169856a1f9e Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Wed, 28 Oct 2026 12:13:25 +0100 Subject: [PATCH 14/36] test(guardrails): cover the buffered ended-stream path and per-choice holdback --- .../test_openai_guardrail_handler.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index f21e22545f5..bc8ffe07b4f 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1519,6 +1519,73 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: ) assert request_data["visited_choices"] == (0, 1) + @pytest.mark.asyncio + async def test_buffered_stream_without_finish_reason_takes_the_ended_path(self) -> None: + """A buffered stream has no finish_reason yet, so the ended path must key off the opt-in. + + Dropping ``deliver_ended_stream_rewrites`` from that condition leaves the buffered + caller on the in-flight path, where rewrites are never written back. + """ + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + class PathProbe(CustomGuardrail): + saw_assembled_response: bool = False + + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object], + input_type: Literal["request", "response"], logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + # Only the ended-stream inspection loop hands the guardrail an assembled + # "response"; the in-flight branch passes raw chunks under "responses" alone. + type(self).saw_assembled_response = isinstance(request_data.get("response"), ModelResponse) + return inputs + + chunks: Final = [ + ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content="secret"), finish_reason=None)] + ) + ] + assert OpenAIChatCompletionsHandler()._first_choice_has_finished(chunks) is False + await OpenAIChatCompletionsHandler().process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=PathProbe(guardrail_name="probe"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + assert PathProbe.saw_assembled_response is True + + @pytest.mark.asyncio + async def test_diff_mode_keeps_each_choice_holdback_with_two_choices(self) -> None: + """The per-choice sinks are merged, so an n=2 diff round must keep both holdbacks. + + Reporting an empty holdback here would release text the guardrail asked to withhold. + """ + from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + class HoldingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object], + input_type: Literal["request", "response"], logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + texts: Final = inputs.get("texts") or () + return {**inputs, "stream_holdback_chars": [len(text) for text in texts]} + + chunks: Final = [ModelResponseStream(choices=[StreamingChoices( + index=index, delta=Delta(content=text, tool_calls=[{ + "index": 0, "id": f"call_{index}", "type": "function", + "function": {"name": "contact", "arguments": "{}"}, + }]), finish_reason="tool_calls", + ) for index, text in enumerate(("alpha", "bravo-long"))])] + sink: Final = StreamTransformSink() + await OpenAIChatCompletionsHandler().process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=HoldingGuardrail(guardrail_name="holding"), + litellm_logging_obj=None, + stream_transform_sink=sink, + ) + assert sink.holdback_per_choice == {0: len("alpha"), 1: len("bravo-long")} + @pytest.mark.asyncio async def test_deliver_ended_stream_tool_rewrites_keep_choice_indices(self) -> None: handler: Final = OpenAIChatCompletionsHandler() From fd3ad67c474b162200e757a0fcdb01926035e4c4 Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Wed, 28 Oct 2026 12:22:00 +0100 Subject: [PATCH 15/36] fix(guardrails): send tool calls to the guardrail only once their arguments are complete --- .../chat/guardrail_translation/handler.py | 19 ++++--- .../test_openai_guardrail_handler.py | 49 +++++++++++++++++++ 2 files changed, 61 insertions(+), 7 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 9ad42b18c70..ef9c1eb34ab 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -909,14 +909,19 @@ class OpenAIChatCompletionsHandler(BaseTranslation): ) sink.record(merged_texts, merged_holdback) return - tool_calls: Final = chain.from_iterable(choice.message.tool_calls or () for choice in assembled.choices) - inputs["tool_calls"] = TypeAdapter(list[ChatCompletionToolCallChunk]).validate_python( - tuple( - MappingProxyType({"index": index, **converted}) - for index, tool_call in enumerate(tool_calls) - if (converted := self._convert_tool_call_to_dict(tool_call)) is not None + # Mid-round the assembled arguments are a truncated JSON fragment such as `{"ci`, which + # main never handed to a guardrail: one that parses them either errors or judges + # arguments that do not exist yet. Send tool calls only once the stream has finished and + # the arguments are complete; the text of those rounds is still inspected above. + if self._first_choice_has_finished(responses_so_far): + tool_calls: Final = chain.from_iterable(choice.message.tool_calls or () for choice in assembled.choices) + inputs["tool_calls"] = TypeAdapter(list[ChatCompletionToolCallChunk]).validate_python( + tuple( + MappingProxyType({"index": index, **converted}) + for index, tool_call in enumerate(tool_calls) + if (converted := self._convert_tool_call_to_dict(tool_call)) is not None + ) ) - ) if responses_so_far and getattr(responses_so_far[0], "model", None): inputs["model"] = responses_so_far[0].model guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index bc8ffe07b4f..a503bbf386e 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1586,6 +1586,55 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: ) assert sink.holdback_per_choice == {0: len("alpha"), 1: len("bravo-long")} + @pytest.mark.asyncio + async def test_diff_mode_withholds_tool_calls_until_their_arguments_are_complete(self) -> None: + """Mid-round the assembled arguments are a truncated JSON fragment, so they are not sent. + + Sending them would hand the guardrail something like `{"ci`, which main never did. + """ + from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + class RecordingGuardrail(CustomGuardrail): + seen: list[object] = [] + + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object], + input_type: Literal["request", "response"], logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + type(self).seen.append(inputs.get("tool_calls")) + return inputs + + def chunk(text: str, arguments: str, finish_reason: str | None) -> ModelResponseStream: + # The text keeps the guardrail in play for the partial round; without any text the + # transform path has nothing to inspect and never calls it at all. + return ModelResponseStream(choices=[StreamingChoices( + index=0, delta=Delta(content=text, tool_calls=[{ + "index": 0, "id": "call_0", "type": "function", + "function": {"name": "weather", "arguments": arguments}, + }]), finish_reason=finish_reason, + )]) + + RecordingGuardrail.seen = [] + handler: Final = OpenAIChatCompletionsHandler() + guardrail: Final = RecordingGuardrail(guardrail_name="recording") + partial: Final = [chunk("looking", '{"ci', None)] + await handler.process_output_streaming_response( + responses_so_far=partial, guardrail_to_apply=guardrail, + litellm_logging_obj=None, stream_transform_sink=StreamTransformSink(), + ) + assert RecordingGuardrail.seen == [None] + + RecordingGuardrail.seen = [] + complete: Final = [chunk("looking", '{"ci', None), chunk(" it up", 'ty": "Paris"}', "tool_calls")] + await handler.process_output_streaming_response( + responses_so_far=complete, guardrail_to_apply=guardrail, + litellm_logging_obj=None, stream_transform_sink=StreamTransformSink(), + ) + sent: Final = RecordingGuardrail.seen[-1] + assert sent is not None + assert [call["function"]["arguments"] for call in sent] == ['{"city": "Paris"}'] + @pytest.mark.asyncio async def test_deliver_ended_stream_tool_rewrites_keep_choice_indices(self) -> None: handler: Final = OpenAIChatCompletionsHandler() From ff5644b3d873bfce95e6ff300432631d135a2005 Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Wed, 28 Oct 2026 12:32:06 +0100 Subject: [PATCH 16/36] perf(guardrails): rebuild the stream only when a round consumes it --- .../chat/guardrail_translation/handler.py | 22 +++++++-- .../test_openai_guardrail_handler.py | 47 +++++++++++++++++++ 2 files changed, 65 insertions(+), 4 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index ef9c1eb34ab..a43c6b193b6 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -850,9 +850,23 @@ class OpenAIChatCompletionsHandler(BaseTranslation): inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if self._streamed_tool_call_fingerprints(responses_so_far): - assembled: Final = self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj) - resolved_request_data["response"] = assembled - if len(assembled.choices) > 1: + # Rebuilding folds every chunk accumulated so far, so doing it on every sampled round is + # quadratic in stream length: a 400-chunk stream costs ~44x one end-of-stream rebuild. + # main rebuilds only on the ended path, so rebuild here only when this round consumes the + # assembled response -- to isolate choices, or to hand over complete tool calls once the + # stream has finished. A single-choice round in flight needs neither. + stream_ended: Final = self._first_choice_has_finished(responses_so_far) + streamed_choice_indices: Final = frozenset( + choice.index for chunk in responses_so_far for choice in chunk.choices + ) + assembled: Final = ( + None + if len(streamed_choice_indices) == 1 and not stream_ended + else self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj) + ) + if assembled is not None: + resolved_request_data["response"] = assembled + if assembled is not None and len(assembled.choices) > 1: choice_rounds: Final = tuple( ( choice, @@ -913,7 +927,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # main never handed to a guardrail: one that parses them either errors or judges # arguments that do not exist yet. Send tool calls only once the stream has finished and # the arguments are complete; the text of those rounds is still inspected above. - if self._first_choice_has_finished(responses_so_far): + if assembled is not None and stream_ended: tool_calls: Final = chain.from_iterable(choice.message.tool_calls or () for choice in assembled.choices) inputs["tool_calls"] = TypeAdapter(list[ChatCompletionToolCallChunk]).validate_python( tuple( diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index a503bbf386e..137383c3483 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1635,6 +1635,53 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert sent is not None assert [call["function"]["arguments"] for call in sent] == ['{"city": "Paris"}'] + @pytest.mark.asyncio + async def test_diff_mode_does_not_rebuild_the_stream_on_in_flight_rounds(self) -> None: + """Rebuilding folds every chunk so far, so a per-round rebuild is quadratic in length. + + main rebuilds only on the ended path; an in-flight single-choice round consumes nothing + from the assembled response, so it must not pay for one. + """ + from unittest.mock import patch + + from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + def chunk(text: str, finish_reason: str | None) -> ModelResponseStream: + return ModelResponseStream(choices=[StreamingChoices( + index=0, delta=Delta(content=text, tool_calls=[{ + "index": 0, "id": "call_0", "type": "function", + "function": {"name": "weather", "arguments": '{"c'}, + }]), finish_reason=finish_reason, + )]) + + handler: Final = OpenAIChatCompletionsHandler() + rebuilds: Final[list[int]] = [] + original: Final = OpenAIChatCompletionsHandler._rebuild_ended_stream_per_choice + + def counting_rebuild(responses_so_far, litellm_logging_obj): # type: ignore[no-untyped-def] # test spy + rebuilds.append(len(responses_so_far)) + return original(responses_so_far, litellm_logging_obj) + + # The real attribute is a staticmethod, so the spy has to be one too: a plain function + # would be bound and receive the handler as its first argument. + with patch.object( + OpenAIChatCompletionsHandler, "_rebuild_ended_stream_per_choice", staticmethod(counting_rebuild) + ): + await handler.process_output_streaming_response( + responses_so_far=[chunk("looking", None)], + guardrail_to_apply=MockGuardrail(guardrail_name="test"), + litellm_logging_obj=None, stream_transform_sink=StreamTransformSink(), + ) + assert rebuilds == [] + + await handler.process_output_streaming_response( + responses_so_far=[chunk("looking", None), chunk(" it up", "tool_calls")], + guardrail_to_apply=MockGuardrail(guardrail_name="test"), + litellm_logging_obj=None, stream_transform_sink=StreamTransformSink(), + ) + assert rebuilds == [2] + @pytest.mark.asyncio async def test_deliver_ended_stream_tool_rewrites_keep_choice_indices(self) -> None: handler: Final = OpenAIChatCompletionsHandler() From f2342b832d051a4800c03e737981bb932bd8d2f3 Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Wed, 28 Oct 2026 13:09:43 +0100 Subject: [PATCH 17/36] feat(guardrails): label per-choice log entries and gate tool calls on complete arguments --- litellm/integrations/custom_guardrail.py | 3 ++ .../chat/guardrail_translation/handler.py | 47 +++++++++++++++++-- litellm/types/utils.py | 3 ++ .../test_openai_guardrail_handler.py | 30 ++++++++++++ 4 files changed, 79 insertions(+), 4 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9e2dc902d55..7b1eaca13af 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1414,6 +1414,9 @@ class CustomGuardrail(CustomLogger): end_time=end_time, duration=duration, masked_entity_count=masked_entity_count, + # Set by callers that inspect one choice at a time, so the per-choice entries of an + # n > 1 response can be told apart. + choice_index=request_data.get("guardrail_choice_index"), **(tracing_detail or {}), ) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index a43c6b193b6..be698093fec 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -676,6 +676,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): missing: Final = object() prior_response: Final = inspection_request_data.get("response", missing) prior_responses: Final = inspection_request_data.get("responses", missing) + prior_choice_index: Final = inspection_request_data.get("guardrail_choice_index", missing) # An undeliverable rewrite on one choice must not end the loop: the caller releases the # original chunks when it catches this, so leaving later choices uninspected would let a # request with n > 1 pair a redactable choice with a prohibited one and have the whole @@ -684,6 +685,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): try: for inspection_response in inspection_responses: inspection_request_data["response"] = inspection_response + # Labels this round's guardrail log entry and span with the choice it inspected. + inspection_request_data["guardrail_choice_index"] = ( + inspection_response.choices[0].index if len(inspection_response.choices) == 1 else None + ) inspection_request_data["responses"] = ( [ self._narrowed_to_choice(chunk, inspection_response.choices[0].index) @@ -715,6 +720,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): inspection_request_data.pop("responses", None) else: inspection_request_data["responses"] = prior_responses + if prior_choice_index is missing: + inspection_request_data.pop("guardrail_choice_index", None) + else: + inspection_request_data["guardrail_choice_index"] = prior_choice_index if undeliverable is not None: raise undeliverable if not deliver_ended_stream_rewrites: @@ -731,6 +740,31 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_name=guardrail_to_apply.guardrail_name or "unknown", ) + @staticmethod + def _streamed_tool_arguments_are_complete(responses_so_far: Sequence["ModelResponseStream"]) -> bool: + """Whether every streamed tool call now carries arguments that parse as JSON. + + Streamed argument fragments are only valid JSON once the last one has arrived, so this is + the completeness test: it holds back `{"ci` and lets `{"city": "Paris"}` through whether + the stream ended normally or was cut short by a disconnect. It reads the raw chunks so the + answer is available before deciding whether to rebuild. + """ + fragments: Final[dict[tuple[int, int], list[str]]] = {} # mutable-ok: local accumulator + for chunk in responses_so_far: + for choice in chunk.choices: + for call in getattr(choice.delta, "tool_calls", None) or (): + arguments = getattr(getattr(call, "function", None), "arguments", None) + if isinstance(arguments, str) and arguments: + fragments.setdefault((choice.index, getattr(call, "index", 0)), []).append(arguments) + if not fragments: + return False + for parts in fragments.values(): + try: + json.loads("".join(parts)) + except ValueError: + return False + return True + @staticmethod def _rebuild_ended_stream_per_choice( responses_so_far: Sequence["ModelResponseStream"], @@ -859,9 +893,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): streamed_choice_indices: Final = frozenset( choice.index for chunk in responses_so_far for choice in chunk.choices ) + tool_arguments_complete: Final = self._streamed_tool_arguments_are_complete(responses_so_far) assembled: Final = ( None - if len(streamed_choice_indices) == 1 and not stream_ended + if len(streamed_choice_indices) == 1 and not stream_ended and not tool_arguments_complete else self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj) ) if assembled is not None: @@ -890,6 +925,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): update=MappingProxyType({"choices": [choice]}) ) resolved_request_data["responses"] = choice_chunks + resolved_request_data["guardrail_choice_index"] = choice.index try: await self._process_streaming_transform( responses_so_far=choice_chunks, @@ -905,6 +941,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): finally: resolved_request_data["response"] = assembled resolved_request_data["responses"] = responses_so_far + resolved_request_data.pop("guardrail_choice_index", None) if transform_undeliverable is not None: raise transform_undeliverable merged_texts: Final = MappingProxyType( @@ -925,9 +962,11 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return # Mid-round the assembled arguments are a truncated JSON fragment such as `{"ci`, which # main never handed to a guardrail: one that parses them either errors or judges - # arguments that do not exist yet. Send tool calls only once the stream has finished and - # the arguments are complete; the text of those rounds is still inspected above. - if assembled is not None and stream_ended: + # arguments that do not exist yet. Completeness is what matters, not the finish reason: + # a stream cut short by a client disconnect carries whole arguments and no finish + # reason, and that call still has to be inspected. Send the calls once every one of + # them parses; the text of the earlier rounds is still inspected above. + if assembled is not None and tool_arguments_complete: tool_calls: Final = chain.from_iterable(choice.message.tool_calls or () for choice in assembled.choices) inputs["tool_calls"] = TypeAdapter(list[ChatCompletionToolCallChunk]).validate_python( tuple( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index cc7ef851111..c6eeb41a9d1 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3551,6 +3551,9 @@ class StandardLoggingGuardrailInformation(TypedDict, total=False): } """ + choice_index: ReadOnly[int | None] + """Which choice of an n > 1 response this entry inspected, when the guardrail ran per choice.""" + guardrail_id: str | None """Unique identifier for the guardrail configuration, e.g. 'gd-eu-pii-001'""" diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 137383c3483..963b7dde07b 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1682,6 +1682,36 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: ) assert rebuilds == [2] + @pytest.mark.asyncio + async def test_per_choice_inspection_labels_its_log_entries_with_the_choice_index(self) -> None: + """Without a label the per-choice entries of an n > 1 response cannot be told apart.""" + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + class IndexRecorder(CustomGuardrail): + seen: list[object] = [] + + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object], + input_type: Literal["request", "response"], logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + type(self).seen.append(request_data.get("guardrail_choice_index")) + return inputs + + IndexRecorder.seen = [] + chunks: Final = [ModelResponseStream(choices=[StreamingChoices( + index=index, delta=Delta(content=text, tool_calls=[{ + "index": 0, "id": f"call_{index}", "type": "function", + "function": {"name": "contact", "arguments": "{}"}, + }]), finish_reason="tool_calls", + ) for index, text in enumerate(("first", "second"))])] + request_data: Final = {} + await OpenAIChatCompletionsHandler().process_output_streaming_response( + responses_so_far=chunks, guardrail_to_apply=IndexRecorder(guardrail_name="indexed"), + request_data=request_data, deliver_ended_stream_rewrites=True, + ) + assert IndexRecorder.seen == [0, 1] + assert "guardrail_choice_index" not in request_data + @pytest.mark.asyncio async def test_deliver_ended_stream_tool_rewrites_keep_choice_indices(self) -> None: handler: Final = OpenAIChatCompletionsHandler() From 5a2e78e55052c1ce8ef40d5bfebe86c65d46e4f5 Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Wed, 28 Oct 2026 18:33:51 +0100 Subject: [PATCH 18/36] fix(types): include choice_index in the guardrail audit fields --- litellm/types/utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index c6eeb41a9d1..d27c4bb62a6 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3511,6 +3511,7 @@ AUDIT_GUARDRAIL_FIELDS: Final[frozenset[str]] = frozenset( "end_time", "duration", "masked_entity_count", + "choice_index", "guardrail_id", "policy_template", "detection_method", From 55545c693fe5f6183d74b1ca21c3eb00ce57a868 Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Wed, 28 Oct 2026 23:25:18 +0100 Subject: [PATCH 19/36] fix(guardrails): treat tool calls with no arguments yet as incomplete --- .../chat/guardrail_translation/handler.py | 81 ++++++-------- .../test_openai_guardrail_handler.py | 105 ++++++------------ 2 files changed, 71 insertions(+), 115 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index be698093fec..18c8b5aca9e 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -669,7 +669,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if pre_guardrail_tool_calls and len(model_response.choices) > 1 else (model_response,) ) - # Local import avoids a module-level cyclic import with the policy engine. from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite inspection_request_data: Final = request_data if request_data is not None else {} @@ -677,15 +676,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): prior_response: Final = inspection_request_data.get("response", missing) prior_responses: Final = inspection_request_data.get("responses", missing) prior_choice_index: Final = inspection_request_data.get("guardrail_choice_index", missing) - # An undeliverable rewrite on one choice must not end the loop: the caller releases the - # original chunks when it catches this, so leaving later choices uninspected would let a - # request with n > 1 pair a redactable choice with a prohibited one and have the whole - # response delivered unchecked. Inspect every choice, then raise the first failure. undeliverable: UndeliverableStreamRewrite | None = None # rebind-ok: keeps the first failure try: for inspection_response in inspection_responses: inspection_request_data["response"] = inspection_response - # Labels this round's guardrail log entry and span with the choice it inspected. inspection_request_data["guardrail_choice_index"] = ( inspection_response.choices[0].index if len(inspection_response.choices) == 1 else None ) @@ -709,9 +703,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if undeliverable is None: undeliverable = exc finally: - # Restore the caller's keys rather than overwriting them: the inspection loop owns - # these entries only for its own duration, and "responses" is not a key the - # non-streaming path sets at all. if prior_response is missing: inspection_request_data.pop("response", None) else: @@ -740,30 +731,49 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_name=guardrail_to_apply.guardrail_name or "unknown", ) + @staticmethod + def _streamed_choice_indices(responses_so_far: Sequence["ModelResponseStream"]) -> Iterator[int]: + for chunk in responses_so_far: + for choice in chunk.choices: + yield choice.index + + @staticmethod + def _streamed_argument_fragments( + responses_so_far: Sequence["ModelResponseStream"], + ) -> Iterator[tuple[tuple[int, int], str]]: + for chunk in responses_so_far: + for choice in chunk.choices: + for call in choice.delta.tool_calls or (): + arguments = getattr(getattr(call, "function", None), "arguments", None) + yield (choice.index, call.index or 0), arguments if isinstance(arguments, str) else "" + + @staticmethod + def _parses_as_json(arguments: str) -> bool: + if not arguments: + return False + try: + json.loads(arguments) + except ValueError: + return False + return True + @staticmethod def _streamed_tool_arguments_are_complete(responses_so_far: Sequence["ModelResponseStream"]) -> bool: """Whether every streamed tool call now carries arguments that parse as JSON. - Streamed argument fragments are only valid JSON once the last one has arrived, so this is - the completeness test: it holds back `{"ci` and lets `{"city": "Paris"}` through whether - the stream ended normally or was cut short by a disconnect. It reads the raw chunks so the - answer is available before deciding whether to rebuild. + A call that has announced itself but emitted no argument text yet is still in flight, so a + parallel call finishing first does not make the whole set complete. """ - fragments: Final[dict[tuple[int, int], list[str]]] = {} # mutable-ok: local accumulator - for chunk in responses_so_far: - for choice in chunk.choices: - for call in getattr(choice.delta, "tool_calls", None) or (): - arguments = getattr(getattr(call, "function", None), "arguments", None) - if isinstance(arguments, str) and arguments: - fragments.setdefault((choice.index, getattr(call, "index", 0)), []).append(arguments) + fragments: Final = tuple(OpenAIChatCompletionsHandler._streamed_argument_fragments(responses_so_far)) if not fragments: return False - for parts in fragments.values(): - try: - json.loads("".join(parts)) - except ValueError: - return False - return True + joined: Final = MappingProxyType( + { + call: "".join(text for other, text in fragments if other == call) + for call in frozenset(call for call, _ in fragments) + } + ) + return all(OpenAIChatCompletionsHandler._parses_as_json(arguments) for arguments in joined.values()) @staticmethod def _rebuild_ended_stream_per_choice( @@ -884,15 +894,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if self._streamed_tool_call_fingerprints(responses_so_far): - # Rebuilding folds every chunk accumulated so far, so doing it on every sampled round is - # quadratic in stream length: a 400-chunk stream costs ~44x one end-of-stream rebuild. - # main rebuilds only on the ended path, so rebuild here only when this round consumes the - # assembled response -- to isolate choices, or to hand over complete tool calls once the - # stream has finished. A single-choice round in flight needs neither. stream_ended: Final = self._first_choice_has_finished(responses_so_far) - streamed_choice_indices: Final = frozenset( - choice.index for chunk in responses_so_far for choice in chunk.choices - ) + streamed_choice_indices: Final = frozenset(self._streamed_choice_indices(responses_so_far)) tool_arguments_complete: Final = self._streamed_tool_arguments_are_complete(responses_so_far) assembled: Final = ( None @@ -910,14 +913,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): ) for choice in assembled.choices ) - # Local import avoids a module-level cyclic import with the policy engine. from litellm.proxy.policy_engine.pipeline_executor import ( UndeliverableStreamRewrite, ) - # Same rule as the ended-stream inspection loop: an undeliverable rewrite on one - # choice must not leave the rest uninspected, because the caller releases the - # original chunks when it catches this. transform_undeliverable: UndeliverableStreamRewrite | None = None # rebind-ok: keeps the first failure try: for choice, choice_sink, choice_chunks in choice_rounds: @@ -960,12 +959,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation): ) sink.record(merged_texts, merged_holdback) return - # Mid-round the assembled arguments are a truncated JSON fragment such as `{"ci`, which - # main never handed to a guardrail: one that parses them either errors or judges - # arguments that do not exist yet. Completeness is what matters, not the finish reason: - # a stream cut short by a client disconnect carries whole arguments and no finish - # reason, and that call still has to be inspected. Send the calls once every one of - # them parses; the text of the earlier rounds is still inspected above. if assembled is not None and tool_arguments_complete: tool_calls: Final = chain.from_iterable(choice.message.tool_calls or () for choice in assembled.choices) inputs["tool_calls"] = TypeAdapter(list[ChatCompletionToolCallChunk]).validate_python( diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 963b7dde07b..0a0bccceb71 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1434,15 +1434,11 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert request_data["visited_choices"] == (0, 1) assert request_data["metadata"]["trace"] == "retained" if transform: - # The transform path seeds "response" and "responses" up front as the context the - # guardrail reads, so they are its own values rather than leftovers. assert request_data["responses"] is chunks restored: Final = request_data["response"] assert isinstance(restored, ModelResponse) assert tuple(choice.index for choice in restored.choices) == (0, 1) else: - # The ended-stream inspection loop owns these keys only while it runs; neither was in - # request_data beforehand, and the non-streaming path never sets "responses" at all. assert "response" not in request_data assert "responses" not in request_data @@ -1535,8 +1531,6 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object], input_type: Literal["request", "response"], logging_obj: object = None, ) -> GenericGuardrailAPIInputs: - # Only the ended-stream inspection loop hands the guardrail an assembled - # "response"; the in-flight branch passes raw chunks under "responses" alone. type(self).saw_assembled_response = isinstance(request_data.get("response"), ModelResponse) return inputs @@ -1606,8 +1600,6 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: return inputs def chunk(text: str, arguments: str, finish_reason: str | None) -> ModelResponseStream: - # The text keeps the guardrail in play for the partial round; without any text the - # transform path has nothing to inspect and never calls it at all. return ModelResponseStream(choices=[StreamingChoices( index=0, delta=Delta(content=text, tool_calls=[{ "index": 0, "id": "call_0", "type": "function", @@ -1636,81 +1628,52 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert [call["function"]["arguments"] for call in sent] == ['{"city": "Paris"}'] @pytest.mark.asyncio - async def test_diff_mode_does_not_rebuild_the_stream_on_in_flight_rounds(self) -> None: - """Rebuilding folds every chunk so far, so a per-round rebuild is quadratic in length. - - main rebuilds only on the ended path; an in-flight single-choice round consumes nothing - from the assembled response, so it must not pay for one. - """ - from unittest.mock import patch - + async def test_diff_mode_withholds_parallel_tool_calls_until_every_one_is_complete(self) -> None: from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices - def chunk(text: str, finish_reason: str | None) -> ModelResponseStream: - return ModelResponseStream(choices=[StreamingChoices( - index=0, delta=Delta(content=text, tool_calls=[{ - "index": 0, "id": "call_0", "type": "function", - "function": {"name": "weather", "arguments": '{"c'}, - }]), finish_reason=finish_reason, - )]) - - handler: Final = OpenAIChatCompletionsHandler() - rebuilds: Final[list[int]] = [] - original: Final = OpenAIChatCompletionsHandler._rebuild_ended_stream_per_choice - - def counting_rebuild(responses_so_far, litellm_logging_obj): # type: ignore[no-untyped-def] # test spy - rebuilds.append(len(responses_so_far)) - return original(responses_so_far, litellm_logging_obj) - - # The real attribute is a staticmethod, so the spy has to be one too: a plain function - # would be bound and receive the handler as its first argument. - with patch.object( - OpenAIChatCompletionsHandler, "_rebuild_ended_stream_per_choice", staticmethod(counting_rebuild) - ): - await handler.process_output_streaming_response( - responses_so_far=[chunk("looking", None)], - guardrail_to_apply=MockGuardrail(guardrail_name="test"), - litellm_logging_obj=None, stream_transform_sink=StreamTransformSink(), - ) - assert rebuilds == [] - - await handler.process_output_streaming_response( - responses_so_far=[chunk("looking", None), chunk(" it up", "tool_calls")], - guardrail_to_apply=MockGuardrail(guardrail_name="test"), - litellm_logging_obj=None, stream_transform_sink=StreamTransformSink(), - ) - assert rebuilds == [2] - - @pytest.mark.asyncio - async def test_per_choice_inspection_labels_its_log_entries_with_the_choice_index(self) -> None: - """Without a label the per-choice entries of an n > 1 response cannot be told apart.""" - from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices - - class IndexRecorder(CustomGuardrail): + class RecordingGuardrail(CustomGuardrail): seen: list[object] = [] async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object], input_type: Literal["request", "response"], logging_obj: object = None, ) -> GenericGuardrailAPIInputs: - type(self).seen.append(request_data.get("guardrail_choice_index")) + type(self).seen.append(inputs.get("tool_calls")) return inputs - IndexRecorder.seen = [] - chunks: Final = [ModelResponseStream(choices=[StreamingChoices( - index=index, delta=Delta(content=text, tool_calls=[{ - "index": 0, "id": f"call_{index}", "type": "function", - "function": {"name": "contact", "arguments": "{}"}, - }]), finish_reason="tool_calls", - ) for index, text in enumerate(("first", "second"))])] - request_data: Final = {} - await OpenAIChatCompletionsHandler().process_output_streaming_response( - responses_so_far=chunks, guardrail_to_apply=IndexRecorder(guardrail_name="indexed"), - request_data=request_data, deliver_ended_stream_rewrites=True, + def call(index: int, arguments: str) -> dict[str, object]: + return { + "index": index, "id": f"call_{index}", "type": "function", + "function": {"name": "weather", "arguments": arguments}, + } + + RecordingGuardrail.seen = [] + handler: Final = OpenAIChatCompletionsHandler() + guardrail: Final = RecordingGuardrail(guardrail_name="recording") + one_still_empty: Final = [ModelResponseStream(choices=[StreamingChoices( + index=0, delta=Delta(content="looking", tool_calls=[call(0, '{"city": "Paris"}'), call(1, "")]), + finish_reason=None, + )])] + await handler.process_output_streaming_response( + responses_so_far=one_still_empty, guardrail_to_apply=guardrail, + litellm_logging_obj=None, stream_transform_sink=StreamTransformSink(), ) - assert IndexRecorder.seen == [0, 1] - assert "guardrail_choice_index" not in request_data + assert RecordingGuardrail.seen == [None] + + RecordingGuardrail.seen = [] + both_complete: Final = [ModelResponseStream(choices=[StreamingChoices( + index=0, + delta=Delta(content="looking", tool_calls=[call(0, '{"city": "Paris"}'), call(1, '{"city": "Rome"}')]), + finish_reason="tool_calls", + )])] + await handler.process_output_streaming_response( + responses_so_far=both_complete, guardrail_to_apply=guardrail, + litellm_logging_obj=None, stream_transform_sink=StreamTransformSink(), + ) + sent: Final = RecordingGuardrail.seen[-1] + assert sent is not None + assert [c["function"]["arguments"] for c in sent] == ['{"city": "Paris"}', '{"city": "Rome"}'] @pytest.mark.asyncio async def test_deliver_ended_stream_tool_rewrites_keep_choice_indices(self) -> None: From e3276c9c542a99e106d458a6cfbe03d24706745c Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Thu, 29 Oct 2026 11:39:55 +0100 Subject: [PATCH 20/36] fix(guardrails): stop padding rewritten tool calls with empty-argument deltas An ended-stream rewrite wrote the full arguments onto the first fragment and blanked the rest, so the stream carried deltas that said nothing. Write the arguments onto the last fragment instead, keep the first as the announcing delta with its id and name, and drop the ones in between. --- .../chat/guardrail_translation/handler.py | 38 +++++++-- .../test_openai_guardrail_handler.py | 16 ++-- .../test_unified_guardrail.py | 85 +++++++++++++++++++ 3 files changed, 124 insertions(+), 15 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index a4a3956df94..5b8b31a406f 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -1257,16 +1257,36 @@ class OpenAIChatCompletionsHandler(BaseTranslation): f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls for a stream that carried " f"{len(fragments_by_tool_call)}", ) - for before, (name, arguments), fragments in zip( - pre_guardrail_tool_calls, post_guardrail_tool_calls, fragments_by_tool_call - ): - if (name, arguments) == before: - continue - head, *tail = fragments + rewritten: Final = tuple( + (fragments, shape) + for before, shape, fragments in zip( + pre_guardrail_tool_calls, post_guardrail_tool_calls, fragments_by_tool_call + ) + if shape != before + ) + for fragments, (name, arguments) in rewritten: + head, last = fragments[0], fragments[-1] head.function.name = name - head.function.arguments = arguments - for fragment in tail: - fragment.function.arguments = "" + if head is not last: + head.function.arguments = "" + last.function.arguments = arguments + superseded: Final = chain.from_iterable(fragments[1:-1] for fragments, _ in rewritten) + self._drop_tool_call_fragments(responses_so_far, frozenset(id(fragment) for fragment in superseded)) + + @staticmethod + def _drop_tool_call_fragments( + responses_so_far: list["ModelResponseStream"], # mutable-ok: strips the caller's buffered chunks in place + superseded: frozenset[int], + ) -> None: + """Remove fragments a rewrite superseded, so the stream carries no empty-argument tool deltas.""" + if not superseded: + return + for response in responses_so_far: + for choice in response.choices: + present = choice.delta.tool_calls or () + kept = [call for call in present if id(call) not in superseded] + if len(kept) != len(present): + choice.delta.tool_calls = kept or None async def _apply_guardrail_responses_to_output_streaming( self, diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index bffe6e4c123..b3649c9daaa 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1146,10 +1146,12 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: ) assert result is chunks - fragments = [chunk.choices[0].delta.tool_calls for chunk in chunks[:3]] - assert [fragment[0].function.arguments for fragment in fragments] == ['{"fruit": "PERSIMMON"}', "", ""] - assert fragments[0][0].function.name == "lookup_fruit" - assert fragments[0][0].id == "call_1" + fragments = [call for chunk in chunks for call in chunk.choices[0].delta.tool_calls or ()] + arguments = [fragment.function.arguments for fragment in fragments] + assert "".join(arguments) == '{"fruit": "PERSIMMON"}' + assert arguments.count("") <= 1 + assert fragments[0].function.name == "lookup_fruit" + assert fragments[0].id == "call_1" assert chunks[3].choices[0].delta.tool_calls is None assert chunks[3].choices[0].finish_reason == "tool_calls" @@ -1171,8 +1173,10 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: deliver_ended_stream_rewrites=True, ) - fragments = [chunk.choices[0].delta.tool_calls[0] for chunk in chunks[:3]] - assert [fragment.function.name for fragment in fragments] == ["lookup_fruit_reviewed", None, None] + fragments = [call for chunk in chunks for call in chunk.choices[0].delta.tool_calls or ()] + assert [fragment.function.name for fragment in fragments if fragment.function.name] == [ + "lookup_fruit_reviewed" + ] assert json.loads("".join(fragment.function.arguments for fragment in fragments)) == {"fruit": "persimmon"} assert fragments[0].id == "call_1" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 4dd2f2cb6e3..98897bf2936 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -1171,6 +1171,91 @@ class TestStreamingTransform: assert streamed == "hello world" assert streamed != streamed.upper() + @pytest.mark.asyncio + @pytest.mark.asyncio + async def test_rewritten_tool_arguments_do_not_leave_empty_fragments_on_the_wire(self): + """Blanking the trailing fragments pads the stream with meaningless tool-call deltas.""" + def fragment( + arguments: str, announce: bool = False, finish_reason: str | None = None + ) -> ModelResponseStream: + call: Final[dict[str, object]] = {"index": 0, "function": {"arguments": arguments}} + announced: Final = ( + {**call, "id": "call_0", "type": "function", "function": {"name": "lookup", "arguments": arguments}} + if announce + else call + ) + return ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(content=None, role="assistant", tool_calls=[announced]), + finish_reason=finish_reason, + ) + ], + ) + + chunks: Final = [ + _stream_chunk("looking"), + fragment("", announce=True), + fragment('{"q": "SE'), + fragment('CRET"}', finish_reason="tool_calls"), + ] + out: Final = await _drive_stream(UnifiedLLMGuardrails(), _ToolRedactingGuardrail(), chunks) + + delivered_arguments: Final = tuple( + call.function.arguments + for item in out + for choice in getattr(item, "choices", ()) or () + for call in getattr(choice.delta, "tool_calls", None) or () + ) + assert "".join(delivered_arguments) == '{"q": "MASKED"}' + assert delivered_arguments.count("") <= 1 + + @pytest.mark.asyncio + async def test_a_mixed_text_and_tool_call_chunk_never_finishes_ahead_of_its_text(self): + """A client that stops at finish_reason must not lose rewritten text held behind it.""" + mixed: Final = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content="the code is SECRET", + role="assistant", + tool_calls=[ + { + "index": 0, + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": '{"q": "SECRET"}'}, + } + ], + ), + finish_reason="tool_calls", + ) + ], + ) + leading: Final = _stream_chunk("the code is ") + out: Final = await _drive_stream(UnifiedLLMGuardrails(), _ToolRedactingGuardrail(), [leading, mixed]) + + delivered: Final = "".join(_delta_text(item) for item in out) + assert "MASKED" in delivered + + finished_at: Final = next( + ( + position + for position, item in enumerate(out) + if any(getattr(choice, "finish_reason", None) for choice in getattr(item, "choices", ()) or ()) + ), + None, + ) + text_at: Final = next( + (position for position, item in enumerate(out) if "MASKED" in _delta_text(item)), + None, + ) + assert finished_at is not None + assert text_at is not None + assert finished_at >= text_at + @pytest.mark.asyncio async def test_incremental_diff_emits_uppercased_deltas(self): """incremental_diff: the client receives uppercased deltas whose From 744a31ed4ce4330b5307595ef107f1114b02e531 Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Thu, 29 Oct 2026 12:01:45 +0100 Subject: [PATCH 21/36] fix(guardrails): replay tool calls in the order the model sent them and keep their provider fields The text flush before the tool-call replay was a non-final round, so holdback kept the text back and it reached the client after the tool calls even when the model had sent it first. Flush it fully there, without delivering finish_reason, and only when text actually preceded the first tool call. The passthrough also rebuilt the chunk from id, created and model alone, dropping system_fingerprint and provider_specific_fields; copy the chunk and replace its choices instead. --- .../unified_guardrail/unified_guardrail.py | 22 +++++-- .../test_unified_guardrail.py | 66 +++++++++++++++++++ 2 files changed, 83 insertions(+), 5 deletions(-) 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 04b814c19ee..df7bc13a387 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -541,6 +541,7 @@ class UnifiedLLMGuardrails(CustomLogger): finish_reason_per_choice: dict[int, str | None], held_chars_per_choice: dict[int, int], is_final: bool, + deliver_finish_reason: bool = True, ) -> ModelResponseStream | None: """Build the synthetic chunk carrying the newly-guardrailed deltas. @@ -565,7 +566,7 @@ class UnifiedLLMGuardrails(CustomLogger): # content+tool_calls chunk whose passthrough suppressed it) still # needs to reach the client, even if the guardrail returned no text # to emit. Build a terminator chunk carrying finish_reason per choice. - if is_final and finish_reason_per_choice: + if is_final and deliver_finish_reason and finish_reason_per_choice: terminator_choices: Final[list[StreamingChoices]] = [] for choice_idx, finish_reason in finish_reason_per_choice.items(): if finish_reason is None: @@ -616,7 +617,7 @@ class UnifiedLLMGuardrails(CustomLogger): synthetic_choices: Final[list[StreamingChoices]] = [] for choice_idx in mutated_text_per_choice: delta_text = deltas.get(choice_idx, "") - finish_reason = finish_reason_per_choice.get(choice_idx) if is_final else None + finish_reason = finish_reason_per_choice.get(choice_idx) if is_final and deliver_finish_reason else None # Skip a choice with nothing to say: no new content and no # finish_reason to deliver. This avoids emitting an empty delta for an # already-finished choice (e.g. one that terminated via a passed-through @@ -664,6 +665,7 @@ class UnifiedLLMGuardrails(CustomLogger): held_chars_per_choice: dict[int, int], is_final: bool, terminated: asyncio.Event, + deliver_finish_reason: bool = True, ) -> AsyncGenerator[object, None]: """Run one guardrail processing round and emit the resulting diff chunk. @@ -693,6 +695,7 @@ class UnifiedLLMGuardrails(CustomLogger): finish_reason_per_choice=finish_reason_per_choice, held_chars_per_choice=held_chars_per_choice, is_final=is_final, + deliver_finish_reason=deliver_finish_reason, ) except ModifyResponseException as e: terminated.set() @@ -754,7 +757,9 @@ class UnifiedLLMGuardrails(CustomLogger): last_chunk: object | None = None terminated: Final = asyncio.Event() - def _round(reference_chunk: object, is_final: bool) -> AsyncGenerator[object, None]: + def _round( + reference_chunk: object, is_final: bool, deliver_finish_reason: bool = True + ) -> AsyncGenerator[object, None]: return self._emit_transform_round( endpoint_translation=endpoint_translation, guardrail_to_apply=guardrail_to_apply, @@ -769,14 +774,17 @@ class UnifiedLLMGuardrails(CustomLogger): held_chars_per_choice=held_chars_per_choice, is_final=is_final, terminated=terminated, + deliver_finish_reason=deliver_finish_reason, ) saw_tool_calls = False saw_text_content = False + text_preceded_tool_calls = False # rebind-ok: records text arriving before the first tool call try: async for item in response: if self._chunk_has_tool_calls(item): + text_preceded_tool_calls = text_preceded_tool_calls or (not saw_tool_calls and saw_text_content) saw_tool_calls = True responses_so_far.append(item) last_chunk = item @@ -816,8 +824,8 @@ class UnifiedLLMGuardrails(CustomLogger): ): yield out - if saw_text_content: - async for out in _round(last_chunk, is_final=False): + if text_preceded_tool_calls: + async for out in _round(last_chunk, is_final=True, deliver_finish_reason=False): yield out tool_chunks: Final = tuple( self._tool_call_passthrough_chunk( @@ -1015,6 +1023,10 @@ class UnifiedLLMGuardrails(CustomLogger): finish_reason=passthrough_finish, ) ) + if isinstance(item, ModelResponseStream): + passthrough: Final = copy.deepcopy(item) + passthrough.choices = synthetic_choices + return passthrough return ModelResponseStream( id=getattr(item, "id", None), created=getattr(item, "created", None), diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 98897bf2936..b95f9136255 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -1211,6 +1211,72 @@ class TestStreamingTransform: assert "".join(delivered_arguments) == '{"q": "MASKED"}' assert delivered_arguments.count("") <= 1 + @pytest.mark.asyncio + async def test_a_tool_call_streamed_before_text_is_delivered_before_that_text(self) -> None: + """The replay must keep the order the model chose between tool calls and text.""" + chunks: Final = ( + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(tool_calls=[ + {"index": 0, "id": "call_0", "type": "function", + "function": {"name": "contact", "arguments": '{"contact":"ok"}'}}, + ]))]), + _stream_chunk("hello SECRET", finish_reason="stop"), + ) + out: Final = await _drive_stream(UnifiedLLMGuardrails(), _ToolRedactingGuardrail(), chunks) + kinds: Final = tuple( + "tool" if choice.delta.tool_calls else "text" + for chunk in out + for choice in chunk.choices + if choice.delta.tool_calls or choice.delta.content + ) + assert kinds.index("tool") < kinds.index("text") + + @pytest.mark.asyncio + async def test_text_streamed_before_a_tool_call_is_delivered_before_it(self) -> None: + """The replay must not move a tool call ahead of text the model sent first.""" + chunks: Final = ( + _stream_chunk("hello "), + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(tool_calls=[ + {"index": 0, "id": "call_0", "type": "function", + "function": {"name": "contact", "arguments": '{"contact":"ok"}'}}, + ]))]), + _stream_chunk("SECRET", finish_reason="stop"), + ) + out: Final = await _drive_stream(UnifiedLLMGuardrails(), _ToolRedactingGuardrail(), chunks) + kinds: Final = tuple( + "tool" if choice.delta.tool_calls else "text" + for chunk in out + for choice in chunk.choices + if choice.delta.tool_calls or choice.delta.content + ) + assert kinds[0] == "text" + assert "tool" in kinds + assert "".join(_delta_text(chunk) for chunk in out) == "checked:hello MASKED" + + @pytest.mark.asyncio + async def test_tool_call_replay_keeps_the_provider_fields_of_its_chunk(self) -> None: + """A replayed tool-call chunk must not lose the provider metadata the model sent with it.""" + chunks: Final = ( + _stream_chunk("hello SECRET"), + ModelResponseStream( + id="chatcmpl-provider", + created=1234567890, + model="gpt-4", + system_fingerprint="fp_provider", + provider_specific_fields={"service_tier": "scale"}, + choices=[StreamingChoices(index=0, delta=Delta(tool_calls=[ + {"index": 0, "id": "call_0", "type": "function", + "function": {"name": "contact", "arguments": '{"contact":"ok"}'}}, + ]), finish_reason="tool_calls")], + ), + ) + out: Final = await _drive_stream(UnifiedLLMGuardrails(), _ToolRedactingGuardrail(), chunks) + replayed: Final = tuple( + chunk for chunk in out for choice in chunk.choices if choice.delta.tool_calls + ) + assert replayed + assert all(chunk.system_fingerprint == "fp_provider" for chunk in replayed) + assert all(chunk.provider_specific_fields == {"service_tier": "scale"} for chunk in replayed) + @pytest.mark.asyncio async def test_a_mixed_text_and_tool_call_chunk_never_finishes_ahead_of_its_text(self): """A client that stops at finish_reason must not lose rewritten text held behind it.""" From a9cdd82337015150736ddc1564324fe4ce430538 Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Thu, 29 Oct 2026 12:26:31 +0100 Subject: [PATCH 22/36] fix(guardrails): hold back finish_reason while a choice still owes guardrailed text The tool-call passthrough treated a choice as pending only when a processing round had already withheld characters from it. A stream whose first tool call arrives before any text never runs that round, so finish_reason went out on the passthrough while the rewritten text was still buffered, and a client that stops at finish_reason lost it. Treat a choice that buffered text and has emitted none of it as pending too. --- .../unified_guardrail/unified_guardrail.py | 24 ++++++++++++++++++- .../test_unified_guardrail.py | 22 +++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) 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 df7bc13a387..a718ec0eeba 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -11,6 +11,7 @@ import contextlib import copy import json from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Mapping, Sequence +from itertools import chain from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, cast, runtime_checkable import anyio @@ -142,6 +143,25 @@ def _held_choices(held_chars_per_choice: Mapping[int, int]) -> frozenset[int]: return frozenset(idx for idx, held in held_chars_per_choice.items() if held > 0) +def _choices_owing_text( + buffered: Sequence[object], + emitted_text_per_choice: Mapping[int, str], + held_chars_per_choice: Mapping[int, int], +) -> frozenset[int]: + """Choices whose guardrailed text has not reached the client yet. + + A choice that buffered text but has emitted none of it owes text even when no + processing round has run, so ``held_chars_per_choice`` has no entry for it. + """ + buffered_text: Final = frozenset( + getattr(choice, "index", 0) or 0 + for choice in chain.from_iterable(_chunk_choices(chunk) for chunk in buffered) + if getattr(getattr(choice, "delta", None), "content", None) + ) + unstarted: Final = frozenset(idx for idx in buffered_text if not emitted_text_per_choice.get(idx)) + return unstarted | _held_choices(held_chars_per_choice) + + def _recorded_guardrail_information(request_data: _RequestData) -> tuple[StandardLoggingGuardrailInformation, ...]: _metadata_key, metadata_bucket = get_or_create_metadata_bucket(request_data) entries: Final = metadata_bucket.get("standard_logging_guardrail_information") @@ -831,7 +851,9 @@ class UnifiedLLMGuardrails(CustomLogger): self._tool_call_passthrough_chunk( buffered_item, finish_reason_per_choice=finish_reason_per_choice, - held_choices=_held_choices(held_chars_per_choice), + held_choices=_choices_owing_text( + inspected_responses, emitted_text_per_choice, held_chars_per_choice + ), ) for buffered_item in inspected_responses if self._chunk_has_tool_calls(buffered_item) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index b95f9136255..c0ebe662bcf 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -1277,6 +1277,28 @@ class TestStreamingTransform: assert all(chunk.system_fingerprint == "fp_provider" for chunk in replayed) assert all(chunk.provider_specific_fields == {"service_tier": "scale"} for chunk in replayed) + @pytest.mark.asyncio + async def test_finish_on_a_tool_chunk_never_precedes_the_rewritten_text(self) -> None: + """A client that stops at finish_reason must already have the guardrailed text.""" + def tool(index: int, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(tool_calls=[ + {"index": index, "id": f"call_{index}", "type": "function", + "function": {"name": "contact", "arguments": '{"contact":"ok"}'}}, + ]), finish_reason=finish_reason)]) + + chunks: Final = (tool(0), _stream_chunk("hello SECRET"), tool(1, finish_reason="tool_calls")) + out: Final = await _drive_stream(UnifiedLLMGuardrails(), _ToolRedactingGuardrail(), chunks) + delivered: Final = tuple( + (bool(choice.delta.content), choice.finish_reason is not None) + for chunk in out + for choice in chunk.choices + ) + carries_text: Final = tuple(position for position, (text, _) in enumerate(delivered) if text) + finishes: Final = tuple(position for position, (_, finish) in enumerate(delivered) if finish) + assert carries_text and finishes + assert max(carries_text) <= min(finishes) + assert "".join(_delta_text(chunk) for chunk in out) == "checked:hello MASKED" + @pytest.mark.asyncio async def test_a_mixed_text_and_tool_call_chunk_never_finishes_ahead_of_its_text(self): """A client that stops at finish_reason must not lose rewritten text held behind it.""" From 4cf4510928616722bc852d21d968a4f825a2d1e8 Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Thu, 29 Oct 2026 13:31:24 +0100 Subject: [PATCH 23/36] fix(otel): label guardrail spans with the choice they inspected Per-choice guardrail log entries carry choice_index, but the span mapper never read it, so a trace showed several guardrail spans for one request with nothing telling them apart. Expose it as litellm.guardrail.choice_index. --- litellm/integrations/otel/mappers/genai.py | 1 + litellm/integrations/otel/model/payloads.py | 2 ++ litellm/integrations/otel/model/semconv.py | 1 + .../unit/integrations/otel/test_otel_v2_emitter.py | 13 +++++++++++++ 4 files changed, 17 insertions(+) diff --git a/litellm/integrations/otel/mappers/genai.py b/litellm/integrations/otel/mappers/genai.py index 4e07d090edd..132bbf1f401 100644 --- a/litellm/integrations/otel/mappers/genai.py +++ b/litellm/integrations/otel/mappers/genai.py @@ -141,6 +141,7 @@ class GenAIMapper: LiteLLM.GUARDRAIL_CONFIDENCE_SCORE: lambda d: d.confidence_score, LiteLLM.GUARDRAIL_RISK_SCORE: lambda d: d.risk_score, LiteLLM.GUARDRAIL_MASKED_ENTITY_COUNT: lambda d: d.masked_entity_count, + LiteLLM.GUARDRAIL_CHOICE_INDEX: lambda d: d.choice_index, LiteLLM.GUARDRAIL_DURATION: lambda d: d.duration, LiteLLM.GUARDRAIL_ID: lambda d: d.guardrail_id, LiteLLM.GUARDRAIL_POLICY_TEMPLATE: lambda d: d.policy_template, diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index 9c62c45e489..559fa86ef93 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -248,6 +248,7 @@ class GuardrailSpanData: mode: str | None = None status: str | None = None masked_entity_count: int | None = None + choice_index: int | None = None provider: str | None = None action: str | None = None # The guardrail verdict / provider response (e.g. the moderation result), @@ -308,6 +309,7 @@ class GuardrailSpanData: mode=_guardrail_mode_str(get("guardrail_mode")), status=status, masked_entity_count=_total_masked_entities(get("masked_entity_count")), + choice_index=as_int(get("choice_index")), provider=as_str(get("guardrail_provider")), action=as_str(get("guardrail_action")), response_json=_json_or_none(response) if response is not None else None, diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index b7f955a75ed..41a5d85a0bd 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -316,6 +316,7 @@ class LiteLLM: GUARDRAIL_CONFIDENCE_SCORE: Final = "litellm.guardrail.confidence_score" GUARDRAIL_RISK_SCORE: Final = "litellm.guardrail.risk_score" GUARDRAIL_MASKED_ENTITY_COUNT: Final = "litellm.guardrail.masked_entity_count" + GUARDRAIL_CHOICE_INDEX: Final = "litellm.guardrail.choice_index" GUARDRAIL_DURATION: Final = "litellm.guardrail.duration" GUARDRAIL_ID: Final = "litellm.guardrail.id" GUARDRAIL_POLICY_TEMPLATE: Final = "litellm.guardrail.policy_template" diff --git a/tests/unit/integrations/otel/test_otel_v2_emitter.py b/tests/unit/integrations/otel/test_otel_v2_emitter.py index 6e29b2e9c68..403e8f29974 100644 --- a/tests/unit/integrations/otel/test_otel_v2_emitter.py +++ b/tests/unit/integrations/otel/test_otel_v2_emitter.py @@ -344,6 +344,19 @@ def test_guardrail_block_span_is_error_and_carries_verdict(): assert a[LiteLLM.GUARDRAIL_MASKED_ENTITY_COUNT] == 2 +def test_guardrail_span_carries_the_choice_index(): + """A per-choice guardrail entry must say on the span which choice it inspected.""" + engine, exporter = _engine() + engine.emit( + SpanRole.GUARDRAIL, + GuardrailSpanData.from_logging_entry( + {"guardrail_name": "g", "guardrail_status": "success", "choice_index": 1} + ), + ) + (span,) = exporter.get_finished_spans() + assert span.attributes[LiteLLM.GUARDRAIL_CHOICE_INDEX] == 1 + + def test_guardrail_success_span_is_unset(): """On success the status is left UNSET (semconv default) — not forced OK.""" engine, exporter = _engine() From 4be43403284c62c0350e0e9d01c1ab40eb57833c Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Thu, 29 Oct 2026 13:59:10 +0100 Subject: [PATCH 24/36] fix(guardrails): keep calls still in flight out of the inspected response Rebuilding an ended stream gives a tool call that never streamed any argument text an empty JSON object, and that response went into request_data, so a guardrail reading it saw a real call taking no arguments. Hand it a copy without those calls, and keep the rebuilt response intact for the rewrites. Parsing alone could not tell the two apart, since an empty object is valid JSON. --- .../chat/guardrail_translation/handler.py | 52 +++++++++++++++---- .../test_openai_guardrail_handler.py | 47 +++++++++++++++++ 2 files changed, 90 insertions(+), 9 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 18c8b5aca9e..9f9485692c2 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -14,6 +14,7 @@ Pattern Overview: This pattern can be replicated for other message formats (e.g., Anthropic). """ +import copy import json import time import uuid @@ -757,6 +758,19 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return False return True + @staticmethod + def _joined_streamed_arguments( + responses_so_far: Sequence["ModelResponseStream"], + ) -> Mapping[tuple[int, int], str]: + """The argument text streamed so far per (choice index, tool call index).""" + fragments: Final = tuple(OpenAIChatCompletionsHandler._streamed_argument_fragments(responses_so_far)) + return MappingProxyType( + { + call: "".join(text for other, text in fragments if other == call) + for call in frozenset(call for call, _ in fragments) + } + ) + @staticmethod def _streamed_tool_arguments_are_complete(responses_so_far: Sequence["ModelResponseStream"]) -> bool: """Whether every streamed tool call now carries arguments that parse as JSON. @@ -764,17 +778,33 @@ class OpenAIChatCompletionsHandler(BaseTranslation): A call that has announced itself but emitted no argument text yet is still in flight, so a parallel call finishing first does not make the whole set complete. """ - fragments: Final = tuple(OpenAIChatCompletionsHandler._streamed_argument_fragments(responses_so_far)) - if not fragments: + joined: Final = OpenAIChatCompletionsHandler._joined_streamed_arguments(responses_so_far) + if not joined: return False - joined: Final = MappingProxyType( - { - call: "".join(text for other, text in fragments if other == call) - for call in frozenset(call for call, _ in fragments) - } - ) return all(OpenAIChatCompletionsHandler._parses_as_json(arguments) for arguments in joined.values()) + @staticmethod + def _without_unfinished_tool_calls( + assembled: "ModelResponse", responses_so_far: Sequence["ModelResponseStream"] + ) -> "ModelResponse": + """Copy of the rebuilt response without the calls still in flight when the stream ended. + + Rebuilding gives a call that never streamed any argument text an empty JSON object, which + reads to a guardrail as a real call taking no arguments rather than an unfinished one. + """ + joined: Final = OpenAIChatCompletionsHandler._joined_streamed_arguments(responses_so_far) + pruned: Final = copy.deepcopy(assembled) + for choice in pruned.choices: + streamed = choice.message.tool_calls or () + kept = [ + call + for position, call in enumerate(streamed) + if OpenAIChatCompletionsHandler._parses_as_json(joined.get((choice.index, position), "")) + ] + if len(kept) != len(streamed): + choice.message.tool_calls = kept or None + return pruned + @staticmethod def _rebuild_ended_stream_per_choice( responses_so_far: Sequence["ModelResponseStream"], @@ -903,7 +933,11 @@ class OpenAIChatCompletionsHandler(BaseTranslation): else self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj) ) if assembled is not None: - resolved_request_data["response"] = assembled + resolved_request_data["response"] = ( + assembled + if tool_arguments_complete + else self._without_unfinished_tool_calls(assembled, responses_so_far) + ) if assembled is not None and len(assembled.choices) > 1: choice_rounds: Final = tuple( ( diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 0a0bccceb71..3de123a4241 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1675,6 +1675,53 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert sent is not None assert [c["function"]["arguments"] for c in sent] == ['{"city": "Paris"}', '{"city": "Rome"}'] + @pytest.mark.asyncio + async def test_an_unfinished_parallel_call_is_kept_out_of_what_the_guardrail_reads(self) -> None: + """A call that only announced its name must not reach the guardrail, by either route.""" + from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + class RecordingGuardrail(CustomGuardrail): + seen_calls: list[object] = [] + seen_response: list[object] = [] + + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object], + input_type: Literal["request", "response"], logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + type(self).seen_calls.append(inputs.get("tool_calls")) + type(self).seen_response.append(request_data.get("response")) + return inputs + + def call(index: int, arguments: str) -> dict[str, object]: + return { + "index": index, "id": f"call_{index}", "type": "function", + "function": {"name": "weather", "arguments": arguments}, + } + + RecordingGuardrail.seen_calls = [] + RecordingGuardrail.seen_response = [] + handler: Final = OpenAIChatCompletionsHandler() + request_data: Final[dict[str, object]] = {} + ended_with_one_unfinished: Final = [ModelResponseStream(choices=[StreamingChoices( + index=0, delta=Delta(content="looking", tool_calls=[call(0, '{"city": "Paris"}'), call(1, "")]), + finish_reason="tool_calls", + )])] + await handler.process_output_streaming_response( + responses_so_far=ended_with_one_unfinished, + guardrail_to_apply=RecordingGuardrail(guardrail_name="recording"), + litellm_logging_obj=None, request_data=request_data, + stream_transform_sink=StreamTransformSink(), + ) + assert RecordingGuardrail.seen_calls == [None] + inspected: Final = RecordingGuardrail.seen_response[-1] + arguments: Final = tuple( + tool_call.function.arguments + for choice in getattr(inspected, "choices", None) or () + for tool_call in choice.message.tool_calls or () + ) + assert arguments == ('{"city": "Paris"}',) + @pytest.mark.asyncio async def test_deliver_ended_stream_tool_rewrites_keep_choice_indices(self) -> None: handler: Final = OpenAIChatCompletionsHandler() From c4f7c3f66da347201c1bcbd9c1bc982dfb96db77 Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Sat, 31 Oct 2026 23:50:57 +0100 Subject: [PATCH 25/36] perf(guardrails): group streamed tool arguments in one pass Joining the fragments rescanned all of them once per distinct tool call, so a sampled round cost the number of calls times the length of the stream and at n=2 grew with the square of the streamed arguments. Group them in a single pass, and narrow the pruned response rather than the raw one in the per-choice rounds. --- .../chat/guardrail_translation/handler.py | 40 +++++++++++-------- 1 file changed, 24 insertions(+), 16 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 9f9485692c2..3f0500dbf0a 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -19,7 +19,7 @@ import json import time import uuid from collections.abc import Iterator, Mapping, Sequence -from itertools import chain +from itertools import chain, groupby from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast @@ -762,13 +762,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation): def _joined_streamed_arguments( responses_so_far: Sequence["ModelResponseStream"], ) -> Mapping[tuple[int, int], str]: - """The argument text streamed so far per (choice index, tool call index).""" - fragments: Final = tuple(OpenAIChatCompletionsHandler._streamed_argument_fragments(responses_so_far)) + """The argument text streamed so far per (choice index, tool call index). + + Grouped in one pass. Rescanning every fragment once per distinct call made a sampled + round cost the number of calls times the length of the stream, which at n=2 grows with + the square of the streamed arguments. + """ + fragments: Final = sorted( + OpenAIChatCompletionsHandler._streamed_argument_fragments(responses_so_far), + key=lambda fragment: fragment[0], + ) return MappingProxyType( - { - call: "".join(text for other, text in fragments if other == call) - for call in frozenset(call for call, _ in fragments) - } + {call: "".join(text for _, text in group) for call, group in groupby(fragments, key=lambda f: f[0])} ) @staticmethod @@ -932,20 +937,23 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if len(streamed_choice_indices) == 1 and not stream_ended and not tool_arguments_complete else self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj) ) - if assembled is not None: - resolved_request_data["response"] = ( - assembled - if tool_arguments_complete - else self._without_unfinished_tool_calls(assembled, responses_so_far) - ) - if assembled is not None and len(assembled.choices) > 1: + inspected: Final = ( + None + if assembled is None + else assembled + if tool_arguments_complete + else self._without_unfinished_tool_calls(assembled, responses_so_far) + ) + if inspected is not None: + resolved_request_data["response"] = inspected + if inspected is not None and len(inspected.choices) > 1: choice_rounds: Final = tuple( ( choice, StreamTransformSink(), [self._narrowed_to_choice(chunk, choice.index) for chunk in responses_so_far], ) - for choice in assembled.choices + for choice in inspected.choices ) from litellm.proxy.policy_engine.pipeline_executor import ( UndeliverableStreamRewrite, @@ -954,7 +962,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): transform_undeliverable: UndeliverableStreamRewrite | None = None # rebind-ok: keeps the first failure try: for choice, choice_sink, choice_chunks in choice_rounds: - resolved_request_data["response"] = assembled.model_copy( + resolved_request_data["response"] = inspected.model_copy( update=MappingProxyType({"choices": [choice]}) ) resolved_request_data["responses"] = choice_chunks From 215e0143110fa458cd37ed7b5065d32b23e109bf Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Sat, 31 Oct 2026 23:59:41 +0100 Subject: [PATCH 26/36] refactor(guardrails): match unfinished tool calls by their announced id Pruning matched calls by their position in the rebuilt list, which need not follow the order the stream gave them, and did it by deep-copying the response and writing into the model. Match on the id each call announced, build the pruned copies with model_copy, and return the response untouched when every call finished. --- .../chat/guardrail_translation/handler.py | 44 +++++++++++++------ 1 file changed, 30 insertions(+), 14 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 3f0500dbf0a..1ac7f6ebb0a 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -14,7 +14,6 @@ Pattern Overview: This pattern can be replicated for other message formats (e.g., Anthropic). """ -import copy import json import time import uuid @@ -788,27 +787,44 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return False return all(OpenAIChatCompletionsHandler._parses_as_json(arguments) for arguments in joined.values()) + @staticmethod + def _streamed_call_ids(responses_so_far: Sequence["ModelResponseStream"]) -> Mapping[tuple[int, int], str]: + """The id each (choice index, tool call index) announced itself with.""" + choices: Final = chain.from_iterable(chunk.choices for chunk in responses_so_far) + calls: Final = chain.from_iterable( + ((choice.index, call) for call in choice.delta.tool_calls or ()) for choice in choices + ) + return MappingProxyType({(index, call.index or 0): call.id for index, call in calls if call.id}) + + @staticmethod + def _choice_without(choice: "Choices", unfinished: frozenset[str]) -> "Choices": + kept: Final = [call for call in choice.message.tool_calls or () if call.id not in unfinished] + message: Final = choice.message.model_copy(update=MappingProxyType({"tool_calls": kept or None})) + return choice.model_copy(update=MappingProxyType({"message": message})) + @staticmethod def _without_unfinished_tool_calls( assembled: "ModelResponse", responses_so_far: Sequence["ModelResponseStream"] ) -> "ModelResponse": """Copy of the rebuilt response without the calls still in flight when the stream ended. - Rebuilding gives a call that never streamed any argument text an empty JSON object, which - reads to a guardrail as a real call taking no arguments rather than an unfinished one. + Rebuilding gives a call that never streamed complete arguments an empty JSON object, which + reads to a guardrail as a real call taking no arguments rather than an unfinished one. The + calls are matched by the id they announced, since the rebuilt order need not follow the + order the stream gave them. """ joined: Final = OpenAIChatCompletionsHandler._joined_streamed_arguments(responses_so_far) - pruned: Final = copy.deepcopy(assembled) - for choice in pruned.choices: - streamed = choice.message.tool_calls or () - kept = [ - call - for position, call in enumerate(streamed) - if OpenAIChatCompletionsHandler._parses_as_json(joined.get((choice.index, position), "")) - ] - if len(kept) != len(streamed): - choice.message.tool_calls = kept or None - return pruned + unfinished: Final = frozenset( + call_id + for key, call_id in OpenAIChatCompletionsHandler._streamed_call_ids(responses_so_far).items() + if not OpenAIChatCompletionsHandler._parses_as_json(joined.get(key, "")) + ) + if not unfinished: + return assembled + pruned: Final = [ + OpenAIChatCompletionsHandler._choice_without(choice, unfinished) for choice in assembled.choices + ] + return assembled.model_copy(update=MappingProxyType({"choices": pruned})) @staticmethod def _rebuild_ended_stream_per_choice( From 175ddd4faf21ce61e03deedfb896402bd2aa26d4 Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Sun, 1 Nov 2026 22:08:39 +0100 Subject: [PATCH 27/36] fix(guardrails): leave request_data as the round found it A diff-mode round wrote response, responses and the choice index into the caller's request_data and put back its own values instead of what was there, so a stale response outlived the round and later hooks read it. Record what each path found and restore it, including when the round raises. --- .../chat/guardrail_translation/handler.py | 250 ++++++++++-------- .../test_openai_guardrail_handler.py | 37 ++- 2 files changed, 167 insertions(+), 120 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 1ac7f6ebb0a..75d0500324f 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -609,8 +609,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # create a standalone dict (SDK / direct-call path). if request_data is None: request_data = {"responses": responses_so_far} + owes_responses = False # rebind-ok: the standalone dict is ours to keep else: - if "responses" not in request_data: + owes_responses = "responses" not in request_data # rebind-ok: tracks the key this scan adds + if owes_responses: request_data["responses"] = responses_so_far self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) @@ -621,12 +623,16 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Include model information from the first response if available if responses_so_far and hasattr(responses_so_far[0], "model") and responses_so_far[0].model: inputs["model"] = responses_so_far[0].model - guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( - inputs=inputs, - request_data=request_data, - input_type="response", - logging_obj=litellm_logging_obj, - ) + try: + guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + finally: + if owes_responses: + request_data.pop("responses", None) guardrailed_texts: Final = guardrailed_inputs.get("texts", ()) @@ -939,125 +945,145 @@ class OpenAIChatCompletionsHandler(BaseTranslation): resolved_request_data: Final[dict[str, object]] = ( request_data if request_data is not None else {"responses": responses_so_far} ) + missing: Final = object() + prior_response: Final = resolved_request_data.get("response", missing) + prior_responses: Final = resolved_request_data.get("responses", missing) + prior_choice_index: Final = resolved_request_data.get("guardrail_choice_index", missing) if "responses" not in resolved_request_data: resolved_request_data["responses"] = responses_so_far self.merge_user_api_key_metadata_into_request(resolved_request_data, user_api_key_dict) - - inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) - if self._streamed_tool_call_fingerprints(responses_so_far): - stream_ended: Final = self._first_choice_has_finished(responses_so_far) - streamed_choice_indices: Final = frozenset(self._streamed_choice_indices(responses_so_far)) - tool_arguments_complete: Final = self._streamed_tool_arguments_are_complete(responses_so_far) - assembled: Final = ( - None - if len(streamed_choice_indices) == 1 and not stream_ended and not tool_arguments_complete - else self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj) - ) - inspected: Final = ( - None - if assembled is None - else assembled - if tool_arguments_complete - else self._without_unfinished_tool_calls(assembled, responses_so_far) - ) - if inspected is not None: - resolved_request_data["response"] = inspected - if inspected is not None and len(inspected.choices) > 1: - choice_rounds: Final = tuple( - ( - choice, - StreamTransformSink(), - [self._narrowed_to_choice(chunk, choice.index) for chunk in responses_so_far], - ) - for choice in inspected.choices + try: + inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) + if self._streamed_tool_call_fingerprints(responses_so_far): + stream_ended: Final = self._first_choice_has_finished(responses_so_far) + streamed_choice_indices: Final = frozenset(self._streamed_choice_indices(responses_so_far)) + tool_arguments_complete: Final = self._streamed_tool_arguments_are_complete(responses_so_far) + assembled: Final = ( + None + if len(streamed_choice_indices) == 1 and not stream_ended and not tool_arguments_complete + else self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj) ) - from litellm.proxy.policy_engine.pipeline_executor import ( - UndeliverableStreamRewrite, + inspected: Final = ( + None + if assembled is None + else assembled + if tool_arguments_complete + else self._without_unfinished_tool_calls(assembled, responses_so_far) ) - - transform_undeliverable: UndeliverableStreamRewrite | None = None # rebind-ok: keeps the first failure - try: - for choice, choice_sink, choice_chunks in choice_rounds: - resolved_request_data["response"] = inspected.model_copy( - update=MappingProxyType({"choices": [choice]}) + if inspected is not None: + resolved_request_data["response"] = inspected + if inspected is not None and len(inspected.choices) > 1: + choice_rounds: Final = tuple( + ( + choice, + StreamTransformSink(), + [self._narrowed_to_choice(chunk, choice.index) for chunk in responses_so_far], ) - resolved_request_data["responses"] = choice_chunks - resolved_request_data["guardrail_choice_index"] = choice.index - try: - await self._process_streaming_transform( - responses_so_far=choice_chunks, - guardrail_to_apply=guardrail_to_apply, - litellm_logging_obj=litellm_logging_obj, - user_api_key_dict=user_api_key_dict, - request_data=resolved_request_data, - sink=choice_sink, + for choice in inspected.choices + ) + from litellm.proxy.policy_engine.pipeline_executor import ( + UndeliverableStreamRewrite, + ) + + transform_undeliverable: UndeliverableStreamRewrite | None = None # rebind-ok: first failure + try: + for choice, choice_sink, choice_chunks in choice_rounds: + resolved_request_data["response"] = inspected.model_copy( + update=MappingProxyType({"choices": [choice]}) + ) + resolved_request_data["responses"] = choice_chunks + resolved_request_data["guardrail_choice_index"] = choice.index + try: + await self._process_streaming_transform( + responses_so_far=choice_chunks, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=resolved_request_data, + sink=choice_sink, + ) + except UndeliverableStreamRewrite as exc: + if transform_undeliverable is None: + transform_undeliverable = exc + finally: + resolved_request_data["response"] = assembled + resolved_request_data["responses"] = responses_so_far + resolved_request_data.pop("guardrail_choice_index", None) + if transform_undeliverable is not None: + raise transform_undeliverable + merged_texts: Final = MappingProxyType( + dict( + chain.from_iterable( + choice_sink.mutated_text_per_choice.items() for _, choice_sink, _ in choice_rounds ) - except UndeliverableStreamRewrite as exc: - if transform_undeliverable is None: - transform_undeliverable = exc - finally: - resolved_request_data["response"] = assembled - resolved_request_data["responses"] = responses_so_far - resolved_request_data.pop("guardrail_choice_index", None) - if transform_undeliverable is not None: - raise transform_undeliverable - merged_texts: Final = MappingProxyType( - dict( - chain.from_iterable( - choice_sink.mutated_text_per_choice.items() for _, choice_sink, _ in choice_rounds ) ) - ) - merged_holdback: Final = MappingProxyType( - dict( - chain.from_iterable( - choice_sink.holdback_per_choice.items() for _, choice_sink, _ in choice_rounds + merged_holdback: Final = MappingProxyType( + dict( + chain.from_iterable( + choice_sink.holdback_per_choice.items() for _, choice_sink, _ in choice_rounds + ) ) ) - ) - sink.record(merged_texts, merged_holdback) - return - if assembled is not None and tool_arguments_complete: - tool_calls: Final = chain.from_iterable(choice.message.tool_calls or () for choice in assembled.choices) - inputs["tool_calls"] = TypeAdapter(list[ChatCompletionToolCallChunk]).validate_python( - tuple( - MappingProxyType({"index": index, **converted}) - for index, tool_call in enumerate(tool_calls) - if (converted := self._convert_tool_call_to_dict(tool_call)) is not None + sink.record(merged_texts, merged_holdback) + return + if assembled is not None and tool_arguments_complete: + tool_calls: Final = chain.from_iterable( + choice.message.tool_calls or () for choice in assembled.choices ) - ) - if responses_so_far and getattr(responses_so_far[0], "model", None): - inputs["model"] = responses_so_far[0].model - guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( - inputs=inputs, - request_data=resolved_request_data, - input_type="response", - logging_obj=litellm_logging_obj, - ) - - returned_texts = guardrailed_inputs.get("texts") - # No "texts" key means the guardrail made no change (action NONE): the raw - # accumulated text is the guardrailed text. A present-but-shorter list is a - # guardrail contract violation; those choices are omitted below (withheld, - # not emitted raw) so a malformed response fails closed instead of leaking. - if returned_texts is None: - returned_texts = texts_to_check - elif len(returned_texts) < len(texts_to_check): - verbose_proxy_logger.warning( - "OpenAI Chat Completions: guardrail returned %s transformed texts for %s inputs on the " - "streaming transform path; withholding the unmatched choices to fail closed.", - len(returned_texts), - len(texts_to_check), + inputs["tool_calls"] = TypeAdapter(list[ChatCompletionToolCallChunk]).validate_python( + tuple( + MappingProxyType({"index": index, **converted}) + for index, tool_call in enumerate(tool_calls) + if (converted := self._convert_tool_call_to_dict(tool_call)) is not None + ) + ) + if responses_so_far and getattr(responses_so_far[0], "model", None): + inputs["model"] = responses_so_far[0].model + guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=resolved_request_data, + input_type="response", + logging_obj=litellm_logging_obj, ) - holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or () - transformed_texts: Final = MappingProxyType( - {idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts)} - ) - transformed_holdback: Final = MappingProxyType( - {indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback)} - ) - sink.record(transformed_texts, transformed_holdback) + returned_texts = guardrailed_inputs.get("texts") + # No "texts" key means the guardrail made no change (action NONE): the raw + # accumulated text is the guardrailed text. A present-but-shorter list is a + # guardrail contract violation; those choices are omitted below (withheld, + # not emitted raw) so a malformed response fails closed instead of leaking. + if returned_texts is None: + returned_texts = texts_to_check + elif len(returned_texts) < len(texts_to_check): + verbose_proxy_logger.warning( + "OpenAI Chat Completions: guardrail returned %s transformed texts for %s inputs on the " + "streaming transform path; withholding the unmatched choices to fail closed.", + len(returned_texts), + len(texts_to_check), + ) + + holdback: Final = guardrailed_inputs.get("stream_holdback_chars") or () + transformed_texts: Final = MappingProxyType( + {idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts)} + ) + transformed_holdback: Final = MappingProxyType( + { + indices[i]: coerce_stream_holdback_value(holdback[i]) + for i in range(len(indices)) + if i < len(holdback) + } + ) + sink.record(transformed_texts, transformed_holdback) + finally: + for key, prior in ( + ("response", prior_response), + ("responses", prior_responses), + ("guardrail_choice_index", prior_choice_index), + ): + if prior is missing: + resolved_request_data.pop(key, None) + else: + resolved_request_data[key] = prior 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)) diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 3de123a4241..b868e70a902 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1433,14 +1433,8 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: ) assert request_data["visited_choices"] == (0, 1) assert request_data["metadata"]["trace"] == "retained" - if transform: - assert request_data["responses"] is chunks - restored: Final = request_data["response"] - assert isinstance(restored, ModelResponse) - assert tuple(choice.index for choice in restored.choices) == (0, 1) - else: - assert "response" not in request_data - assert "responses" not in request_data + assert "response" not in request_data + assert "responses" not in request_data @pytest.mark.asyncio async def test_undeliverable_rewrite_still_inspects_the_remaining_choices(self) -> None: @@ -1675,6 +1669,33 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert sent is not None assert [c["function"]["arguments"] for c in sent] == ['{"city": "Paris"}', '{"city": "Rome"}'] + @pytest.mark.asyncio + async def test_a_transform_round_leaves_the_callers_request_data_as_it_found_it(self) -> None: + """The round owns these keys only while it runs; what it leaves behind is read by later hooks.""" + from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + handler: Final = OpenAIChatCompletionsHandler() + request_data: Final[dict[str, object]] = {"model": "gpt-4"} + chunks: Final = [ModelResponseStream(choices=[StreamingChoices( + index=0, + delta=Delta(content="looking", tool_calls=[ + {"index": 0, "id": "call_0", "type": "function", + "function": {"name": "weather", "arguments": '{"city": "Paris"}'}}, + ]), + finish_reason="tool_calls", + )])] + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MockGuardrail(guardrail_name="test"), + litellm_logging_obj=None, + request_data=request_data, + stream_transform_sink=StreamTransformSink(), + ) + assert "response" not in request_data + assert "responses" not in request_data + assert "guardrail_choice_index" not in request_data + @pytest.mark.asyncio async def test_an_unfinished_parallel_call_is_kept_out_of_what_the_guardrail_reads(self) -> None: """A call that only announced its name must not reach the guardrail, by either route.""" From e53c6fc8bf35c4cd7fea0384efcba49aa8ca67bd Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Sun, 1 Nov 2026 22:17:24 +0100 Subject: [PATCH 28/36] fix(guardrails): let only the round set the choice index label The label was read straight from request_data, so a caller could put guardrail_choice_index in the request body and have it logged as the choice a guardrail inspected. Drop any incoming value when a round starts, and restore the caller's own on the way out. --- .../chat/guardrail_translation/handler.py | 2 ++ .../test_openai_guardrail_handler.py | 36 +++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 75d0500324f..01cf518e04b 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -682,6 +682,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): prior_response: Final = inspection_request_data.get("response", missing) prior_responses: Final = inspection_request_data.get("responses", missing) prior_choice_index: Final = inspection_request_data.get("guardrail_choice_index", missing) + inspection_request_data.pop("guardrail_choice_index", None) undeliverable: UndeliverableStreamRewrite | None = None # rebind-ok: keeps the first failure try: for inspection_response in inspection_responses: @@ -949,6 +950,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): prior_response: Final = resolved_request_data.get("response", missing) prior_responses: Final = resolved_request_data.get("responses", missing) prior_choice_index: Final = resolved_request_data.get("guardrail_choice_index", missing) + resolved_request_data.pop("guardrail_choice_index", None) if "responses" not in resolved_request_data: resolved_request_data["responses"] = responses_so_far self.merge_user_api_key_metadata_into_request(resolved_request_data, user_api_key_dict) diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index b868e70a902..73abb328848 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1669,6 +1669,42 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert sent is not None assert [c["function"]["arguments"] for c in sent] == ['{"city": "Paris"}', '{"city": "Rome"}'] + @pytest.mark.asyncio + async def test_a_caller_cannot_set_the_choice_index_label(self) -> None: + """The label says which choice the round inspected, so only the round may set it.""" + from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + class LabelReadingGuardrail(CustomGuardrail): + seen: list[object] = [] + + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object], + input_type: Literal["request", "response"], logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + type(self).seen.append(request_data.get("guardrail_choice_index")) + return inputs + + LabelReadingGuardrail.seen = [] + request_data: Final[dict[str, object]] = {"model": "gpt-4", "guardrail_choice_index": 7} + chunks: Final = [ModelResponseStream(choices=[StreamingChoices( + index=0, + delta=Delta(content="looking", tool_calls=[ + {"index": 0, "id": "call_0", "type": "function", + "function": {"name": "weather", "arguments": '{"city": "Paris"}'}}, + ]), + finish_reason="tool_calls", + )])] + await OpenAIChatCompletionsHandler().process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=LabelReadingGuardrail(guardrail_name="label"), + litellm_logging_obj=None, + request_data=request_data, + stream_transform_sink=StreamTransformSink(), + ) + assert LabelReadingGuardrail.seen == [None] + assert request_data["guardrail_choice_index"] == 7 + @pytest.mark.asyncio async def test_a_transform_round_leaves_the_callers_request_data_as_it_found_it(self) -> None: """The round owns these keys only while it runs; what it leaves behind is read by later hooks.""" From e54d4b8d12053d8dbfa72a9024e20dbdb2c8d644 Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Sun, 1 Nov 2026 22:24:57 +0100 Subject: [PATCH 29/36] fix(guardrails): label the diff-mode round with the choice it inspected The ended-stream path labelled every inspection, including a single choice, while a diff-mode round only labelled the per-choice rounds of an n>1 stream. A single-choice round logged no choice at all. Label it with the choice its texts came from. --- litellm/llms/openai/chat/guardrail_translation/handler.py | 2 ++ .../chat/guardrail_translation/test_openai_guardrail_handler.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 01cf518e04b..bd76a2ae9f2 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -951,6 +951,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): prior_responses: Final = resolved_request_data.get("responses", missing) prior_choice_index: Final = resolved_request_data.get("guardrail_choice_index", missing) resolved_request_data.pop("guardrail_choice_index", None) + if len(indices) == 1: + resolved_request_data["guardrail_choice_index"] = indices[0] if "responses" not in resolved_request_data: resolved_request_data["responses"] = responses_so_far self.merge_user_api_key_metadata_into_request(resolved_request_data, user_api_key_dict) diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 73abb328848..d8d06f34519 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1702,7 +1702,7 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: request_data=request_data, stream_transform_sink=StreamTransformSink(), ) - assert LabelReadingGuardrail.seen == [None] + assert LabelReadingGuardrail.seen == [0] assert request_data["guardrail_choice_index"] == 7 @pytest.mark.asyncio From f8a9c28108057077d09b5c692df131182eb587cd Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Sun, 1 Nov 2026 22:32:33 +0100 Subject: [PATCH 30/36] test(guardrails): cover partial tool arguments at n>1 An argument fragment that is still half an object cannot be judged, and the id-based pruning already keeps it away from the guardrail. Lock that in for a stream with more than one choice, where the per-choice rounds narrow the inspected response. --- .../test_openai_guardrail_handler.py | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index d8d06f34519..57931dcad67 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1669,6 +1669,49 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert sent is not None assert [c["function"]["arguments"] for c in sent] == ['{"city": "Paris"}', '{"city": "Rome"}'] + @pytest.mark.asyncio + async def test_partial_arguments_stay_out_of_what_each_choice_inspects(self) -> None: + """Half a JSON object is not a call a guardrail can judge, at any number of choices.""" + from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + class RecordingGuardrail(CustomGuardrail): + seen: list[object] = [] + + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object], + input_type: Literal["request", "response"], logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + type(self).seen.append(request_data.get("response")) + return inputs + + def choice(index: int, arguments: str, label: str) -> StreamingChoices: + return StreamingChoices( + index=index, + delta=Delta(content="looking", tool_calls=[ + {"index": 0, "id": "call_" + label, "type": "function", + "function": {"name": "weather", "arguments": arguments}}, + ]), + finish_reason="tool_calls", + ) + + RecordingGuardrail.seen = [] + await OpenAIChatCompletionsHandler().process_output_streaming_response( + responses_so_far=[ModelResponseStream(choices=[ + choice(0, '{"city": "Pa', "partial"), choice(1, '{"city": "Rome"}', "rome"), + ])], + guardrail_to_apply=RecordingGuardrail(guardrail_name="recording"), + litellm_logging_obj=None, request_data={}, + stream_transform_sink=StreamTransformSink(), + ) + arguments: Final = tuple( + tool_call.function.arguments + for response in RecordingGuardrail.seen + for single in getattr(response, "choices", None) or () + for tool_call in single.message.tool_calls or () + ) + assert arguments == ('{"city": "Rome"}',) + @pytest.mark.asyncio async def test_a_caller_cannot_set_the_choice_index_label(self) -> None: """The label says which choice the round inspected, so only the round may set it.""" From 0ea3b898c1213aa430b7169d6a70cbb13d52b0b4 Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Sun, 1 Nov 2026 22:43:43 +0100 Subject: [PATCH 31/36] perf(guardrails): buy per-choice context only when there is a call to judge At n>1 every sampled round rebuilt the stream and ran one guardrail call per choice, even while no tool call had complete arguments, so a vendor was billed n times a round for context that carried nothing to judge. Rebuild once the stream ends or the arguments parse; until then the choices share one round, with their texts still kept apart. --- .../chat/guardrail_translation/handler.py | 9 +---- .../test_openai_guardrail_handler.py | 36 +++++++++++++++++++ 2 files changed, 37 insertions(+), 8 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index bd76a2ae9f2..3d45feb0eab 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -738,12 +738,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_name=guardrail_to_apply.guardrail_name or "unknown", ) - @staticmethod - def _streamed_choice_indices(responses_so_far: Sequence["ModelResponseStream"]) -> Iterator[int]: - for chunk in responses_so_far: - for choice in chunk.choices: - yield choice.index - @staticmethod def _streamed_argument_fragments( responses_so_far: Sequence["ModelResponseStream"], @@ -960,11 +954,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if self._streamed_tool_call_fingerprints(responses_so_far): stream_ended: Final = self._first_choice_has_finished(responses_so_far) - streamed_choice_indices: Final = frozenset(self._streamed_choice_indices(responses_so_far)) tool_arguments_complete: Final = self._streamed_tool_arguments_are_complete(responses_so_far) assembled: Final = ( None - if len(streamed_choice_indices) == 1 and not stream_ended and not tool_arguments_complete + if not stream_ended and not tool_arguments_complete else self._rebuild_ended_stream_per_choice(responses_so_far, litellm_logging_obj) ) inspected: Final = ( diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 57931dcad67..f139e2c90fc 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1669,6 +1669,42 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert sent is not None assert [c["function"]["arguments"] for c in sent] == ['{"city": "Paris"}', '{"city": "Rome"}'] + @pytest.mark.asyncio + async def test_a_round_with_nothing_to_judge_costs_one_call_at_n2(self) -> None: + """Per-choice rounds buy tool-call context. While no call is complete there is none to buy, + and charging a vendor once per choice for every sampled round buys nothing.""" + from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + class CountingGuardrail(CustomGuardrail): + calls: list[object] = [] + + async def apply_guardrail( + self, inputs: GenericGuardrailAPIInputs, request_data: dict[str, object], + input_type: Literal["request", "response"], logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + type(self).calls.append(inputs.get("texts")) + return inputs + + def choice(index: int, arguments: str) -> StreamingChoices: + return StreamingChoices( + index=index, + delta=Delta(content="looking", tool_calls=[ + {"index": 0, "id": f"call_{index}", "type": "function", + "function": {"name": "weather", "arguments": arguments}}, + ]), + finish_reason=None, + ) + + CountingGuardrail.calls = [] + await OpenAIChatCompletionsHandler().process_output_streaming_response( + responses_so_far=[ModelResponseStream(choices=[choice(0, '{"city": "Pa'), choice(1, '{"ci')])], + guardrail_to_apply=CountingGuardrail(guardrail_name="counting"), + litellm_logging_obj=None, request_data={}, + stream_transform_sink=StreamTransformSink(), + ) + assert len(CountingGuardrail.calls) == 1 + @pytest.mark.asyncio async def test_partial_arguments_stay_out_of_what_each_choice_inspects(self) -> None: """Half a JSON object is not a call a guardrail can judge, at any number of choices.""" From 1a4f6c5542572deb2cdd268e8f8e54f62eb614ab Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Sun, 1 Nov 2026 22:53:39 +0100 Subject: [PATCH 32/36] fix(guardrails): rebuild a stream that carried no choices instead of failing Rebuilding one choice index at a time unpacked the first rebuilt choice, so a stream that ended carrying only usage left nothing to unpack and a policy step over it raised a ValueError. Fall back to rebuilding the chunks as they are when no choice ever appeared. --- .../openai/chat/guardrail_translation/handler.py | 3 +++ .../test_openai_guardrail_handler.py | 16 ++++++++++++++++ 2 files changed, 19 insertions(+) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 3d45feb0eab..d79d2ba5e3c 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -856,6 +856,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): ) for index in choice_indices ) + if not rebuilt_by_index: + whole: Final = stream_chunk_builder(chunks=list(responses_so_far), logging_obj=litellm_logging_obj) + return cast(ModelResponse, whole) if whole is not None else ModelResponse(choices=[]) (_, base_response), *_ = rebuilt_by_index stitched_choices: Final = [ rebuilt.choices[0].model_copy(update=MappingProxyType({"index": index})) diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index f139e2c90fc..0e4ca539165 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1669,6 +1669,22 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert sent is not None assert [c["function"]["arguments"] for c in sent] == ['{"city": "Paris"}', '{"city": "Rome"}'] + @pytest.mark.asyncio + async def test_a_buffered_stream_that_carried_no_choices_is_not_an_error(self) -> None: + """A stream can end carrying only usage, and a policy step still runs over it.""" + from litellm.types.utils import ModelResponseStream + + chunks: Final = [ModelResponseStream(choices=[], usage={"prompt_tokens": 3, "completion_tokens": 0, + "total_tokens": 3})] + result: Final = await OpenAIChatCompletionsHandler().process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MockGuardrail(guardrail_name="test"), + litellm_logging_obj=None, + request_data={}, + deliver_ended_stream_rewrites=True, + ) + assert result is chunks + @pytest.mark.asyncio async def test_a_round_with_nothing_to_judge_costs_one_call_at_n2(self) -> None: """Per-choice rounds buy tool-call context. While no call is complete there is none to buy, From 022612ebfcabe134f6cd3e9ab8c2b76c5d69803a Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Sun, 1 Nov 2026 23:42:38 +0100 Subject: [PATCH 33/36] refactor(guardrails): borrow the inspection keys through one helper Both the ended-stream and diff-mode paths saved and restored response, responses and the choice label by hand, which duplicated the restore and pushed the ended-stream path past the complexity gate after the merge. Lend them through a single context manager instead. --- .../chat/guardrail_translation/handler.py | 74 +++++++++---------- 1 file changed, 33 insertions(+), 41 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 4ebca0ee62a..db2f8005cf4 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -17,7 +17,8 @@ This pattern can be replicated for other message formats (e.g., Anthropic). import json import time import uuid -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Generator, Iterator, Mapping, Sequence +from contextlib import contextmanager from itertools import chain, groupby from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast @@ -78,6 +79,30 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +_INSPECTION_KEYS: Final[tuple[str, ...]] = ("response", "responses", "guardrail_choice_index") + + +@contextmanager +def _borrowed_inspection_keys(request_data: dict[str, object]) -> Generator[None]: + """Lend a round the keys it inspects through, and give the caller's own back. + + The round owns ``response``, ``responses`` and the choice label only while it runs. What it + leaves behind is read by later hooks, and a label the caller supplied is not one a round may + report, so the incoming label is dropped for the duration. + """ + missing: Final = object() + priors: Final = tuple((key, request_data.get(key, missing)) for key in _INSPECTION_KEYS) + request_data.pop("guardrail_choice_index", None) + try: + yield + finally: + for key, prior in priors: + if prior is missing: + request_data.pop(key, None) + else: + request_data[key] = prior + + class OpenAIChatCompletionsHandler(BaseTranslation): """ Handler for processing OpenAI chat completions messages with guardrails. @@ -678,13 +703,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite inspection_request_data: Final = request_data if request_data is not None else {} - missing: Final = object() - prior_response: Final = inspection_request_data.get("response", missing) - prior_responses: Final = inspection_request_data.get("responses", missing) - prior_choice_index: Final = inspection_request_data.get("guardrail_choice_index", missing) - inspection_request_data.pop("guardrail_choice_index", None) undeliverable: UndeliverableStreamRewrite | None = None # rebind-ok: keeps the first failure - try: + with _borrowed_inspection_keys(inspection_request_data): for inspection_response in inspection_responses: inspection_request_data["response"] = inspection_response inspection_request_data["guardrail_choice_index"] = ( @@ -709,19 +729,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation): except UndeliverableStreamRewrite as exc: if undeliverable is None: undeliverable = exc - finally: - if prior_response is missing: - inspection_request_data.pop("response", None) - else: - inspection_request_data["response"] = prior_response - if prior_responses is missing: - inspection_request_data.pop("responses", None) - else: - inspection_request_data["responses"] = prior_responses - if prior_choice_index is missing: - inspection_request_data.pop("guardrail_choice_index", None) - else: - inspection_request_data["guardrail_choice_index"] = prior_choice_index if undeliverable is not None: raise undeliverable if not deliver_ended_stream_rewrites: @@ -938,17 +945,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation): resolved_request_data: Final[dict[str, object]] = ( request_data if request_data is not None else {"responses": responses_so_far} ) - missing: Final = object() - prior_response: Final = resolved_request_data.get("response", missing) - prior_responses: Final = resolved_request_data.get("responses", missing) - prior_choice_index: Final = resolved_request_data.get("guardrail_choice_index", missing) - resolved_request_data.pop("guardrail_choice_index", None) - if len(indices) == 1: - resolved_request_data["guardrail_choice_index"] = indices[0] - if "responses" not in resolved_request_data: - resolved_request_data["responses"] = responses_so_far - self.merge_user_api_key_metadata_into_request(resolved_request_data, user_api_key_dict) - try: + with _borrowed_inspection_keys(resolved_request_data): + if len(indices) == 1: + resolved_request_data["guardrail_choice_index"] = indices[0] + if "responses" not in resolved_request_data: + resolved_request_data["responses"] = responses_so_far + self.merge_user_api_key_metadata_into_request(resolved_request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if self._streamed_tool_call_fingerprints(responses_so_far): stream_ended: Final = self._first_choice_has_finished(responses_so_far) @@ -1069,16 +1071,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation): } ) sink.record(transformed_texts, transformed_holdback) - finally: - for key, prior in ( - ("response", prior_response), - ("responses", prior_responses), - ("guardrail_choice_index", prior_choice_index), - ): - if prior is missing: - resolved_request_data.pop(key, None) - else: - resolved_request_data[key] = prior 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)) From e3722a1a02219093649d1924ceca683ca9dcc764 Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Tue, 3 Nov 2026 12:46:41 +0100 Subject: [PATCH 34/36] fix(guardrails): prune unfinished tool calls per choice The ids of the calls still in flight were flattened into one set and matched against every rebuilt choice, so an n>1 stream where two choices announced the same id lost the finished call along with the unfinished one, and its complete arguments never reached the guardrail. Match on the choice and the id together. --- .../chat/guardrail_translation/handler.py | 11 +++++---- .../test_openai_guardrail_handler.py | 23 +++++++++++++++++++ 2 files changed, 29 insertions(+), 5 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index d79d2ba5e3c..1d604d49faf 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -798,8 +798,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return MappingProxyType({(index, call.index or 0): call.id for index, call in calls if call.id}) @staticmethod - def _choice_without(choice: "Choices", unfinished: frozenset[str]) -> "Choices": - kept: Final = [call for call in choice.message.tool_calls or () if call.id not in unfinished] + def _choice_without(choice: "Choices", unfinished: frozenset[tuple[int, str]]) -> "Choices": + kept: Final = [call for call in choice.message.tool_calls or () if (choice.index, call.id) not in unfinished] message: Final = choice.message.model_copy(update=MappingProxyType({"tool_calls": kept or None})) return choice.model_copy(update=MappingProxyType({"message": message})) @@ -811,12 +811,13 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Rebuilding gives a call that never streamed complete arguments an empty JSON object, which reads to a guardrail as a real call taking no arguments rather than an unfinished one. The - calls are matched by the id they announced, since the rebuilt order need not follow the - order the stream gave them. + calls are matched by the choice they belong to and the id they announced, since the rebuilt + order need not follow the order the stream gave them, and providers number tool calls per + choice, so the same id on another choice is another call. """ joined: Final = OpenAIChatCompletionsHandler._joined_streamed_arguments(responses_so_far) unfinished: Final = frozenset( - call_id + (key[0], call_id) for key, call_id in OpenAIChatCompletionsHandler._streamed_call_ids(responses_so_far).items() if not OpenAIChatCompletionsHandler._parses_as_json(joined.get(key, "")) ) diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 0e4ca539165..63267791ad9 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1721,6 +1721,29 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: ) assert len(CountingGuardrail.calls) == 1 + def test_an_id_in_flight_on_one_choice_does_not_prune_another(self) -> None: + """Providers number tool calls per choice, so the same id on another choice is another call, + and one still in flight must not take a finished call down with it.""" + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + def choice(index: int, arguments: str) -> StreamingChoices: + return StreamingChoices( + index=index, + delta=Delta(content="looking", tool_calls=[ + {"index": 0, "id": "call_shared", "type": "function", + "function": {"name": "weather", "arguments": arguments}}, + ]), + finish_reason="tool_calls", + ) + + handler: Final = OpenAIChatCompletionsHandler() + chunks: Final = [ModelResponseStream(choices=[choice(0, '{"city": "Pa'), choice(1, '{"city": "Rome"}')])] + assembled: Final = handler._rebuild_ended_stream_per_choice(chunks, None) + pruned: Final = handler._without_unfinished_tool_calls(assembled, chunks) + assert [ + [call.function.arguments for call in single.message.tool_calls or ()] for single in pruned.choices + ] == [[], ['{"city": "Rome"}']] + @pytest.mark.asyncio async def test_partial_arguments_stay_out_of_what_each_choice_inspects(self) -> None: """Half a JSON object is not a call a guardrail can judge, at any number of choices.""" From b3dbf6e097a43d81052686128af995a297421758 Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Tue, 3 Nov 2026 12:58:35 +0100 Subject: [PATCH 35/36] perf(guardrails): stop fingerprinting tool calls on every diff-mode round A round fingerprinted every streamed tool-call fragment, dumping and serializing each one, although it only needed to know whether tool calls were present. Repeating that as the stream grew made long tool arguments cost far more CPU than the arguments themselves. Ask whether the stream carries tool calls, and fingerprint only at the end of the stream, where the scan key compares them. On an 800-chunk tool stream this drops the per-round cost from 1083 ms to 1.5 ms across the stream, and the growth from quadratic to linear. --- .../chat/guardrail_translation/handler.py | 21 +++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 1d604d49faf..a4c9d8502a5 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -956,7 +956,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): self.merge_user_api_key_metadata_into_request(resolved_request_data, user_api_key_dict) try: inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) - if self._streamed_tool_call_fingerprints(responses_so_far): + if self._stream_carries_tool_calls(responses_so_far): stream_ended: Final = self._first_choice_has_finished(responses_so_far) tool_arguments_complete: Final = self._streamed_tool_arguments_are_complete(responses_so_far) assembled: Final = ( @@ -1089,12 +1089,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation): 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) - tool_call_fingerprints: Final = self._streamed_tool_call_fingerprints(responses_so_far) + carries_tool_calls: Final = self._stream_carries_tool_calls(responses_so_far) return StreamingScanKey( texts=tuple(self._combine_streaming_texts(chunks).values()), - tool_calls=tool_call_fingerprints if stream_ended else (), + tool_calls=self._streamed_tool_call_fingerprints(responses_so_far) if stream_ended else (), stream_ended=stream_ended, - tool_calls_in_flight=bool(tool_call_fingerprints) and not stream_ended, + tool_calls_in_flight=carries_tool_calls and not stream_ended, ) def released_stream_as_ended(self, responses_so_far: Sequence[object]) -> tuple[object, ...]: @@ -1118,6 +1118,19 @@ class OpenAIChatCompletionsHandler(BaseTranslation): for tool_call in _streamed_delta_tool_calls(stream_item_field(choice, "delta")) ) + @staticmethod + def _stream_carries_tool_calls(responses_so_far: Sequence[object]) -> bool: + """Whether any chunk so far carried a tool call, without fingerprinting them. + + A round only needs to know that tool calls are present; fingerprinting every fragment + dumps and serializes each one, and repeating that as the stream grows is what made a + long tool argument cost far more CPU than the arguments themselves. + """ + return any( + _streamed_delta_tool_calls(stream_item_field(choice, "delta")) + for choice in chain.from_iterable(_stream_chunk_choices(chunk) for chunk in responses_so_far) + ) + @staticmethod def _first_choice_has_finished(responses_so_far: Sequence[object]) -> bool: first_choices: Final = tuple( From 396fc787901a9ebc1c3e1e558d05e66b6a01435e Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Tue, 3 Nov 2026 13:03:38 +0100 Subject: [PATCH 36/36] style(guardrails): drop the explanatory comment on the choice index field AGENTS.md keeps explanatory comments out of the code, and the field's own docstring already says what it is for. --- litellm/integrations/custom_guardrail.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 7b1eaca13af..502c3cc95e0 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1414,8 +1414,6 @@ class CustomGuardrail(CustomLogger): end_time=end_time, duration=duration, masked_entity_count=masked_entity_count, - # Set by callers that inspect one choice at a time, so the per-choice entries of an - # n > 1 response can be told apart. choice_index=request_data.get("guardrail_choice_index"), **(tracing_detail or {}), )