olivv-cs-ComfyUI-FunPack/tests/test_prompt_enhancer.py

"""The prompt enhancer: rewrite the prompt with an LLM BEFORE anything is generated from it.

Not the repair advisor. That one is triggered by a rating and fixes a named fault; this runs
on every generation and only elaborates, so it needs no refinement key, no rating and no
previous run. It sits after shortcut expansion and after `$variables` resolve, which is the
last moment the text is still text — so ComfyUI is fed the enhanced prompt and nothing
downstream has to know it happened.

The invariant everything here defends: a failure must return the ORIGINAL prompt. An enhancer
that empties the prompt on a bad generation would render a blank video and blame the model.
"""
import sys
import types
from pathlib import Path

import pytest

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import conditioning as C


@pytest.fixture
def studio():
    return C.FunPackVideoRefinerV2.__new__(C.FunPackVideoRefinerV2)


class FakeClip:
    """A CLIP that can generate, the way ComfyUI's own TextGenerate node expects."""

    def __init__(self, reply="a detailed prompt", fail=False):
        self.reply = reply
        self.fail = fail
        self.calls = []

    def tokenize(self, text, **kwargs):
        self.calls.append((text, kwargs))
        return {"input_ids": [[1, 2, 3]]}

    def generate(self, tokens, **kwargs):
        if self.fail:
            raise RuntimeError("no lm_head")
        self.last_generate = kwargs
        return [[4, 5, 6]]

    def decode(self, ids, skip_special_tokens=True):
        return self.reply


class EncoderOnlyClip:
    """A text encoder with no generation lane at all — the common wiring mistake."""

    def tokenize(self, text, **kwargs):
        return {}


# ── the failure modes, which all mean "leave the prompt alone" ───────────────

def test_a_clip_that_cannot_generate_leaves_the_prompt_untouched(studio):
    out, status = studio._v2_enhance_prompt(EncoderOnlyClip(), "a cat", "sys")
    assert out == "a cat"
    assert "does not expose text generation" in status


def test_a_clip_that_cannot_generate_says_where_to_wire_one(studio):
    _out, status = studio._v2_enhance_prompt(EncoderOnlyClip(), "a cat", "sys")
    assert "advisor_clip" in status


def test_a_generation_failure_returns_the_original(studio):
    out, status = studio._v2_enhance_prompt(FakeClip(fail=True), "a cat", "sys")
    assert out == "a cat"
    assert "failed" in status


def test_an_empty_generation_returns_the_original(studio):
    """The one that would render a blank video if it were allowed through."""
    out, status = studio._v2_enhance_prompt(FakeClip(reply="   "), "a cat", "sys")
    assert out == "a cat"
    assert "unchanged" in status


def test_an_empty_prompt_is_not_sent_to_the_model(studio):
    clip = FakeClip()
    out, status = studio._v2_enhance_prompt(clip, "   ", "sys")
    assert clip.calls == []
    assert "empty" in status


# ── the happy path ──────────────────────────────────────────────────────────

def test_the_enhanced_text_replaces_the_prompt(studio):
    out, status = studio._v2_enhance_prompt(FakeClip(reply="a ginger cat on a sunlit sill"),
                                            "a cat", "sys")
    assert out == "a ginger cat on a sunlit sill"
    assert "->" in status


def test_the_model_is_handed_the_prompt_it_will_replace(studio):
    clip = FakeClip()
    studio._v2_enhance_prompt(clip, "a cat", "my instructions")
    sent = clip.calls[0]
    assert "a cat" in str(sent)


def test_the_sampling_parameters_reach_the_model(studio):
    clip = FakeClip()
    studio._v2_enhance_prompt(clip, "a cat", "sys", temperature=1.3, top_p=0.5, max_length=128)
    assert clip.last_generate["temperature"] == pytest.approx(1.3)
    assert clip.last_generate["top_p"] == pytest.approx(0.5)
    assert clip.last_generate["max_length"] == 128


def test_a_short_length_limit_is_honoured_not_floored(studio):
    """The repair advisor floors at 128 tokens; a prompt enhancer must be allowed to be brief."""
    clip = FakeClip()
    studio._v2_enhance_prompt(clip, "a cat", "sys", max_length=48)
    assert clip.last_generate["max_length"] == 48


