Skip to content

Commit f5a5b73

Browse files
DouweMDouweM
authored
Add http_client option to GoogleProvider, support api_key for Vertex AI, and use Pydantic AI's cached httpx client by default (#3217)
1 parent f5cc9e3 commit f5a5b73

File tree

6 files changed

+165
-109
lines changed

6 files changed

+165
-109
lines changed

‎docs/models/google.md‎

Lines changed: 76 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@ To use `GoogleModel`, you need to either install `pydantic-ai`, or install `pyda
1111
pip/uv-add "pydantic-ai-slim[google]"
1212
```
1313

14-
---
1514

1615
## Configuration
1716

@@ -27,54 +26,70 @@ Once you have the API key, set it as an environment variable:
2726
export GOOGLE_API_KEY=your-api-key
2827
```
2928

30-
You can then use `GoogleModel` by explicitly creating a provider:
29+
You can then use `GoogleModel` by name (where GLA stands for Generative Language API):
30+
31+
```python
32+
from pydantic_ai import Agent
33+
34+
agent = Agent('google-gla:gemini-2.5-pro')
35+
...
36+
```
37+
38+
Or you can explicitly create the provider:
3139

3240
```python
3341
from pydantic_ai import Agent
3442
from pydantic_ai.models.google import GoogleModel
3543
from pydantic_ai.providers.google import GoogleProvider
3644

3745
provider = GoogleProvider(api_key='your-api-key')
38-
model = GoogleModel('gemini-1.5-flash', provider=provider)
46+
model = GoogleModel('gemini-2.5-pro', provider=provider)
3947
agent = Agent(model)
4048
...
4149
```
4250

4351
### Vertex AI (Enterprise/Cloud)
4452

45-
If you are an enterprise user, you can use the `google-vertex` provider with `GoogleModel` to access Gemini via Vertex AI.
53+
If you are an enterprise user, you can also use `GoogleModel` to access Gemini via Vertex AI.
4654

4755
This interface has a number of advantages over the Generative Language API:
4856

4957
1. The VertexAI API comes with more enterprise readiness guarantees.
50-
2. You can [purchase provisioned throughput](https://cloud.google.com/vertex-ai/generative-ai/docs/provisioned-throughput#purchase-provisioned-throughput) with VertexAI to guarantee capacity.
58+
2. You can [purchase provisioned throughput](https://cloud.google.com/vertex-ai/generative-ai/docs/provisioned-throughput#purchase-provisioned-throughput) with Vertex AI to guarantee capacity.
5159
3. If you're running Pydantic AI inside GCP, you don't need to set up authentication, it should "just work".
5260
4. You can decide which region to use, which might be important from a regulatory perspective, and might improve latency.
5361

54-
The big disadvantage is that for local development you may need to create and configure a "service account", which can be challenging to get right.
62+
You can authenticate using [application default credentials](https://cloud.google.com/docs/authentication/application-default-credentials), a service account, or an [API key](https://cloud.google.com/vertex-ai/generative-ai/docs/start/api-keys?usertype=expressmode).
5563

56-
Whichever way you authenticate, you'll need to have VertexAI enabled in your GCP account.
57-
58-
To use Vertex AI, you may need to set up [application default credentials](https://cloud.google.com/docs/authentication/application-default-credentials) or use a service account. You can also specify the region.
64+
Whichever way you authenticate, you'll need to have Vertex AI enabled in your GCP account.
5965

6066
#### Application Default Credentials
6167

62-
If you have the [`gcloud` CLI](https://cloud.google.com/sdk/gcloud) installed and configured, you can use:
68+
If you have the [`gcloud` CLI](https://cloud.google.com/sdk/gcloud) installed and configured, you can use `GoogleProvider` in Vertex AI mode by name:
69+
70+
```python {test="ci_only"}
71+
from pydantic_ai import Agent
72+
73+
agent = Agent('google-vertex:gemini-2.5-pro')
74+
...
75+
```
76+
77+
Or you can explicitly create the provider and model:
6378

6479
```python {test="ci_only"}
6580
from pydantic_ai import Agent
6681
from pydantic_ai.models.google import GoogleModel
6782
from pydantic_ai.providers.google import GoogleProvider
6883

