gpu-rent/src/gpu_rent/llm_runtime.py

"""Optional LLM runtime (Ollama) beside SwarmUI."""

from __future__ import annotations

import time
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import httpx
import yaml

from gpu_rent.paths import (
    ollama_models_example_path,
    ollama_models_manifest_path,
)

VALID_RUNTIMES = frozenset({"none", "ollama"})

MEMORY_EMBED_MODEL = "nomic-embed-text"

OLLAMA_PRESETS: dict[str, list[str]] = {
    # Requirement: uncensored (abliterated) + solid Russian. Qwen2.5 family.
    # Vision tags preferred for SwarmUI prompt help with images.
    # Memory embed (nomic) is appended separately with use: memory.
    "recommended": ["huihui_ai/qwen2.5-vl-abliterated:7b"],  # ~6 GB
    "light": ["huihui_ai/qwen2.5-vl-abliterated:3b"],  # ~3 GB
    "text": ["huihui_ai/qwen2.5-abliterate:7b"],  # ~5 GB, no vision
    "big": ["huihui_ai/qwen2.5-vl-abliterated:32b"],  # ~21 GB
    "empty": [],
}

OLLAMA_PRESET_LABELS: dict[str, str] = {
    "recommended": "VL 7B abliterate — RU + картинки (~6GB)",
    "light": "VL 3B abliterate — RU + картинки, мало VRAM (~3GB)",
    "text": "7B abliterate text — RU, без vision (~5GB)",
    "big": "VL 32B abliterate — RU + картинки (~21GB)",
    "empty": "только runtime, без pull",
    "keep": "не менять ollama-models.yaml",
}

# Deprecated text blob — prefer menu helpers below.
PRESET_HELP = "\n".join(
    f"{k} — {v}" for k, v in OLLAMA_PRESET_LABELS.items() if k != "keep"
)

LLM_RUNTIME_LABELS: dict[str, str] = {
    "none": "только SwarmUI",
    "ollama": "Ollama (+ pull моделей)",
}

WORKLOAD_LABELS: dict[str, str] = {
    "swarm": "только SwarmUI",
    "both": "SwarmUI + LLM",
    "llm": "только LLM (без SwarmUI)",
}


def llm_runtime_menu() -> list:
    from gpu_rent.prompts import MenuItem

    return [MenuItem(k, f"{k} — {LLM_RUNTIME_LABELS[k]}") for k in ("none", "ollama")]


def workload_menu() -> list:
    from gpu_rent.prompts import MenuItem

    return [MenuItem(k, WORKLOAD_LABELS[k]) for k in ("swarm", "both", "llm")]


def ollama_preset_menu(*, include_keep: bool = False) -> list:
    from gpu_rent.prompts import MenuItem

    keys = list(OLLAMA_PRESETS.keys())
    if include_keep:
        keys.append("keep")
    return [MenuItem(k, OLLAMA_PRESET_LABELS.get(k, k)) for k in keys]


@dataclass(frozen=True)
class OllamaModelEntry:
    name: str
    default: bool = False
    use: str = "chat"  # chat | memory


def _normalize_use(raw: object) -> str:
    s = str(raw or "chat").strip().lower()
    if s in {"memory", "embed", "embedding"}:
        return "memory"
    return "chat"


def normalize_runtime(value: str | None) -> str:
    raw = (value or "none").strip().lower().replace("-", "").replace("_", "")
    if raw in {"", "none", "off", "no", "0"}:
        return "none"
    if raw in {"ollama"}:
        return "ollama"
    raise ValueError(f"неизвестный LLM_RUNTIME={value!r}; жду none|ollama")


def decide_runtime(
    *,
    flag: str | None,
    ollama_flag: bool,
    from_config: str,
) -> str:
    if ollama_flag:
        return "ollama"
    if flag is not None and str(flag).strip() != "":
        return normalize_runtime(flag)
    return normalize_runtime(from_config)


def parse_ollama_models(path: Path) -> list[OllamaModelEntry]:
    if not path.is_file():
        return []
    raw = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
    if not isinstance(raw, dict):
        return []
    items = raw.get("models")
    if items is None:
        return []
    if not isinstance(items, list):
        raise ValueError(f"{path}: models должен быть списком")
    out: list[OllamaModelEntry] = []
    for item in items:
        if isinstance(item, str):
            name = item.strip()
            if name:
                out.append(OllamaModelEntry(name=name, use="chat"))
            continue
        if not isinstance(item, dict):
            continue
        name = str(item.get("name") or "").strip()
        if not name:
            continue
        use = _normalize_use(item.get("use") or item.get("role"))
        out.append(
            OllamaModelEntry(
                name=name,
                default=bool(item.get("default")) and use == "chat",
                use=use,
            )
        )
    return out


def ensure_memory_model_entries(entries: list[OllamaModelEntry]) -> list[OllamaModelEntry]:
    """Append default memory embed if the manifest has chat models but no memory."""
    if not entries:
        return entries
    if any(e.use == "memory" for e in entries):
        return entries
    if all(e.use == "memory" for e in entries):
        return entries
    return list(entries) + [
        OllamaModelEntry(name=MEMORY_EMBED_MODEL, default=False, use="memory")
    ]


