claude-code-workflow 6.3.18 → 6.3.20

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

@@ -9,89 +9,96 @@ Optimized for high-performance similarity search using:
 
 from __future__ import annotations
 
-import json
-import logging
-import sys
-import sqlite3
-import threading
-from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple
-
-from codexlens.entities import SearchResult, SemanticChunk
-from codexlens.errors import StorageError
-
-try:
-    import numpy as np
-    NUMPY_AVAILABLE = True
-except ImportError:
-    np = None  # type: ignore[assignment]
-    NUMPY_AVAILABLE = False
-
-# Try to import ANN index (optional hnswlib dependency)
-try:
-    from codexlens.semantic.ann_index import ANNIndex, HNSWLIB_AVAILABLE
+import json
+import logging
+import sys
+import sqlite3
+import threading
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+from codexlens.entities import SearchResult, SemanticChunk
+from codexlens.errors import StorageError
+
+try:
+    import numpy as np
+    NUMPY_AVAILABLE = True
+except ImportError:
+    np = None  # type: ignore[assignment]
+    NUMPY_AVAILABLE = False
+
+# Try to import ANN index (optional hnswlib dependency)
+try:
+    from codexlens.semantic.ann_index import (
+        ANNIndex,
+        BinaryANNIndex,
+        create_ann_index,
+        HNSWLIB_AVAILABLE,
+    )
 except ImportError:
     HNSWLIB_AVAILABLE = False
     ANNIndex = None
+    BinaryANNIndex = None
+    create_ann_index = None
+
+
+logger = logging.getLogger(__name__)
+
+# Epsilon used to guard against floating point precision edge cases (e.g., near-zero norms).
+EPSILON = 1e-10
+
+# SQLite INTEGER PRIMARY KEY uses signed 64-bit rowids.
+SQLITE_INTEGER_MAX = (1 << 63) - 1
+
+
+def _validate_chunk_id_range(start_id: int, count: int) -> None:
+    """Validate that a batch insert can safely generate sequential chunk IDs."""
+    if count <= 0:
+        return
+
+    last_id = start_id + count - 1
+    if last_id > sys.maxsize or last_id > SQLITE_INTEGER_MAX:
+        raise ValueError(
+            "Chunk ID range overflow: "
+            f"start_id={start_id}, count={count} would allocate up to {last_id}, "
+            f"exceeding limits (sys.maxsize={sys.maxsize}, sqlite_max={SQLITE_INTEGER_MAX}). "
+            "Consider cleaning up the index database or creating a new index database."
+        )
+
+
+def _validate_sql_placeholders(placeholders: str, expected_count: int) -> None:
+    """Validate the placeholder string used for a parameterized SQL IN clause."""
+    expected = ",".join("?" * expected_count)
+    if placeholders != expected:
+        raise ValueError(
+            "Invalid SQL placeholders for IN clause. "
+            f"Expected {expected_count} '?' placeholders."
+        )
+
+
+def _cosine_similarity(a: List[float], b: List[float]) -> float:
+    """Compute cosine similarity between two vectors."""
+    if not NUMPY_AVAILABLE:
+        raise ImportError("numpy required for vector operations")
+
+    a_arr = np.array(a)
+    b_arr = np.array(b)
+
+    norm_a = np.linalg.norm(a_arr)
+    norm_b = np.linalg.norm(b_arr)
+
+    # Use epsilon tolerance to avoid division by (near-)zero due to floating point precision.
+    if norm_a < EPSILON or norm_b < EPSILON:
+        return 0.0
+
+    denom = norm_a * norm_b
+    if denom < EPSILON:
+        return 0.0
 
