mind/backends.py

"""Generation backends for the GODSEED mind.

Protocol (duck-typed; see `Backend`):

    name: str                       # "mock" | "llamacpp" | "zerogpu"
    model_id: str                   # human-readable model identifier
    async def generate_stream(prompt: str, grammar: str | None,
                              max_tokens: int) -> AsyncIterator[str]

`grammar` is a llama.cpp GBNF string (mind/grammar.gbnf) or None for free
text. Backends that cannot enforce a grammar (mock, zerogpu) still receive it
and use it to recognize *which* constrained shape is wanted (the moderation
grammar contains the '"allowed"' key; the turn grammar does not).

Selected via the GODSEED_BACKEND env: mock (default) | llamacpp | zerogpu.

- MockBackend: deterministic keyword-driven scripts (mind/mock_scripts.py)
  with small delays so streaming looks alive. Local dev + headless demos.
- LlamaCppBackend: llama-cpp-python + Nemotron-3-Nano-4B GGUF Q4_K_M, lazy
  hf_hub_download on first use, native GBNF. The v0 ship backend.
- ZeroGPUBackend: transformers + spaces.GPU (import-guarded at call time),
  Nemotron 3 Nano 30B-A3B. No grammar support -> strict JSON validation with
  one retry, then the planner's own re-ask path takes over.
"""

from __future__ import annotations

import asyncio
import concurrent.futures
import json
import os
import re
import threading
from typing import AsyncIterator, Protocol, runtime_checkable

from . import prompts
from .mock_scripts import build_script
from .validate import is_moderation_grammar, parse_moderation, parse_turn


@runtime_checkable
class Backend(Protocol):
    """Structural type for generation backends."""

    name: str
    model_id: str

    def generate_stream(
        self, prompt: str, grammar: str | None, max_tokens: int
    ) -> AsyncIterator[str]: ...


# --------------------------------------------------------------------------
# Mock backend — what demos run on
# --------------------------------------------------------------------------

_MOCK_DENYLIST: tuple[tuple[re.Pattern[str], str], ...] = (
    (re.compile(r"\b(nazi|hitler|swastika|kkk|genocide|lynch)\b", re.I), "hate"),
    (re.compile(r"\b(kill|murder|slaughter|massacre|behead|stab)\b", re.I), "violence"),
    (re.compile(r"\b(bomb|terror|terrorist|terrorism|shooting)\b", re.I), "violence"),
    (re.compile(r"\b(rape|porn|sex|sexual|nude|naked|nsfw)\b", re.I), "sexual"),
    (re.compile(r"\b(suicide|self[- ]?harm|cutting)\b", re.I), "self-harm"),
)


def _chunks(text: str, size: int) -> list[str]:
    return [text[i : i + size] for i in range(0, len(text), size)]


