Fix GoogleModel thinking signature not stored on tool calls when streaming

[jerry-heygen]

Summary

• Fix regression where ThinkingPart.signature was never populated for Google model responses (introduced in Don't insert empty ThinkingPart when Google response ends in text with thought_signature #3516)
• This caused thinking context to be lost in follow-up requests, as Google requires the thought_signature to be sent back to maintain reasoning continuity
• Google's API returns thought_signature on the part following the thinking content (e.g., TextPart or ToolCallPart), not on the ThinkingPart itself
• The fix applies the signature from the following part back to the preceding ThinkingPart.signature

Problem

After #3516, ThinkingPart.signature was always null. The signature was being stored in provider_details on the TextPart/ToolCallPart, but never preserved on the ThinkingPart.

This is critical because when sending message history back to Google, _content_model_response() reads from ThinkingPart.signature to include the signature on the next part. Without it, thinking context is lost in multi-turn conversations.

Per Google's documentation:

"Signatures are returned from the model within other parts in the response, for example function calling or text parts."
"Always send the thought_signature back to the model inside its original Part."

Solution

Non-streaming path: Track last_thinking_part and when a non-thinking part arrives with thought_signature, apply it back to ThinkingPart.signature.

Streaming path: When a non-thinking part with thought_signature arrives, emit a handle_thinking_delta event with the signature to update the previous thinking part.

Test plan

• All existing Google model tests pass (88 passed, 14 skipped)
• test_google_thought_signature_on_thinking_part verifies round-trip behavior
• Updated snapshots to include signature=IsStr() on ThinkingParts

[DouweM]

@jerry-heygen Do you have example code that is failing on main but works on this branch?

Because this part is incorrect:

This is critical because when sending message history back to Google, _content_model_response() reads from ThinkingPart.signature to include the signature on the next part. Without it, thinking context is lost in multi-turn conversations.

The logic here prefers reading thought_signature from the provider_details of the part itself:

pydantic-ai/pydantic_ai_slim/pydantic_ai/models/google.py

Lines 793 to 803 in f2bb374

	for item in m.parts:
	part: PartDict = {}
	if (
	item.provider_details
	and (thought_signature := item.provider_details.get('thought_signature'))
	and m.provider_name == provider_name
	):
	part['thought_signature'] = base64.b64decode(thought_signature)
	elif thinking_part_signature:
	part['thought_signature'] = base64.b64decode(thinking_part_signature)
	thinking_part_signature = None

So it shouldn't be necessary for the signature to (also) be on ThinkingPart.signature.

As you can see in the PR diff, all of the ThinkingParts that now have a signature were followed by parts that already had provider_details['thought_signature'], so that should have been getting sent to Google correctly already.

I'd like to understand the specific scenario in which that wasn't happening, so we can fix it -- adding ThinkingPart.signature shouldn't be necessary.

[jerry-heygen]

Hi @DouweM , so what we have observed on the last release version (v1.26.0) was that the thought signature had been both missing from the ThinkingPart.signature and the provider_details['thought_signature'] (from relevant parts, like ToolCallPart). I have searched for the serialized list[ModelMessage] (where we obtained from agent_run.result.all_messages()) and there was no thought signature strings in it. This has caused the tool call quality in the multi-turn conversation to degrade, but indeed the strange part of the story is Google didn't reject these requests (we use Vertex AI).

We had to downgrade to v1.21.0 where the thought signature is preserved in the ThinkingPart and we have extensively tested that, the tool call quality is higher and more following to the previous thinking when either we use v1.21.0 or this modified branch as in this PR, compared to v1.26.0.

[DouweM]

@jerry-heygen Thanks for the additional context.

I think the bug I introduced in #3516 is that this logic passes provider_details (with the thought_signature) to handle_tool_call_delta correctly:

pydantic-ai/pydantic_ai_slim/pydantic_ai/models/google.py

Lines 734 to 743 in f2bb374

	elif part.function_call:
	maybe_event = self._parts_manager.handle_tool_call_delta(
	vendor_part_id=uuid4(),
	tool_name=part.function_call.name,
	args=part.function_call.args,
	tool_call_id=part.function_call.id,
	provider_details=provider_details,
	)
	if maybe_event is not None: # pragma: no branch
	yield maybe_event

Which then passes it into ToolCallPartDelta correctly as well:

pydantic-ai/pydantic_ai_slim/pydantic_ai/_parts_manager.py

Lines 295 to 297 in f2bb374

	delta = ToolCallPartDelta(
	tool_name_delta=tool_name, args_delta=args, tool_call_id=tool_call_id, provider_details=provider_details
	)

pydantic-ai/pydantic_ai_slim/pydantic_ai/_parts_manager.py

Lines 306 to 308 in f2bb374

	delta = ToolCallPartDelta(
	tool_name_delta=tool_name, args_delta=args, tool_call_id=tool_call_id, provider_details=provider_details
	)

But then the provider_details are not passed along to the ToolCallPart here:

pydantic-ai/pydantic_ai_slim/pydantic_ai/messages.py

Lines 1646 to 1655 in f2bb374

	def as_part(self) -> ToolCallPart | None:
	"""Convert this delta to a fully formed `ToolCallPart` if possible, otherwise return `None`.
	Returns:
	A `ToolCallPart` if `tool_name_delta` is set, otherwise `None`.
	"""
	if self.tool_name_delta is None:
	return None
	return ToolCallPart(self.tool_name_delta, self.args_delta, self.tool_call_id or _generate_tool_call_id())

And here:

pydantic-ai/pydantic_ai_slim/pydantic_ai/messages.py

Lines 1690 to 1741 in f2bb374

	def _apply_to_delta(self, delta: ToolCallPartDelta) -> ToolCallPart | BuiltinToolCallPart | ToolCallPartDelta:
	"""Internal helper to apply this delta to another delta."""
	if self.tool_name_delta:
	# Append incremental text to the existing tool_name_delta
	updated_tool_name_delta = (delta.tool_name_delta or '') + self.tool_name_delta
	delta = replace(delta, tool_name_delta=updated_tool_name_delta)
	if isinstance(self.args_delta, str):
	if isinstance(delta.args_delta, dict):
	raise UnexpectedModelBehavior(
	f'Cannot apply JSON deltas to non-JSON tool arguments ({delta=}, {self=})'
	)
	updated_args_delta = (delta.args_delta or '') + self.args_delta
	delta = replace(delta, args_delta=updated_args_delta)
	elif isinstance(self.args_delta, dict):
	if isinstance(delta.args_delta, str):
	raise UnexpectedModelBehavior(
	f'Cannot apply dict deltas to non-dict tool arguments ({delta=}, {self=})'
	)
	updated_args_delta = {**(delta.args_delta or {}), **self.args_delta}
	delta = replace(delta, args_delta=updated_args_delta)
	if self.tool_call_id:
	delta = replace(delta, tool_call_id=self.tool_call_id)
	# If we now have enough data to create a full ToolCallPart, do so
	if delta.tool_name_delta is not None:
	return ToolCallPart(delta.tool_name_delta, delta.args_delta, delta.tool_call_id or _generate_tool_call_id())
	return delta
	def _apply_to_part(self, part: ToolCallPart | BuiltinToolCallPart) -> ToolCallPart | BuiltinToolCallPart:
	"""Internal helper to apply this delta directly to a `ToolCallPart` or `BuiltinToolCallPart`."""
	if self.tool_name_delta:
	# Append incremental text to the existing tool_name
	tool_name = part.tool_name + self.tool_name_delta
	part = replace(part, tool_name=tool_name)
	if isinstance(self.args_delta, str):
	if isinstance(part.args, dict):
	raise UnexpectedModelBehavior(f'Cannot apply JSON deltas to non-JSON tool arguments ({part=}, {self=})')
	updated_json = (part.args or '') + self.args_delta
	part = replace(part, args=updated_json)
	elif isinstance(self.args_delta, dict):
	if isinstance(part.args, str):
	raise UnexpectedModelBehavior(f'Cannot apply dict deltas to non-dict tool arguments ({part=}, {self=})')
	updated_dict = {**(part.args or {}), **self.args_delta}
	part = replace(part, args=updated_dict)
	if self.tool_call_id:
	part = replace(part, tool_call_id=self.tool_call_id)
	return part

Can you please look into fixing it at that level, instead of by restoring the ThinkingPart.signature?

[jerry-heygen]

@DouweM sure, I've reverted previous changes and applied the correct fix. PTAL