Support enhanced JSON Schema features in Google Gemini 2.5+ models

pydantic_ai_slim/pydantic_ai/_json_schema.py

@@ -25,15 +25,13 @@ def __init__(
         *,
         strict: bool | None = None,
         prefer_inlined_defs: bool = False,
-        simplify_nullable_unions: bool = False,
     ):
         self.schema = schema
 
         self.strict = strict
         self.is_strict_compatible = True  # Can be set to False by subclasses to set `strict` on `ToolDefinition` when set not set by user explicitly
 
         self.prefer_inlined_defs = prefer_inlined_defs
-        self.simplify_nullable_unions = simplify_nullable_unions
 
         self.defs: dict[str, JsonSchema] = self.schema.get('$defs', {})
         self.refs_stack: list[str] = []
@@ -146,39 +144,11 @@ def _handle_union(self, schema: JsonSchema, union_kind: Literal['anyOf', 'oneOf'
 
         handled = [self._handle(member) for member in members]
 
-        # convert nullable unions to nullable types
-        if self.simplify_nullable_unions:
-            handled = self._simplify_nullable_union(handled)
-
-        if len(handled) == 1:
-            # In this case, no need to retain the union
-            return handled[0] | schema
-
         # If we have keys besides the union kind (such as title or discriminator), keep them without modifications
         schema = schema.copy()
         schema[union_kind] = handled
         return schema
 
-    @staticmethod
-    def _simplify_nullable_union(cases: list[JsonSchema]) -> list[JsonSchema]:
-        # TODO: Should we move this to relevant subclasses? Or is it worth keeping here to make reuse easier?
-        if len(cases) == 2 and {'type': 'null'} in cases:
-            # Find the non-null schema
-            non_null_schema = next(
-                (item for item in cases if item != {'type': 'null'}),
-                None,
-            )
-            if non_null_schema:
-                # Create a new schema based on the non-null part, mark as nullable
-                new_schema = deepcopy(non_null_schema)
-                new_schema['nullable'] = True
-                return [new_schema]
-            else:  # pragma: no cover
-                # they are both null, so just return one of them
-                return [cases[0]]
-
-        return cases
-
 
 class InlineDefsJsonSchemaTransformer(JsonSchemaTransformer):
     """Transforms the JSON Schema to inline $defs."""

pydantic_ai_slim/pydantic_ai/models/google.py

@@ -267,7 +267,7 @@ async def count_tokens(
             messages, model_settings, model_request_parameters
         )
 
-        # Annoyingly, the type of `GenerateContentConfigDict.get` is "partially `Unknown`" because `response_schema` includes `typing._UnionGenericAlias`,
+        # Annoyingly, the type of `GenerateContentConfigDict.get` is "partially `Unknown`" because `response_json_schema` includes `typing._UnionGenericAlias`,
         # so without this we'd need `pyright: ignore[reportUnknownMemberType]` on every line and wouldn't get type checking anyway.
         generation_config = cast(dict[str, Any], generation_config)
 
@@ -291,7 +291,7 @@ async def count_tokens(
                     thinking_config=generation_config.get('thinking_config'),
                     media_resolution=generation_config.get('media_resolution'),
                     response_mime_type=generation_config.get('response_mime_type'),
-                    response_schema=generation_config.get('response_schema'),
+                    response_json_schema=generation_config.get('response_json_schema'),
                 ),
             )
 
@@ -455,7 +455,7 @@ async def _build_content_and_config(
             tools=cast(ToolListUnionDict, tools),
             tool_config=tool_config,
             response_mime_type=response_mime_type,
-            response_schema=response_schema,
+            response_json_schema=response_schema,
             response_modalities=modalities,
         )
         return contents, config

pydantic_ai_slim/pydantic_ai/profiles/google.py

@@ -1,9 +1,5 @@
 from __future__ import annotations as _annotations
 
-import warnings
-
-from pydantic_ai.exceptions import UserError
-
 from .._json_schema import JsonSchema, JsonSchemaTransformer
 from . import ModelProfile
 
@@ -23,84 +19,28 @@ def google_model_profile(model_name: str) -> ModelProfile | None:
 class GoogleJsonSchemaTransformer(JsonSchemaTransformer):
     """Transforms the JSON Schema from Pydantic to be suitable for Gemini.
 
-    Gemini which [supports](https://ai.google.dev/gemini-api/docs/function-calling#function_declarations)
-    a subset of OpenAPI v3.0.3.
-
-    Specifically:
-    * gemini doesn't allow the `title` keyword to be set
-    * gemini doesn't allow `$defs` — we need to inline the definitions where possible
+    Gemini supports [a subset of OpenAPI v3.0.3](https://ai.google.dev/gemini-api/docs/function-calling#function_declarations).
     """
 
-    def __init__(self, schema: JsonSchema, *, strict: bool | None = None):
-        super().__init__(schema, strict=strict, prefer_inlined_defs=True, simplify_nullable_unions=True)
-
     def transform(self, schema: JsonSchema) -> JsonSchema:
-        # Note: we need to remove `additionalProperties: False` since it is currently mishandled by Gemini
-        additional_properties = schema.pop(
-            'additionalProperties', None
-        )  # don't pop yet so it's included in the warning
-        if additional_properties:
-            original_schema = {**schema, 'additionalProperties': additional_properties}
-            warnings.warn(
-                '`additionalProperties` is not supported by Gemini; it will be removed from the tool JSON schema.'
-                f' Full schema: {self.schema}\n\n'
-                f'Source of additionalProperties within the full schema: {original_schema}\n\n'
-                'If this came from a field with a type like `dict[str, MyType]`, that field will always be empty.\n\n'
-                "If Google's APIs are updated to support this properly, please create an issue on the Pydantic AI GitHub"
-                ' and we will fix this behavior.',
-                UserWarning,
-            )
-
-        schema.pop('title', None)
+        # Remove properties not supported by Gemini
         schema.pop('$schema', None)
         if (const := schema.pop('const', None)) is not None:
             # Gemini doesn't support const, but it does support enum with a single value
             schema['enum'] = [const]
         schema.pop('discriminator', None)
         schema.pop('examples', None)
 
-        # TODO: Should we use the trick from pydantic_ai.models.openai._OpenAIJsonSchema
-        #   where we add notes about these properties to the field description?
-        schema.pop('exclusiveMaximum', None)
-        schema.pop('exclusiveMinimum', None)
-
-        # Gemini only supports string enums, so we need to convert any enum values to strings.
-        # Pydantic will take care of transforming the transformed string values to the correct type.
-        if enum := schema.get('enum'):
-            schema['type'] = 'string'
-            schema['enum'] = [str(val) for val in enum]
-
         type_ = schema.get('type')
-        if 'oneOf' in schema and 'type' not in schema:  # pragma: no cover
-            # This gets hit when we have a discriminated union
-            # Gemini returns an API error in this case even though it says in its error message it shouldn't...
-            # Changing the oneOf to an anyOf prevents the API error and I think is functionally equivalent
-            schema['anyOf'] = schema.pop('oneOf')
-
         if type_ == 'string' and (fmt := schema.pop('format', None)):
             description = schema.get('description')
             if description:
                 schema['description'] = f'{description} (format: {fmt})'
             else:
                 schema['description'] = f'Format: {fmt}'
 
-        if '$ref' in schema:
-            raise UserError(f'Recursive `$ref`s in JSON Schema are not supported by Gemini: {schema["$ref"]}')
-
-        if 'prefixItems' in schema:
-            # prefixItems is not currently supported in Gemini, so we convert it to items for best compatibility
-            prefix_items = schema.pop('prefixItems')
-            items = schema.get('items')
-            unique_items = [items] if items is not None else []
-            for item in prefix_items:
-                if item not in unique_items:
-                    unique_items.append(item)
-            if len(unique_items) > 1:  # pragma: no cover
-                schema['items'] = {'anyOf': unique_items}
-            elif len(unique_items) == 1:  # pragma: no branch
-                schema['items'] = unique_items[0]
-            schema.setdefault('minItems', len(prefix_items))
-            if items is None:  # pragma: no branch
-                schema.setdefault('maxItems', len(prefix_items))
+        # Note: exclusiveMinimum/exclusiveMaximum are NOT yet supported
+        schema.pop('exclusiveMinimum', None)
+        schema.pop('exclusiveMaximum', None)
 
         return schema

tests/json_body_serializer.py

@@ -76,7 +76,7 @@ def serialize(cassette_dict: Any):  # pragma: lax no cover
                     del data['body']
             if content_type == ['application/x-www-form-urlencoded']:
                 query_params = urllib.parse.parse_qs(data['body'])
-                for key in ['client_secret', 'refresh_token']:  # pragma: no cover
+                for key in ['client_id', 'client_secret', 'refresh_token']:  # pragma: no cover
                     if key in query_params:
                         query_params[key] = ['scrubbed']
                         data['body'] = urllib.parse.urlencode(query_params)

tests/models/cassettes/test_google/test_google_dict_with_additional_properties_native_output.yaml

@@ -0,0 +1,78 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '519'
+      content-type:
+      - application/json
+      host:
+      - generativelanguage.googleapis.com
+    method: POST
+    parsed_body:
+      contents:
+      - parts:
+        - text: Create a config named "api-config" with metadata author="Alice" and version="1.0"
+        role: user
+      generationConfig:
+        responseJsonSchema:
+          description: A response with configuration metadata.
+          properties:
+            metadata:
+              additionalProperties:
+                type: string
+              type: object
+            name:
+              type: string
+          required:
+          - name
+          - metadata
+          title: ConfigResponse
+          type: object
+        responseMimeType: application/json
+        responseModalities:
+        - TEXT
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent
+  response:
+    headers:
+      alt-svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      content-length:
+      - '631'
+      content-type:
+      - application/json; charset=UTF-8
+      server-timing:
+      - gfet4t7; dur=1379
+      transfer-encoding:
+      - chunked
+      vary:
+      - Origin
+      - X-Origin
+      - Referer
+    parsed_body:
+      candidates:
+      - content:
+          parts:
+          - text: '{"name": "api-config", "metadata": {"author": "Alice", "version": "1.0"}}'
+          role: model
+        finishReason: STOP
+        index: 0
+      modelVersion: gemini-2.5-flash
+      responseId: CZMUacOtKv2SxN8Pi7TrsAs
+      usageMetadata:
+        candidatesTokenCount: 25
+        promptTokenCount: 23
+        promptTokensDetails:
+        - modality: TEXT
+          tokenCount: 23
+        thoughtsTokenCount: 158
+        totalTokenCount: 206
+    status:
+      code: 200
+      message: OK
+version: 1

tests/models/cassettes/test_google/test_google_discriminated_union_native_output.yaml

@@ -0,0 +1,102 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '807'
+      content-type:
+      - application/json
+      host:
+      - generativelanguage.googleapis.com
+    method: POST
+    parsed_body:
+      contents:
+      - parts:
+        - text: Tell me about a cat with a meow volume of 5
+        role: user
+      generationConfig:
+        responseJsonSchema:
+          $defs:
+            Cat:
+              properties:
+                meow_volume:
+                  type: integer
+                pet_type:
+                  default: cat
+                  enum:
+                  - cat
+                  type: string
+              required:
+              - meow_volume
+              title: Cat
+              type: object
+            Dog:
+              properties:
+                bark_volume:
+                  type: integer
+                pet_type:
+                  default: dog
+                  enum:
+                  - dog
+                  type: string
+              required:
+              - bark_volume
+              title: Dog
+              type: object
+          description: A response containing a pet.
+          properties:
+            pet:
+              oneOf:
+              - $ref: '#/$defs/Cat'
+              - $ref: '#/$defs/Dog'
+          required:
+          - pet
+          title: PetResponse
+          type: object
+        responseMimeType: application/json
+        responseModalities:
+        - TEXT
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent
+  response:
+    headers:
+      alt-svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      content-length:
+      - '594'
+      content-type:
+      - application/json; charset=UTF-8
+      server-timing:
+      - gfet4t7; dur=1682
+      transfer-encoding:
+      - chunked
+      vary:
+      - Origin
+      - X-Origin
+      - Referer
+    parsed_body:
+      candidates:
+      - content:
+          parts:
+          - text: '{"pet":{"pet_type":"cat","meow_volume":5}}'
+          role: model
+        finishReason: STOP
+        index: 0
+      modelVersion: gemini-2.5-flash
+      responseId: B5MUaePHJLHd7M8PqfX5qAQ
+      usageMetadata:
+        candidatesTokenCount: 16
+        promptTokenCount: 14
+        promptTokensDetails:
+        - modality: TEXT
+          tokenCount: 14
+        thoughtsTokenCount: 181
+        totalTokenCount: 211
+    status:
+      code: 200
+      message: OK
+version: 1