Cortex-Inara/cortex/routers/local_llm.py

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

Routes:
  GET  /settings/models                         → settings page (canonical)
  GET  /settings/local                          → redirect to /settings/models
  POST /settings/local/host                     → save/create a local host
  POST /settings/local/host/{id}/remove         → remove a host (and its models)
  POST /settings/local/google-account           → save/create a Google account
  POST /settings/local/google-account/{id}/remove → remove a Google account
  POST /settings/local/models/add               → add a model (any provider)
  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 json as _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}
    goog_accts  = registry.get("providers", {}).get("google", {}).get("accounts", [])

    # ── Google account rows ───────────────────────────────────────────────────
    google_account_rows = ""
    for a in goog_accts:
        hint = (a.get("api_key") or "")[:10] + "…" if a.get("api_key") else "no key"
        google_account_rows += f'''
        <div class="account-row">
          <div>
            <span class="account-label">{a.get("label") or "Unnamed"}</span>
            <span class="account-hint">{hint}</span>
          </div>
          <form method="POST" action="/settings/local/google-account/{a["id"]}/remove"
                onsubmit="return confirm('Remove this Google account?')">
            <button type="submit" class="btn-link danger">Remove</button>
          </form>
        </div>'''
    if not google_account_rows:
        google_account_rows = '<p class="empty-note">No accounts configured yet.</p>'

    # ── Local host rows ───────────────────────────────────────────────────────
    host_rows = ""
    for h in hosts:
        key_hint = f"…{h['api_key'][-4:]}" if h.get("api_key") else "not set"
        ht    = h.get("host_type", "openwebui")
        ow    = ' selected' if ht == "openwebui" else ''
        ai    = ' selected' if ht == "openai"    else ''
        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="Gaming 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-row">
              <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="field" style="flex:0 0 auto">
                <label>Type</label>
                <select name="host_type">
                  <option value="openwebui"{ow}>Open WebUI / Ollama</option>
                  <option value="openai"{ai}>OpenAI-compatible (OpenRouter, etc.)</option>
                </select>
              </div>
            </div>
            <div class="btn-row">
              <button type="submit" class="btn btn-secondary btn-sm">Save</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 = "".join(
        f'<option value="{h["id"]}">{h.get("label") or h["api_url"]}</option>'
        for h in hosts
    )

    # ── Model rows (all providers) ────────────────────────────────────────────
    _PROVIDER_BADGE = {
        "claude_cli":   ('<span class="pbadge pb-anthropic">Anthropic</span>', "Claude CLI"),
        "gemini_api":   ('<span class="pbadge pb-google">Google</span>', ""),
        "local_openai": ('<span class="pbadge pb-local">Local</span>', ""),
    }
    model_rows = ""
    for m in models:
        resolved = reg._resolve_model(registry, m["id"])
        if not resolved:
            continue
        mtype = m.get("type", "local_openai")
        badge, default_secondary = _PROVIDER_BADGE.get(mtype, ("", ""))

        if mtype == "local_openai":
            h = host_by_id.get(m.get("host_id", ""), {})
            secondary = h.get("label") or h.get("api_url", "")
        elif mtype == "gemini_api":
            acct = next((a for a in goog_accts if a["id"] == m.get("account_id")), None)
            secondary = acct["label"] if acct else ""
        else:
            secondary = default_secondary

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

        model_rows += f'''
        <div class="model-row">
          <div class="model-info">
            <div>{badge}<span class="model-label">{m.get("label") or m.get("model_name","")}</span>{ctx}</div>
            <span class="model-name">{m.get("model_name","")}</span>
            {sec}
            <div class="tag-row">{tags}</div>
          </div>
          <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>'''
    if not model_rows:
        model_rows = '<p class="empty-note">No models added yet.</p>'

    # ── Role assignment rows ──────────────────────────────────────────────────
    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="Configured 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}">'
            f'<span class="role-name">{role.title()}</span>'
            f'<div class="role-slots">'
        )
        for slot in reg.PRIORITY_KEYS[:3]:
            slot_label = slot.replace("_", " ").title()
            sel = (
                f'<select class="role-select" data-role="{role}" '
                f'data-slot="{slot}" title="{slot_label}">\n{model_opts}\n</select>'
            )
            role_rows += f'<div class="role-slot"><span class="slot-label">{slot_label}</span>{sel}</div>'
        role_rows += '</div></div>'

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

    # ── Catalog data + Google accounts for JS ─────────────────────────────────
    google_accounts_js  = _json.dumps(reg.get_google_accounts(username))
    google_catalog_js   = _json.dumps(reg.get_catalog("google"))
    anthropic_catalog_js = _json.dumps(reg.get_catalog("anthropic"))
    has_hosts = "true" if hosts else "false"

    html = (_STATIC / "local_llm.html").read_text()
    replacements = {
        "{{ username }}":           username,
        "{{ google_account_rows }}": google_account_rows,
        "{{ host_rows }}":          host_rows,
        "{{ model_rows }}":         model_rows,
        "{{ host_options }}":       host_options,
        "{{ role_rows }}":          role_rows,
        "{{ role_data_js }}":       role_data_js,
        "{{ google_accounts_js }}": google_accounts_js,
        "{{ google_catalog_js }}":  google_catalog_js,
        "{{ anthropic_catalog_js }}": anthropic_catalog_js,
        "{{ has_hosts }}":          has_hosts,
    }
    for key, val in replacements.items():
        html = html.replace(key, val)
    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/models", include_in_schema=False)
