Cortex-Inara/cortex/routers/local_llm.py

"""
Local LLM settings — per-user host and model configuration.

Routes:
  GET  /settings/local                      → settings page
  POST /settings/local/host                 → save/create host
  POST /settings/local/models/add           → add model entry
  POST /settings/local/models/{id}/activate → set active model
  POST /settings/local/models/{id}/remove   → remove model entry
  GET  /api/local-llm/fetch-models          → proxy to host /api/models (JSON)
"""
import logging
from pathlib import Path

import httpx
import jwt
from fastapi import APIRouter, Form, Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse

from auth_utils import COOKIE_NAME, decode_token
from config import settings as app_settings
import user_settings as us

logger = logging.getLogger(__name__)
router = APIRouter()

_STATIC = Path(__file__).parent.parent / "static"


# ── Auth helper ───────────────────────────────────────────────────────────────

def _get_user(request: Request) -> str | None:
    token = request.cookies.get(COOKIE_NAME)
    if not token:
        return None
    try:
        return decode_token(token)
    except jwt.InvalidTokenError:
        return None


# ── Page renderer ─────────────────────────────────────────────────────────────

def _render(username: str, success: str = "", error: str = "") -> str:
    cfg     = us.get_config(username)
    hosts   = cfg["hosts"]
    models  = cfg["models"]
    active  = cfg.get("active_model_id")

    # Build a host lookup for model rows
    host_by_id = {h["id"]: h for h in hosts}

    # ── Host section ──────────────────────────────────────────────────────────
    if hosts:
        h = hosts[0]   # one host for now
        host_id_val  = h["id"]
        host_label   = h.get("label", "")
        host_url     = h.get("api_url", "")
        host_key_hint = f"…{h['api_key'][-4:]}" if h.get("api_key") else "not set"
    else:
        host_id_val  = ""
        host_label   = ""
        host_url     = app_settings.local_api_url
        host_key_hint = f"server default (…{app_settings.local_api_key[-4:]})" \
                        if app_settings.local_api_key else "not set"

    # ── Model rows ────────────────────────────────────────────────────────────
    model_rows = ""
    for m in models:
        is_active = m["id"] == active
        host      = host_by_id.get(m["host_id"], {})
        host_name = host.get("label") or host.get("api_url") or "unknown host"
        badge     = '<span class="active-badge">active</span>' if is_active else ""
        activate_btn = (
            '<span class="active-label">✓ Active</span>'
            if is_active else
            f'''<form method="POST" action="/settings/local/models/{m["id"]}/activate" style="display:inline">
                  <button type="submit" class="row-btn">Set active</button>
                </form>'''
        )
        model_rows += f'''
        <div class="model-row{"  model-active" if is_active else ""}">
          <div class="model-info">
            <span class="model-label">{m.get("label") or m["model_name"]}</span>{badge}
            <span class="model-name">{m["model_name"]}</span>
            <span class="model-host">{host_name}</span>
          </div>
          <div class="model-actions">
            {activate_btn}
            <form method="POST" action="/settings/local/models/{m["id"]}/remove" style="display:inline"
                  onsubmit="return confirm('Remove {m.get('label') or m['model_name']}?')">
              <button type="submit" class="row-btn danger">Remove</button>
            </form>
          </div>
        </div>'''

    if not model_rows:
        model_rows = '<p class="empty-note">No models added yet. Use "Add Model" below.</p>'

    # ── Host select for Add Model ─────────────────────────────────────────────
    host_options = "".join(
        f'<option value="{h["id"]}">{h.get("label") or h["api_url"]}</option>'
        for h in hosts
    )
    add_section_hidden = "" if hosts else ' style="display:none"'

    html = (_STATIC / "local_llm.html").read_text()
    first_host_id = hosts[0]["id"] if hosts else ""

    html = html.replace("{{ username }}",         username)
    html = html.replace("{{ host_id }}",          host_id_val)
    html = html.replace("{{ host_label }}",       host_label)
    html = html.replace("{{ host_url }}",         host_url)
    html = html.replace("{{ host_key_hint }}",    host_key_hint)
    html = html.replace("{{ model_rows }}",       model_rows)
    html = html.replace("{{ host_options }}",     host_options)
    html = html.replace("{{ first_host_id }}",    first_host_id)
    html = html.replace("{{ add_section_hidden }}", add_section_hidden)
    html = html.replace("{{ has_host }}",         "true" if hosts else "false")
    if success:
        html = html.replace("<!-- SUCCESS -->", f'<p class="msg success">{success}</p>')
    if error:
        html = html.replace("<!-- ERROR -->",   f'<p class="msg error">{error}</p>')
    return html


