Fix StructuredDict with nested JSON schemas using $ref (#2570)

Co-authored-by: Douwe Maan <[email protected]>

Author: ChiaXinLiang

pydantic_ai_slim/pydantic_ai/profiles/_json_schema.py renamed to pydantic_ai_slim/pydantic_ai/_json_schema.py

@@ -6,7 +6,7 @@
 from dataclasses import dataclass
 from typing import Any, Literal
 
-from pydantic_ai.exceptions import UserError
+from .exceptions import UserError
 
 JsonSchema = dict[str, Any]
 

pydantic_ai_slim/pydantic_ai/_utils.py

@@ -6,7 +6,7 @@
 import re
 import time
 import uuid
-from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator
+from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterable, Iterator
 from contextlib import asynccontextmanager, suppress
 from dataclasses import dataclass, fields, is_dataclass
 from datetime import datetime, timezone
@@ -70,16 +70,33 @@ def check_object_json_schema(schema: JsonSchemaValue) -> ObjectJsonSchema:
 
     if schema.get('type') == 'object':
         return schema
-    elif schema.get('$ref') is not None:
-        maybe_result = schema.get('$defs', {}).get(schema['$ref'][8:])  # This removes the initial "#/$defs/".
-
-        if "'$ref': '#/$defs/" in str(maybe_result):
-            return schema  # We can't remove the $defs because the schema contains other references
-        return maybe_result
+    elif ref := schema.get('$ref'):
+        prefix = '#/$defs/'
+        # Return the referenced schema unless it contains additional nested references.
+        if (
+            ref.startswith(prefix)
+            and (resolved := schema.get('$defs', {}).get(ref[len(prefix) :]))
+            and resolved.get('type') == 'object'
+            and not _contains_ref(resolved)
+        ):
+            return resolved
+        return schema
     else:
         raise UserError('Schema must be an object')
 
 
+def _contains_ref(obj: JsonSchemaValue | list[JsonSchemaValue]) -> bool:
+    """Recursively check if an object contains any $ref keys."""
+    items: Iterable[JsonSchemaValue]
+    if isinstance(obj, dict):
+        if '$ref' in obj:
+            return True
+        items = obj.values()
+    else:
+        items = obj
+    return any(isinstance(item, dict | list) and _contains_ref(item) for item in items)  # pyright: ignore[reportUnknownArgumentType]
+
+
 T = TypeVar('T')
 
 

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -20,6 +20,7 @@
 from typing_extensions import TypeAliasType, TypedDict
 
 from .. import _utils
+from .._json_schema import JsonSchemaTransformer
 from .._output import OutputObjectDefinition
 from .._parts_manager import ModelResponsePartsManager
 from .._run_context import RunContext
@@ -40,7 +41,6 @@
 )
 from ..output import OutputMode
 from ..profiles import DEFAULT_PROFILE, ModelProfile, ModelProfileSpec
-from ..profiles._json_schema import JsonSchemaTransformer
 from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from ..usage import RequestUsage

pydantic_ai_slim/pydantic_ai/output.py

@@ -10,6 +10,7 @@
 from typing_extensions import TypeAliasType, TypeVar, deprecated
 
 from . import _utils
+from ._json_schema import InlineDefsJsonSchemaTransformer
 from .messages import ToolCallPart
 from .tools import DeferredToolRequests, ObjectJsonSchema, RunContext, ToolDefinition
 
@@ -311,6 +312,11 @@ def StructuredDict(
     """
     json_schema = _utils.check_object_json_schema(json_schema)
 
+    # Pydantic `TypeAdapter` fails when `object.__get_pydantic_json_schema__` has `$defs`, so we inline them
+    # See https://github.com/pydantic/pydantic/issues/12145
+    if '$defs' in json_schema:
+        json_schema = InlineDefsJsonSchemaTransformer(json_schema).walk()
+
     if name:
         json_schema['title'] = name
 

pydantic_ai_slim/pydantic_ai/profiles/__init__.py

@@ -6,8 +6,8 @@
 
 from typing_extensions import Self
 
+from .._json_schema import InlineDefsJsonSchemaTransformer, JsonSchemaTransformer
 from ..output import StructuredOutputMode
-from ._json_schema import InlineDefsJsonSchemaTransformer, JsonSchemaTransformer
 
 __all__ = [
     'ModelProfile',

pydantic_ai_slim/pydantic_ai/profiles/google.py

@@ -4,8 +4,8 @@
 
 from pydantic_ai.exceptions import UserError
 
+from .._json_schema import JsonSchema, JsonSchemaTransformer
 from . import ModelProfile
-from ._json_schema import JsonSchema, JsonSchemaTransformer
 
 
 def google_model_profile(model_name: str) -> ModelProfile | None:

pydantic_ai_slim/pydantic_ai/profiles/openai.py

@@ -6,8 +6,8 @@
 from dataclasses import dataclass
 from typing import Any, Literal
 
+from .._json_schema import JsonSchema, JsonSchemaTransformer
 from . import ModelProfile
-from ._json_schema import JsonSchema, JsonSchemaTransformer
 
 OpenAISystemPromptRole = Literal['system', 'developer', 'user']
 

tests/models/test_openai.py

@@ -34,10 +34,10 @@
     UserError,
     UserPromptPart,
 )
+from pydantic_ai._json_schema import InlineDefsJsonSchemaTransformer
 from pydantic_ai.builtin_tools import WebSearchTool
 from pydantic_ai.models import ModelRequestParameters
 from pydantic_ai.output import NativeOutput, PromptedOutput, TextOutput, ToolOutput
-from pydantic_ai.profiles._json_schema import InlineDefsJsonSchemaTransformer
 from pydantic_ai.profiles.openai import OpenAIModelProfile, openai_model_profile
 from pydantic_ai.result import RunUsage
 from pydantic_ai.settings import ModelSettings

tests/providers/test_azure.py

@@ -4,8 +4,8 @@
 from inline_snapshot import snapshot
 from pytest_mock import MockerFixture
 
+from pydantic_ai._json_schema import InlineDefsJsonSchemaTransformer
 from pydantic_ai.agent import Agent
-from pydantic_ai.profiles._json_schema import InlineDefsJsonSchemaTransformer
 from pydantic_ai.profiles.cohere import cohere_model_profile
 from pydantic_ai.profiles.deepseek import deepseek_model_profile
 from pydantic_ai.profiles.grok import grok_model_profile

tests/providers/test_bedrock.py

@@ -3,7 +3,7 @@
 import pytest
 from pytest_mock import MockerFixture
 
-from pydantic_ai.profiles._json_schema import InlineDefsJsonSchemaTransformer
+from pydantic_ai._json_schema import InlineDefsJsonSchemaTransformer
 from pydantic_ai.profiles.amazon import amazon_model_profile
 from pydantic_ai.profiles.anthropic import anthropic_model_profile
 from pydantic_ai.profiles.cohere import cohere_model_profile