visualnovel/memory.py

"""Memory is a *budget*, not an archive.

Every turn we assemble the user-side context from (in priority order): the rolling summary,
the current scene, the present characters' sheets, the last k verbatim turns, and the player's
new line — trimmed to a conservative token budget. Older turns get folded into `summary` by
`orchestrator.compact_memory`. This is the literal "Thousand Token Wood".
"""

from __future__ import annotations

from . import characters, config
from .schemas import GameState

# Flags that drive the engine/UI, not story facts — never surfaced in the FACTS: line.
_INTERNAL_FLAGS = {
    "current_music",
    "situation_intro",
    "beat_started_turn",
    "ending_kind",
    "ending_text",
    "player_name",
}
_INTERNAL_PREFIXES = ("milestone50_",)


def assemble_context(state: GameState, player_input: str, action_note: str = "") -> str:
    parts: list[str] = []  # the head: never trimmed (style, scene, sheets, milestones)
    parts.append(f"WORLD STYLE: {state.style_guide}")
    parts.append(f"VIBE: {state.vibe}  |  BEAT: {state.beat}")
    player_name = state.flags.get("player_name", "")
    if player_name:
        parts.append(f'The player is called "{player_name}" - characters address them by name.')
    if state.summary:
        parts.append(f"THE TALE SO FAR:\n{state.summary}")
    parts.append(
        f"WHERE WE ARE:\n- place: {state.scene.place}\n"
        f"- description: {state.scene.description}\n- mood: {state.scene.mood}"
    )
    parts.append(f"SPIRITS ON STAGE:\n{characters.present_sheets_block(state)}")

    # Characters the player has already met but who are currently off-stage.
    # The LLM can bring them back via new_character (same id) — data is preserved.
    off_stage = [
        f"  - {ch.name} (id: {ch.id}, rel: {ch.relationship:+d})"
        for cid, ch in state.characters.items()
        if cid not in state.scene.present
    ]
    if off_stage:
        parts.append(
            "KNOWN BUT OFF STAGE (can return — use their existing id):\n" + "\n".join(off_stage)
        )

    # ── Milestone-50 special moment — one-shot hint armed by apply_directives ──
    for ch in state.characters.values():
        if ch.id in state.scene.present and state.flags.get(f"milestone50_{ch.id}") == "pending":
            parts.append(
                f"*** GROWING CLOSE: {ch.name}'s affection just reached 50+/100. ***\n"
                f"This turn, have {ch.name} initiate a small special moment: suggest going "
                f"somewhere together (emit scene_change), confide a small secret, or give "
                f"the player an affectionate nickname. Make it feel earned, not sudden."
            )

    # ── Pacing nudge — deterministic, fires when a beat lingers too long ──────
    turns_in_beat = state.turn_index - int(state.flags.get("beat_started_turn", "0") or 0)
    if turns_in_beat >= 7 and state.beat in ("opening", "rising", "turn"):
        parts.append(
            f"PACING: {turns_in_beat} turns in beat '{state.beat}' - the scene is "
            f"lingering. Consider advance_beat=true this turn."
        )

    # ── Ending milestone hints — injected when a character hits ±100 ──────────
    for ch in state.characters.values():
        if ch.relationship >= 100 and ch.id in state.scene.present:
            parts.append(
                f"*** ROMANCE PEAK: {ch.name}'s affection is at MAX (100/100). ***\n"
                f'This turn MUST be their heartfelt confession. Emit ending(kind="warm") '
                f'with a poetic 2-4 sentence epilogue. Also emit set_music="romantic".'
            )
        elif ch.relationship <= -100 and ch.id in state.scene.present:
            parts.append(
                f"*** RELATIONSHIP BROKEN: {ch.name} is at MIN (-100/100). ***\n"
                f'They storm off with a final cutting line. MUST emit exit_character="{ch.id}". '
                f'If the stage becomes empty after this, also emit ending(kind="defeat") '
                f'with a dark 2-4 sentence epilogue. Emit set_music="sad".'
            )
    # If scene already empty and not yet ended — nudge for defeat
    if not state.scene.present and state.beat != "ended" and state.characters:
        parts.append(
            "*** THE STAGE IS EMPTY — the wanderer is completely alone. ***\n"
            'Emit ending(kind="defeat") with a melancholic epilogue. Emit set_music="sad".'
        )

    # Surface the current music track prominently so the LLM doesn't switch it needlessly.
    current_music = state.flags.get("current_music")
    if current_music:
        parts.append(
            f'CURRENT BGM: "{current_music}" — leave set_music null unless a MAJOR tonal '
            f"shift justifies a change (max once every 6-8 turns)."
        )

    other_flags = {
        k: v
        for k, v in state.flags.items()
        if k not in _INTERNAL_FLAGS and not k.startswith(_INTERNAL_PREFIXES)
    }
    if other_flags:
        parts.append("FACTS: " + ", ".join(f"{k}={v}" for k, v in other_flags.items()))

    # All turns since the last compaction — bounded by compaction (oldest get folded
    # into `summary`) and by the budget below (oldest lines dropped first).
    recent_lines = [
        f'wanderer: "{t.player}"\n{t.speaker}: "{t.dialogue}"' for t in state.recent_turns
    ]

    tail_parts = []
    if action_note:
        # Kept OUTSIDE the quoted player line — inside the quotes the model reads
        # directives as player speech and ignores them.
        tail_parts.append(f"DIRECTOR NOTE (mandatory):\n{action_note}")
    tail_parts += [
        f'{(player_name or "the wanderer").upper()} NOW SAYS: "{player_input.strip() or "…"}"',
        "Reply as one present spirit, then emit the directives JSON. "
        f"You may speak as one of: {characters.actor_roster(state)}.",
    ]

    return _fit_to_budget(parts, recent_lines, tail_parts, config.CONTEXT_TOKEN_BUDGET)


def should_compact(state: GameState) -> bool:
    # Compaction keeps RECENT_TURNS_K turns verbatim; trigger once at least as many
    # again have piled up, so each compaction call folds a meaningful chunk.
    return len(state.recent_turns) > 2 * config.RECENT_TURNS_K


def _fit_to_budget(
    head_parts: list[str],
    recent_lines: list[str],
    tail_parts: list[str],
    token_budget: int,
) -> str:
    """Rough char≈token/4 budget. Drops the oldest RECENT EXCHANGE lines first; the
    system-critical head (style, scene, sheets) and the player's line are preserved.
    Hard-clamps keeping the tail only as a last resort (head alone over budget)."""
    char_budget = token_budget * 4

    def _join(lines: list[str]) -> str:
        parts = list(head_parts)
        if lines:
            parts.append("RECENT EXCHANGE:\n" + "\n".join(lines))
        parts.extend(tail_parts)
        return "\n\n".join(parts)

    ctx = _join(recent_lines)
    if len(ctx) <= char_budget:
        return ctx
    lines = list(recent_lines)
    while lines:
        lines = lines[1:]
        ctx = _join(["…(earlier exchange trimmed)…", *lines])
        if len(ctx) <= char_budget:
            return ctx
    return "…(earlier dream trimmed)…\n\n" + ctx[-char_budget:]