"""Tests for hybrid search.""" from __future__ import annotations import numpy as np import pytest from wikikb import config, embed, indexer, search _TRAITS = """--- title: "Traits" url: https://x/Traits revid: 1 categories: [] --- # Traits Genius gives plus five diplomacy. Domain limit rises with stewardship. """ _ARMY = """--- title: "Army" url: https://x/Army revid: 2 categories: [] --- # Army Levies and knights fight battles. """ @pytest.fixture def indexed(tmp_kb, fake_model): (config.CORPUS_DIR / "traits.md").write_text(_TRAITS, encoding="utf-8") (config.CORPUS_DIR / "army.md").write_text(_ARMY, encoding="utf-8") indexer.build_index( embed_fn=lambda texts: embed.embed_passages(texts, model=fake_model), log=lambda _: None, ) return fake_model def test_fts_match_query_ors_terms() -> None: assert search._fts_match_query("Domain limit") == '"domain" OR "limit"' assert search._fts_match_query("!!!") == "" def test_reciprocal_rank_fusion_orders_and_dedupes() -> None: fused = search.reciprocal_rank_fusion([[1, 2, 3], [2, 1]]) assert set(fused) == {1, 2, 3} assert fused[-1] == 3 # only appears once, at rank 2 in one list def test_vector_ids_empty() -> None: empty = np.zeros((0, config.EMBED_DIM), dtype=np.float32) query = np.zeros(config.EMBED_DIM, dtype=np.float32) assert search._vector_ids(empty, query, 5, None) == [] def test_search_ranks_relevant_page_first(indexed, fake_model) -> None: results = search.search("domain limit stewardship", model=fake_model) assert results assert results[0].title == "Traits" assert "stewardship" in results[0].text.lower() def test_search_no_fts_match_still_returns_vector_hits(indexed, fake_model) -> None: results = search.search("zzznonexistentxyz", model=fake_model) assert isinstance(results, list) assert results # vector retrieval still yields candidates def test_search_punctuation_query_skips_fts(indexed, fake_model) -> None: # punctuation-only query -> empty FTS match -> _fts_ids returns []; vector still runs results = search.search("!!!", model=fake_model) assert isinstance(results, list) def test_search_empty_index_returns_nothing(tmp_kb, fake_model) -> None: indexer.build_index( embed_fn=lambda texts: embed.embed_passages(texts, model=fake_model), log=lambda _: None, ) assert search.search("anything", model=fake_model) == [] def test_format_empty() -> None: assert search._format([]) == "(no results)" def test_format_truncates_long_text() -> None: result = search.Result( title="T", url="u", source="mediawiki", variant="kcd2", section="S", text="x" * 600, score=1.0, ) out = search._format([result]) assert "…" in out assert "[kcd2] T > S" in out def test_main_requires_query() -> None: assert search.main([]) == 2 def test_main_runs(monkeypatch, capsys) -> None: monkeypatch.setattr( search, "search", lambda *_a, **_k: [search.Result("T", "u", "mediawiki", "kcd2", "S", "body", 1.0)], ) assert search.main(["domain", "limit"]) == 0 assert "[kcd2] T > S" in capsys.readouterr().out def test_vector_ids_filters_before_truncating() -> None: # Three vectors ranked 3,2,1 by similarity; only rowid 1 (the WORST match) is in # scope. Truncating to one candidate first returns nothing; filtering during # selection returns the in-scope hit. vectors = np.array([[1.0], [2.0], [3.0]], dtype=np.float32) query = np.array([1.0], dtype=np.float32) assert search._vector_ids(vectors, query, 1, {1}) == [1] assert search._vector_ids(vectors, query, 1, None) == [3] def _variant_doc(title: str, name: str, body: str) -> str: return ( f'---\ntitle: "{title}"\nurl: https://x/{title}\nrevid: 1\n' f"source: mediawiki\nvariant: {name}\ncategories: []\n---\n\n# {title}\n\n{body}\n" ) @pytest.fixture def indexed_variants(tmp_kb, fake_model, monkeypatch): monkeypatch.setattr(config, "VARIANTS", (("kcd2", "KCD2"), ("kcd1", "KCD"))) (config.CORPUS_DIR / "new.md").write_text( _variant_doc("Alchemy KCD2", "kcd2", "Brewing potions with herbs."), encoding="utf-8" ) (config.CORPUS_DIR / "old.md").write_text( _variant_doc("Alchemy", "kcd1", "Brewing potions with herbs in the first game."), encoding="utf-8", ) (config.CORPUS_DIR / "shared.md").write_text( _variant_doc("Moravia", "unknown", "Brewing herbs region shared by both."), encoding="utf-8", ) indexer.build_index( embed_fn=lambda texts: embed.embed_passages(texts, model=fake_model), log=lambda _: None ) return fake_model def test_search_variant_scope_excludes_the_other_edition(indexed_variants, fake_model) -> None: titles = { r.title for r in search.search("brewing potions herbs", k=10, variant="kcd2", model=fake_model) } assert "Alchemy KCD2" in titles assert "Moravia" in titles # unknown is always in scope assert "Alchemy" not in titles def test_search_any_scope_returns_every_edition(indexed_variants, fake_model) -> None: titles = { r.title for r in search.search("brewing potions herbs", k=10, variant="any", model=fake_model) } assert {"Alchemy KCD2", "Alchemy", "Moravia"} <= titles def test_search_results_carry_source_and_variant(indexed_variants, fake_model) -> None: results = search.search("brewing potions herbs", k=10, variant="any", model=fake_model) assert {r.source for r in results} == {"mediawiki"} assert {r.title: r.variant for r in results}["Alchemy"] == "kcd1"