mission-control-v2/backend/services/llamaswap.py

"""
Engine-Service: liest/schreibt die llama-swap config.yaml und spricht die
llama-swap-API. Portiert & erweitert aus Mission Control v1.

NEU in 2.0: `groups` für Ko-Residenz (schnell + schwer gleichzeitig geladen,
`swap:false`) → Multi-Model-Delegation ohne Nachlade-Latenz.
"""

import os
import re

import httpx
from ruamel.yaml.scalarstring import LiteralScalarString

from config import CMD_TEMPLATE, CONFIG_PATH, DEFAULT_TTL, LLAMA_SWAP_URL, SPEC_DRAFT_MODEL_PATH

# Kanonische Rollen (vereinheitlicht ggü. v1: kein manager/reviewer mehr).
ROLE_IDS = {"vision", "coder", "reasoning", "agent", "scout"}

_CTX_RE = re.compile(r"-(?:c|-ctx-size)\s+(\d+)")
_PATH_RE = re.compile(r"-(?:m|-model)\s+([^\s]+)")
_QUANT_RE = re.compile(r"(Q\d_[A-Z0-9_]+|IQ\d_[A-Z0-9_]+|fp16|bf16)\.gguf", re.IGNORECASE)


# --- Lesen -------------------------------------------------------------------
def read_config() -> dict:
    if not CONFIG_PATH.exists():
        return {"models": {}}
    from ruamel.yaml import YAML
    r_yaml = YAML()
    r_yaml.preserve_quotes = True
    with CONFIG_PATH.open("r", encoding="utf-8") as f:
        data = r_yaml.load(f) or {}
    if not data.get("models"):
        data["models"] = {}
    return data


def _parse_model(name: str, spec: dict) -> dict:
    spec = spec or {}
    cmd = str(spec.get("cmd", "")).strip()
    ctx = int(m.group(1)) if (m := _CTX_RE.search(cmd)) else None

    path = filename = quant = ""
    size_bytes = None
    if (m := _PATH_RE.search(cmd)):
        path = m.group(1).replace("'", "").replace('"', "")
        filename = os.path.basename(path)
        if os.path.exists(path):
            size_bytes = os.path.getsize(path)
        if (q := _QUANT_RE.search(path)):
            quant = q.group(1).upper()

    aliases = spec.get("aliases") or []
    if isinstance(aliases, str):
        aliases = [aliases]
    aliases = [str(a) for a in aliases]
    role = aliases[0].lower() if aliases else (name.lower() if name.lower() in ROLE_IDS else None)

    prompt_cache = "--prompt-cache " in cmd or cmd.endswith("--prompt-cache") or "--prompt-cache-all" in cmd
    spec_draft = None
    if "--spec-draft-model" in cmd:
        m_draft = re.search(r"--spec-draft-model\s+([^\s]+)", cmd)
        if m_draft:
            spec_draft = os.path.basename(m_draft.group(1).replace("'", "").replace('"', ""))
    parallel_match = re.search(r"--parallel\s+(\d+)", cmd)
    parallel_slots = int(parallel_match.group(1)) if parallel_match else 1

    from services.caps import capabilities
    return {
        "name": name,
        "role": role,
        "aliases": aliases,
        "api_ids": [name] + aliases,
        "ctx": ctx,
        "ttl": spec.get("ttl"),
        "cmd": cmd,
        "gguf_path": path,
        "filename": filename,
        "quant": quant,
        "size_bytes": size_bytes,
        "incomplete": not path,
        "prompt_cache": prompt_cache,
        "spec_draft_model": spec_draft,
        "parallel_slots": parallel_slots,
        "capabilities": capabilities(
            name=filename or name, cmd=cmd,
            gguf_path=(path if (path and os.path.exists(path)) else ""),
        ),
    }


def list_models() -> list[dict]:
    cfg = read_config()
    return [_parse_model(name, spec) for name, spec in (cfg.get("models") or {}).items()]


def engine_reachable() -> bool:
    try:
        with httpx.Client(timeout=3.0) as c:
            return c.get(f"{LLAMA_SWAP_URL}/v1/models").status_code == 200
    except Exception:
        return False


