diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 53252df0a28..76c7246b87e 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1,6 +1,6 @@ import base64 import time -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast from litellm.types.llms.openai import ( ChatCompletionAssistantContentValue, @@ -326,10 +326,22 @@ class ChunkProcessor: thinking_blocks: List[ Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"] ] = [] - combined_thinking_text: Optional[str] = None - data: Optional[str] = None - signature: Optional[str] = None - type: Literal["thinking", "redacted_thinking"] = "thinking" + current_thinking_text_parts: List[str] = [] + current_signature: Optional[str] = None + + def _flush_thinking_block() -> None: + nonlocal current_thinking_text_parts, current_signature + if len(current_thinking_text_parts) > 0 and 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: @@ -339,33 +351,25 @@ class ChunkProcessor: for thinking_block in thinking: thinking_type = thinking_block.get("type", None) if thinking_type and thinking_type == "redacted_thinking": - type = "redacted_thinking" - data = thinking_block.get("data", None) + _flush_thinking_block() + redacted_data = thinking_block.get("data", None) + if redacted_data: + thinking_blocks.append( + ChatCompletionRedactedThinkingBlock( + type="redacted_thinking", + data=redacted_data, + ) + ) else: - type = "thinking" thinking_text = thinking_block.get("thinking", None) if thinking_text: - if combined_thinking_text is None: - combined_thinking_text = "" - - combined_thinking_text += thinking_text + current_thinking_text_parts.append(thinking_text) signature = thinking_block.get("signature", None) + if signature: + current_signature = signature + _flush_thinking_block() - if combined_thinking_text and type == "thinking" and signature: - thinking_blocks.append( - ChatCompletionThinkingBlock( - type=type, - thinking=combined_thinking_text, - signature=signature, - ) - ) - elif data and type == "redacted_thinking": - thinking_blocks.append( - ChatCompletionRedactedThinkingBlock( - type=type, - data=data, - ) - ) + _flush_thinking_block() if len(thinking_blocks) > 0: return thinking_blocks 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 eef206ca667..34c9efd8382 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 @@ -158,6 +158,78 @@ def test_get_combined_tool_content(): ] +def test_get_combined_thinking_content_preserves_interleaved_blocks(): + base_chunk = { + "id": "chatcmpl-123", + "object": "chat.completion.chunk", + "created": 1234567890, + "model": "claude-sonnet-4-20250514", + } + + def make_chunk(**delta_kwargs): + return ModelResponseStream( + **{ + **base_chunk, + "choices": [ + { + "index": 0, + "delta": delta_kwargs, + "finish_reason": None, + } + ], + } + ) + + chunks = [ + make_chunk(role="assistant", content=None), + make_chunk( + thinking_blocks=[ + {"type": "thinking", "thinking": "Step 1 analysis...", "signature": None} + ] + ), + make_chunk( + thinking_blocks=[ + {"type": "thinking", "thinking": None, "signature": "sig_block1"} + ] + ), + make_chunk( + thinking_blocks=[ + { + "type": "redacted_thinking", + "data": "EuoBCoYBGAIi...encrypted...", + } + ] + ), + make_chunk( + thinking_blocks=[ + {"type": "thinking", "thinking": "Step 2 analysis...", "signature": None} + ] + ), + make_chunk( + thinking_blocks=[ + {"type": "thinking", "thinking": None, "signature": "sig_block2"} + ] + ), + ] + + 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) + + assert result is not None + assert len(result) == 3 + assert result[0]["type"] == "thinking" + assert result[0]["thinking"] == "Step 1 analysis..." + assert result[0]["signature"] == "sig_block1" + assert result[1]["type"] == "redacted_thinking" + assert result[1]["data"] == "EuoBCoYBGAIi...encrypted..." + assert result[2]["type"] == "thinking" + assert result[2]["thinking"] == "Step 2 analysis..." + assert result[2]["signature"] == "sig_block2" + + def test_cache_read_input_tokens_retained(): chunk1 = ModelResponseStream( id="chatcmpl-95aabb85-c39f-443d-ae96-0370c404d70c", @@ -441,4 +513,4 @@ def test_stream_chunk_builder_anthropic_web_search(): assert usage.prompt_tokens == 50 assert usage.completion_tokens == 27 assert usage.total_tokens == 77 - assert usage.server_tool_use['web_search_requests'] == 2 \ No newline at end of file + assert usage.server_tool_use['web_search_requests'] == 2