Spaces:
Running on Zero
Running on Zero
File size: 5,066 Bytes
0828c2c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | """Tests for BM25 store."""
from pathlib import Path
import pytest
from langchain_core.documents import Document
from src.retrieval.stores.bm25_store import BM25Store
@pytest.fixture
def sample_documents():
"""Create sample documents for testing."""
return [
Document(
page_content="Python is a programming language known for its simplicity.",
metadata={"source": "doc1.txt"},
),
Document(
page_content="Machine learning is a subset of artificial intelligence.",
metadata={"source": "doc2.txt"},
),
Document(
page_content="Deep learning uses neural networks with many layers.",
metadata={"source": "doc3.txt"},
),
Document(
page_content="Natural language processing helps computers understand text.",
metadata={"source": "doc4.txt"},
),
Document(
page_content="Python is widely used for machine learning and data science.",
metadata={"source": "doc5.txt"},
),
]
@pytest.fixture
def temp_bm25_path(tmp_path):
"""Create temporary path for BM25 index."""
return str(tmp_path / "bm25_test")
class TestBM25Store:
"""Tests for BM25Store class."""
def test_init(self, temp_bm25_path):
"""Test BM25Store initialization."""
store = BM25Store(persist_path=temp_bm25_path, tokenizer="simple")
assert store.persist_path == Path(temp_bm25_path)
assert store.tokenizer_type == "simple"
assert not store.is_built()
def test_simple_tokenize(self, temp_bm25_path):
"""Test simple tokenization."""
store = BM25Store(persist_path=temp_bm25_path)
tokens = store._simple_tokenize("Hello, World! This is a test.")
assert tokens == ["hello", "world", "this", "is", "a", "test"]
def test_build_index(self, sample_documents, temp_bm25_path):
"""Test building BM25 index."""
store = BM25Store(persist_path=temp_bm25_path)
store.build_index(sample_documents)
assert store.is_built()
assert len(store.documents) == 5
assert len(store.tokenized_corpus) == 5
def test_build_index_empty_raises(self, temp_bm25_path):
"""Test that building with empty documents raises error."""
store = BM25Store(persist_path=temp_bm25_path)
with pytest.raises(ValueError, match="No documents provided"):
store.build_index([])
def test_search(self, sample_documents, temp_bm25_path):
"""Test BM25 search."""
store = BM25Store(persist_path=temp_bm25_path)
store.build_index(sample_documents)
# Search for Python
results = store.search("Python programming", k=3)
assert len(results) <= 3
# Should find the Python-related documents
assert any("Python" in doc.page_content for doc in results)
def test_search_with_scores(self, sample_documents, temp_bm25_path):
"""Test BM25 search with scores."""
store = BM25Store(persist_path=temp_bm25_path)
store.build_index(sample_documents)
results = store.search_with_scores("machine learning", k=3)
assert len(results) <= 3
# Results should be tuples of (doc, score)
for doc, score in results:
assert isinstance(doc, Document)
assert isinstance(score, float)
assert score >= 0
def test_save_and_load(self, sample_documents, temp_bm25_path):
"""Test saving and loading BM25 index."""
# Build and save
store1 = BM25Store(persist_path=temp_bm25_path)
store1.build_index(sample_documents)
store1.save()
# Load in new instance
store2 = BM25Store(persist_path=temp_bm25_path)
assert store2.load()
assert store2.is_built()
assert len(store2.documents) == 5
# Search should work on loaded store
results = store2.search("Python", k=2)
assert len(results) > 0
def test_load_nonexistent_returns_false(self, temp_bm25_path):
"""Test that loading nonexistent index returns False."""
store = BM25Store(persist_path=temp_bm25_path)
assert not store.load()
def test_delete(self, sample_documents, temp_bm25_path):
"""Test deleting BM25 index."""
store = BM25Store(persist_path=temp_bm25_path)
store.build_index(sample_documents)
store.save()
# Verify saved
assert Path(temp_bm25_path).exists()
# Delete
store.delete()
assert not Path(temp_bm25_path).exists()
assert not store.is_built()
def test_get_stats(self, sample_documents, temp_bm25_path):
"""Test getting index statistics."""
store = BM25Store(persist_path=temp_bm25_path)
store.build_index(sample_documents)
stats = store.get_stats()
assert stats["num_documents"] == 5
assert stats["tokenizer"] == "simple"
assert stats["is_built"]
assert "doc_hash" in stats
|