Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
6 changes: 3 additions & 3 deletions docs/models/google.md
Original file line number Diff line number Diff line change
Expand Up @@ -214,22 +214,22 @@ from pydantic_ai.models.google import GoogleModel, GoogleModelSettings
settings = GoogleModelSettings(
temperature=0.2,
max_tokens=1024,
google_thinking_config={'thinking_budget': 2048},
google_thinking_config={'thinking_level': 'low'},
google_safety_settings=[
{
'category': HarmCategory.HARM_CATEGORY_HATE_SPEECH,
'threshold': HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
}
]
)
model = GoogleModel('gemini-2.5-flash')
model = GoogleModel('gemini-2.5-pro')
agent = Agent(model, model_settings=settings)
...
```

### Disable thinking

You can disable thinking by setting the `thinking_budget` to `0` on the `google_thinking_config`:
On models older than Gemini 2.5 Pro, you can disable thinking by setting the `thinking_budget` to `0` on the `google_thinking_config`:

```python
from pydantic_ai import Agent
Expand Down
16 changes: 10 additions & 6 deletions pydantic_ai_slim/pydantic_ai/models/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -145,24 +145,28 @@
'cohere:command-r7b-12-2024',
'deepseek:deepseek-chat',
'deepseek:deepseek-reasoner',
'google-gla:gemini-flash-latest',
'google-gla:gemini-flash-lite-latest',
'google-gla:gemini-2.0-flash',
'google-gla:gemini-2.0-flash-lite',
'google-gla:gemini-2.5-flash',
'google-gla:gemini-2.5-flash-preview-09-2025',
'google-gla:gemini-2.5-flash-image',
'google-gla:gemini-2.5-flash-lite',
'google-gla:gemini-2.5-flash-lite-preview-09-2025',
'google-gla:gemini-2.5-flash-preview-09-2025',
'google-gla:gemini-2.5-pro',
'google-gla:gemini-flash-latest',
'google-gla:gemini-flash-lite-latest',
'google-gla:gemini-3-pro-preview',
'google-vertex:gemini-flash-latest',
'google-vertex:gemini-flash-lite-latest',
'google-vertex:gemini-2.0-flash',
'google-vertex:gemini-2.0-flash-lite',
'google-vertex:gemini-2.5-flash',
'google-vertex:gemini-2.5-flash-preview-09-2025',
'google-vertex:gemini-2.5-flash-image',
'google-vertex:gemini-2.5-flash-lite',
'google-vertex:gemini-2.5-flash-lite-preview-09-2025',
'google-vertex:gemini-2.5-flash-preview-09-2025',
'google-vertex:gemini-2.5-pro',
'google-vertex:gemini-flash-latest',
'google-vertex:gemini-flash-lite-latest',
'google-vertex:gemini-3-pro-preview',
'grok:grok-2-image-1212',
'grok:grok-2-vision-1212',
'grok:grok-3',
Expand Down
48 changes: 34 additions & 14 deletions pydantic_ai_slim/pydantic_ai/models/google.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
VideoUrl,
)
from ..profiles import ModelProfileSpec
from ..profiles.google import GoogleModelProfile
from ..providers import Provider, infer_provider
from ..settings import ModelSettings
from ..tools import ToolDefinition
Expand Down Expand Up @@ -93,15 +94,17 @@
) from _import_error

LatestGoogleModelNames = Literal[
'gemini-flash-latest',
'gemini-flash-lite-latest',
'gemini-2.0-flash',
'gemini-2.0-flash-lite',
'gemini-2.5-flash',
'gemini-2.5-flash-preview-09-2025',
'gemini-flash-latest',
'gemini-2.5-flash-image',
'gemini-2.5-flash-lite',
'gemini-2.5-flash-lite-preview-09-2025',
'gemini-flash-lite-latest',
'gemini-2.5-pro',
'gemini-3-pro-preview',
]
"""Latest Gemini models."""

