fix(anthropic): emit replayable streaming thinking blocks (#31022)

litellm/llms/anthropic/chat/handler.py

@@ -629,6 +629,7 @@ class ModelResponseIterator:
         Optional[ChatCompletionToolCallChunk],
         List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]],
         Dict[str, Any],
+        Optional[str],
     ]:
         """
         Helper function to handle the content block delta
@@ -636,6 +637,7 @@ class ModelResponseIterator:
         text = ""
         tool_use: Optional[ChatCompletionToolCallChunk] = None
         provider_specific_fields = {}
+        reasoning_content: Optional[str] = None
         content_block = ContentBlockDelta(**chunk)  # type: ignore
         thinking_blocks: List[
             Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]
@@ -670,14 +672,24 @@ class ModelResponseIterator:
             thinking_content = content_block["delta"].get("thinking")
             if isinstance(thinking_content, str) and thinking_content:
                 self.reasoning_content_chunks.append(thinking_content)
-            thinking_blocks = [
-                ChatCompletionThinkingBlock(
-                    type="thinking",
-                    thinking=thinking_content or "",
-                    signature=str(content_block["delta"].get("signature") or ""),
-                )
-            ]
-            provider_specific_fields["thinking_blocks"] = thinking_blocks
+                reasoning_content = thinking_content
+
+            signature = content_block["delta"].get("signature")
+            if isinstance(signature, str) and signature:
+                thinking_blocks = [
+                    ChatCompletionThinkingBlock(
+                        type="thinking",
+                        thinking="".join(
+                            cast(str, block["delta"].get("thinking"))
+                            for block in self.content_blocks
+                            if isinstance(block["delta"].get("thinking"), str)
+                        ),
+                        signature=signature,
+                    )
+                ]
+                provider_specific_fields["thinking_blocks"] = thinking_blocks
+                if reasoning_content is None:
+                    reasoning_content = ""
         elif (
             "content" in content_block["delta"]
             and content_block["delta"].get("type") == "compaction_delta"
@@ -688,25 +700,13 @@ class ModelResponseIterator:
                 "content": content_block["delta"]["content"],
             }
 
-        return text, tool_use, thinking_blocks, provider_specific_fields
-
-    def _handle_reasoning_content(
-        self,
-        thinking_blocks: List[
-            Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]
-        ],
-    ) -> Optional[str]:
-        """
-        Handle the reasoning content
-        """
-        reasoning_content = None
-        for block in thinking_blocks:
-            thinking_content = cast(Optional[str], block.get("thinking"))
-            if reasoning_content is None:
-                reasoning_content = ""
-            if thinking_content is not None:
-                reasoning_content += thinking_content
-        return reasoning_content
+        return (
+            text,
+            tool_use,
+            thinking_blocks,
+            provider_specific_fields,
+            reasoning_content,
+        )
 
     def _handle_redacted_thinking_content(
         self,
@@ -802,11 +802,8 @@ class ModelResponseIterator:
                     tool_use,
                     thinking_blocks,
                     provider_specific_fields,
+                    reasoning_content,
                 ) = self._content_block_delta_helper(chunk=chunk)
-                if thinking_blocks:
-                    reasoning_content = self._handle_reasoning_content(
-                        thinking_blocks=thinking_blocks
-                    )
             elif type_chunk == "content_block_start":
                 """
                 event: content_block_start

tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py

@@ -74,6 +74,137 @@ def test_redacted_thinking_content_block_delta():
     assert "thinking_blocks" in model_response.choices[0].delta.provider_specific_fields
 
 