-
-logger = logging.getLogger(__name__)
-
-# Epsilon used to guard against floating point precision edge cases (e.g., near-zero norms).
-EPSILON = 1e-10
-
-# SQLite INTEGER PRIMARY KEY uses signed 64-bit rowids.
-SQLITE_INTEGER_MAX = (1 << 63) - 1
-
-
-def _validate_chunk_id_range(start_id: int, count: int) -> None:
-    """Validate that a batch insert can safely generate sequential chunk IDs."""
-    if count <= 0:
-        return
-
-    last_id = start_id + count - 1
-    if last_id > sys.maxsize or last_id > SQLITE_INTEGER_MAX:
-        raise ValueError(
-            "Chunk ID range overflow: "
-            f"start_id={start_id}, count={count} would allocate up to {last_id}, "
-            f"exceeding limits (sys.maxsize={sys.maxsize}, sqlite_max={SQLITE_INTEGER_MAX}). "
-            "Consider cleaning up the index database or creating a new index database."
-        )
-
-
-def _validate_sql_placeholders(placeholders: str, expected_count: int) -> None:
-    """Validate the placeholder string used for a parameterized SQL IN clause."""
-    expected = ",".join("?" * expected_count)
-    if placeholders != expected:
-        raise ValueError(
-            "Invalid SQL placeholders for IN clause. "
-            f"Expected {expected_count} '?' placeholders."
-        )
-
-
-def _cosine_similarity(a: List[float], b: List[float]) -> float:
-    """Compute cosine similarity between two vectors."""
-    if not NUMPY_AVAILABLE:
-        raise ImportError("numpy required for vector operations")
-
-    a_arr = np.array(a)
-    b_arr = np.array(b)
-
-    norm_a = np.linalg.norm(a_arr)
-    norm_b = np.linalg.norm(b_arr)
-
-    # Use epsilon tolerance to avoid division by (near-)zero due to floating point precision.
-    if norm_a < EPSILON or norm_b < EPSILON:
-        return 0.0
-
-    denom = norm_a * norm_b
-    if denom < EPSILON:
-        return 0.0
-
-    return float(np.dot(a_arr, b_arr) / denom)
-
-
-class VectorStore:
+    return float(np.dot(a_arr, b_arr) / denom)
+
+
+class VectorStore:
     """SQLite-based vector storage with HNSW-accelerated similarity search.
 
     Performance optimizations:
@@ -106,12 +113,12 @@ class VectorStore:
     # Default embedding dimension (used when creating new index)
     DEFAULT_DIM = 768
 
-    def __init__(self, db_path: str | Path) -> None:
-        if not NUMPY_AVAILABLE:
-            raise ImportError(
-                "Semantic search dependencies not available. "
-                "Install with: pip install codexlens[semantic]"
-            )
+    def __init__(self, db_path: str | Path) -> None:
+        if not NUMPY_AVAILABLE:
+            raise ImportError(
+                "Semantic search dependencies not available. "
+                "Install with: pip install codexlens[semantic]"
+            )
 
         self.db_path = Path(db_path)
         self.db_path.parent.mkdir(parents=True, exist_ok=True)
@@ -148,6 +155,7 @@ class VectorStore:
                     content TEXT NOT NULL,
                     embedding BLOB NOT NULL,
                     metadata TEXT,
+                    category TEXT DEFAULT 'code',
                     created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                 )
             """)
