diff --git a/pyproject.toml b/pyproject.toml index b27fe84..2beda2f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,7 +48,9 @@ addopts = [ "--cov=tests", "--cov-report=term-missing:skip-covered", "--cov-branch", + "-m", "not slow", ] +markers = ["slow: slow tests that build real search indices"] [tool.mypy] python_version = "3.12" @@ -64,6 +66,7 @@ select = ["E", "F", "I", "UP", "B", "SIM"] [tool.coverage.run] source = ["src/storied"] branch = true +omit = ["src/storied/cli.py", "src/storied/srd/*"] [tool.coverage.report] -fail_under = 70 +fail_under = 80 diff --git a/src/storied/search.py b/src/storied/search.py index 8bfebe6..c4f0bb9 100644 --- a/src/storied/search.py +++ b/src/storied/search.py @@ -17,7 +17,15 @@ from fastembed import TextEmbedding EMBED_MODEL = "BAAI/bge-small-en-v1.5" EMBED_DIM = 384 -CHUNK_CHAR_THRESHOLD = 4_000 +CHUNK_CHAR_THRESHOLD = 1_500 + +# Splitting patterns, tried in order from coarsest to finest. +_SPLIT_PATTERNS = [ + r"\n(?=## )", # ## headings + r"\n(?=### )", # ### headings + r"\n(?=#### )", # #### headings (class features, combat sub-topics) + r"\n(?=\*\*[A-Z])", # bold definitions (glossary entries, level features) +] @dataclass @@ -42,56 +50,68 @@ def age_decay(current_day: int, doc_day: int, half_life: int = 3) -> float: return 0.5 ** (age / half_life) +def _split_oversized( + sections: list[str], threshold: int, patterns: list[str], +) -> list[str]: + """Recursively split sections that exceed threshold using finer patterns.""" + if not patterns: + return sections + + pattern, *remaining = patterns + result: list[str] = [] + + for section in sections: + if len(section) <= threshold: + result.append(section) + continue + + parts = re.split(pattern, section) + if len(parts) <= 1: + result.extend(_split_oversized([section], threshold, remaining)) + else: + result.extend( + _split_oversized( + [p for p in parts if p.strip()], threshold, remaining, + ) + ) + + return result + + def chunk_document(path: Path, content: str) -> list[tuple[int, str]]: """Split a document into indexed chunks for embedding. - Small files return a single chunk. Larger files split on ## headers, - each prefixed with the document title for context. + Uses cascading splits (## → ### → #### → **Bold**) to break large + documents into chunks small enough for useful embeddings. Each chunk + gets the document title prepended for context. Returns list of (chunk_index, chunk_text) pairs. """ if not content.strip(): return [(0, path.stem)] - # Extract title from first # heading title_match = re.match(r"^#\s+(.+)", content) title = title_match.group(1).strip() if title_match else path.stem if len(content) < CHUNK_CHAR_THRESHOLD: return [(0, content)] - # Split on ## headers - sections = re.split(r"\n(?=## )", content) - - # If only one section (no ## headers), return as single chunk - if len(sections) <= 1: - return [(0, content)] + raw_sections = _split_oversized( + [content], CHUNK_CHAR_THRESHOLD, _SPLIT_PATTERNS, + ) chunks: list[tuple[int, str]] = [] - chunk_idx = 0 - - for section in sections: + for idx, section in enumerate(raw_sections): section = section.strip() if not section: continue - # Prepend title to non-first chunks for context - if chunk_idx > 0: + if idx > 0 and not section.startswith(f"# {title}"): text = f"# {title}\n\n{section}" else: text = section - # Sub-split oversized sections on ### headers - if len(text) > CHUNK_CHAR_THRESHOLD: - subsections = re.split(r"\n(?=### )", text) - for sub in subsections: - sub = sub.strip() - if sub: - chunks.append((chunk_idx, sub)) - chunk_idx += 1 - else: - chunks.append((chunk_idx, text)) - chunk_idx += 1 + chunks.append((idx, text)) return chunks if chunks else [(0, content)] @@ -110,7 +130,7 @@ def _default_embed(texts: list[str]) -> list[list[float]]: def _connect(db_path: str) -> sqlite3.Connection: """Open a SQLite connection with sqlite-vec loaded.""" - conn = sqlite3.connect(db_path) + conn = sqlite3.connect(db_path, check_same_thread=False) conn.enable_load_extension(True) sqlite_vec.load(conn) conn.enable_load_extension(False) diff --git a/tests/test_search.py b/tests/test_search.py index 7539d54..c1f62e7 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -99,6 +99,45 @@ class TestChunkDocument: assert len(chunks) == 1 assert "some content" in chunks[0][1] + def test_splits_on_h4_headings(self): + content = "# Rogue\n\n" + "\n".join( + f"#### **Feature {i}**\n\n{'x ' * 400}\n" for i in range(5) + ) + chunks = chunk_document(Path("rogue.md"), content) + assert len(chunks) >= 5 + + def test_splits_on_bold_definitions(self): + content = "# Rules Glossary\n\n## Definitions\n\n" + "\n".join( + f"**Term {i}**\nDefinition of term {i}. {'y ' * 400}\n" + for i in range(5) + ) + chunks = chunk_document(Path("glossary.md"), content) + assert len(chunks) >= 5 + + def test_class_features_split_by_level(self): + content = "# Fighter\n\n#### **Fighter Class Features**\n\n" + content += "**Level 1: Fighting Style**\n" + "Choose a style. " * 100 + "\n\n" + content += "**Level 1: Second Wind**\n" + "Heal yourself. " * 100 + "\n\n" + content += "**Level 2: Action Surge**\n" + "Extra action. " * 100 + "\n" + chunks = chunk_document(Path("fighter.md"), content) + texts = [t for _, t in chunks] + assert any("Second Wind" in t for t in texts) + assert any("Action Surge" in t for t in texts) + + def test_oversized_with_no_splittable_headings(self): + content = "# Blob\n\n" + "word " * 1_000 + chunks = chunk_document(Path("blob.md"), content) + assert len(chunks) >= 1 + assert "word" in chunks[0][1] + + def test_chunks_get_title_context(self): + content = "# Rogue\n\n" + "\n".join( + f"#### **Section {i}**\n\n{'z ' * 400}\n" for i in range(5) + ) + chunks = chunk_document(Path("rogue.md"), content) + for _, text in chunks[1:]: + assert text.startswith("# Rogue") + # --- Age Decay --- @@ -346,3 +385,30 @@ class TestSeedFrom: assert seeded.stats()["by_source"]["srd"] == 1 assert seeded.stats()["by_source"]["world"] == 1 seeded.close() + + +class TestThreadSafety: + """Tests for cross-thread access (MCP server runs on a background thread).""" + + def test_search_from_different_thread(self, index: VectorIndex): + index.upsert("world:npcs/vex.md:0", "Captain Vex, harbor master", + {"source": "world", "content_type": "npcs", + "path": "/tmp/vex.md", "title": "Captain Vex"}) + + import threading + + results: list[list[SearchHit]] = [] + error: list[Exception] = [] + + def search_on_thread(): + try: + results.append(index.search("harbor master")) + except Exception as e: + error.append(e) + + t = threading.Thread(target=search_on_thread) + t.start() + t.join() + + assert not error, f"Cross-thread search failed: {error[0]}" + assert len(results[0]) == 1 diff --git a/tests/test_srd_recall.py b/tests/test_srd_recall.py new file mode 100644 index 0000000..4ac0eee --- /dev/null +++ b/tests/test_srd_recall.py @@ -0,0 +1,121 @@ +"""Empirical tests: can recall find key SRD content? + +These test against the real SRD files to verify that chunking and search +produce useful results for the queries a DM actually needs. + +Slow (~60s) because they build a real embedding index. Run with: + pytest tests/test_srd_recall.py -v +""" + +from pathlib import Path + +import pytest + +from storied.search import VectorIndex + +SRD_DIR = Path("rules/srd-5.2.1/sections") + +pytestmark = pytest.mark.slow + + +@pytest.fixture(scope="module") +def srd_index(tmp_path_factory: pytest.TempPathFactory) -> VectorIndex: + """Build a real SRD index (once per module, reused across tests).""" + if not SRD_DIR.exists(): + pytest.skip("SRD not available") + + db_path = tmp_path_factory.mktemp("search") / "srd.db" + index = VectorIndex(db_path) + count = index.reindex_directory(SRD_DIR, source="srd") + assert count > 0, "No documents indexed" + yield index + index.close() + + +def _top_titles(index: VectorIndex, query: str, n: int = 5) -> list[str]: + """Return titles of top N results for a query.""" + hits = index.search(query, limit=n) + return [h.doc_id.split(":")[1] for h in hits] + + +def _assert_any_hit_contains( + index: VectorIndex, query: str, needle: str, top_n: int = 3, +): + """Assert that at least one of the top N results contains needle.""" + hits = index.search(query, limit=top_n) + snippets = [h.snippet for h in hits] + doc_ids = [h.doc_id for h in hits] + assert any( + needle.lower() in s.lower() for s in snippets + ) or any( + needle.lower() in d.lower() for d in doc_ids + ), ( + f"'{needle}' not found in top {top_n} for query '{query}'.\n" + f"Got: {list(zip(doc_ids, [s[:80] for s in snippets]))}" + ) + + +# --- Class Features --- + +class TestClassFeatureRecall: + + def test_sneak_attack(self, srd_index: VectorIndex): + _assert_any_hit_contains(srd_index, "sneak attack", "Sneak Attack") + + def test_second_wind(self, srd_index: VectorIndex): + _assert_any_hit_contains(srd_index, "second wind", "Second Wind") + + def test_action_surge(self, srd_index: VectorIndex): + _assert_any_hit_contains(srd_index, "action surge", "Action Surge") + + def test_cunning_action(self, srd_index: VectorIndex): + _assert_any_hit_contains(srd_index, "cunning action", "Cunning Action") + + def test_uncanny_dodge(self, srd_index: VectorIndex): + _assert_any_hit_contains(srd_index, "uncanny dodge", "Uncanny Dodge") + + def test_wild_shape(self, srd_index: VectorIndex): + _assert_any_hit_contains(srd_index, "wild shape", "Wild Shape") + + def test_lay_on_hands(self, srd_index: VectorIndex): + _assert_any_hit_contains(srd_index, "lay on hands", "Lay on Hands") + + +# --- Conditions --- + +class TestConditionRecall: + + def test_grappled(self, srd_index: VectorIndex): + _assert_any_hit_contains(srd_index, "grappled condition", "Grappled") + + def test_frightened(self, srd_index: VectorIndex): + _assert_any_hit_contains(srd_index, "frightened condition", "Frightened") + + def test_prone(self, srd_index: VectorIndex): + _assert_any_hit_contains(srd_index, "prone condition", "Prone") + + def test_paralyzed(self, srd_index: VectorIndex): + _assert_any_hit_contains(srd_index, "paralyzed condition", "Paralyzed") + + +# --- Core Mechanics --- + +class TestCoreMechanicsRecall: + + def test_opportunity_attack(self, srd_index: VectorIndex): + _assert_any_hit_contains( + srd_index, "opportunity attack", "Opportunity Attack", + ) + + def test_death_saving_throw(self, srd_index: VectorIndex): + _assert_any_hit_contains( + srd_index, "death saving throw", "Death", + ) + + def test_short_rest(self, srd_index: VectorIndex): + _assert_any_hit_contains(srd_index, "short rest", "Short Rest") + + def test_concentration(self, srd_index: VectorIndex): + _assert_any_hit_contains( + srd_index, "concentration spell", "Concentration", + )