Don't insert empty ThinkingPart when Google response ends in text with thought_signature (#3516)

Co-authored-by: Anibal Angulo <[email protected]>

Author: DouweM

pydantic_ai_slim/pydantic_ai/_parts_manager.py

@@ -73,6 +73,7 @@ def handle_text_delta(
         vendor_part_id: VendorId | None,
         content: str,
         id: str | None = None,
+        provider_details: dict[str, Any] | None = None,
         thinking_tags: tuple[str, str] | None = None,
         ignore_leading_whitespace: bool = False,
     ) -> ModelResponseStreamEvent | None:
@@ -88,6 +89,7 @@ def handle_text_delta(
                 a TextPart.
             content: The text content to append to the appropriate TextPart.
             id: An optional id for the text part.
+            provider_details: An optional dictionary of provider-specific details for the text part.
             thinking_tags: If provided, will handle content between the thinking tags as thinking parts.
             ignore_leading_whitespace: If True, will ignore leading whitespace in the content.
 
@@ -121,7 +123,9 @@ def handle_text_delta(
                         self._vendor_id_to_part_index.pop(vendor_part_id)
                         return None
                     else:
-                        return self.handle_thinking_delta(vendor_part_id=vendor_part_id, content=content)
+                        return self.handle_thinking_delta(
+                            vendor_part_id=vendor_part_id, content=content, provider_details=provider_details
+                        )
                 elif isinstance(existing_part, TextPart):
                     existing_text_part_and_index = existing_part, part_index
                 else:
@@ -130,7 +134,9 @@ def handle_text_delta(
         if thinking_tags and content == thinking_tags[0]:
             # When we see a thinking start tag (which is a single token), we'll build a new thinking part instead
             self._vendor_id_to_part_index.pop(vendor_part_id, None)
-            return self.handle_thinking_delta(vendor_part_id=vendor_part_id, content='')
+            return self.handle_thinking_delta(
+                vendor_part_id=vendor_part_id, content='', provider_details=provider_details
+            )
 
         if existing_text_part_and_index is None:
             # This is a workaround for models that emit `<think>\n</think>\n\n` or an empty text part ahead of tool calls (e.g. Ollama + Qwen3),
@@ -140,15 +146,15 @@ def handle_text_delta(
 
             # There is no existing text part that should be updated, so create a new one
             new_part_index = len(self._parts)
-            part = TextPart(content=content, id=id)
+            part = TextPart(content=content, id=id, provider_details=provider_details)
             if vendor_part_id is not None:
                 self._vendor_id_to_part_index[vendor_part_id] = new_part_index
             self._parts.append(part)
             return PartStartEvent(index=new_part_index, part=part)
         else:
             # Update the existing TextPart with the new content delta
             existing_text_part, part_index = existing_text_part_and_index
-            part_delta = TextPartDelta(content_delta=content)
+            part_delta = TextPartDelta(content_delta=content, provider_details=provider_details)
             self._parts[part_index] = part_delta.apply(existing_text_part)
             return PartDeltaEvent(index=part_index, delta=part_delta)
 
@@ -160,6 +166,7 @@ def handle_thinking_delta(
         id: str | None = None,
         signature: str | None = None,
         provider_name: str | None = None,
+        provider_details: dict[str, Any] | None = None,
     ) -> ModelResponseStreamEvent:
         """Handle incoming thinking content, creating or updating a ThinkingPart in the manager as appropriate.
 
@@ -175,6 +182,7 @@ def handle_thinking_delta(
             id: An optional id for the thinking part.
             signature: An optional signature for the thinking content.
             provider_name: An optional provider name for the thinking part.
+            provider_details: An optional dictionary of provider-specific details for the thinking part.
 
         Returns:
             A `PartStartEvent` if a new part was created, or a `PartDeltaEvent` if an existing part was updated.
@@ -204,7 +212,13 @@ def handle_thinking_delta(
             if content is not None or signature is not None:
                 # There is no existing thinking part that should be updated, so create a new one
                 new_part_index = len(self._parts)
-                part = ThinkingPart(content=content or '', id=id, signature=signature, provider_name=provider_name)
+                part = ThinkingPart(
+                    content=content or '',
+                    id=id,
+                    signature=signature,
+                    provider_name=provider_name,
+                    provider_details=provider_details,
+                )
                 if vendor_part_id is not None:  # pragma: no branch
                     self._vendor_id_to_part_index[vendor_part_id] = new_part_index
                 self._parts.append(part)
@@ -216,7 +230,10 @@ def handle_thinking_delta(
                 # Update the existing ThinkingPart with the new content and/or signature delta
                 existing_thinking_part, part_index = existing_thinking_part_and_index
                 part_delta = ThinkingPartDelta(
-                    content_delta=content, signature_delta=signature, provider_name=provider_name
+                    content_delta=content,
+                    signature_delta=signature,
+                    provider_name=provider_name,
+                    provider_details=provider_details,
                 )
                 self._parts[part_index] = part_delta.apply(existing_thinking_part)
                 return PartDeltaEvent(index=part_index, delta=part_delta)
@@ -230,6 +247,7 @@ def handle_tool_call_delta(
         tool_name: str | None = None,
         args: str | dict[str, Any] | None = None,
         tool_call_id: str | None = None,
+        provider_details: dict[str, Any] | None = None,
     ) -> ModelResponseStreamEvent | None:
         """Handle or update a tool call, creating or updating a `ToolCallPart`, `BuiltinToolCallPart`, or `ToolCallPartDelta`.
 
@@ -246,6 +264,7 @@ def handle_tool_call_delta(
                 a name match when `vendor_part_id` is None.
             args: Arguments for the tool call, either as a string, a dictionary of key-value pairs, or None.
             tool_call_id: An optional string representing an identifier for this tool call.
+            provider_details: An optional dictionary of provider-specific details for the tool call part.
 
         Returns:
             - A `PartStartEvent` if a new ToolCallPart or BuiltinToolCallPart is created.
@@ -280,7 +299,9 @@ def handle_tool_call_delta(
 
         if existing_matching_part_and_index is None:
             # No matching part/delta was found, so create a new ToolCallPartDelta (or ToolCallPart if fully formed)
-            delta = ToolCallPartDelta(tool_name_delta=tool_name, args_delta=args, tool_call_id=tool_call_id)
+            delta = ToolCallPartDelta(
+                tool_name_delta=tool_name, args_delta=args, tool_call_id=tool_call_id, provider_details=provider_details
+            )
             part = delta.as_part() or delta
             if vendor_part_id is not None:
                 self._vendor_id_to_part_index[vendor_part_id] = len(self._parts)
@@ -292,7 +313,9 @@ def handle_tool_call_delta(
         else:
             # Update the existing part or delta with the new information
             existing_part, part_index = existing_matching_part_and_index
-            delta = ToolCallPartDelta(tool_name_delta=tool_name, args_delta=args, tool_call_id=tool_call_id)
+            delta = ToolCallPartDelta(
+                tool_name_delta=tool_name, args_delta=args, tool_call_id=tool_call_id, provider_details=provider_details
+            )
             updated_part = delta.apply(existing_part)
             self._parts[part_index] = updated_part
             if isinstance(updated_part, ToolCallPart | BuiltinToolCallPart):
@@ -313,6 +336,7 @@ def handle_tool_call_part(
         args: str | dict[str, Any] | None,
         tool_call_id: str | None = None,
         id: str | None = None,
+        provider_details: dict[str, Any] | None = None,
     ) -> ModelResponseStreamEvent:
         """Immediately create or fully-overwrite a ToolCallPart with the given information.
 
@@ -325,6 +349,7 @@ def handle_tool_call_part(
             args: The arguments for the tool call, either as a string, a dictionary, or None.
             tool_call_id: An optional string identifier for this tool call.
             id: An optional identifier for this tool call part.
+            provider_details: An optional dictionary of provider-specific details for the tool call part.
 
         Returns:
             ModelResponseStreamEvent: A `PartStartEvent` indicating that a new tool call part
@@ -335,6 +360,7 @@ def handle_tool_call_part(
             args=args,
             tool_call_id=tool_call_id or _generate_tool_call_id(),
             id=id,
+            provider_details=provider_details,
         )
         if vendor_part_id is None:
             # vendor_part_id is None, so we unconditionally append a new ToolCallPart to the end of the list

pydantic_ai_slim/pydantic_ai/messages.py

@@ -880,6 +880,11 @@ class BuiltinToolReturnPart(BaseToolReturnPart):
     provider_name: str | None = None
     """The name of the provider that generated the response."""
 
+    provider_details: dict[str, Any] | None = None
+    """Additional data returned by the provider that can't be mapped to standard fields.
+
+    This is used for data that is required to be sent back to APIs, as well as data users may want to access programmatically."""
+
     part_kind: Literal['builtin-tool-return'] = 'builtin-tool-return'
     """Part type identifier, this is available on all parts as a discriminator."""
 
@@ -1021,6 +1026,11 @@ class TextPart:
     id: str | None = None
     """An optional identifier of the text part."""
 
+    provider_details: dict[str, Any] | None = None
+    """Additional data returned by the provider that can't be mapped to standard fields.
+
+    This is used for data that is required to be sent back to APIs, as well as data users may want to access programmatically."""
+
     part_kind: Literal['text'] = 'text'
     """Part type identifier, this is available on all parts as a discriminator."""
 
@@ -1060,14 +1070,13 @@ class ThinkingPart:
     Signatures are only sent back to the same provider.
     """
 
-    part_kind: Literal['thinking'] = 'thinking'
-    """Part type identifier, this is available on all parts as a discriminator."""
-
     provider_details: dict[str, Any] | None = None
-    """Additional provider-specific details in a serializable format.
+    """Additional data returned by the provider that can't be mapped to standard fields.
 
-    This allows storing selected vendor-specific data that isn't mapped to standard ThinkingPart fields.
-    """
+    This is used for data that is required to be sent back to APIs, as well as data users may want to access programmatically."""
+
+    part_kind: Literal['thinking'] = 'thinking'
+    """Part type identifier, this is available on all parts as a discriminator."""
 
     def has_content(self) -> bool:
         """Return `True` if the thinking content is non-empty."""
@@ -1092,6 +1101,11 @@ class FilePart:
     """The name of the provider that generated the response.
     """
 
+    provider_details: dict[str, Any] | None = None
+    """Additional data returned by the provider that can't be mapped to standard fields.
+
+    This is used for data that is required to be sent back to APIs, as well as data users may want to access programmatically."""
+
     part_kind: Literal['file'] = 'file'
     """Part type identifier, this is available on all parts as a discriminator."""
 
@@ -1128,6 +1142,11 @@ class BaseToolCallPart:
 
     This is used by some APIs like OpenAI Responses."""
 
+    provider_details: dict[str, Any] | None = None
+    """Additional data returned by the provider that can't be mapped to standard fields.
+
+    This is used for data that is required to be sent back to APIs, as well as data users may want to access programmatically."""
+
     def args_as_dict(self) -> dict[str, Any]:
         """Return the arguments as a Python dictionary.
 
@@ -1232,11 +1251,7 @@ class ModelResponse:
         # `vendor_details` is deprecated, but we still want to support deserializing model responses stored in a DB before the name was changed
         pydantic.Field(validation_alias=pydantic.AliasChoices('provider_details', 'vendor_details')),
     ] = None
-    """Additional provider-specific details in a serializable format.
-
-    This allows storing selected vendor-specific data that isn't mapped to standard ModelResponse fields.
-    For OpenAI models, this may include 'logprobs', 'finish_reason', etc.
-    """
+    """Additional data returned by the provider that can't be mapped to standard fields."""
 
     provider_response_id: Annotated[
         str | None,
@@ -1460,6 +1475,11 @@ class TextPartDelta:
 
     _: KW_ONLY
 
+    provider_details: dict[str, Any] | None = None
+    """Additional data returned by the provider that can't be mapped to standard fields.
+
+    This is used for data that is required to be sent back to APIs, as well as data users may want to access programmatically."""
+
     part_delta_kind: Literal['text'] = 'text'
     """Part delta type identifier, used as a discriminator."""
 
@@ -1477,7 +1497,11 @@ def apply(self, part: ModelResponsePart) -> TextPart:
         """
         if not isinstance(part, TextPart):
             raise ValueError('Cannot apply TextPartDeltas to non-TextParts')  # pragma: no cover
-        return replace(part, content=part.content + self.content_delta)
+        return replace(
+            part,
+            content=part.content + self.content_delta,
+            provider_details={**(part.provider_details or {}), **(self.provider_details or {})} or None,
+        )
 
     __repr__ = _utils.dataclasses_no_defaults_repr
 
@@ -1501,6 +1525,11 @@ class ThinkingPartDelta:
     Signatures are only sent back to the same provider.
     """
 
+    provider_details: dict[str, Any] | None = None
+    """Additional data returned by the provider that can't be mapped to standard fields.
+
+    This is used for data that is required to be sent back to APIs, as well as data users may want to access programmatically."""
+
     part_delta_kind: Literal['thinking'] = 'thinking'
     """Part delta type identifier, used as a discriminator."""
 
@@ -1526,7 +1555,14 @@ def apply(self, part: ModelResponsePart | ThinkingPartDelta) -> ThinkingPart | T
             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
             new_provider_name = self.provider_name if self.provider_name is not None else part.provider_name
-            return replace(part, content=new_content, signature=new_signature, provider_name=new_provider_name)
+            new_provider_details = {**(part.provider_details or {}), **(self.provider_details or {})} or None
+            return replace(
+                part,
+                content=new_content,
+                signature=new_signature,
+                provider_name=new_provider_name,
+                provider_details=new_provider_details,
+            )
         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')
@@ -1536,6 +1572,8 @@ def apply(self, part: ModelResponsePart | ThinkingPartDelta) -> ThinkingPart | T
                 part = replace(part, signature_delta=self.signature_delta)
             if self.provider_name is not None:
                 part = replace(part, provider_name=self.provider_name)
+            if self.provider_details is not None:
+                part = replace(part, provider_details={**(part.provider_details or {}), **self.provider_details})
             return part
         raise ValueError(  # pragma: no cover
             f'Cannot apply ThinkingPartDeltas to non-ThinkingParts or non-ThinkingPartDeltas ({part=}, {self=})'
@@ -1564,6 +1602,11 @@ class ToolCallPartDelta:
     Note this is never treated as a delta — it can replace None, but otherwise if a
     non-matching value is provided an error will be raised."""
 
+    provider_details: dict[str, Any] | None = None
+    """Additional data returned by the provider that can't be mapped to standard fields.
+
+    This is used for data that is required to be sent back to APIs, as well as data users may want to access programmatically."""
+
     part_delta_kind: Literal['tool_call'] = 'tool_call'
     """Part delta type identifier, used as a discriminator."""