diff --git a/chat/services/prompt.py b/chat/services/prompt.py index 1df074b..48a0a7f 100644 --- a/chat/services/prompt.py +++ b/chat/services/prompt.py @@ -37,6 +37,7 @@ import tiktoken from chat.llm.client import Message from chat.state.edges import get_edge, list_edges_for from chat.state.entities import get_bot, get_you +from chat.state.group_node import get_group_node from chat.state.memory import search_memories from chat.state.world import ( active_scene, @@ -206,6 +207,26 @@ def _build_previous_scene_block(pov_summary: str | None) -> str | None: return "PREVIOUS SCENE SUMMARY:\n" + pov_summary +def _build_group_node_block(group_node: dict | None) -> str | None: + """Render the group-node summary + dynamic as a SHOULD-tier block. + + Used only in 3-entity scenes (you + host + guest). Returns None when + the row is missing or both summary and dynamic are empty. + """ + if not group_node: + return None + summary = (group_node.get("summary") or "").strip() + dynamic = (group_node.get("dynamic") or "").strip() + if not summary and not dynamic: + return None + lines = ["Group dynamic:"] + if summary: + lines.append(f"- Summary: {summary}") + if dynamic: + lines.append(f"- Dynamic: {dynamic}") + return "\n".join(lines) + + def _closing_instruction(speaker_name: str, addressee_name: str) -> str: return ( f"Continue the scene as {speaker_name}, in their voice, responding " @@ -287,6 +308,7 @@ def assemble_narrative_prompt( budget_soft: int = 6000, budget_hard: int = 8000, encoding_name: str = "cl100k_base", + guest_id: str | None = None, ) -> list[Message]: """Assemble the narrative prompt for ``speaker_bot_id`` to respond. @@ -313,6 +335,15 @@ def assemble_narrative_prompt( if chat is None: raise ValueError(f"chat_id {chat_id!r} not found") + # Auto-detect guest from chat state when caller didn't pass one. + # Phase 1 chats have ``guest_bot_id is None``; the auto-detect is a + # no-op there and the function behaves exactly as before. + if guest_id is None: + guest_id = chat.get("guest_bot_id") + # A speaker addressing themself as guest doesn't add a third party. + if guest_id is not None and guest_id == speaker_bot_id: + guest_id = None + you = get_you(conn) addressee_id, addressee_name = _resolve_addressee(conn, addressee, you) @@ -325,9 +356,10 @@ def assemble_narrative_prompt( addressee_name, ) - # Activity for present entities. Phase 1: you + speaker bot. (When a - # guest is added in Phase 1+, callers that know about it can pass - # extra activities via a future hook; for now we keep it strict.) + # Activity for present entities. Core (MUST): you + speaker bot. + # Phase 2 (SHOULD-tier): when a third party (guest) is present in + # the chat, append their activity in a separate block so it can be + # trimmed independently under tight budget. activities: list[dict] = [] you_act = get_activity(conn, "you") if you_act is not None: @@ -341,6 +373,34 @@ def assemble_narrative_prompt( activities.append(bot_act) activity_block = _build_activity_block(activities) + # SHOULD-tier guest activity extension (Phase 2 / Task 43). + guest_activity_block: str | None = None + if guest_id is not None: + guest_act = get_activity(conn, guest_id) + if guest_act is not None: + guest_act = dict(guest_act) + guest_bot = get_bot(conn, guest_id) + guest_act["_display_name"] = ( + guest_bot["name"] if guest_bot else guest_id + ) + guest_activity_block = _build_activity_block([guest_act]) + + # SHOULD-tier group-node block (Phase 2 / Task 43): rendered only + # when the group_node row is present AND it covers all three of + # you + host + guest (per the Task 43 spec). + group_node_block: str | None = None + if guest_id is not None: + gn = get_group_node(conn, chat_id) + if gn is not None: + members = set(gn.get("members") or []) + host_id = chat.get("host_bot_id") + required = {"you"} + if host_id is not None: + required.add(host_id) + required.add(guest_id) + if required.issubset(members): + group_node_block = _build_group_node_block(gn) + container = None if chat.get("active_scene_id"): scene = get_scene(conn, chat["active_scene_id"]) @@ -421,6 +481,8 @@ def assemble_narrative_prompt( include_previous_scene: bool, include_memories_top_k: int, dialogue_keep: int, + include_guest_activity: bool = True, + include_group_node: bool = True, ) -> tuple[str, int, list[dict]]: # dialogue: keep the last `dialogue_keep` turns verbatim; older # turns become an "earlier:" placeholder line. @@ -447,6 +509,8 @@ def assemble_narrative_prompt( other_edges_block if include_other_edges else None, scene_block, activity_block, + guest_activity_block if include_guest_activity else None, + group_node_block if include_group_node else None, prev_block, memories_block, dialogue_block, @@ -463,12 +527,25 @@ def assemble_narrative_prompt( nice_memories_k = min(4, len(memory_summaries)) include_prev = previous_scene_summary is not None include_other = other_edges_block is not None + include_guest_activity = guest_activity_block is not None + include_group_node = group_node_block is not None - body, total, _ = assemble( - include_other_edges=include_other, - include_previous_scene=include_prev, - include_memories_top_k=nice_memories_k, - dialogue_keep=nice_dialogue_keep, + def _build(*, prev: bool, mem_k: int, dlg: int, other: bool, + guest_act: bool, group: bool) -> tuple[str, int]: + body, total, _ = assemble( + include_other_edges=other, + include_previous_scene=prev, + include_memories_top_k=mem_k, + dialogue_keep=dlg, + include_guest_activity=guest_act, + include_group_node=group, + ) + return body, total + + body, total = _build( + prev=include_prev, mem_k=nice_memories_k, dlg=nice_dialogue_keep, + other=include_other, guest_act=include_guest_activity, + group=include_group_node, ) # If under soft, we're done. @@ -478,34 +555,31 @@ def assemble_narrative_prompt( # Drop NICE in order: previous scene → memories beyond top-2 → # older dialogue turns (collapse to 4). if include_prev: - body, total, _ = assemble( - include_other_edges=include_other, - include_previous_scene=False, - include_memories_top_k=nice_memories_k, - dialogue_keep=nice_dialogue_keep, - ) include_prev = False + body, total = _build( + prev=include_prev, mem_k=nice_memories_k, dlg=nice_dialogue_keep, + other=include_other, guest_act=include_guest_activity, + group=include_group_node, + ) if total <= budget_soft: return _emit(body, user_turn_prose) if nice_memories_k > 2: nice_memories_k = 2 - body, total, _ = assemble( - include_other_edges=include_other, - include_previous_scene=False, - include_memories_top_k=nice_memories_k, - dialogue_keep=nice_dialogue_keep, + body, total = _build( + prev=include_prev, mem_k=nice_memories_k, dlg=nice_dialogue_keep, + other=include_other, guest_act=include_guest_activity, + group=include_group_node, ) if total <= budget_soft: return _emit(body, user_turn_prose) if nice_dialogue_keep > baseline_keep: nice_dialogue_keep = baseline_keep - body, total, _ = assemble( - include_other_edges=include_other, - include_previous_scene=False, - include_memories_top_k=nice_memories_k, - dialogue_keep=nice_dialogue_keep, + body, total = _build( + prev=include_prev, mem_k=nice_memories_k, dlg=nice_dialogue_keep, + other=include_other, guest_act=include_guest_activity, + group=include_group_node, ) if total <= budget_soft: return _emit(body, user_turn_prose) @@ -513,21 +587,37 @@ def assemble_narrative_prompt( # Drop more NICE until we're under hard: memories all the way to 0. while nice_memories_k > 0 and total > budget_hard: nice_memories_k = max(0, nice_memories_k - 1) - body, total, _ = assemble( - include_other_edges=include_other, - include_previous_scene=False, - include_memories_top_k=nice_memories_k, - dialogue_keep=nice_dialogue_keep, + body, total = _build( + prev=include_prev, mem_k=nice_memories_k, dlg=nice_dialogue_keep, + other=include_other, guest_act=include_guest_activity, + group=include_group_node, + ) + + # Drop SHOULD-tier blocks in order: guest activity → group node → + # other edges. (Guest activity goes first per Task 43 spec — it's + # the most expendable additive context.) + if include_guest_activity and total > budget_hard: + include_guest_activity = False + body, total = _build( + prev=include_prev, mem_k=nice_memories_k, dlg=nice_dialogue_keep, + other=include_other, guest_act=include_guest_activity, + group=include_group_node, + ) + + if include_group_node and total > budget_hard: + include_group_node = False + body, total = _build( + prev=include_prev, mem_k=nice_memories_k, dlg=nice_dialogue_keep, + other=include_other, guest_act=include_guest_activity, + group=include_group_node, ) - # Drop SHOULD: other edges. if include_other and total > budget_hard: include_other = False - body, total, _ = assemble( - include_other_edges=False, - include_previous_scene=False, - include_memories_top_k=nice_memories_k, - dialogue_keep=nice_dialogue_keep, + body, total = _build( + prev=include_prev, mem_k=nice_memories_k, dlg=nice_dialogue_keep, + other=include_other, guest_act=include_guest_activity, + group=include_group_node, ) if total > budget_hard: diff --git a/tests/test_prompt.py b/tests/test_prompt.py index bef8dc0..2ef4641 100644 --- a/tests/test_prompt.py +++ b/tests/test_prompt.py @@ -253,3 +253,244 @@ def test_must_exceeds_budget_hard_raises_value_error(tmp_path): 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