def already_have_ollama_tag(have: set[str], wanted: str) -> bool:
    """Exact tag match only — qwen2.5:3b must not satisfy qwen2.5:7b."""
    if wanted in have:
        return True
    if ":" not in wanted and f"{wanted}:latest" in have:
        return True
    if wanted.endswith(":latest") and wanted.rsplit(":", 1)[0] in have:
        return True
    return False


def preferred_ollama_model(path: Path) -> str | None:
    """Manifest default chat model, else first chat tag (never memory/embed)."""
    entries = [e for e in parse_ollama_models(path) if e.use == "chat"]
    for entry in entries:
        if entry.default:
            return entry.name
    return entries[0].name if entries else None


def ollama_roles_payload(entries: list[OllamaModelEntry]) -> dict[str, list[str]]:
    chat = [e.name for e in entries if e.use == "chat"]
    memory = [e.name for e in entries if e.use == "memory"]
    return {"chat": chat, "memory": memory}


def _ollama_ps_names(payload: object) -> set[str]:
    names: set[str] = set()
    if not isinstance(payload, dict):
        return names
    for item in payload.get("models") or []:
        if isinstance(item, dict):
            for key in ("name", "model"):
                val = item.get(key)
                if val:
                    names.add(str(val))
        elif isinstance(item, str) and item.strip():
            names.add(item.strip())
    return names


def warmup_ollama_http(
    base_url: str,
    model: str,
    *,
    keep_alive: str = "15m",
    num_ctx: int = 16384,
    timeout: float = 180.0,
) -> str:
    """Load ``model`` into VRAM via a 1-token /api/chat. Skip if /api/ps already has it."""
    base = str(base_url).rstrip("/")
    t0 = time.monotonic()
    timeout_cfg = httpx.Timeout(timeout, connect=8.0)
    with httpx.Client(timeout=timeout_cfg) as client:
        try:
            ps = client.get(f"{base}/api/ps", timeout=8.0)
            ps.raise_for_status()
            if already_have_ollama_tag(_ollama_ps_names(ps.json()), model):
                return f"Ollama warmup skip — {model} уже в VRAM"
        except (httpx.HTTPError, ValueError, TypeError):
            pass
        body = {
            "model": model,
            "stream": False,
            "keep_alive": keep_alive,
            "options": {
                "num_predict": 1,
                "num_ctx": int(num_ctx),
                "temperature": 0,
            },
            "messages": [{"role": "user", "content": "ok"}],
        }
        resp = client.post(f"{base}/api/chat", json=body)
        resp.raise_for_status()
    elapsed = time.monotonic() - t0
    return f"Ollama warmup {model} ok ({elapsed:.1f}s)"


def maybe_warmup_ollama_local(cfg: Any, log: Callable[[str], None]) -> None:
    """After the laptop tunnel is up: load preferred tag if VRAM is cold."""
    try:
        if normalize_runtime(getattr(cfg, "llm_runtime", "none")) != "ollama":
            return
    except ValueError:
        return
    manifest = Path(getattr(cfg, "ollama_models_manifest", "") or "")
    model = preferred_ollama_model(manifest)
    if not model:
        return
    port = int(getattr(cfg, "ollama_local_port", 17811) or 17811)
    log(f"Ollama warmup {model} (VRAM)")
    try:
        log(
            warmup_ollama_http(
                f"http://127.0.0.1:{port}",
                model,
            )
        )
    except Exception as exc:
        log(f"⚠ Ollama warmup: {exc}")


def write_ollama_models_preset(path: Path, preset: str) -> None:
    key = (preset or "recommended").strip().lower()
    if key not in OLLAMA_PRESETS:
        raise ValueError(f"пресет {preset!r}; варианты: {', '.join(OLLAMA_PRESETS)}")
    names = OLLAMA_PRESETS[key]
    lines = [
        "# Локальный манифест Ollama (не коммить). Пример: ollama-models.example.yaml",
        "# name = точный тег для `ollama pull`.",
        "# use: chat — селект Assistent; use: memory — модель памяти (⚙).",
        "# Пустой models: [] — без pull.",
        "models:",
    ]
    if not names:
        lines.append("  []")
    else:
        for i, name in enumerate(names):
            lines.append(f"  - name: {name}")
            lines.append("    use: chat")
            if i == 0:
                lines.append("    default: true")
        lines.append(f"  - name: {MEMORY_EMBED_MODEL}")
        lines.append("    use: memory")
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def ensure_ollama_manifest_from_example() -> Path:
    dest = ollama_models_manifest_path()
    if dest.is_file():
        return dest
    example = ollama_models_example_path()
    if example.is_file():
        dest.write_text(example.read_text(encoding="utf-8"), encoding="utf-8")
    else:
        write_ollama_models_preset(dest, "recommended")
    return dest


def llm_local_port(cfg: Any) -> int | None:
    runtime = normalize_runtime(getattr(cfg, "llm_runtime", "none"))
    if runtime == "ollama":
        return int(getattr(cfg, "ollama_local_port", 17811))
    return None


def llm_remote_port(runtime: str) -> int | None:
    runtime = normalize_runtime(runtime)
    if runtime == "ollama":
        return 11434
    return None


def append_vars_llm_runtime(vars_file: Path, runtime: str) -> None:
    from gpu_rent.varsfile import upsert_vars

    upsert_vars(vars_file, {"LLM_RUNTIME": normalize_runtime(runtime)})