Skip to content

Commit 757d409

Browse files
conradleeclaudeDouweM
authored
Support Gemini enhanced JSON Schema features (#3357)
Co-authored-by: Claude <[email protected]> Co-authored-by: Douwe Maan <[email protected]>
1 parent e7b2f82 commit 757d409

20 files changed

+1698
-216
lines changed

‎pydantic_ai_slim/pydantic_ai/_json_schema.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ def __init__(
2525
*,
2626
strict: bool | None = None,
2727
prefer_inlined_defs: bool = False,
28-
simplify_nullable_unions: bool = False,
28+
simplify_nullable_unions: bool = False, # TODO (v2): Remove this, no longer used
2929
):
3030
self.schema = schema
3131

@@ -146,10 +146,9 @@ def _handle_union(self, schema: JsonSchema, union_kind: Literal['anyOf', 'oneOf'
146146

147147
handled = [self._handle(member) for member in members]
148148

149-
# convert nullable unions to nullable types
149+
# TODO (v2): Remove this feature, no longer used
150150
if self.simplify_nullable_unions:
151151
handled = self._simplify_nullable_union(handled)
152-
153152
if len(handled) == 1:
154153
# In this case, no need to retain the union
155154
return handled[0] | schema
@@ -161,7 +160,7 @@ def _handle_union(self, schema: JsonSchema, union_kind: Literal['anyOf', 'oneOf'
161160

162161
@staticmethod
163162
def _simplify_nullable_union(cases: list[JsonSchema]) -> list[JsonSchema]:
164-
# TODO: Should we move this to relevant subclasses? Or is it worth keeping here to make reuse easier?
163+
# TODO (v2): Remove this method, no longer used
165164
if len(cases) == 2 and {'type': 'null'} in cases:
166165
# Find the non-null schema
167166
non_null_schema = next(

‎pydantic_ai_slim/pydantic_ai/models/google.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -292,7 +292,7 @@ async def count_tokens(
292292
thinking_config=generation_config.get('thinking_config'),
293293
media_resolution=generation_config.get('media_resolution'),
294294
response_mime_type=generation_config.get('response_mime_type'),
295-
response_schema=generation_config.get('response_schema'),
295+
response_json_schema=generation_config.get('response_json_schema'),
296296
),
297297
)
298298

@@ -456,7 +456,7 @@ async def _build_content_and_config(
456456
tools=cast(ToolListUnionDict, tools),
457457
tool_config=tool_config,
458458
response_mime_type=response_mime_type,
459-
response_schema=response_schema,
459+
response_json_schema=response_schema,
460460
response_modalities=modalities,
461461
)
462462
return contents, config

‎pydantic_ai_slim/pydantic_ai/profiles/google.py‎

Lines changed: 5 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,5 @@
11
from __future__ import annotations as _annotations
22

3-
import warnings
4-
5-
from pydantic_ai.exceptions import UserError
6-
73
from .._json_schema import JsonSchema, JsonSchemaTransformer
84
from . import ModelProfile
95

@@ -23,84 +19,28 @@ def google_model_profile(model_name: str) -> ModelProfile | None:
2319
class GoogleJsonSchemaTransformer(JsonSchemaTransformer):
2420
"""Transforms the JSON Schema from Pydantic to be suitable for Gemini.
2521
26-
Gemini which [supports](https://ai.google.dev/gemini-api/docs/function-calling#function_declarations)
27-
a subset of OpenAPI v3.0.3.
28-
29-
Specifically:
30-
* gemini doesn't allow the `title` keyword to be set
31-
* gemini doesn't allow `$defs` — we need to inline the definitions where possible
22+
Gemini supports [a subset of OpenAPI v3.0.3](https://ai.google.dev/gemini-api/docs/function-calling#function_declarations).
3223
"""
3324

34-
def __init__(self, schema: JsonSchema, *, strict: bool | None = None):
35-
super().__init__(schema, strict=strict, prefer_inlined_defs=True, simplify_nullable_unions=True)
36-
3725
def transform(self, schema: JsonSchema) -> JsonSchema:
38-
# Note: we need to remove `additionalProperties: False` since it is currently mishandled by Gemini
39-
additional_properties = schema.pop(
40-
'additionalProperties', None
41-
) # don't pop yet so it's included in the warning
42-
if additional_properties:
43-
original_schema = {**schema, 'additionalProperties': additional_properties}
44-
warnings.warn(
45-
'`additionalProperties` is not supported by Gemini; it will be removed from the tool JSON schema.'
46-
f' Full schema: {self.schema}\n\n'
47-
f'Source of additionalProperties within the full schema: {original_schema}\n\n'
48-
'If this came from a field with a type like `dict[str, MyType]`, that field will always be empty.\n\n'
49-
"If Google's APIs are updated to support this properly, please create an issue on the Pydantic AI GitHub"
50-
' and we will fix this behavior.',
51-
UserWarning,
52-
)
53-
54-
schema.pop('title', None)
26+
# Remove properties not supported by Gemini
5527
schema.pop('$schema', None)
5628
if (const := schema.pop('const', None)) is not None:
5729
# Gemini doesn't support const, but it does support enum with a single value
5830
schema['enum'] = [const]
5931
schema.pop('discriminator', None)
6032
schema.pop('examples', None)
6133

62-
# TODO: Should we use the trick from pydantic_ai.models.openai._OpenAIJsonSchema
63-
# where we add notes about these properties to the field description?
64-
schema.pop('exclusiveMaximum', None)
65-
schema.pop('exclusiveMinimum', None)
66-
67-
# Gemini only supports string enums, so we need to convert any enum values to strings.
68-
# Pydantic will take care of transforming the transformed string values to the correct type.
69-
if enum := schema.get('enum'):
70-
schema['type'] = 'string'
71-
schema['enum'] = [str(val) for val in enum]
72-
7334
type_ = schema.get('type')
74-
if 'oneOf' in schema and 'type' not in schema: # pragma: no cover
75-
# This gets hit when we have a discriminated union
76-
# Gemini returns an API error in this case even though it says in its error message it shouldn't...
77-
# Changing the oneOf to an anyOf prevents the API error and I think is functionally equivalent
78-
schema['anyOf'] = schema.pop('oneOf')
79-
8035
if type_ == 'string' and (fmt := schema.pop('format', None)):
8136
description = schema.get('description')
8237
if description:
8338
schema['description'] = f'{description} (format: {fmt})'
8439
else:
8540
schema['description'] = f'Format: {fmt}'
8641

87-
if '$ref' in schema:
88-
raise UserError(f'Recursive `$ref`s in JSON Schema are not supported by Gemini: {schema["$ref"]}')
89-
90-
if 'prefixItems' in schema:
91-
# prefixItems is not currently supported in Gemini, so we convert it to items for best compatibility
92-
prefix_items = schema.pop('prefixItems')
93-
items = schema.get('items')
94-
unique_items = [items] if items is not None else []
95-
for item in prefix_items:
96-
if item not in unique_items:
97-
unique_items.append(item)
98-
if len(unique_items) > 1: # pragma: no cover
99-
schema['items'] = {'anyOf': unique_items}
100-
elif len(unique_items) == 1: # pragma: no branch
101-
schema['items'] = unique_items[0]
102-
schema.setdefault('minItems', len(prefix_items))
103-
if items is None: # pragma: no branch
104-
schema.setdefault('maxItems', len(prefix_items))
42+
# Note: exclusiveMinimum/exclusiveMaximum are NOT yet supported
43+
schema.pop('exclusiveMinimum', None)
44+
schema.pop('exclusiveMaximum', None)
10545

10646
return schema

‎tests/json_body_serializer.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ def serialize(cassette_dict: Any): # pragma: lax no cover
7676
del data['body']
7777
if content_type == ['application/x-www-form-urlencoded']:
7878
query_params = urllib.parse.parse_qs(data['body'])
79-
for key in ['client_secret', 'refresh_token']: # pragma: no cover
79+
for key in ['client_id', 'client_secret', 'refresh_token']: # pragma: no cover
8080
if key in query_params:
8181
query_params[key] = ['scrubbed']
8282
data['body'] = urllib.parse.urlencode(query_params)

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

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
interactions:
2+
- request:
3+
headers:
4+
accept:
5+
- '*/*'
6+
accept-encoding:
7+
- gzip, deflate
8+
connection:
9+
- keep-alive
10+
content-length:
11+
- '519'
12+
content-type:
13+
- application/json
14+
host:
15+
- generativelanguage.googleapis.com
16+
method: POST
17+
parsed_body:
18+
contents:
19+
- parts:
20+
- text: Create a config named "api-config" with metadata author="Alice" and version="1.0"
21+
role: user
22+
generationConfig:
23+
responseJsonSchema:
24+
description: A response with configuration metadata.
25+
properties:
26+
metadata:
27+
additionalProperties:
28+
type: string
29+
type: object
30+
name:
31+
type: string
32+
required:
33+
- name
34+
- metadata
35+
title: ConfigResponse
36+
type: object
37+
responseMimeType: application/json
38+
responseModalities:
39+
- TEXT
40+
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent
41+
response:
42+
headers:
43+
alt-svc:
44+
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
45+
content-length:
46+
- '631'
47+
content-type:
48+
- application/json; charset=UTF-8
49+
server-timing:
50+
- gfet4t7; dur=1379
51+
transfer-encoding:
52+
- chunked
53+
vary:
54+
- Origin
55+
- X-Origin
56+
- Referer
57+
parsed_body:
58+
candidates:
59+
- content:
60+
parts:
61+
- text: '{"name": "api-config", "metadata": {"author": "Alice", "version": "1.0"}}'
62+
role: model
63+
finishReason: STOP
64+
index: 0
65+
modelVersion: gemini-2.5-flash
66+
responseId: CZMUacOtKv2SxN8Pi7TrsAs
67+
usageMetadata:
68+
candidatesTokenCount: 25
69+
promptTokenCount: 23
70+
promptTokensDetails:
71+
- modality: TEXT
72+
tokenCount: 23
73+
thoughtsTokenCount: 158
74+
totalTokenCount: 206
75+
status:
76+
code: 200
77+
message: OK
78+
version: 1

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

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
interactions:
2+
- request:
3+
headers:
4+
accept:
5+
- '*/*'
6+
accept-encoding:
7+
- gzip, deflate
8+
connection:
9+
- keep-alive
10+
content-length:
11+
- '519'
12+
content-type:
13+
- application/json
14+
host:
15+
- generativelanguage.googleapis.com
16+
method: POST
17+
parsed_body:
18+
contents:
19+
- parts:
20+
- text: Create a config named "api-config" with metadata author="Alice" and version="1.0"
21+
role: user
22+
generationConfig:
23+
responseJsonSchema:
24+
description: A response with configuration metadata.
25+
properties:
26+
metadata:
27+
additionalProperties:
28+
type: string
29+
type: object
30+
name:
31+
type: string
32+
required:
33+
- name
34+
- metadata
35+
title: ConfigResponse
36+
type: object
37+
responseMimeType: application/json
38+
responseModalities:
39+
- TEXT
40+
uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent
41+
response:
42+
headers:
43+
alt-svc:
44+
- h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
45+
content-length:
46+
- '760'
47+
content-type:
48+
- application/json; charset=UTF-8
49+
server-timing:
50+
- gfet4t7; dur=818
51+
transfer-encoding:
52+
- chunked
53+
vary:
54+
- Origin
55+
- X-Origin
56+
- Referer
57+
parsed_body:
58+
candidates:
59+
- avgLogprobs: -1.4634492981713265e-05
60+
content:
61+
parts:
62+
- text: |-
63+
{
64+
"name": "api-config",
65+
"metadata": {
66+
"author": "Alice",
67+
"version": "1.0"
68+
}
69+
}
70+
role: model
71+
finishReason: STOP
72+
modelVersion: gemini-2.0-flash
73+
responseId: 8pMcab_EMqWd28oP46bOiAk
74+
usageMetadata:
75+
candidatesTokenCount: 40
76+
candidatesTokensDetails:
77+
- modality: TEXT
78+
tokenCount: 40
79+
promptTokenCount: 22
80+
promptTokensDetails:
81+
- modality: TEXT
82+
tokenCount: 22
83+
totalTokenCount: 62
84+
status:
85+
code: 200
86+
message: OK
87+
version: 1

0 commit comments

Comments
 (0)