tests/test_bm25.py

"""
tests/test_bm25.py
──────────────────
Unit tests for the BM25 keyword retriever.
"""

from __future__ import annotations

from llama_index.core.schema import TextNode

from docling_pdf_processor.retrieval.bm25 import BM25, KeywordSearchRetriever


def test_bm25_tokenize():
    bm25 = BM25()
    assert bm25.tokenize("Hello, world!") == ["hello", "world"]


def test_bm25_search_ranking():
    bm25 = BM25()
    docs = [
        "the quick brown fox",
        "the lazy dog sleeps",
        "foxes are quick and brown",
    ]
    bm25.fit(docs)
    results = bm25.search("quick brown fox", top_k=3)
    assert len(results) == 3
    # The most relevant doc should be the one containing quick + brown + fox
    top_idx = results[0][0]
    assert top_idx in (0, 2)


def test_keyword_search_retriever():
    nodes = [
        TextNode(text="Python is great for data science."),
        TextNode(text="JavaScript runs in the browser."),
        TextNode(text="Python and machine learning go hand in hand."),
    ]
    retriever = KeywordSearchRetriever(nodes)
    results = retriever.retrieve("machine learning", top_k=2)
    assert len(results) >= 1
    assert "machine learning" in results[0].node.text.lower()