llm-ingestion-okf/tests/test_bm25_ranking.py

"""The BM25 ranking (v1.1 order C, C1): the default reading of `okf consume`.

Each test states one property of the mechanism on a synthetic bundle small
enough to reason about by hand. The measurement against a real collection is
the search gate's (`tools/okf_soek_gate.py`), not this file's.
"""

from __future__ import annotations

from pathlib import Path

from llm_ingestion_okf import bm25, consume


def _concept(concept_id: str, title: str, body: str) -> consume.Concept:
    return consume.Concept(
        path=Path(f"{concept_id}.md"),
        concept_id=concept_id,
        bundle_id="b",
        bundle_id_inherited=True,
        sha256="0" * 64,
        okf_type="concept",
        title=title,
        source_file=f"{concept_id.split('/')[0]}.md",
        adjudication="unknown",
        adjudication_present=False,
        req_number="",
        sources=(),
        sources_present=False,
        locators={},
        frontmatter={},
        body=body,
    )


def test_the_tokeniser_drops_stopwords_and_single_characters_and_folds_case() -> None:
    assert bm25.tokens("What IS the Default model for a teammate?") == [
        "default",
        "model",
        "teammate",
    ]


def test_the_tokeniser_keeps_a_norwegian_word_whole() -> None:
    # A letter outside ASCII must not split a word into fragments that can
    # match something unrelated in an English collection.
    assert bm25.tokens("første") == ["første"]


def test_a_word_the_collection_does_not_hold_lifts_nothing() -> None:
    concepts = [
        _concept("a/one", "One", "the alpha feature is described here"),
        _concept("b/two", "Two", "the beta feature is described here"),
    ]
    plain = bm25.rank(concepts, "alpha feature")
    padded = bm25.rank(concepts, "alpha feature zzqqxx")
    assert [c.concept_id for c, _, _ in plain.ranked] == [c.concept_id for c, _, _ in padded.ranked]
    assert [score for _, score, _ in plain.ranked] == [score for _, score, _ in padded.ranked]


def test_length_normalisation_prefers_the_short_concept_on_one_shared_term() -> None:
    filler = " ".join(f"word{i}" for i in range(400))
    concepts = [
        _concept("a/long", "Long", f"rotation {filler}"),
        _concept("b/short", "Short", "rotation of keys"),
    ]
    ranked = bm25.rank(concepts, "rotation").ranked
    assert ranked[0][0].concept_id == "b/short"


def test_the_title_field_separates_two_equal_bodies() -> None:
    # The passage signal reads bodies only, so two equal bodies tie there and
    # the field signal -- title and path weighted up -- decides.
    concepts = [
        _concept("a/body", "Unrelated", "notes about the sandbox and its settings"),
        _concept("b/title", "Sandbox", "notes about the sandbox and its settings"),
    ]
    assert bm25.rank(concepts, "sandbox").ranked[0][0].concept_id == "b/title"


def test_the_best_window_wins_rather_than_the_sum() -> None:
    # Ten scattered mentions sum to more than one dense window; the rule is
    # `max`, so the concept whose ONE window answers the question leads.
    scattered = " ".join(["hooks"] + [f"pad{i}" for i in range(120)]) * 10
    dense = "hooks configure hooks per event, hooks run commands"
    concepts = [
        _concept("a/scattered", "Scattered", scattered),
        _concept("b/dense", "Dense", dense),
    ]
    ranking = bm25.rank(concepts, "hooks configure event commands")
    assert ranking.ranked[0][0].concept_id == "b/dense"


def test_the_best_window_offset_points_at_the_answering_text() -> None:
    body = "x " * 2000 + "the answer about retention lives here " + "y " * 2000
    ranking = bm25.rank([_concept("a/doc", "Doc", body)], "retention answer")
    start = ranking.best_window["a/doc"]
    assert "retention" in body[start : start + bm25.WINDOW_CHARS]


