Add entity path index and write-through cache for faster tool calls

+14 -1

src/storied/engine.py

··· 297 297 return resolver.load(name, content_type=content_type) 298 298 299 299 def _find_entity(self, name: str) -> dict | None: 300 - """Find an entity by name, searching entity directories in priority order.""" 300 + """Find an entity by name, using the entity index for O(1) lookups.""" 301 301 if not self.world_id: 302 302 return None 303 303 304 + # Fast path: use the entity index from the MCP server 305 + entity_index = self._mcp.entity_index 306 + if entity_index: 307 + path = entity_index.resolve(name) 308 + if path and path.exists(): 309 + return { 310 + "name": name, 311 + "body": path.read_text(), 312 + "entity_type": path.parent.name, 313 + } 314 + return None 315 + 316 + # Fallback: filesystem scan 304 317 entity = load_entity_content(name, self.world_id, self.base_path) 305 318 if entity: 306 319 return {

+11 -2

src/storied/mcp_server.py

··· 19 19 20 20 from storied.log import CampaignLog 21 21 from storied.tools import ( 22 + EntityIndex, 22 23 PLANNER_TOOL_DEFINITIONS, 23 24 SEEDER_TOOL_DEFINITIONS, 24 25 TOOL_DEFINITIONS, ··· 59 60 class MCPServerHandle: 60 61 """Handle to a running in-process MCP server.""" 61 62 62 - def __init__(self, port: int, thread: threading.Thread): 63 + def __init__( 64 + self, port: int, thread: threading.Thread, 65 + entity_index: EntityIndex | None = None, 66 + ): 63 67 self.port = port 64 68 self.url = f"http://127.0.0.1:{port}/sse" 69 + self.entity_index = entity_index 65 70 self._thread = thread 66 71 67 72 def stop(self) -> None: ··· 85 90 if campaign_log is None: 86 91 campaign_log = CampaignLog(world_id, base_path) 87 92 93 + world_dir = base_path / "worlds" / world_id 94 + entity_index = EntityIndex(world_dir) 95 + 88 96 definitions = TOOL_SETS.get(tool_set, TOOL_DEFINITIONS) 89 97 executor = EXECUTORS.get(tool_set, execute_tool) 90 98 mcp_tools = [_to_mcp_tool(d) for d in definitions] ··· 104 112 player_id=player_id, 105 113 base_path=base_path, 106 114 campaign_log=campaign_log, 115 + entity_index=entity_index, 107 116 ) 108 117 return [TextContent(type="text", text=str(result))] 109 118 ··· 143 152 except OSError: 144 153 time.sleep(0.1) 145 154 146 - return MCPServerHandle(port, thread) 155 + return MCPServerHandle(port, thread, entity_index=entity_index)

+115 -25

src/storied/tools.py

