chatbot_v3/tests/unit/infrastructure/embedding/test_bm25.py

"""The `bm25-fa-norm-stop` sparse embedder (ADR-0001, ADR-0005).

These assert the analyzer/weighting behaviour benchmarked in the `emet`
evaluation lab. A change that makes one of these fail is a change that
invalidates that benchmark, not just a failing test.
"""

import pytest

from src.config import SparseEmbeddingSettings
from src.infrastructure.embedding.analyzers import analyze
from src.infrastructure.embedding.bm25 import Bm25SparseEmbedder, _token_index

pytestmark = pytest.mark.unit


@pytest.fixture
def settings() -> SparseEmbeddingSettings:
    return SparseEmbeddingSettings()


@pytest.fixture
def embedder(settings: SparseEmbeddingSettings) -> Bm25SparseEmbedder:
    return Bm25SparseEmbedder(settings)


# --- analyzer ------------------------------------------------------------


def test_analyze_keeps_digits_as_tokens() -> None:
    """Policy numbers, dates, and amounts are exactly what a lexical index
    should match on. An earlier implementation dropped every digit.
    """
    tokens = analyze("بیمه‌نامه شماره ۱۲۳۴۵ صادر شد")
    assert "12345" in tokens


def test_analyze_folds_eastern_digits_to_ascii() -> None:
    """The same number must match however it was authored."""
    assert analyze("۹۹۸۸۷۷") == analyze("٩٩٨٨٧٧") == analyze("998877")


def test_analyze_splits_on_zwnj() -> None:
    """ZWNJ joins compounds visually but they are separate index terms."""
    assert analyze("آتش‌سوزی") == ["آتش", "سوزی"]


def test_analyze_folds_arabic_letterforms_to_persian() -> None:
    # Arabic kaf (U+0643) vs. Persian keheh (U+06A9): the same word typed on
    # two different keyboards must produce the same term.
    assert analyze("كتاب") == analyze("کتاب")


def test_analyze_drops_persian_stopwords() -> None:
    assert analyze("این کتاب و آن مداد") == analyze("کتاب مداد")


def test_analyze_drops_english_stopwords() -> None:
    """The corpus is mixed-script, so the list carries English too."""
    assert analyze("the policy is valid") == ["policy", "valid"]


def test_analyze_lowercases_latin() -> None:
    assert analyze("POLICY Number") == ["policy", "number"]


def test_analyze_rejects_unknown_analyzer() -> None:
    with pytest.raises(ValueError, match="Unknown analyzer"):
        analyze("متن", "fa_norm_stem")


# --- token indexing ------------------------------------------------------


def test_token_index_is_stable_across_calls() -> None:
    assert _token_index("کتاب") == _token_index("کتاب")


def test_token_index_is_pinned_to_known_values() -> None:
    """A golden test. These indices are baked into every stored sparse vector,
    so changing the hash silently orphans the whole index -- a re-ingestion,
    not a deploy. Ingest-time and query-time encoding must agree forever.
    """
    assert _token_index("کتاب") == 1701064151
    assert _token_index("policy") == 741331709
    assert _token_index("12345") == 1232178634


def test_token_index_fits_signed_int32() -> None:
    for token in ("کتاب", "policy", "12345", "بیمه", "x" * 200):
        assert 0 <= _token_index(token) < 2**31 - 1


def test_token_index_distinguishes_different_tokens() -> None:
    assert _token_index("کتاب") != _token_index("مداد")


# --- vector construction -------------------------------------------------


def test_embed_batch_returns_one_vector_per_text(embedder: Bm25SparseEmbedder) -> None:
    assert len(embedder.embed_batch(["سلام دنیا", "یک تست دیگر"])) == 2


def test_embed_batch_empty_text_returns_empty_vector(embedder: Bm25SparseEmbedder) -> None:
    (vector,) = embedder.embed_batch([""])
    assert vector.indices == []
    assert vector.values == []


def test_embed_batch_all_stopwords_returns_empty_vector(embedder: Bm25SparseEmbedder) -> None:
    (vector,) = embedder.embed_batch(["و در به از که"])
    assert vector.indices == []


def test_embed_batch_emits_tokens_in_sorted_order(embedder: Bm25SparseEmbedder) -> None:
    """Deterministic output keeps re-ingestion byte-stable."""
    text = "مداد کتاب دفتر"
    expected = [_token_index(token) for token in sorted(analyze(text))]
    (vector,) = embedder.embed_batch([text])
    assert vector.indices == expected


def test_embed_batch_applies_bm25_saturation_not_raw_counts(
    embedder: Bm25SparseEmbedder,
) -> None:
    """Weight must be sublinear in term frequency: tripling a term must not
    triple its weight, which is the whole point of the `k` parameter.
    """
    (once,) = embedder.embed_batch(["کتاب"])
    (thrice,) = embedder.embed_batch(["کتاب کتاب کتاب"])
    assert thrice.values[0] > once.values[0]
    assert thrice.values[0] < 3 * once.values[0]


def test_embed_batch_query_side_omits_length_normalization(
    embedder: Bm25SparseEmbedder, settings: SparseEmbeddingSettings
) -> None:
    text = "کتاب مداد دفتر خودکار"
    (document,) = embedder.embed_batch([text], query=False)
    (query,) = embedder.embed_batch([text], query=True)

    assert document.indices == query.indices  # same terms, same hashing
    assert document.values != query.values

    k = settings.k
    assert query.values[0] == pytest.approx(1.0 * (k + 1.0) / (1.0 + k))


def test_embed_batch_short_document_outweighs_long_one(
    embedder: Bm25SparseEmbedder,
) -> None:
    """The `b` term discounts a term appearing in a longer document."""
    (short,) = embedder.embed_batch(["کتاب"])
    (long,) = embedder.embed_batch(["کتاب " + " ".join(f"واژه{i}" for i in range(200))])

    short_weight = short.values[short.indices.index(_token_index("کتاب"))]
    long_weight = long.values[long.indices.index(_token_index("کتاب"))]
    assert short_weight > long_weight


def test_embed_batch_respects_configured_parameters() -> None:
    """k/b/avg_len come from settings, so they can be retuned without a code
    change -- and so a retune is visibly a config decision.
    """
    default = Bm25SparseEmbedder(SparseEmbeddingSettings())
    tuned = Bm25SparseEmbedder(SparseEmbeddingSettings(k=2.5, b=0.2, avg_len=64.0))
    text = "کتاب کتاب مداد"

    assert default.embed_batch([text])[0].values != tuned.embed_batch([text])[0].values