# ── Routes ────────────────────────────────────────────────────────────────────

@router.get("/settings/local", include_in_schema=False)
async def local_llm_page(request: Request):
    username = _get_user(request)
    if not username:
        return RedirectResponse("/login", status_code=302)
    return HTMLResponse(_render(username))


@router.post("/settings/local/host", include_in_schema=False)
async def save_host(
    request: Request,
    host_id:  str = Form(""),
    label:    str = Form(""),
    api_url:  str = Form(""),
    api_key:  str = Form(""),
):
    username = _get_user(request)
    if not username:
        return RedirectResponse("/login", status_code=302)

    if not api_url.strip():
        return HTMLResponse(_render(username, error="API URL is required."))

    us.save_host(username, host_id or None, label, api_url, api_key)
    logger.info("local LLM host saved: %s", username)
    return HTMLResponse(_render(username, success="Host saved."))


@router.post("/settings/local/models/add", include_in_schema=False)
async def add_model(
    request:    Request,
    host_id:    str = Form(...),
    label:      str = Form(""),
    model_name: str = Form(...),
):
    username = _get_user(request)
    if not username:
        return RedirectResponse("/login", status_code=302)

    if not model_name.strip():
        return HTMLResponse(_render(username, error="Model name is required."))

    us.add_model(username, host_id, label, model_name)
    logger.info("local model added: %s / %s", username, model_name)
    return HTMLResponse(_render(username, success=f"Model \"{label or model_name}\" added."))


@router.post("/settings/local/models/{model_id}/activate", include_in_schema=False)
async def activate_model(request: Request, model_id: str):
    username = _get_user(request)
    if not username:
        return RedirectResponse("/login", status_code=302)

    if not us.set_active_model(username, model_id):
        return HTMLResponse(_render(username, error="Model not found."))

    logger.info("active local model set: %s / %s", username, model_id)
    return HTMLResponse(_render(username, success="Active model updated."))


@router.post("/settings/local/models/{model_id}/remove", include_in_schema=False)
async def remove_model(request: Request, model_id: str):
    username = _get_user(request)
    if not username:
        return RedirectResponse("/login", status_code=302)

    us.remove_model(username, model_id)
    logger.info("local model removed: %s / %s", username, model_id)
    return HTMLResponse(_render(username, success="Model removed."))


@router.get("/api/local-llm/fetch-models")
async def fetch_models(request: Request) -> JSONResponse:
    """Proxy to the configured host's /api/models endpoint.

    Returns [{id, name}] sorted by name, or an error dict.
    """
    username = _get_user(request)
    if not username:
        return JSONResponse({"error": "Not authenticated"}, status_code=401)

    cfg   = us.get_config(username)
    hosts = cfg.get("hosts", [])

    # Fall back to .env if no host configured yet
    if hosts:
        h       = hosts[0]
        api_url = h.get("api_url", "")
        api_key = h.get("api_key", "")
    else:
        api_url = app_settings.local_api_url
        api_key = app_settings.local_api_key

    if not api_url:
        return JSONResponse({"error": "No host configured."}, status_code=400)

    url     = api_url.rstrip("/") + "/api/models"
    headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}

    try:
        async with httpx.AsyncClient(timeout=8) as client:
            resp = await client.get(url, headers=headers)
        resp.raise_for_status()
        data   = resp.json()
        models = [
            {"id": m["id"], "name": m.get("name") or m["id"]}
            for m in data.get("data", [])
        ]
        models.sort(key=lambda m: m["name"].lower())
        return JSONResponse({"models": models})
    except httpx.HTTPStatusError as e:
        return JSONResponse({"error": f"Host returned {e.response.status_code}"}, status_code=502)
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=502)