From d20535d366f2f0ab8a16d3c5c06a39427243e11a Mon Sep 17 00:00:00 2001 From: Alessandro Pignati Date: Fri, 6 Nov 2026 00:31:02 +0100 Subject: [PATCH] fix(guardrails): leave the tool calls of a one-choice stream to the block inspection The round that saw the stream end still read the assembled tool calls, so a single-choice stream paid two tool rounds where main pays one, the end-of-stream block inspection. Read them only where the per-choice rounds need them, and cover the ended round so the count cannot drift back. The choice index regressions move into the test files their fixes are mapped to, and the comments this branch added go with them. --- litellm/integrations/custom_guardrail.py | 5 - .../chat/guardrail_translation/handler.py | 4 +- litellm/proxy/litellm_pre_call_utils.py | 4 - .../integrations/test_custom_guardrail.py | 23 ++++ .../test_openai_guardrail_handler.py | 108 +++++++++++------- .../test_guardrail_choice_index_label.py | 73 ------------ .../unit/proxy/test_litellm_pre_call_utils.py | 29 +++++ 7 files changed, 118 insertions(+), 128 deletions(-) delete mode 100644 tests/unit/proxy/guardrails/test_guardrail_choice_index_label.py diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index cb000ad0062..7e58625e6e7 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -150,11 +150,6 @@ def _unified_hook_fields(guardrail: "CustomGuardrail", request_data: Mapping[str def _inspected_choice_index(request_data: Mapping[str, object]) -> int | None: - """The choice this round inspected, as the OpenAI guardrail handler labelled it. - - Read behind a type guard: the key rides on the request body, so anything reaching - here that is not a plain int was not written by that handler. - """ labelled: Final = request_data.get("guardrail_choice_index") return labelled if isinstance(labelled, int) and not isinstance(labelled, bool) else None diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index a5278c333b6..60bc7a639dc 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -1047,9 +1047,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): try: inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) stream_ended: Final = fold.ended - inspects_tool_calls: Final = ( - inspecting_one_choice or stream_ended or len(fold.choices) > 1 - ) and fold.carries_tool_calls + inspects_tool_calls: Final = (inspecting_one_choice or len(fold.choices) > 1) and fold.carries_tool_calls if inspects_tool_calls: tool_arguments_complete: Final = self._arguments_are_complete(fold.arguments) assembled: Final = ( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index f7a0dc7449f..9b5b692f0ca 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -858,16 +858,12 @@ def strip_client_guardrail_information(data: dict[str, object]) -> None: # muta ) -# ``guardrail_choice_index`` is written by the OpenAI guardrail handler around each -# per-choice inspection round and read by guardrail spend logs as fact; a client value -# has no legitimate meaning and no key or team setting keeps it. _GUARDRAIL_RESERVED_REQUEST_FIELDS: Final = frozenset({"guardrail_choice_index"}) def _strip_guardrail_reserved_fields( data: dict[str, object], # mutable-ok: strips in place on the request body the pre-call pipeline threads through ) -> None: - """Drop the guardrail-owned per-choice label from the client request body.""" stripped: Final = tuple(sorted(field for field in _GUARDRAIL_RESERVED_REQUEST_FIELDS if field in data)) for field in stripped: data.pop(field, None) diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index c32c1b71f75..f5cce974c1a 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -3813,3 +3813,26 @@ def test_merged_guardrail_entries_names_each_guardrail_once_preferring_its_first {"pii": {"extra_body": {"threshold": 1}}}, "model-guard", ] + + +def _logged_choice_index(request_data: dict[str, object]) -> object: + guardrail: Final = CustomGuardrail(guardrail_name="choice-label-test") + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={}, + request_data=request_data, + guardrail_status="success", + ) + metadata: Final = request_data["metadata"] + assert isinstance(metadata, dict) + rows: Final = metadata["standard_logging_guardrail_information"] + assert isinstance(rows, list) + return rows[0]["choice_index"] + + +def test_a_non_integer_choice_index_is_not_logged() -> None: + """The label rides on the request body, so only one the handler wrote may be persisted.""" + assert _logged_choice_index({"guardrail_choice_index": "1; drop"}) is None + + +def test_an_integer_choice_index_is_logged() -> None: + assert _logged_choice_index({"guardrail_choice_index": 1}) == 1 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 2024c2fecfa..c302e052b83 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 @@ -1591,7 +1591,8 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: 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. + Sending them would hand the guardrail something like `{"ci`, which main never did. Shown + at n>1, the only shape where a sampled round reads tool calls at all. """ from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices @@ -1607,12 +1608,15 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: return inputs def chunk(text: str, arguments: 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": arguments}, - }]), finish_reason=finish_reason, - )]) + return ModelResponseStream(choices=[ + StreamingChoices( + index=index, delta=Delta(content=text, tool_calls=[{ + "index": 0, "id": f"call_{index}", "type": "function", + "function": {"name": "weather", "arguments": arguments}, + }]), finish_reason=finish_reason, + ) + for index in (0, 1) + ]) RecordingGuardrail.seen = [] handler: Final = OpenAIChatCompletionsHandler() @@ -1630,9 +1634,9 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: 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"}'] + assert [ + call["function"]["arguments"] for calls in RecordingGuardrail.seen for call in calls or () + ] == ['{"city": "Paris"}', '{"city": "Paris"}'] @pytest.mark.asyncio async def test_diff_mode_withholds_parallel_tool_calls_until_every_one_is_complete(self) -> None: @@ -1658,10 +1662,13 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: 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, - )])] + one_still_empty: Final = [ModelResponseStream(choices=[ + StreamingChoices( + index=index, delta=Delta(content="looking", tool_calls=[call(0, '{"city": "Paris"}'), call(1, "")]), + finish_reason=None, + ) + for index in (0, 1) + ])] await handler.process_output_streaming_response( responses_so_far=one_still_empty, guardrail_to_apply=guardrail, litellm_logging_obj=None, stream_transform_sink=StreamTransformSink(), @@ -1669,18 +1676,22 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: 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", - )])] + both_complete: Final = [ModelResponseStream(choices=[ + StreamingChoices( + index=index, + delta=Delta(content="looking", + tool_calls=[call(0, '{"city": "Paris"}'), call(1, '{"city": "Rome"}')]), + finish_reason="tool_calls", + ) + for index in (0, 1) + ])] 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"}'] + assert [c["function"]["arguments"] for calls in RecordingGuardrail.seen for c in calls or ()] == [ + '{"city": "Paris"}', '{"city": "Rome"}', '{"city": "Paris"}', '{"city": "Rome"}' + ] @pytest.mark.asyncio async def test_a_legacy_post_call_hook_sees_the_whole_response_once(self) -> None: @@ -1969,7 +1980,11 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: @pytest.mark.asyncio async def test_a_one_choice_round_keeps_the_shape_it_always_had(self) -> None: - """One choice has nothing to keep apart, so a sampled round still sends only its text.""" + """One choice has nothing to keep apart, so a sampled round still sends only its text. + + Including the round that sees the stream end, which is where the block inspection, not + this path, hands the guardrail the assembled tool calls, as it does on main. + """ from litellm.llms.base_llm.guardrail_translation.base_translation import StreamTransformSink from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices @@ -1984,21 +1999,25 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: return inputs RecordingGuardrail.seen = [] - 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"}'}}, - ]), - )])] - await OpenAIChatCompletionsHandler().process_output_streaming_response( - responses_so_far=chunks, - guardrail_to_apply=RecordingGuardrail(guardrail_name="one-choice"), - litellm_logging_obj=None, - request_data={"model": "gpt-4"}, - stream_transform_sink=StreamTransformSink(), - ) - assert RecordingGuardrail.seen == [(None, None)] + def chunk(finish_reason: str | None) -> ModelResponseStream: + return 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=finish_reason, + )]) + + for finish_reason in (None, "tool_calls"): + await OpenAIChatCompletionsHandler().process_output_streaming_response( + responses_so_far=[chunk(finish_reason)], + guardrail_to_apply=RecordingGuardrail(guardrail_name="one-choice"), + litellm_logging_obj=None, + request_data={"model": "gpt-4"}, + stream_transform_sink=StreamTransformSink(), + ) + assert RecordingGuardrail.seen == [(None, None), (None, None)] @pytest.mark.asyncio async def test_a_parallel_round_still_inspects_each_choice(self) -> None: @@ -2092,17 +2111,20 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: 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", - )])] + ended_with_one_unfinished: Final = [ModelResponseStream(choices=[ + StreamingChoices( + index=index, delta=Delta(content="looking", tool_calls=[call(0, '{"city": "Paris"}'), call(1, "")]), + finish_reason="tool_calls", + ) + for index in (0, 1) + ])] 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] + assert RecordingGuardrail.seen_calls == [None, None] inspected: Final = RecordingGuardrail.seen_response[-1] arguments: Final = tuple( tool_call.function.arguments diff --git a/tests/unit/proxy/guardrails/test_guardrail_choice_index_label.py b/tests/unit/proxy/guardrails/test_guardrail_choice_index_label.py deleted file mode 100644 index 45f69ad03d8..00000000000 --- a/tests/unit/proxy/guardrails/test_guardrail_choice_index_label.py +++ /dev/null @@ -1,73 +0,0 @@ -"""The per-choice label on guardrail spend logs is proxy-owned, never caller-supplied. - -`guardrail_choice_index` names which choice a guardrail round inspected. The OpenAI -handler writes it around each round; `StandardLoggingGuardrailInformation` reads it and -spend logs persist it. A caller putting the key in its own request body would otherwise -have that value stored verbatim, with any type or size, on every guardrail row of the -request. -""" - -from typing import Final -from unittest.mock import MagicMock - -import pytest -from fastapi import Request - -from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request - - -def _request_mock() -> Request: - request_mock = MagicMock(spec=Request) - request_mock.url = MagicMock() - request_mock.url.path = "/v1/chat/completions" - request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" - request_mock.method = "POST" - request_mock.query_params = {} - request_mock.headers = {"Content-Type": "application/json"} - request_mock.client = MagicMock() - request_mock.client.host = "127.0.0.1" - return request_mock - - -def _logged_choice_index(request_data: dict[str, object]) -> object: - guardrail: Final = CustomGuardrail(guardrail_name="choice-label-test") - guardrail.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response={}, - request_data=request_data, - guardrail_status="success", - ) - metadata: Final = request_data["metadata"] - assert isinstance(metadata, dict) - rows: Final = metadata["standard_logging_guardrail_information"] - assert isinstance(rows, list) - return rows[0]["choice_index"] - - -@pytest.mark.asyncio -async def test_proxy_drops_a_caller_supplied_choice_index() -> None: - data: Final[dict[str, object]] = { - "model": "gpt-4", - "messages": [{"role": "user", "content": "hi"}], - "guardrail_choice_index": "x" * 4096, - } - - updated: Final = await add_litellm_data_to_request( - data=data, - request=_request_mock(), - user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), - proxy_config=MagicMock(), - general_settings={}, - version="test-version", - ) - - assert "guardrail_choice_index" not in updated - - -def test_a_non_integer_label_is_not_logged() -> None: - assert _logged_choice_index({"guardrail_choice_index": "1; drop"}) is None - - -def test_an_integer_label_is_logged() -> None: - assert _logged_choice_index({"guardrail_choice_index": 1}) == 1 diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 91c87c63067..3b47bca76f2 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -11444,3 +11444,32 @@ def test_caller_tags_reach_litellm_metadata_whether_metadata_is_a_dict_or_a_json LiteLLMProxyRequestSetup.merge_caller_body_tags(request_data=data) assert set(data["litellm_metadata"]["tags"]) == {"key-tag", "client-tag"} assert data["metadata"] == caller_metadata + + +@pytest.mark.asyncio +async def test_proxy_drops_a_caller_supplied_guardrail_choice_index(): + request_mock = MagicMock(spec=Request) + request_mock.url = MagicMock() + request_mock.url.path = "/v1/chat/completions" + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + data: Final[dict[str, object]] = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "guardrail_choice_index": "x" * 4096, + } + + updated: Final = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert "guardrail_choice_index" not in updated