# --- Schreiben ---------------------------------------------------------------
def model_id_from_path(model_path: str) -> str:
    """Sprechende Modell-ID (= API-Name) aus dem GGUF-Pfad: Repo-Ordnername ohne
    '-GGUF'. Fallback: Dateiname ohne Quant-Suffix.
    Split-GGUFs liegen oft in einem Quant-Unterordner (…/Q4_K_M/file-00001-of-…) →
    dann eine Ebene höher (Repo-Ordner) nehmen, sonst hieße das Modell 'Q4_K_M'."""
    d = os.path.basename(os.path.dirname(model_path))
    if re.fullmatch(r"(I?Q\d[\w]*|UD-Q\d[\w]*|F16|BF16|FP16|F32)", d, flags=re.I):
        d = os.path.basename(os.path.dirname(os.path.dirname(model_path)))
    name = re.sub(r"[-_]?GGUF$", "", d, flags=re.I).strip("-_")
    if not name:
        fn = re.sub(r"\.gguf$", "", os.path.basename(model_path), flags=re.I)
        fn = re.sub(r"-\d+-of-\d+$", "", fn)
        name = re.sub(r"[-_](Q\d[\w]*|IQ\d[\w]*|F16|BF16|FP16|F32)$", "", fn, flags=re.I)
    return name or "modell"


def set_role_alias(cfg: dict, model_id: str, role: str | None) -> None:
    """Rolle als eindeutigen llama-swap-`aliases`-Eintrag setzen (vorher bei allen
    anderen Modellen entfernen). role=None/leer entfernt den Alias."""
    models = cfg.get("models") or {}
    role = (role or "").strip().lower()
    if role:
        for mid, spec in models.items():
            if mid == model_id or not isinstance(spec, dict):
                continue
            al = [a for a in (spec.get("aliases") or []) if str(a).lower() != role]
            if al:
                spec["aliases"] = al
            else:
                spec.pop("aliases", None)
    spec = models.get(model_id)
    if isinstance(spec, dict):
        if role and role != model_id.lower():
            spec["aliases"] = [role]
        else:
            spec.pop("aliases", None)


def _augment_vision(cmd: str, model_path: str, mmproj_path: str | None) -> str:
    """Vision-Modelle brauchen --mmproj <projektor> und --jinja."""
    if mmproj_path:
        if "--mmproj" not in cmd:
            cmd += f" --mmproj {mmproj_path}"
        if "--jinja" not in cmd:
            cmd += " --jinja"
    return cmd


def write_config(cfg: dict) -> None:
    """Atomar schreiben (tmp + os.replace), damit llama-swap mit -watch-config nie
    eine halbe Datei sieht. Fehlende Schreibrechte → klare Meldung."""
    try:
        CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
        tmp = CONFIG_PATH.with_name(CONFIG_PATH.name + ".tmp")
        from ruamel.yaml import YAML
        r_yaml = YAML()
        r_yaml.preserve_quotes = True
        with tmp.open("w", encoding="utf-8") as f:
            r_yaml.dump(cfg, f)
        os.replace(tmp, CONFIG_PATH)
    except PermissionError as exc:
        raise PermissionError(
            f"Mission Control darf '{CONFIG_PATH}' nicht schreiben. "
            f"Einmalig: sudo chown -R hitonabi:hitonabi {CONFIG_PATH.parent}"
        ) from exc



def register_model(model_path: str, role: str | None = None, ctx: int = 8192,
                   ttl: int | None = None, mmproj_path: str | None = None,
                   jinja: bool = False) -> str:
    """Ein GGUF als llama-swap-Modell eintragen (cmd + Rolle-Alias). Gibt die
    Modell-ID zurück. jinja=True erzwingt --jinja (Tool-Calling, z.B. fürs Agent-Hirn)."""
    cfg = read_config()
    model_id = model_id_from_path(model_path)
    cmd = CMD_TEMPLATE.replace("{model}", model_path).replace("{ctx}", str(ctx))
    cmd = _augment_vision(cmd, model_path, mmproj_path)
    if jinja and "--jinja" not in cmd:
        cmd += " --jinja"
        
    role_lower = (role or "").strip().lower()
    if role_lower in ("fast", "coder"):
        if "--parallel" not in cmd:
            cmd += " --parallel 2"
        if os.path.exists(SPEC_DRAFT_MODEL_PATH) and "--spec-draft-model" not in cmd:
            cmd += f" --spec-draft-model {SPEC_DRAFT_MODEL_PATH}"
            
    cfg.setdefault("models", {})[model_id] = {
        "cmd": LiteralScalarString(cmd + "\n"),
        "ttl": ttl if ttl is not None else DEFAULT_TTL,
    }
    set_role_alias(cfg, model_id, role)
    write_config(cfg)
    return model_id


