mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
fix(model_armor): scan responses delta fields apart from each other
A Responses turn spells out its reasoning summary, its visible answer and its tool-call arguments in separate delta events. Joining every delta into one string let a finding form across the boundary between two fields that each carry nothing to find, so a safe stream could be blocked. Group the deltas by the field they belong to, join a field's own deltas as they streamed, and keep the fields apart.
This commit is contained in:
parent
efa6c87a95
commit
a0b0fec5f7
2 changed files with 133 additions and 3 deletions
|
|
@ -69,6 +69,10 @@ _RESPONSES_DELTA_EVENT_TYPES: Final = frozenset(
|
|||
event.value for event in ResponsesAPIStreamEvents if event.value.endswith(".delta")
|
||||
)
|
||||
|
||||
# What makes two delta events part of the same field of the turn, rather than two fields that merely
|
||||
# streamed next to each other
|
||||
_RESPONSES_DELTA_FIELD_ATTRS: Final = ("type", "item_id", "output_index", "content_index", "summary_index")
|
||||
|
||||
|
||||
class _StreamSurface(Enum):
|
||||
"""Wire format of a buffered streaming response, which decides how it is read and how it is refused."""
|
||||
|
|
@ -974,15 +978,29 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
return self._responses_api_response_text(assembled_response)
|
||||
return self._extract_content_from_response(assembled_response)
|
||||
|
||||
@staticmethod
|
||||
def _responses_delta_field(chunk: object) -> tuple[str, ...]:
|
||||
"""Which field of the turn a delta event belongs to."""
|
||||
return tuple(str(getattr(chunk, attr, None)) for attr in _RESPONSES_DELTA_FIELD_ATTRS)
|
||||
|
||||
@staticmethod
|
||||
def _responses_delta_text(all_chunks: Sequence[object]) -> str:
|
||||
"""Text a ``/v1/responses`` stream has already spelled out in its delta events."""
|
||||
return "".join(
|
||||
delta
|
||||
"""Text a ``/v1/responses`` stream has already spelled out in its delta events.
|
||||
|
||||
One field's deltas are joined as they streamed, since a finding can be split across them,
|
||||
and separate fields stay apart, so a reasoning summary running into the visible answer
|
||||
cannot spell out a finding that neither of them carries.
|
||||
"""
|
||||
deltas: Final = tuple(
|
||||
(ModelArmorGuardrail._responses_delta_field(chunk), delta)
|
||||
for chunk in all_chunks
|
||||
if getattr(chunk, "type", None) in _RESPONSES_DELTA_EVENT_TYPES
|
||||
and isinstance(delta := getattr(chunk, "delta", None), str)
|
||||
)
|
||||
return "\n".join(
|
||||
"".join(delta for field, delta in deltas if field == streamed_field)
|
||||
for streamed_field in dict.fromkeys(field for field, _ in deltas)
|
||||
)
|
||||
|
||||
def _streaming_content_to_scan(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -4743,6 +4743,118 @@ async def test_streaming_responses_reasoning_summary_deltas_are_scanned_alongsid
|
|||
assert "Streaming response blocked by Model Armor" in rendered
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_responses_deltas_of_separate_fields_do_not_form_a_finding_across_their_boundary():
|
||||
"""Two fields of a turn are separate text, so what runs across their boundary is not model output.
|
||||
|
||||
A reasoning summary ending in half a card number and an answer opening with the other half
|
||||
each carry nothing to find, and joining them without a break would invent one.
|
||||
"""
|
||||
from litellm.types.llms.openai import (
|
||||
OutputTextDeltaEvent,
|
||||
ReasoningSummaryTextDeltaEvent,
|
||||
ResponseCompletedEvent,
|
||||
ResponsesAPIResponse,
|
||||
ResponsesAPIStreamEvents,
|
||||
)
|
||||
|
||||
answer = "1111-1111 is not a full card"
|
||||
summary_delta = ReasoningSummaryTextDeltaEvent(
|
||||
type=ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA,
|
||||
item_id="rs_1",
|
||||
output_index=0,
|
||||
delta="the prefix they gave me is 4111-1111-",
|
||||
)
|
||||
text_delta = OutputTextDeltaEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
|
||||
item_id="msg_1",
|
||||
output_index=1,
|
||||
content_index=0,
|
||||
delta=answer,
|
||||
)
|
||||
completed = ResponseCompletedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=ResponsesAPIResponse(
|
||||
id="resp_1",
|
||||
created_at=0,
|
||||
model="gpt-5-mini",
|
||||
object="response",
|
||||
output=[
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_1",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{"type": "output_text", "text": answer, "annotations": []}],
|
||||
}
|
||||
],
|
||||
parallel_tool_calls=False,
|
||||
tool_choice="auto",
|
||||
tools=[],
|
||||
),
|
||||
)
|
||||
guardrail = _surface_guardrail()
|
||||
post = _armor_post_mock(_MODEL_ARMOR_CLEAN)
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", post):
|
||||
delivered = await _drain_surface_hook(guardrail, (summary_delta, text_delta, completed))
|
||||
|
||||
post.assert_called_once()
|
||||
scanned = post.call_args.kwargs["json"]["modelResponseData"]["text"]
|
||||
assert "4111-1111-" in scanned
|
||||
assert answer in scanned
|
||||
assert "4111-1111-1111-1111" not in scanned
|
||||
rendered = "".join(str(item) for item in delivered)
|
||||
assert "Streaming response blocked by Model Armor" not in rendered
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_responses_one_fields_deltas_still_join_into_a_single_finding():
|
||||
"""A card number split across two deltas of one field is still one card number to scan."""
|
||||
from litellm.types.llms.openai import (
|
||||
OutputTextDeltaEvent,
|
||||
ResponseCompletedEvent,
|
||||
ResponsesAPIResponse,
|
||||
ResponsesAPIStreamEvents,
|
||||
)
|
||||
|
||||
halves = ("my card is 4111-1111-", "1111-1111")
|
||||
text_deltas = tuple(
|
||||
OutputTextDeltaEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
|
||||
item_id="msg_1",
|
||||
output_index=0,
|
||||
content_index=0,
|
||||
delta=half,
|
||||
)
|
||||
for half in halves
|
||||
)
|
||||
completed = ResponseCompletedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=ResponsesAPIResponse(
|
||||
id="resp_1",
|
||||
created_at=0,
|
||||
model="gpt-5-mini",
|
||||
object="response",
|
||||
output=[],
|
||||
parallel_tool_calls=False,
|
||||
tool_choice="auto",
|
||||
tools=[],
|
||||
),
|
||||
)
|
||||
guardrail = _surface_guardrail()
|
||||
post = _armor_post_mock(_MODEL_ARMOR_BLOCK)
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", post):
|
||||
delivered = await _drain_surface_hook(guardrail, (*text_deltas, completed))
|
||||
|
||||
post.assert_called_once()
|
||||
assert "4111-1111-1111-1111" in post.call_args.kwargs["json"]["modelResponseData"]["text"]
|
||||
rendered = "".join(str(item) for item in delivered)
|
||||
assert "4111-1111-1111-1111" not in rendered
|
||||
assert "Streaming response blocked by Model Armor" in rendered
|
||||
|
||||
|
||||
def test_every_responses_delta_event_is_in_the_scanned_set():
|
||||
"""Every ``.delta`` the Responses event enum defines is model output on its way to the client."""
|
||||
from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import (
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue