visualnovel/orchestrator.py

"""The Weaver — director / game master.

Three jobs:
  - init_world(setup)          -> (GameState, opening DirectorOutput)
  - direct_turn(state, input)  -> DirectorOutput   (one grammar-constrained call per turn)
  - compact_memory(state)      -> None             (fold old turns into the rolling summary)

In MOCK mode (config.USE_MOCK) these build nice, deterministic, themed output WITHOUT a model,
so the whole loop runs offline. Flip VN_MOCK=0 to use the real LLM via `llm.py`.
"""

from __future__ import annotations

import random

from . import config, memory, prompts
from .llm import LLMBackend
from .schemas import (
    Character,
    Directives,
    DirectorOutput,
    GameState,
    InitOutput,
    NewCharacter,
    Scene,
    SceneChange,
    SetupForm,
)
from .utils import clip_words, strip_think

# crude keyword cue for the mock director to trigger a move
_MOVE_WORDS = (
    "go",
    "leave",
    "enter",
    "into",
    "toward",
    "outside",
    "deeper",
    "follow",
    "walk",
    "open",
)


# --------------------------------------------------------------------------- #
#  init_world
# --------------------------------------------------------------------------- #
def init_world(llm: LLMBackend, setup: SetupForm) -> tuple[GameState, DirectorOutput]:
    seed = setup.seed if setup.seed is not None else random.randint(1, 2**31 - 1)
    vibe = prompts.build_vibe(setup)

    if config.USE_MOCK:
        init = _mock_init(setup)
    else:
        raw = llm.complete_json(
            messages=[
                {"role": "system", "content": prompts.INIT_SYSTEM},
                {"role": "user", "content": prompts.init_user_prompt(setup)},
            ],
            schema=prompts.init_schema(),
            temperature=0.9,
            top_p=0.95,
            max_tokens=1024,
        )
        raw = _repair_init(raw)
        init = InitOutput.model_validate(raw)

    first = init.first_character
    first.id = _slug(first.id) or _slug(first.name)  # sanitize LLM-supplied id (Windows-safe)
    scene = Scene(
        id=_slug(init.scene.place),
        place=init.scene.place,
        description=init.scene.description,
        mood=init.scene.mood,
        present=[first.id],
        backdrop_seed=_seed_for(seed, init.scene.place),
    )
    flags: dict[str, str] = {}
    if init.situation_intro:
        flags["situation_intro"] = init.situation_intro
    if setup.player_name and setup.player_name != "the wanderer":
        flags["player_name"] = setup.player_name

    state = GameState(
        seed=seed,
        vibe=vibe,
        style_guide=init.style_guide,
        scene=scene,
        characters={
            first.id: Character(
                id=first.id,
                name=first.name,
                one_line=first.one_line,
                appearance=first.appearance,
                voice=first.voice,
                traits=first.traits,
                goals=first.goals,
                sprite_seed=_seed_for(seed, first.id),
                tts_voice_description=first.voice,  # frozen at creation
            )
        },
        flags=flags,
    )
    opening = DirectorOutput(speaker=first.id, dialogue=init.opening_line, emotion="tender")
    return state, opening


# --------------------------------------------------------------------------- #
#  direct_turn
# --------------------------------------------------------------------------- #
def direct_turn(
    llm: LLMBackend,
    state: GameState,
    player_input: str,
    action: str = "talk",
    target: str = "",
) -> DirectorOutput:
    if config.USE_MOCK:
        return _mock_turn(state, player_input)

    # "look around" with room on stage uses a schema where new_character is REQUIRED:
    # the introduction is enforced by the grammar (prompting alone proved unreliable).
    if action == "scout" and len(state.scene.present) >= config.MAX_PRESENT:
        action = "scout_full"  # stage full: milder note, normal schema
    schema = prompts.scout_schema() if action == "scout" else prompts.directive_schema()

    note = prompts.action_note(action, target)
    context = memory.assemble_context(state, player_input, action_note=note)
    try:
        raw = llm.complete_json(
            messages=[
                {"role": "system", "content": prompts.DIRECTOR_SYSTEM},
                {"role": "user", "content": context},
            ],
            schema=schema,
            temperature=0.7,
            top_p=0.9,
            max_tokens=1536,
        )
        out = _repair(state, DirectorOutput.model_validate(raw))
    except Exception as exc:
        # A bad sample must cost the player one bland line, never a crashed turn.
        print(f"[director] turn failed ({type(exc).__name__}: {exc}) - using fallback line")
        speaker = state.scene.present[0] if state.scene.present else "narrator"
        return DirectorOutput(
            speaker=speaker,
            dialogue="Hm? Sorry - I lost my train of thought for a moment. What were you saying?",
            emotion="thoughtful",
        )

    # Anti-repetition guard: the prompt forbids verbatim repeats, but small models still
    # produce them. One hotter retry quoting the forbidden line; accept whatever comes back.
    if _is_repeat(out.dialogue, state):
        try:
            raw = llm.complete_json(
                messages=[
                    {"role": "system", "content": prompts.DIRECTOR_SYSTEM},
                    {
                        "role": "user",
                        "content": context
                        + f'\n\nNOTE: You already said this line: "{out.dialogue}"\n'
                        "Repeating it is FORBIDDEN. Do not reuse any sentence from "
                        "RECENT EXCHANGE - react freshly to the player's latest input "
                        "and to whatever just changed on stage.",
                    },
                ],
                schema=schema,
                temperature=0.9,
                top_p=0.9,
                max_tokens=1536,
                presence_penalty=0.8,  # llama.cpp honours it; other backends ignore it
            )
            out = _repair(state, DirectorOutput.model_validate(raw))
        except Exception:
            pass  # a failed retry must never kill the turn — keep the first output
    return out


