crabcode 0.1.6__py3-none-win_amd64.whl

crabcode_search/indexer.py

@@ -0,0 +1,215 @@
+"""Codebase indexer — orchestrates scanning, chunking, embedding, and storage."""
+
+from __future__ import annotations
+
+import json
+import subprocess
+from dataclasses import dataclass
+from pathlib import Path
+from typing import AsyncGenerator
+
+from crabcode_core.subprocess_utils import subprocess_group_options
+from crabcode_search.chunker import ChunkMeta, chunk_file, is_indexable
+from crabcode_search.embedder import Embedder
+from crabcode_search.store import VectorStore
+
+BATCH_SIZE = 5
+INDEX_DIR_NAME = ".crabcode/search"
+
+
+@dataclass
+class IndexProgress:
+    done: int
+    total: int
+
+
+class CodebaseIndexer:
+    """Coordinates file scanning, chunking, embedding, and vector storage."""
+
+    def __init__(self, cwd: str, embedder: Embedder):
+        self.cwd = Path(cwd).resolve()
+        self.embedder = embedder
+        self._index_dir = self.cwd / INDEX_DIR_NAME
+        self.store = VectorStore(self._index_dir, dimension=embedder.dimension)
+        self._file_index: dict[str, float] = {}
+        self._file_index_path = self._index_dir / "file_index.json"
+        self.total_files: int = 0
+
+        self._load_file_index()
+        self.store.load()
+
+    @property
+    def progress(self) -> IndexProgress:
+        indexed = len(self._file_index)
+        return IndexProgress(done=indexed, total=max(self.total_files, indexed))
+
+    async def build_or_update(self) -> AsyncGenerator[IndexProgress, None]:
+        """Full/incremental index build, yielding progress after each batch."""
+        all_files = self._scan_files()
+        self.total_files = len(all_files)
+        changed = self._detect_changes(all_files)
+        deleted = self._detect_deleted(all_files)
+
+        await self._purge_deleted(deleted)
+
+        if not changed:
+            yield IndexProgress(done=self.total_files, total=self.total_files)
+            return
+
+        total_changed = len(changed)
+        done = 0
+
+        for batch in _batched(changed, BATCH_SIZE):
+            all_chunks: list[ChunkMeta] = []
+            for file_path in batch:
+                rel = str(file_path.relative_to(self.cwd))
+                await self.store.remove_by_file(rel)
+                try:
+                    content = file_path.read_text(encoding="utf-8", errors="replace")
+                except OSError:
+                    continue
+                chunks = chunk_file(rel, content)
+                all_chunks.extend(chunks)
+
+            if all_chunks:
+                texts = [c.content for c in all_chunks]
+                vectors = await self.embedder.embed(texts, task_type="RETRIEVAL_DOCUMENT")
+                await self.store.add(vectors, all_chunks)
+
+            for file_path in batch:
+                rel = str(file_path.relative_to(self.cwd))
+                try:
+                    self._file_index[rel] = file_path.stat().st_mtime
+                except OSError:
+                    pass
+
+            done += len(batch)
+            yield IndexProgress(done=done, total=total_changed)
+
+        self._save_file_index()
+        self.store.save()
+
+    async def incremental_update(self) -> None:
+        """Quick mtime-based update for changed files (typically < 1s)."""
+        all_files = self._scan_files()
+        changed = self._detect_changes(all_files)
+        deleted = self._detect_deleted(all_files)
+
+        await self._purge_deleted(deleted)
+
+        if not changed:
+            if deleted:
+                self._save_file_index()
+                self.store.save()
+            return
+
+        for file_path in changed:
+            rel = str(file_path.relative_to(self.cwd))
+            await self.store.remove_by_file(rel)
+
+        all_chunks: list[ChunkMeta] = []
+        for file_path in changed:
+            rel = str(file_path.relative_to(self.cwd))
+            try:
+                content = file_path.read_text(encoding="utf-8", errors="replace")
+            except OSError:
+                continue
+            chunks = chunk_file(rel, content)
+            all_chunks.extend(chunks)
+
+        if all_chunks:
+            texts = [c.content for c in all_chunks]
+            vectors = await self.embedder.embed(texts, task_type="RETRIEVAL_DOCUMENT")
+            await self.store.add(vectors, all_chunks)
+
+        for file_path in changed:
+            rel = str(file_path.relative_to(self.cwd))
+            try:
+                self._file_index[rel] = file_path.stat().st_mtime
+            except OSError:
+                pass
+
+        self._save_file_index()
+        self.store.save()
+
+    def _scan_files(self) -> list[Path]:
+        """Scan the codebase for indexable files, respecting .gitignore."""
+        tracked = self._git_tracked_files()
+        if tracked is not None:
+            return [f for f in tracked if is_indexable(f)]
+
+        files: list[Path] = []
+        skip_dirs = {".git", "node_modules", "__pycache__", ".venv", "venv",
+                     ".tox", ".mypy_cache", ".ruff_cache", "dist", "build",
+                     ".crabcode", ".cursor"}
+        for item in self.cwd.rglob("*"):
+            if any(part in skip_dirs for part in item.parts):
+                continue
+            if item.is_file() and is_indexable(item):
+                files.append(item)
+        return files
+
+    def _git_tracked_files(self) -> list[Path] | None:
+        """Use git ls-files to get tracked files (respects .gitignore)."""
+        try:
+            result = subprocess.run(
+                ["git", "ls-files", "-z", "--cached", "--others", "--exclude-standard"],
+                cwd=self.cwd,
+                capture_output=True,
+                text=True,
+                encoding="utf-8",
+                errors="replace",
+                timeout=10,
+                **subprocess_group_options(),
+            )
+            if result.returncode != 0:
+                return None
+            files: list[Path] = []
+            for line in result.stdout.split("\0"):
+                if line:
+                    files.append(self.cwd / line)
+            return files
+        except (OSError, subprocess.TimeoutExpired):
+            return None
+
+    def _detect_changes(self, files: list[Path]) -> list[Path]:
+        """Return files whose mtime changed since last index."""
+        changed: list[Path] = []
+        for f in files:
+            rel = str(f.relative_to(self.cwd))
+            try:
+                mtime = f.stat().st_mtime
+            except OSError:
+                continue
+            if rel not in self._file_index or self._file_index[rel] != mtime:
+                changed.append(f)
+        return changed
+
+    def _detect_deleted(self, current_files: list[Path]) -> list[str]:
+        """Return relative paths that are in file_index but no longer on disk."""
+        current_rels = {str(f.relative_to(self.cwd)) for f in current_files}
+        return [rel for rel in self._file_index if rel not in current_rels]
+
+    async def _purge_deleted(self, deleted: list[str]) -> None:
+        """Remove vectors and file_index entries for deleted files."""
+        for rel in deleted:
+            await self.store.remove_by_file(rel)
+            del self._file_index[rel]
+
+    def _load_file_index(self) -> None:
+        if self._file_index_path.exists():
+            try:
+                with open(self._file_index_path, encoding="utf-8") as f:
+                    self._file_index = json.load(f)
+            except (json.JSONDecodeError, OSError):
+                self._file_index = {}
+
+    def _save_file_index(self) -> None:
+        self._index_dir.mkdir(parents=True, exist_ok=True)
+        with open(self._file_index_path, "w", encoding="utf-8") as f:
+            json.dump(self._file_index, f)
+
+
+def _batched(items: list, size: int):
+    for i in range(0, len(items), size):
+        yield items[i : i + size]