··· 10 10 11 11 import yaml 12 12 13 + from storied.character import create_character as char_create 14 + from storied.character import update_character as char_update 15 + from storied.content import ContentResolver 16 + from storied.dice import roll as dice_roll 17 + from storied.log import CampaignLog, load_log 18 + from storied.session import name_to_slug 19 + from storied.session import update_session as session_update 20 + 13 21 # Per-file locks for thread-safe entity writes (establish, mark) 14 22 _file_locks: dict[Path, threading.Lock] = {} 15 23 _file_locks_lock = threading.Lock() ··· 22 30 _file_locks[path] = threading.Lock() 23 31 return _file_locks[path] 24 32 25 - from storied.character import create_character as char_create 26 - from storied.character import update_character as char_update 27 - from storied.content import ContentResolver 28 - from storied.dice import roll as dice_roll 29 - from storied.log import CampaignLog, load_log 30 - from storied.session import name_to_slug 31 - from storied.session import update_session as session_update 33 + 34 + class EntityIndex: 35 + """Name→path index with write-through entity cache. 36 + 37 + Built once at startup by globbing the world directory. Lookups are 38 + dict lookups instead of 6 sequential Path.exists() calls. The cache 39 + stores parsed entity dicts so repeated loads within a turn skip disk I/O. 40 + establish() and mark() update both the index and cache on write. 41 + """ 42 + 43 + def __init__(self, world_dir: Path | None = None): 44 + self._paths: dict[str, Path] = {} 45 + self._cache: dict[Path, dict] = {} 46 + if world_dir and world_dir.exists(): 47 + for md in world_dir.rglob("*.md"): 48 + self._paths[md.stem] = md 49 + 50 + def resolve(self, name: str) -> Path | None: 51 + """Look up an entity's file path by name.""" 52 + return self._paths.get(name) 53 + 54 + def register(self, name: str, path: Path) -> None: 55 + """Register or update an entity's path in the index.""" 56 + self._paths[name] = path 57 + 58 + def cache_get(self, path: Path) -> dict | None: 59 + """Get a cached parsed entity, or None if not cached.""" 60 + return self._cache.get(path) 61 + 62 + def cache_put(self, path: Path, data: dict) -> None: 63 + """Store a parsed entity in the cache.""" 64 + self._cache[path] = data 32 65 33 66 34 67 def roll(notation: str, reason: str | None = None) -> dict: ··· 240 273 base_path: Path | None = None, 241 274 campaign_log: CampaignLog | None = None, 242 275 world_id: str = "default", 276 + entity_index: EntityIndex | None = None, 243 277 ) -> str: 244 278 """Call this after every response. Logs what happened, advances the clock, 245 279 and updates the scene state. ··· 297 331 if event and present and world_id: 298 332 marked = _auto_mark_present( 299 333 present, event, world_id, base_path, campaign_log, 334 + entity_index=entity_index, 300 335 ) 301 336 if marked: 302 337 parts.append(f"Auto-marked: {', '.join(marked)}") ··· 314 349 will: list[str] | None = None, 315 350 world_id: str | None = None, 316 351 base_path: Path | None = None, 352 + entity_index: EntityIndex | None = None, 317 353 ) -> str: 318 354 """Establish or update an entity in the world. 319 355 ··· 361 397 lock = _get_file_lock(file_path) 362 398 with lock: 363 399 # Load existing content if file exists (for partial updates) 364 - existing = _load_entity(file_path) 400 + existing = _load_entity(file_path, entity_index) 365 401 366 402 # Merge with existing content (new values override) 367 403 if description is None: ··· 380 416 file_content = _format_entity(name, description, location, knows, wants, will, was) 381 417 file_path.write_text(file_content) 382 418 419 + # Write-through: update index and cache 420 + updated_data = { 421 + "description": description, "location": location, 422 + "knows": knows, "wants": wants, "will": will, "was": was, 423 + } 424 + if entity_index: 425 + entity_index.register(name, file_path) 426 + entity_index.cache_put(file_path, updated_data) 427 + 383 428 action = "Updated" if existing else "Established" 384 429 return f"{action} {entity_type.rstrip('s')} '{name}'" 385 430 386 431 387 - def _load_entity(file_path: Path) -> dict: 432 + def _load_entity( 433 + file_path: Path, entity_index: EntityIndex | None = None, 434 + ) -> dict: 388 435 """Load an existing entity file and parse its structure.""" 436 + if entity_index: 437 + cached = entity_index.cache_get(file_path) 438 + if cached is not None: 439 + return cached 440 + 389 441 if not file_path.exists(): 390 442 return {} 391 443 ··· 430 482 if was_match: 431 483 result["was"] = _parse_list_items(was_match.group(1)) 432 484 485 + if entity_index: 486 + entity_index.cache_put(file_path, result) 487 + 433 488 return result 434 489 435 490 ··· 499 554 world_id: str, 500 555 base_path: Path | None, 501 556 campaign_log: CampaignLog | None, 557 + entity_index: EntityIndex | None = None, 502 558 ) -> list[str]: 503 559 """Auto-mark present entities with the current event. 504 560 ··· 516 572 continue 517 573 name = link_match.group(1) 518 574 519 - # Try each entity type directory 520 - for etype in ("npcs", "locations", "items", "factions"): 521 - file_path = base_path / "worlds" / world_id / etype / f"{name}.md" 522 - if file_path.exists(): 523 - mark( 524 - entity_type=etype, 525 - name=name, 526 - event=event, 527 - world_id=world_id, 528 - base_path=base_path, 529 - campaign_log=campaign_log, 530 - ) 531 - marked.append(name) 532 - break 575 + # Resolve via index (O(1)) or fall back to directory scan 576 + file_path = entity_index.resolve(name) if entity_index else None 577 + if file_path is None: 578 + for etype in ("npcs", "locations", "items", "factions"): 579 + candidate = base_path / "worlds" / world_id / etype / f"{name}.md" 580 + if candidate.exists(): 581 + file_path = candidate 582 + break 583 + 584 + if file_path and file_path.exists(): 585 + entity_type = file_path.parent.name 586 + mark( 587 + entity_type=entity_type, 588 + name=name, 589 + event=event, 590 + world_id=world_id, 591 + base_path=base_path, 592 + campaign_log=campaign_log, 593 + entity_index=entity_index, 594 + ) 595 + marked.append(name) 533 596 534 597 return marked 535 598 ··· 542 605 world_id: str | None = None, 543 606 base_path: Path | None = None, 544 607 campaign_log: CampaignLog | None = None, 608 + entity_index: EntityIndex | None = None, 545 609 ) -> str: 546 610 """Record an event in an entity's history (## Was section). 547 611 ··· 570 634 if base_path is None: 571 635 base_path = Path.cwd() 572 636 573 - file_path = base_path / "worlds" / world_id / entity_type / f"{name}.md" 637 + # Resolve path via index or filesystem 638 + if entity_index: 639 + file_path = entity_index.resolve(name) 640 + if file_path is None: 641 + file_path = base_path / "worlds" / world_id / entity_type / f"{name}.md" 642 + else: 643 + file_path = base_path / "worlds" / world_id / entity_type / f"{name}.md" 644 + 574 645 if not file_path.exists(): 575 646 return f"Error: Entity '{name}' not found in {entity_type}" 576 647 ··· 584 655 lock = _get_file_lock(file_path) 585 656 with lock: 586 657 # Load existing entity 587 - existing = _load_entity(file_path) 658 + existing = _load_entity(file_path, entity_index) 588 659 589 660 # Append to Was section 590 661 was = existing.get("was", []) ··· 609 680 was, 610 681 ) 611 682 file_path.write_text(file_content) 683 + 684 + # Write-through cache 685 + if entity_index: 686 + entity_index.cache_put(file_path, { 687 + "description": existing.get("description", ""), 688 + "location": existing.get("location", ""), 689 + "knows": existing.get("knows", []), 690 + "wants": existing.get("wants", []), 691 + "will": will, 692 + "was": was, 693 + }) 612 694 613 695 result = f"Marked: {event}" 614 696 if resolved: ··· 985 1067 player_id: str = "default", 986 1068 base_path: Path | None = None, 987 1069 campaign_log: CampaignLog | None = None, 1070 + entity_index: EntityIndex | None = None, 988 1071 ) -> str: 989 1072 """Execute a tool by name with the given input. 990 1073 ··· 1060 1143 base_path=base_path, 1061 1144 campaign_log=campaign_log, 1062 1145 world_id=world_id or "default", 1146 + entity_index=entity_index, 1063 1147 ) 1064 1148 1065 1149 elif tool_name == "establish": ··· 1073 1157 will=tool_input.get("will"), 1074 1158 world_id=world_id, 1075 1159 base_path=base_path, 1160 + entity_index=entity_index, 1076 1161 ) 1077 1162 1078 1163 elif tool_name == "mark": ··· 1084 1169 world_id=world_id, 1085 1170 base_path=base_path, 1086 1171 campaign_log=campaign_log, 1172 + entity_index=entity_index, 1087 1173 ) 1088 1174 1089 1175 elif tool_name == "note_discovery": ··· 1120 1206 world_id: str | None = None, 1121 1207 base_path: Path | None = None, 1122 1208 campaign_log: CampaignLog | None = None, 1209 + entity_index: EntityIndex | None = None, 1123 1210 ) -> str: 1124 1211 """Execute a planner-allowed tool. Rejects anything outside the allowed set.""" 1125 1212 if tool_name not in PLANNER_TOOLS: ··· 1130 1217 world_id=world_id, 1131 1218 base_path=base_path, 1132 1219 campaign_log=campaign_log, 1220 + entity_index=entity_index, 1133 1221 ) 1134 1222 1135 1223 ··· 1145 1233 player_id: str | None = None, 1146 1234 base_path: Path | None = None, 1147 1235 campaign_log: CampaignLog | None = None, 1236 + entity_index: EntityIndex | None = None, 1148 1237 ) -> str: 1149 1238 """Execute a seeder-allowed tool. Rejects anything outside the allowed set.""" 1150 1239 if tool_name not in SEEDER_TOOLS: ··· 1156 1245 player_id=player_id or "default", 1157 1246 base_path=base_path, 1158 1247 campaign_log=campaign_log, 1248 + entity_index=entity_index, 1159 1249 )

