app/registry.py

"""Which models DecisionLab compares.

First the operator's Hugging Face models, in this order (rulings 2026-09-28 and 2026-09-30):
  LightDec_Arthur  Falconsai/LightDec_Arthur  (Arthur; code ships with DecisionLab)
  LightDec_V2      Falconsai/LightDec_V2      (FalconDec; its modeling code must be on the allowlist)
  Enterprise Reflux Laya V2.1  yasserrmd/enterprise-reflux-laya-v21  (a Laya fine-tune; loaded with laya 0.3.20)
  Laya             convaiinnovations/laya
Then every model folder found in the models directory, named "... (local)".

The models directory (MODELS_DIR, default /models, mounted from ./models on the host) is scanned once at
start-up. Two kinds of folder are recognised, in folder-name order:
  - FalconDec (LightDec and variants): a falcondec_config.json; named from that config.
  - Arthur: a config.json with Arthur's architecture keys plus model.safetensors; named "Arthur <tier>".
    Arthur folders hold data only: the code that runs them ships with DecisionLab (app/arthur.py).
Laya is always last and still comes from the Hugging Face Hub.

Pure Python, no torch: discovery and its settings can be tested anywhere.
"""
from __future__ import annotations

import hashlib
import json
import re
from pathlib import Path
from typing import Mapping

VARIANTS = ("fp16", "int8")
CONFIG_FILE = "falcondec_config.json"
DEFAULT_MODELS_DIR = "/models"
PALETTE = ("c0", "c1", "c2", "c3", "c4", "c5", "c6", "c7")   # colour slots; CSS defines each


def _variant(env: Mapping[str, str], name: str) -> str:
    v = env.get(name) or "fp16"
    if v not in VARIANTS:
        raise ValueError(f"{name} must be fp16 or int8 (got '{v}').")
    return v


def _slug(folder: str) -> str:
    return re.sub(r"[^a-z0-9]+", "_", folder.lower()).strip("_") or "model"


def _display_name(folder: Path) -> str:
    try:
        cfg = json.loads((folder / CONFIG_FILE).read_text(encoding="utf-8"))
        name, version = cfg.get("name"), cfg.get("version")
        if name:
            return f"{name} {version}" if version else str(name)
    except (OSError, ValueError, AttributeError):
        pass
    return folder.name


ARTHUR_KEYS = frozenset({"tier", "d", "heads", "e", "buckets", "recursions", "interact", "mlp", "layout", "temperatures"})


