thoxmesh_node/backends.py

"""Local inference backends behind the OpenAI-compatible surface.

``LlamaCppBackend`` is the real one: an in-process ``llama-cpp-python`` model
loaded from a GGUF pulled off the Hub. ``EchoBackend`` is a deterministic stand-in
that lets the server, registration, heartbeat and routing paths be tested in CI
without downloading weights or needing a GPU — the mesh contract is what those
tests are about, and a real model would only make them slow and flaky.
"""

from __future__ import annotations

import logging
import os
import threading
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any, Iterable, Mapping, Sequence

from .errors import BackendError
from .gpu import resolve_gpu_layers

logger = logging.getLogger("thox.backend")


@dataclass
class Completion:
    """One generated completion plus its token accounting."""

    text: str
    prompt_tokens: int
    completion_tokens: int

    @property
    def total_tokens(self) -> int:
        return self.prompt_tokens + self.completion_tokens


class Backend(ABC):
    """A loaded model that can answer chat completions."""

    model_id: str = "unknown"

    @abstractmethod
    def generate(
        self,
        messages: Sequence[Mapping[str, Any]],
        *,
        max_tokens: int,
        temperature: float,
        stop: Sequence[str] | None = None,
    ) -> Completion:
        """Produce a single completion. Must be safe to call concurrently."""

    def close(self) -> None:
        """Release model resources. Idempotent."""


class EchoBackend(Backend):
    """Deterministic backend used by tests and by ``--dry-run``.

    It is intentionally not a language model: it returns a stable, inspectable
    string so a test can assert that a prompt travelled router -> node -> back
    without asserting anything about model quality.
    """

    def __init__(self, model_id: str = "thox-echo") -> None:
        self.model_id = model_id

    def generate(
        self,
        messages: Sequence[Mapping[str, Any]],
        *,
        max_tokens: int,
        temperature: float,
        stop: Sequence[str] | None = None,
    ) -> Completion:
        if not messages:
            raise BackendError("at least one message is required")
        last_user = next(
            (m.get("content", "") for m in reversed(messages) if m.get("role") == "user"),
            messages[-1].get("content", ""),
        )
        text = f"[{self.model_id}] {last_user}"
        words = text.split()
        if max_tokens > 0:
            words = words[:max_tokens]
        rendered = " ".join(words)
        return Completion(
            text=rendered,
            prompt_tokens=sum(len(str(m.get("content", "")).split()) for m in messages),
            completion_tokens=len(words),
        )


class LlamaCppBackend(Backend):
    """GGUF model served in-process via ``llama-cpp-python``.

    A single process-wide lock serialises generation. ``llama.cpp``'s context is
    not safe to use from several threads at once, and on the free CPU tier there
    is no headroom for parallel decoding anyway — one queue with honest latency
    telemetry ranks better in the mesh than several that all thrash.
    """

    def __init__(
        self,
        model_path: str,
        *,
        model_id: str,
        n_ctx: int = 2048,
        n_threads: int = 0,
        n_gpu_layers: str | int | None = "auto",
        chat_format: str | None = None,
    ) -> None:
        try:
            from llama_cpp import Llama
        except ImportError as exc:  # pragma: no cover - env dependent
            raise BackendError(
                "llama-cpp-python is not installed. Install with:\n"
                "  pip install llama-cpp-python\n"
                "There is no manylinux wheel, so it compiles from source and needs a "
                "C/C++ toolchain (build-essential + cmake)."
            ) from exc
        if not os.path.exists(model_path):
            raise BackendError(f"GGUF not found at {model_path}")
        try:
            size = os.path.getsize(model_path)
        except OSError as exc:
            raise BackendError(f"cannot stat GGUF at {model_path}: {exc}") from exc
        if size < 1_000_000:
            # A truncated download or a saved error page presents as a tiny
            # "GGUF" and otherwise fails much later inside the loader with an
            # opaque message.
            raise BackendError(
                f"GGUF at {model_path} is only {size} bytes - the download was "
                "truncated or returned an error page. Delete it and retry."
            )

        self.model_id = model_id
        self._lock = threading.Lock()
        threads = n_threads or _default_threads()
        # Resolved here rather than by callers so every entry point (agent,
        # Space, notebook, CLI) gets the same CPU-safe default.
        gpu_layers = resolve_gpu_layers(n_gpu_layers)
        self.n_gpu_layers = gpu_layers
        logger.info(
            "loading %s (n_ctx=%d, threads=%d, n_gpu_layers=%d)",
            model_path,
            n_ctx,
            threads,
            gpu_layers,
        )
        try:
            self._llama = self._load(Llama, model_path, n_ctx, threads, gpu_layers, chat_format)
        except Exception as exc:  # noqa: BLE001
            message = str(exc)
            if gpu_layers != 0 and _looks_like_gpu_failure(message):
                # A CPU-only llama.cpp build raises when asked to offload. Falling
                # back beats refusing to serve: the node is still a useful CPU
                # member of the mesh, just slower.
                logger.warning(
                    "GPU offload failed (%s); retrying CPU-only. Rebuild "
                    "llama-cpp-python with CMAKE_ARGS='-DGGML_CUDA=ON' for GPU support.",
                    message.splitlines()[0] if message else exc,
                )
                self.n_gpu_layers = 0
                try:
                    self._llama = self._load(Llama, model_path, n_ctx, threads, 0, chat_format)
                except Exception as retry_exc:  # noqa: BLE001
                    raise BackendError(
                        f"failed to load GGUF {model_path} even CPU-only: {retry_exc}"
                    ) from retry_exc
            else:
                raise BackendError(f"failed to load GGUF {model_path}: {exc}") from exc

    @staticmethod
    def _load(llama_cls, model_path, n_ctx, threads, gpu_layers, chat_format):
        """Instantiate llama.cpp. Split out so the CPU retry cannot drift."""
        return llama_cls(
            model_path=model_path,
            n_ctx=n_ctx,
            n_threads=threads,
            n_gpu_layers=gpu_layers,
            verbose=False,
            **({"chat_format": chat_format} if chat_format else {}),
        )

    def generate(
        self,
        messages: Sequence[Mapping[str, Any]],
        *,
        max_tokens: int,
        temperature: float,
        stop: Sequence[str] | None = None,
    ) -> Completion:
        if not messages:
            raise BackendError("at least one message is required")
        payload = [
            {"role": str(m.get("role", "user")), "content": str(m.get("content", ""))}
            for m in messages
        ]
        try:
            with self._lock:
                result = self._llama.create_chat_completion(
                    messages=payload,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    stop=list(stop) if stop else None,
                )
        except Exception as exc:  # noqa: BLE001
            raise BackendError(f"generation failed: {exc}") from exc

        choices = result.get("choices") or []
        if not choices:
            raise BackendError("model returned no choices")
        text = (choices[0].get("message") or {}).get("content") or ""
        usage = result.get("usage") or {}
        return Completion(
            text=text,
            prompt_tokens=int(usage.get("prompt_tokens", 0)),
            completion_tokens=int(usage.get("completion_tokens", 0)),
        )

    def close(self) -> None:
        llama = getattr(self, "_llama", None)
        if llama is not None and hasattr(llama, "close"):
            try:
                llama.close()
            except Exception:  # noqa: BLE001 - best effort
                logger.debug("llama close failed", exc_info=True)


def _default_threads() -> int:
    """Thread count honouring the container's CPU quota, not the host's cores.

    ``os.cpu_count()`` reports the *host* CPUs inside a cgroup-limited container,
    which on a 2-vCPU free Space means spawning many more threads than the quota
    allows and losing throughput to contention. The cgroup v2 quota is read first
    and the host count is only a fallback.
    """
    try:
        with open("/sys/fs/cgroup/cpu.max", "r", encoding="utf-8") as handle:
            quota_text, period_text = handle.read().split()
        if quota_text != "max":
            quota = int(quota_text) / int(period_text)
            if quota >= 1:
                return max(1, int(quota))
    except (OSError, ValueError):
        pass
    return max(1, (os.cpu_count() or 2))


def download_gguf(repo_id: str, filename: str, *, token: str | None = None) -> str:
    """Fetch a GGUF from the Hub and return its local path.

    The token is read from the environment by the caller and passed through; it
    is never written to disk or logged.
    """
    try:
        from huggingface_hub import hf_hub_download
    except ImportError as exc:  # pragma: no cover - env dependent
        raise BackendError("huggingface_hub is required to download a GGUF") from exc
    if not repo_id or not filename:
        raise BackendError("both THOX_MODEL_REPO and THOX_MODEL_FILE must be set")
    logger.info("downloading %s/%s", repo_id, filename)
    try:
        return hf_hub_download(repo_id=repo_id, filename=filename, token=token or None)
    except Exception as exc:  # noqa: BLE001
        # The raw hub error is accurate but rarely actionable. The two failures
        # that actually happen here are a missing token on a private repo and a
        # mistyped GGUF filename, so name both explicitly.
        detail = str(exc)
        lowered = detail.lower()
        hint = ""
        if any(k in lowered for k in ("401", "403", "gated", "unauthorized", "authentication")):
            hint = (
                "\nTHOX model repos are private/gated. Set HF_TOKEN to a token with read "
                "access (Colab: Secrets panel; Space: Settings -> Secrets)."
            )
        elif any(k in lowered for k in ("404", "not found", "entrynotfound")):
            hint = (
                f"\nCheck that '{filename}' exists in '{repo_id}'. GGUF filenames are "
                "case-sensitive and differ between repos "
                "(e.g. thoxmini-3b-Q4_K_M.gguf vs thoxmythos-9b-q4_k_m.gguf)."
            )
        raise BackendError(f"could not download {repo_id}/{filename}: {detail}{hint}") from exc


def build_backend(
    *,
    model_id: str,
    repo_id: str,
    filename: str,
    n_ctx: int,
    n_threads: int,
    n_gpu_layers: str | int | None = "auto",
    token: str | None = None,
    echo: bool = False,
) -> Backend:
    """Construct the configured backend, falling back to echo when asked."""
    if echo or not repo_id:
        if not echo:
            logger.warning("no THOX_MODEL_REPO set; serving the echo backend")
        return EchoBackend(model_id=model_id)
    path = download_gguf(repo_id, filename, token=token)
    return LlamaCppBackend(
        path,
        model_id=model_id,
        n_ctx=n_ctx,
        n_threads=n_threads,
        n_gpu_layers=n_gpu_layers,
    )


def _looks_like_gpu_failure(message: str) -> bool:
    """Whether a loader error is plausibly about GPU offload.

    Matched on text because llama.cpp raises a generic exception for these. Kept
    narrow on purpose: misclassifying a genuine model problem as a GPU problem
    would hide it behind a confusing CPU retry.
    """
    lowered = message.lower()
    return any(
        token in lowered
        for token in (
            "cuda",
            "cublas",
            "gpu",
            "vram",
            "out of memory",
            "no kernel image",
            "device",
        )
    )