# ── the cache: one generation per distinct text, not per scene ───────────────

def test_the_same_text_is_only_generated_once(studio):
    """A multi-scene chain repeats the anchor line across scenes; each call is a full LLM run."""
    clip = FakeClip()
    cache = {}
    first, _ = studio._v2_enhance_prompt(clip, "a cat", "sys", cache=cache)
    second, status = studio._v2_enhance_prompt(clip, "a cat", "sys", cache=cache)
    assert first == second
    assert len(clip.calls) == 1
    assert "reused" in status


def test_different_instructions_are_not_shared(studio):
    clip = FakeClip()
    cache = {}
    studio._v2_enhance_prompt(clip, "a cat", "sys one", cache=cache)
    studio._v2_enhance_prompt(clip, "a cat", "sys two", cache=cache)
    assert len(clip.calls) == 2


def test_a_failed_generation_is_not_cached(studio):
    """Caching the fallback would make one transient failure stick for the whole run."""
    clip = FakeClip(reply="")
    cache = {}
    studio._v2_enhance_prompt(clip, "a cat", "sys", cache=cache)
    assert cache == {}


# ── cleaning what a chat model wraps around the answer ───────────────────────

def test_thinking_traces_are_stripped():
    out = C.FunPackVideoRefinerV2._v2_clean_enhanced_prompt(
        "<think>the user wants a cat</think>A ginger cat sits on a sill.")
    assert out == "A ginger cat sits on a sill."


def test_an_unterminated_thinking_trace_does_not_become_the_prompt():
    """A model that runs out of budget mid-reasoning must not have its reasoning rendered."""
    out = C.FunPackVideoRefinerV2._v2_clean_enhanced_prompt("A cat.<think>hmm, maybe I should")
    assert out == "A cat."


def test_a_leading_preamble_is_dropped():
    out = C.FunPackVideoRefinerV2._v2_clean_enhanced_prompt(
        "Here is the enhanced prompt: A ginger cat sits on a sill.")
    assert out == "A ginger cat sits on a sill."


def test_a_bare_output_label_is_dropped():
    assert C.FunPackVideoRefinerV2._v2_clean_enhanced_prompt("Output: A cat.") == "A cat."


def test_code_fences_are_dropped():
    out = C.FunPackVideoRefinerV2._v2_clean_enhanced_prompt("```\nA cat sits.\n```")
    assert out == "A cat sits."


def test_a_fully_quoted_answer_is_unwrapped():
    assert C.FunPackVideoRefinerV2._v2_clean_enhanced_prompt('"A cat sits."') == "A cat sits."


def test_quotes_inside_the_prompt_survive():
    """Speech is quoted on purpose — unwrapping must not eat it."""
    out = C.FunPackVideoRefinerV2._v2_clean_enhanced_prompt(
        'She says "hello" and waves.')
    assert out == 'She says "hello" and waves.'


# ── where it sits in refine_v2 ──────────────────────────────────────────────

def test_it_runs_before_the_prompt_is_encoded():
    """The whole point: ComfyUI must be fed the enhanced text, so the call has to come before
    _v2_conditioning_source, not after it and not somewhere downstream."""
    import inspect
    src = inspect.getsource(C.FunPackVideoRefinerV2.refine_v2)
    assert src.index("_v2_enhance_prompt(") < src.index("_v2_conditioning_source(")


def test_it_runs_after_variables_resolve():
    """`fully expanded` means shortcuts AND $variables — otherwise the model is handed a $name
    it cannot know the value of."""
    import inspect
    src = inspect.getsource(C.FunPackVideoRefinerV2.refine_v2)
    assert src.index("_resolve_variables(prompt_to_encode") < src.index("_v2_enhance_prompt(")


def test_multi_scene_enhances_each_scene_separately():
    """The editor's scene list is authoritative on COUNT. Rewriting the scenes as one
    paragraph and re-splitting could return a different number and desync every anchor."""
    import inspect
    src = inspect.getsource(C.FunPackVideoRefinerV2.refine_v2)
    assert "for t in split_scene_texts" in src


def test_it_is_off_by_default():
    import inspect
    sig = inspect.signature(C.FunPackVideoRefinerV2.refine_v2)
    assert sig.parameters["prompt_enhance"].default is False