crabcode_search/store.py

@@ -0,0 +1,222 @@
+"""USearch-backed vector store with async-safe concurrent access.
+
+Uses USearch HNSW for all dataset sizes:
+  - count < HNSW_THRESHOLD  → exact brute-force search (exact=True, O(N))
+  - count >= HNSW_THRESHOLD → approximate HNSW search (exact=False, O(log N))
+
+USearch uses SimSIMD instead of OpenMP, which avoids the libomp conflict with
+PyTorch on macOS that made FAISS unusable in the same process.
+
+Persistence:
+  index.usearch  — USearch HNSW index
+  hnsw_state.json — key→metadata-index mapping and next-key counter
+  metadata.json  — list of chunk dicts
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from crabcode_search.chunker import ChunkMeta
+
+logger = logging.getLogger(__name__)
+
+HNSW_THRESHOLD = 100_000
+HNSW_CONNECTIVITY = 16
+HNSW_EXPANSION_ADD = 128
+HNSW_EXPANSION_SEARCH = 64
+
+
+@dataclass
+class SearchResult:
+    chunk: ChunkMeta
+    score: float
+
+
+class VectorStore:
+    """USearch-backed vector store with metadata tracking.
+
+    Small codebases (count < HNSW_THRESHOLD) use exact inner-product search;
+    large codebases use approximate HNSW — both via the same USearch index.
+    """
+
+    def __init__(self, index_dir: Path, dimension: int = 768):
+        self._index_dir = index_dir
+        self._dimension = dimension
+
+        self._index: Any = None          # usearch.index.Index (lazy init)
+        self._metadata: list[dict[str, Any]] = []
+        self._key_to_idx: dict[int, int] = {}   # usearch key → metadata list index
+        self._next_key: int = 0
+        self._lock = asyncio.Lock()
+
+    @property
+    def count(self) -> int:
+        if self._index is None:
+            return 0
+        return len(self._index)
+
+    def _get_or_create_index(self) -> Any:
+        if self._index is None:
+            from usearch.index import Index
+
+            self._index = Index(
+                ndim=self._dimension,
+                metric="ip",
+                dtype="f32",
+                connectivity=HNSW_CONNECTIVITY,
+                expansion_add=HNSW_EXPANSION_ADD,
+                expansion_search=HNSW_EXPANSION_SEARCH,
+            )
+        return self._index
+
+    async def add(self, vectors: np.ndarray, chunks: list[ChunkMeta]) -> None:
+        if vectors.shape[0] == 0:
+            return
+        async with self._lock:
+            idx = self._get_or_create_index()
+            base = len(self._metadata)
+            keys = np.arange(
+                self._next_key, self._next_key + vectors.shape[0], dtype=np.int64
+            )
+            idx.add(keys, vectors)
+            for i, chunk in enumerate(chunks):
+                self._key_to_idx[int(keys[i])] = base + i
+                self._metadata.append({
+                    "file_path": chunk.file_path,
+                    "start_line": chunk.start_line,
+                    "end_line": chunk.end_line,
+                    "signature": chunk.signature,
+                    "content": chunk.content,
+                })
+            self._next_key += vectors.shape[0]
+
+    async def remove_by_file(self, file_path: str) -> None:
+        async with self._lock:
+            if not self._metadata or self._index is None:
+                return
+
+            remove_set = {
+                i for i, m in enumerate(self._metadata) if m["file_path"] == file_path
+            }
+            if not remove_set:
+                return
+
+            keys_to_remove = [k for k, idx in self._key_to_idx.items() if idx in remove_set]
+            for k in keys_to_remove:
+                self._index.remove(k)
+                del self._key_to_idx[k]
+
+            keep = [i for i in range(len(self._metadata)) if i not in remove_set]
+            new_meta = [self._metadata[i] for i in keep]
+            old_to_new = {old: new for new, old in enumerate(keep)}
+            self._key_to_idx = {
+                k: old_to_new[v]
+                for k, v in self._key_to_idx.items()
+                if v in old_to_new
+            }
+            self._metadata = new_meta
+
+    async def search(self, query_vec: np.ndarray, top_k: int = 10) -> list[SearchResult]:
+        async with self._lock:
+            if self._index is None or len(self._index) == 0:
+                return []
+            if query_vec.ndim == 1:
+                query_vec = query_vec.reshape(1, -1)
+            k = min(top_k, len(self._index))
+            exact = len(self._index) < HNSW_THRESHOLD
+            matches = self._index.search(query_vec, k, exact=exact)
+            keys = matches.keys.ravel()
+            distances = matches.distances.ravel()
+
+            results: list[SearchResult] = []
+            for key, dist in zip(keys, distances):
+                key = int(key)
+                if key not in self._key_to_idx:
+                    continue
+                meta_idx = self._key_to_idx[key]
+                if meta_idx >= len(self._metadata):
+                    continue
+                meta = self._metadata[meta_idx]
+                results.append(SearchResult(
+                    chunk=ChunkMeta(**meta),
+                    score=1.0 - float(dist),
+                ))
+            return results
+
+    def save(self) -> None:
+        self._index_dir.mkdir(parents=True, exist_ok=True)
+        if self._index is not None:
+            self._index.save(str(self._index_dir / "index.usearch"))
+        state = {
+            "key_to_idx": {str(k): v for k, v in self._key_to_idx.items()},
+            "next_key": self._next_key,
+        }
+        with open(self._index_dir / "hnsw_state.json", "w", encoding="utf-8") as f:
+            json.dump(state, f)
+        with open(self._index_dir / "metadata.json", "w", encoding="utf-8") as f:
+            json.dump(self._metadata, f)
+        (self._index_dir / "index.npy").unlink(missing_ok=True)
+
+    def load(self) -> bool:
+        meta_path = self._index_dir / "metadata.json"
+        if not meta_path.exists():
+            return False
+        try:
+            with open(meta_path, encoding="utf-8") as f:
+                self._metadata = json.load(f)
+
+            usearch_path = self._index_dir / "index.usearch"
+            if usearch_path.exists():
+                from usearch.index import Index
+
+                idx = Index(
+                    ndim=self._dimension,
+                    metric="ip",
+                    dtype="f32",
+                    connectivity=HNSW_CONNECTIVITY,
+                    expansion_add=HNSW_EXPANSION_ADD,
+                    expansion_search=HNSW_EXPANSION_SEARCH,
+                )
+                idx.load(str(usearch_path))
+                self._index = idx
+                if len(idx) > 0:
+                    self._dimension = idx.ndim
+
+                state_path = self._index_dir / "hnsw_state.json"
+                if state_path.exists():
+                    with open(state_path, encoding="utf-8") as f:
+                        state = json.load(f)
+                    self._key_to_idx = {int(k): v for k, v in state["key_to_idx"].items()}
+                    self._next_key = state["next_key"]
+                else:
+                    n = len(idx)
+                    self._key_to_idx = {k: k for k in range(n)}
+                    self._next_key = n
+                return True
+
+            # Legacy: plain numpy index — migrate to usearch on next save
+            npy_path = self._index_dir / "index.npy"
+            if npy_path.exists():
+                vectors = np.load(npy_path).astype(np.float32)
+                if vectors.shape[0] > 0:
+                    self._dimension = vectors.shape[1]
+                    idx = self._get_or_create_index()
+                    n = vectors.shape[0]
+                    keys = np.arange(n, dtype=np.int64)
+                    idx.add(keys, vectors)
+                    self._key_to_idx = {k: k for k in range(n)}
+                    self._next_key = n
+                return True
+
+            return False
+        except Exception:
+            logger.exception("Failed to load vector store")
+            return False