dashboard/app/plugins/llamacpp.py

from typing import Annotated, Any

import httpx
from fastapi import APIRouter, Form
from fastapi.responses import HTMLResponse

from app.config import Settings, get_settings
from app.plugins.base import Plugin
from app.render import render

router = APIRouter(prefix="/api/plugins/llamacpp", tags=["plugins"])


def _headers(settings: Settings) -> dict[str, str]:
    h: dict[str, str] = {}
    if settings.llama_api_key:
        h["Authorization"] = f"Bearer {settings.llama_api_key}"
    return h


def _client() -> httpx.AsyncClient:
    settings = get_settings()
    return httpx.AsyncClient(
        base_url=settings.llama_base_url.rstrip("/"),
        timeout=settings.llama_timeout,
        headers=_headers(settings),
    )


async def gather_status() -> dict[str, Any]:
    """Query the llama-server router. Never raises; returns status dict."""
    settings = get_settings()
    status: dict[str, Any] = {
        "base_url": settings.llama_base_url,
        "reachable": False,
        "health": None,
        "models": [],
        "error": None,
    }
    models: list[dict[str, str | bool | float]] = []
    try:
        async with _client() as client:
            try:
                r = await client.get("/health")
                status["health"] = r.json().get("status") if r.status_code == 200 else f"http {r.status_code}"
            except httpx.HTTPError:
                pass
            r = await client.get("/models")
            _ = r.raise_for_status()
            status["reachable"] = True
            data = r.json()
            for m in data.get("data", []):
                st: dict[str, Any] = m.get("status") or {}
                item: dict[str, str | bool | float] = {
                    "id": str(m.get("id", "?")),
                    "state": str(st.get("value", "unknown")),
                    "failed": bool(st.get("failed")),
                    "exit_code": str(st.get("exit_code")),
                    "path": m.get("path", ""),
                }
                prog: dict[str, Any] = st.get("progress") or {}
                if prog:
                    done = sum(p.get("done", 0) for p in prog.values())
                    total = sum(p.get("total", 1) for p in prog.values())
                    item["progress"] = round(done / total * 100, 1) if total else 0.0
                models.append(item)
            models.sort(key=lambda m: m["id"])
    except httpx.HTTPError as e:
        status["error"] = f"unreachable: {e.__class__.__name__}"
    except Exception as e: # noqa
        status["error"] = str(e)[:200]
    status["models"] = models
    return status


async def _action(endpoint: str, model: str) -> tuple[bool, str]:
    try:
        async with _client() as client:
            r = await client.post(endpoint, json={"model": model})
            if r.status_code < 300:
                return True, ""
            try:
                detail = r.json()
                msg = detail.get("error") or str(detail)
            except Exception: # noqa
                msg = r.text[:200]
            return False, f"http {r.status_code}: {msg}"
    except httpx.HTTPError as e:
        return False, f"unreachable: {e.__class__.__name__}"


def _with_lists(status: dict[str, Any]) -> dict[str, Any]:
    active = {m["id"] for m in status["models"] if m["state"] in ("loaded", "sleeping", "loading")}
    status["loaded"] = [m for m in status["models"] if m["id"] in active]
    status["available"] = [m for m in status["models"] if m["id"] not in active]
    status["loaded"].sort(key=lambda m: (m["state"] != "loaded", m["state"] != "sleeping", m["id"]))
    status["available"].sort(key=lambda m: m["id"])
    return status


async def _status(message: str, error: str) -> dict[str, Any]:
    status = _with_lists(await gather_status())
    status["message"] = message
    status["error_msg"] = error
    return status


async def _state(message: str = "", error: str = "") -> str:
    return render("plugins/llamacpp_state.html", **await _status(message, error))


async def _skeleton(message: str = "", error: str = "") -> str:
    return render("plugins/llamacpp_skeleton.html", **await _status(message, error))


@router.get("/state")
async def state():
    return HTMLResponse(await _state())


@router.post("/load")
async def load(model: Annotated[str, Form()]):
    ok, err = await _action("/models/load", model)
    return HTMLResponse(
        await _skeleton(
            message=f"loading {model}" if ok else "",
            error="" if ok else err,
        )
    )


@router.post("/unload")
async def unload(model: Annotated[str, Form()]):
    ok, err = await _action("/models/unload", model)
    return HTMLResponse(
        await _skeleton(
            message=f"unloading {model}" if ok else "",
            error="" if ok else err,
        )
    )


@router.post("/rescan")
async def rescan():
    try:
        async with _client() as client:
            r = await client.get("/models", params={"reload": "1"})
            if r.status_code < 300:
                msg = "model list refreshed"
                err = ""
            else:
                msg, err = "", f"http {r.status_code}"
    except httpx.HTTPError as e:
        msg, err = "", f"unreachable: {e.__class__.__name__}"
    return HTMLResponse(await _skeleton(message=msg, error=err))


plugin = Plugin(
    id="llamacpp",
    title="llama.cpp",
    description="Loaded model status for a llama-server in router mode.",
    skeleton_fn=_skeleton,
)