tests/test_llm_helpers.py

"""LLM helper logic that must work without any model: <think> stripping, the
transformers sampling-kwargs mapping, and compaction's defence against thinking
text polluting the rolling summary."""

from __future__ import annotations

import json

from visualnovel import config, orchestrator
from visualnovel.llm import _hf_gen_kwargs
from visualnovel.schemas import Character, GameState, Scene, Turn
from visualnovel.utils import close_truncated_json, strip_think


def test_strip_think_closed_block():
    assert strip_think("<think>reasoning here</think>The summary.") == "The summary."


def test_strip_think_unclosed_block():
    assert strip_think("Partial text <think>truncated reasoning") == "Partial text"


def test_strip_think_passthrough():
    assert strip_think("Plain prose, no tags.") == "Plain prose, no tags."


def test_strip_think_multiple_blocks():
    assert strip_think("<think>a</think>One. <think>b</think>Two.") == "One. Two."


def test_close_truncated_json_mid_string():
    repaired = close_truncated_json('{"speaker": "moth", "dialogue": "Hello the')
    assert json.loads(repaired)["dialogue"] == "Hello the"


def test_close_truncated_json_after_colon_and_comma():
    assert json.loads(close_truncated_json('{"a": "x", "b":')) == {"a": "x", "b": None}
    assert json.loads(close_truncated_json('{"a": "x",')) == {"a": "x"}


def test_close_truncated_json_nested():
    repaired = close_truncated_json('{"d": {"new_character": {"id": "kaaris", "traits": ["bo')
    parsed = json.loads(repaired)
    assert parsed["d"]["new_character"]["traits"] == ["bo"]


def test_close_truncated_json_dangling_escape():
    repaired = close_truncated_json('{"a": "say \\"hi\\" no\\')
    assert json.loads(repaired)["a"] == 'say "hi" no'


def test_close_truncated_json_valid_passthrough():
    assert json.loads(close_truncated_json('{"a": 1}')) == {"a": 1}


def test_hf_gen_kwargs_sampling():
    kw = _hf_gen_kwargs({"temperature": 0.7, "top_p": 0.9, "max_tokens": 1024})
    assert kw == {"max_new_tokens": 1024, "do_sample": True, "temperature": 0.7, "top_p": 0.9}


def test_hf_gen_kwargs_greedy():
    assert _hf_gen_kwargs({}) == {"max_new_tokens": 512, "do_sample": False}
    assert _hf_gen_kwargs({"temperature": 0}) == {"max_new_tokens": 512, "do_sample": False}


def _state(n_turns: int) -> GameState:
    s = GameState(
        seed=1,
        vibe="v",
        style_guide="s",
        scene=Scene(id="a", place="A", description="d", present=["moth"]),
        characters={"moth": Character(id="moth", name="Moth", one_line="a moth")},
        summary="The tale began.",
    )
    for i in range(n_turns):
        s.recent_turns.append(Turn(player=f"p{i}", speaker="moth", dialogue=f"d{i}"))
    return s


class _StubLLM:
    def __init__(self, reply: str) -> None:
        self.reply = reply

    def complete(self, messages, **kw):
        return self.reply

    def complete_json(self, messages, schema, **kw):  # pragma: no cover - unused
        return {}


def test_compact_memory_strips_think(monkeypatch):
    monkeypatch.setattr(config, "USE_MOCK", False)
    state = _state(2 * config.RECENT_TURNS_K + 1)
    orchestrator.compact_memory(_StubLLM("<think>blah blah</think>Real summary."), state)
    assert state.summary == "Real summary."
    assert len(state.recent_turns) == config.RECENT_TURNS_K


def test_compact_memory_all_thinking_falls_back(monkeypatch):
    monkeypatch.setattr(config, "USE_MOCK", False)
    state = _state(2 * config.RECENT_TURNS_K + 1)
    orchestrator.compact_memory(_StubLLM("<think>only truncated thinking"), state)
    assert state.summary  # never empty
    assert "<think>" not in state.summary