backend/real.py

"""Real backends — small GGUF LLMs via llama-cpp-python.

Image generation was dropped (see backend/factory.py docstring). Only the
LLM path remains real on Spaces.

LLM loader follows the verified Build Small Discord recipe (Dean [UNRL],
2026-05-06): `Llama(...)` is constructed INSIDE the `@spaces.GPU` boundary,
not at module level.

Each `RealLLMBackend(model_id)` loads a different GGUF from the roster in
`game/models.py`. Per-model_id instances are cached by `backend/factory.py`.

Model card references (per CLAUDE.md vendor-citation rule):
  https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-4B-GGUF
  Additional roster URLs live in `game/models.py`.
"""

from __future__ import annotations

import os
from typing import TYPE_CHECKING

from game.models import get_spec

from . import config

if TYPE_CHECKING:
    from llama_cpp import Llama


def force_cpu_mode() -> bool:
    """When set, RealLLMBackend loads with `n_gpu_layers=0` and reduced
    context — usable on the Space's 2 vCPU when ZeroGPU quota is gone.
    Toggled by app.py's tick_decisions fallback path."""
    return os.getenv("TOWNLET_FORCE_CPU") == "1"


class RealLLMBackend:
    """A small GGUF LLM via llama-cpp-python, keyed by `model_id` in the roster.

    First call downloads the GGUF (cached after that). On Spaces with
    `preload_from_hub` in README.md, the cache is warm at boot.
    """

    # Module-scope per-model GGUF paths populated by app.py at boot via
    # huggingface_hub.hf_hub_download — matches the canonical Dean
    # (Rizz Therapy) and gokaygokay patterns. None means "resolve lazily
    # via hf_hub_download on first call", which is the local-dev path.
    _model_paths: dict[str, str] = {}

    @classmethod
    def register_model_path(cls, model_id: str, path: str) -> None:
        cls._model_paths[model_id] = path

    def __init__(self, model_id: str) -> None:
        self.model_id = model_id
        self._spec = get_spec(model_id)
        self._llm: Llama | None = None

    def _ensure_loaded(self) -> "Llama":
        if self._llm is None:
            try:
                from llama_cpp import Llama
            except ImportError as e:
                raise RuntimeError(
                    "llama-cpp-python is not installed. To run with real models "
                    "locally, run `.venv/bin/pip install llama-cpp-python` "
                    "(first time compiles from source, takes 5-15 min on Mac). "
                    "Or run with mocks (no FORCE_REAL_MODELS env var set)."
                ) from e

            # Resolve the GGUF path. Prefer the module-scope-cached path
            # that app.py populates at boot (parent process, no GPU
            # contention). Fall back to hf_hub_download here for the
            # local-dev path where the cache isn't pre-populated.
            path = self._model_paths.get(self.model_id)
            if path is None:
                from huggingface_hub import hf_hub_download
                path = hf_hub_download(
                    repo_id=self._spec.repo_id,
                    filename=self._spec.gguf_filename,
                )

            # Five-kwarg Llama() — matches Dean (Rizz Therapy, the Discord
            # canonical recipe) and 1000-Rooms. Earlier we added type_k/
            # type_v/use_mlock/use_mmap/n_batch overrides; in practice
            # they triggered version-gated TypeError retries on every
            # fork and added latency. Stripping back to the minimal set
            # eliminated that whole class of slowdown.
            # n_ctx=8192 fits the 2-step smolagents memory (~5000 prompt
            # + room for response). CPU-only fallback uses smaller n_ctx
            # so 2 vCPU inference stays usable.
            on_spaces = config.IS_SPACES
            cpu_only = force_cpu_mode()
            if cpu_only:
                gpu_layers = 0
                n_ctx = 4096
                use_flash = False
            elif on_spaces:
                gpu_layers = -1
                n_ctx = 8192
                use_flash = True
            else:
                gpu_layers = 0
                n_ctx = 6144
                use_flash = False

            self._llm = Llama(
                model_path=path,
                n_gpu_layers=gpu_layers,
                n_ctx=n_ctx,
                flash_attn=use_flash,
                verbose=False,
            )
        return self._llm

    def generate(self, system_prompt: str, user_prompt: str) -> str:
        return self.generate_messages(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            max_tokens=256,
        )

    def generate_messages(
        self,
        messages: list[dict],
        stop_sequences: list[str] | None = None,
        max_tokens: int = 1024,
        grammar: str | None = None,
    ) -> str:
        # When `grammar` is supplied, llama.cpp constrains output to that
        # GBNF. We use it for CodeAgent decisions to force the response into
        # a single fenced ```py … ``` block — small chat models otherwise
        # ramble and smolagents' parser fails.
        # https://github.com/ggml-org/llama.cpp/blob/master/grammars/README.md
        llm = self._ensure_loaded()
        kwargs: dict = {
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": 0.8,
        }
        if stop_sequences:
            kwargs["stop"] = stop_sequences
        if grammar:
            from llama_cpp import LlamaGrammar

            kwargs["grammar"] = LlamaGrammar.from_string(grammar)
        result = llm.create_chat_completion(**kwargs)
        return result["choices"][0]["message"]["content"].strip()