gpu-rent/src/gpu_rent/llm_runtime.py

"""Optional LLM runtimes (Ollama / llama.cpp) beside SwarmUI."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

import yaml

from gpu_rent.paths import ollama_models_example_path, ollama_models_manifest_path

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

# Presets for setup / interactive up (prompt help: RU + low refusal).
OLLAMA_PRESETS: dict[str, list[str]] = {
    "recommended": ["huihui_ai/qwen2.5-abliterate:7b"],
    "light": ["qwen2.5:3b"],
    "stock": ["qwen2.5:7b"],
    "alt": ["richardyoung/qwen2.5-7b-instruct-abliterated"],
    "empty": [],
}

PRESET_HELP = (
    "recommended — Qwen2.5 7B abliterate (RU/EN, мало отказов, ~5GB)\n"
    "light — qwen2.5:3b (быстрее, слабее)\n"
    "stock — официальный qwen2.5:7b (больше цензуры)\n"
    "alt — другой abliterate-пак 7B\n"
    "empty — только runtime, без pull"
)


@dataclass(frozen=True)
class OllamaModelEntry:
    name: str
    default: bool = False


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"
    if raw in {"llamacpp", "llama", "llamacppserver"}:
        return "llamacpp"
    raise ValueError(f"неизвестный LLM_RUNTIME={value!r}; жду none|ollama|llamacpp")


def decide_runtime(
    *,
    flag: str | None,
    ollama_flag: bool,
    llamacpp_flag: bool,
    from_config: str,
) -> str:
    """CLI flags win over config/vars."""
    if ollama_flag and llamacpp_flag:
        raise ValueError("укажи только --ollama или --llamacpp, не оба")
    if ollama_flag:
        return "ollama"
    if llamacpp_flag:
        return "llamacpp"
    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))
            continue
        if not isinstance(item, dict):
            continue
        name = str(item.get("name") or "").strip()
        if not name:
            continue
        out.append(OllamaModelEntry(name=name, default=bool(item.get("default"))))
    return out


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`. Пустой models: [] — без pull.",
        "models:",
    ]
    if not names:
        lines.append("  []")
    else:
        for i, name in enumerate(names):
            lines.append(f"  - name: {name}")
            if i == 0:
                lines.append("    default: true")
    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))
    if runtime == "llamacpp":
        return int(getattr(cfg, "llamacpp_local_port", 17812))
    return None


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


def append_vars_llm_runtime(vars_file: Path, runtime: str) -> None:
    runtime = normalize_runtime(runtime)
    line = f"LLM_RUNTIME={runtime}\n"
    if not vars_file.is_file():
        vars_file.write_text("# gpu-rent.vars — несекреты\n" + line, encoding="utf-8")
        return
    text = vars_file.read_text(encoding="utf-8")
    lines = text.splitlines(keepends=True)
    out: list[str] = []
    replaced = False
    for row in lines:
        if row.lstrip().startswith("LLM_RUNTIME="):
            out.append(line if row.endswith("\n") else line.rstrip("\n"))
            replaced = True
        else:
            out.append(row)
    if not replaced:
        if out and not out[-1].endswith("\n"):
            out[-1] = out[-1] + "\n"
        out.append(line)
    vars_file.write_text("".join(out), encoding="utf-8")