fix(anthropic): split mixed reasoning stream chunks

This commit is contained in:
Napuh 2026-07-26 12:30:02 +02:00
parent 24123269cc
commit 2a8e29cc0b
2 changed files with 85 additions and 16 deletions

View file

@ -96,6 +96,38 @@ class _CombinedChunkSplitter:
or getattr(delta, "thinking_blocks", None)
)
@staticmethod
def _has_mixed_reasoning_and_text(chunk: Any) -> bool:
choices = getattr(chunk, "choices", None)
if not choices:
return False
delta = getattr(choices[0], "delta", None)
if delta is None:
return False
return bool(getattr(delta, "content", None) and getattr(delta, "reasoning_content", None))
@staticmethod
def _clear_usage(chunk: Any) -> None:
if hasattr(chunk, "usage"):
chunk.usage = None
hidden_params = getattr(chunk, "_hidden_params", None)
if isinstance(hidden_params, dict) and "usage" in hidden_params:
chunk._hidden_params = {key: value for key, value in hidden_params.items() if key != "usage"}
@staticmethod
def _split_mixed_reasoning_and_text(chunk: Any) -> List[Any]:
if not _CombinedChunkSplitter._has_mixed_reasoning_and_text(chunk):
return [chunk]
reasoning_chunk = copy.deepcopy(chunk)
reasoning_chunk.choices[0].finish_reason = None
reasoning_chunk.choices[0].delta.content = None
_CombinedChunkSplitter._clear_usage(reasoning_chunk)
text_chunk = copy.deepcopy(chunk)
text_chunk.choices[0].delta.reasoning_content = None
return [reasoning_chunk, text_chunk]
@staticmethod
def _split(chunk: Any) -> List[Any]:
"""Return ``[chunk]``, or ``[content_chunk, finish_chunk]`` if combined."""
@ -105,6 +137,7 @@ class _CombinedChunkSplitter:
# Content chunk: keep the delta payload, clear the finish_reason.
content_chunk = copy.deepcopy(chunk)
content_chunk.choices[0].finish_reason = None
_CombinedChunkSplitter._clear_usage(content_chunk)
# Finish chunk: keep finish_reason (and usage), clear the delta payload.
finish_chunk = copy.deepcopy(chunk)
@ -127,7 +160,11 @@ class _CombinedChunkSplitter:
if self._sync_iter is None:
self._sync_iter = iter(self._stream)
chunk = next(self._sync_iter) # propagates StopIteration when exhausted
self._buffer.extend(self._split(chunk))
self._buffer.extend(
split_chunk
for combined_chunk in self._split(chunk)
for split_chunk in self._split_mixed_reasoning_and_text(combined_chunk)
)
return self._buffer.popleft()
def __aiter__(self) -> "AsyncIterator[Any]":
@ -139,7 +176,11 @@ class _CombinedChunkSplitter:
if self._async_iter is None:
self._async_iter = self._stream.__aiter__()
chunk = await self._async_iter.__anext__() # propagates StopAsyncIteration
self._buffer.extend(self._split(chunk))
self._buffer.extend(
split_chunk
for combined_chunk in self._split(chunk)
for split_chunk in self._split_mixed_reasoning_and_text(combined_chunk)
)
return self._buffer.popleft()

View file

@ -65,9 +65,7 @@ def _thinking_chunk(thinking: str, signature: str = "") -> MagicMock:
return _make_chunk(Delta(content=None, thinking_blocks=[block]))
def _tool_chunk(
call_id: str, name: Optional[str], arguments: Optional[str]
) -> MagicMock:
def _tool_chunk(call_id: str, name: Optional[str], arguments: Optional[str]) -> MagicMock:
return _make_chunk(
Delta(
content=None,
@ -109,8 +107,7 @@ def _text_deltas(events: List[dict]) -> List[str]:
return [
e["delta"]["text"]
for e in events
if e.get("type") == "content_block_delta"
and e["delta"].get("type") == "text_delta"
if e.get("type") == "content_block_delta" and e["delta"].get("type") == "text_delta"
]
@ -118,8 +115,7 @@ def _input_json_deltas(events: List[dict]) -> List[str]:
return [
e["delta"]["partial_json"]
for e in events
if e.get("type") == "content_block_delta"
and e["delta"].get("type") == "input_json_delta"
if e.get("type") == "content_block_delta" and e["delta"].get("type") == "input_json_delta"
]
@ -127,8 +123,7 @@ def _thinking_deltas(events: List[dict]) -> List[str]:
return [
e["delta"]["thinking"]
for e in events
if e.get("type") == "content_block_delta"
and e["delta"].get("type") == "thinking_delta"
if e.get("type") == "content_block_delta" and e["delta"].get("type") == "thinking_delta"
]
@ -136,8 +131,7 @@ def _signature_deltas(events: List[dict]) -> List[str]:
return [
e["delta"]["signature"]
for e in events
if e.get("type") == "content_block_delta"
and e["delta"].get("type") == "signature_delta"
if e.get("type") == "content_block_delta" and e["delta"].get("type") == "signature_delta"
]
@ -228,9 +222,7 @@ async def test_first_text_delta_after_tool_use_is_not_dropped_async():
_make_chunk(Delta(content=" Bye.")),
_make_chunk(Delta(content=None), finish_reason="stop"),
]
wrapper = AnthropicStreamWrapper(
completion_stream=_AsyncStream(chunks), model="claude-x"
)
wrapper = AnthropicStreamWrapper(completion_stream=_AsyncStream(chunks), model="claude-x")
events = await _drain_async(wrapper)
assert _input_json_deltas(events) == ['{"city": "NY"}']
@ -506,3 +498,39 @@ def test_empty_content_chunk_mid_text_block_is_suppressed_sync():
assert _text_deltas(events) == ["Hi", " there"]
_assert_deltas_match_their_block_type(events)
def _mixed_reasoning_and_text_chunks() -> List[MagicMock]:
return [
_make_chunk(Delta(content=None, reasoning_content="First thought.")),
_make_chunk(
Delta(content="Answer.", reasoning_content=" Last thought."),
finish_reason="stop",
),
]
def _assert_mixed_reasoning_and_text_chunk_is_split(events: List[dict]) -> None:
_assert_deltas_match_their_block_type(events)
assert _thinking_deltas(events) == ["First thought.", " Last thought."]
assert _text_deltas(events) == ["Answer."]
assert [event["type"] for event in events].count("message_delta") == 1
def test_mixed_reasoning_and_text_chunk_is_split_sync():
wrapper = AnthropicStreamWrapper(
completion_stream=iter(_mixed_reasoning_and_text_chunks()),
model="claude-x",
)
_assert_mixed_reasoning_and_text_chunk_is_split(_drain_sync(wrapper))
@pytest.mark.asyncio
async def test_mixed_reasoning_and_text_chunk_is_split_async():
wrapper = AnthropicStreamWrapper(
completion_stream=_AsyncStream(_mixed_reasoning_and_text_chunks()),
model="claude-x",
)
_assert_mixed_reasoning_and_text_chunk_is_split(await _drain_async(wrapper))