diff --git a/src/outlines/types/json_schema_utils.py b/src/outlines/types/json_schema_utils.py index 1b65d1894..255983ffb 100644 --- a/src/outlines/types/json_schema_utils.py +++ b/src/outlines/types/json_schema_utils.py @@ -14,7 +14,9 @@ def schema_type_to_python( schema: dict, - caller_target_type: Literal["pydantic", "typeddict", "dataclass"] + caller_target_type: Literal["pydantic", "typeddict", "dataclass"], + defs: Optional[dict] = None, + seen: frozenset = frozenset(), ) -> Any: """Get a Python type from a JSON Schema dict. @@ -24,6 +26,12 @@ def schema_type_to_python( The JSON Schema dict to convert to a Python type caller_target_type: Literal["pydantic", "typeddict", "dataclass"] The type of the caller + defs: Optional[dict] + The ``$defs``/``definitions`` block of the root schema, used to resolve + ``$ref`` references to nested definitions. + seen: frozenset + The ``$ref`` names already being resolved on the current path, used to + break reference cycles (e.g. self-referential models). Returns ------- @@ -31,6 +39,23 @@ def schema_type_to_python( The Python type """ + if "$ref" in schema: + # Pydantic emits nested models as ``{"$ref": "#/$defs/Name"}`` with the + # referenced schema stored under the root ``$defs``. Resolve it so the + # nested structure is preserved instead of silently widening to ``Any``. + # Only local ``$defs``/``definitions`` refs are handled (we take the last + # path segment); external documents and deeper ``$defs`` paths aren't + # what Pydantic emits, so they fall through to ``Any`` below. + ref_name = schema["$ref"].split("/")[-1] + resolved = (defs or {}).get(ref_name) + if resolved is None or ref_name in seen: + # Unknown ref, or a cycle back to a def already being resolved + # (recursive model): degrade to ``Any`` rather than recurse forever. + return Any + return schema_type_to_python( + resolved, caller_target_type, defs, seen | {ref_name} + ) + if "enum" in schema: values = schema["enum"] return Literal[tuple(values)] @@ -50,7 +75,7 @@ def schema_type_to_python( # Python type and combine them into a Union (mirroring the ``anyOf`` # the regex backend uses for type arrays). members = tuple( - schema_type_to_python({**schema, "type": member}, caller_target_type) + schema_type_to_python({**schema, "type": member}, caller_target_type, defs, seen) for member in t ) return Union[members] if members else Any # type: ignore @@ -68,25 +93,27 @@ def schema_type_to_python( elif t == "array": items = schema.get("items", {}) if items: - item_type = schema_type_to_python(items, caller_target_type) + item_type = schema_type_to_python(items, caller_target_type, defs, seen) else: item_type = Any return List[item_type] # type: ignore elif t == "object": name = schema.get("title") if caller_target_type == "pydantic": - return json_schema_dict_to_pydantic(schema, name) + return json_schema_dict_to_pydantic(schema, name, defs, seen) elif caller_target_type == "typeddict": - return json_schema_dict_to_typeddict(schema, name) + return json_schema_dict_to_typeddict(schema, name, defs, seen) elif caller_target_type == "dataclass": - return json_schema_dict_to_dataclass(schema, name) + return json_schema_dict_to_dataclass(schema, name, defs, seen) return Any def json_schema_dict_to_typeddict( schema: dict, - name: Optional[str] = None + name: Optional[str] = None, + defs: Optional[dict] = None, + seen: frozenset = frozenset(), ) -> _TypedDictMeta: """Convert a JSON Schema dict into a TypedDict class. @@ -96,6 +123,12 @@ def json_schema_dict_to_typeddict( The JSON Schema dict to convert to a TypedDict name: Optional[str] The name of the TypedDict + defs: Optional[dict] + The root schema's ``$defs`` used to resolve ``$ref`` references. When + ``None``, it is read from this schema (the top-level call). + seen: frozenset + The ``$ref`` names already being resolved on the current path, used to + break reference cycles. Returns ------- @@ -103,13 +136,15 @@ def json_schema_dict_to_typeddict( The TypedDict class """ + if defs is None: + defs = schema.get("$defs") or schema.get("definitions") or {} required = set(schema.get("required", [])) properties = schema.get("properties", {}) annotations: Dict[str, Any] = {} for property, details in properties.items(): - typ = schema_type_to_python(details, "typeddict") + typ = schema_type_to_python(details, "typeddict", defs, seen) if property not in required: # NotRequired (PEP 655) marks the KEY optional; Optional only makes the # value nullable, leaving the key required on a total=True TypedDict. @@ -121,7 +156,9 @@ def json_schema_dict_to_typeddict( def json_schema_dict_to_pydantic( schema: dict, - name: Optional[str] = None + name: Optional[str] = None, + defs: Optional[dict] = None, + seen: frozenset = frozenset(), ) -> type[BaseModel]: """Convert a JSON Schema dict into a Pydantic BaseModel class. @@ -131,6 +168,12 @@ def json_schema_dict_to_pydantic( The JSON Schema dict to convert to a Pydantic BaseModel name: Optional[str] The name of the Pydantic BaseModel + defs: Optional[dict] + The root schema's ``$defs`` used to resolve ``$ref`` references. When + ``None``, it is read from this schema (the top-level call). + seen: frozenset + The ``$ref`` names already being resolved on the current path, used to + break reference cycles. Returns ------- @@ -138,13 +181,15 @@ def json_schema_dict_to_pydantic( The Pydantic BaseModel class """ + if defs is None: + defs = schema.get("$defs") or schema.get("definitions") or {} required = set(schema.get("required", [])) properties = schema.get("properties", {}) field_definitions: Dict[str, Any] = {} for property, details in properties.items(): - typ = schema_type_to_python(details, "pydantic") + typ = schema_type_to_python(details, "pydantic", defs, seen) if property not in required: field_definitions[property] = (Optional[typ], None) else: @@ -155,7 +200,9 @@ def json_schema_dict_to_pydantic( def json_schema_dict_to_dataclass( schema: dict, - name: Optional[str] = None + name: Optional[str] = None, + defs: Optional[dict] = None, + seen: frozenset = frozenset(), ) -> type: """Convert a JSON Schema dict into a dataclass. @@ -165,6 +212,12 @@ def json_schema_dict_to_dataclass( The JSON Schema dict to convert to a dataclass name: Optional[str] The name of the dataclass + defs: Optional[dict] + The root schema's ``$defs`` used to resolve ``$ref`` references. When + ``None``, it is read from this schema (the top-level call). + seen: frozenset + The ``$ref`` names already being resolved on the current path, used to + break reference cycles. Returns ------- @@ -172,6 +225,8 @@ def json_schema_dict_to_dataclass( The dataclass """ + if defs is None: + defs = schema.get("$defs") or schema.get("definitions") or {} required = set(schema.get("required", [])) properties = schema.get("properties", {}) @@ -179,7 +234,7 @@ def json_schema_dict_to_dataclass( defaults: Dict[str, Any] = {} for property, details in properties.items(): - typ = schema_type_to_python(details, "dataclass") + typ = schema_type_to_python(details, "dataclass", defs, seen) annotations[property] = typ if property not in required: diff --git a/tests/types/test_json_schema_utils.py b/tests/types/test_json_schema_utils.py index 52895da59..8c4011e4a 100644 --- a/tests/types/test_json_schema_utils.py +++ b/tests/types/test_json_schema_utils.py @@ -1,6 +1,6 @@ import sys from dataclasses import is_dataclass -from typing import Any, List, Literal, Optional, Union +from typing import Any, List, Literal, Optional, Union, get_args, get_origin from pydantic import BaseModel, TypeAdapter from pydantic_core import PydanticUndefined @@ -321,6 +321,96 @@ def test_json_schema_dict_to_pydantic_nested_object(): assert field.model_fields["age"].default is None +def test_schema_type_to_python_ref(): + # A $ref is resolved against the provided defs instead of widening to Any. + defs = { + "Address": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + } + } + result = schema_type_to_python({"$ref": "#/$defs/Address"}, "pydantic", defs) + assert issubclass(result, BaseModel) + assert result.model_fields["city"].annotation is str + + # A $ref that cannot be resolved falls back to Any rather than crashing. + assert schema_type_to_python({"$ref": "#/$defs/Missing"}, "pydantic", {}) is Any + assert schema_type_to_python({"$ref": "#/$defs/Missing"}, "pydantic") is Any + + +def test_json_schema_dict_to_pydantic_ref(): + # This is the shape Pydantic emits for nested models: the nested schema is + # stored under $defs and referenced with $ref. Before $ref resolution these + # fields silently degraded to Any, dropping the nested constraint. + class Address(BaseModel): + city: str + zip: str + + class Person(BaseModel): + name: str + address: Address + tags: List[Address] + + result = json_schema_dict_to_pydantic(Person.model_json_schema()) + + address = result.model_fields["address"].annotation + assert issubclass(address, BaseModel) + assert address.model_fields["city"].annotation is str + assert address.model_fields["zip"].annotation is str + + tags = result.model_fields["tags"].annotation + assert get_origin(tags) is list + item = get_args(tags)[0] + assert issubclass(item, BaseModel) + assert set(item.model_fields) == {"city", "zip"} + + # The reconstructed model enforces the nested structure end to end. + result(name="x", address={"city": "NY", "zip": "1"}, tags=[{"city": "A", "zip": "2"}]) + + +def test_json_schema_dict_to_pydantic_recursive_ref(): + # A self-referential model must not send $ref resolution into infinite + # recursion; the cyclic edge degrades to Any while the rest resolves. + class Node(BaseModel): + value: int + children: List["Node"] = [] + + class Tree(BaseModel): + root: Node + + result = json_schema_dict_to_pydantic(Tree.model_json_schema()) + node = result.model_fields["root"].annotation + assert issubclass(node, BaseModel) + assert node.model_fields["value"].annotation is int + + +def test_json_schema_dict_to_typeddict_ref(): + class Address(BaseModel): + city: str + + class Person(BaseModel): + address: Address + + result = json_schema_dict_to_typeddict(Person.model_json_schema()) + address = result.__annotations__["address"] + assert isinstance(address, _TypedDictMeta) + assert address.__annotations__["city"] is str + + +def test_json_schema_dict_to_dataclass_ref(): + class Address(BaseModel): + city: str + + class Person(BaseModel): + address: Address + + result = json_schema_dict_to_dataclass(Person.model_json_schema()) + address = result.__annotations__["address"] + assert is_dataclass(address) + assert address.__annotations__["city"] is str + + def test_json_schema_dict_to_dataclass_basic(): schema = { "type": "object",