ai-coding-assistant 0.5.0__py3-none-any.whl

coding_assistant/context/enhanced_retriever.py

@@ -0,0 +1,322 @@
+"""Enhanced semantic retriever with hybrid search and context ranking."""
+
+from pathlib import Path
+from typing import List, Dict, Optional
+from datetime import datetime
+
+from coding_assistant.codebase.crawler import CodebaseCrawler
+from coding_assistant.codebase.parser import CodeParser
+from coding_assistant.embeddings.generator import EmbeddingGenerator
+from coding_assistant.embeddings.store import VectorStore
+from coding_assistant.context.chunker import SmartChunker
+from coding_assistant.context.hybrid_search import HybridSearch
+from coding_assistant.context.ranker import ContextRanker
+from coding_assistant.vcs.git import GitIntegration
+
+
+class EnhancedSemanticRetriever:
+    """
+    Enhanced retriever with hybrid search and intelligent ranking.
+
+    Combines:
+    - Vector similarity search
+    - BM25 keyword search
+    - Smart code chunking
+    - Context-aware ranking
+    """
+
+    def __init__(self, project_path: Path, dependency_graph=None):
+        """
+        Initialize the enhanced retriever.
+
+        Args:
+            project_path: Path to the project root
+            dependency_graph: Optional dependency graph for file relationships
+        """
+        self.project_path = Path(project_path).resolve()
+        self.crawler = CodebaseCrawler(self.project_path)
+        self.parser = CodeParser()
+        self.embedder = EmbeddingGenerator()
+
+        # Vector store
+        persist_dir = self.project_path / ".coding_assistant" / "chroma_db"
+        self.store = VectorStore(persist_dir=persist_dir)
+
+        # Smart chunker
+        self.chunker = SmartChunker(max_chunk_tokens=500)
+
+        # Hybrid search (will be initialized after indexing)
+        self.hybrid_search = None
+
+        # Context ranker
+        self.ranker = ContextRanker(dependency_graph=dependency_graph)
+
+        # Git integration (optional - won't fail if not a git repo)
+        try:
+            self.git = GitIntegration(str(self.project_path))
+        except Exception:
+            self.git = None
+
+        # Cache for indexed chunks
+        self._indexed_chunks = []
+
+    def clear_index(self):
+        """Clear the existing index."""
+        self.store.clear()
+        self._indexed_chunks = []
+        self.hybrid_search = None
+
+    def index_codebase(self, max_files: int = 100, force_reindex: bool = False):
+        """
+        Index the codebase with smart chunking.
+
+        Args:
+            max_files: Maximum number of files to index
+            force_reindex: Force re-indexing even if already indexed
+        """
+        if self._indexed_chunks and not force_reindex:
+            print("Codebase already indexed. Use force_reindex=True to re-index.")
+            return
+
+        print("Indexing codebase with smart chunking...")
+
+        # Scan files
+        files = self.crawler.scan(max_files=max_files)
+
+        all_chunks = []
+
+        # Parse and chunk each file
+        for file_info in files:
+            try:
+                content = self.crawler.read_file(file_info['path'])
+                language = self._detect_language(file_info['extension'])
+
+                if language in ('python', 'javascript', 'typescript', 'jsx', 'tsx'):
+                    # Use smart chunking for supported languages
+                    chunks = self.chunker.chunk_code(
+                        content,
+                        file_info['path'],
+                        language
+                    )
+
+                    # Convert to dict format
+                    for chunk in chunks:
+                        all_chunks.append({
+                            'type': chunk.type,
+                            'file_path': chunk.file_path,
+                            'content': chunk.content,
+                            'start_line': chunk.start_line,
+                            'end_line': chunk.end_line,
+                            'language': chunk.language,
+                            'name': chunk.name,
+                            'docstring': chunk.docstring,
+                            'last_modified': datetime.fromtimestamp(
+                                Path(chunk.file_path).stat().st_mtime
+                            ).isoformat()
+                        })
+                else:
+                    # For other files, create simple file-level chunk
+                    all_chunks.append({
+                        'type': 'file',
+                        'file_path': file_info['path'],
+                        'content': content[:5000],
+                        'start_line': 0,
+                        'end_line': len(content.split('\n')),
+                        'language': language,
+                        'name': Path(file_info['path']).name,
+                        'last_modified': datetime.fromtimestamp(
+                            Path(file_info['path']).stat().st_mtime
+                        ).isoformat()
+                    })
+
+            except Exception as e:
+                print(f"Warning: Could not process {file_info['path']}: {e}")
+                continue
+
+        if not all_chunks:
+            print("No chunks to index")
+            return
+
+        print(f"Generating embeddings for {len(all_chunks)} chunks...")
+
+        # Generate embeddings
+        embedded_chunks = self.embedder.embed_code_chunks(all_chunks)
+
+        # Store in vector database
+        print("Storing in vector database...")
+        self.store.add_chunks(embedded_chunks)
+
+        # Cache chunks for BM25
+        self._indexed_chunks = embedded_chunks
+
+        # Initialize hybrid search
+        print("Building BM25 index...")
+        self.hybrid_search = HybridSearch(self.store, self.embedder)
+        self.hybrid_search.index_corpus(embedded_chunks)
+
+        print(f"✓ Indexed {len(all_chunks)} code chunks")
+
+    def retrieve(self, query: str, k: int = 5,
+                current_file: Optional[str] = None,
+                language: Optional[str] = None,
+                use_hybrid: bool = True,
+                use_ranking: bool = True) -> List[Dict]:
+        """
+        Retrieve relevant code chunks with hybrid search and ranking.
+
+        Args:
+            query: User query/question
+            k: Number of results to return
+            current_file: Current file for proximity scoring
+            language: Programming language for language-aware ranking
+            use_hybrid: Whether to use hybrid search (vector + BM25)
+            use_ranking: Whether to apply context ranking
+
+        Returns:
+            List of relevant chunks with metadata
+        """
+        if use_hybrid and self.hybrid_search:
+            # Hybrid search (vector + BM25)
+            results = self.hybrid_search.search(query, n_results=k * 2)
+
+            # Convert hybrid results to full chunk data
+            enriched_results = []
+            for result in results:
+                chunk_id = result['id']
+                # Find full chunk data
+                for chunk in self._indexed_chunks:
+                    if chunk.get('id') == chunk_id:
+                        chunk_copy = chunk.copy()
+                        chunk_copy['similarity'] = result.get('vector_score', 0.0)
+                        chunk_copy['bm25_score'] = result.get('bm25_score', 0.0)
+                        chunk_copy['rrf_score'] = result.get('rrf_score', 0.0)
+                        enriched_results.append(chunk_copy)
+                        break
+        else:
+            # Fallback to vector-only search
+            query_embedding = self.embedder.generate_embedding(query)
+            results = self.store.search(query_embedding, n_results=k * 2)
+            enriched_results = results
+
+        if not enriched_results:
+            return []
+
+        # Apply context ranking
+        if use_ranking:
+            ranked_results = self.ranker.rank(
+                enriched_results,
+                query,
+                current_file=current_file,
+                language=language
+            )
+        else:
+            ranked_results = enriched_results
+
+        # Format results
+        formatted_results = []
+        for result in ranked_results[:k]:
+            metadata = result.get('metadata', {})
+            formatted_results.append({
+                'path': result.get('file_path', metadata.get('file_path', 'unknown')),
+                'type': result.get('type', metadata.get('type', 'file')),
+                'name': result.get('name', metadata.get('name', '')),
+                'similarity': result.get('similarity', 0.0),
+                'rank_score': result.get('rank_score', result.get('similarity', 0.0)),
+                'start_line': result.get('start_line', metadata.get('start_line', 0)),
+                'end_line': result.get('end_line', metadata.get('end_line', 0)),
+                'content': result.get('content', ''),
+                'language': result.get('language', metadata.get('language', 'unknown')),
+                'bm25_score': result.get('bm25_score', 0.0),
+                'rrf_score': result.get('rrf_score', 0.0)
+            })
+
+        return formatted_results
+
+    def get_stats(self) -> Dict:
+        """
+        Get statistics about the indexed codebase.
+
+        Returns:
+            Dictionary with stats
+        """
+        stats = {
+            'total_chunks': self.store.count(),
+            'embedding_dimension': self.embedder.dimension,
+            'indexed_chunks': len(self._indexed_chunks),
+            'hybrid_search_enabled': self.hybrid_search is not None
+        }
+
+        if self.hybrid_search:
+            stats.update(self.hybrid_search.get_stats())
+
+        return stats
+
+    def explain_retrieval(self, query: str, chunk_id: str) -> Dict:
+        """
+        Explain why a specific chunk was retrieved for a query.
+
+        Args:
+            query: The search query
+            chunk_id: ID of the chunk to explain
+
+        Returns:
+            Dict with retrieval explanation
+        """
+        explanation = {}
+
+        if self.hybrid_search:
+            explanation['hybrid_search'] = self.hybrid_search.explain_ranking(query, chunk_id)
+
+        # Find the chunk
+        for chunk in self._indexed_chunks:
+            if chunk.get('id') == chunk_id:
+                if self.ranker:
+                    explanation['ranking'] = self.ranker.explain_ranking(chunk, query)
+                break
+
+        return explanation
+
+    def _detect_language(self, extension: str) -> str:
+        """Detect language from file extension."""
+        extension = extension.lower()
+        language_map = {
+            '.py': 'python',
+            '.js': 'javascript',
+            '.jsx': 'jsx',
+            '.ts': 'typescript',
+            '.tsx': 'tsx',
+            '.mjs': 'javascript',
+            '.cjs': 'javascript',
+        }
+        return language_map.get(extension, 'text')
+
+    def update_chunk_metadata(self, file_path: str, **metadata):
+        """
+        Update metadata for chunks from a specific file.
+
+        Useful for updating reference counts, popularity, etc.
+        """
+        for chunk in self._indexed_chunks:
+            if chunk.get('file_path') == file_path:
+                chunk.update(metadata)
+
+    def configure_ranker(self, **weights):
+        """
+        Configure context ranker weights.
+
+        Args:
+            **weights: New weights (similarity, file_proximity, etc.)
+        """
+        self.ranker.update_weights(**weights)
+
+    def configure_hybrid_search(self, vector_weight: float = 0.5,
+                                keyword_weight: float = 0.5):
+        """
+        Configure hybrid search weights.
+
+        Args:
+            vector_weight: Weight for vector search
+            keyword_weight: Weight for keyword search
+        """
+        if self.hybrid_search:
+            self.hybrid_search.update_weights(vector_weight, keyword_weight)

