visualnovel/llm.py

"""One thin LLM wrapper, two methods:

    complete(messages)               -> str            (free text; used by compact_memory)
    complete_json(messages, schema)  -> dict           (grammar-constrained structured output)

Four backends behind one interface:
    MockLLM          - zero deps, returns schema-shaped defaults (unit tests + smoke)
    LlamaCppLLM      - llama.cpp via llama-cpp-python + GBNF/JSON-schema grammar  (the bonus)
    TransformersLLM  - HF transformers fallback for the ZeroGPU Space
    ModalLLM         - cloud GPU via Modal (set VN_LLM_BACKEND=modal)

Heavy imports are done lazily inside each real backend so a MOCK checkout needs nothing.
"""

from __future__ import annotations

from typing import Any, Protocol

from . import config
from .utils import _quiet_stderr, _safe_print, close_truncated_json, strip_think


def _hf_gen_kwargs(kw: dict[str, Any]) -> dict[str, Any]:
    """Map our complete()/complete_json() kwargs to transformers generate() kwargs.

    Greedy decode when temperature is absent or 0; otherwise enable sampling so
    temperature/top_p actually apply (generate() silently ignores them without
    do_sample=True) and retries produce different outputs.
    """
    out: dict[str, Any] = {"max_new_tokens": kw.get("max_tokens", 512)}
    temperature = kw.get("temperature")
    if temperature:
        out["do_sample"] = True
        out["temperature"] = temperature
        if "top_p" in kw:
            out["top_p"] = kw["top_p"]
    else:
        out["do_sample"] = False
    return out


class LLMBackend(Protocol):
    def complete(self, messages: list[dict[str, str]], **kw: Any) -> str: ...
    def complete_json(self, messages: list[dict[str, str]], schema: dict, **kw: Any) -> dict: ...


# --------------------------------------------------------------------------- #
#  Mock — only used by unit tests / when orchestrator isn't short-circuiting.
#  (In MOCK mode the orchestrator builds nice themed output itself; see orchestrator.py.)
# --------------------------------------------------------------------------- #
class MockLLM:
    def complete(self, messages: list[dict[str, str]], **kw: Any) -> str:
        return "The wood remembers little, and dreams the rest."

    def complete_json(self, messages: list[dict[str, str]], schema: dict, **kw: Any) -> dict:
        return _empty_from_schema(schema)


# --------------------------------------------------------------------------- #
#  llama.cpp  (Llama-Champion bonus + Off-the-Grid local-first)
# --------------------------------------------------------------------------- #
class LlamaCppLLM:
    def __init__(self) -> None:
        from llama_cpp import Llama  # noqa: PLC0415

        # Pull the GGUF once into ./models (see scripts/download_models.py), then load it.
        # _quiet_stderr suppresses C-level messages (e.g. n_ctx_seq < n_ctx_train)
        # that llama.cpp emits before verbose=False takes effect.
        with _quiet_stderr():
            self.llm = Llama(
                model_path=str(config.MODELS_DIR / config.LLM_GGUF_FILE),
                n_ctx=8192,
                n_gpu_layers=-1,  # offload all layers (Metal on Mac / ROCm-HIP on AMD)
                verbose=False,
            )

    def complete(self, messages: list[dict[str, str]], **kw: Any) -> str:
        out = self.llm.create_chat_completion(messages=messages, **kw)
        return out["choices"][0]["message"]["content"]

    def complete_json(self, messages: list[dict[str, str]], schema: dict, **kw: Any) -> dict:
        import json  # noqa: PLC0415

        # llama-cpp-python builds the grammar from the JSON schema for you:
        out = self.llm.create_chat_completion(
            messages=messages,
            response_format={"type": "json_object", "schema": schema},
            temperature=kw.get("temperature", 0.7),
            top_p=kw.get("top_p", 0.9),
            max_tokens=kw.get("max_tokens", 512),
            presence_penalty=kw.get("presence_penalty", 0.0),
        )
        content = out["choices"][0]["message"]["content"]
        try:
            return json.loads(content)
        except json.JSONDecodeError:
            # max_tokens hit mid-object — the grammar can't close it, so we do.
            print(f"[llm] truncated JSON ({len(content)} chars), repairing")
            return json.loads(close_truncated_json(content))


