diff --git a/pyproject.toml b/pyproject.toml index bcfcf8bb2..13258cabb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,6 @@ dependencies = [ "model2vec>=0.4.0", "vicinity>=0.4.4", "numpy>=1.24.0", - "bm25s>=0.2.0", "pathspec>=0.12", "tree-sitter>=0.25,<0.26", "tree-sitter-language-pack>=1.0,<1.8.0,!=1.6.3", diff --git a/src/semble/cache.py b/src/semble/cache.py index 31a73bbef..4b4ff803e 100644 --- a/src/semble/cache.py +++ b/src/semble/cache.py @@ -8,10 +8,14 @@ from pathlib import Path from typing import TYPE_CHECKING +import orjson + +from semble.index.bm25 import BM25 +from semble.index.dense import SelectableBasicBackend from semble.index.file_walker import walk_files from semble.index.files import FileStatus, get_extensions, get_file_status -from semble.index.types import PersistencePath -from semble.types import ContentType +from semble.index.types import CACHE_FORMAT_VERSION, FileManifestEntry, PersistencePath, PreviousIndex, make_chunk_id +from semble.types import Chunk, ContentType from semble.utils import is_git_url, resolve_model_name logger = logging.getLogger(__name__) @@ -99,10 +103,13 @@ def _metadata_matches(metadata: dict, model_path: str, content: Sequence[Content try: content_type = tuple(ContentType(s) for s in metadata["content_type"]) - # chunk_size is absent in indexes built before this field was added; treat None as mismatch - # so old caches are transparently rebuilt with the current chunk size. + # chunk_size and cache_version are absent in indexes built before those fields were added; + # treat that as a mismatch so old caches are transparently rebuilt in the current format. chunk_size_ok = metadata.get("chunk_size") == _DESIRED_CHUNK_LENGTH_CHARS - return metadata["model_path"] == model_path and set(content_type) == set(content) and chunk_size_ok + version_ok = metadata.get("cache_version") == CACHE_FORMAT_VERSION + return ( + metadata["model_path"] == model_path and set(content_type) == set(content) and chunk_size_ok and version_ok + ) except (KeyError, ValueError): return False @@ -131,7 +138,7 @@ def get_validated_cache(path: str, model_path: str | None, content: Sequence[Con extensions = get_extensions(content) path_as_path = Path(path).resolve() - stored_files: list[str] = metadata.get("file_paths", []) + stored_files = metadata.get("files", {}) current_files = [] for file_path in walk_files(path_as_path, extensions=extensions): file_status = get_file_status(file_path, write_time) @@ -145,3 +152,58 @@ def get_validated_cache(path: str, model_path: str | None, content: Sequence[Con return None return index_path + + +def load_previous_for_incremental( + path: str, model_path: str | None, content: Sequence[ContentType] +) -> PreviousIndex | None: + """Load compatible index state for incremental reuse. + + :param path: Source path used to locate the cached index. + :param model_path: Requested model, or None to use the default. + :param content: Content types the cached index must support. + :return: Previous index state, or None if the cache is unavailable or invalid. + """ + try: + index_path = find_index_from_cache_folder(path) + persistence_path = PersistencePath.from_path(index_path) + if persistence_path.non_existing(): + return None + + if model_path is None: + model_path = resolve_model_name() + with open(persistence_path.metadata, encoding="utf-8") as f: + metadata = json.load(f) + if not _metadata_matches(metadata, model_path, content): + return None + + raw_manifest = metadata.get("files") + if not raw_manifest: + return None + manifest = {indexed_path: FileManifestEntry(**entry) for indexed_path, entry in raw_manifest.items()} + + with open(persistence_path.chunks, "rb") as f: + chunks = [Chunk.from_dict(item) for item in orjson.loads(f.read())] + + vectors = SelectableBasicBackend.load(persistence_path.semantic_index).vectors + bm25_index = BM25.load(persistence_path.bm25_index) + chunk_count = len(chunks) + if vectors.shape[0] != chunk_count or len(bm25_index.doc_order) != chunk_count: + return None + expected_ids: list[str] = [] + next_start = 0 + for indexed_path, entry in manifest.items(): + if ( + entry.start != next_start + or entry.count < 0 + or any(chunk.file_path != indexed_path for chunk in chunks[entry.start : entry.start + entry.count]) + ): + return None + expected_ids.extend(make_chunk_id(indexed_path, slot) for slot in range(entry.count)) + next_start += entry.count + if next_start != chunk_count or bm25_index.doc_order != expected_ids: + return None + + return PreviousIndex(chunks=chunks, vectors=vectors, manifest=manifest, bm25_index=bm25_index) + except (OSError, orjson.JSONDecodeError, json.JSONDecodeError, KeyError, TypeError, ValueError): + return None diff --git a/src/semble/index/bm25.py b/src/semble/index/bm25.py new file mode 100644 index 000000000..311dc6d4a --- /dev/null +++ b/src/semble/index/bm25.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +import math +from collections import Counter +from pathlib import Path + +import numpy as np +import numpy.typing as npt +import orjson + +_K1 = 1.5 # Term-frequency saturation +_B = 0.75 # Document length normalization + + +class BM25: + """BM25 inverted index supporting incremental document updates.""" + + def __init__(self) -> None: + """Create an empty index.""" + self._documents: dict[str, Counter[str]] = {} + self._doc_lengths: dict[str, int] = {} + self._total_doc_length = 0 + self.postings: dict[str, dict[str, int]] = {} + self.doc_order: list[str] = [] + self._positions: dict[str, int] = {} + + def add_document(self, chunk_id: str, tokens: list[str]) -> None: + """Index one document, rejecting duplicate IDs.""" + if chunk_id in self._documents: + raise ValueError(f"chunk_id already indexed: {chunk_id}") + counts = Counter(tokens) + self._documents[chunk_id] = counts + self._doc_lengths[chunk_id] = len(tokens) + self._total_doc_length += len(tokens) + for term, count in counts.items(): + self.postings.setdefault(term, {})[chunk_id] = count + + def remove_document(self, chunk_id: str) -> None: + """Remove a document's postings; no-op if chunk_id is not indexed.""" + counts = self._documents.pop(chunk_id, None) + if counts is None: + return + self._total_doc_length -= self._doc_lengths.pop(chunk_id) + for term in counts: + docs = self.postings[term] + docs.pop(chunk_id, None) + if not docs: + del self.postings[term] + + def set_doc_order(self, chunk_ids: list[str]) -> None: + """Set the current global chunk-list order that get_scores' output is aligned to.""" + self.doc_order = chunk_ids + self._positions = {chunk_id: i for i, chunk_id in enumerate(chunk_ids)} + + def get_scores( + self, tokens: list[str], weight_mask: npt.NDArray[np.bool_] | None = None + ) -> npt.NDArray[np.float32]: + """Calculate BM25 scores for a tokenized query. + + :param tokens: Tokenized search query. + :param weight_mask: Optional boolean mask aligned with doc_order. + :return: Scores aligned with doc_order. + """ + output_size = len(self.doc_order) + corpus_size = len(self._documents) + scores: npt.NDArray[np.float32] = np.zeros(output_size, dtype=np.float32) + if not tokens or corpus_size == 0: + return scores + + avgdl = self._total_doc_length / corpus_size + for term, query_tf in Counter(tokens).items(): + docs = self.postings.get(term) + if not docs: + continue + df = len(docs) + idf = math.log(1 + (corpus_size - df + 0.5) / (df + 0.5)) + for chunk_id, tf in docs.items(): + idx = self._positions.get(chunk_id) + if idx is None: + continue + dl = self._doc_lengths[chunk_id] + tfc = tf / (_K1 * (1 - _B + _B * dl / avgdl) + tf) + scores[idx] += query_tf * idf * tfc + + if weight_mask is not None: + scores = scores * weight_mask + return scores + + def save(self, path: Path) -> None: + """Persist the index to path/index.json.""" + path.mkdir(parents=True, exist_ok=True) + payload = {"documents": self._documents, "doc_order": self.doc_order} + (path / "index.json").write_bytes(orjson.dumps(payload)) + + @classmethod + def load(cls, path: Path) -> "BM25": + """Load an index from path/index.json, reconstructing its postings.""" + data = orjson.loads((path / "index.json").read_bytes()) + index = cls() + doc_order = data["doc_order"] + documents = data["documents"] + if len(doc_order) != len(set(doc_order)) or set(documents) != set(doc_order): + raise ValueError("Persisted BM25 document state is inconsistent") + index._documents = {chunk_id: Counter(counts) for chunk_id, counts in documents.items()} + for chunk_id, counts in index._documents.items(): + for term, count in counts.items(): + index.postings.setdefault(term, {})[chunk_id] = count + index._doc_lengths = {chunk_id: sum(counts.values()) for chunk_id, counts in index._documents.items()} + index._total_doc_length = sum(index._doc_lengths.values()) + index.set_doc_order(doc_order) + return index diff --git a/src/semble/index/create.py b/src/semble/index/create.py index b4dd189c8..41db52acb 100644 --- a/src/semble/index/create.py +++ b/src/semble/index/create.py @@ -2,17 +2,39 @@ from collections.abc import Sequence from pathlib import Path -import bm25s +import numpy as np from model2vec.model import StaticModel from vicinity.backends.basic import BasicArgs from semble.chunking import chunk_source +from semble.index.bm25 import BM25 from semble.index.dense import SelectableBasicBackend, embed_chunks from semble.index.file_walker import walk_files -from semble.index.files import FileStatus, detect_language, get_extensions, get_file_status, read_file_text +from semble.index.files import ( + FileStatus, + detect_language, + get_extensions, + get_file_status, + read_file_text, +) from semble.index.sparse import enrich_for_bm25 +from semble.index.types import FileManifestEntry, PreviousIndex, make_chunk_id from semble.tokens import tokenize -from semble.types import Chunk, ContentType +from semble.types import Chunk, ContentType, EmbeddingMatrix + + +def _reindex_file( + bm25_index: BM25, + indexed_path: str, + file_chunks: list[Chunk], + previous_entry: FileManifestEntry | None, +) -> None: + """Replace a file's BM25 postings: remove its old slots (if any), then add its new ones.""" + if previous_entry is not None: + for slot in range(previous_entry.count): + bm25_index.remove_document(make_chunk_id(indexed_path, slot)) + for slot, chunk in enumerate(file_chunks): + bm25_index.add_document(make_chunk_id(indexed_path, slot), tokenize(enrich_for_bm25(chunk))) def create_index_from_path( @@ -20,39 +42,89 @@ def create_index_from_path( model: StaticModel, content: ContentType | Sequence[ContentType] = (ContentType.CODE,), display_root: Path | None = None, -) -> tuple[bm25s.BM25, SelectableBasicBackend, list[Chunk]]: - """Create an index from a resolved directory, optionally storing chunk paths relative to display_root. + previous: PreviousIndex | None = None, +) -> tuple[BM25, SelectableBasicBackend, list[Chunk], dict[str, FileManifestEntry]]: + """Create an index from a resolved directory, optionally reusing a previous index's unchanged files. :param path: Resolved absolute path to index. :param model: The model to use for indexing. :param content: Content types to index. :param display_root: If set, chunk file paths are stored relative to this root. + :param previous: A previously built index to reuse unchanged files' chunks/embeddings/postings from. :raises ValueError: if no items were found, no index can be created. - :return: A bm25 index, vicinity index and list of chunks + :return: A BM25 index, semantic index, list of chunks, and file manifest. """ - chunks: list[Chunk] = [] + # PreviousIndex is consumed; mutate BM25 in place to avoid a copy. + bm25_index = previous.bm25_index if previous is not None else BM25() + previous_manifest = previous.manifest if previous is not None else {} + normalized = (content,) if isinstance(content, ContentType) else content resolved_extensions = get_extensions(normalized) + + chunks: list[Chunk] = [] + chunk_ids: list[str] = [] + vector_parts: list[EmbeddingMatrix] = [] + manifest: dict[str, FileManifestEntry] = {} + embedding_parts: list[tuple[int, int, int]] = [] + changed_chunks: list[Chunk] = [] + for file_path in walk_files(path, resolved_extensions): language = detect_language(file_path) with contextlib.suppress(OSError): file_status = get_file_status(file_path, None) if file_status != FileStatus.VALID: continue - source = read_file_text(file_path) - chunk_path = file_path.relative_to(display_root) if display_root else file_path - chunks.extend(chunk_source(source, str(chunk_path), language)) - if chunks: + indexed_path = str(file_path.relative_to(display_root) if display_root else file_path) + mtime = file_path.stat().st_mtime + previous_entry = previous_manifest.get(indexed_path) + + if previous is not None and previous_entry is not None and previous_entry.mtime == mtime: + file_chunks = previous.chunks[previous_entry.start : previous_entry.start + previous_entry.count] + vector_parts.append( + previous.vectors[previous_entry.start : previous_entry.start + previous_entry.count] + ) + else: + source = read_file_text(file_path) + file_chunks = chunk_source(source, indexed_path, language) + _reindex_file(bm25_index, indexed_path, file_chunks, previous_entry) + + embedding_parts.append((len(vector_parts), len(chunks), len(file_chunks))) + vector_parts.append(np.empty((0, model.dim), dtype=np.float32)) + changed_chunks.extend(file_chunks) + + start = len(chunks) + chunks.extend(file_chunks) + chunk_ids.extend(make_chunk_id(indexed_path, slot) for slot in range(len(file_chunks))) + manifest[indexed_path] = FileManifestEntry(mtime=mtime, start=start, count=len(file_chunks)) + + for indexed_path in previous_manifest.keys() - manifest.keys(): + _reindex_file(bm25_index, indexed_path, [], previous_manifest[indexed_path]) + + if not chunks: + raise ValueError(f"No supported files found under {path}.") + + if previous is None: embeddings = embed_chunks(model, chunks) - bm25_index = bm25s.BM25() - bm25_index.index( - [tokenize(enrich_for_bm25(chunk)) for chunk in chunks], - show_progress=False, - ) - args = BasicArgs() - semantic_index = SelectableBasicBackend(embeddings, args) else: - raise ValueError(f"No supported files found under {path}.") + changed_embeddings = embed_chunks(model, changed_chunks) + offset = 0 + for index, _, count in embedding_parts: + vector_parts[index] = changed_embeddings[offset : offset + count] + offset += count + same_vector_layout = len(manifest) == len(previous_manifest) and all( + (previous_entry := previous_manifest.get(indexed_path)) is not None + and entry.start == previous_entry.start + and entry.count == previous_entry.count + for indexed_path, entry in manifest.items() + ) + if same_vector_layout: + embeddings = previous.vectors + for vector_part, start, count in embedding_parts: + embeddings[start : start + count] = vector_parts[vector_part] + else: + embeddings = np.vstack(vector_parts) + bm25_index.set_doc_order(chunk_ids) + semantic_index = SelectableBasicBackend(embeddings, BasicArgs()) - return bm25_index, semantic_index, chunks + return bm25_index, semantic_index, chunks, manifest diff --git a/src/semble/index/index.py b/src/semble/index/index.py index f51eac81f..3e11c3784 100644 --- a/src/semble/index/index.py +++ b/src/semble/index/index.py @@ -12,14 +12,14 @@ import numpy as np import numpy.typing as npt import orjson -from bm25s import BM25 from model2vec.model import StaticModel -from semble.cache import get_validated_cache +from semble.cache import get_validated_cache, load_previous_for_incremental +from semble.index.bm25 import BM25 from semble.index.create import create_index_from_path from semble.index.dense import SelectableBasicBackend, load_model from semble.index.files import read_file_text -from semble.index.types import PersistencePath +from semble.index.types import CACHE_FORMAT_VERSION, FileManifestEntry, PersistencePath from semble.search import _search_semantic, search from semble.stats import save_search_stats from semble.types import CallType, Chunk, ContentType, IndexStats, SearchResult @@ -60,6 +60,7 @@ def __init__( root: Path | None = None, content: ContentType | Sequence[ContentType] = _DEFAULT_CONTENT, loaded_from_disk: bool = False, + manifest: dict[str, FileManifestEntry] | None = None, ) -> None: """Initialize a SembleIndex. Should be created with from_path or from_git. @@ -71,6 +72,7 @@ def __init__( :param root: Root directory used to read file sizes for token-savings stats. :param content: Content type used when indexing; controls the search pipeline. :param loaded_from_disk: Whether the index was loaded from disk (cache hit); controls CLI messaging. + :param manifest: File modification times and chunk ranges used for incremental reindexing. """ self.model = model self.chunks: list[Chunk] = chunks @@ -82,6 +84,7 @@ def __init__( self._file_sizes: dict[str, int] = self._compute_file_sizes(root) if root else {} self._file_mapping, self._language_mapping = self._populate_mapping() self.loaded_from_disk: bool = loaded_from_disk + self._manifest: dict[str, FileManifestEntry] = manifest or {} def _populate_mapping(self) -> tuple[dict[str, list[int]], dict[str, list[int]]]: """Build (file → chunk indices, language → chunk indices) mappings, in that order.""" @@ -152,14 +155,18 @@ def from_path( model, model_path = load_model(model_path) path = path.resolve() - bm25, vicinity, chunks = create_index_from_path( + previous = load_previous_for_incremental(str(path), model_path, normalized) + bm25_index, semantic_index, chunks, manifest = create_index_from_path( path, model=model, content=normalized, display_root=path, + previous=previous, ) - return SembleIndex(model, bm25, vicinity, chunks, model_path, root=path, content=normalized) + return SembleIndex( + model, bm25_index, semantic_index, chunks, model_path, root=path, content=normalized, manifest=manifest + ) @classmethod def from_git( @@ -207,7 +214,7 @@ def from_git( model, model_path = load_model(model_path) resolved_path = Path(tmp_dir).resolve() - bm25, vicinity, chunks = create_index_from_path( + bm25_index, semantic_index, chunks, manifest = create_index_from_path( resolved_path, model=model, content=normalized, @@ -216,12 +223,13 @@ def from_git( return SembleIndex( model, - bm25, - vicinity, + bm25_index, + semantic_index, chunks, model_path, root=resolved_path, content=normalized, + manifest=manifest, ) def find_related( @@ -310,19 +318,30 @@ def load_from_disk(cls: type[SembleIndex], path: Path | str) -> SembleIndex: missing = ", ".join(str(p) for p in non_existent) raise FileNotFoundError(f"Index not found at {path}. Missing: {missing}") - bm_25_index = BM25.load(persistence_paths.bm25_index) - semantic_index = SelectableBasicBackend.load(persistence_paths.semantic_index) with open(persistence_paths.metadata, "rb") as f: metadata = orjson.loads(f.read()) + found_version = metadata.get("cache_version") + if found_version != CACHE_FORMAT_VERSION: + raise ValueError( + f"Unsupported index format {found_version!r}; expected {CACHE_FORMAT_VERSION}. Rebuild the index." + ) + + bm25_index = BM25.load(persistence_paths.bm25_index) + semantic_index = SelectableBasicBackend.load(persistence_paths.semantic_index) with open(persistence_paths.chunks, "rb") as f: chunk_data = orjson.loads(f.read()) chunks = [] for chunk_item in chunk_data: chunks.append(Chunk.from_dict(chunk_item)) + if len(chunks) != len(bm25_index.doc_order) or len(chunks) != semantic_index.vectors.shape[0]: + raise ValueError("Persisted index components have inconsistent document counts") root_path = metadata["root_path"] model_path = metadata["model_path"] content = tuple(ContentType(s) for s in metadata.get("content_type", ["code"])) + manifest = { + indexed_path: FileManifestEntry(**entry) for indexed_path, entry in metadata.get("files", {}).items() + } if root_path: root_path = Path(root_path) @@ -330,13 +349,14 @@ def load_from_disk(cls: type[SembleIndex], path: Path | str) -> SembleIndex: return cls( model, - bm_25_index, + bm25_index, semantic_index, chunks, model_path, root=root_path, content=content, loaded_from_disk=True, + manifest=manifest, ) def save(self, path: Path | str) -> None: @@ -348,9 +368,8 @@ def save(self, path: Path | str) -> None: self._bm25_index.save(persistence_paths.bm25_index) self._semantic_index.save(persistence_paths.semantic_index) - chunks_as_dict = [chunk.to_dict() for chunk in self.chunks] with open(persistence_paths.chunks, "wb") as f: - data = orjson.dumps(chunks_as_dict) + data = orjson.dumps(self.chunks) f.write(data) from semble.chunking.chunking import _DESIRED_CHUNK_LENGTH_CHARS # avoid circular import at module level @@ -360,8 +379,9 @@ def save(self, path: Path | str) -> None: "time": datetime.now().timestamp(), "model_path": self._model_path, "content_type": list(x.value for x in self._content), - "file_paths": sorted(self._file_mapping), "chunk_size": _DESIRED_CHUNK_LENGTH_CHARS, + "cache_version": CACHE_FORMAT_VERSION, + "files": self._manifest, } with open(persistence_paths.metadata, "wb") as f: data = orjson.dumps(metadata) diff --git a/src/semble/index/types.py b/src/semble/index/types.py index fc3c7f130..b40f1098d 100644 --- a/src/semble/index/types.py +++ b/src/semble/index/types.py @@ -3,6 +3,16 @@ from dataclasses import dataclass from pathlib import Path +from semble.index.bm25 import BM25 +from semble.types import Chunk, EmbeddingMatrix + +CACHE_FORMAT_VERSION = 1 # Bump when the persisted index schema changes. + + +def make_chunk_id(indexed_path: str, slot: int) -> str: + """Return the stable document ID for a file chunk.""" + return f"{indexed_path}:{slot}" + @dataclass class PersistencePath: @@ -28,3 +38,22 @@ def from_path(cls: type[PersistencePath], path: Path) -> PersistencePath: semantic_index=path / "semantic_index", metadata=path / "metadata.json", ) + + +@dataclass +class FileManifestEntry: + """Record a file's modification time and chunk range within the global chunk list.""" + + mtime: float + start: int + count: int + + +@dataclass +class PreviousIndex: + """A previously built index, loaded for reuse during incremental reindexing.""" + + chunks: list[Chunk] + vectors: EmbeddingMatrix + manifest: dict[str, FileManifestEntry] + bm25_index: BM25 diff --git a/src/semble/search.py b/src/semble/search.py index 238d9eb40..cab61da59 100644 --- a/src/semble/search.py +++ b/src/semble/search.py @@ -1,8 +1,8 @@ -import bm25s import numpy as np import numpy.typing as npt from model2vec import StaticModel +from semble.index.bm25 import BM25 from semble.index.dense import SelectableBasicBackend from semble.index.sparse import selector_to_mask from semble.ranking import apply_query_boost, boost_multi_chunk_files, rerank_topk, resolve_alpha @@ -46,7 +46,7 @@ def _sort_top_k(arr: npt.NDArray, top_k: int) -> npt.NDArray[np.int_]: def _search_bm25( query: str, - bm25_index: bm25s.BM25, + bm25_index: BM25, chunks: list[Chunk], top_k: int, selector: npt.NDArray[np.int_] | None, @@ -67,7 +67,7 @@ def search( query: str, model: StaticModel, semantic_index: SelectableBasicBackend, - bm25_index: bm25s.BM25, + bm25_index: BM25, chunks: list[Chunk], top_k: int, alpha: float | None = None, diff --git a/tests/index/test_bm25.py b/tests/index/test_bm25.py new file mode 100644 index 000000000..17f980427 --- /dev/null +++ b/tests/index/test_bm25.py @@ -0,0 +1,90 @@ +import math +from pathlib import Path + +import numpy as np +import orjson +import pytest + +from semble.index.bm25 import BM25 + + +def _build(docs: dict[str, list[str]]) -> BM25: + index = BM25() + for chunk_id, tokens in docs.items(): + index.add_document(chunk_id, tokens) + index.set_doc_order(list(docs)) + return index + + +def test_scoring_matches_lucene_formula() -> None: + """BM25 scores use the Lucene term-frequency formula.""" + index = _build({"a": ["authenticate", "token"], "b": ["login", "password"]}) + scores = index.get_scores(["authenticate"]) + np.testing.assert_allclose(scores[0], math.log(1 + 1.5 / 1.5) / 2.5) + assert scores[1] == 0 + + +def test_removed_and_unordered_documents_stop_scoring() -> None: + """Only documents retained in the current order contribute scores.""" + index = _build({"a": ["authenticate"], "b": ["login"]}) + index.remove_document("missing") + index.set_doc_order(["b"]) + assert np.all(index.get_scores(["authenticate"]) == 0) + + index.remove_document("a") + index.set_doc_order(["a", "b"]) + assert np.all(index.get_scores(["authenticate"]) == 0) + + +def test_duplicate_add_document_raises() -> None: + """Re-adding an already-indexed chunk_id raises, catching caller bugs.""" + index = _build({"a": ["x"]}) + with pytest.raises(ValueError, match="already indexed"): + index.add_document("a", ["y"]) + + +@pytest.mark.parametrize( + ("mask", "expected_nonzero"), + [ + (None, [0, 1]), + (np.array([True, False]), [0]), + ], +) +def test_weight_mask_zeroes_masked_docs(mask: np.ndarray | None, expected_nonzero: list[int]) -> None: + """weight_mask zeroes out scores for masked-out positions, by global chunk order.""" + index = _build({"a": ["shared"], "b": ["shared"]}) + scores = index.get_scores(["shared"], weight_mask=mask) + nonzero = [i for i, s in enumerate(scores) if s > 0] + assert nonzero == expected_nonzero + + +@pytest.mark.parametrize("query", [[], ["zzznonexistent"]]) +def test_unmatched_queries_return_all_zero(query: list[str]) -> None: + """Empty and unknown queries return an all-zero array sized to the corpus.""" + index = _build({"a": ["foo"], "b": ["bar"]}) + scores = index.get_scores(query) + assert scores.shape == (2,) + assert np.all(scores == 0) + + +def test_save_load_preserves_scores_and_doc_order(tmp_path: Path) -> None: + """save/load roundtrips postings and doc_order, producing identical scores for a fixed query.""" + index = _build({"empty": [], "a": ["authenticate", "token"], "b": ["login", "password"]}) + index.save(tmp_path) + + loaded = BM25.load(tmp_path) + assert loaded.doc_order == index.doc_order + np.testing.assert_array_equal(loaded.get_scores(["authenticate"]), index.get_scores(["authenticate"])) + + +def test_load_rejects_inconsistent_document_order(tmp_path: Path) -> None: + """Persisted document order must describe the same documents as the postings.""" + index = _build({"a": ["authenticate"]}) + index.save(tmp_path) + index_path = tmp_path / "index.json" + data = orjson.loads(index_path.read_bytes()) + data["doc_order"] = ["other"] + index_path.write_bytes(orjson.dumps(data)) + + with pytest.raises(ValueError, match="document state"): + BM25.load(tmp_path) diff --git a/tests/index/test_create.py b/tests/index/test_create.py new file mode 100644 index 000000000..d06b3350d --- /dev/null +++ b/tests/index/test_create.py @@ -0,0 +1,163 @@ +from pathlib import Path +from typing import Any +from unittest.mock import patch + +import numpy as np +import orjson +import pytest + +from semble.cache import load_previous_for_incremental +from semble.index.create import create_index_from_path +from semble.index.index import SembleIndex +from semble.index.types import PreviousIndex, make_chunk_id +from semble.types import ContentType + + +def _write_files(root: Path, files: dict[str, str]) -> None: + for rel, content in files.items(): + path = root / rel + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content) + + +def test_incremental_reindex_reuses_updates_and_prunes(mock_model: Any, tmp_path: Path) -> None: + """One incremental pass reuses unchanged vectors, re-embeds changes, and keeps BM25 slots current.""" + _write_files( + tmp_path, + { + "a.py": "def stable_anchor():\n return 1\n", + "b.py": "def changed_value():\n return 2\n", + "c.py": "def unique_gone():\n return 3\n", + "emptying.py": "def becomes_empty():\n return 4\n", + }, + ) + bm25_before, semantic_before, chunks_before, manifest_before = create_index_from_path( + tmp_path, mock_model, display_root=tmp_path + ) + a_entry = manifest_before["a.py"] + b_entry = manifest_before["b.py"] + a_vectors_before = semantic_before.vectors[a_entry.start : a_entry.start + a_entry.count].copy() + b_vectors_before = semantic_before.vectors[b_entry.start : b_entry.start + b_entry.count].copy() + previous = PreviousIndex( + chunks=chunks_before, + vectors=semantic_before.vectors, + manifest=manifest_before, + bm25_index=bm25_before, + ) + _, semantic_unchanged, _, _ = create_index_from_path(tmp_path, mock_model, display_root=tmp_path, previous=previous) + assert semantic_unchanged.vectors is semantic_before.vectors + + (tmp_path / "b.py").write_text("def changed_value():\n return 999\n") + bm25_before, semantic_before, chunks_before, manifest_before = create_index_from_path( + tmp_path, mock_model, display_root=tmp_path, previous=previous + ) + assert semantic_before.vectors is previous.vectors + previous = PreviousIndex( + chunks=chunks_before, + vectors=semantic_before.vectors, + manifest=manifest_before, + bm25_index=bm25_before, + ) + + (tmp_path / "c.py").unlink() + (tmp_path / "emptying.py").write_text(" " * 128) + _write_files(tmp_path, {"d.py": "def brand_new_term():\n return 4\n"}) + bm25_after, semantic_after, _, manifest_after = create_index_from_path( + tmp_path, mock_model, display_root=tmp_path, previous=previous + ) + + a_entry_after = manifest_after["a.py"] + b_entry_after = manifest_after["b.py"] + np.testing.assert_array_equal( + semantic_after.vectors[a_entry_after.start : a_entry_after.start + a_entry_after.count], a_vectors_before + ) + assert not np.array_equal( + b_vectors_before, + semantic_after.vectors[b_entry_after.start : b_entry_after.start + b_entry_after.count], + ) + assert "c.py" not in manifest_after + assert "d.py" in manifest_after + assert manifest_after["emptying.py"].count == 0 + assert bm25_after.get_scores(["unique_gone"]).sum() == 0 + assert bm25_after.get_scores(["becomes_empty"]).sum() == 0 + assert bm25_after.get_scores(["brand", "new", "term"]).sum() > 0 + expected_ids = { + make_chunk_id(indexed_path, slot) + for indexed_path, entry in manifest_after.items() + for slot in range(entry.count) + } + assert set(bm25_after.doc_order) == expected_ids + + +def _build_valid_cache(index_path: Path, mock_model: Any) -> dict: + """Build a real, well-formed on-disk index and return its metadata dict for mutation.""" + src = index_path.parent / "src" + _write_files(src, {"a.py": "def a():\n return 1\n", "b.py": "def b():\n return 2\n"}) + + with patch("semble.index.index.load_model", return_value=(mock_model, "my/model")): + SembleIndex.from_path(src).save(index_path) + return orjson.loads((index_path / "metadata.json").read_bytes()) + + +@pytest.mark.parametrize( + "corrupt", + [ + "missing_cache", + "missing_files_key", + "metadata_mismatch", + "component_length_mismatch", + "length_mismatch", + "overlapping_entries", + "bm25_order_mismatch", + "corrupt_json", + ], +) +def test_load_previous_for_incremental_fails_closed(corrupt: str, tmp_path: Path, mock_model: Any) -> None: + """Any structurally invalid or missing cache state yields None instead of raising.""" + index_path = tmp_path / "index" + + if corrupt != "missing_cache": + metadata = _build_valid_cache(index_path, mock_model) + if corrupt == "missing_files_key": + del metadata["files"] + elif corrupt == "metadata_mismatch": + metadata["model_path"] = "other/model" + elif corrupt == "component_length_mismatch": + chunks_path = index_path / "chunks.json" + chunks = orjson.loads(chunks_path.read_bytes()) + chunks_path.write_bytes(orjson.dumps(chunks[:-1])) + elif corrupt == "length_mismatch": + metadata["files"]["a.py"]["count"] += 5 + elif corrupt == "overlapping_entries": + metadata["files"]["b.py"]["start"] = metadata["files"]["a.py"]["start"] + elif corrupt == "bm25_order_mismatch": + bm25_path = index_path / "bm25_index" / "index.json" + bm25 = orjson.loads(bm25_path.read_bytes()) + bm25["doc_order"].reverse() + bm25_path.write_bytes(orjson.dumps(bm25)) + elif corrupt == "corrupt_json": + (index_path / "metadata.json").write_bytes(b"{not json") + with patch("semble.cache.find_index_from_cache_folder", return_value=index_path): + assert load_previous_for_incremental("/some/path", "my/model", [ContentType.CODE]) is None + return + (index_path / "metadata.json").write_bytes(orjson.dumps(metadata)) + + with patch("semble.cache.find_index_from_cache_folder", return_value=index_path): + result = load_previous_for_incremental("/some/path", "my/model", [ContentType.CODE]) + assert result is None + + +def test_load_previous_for_incremental_happy_path(mock_model: Any, tmp_path: Path) -> None: + """A well-formed cache round-trips into a usable PreviousIndex.""" + index_path = tmp_path / "cache" / "index" + _build_valid_cache(index_path, mock_model) + + with ( + patch("semble.cache.find_index_from_cache_folder", return_value=index_path), + patch("semble.cache.resolve_model_name", return_value="my/model"), + ): + previous = load_previous_for_incremental(str(index_path.parent / "src"), None, [ContentType.CODE]) + + assert previous is not None + assert len(previous.chunks) == previous.vectors.shape[0] == len(previous.bm25_index.doc_order) + assert "a.py" in previous.manifest diff --git a/tests/index/test_index.py b/tests/index/test_index.py index 76a9d040b..08994cde5 100644 --- a/tests/index/test_index.py +++ b/tests/index/test_index.py @@ -2,6 +2,7 @@ from typing import Any from unittest.mock import MagicMock, patch +import orjson import pytest from model2vec import StaticModel @@ -31,7 +32,7 @@ def test_index_markdown_inclusion( mock_model: StaticModel, tmp_project: Path, content: list[ContentType], md_in_results: bool ) -> None: """Markdown files are excluded for code-only and included when docs is requested.""" - _, _, chunks = create_index_from_path(tmp_project, mock_model, content=content) + _, _, chunks, _ = create_index_from_path(tmp_project, mock_model, content=content) has_md = ".md" in {Path(c.file_path).suffix for c in chunks} assert has_md is md_in_results @@ -57,7 +58,7 @@ def test_from_git_include_text_files_deprecated(mock_model: Any, tmp_project: Pa with patch("semble.index.index.load_model", return_value=(mock_model, "")): with patch("subprocess.run", return_value=fake_result): with patch("semble.index.index.create_index_from_path") as mock_create: - mock_create.return_value = (MagicMock(), MagicMock(), [make_chunk("x = 1", "f.py")]) + mock_create.return_value = (MagicMock(), MagicMock(), [make_chunk("x = 1", "f.py")], {}) with pytest.warns(DeprecationWarning, match="include_text_files is deprecated"): SembleIndex.from_git("https://example.com/repo", include_text_files=True) @@ -186,7 +187,9 @@ def test_find_related(indexed_index: SembleIndex) -> None: def test_roundtrip(tmp_path: Path, indexed_index: SembleIndex) -> None: """Test that saving and loading a folder leads to the same data.""" + assert indexed_index.chunks[0].to_dict()["location"] == indexed_index.chunks[0].location indexed_index.save(tmp_path) + assert "location" not in orjson.loads((tmp_path / "chunks.json").read_bytes())[0] with patch.object(StaticModel, "from_pretrained"): index_2 = SembleIndex.load_from_disk(tmp_path) assert index_2.chunks == indexed_index.chunks @@ -194,7 +197,7 @@ def test_roundtrip(tmp_path: Path, indexed_index: SembleIndex) -> None: def test_load_save_roundtrip_preserves_manifest(tmp_path: Path, indexed_index: SembleIndex) -> None: - """load_from_disk followed by save must not clobber file_paths with an empty list.""" + """load_from_disk followed by save must preserve the incremental manifest.""" save_a = tmp_path / "a" save_b = tmp_path / "b" indexed_index.save(save_a) @@ -203,8 +206,8 @@ def test_load_save_roundtrip_preserves_manifest(tmp_path: Path, indexed_index: S loaded.save(save_b) import json - manifest_a = json.loads((save_a / "metadata.json").read_text())["file_paths"] - manifest_b = json.loads((save_b / "metadata.json").read_text())["file_paths"] + manifest_a = json.loads((save_a / "metadata.json").read_text())["files"] + manifest_b = json.loads((save_b / "metadata.json").read_text())["files"] assert manifest_b == manifest_a assert len(manifest_b) > 0 @@ -234,6 +237,28 @@ def test_load_from_disk_missing_files_reports_them(tmp_path: Path) -> None: assert "chunks.json" not in error_msg +@pytest.mark.parametrize( + ("corruption", "message"), + [("version", "Unsupported index format"), ("counts", "inconsistent document counts")], +) +def test_load_from_disk_rejects_incompatible_state( + corruption: str, message: str, tmp_path: Path, indexed_index: SembleIndex +) -> None: + """Incompatible persistence metadata and component counts are rejected.""" + indexed_index.save(tmp_path) + if corruption == "version": + path = tmp_path / "metadata.json" + data = orjson.loads(path.read_bytes()) + del data["cache_version"] + else: + path = tmp_path / "chunks.json" + data = orjson.loads(path.read_bytes())[:-1] + path.write_bytes(orjson.dumps(data)) + + with pytest.raises(ValueError, match=message): + SembleIndex.load_from_disk(tmp_path) + + def test_from_path_uses_cache_when_valid(tmp_project: Path) -> None: """from_path returns the cached index directly when get_validated_cache hits.""" fake_cached = MagicMock(spec=SembleIndex) diff --git a/tests/test_cache.py b/tests/test_cache.py index 52b52a49c..294ebfa52 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -18,6 +18,8 @@ resolve_cache_folder, save_index_to_cache, ) +from semble.chunking.chunking import _DESIRED_CHUNK_LENGTH_CHARS +from semble.index.types import CACHE_FORMAT_VERSION from semble.types import ContentType @@ -134,9 +136,8 @@ def _write_metadata( write_time: float, file_paths: list[str] | None = None, chunk_size: int | None = None, + cache_version: int | None = None, ) -> None: - from semble.chunking.chunking import _DESIRED_CHUNK_LENGTH_CHARS - path.mkdir(parents=True, exist_ok=True) (path / "chunks.json").write_text("[]") (path / "bm25_index").write_text("") @@ -147,8 +148,9 @@ def _write_metadata( "model_path": model_path, "content_type": content_type, "time": write_time, - "file_paths": file_paths if file_paths is not None else [], + "files": {file_path: {} for file_path in file_paths or []}, "chunk_size": chunk_size if chunk_size is not None else _DESIRED_CHUNK_LENGTH_CHARS, + "cache_version": cache_version if cache_version is not None else CACHE_FORMAT_VERSION, } ) ) @@ -189,8 +191,6 @@ def test_get_validated_cache_metadata_mismatch( def test_get_validated_cache_reads_utf8_metadata_with_non_ascii_file_paths(tmp_path: Path) -> None: """Cache metadata is always UTF-8, even when the system default encoding is not.""" - from semble.chunking.chunking import _DESIRED_CHUNK_LENGTH_CHARS - index_path = tmp_path / "index" index_path.mkdir(parents=True) (index_path / "chunks.json").write_text("[]") @@ -207,8 +207,9 @@ def test_get_validated_cache_reads_utf8_metadata_with_non_ascii_file_paths(tmp_p "model_path": "my/model", "content_type": ["docs"], "time": 0.0, - "file_paths": [non_ascii_path], + "files": {non_ascii_path: {}}, "chunk_size": _DESIRED_CHUNK_LENGTH_CHARS, + "cache_version": CACHE_FORMAT_VERSION, }, ensure_ascii=False, ), @@ -228,36 +229,26 @@ def open_with_cp936_default(file: object, mode: str = "r", *args: object, **kwar assert result == index_path -def test_get_validated_cache_chunk_size_mismatch_returns_none(tmp_path: Path) -> None: - """Cache built with a different chunk_size is not reused.""" - from semble.chunking.chunking import _DESIRED_CHUNK_LENGTH_CHARS - - index_path = tmp_path / "index" - _write_metadata(index_path, "my/model", ["code"], float("inf"), chunk_size=_DESIRED_CHUNK_LENGTH_CHARS + 100) - with patch("semble.cache.find_index_from_cache_folder", return_value=index_path): - assert get_validated_cache("/path", "my/model", [ContentType.CODE]) is None - - -def test_get_validated_cache_missing_chunk_size_returns_none(tmp_path: Path) -> None: - """Old cache metadata without chunk_size field is not reused (transparent rebuild).""" +@pytest.mark.parametrize( + ("field", "value"), + [ + ("chunk_size", None), + ("chunk_size", _DESIRED_CHUNK_LENGTH_CHARS + 100), + ("cache_version", None), + ("cache_version", CACHE_FORMAT_VERSION - 1), + ], +) +def test_get_validated_cache_format_mismatch_returns_none(field: str, value: int | None, tmp_path: Path) -> None: + """Missing or stale persistence-format metadata forces a rebuild.""" index_path = tmp_path / "index" - # Write metadata as old semble would — no chunk_size field - index_path.mkdir(parents=True, exist_ok=True) - (index_path / "chunks.json").write_text("[]") - (index_path / "bm25_index").write_text("") - (index_path / "semantic_index").write_text("") - import json as _json - - (index_path / "metadata.json").write_text( - _json.dumps( - { - "model_path": "my/model", - "content_type": ["code"], - "time": float("inf"), - "file_paths": [], - } - ) - ) + _write_metadata(index_path, "my/model", ["code"], float("inf")) + metadata_path = index_path / "metadata.json" + metadata = json.loads(metadata_path.read_text()) + if value is None: + del metadata[field] + else: + metadata[field] = value + metadata_path.write_text(json.dumps(metadata)) with patch("semble.cache.find_index_from_cache_folder", return_value=index_path): assert get_validated_cache("/path", "my/model", [ContentType.CODE]) is None diff --git a/tests/test_search.py b/tests/test_search.py index ed605fbf6..b8bc7a250 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -1,13 +1,13 @@ from typing import Any from unittest.mock import MagicMock, patch -import bm25s import numpy as np import numpy.typing as npt import pytest from model2vec import StaticModel from vicinity.backends.basic import BasicArgs +from semble.index.bm25 import BM25 from semble.index.dense import SelectableBasicBackend, embed_chunks, load_model from semble.search import _search_bm25, _search_semantic, _sort_top_k, search from semble.tokens import tokenize @@ -15,6 +15,16 @@ from tests.conftest import make_chunk +def _build_bm25(chunks: list[Chunk]) -> BM25: + """Build a BM25 index over chunks, keyed by their position.""" + index = BM25() + doc_ids = [f"c{i}" for i in range(len(chunks))] + for doc_id, chunk in zip(doc_ids, chunks): + index.add_document(doc_id, tokenize(chunk.content)) + index.set_doc_order(doc_ids) + return index + + @pytest.fixture def chunks() -> list[Chunk]: """Four small code chunks covering authentication, login, user service, and utils.""" @@ -37,11 +47,9 @@ def embeddings(chunks: list[Chunk]) -> npt.NDArray[np.float32]: @pytest.fixture -def bm25(chunks: list[Chunk]) -> bm25s.BM25: +def bm25(chunks: list[Chunk]) -> BM25: """Pre-built BM25 index over the chunks fixture.""" - index = bm25s.BM25() - index.index([tokenize(chunk.content) for chunk in chunks], show_progress=False) - return index + return _build_bm25(chunks) @pytest.fixture @@ -50,7 +58,7 @@ def semantic(embeddings: npt.NDArray[np.float32]) -> SelectableBasicBackend: return SelectableBasicBackend(embeddings, BasicArgs()) -def test_search_bm25(bm25: bm25s.BM25, chunks: list[Chunk]) -> None: +def test_search_bm25(bm25: BM25, chunks: list[Chunk]) -> None: """search_bm25: returns most relevant chunk first; selector restricts to given indices.""" results = _search_bm25("authenticate token", bm25, chunks, top_k=4, selector=None) assert len(results) > 0 @@ -62,8 +70,8 @@ def test_search_bm25(bm25: bm25s.BM25, chunks: list[Chunk]) -> None: @pytest.mark.parametrize("query", ["", " ", "\n\n", "zzzznonexistentterm"]) -def test_bm25_returns_empty_for_no_match(bm25: bm25s.BM25, chunks: list[Chunk], query: str) -> None: - """Empty / whitespace-only / token-less queries return [] instead of crashing bm25s.""" +def test_bm25_returns_empty_for_no_match(bm25: BM25, chunks: list[Chunk], query: str) -> None: + """Empty / whitespace-only / token-less queries return [] instead of crashing.""" assert _search_bm25(query, bm25, chunks, top_k=3, selector=None) == [] @@ -74,9 +82,7 @@ def test_semantic_search(semantic: SelectableBasicBackend, chunks: list[Chunk], assert all(-1.0 <= r.score <= 1.0 for r in results) -def test_search_hybrid( - chunks: list[Chunk], semantic: SelectableBasicBackend, bm25: bm25s.BM25, mock_model: Any -) -> None: +def test_search_hybrid(chunks: list[Chunk], semantic: SelectableBasicBackend, bm25: BM25, mock_model: Any) -> None: """search_hybrid: returns combined results; identical content in different files produces separate results.""" results = search("authenticate token", mock_model, semantic, bm25, chunks, top_k=3) assert len(results) > 0 @@ -91,8 +97,7 @@ def test_search_hybrid( embs /= np.linalg.norm(embs, axis=1, keepdims=True) + 1e-8 sem_index = SelectableBasicBackend(embs, BasicArgs()) - bm25_index = bm25s.BM25() - bm25_index.index([tokenize(c.content) for c in all_chunks], show_progress=False) + bm25_index = _build_bm25(all_chunks) deduped = search("helper", mock_model, sem_index, bm25_index, all_chunks, top_k=5) result_locations = {r.chunk.file_path for r in deduped} @@ -114,7 +119,7 @@ def test_search_source_labels( top_k: int, chunks: list[Chunk], semantic: SelectableBasicBackend, - bm25: bm25s.BM25, + bm25: BM25, mock_model: Any, ) -> None: """Each result carries a source label matching the search mode used.""" diff --git a/uv.lock b/uv.lock index f04d43a3e..d0a51dbde 100644 --- a/uv.lock +++ b/uv.lock @@ -91,19 +91,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" }, ] -[[package]] -name = "bm25s" -version = "0.3.8" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/98/3f/88010a24298e2ca4b0807ca5a59ee69e867551fc289c3b062d41cefda59b/bm25s-0.3.8.tar.gz", hash = "sha256:d165f096d64961de0b25a0ad4989d42dfe529c94ca12a9a070eef173c62b8612", size = 80495, upload-time = "2026-04-29T02:18:45.493Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/82/c2/e5dbdbf8e99ec504432f9bd82cdf990faff1ffcdd5316272cb7fccce31ed/bm25s-0.3.8-py3-none-any.whl", hash = "sha256:b5ff52a0633ee5add26ca89619f5c5fff66127dcce0984976a2a09e1871a33f9", size = 74452, upload-time = "2026-04-29T02:18:43.892Z" }, -] - [[package]] name = "certifi" version = "2026.4.22" @@ -3143,7 +3130,6 @@ wheels = [ name = "semble" source = { editable = "." } dependencies = [ - { name = "bm25s" }, { name = "model2vec" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, @@ -3179,7 +3165,6 @@ mcp = [ [package.metadata] requires-dist = [ - { name = "bm25s", specifier = ">=0.2.0" }, { name = "einops", marker = "extra == 'benchmark'", specifier = ">=0.8.2" }, { name = "matplotlib", marker = "extra == 'benchmark'", specifier = ">=3.7" }, { name = "mcp", marker = "extra == 'mcp'", specifier = ">=1.0,<2.0" },