app/agents/d3/registry.py

from __future__ import annotations
from dataclasses import dataclass
from typing import Literal
from pydantic import BaseModel

Family = Literal["chart", "hierarchy", "network", "flow"]


@dataclass
class D3Template:
    id: str
    family: Family
    title: str
    when_to_use: str
    data_requirements: str
    schema: type[BaseModel]
    html_template: str
    golden_sample: BaseModel

    def render(self, data: BaseModel) -> str:
        json_payload = data.model_dump_json().replace("</", "<\\/")
        return self.html_template.replace("__DATA__", json_payload)


TEMPLATES: dict[str, D3Template] = {}


def register(t: D3Template) -> None:
    TEMPLATES[t.id] = t


def catalog_for_prompt() -> str:
    return "\n".join(
        f"- {t.id} [{t.family}]: {t.when_to_use} | needs: {t.data_requirements}"
        for t in TEMPLATES.values()
    )