scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 182

fix(model_armor): skip stream scanning for Responses API event streams

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
milan 2026-08-25 15:07:43 +00:00
parent 429b2c3546
commit 7f1d529a6d
2 changed files with 33 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py
View file

@ -849,16 +849,20 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
)
# Collect all chunks
all_chunks: Final[list[ModelResponseStream]] = []
all_chunks: Final[list[object]] = []
async for chunk in response:
all_chunks.append(chunk)
# /v1/messages arrives as raw SSE frames, which stream_chunk_builder cannot assemble
# /v1/messages arrives as raw SSE frames and /v1/responses as Responses API event
# objects; stream_chunk_builder can assemble neither
raw_sse: Final = is_raw_sse_stream(all_chunks)
chat_stream: Final = bool(all_chunks) and all(isinstance(chunk, ModelResponseStream) for chunk in all_chunks)
assembled_response: Final = (
assemble_anthropic_sse_stream(all_chunks, restore_identity=True)
if raw_sse
else stream_chunk_builder(chunks=all_chunks)
if chat_stream
else None
)
if isinstance(assembled_response, ModelResponse):
@ -963,6 +967,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
):
yield error_frame
return
elif not chat_stream:
verbose_proxy_logger.warning(
"Model Armor: streaming response contained unsupported event objects "
"(e.g. /v1/responses events); output scanning was skipped for this response."
)
# Return original chunks if no sanitization needed
for chunk in all_chunks:

23
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py
View file

@ -652,7 +652,7 @@ def _sse_armor_guardrail(**kwargs: object) -> ModelArmorGuardrail:
async def _drain_armor_streaming_hook(
guardrail: ModelArmorGuardrail, chunks: tuple[bytes, ...] = _ANTHROPIC_SSE_CHUNKS
guardrail: ModelArmorGuardrail, chunks: tuple[object, ...] = _ANTHROPIC_SSE_CHUNKS
) -> list[object]:
async def _stream():
for chunk in chunks:
@ -764,6 +764,27 @@ async def test_streaming_hook_masks_raw_anthropic_sse():
assert b"123-45-6789" not in body
@pytest.mark.asyncio
async def test_streaming_hook_passes_through_responses_api_events():
"""/v1/responses streams deliver event objects stream_chunk_builder cannot assemble.
Regression for `500 Error building chunks for logging/streaming usage calculation`.
"""
from litellm.types.llms.openai import GenericEvent, ResponsesAPIStreamEvents
guardrail = _sse_armor_guardrail()
events = (
GenericEvent(type=ResponsesAPIStreamEvents.RESPONSE_CREATED),
GenericEvent(type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED),
)
with patch.object(guardrail.async_handler, "post", AsyncMock()) as mock_post:
delivered = await _drain_armor_streaming_hook(guardrail, chunks=events)
mock_post.assert_not_called()
assert tuple(delivered) == events
@pytest.mark.asyncio
async def test_streaming_hook_fails_closed_on_unparseable_raw_sse():
guardrail = _sse_armor_guardrail()