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

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

import psutil
from config import HF_DOWNLOAD_ENV, MODELS_DIR
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from services import budget, discover, hf, jobengine, llamaswap
from services.fit import evaluate_fit, 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), "running": llamaswap.get_running_models()}


@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 = 0, quant: str = "Q4_K_M", ctx: int = 8192,
        name: str = "", role: str = "") -> dict:
    """Hardware-Fit-Vorschau. params_b<=0 → aus KATALOG (echte Metadaten, MoE-bewusst)
    oder sonst aus dem Namen geschätzt. assigned_ctx = der ctx, der TATSÄCHLICH vergeben
    würde: SETUP-BEWUSST (neben Hirn/warmem Set), nicht nur gegen den Gesamt-RAM.
    So sieht die 'Erweiterte Ansicht' vor dem Download Ampel + echten ctx."""
    ram = _ram_gb()
    pb = params_b if params_b > 0 else budget.params_b_for(name)
    saw = budget.setup_aware_ctx(pb, quant, role=role or None)
    return {
        "params_b": round(pb, 1),
        "fit": evaluate_fit(pb, quant, ctx, ram, name=name),
        "optimal_ctx": max_ctx_for(pb, quant, ram),   # Roh-Obergrenze (Modell allein)
        "assigned_ctx": saw["ctx"],                     # setup-bewusst vergeben
        "budget": {"gtt_gb": saw["gtt_gb"], "reserved_gb": saw["reserved_gb"],
                   "budget_gb": saw["budget_gb"], "mode": saw["mode"]},
        "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.get("/hf/search")
def hf_search(q: str = "") -> dict:
    return {"results": hf.search(q)}


@router.get("/hf/quants")
def hf_quants(repo: str) -> dict:
    repo = hf.normalize_repo(repo)
    return {"repo": repo, "quants": hf.list_quants(repo)}


@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-…). Akzeptiert volle HF-URL ODER org/repo."""
    repo = hf.normalize_repo(req.repo)
    info = hf.resolve_gguf(repo, req.quant)
    if not info["first"]:
        raise HTTPException(404, f"Keine GGUF-Datei für Quant '{req.quant}' in {repo} gefunden.")

    subdir = 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:
        # SETUP-BEWUSST: größter ctx, der neben Hirn/warmem Set passt (nicht nur Modell allein).
        ctx = budget.setup_aware_ctx(budget.params_b_for(repo), req.quant, role=req.role)["ctx"]

    # Sofort registrieren (robust gegen -watch-config) — aber OHNE den Rollen-Alias
    # umzuhängen: der zeigte sonst minutenlang auf eine noch ladende Datei (Lane kalt;
    # bei role=hermes wäre Lucy bis Download-Ende tot gewesen — Review 16.07.).
    try:
        model_id = llamaswap.register_model(
            model_path, role=req.role, ctx=ctx, mmproj_path=mmproj_path, jinja=req.jinja,
            set_alias=False)
    except PermissionError as exc:
        raise HTTPException(500, str(exc))

    # Rolle erst NACH erfolgreichem Download übernehmen. hermes geht dabei durch den
    # warm-bewussten Hirn-Flow (brains-Gruppe, ttl 0, Hermes-Config, Gateway-Restart) —
    # der rohe Alias-Move hatte diesen Flow bisher umgangen.
    role = (req.role or "").strip().lower()

    def _apply_role() -> None:
        if role == "hermes":
            from services.agent import set_agent_brain
            res = set_agent_brain(model_id)
            if not res.get("ok"):
                raise RuntimeError(res.get("reason", "Hirn-Wechsel fehlgeschlagen"))
        else:
            llamaswap.set_role(model_id, role)

    # Download-Job: alle GGUF-Teile (+ mmproj) per --include holen.
    args = [hf.hf_bin(), "download", repo]
    args.extend(info["files"])
    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,
                                 on_done=_apply_role if role else None)
    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)}


class RoleReq(BaseModel):
    role: str | None = None


@router.get("/roles/{role}/recommend")
def recommend_role(role: str) -> dict:
    """Welches installierte Modell passt am besten auf diese Rolle? (Capability + setup-
    bewusster Fit). Basis für 'Empfohlen'-Hinweis + Auto-Pick im Rollen-Zuweisungs-Modal."""
    from services import roles
    return roles.recommend_for_role(role)


@router.post("/models/{model_id}/role")
def set_model_role(model_id: str, body: RoleReq) -> dict:
    new_role = (body.role or "").strip().lower()
    # Guard: Hält das Modell einen LEBENSWICHTIGEN Alias (hermes/embed), darf die Rolle
    # hier nicht weggeklickt werden — sonst verliert Lucy Hirn/Gedächtnis mit einem Klick.
    # Weg: die geschützte Rolle zuerst einem ANDEREN Modell zuweisen (Alias zieht um).
    prot = llamaswap.protected_alias_of(model_id)
    if prot and new_role != prot:
        raise HTTPException(400,
            f"Dieses Modell hält die lebenswichtige Rolle '{prot}'. Weise '{prot}' zuerst "
            f"einem anderen Modell zu — danach lässt sich die Rolle hier ändern.")
    # Das Agent-Hirn (Rolle 'hermes') braucht den warm-bewussten Flow (Alias + brains-Gruppe +
    # ttl 0 + Hermes config.default + Gateway-Restart) — Single Source of Truth UI ↔ Hermes.
    if new_role == "hermes":
        from services.agent import set_agent_brain
        res = set_agent_brain(model_id)
        if not res.get("ok"):
            raise HTTPException(400, res.get("reason", "Fehler beim Setzen des Agent-Hirns"))
        return res
    if not llamaswap.set_role(model_id, body.role):
        raise HTTPException(404, "Modell nicht gefunden")
    return {"ok": True}


class CtxReq(BaseModel):
    ctx: int


@router.get("/models/{model_id}/ctx/auto")
def auto_ctx(model_id: str) -> dict:
    """Setup-bewusster Optimal-ctx für ein bestehendes Modell (Rolle/Params/Quant +
    aktuelles Setup). Basis für den 'Auto'-Button an der Modellkarte."""
    m = next((x for x in llamaswap.list_models() if x["name"] == model_id), None)
    if not m:
        raise HTTPException(404, "Modell nicht gefunden")
    saw = budget.setup_aware_ctx_for_model(m)
    return {"model_id": model_id, "current_ctx": m.get("ctx"),
            "params_b": round(budget.params_of_model(m), 1), "quant": m.get("quant"),
            "role": m.get("role"), **saw}


@router.post("/models/{model_id}/ctx")
def set_model_ctx(model_id: str, body: CtxReq) -> dict:
    if not llamaswap.set_ctx(model_id, body.ctx):
        raise HTTPException(404, "Modell nicht gefunden")
    return {"ok": True}


@router.get("/models/drafts")
def list_drafts(target: str = "") -> dict:
    """Verfügbare Draft-Modelle + ihre Vocab-Kompatibilität zum Ziel-Modell
    (target = GGUF-Pfad). Basis für die idiotensichere Spec-Draft-Auswahl im UI."""
    return llamaswap.drafts_for(target)


class DraftReq(BaseModel):
    draft_path: str | None = None


@router.post("/models/{model_id}/draft")
def set_model_draft(model_id: str, body: DraftReq) -> dict:
    """Setzt/entfernt den Speculative-Decoding-Draft eines Modells. Inkompatible
    (oder nicht prüfbare) Drafts werden serverseitig abgelehnt."""
    try:
        res = llamaswap.set_spec_draft(model_id, body.draft_path)
    except PermissionError as exc:
        raise HTTPException(500, str(exc))
    if not res["ok"]:
        raise HTTPException(400 if "kompatib" in res["reason"].lower() else 404, res["reason"])
    return res


@router.post("/models/unload")
def unload_all_models() -> dict:
    import httpx
    from config import LLAMA_SWAP_URL
    try:
        with httpx.Client(timeout=10.0) as c:
            r = c.post(f"{LLAMA_SWAP_URL}/api/models/unload")
            return {"ok": r.status_code == 200}
    except Exception as exc:
        raise HTTPException(500, str(exc))


@router.post("/models/{model_id}/unload")
def unload_model(model_id: str) -> dict:
    import httpx
    from config import LLAMA_SWAP_URL
    try:
        with httpx.Client(timeout=10.0) as c:
            r = c.post(f"{LLAMA_SWAP_URL}/api/models/unload/{model_id}")
            return {"ok": r.status_code == 200}
    except Exception as exc:
        raise HTTPException(500, str(exc))


@router.post("/models/{model_id}/load")
def load_model(model_id: str) -> dict:
    import httpx
    from config import LLAMA_SWAP_URL
    try:
        # Trigger load by sending a lightweight completion request.
        body = {
            "model": model_id,
            "messages": [{"role": "user", "content": "ping"}],
            "max_tokens": 1
        }
        # High timeout because model loading might take time
        with httpx.Client(timeout=60.0) as c:
            c.post(f"{LLAMA_SWAP_URL}/v1/chat/completions", json=body)
            return {"ok": True}
    except Exception as exc:
        raise HTTPException(500, str(exc))


@router.delete("/models/{model_id}")
def delete(model_id: str) -> dict:
    if not llamaswap.delete_model(model_id):
        raise HTTPException(404, "Modell nicht gefunden")
    return {"ok": True}


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