Joseph Doherty
497 lines
18 KiB
Python
497 lines
18 KiB
Python
"""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,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Task 43: multi-entity prompt assembly (guest_id support)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _seed_with_guest(conn) -> None:
|
|
"""Seed a 3-entity scene: you (Sam) + host (Aria, bot_a) + guest (Iris, bot_b).
|
|
|
|
Group node row is initialized with summary + dynamic, edges in all
|
|
relevant directions are seeded, and activities are recorded for all
|
|
three entities.
|
|
"""
|
|
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="bot_authored", payload={
|
|
"id": "bot_b",
|
|
"name": "Iris",
|
|
"persona": "wry transplant from the Boston office",
|
|
"voice_samples": ["Oh, please.", "Don't make me say it twice."],
|
|
"traits": ["sardonic", "loyal"],
|
|
"backstory": "Met Aria at a conference two years back.",
|
|
"initial_relationship_to_you": "stranger; curious",
|
|
"kickoff_prose": "",
|
|
})
|
|
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": "bot_b",
|
|
"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"},
|
|
})
|
|
# Edges: host -> you, guest -> you, host -> guest, guest -> host.
|
|
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"],
|
|
})
|
|
append_event(conn, kind="edge_update", payload={
|
|
"source_id": "bot_a",
|
|
"target_id": "bot_b",
|
|
"affinity_delta": 20,
|
|
"trust_delta": 15,
|
|
"knowledge_facts": ["studied physics together"],
|
|
})
|
|
append_event(conn, kind="edge_update", payload={
|
|
"source_id": "bot_b",
|
|
"target_id": "you",
|
|
"affinity_delta": 4,
|
|
"trust_delta": 0,
|
|
"knowledge_facts": ["Aria's coworker"],
|
|
})
|
|
append_event(conn, kind="edge_update", payload={
|
|
"source_id": "bot_b",
|
|
"target_id": "bot_a",
|
|
"affinity_delta": 18,
|
|
"trust_delta": 12,
|
|
"knowledge_facts": ["former roommate"],
|
|
})
|
|
# Activity for all three entities — note distinct verbs so we can
|
|
# check whose activity got dropped under tight budget.
|
|
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="activity_change", payload={
|
|
"entity_id": "bot_b",
|
|
"container_id": 1,
|
|
"posture": "leaning against the doorframe",
|
|
"action": {"verb": "smirking-distinctively"},
|
|
"attention": "Aria",
|
|
})
|
|
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", "bot_b"],
|
|
})
|
|
append_event(conn, kind="group_node_initialized", payload={
|
|
"chat_id": "chat_bot_a",
|
|
"members": ["you", "bot_a", "bot_b"],
|
|
"summary": "Three coworkers catching up after hours UNIQUE-GROUP-SUMMARY.",
|
|
"dynamic": "warm-but-prickly UNIQUE-GROUP-DYNAMIC",
|
|
})
|
|
project(conn)
|
|
|
|
|
|
def test_assemble_with_no_guest_matches_phase1(tmp_path):
|
|
"""Regression: 2-entity scenario without guest_id behaves exactly as Phase 1."""
|
|
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=[],
|
|
)
|
|
body = msgs[0].content
|
|
# Phase 1 must blocks present.
|
|
assert "Aria" in body
|
|
assert "PERSONA" in body
|
|
assert "Sam" in body
|
|
assert "ACTIVITIES" in body
|
|
assert "62/100" in body # speaker → addressee edge intact
|
|
# No guest content leaks in.
|
|
assert "Group dynamic" not in body
|
|
assert "Iris" not in body
|
|
|
|
|
|
def test_assemble_with_guest_includes_group_node_summary(tmp_path):
|
|
"""When guest is present (auto-detected via chat.guest_bot_id) and a
|
|
group_node row exists, its summary + dynamic are rendered."""
|
|
db = tmp_path / "t.db"
|
|
apply_migrations(db)
|
|
with open_db(db) as conn:
|
|
_seed_with_guest(conn)
|
|
msgs = assemble_narrative_prompt(
|
|
conn,
|
|
chat_id="chat_bot_a",
|
|
speaker_bot_id="bot_a",
|
|
recent_dialogue=[],
|
|
retrieved_memory_summaries=[],
|
|
)
|
|
body = msgs[0].content
|
|
assert "Group dynamic" in body
|
|
assert "UNIQUE-GROUP-SUMMARY" in body
|
|
assert "UNIQUE-GROUP-DYNAMIC" in body
|
|
# Guest activity also present (SHOULD-tier, fits at default budget).
|
|
assert "smirking-distinctively" in body
|
|
# Speaker's other edges include the host -> guest direction.
|
|
assert "Iris" in body
|
|
|
|
|
|
def test_assemble_when_speaker_is_guest_orients_edges_correctly(tmp_path):
|
|
"""When the guest is the speaker, identity is the guest, the
|
|
addressee edge is guest → you, and other edges include guest → host."""
|
|
db = tmp_path / "t.db"
|
|
apply_migrations(db)
|
|
with open_db(db) as conn:
|
|
_seed_with_guest(conn)
|
|
msgs = assemble_narrative_prompt(
|
|
conn,
|
|
chat_id="chat_bot_a",
|
|
speaker_bot_id="bot_b", # guest as speaker
|
|
recent_dialogue=[],
|
|
retrieved_memory_summaries=[],
|
|
)
|
|
body = msgs[0].content
|
|
# Speaker identity is the guest's persona.
|
|
assert "You are Iris." in body
|
|
assert "wry transplant from the Boston office" in body
|
|
# Edge to addressee is guest → you (Sam) with the seeded values
|
|
# (default 50 + 4 affinity = 54).
|
|
assert "YOUR EDGE TO Sam" in body
|
|
assert "54/100" in body
|
|
# Other edges include guest → host (Aria) with seeded value
|
|
# (default 50 + 18 = 68).
|
|
assert "OTHER EDGES" in body
|
|
assert "Aria" in body
|
|
assert "68/100" in body
|
|
|
|
|
|
def test_assemble_with_tight_budget_drops_guest_activity_first(tmp_path):
|
|
"""Under tight budget MUST blocks survive but SHOULD-tier guest
|
|
activity is dropped first."""
|
|
db = tmp_path / "t.db"
|
|
apply_migrations(db)
|
|
with open_db(db) as conn:
|
|
_seed_with_guest(conn)
|
|
# Short dialogue so MUST core (speaker identity + edge + last 4
|
|
# turns + closing) sits comfortably under the hard budget while
|
|
# SHOULD-tier additions (guest activity, group node, other edges)
|
|
# would push over.
|
|
dialogue = [
|
|
{"speaker": "you", "text": "line-16 hi there"},
|
|
{"speaker": "bot_a", "text": "line-17 hey"},
|
|
{"speaker": "you", "text": "line-18 quiet night"},
|
|
{"speaker": "bot_a", "text": "line-19 indeed"},
|
|
]
|
|
msgs = assemble_narrative_prompt(
|
|
conn,
|
|
chat_id="chat_bot_a",
|
|
speaker_bot_id="bot_a",
|
|
recent_dialogue=dialogue,
|
|
retrieved_memory_summaries=[],
|
|
# MUST core ~310 tokens; SHOULD additions (guest activity +
|
|
# group node + other edges) push it well over 380. budget_hard
|
|
# is set just above MUST core so SHOULD-tier blocks must be
|
|
# trimmed away.
|
|
budget_soft=250,
|
|
budget_hard=340,
|
|
)
|
|
body = msgs[0].content
|
|
# MUST: speaker identity, edge to addressee, last 4 dialogue turns.
|
|
assert "Aria" in body
|
|
assert "YOUR EDGE TO Sam" in body
|
|
for i in range(16, 20):
|
|
assert f"line-{i:02d}" in body
|
|
# Guest activity (SHOULD-tier) must be dropped under tight budget.
|
|
assert "smirking-distinctively" not in body
|
|
# Token budget honoured.
|
|
import tiktoken
|
|
enc = tiktoken.get_encoding("cl100k_base")
|
|
assert len(enc.encode(body)) <= 340
|