"""Context assembly per AI Dungeon's memory system (help.aidungeon.com/faq/the-memory-system): [AI Instructions] always included [Plot Essentials] always included (classic "Memory") [Story Summary] always included (manual in Phase 3, auto in Phase 6) [Used Memories] top-K memory-bank retrievals (Phase 6, when enabled) [Triggered Story Cards] "World Lore: ", conditional; first dropped when over budget [Story history] newest actions that fit the remaining token budget [Author's Note] injected AUTHORS_NOTE_DEPTH actions before the end of history [Latest player action] (+ script frontMemory right after it, Phase 4) """ import functools from dataclasses import dataclass import tiktoken from .. import models, worldstate AUTHORS_NOTE_DEPTH = 3 # actions from the end of history CARD_BUDGET_SHARE = 0.4 # max share of non-reserved budget that story cards may take SEPARATOR = "\n\n" @functools.lru_cache(maxsize=1) def _encoding() -> tiktoken.Encoding: return tiktoken.get_encoding("cl100k_base") def count_tokens(text: str) -> int: return len(_encoding().encode(text)) def truncate_to_last_tokens(text: str, budget: int) -> str: tokens = _encoding().encode(text) if len(tokens) <= budget: return text return _encoding().decode(tokens[-budget:]) @dataclass class Section: label: str text: str @property def tokens(self) -> int: return count_tokens(self.text) def _script_memory(adventure: models.Adventure) -> dict: """Script-provided memory overrides (populated by Phase 4 scripting).""" state = adventure.script_state if isinstance(adventure.script_state, dict) else {} memory = state.get("memory") return memory if isinstance(memory, dict) else {} def _history_text(action: models.Action) -> str: """An AI turn as the model should see it in replayed history: its narration with the state block it emitted re-appended (reconstructed from the stored delta). The block is stripped before storage/UI, so without this every past AI turn would look like one that emitted nothing — biasing the model, by imitation, to stop emitting too. Player turns and blockless turns are returned unchanged.""" text = action.text snap = action.context_snapshot if isinstance(action.context_snapshot, dict) else None if snap: ws = snap.get("world_state") if isinstance(ws, dict): block = worldstate.render_delta_block(ws.get("delta") or {}) if block: text = f"{text}\n{block}" return text def story_actions( adventure: models.Adventure, exclude_action_id: int | None = None ) -> list[models.Action]: """The adventure's non-empty actions, as the model should see them. `exclude_action_id` drops one action from the story — used by retry, where the row being regenerated is still attached to the adventure (it holds the variant history) but must not appear in the context assembled to replace it. """ return [ a for a in adventure.actions if a.text.strip() and (exclude_action_id is None or a.id != exclude_action_id) ] def _visible_npcs(actions: list[models.Action], stat_schema: dict) -> dict[str, str]: """Defined NPCs whose trigger words appear in the recent story — the ones "in scene", so only their stats get injected. Maps npc id -> display name.""" recent = SEPARATOR.join(a.text for a in actions[-6:]).lower() visible: dict[str, str] = {} for npc_key, ndef in (stat_schema.get("npcs") or {}).items(): if not isinstance(ndef, dict): continue if any(trigger in recent for trigger in worldstate.npc_triggers(ndef, npc_key)): visible[npc_key] = worldstate.npc_name(ndef, npc_key) return visible def _match_cards(cards: list[models.StoryCard], window_text: str) -> list[dict]: """AI Dungeon trigger rules: case-insensitive, space-sensitive, partial-word ('boat' triggers on 'boats'). Returns one record per card with the keyword that fired.""" haystack = window_text.lower() matched = [] for card in cards: for key in (k.strip().lower() for k in card.keys.split(",")): if key and key in haystack: matched.append( {"id": card.id, "name": card.name, "keyword": key, "entry": card.entry} ) break return matched def build_context( adventure: models.Adventure, settings: models.Settings, memory_bank: dict | None = None, exclude_action_id: int | None = None, ) -> tuple[str, str, dict]: """Returns (system_text, story_text, context_report). `memory_bank` is the result of memorybank.retrieve_memories (None when the bank is off); `exclude_action_id` omits one action from the story (see story_actions).""" script_mem = _script_memory(adventure) actions = story_actions(adventure, exclude_action_id) # ----- Always-included components ----- system_sections: list[Section] = [Section("narrator", settings.narrator_prompt.strip())] # RPG world state (Phase 12): current stats/milestones + how to report changes. stat_schema = adventure.scenario.stat_schema if adventure.scenario else None has_ws = worldstate.has_schema(stat_schema) if has_ws: guide = worldstate.render_reference(stat_schema) if guide: system_sections.append(Section("world_state_guide", guide)) block = worldstate.render_state_section( adventure.world_state, stat_schema, _visible_npcs(actions, stat_schema) ) if block: system_sections.append(Section("world_state", block)) system_sections.append(Section("world_state_rule", worldstate.EMIT_RULE)) if isinstance(script_mem.get("context"), str) and script_mem["context"].strip(): system_sections.append(Section("script_context", script_mem["context"].strip())) if adventure.ai_instructions.strip(): system_sections.append(Section("ai_instructions", adventure.ai_instructions.strip())) if adventure.memory.strip(): system_sections.append( Section("plot_essentials", f"Plot essentials:\n{adventure.memory.strip()}") ) if adventure.story_summary.strip(): system_sections.append( Section("story_summary", f"Story summary:\n{adventure.story_summary.strip()}") ) if memory_bank and memory_bank.get("used"): lines = "\n".join(f"- {m['text']}" for m in memory_bank["used"]) system_sections.append(Section("used_memories", f"Memories:\n{lines}")) authors_note_text = adventure.authors_note.strip() if isinstance(script_mem.get("authorsNote"), str) and script_mem["authorsNote"].strip(): authors_note_text = script_mem["authorsNote"].strip() authors_note = f"[Author's note: {authors_note_text}]" if authors_note_text else "" front_memory = "" if isinstance(script_mem.get("frontMemory"), str): front_memory = script_mem["frontMemory"].strip() reserved = ( sum(s.tokens for s in system_sections) + count_tokens(authors_note) + count_tokens(front_memory) + (count_tokens(worldstate.EMIT_REMINDER) if has_ws else 0) ) available = max(256, settings.context_token_budget - reserved) # ----- Story cards: triggered by recent story text (the window history could fill) ----- trigger_window = truncate_to_last_tokens(SEPARATOR.join(a.text for a in actions), available) triggered = _match_cards(adventure.story_cards, trigger_window) card_budget = int(available * CARD_BUDGET_SHARE) card_records = [] lore_lines: list[str] = [] used = 0 for match in triggered: line = f"World Lore: {match['entry'].strip()}" tokens = count_tokens(line) included = used + tokens <= card_budget if included: lore_lines.append(line) used += tokens card_records.append( {"id": match["id"], "name": match["name"], "keyword": match["keyword"], "included": included} ) if lore_lines: system_sections.append(Section("world_lore", "\n".join(lore_lines))) # ----- Story history: newest first until the remaining budget is spent ----- history_budget = available - used included_actions: list[models.Action] = [] spent = 0 oldest_truncated = False for action in reversed(actions): # Budget on the text as it will actually appear — with the re-attached # state block (B) when this adventure tracks world state. rendered = _history_text(action) if has_ws else action.text tokens = count_tokens(rendered) + count_tokens(SEPARATOR) if spent + tokens > history_budget: if not included_actions: # Even the newest action alone is over budget: hard-truncate it. included_actions.append( models.Action( adventure_id=action.adventure_id, index=action.index, type=action.type, text=truncate_to_last_tokens(action.text, history_budget), ) ) oldest_truncated = True break included_actions.append(action) spent += tokens included_actions.reverse() # ----- Assemble story text with author's note near the end ----- # Re-attach each AI turn's state block (stripped before storage) so recent # history shows the model its own emit pattern to imitate. texts = [_history_text(a) if has_ws else a.text for a in included_actions] note_sections: list[Section] = [] if authors_note: pos = max(0, len(texts) - AUTHORS_NOTE_DEPTH) before, after = texts[:pos], texts[pos:] if before: note_sections.append(Section("history", SEPARATOR.join(before))) note_sections.append(Section("authors_note", authors_note)) note_sections.append(Section("recent_history", SEPARATOR.join(after))) else: note_sections.append(Section("history", SEPARATOR.join(texts))) if front_memory: note_sections.append(Section("front_memory", front_memory)) if has_ws: # Terminal reminder: the emit rule sits up in the system block, far from # where the model generates; repeat it last, in the strongest recency slot. note_sections.append(Section("world_state_reminder", worldstate.EMIT_REMINDER)) story_sections = [s for s in note_sections if s.text] system_text = SEPARATOR.join(s.text for s in system_sections if s.text) story_text = SEPARATOR.join(s.text for s in story_sections) all_sections = [s for s in system_sections if s.text] + story_sections report = { "sections": [ {"label": s.label, "text": s.text, "tokens": s.tokens} for s in all_sections ], "prompt": {"system": system_text, "story": story_text}, "tokens": { "total": count_tokens(system_text) + count_tokens(story_text), "budget": settings.context_token_budget, }, "cards": card_records, "memories": memory_bank, "history": { "included": len(included_actions), "total": len(actions), "oldest_truncated": oldest_truncated, }, "settings": { "model": settings.model, "api_mode": settings.api_mode, "temperature": settings.temperature, "max_output_tokens": settings.max_output_tokens, }, } return system_text, story_text, report