cognify-code 0.2.0__py3-none-any.whl

ai_code_assistant/retrieval/chunker.py

@@ -0,0 +1,215 @@
+"""
+Code chunking strategies for semantic indexing.
+
+Chunks code into meaningful segments for better retrieval.
+"""
+
+import re
+from dataclasses import dataclass, field
+from typing import List, Optional
+from pathlib import Path
+
+
+@dataclass
+class CodeChunk:
+    """Represents a chunk of code with metadata."""
+    content: str
+    file_path: str
+    start_line: int
+    end_line: int
+    chunk_type: str  # function, class, module, block
+    name: Optional[str] = None
+    language: Optional[str] = None
+    
+    @property
+    def id(self) -> str:
+        """Generate unique ID for this chunk."""
+        return f"{self.file_path}:{self.start_line}-{self.end_line}"
+    
+    def to_dict(self) -> dict:
+        """Convert to dictionary for storage."""
+        return {
+            "content": self.content,
+            "file_path": self.file_path,
+            "start_line": self.start_line,
+            "end_line": self.end_line,
+            "chunk_type": self.chunk_type,
+            "name": self.name or "",
+            "language": self.language or "",
+        }
+
+
+@dataclass
+class CodeChunker:
+    """Splits code files into semantic chunks."""
+    
+    chunk_size: int = 50  # Max lines per chunk
+    chunk_overlap: int = 10  # Overlap between chunks
+    
+    # Language detection by extension
+    LANGUAGE_MAP: dict = field(default_factory=lambda: {
+        ".py": "python",
+        ".js": "javascript",
+        ".ts": "typescript",
+        ".jsx": "javascript",
+        ".tsx": "typescript",
+        ".java": "java",
+        ".go": "go",
+        ".rs": "rust",
+        ".rb": "ruby",
+        ".php": "php",
+        ".c": "c",
+        ".cpp": "cpp",
+        ".h": "c",
+        ".hpp": "cpp",
+        ".cs": "csharp",
+        ".swift": "swift",
+        ".kt": "kotlin",
+        ".scala": "scala",
+        ".sh": "bash",
+        ".yaml": "yaml",
+        ".yml": "yaml",
+        ".json": "json",
+        ".md": "markdown",
+        ".sql": "sql",
+    })
+    
+    # Patterns for detecting code boundaries
+    PYTHON_PATTERNS = {
+        "class": r"^class\s+(\w+)",
+        "function": r"^(?:async\s+)?def\s+(\w+)",
+    }
+    
+    def detect_language(self, file_path: str) -> Optional[str]:
+        """Detect programming language from file extension."""
+        ext = Path(file_path).suffix.lower()
+        return self.LANGUAGE_MAP.get(ext)
+    
+    def chunk_file(self, file_path: str, content: str) -> List[CodeChunk]:
+        """Split a file into semantic chunks."""
+        language = self.detect_language(file_path)
+        lines = content.split("\n")
+        
+        if language == "python":
+            return self._chunk_python(file_path, lines, language)
+        else:
+            return self._chunk_generic(file_path, lines, language)
+    
+    def _chunk_python(self, file_path: str, lines: List[str], language: str) -> List[CodeChunk]:
+        """Chunk Python code by class/function boundaries."""
+        chunks = []
+        current_chunk_start = 0
+        current_chunk_lines = []
+        current_name = None
+        current_type = "module"
+        
+        for i, line in enumerate(lines):
+            # Check for class definition
+            class_match = re.match(self.PYTHON_PATTERNS["class"], line)
+            func_match = re.match(self.PYTHON_PATTERNS["function"], line)
+            
+            if class_match or func_match:
+                # Save previous chunk if exists
+                if current_chunk_lines:
+                    chunks.append(CodeChunk(
+                        content="\n".join(current_chunk_lines),
+                        file_path=file_path,
+                        start_line=current_chunk_start + 1,
+                        end_line=i,
+                        chunk_type=current_type,
+                        name=current_name,
+                        language=language,
+                    ))
+                
+                # Start new chunk
+                current_chunk_start = i
+                current_chunk_lines = [line]
+                if class_match:
+                    current_name = class_match.group(1)
+                    current_type = "class"
+                else:
+                    current_name = func_match.group(1)
+                    current_type = "function"
+            else:
+                current_chunk_lines.append(line)
+                
+                # Check if chunk is too large
+                if len(current_chunk_lines) >= self.chunk_size:
+                    chunks.append(CodeChunk(
+                        content="\n".join(current_chunk_lines),
+                        file_path=file_path,
+                        start_line=current_chunk_start + 1,
+                        end_line=i + 1,
+                        chunk_type=current_type,
+                        name=current_name,
+                        language=language,
+                    ))
+                    # Start new chunk with overlap
+                    overlap_start = max(0, len(current_chunk_lines) - self.chunk_overlap)
+                    current_chunk_lines = current_chunk_lines[overlap_start:]
+                    current_chunk_start = i - len(current_chunk_lines) + 1
+        
+        # Add final chunk
+        if current_chunk_lines:
+            chunks.append(CodeChunk(
+                content="\n".join(current_chunk_lines),
+                file_path=file_path,
+                start_line=current_chunk_start + 1,
+                end_line=len(lines),
+                chunk_type=current_type,
+                name=current_name,
+                language=language,
+            ))
+        
+        return chunks
+
+    def _chunk_generic(self, file_path: str, lines: List[str], language: Optional[str]) -> List[CodeChunk]:
+        """Chunk code using sliding window approach."""
+        chunks = []
+        total_lines = len(lines)
+
+        if total_lines == 0:
+            return chunks
+
+        # If file is small enough, return as single chunk
+        if total_lines <= self.chunk_size:
+            chunks.append(CodeChunk(
+                content="\n".join(lines),
+                file_path=file_path,
+                start_line=1,
+                end_line=total_lines,
+                chunk_type="module",
+                name=Path(file_path).stem,
+                language=language,
+            ))
+            return chunks
+
+        # Use sliding window
+        start = 0
+        while start < total_lines:
+            end = min(start + self.chunk_size, total_lines)
+            chunk_lines = lines[start:end]
+
+            chunks.append(CodeChunk(
+                content="\n".join(chunk_lines),
+                file_path=file_path,
+                start_line=start + 1,
+                end_line=end,
+                chunk_type="block",
+                name=f"{Path(file_path).stem}:{start+1}-{end}",
+                language=language,
+            ))
+
+            # Move window with overlap
+            start += self.chunk_size - self.chunk_overlap
+
+            # Avoid tiny final chunks
+            if total_lines - start < self.chunk_overlap:
+                break
+
+        return chunks
+
+    def chunk_text(self, text: str, file_path: str = "unknown") -> List[CodeChunk]:
+        """Convenience method to chunk text directly."""
+        return self.chunk_file(file_path, text)
+