def test_the_built_in_instructions_forbid_inventing_speech():
    """A video model will say whatever the prompt quotes, so an enhancer that invents dialogue
    puts words in a character's mouth the user never asked for."""
    text = C.V2_PROMPT_ENHANCER_SYSTEM_PROMPT.lower()
    assert "do not invent" in text and "speech" in text


# ── the model that will not stop ────────────────────────────────────────────
#
# Reported from a rental: with the cap raised to 8192 the enhancer used all of it. The H3
# encoder has no trained decoder head, so it may never emit a stop token — `max_length` stops
# being a ceiling and becomes the actual length. Nothing downstream would notice: it would
# just encode 8000 tokens of rambling as the prompt.

def _trim(text, **kw):
    return C.FunPackVideoRefinerV2._v2_trim_runaway_prompt(text, **kw)


def test_a_looping_model_is_cut_at_the_repeat(studio):
    """Repetition, not length, is what distinguishes a stuck model from a detailed one."""
    out, cut = _trim("A ginger cat sits on the sunlit sill. " * 40)
    assert cut is True
    assert out.count("A ginger cat sits on the sunlit sill.") < 3


def test_a_long_prompt_that_never_repeats_survives(studio):
    """H3 has no padding — every prompt token is a real sequence row, so a long prompt is a
    legitimate thing to want and must not be cut for being long."""
    text = " ".join(f"Shot {i}: a different thing happens in this room." for i in range(200))
    out, cut = _trim(text)
    assert cut is False and out == text


def test_length_alone_never_cuts_without_a_token_limit(studio):
    out, cut = _trim("Unique sentence number %d here. " % 0 + "".join(
        f"Sentence {i} describes something new. " for i in range(500)))
    assert cut is False


def test_the_length_bound_comes_from_the_users_token_limit(studio):
    """The limit the user set IS the control; no character cap is invented next to it."""
    text = "".join(f"Sentence {i} describes something new and distinct. " for i in range(500))
    out, cut = _trim(text, max_length=64)
    assert cut is True
    assert len(out) <= 64 * C.FunPackVideoRefinerV2.V2_ENHANCED_PROMPT_CHARS_PER_TOKEN


def test_a_generous_token_limit_keeps_a_generous_prompt(studio):
    text = "".join(f"Sentence {i} describes something new and distinct. " for i in range(200))
    _out, cut = _trim(text, max_length=8192)
    assert cut is False


def test_the_cut_lands_on_a_finished_sentence(studio):
    out, _cut = _trim("A ginger cat sits on the sunlit sill. " * 40)
    assert out.endswith(".")


def test_a_normal_prompt_is_not_cut(studio):
    text = "A ginger cat sits on a sunlit sill."
    out, cut = _trim(text, max_length=400)
    assert out == text and cut is False


def test_a_short_recurring_line_is_not_treated_as_a_loop(studio):
    """"She smiles." legitimately recurs in a long shot list."""
    text = " ".join(f"The {w} moves across the room. She smiles." for w in
                    ("cat", "dog", "bird", "fox", "owl"))
    _out, cut = _trim(text)
    assert cut is False


def test_a_runaway_never_ends_mid_word(studio):
    out, _cut = _trim("supercalifragilistic " * 500, max_length=64)
    assert not out.endswith("supercalifragilis")
    assert out == out.strip()


def test_the_user_is_told_the_model_repeated_itself(studio):
    clip = FakeClip(reply="A ginger cat sits on the sunlit sill. " * 40)
    _out, status = studio._v2_enhance_prompt(clip, "a cat", "sys")
    assert "repeated itself" in status


# ── where it ran ────────────────────────────────────────────────────────────

def test_the_status_says_which_device_generated(studio):
    """A 32B encoder sharing a card with the DiT can execute on the CPU — minutes per call,
    and nothing in the returned text says so."""
    clip = FakeClip()
    clip.patcher = types.SimpleNamespace(load_device="cuda:0")
    _out, status = studio._v2_enhance_prompt(clip, "a cat", "sys")
    assert "cuda:0" in status


