mind/backends.py

"""Generation backends for the NEMOCITY mind.

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

    name: str                       # "mock" | "llamacpp" | "zerogpu"
    model_id: str                   # human-readable model identifier
    oneshot: bool                   # always True — NEMOCITY is one-shot only
    async def generate_stream(prompt: str, grammar: str | None,
                              max_tokens: int) -> AsyncIterator[str]

`grammar` is a plain string TAG (mind/validate.py), never a real grammar:
"build" | "fix" | "moderation" select the wanted JSON shape (validation +
temperature 0.4), a "!strict" suffix means a retry pass at temperature 0,
None means free text. The protocol signature is byte-compatible with
godseed's so the copied plumbing keeps working.

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

- MockBackend: deterministic keyword-driven scripts (mind/mock_scripts.py)
  with small delays so streaming looks alive. Local dev + headless demos.
- ZeroGPUBackend: transformers + spaces.GPU, Nemotron-Nano-9B-v2 resident on
  cuda. ONE @spaces.GPU call per petition (the June 12 lesson). Ship path.
- LlamaCppBackend: llama-cpp-python + a local GGUF; offgrid stretch, not the
  ship path.
"""

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, fix_script
from .validate import (
    FIX_TAG,
    base_tag,
    is_moderation_grammar,
    is_strict,
    parse_build,
    parse_fix,
    parse_moderation,
)


@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.

    Routing (marker constants from mind/prompts.py):
      - moderation tag -> keyword-denylist verdict on the text after
        CANDIDATE_MARKER;
      - prompts containing FIX_MARKER -> fix script (reads the rendered
        candidates/numbers back out of the prompt);
      - everything else -> BUILD JSON keyed on the petition line after the
        last PETITION_MARKER.

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

    name = "mock"
    model_id = "nemocity-mock-scripts"
    oneshot = True

    def __init__(self, delay: float | None = None):
        if delay is None:
            try:
                delay = float(os.environ.get("CITY_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 _petition(self, prompt: str) -> str:
        tail = self._tail(prompt, prompts.PETITION_MARKER)
        return (tail.splitlines() or [""])[0].strip()

    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 = 700
    ) -> 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

        if prompts.FIX_MARKER in prompt:
            payload = fix_script(prompt)
        else:
            payload = build_script(self._petition(prompt))
        for chunk in _chunks(payload, 12):
            await self._tick()
            yield chunk


# --------------------------------------------------------------------------
# llama.cpp backend — local GGUF (offgrid stretch; NOT the ship path)
# --------------------------------------------------------------------------

_THREAD_DONE = object()


class LlamaCppBackend:
    """llama-cpp-python backend (CPU). No grammars — tags only pick the
    temperature; the lenient parsers + planner fallback cover messy output.

    Lazy: nothing is imported or downloaded until the first generate call.
    Model resolution order:
      1. CITY_GGUF env — local path to a .gguf file;
      2. hf_hub_download(CITY_GGUF_REPO, CITY_GGUF_FILE).

    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

    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("CITY_GGUF") or None
        self.repo_id = os.environ.get("CITY_GGUF_REPO", self.DEFAULT_REPO)
        self.filename = os.environ.get("CITY_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()

    # -- 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}
            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

    async def generate_stream(
        self, prompt: str, grammar: str | None = None, max_tokens: int = 700
    ) -> AsyncIterator[str]:
        loop = asyncio.get_running_loop()
        queue: asyncio.Queue = asyncio.Queue(maxsize=64)
        cancelled = threading.Event()
        if grammar is None:
            temperature = 0.7
        else:
            temperature = 0.0 if is_strict(grammar) 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()
                stream = llm.create_completion(
                    prompt=prompt,
                    max_tokens=max_tokens,
                    stream=True,
                    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-Nano-9B-v2 (the ship 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.

    No grammar support, so when a tag is supplied the output is validated
    (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 retry/fallback 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.
    """

    name = "zerogpu"
    oneshot = True  # a petition is ONE @spaces.GPU call (ZeroGPU detaches the
    # GPU between calls; godseed's per-turn loop of ~8 calls crashed 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 (June 12).
    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(
            "CITY_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 guarded by
        # `import spaces` below, which virtualizes the placement.
        self._ensure()

    # -- setup (model resident; cuda placement guarded by spaces) -----------

    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 petition 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 budget and left no time to generate (June 12).
            # Resident + one call = ~14s 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
        if base_tag(grammar) == FIX_TAG:
            return parse_fix(text)[1] is None
        return parse_build(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 = 700
    ) -> 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)
        else:
            temperature = 0.0 if is_strict(grammar) else 0.4
            strict_prompt = (
                prompt + "\nReply with exactly one JSON object and nothing else."
            )
            text = await self._gen(generate, strict_prompt, max_tokens, temperature)
            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 retries then falls back.

        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 CITY_BACKEND env.

    Defaults to mock (always works, zero deps beyond stdlib).
    """
    raw = (name or os.environ.get("CITY_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 CITY_BACKEND {raw!r} (expected mock | zerogpu | llamacpp)"
    )