Cortex-Inara/cortex/user_settings.py

"""
Per-user settings stored in home/{user}/local_llm.json.

Structure:
  {
    "hosts": [{"id", "label", "api_url", "api_key"}, ...],
    "models": [{"id", "host_id", "label", "model_name"}, ...],
    "active_model_id": "<model id>" | null
  }

Values not configured here fall back to .env server defaults.
"""
import json
import logging
import secrets
from pathlib import Path

from config import settings as app_settings

logger = logging.getLogger(__name__)


def _llm_path(username: str) -> Path:
    return app_settings.home_root() / username / "local_llm.json"


def _empty() -> dict:
    return {"hosts": [], "models": [], "active_model_id": None}


def _load(username: str) -> dict:
    path = _llm_path(username)
    if not path.exists():
        return _empty()
    try:
        data = json.loads(path.read_text())
    except (json.JSONDecodeError, OSError):
        logger.warning("local_llm.json for %s is unreadable — starting fresh", username)
        return _empty()

    # Migrate old single-model format {api_url, api_key, model} → new format
    if "hosts" not in data:
        return _migrate_v0(data)

    return data


def _migrate_v0(old: dict) -> dict:
    """Migrate flat {api_url, api_key, model} → hosts/models structure."""
    data = _empty()
    api_url    = old.get("api_url")    or app_settings.local_api_url
    api_key    = old.get("api_key")    or app_settings.local_api_key
    model_name = old.get("model")      or app_settings.local_model

    if not api_url:
        return data

    host_id = secrets.token_hex(4)
    data["hosts"].append({
        "id":      host_id,
        "label":   "Local Model Server",
        "api_url": api_url,
        "api_key": api_key,
    })

    if model_name:
        model_id = secrets.token_hex(4)
        data["models"].append({
            "id":         model_id,
            "host_id":    host_id,
            "label":      model_name,
            "model_name": model_name,
        })
        data["active_model_id"] = model_id

    logger.info("migrated local_llm.json v0 → v1 for user (host=%s)", host_id)
    return data


def _save(username: str, data: dict) -> None:
    _llm_path(username).write_text(json.dumps(data, indent=2))


# ── Public read API ───────────────────────────────────────────────────────────

def get_config(username: str) -> dict:
    """Return the full local LLM config for the user."""
    return _load(username)


def get_active_local_model(username: str) -> dict | None:
    """Return effective {api_url, api_key, model_name, label} for the active model.

    Resolution order:
      1. User's active model + its host config
      2. .env server defaults (LOCAL_API_URL / LOCAL_API_KEY / LOCAL_MODEL)
      3. None — caller should raise a helpful error
    """
    data = _load(username)

    active_id = data.get("active_model_id")
    model = next((m for m in data["models"] if m["id"] == active_id), None)

    if model:
        host = next((h for h in data["hosts"] if h["id"] == model["host_id"]), None)
        if host:
            return {
                "api_url":    host.get("api_url", ""),
                "api_key":    host.get("api_key", ""),
                "model_name": model["model_name"],
                "label":      model.get("label") or model["model_name"],
            }

    # Fall back to .env defaults
    if app_settings.local_api_url and app_settings.local_model:
        return {
            "api_url":    app_settings.local_api_url,
            "api_key":    app_settings.local_api_key,
            "model_name": app_settings.local_model,
            "label":      app_settings.local_model,
        }

    return None


# ── Host management ───────────────────────────────────────────────────────────

def save_host(username: str, host_id: str | None,
              label: str, api_url: str, api_key: str) -> str:
    """Create or update a host. Returns the host ID.

    api_key is only written when non-empty, so submitting a masked placeholder
    with a blank key field leaves the stored key unchanged.
    """
    data = _load(username)

    if host_id:
        for h in data["hosts"]:
            if h["id"] == host_id:
                h["label"]   = label.strip()
                h["api_url"] = api_url.strip()
                if api_key.strip():
                    h["api_key"] = api_key.strip()
                break
        else:
            host_id = None  # ID not found — fall through to create

    if not host_id:
        host_id = secrets.token_hex(4)
        data["hosts"].append({
            "id":      host_id,
            "label":   label.strip(),
            "api_url": api_url.strip(),
            "api_key": api_key.strip(),
        })

    _save(username, data)
    return host_id


# ── Model management ──────────────────────────────────────────────────────────

def add_model(username: str, host_id: str, label: str, model_name: str) -> str:
    """Add a model entry. Auto-activates if it is the first model. Returns the model ID."""
    data = _load(username)
    model_id = secrets.token_hex(4)
    data["models"].append({
        "id":         model_id,
        "host_id":    host_id,
        "label":      label.strip() or model_name.strip(),
        "model_name": model_name.strip(),
    })
    if not data.get("active_model_id"):
        data["active_model_id"] = model_id
    _save(username, data)
    return model_id


def remove_model(username: str, model_id: str) -> None:
    data = _load(username)
    data["models"] = [m for m in data["models"] if m["id"] != model_id]
    if data.get("active_model_id") == model_id:
        data["active_model_id"] = data["models"][0]["id"] if data["models"] else None
    _save(username, data)


def set_active_model(username: str, model_id: str) -> bool:
    """Set the active model. Returns False if the model ID is not found."""
    data = _load(username)
    if not any(m["id"] == model_id for m in data["models"]):
        return False
    data["active_model_id"] = model_id
    _save(username, data)
    return True