From 95348ba610fcd21b742c64786d1fa0eed7d4b9f4 Mon Sep 17 00:00:00 2001 From: albertbausili Date: Fri, 16 Oct 2026 09:36:57 +0200 Subject: [PATCH] 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):