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

coding_assistant/context/ranker.py

@@ -0,0 +1,355 @@
+"""Context ranking with language-aware scoring."""
+
+from typing import List, Dict, Optional
+from datetime import datetime, timedelta
+from pathlib import Path
+
+
+class ContextRanker:
+    """
+    Rank retrieved code chunks based on multiple factors.
+
+    Scoring factors:
+    - Semantic similarity (40%)
+    - File proximity (20%)
+    - Dependency distance (20%)
+    - Recency (10%)
+    - Code popularity/references (10%)
+    """
+
+    def __init__(self, dependency_graph=None):
+        """
+        Initialize the context ranker.
+
+        Args:
+            dependency_graph: Optional dependency graph for file relationships
+        """
+        self.dependency_graph = dependency_graph
+
+        # Scoring weights
+        self.weights = {
+            'similarity': 0.4,
+            'file_proximity': 0.2,
+            'dependency': 0.2,
+            'recency': 0.1,
+            'popularity': 0.1
+        }
+
+    def rank(self, chunks: List[Dict], query: str,
+             current_file: Optional[str] = None,
+             language: Optional[str] = None) -> List[Dict]:
+        """
+        Rank chunks based on relevance to query and context.
+
+        Args:
+            chunks: List of chunks with metadata (must have 'id', 'similarity', etc.)
+            query: The user's query
+            current_file: Current file being edited (for proximity bonus)
+            language: Programming language (for language-specific heuristics)
+
+        Returns:
+            Ranked list of chunks with scores
+        """
+        scored_chunks = []
+
+        for chunk in chunks:
+            score = self._calculate_score(chunk, query, current_file, language)
+            chunk['rank_score'] = score
+            scored_chunks.append(chunk)
+
+        # Sort by rank score (descending)
+        ranked = sorted(scored_chunks, key=lambda x: x['rank_score'], reverse=True)
+
+        return ranked
+
+    def _calculate_score(self, chunk: Dict, query: str,
+                        current_file: Optional[str],
+                        language: Optional[str]) -> float:
+        """Calculate overall relevance score for a chunk."""
+        score = 0.0
+
+        # 1. Semantic similarity (from vector/hybrid search)
+        similarity = chunk.get('similarity', chunk.get('vector_score', 0.0))
+        score += similarity * self.weights['similarity']
+
+        # 2. File proximity
+        if current_file and 'file_path' in chunk:
+            proximity_score = self._file_proximity_score(chunk['file_path'], current_file)
+            score += proximity_score * self.weights['file_proximity']
+
+        # 3. Dependency distance
+        if current_file and self.dependency_graph and 'file_path' in chunk:
+            dep_score = self._dependency_score(chunk['file_path'], current_file)
+            score += dep_score * self.weights['dependency']
+
+        # 4. Recency
+        if 'last_modified' in chunk:
+            recency_score = self._recency_score(chunk['last_modified'])
+            score += recency_score * self.weights['recency']
+
+        # 5. Popularity
+        if 'reference_count' in chunk:
+            popularity_score = self._popularity_score(chunk['reference_count'])
+            score += popularity_score * self.weights['popularity']
+
+        # 6. Language-specific boosts
+        if language:
+            lang_boost = self._language_specific_boost(chunk, query, language)
+            score += lang_boost
+
+        return score
+
+    def _file_proximity_score(self, file_path: str, current_file: str) -> float:
+        """
+        Score based on how close two files are in the directory structure.
+
+        Returns:
+            Score between 0 and 1
+        """
+        if file_path == current_file:
+            return 1.0
+
+        # Convert to Path objects
+        path1 = Path(file_path)
+        path2 = Path(current_file)
+
+        # Same directory: high score
+        if path1.parent == path2.parent:
+            return 0.8
+
+        # Check how many directory levels apart
+        try:
+            # Get relative path
+            rel_path = path1.relative_to(path2.parent)
+            levels = len(rel_path.parts) - 1
+            # Closer = higher score
+            return max(0.0, 1.0 - (levels * 0.2))
+        except ValueError:
+            # Not in same tree, check common parent
+            common_parts = 0
+            for p1, p2 in zip(path1.parts, path2.parts):
+                if p1 == p2:
+                    common_parts += 1
+                else:
+                    break
+
+            # More common directories = higher score
+            return min(1.0, common_parts * 0.15)
+
+    def _dependency_score(self, file_path: str, current_file: str) -> float:
+        """
+        Score based on dependency distance in import graph.
+
+        Returns:
+            Score between 0 and 1
+        """
+        if not self.dependency_graph:
+            return 0.0
+
+        try:
+            # Get shortest path distance in dependency graph
+            distance = self.dependency_graph.get_distance(current_file, file_path)
+
+            if distance == 0:
+                return 1.0
+            elif distance == 1:
+                # Direct dependency
+                return 0.8
+            elif distance == 2:
+                # Second-degree dependency
+                return 0.5
+            else:
+                # Further away
+                return max(0.0, 1.0 - (distance * 0.2))
+        except:
+            return 0.0
+
+    def _recency_score(self, last_modified: datetime) -> float:
+        """
+        Score based on how recently the file was modified.
+
+        Returns:
+            Score between 0 and 1
+        """
+        if isinstance(last_modified, str):
+            try:
+                last_modified = datetime.fromisoformat(last_modified)
+            except:
+                return 0.5  # Default if can't parse
+
+        now = datetime.now()
+        age = now - last_modified
+
+        # Files modified in last day: 1.0
+        if age < timedelta(days=1):
+            return 1.0
+        # Last week: 0.8
+        elif age < timedelta(weeks=1):
+            return 0.8
+        # Last month: 0.5
+        elif age < timedelta(days=30):
+            return 0.5
+        # Last 3 months: 0.3
+        elif age < timedelta(days=90):
+            return 0.3
+        # Older: decay linearly
+        else:
+            days_old = age.days
+            return max(0.0, 1.0 - (days_old / 365))
+
+    def _popularity_score(self, reference_count: int) -> float:
+        """
+        Score based on how many times the code is referenced.
+
+        Returns:
+            Score between 0 and 1
+        """
+        # Logarithmic scaling
+        # 0 refs: 0.0, 1 ref: 0.2, 10 refs: 0.5, 100 refs: 0.7, 1000 refs: 1.0
+        if reference_count == 0:
+            return 0.0
+
+        import math
+        return min(1.0, math.log10(reference_count + 1) / 3)
+
+    def _language_specific_boost(self, chunk: Dict, query: str, language: str) -> float:
+        """
+        Apply language-specific heuristics for ranking.
+
+        Returns:
+            Boost score (can be positive or negative, typically 0.0-0.2)
+        """
+        boost = 0.0
+        file_path = chunk.get('file_path', '').lower()
+        query_lower = query.lower()
+
+        if language == 'python':
+            # Boost test files if query mentions "test"
+            if 'test' in query_lower and 'test_' in file_path:
+                boost += 0.15
+
+            # Boost __init__.py for module/package questions
+            if '__init__.py' in file_path:
+                if any(word in query_lower for word in ['module', 'package', 'import']):
+                    boost += 0.1
+
+            # Boost setup.py, pyproject.toml for dependency questions
+            if 'setup.py' in file_path or 'pyproject.toml' in file_path:
+                if any(word in query_lower for word in ['dependency', 'install', 'package']):
+                    boost += 0.15
+
+            # Boost main.py, app.py for entry point questions
+            if 'main.py' in file_path or 'app.py' in file_path:
+                if any(word in query_lower for word in ['start', 'entry', 'run', 'main']):
+                    boost += 0.1
+
+        elif language in ('javascript', 'typescript'):
+            # Boost index.js, main.js, app.js for entry points
+            if any(name in file_path for name in ['index.js', 'main.js', 'app.js', 'index.ts', 'main.ts']):
+                if any(word in query_lower for word in ['start', 'entry', 'run', 'main']):
+                    boost += 0.1
+
+            # Boost .tsx/.jsx for component questions
+            if file_path.endswith(('.tsx', '.jsx')):
+                if any(word in query_lower for word in ['component', 'render', 'ui', 'view']):
+                    boost += 0.15
+
+            # Boost package.json for dependency questions
+            if 'package.json' in file_path:
+                if any(word in query_lower for word in ['dependency', 'install', 'package', 'npm']):
+                    boost += 0.15
+
+            # Boost test files
+            if any(pattern in file_path for pattern in ['.test.', '.spec.', '__tests__']):
+                if 'test' in query_lower:
+                    boost += 0.15
+
+            # Boost config files for configuration questions
+            if any(pattern in file_path for pattern in ['config.', '.config.', 'webpack', 'vite']):
+                if any(word in query_lower for word in ['config', 'setup', 'build']):
+                    boost += 0.1
+
+        # Generic boosts (language-agnostic)
+
+        # Boost README files for overview questions
+        if 'readme' in file_path:
+            if any(word in query_lower for word in ['what', 'overview', 'about', 'intro']):
+                boost += 0.15
+
+        # Boost documentation files
+        if any(pattern in file_path for pattern in ['docs/', 'documentation/']):
+            if any(word in query_lower for word in ['how', 'guide', 'tutorial', 'example']):
+                boost += 0.1
+
+        # Penalize non-code files for code-specific questions
+        if any(word in query_lower for word in ['function', 'class', 'implement', 'code']):
+            if file_path.endswith(('.md', '.txt', '.json', '.yaml', '.yml')):
+                boost -= 0.1
+
+        return boost
+
+    def update_weights(self, **kwargs):
+        """
+        Update scoring weights.
+
+        Args:
+            **kwargs: New weights (similarity, file_proximity, dependency, recency, popularity)
+        """
+        for key, value in kwargs.items():
+            if key in self.weights:
+                self.weights[key] = value
+
+        # Normalize weights to sum to 1.0
+        total = sum(self.weights.values())
+        if total > 0:
+            for key in self.weights:
+                self.weights[key] /= total
+
+    def get_weights(self) -> Dict[str, float]:
+        """Get current scoring weights."""
+        return self.weights.copy()
+
+    def explain_ranking(self, chunk: Dict, query: str,
+                       current_file: Optional[str] = None,
+                       language: Optional[str] = None) -> Dict:
+        """
+        Explain the ranking score for a specific chunk.
+
+        Returns:
+            Dict with breakdown of score components
+        """
+        similarity = chunk.get('similarity', chunk.get('vector_score', 0.0))
+
+        explanation = {
+            'total_score': chunk.get('rank_score', 0.0),
+            'components': {
+                'similarity': similarity * self.weights['similarity'],
+                'file_proximity': 0.0,
+                'dependency': 0.0,
+                'recency': 0.0,
+                'popularity': 0.0,
+                'language_boost': 0.0
+            }
+        }
+
+        if current_file and 'file_path' in chunk:
+            prox = self._file_proximity_score(chunk['file_path'], current_file)
+            explanation['components']['file_proximity'] = prox * self.weights['file_proximity']
+
+        if current_file and self.dependency_graph and 'file_path' in chunk:
+            dep = self._dependency_score(chunk['file_path'], current_file)
+            explanation['components']['dependency'] = dep * self.weights['dependency']
+
+        if 'last_modified' in chunk:
+            rec = self._recency_score(chunk['last_modified'])
+            explanation['components']['recency'] = rec * self.weights['recency']
+
+        if 'reference_count' in chunk:
+            pop = self._popularity_score(chunk['reference_count'])
+            explanation['components']['popularity'] = pop * self.weights['popularity']
+
+        if language:
+            boost = self._language_specific_boost(chunk, query, language)
+            explanation['components']['language_boost'] = boost
+
+        return explanation

