LLeMbas/src/lembas/api/admin_models.py

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

from __future__ import annotations

import contextlib
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.services.llm.openai_client import MAX_CONTEXT
from lembas.web.templating import render

log = logging.getLogger(__name__)

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

# What the endpoint can do. Endpoints do not advertise any of this reliably, so
# these are an administrator's assertion.
PROTOCOL_CAPABILITIES = ("reasoning", "vision", "tools")

# Which built-in tools this model is given. Distinct from the above: `tools` is
# whether a tools array may be sent at all, these are what goes in it. Every one
# of them is meaningless unless `tools` is on.
TOOL_CAPABILITIES = (
    ("tool_web_search", "Web search"),
    ("tool_knowledge", "Knowledge"),
    ("tool_notes", "Notes"),
    ("tool_memory", "Memory"),
    ("tool_skills", "Skills"),
)

CAPABILITIES = PROTOCOL_CAPABILITIES + tuple(key for key, _ in TOOL_CAPABILITIES)


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()


# --- Listing -----------------------------------------------------------------
PAGE_SIZE = 40

# Filters offered as tabs above the list. Each is a predicate over a Model.
FILTERS: dict[str, tuple[str, object]] = {
    "all": ("All", lambda m: True),
    "enabled": ("Enabled", lambda m: m.enabled),
    "disabled": ("Disabled", lambda m: not m.enabled),
    "pinned": ("Pinned", lambda m: m.pinned),
    "restricted": ("Restricted", lambda m: not m.public),
}


@router.get("/admin/models")
async def models_page(
    request: Request,
    db: Db,
    user: AdminUser,
    saved: str = "",
    q: str = "",
    filter: str = "all",
    connection: str = "",
    page: int = 1,
):
    """The model list.

    Compact rows only -- editing happens on a page of its own. A connection can
    advertise a hundred models, and a list that renders a full form for each of
    them is unusable at that size.
    """
    everything = _ordered(db)

    predicate = FILTERS.get(filter, FILTERS["all"])[1]
    needle = q.strip().lower()

    matching = [
        model
        for model in everything
        if predicate(model)
        and (not connection or model.connection_id == connection)
        and (
            not needle
            or needle in model.model_id.lower()
            or needle in (model.display_name or "").lower()
        )
    ]

    pages = max(1, -(-len(matching) // PAGE_SIZE))
    page = max(1, min(page, pages))
    start = (page - 1) * PAGE_SIZE
    visible = matching[start : start + PAGE_SIZE]

    return render(
        request,
        "admin/models.html",
        {
            "models": visible,
            "total": len(everything),
            "matched": len(matching),
            "page": page,
            "pages": pages,
            "page_start": start,
            "connections": list(db.scalars(select(Connection).order_by(Connection.name))),
            "default_model": settings_store.get(db, "default_model") or "",
            "counts": {
                key: sum(1 for m in everything if test(m)) for key, (_, test) in FILTERS.items()
            },
            "filters": {key: label for key, (label, _) in FILTERS.items()},
            "active_filter": filter if filter in FILTERS else "all",
            "q": q,
            "connection_id": connection,
            "saved": saved,
        },
    )


@router.get("/admin/models/{model_id}/edit")
async def model_detail(
    request: Request, db: Db, user: AdminUser, model_id: str, saved: str = ""
):
    """Everything about one model, on its own page."""
    model = _model(db, model_id)
    ordered = _ordered(db)
    index = next((i for i, m in enumerate(ordered) if m.id == model.id), 0)

    return render(
        request,
        "admin/model_detail.html",
        {
            "model": model,
            "groups": list(db.scalars(select(Group).order_by(Group.name))),
            "capabilities": PROTOCOL_CAPABILITIES,
            "tool_capabilities": TOOL_CAPABILITIES,
            # Rows predating the split have no tool_* keys at all. Showing them
            # unticked would be a lie: tools.enabled_tools treats absent as on
            # when `tools` is on, so that an upgrade does not silently take web
            # search away from every model already configured for it.
            "tool_default": bool((model.capabilities_json or {}).get("tools")),
            "default_model": settings_store.get(db, "default_model") or "",
            "instance_prompt": settings_store.get(db, "system_prompt") or "",
            "position_of": index + 1,
            "total": len(ordered),
            "previous": ordered[index - 1] if index > 0 else None,
            "next": ordered[index + 1] if index + 1 < len(ordered) else None,
            "saved": saved,
        },
    )


# Registered BEFORE /{model_id}: FastAPI matches in registration order, so
# with the parameterised route first, "bulk" is captured as a model id and
# the handler 404s on a model that does not exist.
@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
    )


@router.post("/admin/models/{model_id}")
async def update_model(
    db: Db,
    user: AdminUser,
    model_id: str,
    display_name: str = Form(""),
    description: str = Form(""),
    system_prompt: str = Form(""),
    enabled: bool = Form(False),
    pinned: bool = Form(False),
    public: bool = Form(False),
    position: str = Form(""),
    context_length: str = Form(""),
    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.system_prompt = system_prompt.strip()[:8000]
    # A string, so an emptied field is distinguishable and junk can be ignored
    # rather than becoming a 422 -- the same shape `position` uses below.
    if context_length.strip():
        with contextlib.suppress(ValueError):
            model.context_length = min(max(int(context_length), 0), MAX_CONTEXT)
    else:
        model.context_length = 0
    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()

    # Typing a position is the only workable way to reorder a long list; the
    # up/down buttons are for nudging a model one place.
    if position.strip():
        try:
            wanted = max(1, int(position)) - 1
        except ValueError:
            wanted = None
        if wanted is not None:
            ordered = [m for m in _ordered(db) if m.id != model.id]
            ordered.insert(min(wanted, len(ordered)), model)
            for index, item in enumerate(ordered):
                item.position = index
            db.commit()

    log.info("model %s updated by %s", model.model_id, user.email)
    return RedirectResponse(
        f"/admin/models/{model.id}/edit?saved=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(...),
    back: 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()

    # Back to whichever filtered, paginated view the button was pressed on.
    return RedirectResponse(back or "/admin/models", status_code=303)


@router.post("/admin/models/{model_id}/default")
async def set_default_model(
    db: Db, user: AdminUser, model_id: str, back: str = Form("")
) -> 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(
        back or f"/admin/models/{model.id}/edit?saved=Now+the+default+model.",
        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/{model.id}/edit?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(
        f"/admin/models/{model.id}/edit?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(
        f"/admin/models/{model.id}/edit?saved=Image+removed.", 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"},
    )