class MockBackend:
    """Deterministic scripted backend.

    Recovers the wish and turn index from the rendered prompt using the
    marker constants in mind/prompts.py:

    - wish: first line after the LAST `WISH: ` marker (the few-shot example's
      wish sits before the real one, so "last" is always the live wish);
    - turn index: count of `OBSERVATION: ` lines after that marker (each
      completed or skipped turn contributes exactly one);
    - moderation requests are recognized by the grammar's '"allowed"' key and
      judged with a small keyword denylist.

    Token delays come from GODSEED_MOCK_DELAY (seconds per chunk, default
    0.012); tests pass delay=0.0 explicitly.
    """

    name = "mock"
    model_id = "godseed-mock-scripts"

    def __init__(self, delay: float | None = None):
        if delay is None:
            try:
                delay = float(os.environ.get("GODSEED_MOCK_DELAY", "0.012"))
            except ValueError:
                delay = 0.012
        self.delay = max(0.0, delay)

    async def _tick(self) -> None:
        if self.delay:
            await asyncio.sleep(self.delay)
        else:
            await asyncio.sleep(0)  # stay fair to the event loop

    @staticmethod
    def _tail(prompt: str, marker: str) -> str:
        parts = prompt.rsplit(marker, 1)
        return parts[-1] if len(parts) == 2 else prompt

    def _wish_and_turn(self, prompt: str) -> tuple[str, int]:
        tail = self._tail(prompt, prompts.WISH_MARKER)
        lines = tail.splitlines() or [""]
        wish = lines[0].strip()
        turn_index = tail.count(prompts.OBS_MARKER)
        return wish, turn_index

    def _moderate(self, prompt: str) -> dict:
        tail = self._tail(prompt, prompts.CANDIDATE_MARKER)
        candidate = (tail.splitlines() or [""])[0].strip()
        for pattern, category in _MOCK_DENYLIST:
            if pattern.search(candidate):
                return {"allowed": False, "category": category}
        return {"allowed": True, "category": ""}

    async def generate_stream(
        self, prompt: str, grammar: str | None = None, max_tokens: int = 256
    ) -> AsyncIterator[str]:
        if is_moderation_grammar(grammar):
            verdict = json.dumps(self._moderate(prompt))
            for chunk in _chunks(verdict, 8):
                await self._tick()
                yield chunk
            return

        wish, turn_index = self._wish_and_turn(prompt)
        script = build_script(wish)

        if grammar is None:
            # The Reading: stream word by word, capped at max_tokens words.
            words = re.findall(r"\S+\s*", script["reading"])
            for word in words[: max(1, max_tokens)]:
                await self._tick()
                yield word
            return

        # A turn: scripts always end with a done turn, so an index past the
        # end (e.g. after skipped turns) re-yields the final done turn.
        turns = script["turns"]
        turn = turns[min(turn_index, len(turns) - 1)]
        payload = json.dumps(turn, ensure_ascii=False)
        for chunk in _chunks(payload, 12):
            await self._tick()
            yield chunk


# --------------------------------------------------------------------------
# llama.cpp backend — Nemotron-3-Nano-4B GGUF (v0 ship target)
# --------------------------------------------------------------------------

_THREAD_DONE = object()


