mission-control-v2/backend/routers/models.py

"""Modelle-Endpoints: Liste (mit Caps), Discover, Fit, Register, Groups."""

import psutil
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel

from config import HF_DOWNLOAD_ENV, MODELS_DIR
from services import discover, hf, jobengine, llamaswap
from services.fit import evaluate_fit, extract_params_b, max_ctx_for

router = APIRouter(prefix="/api")


def _ram_gb() -> float:
    return psutil.virtual_memory().total / (1024 ** 3)


@router.get("/models")
def models() -> dict:
    items = llamaswap.list_models()
    return {"models": items, "count": len(items)}


@router.get("/discover")
def discover_models(force: bool = False) -> dict:
    ram = _ram_gb()
    data = discover.refresh_discover(ram) if force else discover.safe_discover(ram)
    if not data:
        raise HTTPException(502, "Modell-Quellen gerade nicht erreichbar — später erneut.")
    return {**data, "sys_ram_gb": round(ram, 1)}


@router.get("/fit")
def fit(params_b: float, quant: str = "Q4_K_M", ctx: int = 8192, name: str = "") -> dict:
    ram = _ram_gb()
    return {
        "fit": evaluate_fit(params_b, quant, ctx, ram, name=name),
        "optimal_ctx": max_ctx_for(params_b, quant, ram),
        "sys_ram_gb": round(ram, 1),
    }


class RegisterReq(BaseModel):
    model_path: str
    role: str | None = None
    ctx: int = 8192
    ttl: int | None = None
    mmproj_path: str | None = None
    jinja: bool = False


@router.post("/models/register")
def register(req: RegisterReq) -> dict:
    try:
        model_id = llamaswap.register_model(
            req.model_path, role=req.role, ctx=req.ctx, ttl=req.ttl,
            mmproj_path=req.mmproj_path, jinja=req.jinja,
        )
    except PermissionError as exc:
        raise HTTPException(500, str(exc))
    return {"ok": True, "model_id": model_id}


class InstallReq(BaseModel):
    repo: str
    role: str | None = None
    quant: str = "Q4_K_M"
    ctx: int | None = None
    jinja: bool = False
    hf_token: str | None = None


@router.post("/models/install")
def install(req: InstallReq) -> dict:
    """Lädt ein Modell von HuggingFace (Hintergrund-Job) UND trägt es sofort in
    llama-swap ein (cmd + Rolle-Alias). llama-swap (-watch-config) lädt es, sobald
    die Datei da ist. Split-GGUFs werden komplett geladen, registriert wird der
    erste Teil (-00001-of-…)."""
    info = hf.resolve_gguf(req.repo, req.quant)
    if not info["first"]:
        raise HTTPException(404, f"Keine GGUF-Datei für Quant '{req.quant}' in {req.repo} gefunden.")

    subdir = req.repo.split("/")[-1]
    target = MODELS_DIR / subdir
    target.mkdir(parents=True, exist_ok=True)
    model_path = str(target / info["first"])
    mmproj_path = str(target / info["mmproj"]) if info["mmproj"] else None

    ctx = req.ctx
    if ctx is None:
        ram = _ram_gb()
        ctx = max_ctx_for(extract_params_b(req.repo), req.quant, ram)

    # Sofort registrieren (Datei kommt gleich) — robust gegen -watch-config.
    try:
        model_id = llamaswap.register_model(
            model_path, role=req.role, ctx=ctx, mmproj_path=mmproj_path, jinja=req.jinja)
    except PermissionError as exc:
        raise HTTPException(500, str(exc))

    # Download-Job: alle GGUF-Teile (+ mmproj) per --include holen.
    args = [hf.hf_bin(), "download", req.repo]
    for f in info["files"]:
        args.append(f)
    if info["mmproj"]:
        args.append(info["mmproj"])
    args += ["--local-dir", str(target)]
    env = dict(HF_DOWNLOAD_ENV)
    if req.hf_token:
        env["HF_TOKEN"] = req.hf_token
    job_id = jobengine.start_job(args, f"download {req.repo}", env=env)
    jobengine.attach_download_progress(job_id, str(target), info["total_bytes"])
    return {"ok": True, "job_id": job_id, "model_id": model_id, "model_path": model_path,
            "total_bytes": info["total_bytes"], "files": len(info["files"])}


@router.get("/jobs")
def jobs() -> dict:
    return {"jobs": jobengine.public_jobs()}


@router.post("/jobs/{job_id}/cancel")
def cancel(job_id: str) -> dict:
    return {"ok": jobengine.cancel_job(job_id)}


@router.get("/groups")
def groups() -> dict:
    return {"groups": llamaswap.list_groups()}


class GroupReq(BaseModel):
    group: str
    members: list[str]
    swap: bool = False
    persist: bool = False


@router.put("/groups")
def set_group(req: GroupReq) -> dict:
    try:
        llamaswap.set_group(req.group, req.members, swap=req.swap, persist=req.persist)
    except PermissionError as exc:
        raise HTTPException(500, str(exc))
    return {"ok": True}