Expand Down Expand Up @@ -228,12 +231,17 @@ def system(self) -> str:
def prepare_request(
self, model_settings: ModelSettings | None, model_request_parameters: ModelRequestParameters
) -> tuple[ModelSettings | None, ModelRequestParameters]:
supports_native_output_with_builtin_tools = GoogleModelProfile.from_profile(
self.profile
).google_supports_native_output_with_builtin_tools
if model_request_parameters.builtin_tools and model_request_parameters.output_tools:
if model_request_parameters.output_mode == 'auto':
model_request_parameters = replace(model_request_parameters, output_mode='prompted')
output_mode = 'native' if supports_native_output_with_builtin_tools else 'prompted'
model_request_parameters = replace(model_request_parameters, output_mode=output_mode)
else:
output_mode = 'NativeOutput' if supports_native_output_with_builtin_tools else 'PromptedOutput'
raise UserError(
'Google does not support output tools and built-in tools at the same time. Use `output_type=PromptedOutput(...)` instead.'
f'Google does not support output tools and built-in tools at the same time. Use `output_type={output_mode}(...)` instead.'
)
return super().prepare_request(model_settings, model_request_parameters)

Expand Down Expand Up @@ -418,9 +426,9 @@ async def _build_content_and_config(
response_mime_type = None
response_schema = None
if model_request_parameters.output_mode == 'native':
if tools:
if model_request_parameters.function_tools:
raise UserError(
'Google does not support `NativeOutput` and tools at the same time. Use `output_type=ToolOutput(...)` instead.'
'Google does not support `NativeOutput` and function tools at the same time. Use `output_type=ToolOutput(...)` instead.'
)
response_mime_type = 'application/json'
output_object = model_request_parameters.output_object
Expand Down Expand Up @@ -685,23 +693,25 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:

for part in parts:
if part.thought_signature:
# Per https://ai.google.dev/gemini-api/docs/function-calling?example=meeting#thought-signatures:
# - Always send the thought_signature back to the model inside its original Part.
# - Don't merge a Part containing a signature with one that does not. This breaks the positional context of the thought.
# - Don't combine two Parts that both contain signatures, as the signature strings cannot be merged.

signature = base64.b64encode(part.thought_signature).decode('utf-8')
# Attach signature to most recent thinking part, if there was one
yield self._parts_manager.handle_thinking_delta(
vendor_part_id='thinking',
vendor_part_id=None,
signature=signature,
provider_name=self.provider_name,
)

if part.text is not None:
if len(part.text) > 0:
if part.thought:
yield self._parts_manager.handle_thinking_delta(
vendor_part_id='thinking', content=part.text
)
yield self._parts_manager.handle_thinking_delta(vendor_part_id=None, content=part.text)
else:
maybe_event = self._parts_manager.handle_text_delta(
vendor_part_id='content', content=part.text
)
maybe_event = self._parts_manager.handle_text_delta(vendor_part_id=None, content=part.text)
if maybe_event is not None: # pragma: no branch
yield maybe_event
elif part.function_call:
Expand Down Expand Up @@ -760,6 +770,7 @@ def timestamp(self) -> datetime:
def _content_model_response(m: ModelResponse, provider_name: str) -> ContentDict: # noqa: C901
parts: list[PartDict] = []
thought_signature: bytes | None = None
function_call_requires_signature: bool = True
for item in m.parts:
part: PartDict = {}
if thought_signature:
Expand All @@ -769,6 +780,15 @@ def _content_model_response(m: ModelResponse, provider_name: str) -> ContentDict
if isinstance(item, ToolCallPart):
function_call = FunctionCallDict(name=item.tool_name, args=item.args_as_dict(), id=item.tool_call_id)
part['function_call'] = function_call
if function_call_requires_signature and not part.get('thought_signature'):
# Per https://ai.google.dev/gemini-api/docs/gemini-3?thinking=high#migrating_from_other_models:
# > If you are transferring a conversation trace from another model (e.g., Gemini 2.5) or injecting
# > a custom function call that was not generated by Gemini 3, you will not have a valid signature.
# > To bypass strict validation in these specific scenarios, populate the field with this specific
# > dummy string: "thoughtSignature": "context_engineering_is_the_way_to_go"
part['thought_signature'] = b'context_engineering_is_the_way_to_go'
# Only the first function call requires a signature
function_call_requires_signature = False
elif isinstance(item, TextPart):
part['text'] = item.content
elif isinstance(item, ThinkingPart):
Expand Down Expand Up @@ -881,7 +901,7 @@ def _function_declaration_from_tool(tool: ToolDefinition) -> FunctionDeclaration
f = FunctionDeclarationDict(
name=tool.name,
description=tool.description or '',
parameters=json_schema, # type: ignore
parameters_json_schema=json_schema,
)
return f

Expand Down
18 changes: 17 additions & 1 deletion pydantic_ai_slim/pydantic_ai/profiles/google.py
Original file line number Diff line number Diff line change
@@ -1,18 +1,34 @@
from __future__ import annotations as _annotations

from dataclasses import dataclass

from .._json_schema import JsonSchema, JsonSchemaTransformer
from . import ModelProfile


@dataclass(kw_only=True)
class GoogleModelProfile(ModelProfile):
"""Profile for models used with `GoogleModel`.
ALL FIELDS MUST BE `google_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.
"""

google_supports_native_output_with_builtin_tools: bool = False
"""Whether the model supports native output with builtin tools.
See https://ai.google.dev/gemini-api/docs/structured-output?example=recipe#structured_outputs_with_tools"""


def google_model_profile(model_name: str) -> ModelProfile | None:
"""Get the model profile for a Google model."""
is_image_model = 'image' in model_name
return ModelProfile(
is_3_or_newer = 'gemini-3' in model_name
return GoogleModelProfile(
json_schema_transformer=GoogleJsonSchemaTransformer,
supports_image_output=is_image_model,
supports_json_schema_output=not is_image_model,
supports_json_object_output=not is_image_model,
supports_tools=not is_image_model,
google_supports_native_output_with_builtin_tools=is_3_or_newer,
)


Expand Down
2 changes: 1 addition & 1 deletion pydantic_ai_slim/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ logfire = ["logfire[httpx]>=3.14.1"]
openai = ["openai>=1.107.2"]
cohere = ["cohere>=5.18.0; platform_system != 'Emscripten'"]
vertexai = ["google-auth>=2.36.0", "requests>=2.32.2"]
google = ["google-genai>=1.50.1"]
google = ["google-genai>=1.51.0"]
anthropic = ["anthropic>=0.70.0"]
groq = ["groq>=0.25.0"]
mistral = ["mistralai>=1.9.10"]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ interactions:
connection:
- keep-alive
content-length:
- '526'
- '560'
content-type:
- application/json
host:
Expand All @@ -19,10 +19,13 @@ interactions:
- parts:
- text: What is the largest city in Mexico?
role: user
generationConfig: {}
generationConfig:
responseModalities:
- TEXT
systemInstruction:
parts:
- text: |-
- text: |2
Always respond with a JSON object that's compatible with this schema:
{"properties": {"city": {"type": "string"}, "country": {"type": "string"}}, "required": ["city", "country"], "title": "CityLocation", "type": "object"}
Expand All @@ -41,7 +44,7 @@ interactions:
content-type:
- application/json; charset=UTF-8
server-timing:
- gfet4t7; dur=780
- gfet4t7; dur=873
transfer-encoding:
- chunked
vary:
Expand All @@ -58,15 +61,15 @@ interactions:
groundingMetadata: {}
index: 0
modelVersion: gemini-2.5-flash
responseId: 6Xq3aPnXNtqKqtsP8ZuDyAc
responseId: 0f4caf7iM-mgz7IPtLeryAc
usageMetadata:
candidatesTokenCount: 13
promptTokenCount: 83
promptTokenCount: 85
promptTokensDetails:
- modality: TEXT
tokenCount: 83
thoughtsTokenCount: 33
totalTokenCount: 129
tokenCount: 85
thoughtsTokenCount: 32
totalTokenCount: 130
status:
code: 200
message: OK
Expand Down
Loading