claude-code-workflow 6.3.18 → 6.3.20

package/codex-lens/src/codexlens/semantic/ann_index.py

@@ -9,6 +9,7 @@ Key features:
 - Incremental vector addition and deletion
 - Thread-safe operations
 - Cosine similarity metric
+- Support for centralized storage mode (single index at project root)
 """
 
 from __future__ import annotations
@@ -19,6 +20,7 @@ from pathlib import Path
 from typing import List, Optional, Tuple
 
 from codexlens.errors import StorageError
+from codexlens.config import VECTORS_HNSW_NAME
 
 from . import SEMANTIC_AVAILABLE
 
@@ -127,6 +129,94 @@ class ANNIndex:
             f"auto_save={auto_save}, expansion_threshold={expansion_threshold}"
         )
 
+    @classmethod
+    def create_central(
+        cls,
+        index_root: Path,
+        dim: int,
+        initial_capacity: int = 50000,
+        auto_save: bool = False,
+        expansion_threshold: float = 0.8,
+    ) -> "ANNIndex":
+        """Create a centralized ANN index at the project index root.
+
+        This method creates a single shared HNSW index file at the project root,
+        rather than per-directory indexes. Use this for projects that want all
+        dense vectors stored in one central location.
+
+        Args:
+            index_root: Root directory for the index (e.g., .codexlens/<project_hash>/)
+            dim: Dimension of embedding vectors
+            initial_capacity: Initial maximum elements capacity (default: 50000)
+            auto_save: Whether to automatically save index after operations (default: False)
+            expansion_threshold: Capacity threshold to trigger auto-expansion (default: 0.8)
+
+        Returns:
+            ANNIndex instance configured for centralized storage
+
+        Example:
+            >>> index = ANNIndex.create_central(Path(".codexlens/abc123"), dim=768)
+            >>> index.hnsw_path  # Returns: .codexlens/abc123/_vectors.hnsw
+        """
+        # Create a dummy index_path that will result in the central hnsw_path
+        # The index_path is used to derive hnsw_path, so we create a virtual path
+        # such that self.hnsw_path = index_root / VECTORS_HNSW_NAME
+        instance = cls.__new__(cls)
+
+        if not SEMANTIC_AVAILABLE:
+            raise ImportError(
+                "Semantic search dependencies not available. "
+                "Install with: pip install codexlens[semantic]"
+            )
+
+        if not HNSWLIB_AVAILABLE:
+            raise ImportError(
+                "hnswlib is required for ANN index. "
+                "Install with: pip install hnswlib"
+            )
+
+        if dim <= 0:
+            raise ValueError(f"Invalid dimension: {dim}")
+
+        if initial_capacity <= 0:
+            raise ValueError(f"Invalid initial capacity: {initial_capacity}")
+
+        if not 0.0 < expansion_threshold < 1.0:
+            raise ValueError(
+                f"Invalid expansion threshold: {expansion_threshold}. Must be between 0 and 1."
+            )
+
+        instance.index_path = index_root
+        instance.dim = dim
+
+        # Centralized mode: use VECTORS_HNSW_NAME directly at index_root
+        instance.hnsw_path = index_root / VECTORS_HNSW_NAME
+
+        # HNSW parameters
+        instance.space = "cosine"
+        instance.M = 16
+        instance.ef_construction = 200
+        instance.ef = 50
+
+        # Memory management parameters
+        instance._auto_save = auto_save
+        instance._expansion_threshold = expansion_threshold
+
+        # Thread safety
+        instance._lock = threading.RLock()
+
+        # HNSW index instance
+        instance._index: Optional[hnswlib.Index] = None
+        instance._max_elements = initial_capacity
+        instance._current_count = 0
+
+        logger.info(
+            f"Initialized centralized ANNIndex at {instance.hnsw_path} with "
+            f"capacity={initial_capacity}, auto_save={auto_save}"
+        )
+
+        return instance
+
     def _ensure_index(self) -> None:
         """Ensure HNSW index is initialized (lazy initialization)."""
         if self._index is None:
@@ -412,3 +502,567 @@ class ANNIndex:
         """
         with self._lock:
             return self._index is not None and self._current_count > 0
