Cortex-Inara/cortex/routers/local_llm.py

"""
Model Registry settings — hosts, models, and role assignments.

Routes:
  GET  /settings/local                        → settings page
  POST /settings/local/host                   → save/create a host
  POST /settings/local/host/{id}/remove       → remove a host (and its models)
  POST /settings/local/models/add             → add a model entry
  POST /settings/local/models/{id}/remove     → remove a model
  POST /api/models/role                       → AJAX: set a role assignment
  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 model_registry as reg

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:
    registry = reg.get_registry(username)
    hosts    = registry.get("hosts", [])
    models   = registry.get("models", [])
    roles    = registry.get("roles", {})
    builtins = reg._builtins()

    host_by_id = {h["id"]: h for h in hosts}

    # ── Host rows ─────────────────────────────────────────────────────────────
    host_rows = ""
    for h in hosts:
        key_hint = f"…{h['api_key'][-4:]}" if h.get("api_key") else "not set"
        host_rows += f'''
        <div class="host-row">
          <form method="POST" action="/settings/local/host" class="host-form">
            <input type="hidden" name="host_id" value="{h["id"]}">
            <div class="field-row">
              <div class="field">
                <label>Label</label>
                <input type="text" name="label" value="{h.get("label","")}"
                       placeholder="Home ML Laptop" autocomplete="off" data-form-type="other">
              </div>
              <div class="field" style="flex:2">
                <label>API URL</label>
                <input type="text" name="api_url" value="{h.get("api_url","")}"
                       placeholder="http://192.168.x.x:3000"
                       autocomplete="off" spellcheck="false" data-form-type="other">
              </div>
            </div>
            <div class="field">
              <label>API Key</label>
              <input type="password" name="api_key" placeholder="Leave blank to keep existing"
                     autocomplete="new-password" data-1p-ignore data-lpignore="true" data-form-type="other">
              <p class="key-status">Current: {key_hint}</p>
            </div>
            <div class="btn-row">
              <button type="submit" class="btn btn-secondary btn-sm">Save host</button>
              <button type="button" class="btn btn-secondary btn-sm fetch-btn"
                      data-host-id="{h["id"]}">Fetch models</button>
              <span class="fetch-status" id="fetch-{h["id"]}"></span>
            </div>
          </form>
          <form method="POST" action="/settings/local/host/{h["id"]}/remove"
                onsubmit="return confirm('Remove host and all its models?')" style="margin-top:0.5rem">
            <button type="submit" class="btn-link danger">Remove host</button>
          </form>
        </div>'''

    if not host_rows:
        host_rows = '<p class="empty-note">No hosts configured yet. Add one below.</p>'

    # ── Host options for add-model form ───────────────────────────────────────
    host_options = "".join(
        f'<option value="{h["id"]}">{h.get("label") or h["api_url"]}</option>'
        for h in hosts
    )
    add_model_hidden = "" if hosts else ' style="display:none"'

    # ── Model rows ────────────────────────────────────────────────────────────
    model_rows = ""
    for m in models:
        resolved = reg._resolve_model(registry, m["id"])
        if not resolved:
            continue
        host_name = ""
        if m.get("type") == "local_openai" and m.get("host_id"):
            h = host_by_id.get(m["host_id"], {})
            host_name = h.get("label") or h.get("api_url", "")

        ctx_badge = f'<span class="ctx-badge">{m.get("context_k",0)}k ctx</span>' if m.get("context_k") else ""
        tags_html = " ".join(
            f'<span class="tag">{t}</span>' for t in (m.get("tags") or [])
        )
        host_html = f'<span class="model-host">{host_name}</span>' if host_name else ""

        model_rows += f'''
        <div class="model-row" id="model-{m["id"]}">
          <div class="model-info">
            <span class="model-label">{m.get("label") or m.get("model_name","")}</span>
            <span class="model-name">{m.get("model_name","")}</span>
            {host_html}{ctx_badge}
            <div class="tag-row">{tags_html}</div>
          </div>
          <div class="model-actions">
            <form method="POST" action="/settings/local/models/{m["id"]}/remove"
                  onsubmit="return confirm('Remove this model?')" style="display:inline">
              <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.</p>'

    # ── Role assignment rows ──────────────────────────────────────────────────
    # Build option list: (none) + built-ins + user models
    model_opts = '<option value="">— .env default —</option>\n'
    model_opts += '<optgroup label="Built-in">\n'
    for bid, bm in builtins.items():
        model_opts += f'  <option value="{bid}">{bm["label"]}</option>\n'
    model_opts += '</optgroup>\n'
    if models:
        model_opts += '<optgroup label="Local models">\n'
        for m in models:
            lbl = m.get("label") or m.get("model_name", m["id"])
            model_opts += f'  <option value="{m["id"]}">{lbl}</option>\n'
        model_opts += '</optgroup>\n'

    role_rows = ""
    for role in app_settings.get_defined_roles():
        role_cfg = roles.get(role, {})
        role_rows += f'<div class="role-row" data-role="{role}"><span class="role-name">{role.title()}</span><div class="role-slots">'
        for slot in reg.PRIORITY_KEYS[:3]:  # primary + backup_1 + backup_2
            current = role_cfg.get(slot) or ""
            slot_label = slot.replace("_", " ").title()
            sel_html = f'<select class="role-select" data-role="{role}" data-slot="{slot}" title="{slot_label}">\n{model_opts}\n</select>'
            # Pre-select current value via JS (simpler than string-building selected attrs)
            role_rows += f'<div class="role-slot"><span class="slot-label">{slot_label}</span>{sel_html}</div>'
        role_rows += '</div></div>'

    # JS data for pre-selecting current role values
    import json as _json
    role_data_js = _json.dumps({
        role: {slot: (roles.get(role, {}).get(slot) or "") for slot in reg.PRIORITY_KEYS[:3]}
        for role in app_settings.get_defined_roles()
    })

    html = (_STATIC / "local_llm.html").read_text()
    html = html.replace("{{ username }}",         username)
    html = html.replace("{{ host_rows }}",         host_rows)
    html = html.replace("{{ model_rows }}",        model_rows)
    html = html.replace("{{ host_options }}",      host_options)
    html = html.replace("{{ add_model_hidden }}",  add_model_hidden)
    html = html.replace("{{ role_rows }}",         role_rows)
    html = html.replace("{{ role_data_js }}",      role_data_js)
    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 models_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."))
    reg.save_host(username, host_id or None, label, api_url, api_key)
    logger.info("model registry host saved: %s", username)
    return HTMLResponse(_render(username, success="Host saved."))


@router.post("/settings/local/host/{host_id}/remove", include_in_schema=False)
async def remove_host(request: Request, host_id: str):
    username = _get_user(request)
    if not username:
        return RedirectResponse("/login", status_code=302)
    reg.remove_host(username, host_id)
    return HTMLResponse(_render(username, success="Host removed."))


@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(...),
    context_k:  int = Form(0),
    tags:       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."))
    tag_list = [t.strip() for t in tags.split(",") if t.strip()]
    reg.save_model(username, None, host_id, label, model_name, context_k, tag_list)
    logger.info("model added to registry: %s / %s", username, model_name)
    return HTMLResponse(_render(username, success=f'Model "{label or model_name}" added.'))


@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)
    reg.remove_model(username, model_id)
    return HTMLResponse(_render(username, success="Model removed."))


@router.post("/api/models/role")
async def set_role(request: Request) -> JSONResponse:
    """AJAX: assign a model to a role priority slot.

    Body: {"role": "chat", "slot": "primary", "model_id": "abc123" | ""}
    """
    username = _get_user(request)
    if not username:
        return JSONResponse({"error": "Not authenticated"}, status_code=401)
    try:
        body = await request.json()
    except Exception:
        return JSONResponse({"error": "Invalid JSON"}, status_code=400)

    role     = body.get("role", "").strip()
    slot     = body.get("slot", "").strip()
    model_id = body.get("model_id", "").strip() or None

    if not role or not slot:
        return JSONResponse({"error": "role and slot are required"}, status_code=400)

    ok = reg.set_role(username, role, slot, model_id)
    if not ok:
        return JSONResponse({"error": f"Invalid slot or model_id not found"}, status_code=400)

    logger.info("role set: %s %s.%s = %s", username, role, slot, model_id)
    return JSONResponse({"ok": True})


@router.get("/api/local-llm/fetch-models")
async def fetch_models(request: Request, host_id: str = "") -> JSONResponse:
    """Proxy to the host's /api/models endpoint. host_id selects which host."""
    username = _get_user(request)
    if not username:
        return JSONResponse({"error": "Not authenticated"}, status_code=401)

    registry = reg.get_registry(username)
    hosts = registry.get("hosts", [])

    if host_id:
        host = next((h for h in hosts if h["id"] == host_id), None)
    else:
        host = hosts[0] if hosts else None

    # Fall back to .env
    if host:
        api_url = host.get("api_url", "")
        api_key = host.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)