mcp-code-indexer 4.2.15__py3-none-any.whl → 4.2.16__py3-none-any.whl

mcp_code_indexer/vector_mode/providers/voyage_client.py

@@ -6,25 +6,29 @@ high-quality code embeddings using the voyage-code-2 model.
 """
 
 import logging
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Tuple
 import voyageai
 
-from ..config import VectorConfig
+
+
+from ..config import VectorConfig, DEFAULT_EMBEDDING_MODEL
+from ..const import MODEL_DIMENSIONS
 
 logger = logging.getLogger(__name__)
 
+
 class VoyageClient:
     """Clean Voyage AI client using official SDK."""
-    
-    def __init__(self, api_key: str, model: str = "voyage-code-2"):
+
+    def __init__(self, api_key: str, model: str = DEFAULT_EMBEDDING_MODEL):
         self.api_key = api_key
         self.model = model
         self._embedding_dimension: int | None = None
-        
+
         # Initialize official Voyage AI client
         self.client = voyageai.Client(api_key=api_key)
         logger.info(f"Initialized Voyage AI client with model {model}")
-    
+
     def health_check(self) -> bool:
         """Check if Voyage AI service is healthy."""
         try:
@@ -33,74 +37,108 @@ class VoyageClient:
         except Exception as e:
             logger.warning(f"Voyage AI health check failed: {e}")
             return False
-    
+
+    def validate_api_access(self) -> None:
+        """
+        Validate API key and access to Voyage AI service.
+
+        Raises:
+            RuntimeError: If API access validation fails with specific error details
+        """
+        logger.info("Validating Voyage AI API access...")
+        try:
+            result = self.client.embed(["test"], model=self.model, input_type="query")
+            if not result or not result.embeddings:
+                raise RuntimeError("Voyage AI API returned empty response")
+            logger.debug("Voyage AI API access validated successfully")
+        except Exception as e:
+            error_msg = str(e).lower()
+
+            if (
+                "401" in error_msg
+                or "unauthorized" in error_msg
+                or "api key" in error_msg
+            ):
+                raise RuntimeError(
+                    f"Voyage AI API authentication failed: Invalid or expired API key. "
+                    f"Please check your VOYAGE_API_KEY. Error: {e}"
+                )
+            elif "403" in error_msg or "forbidden" in error_msg:
+                raise RuntimeError(
+                    f"Voyage AI API access denied: API key lacks required permissions. Error: {e}"
+                )
+            elif "429" in error_msg or "rate limit" in error_msg:
+                raise RuntimeError(
+                    f"Voyage AI API rate limit exceeded: Too many requests. Error: {e}"
+                )
+            elif "quota" in error_msg or "usage" in error_msg:
+                raise RuntimeError(
+                    f"Voyage AI API quota exceeded: Usage limit reached. Error: {e}"
+                )
+            elif "5" in error_msg and ("error" in error_msg or "server" in error_msg):
+                raise RuntimeError(
+                    f"Voyage AI service unavailable: Server error. Error: {e}"
+                )
+            else:
+                raise RuntimeError(f"Voyage AI API access validation failed: {e}")
+
+        logger.info("Voyage AI API access validated successfully")
+
     def generate_embeddings(
-        self,
-        texts: List[str],
-        input_type: str = "document",
-        **kwargs
+        self, texts: List[str], input_type: str = "document", **kwargs
     ) -> List[List[float]]:
         """Generate embeddings for texts using official SDK."""
         if not texts:
             return []
-        
+
         logger.info(f"Generating embeddings for {len(texts)} texts using {self.model}")
-        
+
         try:
             result = self.client.embed(
-                texts=texts,
-                model=self.model,
-                input_type=input_type,
-                truncation=True
+                texts=texts, model=self.model, input_type=input_type, truncation=True
             )
-            
+
             # Log usage if available
-            if hasattr(result, 'usage') and result.usage:
+            if hasattr(result, "usage") and result.usage:
                 logger.debug(f"Token usage: {result.usage.total_tokens}")
-            
+
             logger.info(f"Successfully generated {len(result.embeddings)} embeddings")
             return result.embeddings
-            
+
         except Exception as e:
             logger.error(f"Failed to generate embeddings: {e}")
             raise RuntimeError(f"Embedding generation failed: {e}")
-    
+
     def get_embedding_dimension(self) -> int:
         """Get the dimension of embeddings produced by this model."""
         if self._embedding_dimension is not None:
             return self._embedding_dimension
-        
+
         # Generate a test embedding to determine dimension
         try:
             test_embeddings = self.generate_embeddings(["test"], input_type="query")
             if test_embeddings:
                 self._embedding_dimension = len(test_embeddings[0])
-                logger.info(f"Detected embedding dimension: {self._embedding_dimension}")
+                logger.info(
+                    f"Detected embedding dimension: {self._embedding_dimension}"
+                )
                 return self._embedding_dimension
         except Exception as e:
             logger.warning(f"Could not determine embedding dimension: {e}")
-        
-        # Default dimensions for known Voyage models
-        model_dimensions = {
-            "voyage-code-2": 1536,
-            "voyage-2": 1024,
-            "voyage-large-2": 1536,
-            "voyage-3": 1024,
-        }
-        
-        self._embedding_dimension = model_dimensions.get(self.model, 1536)
+
+        self._embedding_dimension = MODEL_DIMENSIONS[self.model]
         logger.info(f"Using default embedding dimension: {self._embedding_dimension}")
         return self._embedding_dimension
-    
+
     def estimate_cost(self, texts: List[str]) -> Dict[str, Any]:
         """Estimate the cost of embedding generation."""
         # Rough token estimation (4 chars per token)
         total_tokens = sum(len(text) // 4 for text in texts)
-        
+
         # Voyage AI pricing (approximate, may change)
         cost_per_1k_tokens = 0.00013  # voyage-code-2 pricing
         estimated_cost = (total_tokens / 1000) * cost_per_1k_tokens
-        
+
         return {
             "total_tokens": total_tokens,
             "total_texts": len(texts),
@@ -108,11 +146,75 @@ class VoyageClient:
             "model": self.model,
         }
 
+    def count_tokens(self, texts: List[str]) -> int:
+        return self.client.count_tokens(texts, self.model)
+
+    def generate_embeddings_batch(
+        self,
+        all_texts: List[str],
+        file_boundaries: List[Tuple[str, int, int]],
+        input_type: str = "document",
+        **kwargs,
+    ) -> Dict[str, List[List[float]]]:
+        """
+        Generate embeddings for texts from multiple files in a single batch call.
+
+        Note: Token limits should be handled at the calling level. This method
+        assumes the provided texts fit within API limits.
+
+        Args:
+            all_texts: Flattened list of all text chunks from all files
+            file_boundaries: List of (file_path, start_idx, end_idx) tuples indicating
+                           which embeddings belong to which file
+            input_type: Type of input for embedding generation
+            **kwargs: Additional arguments for embedding generation
+
+        Returns:
+            Dictionary mapping file paths to their corresponding embeddings
+
+        Raises:
+            RuntimeError: If embedding generation fails
+        """
+        if not all_texts:
+            return {}
+
+        if not file_boundaries:
+            raise ValueError(
+                "file_boundaries cannot be empty when all_texts is provided"
+            )
+
+        logger.info(
+            f"Generating batch embeddings for {len(all_texts)} texts from {len(file_boundaries)} files using {self.model}"
+        )
+
+        try:
+            # Generate embeddings for all texts in a single API call
+            all_embeddings = self.generate_embeddings(
+                all_texts, input_type=input_type, **kwargs
+            )
+
+            # Group embeddings by file using boundaries
+            file_embeddings = {}
+            for file_path, start_idx, end_idx in file_boundaries:
+                file_embeddings[file_path] = all_embeddings[start_idx:end_idx]
+
+            logger.info(
+                f"Successfully generated batch embeddings for {len(file_boundaries)} files "
+                f"({len(all_embeddings)} total embeddings)"
+            )
+
+            return file_embeddings
+
+        except Exception as e:
+            logger.error(f"Failed to generate batch embeddings: {e}")
+            raise RuntimeError(f"Batch embedding generation failed: {e}")
+
+
 def create_voyage_client(config: VectorConfig) -> VoyageClient:
     """Create a Voyage client from configuration."""
     if not config.voyage_api_key:
         raise ValueError("VOYAGE_API_KEY is required for embedding generation")
-    
+
     return VoyageClient(
         api_key=config.voyage_api_key,
         model=config.embedding_model,

mcp_code_indexer/vector_mode/services/__init__.py

@@ -0,0 +1,9 @@
+"""
+Vector mode services module.
+
+Contains service classes that handle specialized vector operations.
+"""
+
+from .embedding_service import EmbeddingService
+
+__all__ = ["EmbeddingService"]

mcp_code_indexer/vector_mode/services/embedding_service.py

@@ -0,0 +1,389 @@
+"""
+Embedding Service for converting code chunks to vector embeddings.
+
+Provides a clean abstraction layer between chunked code and embedding providers,
+handling text preparation, batching, and provider communication.
+"""
+
+import asyncio
+import logging
+from pathlib import Path
+from typing import List, Dict, Tuple
+
+
+
+from ..chunking.ast_chunker import CodeChunk
+from ..providers.voyage_client import VoyageClient
+from ..config import VectorConfig
+
+logger = logging.getLogger(__name__)
+
+
+class EmbeddingService:
+    """
+    Service for converting code chunks to embeddings.
+
+    Handles text preparation, context enhancement, batching, and provider
+    communication while maintaining separation from daemon orchestration logic.
+    """
+
+    def __init__(self, embedding_client: VoyageClient, config: VectorConfig):
+        """Initialize embedding service with client and configuration."""
+        self.embedding_client = embedding_client
+        self.config = config
+
+        # Validate API access immediately during initialization
+        self.embedding_client.validate_api_access()
+
+    async def generate_embeddings_for_chunks(
+        self, chunks: List[CodeChunk], project_name: str, file_path: Path
+    ) -> List[List[float]]:
+        """
+        Generate embeddings for a list of code chunks.
+
+        Args:
+            chunks: List of code chunks to embed
+            project_name: Name of the project (for logging)
+            file_path: Path to source file (for logging)
+
+        Returns:
+            List of embedding vectors (one per chunk)
+
+        Raises:
+            RuntimeError: If embedding generation fails
+        """
+        if not chunks:
+            logger.debug(f"No chunks provided for {file_path}")
+            return []
+
+        logger.info(
+            f"Generating embeddings for {len(chunks)} chunks from {file_path}",
+            extra={
+                "structured_data": {
+                    "project_name": project_name,
+                    "file_path": str(file_path),
+                    "chunk_count": len(chunks),
+                    "embedding_model": self.config.embedding_model,
+                }
+            },
+        )
+
+        try:
+            # Extract text content from chunks with context enhancement
+            texts = self._prepare_chunk_texts(chunks)
+
+            # Process chunks in batches to respect API limits
+            all_embeddings = []
+            batch_size = self.config.batch_size
+
+            for i in range(0, len(texts), batch_size):
+                batch_texts = texts[i : i + batch_size]
+                batch_chunks = chunks[i : i + batch_size]
+
+                logger.debug(
+                    f"Processing embedding batch {i//batch_size + 1} "
+                    f"({len(batch_texts)} chunks) for {file_path}"
+                )
+
+                # Generate embeddings using async/sync bridge
+                embeddings = await asyncio.get_event_loop().run_in_executor(
+                    None,
+                    lambda: self.embedding_client.generate_embeddings(
+                        batch_texts, input_type="document"  # Code chunks are documents
+                    ),
+                )
+
+                all_embeddings.extend(embeddings)
+
+                # Log batch statistics
+                self._log_batch_stats(
+                    batch_chunks, i // batch_size + 1, len(embeddings)
+                )
+
+            logger.info(
+                f"Successfully generated {len(all_embeddings)} embeddings for {file_path}",
+                extra={
+                    "structured_data": {
+                        "project_name": project_name,
+                        "file_path": str(file_path),
+                        "embedding_count": len(all_embeddings),
+                    }
+                },
+            )
+
+            return all_embeddings
+
+        except Exception as e:
+            logger.error(
+                f"Failed to generate embeddings for {file_path}: {e}",
+                extra={
+                    "structured_data": {
+                        "project_name": project_name,
+                        "file_path": str(file_path),
+                        "chunk_count": len(chunks),
+                        "error": str(e),
+                    }
+                },
+                exc_info=True,
+            )
+            raise
+
+    def _prepare_chunk_texts(self, chunks: List[CodeChunk]) -> List[str]:
+        """
+        Prepare text content from chunks with context enhancement.
+
+        Args:
+            chunks: List of code chunks
+
+        Returns:
+            List of prepared text strings ready for embedding
+        """
+        texts = []
+        for chunk in chunks:
+            # Include chunk context for better embeddings
+            text_content = chunk.content
+            if chunk.name:
+                # Prefix with chunk name for context
+                text_content = f"# {chunk.name}\n{chunk.content}"
+            texts.append(text_content)
+        return texts
+
+    def _log_batch_stats(
+        self, batch_chunks: List[CodeChunk], batch_num: int, embedding_count: int
+    ) -> None:
+        """Log statistics for a processed batch."""
+        chunk_types = {}
+        redacted_count = 0
+
+        for chunk in batch_chunks:
+            chunk_type = chunk.chunk_type.value
+            chunk_types[chunk_type] = chunk_types.get(chunk_type, 0) + 1
+            if chunk.redacted:
+                redacted_count += 1
+
+        logger.debug(
+            f"Batch {batch_num} complete: "
+            f"{embedding_count} embeddings generated, "
+            f"chunk types: {chunk_types}, "
+            f"redacted: {redacted_count}"
+        )
+
+    async def generate_embeddings_for_multiple_files(
+        self, file_chunks: Dict[str, List[CodeChunk]], project_name: str
+    ) -> Dict[str, List[List[float]]]:
+        """
+        Generate embeddings for chunks from multiple files in a single batch operation.
+
+        Args:
+            file_chunks: Dictionary mapping file paths to their code chunks
+            project_name: Name of the project (for logging)
+
+        Returns:
+            Dictionary mapping file paths to their corresponding embeddings
+
+        Raises:
+            RuntimeError: If embedding generation fails
+        """
+        if not file_chunks:
+            logger.debug("No file chunks provided for batch processing")
+            return {}
+
+        total_chunks = sum(len(chunks) for chunks in file_chunks.values())
+        logger.info(
+            f"Generating batch embeddings for {len(file_chunks)} files "
+            f"({total_chunks} total chunks)",
+            extra={
+                "structured_data": {
+                    "project_name": project_name,
+                    "file_count": len(file_chunks),
+                    "chunk_count": total_chunks,
+                    "embedding_model": self.config.embedding_model,
+                }
+            },
+        )
+
+        try:
+            # Build batches from the start based on actual token counts and text limits
+            batches = await self._build_token_aware_batches(file_chunks, project_name)
+
+            if not batches:
+                logger.debug("No valid chunks found after text preparation")
+                return {}
+
+            # Process each batch
+            logger.info(
+                f"Processing {len(batches)} token-aware batches for project {project_name}"
+            )
+            all_file_embeddings = {}
+
+            for i, (batch_texts, batch_boundaries) in enumerate(batches):
+                logger.debug(
+                    f"Processing batch {i + 1}/{len(batches)}: {len(batch_texts)} texts"
+                )
+
+                # Generate embeddings for this batch
+                batch_file_embeddings = await asyncio.get_event_loop().run_in_executor(
+                    None,
+                    lambda texts=batch_texts, boundaries=batch_boundaries: self.embedding_client.generate_embeddings_batch(
+                        all_texts=texts,
+                        file_boundaries=boundaries,
+                        input_type="document",
+                    ),
+                )
+
+                # Merge results
+                all_file_embeddings.update(batch_file_embeddings)
+
+            file_embeddings = all_file_embeddings
+
+            # Log batch statistics
+            self._log_batch_embedding_stats(file_chunks, file_embeddings)
+
+            logger.info(
+                f"Successfully generated batch embeddings for {len(file_embeddings)} files",
+                extra={
+                    "structured_data": {
+                        "project_name": project_name,
+                        "files_processed": len(file_embeddings),
+                        "total_embeddings": sum(
+                            len(embs) for embs in file_embeddings.values()
+                        ),
+                    }
+                },
+            )
+
+            return file_embeddings
+
+        except Exception as e:
+            logger.error(
+                f"Failed to generate batch embeddings: {e}",
+                extra={
+                    "structured_data": {
+                        "project_name": project_name,
+                        "file_count": len(file_chunks),
+                        "chunk_count": total_chunks,
+                        "error": str(e),
+                    }
+                },
+                exc_info=True,
+            )
+            raise
+
+    async def _build_token_aware_batches(
+        self, file_chunks: Dict[str, List[CodeChunk]], project_name: str
+    ) -> List[Tuple[List[str], List[Tuple[str, int, int]]]]:
+        """
+        Build batches from file chunks respecting both token and text count limits.
+
+        Args:
+            file_chunks: Dictionary mapping file paths to their code chunks
+            project_name: Name of the project (for logging)
+
+        Returns:
+            List of tuples: (batch_texts, batch_file_boundaries)
+        """
+        batches = []  # List of (batch_texts, batch_file_boundaries)
+        current_batch_texts = []
+        current_batch_boundaries = []
+        current_batch_tokens = 0
+        batch_idx = 0
+
+        for file_path, chunks in file_chunks.items():
+            if not chunks:
+                continue
+
+            # Prepare texts for this file
+            file_texts = self._prepare_chunk_texts(chunks)
+            if not file_texts:
+                continue
+
+            # Count tokens for this file using accurate Voyage API
+            file_tokens = await asyncio.get_event_loop().run_in_executor(
+                None, lambda texts=file_texts: self.embedding_client.count_tokens(texts)
+            )
+
+            logger.debug(
+                f"File {file_path}: {len(file_texts)} texts, {file_tokens} tokens"
+            )
+
+            # If adding this file would exceed token limit OR text count limit, finalize current batch
+            if (
+                current_batch_tokens + file_tokens
+                > self.config.voyage_max_tokens_per_batch
+                or len(current_batch_texts) + len(file_texts)
+                > self.config.voyage_batch_size_limit
+            ) and current_batch_texts:
+
+                # Determine which limit was exceeded for logging
+                token_exceeded = (
+                    current_batch_tokens + file_tokens
+                    > self.config.voyage_max_tokens_per_batch
+                )
+                count_exceeded = (
+                    len(current_batch_texts) + len(file_texts)
+                    > self.config.voyage_batch_size_limit
+                )
+
+                logger.info(
+                    f"Finalizing batch {len(batches) + 1}: {len(current_batch_texts)} texts, "
+                    f"{current_batch_tokens} tokens (limit exceeded: "
+                    f"tokens={token_exceeded}, count={count_exceeded})"
+                )
+
+                batches.append((current_batch_texts, current_batch_boundaries))
+
+                # Start new batch
+                current_batch_texts = []
+                current_batch_boundaries = []
+                current_batch_tokens = 0
+                batch_idx = 0
+
+            # Add this file to current batch
+            start_idx = batch_idx
+            end_idx = batch_idx + len(file_texts)
+            current_batch_texts.extend(file_texts)
+            current_batch_boundaries.append((file_path, start_idx, end_idx))
+            current_batch_tokens += file_tokens
+            batch_idx = end_idx
+
+        # Add final batch if it has content
+        if current_batch_texts:
+            logger.info(
+                f"Finalizing final batch {len(batches) + 1}: {len(current_batch_texts)} texts, "
+                f"{current_batch_tokens} tokens"
+            )
+            batches.append((current_batch_texts, current_batch_boundaries))
+
+        return batches
+
+    def _log_batch_embedding_stats(
+        self,
+        file_chunks: Dict[str, List[CodeChunk]],
+        file_embeddings: Dict[str, List[List[float]]],
+    ) -> None:
+        """Log statistics for batch embedding processing."""
+        total_chunks = 0
+        chunk_types = {}
+        redacted_count = 0
+        languages = set()
+
+        for file_path, chunks in file_chunks.items():
+            total_chunks += len(chunks)
+            for chunk in chunks:
+                chunk_type = chunk.chunk_type.value
+                chunk_types[chunk_type] = chunk_types.get(chunk_type, 0) + 1
+                if chunk.redacted:
+                    redacted_count += 1
+                languages.add(chunk.language)
+
+        total_embeddings = sum(len(embs) for embs in file_embeddings.values())
+
+        logger.debug(
+            f"Batch embedding complete: "
+            f"{len(file_chunks)} files, "
+            f"{total_chunks} chunks, "
+            f"{total_embeddings} embeddings generated, "
+            f"chunk types: {chunk_types}, "
+            f"languages: {sorted(languages)}, "
+            f"redacted: {redacted_count}"
+        )