Update json schema creation to dereference pydantic complex data types

Author: joshuawootonn

py/src/braintrust/parameters.py

@@ -1,5 +1,6 @@
 """Evaluation parameters support for Python SDK."""
 
+from copy import deepcopy
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Literal, TypedDict
 
@@ -32,6 +33,7 @@ class ModelParameter(TypedDict):
     description: NotRequired[str | None]
 
 
+JSONValue = None | bool | int | float | str | list["JSONValue"] | dict[str, "JSONValue"]
 ValidatedParameters = dict[str, object]
 ParameterSchema = PromptParameter | ModelParameter | type[object] | None
 EvalParameters = dict[str, ParameterSchema]
@@ -97,8 +99,61 @@ def _get_pydantic_fields(schema: Any) -> dict[str, Any]:
     return getattr(schema, "__fields__", {})
 
 
+def _resolve_json_pointer(document: dict[str, JSONValue], pointer: str) -> JSONValue:
+    if pointer == "#":
+        return document
+    if not pointer.startswith("#/"):
+        raise ValueError(f"Unsupported JSON schema ref '{pointer}'")
+
+    current: JSONValue = document
+    for raw_part in pointer[2:].split("/"):
+        part = raw_part.replace("~1", "/").replace("~0", "~")
+        if not isinstance(current, dict) or part not in current:
+            raise ValueError(f"JSON schema ref '{pointer}' could not be resolved")
+        current = current[part]
+    return current
+
+
+def _resolve_local_json_schema_refs(
+    node: JSONValue,
+    root: dict[str, JSONValue],
+    resolving: tuple[str, ...] = (),
+) -> JSONValue:
+    if isinstance(node, list):
+        return [_resolve_local_json_schema_refs(item, root, resolving) for item in node]
+
+    if not isinstance(node, dict):
+        return node
+
+    ref = node.get("$ref")
+    if isinstance(ref, str):
+        if ref in resolving:
+            raise ValueError(f"Cyclic JSON schema ref '{ref}'")
+
+        resolved = deepcopy(_resolve_json_pointer(root, ref))
+        resolved = _resolve_local_json_schema_refs(resolved, root, resolving + (ref,))
+
+        siblings = {
+            key: _resolve_local_json_schema_refs(value, root, resolving)
+            for key, value in node.items()
+            if key != "$ref"
+        }
+        if siblings:
+            if not isinstance(resolved, dict):
+                raise ValueError(f"Cannot merge sibling keys into non-object JSON schema ref '{ref}'")
+            merged = dict(resolved)
+            merged.update(siblings)
+            return merged
+        return resolved
+
+    return {key: _resolve_local_json_schema_refs(value, root, resolving) for key, value in node.items()}
+
+
 def _serialize_pydantic_parameter_schema(schema: Any) -> dict[str, Any]:
     schema_json = _pydantic_to_json_schema(schema)
+    schema_json = _resolve_local_json_schema_refs(schema_json, schema_json)
+    schema_json.pop("$defs", None)
+    schema_json.pop("definitions", None)
     fields = _get_pydantic_fields(schema)
     if len(fields) == 1 and "value" in fields:
         properties = schema_json.get("properties")

py/src/braintrust/test_parameters.py

@@ -4,13 +4,24 @@
 from braintrust.parameters import (
     RemoteEvalParameters,
     parameters_to_json_schema,
+    serialize_eval_parameters,
     validate_parameters,
 )
 
 
 HAS_PYDANTIC = importlib.util.find_spec("pydantic") is not None
 
 
+def _contains_json_schema_ref(node):
+    if isinstance(node, dict):
+        if "$ref" in node or "$defs" in node or "definitions" in node:
+            return True
+        return any(_contains_json_schema_ref(value) for value in node.values())
+    if isinstance(node, list):
+        return any(_contains_json_schema_ref(value) for value in node)
+    return False
+
+
 @pytest.mark.skipif(not HAS_PYDANTIC, reason="pydantic not installed")
 def test_validate_local_parameters_with_prompt_and_model_defaults():
     from pydantic import BaseModel
@@ -130,6 +141,272 @@ class PrefixParam(BaseModel):
     }
 
 
