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

mcp_vector_search/core/directory_index.py

@@ -0,0 +1,318 @@
+"""Directory index for tracking project structure and file relationships."""
+
+import json
+from collections import defaultdict
+from pathlib import Path
+from typing import Any
+
+from loguru import logger
+
+from .models import Directory
+
+
+class DirectoryIndex:
+    """Manages directory structure and file-directory relationships."""
+
+    def __init__(self, index_path: Path) -> None:
+        """Initialize directory index.
+
+        Args:
+            index_path: Path to directory index file (JSON)
+        """
+        self.index_path = index_path
+        self.directories: dict[str, Directory] = {}  # path -> Directory
+        self.file_to_directory: dict[str, str] = {}  # file_path -> directory_path
+        self.directory_files: dict[str, list[str]] = defaultdict(
+            list
+        )  # dir_path -> [file_paths]
+
+    def load(self) -> None:
+        """Load directory index from disk."""
+        if not self.index_path.exists():
+            logger.debug("No directory index found, starting fresh")
+            return
+
+        try:
+            with open(self.index_path) as f:
+                data = json.load(f)
+
+            # Load directories
+            for dir_data in data.get("directories", []):
+                directory = Directory.from_dict(dir_data)
+                self.directories[str(directory.path)] = directory
+
+            # Load file mappings
+            self.file_to_directory = data.get("file_to_directory", {})
+
+            # Rebuild directory_files from file_to_directory
+            self.directory_files = defaultdict(list)
+            for file_path, dir_path in self.file_to_directory.items():
+                self.directory_files[dir_path].append(file_path)
+
+            logger.info(f"Loaded {len(self.directories)} directories from index")
+
+        except Exception as e:
+            logger.error(f"Failed to load directory index: {e}")
+            self.directories = {}
+            self.file_to_directory = {}
+            self.directory_files = defaultdict(list)
+
+    def save(self) -> None:
+        """Save directory index to disk."""
+        try:
+            # Ensure parent directory exists
+            self.index_path.parent.mkdir(parents=True, exist_ok=True)
+
+            data = {
+                "directories": [d.to_dict() for d in self.directories.values()],
+                "file_to_directory": self.file_to_directory,
+            }
+
+            with open(self.index_path, "w") as f:
+                json.dump(data, f, indent=2)
+
+            logger.debug(f"Saved {len(self.directories)} directories to index")
+
+        except Exception as e:
+            logger.error(f"Failed to save directory index: {e}")
+            raise
+
+    def add_directory(self, directory: Directory) -> None:
+        """Add or update a directory in the index.
+
+        Args:
+            directory: Directory to add
+        """
+        dir_path = str(directory.path)
+        self.directories[dir_path] = directory
+
+    def add_file(self, file_path: Path, directory_path: Path) -> None:
+        """Associate a file with its directory.
+
+        Args:
+            file_path: Path to the file
+            directory_path: Path to the directory containing the file
+        """
+        file_path_str = str(file_path)
+        dir_path_str = str(directory_path)
+
+        self.file_to_directory[file_path_str] = dir_path_str
+        if file_path_str not in self.directory_files[dir_path_str]:
+            self.directory_files[dir_path_str].append(file_path_str)
+
+        # Update directory file count
+        if dir_path_str in self.directories:
+            self.directories[dir_path_str].file_count = len(
+                self.directory_files[dir_path_str]
+            )
+
+    def get_directory(self, directory_path: Path) -> Directory | None:
+        """Get directory by path.
+
+        Args:
+            directory_path: Path to directory
+
+        Returns:
+            Directory object or None if not found
+        """
+        return self.directories.get(str(directory_path))
+
+    def get_files_in_directory(self, directory_path: Path) -> list[str]:
+        """Get all files in a directory.
+
+        Args:
+            directory_path: Path to directory
+
+        Returns:
+            List of file paths in the directory
+        """
+        return self.directory_files.get(str(directory_path), [])
+
+    def get_subdirectories(self, directory_path: Path) -> list[Directory]:
+        """Get all immediate subdirectories.
+
+        Args:
+            directory_path: Path to parent directory
+
+        Returns:
+            List of subdirectory objects
+        """
+        parent_path_str = str(directory_path)
+        subdirs = []
+
+        for _dir_path_str, directory in self.directories.items():
+            if directory.parent_path and str(directory.parent_path) == parent_path_str:
+                subdirs.append(directory)
+
+        return subdirs
+
+    def get_root_directories(self) -> list[Directory]:
+        """Get all root-level directories (no parent).
+
+        Returns:
+            List of root directory objects
+        """
+        return [d for d in self.directories.values() if d.parent_path is None]
+
+    def delete_directory(self, directory_path: Path) -> None:
+        """Remove directory and its file associations.
+
+        Args:
+            directory_path: Path to directory to remove
+        """
+        dir_path_str = str(directory_path)
+
+        # Remove directory
+        if dir_path_str in self.directories:
+            del self.directories[dir_path_str]
+
+        # Remove file associations
+        if dir_path_str in self.directory_files:
+            for file_path in self.directory_files[dir_path_str]:
+                if file_path in self.file_to_directory:
+                    del self.file_to_directory[file_path]
+            del self.directory_files[dir_path_str]
+
+    def delete_file(self, file_path: Path) -> None:
+        """Remove file from directory associations.
+
+        Args:
+            file_path: Path to file to remove
+        """
+        file_path_str = str(file_path)
+
+        if file_path_str in self.file_to_directory:
+            dir_path = self.file_to_directory[file_path_str]
+            del self.file_to_directory[file_path_str]
+
+            # Remove from directory_files
+            if dir_path in self.directory_files:
+                self.directory_files[dir_path] = [
+                    f for f in self.directory_files[dir_path] if f != file_path_str
+                ]
+
+                # Update directory file count
+                if dir_path in self.directories:
+                    self.directories[dir_path].file_count = len(
+                        self.directory_files[dir_path]
+                    )
+
+    def rebuild_from_files(
+        self,
+        file_paths: list[Path],
+        root_path: Path,
+        chunk_stats: dict[str, dict] | None = None,
+    ) -> None:
+        """Rebuild directory index from list of files with statistics from chunks.
+
+        Args:
+            file_paths: List of file paths to index
+            root_path: Project root path
+            chunk_stats: Optional dict mapping file_path -> {'chunks': count, 'language': str}
+        """
+        self.directories = {}
+        self.file_to_directory = {}
+        self.directory_files = defaultdict(list)
+
+        # Track all unique directories and their statistics
+        dir_set = set()
+        dir_chunks = defaultdict(int)  # directory -> total chunks
+        dir_languages = defaultdict(
+            lambda: defaultdict(int)
+        )  # directory -> {language: count}
+        dir_modified = defaultdict(float)  # directory -> most recent modification time
+
+        for file_path in file_paths:
+            # Get relative path from root
+            try:
+                rel_path = file_path.relative_to(root_path)
+                parent_dir = rel_path.parent
+
+                # Add all parent directories up to root
+                current = parent_dir
+                while current != Path("."):
+                    dir_set.add(current)
+
+                    # Accumulate statistics up the directory tree
+                    if chunk_stats and str(file_path) in chunk_stats:
+                        stats = chunk_stats[str(file_path)]
+                        dir_chunks[current] += stats.get("chunks", 0)
+                        if "language" in stats:
+                            dir_languages[current][stats["language"]] += stats.get(
+                                "chunks", 0
+                            )
+                        # Track most recent modification time
+                        if "modified" in stats:
+                            dir_modified[current] = max(
+                                dir_modified.get(current, 0), stats["modified"]
+                            )
+
+                    current = current.parent
+
+                # Associate file with its direct parent
+                if parent_dir != Path("."):
+                    self.add_file(rel_path, parent_dir)
+
+            except ValueError:
+                # File not relative to root, skip
+                logger.warning(f"File {file_path} not under root {root_path}")
+                continue
+
+        # Create Directory objects for all directories
+        for dir_path in sorted(dir_set):
+            # Determine parent
+            parent_path = dir_path.parent if dir_path.parent != Path(".") else None
+
+            # Check if it's a package
+            is_package = False
+            full_dir_path = root_path / dir_path
+            if (full_dir_path / "__init__.py").exists():
+                is_package = True
+            elif (full_dir_path / "package.json").exists():
+                is_package = True
+
+            directory = Directory(
+                path=dir_path,
+                name=dir_path.name,
+                parent_path=parent_path,
+                depth=len(dir_path.parts),
+                is_package=is_package,
+                total_chunks=dir_chunks.get(dir_path, 0),
+                languages=dict(dir_languages.get(dir_path, {})),
+                last_modified=dir_modified.get(dir_path),
+            )
+
+            self.add_directory(directory)
+
+        # Update subdirectory counts
+        for directory in self.directories.values():
+            subdirs = self.get_subdirectories(directory.path)
+            directory.subdirectory_count = len(subdirs)
+
+        logger.info(
+            f"Rebuilt directory index with {len(self.directories)} directories, {sum(dir_chunks.values())} total chunks"
+        )
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get directory index statistics.
+
+        Returns:
+            Dictionary with statistics
+        """
+        return {
+            "total_directories": len(self.directories),
+            "total_files": len(self.file_to_directory),
+            "root_directories": len(self.get_root_directories()),
+            "packages": sum(1 for d in self.directories.values() if d.is_package),
+        }
+
+    def reset(self) -> None:
+        """Clear all directory data."""
+        self.directories = {}
+        self.file_to_directory = {}
+        self.directory_files = defaultdict(list)
+
+        if self.index_path.exists():
+            self.index_path.unlink()
+
+        logger.info("Directory index reset")

mcp_vector_search/core/embeddings.py

@@ -0,0 +1,294 @@
+"""Embedding generation for MCP Vector Search."""
+
+import hashlib
+import json
+from pathlib import Path
+
+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]] = {}
+        self._access_order: list[str] = []  # For LRU eviction
+        self._cache_hits = 0
+        self._cache_misses = 0
+
+    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) -> list[float] | None:
+        """Get cached embedding for content."""
+        cache_key = self._hash_content(content)
+
+        # Check memory cache first
+        if cache_key in self._memory_cache:
+            self._cache_hits += 1
+            # Move to end for LRU
+            self._access_order.remove(cache_key)
+            self._access_order.append(cache_key)
+            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) as f:
+                    content_str = await f.read()
+                    embedding = json.loads(content_str)
+
+                    # Add to memory cache with LRU management
+                    self._add_to_memory_cache(cache_key, embedding)
+                    self._cache_hits += 1
+                    return embedding
+            except Exception as e:
+                logger.warning(f"Failed to load cached embedding: {e}")
+
+        self._cache_misses += 1
+        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 with LRU management
+        self._add_to_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 _add_to_memory_cache(self, cache_key: str, embedding: list[float]) -> None:
+        """Add embedding to memory cache with LRU eviction.
+
+        Args:
+            cache_key: Cache key for the embedding
+            embedding: Embedding vector to cache
+        """
+        # If already in cache, update and move to end
+        if cache_key in self._memory_cache:
+            self._access_order.remove(cache_key)
+            self._access_order.append(cache_key)
+            self._memory_cache[cache_key] = embedding
+            return
+
+        # If cache is full, evict least recently used
+        if len(self._memory_cache) >= self.max_size:
+            lru_key = self._access_order.pop(0)
+            del self._memory_cache[lru_key]
+
+        # Add new embedding
+        self._memory_cache[cache_key] = embedding
+        self._access_order.append(cache_key)
+
+    def clear_memory_cache(self) -> None:
+        """Clear the in-memory cache."""
+        self._memory_cache.clear()
+        self._access_order.clear()
+
+    def get_cache_stats(self) -> dict[str, any]:
+        """Get cache performance statistics.
+
+        Returns:
+            Dictionary with cache statistics
+        """
+        total_requests = self._cache_hits + self._cache_misses
+        hit_rate = self._cache_hits / total_requests if total_requests > 0 else 0.0
+        disk_files = (
+            len(list(self.cache_dir.glob("*.json"))) if self.cache_dir.exists() else 0
+        )
+
+        return {
+            "memory_cache_size": len(self._memory_cache),
+            "memory_cached": len(self._memory_cache),  # Alias for compatibility
+            "max_cache_size": self.max_size,
+            "memory_limit": self.max_size,  # Alias for compatibility
+            "cache_hits": self._cache_hits,
+            "cache_misses": self._cache_misses,
+            "hit_rate": round(hit_rate, 3),
+            "disk_cache_files": disk_files,
+            "disk_cached": disk_files,  # Alias for compatibility
+        }
+
+
+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: EmbeddingCache | None = 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, strict=False)
+                ):
+                    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: Path | None = 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.debug(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

mcp_vector_search/core/exceptions.py

@@ -0,0 +1,89 @@
+"""Custom exception hierarchy for MCP Vector Search."""
+
+from typing import Any
+
+
+class MCPVectorSearchError(Exception):
+    """Base exception for MCP Vector Search."""
+
+    def __init__(self, message: str, context: dict[str, Any] | None = None) -> None:
+        super().__init__(message)
+        self.context = context or {}
+
+
+class DatabaseError(MCPVectorSearchError):
+    """Database-related errors."""
+
+    pass
+
+
+class DatabaseInitializationError(DatabaseError):
+    """Database initialization failed."""
+
+    pass
+
+
+class DatabaseNotInitializedError(DatabaseError):
+    """Operation attempted on uninitialized database."""
+
+    pass
+
+
+class ConnectionPoolError(DatabaseError):
+    """Connection pool operation failed."""
+
+    pass
+
+
+class DocumentAdditionError(DatabaseError):
+    """Failed to add documents to database."""
+
+    pass
+
+
+class SearchError(DatabaseError):
+    """Search operation failed."""
+
+    pass
+
+
+class IndexCorruptionError(DatabaseError):
+    """Index corruption detected."""
+
+    pass
+
+
+class ParsingError(MCPVectorSearchError):
+    """Code parsing errors."""
+
+    pass
+
+
+class EmbeddingError(MCPVectorSearchError):
+    """Embedding generation errors."""
+
+    pass
+
+
+class ConfigurationError(MCPVectorSearchError):
+    """Configuration validation errors."""
+
+    pass
+
+
+class ProjectError(MCPVectorSearchError):
+    """Project management errors."""
+
+    pass
+
+
+class ProjectNotFoundError(ProjectError):
+    """Project directory or configuration not found."""
+
+    pass
+
+
+class ProjectInitializationError(ProjectError):
+    """Failed to initialize project."""
+
+    pass