File size: 4,366 Bytes
87149fa d7d27a3 87149fa d7d27a3 87149fa d7d27a3 87149fa d7d27a3 87149fa | 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 | """Chunking and retrieval, with a bag-of-words embedder so tests stay offline."""
import numpy as np
import pytest
from dataloader import SEPARATOR
from retrieval import ProfileIndex, chunk_documents
VOCAB = ["faiss", "kubernetes", "volunteer", "python", "rust"]
def fake_embed(texts):
vectors = []
for text in texts:
lowered = text.lower()
vectors.append([float(word in lowered) for word in VOCAB])
return np.array(vectors, dtype=np.float32)
def profile(*documents):
return SEPARATOR.join(documents)
def test_chunks_carry_their_source_label():
text = profile("# Source: resume\n\nWorked with FAISS.", "# Source: volunteer\n\nTaught Python.")
chunks = chunk_documents(text)
assert any(c.startswith("[resume]") for c in chunks)
assert any(c.startswith("[volunteer]") for c in chunks)
def test_content_survives_chunking():
chunks = chunk_documents(profile("# Source: resume\n\nWorked with FAISS."))
assert "Worked with FAISS." in "\n".join(chunks)
def test_long_documents_split_into_several_chunks():
body = " ".join(f"word{i}" for i in range(40))
chunks = chunk_documents(
profile(f"# Source: resume\n\n{body}"),
words_per_chunk=10,
overlap=3,
)
assert len(chunks) > 1
def test_adjacent_chunks_overlap_by_the_requested_words():
body = " ".join(f"word{i}" for i in range(23))
chunks = chunk_documents(
profile(f"# Source: resume\n\n{body}"),
words_per_chunk=10,
overlap=3,
)
chunk_words = [chunk.split("]\n", 1)[1].split() for chunk in chunks]
assert all(
previous[-3:] == current[:3]
for previous, current in zip(chunk_words, chunk_words[1:])
)
def test_chunks_do_not_cross_source_document_boundaries():
chunks = chunk_documents(
profile(
"# Source: resume\n\n" + " ".join(f"resume{i}" for i in range(12)),
"# Source: projects\n\n" + " ".join(f"project{i}" for i in range(12)),
),
words_per_chunk=8,
overlap=2,
)
assert all(
("resume" in chunk) != ("project" in chunk)
for chunk in chunks
)
@pytest.mark.parametrize(
"words_per_chunk,overlap",
[(0, 0), (10, -1), (10, 10), (10, 11)],
)
def test_invalid_word_chunk_settings_are_rejected(words_per_chunk, overlap):
with pytest.raises(ValueError, match="words_per_chunk > overlap"):
chunk_documents(
"text",
words_per_chunk=words_per_chunk,
overlap=overlap,
)
def test_short_documents_stay_whole():
assert len(chunk_documents(profile("# Source: resume\n\nShort."))) == 1
def test_empty_input_yields_no_chunks():
assert chunk_documents("") == []
def test_documents_without_a_source_header_still_chunk():
chunks = chunk_documents("Just a body with no header.")
assert len(chunks) == 1
assert "Just a body" in chunks[0]
class TestSearch:
@pytest.fixture
def index(self):
return ProfileIndex(
["[resume] Built a FAISS index", "[volunteer] Volunteer teaching", "[projects] Python tooling"],
fake_embed,
)
def test_search_ranks_the_matching_chunk_first(self, index):
assert "FAISS" in index.search("faiss", k=1)[0]
assert "Volunteer" in index.search("volunteer", k=1)[0]
def test_search_respects_k(self, index):
assert len(index.search("python", k=2)) == 2
def test_hybrid_search_returns_bm25_and_dense_matches(self, index):
results = index.search("volunteer", k=3)
assert "Volunteer" in results[0]
def test_k_larger_than_the_corpus_is_safe(self, index):
assert len(index.search("python", k=99)) == 3
def test_empty_query_returns_nothing(self, index):
assert index.search(" ") == []
def test_non_positive_k_returns_nothing(self, index):
assert index.search("python", k=0) == []
def test_search_text_joins_results(self, index):
assert "---" in index.search_text("python", k=2)
def test_search_text_reports_when_nothing_is_found(self, index):
assert "No relevant excerpts" in index.search_text("", k=3)
def test_an_empty_index_does_not_crash():
empty = ProfileIndex([], fake_embed)
assert empty.search("anything") == []
assert "No relevant excerpts" in empty.search_text("anything")
|