wiki-kb/tests/test_search.py

"""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"