mcp-vector-search 0.12.6__py3-none-any.whl

mcp_vector_search/core/database.py

@@ -0,0 +1,1214 @@
+"""Database abstraction and ChromaDB implementation for MCP Vector Search."""
+
+import asyncio
+import json
+import shutil
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, Protocol, runtime_checkable
+
+from loguru import logger
+
+from .connection_pool import ChromaConnectionPool
+from .exceptions import (
+    DatabaseError,
+    DatabaseInitializationError,
+    DatabaseNotInitializedError,
+    DocumentAdditionError,
+    IndexCorruptionError,
+    SearchError,
+)
+from .models import CodeChunk, IndexStats, SearchResult
+
+
+@runtime_checkable
+class EmbeddingFunction(Protocol):
+    """Protocol for embedding functions."""
+
+    def __call__(self, texts: list[str]) -> list[list[float]]:
+        """Generate embeddings for input texts."""
+        ...
+
+
+class VectorDatabase(ABC):
+    """Abstract interface for vector database operations."""
+
+    @abstractmethod
+    async def initialize(self) -> None:
+        """Initialize the database connection and collections."""
+        ...
+
+    @abstractmethod
+    async def close(self) -> None:
+        """Close database connections and cleanup resources."""
+        ...
+
+    @abstractmethod
+    async def add_chunks(self, chunks: list[CodeChunk]) -> None:
+        """Add code chunks to the database.
+
+        Args:
+            chunks: List of code chunks to add
+        """
+        ...
+
+    @abstractmethod
+    async def search(
+        self,
+        query: str,
+        limit: int = 10,
+        filters: dict[str, Any] | None = None,
+        similarity_threshold: float = 0.7,
+    ) -> list[SearchResult]:
+        """Search for similar code chunks.
+
+        Args:
+            query: Search query
+            limit: Maximum number of results
+            filters: Optional filters to apply
+            similarity_threshold: Minimum similarity score
+
+        Returns:
+            List of search results
+        """
+        ...
+
+    @abstractmethod
+    async def delete_by_file(self, file_path: Path) -> int:
+        """Delete all chunks for a specific file.
+
+        Args:
+            file_path: Path to the file
+
+        Returns:
+            Number of deleted chunks
+        """
+        ...
+
+    @abstractmethod
+    async def get_stats(self) -> IndexStats:
+        """Get database statistics.
+
+        Returns:
+            Index statistics
+        """
+        ...
+
+    @abstractmethod
+    async def reset(self) -> None:
+        """Reset the database (delete all data)."""
+        ...
+
+    @abstractmethod
+    async def get_all_chunks(self) -> list[CodeChunk]:
+        """Get all chunks from the database.
+
+        Returns:
+            List of all code chunks with metadata
+        """
+        ...
+
+    @abstractmethod
+    async def health_check(self) -> bool:
+        """Check database health and integrity.
+
+        Returns:
+            True if database is healthy, False otherwise
+        """
+        ...
+
+    async def __aenter__(self) -> "VectorDatabase":
+        """Async context manager entry."""
+        await self.initialize()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Async context manager exit."""
+        await self.close()
+
+
+class ChromaVectorDatabase(VectorDatabase):
+    """ChromaDB implementation of vector database."""
+
+    def __init__(
+        self,
+        persist_directory: Path,
+        embedding_function: EmbeddingFunction,
+        collection_name: str = "code_search",
+    ) -> None:
+        """Initialize ChromaDB vector database.
+
+        Args:
+            persist_directory: Directory to persist database
+            embedding_function: Function to generate embeddings
+            collection_name: Name of the collection
+        """
+        self.persist_directory = persist_directory
+        self.embedding_function = embedding_function
+        self.collection_name = collection_name
+        self._client = None
+        self._collection = None
+
+    async def initialize(self) -> None:
+        """Initialize ChromaDB client and collection with corruption recovery."""
+        try:
+            import chromadb
+
+            # Ensure directory exists
+            self.persist_directory.mkdir(parents=True, exist_ok=True)
+
+            # Check for corruption before initializing
+            await self._detect_and_recover_corruption()
+
+            # Create client with new API
+            self._client = chromadb.PersistentClient(
+                path=str(self.persist_directory),
+                settings=chromadb.Settings(
+                    anonymized_telemetry=False,
+                    allow_reset=True,
+                ),
+            )
+
+            # Create or get collection
+            self._collection = self._client.get_or_create_collection(
+                name=self.collection_name,
+                embedding_function=self.embedding_function,
+                metadata={
+                    "description": "Semantic code search collection",
+                },
+            )
+
+            logger.debug(f"ChromaDB initialized at {self.persist_directory}")
+
+        except Exception as e:
+            # Check if this is a corruption error
+            error_msg = str(e).lower()
+            if any(
+                indicator in error_msg
+                for indicator in [
+                    "pickle",
+                    "unpickling",
+                    "eof",
+                    "ran out of input",
+                    "hnsw",
+                    "index",
+                    "deserialize",
+                    "corrupt",
+                ]
+            ):
+                logger.warning(f"Detected index corruption: {e}")
+                # Try to recover
+                await self._recover_from_corruption()
+                # Retry initialization
+                await self.initialize()
+            else:
+                logger.error(f"Failed to initialize ChromaDB: {e}")
+                raise DatabaseInitializationError(
+                    f"ChromaDB initialization failed: {e}"
+                ) from e
+
+    async def remove_file_chunks(self, file_path: str) -> int:
+        """Remove all chunks for a specific file.
+
+        Args:
+            file_path: Relative path to the file
+
+        Returns:
+            Number of chunks removed
+        """
+        if not self._collection:
+            raise DatabaseNotInitializedError("Database not initialized")
+
+        try:
+            # Get all chunks for this file
+            results = self._collection.get(where={"file_path": file_path})
+
+            if not results["ids"]:
+                return 0
+
+            # Delete the chunks
+            self._collection.delete(ids=results["ids"])
+
+            removed_count = len(results["ids"])
+            logger.debug(f"Removed {removed_count} chunks for file: {file_path}")
+            return removed_count
+
+        except Exception as e:
+            logger.error(f"Failed to remove chunks for file {file_path}: {e}")
+            return 0
+
+    async def close(self) -> None:
+        """Close database connections."""
+        if self._client:
+            # ChromaDB doesn't require explicit closing
+            self._client = None
+            self._collection = None
+            logger.debug("ChromaDB connections closed")
+
+    async def add_chunks(self, chunks: list[CodeChunk]) -> None:
+        """Add code chunks to the database."""
+        if not self._collection:
+            raise DatabaseNotInitializedError("Database not initialized")
+
+        if not chunks:
+            return
+
+        try:
+            documents = []
+            metadatas = []
+            ids = []
+
+            for chunk in chunks:
+                # Create searchable text
+                searchable_text = self._create_searchable_text(chunk)
+                documents.append(searchable_text)
+
+                # Create metadata
+                metadata = {
+                    "file_path": str(chunk.file_path),
+                    "start_line": chunk.start_line,
+                    "end_line": chunk.end_line,
+                    "language": chunk.language,
+                    "chunk_type": chunk.chunk_type,
+                    "function_name": chunk.function_name or "",
+                    "class_name": chunk.class_name or "",
+                    "docstring": chunk.docstring or "",
+                    "complexity_score": chunk.complexity_score,
+                    # Hierarchy fields (convert lists to JSON strings for ChromaDB)
+                    "chunk_id": chunk.chunk_id or "",
+                    "parent_chunk_id": chunk.parent_chunk_id or "",
+                    "child_chunk_ids": json.dumps(chunk.child_chunk_ids or []),
+                    "chunk_depth": chunk.chunk_depth,
+                    # Additional metadata (convert lists/dicts to JSON strings)
+                    "decorators": json.dumps(chunk.decorators or []),
+                    "parameters": json.dumps(chunk.parameters or []),
+                    "return_type": chunk.return_type or "",
+                    "type_annotations": json.dumps(chunk.type_annotations or {}),
+                    # Monorepo support
+                    "subproject_name": chunk.subproject_name or "",
+                    "subproject_path": chunk.subproject_path or "",
+                }
+                metadatas.append(metadata)
+
+                # Use chunk ID
+                ids.append(chunk.id)
+
+            # Add to collection
+            self._collection.add(
+                documents=documents,
+                metadatas=metadatas,
+                ids=ids,
+            )
+
+            logger.debug(f"Added {len(chunks)} chunks to database")
+
+        except Exception as e:
+            logger.error(f"Failed to add chunks: {e}")
+            raise DocumentAdditionError(f"Failed to add chunks: {e}") from e
+
+    async def search(
+        self,
+        query: str,
+        limit: int = 10,
+        filters: dict[str, Any] | None = None,
+        similarity_threshold: float = 0.7,
+    ) -> list[SearchResult]:
+        """Search for similar code chunks."""
+        if not self._collection:
+            raise DatabaseNotInitializedError("Database not initialized")
+
+        try:
+            # Build where clause
+            where_clause = self._build_where_clause(filters) if filters else None
+
+            # Perform search
+            results = self._collection.query(
+                query_texts=[query],
+                n_results=limit,
+                where=where_clause,
+                include=["documents", "metadatas", "distances"],
+            )
+
+            # Process results
+            search_results = []
+
+            if results["documents"] and results["documents"][0]:
+                for i, (doc, metadata, distance) in enumerate(
+                    zip(
+                        results["documents"][0],
+                        results["metadatas"][0],
+                        results["distances"][0],
+                        strict=False,
+                    )
+                ):
+                    # Convert distance to similarity (ChromaDB uses cosine distance)
+                    # For cosine distance, use a more permissive conversion that handles distances > 1.0
+                    # Convert to a 0-1 similarity score where lower distances = higher similarity
+                    similarity = max(0.0, 1.0 / (1.0 + distance))
+
+                    if similarity >= similarity_threshold:
+                        result = SearchResult(
+                            content=doc,
+                            file_path=Path(metadata["file_path"]),
+                            start_line=metadata["start_line"],
+                            end_line=metadata["end_line"],
+                            language=metadata["language"],
+                            similarity_score=similarity,
+                            rank=i + 1,
+                            chunk_type=metadata.get("chunk_type", "code"),
+                            function_name=metadata.get("function_name") or None,
+                            class_name=metadata.get("class_name") or None,
+                        )
+                        search_results.append(result)
+
+            logger.debug(f"Found {len(search_results)} results for query: {query}")
+            return search_results
+
+        except Exception as e:
+            logger.error(f"Search failed: {e}")
+            raise SearchError(f"Search failed: {e}") from e
+
+    async def delete_by_file(self, file_path: Path) -> int:
+        """Delete all chunks for a specific file."""
+        if not self._collection:
+            raise DatabaseNotInitializedError("Database not initialized")
+
+        try:
+            # Get all chunks for this file
+            results = self._collection.get(
+                where={"file_path": str(file_path)},
+                include=["metadatas"],
+            )
+
+            if results["ids"]:
+                self._collection.delete(ids=results["ids"])
+                count = len(results["ids"])
+                logger.debug(f"Deleted {count} chunks for {file_path}")
+                return count
+
+            return 0
+
+        except Exception as e:
+            logger.error(f"Failed to delete chunks for {file_path}: {e}")
+            raise DatabaseError(f"Failed to delete chunks: {e}") from e
+
+    async def get_stats(self) -> IndexStats:
+        """Get database statistics with optimized chunked queries."""
+        if not self._collection:
+            raise DatabaseNotInitializedError("Database not initialized")
+
+        try:
+            # Get total count (fast operation)
+            count = self._collection.count()
+
+            if count == 0:
+                return IndexStats(
+                    total_files=0,
+                    total_chunks=0,
+                    languages={},
+                    file_types={},
+                    index_size_mb=0.0,
+                    last_updated="N/A",
+                    embedding_model="unknown",
+                )
+
+            # Process in chunks to avoid loading everything at once
+            batch_size_limit = 1000
+
+            files = set()
+            language_counts: dict[str, int] = {}
+            file_type_counts: dict[str, int] = {}
+
+            offset = 0
+            while offset < count:
+                # Fetch batch
+                batch_size = min(batch_size_limit, count - offset)
+                logger.debug(
+                    f"Processing database stats: batch {offset // batch_size_limit + 1}, "
+                    f"{offset}-{offset + batch_size} of {count} chunks"
+                )
+
+                results = self._collection.get(
+                    include=["metadatas"],
+                    limit=batch_size,
+                    offset=offset,
+                )
+
+                # Process batch metadata
+                for metadata in results.get("metadatas", []):
+                    # Language stats
+                    lang = metadata.get("language", "unknown")
+                    language_counts[lang] = language_counts.get(lang, 0) + 1
+
+                    # File stats
+                    file_path = metadata.get("file_path", "")
+                    if file_path:
+                        files.add(file_path)
+                        ext = Path(file_path).suffix or "no_extension"
+                        file_type_counts[ext] = file_type_counts.get(ext, 0) + 1
+
+                offset += batch_size
+
+                # Yield to event loop periodically to prevent blocking
+                await asyncio.sleep(0)
+
+            # Estimate index size (rough approximation: ~1KB per chunk)
+            index_size_mb = count * 0.001
+
+            return IndexStats(
+                total_files=len(files),
+                total_chunks=count,
+                languages=language_counts,
+                file_types=file_type_counts,
+                index_size_mb=index_size_mb,
+                last_updated="unknown",
+                embedding_model="unknown",
+            )
+
+        except Exception as e:
+            logger.error(f"Failed to get database statistics: {e}")
+            # Return empty stats instead of raising
+            return IndexStats(
+                total_files=0,
+                total_chunks=0,
+                languages={},
+                file_types={},
+                index_size_mb=0.0,
+                last_updated="error",
+                embedding_model="unknown",
+            )
+
+    async def reset(self) -> None:
+        """Reset the database."""
+        if self._client:
+            try:
+                self._client.reset()
+                # Recreate collection
+                await self.initialize()
+                logger.info("Database reset successfully")
+            except Exception as e:
+                logger.error(f"Failed to reset database: {e}")
+                raise DatabaseError(f"Failed to reset database: {e}") from e
+
+    async def get_all_chunks(self) -> list[CodeChunk]:
+        """Get all chunks from the database.
+
+        Returns:
+            List of all code chunks with metadata
+        """
+        if not self._collection:
+            raise DatabaseNotInitializedError("Database not initialized")
+
+        try:
+            # Get all documents from collection
+            results = self._collection.get(include=["metadatas", "documents"])
+
+            chunks = []
+            if results and results.get("ids"):
+                for i, _chunk_id in enumerate(results["ids"]):
+                    metadata = results["metadatas"][i]
+                    content = results["documents"][i]
+
+                    # Parse JSON strings back to lists/dicts
+                    child_chunk_ids = metadata.get("child_chunk_ids", "[]")
+                    if isinstance(child_chunk_ids, str):
+                        child_chunk_ids = json.loads(child_chunk_ids)
+
+                    decorators = metadata.get("decorators", "[]")
+                    if isinstance(decorators, str):
+                        decorators = json.loads(decorators)
+
+                    parameters = metadata.get("parameters", "[]")
+                    if isinstance(parameters, str):
+                        parameters = json.loads(parameters)
+
+                    type_annotations = metadata.get("type_annotations", "{}")
+                    if isinstance(type_annotations, str):
+                        type_annotations = json.loads(type_annotations)
+
+                    chunk = CodeChunk(
+                        content=content,
+                        file_path=Path(metadata["file_path"]),
+                        start_line=metadata["start_line"],
+                        end_line=metadata["end_line"],
+                        language=metadata["language"],
+                        chunk_type=metadata.get("chunk_type", "code"),
+                        function_name=metadata.get("function_name"),
+                        class_name=metadata.get("class_name"),
+                        docstring=metadata.get("docstring"),
+                        imports=metadata.get("imports", []),
+                        complexity_score=metadata.get("complexity_score", 0.0),
+                        chunk_id=metadata.get("chunk_id"),
+                        parent_chunk_id=metadata.get("parent_chunk_id"),
+                        child_chunk_ids=child_chunk_ids,
+                        chunk_depth=metadata.get("chunk_depth", 0),
+                        decorators=decorators,
+                        parameters=parameters,
+                        return_type=metadata.get("return_type"),
+                        type_annotations=type_annotations,
+                        subproject_name=metadata.get("subproject_name"),
+                        subproject_path=metadata.get("subproject_path"),
+                    )
+                    chunks.append(chunk)
+
+            logger.debug(f"Retrieved {len(chunks)} chunks from database")
+            return chunks
+
+        except Exception as e:
+            logger.error(f"Failed to get all chunks: {e}")
+            raise DatabaseError(f"Failed to get all chunks: {e}") from e
+
+    def _create_searchable_text(self, chunk: CodeChunk) -> str:
+        """Create optimized searchable text from code chunk."""
+        parts = [chunk.content]
+
+        # Add contextual information
+        if chunk.function_name:
+            parts.append(f"Function: {chunk.function_name}")
+
+        if chunk.class_name:
+            parts.append(f"Class: {chunk.class_name}")
+
+        if chunk.docstring:
+            parts.append(f"Documentation: {chunk.docstring}")
+
+        # Add language and file context
+        parts.append(f"Language: {chunk.language}")
+        parts.append(f"File: {chunk.file_path.name}")
+
+        return "\n".join(parts)
+
+    def _build_where_clause(self, filters: dict[str, Any]) -> dict[str, Any]:
+        """Build ChromaDB where clause from filters."""
+        where = {}
+
+        for key, value in filters.items():
+            if isinstance(value, list):
+                where[key] = {"$in": value}
+            elif isinstance(value, str) and value.startswith("!"):
+                where[key] = {"$ne": value[1:]}
+            else:
+                where[key] = value
+
+        return where
+
+    async def _detect_and_recover_corruption(self) -> None:
+        """Detect and recover from index corruption proactively."""
+        # Check for common corruption indicators in ChromaDB files
+        chroma_db_path = self.persist_directory / "chroma.sqlite3"
+
+        # If database doesn't exist yet, nothing to check
+        if not chroma_db_path.exists():
+            return
+
+        # Check for HNSW index files that might be corrupted
+        self.persist_directory / "chroma-collections.parquet"
+        index_path = self.persist_directory / "index"
+
+        if index_path.exists():
+            # Look for pickle files in the index
+            pickle_files = list(index_path.glob("**/*.pkl"))
+            pickle_files.extend(list(index_path.glob("**/*.pickle")))
+
+            for pickle_file in pickle_files:
+                try:
+                    # Try to read the pickle file to detect corruption
+                    import pickle
+
+                    with open(pickle_file, "rb") as f:
+                        pickle.load(f)
+                except (EOFError, pickle.UnpicklingError, Exception) as e:
+                    logger.warning(
+                        f"Corrupted index file detected: {pickle_file} - {e}"
+                    )
+                    await self._recover_from_corruption()
+                    return
+
+    async def _recover_from_corruption(self) -> None:
+        """Recover from index corruption by rebuilding the index."""
+        logger.info("Attempting to recover from index corruption...")
+
+        # Create backup directory
+        backup_dir = (
+            self.persist_directory.parent / f"{self.persist_directory.name}_backup"
+        )
+        backup_dir.mkdir(exist_ok=True)
+
+        # Backup current state (in case we need it)
+        import time
+
+        timestamp = int(time.time())
+        backup_path = backup_dir / f"backup_{timestamp}"
+
+        if self.persist_directory.exists():
+            try:
+                shutil.copytree(self.persist_directory, backup_path)
+                logger.info(f"Created backup at {backup_path}")
+            except Exception as e:
+                logger.warning(f"Could not create backup: {e}")
+
+        # Clear the corrupted index
+        if self.persist_directory.exists():
+            try:
+                shutil.rmtree(self.persist_directory)
+                logger.info(f"Cleared corrupted index at {self.persist_directory}")
+            except Exception as e:
+                logger.error(f"Failed to clear corrupted index: {e}")
+                raise IndexCorruptionError(
+                    f"Could not clear corrupted index: {e}"
+                ) from e
+
+        # Recreate the directory
+        self.persist_directory.mkdir(parents=True, exist_ok=True)
+        logger.info("Index directory recreated. Please re-index your codebase.")
+
+    async def health_check(self) -> bool:
+        """Check database health and integrity.
+
+        Returns:
+            True if database is healthy, False otherwise
+        """
+        try:
+            # First check if client is initialized
+            if not self._client or not self._collection:
+                logger.warning("Database not initialized")
+                return False
+
+            # Try a simple operation to test the connection
+            try:
+                # Attempt to get count - this will fail if index is corrupted
+                count = self._collection.count()
+                logger.debug(f"Health check passed: {count} chunks in database")
+
+                # Try a minimal query to ensure search works
+                self._collection.query(
+                    query_texts=["test"], n_results=1, include=["metadatas"]
+                )
+
+                return True
+
+            except Exception as e:
+                error_msg = str(e).lower()
+                if any(
+                    indicator in error_msg
+                    for indicator in [
+                        "pickle",
+                        "unpickling",
+                        "eof",
+                        "ran out of input",
+                        "hnsw",
+                        "index",
+                        "deserialize",
+                        "corrupt",
+                    ]
+                ):
+                    logger.error(f"Index corruption detected during health check: {e}")
+                    return False
+                else:
+                    # Some other error
+                    logger.warning(f"Health check failed: {e}")
+                    return False
+
+        except Exception as e:
+            logger.error(f"Health check error: {e}")
+            return False
+
+
+class PooledChromaVectorDatabase(VectorDatabase):
+    """ChromaDB implementation with connection pooling for improved performance."""
+
+    def __init__(
+        self,
+        persist_directory: Path,
+        embedding_function: EmbeddingFunction,
+        collection_name: str = "code_search",
+        max_connections: int = 10,
+        min_connections: int = 2,
+        max_idle_time: float = 300.0,
+        max_connection_age: float = 3600.0,
+    ) -> None:
+        """Initialize pooled ChromaDB vector database.
+
+        Args:
+            persist_directory: Directory to persist database
+            embedding_function: Function to generate embeddings
+            collection_name: Name of the collection
+            max_connections: Maximum number of connections in pool
+            min_connections: Minimum number of connections to maintain
+            max_idle_time: Maximum time a connection can be idle (seconds)
+            max_connection_age: Maximum age of a connection (seconds)
+        """
+        self.persist_directory = persist_directory
+        self.embedding_function = embedding_function
+        self.collection_name = collection_name
+
+        self._pool = ChromaConnectionPool(
+            persist_directory=persist_directory,
+            embedding_function=embedding_function,
+            collection_name=collection_name,
+            max_connections=max_connections,
+            min_connections=min_connections,
+            max_idle_time=max_idle_time,
+            max_connection_age=max_connection_age,
+        )
+
+    async def initialize(self) -> None:
+        """Initialize the connection pool."""
+        await self._pool.initialize()
+        logger.debug(f"Pooled ChromaDB initialized at {self.persist_directory}")
+
+    async def close(self) -> None:
+        """Close the connection pool."""
+        await self._pool.close()
+        logger.debug("Pooled ChromaDB connections closed")
+
+    async def add_chunks(self, chunks: list[CodeChunk]) -> None:
+        """Add code chunks to the database using pooled connection."""
+        if not chunks:
+            return
+
+        # Ensure pool is initialized
+        if not self._pool._initialized:
+            await self._pool.initialize()
+
+        try:
+            async with self._pool.get_connection() as conn:
+                # Prepare data for ChromaDB
+                documents = []
+                metadatas = []
+                ids = []
+
+                for chunk in chunks:
+                    documents.append(chunk.content)
+                    metadatas.append(
+                        {
+                            "file_path": str(chunk.file_path),
+                            "start_line": chunk.start_line,
+                            "end_line": chunk.end_line,
+                            "language": chunk.language,
+                            "chunk_type": chunk.chunk_type,
+                            "function_name": chunk.function_name or "",
+                            "class_name": chunk.class_name or "",
+                            "docstring": chunk.docstring or "",
+                            "complexity_score": chunk.complexity_score,
+                            # Hierarchy fields (convert lists to JSON strings for ChromaDB)
+                            "chunk_id": chunk.chunk_id or "",
+                            "parent_chunk_id": chunk.parent_chunk_id or "",
+                            "child_chunk_ids": json.dumps(chunk.child_chunk_ids or []),
+                            "chunk_depth": chunk.chunk_depth,
+                            # Additional metadata (convert lists/dicts to JSON strings)
+                            "decorators": json.dumps(chunk.decorators or []),
+                            "parameters": json.dumps(chunk.parameters or []),
+                            "return_type": chunk.return_type or "",
+                            "type_annotations": json.dumps(
+                                chunk.type_annotations or {}
+                            ),
+                            # Monorepo support
+                            "subproject_name": chunk.subproject_name or "",
+                            "subproject_path": chunk.subproject_path or "",
+                        }
+                    )
+                    ids.append(chunk.id)
+
+                # Add to collection
+                conn.collection.add(documents=documents, metadatas=metadatas, ids=ids)
+
+                logger.debug(f"Added {len(chunks)} chunks to database")
+
+        except Exception as e:
+            logger.error(f"Failed to add chunks: {e}")
+            raise DocumentAdditionError(f"Failed to add chunks: {e}") from e
+
+    async def search(
+        self,
+        query: str,
+        limit: int = 10,
+        filters: dict[str, Any] | None = None,
+        similarity_threshold: float = 0.7,
+    ) -> list[SearchResult]:
+        """Search for similar code chunks using pooled connection."""
+        # Ensure pool is initialized
+        if not self._pool._initialized:
+            await self._pool.initialize()
+
+        try:
+            async with self._pool.get_connection() as conn:
+                # Build where clause
+                where_clause = self._build_where_clause(filters) if filters else None
+
+                # Perform search
+                results = conn.collection.query(
+                    query_texts=[query],
+                    n_results=limit,
+                    where=where_clause,
+                    include=["documents", "metadatas", "distances"],
+                )
+
+                # Process results
+                search_results = []
+
+                if results["documents"] and results["documents"][0]:
+                    for i, (doc, metadata, distance) in enumerate(
+                        zip(
+                            results["documents"][0],
+                            results["metadatas"][0],
+                            results["distances"][0],
+                            strict=False,
+                        )
+                    ):
+                        # Convert distance to similarity (ChromaDB uses cosine distance)
+                        # For cosine distance, use a more permissive conversion that handles distances > 1.0
+                        # Convert to a 0-1 similarity score where lower distances = higher similarity
+                        similarity = max(0.0, 1.0 / (1.0 + distance))
+
+                        if similarity >= similarity_threshold:
+                            result = SearchResult(
+                                content=doc,
+                                file_path=Path(metadata["file_path"]),
+                                start_line=metadata["start_line"],
+                                end_line=metadata["end_line"],
+                                language=metadata["language"],
+                                similarity_score=similarity,
+                                rank=i + 1,
+                                chunk_type=metadata.get("chunk_type", "code"),
+                                function_name=metadata.get("function_name") or None,
+                                class_name=metadata.get("class_name") or None,
+                            )
+                            search_results.append(result)
+
+                logger.debug(f"Found {len(search_results)} results for query: {query}")
+                return search_results
+
+        except Exception as e:
+            logger.error(f"Search failed: {e}")
+            raise SearchError(f"Search failed: {e}") from e
+
+    async def delete_by_file(self, file_path: Path) -> int:
+        """Delete all chunks for a specific file using pooled connection."""
+        try:
+            async with self._pool.get_connection() as conn:
+                # Get all chunks for this file
+                results = conn.collection.get(
+                    where={"file_path": str(file_path)}, include=["metadatas"]
+                )
+
+                if not results["ids"]:
+                    return 0
+
+                # Delete the chunks
+                conn.collection.delete(ids=results["ids"])
+
+                deleted_count = len(results["ids"])
+                logger.debug(f"Deleted {deleted_count} chunks for file: {file_path}")
+                return deleted_count
+
+        except Exception as e:
+            logger.error(f"Failed to delete chunks for file {file_path}: {e}")
+            raise DatabaseError(f"Failed to delete chunks: {e}") from e
+
+    async def get_stats(self) -> IndexStats:
+        """Get database statistics with connection pooling and chunked queries."""
+        try:
+            async with self._pool.get_connection() as conn:
+                # Get total count (fast operation)
+                count = conn.collection.count()
+
+                if count == 0:
+                    return IndexStats(
+                        total_files=0,
+                        total_chunks=0,
+                        languages={},
+                        file_types={},
+                        index_size_mb=0.0,
+                        last_updated="N/A",
+                        embedding_model="unknown",
+                    )
+
+                # Process in chunks to avoid loading everything at once
+                batch_size_limit = 1000
+
+                files = set()
+                language_counts: dict[str, int] = {}
+                file_type_counts: dict[str, int] = {}
+
+                offset = 0
+                while offset < count:
+                    # Fetch batch
+                    batch_size = min(batch_size_limit, count - offset)
+                    logger.debug(
+                        f"Processing database stats: batch {offset // batch_size_limit + 1}, "
+                        f"{offset}-{offset + batch_size} of {count} chunks"
+                    )
+
+                    results = conn.collection.get(
+                        include=["metadatas"],
+                        limit=batch_size,
+                        offset=offset,
+                    )
+
+                    # Process batch metadata
+                    for metadata in results.get("metadatas", []):
+                        # Language stats
+                        lang = metadata.get("language", "unknown")
+                        language_counts[lang] = language_counts.get(lang, 0) + 1
+
+                        # File stats
+                        file_path = metadata.get("file_path", "")
+                        if file_path:
+                            files.add(file_path)
+                            ext = Path(file_path).suffix or "no_extension"
+                            file_type_counts[ext] = file_type_counts.get(ext, 0) + 1
+
+                    offset += batch_size
+
+                    # Yield to event loop periodically to prevent blocking
+                    await asyncio.sleep(0)
+
+                # Estimate index size (rough approximation: ~1KB per chunk)
+                index_size_mb = count * 0.001
+
+                return IndexStats(
+                    total_files=len(files),
+                    total_chunks=count,
+                    languages=language_counts,
+                    file_types=file_type_counts,
+                    index_size_mb=index_size_mb,
+                    last_updated="unknown",
+                    embedding_model="unknown",
+                )
+
+        except Exception as e:
+            logger.error(f"Failed to get database statistics: {e}")
+            # Return empty stats instead of raising
+            return IndexStats(
+                total_files=0,
+                total_chunks=0,
+                languages={},
+                file_types={},
+                index_size_mb=0.0,
+                last_updated="error",
+                embedding_model="unknown",
+            )
+
+    async def remove_file_chunks(self, file_path: str) -> int:
+        """Remove all chunks for a specific file using pooled connection."""
+        try:
+            async with self._pool.get_connection() as conn:
+                # Get all chunks for this file
+                results = conn.collection.get(where={"file_path": file_path})
+
+                if not results["ids"]:
+                    return 0
+
+                # Delete the chunks
+                conn.collection.delete(ids=results["ids"])
+
+                return len(results["ids"])
+
+        except Exception as e:
+            logger.error(f"Failed to remove chunks for file {file_path}: {e}")
+            return 0
+
+    async def reset(self) -> None:
+        """Reset the database using pooled connection."""
+        try:
+            async with self._pool.get_connection() as conn:
+                conn.client.reset()
+                # Reinitialize the pool after reset
+                await self._pool.close()
+                await self._pool.initialize()
+                logger.info("Database reset successfully")
+        except Exception as e:
+            logger.error(f"Failed to reset database: {e}")
+            raise DatabaseError(f"Failed to reset database: {e}") from e
+
+    async def get_all_chunks(self) -> list[CodeChunk]:
+        """Get all chunks from the database using pooled connection.
+
+        Returns:
+            List of all code chunks with metadata
+        """
+        try:
+            async with self._pool.get_connection() as conn:
+                # Get all documents from collection
+                results = conn.collection.get(include=["metadatas", "documents"])
+
+                chunks = []
+                if results and results.get("ids"):
+                    for i, _chunk_id in enumerate(results["ids"]):
+                        metadata = results["metadatas"][i]
+                        content = results["documents"][i]
+
+                        # Parse JSON strings back to lists/dicts
+                        child_chunk_ids = metadata.get("child_chunk_ids", "[]")
+                        if isinstance(child_chunk_ids, str):
+                            child_chunk_ids = json.loads(child_chunk_ids)
+
+                        decorators = metadata.get("decorators", "[]")
+                        if isinstance(decorators, str):
+                            decorators = json.loads(decorators)
+
+                        parameters = metadata.get("parameters", "[]")
+                        if isinstance(parameters, str):
+                            parameters = json.loads(parameters)
+
+                        type_annotations = metadata.get("type_annotations", "{}")
+                        if isinstance(type_annotations, str):
+                            type_annotations = json.loads(type_annotations)
+
+                        chunk = CodeChunk(
+                            content=content,
+                            file_path=Path(metadata["file_path"]),
+                            start_line=metadata["start_line"],
+                            end_line=metadata["end_line"],
+                            language=metadata["language"],
+                            chunk_type=metadata.get("chunk_type", "code"),
+                            function_name=metadata.get("function_name"),
+                            class_name=metadata.get("class_name"),
+                            docstring=metadata.get("docstring"),
+                            imports=metadata.get("imports", []),
+                            complexity_score=metadata.get("complexity_score", 0.0),
+                            chunk_id=metadata.get("chunk_id"),
+                            parent_chunk_id=metadata.get("parent_chunk_id"),
+                            child_chunk_ids=child_chunk_ids,
+                            chunk_depth=metadata.get("chunk_depth", 0),
+                            decorators=decorators,
+                            parameters=parameters,
+                            return_type=metadata.get("return_type"),
+                            type_annotations=type_annotations,
+                            subproject_name=metadata.get("subproject_name"),
+                            subproject_path=metadata.get("subproject_path"),
+                        )
+                        chunks.append(chunk)
+
+                logger.debug(f"Retrieved {len(chunks)} chunks from database")
+                return chunks
+
+        except Exception as e:
+            logger.error(f"Failed to get all chunks: {e}")
+            raise DatabaseError(f"Failed to get all chunks: {e}") from e
+
+    def _build_where_clause(self, filters: dict[str, Any]) -> dict[str, Any] | None:
+        """Build ChromaDB where clause from filters."""
+        if not filters:
+            return None
+
+        conditions = []
+
+        for key, value in filters.items():
+            if key == "language" and value:
+                conditions.append({"language": {"$eq": value}})
+            elif key == "file_path" and value:
+                if isinstance(value, list):
+                    conditions.append({"file_path": {"$in": [str(p) for p in value]}})
+                else:
+                    conditions.append({"file_path": {"$eq": str(value)}})
+            elif key == "chunk_type" and value:
+                conditions.append({"chunk_type": {"$eq": value}})
+
+        if not conditions:
+            return None
+        elif len(conditions) > 1:
+            return {"$and": conditions}
+        else:
+            return conditions[0]
+
+    def get_pool_stats(self) -> dict[str, Any]:
+        """Get connection pool statistics."""
+        return self._pool.get_stats()
+
+    async def health_check(self) -> bool:
+        """Perform a health check on the database and connection pool."""
+        try:
+            # Check pool health
+            pool_healthy = await self._pool.health_check()
+            if not pool_healthy:
+                return False
+
+            # Try a simple query to verify database integrity
+            try:
+                async with self._pool.get_connection() as conn:
+                    # Test basic operations
+                    conn.collection.count()
+                    conn.collection.query(
+                        query_texts=["test"], n_results=1, include=["metadatas"]
+                    )
+                return True
+            except Exception as e:
+                error_msg = str(e).lower()
+                if any(
+                    indicator in error_msg
+                    for indicator in [
+                        "pickle",
+                        "unpickling",
+                        "eof",
+                        "ran out of input",
+                        "hnsw",
+                        "index",
+                        "deserialize",
+                        "corrupt",
+                    ]
+                ):
+                    logger.error(f"Index corruption detected: {e}")
+                    # Attempt recovery
+                    await self._recover_from_corruption()
+                    return False
+                else:
+                    logger.warning(f"Health check failed: {e}")
+                    return False
+        except Exception as e:
+            logger.error(f"Health check error: {e}")
+            return False
+
+    async def _recover_from_corruption(self) -> None:
+        """Recover from index corruption by rebuilding the index."""
+        logger.info("Attempting to recover from index corruption...")
+
+        # Close the pool first
+        await self._pool.close()
+
+        # Create backup directory
+        backup_dir = (
+            self.persist_directory.parent / f"{self.persist_directory.name}_backup"
+        )
+        backup_dir.mkdir(exist_ok=True)
+
+        # Backup current state
+        import time
+
+        timestamp = int(time.time())
+        backup_path = backup_dir / f"backup_{timestamp}"
+
+        if self.persist_directory.exists():
+            try:
+                shutil.copytree(self.persist_directory, backup_path)
+                logger.info(f"Created backup at {backup_path}")
+            except Exception as e:
+                logger.warning(f"Could not create backup: {e}")
+
+        # Clear the corrupted index
+        if self.persist_directory.exists():
+            try:
+                shutil.rmtree(self.persist_directory)
+                logger.info(f"Cleared corrupted index at {self.persist_directory}")
+            except Exception as e:
+                logger.error(f"Failed to clear corrupted index: {e}")
+                raise IndexCorruptionError(
+                    f"Could not clear corrupted index: {e}"
+                ) from e
+
+        # Recreate the directory
+        self.persist_directory.mkdir(parents=True, exist_ok=True)
+
+        # Reinitialize the pool
+        await self._pool.initialize()
+        logger.info("Index recovered. Please re-index your codebase.")
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        await self.initialize()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.close()