fix: update ThinkingPart when delta contains signature (#2012)

Author: Kludex

pydantic_ai_slim/pydantic_ai/messages.py

@@ -763,7 +763,9 @@ def apply(self, part: ModelResponsePart | ThinkingPartDelta) -> ThinkingPart | T
             ValueError: If `part` is not a `ThinkingPart`.
         """
         if isinstance(part, ThinkingPart):
-            return replace(part, content=part.content + self.content_delta if self.content_delta else None)
+            new_content = part.content + self.content_delta if self.content_delta else part.content
+            new_signature = self.signature_delta if self.signature_delta is not None else part.signature
+            return replace(part, content=new_content, signature=new_signature)
         elif isinstance(part, ThinkingPartDelta):
             if self.content_delta is None and self.signature_delta is None:
                 raise ValueError('Cannot apply ThinkingPartDelta with no content or signature')

tests/test_thinking_part.py

@@ -1,9 +1,10 @@
 from __future__ import annotations as _annotations
 
 import pytest
+from inline_snapshot import snapshot
 
 from pydantic_ai._thinking_part import split_content_into_text_and_thinking
-from pydantic_ai.messages import ModelResponsePart, TextPart, ThinkingPart
+from pydantic_ai.messages import ModelResponsePart, TextPart, ThinkingPart, ThinkingPartDelta
 
 
 @pytest.mark.parametrize(
@@ -26,3 +27,43 @@
 )
 def test_split_content_into_text_and_thinking(content: str, parts: list[ModelResponsePart]):
     assert split_content_into_text_and_thinking(content) == parts
+
+
+def test_thinking_part_delta_applies_both_content_and_signature():
+    thinking_part = ThinkingPart(content='Initial content', signature='initial_sig')
+    delta = ThinkingPartDelta(content_delta=' added', signature_delta='new_sig')
+
+    result = delta.apply(thinking_part)
+
+    # The content is appended, and the signature is updated.
+    assert result == snapshot(ThinkingPart(content='Initial content added', signature='new_sig'))
+
+
+def test_thinking_part_delta_applies_signature_only():
+    thinking_part = ThinkingPart(content='Initial content', signature='initial_sig')
+    delta_sig_only = ThinkingPartDelta(content_delta=None, signature_delta='sig_only')
+
+    result_sig_only = delta_sig_only.apply(thinking_part)
+
+    # The content is unchanged, and the signature is updated.
+    assert result_sig_only == snapshot(ThinkingPart(content='Initial content', signature='sig_only'))
+
+
+def test_thinking_part_delta_applies_content_only_preserves_signature():
+    thinking_part = ThinkingPart(content='Initial content', signature='initial_sig')
+    delta_content_only = ThinkingPartDelta(content_delta=' more', signature_delta=None)
+
+    result_content_only = delta_content_only.apply(thinking_part)
+
+    # The content is appended, and the signature is preserved.
+    assert result_content_only == snapshot(ThinkingPart(content='Initial content more', signature='initial_sig'))
+
+
+def test_thinking_part_delta_applies_to_part_with_none_signature():
+    thinking_part_no_sig = ThinkingPart(content='No sig content', signature=None)
+    delta_to_none_sig = ThinkingPartDelta(content_delta=' extra', signature_delta='added_sig')
+
+    result_none_sig = delta_to_none_sig.apply(thinking_part_no_sig)
+
+    # The content is appended, and the signature is updated.
+    assert result_none_sig == snapshot(ThinkingPart(content='No sig content extra', signature='added_sig'))