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

mcp_vector_search/core/database.py

@@ -0,0 +1,431 @@
+"""Database abstraction and ChromaDB implementation for MCP Vector Search."""
+
+import asyncio
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
+
+from loguru import logger
+
+from .exceptions import (
+    DatabaseError,
+    DatabaseInitializationError,
+    DatabaseNotInitializedError,
+    DocumentAdditionError,
+    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: Optional[Dict[str, Any]] = 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)."""
+        ...
+
+    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."""
+        try:
+            import chromadb
+
+            # Ensure directory exists
+            self.persist_directory.mkdir(parents=True, exist_ok=True)
+
+            # 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.info(f"ChromaDB initialized at {self.persist_directory}")
+
+        except Exception as e:
+            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,
+                }
+                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: Optional[Dict[str, Any]] = 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],
+                    )
+                ):
+                    # Convert distance to similarity (ChromaDB uses cosine distance)
+                    similarity = 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."""
+        if not self._collection:
+            raise DatabaseNotInitializedError("Database not initialized")
+
+        try:
+            # Get total count
+            count = self._collection.count()
+
+            # Get sample for language distribution
+            sample_results = self._collection.get(
+                limit=min(1000, count) if count > 0 else 0,
+                include=["metadatas"],
+            )
+
+            languages = {}
+            file_types = {}
+
+            if sample_results["metadatas"]:
+                for metadata in sample_results["metadatas"]:
+                    # Count languages
+                    lang = metadata.get("language", "unknown")
+                    languages[lang] = languages.get(lang, 0) + 1
+
+                    # Count file types
+                    file_path = metadata.get("file_path", "")
+                    ext = Path(file_path).suffix or "no_extension"
+                    file_types[ext] = file_types.get(ext, 0) + 1
+
+            # Estimate index size (rough approximation)
+            index_size_mb = count * 0.001  # Rough estimate
+
+            return IndexStats(
+                total_files=len(set(m.get("file_path", "") for m in sample_results.get("metadatas", []))),
+                total_chunks=count,
+                languages=languages,
+                file_types=file_types,
+                index_size_mb=index_size_mb,
+                last_updated="unknown",  # TODO: Track this
+                embedding_model="unknown",  # TODO: Track this
+            )
+
+        except Exception as e:
+            logger.error(f"Failed to get stats: {e}")
+            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
+
+    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

mcp_vector_search/core/embeddings.py

@@ -0,0 +1,250 @@
+"""Embedding generation for MCP Vector Search."""
+
+import hashlib
+import json
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import aiofiles
+from loguru import logger
+from sentence_transformers import SentenceTransformer
+
+from .exceptions import EmbeddingError
+
+
+class EmbeddingCache:
+    """LRU cache for embeddings with disk persistence."""
+
+    def __init__(self, cache_dir: Path, max_size: int = 1000) -> None:
+        """Initialize embedding cache.
+        
+        Args:
+            cache_dir: Directory to store cached embeddings
+            max_size: Maximum number of embeddings to keep in memory
+        """
+        self.cache_dir = cache_dir
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self.max_size = max_size
+        self._memory_cache: Dict[str, List[float]] = {}
+
+    def _hash_content(self, content: str) -> str:
+        """Generate cache key from content."""
+        return hashlib.sha256(content.encode()).hexdigest()[:16]
+
+    async def get_embedding(self, content: str) -> Optional[List[float]]:
+        """Get cached embedding for content."""
+        cache_key = self._hash_content(content)
+
+        # Check memory cache first
+        if cache_key in self._memory_cache:
+            return self._memory_cache[cache_key]
+
+        # Check disk cache
+        cache_file = self.cache_dir / f"{cache_key}.json"
+        if cache_file.exists():
+            try:
+                async with aiofiles.open(cache_file, "r") as f:
+                    content_str = await f.read()
+                    embedding = json.loads(content_str)
+
+                    # Add to memory cache if space available
+                    if len(self._memory_cache) < self.max_size:
+                        self._memory_cache[cache_key] = embedding
+
+                    return embedding
+            except Exception as e:
+                logger.warning(f"Failed to load cached embedding: {e}")
+
+        return None
+
+    async def store_embedding(self, content: str, embedding: List[float]) -> None:
+        """Store embedding in cache."""
+        cache_key = self._hash_content(content)
+
+        # Store in memory cache if space available
+        if len(self._memory_cache) < self.max_size:
+            self._memory_cache[cache_key] = embedding
+
+        # Store in disk cache
+        cache_file = self.cache_dir / f"{cache_key}.json"
+        try:
+            async with aiofiles.open(cache_file, "w") as f:
+                await f.write(json.dumps(embedding))
+        except Exception as e:
+            logger.warning(f"Failed to cache embedding: {e}")
+
+    def clear_memory_cache(self) -> None:
+        """Clear the in-memory cache."""
+        self._memory_cache.clear()
+
+    def get_cache_stats(self) -> Dict[str, int]:
+        """Get cache statistics."""
+        disk_files = len(list(self.cache_dir.glob("*.json")))
+        return {
+            "memory_cached": len(self._memory_cache),
+            "disk_cached": disk_files,
+            "memory_limit": self.max_size,
+        }
+
+
+class CodeBERTEmbeddingFunction:
+    """ChromaDB-compatible embedding function using CodeBERT."""
+
+    def __init__(self, model_name: str = "microsoft/codebert-base") -> None:
+        """Initialize CodeBERT embedding function.
+
+        Args:
+            model_name: Name of the sentence transformer model
+        """
+        try:
+            self.model = SentenceTransformer(model_name)
+            self.model_name = model_name
+            logger.info(f"Loaded embedding model: {model_name}")
+        except Exception as e:
+            logger.error(f"Failed to load embedding model {model_name}: {e}")
+            raise EmbeddingError(f"Failed to load embedding model: {e}") from e
+
+    def __call__(self, input: List[str]) -> List[List[float]]:
+        """Generate embeddings for input texts (ChromaDB interface)."""
+        try:
+            embeddings = self.model.encode(input, convert_to_numpy=True)
+            return embeddings.tolist()
+        except Exception as e:
+            logger.error(f"Failed to generate embeddings: {e}")
+            raise EmbeddingError(f"Failed to generate embeddings: {e}") from e
+
+
+class BatchEmbeddingProcessor:
+    """Batch processing for efficient embedding generation with caching."""
+
+    def __init__(
+        self,
+        embedding_function: CodeBERTEmbeddingFunction,
+        cache: Optional[EmbeddingCache] = None,
+        batch_size: int = 32,
+    ) -> None:
+        """Initialize batch embedding processor.
+        
+        Args:
+            embedding_function: Function to generate embeddings
+            cache: Optional embedding cache
+            batch_size: Size of batches for processing
+        """
+        self.embedding_function = embedding_function
+        self.cache = cache
+        self.batch_size = batch_size
+
+    async def process_batch(self, contents: List[str]) -> List[List[float]]:
+        """Process a batch of content for embeddings.
+        
+        Args:
+            contents: List of text content to embed
+            
+        Returns:
+            List of embeddings
+        """
+        if not contents:
+            return []
+
+        embeddings = []
+        uncached_contents = []
+        uncached_indices = []
+
+        # Check cache for each content if cache is available
+        if self.cache:
+            for i, content in enumerate(contents):
+                cached_embedding = await self.cache.get_embedding(content)
+                if cached_embedding:
+                    embeddings.append(cached_embedding)
+                else:
+                    embeddings.append(None)  # Placeholder
+                    uncached_contents.append(content)
+                    uncached_indices.append(i)
+        else:
+            # No cache, process all content
+            uncached_contents = contents
+            uncached_indices = list(range(len(contents)))
+            embeddings = [None] * len(contents)
+
+        # Generate embeddings for uncached content
+        if uncached_contents:
+            logger.debug(f"Generating {len(uncached_contents)} new embeddings")
+
+            try:
+                new_embeddings = []
+                for i in range(0, len(uncached_contents), self.batch_size):
+                    batch = uncached_contents[i : i + self.batch_size]
+                    batch_embeddings = self.embedding_function(batch)
+                    new_embeddings.extend(batch_embeddings)
+
+                # Cache new embeddings and fill placeholders
+                for i, (content, embedding) in enumerate(
+                    zip(uncached_contents, new_embeddings)
+                ):
+                    if self.cache:
+                        await self.cache.store_embedding(content, embedding)
+                    embeddings[uncached_indices[i]] = embedding
+
+            except Exception as e:
+                logger.error(f"Failed to generate embeddings: {e}")
+                raise EmbeddingError(f"Failed to generate embeddings: {e}") from e
+
+        return embeddings
+
+    def get_stats(self) -> Dict[str, any]:
+        """Get processor statistics."""
+        stats = {
+            "model_name": self.embedding_function.model_name,
+            "batch_size": self.batch_size,
+            "cache_enabled": self.cache is not None,
+        }
+
+        if self.cache:
+            stats.update(self.cache.get_cache_stats())
+
+        return stats
+
+
+def create_embedding_function(
+    model_name: str = "microsoft/codebert-base",
+    cache_dir: Optional[Path] = None,
+    cache_size: int = 1000,
+):
+    """Create embedding function and cache.
+
+    Args:
+        model_name: Name of the embedding model
+        cache_dir: Directory for caching embeddings
+        cache_size: Maximum cache size
+
+    Returns:
+        Tuple of (embedding_function, cache)
+    """
+    try:
+        # Use ChromaDB's built-in sentence transformer function
+        from chromadb.utils import embedding_functions
+
+        # Map our model names to sentence-transformers compatible names
+        model_mapping = {
+            "microsoft/codebert-base": "sentence-transformers/all-MiniLM-L6-v2",  # Fallback to working model
+            "microsoft/unixcoder-base": "sentence-transformers/all-MiniLM-L6-v2",  # Fallback to working model
+        }
+
+        actual_model = model_mapping.get(model_name, model_name)
+
+        embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(
+            model_name=actual_model
+        )
+
+        logger.info(f"Created ChromaDB embedding function with model: {actual_model}")
+
+    except Exception as e:
+        logger.warning(f"Failed to create ChromaDB embedding function: {e}")
+        # Fallback to our custom implementation
+        embedding_function = CodeBERTEmbeddingFunction(model_name)
+
+    cache = None
+    if cache_dir:
+        cache = EmbeddingCache(cache_dir, cache_size)
+
+    return embedding_function, cache