def test_a_wrapper_device_is_found_too(studio):
    clip = FakeClip()
    clip._device = "cpu"
    _out, status = studio._v2_enhance_prompt(clip, "a cat", "sys")
    assert "cpu" in status


def test_an_unreported_device_says_so_rather_than_guessing(studio):
    _out, status = studio._v2_enhance_prompt(FakeClip(), "a cat", "sys")
    assert "unreported" in status


def test_the_status_reports_the_cap_that_was_applied(studio):
    _out, status = studio._v2_enhance_prompt(FakeClip(), "a cat", "sys", max_length=256)
    assert "256" in status


# ── one generation per prompt that is actually encoded ──────────────────────
#
# Reported from a rental: the enhancer ran twice. The Movie Editor's builder forces
# split_by_transitions on for EVERY project (builder.py), and the per-scene conditionings
# replace the base entry — so enhancing the base as well was a second full generation whose
# result was then thrown away. At ~59s per call that is a minute of pure waste per run.

def _refine_source():
    import inspect
    return inspect.getsource(C.FunPackVideoRefinerV2.refine_v2)


def test_the_base_prompt_is_not_enhanced_when_the_run_splits():
    src = _refine_source()
    assert "_enhance_base = _enhance_own and not split_by_transitions" in src


def test_the_base_enhancement_is_gated_on_that_flag():
    """Gating the print but not the call would keep paying for the generation."""
    src = _refine_source()
    assert "if _enhance_base:" in src


def test_the_scene_path_is_still_enhanced_when_splitting():
    src = _refine_source()
    assert "if _enhance_own and split_scene_texts:" in src


def test_the_skip_is_reported_not_silent():
    """A user watching the console must see why the base was not enhanced."""
    src = _refine_source()
    assert "base prompt skipped" in src


def test_thinking_is_off_unless_asked_for():
    """The double run looked like thinking mode; it was not. Thinking defaults off and is
    passed explicitly, so it can never be inherited from the advisor's own default of True."""
    import inspect
    sig = inspect.signature(C.FunPackVideoRefinerV2.refine_v2)
    assert sig.parameters["prompt_enhance_thinking"].default is False
    assert "thinking=prompt_enhance_thinking" in _refine_source()


def test_the_enhancer_never_inherits_the_advisors_thinking_default(studio):
    clip = FakeClip()
    studio._v2_enhance_prompt(clip, "a cat", "sys")
    assert clip.calls[0][1].get("thinking") is False


# ── the Studio image reaches the enhancer on BOTH paths ─────────────────────
#
# The Editor forces the per-scene path, which used to call the enhancer with no image at all —
# so only a single-prompt run could ever show it the picture.

def test_the_image_reaches_the_enhancer():
    clip = FakeClip()
    image = object()
    studio = C.FunPackVideoRefinerV2.__new__(C.FunPackVideoRefinerV2)
    studio._v2_enhance_prompt(clip, "a cat", "sys", image=image)
    assert clip.calls[0][1].get("image") is image


def test_both_paths_hand_over_the_enhancer_image():
    src = _refine_source()
    assert src.count("image=prompt_enhance_image") == 3
    assert "thinking=prompt_enhance_thinking, image=source_image" not in src


def test_studio_gates_the_image_on_its_own_switch():
    """Not on vision_conditioning: showing the enhancer the picture is a separate choice."""
    import inspect
    src = inspect.getsource(C.FunPackStudio.run)
    assert ('prompt_enhance_image=source_image if rf.get("prompt_enhance_use_image", True) '
            'else None') in src


def test_every_rewrite_is_reported_for_the_composer():
    src = _refine_source()
    assert src.count("self._v2_note_enhanced(") == 3
    import inspect
    assert '{"ui": {"funpack_enhanced": enhanced}, "result": result}' in inspect.getsource(C.FunPackStudio.run)


# ── every sampling knob in the Composer reaches clip.generate ───────────────

def test_the_sampling_settings_reach_the_generator(studio):
    clip = FakeClip()
    studio._v2_enhance_prompt(clip, "a cat", "sys", seed=7, temperature=0.4, top_p=0.8,
                              top_k=12, min_p=0.1, repetition_penalty=1.1,
                              presence_penalty=0.5, do_sample=False)
    g = clip.last_generate
    assert (g["temperature"], g["top_p"], g["top_k"], g["min_p"]) == (0.4, 0.8, 12, 0.1)
    assert (g["repetition_penalty"], g["presence_penalty"], g["do_sample"], g["seed"]) == (1.1, 0.5, False, 7)


