buzz/scripts/generate-databricks-model-names.py

#!/usr/bin/env python3
"""Generate Databricks model-name registries from models.dev.

Emits two generated files in one invocation:

    crates/buzz-agent/src/databricks_model_names.rs  (Rust)
    desktop/src/features/agents/lib/databricksModelNames.ts  (TypeScript)

Usage (from repo root):
    python3 scripts/generate-databricks-model-names.py

Fetches https://models.dev/api.json, extracts providers.databricks.models,
and emits sorted (id, display_name) tables formatted for each language.
Re-run whenever Databricks ships a new managed endpoint and commit the diff.

Both files are kept in sync by this script — never edit them by hand.
"""

import json
import re
import subprocess
import sys
from pathlib import Path

URL = "https://models.dev/api.json"
REPO_ROOT = Path(__file__).resolve().parent.parent
RUST_OUT = REPO_ROOT / "crates/buzz-agent/src/databricks_model_names.rs"
TS_OUT = REPO_ROOT / "desktop/src/features/agents/lib/databricksModelNames.ts"

# Allowed characters in endpoint IDs and curated names.
SAFE_ID_RE = re.compile(r"^[a-z0-9][a-z0-9.\-]*$")
SAFE_NAME_RE = re.compile(r"^[^\x00-\x1f\"\\<>&]+$")


def fetch(url: str) -> bytes:
    """Fetch URL via curl; raise on HTTP error."""
    result = subprocess.run(
        ["curl", "--fail", "--silent", "--max-time", "30", "-A", "Mozilla/5.0", url],
        capture_output=True,
    )
    if result.returncode != 0:
        raise RuntimeError(
            f"curl failed (exit {result.returncode}): {result.stderr.decode()}"
        )
    return result.stdout


def validate_entries(
    entries: list[tuple[str, str]],
) -> list[tuple[str, str]]:
    """Validate all (id, name) pairs and raise on unexpected shapes."""
    for id_, name in entries:
        if not isinstance(id_, str) or not isinstance(name, str):
            raise ValueError(f"Non-string entry: {id_!r} -> {name!r}")
        if not SAFE_ID_RE.match(id_):
            raise ValueError(f"Unsafe endpoint ID: {id_!r}")
        if not SAFE_NAME_RE.match(name):
            raise ValueError(f"Unsafe display name for {id_!r}: {name!r}")
    return entries


def rust_str(s: str) -> str:
    """Emit a double-quoted Rust string literal (backslash + quote only)."""
    return '"' + s.replace("\\", "\\\\").replace('"', '\\"') + '"'


def ts_str(s: str) -> str:
    """Emit a double-quoted TypeScript string literal."""
    return '"' + s.replace("\\", "\\\\").replace('"', '\\"') + '"'


def write_rust(entries: list[tuple[str, str]]) -> None:
    """Write the Rust generated file, then rustfmt it for byte-for-byte stability."""
    lines = [
        "// GENERATED by scripts/generate-databricks-model-names.py",
        "// Source:  https://models.dev/api.json -- providers.databricks.models",
        "// Refresh: python3 scripts/generate-databricks-model-names.py",
        "//",
        "// Do not hand-edit -- rerun the script to update.",
        "",
        "/// Curated display names for known Databricks AI Gateway endpoints.",
        "///",
        "/// Keys are endpoint IDs returned verbatim by the discovery APIs.",
        "/// Values are human-readable display names sourced from models.dev.",
        "///",
        "/// Unknown endpoint IDs are displayed as their raw ID -- no guessing.",
        "pub(crate) static DATABRICKS_MODEL_NAMES: &[(&str, &str)] = &[",
    ]
    for id_, name in entries:
        lines.append(f"    ({rust_str(id_)}, {rust_str(name)}),")
    lines += ["];", ""]
    RUST_OUT.write_text("\n".join(lines))
    # Run rustfmt so the committed file is always formatter-clean and
    # a subsequent generator run reproduces it byte-for-byte.
    result = subprocess.run(
        ["cargo", "fmt", "--", str(RUST_OUT)],
        cwd=REPO_ROOT,
        capture_output=True,
    )
    if result.returncode != 0:
        raise RuntimeError(
            f"rustfmt failed: {result.stderr.decode()}"
        )
    print(f"Wrote {RUST_OUT.relative_to(REPO_ROOT)}")


