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