+118 -1

tests/test_entities.py

··· 10 10 load_entity_content, 11 11 resolve_wiki_link, 12 12 ) 13 - from storied.tools import establish, mark 13 + from storied.tools import EntityIndex, establish, mark 14 14 15 15 16 16 @pytest.fixture ··· 452 452 def test_load_entity_content_not_found(self, world_base: Path): 453 453 entity = load_entity_content("Nobody", "test-world", world_base) 454 454 assert entity is None 455 + 456 + 457 + # ── EntityIndex ────────────────────────────────────────────────────────── 458 + 459 + 460 + @pytest.fixture 461 + def indexed_world(world_base: Path) -> tuple[Path, EntityIndex]: 462 + """Create a world with entities and build an index.""" 463 + world_dir = world_base / "worlds" / "test-world" 464 + for etype, name in [("npcs", "Vera"), ("locations", "Tavern"), ("items", "Sword")]: 465 + d = world_dir / etype 466 + d.mkdir(parents=True, exist_ok=True) 467 + (d / f"{name}.md").write_text(f"# {name}\n\n## Is\n\nA {etype[:-1]}.\n") 468 + return world_dir, EntityIndex(world_dir) 469 + 470 + 471 + class TestEntityIndex: 472 + """Tests for the entity path index and write-through cache.""" 473 + 474 + def test_resolve_existing(self, indexed_world: tuple[Path, EntityIndex]): 475 + _, index = indexed_world 476 + path = index.resolve("Vera") 477 + assert path is not None 478 + assert path.name == "Vera.md" 479 + assert "npcs" in str(path) 480 + 481 + def test_resolve_missing(self, indexed_world: tuple[Path, EntityIndex]): 482 + _, index = indexed_world 483 + assert index.resolve("Nobody") is None 484 + 485 + def test_register_new_entity(self, indexed_world: tuple[Path, EntityIndex]): 486 + world_dir, index = indexed_world 487 + new_path = world_dir / "npcs" / "Henrik.md" 488 + index.register("Henrik", new_path) 489 + assert index.resolve("Henrik") == new_path 490 + 491 + def test_cache_miss_returns_none(self, indexed_world: tuple[Path, EntityIndex]): 492 + world_dir, index = indexed_world 493 + assert index.cache_get(world_dir / "npcs" / "Vera.md") is None 494 + 495 + def test_cache_put_and_get(self, indexed_world: tuple[Path, EntityIndex]): 496 + world_dir, index = indexed_world 497 + path = world_dir / "npcs" / "Vera.md" 498 + data = {"description": "A tavern keeper.", "knows": ["a secret"]} 499 + index.cache_put(path, data) 500 + assert index.cache_get(path) == data 501 + 502 + def test_cache_survives_register(self, indexed_world: tuple[Path, EntityIndex]): 503 + world_dir, index = indexed_world 504 + path = world_dir / "npcs" / "Vera.md" 505 + data = {"description": "Vera"} 506 + index.cache_put(path, data) 507 + index.register("Vera", path) 508 + assert index.cache_get(path) == data 509 + 510 + def test_empty_world_dir(self, tmp_path: Path): 511 + world_dir = tmp_path / "worlds" / "empty" 512 + world_dir.mkdir(parents=True) 513 + index = EntityIndex(world_dir) 514 + assert index.resolve("anything") is None 515 + 516 + def test_none_world_dir(self): 517 + index = EntityIndex(None) 518 + assert index.resolve("anything") is None 519 + 520 + 521 + class TestEstablishWithIndex: 522 + """Tests that establish writes through to the index and cache.""" 523 + 524 + def test_establish_registers_in_index(self, world_base: Path): 525 + world_dir = world_base / "worlds" / "test-world" 526 + index = EntityIndex(world_dir) 527 + establish( 528 + entity_type="npcs", name="New NPC", 529 + description="A stranger.", 530 + world_id="test-world", base_path=world_base, 531 + entity_index=index, 532 + ) 533 + assert index.resolve("New NPC") is not None 534 + 535 + def test_establish_caches_entity(self, world_base: Path): 536 + world_dir = world_base / "worlds" / "test-world" 537 + index = EntityIndex(world_dir) 538 + establish( 539 + entity_type="npcs", name="Cached NPC", 540 + description="Cached.", 541 + knows=["a secret"], 542 + world_id="test-world", base_path=world_base, 543 + entity_index=index, 544 + ) 545 + path = index.resolve("Cached NPC") 546 + cached = index.cache_get(path) 547 + assert cached is not None 548 + assert cached["description"] == "Cached." 549 + assert cached["knows"] == ["a secret"] 550 + 551 + def test_mark_updates_cache( 552 + self, world_base: Path, campaign_log: CampaignLog, 553 + ): 554 + world_dir = world_base / "worlds" / "test-world" 555 + index = EntityIndex(world_dir) 556 + establish( 557 + entity_type="npcs", name="Markable", 558 + description="An NPC.", 559 + world_id="test-world", base_path=world_base, 560 + entity_index=index, 561 + ) 562 + mark( 563 + entity_type="npcs", name="Markable", 564 + event="Something happened", 565 + world_id="test-world", base_path=world_base, 566 + campaign_log=campaign_log, 567 + entity_index=index, 568 + ) 569 + path = index.resolve("Markable") 570 + cached = index.cache_get(path) 571 + assert any("Something happened" in w for w in cached["was"])

Configure Feed

Configure Feed