From 236df8e316ce6b2e7a932166aac39fbedf987a41 Mon Sep 17 00:00:00 2001 From: jinwukong <134120444+jinwukong@users.noreply.github.com> Date: Fri, 11 Sep 2026 12:16:58 +0800 Subject: [PATCH] refactor(streaming): assemble thinking blocks without mutable state --- .../streaming_chunk_builder_utils.py | 112 ++++++++---------- .../test_streaming_chunk_builder_utils.py | 29 ++++- 2 files changed, 77 insertions(+), 64 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index d6823c7fa14..a98d4dfa063 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1,9 +1,9 @@ import base64 import time from collections.abc import Callable, Iterator, Mapping, Sequence -from itertools import groupby +from itertools import accumulate, chain, groupby, tee from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypedDict, Union, cast +from typing import TYPE_CHECKING, Any, Final, NamedTuple, TypeAlias, TypedDict, Union, cast from typing_extensions import ReadOnly, Required @@ -69,6 +69,11 @@ class _ThinkingChunk(TypedDict): choices: Sequence[_ThinkingChoice] +class _ThinkingStreamFragment(NamedTuple): + block: _ThinkingBlockFragment + is_snapshot: bool + + class _ContentChoice(TypedDict, total=False): delta: Mapping[str, str | None] @@ -662,69 +667,56 @@ class ChunkProcessor: def get_combined_thinking_content( self, chunks: Sequence["_ThinkingChunk"] ) -> list[Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]] | None: + fragments, boundary_fragments = tee(self._iter_thinking_fragments(chunks)) + # Count completed blocks before each fragment, keeping signatures with their preceding text. + closed_blocks: Final = accumulate( + ( + int(fragment.block.get("type") == "redacted_thinking" or bool(fragment.block.get("signature"))) + for fragment in boundary_fragments + ), + initial=0, + ) + grouped: Final = groupby(zip(closed_blocks, fragments, strict=False), key=lambda entry: entry[0]) + groups: Final = (tuple(fragment for _, fragment in group) for _, group in grouped) + blocks: Final = tuple(block for group in groups if (block := self._assemble_thinking_block(group)) is not None) + return list(blocks) if blocks else None # mutable-ok: Message.thinking_blocks requires a list + + @staticmethod + def _iter_thinking_fragments(chunks: Sequence["_ThinkingChunk"]) -> Iterator[_ThinkingStreamFragment]: + for choice in chain.from_iterable(chunk["choices"] for chunk in chunks): + if (delta := choice.get("delta")) is None or not isinstance(blocks := delta.get("thinking_blocks"), list): + continue + for block in blocks: + yield _ThinkingStreamFragment( + block, + isinstance(provider_fields := delta.get("provider_specific_fields"), Mapping) + and provider_fields.get("thinking_blocks") == blocks, + ) + + @staticmethod + def _assemble_thinking_block( + fragments: Sequence[_ThinkingStreamFragment], + ) -> Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock", None]: from litellm.types.llms.openai import ( ChatCompletionRedactedThinkingBlock, ChatCompletionThinkingBlock, ) - thinking_blocks: Final[list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock]] = [] - current_thinking_text_parts: list[str] = [] - current_signature: str | None = None - - def _flush_thinking_block() -> None: - nonlocal current_thinking_text_parts, current_signature - if current_signature: - thinking_blocks.append( - ChatCompletionThinkingBlock( - type="thinking", - thinking="".join(current_thinking_text_parts), - signature=current_signature, - ) - ) - current_thinking_text_parts = [] - current_signature = None - - for chunk in chunks: - choices = chunk["choices"] - for choice in choices: - delta = choice.get("delta", {}) - thinking = delta.get("thinking_blocks", None) - if thinking and isinstance(thinking, list): - for thinking_block in thinking: - thinking_type = thinking_block.get("type", None) - if thinking_type and thinking_type == "redacted_thinking": - _flush_thinking_block() - redacted_data = thinking_block.get("data", None) - if redacted_data: - thinking_blocks.append( - ChatCompletionRedactedThinkingBlock( - type="redacted_thinking", - data=redacted_data, - ) - ) - else: - thinking_text, signature, provider_fields = ( - thinking_block.get("thinking"), - thinking_block.get("signature"), - delta.get("provider_specific_fields"), - ) - if ( - signature - and isinstance(provider_fields, Mapping) - and provider_fields.get("thinking_blocks") == thinking - ): - current_thinking_text_parts.clear() - if thinking_text: - current_thinking_text_parts.append(thinking_text) - if signature: - current_signature = signature - _flush_thinking_block() - - _flush_thinking_block() - - if len(thinking_blocks) > 0: - return thinking_blocks - return None + last: Final = fragments[-1] + if last.block.get("type") == "redacted_thinking": + return ( + ChatCompletionRedactedThinkingBlock(type="redacted_thinking", data=data) + if (data := last.block.get("data")) + else None + ) + if not (signature := last.block.get("signature")): + return None + text: Final = ( + last.block.get("thinking") or "" + if last.is_snapshot + else "".join(fragment.block.get("thinking") or "" for fragment in fragments) + ) + return ChatCompletionThinkingBlock(type="thinking", thinking=text, signature=signature) def get_combined_reasoning_content(self, chunks: Sequence["_ContentChunk"]) -> ChatCompletionAssistantContentValue: return self.get_combined_content(chunks, delta_key="reasoning_content") diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 47f06b24016..3c42221d682 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -218,11 +218,8 @@ def test_get_combined_thinking_content_preserves_interleaved_blocks(): ), ] - thinking_chunks = [ - chunk for chunk in chunks if chunk["choices"][0]["delta"].get("thinking_blocks") - ] processor = ChunkProcessor(chunks=chunks) - result = processor.get_combined_thinking_content(thinking_chunks) + result = processor.get_combined_thinking_content(chunks) assert result is not None assert len(result) == 3 @@ -260,6 +257,30 @@ def test_stream_chunk_builder_distinguishes_thinking_snapshots_from_repeated_del ] +def test_incomplete_thinking_stream_preserves_summary_without_signed_blocks() -> None: + chunk: Final = ModelResponseStream( + id="chatcmpl-incomplete-thinking", + model="claude-opus-5", + choices=[ + StreamingChoices( + index=0, + finish_reason="length", + delta=Delta( + reasoning_content="Unfinished reasoning", + thinking_blocks=[ChatCompletionThinkingBlock(type="thinking", thinking="Unfinished reasoning")], + ), + ) + ], + ) + + response: Final = stream_chunk_builder(chunks=[chunk]) + + assert response is not None + assert response.choices[0].finish_reason == "length" + assert response.choices[0].message.reasoning_content == "Unfinished reasoning" + assert response.choices[0].message.thinking_blocks is None + + def test_cache_read_input_tokens_retained(): chunk1 = ModelResponseStream( id="chatcmpl-95aabb85-c39f-443d-ae96-0370c404d70c",