scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(guardrails): leave the tool calls of a one-choice stream to the block inspection

Browse source 
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.
This commit is contained in:
Alessandro Pignati 2026-11-06 00:31:02 +01:00
parent 277f14c134
commit d20535d366
7 changed files with 118 additions and 128 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/integrations/custom_guardrail.py
View file

@ -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

4
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -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 = (

4
litellm/proxy/litellm_pre_call_utils.py
View file

@ -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)

23
tests/unit/integrations/test_custom_guardrail.py
View file

@ -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

108
tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py
View file

@ -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

73
tests/unit/proxy/guardrails/test_guardrail_choice_index_label.py
View file

@ -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

29
tests/unit/proxy/test_litellm_pre_call_utils.py
View file

@ -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