def test_a_concept_matching_nothing_carries_zero_lexical_and_sorts_by_id() -> None:
    concepts = [
        _concept("c/none", "C", "nothing relevant"),
        _concept("b/none", "B", "nothing relevant"),
        _concept("a/hit", "A", "the keyword appears"),
    ]
    ranked = bm25.rank(concepts, "keyword").ranked
    assert [(c.concept_id, lexical) for c, _, lexical in ranked] == [
        ("a/hit", 1),
        ("b/none", 0),
        ("c/none", 0),
    ]


def test_two_rankings_of_the_same_input_are_identical() -> None:
    concepts = [_concept(f"d{i}/c", f"T{i}", f"shared term {i} " * (i + 1)) for i in range(12)]
    first = bm25.rank(concepts, "shared term")
    second = bm25.rank(list(reversed(concepts)), "shared term")
    assert [(c.concept_id, s) for c, s, _ in first.ranked] == [
        (c.concept_id, s) for c, s, _ in second.ranked
    ]


def test_the_default_ranking_is_bm25_and_the_fusion_is_still_reachable() -> None:
    assert consume.DEFAULT_RANKING == "bm25"
    assert set(consume.RANKINGS) == {"bm25", "fusion"}


def test_a_directory_every_concept_shares_matches_nothing() -> None:
    # K3-20's defect, one signal over: in a one-document bundle every id
    # carries the document's directory, so a question naming the document
    # would otherwise match every concept.
    concepts = [
        _concept("handbook/intro", "Intro", "welcome to the club"),
        _concept("handbook/fees", "Fees", "membership costs money"),
    ]
    ranked = bm25.rank(concepts, "handbook").ranked
    assert [lexical for _, _, lexical in ranked] == [0, 0]
    assert [score for _, score, _ in ranked] == [0.0, 0.0]


def test_a_directory_that_separates_concepts_still_counts() -> None:
    concepts = [
        _concept("billing/overview", "Overview", "general words"),
        _concept("security/overview", "Overview", "general words"),
    ]
    assert bm25.rank(concepts, "billing").ranked[0][0].concept_id == "billing/overview"


def test_an_absent_inflection_reaches_the_form_the_collection_holds() -> None:
    # `vinterberedskapen` occurs nowhere; `vinterberedskap` does, and it is the
    # shared prefix -- a WORD of this collection -- that bridges the two
    # (`consume.tokens_match`, the rule the older ranking measured for
    # Norwegian inflection and compounding).
    concepts = [
        _concept("a/winter", "Vinterberedskap", "vinterberedskap kontrolleres hver host"),
        _concept("b/summer", "Sommer", "sommerdrift og vedlikehold"),
    ]
    ranked = bm25.rank(concepts, "Når kontrolleres vinterberedskapen?").ranked
    assert ranked[0][0].concept_id == "a/winter"
    assert ranked[0][2] == 2


def test_a_word_the_collection_holds_is_matched_as_itself_only() -> None:
    # `mode` is in the collection, so it is never widened to `model`: the
    # bridge is for a word that is absent, never a second reading of one that
    # is present.
    concepts = [
        _concept("a/mode", "Fast mode", "fast mode speeds output"),
        _concept("b/model", "Model", "model selection and model aliases"),
    ]
    ranked = bm25.rank(concepts, "mode").ranked
    assert [(c.concept_id, lexical) for c, _, lexical in ranked] == [("a/mode", 1), ("b/model", 0)]


def test_a_number_a_document_is_known_by_is_kept_whole() -> None:
    # `4.2` split on the dot is two single characters, and single characters
    # are dropped: without the whole token a question naming a section by its
    # number matches nothing in the section it names.
    assert "4.2" in bm25.tokens("Hva staar i punkt 4.2?")
    assert "10.2-2" in bm25.tokens("Krav 10.2—2")


def test_a_concept_named_by_its_number_carries_a_lexical_match() -> None:
    concepts = [
        _concept("r/vakthold-4-2", "Vakthold 4.2", "Vakten gaar fra fredag til soendag."),
        _concept("r/notat", "Notat om ettersyn", "Kontrollen av hytta foeres i skjema."),
    ]
    ranked = bm25.rank(concepts, "Hva staar i punkt 4.2 om kontrollen av hytta?").ranked
    lexical = {concept.concept_id: count for concept, _, count in ranked}
    assert lexical["r/vakthold-4-2"] >= 1