def write_ts(entries: list[tuple[str, str]]) -> None:
    """Write the TypeScript generated file, then biome-format it."""
    lines = [
        "// GENERATED by scripts/generate-databricks-model-names.py",
        "// Source:  https://models.dev/api.json -- providers.databricks.models",
        "// Refresh: python3 scripts/generate-databricks-model-names.py",
        "//",
        "// Do not hand-edit -- rerun the script to update.",
        "",
        "/**",
        " * Curated display names for known Databricks AI Gateway endpoints.",
        " *",
        " * Keys are endpoint IDs returned verbatim by the discovery APIs.",
        " * Values are human-readable display names sourced from models.dev.",
        " *",
        " * Unknown endpoint IDs are resolved by resolveModelLabel() as raw IDs.",
        " * Use a Map to avoid Object.prototype key collisions.",
        " */",
        "export const DATABRICKS_MODEL_NAMES: Map<string, string> = new Map([",
    ]
    for id_, name in entries:
        lines.append(f"  [{ts_str(id_)}, {ts_str(name)}],")
    lines += ["]);\n"]
    TS_OUT.write_text("\n".join(lines))
    # biome format for byte-for-byte stability on regenerate.
    desktop_dir = REPO_ROOT / "desktop"
    biome_bin = desktop_dir / "node_modules/.bin/biome"
    if biome_bin.exists():
        result = subprocess.run(
            [str(biome_bin), "format", "--write", str(TS_OUT)],
            cwd=desktop_dir,
            capture_output=True,
        )
        if result.returncode != 0:
            raise RuntimeError(
                f"biome format failed: {result.stderr.decode()}"
            )
    print(f"Wrote {TS_OUT.relative_to(REPO_ROOT)}")


def extract_entries(data: object) -> list[tuple[str, str]]:
    """Pull sorted (id, name) pairs out of the models.dev payload.

    Every container and leaf shape is checked explicitly so an upstream
    restructure fails with an actionable message instead of a bare KeyError
    or — worse — a silently degraded table where a malformed entry emits
    `id -> id` and permanently masks the real curated name.
    """
    if not isinstance(data, dict):
        raise RuntimeError(
            f"Unexpected models.dev shape: root must be an object, got {type(data).__name__}"
        )
    provider = data.get("databricks")
    if provider is None:
        raise RuntimeError("Unexpected models.dev shape: missing data['databricks']")
    if not isinstance(provider, dict):
        raise RuntimeError(
            "Unexpected models.dev shape: data['databricks'] must be an object, "
            f"got {type(provider).__name__}"
        )
    models = provider.get("models")
    if models is None:
        raise RuntimeError(
            "Unexpected models.dev shape: missing data['databricks']['models']"
        )
    if not isinstance(models, dict):
        raise RuntimeError(
            "Unexpected models.dev shape: data['databricks']['models'] must be an "
            f"object, got {type(models).__name__}"
        )
    if not models:
        raise RuntimeError(
            "Unexpected models.dev shape: data['databricks']['models'] is empty"
        )

    entries: list[tuple[str, str]] = []
    for model_id, model in models.items():
        where = f"data['databricks']['models'][{model_id!r}]"
        if not isinstance(model, dict):
            raise RuntimeError(
                f"Unexpected models.dev shape: {where} must be an object, "
                f"got {type(model).__name__}"
            )
        name = model.get("name")
        if not isinstance(name, str):
            raise RuntimeError(
                f"Unexpected models.dev shape: {where}['name'] must be a string, "
                f"got {type(name).__name__}"
            )
        entries.append((model_id, name))
    return sorted(entries)


def main() -> None:
    raw = fetch(URL)
    try:
        data = json.loads(raw)
    except json.JSONDecodeError as e:
        raise RuntimeError(f"models.dev response is not valid JSON: {e}") from e

    entries = validate_entries(extract_entries(data))

    write_rust(entries)
    write_ts(entries)
    print(f"Done — {len(entries)} Databricks endpoints.")


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:  # noqa: BLE001
        print(f"error: {exc}", file=sys.stderr)
        sys.exit(1)