LLeMbas/src/lembas/api/admin_models.py

"""Model administration: ordering, defaults, images, access and capabilities."""

from __future__ import annotations

import logging

from fastapi import APIRouter, File, Form, HTTPException, Request, Response, UploadFile, status
from fastapi.responses import FileResponse, RedirectResponse
from sqlalchemy import select
from sqlalchemy.orm import Session as DBSession

from lembas.api.deps import AdminUser, Db, RequiredUser
from lembas.db.models import Connection, Group, Model
from lembas.services import settings_store, uploads
from lembas.web.templating import render

log = logging.getLogger(__name__)

router = APIRouter(tags=["admin-models"])

CAPABILITIES = ("reasoning", "vision", "tools")


def _model(db: DBSession, model_id: str) -> Model:
    model = db.get(Model, model_id)
    if model is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, "That model no longer exists.")
    return model


def _ordered(db: DBSession) -> list[Model]:
    return list(
        db.scalars(
            select(Model).join(Connection).order_by(Model.position, Model.model_id)
        )
    )


def _renumber(db: DBSession) -> None:
    """Rewrite positions to 0..n-1.

    Keeps the numbers dense so a move is always a swap with a neighbour, and
    stops repeated reordering drifting into large sparse values.
    """
    for index, model in enumerate(_ordered(db)):
        model.position = index
    db.commit()


# --- The admin page ----------------------------------------------------------
@router.get("/admin/models")
async def models_page(request: Request, db: Db, user: AdminUser, saved: str = ""):
    models = _ordered(db)
    return render(
        request,
        "admin/models.html",
        {
            "models": models,
            "groups": list(db.scalars(select(Group).order_by(Group.name))),
            "connections": list(db.scalars(select(Connection).order_by(Connection.name))),
            "default_model": settings_store.get(db, "default_model") or "",
            "capabilities": CAPABILITIES,
            "saved": saved,
        },
    )


@router.post("/admin/models/{model_id}")
async def update_model(
    db: Db,
    user: AdminUser,
    model_id: str,
    display_name: str = Form(""),
    description: str = Form(""),
    enabled: bool = Form(False),
    pinned: bool = Form(False),
    public: bool = Form(False),
    group_ids: list[str] = Form(default=[]),
    capability: list[str] = Form(default=[]),
) -> Response:
    model = _model(db, model_id)

    model.display_name = display_name.strip()[:300]
    model.description = description.strip()[:2000]
    model.enabled = enabled
    model.pinned = pinned
    model.public = public

    # Absent checkboxes are simply missing from a form post, so the submitted
    # list IS the complete new state -- rebuild rather than merge.
    model.capabilities_json = {name: (name in capability) for name in CAPABILITIES}

    if public:
        # Group rows would be dead weight and misleading in the UI.
        model.groups = []
    else:
        model.groups = list(db.scalars(select(Group).where(Group.id.in_(group_ids or []))))

    db.commit()
    log.info("model %s updated by %s", model.model_id, user.email)
    return RedirectResponse("/admin/models?saved=Model+saved.", status_code=303)


@router.post("/admin/models/{model_id}/move")
async def move_model(
    db: Db, user: AdminUser, model_id: str, direction: str = Form(...)
) -> Response:
    """Swap a model with its neighbour."""
    model = _model(db, model_id)
    ordered = _ordered(db)
    index = next((i for i, m in enumerate(ordered) if m.id == model.id), None)

    if index is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, "That model no longer exists.")

    target = index - 1 if direction == "up" else index + 1
    if 0 <= target < len(ordered):
        ordered[index], ordered[target] = ordered[target], ordered[index]
        for position, item in enumerate(ordered):
            item.position = position
        db.commit()

    return RedirectResponse("/admin/models", status_code=303)


@router.post("/admin/models/{model_id}/default")
async def set_default_model(db: Db, user: AdminUser, model_id: str) -> Response:
    """Make a model the instance default for new chats."""
    model = _model(db, model_id)
    settings_store.update(db, {"default_model": model.model_id})
    log.info("default model set to %s by %s", model.model_id, user.email)
    return RedirectResponse(
        f"/admin/models?saved={model.label}+is+now+the+default.", status_code=303
    )


@router.post("/admin/models/{model_id}/image")
async def upload_model_image(
    db: Db, user: AdminUser, model_id: str, image: UploadFile = File(...)
) -> Response:
    model = _model(db, model_id)
    payload = await image.read()

    try:
        filename = uploads.save_model_image(payload, image.content_type or "")
    except uploads.UploadError as exc:
        return RedirectResponse(f"/admin/models?saved={exc}", status_code=303)

    # Remove the old file rather than orphaning it in the uploads directory.
    if model.image_path:
        uploads.delete_model_image(model.image_path)

    model.image_path = filename
    db.commit()
    return RedirectResponse("/admin/models?saved=Image+updated.", status_code=303)


@router.post("/admin/models/{model_id}/image/delete")
async def delete_model_image(db: Db, user: AdminUser, model_id: str) -> Response:
    model = _model(db, model_id)
    if model.image_path:
        uploads.delete_model_image(model.image_path)
        model.image_path = ""
        db.commit()
    return RedirectResponse("/admin/models?saved=Image+removed.", status_code=303)


@router.post("/admin/models/bulk")
async def bulk_models(
    db: Db, user: AdminUser, action: str = Form(...), model_ids: list[str] = Form(default=[])
) -> Response:
    """Enable or disable several models at once.

    A freshly refreshed connection can advertise dozens of models; turning them
    off one at a time is not a reasonable way to spend an afternoon.
    """
    models = list(db.scalars(select(Model).where(Model.id.in_(model_ids or []))))
    for model in models:
        if action == "enable":
            model.enabled = True
        elif action == "disable":
            model.enabled = False
        elif action == "public":
            model.public = True
            model.groups = []
        elif action == "private":
            model.public = False
    db.commit()
    _renumber(db)
    return RedirectResponse(
        f"/admin/models?saved={len(models)}+model(s)+updated.", status_code=303
    )


# --- Serving model images ----------------------------------------------------
@router.get("/uploads/models/{filename}")
async def model_image(user: RequiredUser, filename: str) -> Response:
    """Serve a stored model avatar.

    Behind the auth guard: these are instance assets, not public files, and
    the path resolution in uploads refuses anything outside the directory.
    """
    path = uploads.model_image_path(filename)
    if path is None:
        raise HTTPException(status.HTTP_404_NOT_FOUND, "No such image.")
    return FileResponse(
        path,
        media_type=uploads.media_type_for(filename),
        # Filenames are random and content-addressed in practice, so a long
        # cache is safe: a new image gets a new name.
        headers={"Cache-Control": "private, max-age=604800"},
    )