def test_the_repair_advisor_keeps_its_old_sampling(studio):
    """New knobs are the enhancer's; the advisor's defaults must not move."""
    clip = FakeClip()
    studio._v2_generate_advisor_text(clip, "sys", "a cat")
    g = clip.last_generate
    assert (g["min_p"], g["presence_penalty"], g["do_sample"], g["top_k"]) == (0.05, 0.0, True, 50)


def test_a_fixed_enhancer_seed_overrides_the_run_seed():
    src = _refine_source()
    assert '_enhance_seed = int(_enhance_sampling.pop("seed", 0) or 0) or seed' in src
    assert src.count("seed=_enhance_seed") == 3


# ── how the instructions reach the model ─────────────────────────────────────

def test_instructions_ride_in_the_models_own_template(studio):
    # ComfyUI's Gemma tokenizer swallows unknown kwargs and defaults to skip_template=True,
    # so a system_prompt kwarg vanished and the raw prompt went in with no chat turn and no
    # image marker. The model's own template must be asked for, with the instructions in it.
    clip = FakeClip()
    studio._v2_enhance_prompt(clip, "a cat", "BE VIVID", image="img")
    text, kwargs = clip.calls[0]
    assert text.startswith("BE VIVID") and text.endswith("a cat")
    assert kwargs["skip_template"] is False and kwargs["min_length"] == 1
    assert kwargs["image"] == "img"
    assert "system_prompt" not in kwargs


def test_the_funpack_advisor_llm_still_gets_a_real_system_turn(studio):
    wrapper = C._FunPackAdvisorLLMWrapper.__new__(C._FunPackAdvisorLLMWrapper)
    seen = {}
    wrapper.tokenize = lambda text, **kw: seen.update(text=text, **kw) or {}
    wrapper.generate = lambda tokens, **kw: [1]
    wrapper.decode = lambda ids, skip_special_tokens=True: "a vivid cat"
    studio._v2_enhance_prompt(wrapper, "a cat", "BE VIVID")
    assert seen["text"] == "a cat" and seen["system_prompt"] == "BE VIVID"


# ── reference: shortcuts + whole lorebooks, appended after the prompt ────────

def _sources(tmp_path, monkeypatch):
    import json as _json
    import templates
    monkeypatch.setattr(templates, "load_shortcut_db", lambda: {"shortcuts": {
        "k1": {"name": "Mira", "replacements": ["red coat, silver hair"], "triggers": ["/mira"]},
        "k2": {"name": "Rain", "replacements": ["heavy rain"], "triggers": ["/rain"]}}})
    lb = tmp_path / "lore.json"
    lb.write_text(_json.dumps({"entries": {
        "0": {"comment": "Castle", "content": "black stone towers", "keys": ["castle"], "order": 2},
        "1": {"comment": "Sea", "content": "grey winter sea", "keys": ["sea"], "order": 1},
        "2": {"comment": "Tone", "content": "melancholy", "keys": [], "constant": True, "order": 3}}}))
    return C.FunPackVideoRefinerV2._v2_enhancer_sources(
        ["k1", "k2", "gone"], [str(lb), str(tmp_path / "nope.json")])


def test_reference_sends_only_what_the_prompt_mentions(tmp_path, monkeypatch):
    ref = C.FunPackVideoRefinerV2._v2_enhancer_reference(
        "Mira walks to the castle", _sources(tmp_path, monkeypatch))
    assert "[Shortcut] Mira\nred coat, silver hair" in ref and "heavy rain" not in ref
    assert "black stone towers" in ref and "grey winter sea" not in ref
    assert "melancholy" in ref                      # constant joins a match


def test_expanded_shortcut_content_counts_as_a_mention(tmp_path, monkeypatch):
    ref = C.FunPackVideoRefinerV2._v2_enhancer_reference(
        "a woman in a red coat, silver hair", _sources(tmp_path, monkeypatch))
    assert "[Shortcut] Mira" in ref and "heavy rain" not in ref


