Skip to content

Resolve $ref/$defs when converting JSON schema to Python types - #1908

Closed
chakshu-dhannawat wants to merge 4 commits into
dottxt-ai:mainfrom
chakshu-dhannawat:fix/json-schema-ref-resolution
Closed

Resolve $ref/$defs when converting JSON schema to Python types#1908
chakshu-dhannawat wants to merge 4 commits into
dottxt-ai:mainfrom
chakshu-dhannawat:fix/json-schema-ref-resolution

Conversation

@chakshu-dhannawat

Copy link
Copy Markdown

What

schema_type_to_python (in outlines/types/json_schema_utils.py) never handled $ref. Pydantic emits nested models as references ({"$ref": "#/$defs/Address"}) with the real schema under the root $defs, so any nested-model field fell through to Any. The nested structure was silently lost when a raw JSON schema was converted to a pydantic model / TypedDict / dataclass.

This matters in practice because JsonSchema.convert_to(schema, ["dataclass", "typeddict", "pydantic"]) is exactly the path the Gemini backend uses when you pass a JSON schema as the output type, so nested objects were not being constrained.

Repro (before)

from pydantic import BaseModel
from typing import List
from outlines.types.dsl import JsonSchema

class Address(BaseModel):
    city: str
    zip: str

class Person(BaseModel):
    name: str
    address: Address
    tags: List[Address]

M = JsonSchema.convert_to(Person.model_json_schema(), ["pydantic"])
print(M.model_fields["address"].annotation)  # typing.Any   <-- structure lost
print(M.model_fields["tags"].annotation)     # typing.List[typing.Any]

Fix

  • Resolve $ref against the root $defs/definitions and thread it through the recursion (array items, type-list unions, and the object -> converter calls).
  • Guard against reference cycles with a seen set of ref names on the current path, so self-referential (recursive) models degrade the cyclic edge to Any instead of recursing forever. Without this a wrapped recursive model would hit RecursionError.
  • Unknown/unresolvable refs fall back to Any (no crash), same as before.

After the fix, address resolves to the rebuilt Address model, tags to List[Address], and the reconstructed model validates the nested structure end to end.

Scope

I deliberately kept this to direct $ref resolution. Optional[NestedModel] is emitted as anyOf: [{$ref}, {"type": "null"}]; anyOf is a separate pre-existing limitation this PR doesn't touch, so that specific case still widens to Any (unchanged behavior, not a regression).

Tests

Added regression tests in tests/types/test_json_schema_utils.py covering nested $ref across pydantic/typeddict/dataclass, a dangling ref, and a recursive model. They fail on the current code and pass with the fix. Full tests/types/ suite is green; ruff and mypy clean.

schema_type_to_python did not handle $ref, so any nested model reference
(the shape Pydantic emits for nested models, e.g. {"$ref": "#/$defs/Address"})
fell through to Any. This silently dropped the nested structure when a raw
JSON schema was converted to a pydantic/typeddict/dataclass, which is the
path Gemini uses for JSON-schema output types.

Resolve $ref against the root $defs/definitions and thread it through the
recursion. A seen-set guards against reference cycles so self-referential
(recursive) models degrade the cyclic edge to Any instead of recursing
forever. Unknown refs also fall back to Any.

Adds regression tests covering nested refs across all three targets, a
dangling ref, and a recursive model.

@ErenAta16 ErenAta16 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Checked this against the actual behavior on main first, then ran the fix. Both hold up.

On main, the bug is exactly as described: a nested pydantic model comes through as {"$ref": "#/$defs/Address"}, schema_type_to_python has no $ref branch, so it falls past every type check and returns Any. Confirmed by converting a Person schema with a nested Address ref, the address field on the resulting model came out as typing.Any, the whole nested structure gone.

With this PR applied, the same conversion resolves address to a real Address model with city: str. The threading of defs/seen through every recursive call site (type-array union members, array items, and all three of the object converters) looks complete, I didn't find a recursion path that drops them.

The part I most wanted to check was the cycle handling, since a self-referential model is the obvious way to make $ref resolution loop forever. Built a Node schema whose child field refs Node itself. It resolves without a RecursionError: the first level materializes as a real Node model, and the repeated ref (now in seen) degrades to Any, so child ends up Optional[Node]. That's the right call, you can't build an infinitely-nested Python type anyway, so breaking after one level and widening is the sensible degradation rather than crashing.