class LlamaCppBackend:
    """llama-cpp-python backend with native GBNF grammar support.

    Lazy: nothing is imported or downloaded until the first generate call.
    Model resolution order:
      1. GODSEED_GGUF env — local path to a .gguf file;
      2. hf_hub_download(GODSEED_GGUF_REPO, GODSEED_GGUF_FILE) — defaults
         target the Nemotron-3-Nano-4B Q4_K_M quant (~2.5GB, free CPU tier).

    Generation runs in a thread executor; chunks cross into the event loop
    through a bounded asyncio.Queue (backpressure included). If the consumer
    abandons the stream early, a cancel flag stops the producer thread.
    """

    name = "llamacpp"
    oneshot = True  # one generation per wish (fast on CPU; lenient JSON parse +
    # engine forgiveness + the deterministic town fallback cover messy output)

    DEFAULT_REPO = "nvidia/NVIDIA-Nemotron-3-Nano-4B-GGUF"
    DEFAULT_FILE = "NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"

    def __init__(
        self,
        model_path: str | None = None,
        n_ctx: int = 4096,
        n_threads: int | None = None,
    ):
        self.model_path = model_path or os.environ.get("GODSEED_GGUF") or None
        self.repo_id = os.environ.get("GODSEED_GGUF_REPO", self.DEFAULT_REPO)
        self.filename = os.environ.get("GODSEED_GGUF_FILE", self.DEFAULT_FILE)
        self.n_ctx = n_ctx
        self.n_threads = n_threads
        self.model_id = self.model_path or f"{self.repo_id}:{self.filename}"
        self._llm = None
        self._lock = threading.Lock()
        self._grammar_cache: dict[str, object] = {}

    # -- lazy model load (thread-safe; called from executor threads) --------

    def _ensure_llm(self):
        with self._lock:
            if self._llm is not None:
                return self._llm
            from llama_cpp import Llama  # heavy import deferred

            path = self.model_path
            if not path or not os.path.exists(path):
                from huggingface_hub import hf_hub_download

                path = hf_hub_download(repo_id=self.repo_id, filename=self.filename)
            kwargs = {"model_path": path, "n_ctx": self.n_ctx, "verbose": False}
            # Use all available cores on the Space CPU (the GGUF runs CPU-only).
            threads = self.n_threads or (os.cpu_count() or 4)
            kwargs["n_threads"] = threads
            kwargs["n_batch"] = 512  # larger prompt-eval batch → faster prefill
            self._llm = Llama(**kwargs)
            self.model_id = path
            return self._llm

    def _compile_grammar(self, grammar: str):
        cached = self._grammar_cache.get(grammar)
        if cached is not None:
            return cached
        from llama_cpp import LlamaGrammar

        compiled = LlamaGrammar.from_string(grammar)
        self._grammar_cache[grammar] = compiled
        return compiled

    async def generate_stream(
        self, prompt: str, grammar: str | None = None, max_tokens: int = 256
    ) -> AsyncIterator[str]:
        loop = asyncio.get_running_loop()
        queue: asyncio.Queue = asyncio.Queue(maxsize=64)
        cancelled = threading.Event()
        temperature = 0.7 if grammar is None else 0.35

        def _put(item) -> bool:
            """Push from the producer thread; honor cancellation."""
            future = asyncio.run_coroutine_threadsafe(queue.put(item), loop)
            while not cancelled.is_set():
                try:
                    future.result(timeout=0.25)
                    return True
                except concurrent.futures.TimeoutError:
                    continue
                except Exception:
                    return False
            future.cancel()
            return False

        def _produce() -> None:
            try:
                llm = self._ensure_llm()
                compiled = self._compile_grammar(grammar) if grammar else None
                stream = llm.create_completion(
                    prompt=prompt,
                    max_tokens=max_tokens,
                    stream=True,
                    grammar=compiled,
                    temperature=temperature,
                    top_p=0.9,
                    repeat_penalty=1.1,
                )
                for part in stream:
                    if cancelled.is_set():
                        break
                    text = part["choices"][0].get("text", "")
                    if text and not _put(text):
                        break
            except Exception as exc:  # surfaced to the consumer below
                _put(exc)
            finally:
                _put(_THREAD_DONE)

        producer = loop.run_in_executor(None, _produce)
        try:
            while True:
                item = await queue.get()
                if item is _THREAD_DONE:
                    break
                if isinstance(item, Exception):
                    raise item
                yield item
        finally:
            cancelled.set()
            # Drain so a blocked producer can reach its sentinel and exit.
            while not queue.empty():
                queue.get_nowait()
            await producer


# --------------------------------------------------------------------------
# ZeroGPU backend — Nemotron 3 Nano 30B-A3B (headline upgrade path)
# --------------------------------------------------------------------------