def test_no_mention_sends_the_whole_source(tmp_path, monkeypatch):
    ref = C.FunPackVideoRefinerV2._v2_enhancer_reference("a dog", _sources(tmp_path, monkeypatch))
    assert "[Shortcut] Mira" in ref and "[Shortcut] Rain" in ref
    assert ref.index("grey winter sea") < ref.index("black stone towers") < ref.index("melancholy")


def test_reference_goes_after_the_prompt_and_not_into_the_readout(studio):
    clip = FakeClip("a vivid cat")
    out, _ = studio._v2_enhance_prompt(clip, "a cat", "SYS", reference="\n\nREF")
    assert clip.calls[0][0].endswith("a cat\n\nREF") and out == "a vivid cat"
    assert studio._v2_enhance_prompt(FakeClip(fail=True), "a cat", "SYS", reference="\n\nREF")[0] == "a cat"


def test_reference_is_empty_when_nothing_is_picked():
    assert C.FunPackVideoRefinerV2._v2_enhancer_reference("a cat", []) == ""


def test_a_shortcuts_file_is_a_source_of_its_own(tmp_path, monkeypatch):
    import json as _json
    import templates
    pack = tmp_path / "pack.json"
    pack.write_text(_json.dumps({"shortcuts": {
        "a": {"name": "Neon", "triggers": ["/neon"], "replacements": ["pink neon glow"]},
        "b": {"name": "Fog", "triggers": ["/fog"], "replacements": ["low fog"], "enabled": False}}}))
    monkeypatch.setattr(templates, "load_shortcut_db", lambda: {"shortcuts": {}})
    src = C.FunPackVideoRefinerV2._v2_enhancer_sources([], [str(pack)])
    ref = C.FunPackVideoRefinerV2._v2_enhancer_reference("a street", src)
    assert "[Shortcut] Neon\npink neon glow" in ref and "low fog" not in ref   # disabled stays out


def test_shortcut_content_with_a_variable_is_recognised(monkeypatch):
    import templates
    monkeypatch.setattr(templates, "load_shortcut_db", lambda: {"shortcuts": {
        "o": {"name": "Outfit", "replacements": ["wearing $dress and boots"], "triggers": ["/outfit"]},
        "r": {"name": "Rain", "replacements": ["heavy rain"], "triggers": ["/rain"]}}})
    src = C.FunPackVideoRefinerV2._v2_enhancer_sources(["o", "r"], [], [{"name": "dress", "value": "a green dress"}])
    ref = C.FunPackVideoRefinerV2._v2_enhancer_reference("she is wearing a green dress and boots", src)
    assert "[Shortcut] Outfit\nwearing a green dress and boots" in ref and "heavy rain" not in ref


def test_names_and_keywords_match_whole_words_only(tmp_path, monkeypatch):
    import json as _json
    import templates
    monkeypatch.setattr(templates, "load_shortcut_db", lambda: {"shortcuts": {
        "r": {"name": "Rain", "replacements": ["heavy rain"], "triggers": ["/rain"]},
        "m": {"name": "Mira", "replacements": ["red coat"], "triggers": ["/mira"]}}})
    lb = tmp_path / "lore.json"
    lb.write_text(_json.dumps({"entries": [
        {"comment": "Sea", "content": "grey sea", "keys": ["sea"]},
        {"comment": "Port", "content": "old port", "keys": ["port"], "matchWholeWords": False},
        {"comment": "Tower", "content": "black tower", "keys": ["/tow(er|ers)/"]}]}))
    src = C.FunPackVideoRefinerV2._v2_enhancer_sources(["r", "m"], [str(lb)])
    ref = C.FunPackVideoRefinerV2._v2_enhancer_reference("Mira takes the train in autumn season to the airport towers", src)
    assert "[Shortcut] Mira" in ref and "heavy rain" not in ref        # "train" is not "rain"
    assert "grey sea" not in ref                                      # "season" is not "sea"
    assert "old port" in ref                                          # entry opted out of whole words
    assert "black tower" in ref                                       # regex key