6984
provider = GoogleProvider(vertexai=True)
70-
model = GoogleModel('gemini-1.5-flash', provider=provider)
85+
model = GoogleModel('gemini-2.5-pro', provider=provider)
7186
agent = Agent(model)
7287
...
7388
```
7489

7590
#### Service Account
7691

77-
To use a service account JSON file:
92+
To use a service account JSON file, explicitly create the provider and model:
7893

7994
```python {title="google_model_service_account.py" test="skip"}
8095
from google.oauth2 import service_account
@@ -93,43 +108,66 @@ agent = Agent(model)
93108
...
94109
```
95110

96-
#### Customizing Location
111+
#### API Key
112+
113+
To use Vertex AI with an API key, [create a key](https://cloud.google.com/vertex-ai/generative-ai/docs/start/api-keys?usertype=expressmode) and set it as an environment variable:
114+
115+
```bash
116+
export GOOGLE_API_KEY=your-api-key
117+
```
118+
119+
You can then use `GoogleModel` in Vertex AI mode by name:
120+
121+
```python {test="ci_only"}
122+
from pydantic_ai import Agent
123+
124+
agent = Agent('google-vertex:gemini-2.5-pro')
125+
...
126+
```
97127

98-
You can specify the location when using Vertex AI:
128+
Or you can explicitly create the provider and model:
129+
130+
```python {test="skip"}
131+
from pydantic_ai import Agent
132+
from pydantic_ai.models.google import GoogleModel
133+
from pydantic_ai.providers.google import GoogleProvider
134+
135+
provider = GoogleProvider(vertexai=True, api_key='your-api-key')
136+
model = GoogleModel('gemini-2.5-pro', provider=provider)
137+
agent = Agent(model)
138+
...
139+
```
140+
141+
#### Customizing Location or Project
142+
143+
You can specify the location and/or projectwhen using Vertex AI:
99144

100145
```python {title="google_model_location.py" test="skip"}
101146
from pydantic_ai import Agent
102147
from pydantic_ai.models.google import GoogleModel
103148
from pydantic_ai.providers.google import GoogleProvider
104149

105-
provider = GoogleProvider(vertexai=True, location='asia-east1')
106-
model = GoogleModel('gemini-1.5-flash', provider=provider)
150+
provider = GoogleProvider(vertexai=True, location='asia-east1', project='your-gcp-project-id')
151+
model = GoogleModel('gemini-2.5-pro', provider=provider)
107152
agent = Agent(model)
108153
...
109154
```
110155

111-
#### Customizing Model
156+
#### Model Garden
112157

113-
You can access models from the [Model Garden](https://cloud.google.com/model-garden?hl=en) that support the generateContent API and are available under your GCP project, including but not limited to Gemini, using one of the following `model_name` patterns:
158+
You can access models from the [Model Garden](https://cloud.google.com/model-garden?hl=en) that support the `generateContent` API and are available under your GCP project, including but not limited to Gemini, using one of the following `model_name` patterns:
114159

115160
- `{model_id}` for Gemini models
116161
- `{publisher}/{model_id}`
117162
- `publishers/{publisher}/models/{model_id}`
118163
- `projects/{project}/locations/{location}/publishers/{publisher}/models/{model_id}`
119164

120165
```python {test="skip"}
121-
from google.oauth2 import service_account
122-
123166
from pydantic_ai import Agent
124167
from pydantic_ai.models.google import GoogleModel
125168
from pydantic_ai.providers.google import GoogleProvider
126169

127-
credentials = service_account.Credentials.from_service_account_file(
128-
'path/to/service-account.json',
129-
scopes=['https://www.googleapis.com/auth/cloud-platform'],
130-
)
131170
provider = GoogleProvider(
132-
credentials=credentials,
133171
project='your-gcp-project-id',
134172
location='us-central1', # the region where the model is available
135173
)
@@ -138,31 +176,32 @@ agent = Agent(model)
138176
...
139177
```
140178

141-
## Provider Argument
142-
143-
You can supply a custom `GoogleProvider` instance using the `provider` argument to configure advanced client options, such as setting a custom `base_url`.
179+
## Custom HTTP Client
144180

145-
This is useful if you're using a custom-compatible endpoint with the Google Generative Language API.
181+
You can customize the `GoogleProvider` with a custom `httpx.AsyncClient`:
146182

147183
```python
148-
from google.genai import Client
149-
from google.genai.types import HttpOptions
184+
from httpx import AsyncClient
150185

151186
from pydantic_ai import Agent
152187
from pydantic_ai.models.google import GoogleModel
153188
from pydantic_ai.providers.google import GoogleProvider
154189

