aigc-home/calliope-backend/scripts/build_widget_orders.py

"""Regenerate ``src/calliope/comfyui/ui_widget_orders.py`` (HEL-681).

Offline widget-order registry for the UI→API converter. Sources, in priority:

1. ``widgets_values_named`` instances across ``example_ComfyUI_workflows``
   (including subgraph inner nodes) — the frontend serializes named widgets
   in widget order, so the key order IS the serialization order;
2. the repo's 10 ground-truth API exports (root ``*_API.json``) — scalar
   input names in export order (control_after_generate entries are absent
   there; the converter's walk skips them positionally).

Run from anywhere inside the repo:
    python scripts/build_widget_orders.py
"""

from __future__ import annotations

import collections
import json
from pathlib import Path

BACKEND = Path(__file__).resolve().parents[1]
REPO = BACKEND.parent
EXAMPLES = REPO / "example_ComfyUI_workflows"
OUT = BACKEND / "src" / "calliope" / "comfyui" / "ui_widget_orders.py"

BUTTONS = {"preview", "audioUI", "compare_view", "ToggleableColorPicker", "imagebrowser"}
MAX_SIGNATURES = 3
MAX_ORDER_LEN = 48

# Hand-curated fallbacks for node classes the corpus never serializes with
# names (classic core/custom widgets whose order is stable and public).
CURATED: dict[str, list[list[str]]] = {
    "PrimitiveNode": [["value"]],
    "Seed (rgthree)": [["seed"]],
}


def _iter_ui_nodes() -> list[dict]:
    nodes: list[dict] = []
    for path in sorted(EXAMPLES.rglob("*.json")):
        try:
            data = json.loads(path.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, OSError):
            continue
        if not (isinstance(data, dict) and isinstance(data.get("nodes"), list)):
            continue
        nodes.extend(n for n in data["nodes"] if isinstance(n, dict))
        for sub in data.get("definitions", {}).get("subgraphs", []) or []:
            if isinstance(sub, dict):
                nodes.extend(n for n in sub.get("nodes", []) or [] if isinstance(n, dict))
    return nodes


def _iter_api_workflows() -> list[dict]:
    out: list[dict] = []
    for path in sorted(EXAMPLES.glob("*_API.json")):
        try:
            data = json.loads(path.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, OSError):
            continue
        if isinstance(data, dict) and "nodes" not in data:
            out.append(data)
    return out


def main() -> None:
    learned: dict[str, collections.Counter[tuple[str, ...]]] = collections.defaultdict(
        collections.Counter
    )
    for node in _iter_ui_nodes():
        named = node.get("widgets_values_named")
        if not isinstance(named, dict) or not named:
            continue
        keys = tuple(
            str(k) for k in named.keys() if k and k not in BUTTONS
        )
        if keys and len(keys) <= MAX_ORDER_LEN:
            learned[str(node.get("type") or "")][keys] += 1

    for workflow in _iter_api_workflows():
        for node in workflow.values():
            if not isinstance(node, dict):
                continue
            names = [
                str(k)
                for k, v in (node.get("inputs") or {}).items()
                if not isinstance(v, list) and k and k not in BUTTONS
            ]
            if len(names) > MAX_ORDER_LEN:
                continue
            # Display nodes (PreviewAny…) carry zero scalar inputs — that
            # empty order is real signal: their positional arrays are pure
            # button leftovers and map to {}.
            learned[str(node.get("class_type") or "")][tuple(names)] += 1

    registry: dict[str, list[list[str]]] = {}
    for node_type, counter in sorted(learned.items()):
        if not node_type:
            continue
        orders: list[list[str]] = []
        for signature, _count in counter.most_common(MAX_SIGNATURES):
            orders.append(list(signature))
        if orders:
            registry[node_type] = orders
    for node_type, orders in CURATED.items():
        registry.setdefault(node_type, orders)

    body = json.dumps(registry, ensure_ascii=False, indent=1, sort_keys=True)
    OUT.write_text(
        '"""GENERATED by scripts/build_widget_orders.py — do not edit by hand.\n'
        "\n"
        "Widget-order registry for the UI→API converter: class_type → candidate\n"
        "widget name orders, learned from the repo's named-widget UI files and\n"
        "ground-truth API exports (HEL-681).\n"
        '"""\n'
        "\n"
        "from __future__ import annotations\n"
        "\n"
        f"WIDGET_ORDERS: dict[str, list[list[str]]] = {body}\n",
        encoding="utf-8",
    )
    total = sum(len(o) for orders in registry.values() for o in orders)
    print(f"wrote {OUT.name}: {len(registry)} node types, {total} orders total")


if __name__ == "__main__":
    main()