# --------------------------------------------------------------------------- #
#  transformers fallback (for the hosted ZeroGPU Space, if llama.cpp won't build there)
# --------------------------------------------------------------------------- #
class TransformersLLM:
    def __init__(self) -> None:
        import torch  # noqa: PLC0415
        from transformers import AutoModelForCausalLM, AutoTokenizer  # noqa: PLC0415
        from transformers.utils import logging as hf_logging  # noqa: PLC0415

        hf_logging.set_verbosity_error()  # advisories bypass stdlib logging config
        print(f"[llm] Loading {config.LLM_REPO}…")
        self.tok = AutoTokenizer.from_pretrained(config.LLM_REPO)
        self.model = AutoModelForCausalLM.from_pretrained(
            config.LLM_REPO, torch_dtype="auto", device_map="auto"
        )
        print(f"[llm] {config.LLM_REPO} ready")
        self.torch = torch

    def _apply_template(self, messages: list[dict[str, str]], enable_thinking: bool = True) -> str:
        """Apply chat template, disabling Qwen3 thinking mode when asked."""
        try:
            return self.tok.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
                enable_thinking=enable_thinking,
            )
        except TypeError:
            # Tokenizer doesn't support enable_thinking (non-Qwen3 models)
            return self.tok.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )

    def complete(self, messages: list[dict[str, str]], **kw: Any) -> str:
        # Thinking off: the only free-text consumer is compact_memory, where a 256-token
        # budget would be eaten by the <think> block. strip_think() catches leftovers.
        text = self._apply_template(messages, enable_thinking=False)
        ids = self.tok(text, return_tensors="pt").to(self.model.device)
        out = self.model.generate(**ids, **_hf_gen_kwargs(kw))
        raw = self.tok.decode(out[0][ids.input_ids.shape[1] :], skip_special_tokens=True)
        return strip_think(raw)

    def complete_json(self, messages: list[dict[str, str]], schema: dict, **kw: Any) -> dict:
        import json  # noqa: PLC0415
        import re  # noqa: PLC0415

        defs = schema.get("$defs", {})

        def _resolve(s: dict) -> Any:
            if "$ref" in s:
                return _resolve(defs.get(s["$ref"].split("/")[-1], {}))
            if "anyOf" in s:
                for opt in s["anyOf"]:
                    if opt.get("type") != "null":
                        return _resolve(opt)
                return None
            if s.get("type") == "object" or "properties" in s:
                return {k: _resolve(v) for k, v in s.get("properties", {}).items()}
            if s.get("type") == "array":
                items = s.get("items", {})
                return [_resolve(items)] if items else []
            t = s.get("type", "string")
            if isinstance(t, list):
                t = next((x for x in t if x != "null"), "string")
            return {"string": "...", "integer": 0, "number": 0.0, "boolean": False}.get(t, None)

        skeleton = json.dumps(_resolve(schema), indent=2, ensure_ascii=False)
        hint = (
            "\n\nRespond with ONLY a valid JSON object. Required structure:\n"
            + skeleton
            + "\nNo markdown fences. No explanation. Only the JSON object."
        )
        augmented = [
            {**m, "content": m["content"] + hint} if m["role"] == "system" else m for m in messages
        ]
        for attempt in range(3):
            # Disable Qwen3 thinking mode: <think> blocks waste tokens and can truncate the JSON
            text = self._apply_template(augmented, enable_thinking=False)
            ids = self.tok(text, return_tensors="pt").to(self.model.device)
            out = self.model.generate(**ids, **_hf_gen_kwargs(kw))
            raw = self.tok.decode(out[0][ids.input_ids.shape[1] :], skip_special_tokens=True)

            # Prefer text AFTER </think>, fall back to text INSIDE <think>.
            think_end = raw.rfind("</think>")
            if think_end != -1:
                candidates = [
                    raw[think_end + len("</think>") :],
                    re.search(r"<think>(.*)</think>", raw, re.DOTALL).group(1),  # inside think
                ]
            else:
                candidates = [raw]

            parsed = False
            for candidate in candidates:
                candidate = re.sub(r"```(?:json)?\s*", "", candidate).strip()
                start, end = candidate.find("{"), candidate.rfind("}") + 1
                if start != -1 and end > start:
                    try:
                        result = json.loads(candidate[start:end])
                        parsed = True
                        return result
                    except json.JSONDecodeError as exc:
                        print(f"[llm] JSON parse error (attempt {attempt + 1}/3): {exc}")
            if not parsed:
                print(f"[llm] attempt {attempt + 1}/3: no JSON found. raw[:300]={raw[:300]!r}")
        print("[llm] WARNING: all 3 retries failed, using empty schema fallback")
        return _empty_from_schema(schema)


# --------------------------------------------------------------------------- #
#  Modal cloud GPU  (VN_LLM_BACKEND=modal)
#  Requires: pip install modal  +  modal deploy modal_app.py
# --------------------------------------------------------------------------- #
class ModalLLM:
    """Thin proxy that calls the deployed Modal LLM backend remotely."""

    def __init__(self) -> None:
        import modal  # noqa: PLC0415

        # Look up the already-deployed class — no local GPU needed.
        self._cls = modal.Cls.from_name("vn-app", "ModalLLMBackend")
        self._backend = self._cls()

    def warmup(self) -> None:
        """Fire-and-forget ping so the GPU container is warm before the first turn."""
        try:
            self._backend.complete.spawn([{"role": "user", "content": "hi"}], max_tokens=1)
        except Exception as exc:  # missing deployment must not kill server start
            _safe_print(f"[ModalLLM] warmup skipped: {exc}")

    def complete(self, messages: list[dict[str, str]], **kw: Any) -> str:
        result = self._backend.complete.remote(messages, **kw)
        if config.DEBUG:
            _safe_print(f"[ModalLLM] complete -> {result[:120]!r}")
        return result

    def complete_json(self, messages: list[dict[str, str]], schema: dict, **kw: Any) -> dict:
        if config.DEBUG:
            _safe_print(f"[ModalLLM] complete_json input messages:\n{messages}\n")
        result = self._backend.complete_json.remote(messages, schema, **kw)
        if config.DEBUG:
            _safe_print(f"[ModalLLM] complete_json -> {result}")
        return result


# --------------------------------------------------------------------------- #
#  Factory
# --------------------------------------------------------------------------- #
def get_llm() -> LLMBackend:
    if config.USE_MOCK or config.LLM_BACKEND == "mock":
        return MockLLM()
    if config.LLM_BACKEND == "modal":
        return ModalLLM()
    if config.LLM_BACKEND == "transformers":
        return TransformersLLM()
    return LlamaCppLLM()


def _empty_from_schema(schema: dict) -> dict:
    """Build a minimal object that satisfies a (simple) JSON schema's required fields."""
    if schema.get("type") == "object":
        obj: dict[str, Any] = {}
        props = schema.get("properties", {})
        for key in schema.get("required", []):
            obj[key] = _empty_from_schema(props.get(key, {}))
        return obj
    t = schema.get("type")
    if isinstance(t, list):
        t = next((x for x in t if x != "null"), "null")
    return {
        "string": "",
        "integer": 0,
        "number": 0,
        "boolean": False,
        "array": [],
        "object": {},
    }.get(t, None)