mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix(streaming): do not raise when a collected stream carries raw SSE frames
The guardrail post-call hooks collect a stream by appending whatever the response iterator yields, so `all_chunks` can hold raw SSE frames (bytes or str) next to parsed chunks. `stream_chunk_builder` indexes every entry, so the first raw frame raises and the hook turns a completion the client had already received into a 500. Six of the nine guardrails that reassemble a collected stream call `stream_chunk_builder` with no guard for this: noma, model_armor, repelloai, panw_prisma_airs, cisco_ai_defense and microsoft_purview. Only the bedrock and tool_permission hooks check `is_raw_sse_stream` first, and presidio filters by type while collecting. Fix it once at the assembler rather than six times at the callers. Assembling only the parsed remainder would hand back a response that silently omits content, which a guardrail would then scan and pass, so treat the whole stream as unassemblable instead, matching the `any(...)` stance `is_raw_sse_stream` already takes. Returning None meets the empty-case contract immediately above it, and every one of those callers already routes a non-ModelResponse result to its pass-through path. Relates to #37873
This commit is contained in:
parent
6a0d03914c
commit
462199de08
2 changed files with 64 additions and 0 deletions
|
|
@ -8548,6 +8548,19 @@ def stream_chunk_builder(
|
|||
if not chunks:
|
||||
return None
|
||||
|
||||
# A collected stream can carry raw SSE frames (bytes/str) alongside parsed
|
||||
# chunks, because the guardrail post-call hooks append whatever the response
|
||||
# iterator yields. Those have no chunk shape, so the first one reached raises
|
||||
# and the caller turns a completion the client already received into a 500.
|
||||
# Assembling only the parsed remainder would instead hand back a response that
|
||||
# silently omits content, so treat the whole stream as unassemblable, matching
|
||||
# the `any(...)` stance `is_raw_sse_stream` already takes for this shape.
|
||||
# Returning None meets the empty-case contract just above, which every caller
|
||||
# handles by passing the stream through unscanned.
|
||||
if any(isinstance(chunk, (str, bytes, bytearray)) for chunk in chunks):
|
||||
verbose_logger.debug("stream_chunk_builder: raw SSE frames in the collected stream, not assembling")
|
||||
return None
|
||||
|
||||
processor: Final = ChunkProcessor(chunks, messages)
|
||||
chunks = processor.chunks
|
||||
|
||||
|
|
|
|||
|
|
@ -1337,3 +1337,54 @@ def test_calculate_usage_fills_unknown_split_from_reasoning_estimate(
|
|||
assert usage.completion_tokens == 100
|
||||
assert usage.completion_tokens_details.reasoning_tokens == expected_reasoning_tokens
|
||||
assert usage.completion_tokens_details.text_tokens == expected_text_tokens
|
||||
|
||||
|
||||
class TestRawSSEFramesInACollectedStream:
|
||||
"""A collected stream can carry raw SSE frames next to parsed chunks.
|
||||
|
||||
The guardrail post-call hooks append whatever the response iterator yields, so
|
||||
`stream_chunk_builder` receives bytes/str entries it cannot index. Reaching one
|
||||
used to raise, and the caller turned a completion the client had already
|
||||
received into a 500 (issue #37873). Only two of the nine guardrails that
|
||||
reassemble a stream carry the `is_raw_sse_stream` guard that avoids this call.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _chunk(content):
|
||||
return ModelResponseStream(
|
||||
id="chatcmpl-a83572f7",
|
||||
created=1,
|
||||
model="nvidia_nim/meta/llama-3.1-8b-instruct",
|
||||
object="chat.completion.chunk",
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
finish_reason=None,
|
||||
index=0,
|
||||
delta=Delta(role="assistant", content=content),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"frame",
|
||||
[b'data: {"id":"x"}\n\n', 'data: {"id":"x"}\n\n', bytearray(b"data: x\n\n")],
|
||||
)
|
||||
def test_a_raw_frame_does_not_raise(self, frame):
|
||||
chunks = [self._chunk("Hello"), frame, self._chunk(" world")]
|
||||
|
||||
# Must not raise: an APIError here becomes a 500 for an already-delivered completion.
|
||||
assert stream_chunk_builder(chunks=chunks) is None
|
||||
|
||||
def test_an_all_raw_stream_does_not_raise(self):
|
||||
assert stream_chunk_builder(chunks=[b"data: a\n\n", b"data: b\n\n"]) is None
|
||||
|
||||
def test_a_stream_of_real_chunks_is_still_assembled(self):
|
||||
"""Control: the ordinary path is untouched, so this passes with or without the fix."""
|
||||
built = stream_chunk_builder(chunks=[self._chunk("Hello"), self._chunk(" world")])
|
||||
|
||||
assert built is not None
|
||||
assert built.choices[0].message.content == "Hello world"
|
||||
|
||||
def test_empty_and_none_contracts_are_unchanged(self):
|
||||
"""Control: the existing empty-case contract still holds."""
|
||||
assert stream_chunk_builder(chunks=[]) is None
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue