src/id/llm/local.py

"""In-process llama.cpp runtime for self-contained (no-API) deployment.

The Space runs a small local model (OpenBMB MiniCPM4.1-8B, GGUF) through the
llama.cpp runtime instead of calling a cloud API. Generation happens inside an
``@spaces.GPU`` function so it runs on Hugging Face ZeroGPU; off ZeroGPU the
decorator is a no-op and it falls back to CPU/Metal.

The public surface mimics the tiny slice of the OpenAI SDK that
``LLMClient`` uses (``client.chat.completions.create(...)`` returning an object
with ``.choices[0].message.content`` and ``.usage``), so the engine and the
client's call/retry logic stay unchanged.
"""

from __future__ import annotations

import os
import threading
from types import SimpleNamespace
from typing import Any

# Model selection (overridable via Space variables).
GGUF_REPO = os.getenv("F_ID_GGUF_REPO", "openbmb/MiniCPM4.1-8B-GGUF")
GGUF_FILE = os.getenv("F_ID_GGUF_FILE", "*[qQ]4_[kK]_[mM]*.gguf")
N_CTX = int(os.getenv("F_ID_CTX", "16384"))
# -1 offloads every layer to the GPU (ZeroGPU); set 0 to force CPU.
N_GPU_LAYERS = int(os.getenv("F_ID_N_GPU_LAYERS", "-1"))
GPU_DURATION = int(os.getenv("F_ID_GPU_DURATION", "120"))


def _env_flag(name: str, default: bool) -> bool:
    raw = os.getenv(name)
    if raw is None:
        return default
    return raw.strip().lower() in ("1", "true", "yes", "on")


# MiniCPM4.1 is a hybrid reasoning model: its chat template emits a `<think>` block
# unless `enable_thinking=False` is passed to the template render. Reasoning is
# disabled by default here because the long/looping chains stall play-time turns.
ENABLE_THINKING = _env_flag("F_ID_ENABLE_THINKING", False)

_LOCK = threading.Lock()
_LLM: Any = None  # llama_cpp.Llama, built lazily

try:  # ZeroGPU decorator; harmless no-op everywhere else.
    import spaces  # type: ignore

    _gpu = spaces.GPU(duration=GPU_DURATION)
except Exception:  # pragma: no cover - exercised only off ZeroGPU

    def _gpu(fn):  # type: ignore
        return fn


def _preload_cuda() -> None:
    """Make the CUDA runtime resolvable for llama.cpp's CUDA-linked .so.

    The prebuilt ``cu124`` ``llama-cpp-python`` wheel is dynamically linked
    against ``libcudart``/``libcublas``, which are not on the loader path in the
    HF Spaces image. We ship them as ``nvidia-*`` pip wheels and ``dlopen`` them
    with ``RTLD_GLOBAL`` (cudart first) so ``libllama.so`` finds the symbols.
    """
    if N_GPU_LAYERS == 0:
        return
    import ctypes
    import glob
    import site

    roots: set[str] = set()
    getsp = getattr(site, "getsitepackages", None)
    if getsp:
        roots.update(getsp())
    import sys

    roots.update(p for p in sys.path if p.endswith("site-packages"))
    # Load order matters: cudart -> cublasLt -> cublas.
    for pattern in (
        "nvidia/cuda_runtime/lib/libcudart.so*",
        "nvidia/cublas/lib/libcublasLt.so*",
        "nvidia/cublas/lib/libcublas.so*",
    ):
        for root in roots:
            hits = glob.glob(os.path.join(root, pattern))
            if hits:
                try:
                    ctypes.CDLL(hits[0], mode=ctypes.RTLD_GLOBAL)
                except OSError:
                    pass
                break


def _install_no_think_handler(llm: Any) -> None:
    """Force the chat template to render with ``enable_thinking=False``.

    ``Llama.create_chat_completion`` does not forward arbitrary kwargs to the
    Jinja chat template, so the only way to flip the hybrid model's reasoning
    switch is to replace the chat handler with one that injects the flag. We
    rebuild the formatter from the GGUF's embedded template the same way
    llama-cpp-python does internally, subclassing it to pin the flag.
    """
    template = (llm.metadata or {}).get("tokenizer.chat_template")
    if not template:
        return  # no embedded template; nothing to override

    from llama_cpp.llama_chat_format import Jinja2ChatFormatter

    eos_id = llm.token_eos()
    bos_id = llm.token_bos()
    eos_token = llm._model.token_get_text(eos_id) if eos_id != -1 else ""
    bos_token = llm._model.token_get_text(bos_id) if bos_id != -1 else ""

    class _NoThinkFormatter(Jinja2ChatFormatter):
        def __call__(self, **kwargs: Any):  # type: ignore[override]
            kwargs.setdefault("enable_thinking", False)
            return super().__call__(**kwargs)

    llm.chat_handler = _NoThinkFormatter(
        template=template,
        eos_token=eos_token,
        bos_token=bos_token,
        stop_token_ids=[eos_id],
    ).to_chat_handler()


def _load_model() -> Any:
    """Build (once) the llama.cpp model from the cached GGUF."""
    global _LLM
    if _LLM is None:
        with _LOCK:
            if _LLM is None:
                _preload_cuda()
                from llama_cpp import Llama

                llm = Llama.from_pretrained(
                    repo_id=GGUF_REPO,
                    filename=GGUF_FILE,
                    n_ctx=N_CTX,
                    n_gpu_layers=N_GPU_LAYERS,
                    verbose=False,
                )
                if not ENABLE_THINKING:
                    try:
                        _install_no_think_handler(llm)
                    except Exception:  # never block startup on the override
                        pass
                _LLM = llm
    return _LLM


def prefetch() -> None:
    """Download the GGUF to the HF cache on CPU (before any GPU allocation)."""
    # huggingface_hub does not glob, so resolve the filename via the repo listing.
    from fnmatch import fnmatch

    from huggingface_hub import hf_hub_download, list_repo_files

    files = [f for f in list_repo_files(GGUF_REPO) if fnmatch(f, GGUF_FILE)]
    if files:
        hf_hub_download(GGUF_REPO, files[0])


@_gpu
def _generate(
    messages: list[dict[str, str]],
    temperature: float,
    top_p: float | None,
    max_tokens: int | None,
    json_mode: bool,
) -> tuple[str, int, int, int]:
    llm = _load_model()
    kwargs: dict[str, Any] = {"messages": messages, "temperature": temperature}
    if top_p is not None:
        kwargs["top_p"] = top_p
    if max_tokens:
        kwargs["max_tokens"] = max_tokens
    if json_mode:
        kwargs["response_format"] = {"type": "json_object"}
    out = llm.create_chat_completion(**kwargs)
    choice = out["choices"][0]["message"]
    text = choice.get("content") or ""
    usage = out.get("usage") or {}
    pt = int(usage.get("prompt_tokens", 0) or 0)
    ct = int(usage.get("completion_tokens", 0) or 0)
    tt = int(usage.get("total_tokens", 0) or (pt + ct))
    return text, pt, ct, tt


class _Completions:
    def create(
        self,
        *,
        messages: list[dict[str, str]],
        model: str | None = None,
        temperature: float = 0.7,
        top_p: float | None = None,
        max_tokens: int | None = None,
        response_format: dict[str, Any] | None = None,
        **_: Any,
    ) -> SimpleNamespace:
        from .client import LLMError

        json_mode = bool(response_format) and response_format.get("type") == "json_object"
        try:
            text, pt, ct, tt = _generate(list(messages), temperature, top_p, max_tokens, json_mode)
        except Exception as exc:  # surface as the engine's expected error type
            raise LLMError(f"local llama.cpp generation failed: {exc}") from exc
        return SimpleNamespace(
            choices=[SimpleNamespace(message=SimpleNamespace(content=text))],
            usage=SimpleNamespace(prompt_tokens=pt, completion_tokens=ct, total_tokens=tt),
        )


class LocalLlamaClient:
    """Drop-in stand-in for ``openai.OpenAI`` backed by in-process llama.cpp."""

    def __init__(self) -> None:
        self.chat = SimpleNamespace(completions=_Completions())