coding_assistant/context/retriever.py

@@ -0,0 +1,119 @@
+"""Semantic retriever for code search using embeddings."""
+from pathlib import Path
+from typing import List, Dict, Optional
+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
+
+
+class SemanticRetriever:
+    """Retrieve relevant code using semantic search."""
+
+    def __init__(self, project_path: Path):
+        """
+        Initialize the semantic retriever.
+
+        Args:
+            project_path: Path to the project root
+        """
+        self.project_path = Path(project_path).resolve()
+        self.crawler = CodebaseCrawler(self.project_path)
+        self.parser = CodeParser()
+        self.embedder = EmbeddingGenerator()
+        self.store = VectorStore(persist_dir=self.project_path / ".coding_assistant" / "chroma_db")
+
+    def clear_index(self):
+        """Clear the existing index."""
+        self.store.clear()
+
+    def index_codebase(self, max_files: int = 100):
+        """
+        Index the codebase for semantic search.
+
+        Args:
+            max_files: Maximum number of files to index
+        """
+        # Scan files
+        files = self.crawler.scan(max_files=max_files)
+
+        all_chunks = []
+
+        # Parse each file
+        for file_info in files:
+            try:
+                content = self.crawler.read_file(file_info['path'])
+                
+                # Only parse Python files for now (parser supports Python)
+                if file_info['extension'] == '.py':
+                    parsed = self.parser.parse_file(file_info['path'], content)
+                    all_chunks.extend(parsed['chunks'])
+                else:
+                    # For non-Python files, create a simple file-level chunk
+                    all_chunks.append({
+                        'type': 'file',
+                        'file_path': file_info['path'],
+                        'content': content[:5000],  # Limit content size
+                        'start_line': 0,
+                        'end_line': len(content.split('\n'))
+                    })
+            except Exception as e:
+                # Skip files that can't be parsed
+                if hasattr(self, '_verbose') and self._verbose:
+                    print(f"Warning: Could not parse {file_info['path']}: {e}")
+                continue
+
+        if not all_chunks:
+            print("No chunks to index")
+            return
+
+        # Generate embeddings
+        embedded_chunks = self.embedder.embed_code_chunks(all_chunks)
+
+        # Store in vector database
+        self.store.add_chunks(embedded_chunks)
+
+    def get_stats(self) -> Dict:
+        """
+        Get statistics about the indexed codebase.
+
+        Returns:
+            Dictionary with stats including total_chunks and embedding_dimension
+        """
+        return {
+            'total_chunks': self.store.count(),
+            'embedding_dimension': self.embedder.dimension
+        }
+
+    def retrieve(self, query: str, k: int = 5) -> List[Dict]:
+        """
+        Retrieve relevant code chunks for a query.
+
+        Args:
+            query: User query/question
+            k: Number of results to return
+
+        Returns:
+            List of relevant chunks with metadata
+        """
+        # Generate query embedding
+        query_embedding = self.embedder.generate_embedding(query)
+
+        # Search vector store
+        results = self.store.search(query_embedding, n_results=k)
+
+        # Format results to match expected structure
+        formatted_results = []
+        for result in results:
+            metadata = result['metadata']
+            formatted_results.append({
+                'path': metadata['file_path'],
+                'type': metadata['type'],
+                'similarity': result['similarity'],
+                'start_line': metadata['start_line'],
+                'end_line': metadata['end_line'],
+                'content': result['content']
+            })
+
+        return formatted_results
+