ai_code_assistant/retrieval/indexer.py

@@ -0,0 +1,236 @@
+"""
+Codebase indexer using ChromaDB and sentence-transformers.
+
+Indexes code files for semantic search.
+"""
+
+import os
+import hashlib
+from pathlib import Path
+from typing import List, Optional, Set
+from dataclasses import dataclass, field
+
+import chromadb
+from chromadb.config import Settings
+from sentence_transformers import SentenceTransformer
+
+from .chunker import CodeChunker, CodeChunk
+
+
+@dataclass
+class IndexStats:
+    """Statistics about the index."""
+    total_files: int = 0
+    total_chunks: int = 0
+    indexed_files: List[str] = field(default_factory=list)
+    skipped_files: List[str] = field(default_factory=list)
+    errors: List[str] = field(default_factory=list)
+
+
+@dataclass  
+class IndexConfig:
+    """Configuration for the indexer."""
+    # Embedding model (runs locally)
+    embedding_model: str = "all-MiniLM-L6-v2"
+    # ChromaDB persistence directory
+    persist_directory: str = ".ai-assistant-index"
+    # Collection name
+    collection_name: str = "codebase"
+    # File extensions to index
+    extensions: Set[str] = field(default_factory=lambda: {
+        ".py", ".js", ".ts", ".jsx", ".tsx", ".java", ".go", ".rs",
+        ".rb", ".php", ".c", ".cpp", ".h", ".hpp", ".cs", ".swift",
+        ".kt", ".scala", ".sh", ".yaml", ".yml", ".json", ".md", ".sql",
+        ".html", ".css", ".scss", ".vue", ".svelte",
+    })
+    # Directories to ignore
+    ignore_dirs: Set[str] = field(default_factory=lambda: {
+        ".git", ".svn", ".hg", "node_modules", "__pycache__", ".pytest_cache",
+        ".mypy_cache", ".tox", ".venv", "venv", "env", ".env", "dist", "build",
+        "target", ".idea", ".vscode", "coverage", ".coverage", "htmlcov",
+        ".eggs", "*.egg-info", ".ai-assistant-index",
+    })
+    # Max file size in bytes (skip large files)
+    max_file_size: int = 1024 * 1024  # 1MB
+
+
+class CodebaseIndexer:
+    """Indexes codebase into ChromaDB for semantic search."""
+    
+    def __init__(self, config: Optional[IndexConfig] = None, root_path: Optional[str] = None):
+        """Initialize the indexer.
+        
+        Args:
+            config: Index configuration
+            root_path: Root directory of the codebase
+        """
+        self.config = config or IndexConfig()
+        self.root_path = Path(root_path) if root_path else Path.cwd()
+        self.chunker = CodeChunker()
+        
+        # Initialize embedding model (lazy load)
+        self._embedder: Optional[SentenceTransformer] = None
+        
+        # Initialize ChromaDB
+        persist_path = self.root_path / self.config.persist_directory
+        self._client = chromadb.PersistentClient(
+            path=str(persist_path),
+            settings=Settings(anonymized_telemetry=False),
+        )
+        self._collection = self._client.get_or_create_collection(
+            name=self.config.collection_name,
+            metadata={"description": "Codebase index for semantic search"},
+        )
+    
+    @property
+    def embedder(self) -> SentenceTransformer:
+        """Lazy load the embedding model."""
+        if self._embedder is None:
+            self._embedder = SentenceTransformer(self.config.embedding_model)
+        return self._embedder
+    
+    def _should_index_file(self, file_path: Path) -> bool:
+        """Check if a file should be indexed."""
+        # Check extension
+        if file_path.suffix.lower() not in self.config.extensions:
+            return False
+        
+        # Check if in ignored directory
+        for part in file_path.parts:
+            if part in self.config.ignore_dirs:
+                return False
+            # Handle glob patterns like *.egg-info
+            for pattern in self.config.ignore_dirs:
+                if "*" in pattern and file_path.match(pattern):
+                    return False
+        
+        # Check file size
+        try:
+            if file_path.stat().st_size > self.config.max_file_size:
+                return False
+        except OSError:
+            return False
+        
+        return True
+    
+    def _get_file_hash(self, content: str) -> str:
+        """Get hash of file content for change detection."""
+        return hashlib.md5(content.encode()).hexdigest()
+    
+    def _read_file(self, file_path: Path) -> Optional[str]:
+        """Safely read a file."""
+        try:
+            return file_path.read_text(encoding="utf-8", errors="ignore")
+        except Exception:
+            return None
+    
+    def index_file(self, file_path: Path, stats: IndexStats) -> bool:
+        """Index a single file.
+        
+        Returns True if file was indexed, False if skipped.
+        """
+        content = self._read_file(file_path)
+        if content is None:
+            stats.errors.append(f"Could not read: {file_path}")
+            return False
+        
+        # Get relative path for storage
+        try:
+            rel_path = str(file_path.relative_to(self.root_path))
+        except ValueError:
+            rel_path = str(file_path)
+        
+        # Chunk the file
+        chunks = self.chunker.chunk_file(rel_path, content)
+        if not chunks:
+            stats.skipped_files.append(rel_path)
+            return False
+        
+        # Generate embeddings
+        texts = [chunk.content for chunk in chunks]
+        embeddings = self.embedder.encode(texts, show_progress_bar=False).tolist()
+        
+        # Prepare data for ChromaDB
+        ids = [chunk.id for chunk in chunks]
+        documents = texts
+        metadatas = [chunk.to_dict() for chunk in chunks]
+
+        # Delete old chunks for this file (if re-indexing)
+        try:
+            existing = self._collection.get(where={"file_path": rel_path})
+            if existing["ids"]:
+                self._collection.delete(ids=existing["ids"])
+        except Exception:
+            pass  # Collection might be empty
+
+        # Add to collection
+        self._collection.add(
+            ids=ids,
+            documents=documents,
+            embeddings=embeddings,
+            metadatas=metadatas,
+        )
+
+        stats.indexed_files.append(rel_path)
+        stats.total_chunks += len(chunks)
+        return True
+
+    def index_directory(self, directory: Optional[Path] = None, verbose: bool = True) -> IndexStats:
+        """Index all files in a directory recursively.
+
+        Args:
+            directory: Directory to index (defaults to root_path)
+            verbose: Print progress
+
+        Returns:
+            IndexStats with results
+        """
+        directory = directory or self.root_path
+        stats = IndexStats()
+
+        # Find all files to index
+        files_to_index = []
+        for file_path in directory.rglob("*"):
+            if file_path.is_file() and self._should_index_file(file_path):
+                files_to_index.append(file_path)
+
+        stats.total_files = len(files_to_index)
+
+        if verbose:
+            print(f"Found {len(files_to_index)} files to index...")
+
+        # Index each file
+        for i, file_path in enumerate(files_to_index):
+            if verbose and (i + 1) % 10 == 0:
+                print(f"  Indexed {i + 1}/{len(files_to_index)} files...")
+
+            try:
+                self.index_file(file_path, stats)
+            except Exception as e:
+                stats.errors.append(f"{file_path}: {str(e)}")
+
+        if verbose:
+            print(f"\n✓ Indexed {len(stats.indexed_files)} files ({stats.total_chunks} chunks)")
+            if stats.errors:
+                print(f"⚠ {len(stats.errors)} errors occurred")
+
+        return stats
+
+    def clear_index(self) -> None:
+        """Clear the entire index."""
+        self._client.delete_collection(self.config.collection_name)
+        self._collection = self._client.create_collection(
+            name=self.config.collection_name,
+            metadata={"description": "Codebase index for semantic search"},
+        )
+
+    def get_stats(self) -> dict:
+        """Get statistics about the current index."""
+        count = self._collection.count()
+        return {
+            "total_chunks": count,
+            "collection_name": self.config.collection_name,
+            "embedding_model": self.config.embedding_model,
+            "root_path": str(self.root_path),
+        }
+

