Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion pydantic_ai_slim/pydantic_ai/messages.py
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand Down
43 changes: 42 additions & 1 deletion tests/test_thinking_part.py
Original file line number Diff line number Diff line change
@@ -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(
Expand All @@ -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'))