dify/api/core/dify_builder/contract_gen.py

"""Generate the JSON Schema + TypeScript contract artifacts from the one
authored Python source (``contract.py``, plus ``EntryMode``/``ConversationItem``
from ``models.py`` and ``SessionView`` from ``services.dify_builder.service``).

Regenerate after any change to those sources::

    uv run --directory api python -m core.dify_builder.contract_gen

This writes the two checked-in derived files:

- ``api/core/dify_builder/contract_schema.json``
- ``web/app/components/dify-builder/contract/types.ts``

Both are DERIVED -- never hand-edit them. ``tests/unit_tests/core/
dify_builder/test_contract.py::test_schema_in_lockstep`` byte-compares a
fresh ``generate()`` call against the checked-in files to guard drift.

Types are *discovered*, not hand-listed, so new enums/dataclasses/cards
added to the authored source are picked up automatically:

- enums: every ``StrEnum`` subclass defined in ``contract.py`` (by
  ``__module__``), plus ``EntryMode`` (defined in ``models.py``).
- dataclasses: every ``@dataclass`` defined in ``contract.py`` (by
  ``__module__``), excluding the private ``_Card`` base.
- the discriminated ``ConversationItem`` union: one branch per
  ``_Card.__subclasses__()``, keyed by that card's ``kind`` ClassVar.
- ``SessionView`` (``services.dify_builder.service``) is generated last,
  by the same dataclass machinery.
"""

import dataclasses
import json
import types
import typing
from enum import StrEnum
from pathlib import Path
from typing import Any

from core.dify_builder import contract
from core.dify_builder.contract import _Card
from core.dify_builder.models import EntryMode
from services.dify_builder.service import SessionView

_JSON_SCHEMA_DIALECT = "https://json-schema.org/draft/2020-12/schema"

_TS_HEADER = (
    "// AUTO-GENERATED by api/core/dify_builder/contract_gen.py — do not edit by hand.",
    "// Regenerate: uv run --directory api python -m core.dify_builder.contract_gen",
)


def _discover_enums() -> list[type]:
    """Every ``StrEnum`` defined (not merely imported) in ``contract.py``,
    in definition order, plus ``EntryMode`` (defined in ``models.py``)."""
    enums = [
        obj
        for _name, obj in vars(contract).items()
        if isinstance(obj, type)
        and issubclass(obj, StrEnum)
        and obj is not StrEnum
        and obj.__module__ == contract.__name__
    ]
    enums.append(EntryMode)
    return enums


def _discover_dataclasses() -> list[type]:
    """Every ``@dataclass`` defined (not merely imported) in ``contract.py``,
    in definition order, excluding the private ``_Card`` discriminant base."""
    return [
        obj
        for _name, obj in vars(contract).items()
        if dataclasses.is_dataclass(obj)
        and isinstance(obj, type)
        and obj.__module__ == contract.__name__
        and obj is not _Card
    ]


def _split_optional(tp: Any) -> tuple[bool, Any]:
    """``X | None`` -> ``(True, X)``; anything else -> ``(False, tp)``."""
    origin = typing.get_origin(tp)
    if origin is types.UnionType or origin is typing.Union:
        args = typing.get_args(tp)
        non_none = [a for a in args if a is not type(None)]
        if len(non_none) == 1 and len(non_none) == len(args) - 1:
            return True, non_none[0]
    return False, tp


def _json_type(tp: Any) -> dict[str, Any]:
    if tp is str:
        return {"type": "string"}
    if tp is bool:
        return {"type": "boolean"}
    if tp is int:
        return {"type": "integer"}
    if tp is float:
        return {"type": "number"}
    origin = typing.get_origin(tp)
    if tp is dict or origin is dict:
        return {"type": "object"}
    if origin is list:
        (item_tp,) = typing.get_args(tp)
        return {"type": "array", "items": _json_type(item_tp)}
    if isinstance(tp, type) and issubclass(tp, StrEnum):
        return {"$ref": f"#/$defs/{tp.__name__}"}
    if dataclasses.is_dataclass(tp):
        return {"$ref": f"#/$defs/{tp.__name__}"}
    raise TypeError(f"contract_gen: no JSON Schema mapping for type {tp!r}")


def _ts_type(tp: Any) -> str:
    if tp is str:
        return "string"
    if tp is bool:
        return "boolean"
    if tp is int:
        return "number"
    if tp is float:
        return "number"
    origin = typing.get_origin(tp)
    if tp is dict or origin is dict:
        return "Record<string, unknown>"
    if origin is list:
        (item_tp,) = typing.get_args(tp)
        return f"{_ts_type(item_tp)}[]"
    if isinstance(tp, type) and issubclass(tp, StrEnum):
        return tp.__name__
    if dataclasses.is_dataclass(tp):
        return tp.__name__
    raise TypeError(f"contract_gen: no TypeScript mapping for type {tp!r}")


