mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41: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)
|
return tuple(str(getattr(chunk, attr, None)) for attr in _RESPONSES_DELTA_FIELD_ATTRS)
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _responses_delta_text(all_chunks: Sequence[object]) -> str:
|
def _responses_delta_field_texts(all_chunks: Sequence[object]) -> tuple[str, ...]:
|
||||||
"""Text a ``/v1/responses`` stream has already spelled out in its delta events.
|
"""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,
|
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
|
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
|
if getattr(chunk, "type", None) in _RESPONSES_DELTA_EVENT_TYPES
|
||||||
and isinstance(delta := getattr(chunk, "delta", None), str)
|
and isinstance(delta := getattr(chunk, "delta", None), str)
|
||||||
)
|
)
|
||||||
return "\n".join(
|
return tuple(
|
||||||
"".join(delta for field, delta in deltas if field == streamed_field)
|
"".join(delta for field, delta in deltas if field == streamed_field)
|
||||||
for streamed_field in dict.fromkeys(field for field, _ in deltas)
|
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.
|
"""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
|
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)
|
content: Final = self._extract_streaming_content(assembled_response)
|
||||||
if surface is not _StreamSurface.RESPONSES:
|
if surface is not _StreamSurface.RESPONSES:
|
||||||
return content
|
return content
|
||||||
delta_text: Final = self._responses_delta_text(all_chunks)
|
unscanned: Final = tuple(text for text in self._responses_delta_field_texts(all_chunks) if text not in content)
|
||||||
if delta_text in content:
|
return "\n".join(part for part in (content, *unscanned) if part)
|
||||||
return content
|
|
||||||
if content in delta_text:
|
|
||||||
return delta_text
|
|
||||||
return f"{content}\n{delta_text}"
|
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _apply_sanitized_content(assembled_response: ModelResponse, sanitized_content: str) -> None:
|
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
|
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():
|
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."""
|
"""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 (
|
from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import (
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue