From 3216bd431a1be87eced35de670f0bd4e8810a634 Mon Sep 17 00:00:00 2001 From: Pringled Date: Thu, 16 Jul 2026 08:09:02 +0200 Subject: [PATCH 1/4] feat: Add incremental indexing --- pyproject.toml | 1 - src/semble/cache.py | 71 +++++++++++++++-- src/semble/index/bm25.py | 110 +++++++++++++++++++++++++++ src/semble/index/create.py | 124 ++++++++++++++++++++++++------ src/semble/index/index.py | 45 ++++++++--- src/semble/index/types.py | 29 +++++++ src/semble/search.py | 6 +- tests/index/test_bm25.py | 90 ++++++++++++++++++++++ tests/index/test_create.py | 152 +++++++++++++++++++++++++++++++++++++ tests/index/test_index.py | 27 ++++++- tests/test_cache.py | 57 ++++++-------- tests/test_search.py | 33 ++++---- uv.lock | 15 ---- 13 files changed, 654 insertions(+), 106 deletions(-) create mode 100644 src/semble/index/bm25.py create mode 100644 tests/index/test_bm25.py create mode 100644 tests/index/test_create.py 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..5ab3bc12c 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 @@ -145,3 +152,57 @@ 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 a prior on-disk index for reuse during incremental reindexing. + + Returns None -- forfeiting the incremental optimization for this run, so the + caller falls back to a full rebuild -- for malformed or inconsistent cache + state. + """ + 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..1a4eb803c --- /dev/null +++ b/src/semble/index/bm25.py @@ -0,0 +1,110 @@ +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 +_B = 0.75 + + +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]: + """Return a BM25 score per position in doc_order, compatible with bm25s.get_scores. + + Unknown query tokens contribute 0 rather than raising. + """ + 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) + documents = {chunk_id: dict(counts) for chunk_id, counts in self._documents.items()} + payload = {"documents": 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..52dc03221 100644 --- a/src/semble/index/create.py +++ b/src/semble/index/create.py @@ -1,18 +1,56 @@ import contextlib +import hashlib 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, +) 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 _embed_changed_chunks( + model: StaticModel, + embedding_parts: list[tuple[int, int]], + changed_chunks: list[Chunk], + vector_parts: list[EmbeddingMatrix], +) -> None: + """Embed changed chunks in one batch and restore their per-file rows.""" + if not changed_chunks: + return + 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 def create_index_from_path( @@ -20,39 +58,81 @@ 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. + + Unchanged files (same content hash as ``previous``) are not chunked, embedded, or re-indexed + into BM25 -- their existing chunks, embeddings, and postings are reused as-is. :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] = [] + 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]] = [] + 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: - 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: + + indexed_path = str(file_path.relative_to(display_root) if display_root else file_path) + raw_source = file_path.read_bytes() + file_hash = hashlib.sha256(raw_source).hexdigest() + previous_entry = previous_manifest.get(indexed_path) + + if previous is not None and previous_entry is not None and previous_entry.file_hash == file_hash: + 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 = raw_source.decode("utf-8", errors="replace") + file_chunks = chunk_source(source, indexed_path, language) + _reindex_file(bm25_index, indexed_path, file_chunks, previous_entry) + + if file_chunks: + embedding_parts.append((len(vector_parts), len(file_chunks))) + vector_parts.append(np.empty((0, 0), dtype=np.float32)) + changed_chunks.extend(file_chunks) + else: + vector_parts.append(np.empty((0, model.dim), dtype=np.float32)) + + 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(file_hash=file_hash, start=start, count=len(file_chunks)) + + if previous is not None: + for indexed_path, previous_entry in previous_manifest.items(): + if indexed_path not in manifest: + for slot in range(previous_entry.count): + bm25_index.remove_document(make_chunk_id(indexed_path, slot)) + + if not chunks: raise ValueError(f"No supported files found under {path}.") - return bm25_index, semantic_index, chunks + _embed_changed_chunks(model, embedding_parts, changed_chunks, vector_parts) + + embeddings = np.vstack(vector_parts) + bm25_index.set_doc_order(chunk_ids) + semantic_index = SelectableBasicBackend(embeddings, BasicArgs()) + + return bm25_index, semantic_index, chunks, manifest diff --git a/src/semble/index/index.py b/src/semble/index/index.py index f51eac81f..ceee3ce9a 100644 --- a/src/semble/index/index.py +++ b/src/semble/index/index.py @@ -6,20 +6,21 @@ import warnings from collections import defaultdict from collections.abc import Sequence +from dataclasses import asdict from datetime import datetime from pathlib import Path 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 +61,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 +73,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 hashes and chunk ranges used for incremental reindexing. """ self.model = model self.chunks: list[Chunk] = chunks @@ -82,6 +85,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 +156,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 +215,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 +224,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 +319,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 +350,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: @@ -362,6 +383,8 @@ def save(self, path: Path | str) -> None: "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": {indexed_path: asdict(entry) for indexed_path, entry in self._manifest.items()}, } 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..199c93ead 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 = 3 + + +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: + """Records a file's content hash and its chunk range within the global chunk list.""" + + file_hash: str + 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..d80baa135 --- /dev/null +++ b/tests/index/test_create.py @@ -0,0 +1,152 @@ +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) + np.testing.assert_array_equal(semantic_unchanged.vectors, semantic_before.vectors) + + (tmp_path / "b.py").write_text("def changed_value():\n return 999\n") + (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..94a6c3400 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) @@ -234,6 +235,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..5ef619676 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("") @@ -149,6 +150,7 @@ def _write_metadata( "time": write_time, "file_paths": file_paths if file_paths is not None else [], "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("[]") @@ -209,6 +209,7 @@ def test_get_validated_cache_reads_utf8_metadata_with_non_ascii_file_paths(tmp_p "time": 0.0, "file_paths": [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" }, From 66f8ad275bf27d1c3a72d2e76d56fade2ecb7507 Mon Sep 17 00:00:00 2001 From: Pringled Date: Thu, 16 Jul 2026 09:12:16 +0200 Subject: [PATCH 2/4] Optimizations --- src/semble/cache.py | 9 ++++--- src/semble/index/bm25.py | 13 +++++----- src/semble/index/create.py | 50 ++++++++++++++++++++++++++------------ src/semble/index/index.py | 6 ++--- src/semble/index/types.py | 2 +- tests/index/test_create.py | 13 +++++++++- tests/index/test_index.py | 2 ++ 7 files changed, 64 insertions(+), 31 deletions(-) diff --git a/src/semble/cache.py b/src/semble/cache.py index 5ab3bc12c..fd08878bc 100644 --- a/src/semble/cache.py +++ b/src/semble/cache.py @@ -157,11 +157,12 @@ def get_validated_cache(path: str, model_path: str | None, content: Sequence[Con def load_previous_for_incremental( path: str, model_path: str | None, content: Sequence[ContentType] ) -> PreviousIndex | None: - """Load a prior on-disk index for reuse during incremental reindexing. + """Load compatible index state for incremental reuse. - Returns None -- forfeiting the incremental optimization for this run, so the - caller falls back to a full rebuild -- for malformed or inconsistent cache - state. + :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) diff --git a/src/semble/index/bm25.py b/src/semble/index/bm25.py index 1a4eb803c..311dc6d4a 100644 --- a/src/semble/index/bm25.py +++ b/src/semble/index/bm25.py @@ -8,8 +8,8 @@ import numpy.typing as npt import orjson -_K1 = 1.5 -_B = 0.75 +_K1 = 1.5 # Term-frequency saturation +_B = 0.75 # Document length normalization class BM25: @@ -55,9 +55,11 @@ def set_doc_order(self, chunk_ids: list[str]) -> None: def get_scores( self, tokens: list[str], weight_mask: npt.NDArray[np.bool_] | None = None ) -> npt.NDArray[np.float32]: - """Return a BM25 score per position in doc_order, compatible with bm25s.get_scores. + """Calculate BM25 scores for a tokenized query. - Unknown query tokens contribute 0 rather than raising. + :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) @@ -87,8 +89,7 @@ def get_scores( def save(self, path: Path) -> None: """Persist the index to path/index.json.""" path.mkdir(parents=True, exist_ok=True) - documents = {chunk_id: dict(counts) for chunk_id, counts in self._documents.items()} - payload = {"documents": documents, "doc_order": self.doc_order} + payload = {"documents": self._documents, "doc_order": self.doc_order} (path / "index.json").write_bytes(orjson.dumps(payload)) @classmethod diff --git a/src/semble/index/create.py b/src/semble/index/create.py index 52dc03221..1cf98cb9e 100644 --- a/src/semble/index/create.py +++ b/src/semble/index/create.py @@ -37,9 +37,21 @@ def _reindex_file( bm25_index.add_document(make_chunk_id(indexed_path, slot), tokenize(enrich_for_bm25(chunk))) +def _remove_deleted_files( + bm25_index: BM25, + manifest: dict[str, FileManifestEntry], + previous_manifest: dict[str, FileManifestEntry], +) -> None: + """Remove BM25 documents belonging to files absent from the new manifest.""" + for indexed_path, previous_entry in previous_manifest.items(): + if indexed_path not in manifest: + for slot in range(previous_entry.count): + bm25_index.remove_document(make_chunk_id(indexed_path, slot)) + + def _embed_changed_chunks( model: StaticModel, - embedding_parts: list[tuple[int, int]], + embedding_parts: list[tuple[int, int, int]], changed_chunks: list[Chunk], vector_parts: list[EmbeddingMatrix], ) -> None: @@ -48,7 +60,7 @@ def _embed_changed_chunks( return changed_embeddings = embed_chunks(model, changed_chunks) offset = 0 - for index, count in embedding_parts: + for index, _, count in embedding_parts: vector_parts[index] = changed_embeddings[offset : offset + count] offset += count @@ -62,9 +74,6 @@ def create_index_from_path( ) -> tuple[BM25, SelectableBasicBackend, list[Chunk], dict[str, FileManifestEntry]]: """Create an index from a resolved directory, optionally reusing a previous index's unchanged files. - Unchanged files (same content hash as ``previous``) are not chunked, embedded, or re-indexed - into BM25 -- their existing chunks, embeddings, and postings are reused as-is. - :param path: Resolved absolute path to index. :param model: The model to use for indexing. :param content: Content types to index. @@ -73,6 +82,7 @@ def create_index_from_path( :raises ValueError: if no items were found, no index can be created. :return: A BM25 index, semantic index, list of chunks, and file manifest. """ + # 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 {} @@ -83,7 +93,7 @@ def create_index_from_path( chunk_ids: list[str] = [] vector_parts: list[EmbeddingMatrix] = [] manifest: dict[str, FileManifestEntry] = {} - embedding_parts: list[tuple[int, int]] = [] + embedding_parts: list[tuple[int, int, int]] = [] changed_chunks: list[Chunk] = [] for file_path in walk_files(path, resolved_extensions): @@ -109,8 +119,8 @@ def create_index_from_path( _reindex_file(bm25_index, indexed_path, file_chunks, previous_entry) if file_chunks: - embedding_parts.append((len(vector_parts), len(file_chunks))) - vector_parts.append(np.empty((0, 0), dtype=np.float32)) + embedding_parts.append((len(vector_parts), len(chunks), len(file_chunks))) + vector_parts.append(np.empty((0, 0), dtype=np.float32)) # Replaced after batched embedding. changed_chunks.extend(file_chunks) else: vector_parts.append(np.empty((0, model.dim), dtype=np.float32)) @@ -121,17 +131,27 @@ def create_index_from_path( manifest[indexed_path] = FileManifestEntry(file_hash=file_hash, start=start, count=len(file_chunks)) if previous is not None: - for indexed_path, previous_entry in previous_manifest.items(): - if indexed_path not in manifest: - for slot in range(previous_entry.count): - bm25_index.remove_document(make_chunk_id(indexed_path, slot)) + _remove_deleted_files(bm25_index, manifest, previous_manifest) if not chunks: raise ValueError(f"No supported files found under {path}.") - _embed_changed_chunks(model, embedding_parts, changed_chunks, vector_parts) - - embeddings = np.vstack(vector_parts) + if previous is None: + embeddings = embed_chunks(model, chunks) + else: + _embed_changed_chunks(model, embedding_parts, changed_chunks, vector_parts) + 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()) diff --git a/src/semble/index/index.py b/src/semble/index/index.py index ceee3ce9a..a1e95bb65 100644 --- a/src/semble/index/index.py +++ b/src/semble/index/index.py @@ -6,7 +6,6 @@ import warnings from collections import defaultdict from collections.abc import Sequence -from dataclasses import asdict from datetime import datetime from pathlib import Path @@ -369,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 @@ -384,7 +382,7 @@ def save(self, path: Path | str) -> None: "file_paths": sorted(self._file_mapping), "chunk_size": _DESIRED_CHUNK_LENGTH_CHARS, "cache_version": CACHE_FORMAT_VERSION, - "files": {indexed_path: asdict(entry) for indexed_path, entry in self._manifest.items()}, + "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 199c93ead..071c036c6 100644 --- a/src/semble/index/types.py +++ b/src/semble/index/types.py @@ -6,7 +6,7 @@ from semble.index.bm25 import BM25 from semble.types import Chunk, EmbeddingMatrix -CACHE_FORMAT_VERSION = 3 +CACHE_FORMAT_VERSION = 1 # Bump when the persisted index schema changes. def make_chunk_id(indexed_path: str, slot: int) -> str: diff --git a/tests/index/test_create.py b/tests/index/test_create.py index d80baa135..d06b3350d 100644 --- a/tests/index/test_create.py +++ b/tests/index/test_create.py @@ -45,9 +45,20 @@ def test_incremental_reindex_reuses_updates_and_prunes(mock_model: Any, tmp_path bm25_index=bm25_before, ) _, semantic_unchanged, _, _ = create_index_from_path(tmp_path, mock_model, display_root=tmp_path, previous=previous) - np.testing.assert_array_equal(semantic_unchanged.vectors, semantic_before.vectors) + 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"}) diff --git a/tests/index/test_index.py b/tests/index/test_index.py index 94a6c3400..0d653df06 100644 --- a/tests/index/test_index.py +++ b/tests/index/test_index.py @@ -187,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 From bbcb40a4378a33dddeb65ff6384e500cda0a82a9 Mon Sep 17 00:00:00 2001 From: Pringled Date: Thu, 16 Jul 2026 09:59:54 +0200 Subject: [PATCH 3/4] CRLF normalization fix --- src/semble/index/create.py | 2 +- tests/index/test_create.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/semble/index/create.py b/src/semble/index/create.py index 1cf98cb9e..cb7416239 100644 --- a/src/semble/index/create.py +++ b/src/semble/index/create.py @@ -114,7 +114,7 @@ def create_index_from_path( previous.vectors[previous_entry.start : previous_entry.start + previous_entry.count] ) else: - source = raw_source.decode("utf-8", errors="replace") + source = raw_source.decode("utf-8", errors="replace").replace("\r\n", "\n").replace("\r", "\n") file_chunks = chunk_source(source, indexed_path, language) _reindex_file(bm25_index, indexed_path, file_chunks, previous_entry) diff --git a/tests/index/test_create.py b/tests/index/test_create.py index d06b3350d..ce0e96226 100644 --- a/tests/index/test_create.py +++ b/tests/index/test_create.py @@ -25,7 +25,7 @@ def test_incremental_reindex_reuses_updates_and_prunes(mock_model: Any, tmp_path _write_files( tmp_path, { - "a.py": "def stable_anchor():\n return 1\n", + "a.py": "def stable_anchor():\r\n return 1\r\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", @@ -36,6 +36,7 @@ def test_incremental_reindex_reuses_updates_and_prunes(mock_model: Any, tmp_path ) a_entry = manifest_before["a.py"] b_entry = manifest_before["b.py"] + assert all("\r" not in chunk.content for chunk in chunks_before[a_entry.start : a_entry.start + a_entry.count]) 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( From c379a7eb5282f09ebae42d71cd4aea3deb2dd618 Mon Sep 17 00:00:00 2001 From: Pringled Date: Thu, 16 Jul 2026 11:20:49 +0200 Subject: [PATCH 4/4] Switch back to mtime, some simplifications --- src/semble/cache.py | 2 +- src/semble/index/create.py | 58 ++++++++++---------------------------- src/semble/index/index.py | 3 +- src/semble/index/types.py | 4 +-- tests/index/test_create.py | 3 +- tests/index/test_index.py | 6 ++-- tests/test_cache.py | 4 +-- 7 files changed, 25 insertions(+), 55 deletions(-) diff --git a/src/semble/cache.py b/src/semble/cache.py index fd08878bc..4b4ff803e 100644 --- a/src/semble/cache.py +++ b/src/semble/cache.py @@ -138,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) diff --git a/src/semble/index/create.py b/src/semble/index/create.py index cb7416239..41db52acb 100644 --- a/src/semble/index/create.py +++ b/src/semble/index/create.py @@ -1,5 +1,4 @@ import contextlib -import hashlib from collections.abc import Sequence from pathlib import Path @@ -16,6 +15,7 @@ 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 @@ -37,34 +37,6 @@ def _reindex_file( bm25_index.add_document(make_chunk_id(indexed_path, slot), tokenize(enrich_for_bm25(chunk))) -def _remove_deleted_files( - bm25_index: BM25, - manifest: dict[str, FileManifestEntry], - previous_manifest: dict[str, FileManifestEntry], -) -> None: - """Remove BM25 documents belonging to files absent from the new manifest.""" - for indexed_path, previous_entry in previous_manifest.items(): - if indexed_path not in manifest: - for slot in range(previous_entry.count): - bm25_index.remove_document(make_chunk_id(indexed_path, slot)) - - -def _embed_changed_chunks( - model: StaticModel, - embedding_parts: list[tuple[int, int, int]], - changed_chunks: list[Chunk], - vector_parts: list[EmbeddingMatrix], -) -> None: - """Embed changed chunks in one batch and restore their per-file rows.""" - if not changed_chunks: - return - 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 - - def create_index_from_path( path: Path, model: StaticModel, @@ -104,34 +76,30 @@ def create_index_from_path( continue indexed_path = str(file_path.relative_to(display_root) if display_root else file_path) - raw_source = file_path.read_bytes() - file_hash = hashlib.sha256(raw_source).hexdigest() + 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.file_hash == file_hash: + 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 = raw_source.decode("utf-8", errors="replace").replace("\r\n", "\n").replace("\r", "\n") + source = read_file_text(file_path) file_chunks = chunk_source(source, indexed_path, language) _reindex_file(bm25_index, indexed_path, file_chunks, previous_entry) - if file_chunks: - embedding_parts.append((len(vector_parts), len(chunks), len(file_chunks))) - vector_parts.append(np.empty((0, 0), dtype=np.float32)) # Replaced after batched embedding. - changed_chunks.extend(file_chunks) - else: - vector_parts.append(np.empty((0, model.dim), dtype=np.float32)) + 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(file_hash=file_hash, start=start, count=len(file_chunks)) + manifest[indexed_path] = FileManifestEntry(mtime=mtime, start=start, count=len(file_chunks)) - if previous is not None: - _remove_deleted_files(bm25_index, manifest, previous_manifest) + 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}.") @@ -139,7 +107,11 @@ def create_index_from_path( if previous is None: embeddings = embed_chunks(model, chunks) else: - _embed_changed_chunks(model, embedding_parts, changed_chunks, vector_parts) + 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 diff --git a/src/semble/index/index.py b/src/semble/index/index.py index a1e95bb65..3e11c3784 100644 --- a/src/semble/index/index.py +++ b/src/semble/index/index.py @@ -72,7 +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 hashes and chunk ranges used for incremental reindexing. + :param manifest: File modification times and chunk ranges used for incremental reindexing. """ self.model = model self.chunks: list[Chunk] = chunks @@ -379,7 +379,6 @@ 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, diff --git a/src/semble/index/types.py b/src/semble/index/types.py index 071c036c6..b40f1098d 100644 --- a/src/semble/index/types.py +++ b/src/semble/index/types.py @@ -42,9 +42,9 @@ def from_path(cls: type[PersistencePath], path: Path) -> PersistencePath: @dataclass class FileManifestEntry: - """Records a file's content hash and its chunk range within the global chunk list.""" + """Record a file's modification time and chunk range within the global chunk list.""" - file_hash: str + mtime: float start: int count: int diff --git a/tests/index/test_create.py b/tests/index/test_create.py index ce0e96226..d06b3350d 100644 --- a/tests/index/test_create.py +++ b/tests/index/test_create.py @@ -25,7 +25,7 @@ def test_incremental_reindex_reuses_updates_and_prunes(mock_model: Any, tmp_path _write_files( tmp_path, { - "a.py": "def stable_anchor():\r\n return 1\r\n", + "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", @@ -36,7 +36,6 @@ def test_incremental_reindex_reuses_updates_and_prunes(mock_model: Any, tmp_path ) a_entry = manifest_before["a.py"] b_entry = manifest_before["b.py"] - assert all("\r" not in chunk.content for chunk in chunks_before[a_entry.start : a_entry.start + a_entry.count]) 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( diff --git a/tests/index/test_index.py b/tests/index/test_index.py index 0d653df06..08994cde5 100644 --- a/tests/index/test_index.py +++ b/tests/index/test_index.py @@ -197,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) @@ -206,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 diff --git a/tests/test_cache.py b/tests/test_cache.py index 5ef619676..294ebfa52 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -148,7 +148,7 @@ 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, } @@ -207,7 +207,7 @@ 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, },