arborist/tests/test_search_fts5.py

"""Unit tests for progressive-AND fallback in arborist.search.fts5."""

from __future__ import annotations

from typing import Iterator

from arborist.document import Document
from arborist.ingest import ingest_source
from arborist.search import FTS5Backend
from arborist.search.fts5 import _progressive_and_token_chains
from arborist.source import Source
from arborist.store import connect


class _FakeSource(Source):
    source_type = "fake"

    def __init__(self, docs: list[Document]):
        self.docs = docs

    def iter_documents(self) -> Iterator[Document]:
        yield from self.docs


def _doc(uri: str, content: str, title: str | None = None) -> Document:
    return Document(
        uri=uri,
        content=content,
        source_type="fake",
        title=title if title is not None else uri.rsplit("/", 1)[-1],
        edges=[],
    )


# ----------------------------------------------------------- helper unit tests

def test_progressive_chains_single_token_passthrough():
    assert _progressive_and_token_chains(["alpha"]) == [["alpha"]]


def test_progressive_chains_empty_input():
    assert _progressive_and_token_chains([]) == [[]]


def test_progressive_chains_drops_shortest_first():
    # Lengths: Gundremmingen=13, located=7, Bavaria=7. Tie between
    # "located" and "Bavaria" — input order ("located" first) breaks it.
    chains = _progressive_and_token_chains(["Gundremmingen", "located", "Bavaria"])
    assert chains == [
        ["Gundremmingen", "located", "Bavaria"],
        ["Gundremmingen", "Bavaria"],
        ["Gundremmingen"],
    ]


def test_progressive_chains_strict_length_ordering():
    # Pure length sort, no ties.
    chains = _progressive_and_token_chains(["aaa", "bb", "ccccc", "dddd"])
    assert chains == [
        ["aaa", "bb", "ccccc", "dddd"],
        ["aaa", "ccccc", "dddd"],   # "bb" (shortest) dropped first
        ["ccccc", "dddd"],          # "aaa" dropped second
        ["ccccc"],                  # "dddd" dropped third
    ]


def test_progressive_chains_always_keeps_one_token():
    # With N tokens, we yield N chains: full, full-1, full-2, ..., 1.
    chains = _progressive_and_token_chains(["x", "yy", "zzz"])
    assert len(chains) == 3
    assert all(c for c in chains)  # no empty chain
    assert len(chains[-1]) == 1


# ----------------------------------------------------------- end-to-end search

def _ingest(tmp_path, docs: list[Document]):
    db_path = tmp_path / "test.db"
    conn = connect(db_path)
    ingest_source(conn, _FakeSource(docs))
    return conn


def test_search_progressive_and_drops_high_df_token(tmp_path):
    """The "located" failure mode: full AND zero-hits; drop "located"
    and the topical AND ("Gundremmingen AND Bavaria") finds the answer.
    The old straight-to-OR fallback would have included every doc with
    "located", "Bavaria", or "Gundremmingen"; progressive-AND keeps the
    intersection tight.
    """
    docs = [
        _doc("test://gundremmingen-bavaria",
             "Gundremmingen is a town in Bavaria.",
             title="Gundremmingen"),
        _doc("test://bavaria",
             "Bavaria is a state in Germany.",
             title="Bavaria"),
        # Many decoy docs that contain "located" but are off-topic.
        # Old OR-mode would surface these alongside the real answer.
        *[_doc(f"test://decoy-{i}",
               f"The widget {i} is located on the third shelf.",
               title=f"Widget {i}")
          for i in range(20)],
    ]
    conn = _ingest(tmp_path, docs)
    try:
        backend = FTS5Backend(conn)
        hits = backend.search("where is Gundremmingen located in Bavaria?", limit=8)
        assert hits, "expected progressive-AND to surface the topical doc"
        # Top hit must be the topical doc, not a decoy.
        assert hits[0].title == "Gundremmingen"
        # Decoys ("located" but not "Gundremmingen") must not appear at all,
        # because progressive-AND dropping "located" still requires
        # "Gundremmingen AND Bavaria".
        assert all("Widget" not in h.title for h in hits)
    finally:
        conn.close()