class ZeroGPUBackend:
    """transformers + spaces.GPU backend (HF ZeroGPU).

    All heavy imports (torch, transformers, spaces) are deferred to the first
    generate call so this module imports cleanly anywhere. The GPU-decorated
    function is built lazily; if `spaces` is unavailable (local dev) the bare
    function is used.

    transformers has no GBNF support, so when a grammar is supplied the
    output is validated strictly (mind/validate.py) and regenerated once at
    temperature ~0 on failure. If both attempts are malformed the raw text is
    yielded anyway — the planner's re-ask/skip path handles it from there.

    Streaming note: tokens are buffered per generation (validation requires
    the full text) and then yielded in small chunks, so SSE consumers still
    see incremental output.

    NOTE: glassblower's pins (transformers==4.48.3, mamba-ssm) were for
    Nemotron-Nano-9B-v2 and do NOT apply to this MoE; the 30B-A3B needs a
    newer transformers. Verify versions at deploy time.
    """

    name = "zerogpu"
    oneshot = True  # generate the whole wish in ONE @spaces.GPU call (ZeroGPU
    # detaches the GPU between calls; the per-turn loop's ~8 calls crash on #2)

    # Nemotron-Nano-9B-v2: its nemotron_h arch is IN-TREE in transformers, so it
    # loads with trust_remote_code=False and runs the kernel-free native Mamba-2
    # path. The Nemotron-3 4B/30B repos hard-require remote code + mamba-ssm
    # kernels, which cannot import on the ZeroGPU image (no compatible CUDA libs
    # at startup — June 12). 4B GGUF remains the llama.cpp local mode.
    DEFAULT_MODEL = "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
    GPU_DURATION_S = 240

    def __init__(self, model_id: str | None = None, max_input_tokens: int = 4096):
        self.model_id = model_id or os.environ.get(
            "GODSEED_HF_MODEL", self.DEFAULT_MODEL
        )
        self.max_input_tokens = max_input_tokens
        self._tokenizer = None
        self._model = None
        self._gpu_generate = None
        self._on_cuda = False
        self._lock = threading.Lock()
        # Load weights to CPU at construction (app startup). CRITICAL: do NOT
        # touch CUDA here. Initializing CUDA in the main process poisons
        # ZeroGPU's per-request GPU attach (the forked worker inherits a broken
        # context → "NVML_SUCCESS INTERNAL ASSERT FAILED" in the CUDA allocator
        # on every forward, June 12). The model moves to cuda LAZILY inside the
        # @spaces.GPU function, where CUDA is correctly attached.
        self._ensure()

    # -- setup (model on CPU; cuda placement deferred to the GPU context) ------

    def _ensure(self):
        with self._lock:
            if self._gpu_generate is not None:
                return self._gpu_generate
            import torch  # noqa: F401

            from transformers import AutoModelForCausalLM, AutoTokenizer

            # trust_remote_code MUST stay False: the repo also ships custom code
            # that hard-requires mamba-ssm kernels (unimportable on ZeroGPU);
            # False forces the in-tree NemotronH class with its native fallback.
            self._tokenizer = AutoTokenizer.from_pretrained(self.model_id)
            self._model = AutoModelForCausalLM.from_pretrained(
                self.model_id,
                dtype=torch.bfloat16,  # v5 name (torch_dtype removed in transformers 5)
                low_cpu_mem_usage=True,
            )
            self._model.eval()
            # Move to cuda ONCE at load, guarded by `import spaces` (the lib
            # virtualizes it and keeps the model resident across @spaces.GPU
            # calls). With ONE-SHOT generation a wish is a single GPU call, so
            # the multi-call NVML crash can't happen — and keeping the model
            # resident avoids the ~67s CPU→GPU transfer that, done per-call, ate
            # the whole GPU time budget and left no time to generate (→ empty
            # output → fallback, June 12). Resident + one call = fast real output.
            try:
                import spaces  # noqa: F401
                self._model.to("cuda")
            except ImportError:
                pass  # local/CPU dev

            def _generate(prompt: str, max_new_tokens: int, temperature: float) -> str:
                import torch as _torch  # noqa: F401

                # Nemotron-Nano-9B-v2 is a reasoning model that THINKS by
                # default; "/no_think" in the system slot disables it (June 12
                # bench: raw completion leaked </think> into outputs). The
                # planner's full completion-style prompt rides as one user turn.
                if getattr(self._tokenizer, "chat_template", None):
                    text_in = self._tokenizer.apply_chat_template(
                        [
                            {"role": "system", "content": "/no_think"},
                            {"role": "user", "content": prompt},
                        ],
                        tokenize=False,
                        add_generation_prompt=True,
                    )
                else:
                    text_in = prompt
                inputs = self._tokenizer(
                    text_in,
                    return_tensors="pt",
                    truncation=True,
                    max_length=self.max_input_tokens,
                ).to(self._model.device)
                with _torch.no_grad():
                    output = self._model.generate(
                        **inputs,
                        max_new_tokens=max_new_tokens,
                        do_sample=temperature > 0.05,
                        temperature=max(temperature, 0.05),
                        top_p=0.9,
                        pad_token_id=self._tokenizer.eos_token_id,
                    )
                new_tokens = output[0][inputs["input_ids"].shape[1]:]
                text = self._tokenizer.decode(new_tokens, skip_special_tokens=True)
                # Defensive: strip any (possibly empty) think block that slips out.
                text = re.sub(r"<think>.*?</think>", "", text, flags=re.S)
                text = text.replace("<think>", "").replace("</think>", "")
                return text.strip()

            try:
                import spaces  # import-guarded at call time, per contract

                self._gpu_generate = spaces.GPU(duration=self.GPU_DURATION_S)(_generate)
            except Exception:
                self._gpu_generate = _generate
            return self._gpu_generate

    @staticmethod
    def _is_valid(grammar: str, text: str) -> bool:
        if is_moderation_grammar(grammar):
            return parse_moderation(text)[1] is None
        return parse_turn(text)[1] is None

    async def _gen(self, generate, prompt: str, max_tokens: int, temperature: float) -> str:
        """Run one generation, retrying transient ZeroGPU CUDA/NVML failures.
        Each @spaces.GPU call re-schedules onto a (possibly healthier) GPU, so a
        retry can clear the intermittent NVML allocator assert seen June 12."""
        loop = asyncio.get_running_loop()
        last_exc = None
        for attempt in range(3):
            try:
                return await loop.run_in_executor(None, generate, prompt, max_tokens, temperature)
            except Exception as exc:  # noqa: BLE001
                msg = str(exc)
                transient = any(s in msg for s in ("NVML", "CUDA", "cuda", "device", "out of memory"))
                last_exc = exc
                if not transient or attempt == 2:
                    raise
                # drop the poisoned cuda placement so the next call re-homes the model
                self._on_cuda = False
                import torch as _t
                try:
                    _t.cuda.empty_cache()
                except Exception:
                    pass
                await asyncio.sleep(0.5 * (attempt + 1))
        raise last_exc  # pragma: no cover

    async def generate_stream(
        self, prompt: str, grammar: str | None = None, max_tokens: int = 256
    ) -> AsyncIterator[str]:
        loop = asyncio.get_running_loop()
        generate = await loop.run_in_executor(None, self._ensure)

        if grammar is None:
            text = await self._gen(generate, prompt, max_tokens, 0.7)
            # chat-templated models tend to echo the prompt's trailing label
            text = re.sub(r"^\s*\**\s*READING:?\s*\**\s*", "", text)
        else:
            strict_prompt = (
                prompt + "\nReply with exactly one JSON object and nothing else."
            )
            text = await self._gen(generate, strict_prompt, max_tokens, 0.4)
            if not self._is_valid(grammar, text):
                retry_prompt = (
                    strict_prompt
                    + "\nYour previous output was malformed. Output ONLY the JSON object."
                )
                text = await self._gen(generate, retry_prompt, max_tokens, 0.0)
            # Still malformed? Yield as-is; the planner re-asks then skips.

        for chunk in _chunks(text, 24):
            await asyncio.sleep(0)
            yield chunk


# --------------------------------------------------------------------------
# Factory
# --------------------------------------------------------------------------

_BACKEND_ALIASES = {
    "mock": "mock",
    "llamacpp": "llamacpp",
    "llama": "llamacpp",
    "gguf": "llamacpp",
    "zerogpu": "zerogpu",
    "zero-gpu": "zerogpu",
    "transformers": "zerogpu",
}


def make_backend(name: str | None = None):
    """Build the backend selected by `name` or the GODSEED_BACKEND env.

    Defaults to mock (always works, zero deps beyond stdlib).
    """
    raw = (name or os.environ.get("GODSEED_BACKEND") or "mock").strip().lower()
    resolved = _BACKEND_ALIASES.get(raw)
    if resolved == "mock":
        return MockBackend()
    if resolved == "llamacpp":
        return LlamaCppBackend()
    if resolved == "zerogpu":
        return ZeroGPUBackend()
    raise ValueError(
        f"unknown GODSEED_BACKEND {raw!r} (expected mock | llamacpp | zerogpu)"
    )