Importable Pydantic models for Open Responses

generate.py

#!/usr/bin/env -S uv run --script
# Run this standalone with `uv run --script generate.py` or `chmod +x generate.py && ./generate.py`.
# /// script
# requires-python = ">=3.13"
# dependencies = [
#   "datamodel-code-generator>=0.58.0",
# ]
# ///

from __future__ import annotations

import sys

# Keep this gist-friendly: do not let neighboring files shadow stdlib modules.
if sys.path:
    sys.path.pop(0)

import argparse
import json
from pathlib import Path
from urllib.request import urlopen

from datamodel_code_generator import (
    DataModelType,
    InputFileType,
    LiteralType,
    PythonVersion,
    generate,
)


SPEC_URL = "https://www.openresponses.org/openapi/openapi.json"
OUTPUT = Path(__file__).with_name("generated.py")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Generate OpenResponses Pydantic models.",
        epilog="Warning: the destination file is overwritten.",
    )
    parser.add_argument(
        "output",
        nargs="?",
        type=Path,
        default=OUTPUT,
        metavar="OUTPUT",
        help=f"destination Python file (default: {OUTPUT})",
    )
    args = parser.parse_args()
    generate_models(args.output)
    print(f"Overwrote {args.output}")


def generate_models(output: Path) -> None:
    """Generate importable Pydantic models from the OpenResponses OpenAPI schema."""

    schema = json.loads(urlopen(SPEC_URL, timeout=30).read())
    fix_nullable_item_reference_discriminator(schema)
    remove_item_param_discriminator(schema)
    remove_websocket_response_create_disallowed_overrides(schema)

    generate(
        schema,
        input_file_type=InputFileType.OpenAPI,
        output=output,
        output_model_type=DataModelType.PydanticV2BaseModel,
        target_python_version=PythonVersion.PY_313,
        use_union_operator=True,
        use_standard_collections=True,
        enum_field_as_literal=LiteralType.All,
        apply_default_values_for_required_fields=True,
        use_annotated=True,
        field_constraints=True,
        custom_file_header=(
            f"# Generated from {SPEC_URL}\n"
            "# Command: uv run python -m openresponses_spec.generate\n"
            "# Tool: datamodel-code-generator\n"
            "#\n"
            "# Pre-generation compatibility transform:\n"
            "# - ItemReferenceParam.type is made non-null item_reference to match the docs\n"
            "#   and avoid datamodel-code-generator emitting Literal['ItemReferenceParam'].\n"
            "# - ItemParam.discriminator is removed before codegen because type='message'\n"
            "#   is shared by several variants. openresponses_spec/types.py defines a\n"
            "#   callable Pydantic discriminator for that union."
        ),
    )


def fix_nullable_item_reference_discriminator(schema: dict) -> None:
    """Make ItemReferenceParam.type a concrete discriminator value.

    The published OpenAPI schema currently allows ItemReferenceParam.type to be
    null or omitted, while also using "type" as the discriminator for ItemParam.
    datamodel-code-generator then emits Literal["ItemReferenceParam"] | None,
    which is both the wrong literal value and invalid for Pydantic
    discriminated unions. The rendered docs describe the value as always
    "item_reference", so we normalize the codegen input to that shape.
    """

    schemas = schema["components"]["schemas"]
    item_reference_type = schemas["ItemReferenceParam"]["properties"]["type"]
    item_reference_type.clear()
    item_reference_type.update(
        {
            "type": "string",
            "enum": ["item_reference"],
            "description": "The type of item to reference. Always `item_reference`.",
            "default": "item_reference",
        }
    )


def remove_item_param_discriminator(schema: dict) -> None:
    """Remove the one OpenAPI discriminator that cannot be modeled by type alone.

    ItemParam includes user, system, developer, and assistant messages. Those
    variants all use type="message" and are distinguished by role, so a normal
    Pydantic type-field discriminator is ambiguous. The replacement callable
    discriminator is defined in openresponses_spec.types.
    """

    schema["components"]["schemas"]["ItemParam"].pop("discriminator", None)


def remove_websocket_response_create_disallowed_overrides(schema: dict) -> None:
    """Remove impossible websocket-only overrides from the allOf wrapper.

    The published schema models WebSocketResponseCreateEvent as CreateResponseBody
    plus fields marked x-openresponses-disallowed. Those fields intentionally
    conflict with the referenced CreateResponseBody field types, which makes
    datamodel-code-generator emit a subclass that mypy correctly rejects.
    """

    websocket_create = schema["components"]["schemas"]["WebSocketResponseCreateEvent"]
    for part in websocket_create.get("allOf", []):
        properties = part.get("properties")
        if not isinstance(properties, dict):
            continue
        for name, value in list(properties.items()):
            if isinstance(value, dict) and value.get("x-openresponses-disallowed") is True:
                properties.pop(name)


if __name__ == "__main__":
    main()