From 79e04e60511e5d5c26d5f797cfcb5f74281a4d20 Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Thu, 5 Nov 2026 00:39:47 +0100 Subject: [PATCH] fix(guardrails): hand a legacy post-call hook every choice in one call The per-choice rounds exist so a guardrail reads one candidate completion at a time instead of a response no provider would return. A legacy post-call hook's contract is the opposite: one call with the response its route would have given it, which on an n=2 tool stream it stopped getting. Let a guardrail say which it is, and hand the hook the whole response again. --- litellm/integrations/custom_guardrail.py | 10 +++++ .../chat/guardrail_translation/handler.py | 7 +++- .../proxy/policy_engine/pipeline_executor.py | 6 +++ .../test_openai_guardrail_handler.py | 42 +++++++++++++++++++ 4 files changed, 63 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 5220dc576fc..a49aa04c709 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1136,6 +1136,16 @@ class CustomGuardrail(CustomLogger): """ return True + def inspects_choices_separately(self) -> bool: + """Whether a stream's choices may be handed over one at a time. + + An n>1 stream carrying tool calls is inspected choice by choice, so a guardrail reads + one candidate completion at a time instead of a response no provider would return. A + guardrail whose contract is the whole response, like a legacy post-call hook, overrides + this to return False and is handed every choice in one call. + """ + return True + def structured_messages_cover_full_request(self) -> bool: """Whether returned ``structured_messages`` span the whole request. diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 8c938ad6c1f..d23b52dc154 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -672,7 +672,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): 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 + if pre_guardrail_tool_calls + and len(model_response.choices) > 1 + and guardrail_to_apply.inspects_choices_separately() else (model_response,) ) from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite @@ -977,7 +979,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): ) if inspected is not None: resolved_request_data["response"] = inspected - if inspected is not None and len(inspected.choices) > 1: + inspects_each: Final = guardrail_to_apply.inspects_choices_separately() + if inspected is not None and len(inspected.choices) > 1 and inspects_each: choice_rounds: Final = tuple( ( choice, diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 4ad29bd3478..8a6b52b1a31 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -139,6 +139,9 @@ class _StreamRewriteObserver(CustomGuardrail): def structured_messages_cover_full_request(self) -> bool: return self.inner.structured_messages_cover_full_request() + def inspects_choices_separately(self) -> bool: + return self.inner.inspects_choices_separately() + @_logged_by_inner_guardrail async def apply_guardrail( self, @@ -225,6 +228,9 @@ class _LegacyHookStreamAdapter(CustomGuardrail): def structured_messages_cover_full_request(self) -> bool: return self.inner.structured_messages_cover_full_request() + def inspects_choices_separately(self) -> bool: + return False + @_logged_by_inner_guardrail async def apply_guardrail( 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 46dead0d59a..1343333100d 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,48 @@ 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_legacy_post_call_hook_sees_the_whole_response_once(self) -> None: + """The legacy hook's contract is one call with the response the route would have given it, + so the per-choice rounds, which exist for guardrails that read one choice at a time, do not + apply to it.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.policy_engine.pipeline_executor import _LegacyHookStreamAdapter + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + class LegacyHookGuardrail(CustomGuardrail): + choice_counts: list[int] = [] + + async def async_post_call_success_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: object + ) -> None: + type(self).choice_counts.append(len(getattr(response, "choices", None) or ())) + + LegacyHookGuardrail.choice_counts = [] + handler: Final = OpenAIChatCompletionsHandler() + adapter: Final = _LegacyHookStreamAdapter( + LegacyHookGuardrail(guardrail_name="legacy"), handler, UserAPIKeyAuth(api_key="hashed-key") + ) + chunks: Final = [ModelResponseStream(choices=[ + StreamingChoices( + index=index, + delta=Delta(content="looking", tool_calls=[ + {"index": 0, "id": f"call_{index}", "type": "function", + "function": {"name": "weather", "arguments": '{"city": "Paris"}'}}, + ]), + finish_reason="tool_calls", + ) + for index in (0, 1) + ])] + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=adapter, + litellm_logging_obj=None, + request_data={}, + deliver_ended_stream_rewrites=True, + ) + assert LegacyHookGuardrail.choice_counts == [2] + @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."""