def test_a_word_from_inside_a_shortcut_finds_it_even_when_another_matches(tmp_path):
    import json as _json
    pack = tmp_path / "library.json"
    pack.write_text(_json.dumps({"shortcuts": {
        "a": {"name": "Mira", "triggers": ["/mira"], "replacements": ["woman in a crimson trench coat"]},
        "b": {"name": "Night", "triggers": ["/night"], "replacements": ["moonlit street"]},
        "c": {"name": "Rain", "triggers": ["/rain"], "replacements": ["heavy rain"]}}}))
    src = C.FunPackVideoRefinerV2._v2_enhancer_sources([], [str(pack)])
    ref = C.FunPackVideoRefinerV2._v2_enhancer_reference("a figure in a trench coat at night", src)
    assert "[Shortcut] Mira" in ref and "[Shortcut] Night" in ref      # "trench" + name "night"
    assert "heavy rain" not in ref
    # filler words never count: "with the" alone mentions nothing, so the whole file goes
    assert "heavy rain" in C.FunPackVideoRefinerV2._v2_enhancer_reference("with the", src)


# ── the model's thinking reaches the readout, never the prompt ───────────────

def test_thinking_is_kept_for_the_readout_and_out_of_the_prompt(studio):
    clip = FakeClip("<think>She wants mood, so add rain.</think>A cat in the rain.")
    out, _ = studio._v2_enhance_prompt(clip, "a cat", "SYS", thinking=True)
    assert out == "A cat in the rain."
    assert studio._v2_last_thinking == "She wants mood, so add rain."


def test_thinking_opened_by_the_template_is_still_found(studio):
    clip = FakeClip("Mood first.\n</think>\nA cat.")
    out, _ = studio._v2_enhance_prompt(clip, "a cat", "SYS", thinking=True)
    assert out == "A cat." and studio._v2_last_thinking == "Mood first."


def test_no_thinking_means_an_empty_readout(studio):
    studio._v2_enhance_prompt(FakeClip("A cat."), "a cat", "SYS")
    assert studio._v2_last_thinking == ""


def test_the_advisor_llm_wrapper_hands_over_its_stripped_reasoning():
    w = C._FunPackAdvisorLLMWrapper.__new__(C._FunPackAdvisorLLMWrapper)
    w._tokenizer = types.SimpleNamespace(decode=lambda ids, skip_special_tokens=False:
                                         "<think>plan it</think>A dog.<|im_end|>")
    text = w.decode({"output_ids": [[1, 2, 3]], "prompt_length": 0})
    assert text == "A dog." and w.last_thinking == "plan it"


# ── chat: the next run reads its earlier rewrite and the user's comments ─────

def test_no_chat_adds_nothing():
    assert C.FunPackVideoRefinerV2._v2_enhancer_chat([]) == ""
    assert C.FunPackVideoRefinerV2._v2_enhancer_chat([{"rewrites": {"whole": "x"}, "comment": "  "}]) == ""


def test_the_chat_shows_the_latest_rewrite_and_every_comment_in_order():
    chat = [{"rewrites": {"whole": "a red fox runs"}, "comment": "make it night"},
            {"rewrites": {}, "comment": "and add snow"},
            {"rewrites": {"whole": "a red fox runs at night"}, "comment": "colder light"}]
    text = C.FunPackVideoRefinerV2._v2_enhancer_chat(chat)
    assert "a red fox runs at night" in text and "a red fox runs\n" not in text
    order = ["make it night", "and add snow", "colder light"]
    assert [text.index(p) for p in order] == sorted(text.index(p) for p in order)
    assert text.index("YOUR LATEST REWRITE") < text.index("USER FEEDBACK")


def test_a_chat_run_labels_the_parts_and_tells_the_model_it_is_revising(studio):
    clip = FakeClip()
    chat = C.FunPackVideoRefinerV2._v2_enhancer_chat(
        [{"rewrites": {"whole": "a fox at dusk"}, "comment": "make it night"}])
    studio._v2_enhance_prompt(clip, "a fox", "my rules", chat=chat)
    sent = str(clip.calls[0])
    assert "ORIGINAL PROMPT:\\na fox" in sent or "ORIGINAL PROMPT:\na fox" in sent
    assert "REVISION" in sent and "my rules" in sent and "make it night" in sent