ai_code_assistant/retrieval/search.py

@@ -0,0 +1,239 @@
+"""
+Semantic search over indexed codebase.
+
+Provides natural language search capabilities.
+"""
+
+from dataclasses import dataclass, field
+from typing import List, Optional
+from pathlib import Path
+
+import chromadb
+from chromadb.config import Settings
+from sentence_transformers import SentenceTransformer
+
+from .indexer import IndexConfig
+
+
+@dataclass
+class SearchResult:
+    """A single search result."""
+    content: str
+    file_path: str
+    start_line: int
+    end_line: int
+    chunk_type: str
+    name: str
+    language: str
+    score: float  # Similarity score (0-1, higher is better)
+    
+    def __str__(self) -> str:
+        """Format result for display."""
+        return (
+            f"📄 {self.file_path}:{self.start_line}-{self.end_line} "
+            f"({self.chunk_type}: {self.name}) [score: {self.score:.3f}]"
+        )
+    
+    def to_dict(self) -> dict:
+        """Convert to dictionary."""
+        return {
+            "content": self.content,
+            "file_path": self.file_path,
+            "start_line": self.start_line,
+            "end_line": self.end_line,
+            "chunk_type": self.chunk_type,
+            "name": self.name,
+            "language": self.language,
+            "score": self.score,
+        }
+
+
+@dataclass
+class SearchResponse:
+    """Response from a search query."""
+    query: str
+    results: List[SearchResult] = field(default_factory=list)
+    total_results: int = 0
+    
+    @property
+    def has_results(self) -> bool:
+        return len(self.results) > 0
+    
+    def format_for_llm(self, max_results: int = 5) -> str:
+        """Format results as context for LLM prompts."""
+        if not self.results:
+            return "No relevant code found."
+        
+        parts = [f"Found {len(self.results)} relevant code sections:\n"]
+        
+        for i, result in enumerate(self.results[:max_results], 1):
+            parts.append(f"\n--- Result {i}: {result.file_path}:{result.start_line}-{result.end_line} ---")
+            parts.append(f"Type: {result.chunk_type} | Name: {result.name}")
+            parts.append(f"```{result.language or ''}")
+            parts.append(result.content)
+            parts.append("```\n")
+        
+        return "\n".join(parts)
+
+
+class CodebaseSearch:
+    """Semantic search over indexed codebase."""
+    
+    def __init__(self, config: Optional[IndexConfig] = None, root_path: Optional[str] = None):
+        """Initialize search.
+        
+        Args:
+            config: Index configuration (must match indexer config)
+            root_path: Root directory of the codebase
+        """
+        self.config = config or IndexConfig()
+        self.root_path = Path(root_path) if root_path else Path.cwd()
+        
+        # Initialize embedding model (lazy load)
+        self._embedder: Optional[SentenceTransformer] = None
+        
+        # Connect to ChromaDB
+        persist_path = self.root_path / self.config.persist_directory
+        if not persist_path.exists():
+            raise FileNotFoundError(
+                f"Index not found at {persist_path}. Run 'ai-assist index' first."
+            )
+        
+        self._client = chromadb.PersistentClient(
+            path=str(persist_path),
+            settings=Settings(anonymized_telemetry=False),
+        )
+        
+        try:
+            self._collection = self._client.get_collection(self.config.collection_name)
+        except Exception:
+            raise FileNotFoundError(
+                f"Collection '{self.config.collection_name}' not found. Run 'ai-assist index' first."
+            )
+    
+    @property
+    def embedder(self) -> SentenceTransformer:
+        """Lazy load the embedding model."""
+        if self._embedder is None:
+            self._embedder = SentenceTransformer(self.config.embedding_model)
+        return self._embedder
+    
+    def search(
+        self,
+        query: str,
+        top_k: int = 10,
+        min_score: float = 0.0,
+        file_filter: Optional[str] = None,
+        language_filter: Optional[str] = None,
+    ) -> SearchResponse:
+        """Search for relevant code.
+        
+        Args:
+            query: Natural language search query
+            top_k: Maximum number of results
+            min_score: Minimum similarity score (0-1)
+            file_filter: Filter by file path (substring match)
+            language_filter: Filter by programming language
+            
+        Returns:
+            SearchResponse with results
+        """
+        # Generate query embedding
+        query_embedding = self.embedder.encode(query, show_progress_bar=False).tolist()
+        
+        # Build where clause for filtering
+        where = None
+        if language_filter:
+            where = {"language": language_filter}
+        
+        # Search ChromaDB
+        results = self._collection.query(
+            query_embeddings=[query_embedding],
+            n_results=top_k,
+            where=where,
+            include=["documents", "metadatas", "distances"],
+        )
+        
+        # Process results
+        search_results = []
+
+        if results["documents"] and results["documents"][0]:
+            documents = results["documents"][0]
+            metadatas = results["metadatas"][0] if results["metadatas"] else [{}] * len(documents)
+            distances = results["distances"][0] if results["distances"] else [0] * len(documents)
+
+            for doc, meta, distance in zip(documents, metadatas, distances):
+                # Convert distance to similarity score (ChromaDB returns L2 distance)
+                # Lower distance = more similar, so we convert
+                score = 1 / (1 + distance)
+
+                # Apply minimum score filter
+                if score < min_score:
+                    continue
+
+                # Apply file filter
+                file_path = meta.get("file_path", "")
+                if file_filter and file_filter.lower() not in file_path.lower():
+                    continue
+
+                search_results.append(SearchResult(
+                    content=doc,
+                    file_path=file_path,
+                    start_line=int(meta.get("start_line", 0)),
+                    end_line=int(meta.get("end_line", 0)),
+                    chunk_type=meta.get("chunk_type", "unknown"),
+                    name=meta.get("name", ""),
+                    language=meta.get("language", ""),
+                    score=score,
+                ))
+
+        return SearchResponse(
+            query=query,
+            results=search_results,
+            total_results=len(search_results),
+        )
+
+    def search_similar(self, code: str, top_k: int = 5) -> SearchResponse:
+        """Find code similar to the given code snippet.
+
+        Args:
+            code: Code snippet to find similar code for
+            top_k: Maximum number of results
+
+        Returns:
+            SearchResponse with similar code
+        """
+        return self.search(query=code, top_k=top_k)
+
+    def get_file_context(self, file_path: str) -> List[SearchResult]:
+        """Get all indexed chunks for a specific file.
+
+        Args:
+            file_path: Path to the file
+
+        Returns:
+            List of SearchResult for the file
+        """
+        results = self._collection.get(
+            where={"file_path": file_path},
+            include=["documents", "metadatas"],
+        )
+
+        search_results = []
+        if results["documents"]:
+            for doc, meta in zip(results["documents"], results["metadatas"]):
+                search_results.append(SearchResult(
+                    content=doc,
+                    file_path=meta.get("file_path", ""),
+                    start_line=int(meta.get("start_line", 0)),
+                    end_line=int(meta.get("end_line", 0)),
+                    chunk_type=meta.get("chunk_type", "unknown"),
+                    name=meta.get("name", ""),
+                    language=meta.get("language", ""),
+                    score=1.0,  # Not a similarity search
+                ))
+
+        # Sort by line number
+        search_results.sort(key=lambda x: x.start_line)
+        return search_results
+

ai_code_assistant/reviewer/__init__.py

@@ -0,0 +1,7 @@
+"""Code review module for AI Code Assistant."""
+
+from ai_code_assistant.reviewer.analyzer import CodeAnalyzer, ReviewResult, ReviewIssue
+from ai_code_assistant.reviewer.prompts import REVIEW_PROMPTS
+
+__all__ = ["CodeAnalyzer", "ReviewResult", "ReviewIssue", "REVIEW_PROMPTS"]
+