Stop incorrectly treating side-by-side JSON schema $refs as recursion (#1697)

Author: DouweM

pydantic_ai_slim/pydantic_ai/models/_json_schema.py

@@ -25,7 +25,7 @@ def __init__(
         self.simplify_nullable_unions = simplify_nullable_unions
 
         self.defs: dict[str, JsonSchema] = self.schema.get('$defs', {})
-        self.refs_stack = tuple[str, ...]()
+        self.refs_stack: list[str] = []
         self.recursive_refs = set[str]()
 
     @abstractmethod
@@ -62,13 +62,16 @@ def walk(self) -> JsonSchema:
         return handled
 
     def _handle(self, schema: JsonSchema) -> JsonSchema:
+        nested_refs = 0
         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,)
+                self.refs_stack.append(key)
+                nested_refs += 1
+
                 def_schema = self.defs.get(key)
                 if def_schema is None:  # pragma: no cover
                     raise UserError(f'Could not find $ref definition for {key}')
@@ -87,6 +90,9 @@ def _handle(self, schema: JsonSchema) -> JsonSchema:
         # Apply the base transform
         schema = self.transform(schema)
 
+        if nested_refs > 0:
+            self.refs_stack = self.refs_stack[:-nested_refs]
+
         return schema
 
     def _handle_object(self, schema: JsonSchema) -> JsonSchema:

tests/models/test_gemini.py

@@ -164,9 +164,17 @@ async def test_require_response_tool(allow_model_requests: None):
 
 
 async def test_json_def_replaced(allow_model_requests: None):
+    class Axis(BaseModel):
+        label: str
+
+    class Chart(BaseModel):
+        x_axis: Axis
+        y_axis: Axis
+
     class Location(BaseModel):
         lat: float
         lng: float = 1.1
+        chart: Chart
 
     class Locations(BaseModel):
         locations: list[Location]
@@ -175,15 +183,28 @@ class Locations(BaseModel):
     assert json_schema == snapshot(
         {
             '$defs': {
+                'Axis': {
+                    'properties': {'label': {'title': 'Label', 'type': 'string'}},
+                    'required': ['label'],
+                    'title': 'Axis',
+                    'type': 'object',
+                },
+                'Chart': {
+                    'properties': {'x_axis': {'$ref': '#/$defs/Axis'}, 'y_axis': {'$ref': '#/$defs/Axis'}},
+                    'required': ['x_axis', 'y_axis'],
+                    'title': 'Chart',
+                    'type': 'object',
+                },
                 'Location': {
                     'properties': {
                         'lat': {'title': 'Lat', 'type': 'number'},
                         'lng': {'default': 1.1, 'title': 'Lng', 'type': 'number'},
+                        'chart': {'$ref': '#/$defs/Chart'},
                     },
-                    'required': ['lat'],
+                    'required': ['lat', 'chart'],
                     'title': 'Location',
                     'type': 'object',
-                }
+                },
             },
             'properties': {'locations': {'items': {'$ref': '#/$defs/Location'}, 'title': 'Locations', 'type': 'array'}},
             'required': ['locations'],
@@ -201,20 +222,36 @@ class Locations(BaseModel):
     mrp = ModelRequestParameters(function_tools=[], allow_text_output=True, output_tools=[output_tool])
     mrp = m.customize_request_parameters(mrp)
     assert m._get_tools(mrp) == snapshot(
-        _GeminiTools(
-            function_declarations=[
-                _GeminiFunction(
-                    name='result',
-                    description='This is the tool for the final Result',
-                    parameters={
+        {
+            'function_declarations': [
+                {
+                    'name': 'result',
+                    'description': 'This is the tool for the final Result',
+                    'parameters': {
                         'properties': {
                             'locations': {
                                 'items': {
                                     'properties': {
                                         'lat': {'type': 'number'},
                                         'lng': {'type': 'number'},
+                                        'chart': {
+                                            'properties': {
+                                                'x_axis': {
+                                                    'properties': {'label': {'type': 'string'}},
+                                                    'required': ['label'],
+                                                    'type': 'object',
+                                                },
+                                                'y_axis': {
+                                                    'properties': {'label': {'type': 'string'}},
+                                                    'required': ['label'],
+                                                    'type': 'object',
+                                                },
+                                            },
+                                            'required': ['x_axis', 'y_axis'],
+                                            'type': 'object',
+                                        },
                                     },
-                                    'required': ['lat'],
+                                    'required': ['lat', 'chart'],
                                     'type': 'object',
                                 },
                                 'type': 'array',
@@ -223,9 +260,9 @@ class Locations(BaseModel):
                         'required': ['locations'],
                         'type': 'object',
                     },
-                )
+                }
             ]
-        )
+        }
     )