mirror of
https://github.com/pydantic/pydantic-ai.git
synced 2026-10-11 21:59:51 +00:00
Fix merge_json_schema_defs corrupting branch schemas on transitive def-name collisions (#9735)
Co-authored-by: pydanty <[email protected]> Co-authored-by: Douwe Maan <[email protected]>
This commit is contained in:
pydanty[bot]
co-authored by
pydanty
Douwe Maan
pydanty
Douwe Maan
parent
512b80b493
commit
90051ef7f6
@@ -988,48 +988,45 @@ def merge_json_schema_defs(schemas: list[dict[str, Any]]) -> tuple[list[dict[str
|
||||
rewritten_schemas.append(schema)
|
||||
continue
|
||||
|
||||
schema = schema.copy()
|
||||
defs = schema.pop('$defs', None)
|
||||
schema_name_mapping: dict[str, str] = {}
|
||||
schema = copy.deepcopy(schema)
|
||||
defs = schema.pop('$defs')
|
||||
renames: dict[str, str] = {}
|
||||
new_names: list[str] = []
|
||||
|
||||
# Process definitions and build mapping
|
||||
# New names are kept; same-named defs with a different body get a unique name. Both point
|
||||
# `all_defs` at this schema's own copy of the def, whose refs are rewritten in place below.
|
||||
for name, def_schema in defs.items():
|
||||
if name not in all_defs:
|
||||
all_defs[name] = def_schema
|
||||
schema_name_mapping[name] = name
|
||||
new_names.append(name)
|
||||
elif def_schema != all_defs[name]:
|
||||
# Different def with same name — assign a unique name
|
||||
schema_name_mapping[name] = _unique_def_name(name, schema, all_defs)
|
||||
all_defs[schema_name_mapping[name]] = def_schema
|
||||
# else: structurally equal — handled below
|
||||
renames[name] = _unique_def_name(name, schema, all_defs)
|
||||
all_defs[renames[name]] = def_schema
|
||||
|
||||
# Defs that are structurally equal (same dict) may still be semantically
|
||||
# different if they contain $refs that point to defs that were renamed in
|
||||
# this schema. E.g. both schemas have Wrapper={$ref Inner}, but their
|
||||
# Inner defs differ, so Schema B's Inner was renamed to Inner_1. The shared
|
||||
# Wrapper is not actually equal — Schema B needs its own copy with updated refs.
|
||||
# Loop until stable, since creating a copy can trigger further copies
|
||||
# in defs that reference it (transitive chains).
|
||||
# A def equal to an earlier schema's def ("shared") can reuse it only if none of the defs it
|
||||
# transitively references were renamed; otherwise its `$ref`s resolve differently and it needs
|
||||
# its own copy. Iterate to a fixpoint before keeping any shared def, so the result does not
|
||||
# depend on the order of `$defs`.
|
||||
shared = [name for name in defs if name not in renames and name not in new_names]
|
||||
changed = True
|
||||
while changed:
|
||||
changed = False
|
||||
for name, def_schema in defs.items():
|
||||
if name not in schema_name_mapping:
|
||||
updated = copy.deepcopy(def_schema)
|
||||
_update_mapped_json_schema_refs(updated, schema_name_mapping)
|
||||
if updated != def_schema:
|
||||
schema_name_mapping[name] = _unique_def_name(name, schema, all_defs)
|
||||
all_defs[schema_name_mapping[name]] = updated
|
||||
changed = True
|
||||
else:
|
||||
schema_name_mapping[name] = name
|
||||
for name in shared:
|
||||
if name in renames:
|
||||
continue
|
||||
updated = copy.deepcopy(defs[name])
|
||||
_update_mapped_json_schema_refs(updated, renames)
|
||||
if updated != defs[name]:
|
||||
renames[name] = _unique_def_name(name, schema, all_defs)
|
||||
all_defs[renames[name]] = defs[name]
|
||||
changed = True
|
||||
|
||||
# Update refs inside definitions so internal cross-references
|
||||
# (e.g. Outer referencing Inner which was renamed to Inner_1) are corrected.
|
||||
for new_name in schema_name_mapping.values():
|
||||
_update_mapped_json_schema_refs(all_defs[new_name], schema_name_mapping)
|
||||
# Rewrite the refs of every def this schema contributes. Shared defs left under their
|
||||
# name only reference unrenamed defs, so the earlier schema's copy is already correct.
|
||||
for name in [*new_names, *renames]:
|
||||
_update_mapped_json_schema_refs(defs[name], renames)
|
||||
|
||||
_update_mapped_json_schema_refs(schema, schema_name_mapping)
|
||||
_update_mapped_json_schema_refs(schema, renames)
|
||||
rewritten_schemas.append(schema)
|
||||
|
||||
return rewritten_schemas, all_defs
|
||||
|
||||
@@ -3,7 +3,7 @@ import json
|
||||
from typing import Annotated, Any, Optional
|
||||
|
||||
import pytest
|
||||
from pydantic import AfterValidator, BaseModel, Field
|
||||
from pydantic import AfterValidator, BaseModel, Field, create_model
|
||||
from typing_extensions import TypeAliasType
|
||||
|
||||
from pydantic_ai import (
|
||||
@@ -98,6 +98,124 @@ async def test_auto_output_json_schema():
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('names', 'expected'),
|
||||
[
|
||||
pytest.param(
|
||||
('AOuter', 'Middle', 'ZLeaf'),
|
||||
snapshot(
|
||||
{
|
||||
'anyOf': [
|
||||
{
|
||||
'properties': {'middle': {'$ref': '#/$defs/Middle'}},
|
||||
'required': ['middle'],
|
||||
'title': 'AOuter',
|
||||
'type': 'object',
|
||||
},
|
||||
{
|
||||
'properties': {'middle': {'$ref': '#/$defs/AOuter_Middle_1'}},
|
||||
'required': ['middle'],
|
||||
'title': 'AOuter',
|
||||
'type': 'object',
|
||||
},
|
||||
],
|
||||
'$defs': {
|
||||
'Middle': {
|
||||
'properties': {'leaf': {'$ref': '#/$defs/ZLeaf'}},
|
||||
'required': ['leaf'],
|
||||
'title': 'Middle',
|
||||
'type': 'object',
|
||||
},
|
||||
'ZLeaf': {
|
||||
'properties': {'value': {'title': 'Value', 'type': 'string'}},
|
||||
'required': ['value'],
|
||||
'title': 'ZLeaf',
|
||||
'type': 'object',
|
||||
},
|
||||
'AOuter_ZLeaf_1': {
|
||||
'properties': {'value': {'title': 'Value', 'type': 'integer'}},
|
||||
'required': ['value'],
|
||||
'title': 'ZLeaf',
|
||||
'type': 'object',
|
||||
},
|
||||
'AOuter_Middle_1': {
|
||||
'properties': {'leaf': {'$ref': '#/$defs/AOuter_ZLeaf_1'}},
|
||||
'required': ['leaf'],
|
||||
'title': 'Middle',
|
||||
'type': 'object',
|
||||
},
|
||||
},
|
||||
}
|
||||
),
|
||||
id='sorted-chain',
|
||||
),
|
||||
pytest.param(
|
||||
('A', 'C', 'B'),
|
||||
snapshot(
|
||||
{
|
||||
'anyOf': [
|
||||
{
|
||||
'properties': {'middle': {'$ref': '#/$defs/C'}},
|
||||
'required': ['middle'],
|
||||
'title': 'A',
|
||||
'type': 'object',
|
||||
},
|
||||
{
|
||||
'properties': {'middle': {'$ref': '#/$defs/A_C_1'}},
|
||||
'required': ['middle'],
|
||||
'title': 'A',
|
||||
'type': 'object',
|
||||
},
|
||||
],
|
||||
'$defs': {
|
||||
'B': {
|
||||
'properties': {'value': {'title': 'Value', 'type': 'string'}},
|
||||
'required': ['value'],
|
||||
'title': 'B',
|
||||
'type': 'object',
|
||||
},
|
||||
'C': {
|
||||
'properties': {'leaf': {'$ref': '#/$defs/B'}},
|
||||
'required': ['leaf'],
|
||||
'title': 'C',
|
||||
'type': 'object',
|
||||
},
|
||||
'A_B_1': {
|
||||
'properties': {'value': {'title': 'Value', 'type': 'integer'}},
|
||||
'required': ['value'],
|
||||
'title': 'B',
|
||||
'type': 'object',
|
||||
},
|
||||
'A_C_1': {
|
||||
'properties': {'leaf': {'$ref': '#/$defs/A_B_1'}},
|
||||
'required': ['leaf'],
|
||||
'title': 'C',
|
||||
'type': 'object',
|
||||
},
|
||||
},
|
||||
}
|
||||
),
|
||||
id='out-of-order-chain',
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_output_json_schema_transitive_collision_preserves_branches(
|
||||
names: tuple[str, str, str], expected: dict[str, Any]
|
||||
):
|
||||
"""Same-named nested defs with different bodies keep per-branch fidelity."""
|
||||
outer_name, middle_name, leaf_name = names
|
||||
first_leaf = create_model(leaf_name, value=(str, ...))
|
||||
first_middle = create_model(middle_name, leaf=(first_leaf, ...))
|
||||
first = create_model(outer_name, middle=(first_middle, ...))
|
||||
second_leaf = create_model(leaf_name, value=(int, ...))
|
||||
second_middle = create_model(middle_name, leaf=(second_leaf, ...))
|
||||
second = create_model(outer_name, middle=(second_middle, ...))
|
||||
|
||||
schema = Agent(TestModel(), output_type=[first, second]).output_json_schema()
|
||||
|
||||
assert schema == expected
|
||||
|
||||
|
||||
async def test_tool_output_json_schema():
|
||||
# one output
|
||||
agent = Agent(
|
||||
|
||||
+145
-1
@@ -3,8 +3,10 @@ from __future__ import annotations as _annotations
|
||||
import asyncio
|
||||
import contextlib
|
||||
import contextvars
|
||||
import copy
|
||||
import functools
|
||||
import importlib.util
|
||||
import itertools
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
@@ -13,7 +15,7 @@ from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass, field
|
||||
from importlib.metadata import distributions
|
||||
from types import ModuleType
|
||||
from typing import Any
|
||||
from typing import Any, Literal, cast
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
@@ -30,6 +32,7 @@ from pydantic_ai._utils import (
|
||||
get_first_param_type,
|
||||
group_by_temporal,
|
||||
is_async_callable,
|
||||
is_str_dict,
|
||||
merge_json_schema_defs,
|
||||
replace_no_init,
|
||||
run_in_executor,
|
||||
@@ -1133,6 +1136,147 @@ def test_merge_json_schema_defs_structurally_equal_with_different_ref_targets():
|
||||
)
|
||||
|
||||
|
||||
def _chain_schema(
|
||||
names: tuple[str, str, str, str],
|
||||
value_type: str,
|
||||
*,
|
||||
defs_layout: Literal['sorted', 'chain'],
|
||||
title: str,
|
||||
) -> dict[str, Any]:
|
||||
d1, d2, d3, leaf = names
|
||||
chain_defs = {
|
||||
d1: {'type': 'object', 'properties': {'next': {'$ref': f'#/$defs/{d2}'}}},
|
||||
d2: {'type': 'object', 'properties': {'next': {'$ref': f'#/$defs/{d3}'}}},
|
||||
d3: {'type': 'object', 'properties': {'next': {'$ref': f'#/$defs/{leaf}'}}},
|
||||
leaf: {'type': 'object', 'properties': {'value': {'type': value_type}}},
|
||||
}
|
||||
if defs_layout == 'sorted':
|
||||
chain_defs = {name: chain_defs[name] for name in sorted(chain_defs)}
|
||||
|
||||
return {
|
||||
'$defs': chain_defs,
|
||||
'title': title,
|
||||
'type': 'object',
|
||||
'properties': {'root': {'$ref': f'#/$defs/{d1}'}},
|
||||
}
|
||||
|
||||
|
||||
def _assert_all_json_schema_refs_resolve(value: Any, defs: dict[str, dict[str, Any]]) -> None:
|
||||
if is_str_dict(value):
|
||||
if ref := value.get('$ref'):
|
||||
assert str(ref).removeprefix('#/$defs/') in defs
|
||||
for nested in value.values():
|
||||
_assert_all_json_schema_refs_resolve(nested, defs)
|
||||
elif isinstance(value, list):
|
||||
for nested in cast(list[Any], value):
|
||||
_assert_all_json_schema_refs_resolve(nested, defs)
|
||||
|
||||
|
||||
def _resolve_branch_leaf_value_type(branch: dict[str, Any], defs: dict[str, dict[str, Any]]) -> str:
|
||||
"""Follow `$ref`s from a merged branch down to its leaf model's `value` field."""
|
||||
node: dict[str, Any] = branch
|
||||
while True:
|
||||
props: dict[str, Any] = node.get('properties', {})
|
||||
next_props = [prop for prop in props.values() if '$ref' in prop]
|
||||
if not next_props:
|
||||
return str(node['properties']['value']['type'])
|
||||
node = defs[str(next_props[0]['$ref']).removeprefix('#/$defs/')]
|
||||
|
||||
|
||||
@pytest.mark.parametrize('names', list(itertools.permutations(('A', 'B', 'C', 'D'))))
|
||||
@pytest.mark.parametrize('defs_layout', ['sorted', 'chain'])
|
||||
def test_merge_json_schema_defs_transitive_rename_all_orderings(
|
||||
names: tuple[str, str, str, str], defs_layout: Literal['sorted', 'chain']
|
||||
) -> None:
|
||||
"""Every 4-level chain ordering keeps branch types; too many cases to run through `Agent.output_json_schema()`."""
|
||||
schemas = [
|
||||
_chain_schema(names, 'string', defs_layout=defs_layout, title='StringRoot'),
|
||||
_chain_schema(names, 'integer', defs_layout=defs_layout, title='IntegerRoot'),
|
||||
]
|
||||
original_schemas = copy.deepcopy(schemas)
|
||||
|
||||
rewritten_schemas, all_defs = merge_json_schema_defs(schemas)
|
||||
|
||||
assert _resolve_branch_leaf_value_type(rewritten_schemas[0], all_defs) == 'string'
|
||||
assert _resolve_branch_leaf_value_type(rewritten_schemas[1], all_defs) == 'integer'
|
||||
_assert_all_json_schema_refs_resolve(rewritten_schemas, all_defs)
|
||||
_assert_all_json_schema_refs_resolve(all_defs, all_defs)
|
||||
assert schemas == original_schemas
|
||||
|
||||
|
||||
def test_merge_json_schema_defs_transitive_rename_snapshot():
|
||||
"""Pins the names for one previously broken ordering; the Agent-level test covers the user-facing path."""
|
||||
names = ('A', 'C', 'B', 'D')
|
||||
schemas = [
|
||||
_chain_schema(names, 'string', defs_layout='sorted', title='StringRoot'),
|
||||
_chain_schema(names, 'integer', defs_layout='sorted', title='IntegerRoot'),
|
||||
]
|
||||
|
||||
rewritten_schemas, all_defs = merge_json_schema_defs(schemas)
|
||||
|
||||
assert (all_defs, rewritten_schemas) == snapshot(
|
||||
(
|
||||
{
|
||||
'A': {'type': 'object', 'properties': {'next': {'$ref': '#/$defs/C'}}},
|
||||
'B': {'type': 'object', 'properties': {'next': {'$ref': '#/$defs/D'}}},
|
||||
'C': {'type': 'object', 'properties': {'next': {'$ref': '#/$defs/B'}}},
|
||||
'D': {'type': 'object', 'properties': {'value': {'type': 'string'}}},
|
||||
'IntegerRoot_D_1': {'type': 'object', 'properties': {'value': {'type': 'integer'}}},
|
||||
'IntegerRoot_B_1': {'type': 'object', 'properties': {'next': {'$ref': '#/$defs/IntegerRoot_D_1'}}},
|
||||
'IntegerRoot_C_1': {'type': 'object', 'properties': {'next': {'$ref': '#/$defs/IntegerRoot_B_1'}}},
|
||||
'IntegerRoot_A_1': {'type': 'object', 'properties': {'next': {'$ref': '#/$defs/IntegerRoot_C_1'}}},
|
||||
},
|
||||
[
|
||||
{'title': 'StringRoot', 'type': 'object', 'properties': {'root': {'$ref': '#/$defs/A'}}},
|
||||
{'title': 'IntegerRoot', 'type': 'object', 'properties': {'root': {'$ref': '#/$defs/IntegerRoot_A_1'}}},
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_merge_json_schema_defs_shared_chain_not_renamed_for_unrelated_collision():
|
||||
"""Shared defs that do not reach a renamed def are reused, not copied; a merge-level detail of the def names."""
|
||||
shared_defs = {
|
||||
'SharedRoot': {'type': 'object', 'properties': {'leaf': {'$ref': '#/$defs/SharedLeaf'}}},
|
||||
'SharedLeaf': {'type': 'object', 'properties': {'value': {'type': 'boolean'}}},
|
||||
}
|
||||
schemas = [
|
||||
{
|
||||
'$defs': {
|
||||
**shared_defs,
|
||||
'Collision': {'type': 'object', 'properties': {'value': {'type': 'string'}}},
|
||||
},
|
||||
'title': 'First',
|
||||
'properties': {'shared': {'$ref': '#/$defs/SharedRoot'}},
|
||||
},
|
||||
{
|
||||
'$defs': {
|
||||
**shared_defs,
|
||||
'Collision': {'type': 'object', 'properties': {'value': {'type': 'integer'}}},
|
||||
},
|
||||
'title': 'Second',
|
||||
'properties': {'shared': {'$ref': '#/$defs/SharedRoot'}},
|
||||
},
|
||||
]
|
||||
|
||||
rewritten_schemas, all_defs = merge_json_schema_defs(schemas)
|
||||
|
||||
assert (all_defs, rewritten_schemas) == snapshot(
|
||||
(
|
||||
{
|
||||
'SharedRoot': {'type': 'object', 'properties': {'leaf': {'$ref': '#/$defs/SharedLeaf'}}},
|
||||
'SharedLeaf': {'type': 'object', 'properties': {'value': {'type': 'boolean'}}},
|
||||
'Collision': {'type': 'object', 'properties': {'value': {'type': 'string'}}},
|
||||
'Second_Collision_1': {'type': 'object', 'properties': {'value': {'type': 'integer'}}},
|
||||
},
|
||||
[
|
||||
{'title': 'First', 'properties': {'shared': {'$ref': '#/$defs/SharedRoot'}}},
|
||||
{'title': 'Second', 'properties': {'shared': {'$ref': '#/$defs/SharedRoot'}}},
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_strip_markdown_fences():
|
||||
assert strip_markdown_fences('{"foo": "bar"}') == '{"foo": "bar"}'
|
||||
assert strip_markdown_fences('```json\n{"foo": "bar"}\n```') == '{"foo": "bar"}'
|
||||
|
||||
Reference in New Issue
Block a user