Something went wrong. Try again.
A 5e storytelling engine with an LLM DM
Something went wrong. Try again.
Python
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455# pyright: reportArgumentType=false# Tests pass dict literals where typed inputs are stricter than runtime needs."""Tests for engine helper functions and context building."""
from pathlib import Pathfrom unittest.mock import patch
import pytest
from storied.engine import ( _extract_roll_reason, _tool_notification, load_prompt,)
class TestLoadPrompt: def test_loads_dm_system(self): result = load_prompt("dm-system") assert "Dungeon Master" in result
def test_custom_prompts_path(self, tmp_path: Path): prompt_file = tmp_path / "test-prompt.md" prompt_file.write_text("You are a test prompt.") result = load_prompt("test-prompt", prompts_path=tmp_path) assert result == "You are a test prompt."
def test_missing_prompt_raises(self, tmp_path: Path): with pytest.raises(FileNotFoundError): load_prompt("nonexistent", prompts_path=tmp_path)
class TestExtractRollReason: def test_extracts_reason(self): assert _extract_roll_reason('{"reason": "Athletics"}') == "Athletics"
def test_no_reason_field(self): assert _extract_roll_reason('{"notation": "1d20"}') is None
def test_invalid_json(self): assert _extract_roll_reason("not json") is None
def test_empty_string(self): assert _extract_roll_reason("") is None
class TestToolNotification: def test_known_tool(self): assert "Rolling" in _tool_notification("roll")
def test_mcp_prefixed_tool(self): assert "Establishing" in _tool_notification("mcp__storied__establish")
def test_unknown_tool(self): result = _tool_notification("something_new") assert "something_new" in result
def test_tune_tool(self): assert "Tuning" in _tool_notification("tune")
class TestDMEngineContext: """Tests for DMEngine context building (mocks MCP server startup)."""
@pytest.fixture def engine(self, tmp_path: Path): from storied.engine import DMEngine
world_dir = tmp_path / "worlds" / "test" world_dir.mkdir(parents=True)
prompts_dir = tmp_path / "prompts" prompts_dir.mkdir() (prompts_dir / "dm-system.md").write_text("You are a DM.")
with patch("storied.engine.start_mcp_server") as mock_mcp: from storied.initiative import InitiativeTracker from storied.log import CampaignLog from storied.tools import EntityIndex
mock_mcp.return_value = type( "Handle", (), { "url": "http://localhost:0/sse", "ctx": type( "Ctx", (), { "campaign_log": CampaignLog("test"), "entity_index": EntityIndex(world_dir), "vector_index": None, "initiative": InitiativeTracker(), }, )(), }, )() return DMEngine( world_id="test", player_id="default", prompt_name="dm-system", )
def test_build_context_no_style(self, engine): engine._build_context() assert "Style" not in engine._context_parts
def test_build_context_with_style(self, engine, tmp_path: Path): style_path = tmp_path / "worlds" / "test" / "style.md" style_path.write_text("# Style\n\nMore intrigue, less combat.\n")
engine._build_context()
assert "Style" in engine._context_parts assert "intrigue" in engine._context_parts["Style"]
def test_build_context_no_arc(self, engine): engine._build_context() assert "Arc" not in engine._context_parts
def test_build_context_with_arc(self, engine, tmp_path: Path): arc_path = tmp_path / "worlds" / "test" / "arc.md" arc_path.parent.mkdir(parents=True, exist_ok=True) arc_path.write_text( "# Campaign Arc\n\n## Premise\nA quiet mystery.\n" "\n## Off the Table\n- the cosmic-prisoner trope\n" )
engine._build_context()
assert "Arc" in engine._context_parts assert "Campaign Arc" in engine._context_parts["Arc"] assert "cosmic-prisoner" in engine._context_parts["Arc"]
def test_arc_appears_after_style_in_context_order( self, engine, tmp_path: Path, ): style_path = tmp_path / "worlds" / "test" / "style.md" style_path.write_text("# Style\n\nDark tone.\n") arc_path = tmp_path / "worlds" / "test" / "arc.md" arc_path.write_text("# Campaign Arc\n\n## Premise\nMystery.\n")
engine._build_context() parts = list(engine._context_parts.keys())
style_idx = parts.index("Style") arc_idx = parts.index("Arc") assert style_idx < arc_idx
def test_time_is_first_context_part(self, engine, tmp_path: Path): """The ambient clock header goes at the top of every turn's context so the DM can't miss it. Style comes immediately after.""" style_path = tmp_path / "worlds" / "test" / "style.md" style_path.write_text("# Style\n\nDark tone.\n")
engine._build_context() parts = list(engine._context_parts.keys())
assert parts[0] == "Time" assert parts[1] == "Style"
def test_time_header_shows_current_clock(self, engine): """The ambient time header always includes the current time anchor and period of day, so the DM's first impression each turn is what time it is.""" engine._build_context() header = engine._context_parts["Time"] assert "⏰ Current Time" in header now = engine._campaign_log.get_current_time() assert now.to_anchor() in header assert now.period_of_day().lower() in header.lower()
def test_estimate_tokens(self): from storied.engine import DMEngine
assert DMEngine._estimate_tokens("a" * 400) == 100
def test_format_entity(self, engine): result = engine._format_entity( "Npc", { "name": "Vera", "body": "Tavern owner.", }, ) assert "## Npc: Vera" in result assert "Tavern owner." in result
def test_parse_knowledge_file_no_frontmatter(self, engine, tmp_path: Path): f = tmp_path / "note.md" f.write_text("Just some text.") result = engine._parse_knowledge_file(f) assert result["body"] == "Just some text."
def test_parse_knowledge_file_with_frontmatter(self, engine, tmp_path: Path): f = tmp_path / "note.md" f.write_text("---\ntype: npc\nname: Vera\n---\n\nTavern owner.") result = engine._parse_knowledge_file(f) assert result["name"] == "Vera" assert result["body"] == "Tavern owner."
def test_parse_knowledge_file_malformed_frontmatter( self, engine, tmp_path: Path, ): f = tmp_path / "broken.md" f.write_text("---\nnot: [valid yaml\n---\n\nBody.") result = engine._parse_knowledge_file(f) # Falls back to body-only when frontmatter is unparseable assert "Body." in result["body"]
def test_parse_knowledge_file_open_frontmatter(self, engine, tmp_path: Path): f = tmp_path / "open.md" f.write_text("---\nnever closed") result = engine._parse_knowledge_file(f) # Falls back to whole-file body when there's no closing --- assert "never closed" in result["body"]
def test_log_transcript_writes_file(self, tmp_path: Path): from unittest.mock import patch
from storied.engine import DMEngine from storied.initiative import InitiativeTracker from storied.tools import EntityIndex
(tmp_path / "worlds" / "test").mkdir(parents=True) (tmp_path / "prompts").mkdir() (tmp_path / "prompts" / "dm-system.md").write_text("DM.")
transcript_path = tmp_path / "transcripts" / "session.jsonl" with patch("storied.engine.start_mcp_server") as mock_mcp: from storied.log import CampaignLog
mock_mcp.return_value = type( "Handle", (), { "url": "http://localhost:0/sse", "ctx": type( "Ctx", (), { "campaign_log": CampaignLog("test"), "entity_index": EntityIndex(tmp_path / "worlds" / "test"), "vector_index": None, "initiative": InitiativeTracker(), }, )(), }, )() engine = DMEngine( world_id="test", player_id="default", prompt_name="dm-system", transcript_path=transcript_path, )
engine._log_transcript("test_event", {"foo": "bar"}) assert transcript_path.exists() content = transcript_path.read_text() assert "test_event" in content assert "bar" in content
def test_log_transcript_no_path_is_noop(self, engine): # Engine constructed without a transcript_path → method returns early engine._log_transcript("noop", {"x": 1}) # No assertion needed; just verifying it doesn't crash
def test_get_context_stats_returns_breakdown(self, engine): engine._build_context() stats = engine.get_context_stats() assert "model_limit" in stats assert "system_prompt" in stats assert "context_parts" in stats assert isinstance(stats["context_parts"], dict) assert stats["context_total"] > 0
def test_get_current_time_returns_string(self, engine): result = engine.get_current_time() assert isinstance(result, str) assert len(result) > 0
def test_reset_clears_session(self, engine): engine._session_id = "abc-123" engine.reset() assert engine._session_id is None
def test_build_context_with_character(self, engine, tmp_path: Path): # Drop a character file in place; _build_context should pick it up from storied.character import create_character
(tmp_path / "players" / "default").mkdir(parents=True) create_character( player_id="default", name="Mira", race="Human", char_class="Rogue", level=3, abilities={ "strength": 10, "dexterity": 18, "constitution": 14, "intelligence": 14, "wisdom": 12, "charisma": 16, }, hp_max=24, ac=16, ) engine._build_context() assert "Character" in engine._context_parts assert "Mira" in engine._context_parts["Character"]
def test_build_context_with_session(self, engine, tmp_path: Path): from storied.session import save_session
(tmp_path / "players" / "default").mkdir(parents=True) save_session( "default", { "location": "The Tavern", "body": "## Present\n- [[Vera]]", "situation": "Resting", }, ) engine._build_context() assert "Session" in engine._context_parts
def test_build_context_loads_present_entities(self, engine, tmp_path: Path): from storied.session import save_session from storied.tools.entities import _do_establish
(tmp_path / "players" / "default").mkdir(parents=True)
# Establish an NPC and put them in the session's present list _do_establish( "npcs", "Vera", "Tavern owner.", "[[The Tavern]]", None, None, None, "test", engine._mcp.ctx.entity_index, type("FakeIdx", (), {"upsert": lambda *a, **k: None})(), ) save_session( "default", { "location": "The Tavern", "body": "## Present\n- [[Vera]]", }, )
engine._build_context() # The Vera entity should have been loaded into the DM context assert any("Vera" in v for v in engine._context_parts.values())
def test_load_player_knowledge_returns_none_without_dir(self, engine): # No knowledge dir created → returns None result = engine._load_player_knowledge() assert result is None
def test_load_player_knowledge_aggregates_files(self, engine, tmp_path: Path): knowledge = tmp_path / "players" / "default" / "worlds" / "test" / "npcs" knowledge.mkdir(parents=True) (knowledge / "vera.md").write_text( "---\nname: Vera Blackwater\n---\n\nA tavern owner." ) result = engine._load_player_knowledge() assert result is not None assert "Vera Blackwater" in result assert "tavern owner" in result
def test_find_entity_via_index(self, engine, tmp_path: Path): # Drop an entity file directly and rebuild the index so the lookup hits npc_path = tmp_path / "worlds" / "test" / "npcs" / "Vera.md" npc_path.parent.mkdir(parents=True, exist_ok=True) npc_path.write_text("# Vera\n\nA tavern owner.") engine._mcp.ctx.entity_index.register("Vera", npc_path)
result = engine._find_entity("Vera") assert result is not None assert result["name"] == "Vera" assert "tavern owner" in result["body"] assert result["entity_type"] == "npcs"
def test_find_entity_returns_none_when_missing(self, engine): assert engine._find_entity("Nobody") is None
def test_build_context_loads_location_and_one_hop_linked( self, engine, tmp_path: Path, ): """When the session points at a location, _build_context should load the location, then one-hop into entities the location wikilinks.""" from storied.session import save_session
(tmp_path / "players" / "default").mkdir(parents=True)
# Location wikilinks to a related NPC loc_path = tmp_path / "worlds" / "test" / "locations" / "Tavern.md" loc_path.parent.mkdir(parents=True, exist_ok=True) loc_path.write_text("# Tavern\n\nA cozy spot where [[Vera]] holds court.") engine._mcp.ctx.entity_index.register("Tavern", loc_path)
# The linked NPC npc_path = tmp_path / "worlds" / "test" / "npcs" / "Vera.md" npc_path.parent.mkdir(parents=True, exist_ok=True) npc_path.write_text("# Vera\n\nTavern owner.") engine._mcp.ctx.entity_index.register("Vera", npc_path)
save_session( "default", { "location": "Tavern", "body": "Player just walked in.", }, )
engine._build_context() # Location should be loaded assert "Location" in engine._context_parts # One-hop linked entity should be picked up via Linked: prefix linked_keys = [k for k in engine._context_parts if k.startswith("Linked:")] assert any("Vera" in k for k in linked_keys)
def test_build_context_includes_notifications(self, engine): from storied import notifications
notifications.append( engine.world_id, "World tick: Vera left the tavern", ) engine._build_context() assert "Notifications" in engine._context_parts assert "Vera left the tavern" in engine._context_parts["Notifications"]
def test_build_context_injects_initiative_when_active(self, engine): from storied.initiative import Combatant
engine._mcp.ctx.initiative.begin( [ Combatant(name="Goblin", initiative=10, hp=7, hp_max=7, ac=15), ] ) engine._build_context() assert "Initiative" in engine._context_parts assert "Goblin" in engine._context_parts["Initiative"]