From 6d33cda010971adc45220842f30033576fe02b5c Mon Sep 17 00:00:00 2001 From: Eli Date: Wed, 2 Sep 2026 10:22:33 -0400 Subject: [PATCH 1/2] Address `$ref`s so a schema survives both Pydantic and the wire (#761) `_inline_refs` followed every `$ref` to its target and deleted every `$defs`. On a recursive type there is no fixed point to reach, so it recurred until the stack ran out -- the `RecursionError` in #761. That reached users through `Encodable`, since `_pydantic_type_tuple` inlines the schema of a model built over the element types: type A = int | list[A] pydantic.TypeAdapter(Encodable[tuple[A, int]]).json_schema() Inlining was only ever a way to keep a `$ref` out of a `WithJsonSchema` value, which pydantic/pydantic#12145 rejects. Re-addressing the references achieves that without expanding anything, and a reference is what the providers want in any case: measured against five of them with the returned arguments validated against the intended schema, every ref-free encoding of a cycle either fails the request or corrupts the data, while `$defs` and `$ref` are accepted and correctly filled. It is also what OpenAI documents. So `_bundle_refs` renames each definition to a content hash, anchors it with an `$id` and points the references at that absolute URI -- the one ref form pydantic's counter tolerates, and the only one meaningful in a fragment, since `#/$defs/X` addresses the root of whatever document the fragment is embedded in rather than its own. `_rebundle` undoes that once the document is whole: definitions hoisted to the root, references back to pointers, and the keywords beside a `$ref` dropped, since a provider reads a reference as the whole subschema rather than composing it with what sits beside it. Recursive types remain unsupported on the `openrouter/google/...` route, which drops a `$ref` and sends the object it named as a string. That is OpenRouter's translation rather than Gemini -- a direct `gemini/` call on a 2.0+ model passes `$defs` through -- and it already affects ordinary repeated field types today, so it is left as its own issue. Co-Authored-By: Claude Opus 5 (1M context) --- .../handlers/llm/harness/serialization.py | 108 +++++++-- .../handlers/llm/harness/synthesis/body.py | 6 +- .../llm/harness/synthesis/function.py | 4 +- ...test_handlers_llm_harness_serialization.py | 220 ++++++++++++++++++ 4 files changed, 308 insertions(+), 30 deletions(-) diff --git a/effectful/handlers/llm/harness/serialization.py b/effectful/handlers/llm/harness/serialization.py index 713de936f..117405807 100644 --- a/effectful/handlers/llm/harness/serialization.py +++ b/effectful/handlers/llm/harness/serialization.py @@ -10,6 +10,7 @@ import contextvars import dataclasses import functools +import hashlib import inspect import io import json @@ -284,28 +285,81 @@ def emit(block: OpenAIMessageContentListBlock) -> None: return [ChatCompletionTextObject(type="text", text=heading), *blocks] -def _inline_refs(schema: dict) -> dict: - """Inline ``$ref`` pointers so ``WithJsonSchema`` never emits orphan refs. +# Base URI for the definitions a fragment carries. An absolute URI is the one ref +# form Pydantic's ref counter tolerates, and an `$id`-bearing subschema is a +# 2020-12 embedded resource, so a ref to it resolves wherever the fragment does. +_DEFS_BASE = "https://effectful.invalid/$defs/" - Workaround for https://github.com/pydantic/pydantic/issues/12145 — - Pydantic's ``GenerateJsonSchema`` does not merge user-provided ``$defs`` - into its internal ref map, so any ``$ref`` in a ``WithJsonSchema`` value - causes a ``KeyError`` when the annotated type is composed into a model. + +def _bundle_refs(schema: dict) -> dict: + """``schema`` with its definitions addressed absolutely rather than by pointer. + + Works around https://github.com/pydantic/pydantic/issues/12145: Pydantic's + ``GenerateJsonSchema`` does not merge a user-provided ``$defs`` into its + internal ref map, so a ``#/$defs/...`` ref in a `pydantic.WithJsonSchema` + value raises ``KeyError`` once the annotated type is composed into a model. """ - defs = schema.get("$defs", {}) + defs = schema.get("$defs") or {} + if not defs: + return schema + uris = { + name: _DEFS_BASE + + hashlib.sha256( + json.dumps(body, sort_keys=True, default=str).encode() + ).hexdigest()[:16] + for name, body in defs.items() + } + + def walk(node: typing.Any) -> typing.Any: + if isinstance(node, list): + return [walk(item) for item in node] + if not isinstance(node, dict): + return node + return { + k: uris.get(v.removeprefix("#/$defs/"), v) + if k == "$ref" and isinstance(v, str) + else walk(v) + for k, v in node.items() + } + + bundled = walk(schema) + bundled["$defs"] = { + uris[name]: {"$id": uris[name], **body} + for name, body in bundled["$defs"].items() + } + return bundled + - def _resolve(obj): - if isinstance(obj, dict): - if "$ref" in obj: - ref_name = obj["$ref"].split("/")[-1] - if ref_name in defs: - return _resolve(defs[ref_name]) - return {k: _resolve(v) for k, v in obj.items() if k != "$defs"} - if isinstance(obj, list): - return [_resolve(item) for item in obj] - return obj +def _rebundle(schema: dict) -> dict: + """``schema`` in the JSON Schema subset the providers implement. - return _resolve(schema) + They take definitions only at the document root, refuse a remote reference + outright, and reject any keyword sitting beside a ``$ref``. + """ + defs: dict[str, dict] = {} + + def walk(node: typing.Any) -> typing.Any: + if isinstance(node, list): + return [walk(item) for item in node] + if not isinstance(node, dict): + return node + for name, body in (node.get("$defs") or {}).items(): + uri = body.get("$id") + defs[uri.removeprefix(_DEFS_BASE) if uri else name] = walk( + {k: v for k, v in body.items() if k != "$id"} + ) + if isinstance(node.get("$ref"), str): + # Siblings go too, which a provider requires: it reads a `$ref` as + # the whole subschema rather than composing it with what sits + # beside it, the way draft 2020-12 does. + ref = node["$ref"] + if ref.startswith(_DEFS_BASE): + ref = f"#/$defs/{ref.removeprefix(_DEFS_BASE)}" + return {"$ref": ref} + return {k: walk(v) for k, v in node.items() if k != "$defs"} + + document = walk(schema) + return {**document, "$defs": defs} if defs else document @dataclasses.dataclass(frozen=True, eq=True) @@ -590,7 +644,7 @@ def _pydantic_type_complex(ty): """Encode ``complex`` as ``{"real": float, "imag": float}``.""" schema = pydantic.TypeAdapter(_ComplexModel).json_schema() - schema = _ensure_strict_json_schema(_inline_refs(schema), path=(), root={}) + schema = _ensure_strict_json_schema(_bundle_refs(schema), path=(), root={}) return typing.Annotated[ ty, @@ -635,7 +689,7 @@ def _nt_serialize(value, info: pydantic.SerializationInfo): return typing.Annotated[ ty, pydantic.PlainSerializer(_nt_serialize), - pydantic.WithJsonSchema(_inline_refs(nt_model.model_json_schema())), + pydantic.WithJsonSchema(_bundle_refs(nt_model.model_json_schema())), ] args = typing.get_args(ty) @@ -675,7 +729,7 @@ def _serialize(value, info: pydantic.SerializationInfo): ty, pydantic.BeforeValidator(_decode), pydantic.PlainSerializer(_serialize), - pydantic.WithJsonSchema(_inline_refs(model.model_json_schema())), + pydantic.WithJsonSchema(_bundle_refs(model.model_json_schema())), ] @@ -722,7 +776,7 @@ def _pydantic_type_image(ty: type[Image.Image]): pydantic.InstanceOf, pydantic.BeforeValidator(_validate_image), pydantic.PlainSerializer(_serialize_image), - pydantic.WithJsonSchema(_inline_refs(adapter.json_schema())), + pydantic.WithJsonSchema(_bundle_refs(adapter.json_schema())), ] @@ -871,7 +925,7 @@ def _serialize_name_and_tool(value: _NameAndTool) -> ChatCompletionToolParam: "function": { "name": name, "description": description, - "parameters": response_format["json_schema"]["schema"], + "parameters": _rebundle(response_format["json_schema"]["schema"]), "strict": True, }, } @@ -888,7 +942,7 @@ def _pydantic_type_name_and_tool(ty: type[_NameAndTool]): to build a `TypeAdapter` for the `Tool` field and fail. """ schema = pydantic.TypeAdapter(ChatCompletionToolParam).json_schema() - schema = _ensure_strict_json_schema(_inline_refs(schema), path=(), root={}) + schema = _ensure_strict_json_schema(_bundle_refs(schema), path=(), root={}) return typing.Annotated[ ty, pydantic.InstanceOf, @@ -969,7 +1023,7 @@ def _pydantic_type_tool_call(ty: type[DecodedToolCall]): # Use OpenAI's ChatCompletionMessageToolCall (has actual fields: id, function, # type) rather than litellm's (empty dict with extra="allow"). schema = OpenAIChatCompletionMessageToolCall.model_json_schema() - schema = _ensure_strict_json_schema(_inline_refs(schema), path=(), root={}) + schema = _ensure_strict_json_schema(_bundle_refs(schema), path=(), root={}) return typing.Annotated[ ty, pydantic.InstanceOf, @@ -1028,3 +1082,7 @@ def base(tool: Tool) -> str: class _BoxedResponse[T](pydantic.BaseModel): value: T + + @classmethod + def model_json_schema(cls, *args, **kwargs) -> dict[str, typing.Any]: + return _rebundle(super().model_json_schema(*args, **kwargs)) diff --git a/effectful/handlers/llm/harness/synthesis/body.py b/effectful/handlers/llm/harness/synthesis/body.py index a690f27b1..c64ba22fd 100644 --- a/effectful/handlers/llm/harness/synthesis/body.py +++ b/effectful/handlers/llm/harness/synthesis/body.py @@ -65,7 +65,7 @@ DecodedToolCall, EncodedFunction, TypeToPydanticType, - _inline_refs, + _bundle_refs, _serialize_callable, ) from effectful.handlers.llm.harness.synthesis.function import ( @@ -286,7 +286,7 @@ def _validate( pydantic.BeforeValidator(_validate), pydantic.PlainSerializer(lambda value: _serialize_callable(value)), pydantic.WithJsonSchema( - _inline_refs(pydantic.TypeAdapter(typed_enc).json_schema()), + _bundle_refs(pydantic.TypeAdapter(typed_enc).json_schema()), mode="validation", ), pydantic.WithJsonSchema( @@ -456,7 +456,7 @@ def _doctest_apply(op, *args, **kwargs): pydantic.BeforeValidator(_validate), pydantic.PlainSerializer(_serialize_callable), pydantic.WithJsonSchema( - _inline_refs(pydantic.TypeAdapter(typed_enc).json_schema()), + _bundle_refs(pydantic.TypeAdapter(typed_enc).json_schema()), mode="validation", ), pydantic.WithJsonSchema( diff --git a/effectful/handlers/llm/harness/synthesis/function.py b/effectful/handlers/llm/harness/synthesis/function.py index b95c47098..6a82f5648 100644 --- a/effectful/handlers/llm/harness/synthesis/function.py +++ b/effectful/handlers/llm/harness/synthesis/function.py @@ -16,7 +16,7 @@ _TYPE_CHECK_ANCHOR_KEY, EncodedFunction, TypeToPydanticType, - _inline_refs, + _bundle_refs, _serialize_callable, ) @@ -398,7 +398,7 @@ def _validate( pydantic.BeforeValidator(_validate), pydantic.PlainSerializer(_serialize_callable), pydantic.WithJsonSchema( - _inline_refs(pydantic.TypeAdapter(typed_enc).json_schema()), + _bundle_refs(pydantic.TypeAdapter(typed_enc).json_schema()), mode="validation", ), pydantic.WithJsonSchema( diff --git a/tests/test_handlers_llm_harness_serialization.py b/tests/test_handlers_llm_harness_serialization.py index 4297bddfa..5508557a8 100644 --- a/tests/test_handlers_llm_harness_serialization.py +++ b/tests/test_handlers_llm_harness_serialization.py @@ -21,6 +21,7 @@ TypedDict, TypeVar, Union, + cast, ) import litellm @@ -34,11 +35,18 @@ RestrictedPythonExecutor, ) from effectful.handlers.llm.harness.serialization import ( + _DEFS_BASE, _NAME2TOOL_KEY, _TYPE_CHECK_ANCHOR_KEY, CONTENT_BLOCK_TYPES, DecodedToolCall, + EncodedFunction, + _BoxedResponse, + _bundle_refs, + _ComplexModel, _NameAndTool, + _rebundle, + _serialize_name_and_tool, to_content_blocks, ) from effectful.handlers.llm.harness.validation.ty import TyTypeChecker @@ -178,6 +186,32 @@ class _PersonWithAddressModel(pydantic.BaseModel): type _RecursiveAlias = int | list[_RecursiveAlias] type _RecursiveGenAlias[T] = T | list[_RecursiveGenAlias[T]] + +class _Tree(pydantic.BaseModel): + """A model that refers to itself, so its schema needs a definition.""" + + label: str + kids: list["_Tree"] = [] + + +class _TreeParams(pydantic.BaseModel): + tree: _Tree + + +@dataclass +class _Place: + city: str + + +@dataclass +class _Resident: + """A definition reached twice without being recursive.""" + + name: str + home: _Place + work: _Place + + # Stands in for `Kernel` in `docs/source/llm_examples/optimization/kernels.py`: # an alias that is in a skill's lexical *scope*, so the alias object itself # reaches the encoding as a value. @@ -942,6 +976,192 @@ def test_recursive_alias_does_not_diverge(ty): assert "$ref" in schema, "recursion should resolve to a reference, not inline" +# --------------------------------------------------------------------------- +# Reference addressing +# --------------------------------------------------------------------------- + + +def _refs(node): + """Every ``$ref`` string anywhere in `node`.""" + if isinstance(node, dict): + for key, value in node.items(): + if key == "$ref": + yield value + else: + yield from _refs(value) + elif isinstance(node, list): + for item in node: + yield from _refs(item) + + +def _defs_paths(node, path=()): + """Where in `node` a ``$defs`` block sits.""" + if isinstance(node, dict): + for key, value in node.items(): + if key == "$defs": + yield path + else: + yield from _defs_paths(value, path + (key,)) + elif isinstance(node, list): + for i, item in enumerate(node): + yield from _defs_paths(item, path + (i,)) + + +@pytest.mark.parametrize( + "ty", [_RecursiveAlias, _RecursiveGenAlias[int]], ids=["plain", "generic"] +) +def test_a_recursive_schema_can_be_bundled(ty): + """A self-referential schema bundles without diverging -- issue #761.""" + bundled = _bundle_refs(pydantic.TypeAdapter(Encodable[ty]).json_schema()) + assert not any(r.startswith("#") for r in _refs(bundled)), ( + "a fragment cannot carry a pointer to a root it does not have" + ) + + +@pytest.mark.parametrize( + "ty", + [tuple[_RecursiveAlias, int], list[_RecursiveAlias]], + ids=["tuple", "list"], +) +def test_a_recursive_type_has_an_encoding(ty): + """`_pydantic_type_tuple` bundles the schema of a model over the element + types, so a recursive element reaches the bundler through `Encodable`.""" + assert pydantic.TypeAdapter(Encodable[ty]).json_schema() + + +def test_a_bundled_fragment_composes_into_a_model(): + """The reason the bundler exists: pydantic/pydantic#12145 rejects a local + ``$ref`` in a `WithJsonSchema` value once the type is composed.""" + fragment = _bundle_refs( + pydantic.TypeAdapter(Encodable[_RecursiveAlias]).json_schema() + ) + + class Composed(pydantic.BaseModel): + recursive: Annotated[Any, pydantic.WithJsonSchema(fragment)] + plain: int + + assert Composed.model_json_schema() + + +@pytest.mark.parametrize( + "schema", + [ + pydantic.TypeAdapter(litellm.ChatCompletionToolParam).json_schema(), + pydantic.TypeAdapter(_ComplexModel).json_schema(), + EncodedFunction.model_json_schema(), + ], + ids=["tool-param", "complex", "encoded-function"], +) +def test_bundling_only_readdresses(schema): + """The definitions and their structure are pydantic's; only how a reference + names one changes.""" + bundled = _bundle_refs(json.loads(json.dumps(schema))) + assert not any(r.startswith("#") for r in _refs(bundled)) + assert sorted(_defs_paths(bundled)) == sorted(_defs_paths(schema)) + assert len(bundled.get("$defs") or {}) == len(schema.get("$defs") or {}) + + +def test_bundling_keeps_the_keywords_beside_a_ref(): + """In 2020-12 a ``$ref`` composes with its siblings rather than replacing + them, and pydantic puts a field's ``description`` there.""" + + class _Inner(pydantic.BaseModel): + x: int + + class _Outer(pydantic.BaseModel): + described: _Inner = pydantic.Field(..., description="what it is for") + plain: _Inner + + described = _bundle_refs(_Outer.model_json_schema())["properties"]["described"] + assert described["description"] == "what it is for" + assert described["$ref"].startswith(_DEFS_BASE) + + +def test_rebundling_narrows_a_document_to_what_providers_accept(): + """Four separate rejections: a remote reference, a keyword beside a + ``$ref``, a ``$defs`` anywhere but the root, and a reference in any other + form are each answered with a ``BadRequestError``.""" + document = litellm.utils.type_to_response_format_param(_TreeParams)["json_schema"][ + "schema" + ] + wire = _rebundle(json.loads(json.dumps(document))) + + assert "$id" not in json.dumps(wire) + assert all(r.startswith("#/$defs/") for r in _refs(wire)) + assert list(_defs_paths(wire)) == [()] + assert list(_refs(wire)), "the recursion should survive as a reference" + + def no_siblings(node): + if isinstance(node, dict): + assert "$ref" not in node or not set(node) - {"$ref"}, node + for value in node.values(): + no_siblings(value) + elif isinstance(node, list): + for item in node: + no_siblings(item) + + no_siblings(wire) + + +def test_rebundling_a_ref_that_carries_a_description_terminates(): + """`_ensure_strict_json_schema` unravels such a ``$ref`` by inlining its + target and has no guard against doing so forever, so the siblings on the + references that survive are dropped before it runs.""" + schema = { + "type": "object", + "$defs": { + "N": { + "type": "object", + "properties": { + "kids": { + "type": "array", + "items": {"$ref": "#/$defs/N", "description": "a child"}, + } + }, + "required": ["kids"], + "additionalProperties": False, + } + }, + "properties": {"tree": {"$ref": "#/$defs/N", "title": "Tree"}}, + "required": ["tree"], + "additionalProperties": False, + } + assert _rebundle(schema) + + +@pytest.mark.parametrize("ty", [_Tree, _Resident], ids=["recursive", "repeated"]) +def test_a_tool_parameter_is_advertised_by_root_pointer(ty): + def plant(subject: ty) -> str: + """Plant {subject}.""" + return "" + + tool = cast(Tool, plant) + tool.__signature__ = inspect.signature(plant) + parameters = _serialize_name_and_tool(_NameAndTool("plant", tool))["function"][ + "parameters" + ] + assert list(_refs(parameters)) + assert all(r.startswith("#/$defs/") for r in _refs(parameters)) + assert list(_defs_paths(parameters)) == [()] + + +def test_a_recursive_response_type_is_advertised_as_a_reference(): + """Carried on the boxing model itself, which is where litellm reads a + request's ``response_format`` from.""" + boxed = pydantic.create_model( + "BoxedResponse", value=Encodable[_Tree], __base__=_BoxedResponse + ) + for schema in ( + boxed.model_json_schema(), + # And through litellm, which runs its own strict pass over what it finds + # here; the references have to survive that too. + litellm.utils.type_to_response_format_param(boxed)["json_schema"]["schema"], + ): + assert list(_refs(schema)), "the recursion should survive as a reference" + assert all(r.startswith("#/$defs/") for r in _refs(schema)) + assert list(_defs_paths(schema)) == [()] + + # --------------------------------------------------------------------------- # An alias as a *value* # --------------------------------------------------------------------------- From 3cef8a5176960e8ed5cccdd1f39f0488cd6da208 Mon Sep 17 00:00:00 2001 From: Eli Date: Wed, 2 Sep 2026 18:55:35 -0400 Subject: [PATCH 2/2] Let Pydantic own the references, per review The reference handling this branch added is unnecessary. Every `WithJsonSchema` it worked around was passing a schema *dict* generated from a type we were already holding, and pydantic cannot see into a dict -- hence pydantic/pydantic#12145, and hence the renaming. Naming the type instead, through `json_schema_input_type` on the validator and `return_type` on the serializer, hands pydantic the same shape in a form it understands. It then owns every definition and reference, including the recursive ones, which is what #761 needed and what the review asked for. So `_inline_refs` goes, and so do the `_bundle_refs`/`_rebundle` pair that replaced it, the definition hashing, and the three `_ensure_strict_json_schema` calls -- litellm applies that pass to the whole document anyway. The tests added here go with them; they covered machinery that no longer exists, and what remains is covered by the suite as it stands on master. Four serializers had to hand back the model rather than a dict, which `return_type` expects and which pydantic otherwise only warns about. That also retires the manual per-field adapter loops in the tuple encodings: the model's own fields carry the element encodings. The `NamedTuple` branch gains a validator, since `WithJsonSchema` with no mode had been covering both directions and `return_type` covers only one. Co-Authored-By: Claude Opus 5 (1M context) --- .../handlers/llm/harness/serialization.py | 221 ++++++------------ .../handlers/llm/harness/synthesis/body.py | 23 +- .../llm/harness/synthesis/function.py | 12 +- ...test_handlers_llm_harness_serialization.py | 220 ----------------- 4 files changed, 71 insertions(+), 405 deletions(-) diff --git a/effectful/handlers/llm/harness/serialization.py b/effectful/handlers/llm/harness/serialization.py index 117405807..c78c97bfc 100644 --- a/effectful/handlers/llm/harness/serialization.py +++ b/effectful/handlers/llm/harness/serialization.py @@ -10,7 +10,6 @@ import contextvars import dataclasses import functools -import hashlib import inspect import io import json @@ -29,7 +28,6 @@ ChatCompletionToolParam, OpenAIMessageContentListBlock, ) -from openai.lib._pydantic import _ensure_strict_json_schema from openai.types.chat import ( ChatCompletionMessageToolCall as OpenAIChatCompletionMessageToolCall, ) @@ -285,83 +283,6 @@ def emit(block: OpenAIMessageContentListBlock) -> None: return [ChatCompletionTextObject(type="text", text=heading), *blocks] -# Base URI for the definitions a fragment carries. An absolute URI is the one ref -# form Pydantic's ref counter tolerates, and an `$id`-bearing subschema is a -# 2020-12 embedded resource, so a ref to it resolves wherever the fragment does. -_DEFS_BASE = "https://effectful.invalid/$defs/" - - -def _bundle_refs(schema: dict) -> dict: - """``schema`` with its definitions addressed absolutely rather than by pointer. - - Works around https://github.com/pydantic/pydantic/issues/12145: Pydantic's - ``GenerateJsonSchema`` does not merge a user-provided ``$defs`` into its - internal ref map, so a ``#/$defs/...`` ref in a `pydantic.WithJsonSchema` - value raises ``KeyError`` once the annotated type is composed into a model. - """ - defs = schema.get("$defs") or {} - if not defs: - return schema - uris = { - name: _DEFS_BASE - + hashlib.sha256( - json.dumps(body, sort_keys=True, default=str).encode() - ).hexdigest()[:16] - for name, body in defs.items() - } - - def walk(node: typing.Any) -> typing.Any: - if isinstance(node, list): - return [walk(item) for item in node] - if not isinstance(node, dict): - return node - return { - k: uris.get(v.removeprefix("#/$defs/"), v) - if k == "$ref" and isinstance(v, str) - else walk(v) - for k, v in node.items() - } - - bundled = walk(schema) - bundled["$defs"] = { - uris[name]: {"$id": uris[name], **body} - for name, body in bundled["$defs"].items() - } - return bundled - - -def _rebundle(schema: dict) -> dict: - """``schema`` in the JSON Schema subset the providers implement. - - They take definitions only at the document root, refuse a remote reference - outright, and reject any keyword sitting beside a ``$ref``. - """ - defs: dict[str, dict] = {} - - def walk(node: typing.Any) -> typing.Any: - if isinstance(node, list): - return [walk(item) for item in node] - if not isinstance(node, dict): - return node - for name, body in (node.get("$defs") or {}).items(): - uri = body.get("$id") - defs[uri.removeprefix(_DEFS_BASE) if uri else name] = walk( - {k: v for k, v in body.items() if k != "$id"} - ) - if isinstance(node.get("$ref"), str): - # Siblings go too, which a provider requires: it reads a `$ref` as - # the whole subschema rather than composing it with what sits - # beside it, the way draft 2020-12 does. - ref = node["$ref"] - if ref.startswith(_DEFS_BASE): - ref = f"#/$defs/{ref.removeprefix(_DEFS_BASE)}" - return {"$ref": ref} - return {k: walk(v) for k, v in node.items() if k != "$defs"} - - document = walk(schema) - return {**document, "$defs": defs} if defs else document - - @dataclasses.dataclass(frozen=True, eq=True) class DecodedToolCall[T]: """ @@ -557,15 +478,13 @@ def _pydantic_type_base(ty: typing.Any) -> typing.Any: return typing.Annotated[ ty, pydantic.InstanceOf, - pydantic.PlainSerializer(_serialize_unencodable), - _NoEncoding(ty), - pydantic.WithJsonSchema( - { - "type": "string", - "description": inspect.formatannotation(ty), - }, - mode="serialization", + pydantic.PlainSerializer( + _serialize_unencodable, + return_type=typing.Annotated[ + str, pydantic.Field(description=inspect.formatannotation(ty)) + ], ), + _NoEncoding(ty), ] @@ -589,16 +508,18 @@ def _pydantic_type_type(ty: typing.Any) -> typing.Any: return typing.Annotated[ ty, pydantic.InstanceOf, - pydantic.PlainSerializer(_best_effort_schema), - pydantic.WithJsonSchema( - { - "type": "object", - "additionalProperties": True, - "description": ( - "A Python type, as the JSON schema of its encoding. Shown for " - "reference; a type cannot be reconstructed from its schema." + pydantic.PlainSerializer( + _best_effort_schema, + return_type=typing.Annotated[ + dict[str, typing.Any], + pydantic.Field( + description=( + "A Python type, as the JSON schema of its encoding. Shown " + "for reference; a type cannot be reconstructed from its " + "schema." + ) ), - } + ], ), ] @@ -643,14 +564,12 @@ def _serialize_complex(value: complex) -> _ComplexModel: def _pydantic_type_complex(ty): """Encode ``complex`` as ``{"real": float, "imag": float}``.""" - schema = pydantic.TypeAdapter(_ComplexModel).json_schema() - schema = _ensure_strict_json_schema(_bundle_refs(schema), path=(), root={}) - return typing.Annotated[ ty, - pydantic.BeforeValidator(_validate_complex), - pydantic.PlainSerializer(_serialize_complex), - pydantic.WithJsonSchema(schema), + pydantic.BeforeValidator( + _validate_complex, json_schema_input_type=_ComplexModel + ), + pydantic.PlainSerializer(_serialize_complex, return_type=_ComplexModel), ] @@ -670,7 +589,6 @@ def _pydantic_type_tuple(ty): hints = typing.get_type_hints(ty) nt_fields: list[str] = list(ty._fields) nt_types = [hints.get(f, typing.Any) for f in nt_fields] - nt_adapters = [pydantic.TypeAdapter(t) for t in nt_types] nt_model = pydantic.create_model( ty.__name__, __config__={"extra": "forbid"}, @@ -678,18 +596,20 @@ def _pydantic_type_tuple(ty): **{f: (t, ...) for f, t in zip(nt_fields, nt_types)}, ) - def _nt_serialize(value, info: pydantic.SerializationInfo): - return { - f: nt_adapters[i].dump_python( - getattr(value, f), mode="json", context=info.context - ) - for i, f in enumerate(nt_fields) - } + def _nt_serialize(value): + return nt_model.model_construct(**{f: getattr(value, f) for f in nt_fields}) + + def _nt_decode(value): + """Reshape the named-field object form back into the positional tuple + Pydantic validates a `NamedTuple` from.""" + if isinstance(value, collections.abc.Mapping): + return tuple(value[f] for f in nt_fields) + return value return typing.Annotated[ ty, - pydantic.PlainSerializer(_nt_serialize), - pydantic.WithJsonSchema(_bundle_refs(nt_model.model_json_schema())), + pydantic.BeforeValidator(_nt_decode, json_schema_input_type=nt_model), + pydantic.PlainSerializer(_nt_serialize, return_type=nt_model), ] args = typing.get_args(ty) @@ -704,8 +624,6 @@ def _nt_serialize(value, info: pydantic.SerializationInfo): # tuple[()] (empty args with origin) maps to zero fields; otherwise use args. effective: list[typing.Any] = list(args) - adapters = [pydantic.TypeAdapter(a) for a in effective] - model = pydantic.create_model( "TupleItems", __config__={"extra": "forbid"}, @@ -719,17 +637,13 @@ def _decode(value): return tuple(value[f"item_{i}"] for i in range(len(effective))) return value - def _serialize(value, info: pydantic.SerializationInfo): - return { - f"item_{i}": adapters[i].dump_python(v, mode="json", context=info.context) - for i, v in enumerate(value) - } + def _serialize(value): + return model.model_construct(**{f"item_{i}": v for i, v in enumerate(value)}) return typing.Annotated[ ty, - pydantic.BeforeValidator(_decode), - pydantic.PlainSerializer(_serialize), - pydantic.WithJsonSchema(_bundle_refs(model.model_json_schema())), + pydantic.BeforeValidator(_decode, json_schema_input_type=model), + pydantic.PlainSerializer(_serialize, return_type=model), ] @@ -770,13 +684,15 @@ def _serialize_image(value: Image.Image) -> ChatCompletionImageObject: @TypeToPydanticType.register(Image.Image) def _pydantic_type_image(ty: type[Image.Image]): - adapter = pydantic.TypeAdapter(ChatCompletionImageObject) return typing.Annotated[ ty, pydantic.InstanceOf, - pydantic.BeforeValidator(_validate_image), - pydantic.PlainSerializer(_serialize_image), - pydantic.WithJsonSchema(_bundle_refs(adapter.json_schema())), + pydantic.BeforeValidator( + _validate_image, json_schema_input_type=ChatCompletionImageObject + ), + pydantic.PlainSerializer( + _serialize_image, return_type=ChatCompletionImageObject + ), ] @@ -784,15 +700,15 @@ def _pydantic_type_image(ty: type[Image.Image]): # when a function is handed to it as a value (e.g. a tool's return) -- just the # source, with none of the synthesis instructions the `SynthesizedFunction` subtype # carries for the generation direction. Its JSON schema (docstring included, since -# pydantic renders it as the schema `description`) is the ``mode="serialization"`` -# schema of every synthesized-callable encoding, so keep the docstring model-facing. +# pydantic renders it as the schema `description`) is what every synthesized-callable +# encoding serializes as, so keep the docstring model-facing. class EncodedFunction(pydantic.BaseModel): """A function, encoded as a string of its complete Python source.""" code: str = pydantic.Field(..., description="Python source defining the function.") -def _serialize_callable(value: collections.abc.Callable) -> dict: +def _serialize_callable(value: collections.abc.Callable) -> EncodedFunction: """Encode a callable back to its ``code`` form (source, or a stub). Emits a plain `EncodedFunction` -- which is exactly what the serialization JSON @@ -813,7 +729,7 @@ def _serialize_callable(value: collections.abc.Callable) -> dict: source = None if source: - return EncodedFunction(code=textwrap.dedent(source)).model_dump() + return EncodedFunction(code=textwrap.dedent(source)) name = getattr(value, "__name__", None) docstring = inspect.getdoc(value) @@ -831,7 +747,7 @@ def _serialize_callable(value: collections.abc.Callable) -> dict: """{docstring}""" ... ''' - return EncodedFunction(code=stub_code).model_dump() + return EncodedFunction(code=stub_code) @TypeToPydanticType.register(collections.abc.Callable) @@ -839,10 +755,7 @@ def _pydantic_callable_serialize_only(ty: typing.Any) -> typing.Any: return typing.Annotated[ ty, pydantic.InstanceOf, - pydantic.PlainSerializer(_serialize_callable), - pydantic.WithJsonSchema( - EncodedFunction.model_json_schema(), mode="serialization" - ), + pydantic.PlainSerializer(_serialize_callable, return_type=EncodedFunction), ] @@ -851,10 +764,7 @@ def _pydantic_type_tool(ty: type[Tool]) -> typing.Any: return typing.Annotated[ ty, pydantic.InstanceOf, - pydantic.PlainSerializer(_serialize_callable), - pydantic.WithJsonSchema( - EncodedFunction.model_json_schema(), mode="serialization" - ), + pydantic.PlainSerializer(_serialize_callable, return_type=EncodedFunction), ] @@ -925,7 +835,7 @@ def _serialize_name_and_tool(value: _NameAndTool) -> ChatCompletionToolParam: "function": { "name": name, "description": description, - "parameters": _rebundle(response_format["json_schema"]["schema"]), + "parameters": response_format["json_schema"]["schema"], "strict": True, }, } @@ -941,14 +851,15 @@ def _pydantic_type_name_and_tool(ty: type[_NameAndTool]): would route to `_pydantic_type_tuple`'s NamedTuple branch, which would try to build a `TypeAdapter` for the `Tool` field and fail. """ - schema = pydantic.TypeAdapter(ChatCompletionToolParam).json_schema() - schema = _ensure_strict_json_schema(_bundle_refs(schema), path=(), root={}) return typing.Annotated[ ty, pydantic.InstanceOf, - pydantic.BeforeValidator(_validate_name_and_tool), - pydantic.PlainSerializer(_serialize_name_and_tool), - pydantic.WithJsonSchema(schema), + pydantic.BeforeValidator( + _validate_name_and_tool, json_schema_input_type=ChatCompletionToolParam + ), + pydantic.PlainSerializer( + _serialize_name_and_tool, return_type=ChatCompletionToolParam + ), ] @@ -988,7 +899,7 @@ def _validate_tool_call( def _serialize_tool_call( value: DecodedToolCall, info: pydantic.SerializationInfo -) -> dict: +) -> OpenAIChatCompletionMessageToolCall: ctx = info.context or {} encoded_args: dict[str, typing.Any] = {} if value.source is not None: @@ -1015,21 +926,23 @@ def _serialize_tool_call( "arguments": json.dumps(encoded_args), }, } - ).model_dump(mode="json") + ) @TypeToPydanticType.register(DecodedToolCall) def _pydantic_type_tool_call(ty: type[DecodedToolCall]): # Use OpenAI's ChatCompletionMessageToolCall (has actual fields: id, function, # type) rather than litellm's (empty dict with extra="allow"). - schema = OpenAIChatCompletionMessageToolCall.model_json_schema() - schema = _ensure_strict_json_schema(_bundle_refs(schema), path=(), root={}) return typing.Annotated[ ty, pydantic.InstanceOf, - pydantic.BeforeValidator(_validate_tool_call), - pydantic.PlainSerializer(_serialize_tool_call), - pydantic.WithJsonSchema(schema), + pydantic.BeforeValidator( + _validate_tool_call, + json_schema_input_type=OpenAIChatCompletionMessageToolCall, + ), + pydantic.PlainSerializer( + _serialize_tool_call, return_type=OpenAIChatCompletionMessageToolCall + ), ] @@ -1082,7 +995,3 @@ def base(tool: Tool) -> str: class _BoxedResponse[T](pydantic.BaseModel): value: T - - @classmethod - def model_json_schema(cls, *args, **kwargs) -> dict[str, typing.Any]: - return _rebundle(super().model_json_schema(*args, **kwargs)) diff --git a/effectful/handlers/llm/harness/synthesis/body.py b/effectful/handlers/llm/harness/synthesis/body.py index c64ba22fd..cc23430c6 100644 --- a/effectful/handlers/llm/harness/synthesis/body.py +++ b/effectful/handlers/llm/harness/synthesis/body.py @@ -65,7 +65,6 @@ DecodedToolCall, EncodedFunction, TypeToPydanticType, - _bundle_refs, _serialize_callable, ) from effectful.handlers.llm.harness.synthesis.function import ( @@ -283,15 +282,8 @@ def _validate( return typing.Annotated[ pydantic.InstanceOf[ty_], # type: ignore - pydantic.BeforeValidator(_validate), - pydantic.PlainSerializer(lambda value: _serialize_callable(value)), - pydantic.WithJsonSchema( - _bundle_refs(pydantic.TypeAdapter(typed_enc).json_schema()), - mode="validation", - ), - pydantic.WithJsonSchema( - EncodedFunction.model_json_schema(), mode="serialization" - ), + pydantic.BeforeValidator(_validate, json_schema_input_type=typed_enc), + pydantic.PlainSerializer(_serialize_callable, return_type=EncodedFunction), ] @@ -453,15 +445,8 @@ def _doctest_apply(op, *args, **kwargs): return typing.Annotated[ pydantic.InstanceOf[ty_], # type: ignore - pydantic.BeforeValidator(_validate), - pydantic.PlainSerializer(_serialize_callable), - pydantic.WithJsonSchema( - _bundle_refs(pydantic.TypeAdapter(typed_enc).json_schema()), - mode="validation", - ), - pydantic.WithJsonSchema( - EncodedFunction.model_json_schema(), mode="serialization" - ), + pydantic.BeforeValidator(_validate, json_schema_input_type=typed_enc), + pydantic.PlainSerializer(_serialize_callable, return_type=EncodedFunction), ] diff --git a/effectful/handlers/llm/harness/synthesis/function.py b/effectful/handlers/llm/harness/synthesis/function.py index 6a82f5648..ce68958b1 100644 --- a/effectful/handlers/llm/harness/synthesis/function.py +++ b/effectful/handlers/llm/harness/synthesis/function.py @@ -16,7 +16,6 @@ _TYPE_CHECK_ANCHOR_KEY, EncodedFunction, TypeToPydanticType, - _bundle_refs, _serialize_callable, ) @@ -395,13 +394,6 @@ def _validate( return typing.Annotated[ ty, pydantic.InstanceOf, - pydantic.BeforeValidator(_validate), - pydantic.PlainSerializer(_serialize_callable), - pydantic.WithJsonSchema( - _bundle_refs(pydantic.TypeAdapter(typed_enc).json_schema()), - mode="validation", - ), - pydantic.WithJsonSchema( - EncodedFunction.model_json_schema(), mode="serialization" - ), + pydantic.BeforeValidator(_validate, json_schema_input_type=typed_enc), + pydantic.PlainSerializer(_serialize_callable, return_type=EncodedFunction), ] diff --git a/tests/test_handlers_llm_harness_serialization.py b/tests/test_handlers_llm_harness_serialization.py index 5508557a8..4297bddfa 100644 --- a/tests/test_handlers_llm_harness_serialization.py +++ b/tests/test_handlers_llm_harness_serialization.py @@ -21,7 +21,6 @@ TypedDict, TypeVar, Union, - cast, ) import litellm @@ -35,18 +34,11 @@ RestrictedPythonExecutor, ) from effectful.handlers.llm.harness.serialization import ( - _DEFS_BASE, _NAME2TOOL_KEY, _TYPE_CHECK_ANCHOR_KEY, CONTENT_BLOCK_TYPES, DecodedToolCall, - EncodedFunction, - _BoxedResponse, - _bundle_refs, - _ComplexModel, _NameAndTool, - _rebundle, - _serialize_name_and_tool, to_content_blocks, ) from effectful.handlers.llm.harness.validation.ty import TyTypeChecker @@ -186,32 +178,6 @@ class _PersonWithAddressModel(pydantic.BaseModel): type _RecursiveAlias = int | list[_RecursiveAlias] type _RecursiveGenAlias[T] = T | list[_RecursiveGenAlias[T]] - -class _Tree(pydantic.BaseModel): - """A model that refers to itself, so its schema needs a definition.""" - - label: str - kids: list["_Tree"] = [] - - -class _TreeParams(pydantic.BaseModel): - tree: _Tree - - -@dataclass -class _Place: - city: str - - -@dataclass -class _Resident: - """A definition reached twice without being recursive.""" - - name: str - home: _Place - work: _Place - - # Stands in for `Kernel` in `docs/source/llm_examples/optimization/kernels.py`: # an alias that is in a skill's lexical *scope*, so the alias object itself # reaches the encoding as a value. @@ -976,192 +942,6 @@ def test_recursive_alias_does_not_diverge(ty): assert "$ref" in schema, "recursion should resolve to a reference, not inline" -# --------------------------------------------------------------------------- -# Reference addressing -# --------------------------------------------------------------------------- - - -def _refs(node): - """Every ``$ref`` string anywhere in `node`.""" - if isinstance(node, dict): - for key, value in node.items(): - if key == "$ref": - yield value - else: - yield from _refs(value) - elif isinstance(node, list): - for item in node: - yield from _refs(item) - - -def _defs_paths(node, path=()): - """Where in `node` a ``$defs`` block sits.""" - if isinstance(node, dict): - for key, value in node.items(): - if key == "$defs": - yield path - else: - yield from _defs_paths(value, path + (key,)) - elif isinstance(node, list): - for i, item in enumerate(node): - yield from _defs_paths(item, path + (i,)) - - -@pytest.mark.parametrize( - "ty", [_RecursiveAlias, _RecursiveGenAlias[int]], ids=["plain", "generic"] -) -def test_a_recursive_schema_can_be_bundled(ty): - """A self-referential schema bundles without diverging -- issue #761.""" - bundled = _bundle_refs(pydantic.TypeAdapter(Encodable[ty]).json_schema()) - assert not any(r.startswith("#") for r in _refs(bundled)), ( - "a fragment cannot carry a pointer to a root it does not have" - ) - - -@pytest.mark.parametrize( - "ty", - [tuple[_RecursiveAlias, int], list[_RecursiveAlias]], - ids=["tuple", "list"], -) -def test_a_recursive_type_has_an_encoding(ty): - """`_pydantic_type_tuple` bundles the schema of a model over the element - types, so a recursive element reaches the bundler through `Encodable`.""" - assert pydantic.TypeAdapter(Encodable[ty]).json_schema() - - -def test_a_bundled_fragment_composes_into_a_model(): - """The reason the bundler exists: pydantic/pydantic#12145 rejects a local - ``$ref`` in a `WithJsonSchema` value once the type is composed.""" - fragment = _bundle_refs( - pydantic.TypeAdapter(Encodable[_RecursiveAlias]).json_schema() - ) - - class Composed(pydantic.BaseModel): - recursive: Annotated[Any, pydantic.WithJsonSchema(fragment)] - plain: int - - assert Composed.model_json_schema() - - -@pytest.mark.parametrize( - "schema", - [ - pydantic.TypeAdapter(litellm.ChatCompletionToolParam).json_schema(), - pydantic.TypeAdapter(_ComplexModel).json_schema(), - EncodedFunction.model_json_schema(), - ], - ids=["tool-param", "complex", "encoded-function"], -) -def test_bundling_only_readdresses(schema): - """The definitions and their structure are pydantic's; only how a reference - names one changes.""" - bundled = _bundle_refs(json.loads(json.dumps(schema))) - assert not any(r.startswith("#") for r in _refs(bundled)) - assert sorted(_defs_paths(bundled)) == sorted(_defs_paths(schema)) - assert len(bundled.get("$defs") or {}) == len(schema.get("$defs") or {}) - - -def test_bundling_keeps_the_keywords_beside_a_ref(): - """In 2020-12 a ``$ref`` composes with its siblings rather than replacing - them, and pydantic puts a field's ``description`` there.""" - - class _Inner(pydantic.BaseModel): - x: int - - class _Outer(pydantic.BaseModel): - described: _Inner = pydantic.Field(..., description="what it is for") - plain: _Inner - - described = _bundle_refs(_Outer.model_json_schema())["properties"]["described"] - assert described["description"] == "what it is for" - assert described["$ref"].startswith(_DEFS_BASE) - - -def test_rebundling_narrows_a_document_to_what_providers_accept(): - """Four separate rejections: a remote reference, a keyword beside a - ``$ref``, a ``$defs`` anywhere but the root, and a reference in any other - form are each answered with a ``BadRequestError``.""" - document = litellm.utils.type_to_response_format_param(_TreeParams)["json_schema"][ - "schema" - ] - wire = _rebundle(json.loads(json.dumps(document))) - - assert "$id" not in json.dumps(wire) - assert all(r.startswith("#/$defs/") for r in _refs(wire)) - assert list(_defs_paths(wire)) == [()] - assert list(_refs(wire)), "the recursion should survive as a reference" - - def no_siblings(node): - if isinstance(node, dict): - assert "$ref" not in node or not set(node) - {"$ref"}, node - for value in node.values(): - no_siblings(value) - elif isinstance(node, list): - for item in node: - no_siblings(item) - - no_siblings(wire) - - -def test_rebundling_a_ref_that_carries_a_description_terminates(): - """`_ensure_strict_json_schema` unravels such a ``$ref`` by inlining its - target and has no guard against doing so forever, so the siblings on the - references that survive are dropped before it runs.""" - schema = { - "type": "object", - "$defs": { - "N": { - "type": "object", - "properties": { - "kids": { - "type": "array", - "items": {"$ref": "#/$defs/N", "description": "a child"}, - } - }, - "required": ["kids"], - "additionalProperties": False, - } - }, - "properties": {"tree": {"$ref": "#/$defs/N", "title": "Tree"}}, - "required": ["tree"], - "additionalProperties": False, - } - assert _rebundle(schema) - - -@pytest.mark.parametrize("ty", [_Tree, _Resident], ids=["recursive", "repeated"]) -def test_a_tool_parameter_is_advertised_by_root_pointer(ty): - def plant(subject: ty) -> str: - """Plant {subject}.""" - return "" - - tool = cast(Tool, plant) - tool.__signature__ = inspect.signature(plant) - parameters = _serialize_name_and_tool(_NameAndTool("plant", tool))["function"][ - "parameters" - ] - assert list(_refs(parameters)) - assert all(r.startswith("#/$defs/") for r in _refs(parameters)) - assert list(_defs_paths(parameters)) == [()] - - -def test_a_recursive_response_type_is_advertised_as_a_reference(): - """Carried on the boxing model itself, which is where litellm reads a - request's ``response_format`` from.""" - boxed = pydantic.create_model( - "BoxedResponse", value=Encodable[_Tree], __base__=_BoxedResponse - ) - for schema in ( - boxed.model_json_schema(), - # And through litellm, which runs its own strict pass over what it finds - # here; the references have to survive that too. - litellm.utils.type_to_response_format_param(boxed)["json_schema"]["schema"], - ): - assert list(_refs(schema)), "the recursion should survive as a reference" - assert all(r.startswith("#/$defs/") for r in _refs(schema)) - assert list(_defs_paths(schema)) == [()] - - # --------------------------------------------------------------------------- # An alias as a *value* # ---------------------------------------------------------------------------