mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(model_armor): scan each responses field once, not twice
Separating delta fields stopped the terminal body from matching the delta text, so a turn with two visible fields sent Model Armor both copies. Only the delta fields the body does not already carry are appended now.
This commit is contained in:
parent
a0b0fec5f7
commit
669e70055e
2 changed files with 67 additions and 10 deletions
|
|
@ -984,8 +984,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
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.
|
||||
def _responses_delta_field_texts(all_chunks: Sequence[object]) -> tuple[str, ...]:
|
||||
"""Text each field of a ``/v1/responses`` turn 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
|
||||
|
|
@ -997,7 +997,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
if getattr(chunk, "type", None) in _RESPONSES_DELTA_EVENT_TYPES
|
||||
and isinstance(delta := getattr(chunk, "delta", None), str)
|
||||
)
|
||||
return "\n".join(
|
||||
return tuple(
|
||||
"".join(delta for field, delta in deltas if field == streamed_field)
|
||||
for streamed_field in dict.fromkeys(field for field, _ in deltas)
|
||||
)
|
||||
|
|
@ -1011,17 +1011,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
"""Text to scan for a buffered stream, which is everything the client is about to receive.
|
||||
|
||||
A ``/v1/responses`` stream also spells out reasoning summaries and tool-call arguments in
|
||||
delta events that its terminal body never repeats, so body and deltas are scanned together.
|
||||
delta events that its terminal body never repeats, so every delta field the body does not
|
||||
already carry is scanned after it.
|
||||
"""
|
||||
content: Final = self._extract_streaming_content(assembled_response)
|
||||
if surface is not _StreamSurface.RESPONSES:
|
||||
return content
|
||||
delta_text: Final = self._responses_delta_text(all_chunks)
|
||||
if delta_text in content:
|
||||
return content
|
||||
if content in delta_text:
|
||||
return delta_text
|
||||
return f"{content}\n{delta_text}"
|
||||
unscanned: Final = tuple(text for text in self._responses_delta_field_texts(all_chunks) if text not in content)
|
||||
return "\n".join(part for part in (content, *unscanned) if part)
|
||||
|
||||
@staticmethod
|
||||
def _apply_sanitized_content(assembled_response: ModelResponse, sanitized_content: str) -> None:
|
||||
|
|
|
|||
|
|
@ -4855,6 +4855,66 @@ async def test_streaming_responses_one_fields_deltas_still_join_into_a_single_fi
|
|||
assert "Streaming response blocked by Model Armor" in rendered
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_responses_fields_the_body_repeats_are_not_scanned_a_second_time():
|
||||
"""A turn whose visible fields all reach the terminal body is scanned once, not twice.
|
||||
|
||||
Two output_text fields stream as deltas and come back in the completed body, so scanning the
|
||||
deltas on top of the body would send Model Armor two copies of everything the client sees.
|
||||
"""
|
||||
from litellm.types.llms.openai import (
|
||||
OutputTextDeltaEvent,
|
||||
ResponseCompletedEvent,
|
||||
ResponsesAPIResponse,
|
||||
ResponsesAPIStreamEvents,
|
||||
)
|
||||
|
||||
paragraphs = ("the first thing to know", "a second and separate point")
|
||||
text_deltas = tuple(
|
||||
OutputTextDeltaEvent(
|
||||
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
|
||||
item_id=f"msg_{index}",
|
||||
output_index=index,
|
||||
content_index=0,
|
||||
delta=paragraph,
|
||||
)
|
||||
for index, paragraph in enumerate(paragraphs)
|
||||
)
|
||||
completed = ResponseCompletedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=ResponsesAPIResponse(
|
||||
id="resp_1",
|
||||
created_at=0,
|
||||
model="gpt-5-mini",
|
||||
object="response",
|
||||
output=[
|
||||
{
|
||||
"type": "message",
|
||||
"id": f"msg_{index}",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{"type": "output_text", "text": paragraph, "annotations": []}],
|
||||
}
|
||||
for index, paragraph in enumerate(paragraphs)
|
||||
],
|
||||
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, (*text_deltas, completed))
|
||||
|
||||
post.assert_called_once()
|
||||
scanned = post.call_args.kwargs["json"]["modelResponseData"]["text"]
|
||||
assert [scanned.count(paragraph) for paragraph in paragraphs] == [1, 1]
|
||||
rendered = "".join(str(item) for item in delivered)
|
||||
assert all(paragraph in rendered for paragraph in paragraphs)
|
||||
|
||||
|
||||
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