# --- Groups (Ko-Residenz) ----------------------------------------------------
def set_group(group: str, members: list[str], swap: bool = False, persist: bool = False) -> None:
    """llama-swap-`groups`-Eintrag setzen. swap=False → alle Mitglieder dürfen
    GLEICHZEITIG laufen (Ko-Residenz, keine Nachlade-Latenz). persist=True →
    Mitglieder werden nie automatisch entladen."""
    cfg = read_config()
    groups = cfg.setdefault("groups", {})
    groups[group] = {"swap": swap, "persist": persist, "members": list(members)}
    write_config(cfg)


def list_groups() -> dict:
    return read_config().get("groups") or {}


def set_role(model_id: str, role: str | None) -> bool:
    """Rolle (llama-swap-Alias) eines bestehenden Modells setzen/ändern. So tauscht man
    z.B. das `fast`-Hirn: Rolle `fast` auf ein anderes Modell legen (Alias wandert)."""
    cfg = read_config()
    if model_id not in (cfg.get("models") or {}):
        return False
    set_role_alias(cfg, model_id, role)
    write_config(cfg)
    return True


def set_ctx(model_id: str, ctx: int) -> bool:
    """Kontextlänge (-c) eines bestehenden Modells ändern."""
    cfg = read_config()
    spec = (cfg.get("models") or {}).get(model_id)
    if not spec:
        return False
    cmd = str(spec.get("cmd", ""))
    if _CTX_RE.search(cmd):
        cmd = re.sub(r"-(?:c|-ctx-size)\s+\d+", f"-c {ctx}", cmd)
    else:
        cmd = cmd.rstrip() + f" -c {ctx}"
    spec["cmd"] = LiteralScalarString(cmd if cmd.endswith("\n") else cmd + "\n")
    write_config(cfg)
    return True


def delete_model(model_id: str) -> bool:
    """Entfernt einen Modell-Eintrag aus der config.yaml, löscht die zugehörigen
    GGUF-Dateien (auch Splits) vom Datenträger und bereinigt leere Ordner.
    """
    cfg = read_config()
    models = cfg.get("models") or {}
    if model_id not in models:
        return False

    model_spec = models[model_id] or {}
    cmd = str(model_spec.get("cmd", "")).strip()
    if (m := _PATH_RE.search(cmd)):
        path = m.group(1).replace("'", "").replace('"', "")
        if path:
            # 1. Haupt-GGUF-Datei löschen
            if os.path.exists(path):
                try:
                    os.remove(path)
                except Exception:
                    pass

            # 2. Split-GGUF-Teile löschen (z.B. dateiname-00001-of-00005.gguf etc.)
            dirname = os.path.dirname(path)
            basename = os.path.basename(path)
            if os.path.isdir(dirname):
                split_idx = basename.find("-00001-of-")
                if split_idx != -1:
                    prefix = basename[:split_idx]
                    for f in os.listdir(dirname):
                        if f.startswith(prefix) and f.endswith(".gguf"):
                            try:
                                os.remove(os.path.join(dirname, f))
                            except Exception:
                                pass

                # mmproj-Datei (Vision adapter) aus dem Befehl parsen & löschen
                if "mmproj" in cmd:
                    mmproj_match = re.search(r'--mmproj\s+[\'"]?([^\s\'"]+)[\'"]?', cmd)
                    if mmproj_match:
                        m_path = mmproj_match.group(1)
                        if os.path.exists(m_path):
                            try:
                                os.remove(m_path)
                            except Exception:
                                pass

                # 3. Eltern-Ordner löschen, falls er leer ist und nicht der Modelle-Wurzelordner selbst ist
                try:
                    if not os.listdir(dirname) and os.path.basename(dirname) != "models":
                        os.rmdir(dirname)
                except Exception:
                    pass

    del models[model_id]
    for g in (cfg.get("groups") or {}).values():
        if isinstance(g, dict) and model_id in (g.get("members") or []):
            g["members"] = [m for m in g["members"] if m != model_id]
    write_config(cfg)
    return True


def get_running_models() -> list[str]:
    """Fragt den /running Endpunkt von llama-swap ab. Gibt geladene Modelle zurück."""
    try:
        with httpx.Client(timeout=2.0) as c:
            r = c.get(f"{LLAMA_SWAP_URL}/running")
            if r.status_code == 200:
                return r.json().get("running") or []
    except Exception:
        pass
    return []