chat/tests/test_prompt.py

"""Tests for chat.services.prompt.assemble_narrative_prompt.

Covers Task 18 — must/should/nice trim tiers (Requirements §3.2) and
the speaker prompt assembly order (§6.3). Tests use direct event-log
seeding so the projector populates state exactly the way the runtime
will at play-time. No LLM is invoked: prompt assembly is deterministic.
"""

from __future__ import annotations

import pytest

from chat.db.connection import open_db
from chat.db.migrate import apply_migrations
from chat.eventlog.log import append_event
from chat.eventlog.projector import project
import chat.state.entities  # noqa: F401  (registers handlers)
import chat.state.edges  # noqa: F401
import chat.state.memory  # noqa: F401
import chat.state.world  # noqa: F401
from chat.llm.client import Message
from chat.services.prompt import assemble_narrative_prompt


def _seed_basic(conn) -> None:
    """Seed bot, you-entity, edge, chat, container, scene, activities."""
    append_event(conn, kind="bot_authored", payload={
        "id": "bot_a",
        "name": "Aria",
        "persona": "reserved coworker who notices things",
        "voice_samples": ["I — sorry, I didn't mean to.", "Right. Of course."],
        "traits": ["introverted", "observant"],
        "backstory": "An archivist who joined the firm last spring.",
        "initial_relationship_to_you": "coworker; mild crush; never voiced",
        "kickoff_prose": "you stay late at the office",
    })
    append_event(conn, kind="you_authored", payload={
        "name": "Sam",
        "pronouns": "they/them",
        "persona": "tired analyst",
    })
    append_event(conn, kind="chat_created", payload={
        "id": "chat_bot_a",
        "host_bot_id": "bot_a",
        "guest_bot_id": None,
        "initial_time": "2026-04-26T20:00:00+00:00",
        "narrative_anchor": "Day 1 evening",
        "weather": "clear",
    })
    append_event(conn, kind="container_created", payload={
        "chat_id": "chat_bot_a",
        "name": "office bullpen",
        "type": "workplace",
        "properties": {"public": False, "moving": False, "audible_range": "room"},
    })
    append_event(conn, kind="edge_update", payload={
        "source_id": "bot_a",
        "target_id": "you",
        "affinity_delta": 12,
        "trust_delta": 5,
        "knowledge_facts": [
            "they work on the same floor",
            "they've stayed late twice this week",
        ],
    })
    append_event(conn, kind="activity_change", payload={
        "entity_id": "you",
        "container_id": 1,
        "posture": "sitting at your desk",
        "action": {"verb": "finishing emails"},
        "attention": "the screen",
        "holding": ["coffee mug"],
    })
    append_event(conn, kind="activity_change", payload={
        "entity_id": "bot_a",
        "container_id": 1,
        "posture": "sitting at her desk",
        "action": {"verb": "pretending to work"},
        "attention": "you, in glances",
    })
    append_event(conn, kind="scene_opened", payload={
        "chat_id": "chat_bot_a",
        "container_id": 1,
        "started_at": "2026-04-26T20:00:00+00:00",
        "participants": ["you", "bot_a"],
    })
    project(conn)


def test_basic_assembly_returns_system_message_with_all_must_blocks(tmp_path):
    db = tmp_path / "t.db"
    apply_migrations(db)
    with open_db(db) as conn:
        _seed_basic(conn)
        msgs = assemble_narrative_prompt(
            conn,
            chat_id="chat_bot_a",
            speaker_bot_id="bot_a",
            recent_dialogue=[],
            retrieved_memory_summaries=[],
        )
    assert isinstance(msgs, list)
    assert len(msgs) == 1
    sys_msg = msgs[0]
    assert isinstance(sys_msg, Message)
    assert sys_msg.role == "system"
    body = sys_msg.content
    # Must-include markers
    assert "Aria" in body
    assert "PERSONA" in body
    assert "ACTIVITIES" in body
    assert "CURRENT SCENE" in body
    # Edge to addressee — name + numeric values (default affinity 50, +12 = 62)
    assert "Sam" in body
    assert "62/100" in body


def test_user_turn_appended_as_user_message(tmp_path):
    db = tmp_path / "t.db"
    apply_migrations(db)
    with open_db(db) as conn:
        _seed_basic(conn)
        msgs = assemble_narrative_prompt(
            conn,
            chat_id="chat_bot_a",
            speaker_bot_id="bot_a",
            user_turn_prose="*looks up* Hey.",
            recent_dialogue=[],
            retrieved_memory_summaries=[],
        )
    assert len(msgs) == 2
    assert msgs[0].role == "system"
    assert msgs[1].role == "user"
    assert msgs[1].content == "*looks up* Hey."