+@pytest.mark.skipif(not HAS_PYDANTIC, reason="pydantic not installed")
+def test_parameters_to_json_schema_keeps_complex_single_value_models_self_contained():
+    from pydantic import BaseModel, Field
+
+    class ComplexValue(BaseModel):
+        a: int = Field(default=1)
+        b: list[int] = Field(default=[2, 3])
+
+    class ComplexParameter(BaseModel):
+        value: ComplexValue = Field(
+            default=ComplexValue(),
+            description="Complex example parameter",
+        )
+
+    schema = parameters_to_json_schema({"complex": ComplexParameter})
+
+    assert schema["properties"]["complex"] == {
+        "type": "object",
+        "properties": {
+            "a": {
+                "default": 1,
+                "title": "A",
+                "type": "integer",
+            },
+            "b": {
+                "default": [2, 3],
+                "items": {
+                    "type": "integer",
+                },
+                "title": "B",
+                "type": "array",
+            },
+        },
+        "default": {
+            "a": 1,
+            "b": [2, 3],
+        },
+        "description": "Complex example parameter",
+        "title": "ComplexValue",
+    }
+
+
+@pytest.mark.skipif(not HAS_PYDANTIC, reason="pydantic not installed")
+def test_serialize_eval_parameters_does_not_emit_dangling_ref_for_complex_single_value_models():
+    from pydantic import BaseModel, Field
+
+    class ComplexValue(BaseModel):
+        a: int = Field(default=1)
+        b: list[int] = Field(default=[2, 3])
+
+    class ComplexParameter(BaseModel):
+        value: ComplexValue = Field(
+            default=ComplexValue(),
+            description="Complex example parameter",
+        )
+
+    serialized = serialize_eval_parameters({"complex": ComplexParameter})
+
+    assert serialized["complex"] == {
+        "type": "data",
+        "schema": {
+            "type": "object",
+            "properties": {
+                "a": {
+                    "default": 1,
+                    "title": "A",
+                    "type": "integer",
+                },
+                "b": {
+                    "default": [2, 3],
+                    "items": {
+                        "type": "integer",
+                    },
+                    "title": "B",
+                    "type": "array",
+                },
+            },
+            "default": {
+                "a": 1,
+                "b": [2, 3],
+            },
+            "description": "Complex example parameter",
+            "title": "ComplexValue",
+        },
+        "default": {
+            "a": 1,
+            "b": [2, 3],
+        },
+        "description": "Complex example parameter",
+    }
+
+
+@pytest.mark.skipif(not HAS_PYDANTIC, reason="pydantic not installed")
+def test_parameters_to_json_schema_keeps_array_of_objects_single_value_models_self_contained():
+    from pydantic import BaseModel, Field
+
+    class ComplexItem(BaseModel):
+        a: int = Field(default=1)
+        b: str = Field(default="x")
+
+    class ComplexParameter(BaseModel):
+        value: list[ComplexItem] = Field(
+            default=[ComplexItem()],
+            description="Array example parameter",
+        )
+
+    schema = parameters_to_json_schema({"complex_array": ComplexParameter})
+
+    assert schema["properties"]["complex_array"] == {
+        "type": "array",
+        "items": {
+            "type": "object",
+            "properties": {
+                "a": {
+                    "default": 1,
+                    "title": "A",
+                    "type": "integer",
+                },
+                "b": {
+                    "default": "x",
+                    "title": "B",
+                    "type": "string",
+                },
+            },
+            "title": "ComplexItem",
+        },
+        "default": [
+            {
+                "a": 1,
+                "b": "x",
+            }
+        ],
+        "description": "Array example parameter",
+        "title": "Value",
+    }
+    assert not _contains_json_schema_ref(schema["properties"]["complex_array"])
+
+
+@pytest.mark.skipif(not HAS_PYDANTIC, reason="pydantic not installed")
+def test_parameters_to_json_schema_inlines_refs_for_multi_field_models():
+    from pydantic import BaseModel, Field
+
+    class Address(BaseModel):
+        street: str = Field(default="Main")
+        zip_code: int = Field(default=12345)
+
+    class ComplexParameter(BaseModel):
+        address: Address = Field(default=Address())
+        enabled: bool = Field(default=True)
+
+    schema = parameters_to_json_schema({"complex": ComplexParameter})
+
+    assert not _contains_json_schema_ref(schema["properties"]["complex"])
+    assert schema["properties"]["complex"]["properties"]["address"] == {
+        "type": "object",
+        "properties": {
+            "street": {
+                "default": "Main",
+                "title": "Street",
+                "type": "string",
+            },
+            "zip_code": {
+                "default": 12345,
+                "title": "Zip Code",
+                "type": "integer",
+            },
+        },
+        "default": {
+            "street": "Main",
+            "zip_code": 12345,
+        },
+        "title": "Address",
+    }
+
+
+def test_parameters_to_json_schema_inlines_legacy_definitions_refs():
+    class _FakeField:
+        required = False
+
+    class LegacyParameter:
+        __fields__ = {"value": _FakeField()}
+
+        @classmethod
+        def parse_obj(cls, value):
+            return value
+
+        @classmethod
+        def schema(cls):
+            return {
+                "type": "object",
+                "properties": {
+                    "value": {
+                        "$ref": "#/definitions/ComplexValue",
+                        "description": "Legacy example parameter",
+                        "default": {"a": 1},
+                    },
+                },
+                "definitions": {
+                    "ComplexValue": {
+                        "type": "object",
+                        "properties": {
+                            "a": {
+                                "type": "integer",
+                                "title": "A",
+                                "default": 1,
+                            },
+                        },
+                        "title": "ComplexValue",
+                    },
+                },
+            }
+
+    schema = parameters_to_json_schema({"legacy": LegacyParameter})
+
+    assert schema["properties"]["legacy"] == {
+        "type": "object",
+        "properties": {
+            "a": {
+                "type": "integer",
+                "title": "A",
+                "default": 1,
+            },
+        },
+        "title": "ComplexValue",
+        "description": "Legacy example parameter",
+        "default": {"a": 1},
+    }
+    assert not _contains_json_schema_ref(schema["properties"]["legacy"])
+
+
+def test_parameters_to_json_schema_raises_for_cyclic_local_refs():
+    class _FakeField:
+        required = False
+
+    class CyclicParameter:
+        __fields__ = {"value": _FakeField()}
+
+        @classmethod
+        def parse_obj(cls, value):
+            return value
+
+        @classmethod
+        def schema(cls):
+            return {
+                "type": "object",
+                "properties": {
+                    "value": {
+                        "$ref": "#/definitions/Node",
+                    },
+                },
+                "definitions": {
+                    "Node": {
+                        "type": "object",
+                        "properties": {
+                            "child": {
+                                "$ref": "#/definitions/Node",
+                            },
+                        },
+                    },
+                },
+            }
+
+    with pytest.raises(ValueError, match="Cyclic JSON schema ref"):
+        parameters_to_json_schema({"cyclic": CyclicParameter})
+
+
 @pytest.mark.skipif(not HAS_PYDANTIC, reason="pydantic not installed")
 def test_parameters_to_json_schema_marks_prompt_and_model_without_defaults_required():
     schema = parameters_to_json_schema(