155-
client = Client(
156-
api_key='gemini-custom-api-key',
157-
http_options=HttpOptions(base_url='gemini-custom-base-url'),
190+
custom_http_client = AsyncClient(timeout=30)
191+
model = GoogleModel(
192+
'gemini-2.5-pro',
193+
provider=GoogleProvider(api_key='your-api-key', http_client=custom_http_client),
158194
)
159-
provider = GoogleProvider(client=client)
160-
model = GoogleModel('gemini-1.5-flash', provider=provider)
161195
agent = Agent(model)
162196
...
163197
```
164198

165-
## Model Settings
199+
200+
## Document, Image, Audio, and Video Input
201+
202+
`GoogleModel` supports multi-modal input, including documents, images, audio, and video. See the [input documentation](../input.md) for details and examples.
203+
204+
## Model settings
166205

167206
You can customize model behavior using [`GoogleModelSettings`][pydantic_ai.models.google.GoogleModelSettings]:
168207

@@ -188,16 +227,6 @@ agent = Agent(model, model_settings=settings)
188227
...
189228
```
190229

191-
See the [Gemini API docs](https://ai.google.dev/gemini-api/docs/safety-settings) for more on safety settings, and [thinking config](https://ai.google.dev/gemini-api/docs/thinking).
192-
193-
## Document, Image, Audio, and Video Input
194-
195-
`GoogleModel` supports multi-modal input, including documents, images, audio, and video. See the [input documentation](../input.md) for details and examples.
196-
197-
## Model settings
198-
199-
You can use the [`GoogleModelSettings`][pydantic_ai.models.google.GoogleModelSettings] class to customize the model request.
200-
201230
### Disable thinking
202231

203232
You can disable thinking by setting the `thinking_budget` to `0` on the `google_thinking_config`:

‎pydantic_ai_slim/pydantic_ai/providers/gateway.py‎

Lines changed: 21 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,13 @@
33
from __future__ import annotations as _annotations
44

55
import os
6+
from collections.abc import Awaitable, Callable
67
from typing import TYPE_CHECKING, Any, Literal, overload
78

89
import httpx
910

1011
from pydantic_ai.exceptions import UserError
11-
from pydantic_ai.models import cached_async_http_client, get_user_agent
12+
from pydantic_ai.models import cached_async_http_client
1213

1314
if TYPE_CHECKING:
1415
from botocore.client import BaseClient
@@ -109,7 +110,7 @@ def gateway_provider(
109110

110111
base_url = base_url or os.getenv('PYDANTIC_AI_GATEWAY_BASE_URL', GATEWAY_BASE_URL)
111112
http_client = http_client or cached_async_http_client(provider=f'gateway/{upstream_provider}')
112-
http_client.event_hooks = {'request': [_request_hook]}
113+
http_client.event_hooks = {'request': [_request_hook(api_key)]}
113114

114115
if upstream_provider in ('openai', 'openai-chat', 'openai-responses'):
115116
from .openai import OpenAIProvider
@@ -140,41 +141,37 @@ def gateway_provider(
140141
region_name='pydantic-ai-gateway', # Fake region name to avoid NoRegionError
141142
)
142143
elif upstream_provider == 'google-vertex':
143-
from google.genai import Client as GoogleClient
144-
145144
from .google import GoogleProvider
146145

147146
return GoogleProvider(
148-
client=GoogleClient(
149-
vertexai=True,
150-
api_key='unset',
151-
http_options={
152-
'base_url': _merge_url_path(base_url, 'google-vertex'),
153-
'headers': {'User-Agent': get_user_agent(), 'Authorization': api_key},
154-
# TODO(Marcelo): Until https://github.com/googleapis/python-genai/issues/1357 is solved.
155-
'async_client_args': {
156-
'transport': httpx.AsyncHTTPTransport(),
157-
'event_hooks': {'request': [_request_hook]},
158-
},
159-
},
160-
)
147+
vertexai=True,
148+
api_key=api_key,
149+
base_url=_merge_url_path(base_url, 'google-vertex'),
150+
http_client=http_client,
161151
)
162152
else:
163153
raise UserError(f'Unknown upstream provider: {upstream_provider}')
164154

165155

166-
async def _request_hook(request: httpx.Request) -> httpx.Request:
156+
def _request_hook(api_key: str) -> Callable[[httpx.Request], Awaitable[httpx.Request]]:
167157
"""Request hook for the gateway provider.
168158
169-
It adds the `"traceparent"` header to the request.
159+
It adds the `"traceparent"` and `"Authorization"` headers to the request.
170160
"""
171-
from opentelemetry.propagate import inject
172161

173-
headers: dict[str, Any] = {}
174-
inject(headers)
175-
request.headers.update(headers)
162+
async def _hook(request: httpx.Request) -> httpx.Request:
163+
from opentelemetry.propagate import inject
164+
165+
headers: dict[str, Any] = {}
166+
inject(headers)
167+
request.headers.update(headers)
168+
169+
if 'Authorization' not in request.headers:
170+
request.headers['Authorization'] = f'Bearer {api_key}'
171+
172+
return request
176173

177-
return request
174+
return _hook
178175

179176

180177
def _merge_url_path(base_url: str, path: str) -> str:

0 commit comments

Comments
 (0)