def _is_required_field(f: "dataclasses.Field[Any]", is_optional_type: bool) -> bool:
    has_default = f.default is not dataclasses.MISSING or f.default_factory is not dataclasses.MISSING
    return not (has_default or is_optional_type)


def _enum_json_schema(enum_cls: type) -> dict[str, Any]:
    return {"type": "string", "enum": [member.value for member in enum_cls]}


def _enum_ts(enum_cls: type) -> str:
    values = " | ".join(f"'{member.value}'" for member in enum_cls)
    return f"export type {enum_cls.__name__} = {values};"


def _dataclass_json_schema(dc: type) -> dict[str, Any]:
    hints = typing.get_type_hints(dc)
    properties: dict[str, Any] = {}
    required: list[str] = []
    for f in dataclasses.fields(dc):
        tp = hints[f.name]
        is_optional_type, inner = _split_optional(tp)
        json_type = _json_type(inner if is_optional_type else tp)
        if is_optional_type:
            json_type = {"anyOf": [json_type, {"type": "null"}]}
        properties[f.name] = json_type
        if _is_required_field(f, is_optional_type):
            required.append(f.name)
    return {
        "type": "object",
        "properties": properties,
        "required": required,
        "additionalProperties": False,
    }


def _dataclass_ts(dc: type) -> str:
    hints = typing.get_type_hints(dc)
    parts: list[str] = []
    for f in dataclasses.fields(dc):
        tp = hints[f.name]
        is_optional_type, inner = _split_optional(tp)
        ts_type = _ts_type(inner if is_optional_type else tp)
        if is_optional_type:
            ts_type = f"{ts_type} | null"
        optional_marker = "?" if not _is_required_field(f, is_optional_type) else ""
        parts.append(f"{f.name}{optional_marker}: {ts_type}")
    if not parts:
        return f"export interface {dc.__name__} {{}}"
    body = " ".join(f"{part};" for part in parts)
    return f"export interface {dc.__name__} {{ {body} }}"


def _conversation_item_json_schema(cards: list[type]) -> dict[str, Any]:
    return {
        "oneOf": [
            {
                "type": "object",
                "properties": {
                    "seq": {"type": "integer"},
                    "at_version": {"type": "integer"},
                    "kind": {"const": card.kind.value},
                    "payload": {"$ref": f"#/$defs/{card.__name__}"},
                },
                "required": ["seq", "at_version", "kind", "payload"],
                "additionalProperties": False,
            }
            for card in cards
        ]
    }


def _conversation_item_ts(cards: list[type]) -> str:
    branch_lines = "\n".join(
        f"  | {{ seq: number; at_version: number; kind: '{card.kind.value}'; payload: {card.__name__} }}"
        for card in cards
    )
    return f"export type ConversationItem =\n{branch_lines};"


def generate() -> tuple[dict[str, Any], str]:
    """Return ``(json_schema_dict, typescript_source)`` -- pure, deterministic."""
    enums = _discover_enums()
    dataclass_types = _discover_dataclasses()
    cards = list(_Card.__subclasses__())

    defs: dict[str, Any] = {}
    for enum_cls in enums:
        defs[enum_cls.__name__] = _enum_json_schema(enum_cls)
    for dc in dataclass_types:
        defs[dc.__name__] = _dataclass_json_schema(dc)
    defs["ConversationItem"] = _conversation_item_json_schema(cards)
    defs["SessionView"] = _dataclass_json_schema(SessionView)

    schema: dict[str, Any] = {"$schema": _JSON_SCHEMA_DIALECT, "$defs": defs}

    ts_lines: list[str] = list(_TS_HEADER)
    ts_lines += [_enum_ts(enum_cls) for enum_cls in enums]
    ts_lines += [_dataclass_ts(dc) for dc in dataclass_types]
    ts_lines.append(_dataclass_ts(SessionView))
    ts_lines.append(_conversation_item_ts(cards))
    ts_source = "\n".join(ts_lines) + "\n"

    return schema, ts_source


def main() -> None:
    """Write ``contract_schema.json`` + ``web/.../contract/types.ts`` to disk."""
    schema, ts_source = generate()

    repo_root = next(p for p in Path(__file__).parents if (p / "web").is_dir() and (p / "api").is_dir())
    schema_path = Path(__file__).with_name("contract_schema.json")
    ts_path = repo_root / "web" / "app" / "components" / "dify-builder" / "contract" / "types.ts"

    schema_path.write_text(json.dumps(schema, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    ts_path.parent.mkdir(parents=True, exist_ok=True)
    ts_path.write_text(ts_source, encoding="utf-8")


if __name__ == "__main__":
    main()