261 lines
8.4 KiB
Python
261 lines
8.4 KiB
Python
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]:
|
|
"""Build the request headers for llama-server calls.
|
|
|
|
Args:
|
|
settings: app settings (provides the optional API key).
|
|
|
|
Returns:
|
|
Headers including a Bearer Authorization only when
|
|
`DASH_LLAMA_API_KEY` is set.
|
|
"""
|
|
h: dict[str, str] = {}
|
|
if settings.llama_api_key:
|
|
h["Authorization"] = f"Bearer {settings.llama_api_key}"
|
|
return h
|
|
|
|
|
|
def _client() -> httpx.AsyncClient:
|
|
"""Create an httpx client pointed at the configured llama-server.
|
|
|
|
Returns:
|
|
An AsyncClient with base URL, timeout, and auth headers from
|
|
settings (callers must use it as an async context manager).
|
|
"""
|
|
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 for health and loaded-model status.
|
|
|
|
Hits /health and /models on the router endpoint. Per model it records
|
|
the router state (loading/loaded/sleeping/...), failure info, path,
|
|
and — when the router reports progress — an aggregate load percentage
|
|
(done/total summed over the progress fields). Never raises: any
|
|
failure is folded into the "error" field so the UI can still render.
|
|
|
|
Returns:
|
|
A dict with base_url, reachable, health, error, and models
|
|
(sorted by model id).
|
|
"""
|
|
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]:
|
|
"""POST a load/unload action to the llama-server router.
|
|
|
|
Args:
|
|
endpoint: router endpoint path, "/models/load" or "/models/unload".
|
|
model: model id to act on.
|
|
|
|
Returns:
|
|
(True, "") on success, else (False, error description) covering
|
|
HTTP errors and unreachable-server cases.
|
|
"""
|
|
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]:
|
|
"""Split the model list into "loaded" and "available" for the UI.
|
|
|
|
A model counts as active while its state is loaded, sleeping, or
|
|
loading. "loaded" is sorted loaded-first, then sleeping, then by id;
|
|
"available" is sorted by id.
|
|
|
|
Args:
|
|
status: dict from gather_status.
|
|
|
|
Returns:
|
|
The same dict, mutated to carry the two extra lists.
|
|
"""
|
|
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]:
|
|
"""Build the template context: live status plus flash message/error.
|
|
|
|
Args:
|
|
message: transient success message to display, or "".
|
|
error: transient error message to display, or "".
|
|
|
|
Returns:
|
|
gather_status() output with loaded/available lists, message, and
|
|
error_msg added.
|
|
"""
|
|
status = _with_lists(await gather_status())
|
|
status["message"] = message
|
|
status["error_msg"] = error
|
|
return status
|
|
|
|
|
|
async def _state(message: str = "", error: str = "") -> str:
|
|
"""Render the compact state fragment (polling view).
|
|
|
|
Args:
|
|
message: transient success message, or "".
|
|
error: transient error message, or "".
|
|
|
|
Returns:
|
|
The rendered llamacpp_state.html fragment.
|
|
"""
|
|
return render("plugins/llamacpp_state.html", **await _status(message, error))
|
|
|
|
|
|
async def _skeleton(message: str = "", error: str = "") -> str:
|
|
"""Render the full skeleton fragment (initial + post-action view).
|
|
|
|
Args:
|
|
message: transient success message, or "".
|
|
error: transient error message, or "".
|
|
|
|
Returns:
|
|
The rendered llamacpp_skeleton.html fragment.
|
|
"""
|
|
return render("plugins/llamacpp_skeleton.html", **await _status(message, error))
|
|
|
|
|
|
@router.get("/state")
|
|
async def state():
|
|
"""Poll endpoint: return the compact state fragment."""
|
|
return HTMLResponse(await _state())
|
|
|
|
|
|
@router.post("/load")
|
|
async def load(model: Annotated[str, Form()]):
|
|
"""Ask the router to load a model, then re-render the skeleton.
|
|
|
|
Args:
|
|
model: model id from the form.
|
|
|
|
Returns:
|
|
The skeleton fragment with a success message or error.
|
|
"""
|
|
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()]):
|
|
"""Ask the router to unload a model, then re-render the skeleton.
|
|
|
|
Args:
|
|
model: model id from the form.
|
|
|
|
Returns:
|
|
The skeleton fragment with a success message or error.
|
|
"""
|
|
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():
|
|
"""Ask the router to rescan its model directory, then re-render.
|
|
|
|
Returns:
|
|
The skeleton fragment with a refresh message or error.
|
|
"""
|
|
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,
|
|
)
|