def test_must_only_succeeds_with_empty_optional_blocks(tmp_path):
    """No dialogue, memories, other edges, or previous scene summary — should not raise."""
    db = tmp_path / "t.db"
    apply_migrations(db)
    with open_db(db) as conn:
        _seed_basic(conn)
        msgs = assemble_narrative_prompt(
            conn,
            chat_id="chat_bot_a",
            speaker_bot_id="bot_a",
            recent_dialogue=None,  # default → nothing
            retrieved_memory_summaries=None,
            user_turn_prose=None,
        )
    assert len(msgs) == 1
    body = msgs[0].content
    # Must blocks present
    assert "PERSONA" in body
    assert "ACTIVITIES" in body
    # Optional blocks not in body (nothing to render)
    assert "OTHER EDGES" not in body
    assert "PREVIOUS SCENE SUMMARY" not in body
    assert "RELEVANT MEMORIES" not in body


def test_long_dialogue_keeps_last_4_verbatim_and_summarizes_earlier(tmp_path):
    """Stuff a huge dialogue history under budget pressure; older turns
    must be elided to a placeholder, the last 4 verbatim, and earlier
    unique markers gone.
    """
    db = tmp_path / "t.db"
    apply_migrations(db)
    with open_db(db) as conn:
        _seed_basic(conn)
        dialogue = []
        for i in range(20):
            speaker = "you" if i % 2 == 0 else "bot_a"
            # Each line ~250 tokens of filler => 20 turns ≈ 5000 tokens,
            # which together with MUST blocks pushes over soft (1500).
            dialogue.append({
                "speaker": speaker,
                "text": f"unique-line-marker-{i:02d} " + ("filler " * 200),
            })
        msgs = assemble_narrative_prompt(
            conn,
            chat_id="chat_bot_a",
            speaker_bot_id="bot_a",
            recent_dialogue=dialogue,
            retrieved_memory_summaries=[],
            # Soft small enough to force NICE trim but hard fits MUST + 4.
            budget_soft=1200,
            budget_hard=8000,
        )
    body = msgs[0].content
    # The last 4 unique markers (16, 17, 18, 19) must be present verbatim.
    for i in range(16, 20):
        assert f"unique-line-marker-{i:02d}" in body, f"expected last-4 marker {i} in body"
    # Older markers must be dropped (replaced by elision placeholder).
    for i in range(0, 16):
        assert f"unique-line-marker-{i:02d}" not in body
    # An "earlier" summary line must be present.
    assert "earlier" in body.lower()
    # Token count of system message respects hard budget.
    import tiktoken
    enc = tiktoken.get_encoding("cl100k_base")
    assert len(enc.encode(body)) <= 8000


def test_memories_drop_to_top_2_under_budget_pressure(tmp_path):
    """4 memory summaries, each large; under tight soft budget only 2 should appear."""
    db = tmp_path / "t.db"
    apply_migrations(db)
    with open_db(db) as conn:
        _seed_basic(conn)
        # Each ~1500 tokens of repeated text; drop tier should kick in.
        long_chunk = "alpha beta gamma delta " * 400
        memories = [
            f"MEMORY-A {long_chunk}",
            f"MEMORY-B {long_chunk}",
            f"MEMORY-C {long_chunk}",
            f"MEMORY-D {long_chunk}",
        ]
        msgs = assemble_narrative_prompt(
            conn,
            chat_id="chat_bot_a",
            speaker_bot_id="bot_a",
            recent_dialogue=[],
            retrieved_memory_summaries=memories,
            # Pressure: budgets that allow MUST + 2 memories but not 4.
            budget_soft=4000,
            budget_hard=5000,
        )
    body = msgs[0].content
    # MEMORY-A and MEMORY-B are the top-2 and should remain; C & D dropped.
    assert "MEMORY-A" in body
    assert "MEMORY-B" in body
    assert "MEMORY-C" not in body
    assert "MEMORY-D" not in body
    # Token count fits the hard budget.
    import tiktoken
    enc = tiktoken.get_encoding("cl100k_base")
    assert len(enc.encode(body)) <= 5000


def test_must_exceeds_budget_hard_raises_value_error(tmp_path):
    db = tmp_path / "t.db"
    apply_migrations(db)
    with open_db(db) as conn:
        _seed_basic(conn)
        with pytest.raises(ValueError):
            assemble_narrative_prompt(
                conn,
                chat_id="chat_bot_a",
                speaker_bot_id="bot_a",
                recent_dialogue=[],
                retrieved_memory_summaries=[],
                budget_soft=5,
                budget_hard=10,
            )