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