# --------------------------------------------------------------------------- #
#  compact_memory
# --------------------------------------------------------------------------- #
def compact_memory(llm: LLMBackend, state: GameState) -> None:
    keep = config.RECENT_TURNS_K
    old = state.recent_turns[:-keep] if keep else state.recent_turns
    if not old:
        return
    if config.USE_MOCK:
        new_summary = (state.summary + " " + " ".join(t.dialogue for t in old)).strip()
        state.summary = new_summary[-800:]
    else:
        who = state.flags.get("player_name", "wanderer")
        recent_blob = "\n".join(f'{who}: "{t.player}"\n{t.speaker}: "{t.dialogue}"' for t in old)
        # Trailing /no_think = Qwen3 soft switch; without it the GGUF chat template enables
        # thinking and the whole token budget is eaten by an unclosed <think> block.
        raw = llm.complete(
            messages=[
                {"role": "system", "content": prompts.COMPACT_SYSTEM},
                {
                    "role": "user",
                    "content": f"SUMMARY SO FAR:\n{state.summary}\n\nRECENT:\n{recent_blob}\n/no_think",
                },
            ],
            temperature=0.3,
            top_p=0.9,
            max_tokens=320,
        )
        # Qwen3 (GGUF and transformers) may wrap or fill the reply with <think> blocks;
        # never let thinking text — or an empty husk of it — become the rolling summary.
        new_summary = strip_think(raw)
        if not new_summary:
            new_summary = (state.summary + " " + " ".join(t.dialogue for t in old)).strip()[-800:]
        state.summary = new_summary
    state.recent_turns = state.recent_turns[-keep:] if keep else []


# --------------------------------------------------------------------------- #
#  Validation / repair of a real LLM output
# --------------------------------------------------------------------------- #
def _repair_init(raw: dict) -> dict:
    """Coerce a loosely-shaped InitOutput dict into the expected structure.

    Models without grammar constraints sometimes return `scene` as a plain string
    and omit required fields like `one_line`/`appearance` in `first_character`.
    """
    scene = raw.get("scene")
    if isinstance(scene, str):
        place = scene.split(",")[0].split(".")[0][:40].strip() or "School Grounds"
        raw["scene"] = {"place": place, "description": scene, "mood": "neutral"}
    elif isinstance(scene, dict):
        if "place" not in scene:
            scene["place"] = scene.get("name") or scene.get("location") or "School Grounds"
        if "description" not in scene:
            scene["description"] = scene.get("setting") or scene.get("place") or "A lovely setting."

    fc = raw.get("first_character")
    if isinstance(fc, dict):
        if not fc.get("one_line"):
            traits = fc.get("traits", [])
            name = fc.get("name", "stranger")
            fc["one_line"] = (
                fc.pop("bio", None)
                or fc.pop("description", None)
                or (f"a {traits[0].lower()} person" if traits else f"a mysterious {name}")
            )
        if not fc.get("appearance"):
            fc["appearance"] = (
                fc.pop("physical_description", None)
                or fc.pop("looks", None)
                or "a young person with distinctive anime-style features"
            )
    return raw


def _norm(s: str) -> str:
    """Normalize dialogue for repetition comparison (case/punctuation/whitespace)."""
    import re  # noqa: PLC0415

    return re.sub(r"[\W_]+", " ", s.lower()).strip()


def _is_repeat(dialogue: str, state: GameState) -> bool:
    """True when the dialogue is a (near-)verbatim copy of one of the last 3 lines."""
    import difflib  # noqa: PLC0415

    cand = _norm(dialogue)
    if not cand:
        return False
    for t in state.recent_turns[-3:]:
        prev = _norm(t.dialogue)
        if cand == prev or difflib.SequenceMatcher(None, cand, prev).ratio() >= 0.95:
            return True
    return False