@@ -155,6 +163,10 @@ class VectorStore:
                 CREATE INDEX IF NOT EXISTS idx_chunks_file
                 ON semantic_chunks(file_path)
             """)
+            conn.execute("""
+                CREATE INDEX IF NOT EXISTS idx_chunks_category
+                ON semantic_chunks(category)
+            """)
             # Model configuration table - tracks which model generated the embeddings
             conn.execute("""
                 CREATE TABLE IF NOT EXISTS embeddings_config (
@@ -170,6 +182,8 @@ class VectorStore:
 
             # Migration: Add backend column to existing tables
             self._migrate_backend_column(conn)
+            # Migration: Add category column
+            self._migrate_category_column(conn)
 
             conn.commit()
 
@@ -190,6 +204,28 @@ class VectorStore:
                 ADD COLUMN backend TEXT NOT NULL DEFAULT 'fastembed'
             """)
 
+    def _migrate_category_column(self, conn: sqlite3.Connection) -> None:
+        """Add category column to existing semantic_chunks table if not present.
+
+        Args:
+            conn: Active SQLite connection
+        """
+        # Check if category column exists
+        cursor = conn.execute("PRAGMA table_info(semantic_chunks)")
+        columns = [row[1] for row in cursor.fetchall()]
+
+        if 'category' not in columns:
+            logger.info("Migrating semantic_chunks table: adding category column")
+            conn.execute("""
+                ALTER TABLE semantic_chunks
+                ADD COLUMN category TEXT DEFAULT 'code'
+            """)
+            # Create index for fast category filtering
+            conn.execute("""
+                CREATE INDEX IF NOT EXISTS idx_chunks_category
+                ON semantic_chunks(category)
+            """)
+
     def _init_ann_index(self) -> None:
         """Initialize ANN index (lazy loading from existing data)."""
         if not HNSWLIB_AVAILABLE:
@@ -338,14 +374,14 @@ class VectorStore:
             ]
             self._embedding_matrix = np.vstack(embeddings)
 
-            # Pre-compute norms for faster similarity calculation
-            self._embedding_norms = np.linalg.norm(
-                self._embedding_matrix, axis=1, keepdims=True
-            )
-            # Avoid division by zero
-            self._embedding_norms = np.where(
-                self._embedding_norms == 0, EPSILON, self._embedding_norms
-            )
+            # Pre-compute norms for faster similarity calculation
+            self._embedding_norms = np.linalg.norm(
+                self._embedding_matrix, axis=1, keepdims=True
+            )
+            # Avoid division by zero
+            self._embedding_norms = np.where(
+                self._embedding_norms == 0, EPSILON, self._embedding_norms
+            )
 
             return True
 
@@ -383,9 +419,16 @@ class VectorStore:
                 self._ann_index = None
                 return False
 
-    def add_chunk(self, chunk: SemanticChunk, file_path: str) -> int:
+    def add_chunk(
+        self, chunk: SemanticChunk, file_path: str, category: str = "code"
+    ) -> int:
         """Add a single chunk with its embedding.
 
+        Args:
+            chunk: SemanticChunk with embedding
+            file_path: Path to the source file
+            category: File category ('code' or 'doc'), default 'code'
+
         Returns:
             The inserted chunk ID.
         """
@@ -399,10 +442,10 @@ class VectorStore:
         with sqlite3.connect(self.db_path) as conn:
             cursor = conn.execute(
                 """
-                INSERT INTO semantic_chunks (file_path, content, embedding, metadata)
-                VALUES (?, ?, ?, ?)
+                INSERT INTO semantic_chunks (file_path, content, embedding, metadata, category)
+                VALUES (?, ?, ?, ?, ?)
                 """,
-                (file_path, chunk.content, embedding_blob, metadata_json)
+                (file_path, chunk.content, embedding_blob, metadata_json, category)
             )
             conn.commit()
             chunk_id = cursor.lastrowid or 0
@@ -420,9 +463,16 @@ class VectorStore:
         self._invalidate_cache()
         return chunk_id
 
-    def add_chunks(self, chunks: List[SemanticChunk], file_path: str) -> List[int]:
+    def add_chunks(
+        self, chunks: List[SemanticChunk], file_path: str, category: str = "code"
+    ) -> List[int]:
         """Add multiple chunks with embeddings (batch insert).
 
+        Args:
+            chunks: List of SemanticChunk objects with embeddings
+            file_path: Path to the source file
+            category: File category ('code' or 'doc'), default 'code'
+
         Returns:
             List of inserted chunk IDs.
         """
@@ -438,7 +488,7 @@ class VectorStore:
             embedding_arr = np.array(chunk.embedding, dtype=np.float32)
             embedding_blob = embedding_arr.tobytes()
             metadata_json = json.dumps(chunk.metadata) if chunk.metadata else None
-            batch_data.append((file_path, chunk.content, embedding_blob, metadata_json))
+            batch_data.append((file_path, chunk.content, embedding_blob, metadata_json, category))
             embeddings_list.append(embedding_arr)
 
         # Batch insert to SQLite
@@ -449,8 +499,8 @@ class VectorStore:
 
             conn.executemany(
                 """
-                INSERT INTO semantic_chunks (file_path, content, embedding, metadata)
-                VALUES (?, ?, ?, ?)
+                INSERT INTO semantic_chunks (file_path, content, embedding, metadata, category)
+                VALUES (?, ?, ?, ?, ?)
                 """,
                 batch_data
             )
@@ -472,11 +522,12 @@ class VectorStore:
         self._invalidate_cache()
         return ids
 
-    def add_chunks_batch(
-        self,
-        chunks_with_paths: List[Tuple[SemanticChunk, str]],
-        update_ann: bool = True,
-        auto_save_ann: bool = True,
+    def add_chunks_batch(
+        self,
+        chunks_with_paths: List[Tuple[SemanticChunk, str]],
+        update_ann: bool = True,
+        auto_save_ann: bool = True,
+        categories: Optional[List[str]] = None,
     ) -> List[int]:
         """Batch insert chunks from multiple files in a single transaction.
 
@@ -487,19 +538,28 @@ class VectorStore:
             update_ann: If True, update ANN index with new vectors (default: True)
             auto_save_ann: If True, save ANN index after update (default: True).
                           Set to False for bulk inserts to reduce I/O overhead.
+            categories: Optional list of categories per chunk. If None, defaults to 'code'.
+                       If provided, must match length of chunks_with_paths.
 
-        Returns:
-            List of inserted chunk IDs
-        """
-        if not chunks_with_paths:
-            return []
-
-        batch_size = len(chunks_with_paths)
-
-        # Prepare batch data
-        batch_data = []
-        embeddings_list = []
-        for chunk, file_path in chunks_with_paths:
+        Returns:
+            List of inserted chunk IDs
+        """
+        if not chunks_with_paths:
+            return []
+
+        batch_size = len(chunks_with_paths)
+
+        # Validate categories if provided
+        if categories is not None and len(categories) != batch_size:
+            raise ValueError(
+                f"categories length ({len(categories)}) must match "
+                f"chunks_with_paths length ({batch_size})"
+            )
+
+        # Prepare batch data
+        batch_data = []
+        embeddings_list = []
+        for i, (chunk, file_path) in enumerate(chunks_with_paths):
             if chunk.embedding is None:
                 raise ValueError("All chunks must have embeddings")
             # Optimize: avoid repeated np.array() if already numpy
@@ -509,55 +569,57 @@ class VectorStore:
                 embedding_arr = np.array(chunk.embedding, dtype=np.float32)
             embedding_blob = embedding_arr.tobytes()
             metadata_json = json.dumps(chunk.metadata) if chunk.metadata else None
-            batch_data.append((file_path, chunk.content, embedding_blob, metadata_json))
+            category = categories[i] if categories else "code"
+            batch_data.append((file_path, chunk.content, embedding_blob, metadata_json, category))
             embeddings_list.append(embedding_arr)
 
-        # Batch insert to SQLite in single transaction
-        with sqlite3.connect(self.db_path) as conn:
-            # Get starting ID before insert
-            row = conn.execute("SELECT MAX(id) FROM semantic_chunks").fetchone()
-            start_id = (row[0] or 0) + 1
-
-            _validate_chunk_id_range(start_id, batch_size)
-
-            conn.executemany(
-                """
-                INSERT INTO semantic_chunks (file_path, content, embedding, metadata)
-                VALUES (?, ?, ?, ?)
+        # Batch insert to SQLite in single transaction
+        with sqlite3.connect(self.db_path) as conn:
+            # Get starting ID before insert
+            row = conn.execute("SELECT MAX(id) FROM semantic_chunks").fetchone()
+            start_id = (row[0] or 0) + 1
+
+            _validate_chunk_id_range(start_id, batch_size)
+
+            conn.executemany(
+                """
+                INSERT INTO semantic_chunks (file_path, content, embedding, metadata, category)
+                VALUES (?, ?, ?, ?, ?)
                 """,
                 batch_data
-            )
-            conn.commit()
-            # Calculate inserted IDs based on starting ID
-            ids = list(range(start_id, start_id + batch_size))
-
-        # Handle ANN index updates
-        if embeddings_list and update_ann and self._ensure_ann_index(len(embeddings_list[0])):
-            with self._ann_write_lock:
-                # In bulk insert mode, accumulate for later batch update
-                if self._bulk_insert_mode:
-                    self._bulk_insert_ids.extend(ids)
-                    self._bulk_insert_embeddings.extend(embeddings_list)
-                else:
-                    # Normal mode: update immediately
-                    try:
-                        embeddings_matrix = np.vstack(embeddings_list)
-                        self._ann_index.add_vectors(ids, embeddings_matrix)
-                        if auto_save_ann:
-                            self._ann_index.save()
-                    except Exception as e:
-                        logger.warning("Failed to add batch to ANN index: %s", e)
+            )
+            conn.commit()
+            # Calculate inserted IDs based on starting ID
+            ids = list(range(start_id, start_id + batch_size))
+
+        # Handle ANN index updates
+        if embeddings_list and update_ann and self._ensure_ann_index(len(embeddings_list[0])):
+            with self._ann_write_lock:
+                # In bulk insert mode, accumulate for later batch update
+                if self._bulk_insert_mode:
+                    self._bulk_insert_ids.extend(ids)
+                    self._bulk_insert_embeddings.extend(embeddings_list)
+                else:
+                    # Normal mode: update immediately
+                    try:
+                        embeddings_matrix = np.vstack(embeddings_list)
+                        self._ann_index.add_vectors(ids, embeddings_matrix)
+                        if auto_save_ann:
+                            self._ann_index.save()
+                    except Exception as e:
+                        logger.warning("Failed to add batch to ANN index: %s", e)
 
         # Invalidate cache after modification
         self._invalidate_cache()
         return ids
 
-    def add_chunks_batch_numpy(
-        self,
-        chunks_with_paths: List[Tuple[SemanticChunk, str]],
-        embeddings_matrix: np.ndarray,
-        update_ann: bool = True,
+    def add_chunks_batch_numpy(
+        self,
+        chunks_with_paths: List[Tuple[SemanticChunk, str]],
+        embeddings_matrix: np.ndarray,
+        update_ann: bool = True,
         auto_save_ann: bool = True,
+        categories: Optional[List[str]] = None,
     ) -> List[int]:
         """Batch insert chunks with pre-computed numpy embeddings matrix.
 
@@ -569,19 +631,27 @@ class VectorStore:
             embeddings_matrix: Pre-computed embeddings as (N, D) numpy array
             update_ann: If True, update ANN index with new vectors (default: True)
             auto_save_ann: If True, save ANN index after update (default: True)
+            categories: Optional list of categories per chunk. If None, defaults to 'code'.
+
+        Returns:
+            List of inserted chunk IDs
+        """
+        if not chunks_with_paths:
+            return []
 
-        Returns:
-            List of inserted chunk IDs
-        """
-        if not chunks_with_paths:
-            return []
-
-        batch_size = len(chunks_with_paths)
-
-        if len(chunks_with_paths) != embeddings_matrix.shape[0]:
-            raise ValueError(
-                f"Mismatch: {len(chunks_with_paths)} chunks but "
-                f"{embeddings_matrix.shape[0]} embeddings"
+        batch_size = len(chunks_with_paths)
+
+        if len(chunks_with_paths) != embeddings_matrix.shape[0]:
+            raise ValueError(
+                f"Mismatch: {len(chunks_with_paths)} chunks but "
+                f"{embeddings_matrix.shape[0]} embeddings"
+            )
+
+        # Validate categories if provided
+        if categories is not None and len(categories) != batch_size:
+            raise ValueError(
+                f"categories length ({len(categories)}) must match "
+                f"chunks_with_paths length ({batch_size})"
             )
 
         # Ensure float32 format
@@ -593,49 +663,50 @@ class VectorStore:
             embedding_arr = embeddings_matrix[i]
             embedding_blob = embedding_arr.tobytes()
             metadata_json = json.dumps(chunk.metadata) if chunk.metadata else None
-            batch_data.append((file_path, chunk.content, embedding_blob, metadata_json))
-
-        # Batch insert to SQLite in single transaction
-        with sqlite3.connect(self.db_path) as conn:
-            # Get starting ID before insert
-            row = conn.execute("SELECT MAX(id) FROM semantic_chunks").fetchone()
-            start_id = (row[0] or 0) + 1
-
-            _validate_chunk_id_range(start_id, batch_size)
-
-            conn.executemany(
-                """
-                INSERT INTO semantic_chunks (file_path, content, embedding, metadata)
-                VALUES (?, ?, ?, ?)
+            category = categories[i] if categories else "code"
+            batch_data.append((file_path, chunk.content, embedding_blob, metadata_json, category))
+
+        # Batch insert to SQLite in single transaction
+        with sqlite3.connect(self.db_path) as conn:
+            # Get starting ID before insert
+            row = conn.execute("SELECT MAX(id) FROM semantic_chunks").fetchone()
+            start_id = (row[0] or 0) + 1
+
+            _validate_chunk_id_range(start_id, batch_size)
+
+            conn.executemany(
+                """
+                INSERT INTO semantic_chunks (file_path, content, embedding, metadata, category)
+                VALUES (?, ?, ?, ?, ?)
                 """,
                 batch_data
-            )
-            conn.commit()
-            # Calculate inserted IDs based on starting ID
-            ids = list(range(start_id, start_id + batch_size))
-
-        # Handle ANN index updates
-        if update_ann and self._ensure_ann_index(embeddings_matrix.shape[1]):
-            with self._ann_write_lock:
-                # In bulk insert mode, accumulate for later batch update
-                if self._bulk_insert_mode:
-                    self._bulk_insert_ids.extend(ids)
-                    # Split matrix into individual arrays for accumulation
-                    self._bulk_insert_embeddings.extend([embeddings_matrix[i] for i in range(len(ids))])
-                else:
-                    # Normal mode: update immediately
-                    try:
-                        self._ann_index.add_vectors(ids, embeddings_matrix)
-                        if auto_save_ann:
-                            self._ann_index.save()
-                    except Exception as e:
-                        logger.warning("Failed to add batch to ANN index: %s", e)
+            )
+            conn.commit()
+            # Calculate inserted IDs based on starting ID
+            ids = list(range(start_id, start_id + batch_size))
+
+        # Handle ANN index updates
+        if update_ann and self._ensure_ann_index(embeddings_matrix.shape[1]):
+            with self._ann_write_lock:
+                # In bulk insert mode, accumulate for later batch update
+                if self._bulk_insert_mode:
+                    self._bulk_insert_ids.extend(ids)
+                    # Split matrix into individual arrays for accumulation
+                    self._bulk_insert_embeddings.extend([embeddings_matrix[i] for i in range(len(ids))])
+                else:
+                    # Normal mode: update immediately
+                    try:
+                        self._ann_index.add_vectors(ids, embeddings_matrix)
+                        if auto_save_ann:
+                            self._ann_index.save()
+                    except Exception as e:
+                        logger.warning("Failed to add batch to ANN index: %s", e)
 
         # Invalidate cache after modification
         self._invalidate_cache()
         return ids
 
-    def begin_bulk_insert(self) -> None:
+    def begin_bulk_insert(self) -> None:
         """Begin bulk insert mode - disable ANN auto-update for better performance.
 
         Usage:
@@ -651,45 +722,45 @@ class VectorStore:
                 for batch in batches:
                     store.add_chunks_batch(batch)
         """
-        with self._ann_write_lock:
-            self._bulk_insert_mode = True
-            self._bulk_insert_ids.clear()
-            self._bulk_insert_embeddings.clear()
-        logger.debug("Entered bulk insert mode")
-
-    def end_bulk_insert(self) -> None:
+        with self._ann_write_lock:
+            self._bulk_insert_mode = True
+            self._bulk_insert_ids.clear()
+            self._bulk_insert_embeddings.clear()
+        logger.debug("Entered bulk insert mode")
+
+    def end_bulk_insert(self) -> None:
         """End bulk insert mode and rebuild ANN index from accumulated data.
 
         This method should be called after all bulk inserts are complete to
         update the ANN index in a single batch operation.
         """
-        with self._ann_write_lock:
-            if not self._bulk_insert_mode:
-                logger.warning("end_bulk_insert called but not in bulk insert mode")
-                return
-
-            self._bulk_insert_mode = False
-            bulk_ids = list(self._bulk_insert_ids)
-            bulk_embeddings = list(self._bulk_insert_embeddings)
-            self._bulk_insert_ids.clear()
-            self._bulk_insert_embeddings.clear()
-
-        # Update ANN index with accumulated data.
-        if bulk_ids and bulk_embeddings:
-            if self._ensure_ann_index(len(bulk_embeddings[0])):
-                with self._ann_write_lock:
-                    try:
-                        embeddings_matrix = np.vstack(bulk_embeddings)
-                        self._ann_index.add_vectors(bulk_ids, embeddings_matrix)
-                        self._ann_index.save()
-                        logger.info(
-                            "Bulk insert complete: added %d vectors to ANN index",
-                            len(bulk_ids),
-                        )
-                    except Exception as e:
-                        logger.error("Failed to update ANN index after bulk insert: %s", e)
-
-        logger.debug("Exited bulk insert mode")
+        with self._ann_write_lock:
+            if not self._bulk_insert_mode:
+                logger.warning("end_bulk_insert called but not in bulk insert mode")
+                return
+
+            self._bulk_insert_mode = False
+            bulk_ids = list(self._bulk_insert_ids)
+            bulk_embeddings = list(self._bulk_insert_embeddings)
+            self._bulk_insert_ids.clear()
+            self._bulk_insert_embeddings.clear()
+
+        # Update ANN index with accumulated data.
+        if bulk_ids and bulk_embeddings:
+            if self._ensure_ann_index(len(bulk_embeddings[0])):
+                with self._ann_write_lock:
+                    try:
+                        embeddings_matrix = np.vstack(bulk_embeddings)
+                        self._ann_index.add_vectors(bulk_ids, embeddings_matrix)
+                        self._ann_index.save()
+                        logger.info(
+                            "Bulk insert complete: added %d vectors to ANN index",
+                            len(bulk_ids),
+                        )
+                    except Exception as e:
+                        logger.error("Failed to update ANN index after bulk insert: %s", e)
+
+        logger.debug("Exited bulk insert mode")
 
     class BulkInsertContext:
         """Context manager for bulk insert operations."""
@@ -752,108 +823,114 @@ class VectorStore:
             self._invalidate_cache()
         return deleted
 
-    def search_similar(
-        self,
-        query_embedding: List[float],
-        top_k: int = 10,
-        min_score: float = 0.0,
-        return_full_content: bool = True,
-    ) -> List[SearchResult]:
-        """Find chunks most similar to query embedding.
+    def search_similar(
+        self,
+        query_embedding: List[float],
+        top_k: int = 10,
+        min_score: float = 0.0,
+        return_full_content: bool = True,
+        category: Optional[str] = None,
+    ) -> List[SearchResult]:
+        """Find chunks most similar to query embedding.
 
         Uses HNSW index for O(log N) search when available, falls back to
         brute-force NumPy search otherwise.
 
-        Args:
-            query_embedding: Query vector.
-            top_k: Maximum results to return.
-            min_score: Minimum cosine similarity score in [0.0, 1.0].
-            return_full_content: If True, return full code block content.
-
-        Returns:
-            List of SearchResult ordered by similarity (highest first).
-        """
-        query_vec = np.array(query_embedding, dtype=np.float32)
-
-        if not 0.0 <= min_score <= 1.0:
-            raise ValueError(
-                f"Invalid min_score: {min_score}. Must be within [0.0, 1.0] for cosine similarity."
-            )
-
-        # Try HNSW search first (O(log N))
-        if (
-            HNSWLIB_AVAILABLE
-            and self._ann_index is not None
-            and self._ann_index.is_loaded
+        Args:
+            query_embedding: Query vector.
+            top_k: Maximum results to return.
+            min_score: Minimum cosine similarity score in [0.0, 1.0].
+            return_full_content: If True, return full code block content.
+            category: Optional category filter ('code' or 'doc'). If None, returns all.
+
+        Returns:
+            List of SearchResult ordered by similarity (highest first).
+        """
+        query_vec = np.array(query_embedding, dtype=np.float32)
+
+        if not 0.0 <= min_score <= 1.0:
+            raise ValueError(
+                f"Invalid min_score: {min_score}. Must be within [0.0, 1.0] for cosine similarity."
+            )
+
+        # Try HNSW search first (O(log N))
+        if (
+            HNSWLIB_AVAILABLE
+            and self._ann_index is not None
+            and self._ann_index.is_loaded
             and self._ann_index.count() > 0
         ):
             try:
                 return self._search_with_ann(
-                    query_vec, top_k, min_score, return_full_content
+                    query_vec, top_k, min_score, return_full_content, category
                 )
             except Exception as e:
                 logger.warning("ANN search failed, falling back to brute-force: %s", e)
 
         # Fallback to brute-force search (O(N))
         return self._search_brute_force(
-            query_vec, top_k, min_score, return_full_content
+            query_vec, top_k, min_score, return_full_content, category
         )
 
-    def _search_with_ann(
-        self,
-        query_vec: np.ndarray,
-        top_k: int,
-        min_score: float,
-        return_full_content: bool,
-    ) -> List[SearchResult]:
-        """Search using HNSW index (O(log N)).
-
-        Args:
-            query_vec: Query vector as numpy array
-            top_k: Maximum results to return
-            min_score: Minimum cosine similarity score in [0.0, 1.0]
-            return_full_content: If True, return full code block content
+    def _search_with_ann(
+        self,
+        query_vec: np.ndarray,
+        top_k: int,
+        min_score: float,
+        return_full_content: bool,
+        category: Optional[str] = None,
+    ) -> List[SearchResult]:
+        """Search using HNSW index (O(log N)).
+
+        Args:
+            query_vec: Query vector as numpy array
+            top_k: Maximum results to return
+            min_score: Minimum cosine similarity score in [0.0, 1.0]
+            return_full_content: If True, return full code block content
+            category: Optional category filter ('code' or 'doc')
 
         Returns:
             List of SearchResult ordered by similarity (highest first)
         """
         # Limit top_k to available vectors to prevent hnswlib error
         ann_count = self._ann_index.count()
-        effective_top_k = min(top_k, ann_count) if ann_count > 0 else 0
+        # When category filtering, fetch more candidates to compensate for filtering
+        fetch_k = top_k * 3 if category else top_k
+        effective_top_k = min(fetch_k, ann_count) if ann_count > 0 else 0
 
         if effective_top_k == 0:
             return []
 
-        # HNSW search returns (ids, distances)
-        # For cosine space: distance = 1 - similarity
-        ids, distances = self._ann_index.search(query_vec, effective_top_k)
-
-        if ids is None or distances is None:
-            logger.debug(
-                "ANN search returned null results (ids=%s, distances=%s)",
-                ids,
-                distances,
-            )
-            return []
-
-        if len(ids) == 0 or len(distances) == 0:
-            logger.debug(
-                "ANN search returned empty results (ids=%s, distances=%s)",
-                ids,
-                distances,
-            )
-            return []
-
-        if len(ids) != len(distances):
-            logger.warning(
-                "ANN search returned mismatched result lengths (%d ids, %d distances)",
-                len(ids),
-                len(distances),
-            )
-            return []
-
-        # Convert distances to similarity scores
-        scores = [1.0 - d for d in distances]
+        # HNSW search returns (ids, distances)
+        # For cosine space: distance = 1 - similarity
+        ids, distances = self._ann_index.search(query_vec, effective_top_k)
+
+        if ids is None or distances is None:
+            logger.debug(
+                "ANN search returned null results (ids=%s, distances=%s)",
+                ids,
+                distances,
+            )
+            return []
+
+        if len(ids) == 0 or len(distances) == 0:
+            logger.debug(
+                "ANN search returned empty results (ids=%s, distances=%s)",
+                ids,
+                distances,
+            )
+            return []
+
+        if len(ids) != len(distances):
+            logger.warning(
+                "ANN search returned mismatched result lengths (%d ids, %d distances)",
+                len(ids),
+                len(distances),
+            )
+            return []
+
+        # Convert distances to similarity scores
+        scores = [1.0 - d for d in distances]
 
         # Filter by min_score
         filtered = [
@@ -868,23 +945,29 @@ class VectorStore:
         top_ids = [f[0] for f in filtered]
         top_scores = [f[1] for f in filtered]
 
-        # Fetch content from SQLite
-        return self._fetch_results_by_ids(top_ids, top_scores, return_full_content)
+        # Fetch content from SQLite with category filtering
+        results = self._fetch_results_by_ids(
+            top_ids, top_scores, return_full_content, category
+        )
+        # Apply final limit after category filtering
+        return results[:top_k]
 
-    def _search_brute_force(
-        self,
-        query_vec: np.ndarray,
-        top_k: int,
-        min_score: float,
-        return_full_content: bool,
-    ) -> List[SearchResult]:
-        """Brute-force search using NumPy (O(N) fallback).
+    def _search_brute_force(
+        self,
+        query_vec: np.ndarray,
+        top_k: int,
+        min_score: float,
+        return_full_content: bool,
+        category: Optional[str] = None,
+    ) -> List[SearchResult]:
+        """Brute-force search using NumPy (O(N) fallback).
 
-        Args:
-            query_vec: Query vector as numpy array
-            top_k: Maximum results to return
-            min_score: Minimum cosine similarity score in [0.0, 1.0]
-            return_full_content: If True, return full code block content
+        Args:
+            query_vec: Query vector as numpy array
+            top_k: Maximum results to return
+            min_score: Minimum cosine similarity score in [0.0, 1.0]
+            return_full_content: If True, return full code block content
+            category: Optional category filter ('code' or 'doc')
 
         Returns:
             List of SearchResult ordered by similarity (highest first)
@@ -919,27 +1002,31 @@ class VectorStore:
             if len(valid_indices) == 0:
                 return []
 
-            # Sort by score descending and take top_k
+            # When category filtering, fetch more candidates to compensate for filtering
+            fetch_k = top_k * 3 if category else top_k
+
+            # Sort by score descending and take top candidates
             valid_scores = scores[valid_indices]
-            sorted_order = np.argsort(valid_scores)[::-1][:top_k]
+            sorted_order = np.argsort(valid_scores)[::-1][:fetch_k]
             top_indices = valid_indices[sorted_order]
             top_scores = valid_scores[sorted_order]
 
             # Get chunk IDs for top results
             top_ids = [self._chunk_ids[i] for i in top_indices]
 
-        # Fetch content only for top-k results (lazy loading)
+        # Fetch content only for top-k results (lazy loading) with category filtering
         results = self._fetch_results_by_ids(
-            top_ids, top_scores.tolist(), return_full_content
+            top_ids, top_scores.tolist(), return_full_content, category
         )
-
-        return results
+        # Apply final limit after category filtering
+        return results[:top_k]
 
     def _fetch_results_by_ids(
         self,
         chunk_ids: List[int],
         scores: List[float],
         return_full_content: bool,
+        category: Optional[str] = None,
     ) -> List[SearchResult]:
         """Fetch full result data for specific chunk IDs.
 
@@ -947,29 +1034,40 @@ class VectorStore:
             chunk_ids: List of chunk IDs to fetch.
             scores: Corresponding similarity scores.
             return_full_content: Whether to include full content.
+            category: Optional category filter ('code' or 'doc').
 
         Returns:
             List of SearchResult objects.
         """
-        if not chunk_ids:
-            return []
-
-        # Build parameterized query for IN clause
-        placeholders = ",".join("?" * len(chunk_ids))
-        _validate_sql_placeholders(placeholders, len(chunk_ids))
-
-        # SQL injection prevention:
-        # - Only a validated placeholders string (commas + '?') is interpolated into the query.
-        # - User-provided values are passed separately via sqlite3 parameters.
-        query = """
-            SELECT id, file_path, content, metadata
-            FROM semantic_chunks
-            WHERE id IN ({placeholders})
-        """.format(placeholders=placeholders)
+        if not chunk_ids:
+            return []
+
+        # Build parameterized query for IN clause
+        placeholders = ",".join("?" * len(chunk_ids))
+        _validate_sql_placeholders(placeholders, len(chunk_ids))
+
+        # SQL injection prevention:
+        # - Only a validated placeholders string (commas + '?') is interpolated into the query.
+        # - User-provided values are passed separately via sqlite3 parameters.
+        # - Category filter is added as a separate parameter
+        if category:
+            query = """
+                SELECT id, file_path, content, metadata
+                FROM semantic_chunks
+                WHERE id IN ({placeholders}) AND category = ?
+            """.format(placeholders=placeholders)
+            params = list(chunk_ids) + [category]
+        else:
+            query = """
+                SELECT id, file_path, content, metadata
+                FROM semantic_chunks
+                WHERE id IN ({placeholders})
+            """.format(placeholders=placeholders)
+            params = chunk_ids
 
         with sqlite3.connect(self.db_path) as conn:
             conn.execute("PRAGMA mmap_size = 30000000000")
-            rows = conn.execute(query, chunk_ids).fetchall()
+            rows = conn.execute(query, params).fetchall()
 
         # Build ID -> row mapping
         id_to_row = {r[0]: r for r in rows}
@@ -1026,6 +1124,28 @@ class VectorStore:
             row = conn.execute("SELECT COUNT(*) FROM semantic_chunks").fetchone()
             return row[0] if row else 0
 
+    def get_all_chunks(self) -> List[SemanticChunk]:
+        """Get all chunks from the store.
+
+        Returns:
+            List of SemanticChunk objects with id and content.
+        """
+        with sqlite3.connect(self.db_path) as conn:
+            conn.row_factory = sqlite3.Row
+            rows = conn.execute(
+                "SELECT id, file_path, content, metadata FROM semantic_chunks"
+            ).fetchall()
+
+            chunks = []
+            for row in rows:
+                chunks.append(SemanticChunk(
+                    id=row["id"],
+                    content=row["content"],
+                    file_path=row["file_path"],
+                    metadata=json.loads(row["metadata"]) if row["metadata"] else None,
+                ))
+            return chunks
+
     def clear_cache(self) -> None:
         """Manually clear the embedding cache."""
         self._invalidate_cache()