+
+
+
+class BinaryANNIndex:
+    """Binary vector ANN index using Hamming distance for fast coarse retrieval.
+
+    .. deprecated::
+        This class is deprecated. Use :class:`codexlens.search.binary_searcher.BinarySearcher`
+        instead, which provides faster memory-mapped search with centralized storage.
+
+    Optimized for binary vectors (256-bit / 32 bytes per vector).
+    Uses packed binary representation for memory efficiency.
+
+    Performance characteristics:
+    - Storage: 32 bytes per vector (vs ~8KB for dense vectors)
+    - Distance: Hamming distance via XOR + popcount (CPU-efficient)
+    - Search: O(N) brute-force with SIMD-accelerated distance computation
+
+    Index parameters:
+    - dim: Binary vector dimension (default: 256)
+    - packed_dim: Packed bytes size (dim / 8 = 32 for 256-bit)
+
+    Usage:
+        index = BinaryANNIndex(index_path, dim=256)
+        index.add_vectors([1, 2, 3], packed_vectors)  # List of 32-byte packed vectors
+        ids, distances = index.search(query_packed, top_k=10)
+    """
+
+    DEFAULT_DIM = 256  # Default binary vector dimension
+
+    def __init__(
+        self,
+        index_path: Path,
+        dim: int = 256,
+        initial_capacity: int = 100000,
+        auto_save: bool = False,
+    ) -> None:
+        """Initialize Binary ANN index.
+
+        Args:
+            index_path: Path to database (index will be saved as _binary_vectors.bin)
+            dim: Dimension of binary vectors (default: 256)
+            initial_capacity: Initial capacity hint (default: 100000)
+            auto_save: Whether to automatically save index after operations
+
+        Raises:
+            ImportError: If required dependencies are not available
+            ValueError: If dimension is invalid
+        """
+        if not SEMANTIC_AVAILABLE:
+            raise ImportError(
+                "Semantic search dependencies not available. "
+                "Install with: pip install codexlens[semantic]"
+            )
+
+        import warnings
+        warnings.warn(
+            "BinaryANNIndex is deprecated. Use codexlens.search.binary_searcher.BinarySearcher "
+            "instead for faster memory-mapped search with centralized storage.",
+            DeprecationWarning,
+            stacklevel=2
+        )
+
+        if dim <= 0 or dim % 8 != 0:
+            raise ValueError(
+                f"Invalid dimension: {dim}. Must be positive and divisible by 8."
+            )
+
+        self.index_path = Path(index_path)
+        self.dim = dim
+        self.packed_dim = dim // 8  # 32 bytes for 256-bit vectors
+
+        # Derive binary index path from database path
+        db_stem = self.index_path.stem
+        self.binary_path = self.index_path.parent / f"{db_stem}_binary_vectors.bin"
+
+        # Memory management
+        self._auto_save = auto_save
+        self._initial_capacity = initial_capacity
+
+        # Thread safety
+        self._lock = threading.RLock()
+
+        # In-memory storage: id -> packed binary vector
+        self._vectors: dict[int, bytes] = {}
+        self._id_list: list[int] = []  # Ordered list for efficient iteration
+
+        # Cached numpy array for vectorized search (invalidated on add/remove)
+        self._vectors_matrix: Optional[np.ndarray] = None
+        self._ids_array: Optional[np.ndarray] = None
+        self._cache_valid: bool = False
+
+        logger.info(
+            f"Initialized BinaryANNIndex with dim={dim}, packed_dim={self.packed_dim}"
+        )
+
+    def add_vectors(self, ids: List[int], vectors: List[bytes]) -> None:
+        """Add packed binary vectors to the index.
+
+        Args:
+            ids: List of vector IDs (must be unique)
+            vectors: List of packed binary vectors (each of size packed_dim bytes)
+
+        Raises:
+            ValueError: If shapes don't match or vectors are invalid
+            StorageError: If index operation fails
+        """
+        if len(ids) == 0:
+            return
+
+        if len(vectors) != len(ids):
+            raise ValueError(
+                f"Number of vectors ({len(vectors)}) must match number of IDs ({len(ids)})"
+            )
+
+        # Validate vector sizes
+        for i, vec in enumerate(vectors):
+            if len(vec) != self.packed_dim:
+                raise ValueError(
+                    f"Vector {i} has size {len(vec)}, expected {self.packed_dim}"
+                )
+
+        with self._lock:
+            try:
+                for vec_id, vec in zip(ids, vectors):
+                    if vec_id not in self._vectors:
+                        self._id_list.append(vec_id)
+                    self._vectors[vec_id] = vec
+
+                # Invalidate cache on modification
+                self._cache_valid = False
+
+                logger.debug(
+                    f"Added {len(ids)} binary vectors to index (total: {len(self._vectors)})"
+                )
+
+                if self._auto_save:
+                    self.save()
+
+            except Exception as e:
+                raise StorageError(f"Failed to add vectors to Binary ANN index: {e}")
+
+    def add_vectors_numpy(self, ids: List[int], vectors: np.ndarray) -> None:
+        """Add unpacked binary vectors (0/1 values) to the index.
+
+        Convenience method that packs the vectors before adding.
+
+        Args:
+            ids: List of vector IDs (must be unique)
+            vectors: Numpy array of shape (N, dim) with binary values (0 or 1)
+
+        Raises:
+            ValueError: If shapes don't match
+            StorageError: If index operation fails
+        """
+        if len(ids) == 0:
+            return
+
+        if vectors.shape[0] != len(ids):
+            raise ValueError(
+                f"Number of vectors ({vectors.shape[0]}) must match number of IDs ({len(ids)})"
+            )
+
+        if vectors.shape[1] != self.dim:
+            raise ValueError(
+                f"Vector dimension ({vectors.shape[1]}) must match index dimension ({self.dim})"
+            )
+
+        # Pack vectors
+        packed_vectors = []
+        for i in range(vectors.shape[0]):
+            packed = np.packbits(vectors[i].astype(np.uint8)).tobytes()
+            packed_vectors.append(packed)
+
+        self.add_vectors(ids, packed_vectors)
+
+    def remove_vectors(self, ids: List[int]) -> None:
+        """Remove vectors from the index.
+
+        Args:
+            ids: List of vector IDs to remove
+
+        Raises:
+            StorageError: If index operation fails
+
+        Note:
+            Optimized for batch deletion using set operations instead of
+            O(N) list.remove() calls for each ID.
+        """
+        if len(ids) == 0:
+            return
+
+        with self._lock:
+            try:
+                # Use set for O(1) lookup during filtering
+                ids_to_remove = set(ids)
+                removed_count = 0
+
+                # Remove from dictionary - O(1) per deletion
+                for vec_id in ids_to_remove:
+                    if vec_id in self._vectors:
+                        del self._vectors[vec_id]
+                        removed_count += 1
+
+                # Rebuild ID list efficiently - O(N) once instead of O(N) per removal
+                if removed_count > 0:
+                    self._id_list = [id_ for id_ in self._id_list if id_ not in ids_to_remove]
+                    # Invalidate cache on modification
+                    self._cache_valid = False
+
+                logger.debug(f"Removed {removed_count}/{len(ids)} vectors from index")
+
+                if self._auto_save and removed_count > 0:
+                    self.save()
+
+            except Exception as e:
+                raise StorageError(
+                    f"Failed to remove vectors from Binary ANN index: {e}"
+                )
+
+    def _build_cache(self) -> None:
+        """Build numpy array cache from vectors dict for vectorized search.
+
+        Pre-computes a contiguous numpy array from all vectors for efficient
+        batch distance computation. Called lazily on first search after modification.
+        """
+        if self._cache_valid:
+            return
+
+        n_vectors = len(self._id_list)
+        if n_vectors == 0:
+            self._vectors_matrix = None
+            self._ids_array = None
+            self._cache_valid = True
+            return
+
+        # Build contiguous numpy array of all packed vectors
+        # Shape: (n_vectors, packed_dim) with uint8 dtype
+        self._vectors_matrix = np.empty((n_vectors, self.packed_dim), dtype=np.uint8)
+        self._ids_array = np.array(self._id_list, dtype=np.int64)
+
+        for i, vec_id in enumerate(self._id_list):
+            vec_bytes = self._vectors[vec_id]
+            self._vectors_matrix[i] = np.frombuffer(vec_bytes, dtype=np.uint8)
+
+        self._cache_valid = True
+        logger.debug(f"Built vectorized cache for {n_vectors} binary vectors")
+
+    def search(
+        self, query: bytes, top_k: int = 10
+    ) -> Tuple[List[int], List[int]]:
+        """Search for nearest neighbors using Hamming distance.
+
+        Uses vectorized batch computation for O(N) search with SIMD acceleration.
+        Pre-computes and caches numpy arrays for efficient repeated queries.
+
+        Args:
+            query: Packed binary query vector (size: packed_dim bytes)
+            top_k: Number of nearest neighbors to return
+
+        Returns:
+            Tuple of (ids, distances) where:
+            - ids: List of vector IDs ordered by Hamming distance (ascending)
+            - distances: List of Hamming distances (lower = more similar)
+
+        Raises:
+            ValueError: If query size is invalid
+            StorageError: If search operation fails
+        """
+        if len(query) != self.packed_dim:
+            raise ValueError(
+                f"Query size ({len(query)}) must match packed_dim ({self.packed_dim})"
+            )
+
+        with self._lock:
+            try:
+                if len(self._vectors) == 0:
+                    return [], []
+
+                # Build cache if needed (lazy initialization)
+                self._build_cache()
+
+                if self._vectors_matrix is None or self._ids_array is None:
+                    return [], []
+
+                # Vectorized Hamming distance computation
+                # 1. Convert query to numpy array
+                query_arr = np.frombuffer(query, dtype=np.uint8)
+
+                # 2. Broadcast XOR: (1, packed_dim) XOR (n_vectors, packed_dim)
+                #    Result shape: (n_vectors, packed_dim)
+                xor_result = np.bitwise_xor(query_arr, self._vectors_matrix)
+
+                # 3. Vectorized popcount using lookup table for efficiency
+                #    np.unpackbits is slow for large arrays, use popcount LUT instead
+                popcount_lut = np.array([bin(i).count('1') for i in range(256)], dtype=np.uint8)
+                bit_counts = popcount_lut[xor_result]
+
+                # 4. Sum across packed bytes to get Hamming distance per vector
+                distances = bit_counts.sum(axis=1)
+
+                # 5. Get top-k using argpartition (O(N) instead of O(N log N) for full sort)
+                n_vectors = len(distances)
+                k = min(top_k, n_vectors)
+
+                if k == n_vectors:
+                    # No partitioning needed, just sort all
+                    sorted_indices = np.argsort(distances)
+                else:
+                    # Use argpartition for O(N) partial sort
+                    partition_indices = np.argpartition(distances, k)[:k]
+                    # Sort only the top-k
+                    top_k_distances = distances[partition_indices]
+                    sorted_order = np.argsort(top_k_distances)
+                    sorted_indices = partition_indices[sorted_order]
+
+                # 6. Return results
+                result_ids = self._ids_array[sorted_indices].tolist()
+                result_dists = distances[sorted_indices].tolist()
+
+                return result_ids, result_dists
+
+            except Exception as e:
+                raise StorageError(f"Failed to search Binary ANN index: {e}")
+
+    def search_numpy(
+        self, query: np.ndarray, top_k: int = 10
+    ) -> Tuple[List[int], List[int]]:
+        """Search with unpacked binary query vector.
+
+        Convenience method that packs the query before searching.
+
+        Args:
+            query: Binary query vector of shape (dim,) with values 0 or 1
+            top_k: Number of nearest neighbors to return
+
+        Returns:
+            Tuple of (ids, distances)
+        """
+        if query.ndim == 2:
+            query = query.flatten()
+
+        if len(query) != self.dim:
+            raise ValueError(
+                f"Query dimension ({len(query)}) must match index dimension ({self.dim})"
+            )
+
+        packed_query = np.packbits(query.astype(np.uint8)).tobytes()
+        return self.search(packed_query, top_k)
+
+    def search_batch(
+        self, queries: List[bytes], top_k: int = 10
+    ) -> List[Tuple[List[int], List[int]]]:
+        """Batch search for multiple queries.
+
+        Args:
+            queries: List of packed binary query vectors
+            top_k: Number of nearest neighbors to return per query
+
+        Returns:
+            List of (ids, distances) tuples, one per query
+        """
+        results = []
+        for query in queries:
+            ids, dists = self.search(query, top_k)
+            results.append((ids, dists))
+        return results
+
+    def save(self) -> None:
+        """Save index to disk.
+
+        Binary format:
+        - 4 bytes: magic number (0x42494E56 = "BINV")
+        - 4 bytes: version (1)
+        - 4 bytes: dim
+        - 4 bytes: packed_dim
+        - 4 bytes: num_vectors
+        - For each vector:
+          - 4 bytes: id
+          - packed_dim bytes: vector data
+
+        Raises:
+            StorageError: If save operation fails
+        """
+        with self._lock:
+            try:
+                if len(self._vectors) == 0:
+                    logger.debug("Skipping save: index is empty")
+                    return
+
+                # Ensure parent directory exists
+                self.binary_path.parent.mkdir(parents=True, exist_ok=True)
+
+                with open(self.binary_path, "wb") as f:
+                    # Header
+                    f.write(b"BINV")  # Magic number
+                    f.write(np.array([1], dtype=np.uint32).tobytes())  # Version
+                    f.write(np.array([self.dim], dtype=np.uint32).tobytes())
+                    f.write(np.array([self.packed_dim], dtype=np.uint32).tobytes())
+                    f.write(
+                        np.array([len(self._vectors)], dtype=np.uint32).tobytes()
+                    )
+
+                    # Vectors
+                    for vec_id in self._id_list:
+                        f.write(np.array([vec_id], dtype=np.uint32).tobytes())
+                        f.write(self._vectors[vec_id])
+
+                logger.debug(
+                    f"Saved binary index to {self.binary_path} "
+                    f"({len(self._vectors)} vectors)"
+                )
+
+            except Exception as e:
+                raise StorageError(f"Failed to save Binary ANN index: {e}")
+
+    def load(self) -> bool:
+        """Load index from disk.
+
+        Returns:
+            True if index was loaded successfully, False if index file doesn't exist
+
+        Raises:
+            StorageError: If load operation fails
+        """
+        with self._lock:
+            try:
+                if not self.binary_path.exists():
+                    logger.debug(f"Binary index file not found: {self.binary_path}")
+                    return False
+
+                with open(self.binary_path, "rb") as f:
+                    # Read header
+                    magic = f.read(4)
+                    if magic != b"BINV":
+                        raise StorageError(
+                            f"Invalid binary index file: bad magic number"
+                        )
+
+                    version = np.frombuffer(f.read(4), dtype=np.uint32)[0]
+                    if version != 1:
+                        raise StorageError(
+                            f"Unsupported binary index version: {version}"
+                        )
+
+                    file_dim = np.frombuffer(f.read(4), dtype=np.uint32)[0]
+                    file_packed_dim = np.frombuffer(f.read(4), dtype=np.uint32)[0]
+                    num_vectors = np.frombuffer(f.read(4), dtype=np.uint32)[0]
+
+                    if file_dim != self.dim or file_packed_dim != self.packed_dim:
+                        raise StorageError(
+                            f"Dimension mismatch: file has dim={file_dim}, "
+                            f"packed_dim={file_packed_dim}, "
+                            f"expected dim={self.dim}, packed_dim={self.packed_dim}"
+                        )
+
+                    # Clear existing data
+                    self._vectors.clear()
+                    self._id_list.clear()
+                    self._cache_valid = False
+
+                    # Read vectors
+                    for _ in range(num_vectors):
+                        vec_id = np.frombuffer(f.read(4), dtype=np.uint32)[0]
+                        vec_data = f.read(self.packed_dim)
+                        self._vectors[int(vec_id)] = vec_data
+                        self._id_list.append(int(vec_id))
+
+                logger.info(
+                    f"Loaded binary index from {self.binary_path} "
+                    f"({len(self._vectors)} vectors)"
+                )
+
+                return True
+
+            except StorageError:
+                raise
+            except Exception as e:
+                raise StorageError(f"Failed to load Binary ANN index: {e}")
+
+    def count(self) -> int:
+        """Get number of vectors in the index.
+
+        Returns:
+            Number of vectors currently in the index
+        """
+        with self._lock:
+            return len(self._vectors)
+
+    @property
+    def is_loaded(self) -> bool:
+        """Check if index has vectors.
+
+        Returns:
+            True if index has vectors, False otherwise
+        """
+        with self._lock:
+            return len(self._vectors) > 0
+
+    def get_vector(self, vec_id: int) -> Optional[bytes]:
+        """Get a specific vector by ID.
+
+        Args:
+            vec_id: Vector ID to retrieve
+
+        Returns:
+            Packed binary vector or None if not found
+        """
+        with self._lock:
+            return self._vectors.get(vec_id)
+
+    def clear(self) -> None:
+        """Clear all vectors from the index."""
+        with self._lock:
+            self._vectors.clear()
+            self._id_list.clear()
+            self._vectors_matrix = None
+            self._ids_array = None
+            self._cache_valid = False
+            logger.debug("Cleared binary index")
+
+
+def create_ann_index(
+    index_path: Path,
+    index_type: str = "hnsw",
+    dim: int = 2048,
+    **kwargs,
+) -> ANNIndex | BinaryANNIndex:
+    """Factory function to create an ANN index.
+
+    Args:
+        index_path: Path to database file
+        index_type: Type of index - "hnsw" for dense vectors, "binary" for binary vectors
+        dim: Vector dimension (default: 2048 for dense, 256 for binary)
+        **kwargs: Additional arguments passed to the index constructor
+
+    Returns:
+        ANNIndex for dense vectors or BinaryANNIndex for binary vectors
+
+    Raises:
+        ValueError: If index_type is invalid
+
+    Example:
+        >>> # Dense vector index (HNSW)
+        >>> dense_index = create_ann_index(path, index_type="hnsw", dim=2048)
+        >>> dense_index.add_vectors(ids, dense_vectors)
+        >>>
+        >>> # Binary vector index (Hamming distance)
+        >>> binary_index = create_ann_index(path, index_type="binary", dim=256)
+        >>> binary_index.add_vectors(ids, packed_vectors)
+    """
+    index_type = index_type.lower()
+
+    if index_type == "hnsw":
+        return ANNIndex(index_path=index_path, dim=dim, **kwargs)
+    elif index_type == "binary":
+        # Default to 256 for binary if not specified
+        if dim == 2048:  # Default dense dim was used
+            dim = 256
+        return BinaryANNIndex(index_path=index_path, dim=dim, **kwargs)
+    else:
+        raise ValueError(
+            f"Invalid index_type: {index_type}. Must be 'hnsw' or 'binary'."
+        )