+def test_streaming_thinking_blocks_are_replayable_after_signature_delta():
+    model_response_iterator = ModelResponseIterator(
+        streaming_response=MagicMock(), sync_stream=True, json_mode=False
+    )
+    chunks = [
+        {
+            "type": "content_block_start",
+            "index": 0,
+            "content_block": {"type": "thinking", "thinking": ""},
+        },
+        {
+            "type": "content_block_delta",
+            "index": 0,
+            "delta": {"type": "thinking_delta", "thinking": "Step 1. "},
+        },
+        {
+            "type": "content_block_delta",
+            "index": 0,
+            "delta": {"type": "thinking_delta", "thinking": "Step 2."},
+        },
+        {
+            "type": "content_block_delta",
+            "index": 0,
+            "delta": {"type": "signature_delta", "signature": "sig-final"},
+        },
+    ]
+
+    parsed_chunks = [
+        model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks
+    ]
+    reasoning_content = "".join(
+        getattr(chunk.choices[0].delta, "reasoning_content", None) or ""
+        for chunk in parsed_chunks
+    )
+    thinking_blocks = tuple(
+        block
+        for chunk in parsed_chunks
+        for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
+    )
+    expected_thinking_block = {
+        "type": "thinking",
+        "thinking": "Step 1. Step 2.",
+        "signature": "sig-final",
+    }
+
+    assert reasoning_content == "Step 1. Step 2."
+    assert thinking_blocks == (expected_thinking_block,)
+    assert parsed_chunks[-1].choices[0].delta.provider_specific_fields == {
+        "thinking_blocks": [expected_thinking_block]
+    }
+
+
+def test_streaming_unsigned_thinking_deltas_keep_reasoning_content():
+    model_response_iterator = ModelResponseIterator(
+        streaming_response=MagicMock(), sync_stream=True, json_mode=False
+    )
+    chunks = [
+        {
+            "type": "content_block_start",
+            "index": 0,
+            "content_block": {"type": "thinking", "thinking": ""},
+        },
+        {
+            "type": "content_block_delta",
+            "index": 0,
+            "delta": {"type": "thinking_delta", "thinking": "Step 1. "},
+        },
+        {
+            "type": "content_block_delta",
+            "index": 0,
+            "delta": {"type": "thinking_delta", "thinking": "Step 2."},
+        },
+        {"type": "content_block_stop", "index": 0},
+    ]
+
+    parsed_chunks = [
+        model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks
+    ]
+    reasoning_content = "".join(
+        getattr(chunk.choices[0].delta, "reasoning_content", None) or ""
+        for chunk in parsed_chunks
+    )
+    thinking_blocks = tuple(
+        block
+        for chunk in parsed_chunks
+        for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
+    )
+
+    assert reasoning_content == "Step 1. Step 2."
+    assert thinking_blocks == ()
+
+
+def test_streaming_truncated_thinking_deltas_keep_reasoning_content():
+    model_response_iterator = ModelResponseIterator(
+        streaming_response=MagicMock(), sync_stream=True, json_mode=False
+    )
+    chunks = [
+        {
+            "type": "content_block_start",
+            "index": 0,
+            "content_block": {"type": "thinking", "thinking": ""},
+        },
+        {
+            "type": "content_block_delta",
+            "index": 0,
+            "delta": {"type": "thinking_delta", "thinking": "Step 1. "},
+        },
+        {
+            "type": "content_block_delta",
+            "index": 0,
+            "delta": {"type": "thinking_delta", "thinking": "Step 2."},
+        },
+    ]
+
+    parsed_chunks = [
+        model_response_iterator.chunk_parser(chunk=chunk) for chunk in chunks
+    ]
+    reasoning_content = "".join(
+        getattr(chunk.choices[0].delta, "reasoning_content", None) or ""
+        for chunk in parsed_chunks
+    )
+    thinking_blocks = tuple(
+        block
+        for chunk in parsed_chunks
+        for block in (getattr(chunk.choices[0].delta, "thinking_blocks", None) or [])
+    )
+
+    assert reasoning_content == "Step 1. Step 2."
+    assert thinking_blocks == ()
+
+
 def test_handle_json_mode_chunk_response_format_tool():
     model_response_iterator = ModelResponseIterator(
         streaming_response=MagicMock(), sync_stream=True, json_mode=True