From a5835bb859ca413cd4d765d8e865a8b44cc3e6c7 Mon Sep 17 00:00:00 2001 From: jinwukong <134120444+jinwukong@users.noreply.github.com> Date: Thu, 10 Sep 2026 18:25:01 +0800 Subject: [PATCH 1/3] fix(responses): preserve signed thinking when replaying streams --- .../streaming_chunk_builder_utils.py | 16 ++- .../streaming_iterator.py | 11 ++ .../transformation.py | 6 +- .../test_streaming_chunk_builder_utils.py | 25 ++++ .../test_reasoning_input_item_preservation.py | 2 +- .../test_streaming_iterator_transformation.py | 118 +++++++++++++++++- 6 files changed, 170 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 90698296142..d6823c7fa14 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -58,6 +58,7 @@ class _ThinkingBlockFragment(TypedDict, total=False): class _ThinkingDelta(TypedDict, total=False): thinking_blocks: Sequence[_ThinkingBlockFragment] + provider_specific_fields: ReadOnly[Mapping[str, object] | None] class _ThinkingChoice(TypedDict, total=False): @@ -672,7 +673,7 @@ class ChunkProcessor: def _flush_thinking_block() -> None: nonlocal current_thinking_text_parts, current_signature - if len(current_thinking_text_parts) > 0 and current_signature: + if current_signature: thinking_blocks.append( ChatCompletionThinkingBlock( type="thinking", @@ -702,10 +703,19 @@ class ChunkProcessor: ) ) else: - thinking_text = thinking_block.get("thinking", None) + 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) - signature = thinking_block.get("signature", None) if signature: current_signature = signature _flush_thinking_block() diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 660dd8f0c92..6a51ac027f0 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -547,6 +547,16 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) return response + def _encoded_thinking_blocks(self) -> str | None: + response: Final = ( + self.litellm_model_response + if isinstance(self.litellm_model_response, ModelResponse) + else self.create_litellm_model_response() + ) + if response is None: + return None + return LiteLLMCompletionResponsesConfig.encode_thinking_blocks(response.choices[0].message) + @staticmethod def _snapshot_chunk_for_stream_chunk_builder( chunk: ModelResponseStream, @@ -746,6 +756,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): **{ "id": reasoning_item_id, "type": "reasoning", + "encrypted_content": self._encoded_thinking_blocks(), "summary": [ { "type": "summary_text", diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index fca5b0d11cf..19a7ab35a2f 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1384,7 +1384,7 @@ class LiteLLMCompletionResponsesConfig: input_item: Mapping[str, object], ) -> tuple[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, ...] | None: """ - Decode ``encrypted_content`` written by ``_encode_thinking_blocks`` back + Decode ``encrypted_content`` written by ``encode_thinking_blocks`` back into the signed thinking blocks it serialized. LiteLLM writes this field itself for providers whose reasoning is signed @@ -2403,7 +2403,7 @@ class LiteLLMCompletionResponsesConfig: return output_items @staticmethod - def _encode_thinking_blocks(message: Message) -> str | None: + def encode_thinking_blocks(message: Message) -> str | None: thinking_blocks: Final[Sequence[Mapping[str, object]]] = getattr(message, "thinking_blocks", None) or () preserved: Final = tuple(block for block in thinking_blocks if block.get("signature") or block.get("data")) return json.dumps(preserved, separators=(",", ":")) if preserved else None @@ -2417,7 +2417,7 @@ class LiteLLMCompletionResponsesConfig: if hasattr(choice, "message") and choice.message: message = choice.message reasoning_content: str = getattr(message, "reasoning_content", None) or "" - encrypted_content = LiteLLMCompletionResponsesConfig._encode_thinking_blocks(message) + encrypted_content = LiteLLMCompletionResponsesConfig.encode_thinking_blocks(message) if reasoning_content or encrypted_content: # Only check the first choice for reasoning content return [ 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 efe4209c1c9..47f06b24016 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 @@ -8,6 +8,7 @@ import pytest from litellm import ChatCompletionUsageBlock, stream_chunk_builder from litellm.types.utils import GenericStreamingChunk from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor +from litellm.types.llms.openai import ChatCompletionThinkingBlock from litellm.types.utils import ( ChatCompletionDeltaToolCall, ChatCompletionMessageToolCall, @@ -235,6 +236,30 @@ def test_get_combined_thinking_content_preserves_interleaved_blocks(): assert result[2]["signature"] == "sig_block2" +@pytest.mark.parametrize("snapshot", [True, False], ids=["provider-snapshot", "genuine-final-delta"]) +def test_stream_chunk_builder_distinguishes_thinking_snapshots_from_repeated_deltas(snapshot: bool) -> None: + signed: Final = ChatCompletionThinkingBlock(type="thinking", thinking="echo", signature="test-signature") + deltas: Final = ( + Delta(thinking_blocks=[ChatCompletionThinkingBlock(type="thinking", thinking="echo")]), + Delta(thinking_blocks=[signed], provider_specific_fields={"thinking_blocks": [signed]} if snapshot else None), + ) + chunks: Final = [ + ModelResponseStream( + id="chatcmpl-thinking", + model="claude-opus-5", + choices=[StreamingChoices(index=0, delta=delta, finish_reason="stop" if index == 1 else None)], + ) + for index, delta in enumerate(deltas) + ] + + response: Final = stream_chunk_builder(chunks=chunks) + + assert response is not None + assert response.choices[0].message.thinking_blocks == [ + {"type": "thinking", "thinking": "echo" if snapshot else "echoecho", "signature": "test-signature"} + ] + + def test_cache_read_input_tokens_retained(): chunk1 = ModelResponseStream( id="chatcmpl-95aabb85-c39f-443d-ae96-0370c404d70c", diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py index 5e001bdbbbb..db7a6238c4f 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py @@ -167,7 +167,7 @@ class TestEncryptedReasoningRoundTrip: {"type": "redacted_thinking", "data": "redacted-payload"}, ] message = Message(role="assistant", content="answer", thinking_blocks=blocks) - encoded = LiteLLMCompletionResponsesConfig._encode_thinking_blocks(message) + encoded = LiteLLMCompletionResponsesConfig.encode_thinking_blocks(message) decoded = LiteLLMCompletionResponsesConfig._decode_thinking_blocks_from_input_item( {"type": "reasoning", "encrypted_content": encoded} ) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 850ee7ba623..61dceb27f6b 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -19,8 +19,16 @@ import pytest from litellm.responses.litellm_completion_transformation.streaming_iterator import ( LiteLLMCompletionStreamingIterator, ) +from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, +) from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.types.llms.openai import ResponsesAPIStreamEvents +from litellm.types.llms.openai import ( + ChatCompletionRedactedThinkingBlock, + ChatCompletionThinkingBlock, + ResponseCompletedEvent, + ResponsesAPIStreamEvents, +) from litellm.types.utils import ( Delta, ModelResponse, @@ -705,3 +713,111 @@ def test_streamed_unrecognized_tool_choice_is_echoed_as_auto() -> None: ] assert [event.response.tool_choice for event in response_events] == ["auto", "auto", "auto"] + + +def _signed_thinking_chunks(cumulative: bool) -> tuple[ModelResponseStream, ...]: + blocks: Final = [ + ChatCompletionThinkingBlock(type="thinking", thinking="One plus one equals two.", signature="test-signature"), + ChatCompletionRedactedThinkingBlock(type="redacted_thinking", data="test-redacted-data"), + ] + deltas: Final = ( + Delta( + reasoning_content="One plus ", + thinking_blocks=[ChatCompletionThinkingBlock(type="thinking", thinking="One plus ")], + ), + Delta( + reasoning_content="one equals two.", + thinking_blocks=[ChatCompletionThinkingBlock(type="thinking", thinking="one equals two.")], + ), + Delta( + thinking_blocks=blocks + if cumulative + else [ + ChatCompletionThinkingBlock(type="thinking", thinking="", signature="test-signature"), + ChatCompletionRedactedThinkingBlock(type="redacted_thinking", data="test-redacted-data"), + ], + provider_specific_fields={"thinking_blocks": blocks} if cumulative else None, + ), + Delta(content="2"), + ) + return tuple( + ModelResponseStream( + id=CHAT_COMPLETION_ID, + model="test-model", + choices=[StreamingChoices(index=0, delta=delta, finish_reason="stop" if index == 3 else None)], + ) + for index, delta in enumerate(deltas) + ) + + +@pytest.mark.parametrize("cumulative", [True, False], ids=["cumulative-provider-blocks", "delta-blocks"]) +@pytest.mark.parametrize("asynchronous", [True, False], ids=["async", "sync"]) +async def test_completed_response_replays_signed_thinking_unchanged(cumulative: bool, asynchronous: bool) -> None: + iterator: Final = _build_iterator(_signed_thinking_chunks(cumulative)) + events: Final = [event async for event in iterator] if asynchronous else list(iterator) + completed: Final = next(event for event in events if isinstance(event, ResponseCompletedEvent)) + reasoning: Final = next(item for item in completed.response.output if item.type == "reasoning") + + assert reasoning.encrypted_content is not None + assert json.loads(reasoning.encrypted_content) == [ + {"type": "thinking", "thinking": "One plus one equals two.", "signature": "test-signature"}, + {"type": "redacted_thinking", "data": "test-redacted-data"}, + ] + messages: Final = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=reasoning.model_dump(exclude_none=True), replay_reasoning=True + ) + assert len(messages) == 1 + assert messages[0]["thinking_blocks"] == [ + {"type": "thinking", "thinking": "One plus one equals two.", "signature": "test-signature"}, + {"type": "redacted_thinking", "data": "test-redacted-data"}, + ] + + +@pytest.mark.parametrize("cumulative", [True, False], ids=["cumulative-provider-blocks", "delta-blocks"]) +async def test_reasoning_done_preserves_the_replay_payload(cumulative: bool) -> None: + iterator: Final = _build_iterator(_signed_thinking_chunks(cumulative)) + events: Final = [event async for event in iterator] + done: Final = next( + event for event in events if event.type == "response.output_item.done" and event.item.type == "reasoning" + ) + completed: Final = next(event for event in events if isinstance(event, ResponseCompletedEvent)) + reasoning: Final = next(item for item in completed.response.output if item.type == "reasoning") + payload: Final = done.item.model_dump().get("encrypted_content") + + assert payload is not None + assert json.loads(payload) == [ + {"type": "thinking", "thinking": "One plus one equals two.", "signature": "test-signature"}, + {"type": "redacted_thinking", "data": "test-redacted-data"}, + ] + assert payload == reasoning.encrypted_content + + +@pytest.mark.parametrize("asynchronous", [True, False], ids=["async", "sync"]) +async def test_streamed_signature_only_thinking_is_replayable(asynchronous: bool) -> None: + block: Final = ChatCompletionThinkingBlock(type="thinking", thinking="", signature="opaque-signature") + chunks: Final = [ + ModelResponseStream( + id=CHAT_COMPLETION_ID, + model="test-model", + choices=[ + StreamingChoices( + index=0, + delta=Delta(thinking_blocks=[block], provider_specific_fields={"thinking_blocks": [block]}), + ) + ], + ), + _tool_call_chunk(finish_reason="tool_calls"), + ] + iterator: Final = _build_iterator(chunks) + events: Final = [event async for event in iterator] if asynchronous else list(iterator) + completed: Final = next(event for event in events if isinstance(event, ResponseCompletedEvent)) + reasoning: Final = next(item for item in completed.response.output if item.type == "reasoning") + assert json.loads(reasoning.encrypted_content) == [block] + + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=[item.model_dump(exclude_none=True) for item in completed.response.output], + responses_api_request={}, + replay_reasoning=True, + ) + tool_message: Final = next(message for message in messages if message.get("tool_calls")) + assert tool_message["thinking_blocks"] == [block] From 6d7c160fec2098cb662a8cc5052ba72c3f155112 Mon Sep 17 00:00:00 2001 From: jinwukong <134120444+jinwukong@users.noreply.github.com> Date: Fri, 11 Sep 2026 11:48:44 +0800 Subject: [PATCH 2/3] test(responses): cover reasoning events before response assembly --- .../test_streaming_iterator_transformation.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 61dceb27f6b..25ad09ece63 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -821,3 +821,21 @@ async def test_streamed_signature_only_thinking_is_replayable(asynchronous: bool ) tool_message: Final = next(message for message in messages if message.get("tool_calls")) assert tool_message["thinking_blocks"] == [block] + + +def test_reasoning_done_without_a_response_snapshot_preserves_summary() -> None: + iterator: Final = _build_iterator([]) + + event: Final = iterator.create_reasoning_output_item_done_event( + reasoning_item_id="rs_pending", + reasoning_content="The response snapshot is not available yet.", + sequence_number=7, + ) + + assert event.type == "response.output_item.done" + assert event.sequence_number == 7 + assert event.item.model_dump(exclude_none=True) == { + "id": "rs_pending", + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "The response snapshot is not available yet."}], + } 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 3/3] 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",