crabcode 0.1.0__py3-none-any.whl

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") as f:
+            json.dump(state, f)
+        with open(self._index_dir / "metadata.json", "w") 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) 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) 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

crabcode_search/tool.py

@@ -0,0 +1,232 @@
+"""CodebaseSearchTool — semantic code search as a CrabCode extra tool."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+import json
+
+# Must be set before ANY native library (FAISS / PyTorch) loads libomp.
+os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
+os.environ.setdefault("OMP_NUM_THREADS", "1")
+os.environ.setdefault("MKL_NUM_THREADS", "1")
+
+import asyncio
+from typing import Any
+
+from crabcode_core.types.tool import Tool, ToolContext, ToolResult
+
+from crabcode_search.background import STATUS_FILE_NAME
+from crabcode_search.embedder import create_embedder
+from crabcode_search.indexer import CodebaseIndexer
+
+
+class CodebaseSearchTool(Tool):
+    """Semantic codebase search powered by embeddings and FAISS."""
+
+    name = "CodebaseSearch"
+    description = "Search the codebase using natural language queries. Returns semantically relevant code snippets with file paths and line numbers. Use this when you need to find code by meaning rather than exact text. For exact text/regex matching, use Grep instead."
+    is_read_only = True
+    is_concurrency_safe = True
+    input_schema = {
+        "type": "object",
+        "properties": {
+            "query": {
+                "type": "string",
+                "description": "Natural language search query about the codebase.",
+            },
+            "target_directory": {
+                "type": "string",
+                "description": "Optional subdirectory to limit search scope.",
+            },
+            "num_results": {
+                "type": "integer",
+                "description": "Number of results to return (default 10).",
+                "default": 10,
+            },
+        },
+        "required": ["query"],
+    }
+
+    def __init__(self) -> None:
+        self._indexer: CodebaseIndexer | None = None
+        self._tool_config: dict[str, Any] = {}
+        self._init_lock = asyncio.Lock()
+        self._status_path: Path | None = None
+
+    async def setup(self, context: ToolContext) -> None:
+        await super().setup(context)
+        self._tool_config = context.tool_config
+        self._status_path = Path(context.cwd) / ".crabcode" / "search" / STATUS_FILE_NAME
+
+        status = self._read_background_status()
+        if status:
+            state = status.get("state")
+            if state == "ready":
+                chunks = status.get("chunks")
+                message = "Background index ready"
+                if chunks is not None:
+                    message += f" ({chunks} chunks)"
+                self.emit_event("ready", {"message": message})
+                return
+            if state in {"starting", "preloading", "loading_index", "indexing"}:
+                done = status.get("done")
+                total = status.get("total")
+                percent = None
+                if isinstance(done, int) and isinstance(total, int) and total > 0:
+                    percent = done / total
+                self.emit_event(
+                    "progress",
+                    {
+                        "message": f"Background {state}...",
+                        "percent": percent,
+                    },
+                )
+                return
+
+        await self._ensure_initialized()
+
+    async def _ensure_initialized(self) -> None:
+        if self._indexer is not None:
+            return
+
+        async with self._init_lock:
+            if self._indexer is not None:
+                return
+
+            context = self._setup_context
+            if context is None:
+                return
+
+            config = self._tool_config
+            embedder = create_embedder(
+            backend=config.get("embedder", "ollama"),
+            model=config.get("model"),
+            dimension=config.get("dimension"),
+            batch_size=config.get("batch_size"),
+            device=config.get("device"),
+            )
+            await embedder.preload()
+            self._indexer = CodebaseIndexer(context.cwd, embedder=embedder)
+
+            threads = config.get("threads")
+            if threads is not None:
+                threads = int(threads)
+                os.environ["OMP_NUM_THREADS"] = str(threads)
+                os.environ["MKL_NUM_THREADS"] = str(threads)
+                try:
+                    import torch
+                    torch.set_num_threads(threads)
+                except Exception:
+                    pass
+
+            all_files = self._indexer._scan_files()
+            self._indexer.total_files = len(all_files)
+            if self._indexer.store.count > 0 and not self._indexer._detect_changes(
+                all_files
+            ):
+                self.emit_event("ready", {
+                    "message": f"Index loaded ({self._indexer.store.count} chunks)",
+                })
+                return
+
+            self.emit_event("progress", {"message": "Starting indexing...", "percent": 0})
+            self._background_task = asyncio.create_task(self._build_index())
+
+    def _read_background_status(self) -> dict[str, Any] | None:
+        if self._status_path is None or not self._status_path.exists():
+            return None
+        try:
+            data = json.loads(self._status_path.read_text(encoding="utf-8"))
+        except (OSError, json.JSONDecodeError):
+            return None
+        return data if isinstance(data, dict) else None
+
+    async def _build_index(self) -> None:
+        assert self._indexer is not None
+        async for progress in self._indexer.build_or_update():
+            if progress.total > 0:
+                pct = progress.done / progress.total
+            else:
+                pct = 1.0
+            self.emit_event("progress", {
+                "message": f"Indexing: {progress.done}/{progress.total} files",
+                "percent": pct,
+            })
+        self.emit_event("ready", {
+            "message": f"Index ready ({self._indexer.store.count} chunks, {self._indexer.total_files} files)",
+        })
+
+    async def get_prompt(self, **kwargs: Any) -> str:
+        return (
+            "Search the codebase using natural language queries. "
+            "Returns semantically relevant code snippets with file paths and line numbers. "
+            "Use this when you need to find code by meaning rather than exact text. "
+            "For example: 'where is authentication handled', 'how does the build system work', "
+            "'find the payment processing logic'. "
+            "For exact text/regex matching, use Grep instead. "
+            "For finding files by name, use Glob instead."
+        )
+
+    async def call(
+        self,
+        tool_input: dict[str, Any],
+        context: ToolContext,
+    ) -> ToolResult:
+        if self._indexer is None:
+            await self._ensure_initialized()
+        if self._indexer is None:
+            return ToolResult(
+                result_for_model="CodebaseSearch is not initialized. Use Grep or Glob instead.",
+                is_error=True,
+            )
+
+        query = tool_input.get("query", "")
+        if not query:
+            return ToolResult(result_for_model="query is required.", is_error=True)
+
+        num_results = tool_input.get("num_results", 10)
+        target_dir = tool_input.get("target_directory")
+
+        building = self._background_task and not self._background_task.done()
+
+        if not building:
+            await self._indexer.incremental_update()
+
+        query_vec = await self._indexer.embedder.embed([query], task_type="CODE_RETRIEVAL_QUERY")
+        results = await self._indexer.store.search(query_vec[0], top_k=num_results * 2)
+
+        if target_dir:
+            results = [r for r in results if r.chunk.file_path.startswith(target_dir)]
+
+        results = results[:num_results]
+
+        if not results:
+            msg = "No results found."
+            if building:
+                progress = self._indexer.progress
+                msg += (
+                    f"\n\n[Note: Index still building ({progress.done}/{progress.total} files). "
+                    "Results may be incomplete. Use Grep for precise searches.]"
+                )
+            return ToolResult(result_for_model=msg)
+
+        output_parts: list[str] = []
+        for i, r in enumerate(results, 1):
+            output_parts.append(
+                f"[{i}] {r.chunk.file_path}:{r.chunk.start_line}-{r.chunk.end_line} "
+                f"(score: {r.score:.3f})\n"
+                f"    {r.chunk.signature}\n"
+                f"```\n{r.chunk.content}\n```"
+            )
+
+        output = "\n\n".join(output_parts)
+
+        if building:
+            progress = self._indexer.progress
+            output += (
+                f"\n\n[Note: Index still building ({progress.done}/{progress.total} files). "
+                "Results may be incomplete. Use Grep for precise searches.]"
+            )
+
+        return ToolResult(result_for_model=output)