Skip to content

Commit f42e523

Browse files
jerry-heygenDouweM
jerry-heygen
and
DouweM
authored
Fix GoogleModel thinking signature not stored on tool calls when streaming (#3647)
Co-authored-by: Douwe Maan <hi@douwe.me>
1 parent eb2023e commit f42e523

File tree

4 files changed

+326
-8
lines changed

4 files changed

+326
-8
lines changed

‎pydantic_ai_slim/pydantic_ai/messages.py‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1652,7 +1652,12 @@ def as_part(self) -> ToolCallPart | None:
16521652
if self.tool_name_delta is None:
16531653
return None
16541654

1655-
return ToolCallPart(self.tool_name_delta, self.args_delta, self.tool_call_id or _generate_tool_call_id())
1655+
return ToolCallPart(
1656+
self.tool_name_delta,
1657+
self.args_delta,
1658+
self.tool_call_id or _generate_tool_call_id(),
1659+
provider_details=self.provider_details,
1660+
)
16561661

16571662
@overload
16581663
def apply(self, part: ModelResponsePart) -> ToolCallPart | BuiltinToolCallPart: ...
@@ -1712,9 +1717,18 @@ def _apply_to_delta(self, delta: ToolCallPartDelta) -> ToolCallPart | BuiltinToo
17121717
if self.tool_call_id:
17131718
delta = replace(delta, tool_call_id=self.tool_call_id)
17141719

1720+
if self.provider_details:
1721+
merged_provider_details = {**(delta.provider_details or {}), **self.provider_details}
1722+
delta = replace(delta, provider_details=merged_provider_details)
1723+
17151724
# If we now have enough data to create a full ToolCallPart, do so
17161725
if delta.tool_name_delta is not None:
1717-
return ToolCallPart(delta.tool_name_delta, delta.args_delta, delta.tool_call_id or _generate_tool_call_id())
1726+
return ToolCallPart(
1727+
delta.tool_name_delta,
1728+
delta.args_delta,
1729+
delta.tool_call_id or _generate_tool_call_id(),
1730+
provider_details=delta.provider_details,
1731+
)
17181732

17191733
return delta
17201734

@@ -1738,6 +1752,11 @@ def _apply_to_part(self, part: ToolCallPart | BuiltinToolCallPart) -> ToolCallPa
17381752

17391753
if self.tool_call_id:
17401754
part = replace(part, tool_call_id=self.tool_call_id)
1755+
1756+
if self.provider_details:
1757+
merged_provider_details = {**(part.provider_details or {}), **self.provider_details}
1758+
part = replace(part, provider_details=merged_provider_details)
1759+
17411760
return part
17421761

17431762
__repr__ = _utils.dataclasses_no_defaults_repr

‎tests/models/cassettes/test_google/test_google_streaming_tool_call_thought_signature.yaml‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
interactions:
2+
- request:
3+
headers:
4+
accept:
5+
- '*/*'
6+
accept-encoding:
7+
- gzip, deflate
8+
connection:
9+
- keep-alive
10+
content-length:
11+
- '343'
12+
content-type:
13+
- application/json
14+
host:
15+
- generativelanguage.googleapis.com
16+
method: POST
17+
parsed_body:
18+
contents:
19+
- parts:
20+
- text: What is the capital of the user country? Call the tool
21+
role: user
22+
generationConfig:
23+
responseModalities:
24+
- TEXT
25+
tools:
26+
- functionDeclarations:
27+
- description: ''
28+
name: get_country
29+
parameters_json_schema:
30+
additionalProperties: false
31+
properties: {}
32+
type: object
33+
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-3-pro-preview:streamGenerateContent?alt=sse
34+
response:
35+
body:
36+
string: "data: {\"candidates\": [{\"content\": {\"parts\": [{\"functionCall\": {\"name\": \"get_country\",\"args\":
37+
{}},\"thoughtSignature\": \"Et4ECtsEAXLI2nyJtljzBkvSzfFwCuDHg1Ubv1K7gmhe/8Oj4BE0bMEnzcpgKQseNP2rR0JUZs8RmtbDWR35oJIKp4xTji+uQwSb7Fui2IxdaP3uLc4HGnQKIIuI+c1I3t0HfFHYN/oY4xj7gCkxy9Sk4OjOjSnEjC6CZs7qYLHX/0abu+hou4rQl0S/+sWdZJfR26eb7W9Ct6RwUAnFYATDu34opPFtAR2KtmrOdffLwYU//6lNk9JbZP9mvz3I9+mYKMwdZ+5iKt6cHPTNZk17yQM4+jXuXXoVC8fwxm/G5VMcB1TGiZRUFrd5ohzNAyCz+/YRfbetVJtqIZglU69iGTtr+5KHMTloxq6TR8QExwz9re3V08GCtG0SlXN937cUNnoAzUOJaJyTha6PA7b5Qrz9E7B3IMsRlvjoza+cLUlrvFFumnyNaMAvwLJfOUvNVFzg/57qhMh5ZYr1IrAW4fCSBVuPW2UH/7GIvORkA0mVSt83DED599s+JsB1dmnu668Xpw8JLKRnpmXnGpUQKNX9+JEOU3EoucvV86H/kgjgCNi5jW9NyxG2DZ6up18SQdOZjGTCVNB5/792gbmu9FqkUQeQW3aZ5ocynLeY5XCtH6NTiMkegw0yKIP1itmaBabshFQ1rn8qNYjLcTobvSgsof9asSyIaa9dgBNv9MwA1txUFN9ybXVONp/7la0osS+xqn8dUSeUYfbVtLCV7hkwWcywETLAA7wP0DL21Hnm42Ln2g2pDIpMaYsI0TJd4O7eoribaoAJsyjV6OODey3zfWKrKT8rYF9eplXe\"}],\"role\":
38+
\"model\"},\"index\": 0}],\"usageMetadata\": {\"promptTokenCount\": 29,\"candidatesTokenCount\": 10,\"totalTokenCount\":
39+
156,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\": 29}],\"thoughtsTokenCount\": 117},\"modelVersion\":
40+
\"gemini-3-pro-preview\",\"responseId\": \"WC44aZ29Erflz7IPp4CZwQ4\"}\r\n\r\ndata: {\"candidates\": [{\"content\":
41+
{\"parts\": [{\"text\": \"\"}],\"role\": \"model\"},\"finishReason\": \"STOP\",\"index\": 0}],\"usageMetadata\": {\"promptTokenCount\":
42+
29,\"candidatesTokenCount\": 10,\"totalTokenCount\": 156,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\":
43+
29}],\"thoughtsTokenCount\": 117},\"modelVersion\": \"gemini-3-pro-preview\",\"responseId\": \"WC44aZ29Erflz7IPp4CZwQ4\"}\r\n\r\n"
44+
headers:
45+
alt-svc:
46+
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
47+
content-disposition:
48+
- attachment
49+
content-type:
50+
- text/event-stream
51+
server-timing:
52+
- gfet4t7; dur=4243
53+
transfer-encoding:
54+
- chunked
55+
vary:
56+
- Origin
57+
- X-Origin
58+
- Referer
59+
status:
60+
code: 200
61+
message: OK
62+
- request:
63+
headers:
64+
accept:
65+
- '*/*'
66+
accept-encoding:
67+
- gzip, deflate
68+
connection:
69+
- keep-alive
70+
content-length:
71+
- '1478'
72+
content-type:
73+
- application/json
74+
host:
75+
- generativelanguage.googleapis.com
76+
method: POST
77+
parsed_body:
78+
contents:
79+
- parts:
80+
- text: What is the capital of the user country? Call the tool
81+
role: user
82+
- parts:
83+
- functionCall:
84+
args: {}
85+
id: pyd_ai_e31595c62bcb4dc4ae439480bd2f0623
86+
name: get_country
87+
thoughtSignature: Et4ECtsEAXLI2nyJtljzBkvSzfFwCuDHg1Ubv1K7gmhe_8Oj4BE0bMEnzcpgKQseNP2rR0JUZs8RmtbDWR35oJIKp4xTji-uQwSb7Fui2IxdaP3uLc4HGnQKIIuI-c1I3t0HfFHYN_oY4xj7gCkxy9Sk4OjOjSnEjC6CZs7qYLHX_0abu-hou4rQl0S_-sWdZJfR26eb7W9Ct6RwUAnFYATDu34opPFtAR2KtmrOdffLwYU__6lNk9JbZP9mvz3I9-mYKMwdZ-5iKt6cHPTNZk17yQM4-jXuXXoVC8fwxm_G5VMcB1TGiZRUFrd5ohzNAyCz-_YRfbetVJtqIZglU69iGTtr-5KHMTloxq6TR8QExwz9re3V08GCtG0SlXN937cUNnoAzUOJaJyTha6PA7b5Qrz9E7B3IMsRlvjoza-cLUlrvFFumnyNaMAvwLJfOUvNVFzg_57qhMh5ZYr1IrAW4fCSBVuPW2UH_7GIvORkA0mVSt83DED599s-JsB1dmnu668Xpw8JLKRnpmXnGpUQKNX9-JEOU3EoucvV86H_kgjgCNi5jW9NyxG2DZ6up18SQdOZjGTCVNB5_792gbmu9FqkUQeQW3aZ5ocynLeY5XCtH6NTiMkegw0yKIP1itmaBabshFQ1rn8qNYjLcTobvSgsof9asSyIaa9dgBNv9MwA1txUFN9ybXVONp_7la0osS-xqn8dUSeUYfbVtLCV7hkwWcywETLAA7wP0DL21Hnm42Ln2g2pDIpMaYsI0TJd4O7eoribaoAJsyjV6OODey3zfWKrKT8rYF9eplXe
88+
role: model
89+
- parts:
90+
- functionResponse:
91+
id: pyd_ai_e31595c62bcb4dc4ae439480bd2f0623
92+
name: get_country
93+
response:
94+
return_value: Mexico
95+
role: user
96+
generationConfig:
97+
responseModalities:
98+
- TEXT
99+
tools:
100+
- functionDeclarations:
101+
- description: ''
102+
name: get_country
103+
parameters_json_schema:
104+
additionalProperties: false
105+
properties: {}
106+
type: object
107+
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-3-pro-preview:streamGenerateContent?alt=sse
108+
response:
109+
body:
110+
string: "data: {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \"The\"}],\"role\": \"model\"},\"index\": 0}],\"usageMetadata\":
111+
{\"promptTokenCount\": 55,\"candidatesTokenCount\": 1,\"totalTokenCount\": 56,\"promptTokensDetails\": [{\"modality\":
112+
\"TEXT\",\"tokenCount\": 55}]},\"modelVersion\": \"gemini-3-pro-preview\",\"responseId\": \"hy44aYu1BO-ez7IP-cTyoAk\"}\r\n\r\ndata:
113+
{\"candidates\": [{\"content\": {\"parts\": [{\"text\": \" capital of Mexico is **Mexico City**.\"}],\"role\": \"model\"},\"index\":
114+
0}],\"usageMetadata\": {\"promptTokenCount\": 55,\"candidatesTokenCount\": 9,\"totalTokenCount\": 64,\"promptTokensDetails\":
115+
[{\"modality\": \"TEXT\",\"tokenCount\": 55}]},\"modelVersion\": \"gemini-3-pro-preview\",\"responseId\": \"hy44aYu1BO-ez7IP-cTyoAk\"}\r\n\r\ndata:
116+
{\"candidates\": [{\"content\": {\"parts\": [{\"text\": \"\"}],\"role\": \"model\"},\"finishReason\": \"STOP\",\"index\":
117+
0}],\"usageMetadata\": {\"promptTokenCount\": 170,\"candidatesTokenCount\": 9,\"totalTokenCount\": 179,\"promptTokensDetails\":
118+
[{\"modality\": \"TEXT\",\"tokenCount\": 170}]},\"modelVersion\": \"gemini-3-pro-preview\",\"responseId\": \"hy44aYu1BO-ez7IP-cTyoAk\"}\r\n\r\n"
119+
headers:
120+
alt-svc:
121+
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
122+
content-disposition:
123+
- attachment
124+
content-type:
125+
- text/event-stream
126+
server-timing:
127+
- gfet4t7; dur=46746
128+
transfer-encoding:
129+
- chunked
130+
vary:
131+
- Origin
132+
- X-Origin
133+
- Referer
134+
status:
135+
code: 200
136+
message: OK
137+
version: 1

‎tests/models/test_google.py‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,9 @@
1515
from typing_extensions import TypedDict
1616

1717
from pydantic_ai import (
18+
AgentRunResult,
19+
AgentRunResultEvent,
20+
AgentStreamEvent,
1821
AudioUrl,
1922
BinaryContent,
2023
BinaryImage,
@@ -4425,3 +4428,127 @@ def test_google_missing_tool_call_thought_signature():
44254428
],
44264429
}
44274430
)
4431+
4432+
4433+
async def test_google_streaming_tool_call_thought_signature(
4434+
allow_model_requests: None, google_provider: GoogleProvider
4435+
):
4436+
model = GoogleModel('gemini-3-pro-preview', provider=google_provider)
4437+
agent = Agent(model=model)
4438+
4439+
@agent.tool_plain
4440+
def get_country() -> str:
4441+
return 'Mexico'
4442+
4443+
events: list[AgentStreamEvent] = []
4444+
result: AgentRunResult | None = None
4445+
async for event in agent.run_stream_events('What is the capital of the user country? Call the tool'):
4446+
if isinstance(event, AgentRunResultEvent):
4447+
result = event.result
4448+
else:
4449+
events.append(event)
4450+
4451+
assert result is not None
4452+
assert result.all_messages() == snapshot(
4453+
[
4454+
ModelRequest(
4455+
parts=[
4456+
UserPromptPart(
4457+
content='What is the capital of the user country? Call the tool',
4458+
timestamp=IsDatetime(),
4459+
)
4460+
],
4461+
run_id=IsStr(),
4462+
),
4463+
ModelResponse(
4464+
parts=[
4465+
ToolCallPart(
4466+
tool_name='get_country',
4467+
args={},
4468+
tool_call_id=IsStr(),
4469+
provider_details={'thought_signature': IsStr()},
4470+
)
4471+
],
4472+
usage=RequestUsage(
4473+
input_tokens=29, output_tokens=127, details={'thoughts_tokens': 117, 'text_prompt_tokens': 29}
4474+
),
4475+
model_name='gemini-3-pro-preview',
4476+
timestamp=IsDatetime(),
4477+
provider_name='google-gla',
4478+
provider_details={'finish_reason': 'STOP'},
4479+
provider_response_id=IsStr(),
4480+
finish_reason='stop',
4481+
run_id=IsStr(),
4482+
),
4483+
ModelRequest(
4484+
parts=[
4485+
ToolReturnPart(
4486+
tool_name='get_country',
4487+
content='Mexico',
4488+
tool_call_id=IsStr(),
4489+
timestamp=IsDatetime(),
4490+
)
4491+
],
4492+
run_id=IsStr(),
4493+
),
4494+
ModelResponse(
4495+
parts=[TextPart(content='The capital of Mexico is **Mexico City**.')],
4496+
usage=RequestUsage(input_tokens=170, output_tokens=9, details={'text_prompt_tokens': 170}),
4497+
model_name='gemini-3-pro-preview',
4498+
timestamp=IsDatetime(),
4499+
provider_name='google-gla',
4500+
provider_details={'finish_reason': 'STOP'},
4501+
provider_response_id=IsStr(),
4502+
finish_reason='stop',
4503+
run_id=IsStr(),
4504+
),
4505+
]
4506+
)
4507+
assert events == snapshot(
4508+
[
4509+
PartStartEvent(
4510+
index=0,
4511+
part=ToolCallPart(
4512+
tool_name='get_country',
4513+
args={},
4514+
tool_call_id=IsStr(),
4515+
provider_details={'thought_signature': IsStr()},
4516+
),
4517+
),
4518+
PartEndEvent(
4519+
index=0,
4520+
part=ToolCallPart(
4521+
tool_name='get_country',
4522+
args={},
4523+
tool_call_id=IsStr(),
4524+
provider_details={'thought_signature': IsStr()},
4525+
),
4526+
),
4527+
FunctionToolCallEvent(
4528+
part=ToolCallPart(
4529+
tool_name='get_country',
4530+
args={},
4531+
tool_call_id=IsStr(),
4532+
provider_details={'thought_signature': IsStr()},
4533+
)
4534+
),
4535+
FunctionToolResultEvent(
4536+
result=ToolReturnPart(
4537+
tool_name='get_country',
4538+
content='Mexico',
4539+
tool_call_id=IsStr(),
4540+
timestamp=IsDatetime(),
4541+
)
4542+
),
4543+
PartStartEvent(index=0, part=TextPart(content='The')),
4544+
FinalResultEvent(tool_name=None, tool_call_id=None),
4545+
PartDeltaEvent(
4546+
index=0,
4547+
delta=TextPartDelta(content_delta=' capital of Mexico is **Mexico City**.'),
4548+
),
4549+
PartEndEvent(
4550+
index=0,
4551+
part=TextPart(content='The capital of Mexico is **Mexico City**.'),
4552+
),
4553+
]
4554+
)

0 commit comments

Comments
 (0)