Skip to content

Capture thought signatures from Google models #2305

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

Capture thought signatures from Google models #2305

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 1 addition & 1 deletion pydantic_ai_slim/pydantic_ai/messages.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -682,7 +682,7 @@ class ThinkingPart:
signature: str | None = None
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I suppose the alternative to base64 encoding/decoding would be do signature: str | bytes | None = None, but I'm not sure what kind of ripple effects that would produce.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's do that

"""The signature of the thinking.

The signature is only available on the Anthropic models.
The signature is only available on the Anthropic and Google models.
"""

part_kind: Literal['thinking'] = 'thinking'
Expand Down
18 changes: 16 additions & 2 deletions pydantic_ai_slim/pydantic_ai/models/google.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,14 @@
"""


def _bytes_to_base64_string(data: bytes) -> str:
return base64.b64encode(data).decode("utf-8")


def _base64_string_to_bytes(string: str) -> bytes:
return base64.b64decode(string)


class GoogleModelSettings(ModelSettings, total=False):
"""Settings used for a Gemini model request."""

Expand Down Expand Up @@ -493,7 +501,13 @@ def _content_model_response(m: ModelResponse) -> ContentDict:
elif isinstance(item, ThinkingPart): # pragma: no cover
# NOTE: We don't send ThinkingPart to the providers yet. If you are unsatisfied with this,
# please open an issue. The below code is the code to send thinking to the provider.
# parts.append({'text': item.content, 'thought': True})
# parts.append(
# PartDict(
# thought=True,
# thought_signature=_base64_string_to_bytes(item.signature),
# text=item.content,
# )
# )
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

NOTE: We don't send ThinkingPart to the providers yet. If you are unsatisfied with this,

Can we enable it here, or is there a broader rule about this in place?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah we can return the signatures, just not the actual thought text.

pass
else:
assert_never(item)
Expand All @@ -511,7 +525,7 @@ def _process_response_from_parts(
for part in parts:
if part.text is not None:
if part.thought:
items.append(ThinkingPart(content=part.text))
items.append(ThinkingPart(content=part.text, signature=_bytes_to_base64_string(part.thought_signature)))
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually, I don't think this is going to work.

The thinking signature in the example (see Appendix 1) comes right after the thinking part.

So, I think support for a signature field would need to be added to ToolCallPart() and TextPart()?

Copy link
Collaborator

@DouweM DouweM Jul 28, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, interesting, I had missed that detail in reading your issue originally.

For consistency with Anthropic, and to not have to modify all the other part classes with this Google-specific detail, I would prefer to store the signature on the most recent thinking part in Pydantic AI representation, and then to include it on the (next) text/tool part when sending it back to Google. Can you see if you can make that work?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note that we'll also need to handle the signature when streaming, which is implemented here:

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

Line 459 in 806d56d

yield self._parts_manager.handle_thinking_delta(vendor_part_id='thinking', content=part.text)

Since the signature will come on a text/tool message, we'll want to check if its set there, and if so call handle_thinking_delta(vendor_part_id='thinking', signature=signature) before we call handle_text_delta or handle_tool_call_delta so that the signature will be added to the most recently seen thinking part.

else:
items.append(TextPart(content=part.text))
elif part.function_call:
Expand Down
Loading