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 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

    if "databricks" not in data or "models" not in data["databricks"]:
        raise RuntimeError(
            "Unexpected models.dev shape: missing data['databricks']['models']"
        )

    models: dict = data["databricks"]["models"]
    raw_entries = sorted(
        (k, v["name"] if isinstance(v, dict) else k) for k, v in models.items()
    )
    entries = validate_entries(raw_entries)

    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)