fix(anthropic): handle partial JSON chunks in streaming responses (#17493)

Fixes #17473 - Anthropic streaming fails with JSONDecodeError when
network fragmentation causes SSE data to arrive in partial chunks.

Changes:
- Add accumulated_json buffer and chunk_type to ModelResponseIterator
- Add _handle_accumulated_json_chunk() to accumulate partial JSON
- Add _parse_sse_data() to handle both complete and partial chunks
- Modify __next__ and __anext__ to use accumulation logic
- Add unit tests for partial chunk handling
This commit is contained in:
Cesar Garcia 2025-12-08 04:34:42 -03:00 • committed by GitHub
parent 0f5694c8eb
commit b6b155d67b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 210 additions and 58 deletions

View file

@ -10,6 +10,7 @@ from typing import (
Callable,
Dict,
List,
Literal,
Optional,
Tuple,
Union,
@ -498,6 +499,11 @@ class ModelResponseIterator:
# Track if we've converted any response_format tools (affects finish_reason)
self.converted_response_format_tool: bool = False
# For handling partial JSON chunks from fragmentation
# See: https://github.com/BerriAI/litellm/issues/17473
self.accumulated_json: str = ""
self.chunk_type: Literal["valid_json", "accumulated_json"] = "valid_json"
def check_empty_tool_call_args(self) -> bool:
"""
Check if the tool call block so far has been an empty string
@ -866,42 +872,105 @@ class ModelResponseIterator:
usage = self._handle_usage(anthropic_usage_chunk=message_delta["usage"])
return finish_reason, usage
def _handle_accumulated_json_chunk(
self, data_str: str
) -> Optional[GenericStreamingChunk]:
"""
Handle partial JSON chunks by accumulating them until valid JSON is received.
This fixes network fragmentation issues where SSE data chunks may be split
across TCP packets. See: https://github.com/BerriAI/litellm/issues/17473
Args:
data_str: The JSON string to parse (without "data:" prefix)
Returns:
GenericStreamingChunk if JSON is complete, None if still accumulating
"""
# Accumulate JSON data
self.accumulated_json += data_str
# Try to parse the accumulated JSON
try:
data_json = json.loads(self.accumulated_json)
self.accumulated_json = "" # Reset after successful parsing
return self.chunk_parser(chunk=data_json)
except json.JSONDecodeError:
# If it's not valid JSON yet, continue to the next chunk
return None
def _parse_sse_data(self, str_line: str) -> Optional[GenericStreamingChunk]:
"""
Parse SSE data line, handling both complete and partial JSON chunks.
Args:
str_line: The SSE line starting with "data:"
Returns:
GenericStreamingChunk if parsing succeeded, None if accumulating partial JSON
"""
data_str = str_line[5:] # Remove "data:" prefix
if self.chunk_type == "accumulated_json":
# Already in accumulation mode, keep accumulating
return self._handle_accumulated_json_chunk(data_str)
# Try to parse as valid JSON first
try:
data_json = json.loads(data_str)
return self.chunk_parser(chunk=data_json)
except json.JSONDecodeError:
# Switch to accumulation mode and start accumulating
self.chunk_type = "accumulated_json"
return self._handle_accumulated_json_chunk(data_str)
# Sync iterator
def __iter__(self):
return self
def __next__(self):
try:
chunk = self.response_iterator.__next__()
except StopIteration:
raise StopIteration
except ValueError as e:
raise RuntimeError(f"Error receiving chunk from stream: {e}")
while True:
try:
chunk = self.response_iterator.__next__()
except StopIteration:
# If we have accumulated JSON when stream ends, try to parse it
if self.accumulated_json:
try:
data_json = json.loads(self.accumulated_json)
self.accumulated_json = ""
return self.chunk_parser(chunk=data_json)
except json.JSONDecodeError:
pass
raise StopIteration
except ValueError as e:
raise RuntimeError(f"Error receiving chunk from stream: {e}")
try:
str_line = chunk
if isinstance(chunk, bytes): # Handle binary data
str_line = chunk.decode("utf-8") # Convert bytes to string
index = str_line.find("data:")
if index != -1:
str_line = str_line[index:]
try:
str_line = chunk
if isinstance(chunk, bytes): # Handle binary data
str_line = chunk.decode("utf-8") # Convert bytes to string
index = str_line.find("data:")
if index != -1:
str_line = str_line[index:]
if str_line.startswith("data:"):
data_json = json.loads(str_line[5:])
return self.chunk_parser(chunk=data_json)
else:
return GenericStreamingChunk(
text="",
is_finished=False,
finish_reason="",
usage=None,
index=0,
tool_use=None,
)
except StopIteration:
raise StopIteration
except ValueError as e:
raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}")
if str_line.startswith("data:"):
result = self._parse_sse_data(str_line)
if result is not None:
return result
# If None, continue loop to get more chunks for accumulation
else:
return GenericStreamingChunk(
text="",
is_finished=False,
finish_reason="",
usage=None,
index=0,
tool_use=None,
)
except StopIteration:
raise StopIteration
except ValueError as e:
raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}")
# Async iterator
def __aiter__(self):
@ -909,37 +978,48 @@ class ModelResponseIterator:
return self
async def __anext__(self):
try:
chunk = await self.async_response_iterator.__anext__()
except StopAsyncIteration:
raise StopAsyncIteration
except ValueError as e:
raise RuntimeError(f"Error receiving chunk from stream: {e}")
while True:
try:
chunk = await self.async_response_iterator.__anext__()
except StopAsyncIteration:
# If we have accumulated JSON when stream ends, try to parse it
if self.accumulated_json:
try:
data_json = json.loads(self.accumulated_json)
self.accumulated_json = ""
return self.chunk_parser(chunk=data_json)
except json.JSONDecodeError:
pass
raise StopAsyncIteration
except ValueError as e:
raise RuntimeError(f"Error receiving chunk from stream: {e}")
try:
str_line = chunk
if isinstance(chunk, bytes): # Handle binary data
str_line = chunk.decode("utf-8") # Convert bytes to string
index = str_line.find("data:")
if index != -1:
str_line = str_line[index:]
try:
str_line = chunk
if isinstance(chunk, bytes): # Handle binary data
str_line = chunk.decode("utf-8") # Convert bytes to string
index = str_line.find("data:")
if index != -1:
str_line = str_line[index:]
if str_line.startswith("data:"):
data_json = json.loads(str_line[5:])
return self.chunk_parser(chunk=data_json)
else:
return GenericStreamingChunk(
text="",
is_finished=False,
finish_reason="",
usage=None,
index=0,
tool_use=None,
)
except StopAsyncIteration:
raise StopAsyncIteration
except ValueError as e:
raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}")
if str_line.startswith("data:"):
result = self._parse_sse_data(str_line)
if result is not None:
return result
# If None, continue loop to get more chunks for accumulation
else:
return GenericStreamingChunk(
text="",
is_finished=False,
finish_reason="",
usage=None,
index=0,
tool_use=None,
)
except StopAsyncIteration:
raise StopAsyncIteration
except ValueError as e:
raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}")
def convert_str_chunk_to_generic_chunk(self, chunk: str) -> ModelResponseStream:
"""

View file

@ -460,3 +460,75 @@ def test_streaming_chunks_have_stable_ids():
response_two = iterator.chunk_parser(chunk=second_chunk)
assert response_one.id == response_two.id == iterator.response_id
def test_partial_json_chunk_accumulation():
"""
Test that partial JSON chunks are accumulated correctly.
This tests the fix for https://github.com/BerriAI/litellm/issues/17473
where network fragmentation can cause SSE data to arrive in partial chunks.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
# Simulate a complete JSON chunk being split into two parts
partial_chunk_1 = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hel'
partial_chunk_2 = 'lo"}}'
# First partial chunk should return None (still accumulating)
result1 = iterator._parse_sse_data(f"data:{partial_chunk_1}")
assert result1 is None, "First partial chunk should return None while accumulating"
assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode"
assert iterator.accumulated_json == partial_chunk_1, "Should have accumulated first part"
# Second partial chunk should complete the JSON and return a parsed result
result2 = iterator._parse_sse_data(f"data:{partial_chunk_2}")
assert result2 is not None, "Second chunk should return parsed result"
assert iterator.accumulated_json == "", "Buffer should be cleared after successful parse"
assert result2.choices[0].delta.content == "Hello", f"Expected 'Hello', got '{result2.choices[0].delta.content}'"
def test_complete_json_chunk_no_accumulation():
"""
Test that complete JSON chunks are parsed immediately without accumulation.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
complete_chunk = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}'
result = iterator._parse_sse_data(f"data:{complete_chunk}")
assert result is not None, "Complete chunk should return parsed result immediately"
assert iterator.chunk_type == "valid_json", "Should remain in valid_json mode"
assert iterator.accumulated_json == "", "Buffer should remain empty"
assert result.choices[0].delta.content == "Hello", f"Expected 'Hello', got '{result.choices[0].delta.content}'"
def test_multiple_partial_chunks_accumulation():
"""
Test that multiple partial chunks can be accumulated across several iterations.
"""
iterator = ModelResponseIterator(
streaming_response=MagicMock(), sync_stream=True, json_mode=False
)
# Split a JSON chunk into three parts
part1 = '{"type":"content_block_del'
part2 = 'ta","index":0,"delta":{"type":"text_del'
part3 = 'ta","text":"Hello"}}'
result1 = iterator._parse_sse_data(f"data:{part1}")
assert result1 is None
assert iterator.accumulated_json == part1
result2 = iterator._parse_sse_data(f"data:{part2}")
assert result2 is None
assert iterator.accumulated_json == part1 + part2
result3 = iterator._parse_sse_data(f"data:{part3}")
assert result3 is not None
assert iterator.accumulated_json == ""
assert result3.choices[0].delta.content == "Hello"