feat: Add partial reindexing

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",

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

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

src/semble/index/create.py

@@ -2,57 +2,129 @@
 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(
     path: 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