Generalize the JSON schema transformations (#1481)

Author: dmontagu

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -1,10 +1,8 @@
 from __future__ import annotations as _annotations
 
 import base64
-import re
 from collections.abc import AsyncIterator, Sequence
 from contextlib import asynccontextmanager
-from copy import deepcopy
 from dataclasses import dataclass, field, replace
 from datetime import datetime
 from typing import Annotated, Any, Literal, Protocol, Union, cast
@@ -46,6 +44,7 @@
     check_allow_model_requests,
     get_user_agent,
 )
+from .json_schema import JsonSchema, WalkJsonSchema
 
 LatestGeminiModelNames = Literal[
     'gemini-1.5-flash',
@@ -156,7 +155,7 @@ async def request_stream(
 
     def customize_request_parameters(self, model_request_parameters: ModelRequestParameters) -> ModelRequestParameters:
         def _customize_tool_def(t: ToolDefinition):
-            return replace(t, parameters_json_schema=_GeminiJsonSchema(t.parameters_json_schema).simplify())
+            return replace(t, parameters_json_schema=_GeminiJsonSchema(t.parameters_json_schema).walk())
 
         return ModelRequestParameters(
             function_tools=[_customize_tool_def(tool) for tool in model_request_parameters.function_tools],
@@ -760,7 +759,7 @@ class _GeminiPromptFeedback(TypedDict):
 _gemini_streamed_response_ta = pydantic.TypeAdapter(list[_GeminiResponse], config=pydantic.ConfigDict(defer_build=True))
 
 
-class _GeminiJsonSchema:
+class _GeminiJsonSchema(WalkJsonSchema):
     """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)
@@ -771,72 +770,58 @@ class _GeminiJsonSchema:
     * gemini doesn't allow `$defs` — we need to inline the definitions where possible
     """
 
-    def __init__(self, schema: _utils.ObjectJsonSchema):
-        self.schema = deepcopy(schema)
-        self.defs = self.schema.pop('$defs', {})
+    def __init__(self, schema: JsonSchema):
+        super().__init__(schema, prefer_inlined_defs=True, simplify_nullable_unions=True)
 
-    def simplify(self) -> dict[str, Any]:
-        self._simplify(self.schema, refs_stack=())
-        return self.schema
-
-    def _simplify(self, schema: dict[str, Any], refs_stack: tuple[str, ...]) -> None:
+    def transform(self, schema: JsonSchema) -> JsonSchema:
         schema.pop('title', None)
         schema.pop('default', None)
         schema.pop('$schema', None)
+        if (const := schema.pop('const', None)) is not None:  # pragma: no cover
+            # 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)
-        if ref := schema.pop('$ref', None):
-            # noinspection PyTypeChecker
-            key = re.sub(r'^#/\$defs/', '', ref)
-            if key in refs_stack:
-                raise UserError('Recursive `$ref`s in JSON Schema are not supported by Gemini')
-            refs_stack += (key,)
-            schema_def = self.defs[key]
-            self._simplify(schema_def, refs_stack)
-            schema.update(schema_def)
-            return
-
-        if any_of := schema.get('anyOf'):
-            for item_schema in any_of:
-                self._simplify(item_schema, refs_stack)
-            if len(any_of) == 2 and {'type': 'null'} in any_of:
-                for item_schema in any_of:
-                    if item_schema != {'type': 'null'}:
-                        schema.clear()
-                        schema.update(item_schema)
-                        schema['nullable'] = True
-                        return
 
         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_ == 'object':
-            self._object(schema, refs_stack)
-        elif type_ == 'array':
-            return self._array(schema, refs_stack)
-        elif type_ == 'string' and (fmt := schema.pop('format', None)):
+        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}'
 
-    def _object(self, schema: dict[str, Any], refs_stack: tuple[str, ...]) -> None:
-        ad_props = schema.pop('additionalProperties', None)
-        if ad_props:
-            raise UserError('Additional properties in JSON Schema are not supported by Gemini')
-
-        if properties := schema.get('properties'):  # pragma: no branch
-            for value in properties.values():
-                self._simplify(value, refs_stack)
-
-    def _array(self, schema: dict[str, Any], refs_stack: tuple[str, ...]) -> None:
-        if prefix_items := schema.get('prefixItems'):
-            # TODO I think this not is supported by Gemini, maybe we should raise an error?
-            for prefix_item in prefix_items:
-                self._simplify(prefix_item, refs_stack)
-
-        if items_schema := schema.get('items'):  # pragma: no branch
-            self._simplify(items_schema, refs_stack)
+        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:
+                schema['items'] = unique_items[0]
+            schema.setdefault('minItems', len(prefix_items))
+            if items is None:
+                schema.setdefault('maxItems', len(prefix_items))
+
+        return schema
 
 
 def _ensure_decodeable(content: bytearray) -> bytearray:

pydantic_ai_slim/pydantic_ai/models/json_schema.py

@@ -0,0 +1,152 @@
+import re
+from abc import ABC, abstractmethod
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Any, Literal
+
+from pydantic_ai.exceptions import UserError
+
+JsonSchema = dict[str, Any]
+
+
+@dataclass(init=False)
+class WalkJsonSchema(ABC):
+    """Walks a JSON schema, applying transformations to it at each level."""
+
+    def __init__(
+        self, schema: JsonSchema, *, prefer_inlined_defs: bool = False, simplify_nullable_unions: bool = False
+    ):
+        self.schema = deepcopy(schema)
+        self.prefer_inlined_defs = prefer_inlined_defs
+        self.simplify_nullable_unions = simplify_nullable_unions
+
+        self.defs: dict[str, JsonSchema] = self.schema.pop('$defs', {})
+        self.refs_stack = tuple[str, ...]()
+        self.recursive_refs = set[str]()
+
+    @abstractmethod
+    def transform(self, schema: JsonSchema) -> JsonSchema:
+        """Make changes to the schema."""
+        return schema
+
+    def walk(self) -> JsonSchema:
+        handled = self._handle(deepcopy(self.schema))
+
+        if not self.prefer_inlined_defs and self.defs:
+            handled['$defs'] = {k: self._handle(v) for k, v in self.defs.items()}
+
+        elif self.recursive_refs:  # pragma: no cover
+            # If we are preferring inlined defs and there are recursive refs, we _have_ to use a $defs+$ref structure
+            # We try to use whatever the original root key was, but if it is already in use,
+            # we modify it to avoid collisions.
+            defs = {key: self.defs[key] for key in self.recursive_refs}
+            root_ref = self.schema.get('$ref')
+            root_key = None if root_ref is None else re.sub(r'^#/\$defs/', '', root_ref)
+            if root_key is None:
+                root_key = self.schema.get('title', 'root')
+                while root_key in defs:
+                    # Modify the root key until it is not already in use
+                    root_key = f'{root_key}_root'
+
+            defs[root_key] = handled
+            return {'$defs': defs, '$ref': f'#/$defs/{root_key}'}
+
+        return handled
+
+    def _handle(self, schema: JsonSchema) -> JsonSchema:
+        if self.prefer_inlined_defs:
+            while ref := schema.get('$ref'):
+                key = re.sub(r'^#/\$defs/', '', ref)
+                if key in self.refs_stack:
+                    self.recursive_refs.add(key)
+                    break  # recursive ref can't be unpacked
+                self.refs_stack += (key,)
+                def_schema = self.defs.get(key)
+                if def_schema is None:  # pragma: no cover
+                    raise UserError(f'Could not find $ref definition for {key}')
+                schema = def_schema
+
+        # Handle the schema based on its type / structure
+        type_ = schema.get('type')
+        if type_ == 'object':
+            schema = self._handle_object(schema)
+        elif type_ == 'array':
+            schema = self._handle_array(schema)
+        elif type_ is None:
+            schema = self._handle_union(schema, 'anyOf')
+            schema = self._handle_union(schema, 'oneOf')
+
+        # Apply the base transform
+        schema = self.transform(schema)
+
+        return schema
+
+    def _handle_object(self, schema: JsonSchema) -> JsonSchema:
+        if properties := schema.get('properties'):
+            handled_properties = {}
+            for key, value in properties.items():
+                handled_properties[key] = self._handle(value)
+            schema['properties'] = handled_properties
+
+        if (additional_properties := schema.get('additionalProperties')) is not None:
+            if isinstance(additional_properties, bool):
+                schema['additionalProperties'] = additional_properties
+            else:  # pragma: no cover
+                schema['additionalProperties'] = self._handle(additional_properties)
+
+        if (pattern_properties := schema.get('patternProperties')) is not None:
+            handled_pattern_properties = {}
+            for key, value in pattern_properties.items():
+                handled_pattern_properties[key] = self._handle(value)
+            schema['patternProperties'] = handled_pattern_properties
+
+        return schema
+
+    def _handle_array(self, schema: JsonSchema) -> JsonSchema:
+        if prefix_items := schema.get('prefixItems'):
+            schema['prefixItems'] = [self._handle(item) for item in prefix_items]
+
+        if items := schema.get('items'):
+            schema['items'] = self._handle(items)
+
+        return schema
+
+    def _handle_union(self, schema: JsonSchema, union_kind: Literal['anyOf', 'oneOf']) -> JsonSchema:
+        members = schema.get(union_kind)
+        if not members:
+            return schema
+
+        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]
+
+        # 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  # pragma: no cover