async def models_page_canonical(request: Request):
    username = _get_user(request)
    if not username:
        return RedirectResponse("/login", status_code=302)
    return HTMLResponse(_render(username))


@router.get("/settings/local", include_in_schema=False)
async def models_page_legacy(request: Request):
    return RedirectResponse("/settings/models", status_code=301)


@router.post("/settings/local/google-account", include_in_schema=False)
async def save_google_account(
    request:    Request,
    account_id: str = Form(""),
    label:      str = Form(""),
    api_key:    str = Form(""),
):
    username = _get_user(request)
    if not username:
        return RedirectResponse("/login", status_code=302)
    if not api_key.strip() and not account_id.strip():
        return HTMLResponse(_render(username, error="API key is required."))
    reg.save_google_account(username, account_id or None, label, api_key)
    return HTMLResponse(_render(username, success="Google account saved."))


@router.post("/settings/local/google-account/{account_id}/remove", include_in_schema=False)
async def remove_google_account(request: Request, account_id: str):
    username = _get_user(request)
    if not username:
        return RedirectResponse("/login", status_code=302)
    reg.remove_google_account(username, account_id)
    return HTMLResponse(_render(username, success="Google account removed."))


@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(""),
    host_type: str = Form("openwebui"),
):
    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, host_type)
    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,
    provider:         str = Form("local"),
    label:            str = Form(""),
    context_k:        int = Form(0),
    tags:             str = Form(""),
    # local-only fields
    host_id:          str = Form(""),
    model_name:       str = Form(""),
    # cloud-only fields
    cloud_model_name: str = Form(""),
    account_id:       str = Form(""),
    credential_id:    str = Form("cli"),
):
    username = _get_user(request)
    if not username:
        return RedirectResponse("/login", status_code=302)

    tag_list = [t.strip() for t in tags.split(",") if t.strip()]

    if provider == "local":
        if not model_name.strip():
            return HTMLResponse(_render(username, error="Model name is required."))
        if not host_id.strip():
            return HTMLResponse(_render(username, error="Select a host."))
        reg.save_model(username, None, host_id, label, model_name, context_k, tag_list)
        display = label or model_name

    elif provider in ("google", "anthropic"):
        if not cloud_model_name.strip():
            return HTMLResponse(_render(username, error="Select a model from the catalog."))
        if provider == "google" and not account_id.strip():
            return HTMLResponse(_render(username, error="Select a Google account."))
        reg.save_cloud_model(
            username, None, provider, cloud_model_name, label,
            account_id=account_id or None,
            credential_id=credential_id or None,
            context_k=context_k, tags=tag_list,
        )
        display = label or cloud_model_name
    else:
        return HTMLResponse(_render(username, error=f"Unknown provider: {provider}"))

    logger.info("model added: %s / %s (%s)", username, display, provider)
    return HTMLResponse(_render(username, success=f'Model "{display}" 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": "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 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", [])

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

    if host:
        api_url, api_key, host_type = host.get("api_url",""), host.get("api_key",""), host.get("host_type","openwebui")
    else:
        api_url, api_key, host_type = app_settings.local_api_url, app_settings.local_api_key, "openwebui"

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

    models_path = "/models" if host_type == "openai" else "/api/models"
    url         = api_url.rstrip("/") + models_path
    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 = sorted(
            [{"id": m["id"], "name": m.get("name") or m["id"]} for m in data.get("data", [])],
            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)