def test_search_full_and_succeeds_no_progressive_drop(tmp_path):
    """When the full AND already returns hits, progressive-AND
    is a no-op — the topical chain wins on chain 0.
    """
    docs = [
        _doc("test://a", "alpha beta gamma delta", title="A"),
        _doc("test://b", "alpha beta gamma", title="B"),
        _doc("test://c", "alpha", title="C"),
    ]
    conn = _ingest(tmp_path, docs)
    try:
        backend = FTS5Backend(conn)
        hits = backend.search("alpha beta gamma delta")
        # Only doc A contains all four; full AND should return exactly it.
        assert len(hits) == 1
        assert hits[0].title == "A"
    finally:
        conn.close()


def test_search_falls_back_to_or_when_no_chain_hits(tmp_path):
    """If even single-token AND returns zero, OR-mode fallback fires.
    OR can match any one of the synonyms / tokens.
    """
    docs = [
        _doc("test://a", "alpha is one fact", title="Alpha"),
        _doc("test://b", "beta is another fact", title="Beta"),
    ]
    conn = _ingest(tmp_path, docs)
    try:
        backend = FTS5Backend(conn)
        # No doc contains "zeta" — every AND chain (including single-token
        # "zeta") returns zero. OR-mode with synonym "alpha" fires.
        hits = backend.search("zeta", extra_or_tokens={"alpha"})
        assert hits, "OR-mode synonym fallback should still find alpha"
        assert hits[0].title == "Alpha"
    finally:
        conn.close()


def test_search_or_fallback_when_synonym_not_in_corpus(tmp_path):
    """OR-mode with no synonyms and no matching tokens returns empty —
    not an exception.
    """
    docs = [_doc("test://a", "completely unrelated content", title="A")]
    conn = _ingest(tmp_path, docs)
    try:
        backend = FTS5Backend(conn)
        hits = backend.search("Gundremmingen")
        assert hits == []
    finally:
        conn.close()


# ----------------------------------------------------------- B: DF filter at OR

def test_or_pool_df_filter_drops_high_df_token(tmp_path, monkeypatch):
    """When OR-mode fires and one of the tokens is high-DF, the DF
    filter drops it before MATCH so it doesn't dominate the candidate
    set. Synonym path: no AND-able tokens exist; OR-fallback fires.
    """
    from arborist.search import fts5 as _fts5_mod
    monkeypatch.setattr(_fts5_mod, "_OR_FALLBACK_MAX_TOKEN_DF", 5)

    # Build a corpus where "common" appears in many docs (above
    # threshold of 5) and "rare" appears in few.
    docs: list[Document] = []
    for i in range(15):
        docs.append(_doc(f"test://common-{i}",
                         f"this document contains common token number {i}",
                         title=f"common-doc-{i}"))
    docs.append(_doc("test://rare",
                     "rare needle is buried somewhere",
                     title="needle-doc"))
    conn = _ingest(tmp_path, docs)
    try:
        backend = FTS5Backend(conn)
        # Query has no AND-success path because "zeta" doesn't exist —
        # progressive-AND will hit zero, OR fires. extra_or_tokens
        # carries both "common" (DF=15, >threshold) and "rare" (DF=1).
        hits = backend.search(
            "zeta",
            extra_or_tokens={"common", "rare"},
        )
        assert hits, "OR-mode should have surfaced the rare needle"
        # The needle doc must be top — "common" got DF-filtered out so
        # only "rare" actually matched.
        assert hits[0].title == "needle-doc"
        # No common-doc-N should appear: filter dropped "common".
        assert all(not h.title.startswith("common-doc-") for h in hits)
    finally:
        conn.close()


def test_or_pool_df_filter_keeps_input_when_all_above_threshold(tmp_path, monkeypatch):
    """If every candidate is above threshold the filter must NOT empty
    the pool — we'd rather have a slow OR than zero hits. The keep-all
    fallback fires.
    """
    from arborist.search import fts5 as _fts5_mod
    monkeypatch.setattr(_fts5_mod, "_OR_FALLBACK_MAX_TOKEN_DF", 0)  # nothing passes

    docs = [_doc(f"test://d{i}", "alpha bravo charlie", title=f"d{i}")
            for i in range(3)]
    conn = _ingest(tmp_path, docs)
    try:
        backend = FTS5Backend(conn)
        # extra_or_tokens carries content tokens; under threshold=0 all
        # candidates would be dropped, but the fallback keeps them so
        # OR-mode still finds the docs.
        hits = backend.search("zeta", extra_or_tokens={"alpha"})
        assert hits, "all-above-threshold case must keep tokens, not empty pool"
    finally:
        conn.close()