def test_no_chat_leaves_the_message_as_it_was(studio):
    clip = FakeClip()
    studio._v2_enhance_prompt(clip, "a fox", "my rules")
    assert "ORIGINAL PROMPT" not in str(clip.calls[0]) and "REVISION" not in str(clip.calls[0])


def test_an_echoed_conversation_is_cut_from_the_answer(studio):
    """What a model did live: repeated its rewrite, then the feedback, as one prompt."""
    echo = ("A fox pads through a moonlit field. USER FEEDBACK (oldest first): - make it night. "
            "Generate a new prompt.")
    out, _ = studio._v2_enhance_prompt(FakeClip(reply=echo), "a fox", "sys",
                                       chat="\n\nUSER FEEDBACK (oldest first):\n- make it night")
    assert out == "A fox pads through a moonlit field."
    labelled = "YOUR LATEST REWRITE: A fox at night, snow falling."
    out, _ = studio._v2_enhance_prompt(FakeClip(reply=labelled), "a fox", "sys",
                                       chat="\n\nUSER FEEDBACK (oldest first):\n- snow")
    assert out == "A fox at night, snow falling."


def test_each_scene_sees_its_own_earlier_rewrite():
    chat = [{"rewrites": {"0": "scene one text", "1": "scene two text"}, "comment": "brighter"}]
    second = C.FunPackVideoRefinerV2._v2_enhancer_chat(chat, 1)
    assert "scene two text" in second and "scene one text" not in second
    # One rewrite on record (a single-prompt run) serves whatever asks.
    assert "only" in C.FunPackVideoRefinerV2._v2_enhancer_chat(
        [{"rewrites": {"whole": "only"}, "comment": "c"}], 3)


def test_the_chat_reaches_the_model_and_a_new_comment_regenerates(studio):
    clip = FakeClip()
    cache = {}
    studio._v2_enhance_prompt(clip, "a cat", "sys", cache=cache)
    studio._v2_enhance_prompt(clip, "a cat", "sys", cache=cache, chat="\n\nUSER FEEDBACK (oldest first):\n- bigger")
    assert len(clip.calls) == 2, "a comment must not be answered from the no-chat cache"
    assert "bigger" in str(clip.calls[1])


def test_every_enhancement_path_passes_the_chat():
    import inspect
    src = inspect.getsource(C.FunPackVideoRefinerV2.refine_v2)
    assert src.count("chat=self._v2_enhancer_chat(prompt_enhance_chat") == 3


def test_every_readout_says_what_chat_the_model_was_sent():
    """Last run ▸ "Chat sent" is how a user checks what the model actually read."""
    import inspect
    src = inspect.getsource(C.FunPackVideoRefinerV2.refine_v2)
    assert src.count('"chat": ') == 3


def test_a_chat_run_records_the_whole_message_for_the_readout(studio):
    chat = C.FunPackVideoRefinerV2._v2_enhancer_chat([{"rewrites": {"whole": "a fox at dusk"},
                                                        "comment": "make it night"}])
    cache = {}
    studio._v2_enhance_prompt(FakeClip(), "a fox", "my rules", chat=chat, cache=cache)
    sent = studio._v2_last_sent
    for part in ("SYSTEM:", "my rules", "REVISION", "MESSAGE:", "ORIGINAL PROMPT:", "a fox at dusk", "make it night"):
        assert part in sent, part
    studio._v2_enhance_prompt(FakeClip(), "a fox", "my rules", chat=chat, cache=cache)
    assert studio._v2_last_sent == sent                    # a reused answer reports the same
    studio._v2_enhance_prompt(FakeClip(), "a fox", "my rules")
    assert studio._v2_last_sent == ""


def test_reference_intro_is_editable_and_blank_means_built_in(tmp_path, monkeypatch):
    src = _sources(tmp_path, monkeypatch)
    R = C.FunPackVideoRefinerV2._v2_enhancer_reference
    assert R("a dog", src).startswith("\n\n" + C.V2_PROMPT_ENHANCER_REFERENCE_INTRO + "\n\n")
    assert R("a dog", src, "  ").startswith("\n\n" + C.V2_PROMPT_ENHANCER_REFERENCE_INTRO)
    mine = R("a dog", src, "Notes you may use:")
    assert mine.startswith("\n\nNotes you may use:\n\n[") and "Reference entries" not in mine