coding_assistant/context/hybrid_search.py

@@ -0,0 +1,311 @@
+"""Hybrid search combining vector similarity and keyword search (BM25)."""
+
+from typing import List, Dict, Optional, Tuple
+from rank_bm25 import BM25Okapi
+import re
+
+
+class HybridSearch:
+    """
+    Hybrid search that combines vector similarity with BM25 keyword search.
+
+    Uses Reciprocal Rank Fusion (RRF) to combine results from both methods.
+    """
+
+    RRF_CONSTANT = 60  # Standard RRF constant
+
+    def __init__(self, vector_store, embedding_generator):
+        """
+        Initialize hybrid search.
+
+        Args:
+            vector_store: VectorStore instance for semantic search
+            embedding_generator: EmbeddingGenerator for query embeddings
+        """
+        self.vector_store = vector_store
+        self.embedding_generator = embedding_generator
+
+        # BM25 index (will be built from corpus)
+        self.bm25 = None
+        self.corpus = []
+        self.chunk_ids = []
+
+    def index_corpus(self, chunks: List[Dict]):
+        """
+        Build BM25 index from code chunks.
+
+        Args:
+            chunks: List of code chunks with 'id' and 'content' keys
+        """
+        self.corpus = []
+        self.chunk_ids = []
+
+        for chunk in chunks:
+            content = chunk.get('content', '')
+            self.corpus.append(content)
+            self.chunk_ids.append(chunk.get('id'))
+
+        # Tokenize corpus for BM25
+        tokenized_corpus = [self._tokenize(doc) for doc in self.corpus]
+
+        # Build BM25 index
+        self.bm25 = BM25Okapi(tokenized_corpus)
+
+    def search(self, query: str, n_results: int = 10,
+               vector_weight: float = 0.5,
+               keyword_weight: float = 0.5) -> List[Dict]:
+        """
+        Hybrid search combining vector and keyword search.
+
+        Args:
+            query: Search query
+            n_results: Number of results to return
+            vector_weight: Weight for vector search (0-1)
+            keyword_weight: Weight for keyword search (0-1)
+
+        Returns:
+            List of ranked results
+        """
+        if not self.bm25 or not self.corpus:
+            # Fallback to vector-only search
+            return self._vector_only_search(query, n_results)
+
+        # 1. Vector similarity search
+        vector_results = self._vector_search(query, n_results * 2)
+
+        # 2. BM25 keyword search
+        bm25_results = self._bm25_search(query, n_results * 2)
+
+        # 3. Combine using Reciprocal Rank Fusion
+        combined = self._reciprocal_rank_fusion(
+            vector_results,
+            bm25_results,
+            vector_weight,
+            keyword_weight
+        )
+
+        # Return top N results
+        return combined[:n_results]
+
+    def _vector_search(self, query: str, k: int) -> List[Tuple[str, float]]:
+        """
+        Perform vector similarity search.
+
+        Returns:
+            List of (chunk_id, score) tuples
+        """
+        # Generate query embedding
+        query_embedding = self.embedding_generator.generate_embedding(query)
+
+        # Search vector store
+        results = self.vector_store.search(query_embedding, n_results=k)
+
+        # Extract chunk IDs and scores (similarity)
+        return [(r['id'], r['similarity']) for r in results]
+
+    def _bm25_search(self, query: str, k: int) -> List[Tuple[str, float]]:
+        """
+        Perform BM25 keyword search.
+
+        Returns:
+            List of (chunk_id, score) tuples
+        """
+        if not self.bm25:
+            return []
+
+        # Tokenize query
+        tokenized_query = self._tokenize(query)
+
+        # Get BM25 scores
+        scores = self.bm25.get_scores(tokenized_query)
+
+        # Sort by score and get top K
+        scored_docs = [(self.chunk_ids[i], scores[i]) for i in range(len(scores))]
+        scored_docs.sort(key=lambda x: x[1], reverse=True)
+
+        return scored_docs[:k]
+
+    def _reciprocal_rank_fusion(
+        self,
+        vector_results: List[Tuple[str, float]],
+        bm25_results: List[Tuple[str, float]],
+        vector_weight: float = 0.5,
+        keyword_weight: float = 0.5
+    ) -> List[Dict]:
+        """
+        Combine results using Reciprocal Rank Fusion.
+
+        RRF formula: score = Σ weight / (rank + constant)
+
+        Args:
+            vector_results: Results from vector search
+            bm25_results: Results from BM25 search
+            vector_weight: Weight for vector results
+            keyword_weight: Weight for keyword results
+
+        Returns:
+            Combined and ranked results
+        """
+        rrf_scores = {}
+
+        # Add scores from vector search
+        for rank, (chunk_id, score) in enumerate(vector_results):
+            rrf_score = vector_weight / (rank + self.RRF_CONSTANT)
+            if chunk_id in rrf_scores:
+                rrf_scores[chunk_id]['rrf_score'] += rrf_score
+                rrf_scores[chunk_id]['vector_score'] = score
+            else:
+                rrf_scores[chunk_id] = {
+                    'id': chunk_id,
+                    'rrf_score': rrf_score,
+                    'vector_score': score,
+                    'bm25_score': 0.0,
+                    'vector_rank': rank + 1,
+                    'bm25_rank': None
+                }
+
+        # Add scores from BM25 search
+        for rank, (chunk_id, score) in enumerate(bm25_results):
+            rrf_score = keyword_weight / (rank + self.RRF_CONSTANT)
+            if chunk_id in rrf_scores:
+                rrf_scores[chunk_id]['rrf_score'] += rrf_score
+                rrf_scores[chunk_id]['bm25_score'] = score
+                rrf_scores[chunk_id]['bm25_rank'] = rank + 1
+            else:
+                rrf_scores[chunk_id] = {
+                    'id': chunk_id,
+                    'rrf_score': rrf_score,
+                    'vector_score': 0.0,
+                    'bm25_score': score,
+                    'vector_rank': None,
+                    'bm25_rank': rank + 1
+                }
+
+        # Sort by RRF score
+        ranked = sorted(
+            rrf_scores.values(),
+            key=lambda x: x['rrf_score'],
+            reverse=True
+        )
+
+        return ranked
+
+    def _vector_only_search(self, query: str, k: int) -> List[Dict]:
+        """Fallback to vector-only search."""
+        query_embedding = self.embedding_generator.generate_embedding(query)
+        results = self.vector_store.search(query_embedding, n_results=k)
+
+        return [{
+            'id': r['id'],
+            'rrf_score': r['similarity'],
+            'vector_score': r['similarity'],
+            'bm25_score': 0.0,
+            'vector_rank': i + 1,
+            'bm25_rank': None
+        } for i, r in enumerate(results)]
+
+    def _tokenize(self, text: str) -> List[str]:
+        """
+        Tokenize text for BM25 search.
+
+        Uses code-aware tokenization that preserves identifiers and keywords.
+        """
+        # Convert to lowercase
+        text = text.lower()
+
+        # Split on non-alphanumeric characters but keep underscores
+        # This preserves Python/JS identifiers like my_function_name
+        tokens = re.findall(r'\b\w+\b', text)
+
+        # Also split camelCase into separate tokens
+        expanded_tokens = []
+        for token in tokens:
+            # Split camelCase: myFunctionName -> my, function, name
+            camel_split = re.sub('([A-Z][a-z]+)', r' \1', token)
+            camel_split = re.sub('([A-Z]+)', r' \1', camel_split)
+            expanded_tokens.extend(camel_split.split())
+
+        # Remove very short tokens (< 2 chars) and common code words
+        stop_words = {'if', 'for', 'while', 'do', 'is', 'to', 'in', 'of', 'and', 'or', 'not'}
+        filtered = [t for t in expanded_tokens if len(t) >= 2 and t not in stop_words]
+
+        return filtered
+
+    def get_stats(self) -> Dict:
+        """Get statistics about the search index."""
+        return {
+            'corpus_size': len(self.corpus),
+            'bm25_indexed': self.bm25 is not None,
+            'vector_store_count': self.vector_store.count() if self.vector_store else 0
+        }
+
+    def update_weights(self, vector_weight: float, keyword_weight: float):
+        """
+        Update the weights for vector vs keyword search.
+
+        Args:
+            vector_weight: Weight for vector search (0-1)
+            keyword_weight: Weight for keyword search (0-1)
+        """
+        # Normalize weights to sum to 1
+        total = vector_weight + keyword_weight
+        if total > 0:
+            self.vector_weight = vector_weight / total
+            self.keyword_weight = keyword_weight / total
+        else:
+            self.vector_weight = 0.5
+            self.keyword_weight = 0.5
+
+    def explain_ranking(self, query: str, chunk_id: str) -> Dict:
+        """
+        Explain why a specific chunk was ranked for a query.
+
+        Args:
+            query: The search query
+            chunk_id: ID of the chunk to explain
+
+        Returns:
+            Dict with ranking explanation
+        """
+        # Get vector score
+        query_embedding = self.embedding_generator.generate_embedding(query)
+        vector_results = self.vector_store.search(query_embedding, n_results=100)
+        vector_rank = None
+        vector_score = 0.0
+
+        for i, result in enumerate(vector_results):
+            if result['id'] == chunk_id:
+                vector_rank = i + 1
+                vector_score = result['similarity']
+                break
+
+        # Get BM25 score
+        tokenized_query = self._tokenize(query)
+        bm25_scores = self.bm25.get_scores(tokenized_query) if self.bm25 else []
+
+        bm25_rank = None
+        bm25_score = 0.0
+
+        try:
+            idx = self.chunk_ids.index(chunk_id)
+            bm25_score = bm25_scores[idx]
+
+            # Find rank
+            sorted_scores = sorted(enumerate(bm25_scores), key=lambda x: x[1], reverse=True)
+            for rank, (i, score) in enumerate(sorted_scores):
+                if i == idx:
+                    bm25_rank = rank + 1
+                    break
+        except (ValueError, IndexError):
+            pass
+
+        return {
+            'chunk_id': chunk_id,
+            'query': query,
+            'vector_rank': vector_rank,
+            'vector_score': vector_score,
+            'bm25_rank': bm25_rank,
+            'bm25_score': bm25_score,
+            'appears_in_vector': vector_rank is not None,
+            'appears_in_bm25': bm25_rank is not None
+        }