Skip to content

Commit ac85205

Browse files
KludexKludex
committed
Add support for Google
1 parent 13d7433 commit ac85205

File tree

5 files changed

+302
-2
lines changed

5 files changed

+302
-2
lines changed

‎pydantic_ai_slim/pydantic_ai/builtin_tools.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from abc import ABC
44
from dataclasses import dataclass
5-
from typing import Literal
5+
from typing import Any, Literal
66

77
from typing_extensions import TypedDict
88

@@ -15,9 +15,11 @@ class AbstractBuiltinTool(ABC):
1515
1616
This class is abstract and cannot be instantiated directly.
1717
18-
The builtin tools are are passed to the model as part of the `ModelRequestParameters`.
18+
The builtin tools are passed to the model as part of the `ModelRequestParameters`.
1919
"""
2020

21+
def handle_custom_tool_definition(self, model: str) -> Any: ...
22+
2123

2224
@dataclass
2325
class WebSearchTool(AbstractBuiltinTool):

‎pydantic_ai_slim/pydantic_ai/models/google.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111

1212
from typing_extensions import assert_never
1313

14+
from pydantic_ai.builtin_tools import CodeExecutionTool, WebSearchTool
1415
from pydantic_ai.providers import Provider
1516

1617
from .. import UnexpectedModelBehavior, UserError, _utils, usage
@@ -55,10 +56,12 @@
5556
FunctionDeclarationDict,
5657
GenerateContentConfigDict,
5758
GenerateContentResponse,
59+
GoogleSearchDict,
5860
Part,
5961
PartDict,
6062
SafetySettingDict,
6163
ThinkingConfigDict,
64+
ToolCodeExecutionDict,
6265
ToolConfigDict,
6366
ToolDict,
6467
ToolListUnionDict,
@@ -192,6 +195,7 @@ def _customize_tool_def(t: ToolDefinition):
192195

193196
return ModelRequestParameters(
194197
function_tools=[_customize_tool_def(tool) for tool in model_request_parameters.function_tools],
198+
builtin_tools=model_request_parameters.builtin_tools,
195199
allow_text_output=model_request_parameters.allow_text_output,
196200
output_tools=[_customize_tool_def(tool) for tool in model_request_parameters.output_tools],
197201
)
@@ -216,6 +220,13 @@ def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[T
216220
ToolDict(function_declarations=[_function_declaration_from_tool(t)])
217221
for t in model_request_parameters.output_tools
218222
]
223+
for tool in model_request_parameters.builtin_tools:
224+
if isinstance(tool, WebSearchTool):
225+
tools.append(ToolDict(google_search=GoogleSearchDict()))
226+
elif isinstance(tool, CodeExecutionTool):
227+
tools.append(ToolDict(code_execution=ToolCodeExecutionDict()))
228+
else:
229+
raise UserError(f'Unsupported builtin tool: {tool}')
219230
return tools or None
220231

221232
def _get_tool_config(

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

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
interactions:
2+
- request:
3+
headers:
4+
accept:
5+
- '*/*'
6+
accept-encoding:
7+
- gzip, deflate
8+
connection:
9+
- keep-alive
10+
content-length:
11+
- '234'
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 day is today in Utrecht?
21+
role: user
22+
generationConfig: {}
23+
systemInstruction:
24+
parts:
25+
- text: You are a helpful chatbot.
26+
role: user
27+
tools:
28+
- codeExecution: {}
29+
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
30+
response:
31+
headers:
32+
alt-svc:
33+
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
34+
content-length:
35+
- '824'
36+
content-type:
37+
- application/json; charset=UTF-8
38+
server-timing:
39+
- gfet4t7; dur=634
40+
transfer-encoding:
41+
- chunked
42+
vary:
43+
- Origin
44+
- X-Origin
45+
- Referer
46+
parsed_body:
47+
candidates:
48+
- content:
49+
parts:
50+
- text: |
51+
To give you the exact day in Utrecht, I need to know the current date. Can you please provide the date? I can then determine the day.
52+
role: model
53+
finishReason: STOP
54+
modelVersion: gemini-2.0-flash
55+
responseId: WN42aLH_CoKZgLUP97W1wQg
56+
usageMetadata:
57+
candidatesTokenCount: 32
58+
candidatesTokensDetails:
59+
- modality: TEXT
60+
tokenCount: 32
61+
promptTokenCount: 13
62+
promptTokensDetails:
63+
- modality: TEXT
64+
tokenCount: 13
65+
toolUsePromptTokensDetails:
66+
- modality: TEXT
67+
totalTokenCount: 45
68+
status:
69+
code: 200
70+
message: OK
71+
version: 1

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

Lines changed: 199 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,199 @@
1+
interactions:
2+
- request:
3+
headers:
4+
accept:
5+
- '*/*'
6+
accept-encoding:
7+
- gzip, deflate
8+
connection:
9+
- keep-alive
10+
content-length:
11+
- '233'
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 day is today in Utrecht?
21+
role: user
22+
generationConfig: {}
23+
systemInstruction:
24+
parts:
25+
- text: You are a helpful chatbot.
26+
role: user
27+
tools:
28+
- googleSearch: {}
29+
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
30+
response:
31+
headers:
32+
alt-svc:
33+
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
34+
content-length:
35+
- '5991'
36+
content-type:
37+
- application/json; charset=UTF-8
38+
server-timing:
39+
- gfet4t7; dur=1340
40+
transfer-encoding:
41+
- chunked
42+
vary:
43+
- Origin
44+
- X-Origin
45+
- Referer
46+
parsed_body:
47+
candidates:
48+
- content:
49+
parts:
50+
- text: |
51+
Today is Wednesday, May 28, 2025, in Utrecht.
52+
role: model
53+
finishReason: STOP
54+
groundingMetadata:
55+
retrievalMetadata: {}
56+
searchEntryPoint:
57+
renderedContent: |
58+
<style>
59+
.container {
60+
align-items: center;
61+
border-radius: 8px;
62+
display: flex;
63+
font-family: Google Sans, Roboto, sans-serif;
64+
font-size: 14px;
65+
line-height: 20px;
66+
padding: 8px 12px;
67+
}
68+
.chip {
69+
display: inline-block;
70+
border: solid 1px;
71+
border-radius: 16px;
72+
min-width: 14px;
73+
padding: 5px 16px;
74+
text-align: center;
75+
user-select: none;
76+
margin: 0 8px;
77+
-webkit-tap-highlight-color: transparent;
78+
}
79+
.carousel {
80+
overflow: auto;
81+
scrollbar-width: none;
82+
white-space: nowrap;
83+
margin-right: -12px;
84+
}
85+
.headline {
86+
display: flex;
87+
margin-right: 4px;
88+
}
89+
.gradient-container {
90+
position: relative;
91+
}
92+
.gradient {
93+
position: absolute;
94+
transform: translate(3px, -9px);
95+
height: 36px;
96+
width: 9px;
97+
}
98+
@media (prefers-color-scheme: light) {
99+
.container {
100+
background-color: #fafafa;
101+
box-shadow: 0 0 0 1px #0000000f;
102+
}
103+
.headline-label {
104+
color: #1f1f1f;
105+
}
106+
.chip {
107+
background-color: #ffffff;
108+
border-color: #d2d2d2;
109+
color: #5e5e5e;
110+
text-decoration: none;
111+
}
112+
.chip:hover {
113+
background-color: #f2f2f2;
114+
}
115+
.chip:focus {
116+
background-color: #f2f2f2;
117+
}
118+
.chip:active {
119+
background-color: #d8d8d8;
120+
border-color: #b6b6b6;
121+
}
122+
.logo-dark {
123+
display: none;
124+
}
125+
.gradient {
126+
background: linear-gradient(90deg, #fafafa 15%, #fafafa00 100%);
127+
}
128+
}
129+
@media (prefers-color-scheme: dark) {
130+
.container {
131+
background-color: #1f1f1f;
132+
box-shadow: 0 0 0 1px #ffffff26;
133+
}
134+
.headline-label {
135+
color: #fff;
136+
}
137+
.chip {
138+
background-color: #2c2c2c;
139+
border-color: #3c4043;
140+
color: #fff;
141+
text-decoration: none;
142+
}
143+
.chip:hover {
144+
background-color: #353536;
145+
}
146+
.chip:focus {
147+
background-color: #353536;
148+
}
149+
.chip:active {
150+
background-color: #464849;
151+
border-color: #53575b;
152+
}
153+
.logo-light {
154+
display: none;
155+
}
156+
.gradient {
157+
background: linear-gradient(90deg, #1f1f1f 15%, #1f1f1f00 100%);
158+
}
159+
}
160+
</style>
161+
<div class="container">
162+
<div class="headline">
163+
<svg class="logo-light" width="18" height="18" viewBox="9 9 35 35" fill="none" xmlns="http://www.w3.org/2000/svg">
164+
<path fill-rule="evenodd" clip-rule="evenodd" d="M42.8622 27.0064C42.8622 25.7839 42.7525 24.6084 42.5487 23.4799H26.3109V30.1568H35.5897C35.1821 32.3041 33.9596 34.1222 32.1258 35.3448V39.6864H37.7213C40.9814 36.677 42.8622 32.2571 42.8622 27.0064V27.0064Z" fill="#4285F4"/>
165+
<path fill-rule="evenodd" clip-rule="evenodd" d="M26.3109 43.8555C30.9659 43.8555 34.8687 42.3195 37.7213 39.6863L32.1258 35.3447C30.5898 36.3792 28.6306 37.0061 26.3109 37.0061C21.8282 37.0061 18.0195 33.9811 16.6559 29.906H10.9194V34.3573C13.7563 39.9841 19.5712 43.8555 26.3109 43.8555V43.8555Z" fill="#34A853"/>
166+
<path fill-rule="evenodd" clip-rule="evenodd" d="M16.6559 29.8904C16.3111 28.8559 16.1074 27.7588 16.1074 26.6146C16.1074 25.4704 16.3111 24.3733 16.6559 23.3388V18.8875H10.9194C9.74388 21.2072 9.06992 23.8247 9.06992 26.6146C9.06992 29.4045 9.74388 32.022 10.9194 34.3417L15.3864 30.8621L16.6559 29.8904V29.8904Z" fill="#FBBC05"/>
167+
<path fill-rule="evenodd" clip-rule="evenodd" d="M26.3109 16.2386C28.85 16.2386 31.107 17.1164 32.9095 18.8091L37.8466 13.8719C34.853 11.082 30.9659 9.3736 26.3109 9.3736C19.5712 9.3736 13.7563 13.245 10.9194 18.8875L16.6559 23.3388C18.0195 19.2636 21.8282 16.2386 26.3109 16.2386V16.2386Z" fill="#EA4335"/>
168+
</svg>
169+
<svg class="logo-dark" width="18" height="18" viewBox="0 0 48 48" xmlns="http://www.w3.org/2000/svg">
170+
<circle cx="24" cy="23" fill="#FFF" r="22"/>
171+
<path d="M33.76 34.26c2.75-2.56 4.49-6.37 4.49-11.26 0-.89-.08-1.84-.29-3H24.01v5.99h8.03c-.4 2.02-1.5 3.56-3.07 4.56v.75l3.91 2.97h.88z" fill="#4285F4"/>
172+
<path d="M15.58 25.77A8.845 8.845 0 0 0 24 31.86c1.92 0 3.62-.46 4.97-1.31l4.79 3.71C31.14 36.7 27.65 38 24 38c-5.93 0-11.01-3.4-13.45-8.36l.17-1.01 4.06-2.85h.8z" fill="#34A853"/>
173+
<path d="M15.59 20.21a8.864 8.864 0 0 0 0 5.58l-5.03 3.86c-.98-2-1.53-4.25-1.53-6.64 0-2.39.55-4.64 1.53-6.64l1-.22 3.81 2.98.22 1.08z" fill="#FBBC05"/>
174+
<path d="M24 14.14c2.11 0 4.02.75 5.52 1.98l4.36-4.36C31.22 9.43 27.81 8 24 8c-5.93 0-11.01 3.4-13.45 8.36l5.03 3.85A8.86 8.86 0 0 1 24 14.14z" fill="#EA4335"/>
175+
</svg>
176+
<div class="gradient-container"><div class="gradient"></div></div>
177+
</div>
178+
<div class="carousel">
179+
<a class="chip" href="https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9wXF8ek0JqgT0M4R24jWmUjYcW0KmEax-fTkbZfncRMDKieXEJSp-nWfV-utSapLCm3Pm5Kb51BX3oPW77gVM6iZwubNdKrk6KxPd2QTTicnu00DeeRp0jY7UOGJ_Awr1opc0QIKCVmSBNAn8-UldCuJjA-Fzn4NI1pDbzFI7kFhHreLWIY5tHh3ukVQo3CZnq6OpSVhb2Nfs">current date Utrecht</a>
180+
</div>
181+
</div>
182+
webSearchQueries:
183+
- current date Utrecht
184+
modelVersion: gemini-2.0-flash
185+
responseId: LN42aKzNCaqvgLUP1Lz-8As
186+
usageMetadata:
187+
candidatesTokenCount: 19
188+
candidatesTokensDetails:
189+
- modality: TEXT
190+
tokenCount: 19
191+
promptTokenCount: 13
192+
promptTokensDetails:
193+
- modality: TEXT
194+
tokenCount: 13
195+
totalTokenCount: 32
196+
status:
197+
code: 200
198+
message: OK
199+
version: 1

‎tests/models/test_google.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from typing_extensions import TypedDict
1414

1515
from pydantic_ai.agent import Agent
16+
from pydantic_ai.builtin_tools import CodeExecutionTool, WebSearchTool
1617
from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior
1718
from pydantic_ai.messages import (
1819
BinaryContent,
@@ -529,3 +530,19 @@ async def test_google_model_safety_settings(allow_model_requests: None, google_p
529530

530531
with pytest.raises(UnexpectedModelBehavior, match='Safety settings triggered'):
531532
await agent.run('Tell me a joke about a Brazilians.')
533+
534+
535+
async def test_google_model_web_search_tool(allow_model_requests: None, google_provider: GoogleProvider):
536+
m = GoogleModel('gemini-2.0-flash', provider=google_provider)
537+
agent = Agent(m, system_prompt='You are a helpful chatbot.', builtin_tools=[WebSearchTool()])
538+
539+
result = await agent.run('What day is today in Utrecht?')
540+
assert result.output == snapshot('Today is Wednesday, May 28, 2025, in Utrecht.\n')
541+
542+
543+
async def test_google_model_code_execution_tool(allow_model_requests: None, google_provider: GoogleProvider):
544+
m = GoogleModel('gemini-2.0-flash', provider=google_provider)
545+
agent = Agent(m, system_prompt='You are a helpful chatbot.', builtin_tools=[CodeExecutionTool()])
546+
547+
result = await agent.run('What day is today in Utrecht?')
548+
assert result.output == snapshot('Today is Wednesday, May 28, 2025, in Utrecht.\n')

0 commit comments

Comments
 (0)