def _arthur_config(folder: Path) -> dict | None:
    """The folder's Arthur config, or None if it is not an Arthur model folder."""
    if not (folder / "config.json").is_file() or not (folder / "model.safetensors").is_file():
        return None
    try:
        cfg = json.loads((folder / "config.json").read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return None
    return cfg if isinstance(cfg, dict) and ARTHUR_KEYS <= set(cfg) else None


def _kind(folder: Path) -> str | None:
    if (folder / CONFIG_FILE).is_file():
        return "lightdec"
    if _arthur_config(folder) is not None:
        return "arthur"
    return None


def discover(models_dir: str) -> list[tuple[Path, str]]:
    """(folder, kind) for every recognised model folder in models_dir, sorted by folder name. Missing dir -> []."""
    root = Path(models_dir)
    if not root.is_dir():
        return []
    found = ((p, _kind(p)) for p in root.iterdir() if p.is_dir())
    return sorted(((p, k) for p, k in found if k), key=lambda pk: pk[0].name)


def model_specs(env: Mapping[str, str]) -> list[dict]:
    """One dict per model, in the order the UI shows them: the four Hub models, then the models folder."""
    variant = _variant(env, "LIGHTDEC_VARIANT")
    specs = [
        {"key": "lightdec_arthur", "name": "LightDec_Arthur", "kind": "arthur", "side": PALETTE[0], "source": "hub",
         "repo": env.get("LIGHTDEC_ARTHUR_REPO") or "Falconsai/LightDec_Arthur"},
        {"key": "lightdec_v2", "name": "LightDec_V2", "kind": "lightdec", "side": PALETTE[1], "source": "hub",
         "repo": env.get("LIGHTDEC_V2_REPO") or "Falconsai/LightDec_V2", "revision": None, "variant": variant},
        {"key": "enterprise_reflux_laya_v21", "name": "Enterprise Reflux Laya V2.1", "kind": "laya", "side": PALETTE[2],
         "repo": env.get("REFLUX_LAYA_REPO") or "yasserrmd/enterprise-reflux-laya-v21", "fallback_repo": ""},
        {"key": "laya", "name": "Laya", "kind": "laya", "side": "laya",
         "repo": env.get("LAYA_REPO") or "convaiinnovations/laya",
         "fallback_repo": env.get("LAYA_FALLBACK_REPO", "")},
    ]
    used = {s["key"] for s in specs}
    for i, (folder, kind) in enumerate(discover(env.get("MODELS_DIR") or DEFAULT_MODELS_DIR), start=3):
        key, n = _slug(folder.name), 2
        while key in used:
            key, n = f"{_slug(folder.name)}_{n}", n + 1
        used.add(key)
        spec = {"key": key, "kind": kind, "side": PALETTE[i % len(PALETTE)], "source": "local", "path": str(folder),
                "path_env": "MODELS_DIR"}
        if kind == "arthur":
            cfg = _arthur_config(folder)
            name = f"Arthur {cfg.get('tier', folder.name)}" + (" (pretrained)" if cfg.get("pretrain") else "")
        else:
            name = _display_name(folder)
            spec["variant"] = variant
        spec["name"] = f"{name} (local)"
        specs.append(spec)
    return specs


def load_order(env: Mapping[str, str]) -> list[str]:
    """Which models load at start-up, in order. Default: comparison order."""
    known = [s["key"] for s in model_specs(env)]
    raw = env.get("LOAD_ORDER")
    wanted = [k.strip() for k in raw.split(",")] if raw else known
    return [k for k in wanted if k in known]


def resolve_local_dir(path: str, variant: str, label: str, path_env: str) -> Path:
    """The folder load_falcondec should read, or a FileNotFoundError that says how to fix it."""
    root = Path(path)
    if not root.is_dir():
        raise FileNotFoundError(f"{label} folder not found at {path}. It was in the models folder at start-up; "
                                "put it back or restart DecisionLab.")
    folder = root / "compact-int8" if variant == "int8" else root
    if not (folder / CONFIG_FILE).is_file():
        raise FileNotFoundError(f"No {CONFIG_FILE} in {folder}. This model has no int8 copy; set LIGHTDEC_VARIANT to fp16.")
    return folder


# ----------------------------------------------------------------------------- modeling-code trust (DL-SA-002)
# falcondec_modeling.py is executed as Python when a model loads. Only files whose sha256 (line endings
# normalised to \n) is known are run. KNOWN: the FalconDec modeling file shipped with LightDec v1.0.2 and
# LightDec_V2_Long v1.0.0 (identical). More hashes can be trusted via TRUSTED_MODELING_SHA256 (comma-separated).
MODELING_FILE = "falcondec_modeling.py"
KNOWN_MODELING_SHA256 = frozenset({"cc211c2d50a1e6946ed01860abb77673bb15f33d1022cd0d9c8739e166ec6b93"})


def modeling_sha256(path: Path) -> str:
    return hashlib.sha256(Path(path).read_bytes().replace(b"\r\n", b"\n")).hexdigest()


def trusted_modeling(folder: Path, env: Mapping[str, str]) -> Path:
    """The modeling file to execute from folder, or PermissionError if its code is not on the allowlist."""
    f = Path(folder) / MODELING_FILE
    if not f.is_file():
        raise FileNotFoundError(f"No {MODELING_FILE} in {folder}.")
    extra = {h.strip().lower() for h in (env.get("TRUSTED_MODELING_SHA256") or "").split(",") if h.strip()}
    digest = modeling_sha256(f)
    if digest not in KNOWN_MODELING_SHA256 | extra:
        raise PermissionError(f"Refusing to run {f}: its sha256 {digest} is not a known FalconDec modeling file. "
                              "If you trust it, add the hash to TRUSTED_MODELING_SHA256 in .env.")
    return f


def warmup_enabled(env: Mapping[str, str]) -> bool:
    """Run each model once after loading? Off with DLAB_WARMUP=0 (the ZeroGPU Space: no real GPU outside
    @spaces.GPU, so a warm-up there would fail)."""
    return env.get("DLAB_WARMUP", "1") != "0"