def _repair(state: GameState, out: DirectorOutput) -> DirectorOutput:
    # A newcomer introduced THIS turn may speak even though they are not yet on stage
    # (apply_directives adds them right after) — e.g. the "look around" action.
    nc_arriving = out.directives.new_character
    nc_id = (_slug(nc_arriving.id) or _slug(nc_arriving.name)) if nc_arriving else None
    # speaker must be present, the arriving newcomer, or narrator; else first on stage
    if out.speaker != "narrator" and out.speaker not in state.scene.present:
        if nc_id and _slug(out.speaker) == nc_id:
            out.speaker = nc_id  # normalize to the id apply_directives will create
        else:
            out.speaker = state.scene.present[0] if state.scene.present else "narrator"
    # exit must reference someone on stage
    if out.directives.exit_character and out.directives.exit_character not in state.scene.present:
        out.directives.exit_character = None
    # clamp delta (range isn't grammar-enforceable)
    out.directives.relationship_delta = max(-25, min(25, out.directives.relationship_delta))
    # new_character must arrive playable: traits feed the discovery system, appearance
    # feeds every sprite render (same fallbacks as _repair_init for the first character)
    nc = out.directives.new_character
    if nc is not None:
        if not nc.one_line.strip():
            nc.one_line = f"a mysterious {nc.name}"
        if not nc.traits:
            nc.traits = ["curious", "reserved"]
        if not nc.appearance.strip():
            nc.appearance = nc.one_line or "a young person with distinctive anime-style features"
        # Cap runaway generations: these fields are re-injected into the context every
        # turn the character is on stage (one observed `goals` ran past 5000 chars).
        nc.one_line = clip_words(nc.one_line, 30)
        nc.appearance = clip_words(nc.appearance, 40)
        nc.voice = clip_words(nc.voice, 25)
        nc.goals = clip_words(nc.goals, 50)
    return out


# --------------------------------------------------------------------------- #
#  Mock generators (deterministic-ish, themed; make the loop fun without a model)
# --------------------------------------------------------------------------- #
def _mock_init(setup: SetupForm) -> InitOutput:
    theme = config.THEMES.get(setup.theme, next(iter(config.THEMES.values())))
    cid = "hana"
    return InitOutput(
        style_guide=f"{config.STYLE_BASE}. Setting: {theme['setting']}. Palette: {theme['palette']}. "
        f"Tone: {setup.tone}, warm and romantic.",
        scene=SceneChange(
            place=theme["setting"].split(",")[0].strip().title(),
            description=f"The scene opens in {theme['setting']}. "
            "A soft breeze carries cherry-blossom petals across the path ahead.",
            mood=setup.tone,
        ),
        first_character=NewCharacter(
            id=cid,
            name="Hana",
            one_line="a warm-hearted girl who always carries too many books and hides a poet's heart",
            appearance="a slender young woman with wavy chestnut hair pinned loosely, warm amber eyes, "
            "school uniform slightly dishevelled, a small notebook tucked under her arm",
            voice="bright and a little breathless, laughs easily, sometimes trails into a half-whispered aside",
            traits=["warm", "secretly poetic", "terrified of being ordinary"],
            goals="to share one of her poems with someone who will truly listen",
        ),
        opening_line="Oh! I wasn't — I mean… hi. Sorry, I was miles away. You're the new transfer student, right? "
        "I'm Hana. Welcome to, um, wherever this is.",
        situation_intro=(
            f"You find yourself at {theme['setting']}. "
            "It's one of those afternoons where the light seems to linger just a little longer than usual, "
            "and the world feels oddly full of possibility. You're about to cross paths with someone "
            "you won't easily forget."
        ),
    )


def _mock_turn(state: GameState, player_input: str) -> DirectorOutput:
    rng = random.Random(state.seed + state.turn_index)
    speaker = state.scene.present[0] if state.scene.present else "narrator"
    name = state.characters[speaker].name if speaker in state.characters else "The wood"
    low = player_input.lower()

    emotion = "joy" if any(w in low for w in ("thank", "kind", "friend", "love")) else "neutral"
    if any(w in low for w in ("no", "stop", "afraid", "run")):
        emotion = "fear"

    directives = Directives(relationship_delta=rng.choice([0, 1, 1, 2]))

    # move on a movement cue (and not too often)
    if any(w in low for w in _MOVE_WORDS) and rng.random() < 0.8:
        place = rng.choice(
            ["The Hollow Stair", "A Clearing of Bells", "The Underbridge", "The Quiet Library"]
        )
        directives.scene_change = SceneChange(
            place=place,
            description=f"The path folds and {place.lower()} dreams itself into being around you.",
            mood=state.scene.mood,
        )
        dialogue = (
            f"This way, then—the wood likes to be followed. {place} was just now imagined for you."
        )
    else:
        dialogue = (
            f"“{player_input.strip() or '...'}”, you say. {name} considers this as if tasting it."
        )

    return DirectorOutput(
        speaker=speaker, dialogue=dialogue, emotion=emotion, directives=directives
    )


# small shared helpers (kept identical to state.py's so seeds line up)
def _slug(text: str) -> str:
    return "".join(c if c.isalnum() else "_" for c in text.lower()).strip("_")[:40] or "x"


def _seed_for(world_seed: int, key: str) -> int:
    return (world_seed * 2654435761 + sum(ord(c) for c in key) * 40503) % (2**31)