Two small things, both non-blocking:

  • The $ref split takes the last path segment (ref.split("/")[-1]), which is right for the pydantic #/$defs/Name shape. It won't resolve refs into external documents or nested $defs paths, but those aren't what pydantic emits, so scoping it to the local-$defs case is reasonable. Might be worth a one-line note in the docstring that only local $defs/definitions refs are handled.
  • defs is read from the root schema only on the top-level call (if defs is None). That's correct for pydantic output, where all definitions sit at the root, just flagging it as an assumption in case a schema nests its own $defs deeper.

Solid fix, and good instinct adding the cycle guard rather than leaving it as a latent stack overflow.

@chakshu-dhannawat

Copy link
Copy Markdown
Author

Thanks for the thorough review, especially checking the self-referential cycle case, that was the part I was most careful about too.

Good call on the docstring. Added a note (1f9768e) that only local $defs/definitions refs are resolved (last path segment), and that external documents / deeper $defs paths fall through to Any since that isn't what Pydantic emits. That covers both points you raised: the local-ref scoping and the root-only defs assumption both come down to the same "we target Pydantic's output shape" decision, so I folded them into one note where the resolution happens.

@ErenAta16

Copy link
Copy Markdown
Contributor

The docstring note covers it, and folding both points into one at the resolution site was the right call. They really are the same decision, and two separate notes would have implied two independent limitations.

Nothing further from me on this one.

@ErenAta16 ErenAta16 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Resolution works for a bare {"$ref": ...}, but Pydantic wraps optional nested models in anyOf, and there's no anyOf branch — so the most common shape falls back to Any.

class Address(BaseModel):
    street: str

class Holder(BaseModel):
    req: Address                    # {"$ref": "#/$defs/Address"}
    opt: Optional[Address] = None   # {"anyOf": [{"$ref": "#/$defs/Address"}, {"type": "null"}]}
    lst: List[Address] = []
    mp:  Dict[str, Address] = {}

json_schema_dict_to_pydantic(Holder.model_json_schema()) on this branch:

field resolved to
req Address
opt Optional[Any]
lst Optional[List[Address]]
mp Optional[Mp]

The type-is-a-list branch handles {"type": ["string", "null"]}, but Optional[Address] isn't that shape — it's an anyOf of two sub-schemas. schema.get("type") is None, nothing matches, and it falls through to the Any at the bottom. So $ref resolution reaches required nested models and misses optional ones, which is where most nested models live.

Second one: a self-referential model produces a model with no fields at all.

class Node(BaseModel):
    v: int
    child: Optional["Node"] = None
Node.model_rebuild()

Node.model_json_schema()
# {"$defs": {"Node": {...}}, "$ref": "#/$defs/Node"}

json_schema_dict_to_pydantic(Node.model_json_schema()).model_fields
# {}

For a recursive model Pydantic makes the root a $ref and puts everything under $defs. json_schema_dict_to_pydantic reads schema.get("properties", {}) off that root directly — the new $ref handling lives in schema_type_to_python, which never sees the root. properties is empty, so create_model gets no fields and returns silently. That's the case the new seen docstring names ("break reference cycles, e.g. self-referential models"), and it doesn't reach the cycle guard at all.

Both are covered by resolving $ref at the entry of the three json_schema_dict_to_* functions after defs is populated, and adding an anyOf/oneOf branch alongside the type-list one that maps each sub-schema through schema_type_to_python and unions the results.

Dict[str, Address] coming back as a generated Mp model rather than Dict[str, Address] is pre-existing, not this PR.

Measured on outlines @ 13fe8dd with the PR applied, Python 3.10.12.

@ErenAta16 ErenAta16 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#1963 is open against the same function and already covers both of the cases above, so the anyOf branch I suggested would be duplicated work.

Same Holder / Node script, each PR applied to 13fe8dd2 on its own:

req: Address opt: Optional[Address] Node fields
this PR Address Optional[Any] []
#1963 Address Optional[Address] ['v', 'child']
#1951 Any Optional[Any] []

#1963 resolves the root-level $ref before reading properties, which is what gets Node its fields back, and it handles the combinator keywords, which is what gets opt its type back.

#1951 and #1949 are also open against schema_type_to_python — four PRs on the same function right now, which is worth someone deciding an order on before any of them lands and makes the others conflict.

Measured on outlines @ 13fe8dd, Python 3.10.12.

@chakshu-dhannawat

Copy link
Copy Markdown
Author

@ErenAta16 thanks for the detailed review. You are right: this PR resolves the bare $ref case but misses anyOf wrapped nested models and root-level $ref for self-referential types. I see @DerMayer1 opened #1963 which covers both of those gaps plus oneOf/allOf more generally, so continuing this smaller PR would just create duplicate work and likely conflicts. I am closing this in favor of #1963 and will move my review/testing attention there.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants