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

coding_assistant/documentation/writers/markdown_writer.py

@@ -0,0 +1,322 @@
+"""Markdown documentation writer.
+
+This module writes generated documentation and diagrams to markdown files
+with proper structure, formatting, and navigation.
+"""
+
+from pathlib import Path
+from typing import Dict, List
+from datetime import datetime
+
+from coding_assistant.documentation.decomposition.partitioner import Partition
+from coding_assistant.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class MarkdownWriter:
+    """
+    Write documentation to markdown files.
+
+    Creates:
+    - README.md (repository overview)
+    - modules/ (per-module documentation)
+    - diagrams/ (visual diagrams)
+    - index.md (navigation)
+    """
+
+    def __init__(self, output_dir: Path):
+        """
+        Initialize markdown writer.
+
+        Args:
+            output_dir: Directory to write documentation
+        """
+        self.output_dir = Path(output_dir)
+        self.modules_dir = self.output_dir / "modules"
+        self.diagrams_dir = self.output_dir / "diagrams"
+
+        # Create directories
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        self.modules_dir.mkdir(exist_ok=True)
+        self.diagrams_dir.mkdir(exist_ok=True)
+
+    def write_documentation(self,
+                             overview: str,
+                             module_docs: Dict[str, str],
+                             diagrams: Dict[str, str],
+                             partitions: List[Partition],
+                             repo_metadata: Dict) -> List[Path]:
+        """
+        Write all documentation to files.
+
+        Args:
+            overview: Repository overview markdown
+            module_docs: Module name -> documentation mapping
+            diagrams: Diagram type -> mermaid code mapping
+            partitions: List of partitions
+            repo_metadata: Repository metadata
+
+        Returns:
+            List of paths to written files
+        """
+        logger.info(f"Writing documentation to {self.output_dir}")
+
+        written_files = []
+
+        # 1. Write README.md (overview)
+        readme_path = self._write_readme(overview, diagrams, repo_metadata)
+        written_files.append(readme_path)
+
+        # 2. Write module documentation
+        module_paths = self._write_modules(module_docs, partitions)
+        written_files.extend(module_paths)
+
+        # 3. Write diagrams
+        diagram_paths = self._write_diagrams(diagrams)
+        written_files.extend(diagram_paths)
+
+        # 4. Write index/navigation
+        index_path = self._write_index(module_docs, diagrams, repo_metadata)
+        written_files.append(index_path)
+
+        logger.info(f"Wrote {len(written_files)} documentation files")
+
+        return written_files
+
+    def _write_readme(self,
+                       overview: str,
+                       diagrams: Dict[str, str],
+                       repo_metadata: Dict) -> Path:
+        """Write README.md with repository overview."""
+
+        readme_path = self.output_dir / "README.md"
+
+        content = []
+
+        # Header
+        repo_name = repo_metadata.get('name', 'Repository')
+        content.append(f"# {repo_name} Documentation")
+        content.append("")
+        content.append(f"*Generated on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*")
+        content.append("")
+
+        # Overview
+        content.append(overview)
+        content.append("")
+
+        # Architecture diagram (if available)
+        if 'architecture' in diagrams:
+            content.append("## Architecture")
+            content.append("")
+            content.append("```mermaid")
+            content.append(diagrams['architecture'])
+            content.append("```")
+            content.append("")
+
+        # Navigation
+        content.append("## Documentation")
+        content.append("")
+        content.append("- [Module Index](index.md) - Browse all modules")
+        content.append("- [Modules](modules/) - Detailed module documentation")
+        content.append("- [Diagrams](diagrams/) - Visual diagrams")
+        content.append("")
+
+        # Repository stats
+        content.append("## Repository Statistics")
+        content.append("")
+        content.append(f"- **Total Files**: {repo_metadata.get('total_files', 'N/A')}")
+        content.append(f"- **Modules Detected**: {repo_metadata.get('total_modules', 'N/A')}")
+        content.append(f"- **Partitions**: {repo_metadata.get('total_partitions', 'N/A')}")
+        content.append("")
+
+        # Footer
+        content.append("---")
+        content.append("")
+        content.append("*This documentation was auto-generated using CodeWiki integration.*")
+        content.append("")
+
+        readme_path.write_text("\\n".join(content))
+
+        logger.debug(f"Wrote README.md ({len(content)} lines)")
+
+        return readme_path
+
+    def _write_modules(self,
+                        module_docs: Dict[str, str],
+                        partitions: List[Partition]) -> List[Path]:
+        """Write individual module documentation files."""
+
+        written_paths = []
+
+        # Create partition lookup
+        partition_map = {p.name: p for p in partitions}
+
+        for module_name, doc_content in module_docs.items():
+            # Sanitize module name for filename
+            safe_name = self._sanitize_filename(module_name)
+            module_path = self.modules_dir / f"{safe_name}.md"
+
+            content = []
+
+            # Header
+            content.append(f"# {module_name}")
+            content.append("")
+
+            # Metadata (if partition available)
+            partition = partition_map.get(module_name)
+            if partition:
+                content.append("## Module Information")
+                content.append("")
+                content.append(f"- **Files**: {len(partition.files)}")
+                content.append(f"- **Lines of Code**: {partition.size_loc}")
+                content.append(f"- **Cohesion Score**: {partition.cohesion_score:.2f}")
+                content.append(f"- **Dependencies**: {len(partition.dependencies)}")
+                content.append("")
+
+                # File list
+                if partition.files:
+                    content.append("### Files in Module")
+                    content.append("")
+                    for file_path in partition.files[:20]:  # Limit to 20
+                        rel_path = Path(file_path).name
+                        content.append(f"- `{rel_path}`")
+
+                    if len(partition.files) > 20:
+                        content.append(f"- *... and {len(partition.files) - 20} more files*")
+
+                    content.append("")
+
+            # Main documentation
+            content.append("## Documentation")
+            content.append("")
+            content.append(doc_content)
+            content.append("")
+
+            # Navigation
+            content.append("---")
+            content.append("")
+            content.append("[← Back to Index](../index.md)")
+            content.append("")
+
+            module_path.write_text("\\n".join(content))
+            written_paths.append(module_path)
+
+        logger.debug(f"Wrote {len(written_paths)} module documentation files")
+
+        return written_paths
+
+    def _write_diagrams(self, diagrams: Dict[str, str]) -> List[Path]:
+        """Write diagram files in mermaid format."""
+
+        written_paths = []
+
+        for diagram_type, diagram_code in diagrams.items():
+            diagram_path = self.diagrams_dir / f"{diagram_type}.mmd"
+
+            diagram_path.write_text(diagram_code)
+            written_paths.append(diagram_path)
+
+            # Also create markdown file with embedded diagram
+            md_path = self.diagrams_dir / f"{diagram_type}.md"
+
+            content = []
+            content.append(f"# {diagram_type.title()} Diagram")
+            content.append("")
+            content.append(f"*Generated diagram for {diagram_type}*")
+            content.append("")
+            content.append("```mermaid")
+            content.append(diagram_code)
+            content.append("```")
+            content.append("")
+            content.append("---")
+            content.append("")
+            content.append(f"[Download .mmd file]({diagram_type}.mmd) | [Back to Index](../index.md)")
+            content.append("")
+
+            md_path.write_text("\\n".join(content))
+            written_paths.append(md_path)
+
+        logger.debug(f"Wrote {len(written_paths)} diagram files")
+
+        return written_paths
+
+    def _write_index(self,
+                      module_docs: Dict[str, str],
+                      diagrams: Dict[str, str],
+                      repo_metadata: Dict) -> Path:
+        """Write index.md with navigation."""
+
+        index_path = self.output_dir / "index.md"
+
+        content = []
+
+        # Header
+        content.append("# Documentation Index")
+        content.append("")
+        content.append(f"*{repo_metadata.get('name', 'Repository')} Documentation*")
+        content.append("")
+
+        # Modules section
+        content.append("## Modules")
+        content.append("")
+        content.append(f"This repository contains {len(module_docs)} documented modules:")
+        content.append("")
+
+        # Sort modules alphabetically
+        sorted_modules = sorted(module_docs.keys())
+
+        for module_name in sorted_modules:
+            safe_name = self._sanitize_filename(module_name)
+            content.append(f"- [{module_name}](modules/{safe_name}.md)")
+
+        content.append("")
+
+        # Diagrams section
+        if diagrams:
+            content.append("## Diagrams")
+            content.append("")
+            content.append("Visual representations of the codebase:")
+            content.append("")
+
+            for diagram_type in sorted(diagrams.keys()):
+                content.append(f"- [{diagram_type.title()}](diagrams/{diagram_type}.md)")
+
+            content.append("")
+
+        # Quick links
+        content.append("## Quick Links")
+        content.append("")
+        content.append("- [Repository Overview](README.md)")
+        content.append("")
+
+        # Footer
+        content.append("---")
+        content.append("")
+        content.append(f"*Total: {len(module_docs)} modules, {len(diagrams)} diagrams*")
+        content.append("")
+
+        index_path.write_text("\\n".join(content))
+
+        logger.debug("Wrote index.md")
+
+        return index_path
+
+    def _sanitize_filename(self, name: str) -> str:
+        """
+        Sanitize module name for use as filename.
+
+        Replaces special characters with underscores.
+        """
+        import re
+
+        # Replace path separators, dots, spaces with underscores
+        sanitized = re.sub(r'[/\\\\.\\s]', '_', name)
+
+        # Remove any remaining special characters
+        sanitized = re.sub(r'[^a-zA-Z0-9_-]', '', sanitized)
+
+        # Remove leading/trailing underscores
+        sanitized = sanitized.strip('_')
+
+        return sanitized or 'unnamed_module'

coding_assistant/embeddings/generator.py

@@ -0,0 +1,89 @@
+"""Embedding generation using sentence-transformers (local, no API needed)."""
+from typing import List, Dict
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from rich.console import Console
+
+console = Console()
+
+
+class EmbeddingGenerator:
+    """Generate embeddings for code using local sentence-transformers."""
+
+    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
+        """
+        Initialize the embedding generator.
+
+        Args:
+            model_name: Model to use. Default is all-MiniLM-L6-v2
+                       (fast, 384 dimensions, good for code)
+        """
+        console.print(f"[dim]Loading embedding model: [bold #06B6D4]{model_name}[/bold #06B6D4]...[/dim]")
+        self.model = SentenceTransformer(model_name)
+        self.dimension = self.model.get_sentence_embedding_dimension()
+        console.print(f"[bold #10B981]✓[/bold #10B981] [dim]Model loaded (dimension: {self.dimension})[/dim]")
+
+    def generate_embedding(self, text: str) -> List[float]:
+        """
+        Generate embedding for a single text.
+
+        Args:
+            text: Text to embed
+
+        Returns:
+            Embedding vector as list of floats
+        """
+        embedding = self.model.encode(text, convert_to_numpy=True)
+        return embedding.tolist()
+
+    def generate_embeddings_batch(self, texts: List[str]) -> List[List[float]]:
+        """
+        Generate embeddings for multiple texts (more efficient).
+
+        Args:
+            texts: List of texts to embed
+
+        Returns:
+            List of embedding vectors
+        """
+        embeddings = self.model.encode(
+            texts,
+            convert_to_numpy=True,
+            show_progress_bar=len(texts) > 10
+        )
+        return embeddings.tolist()
+
+    def embed_code_chunks(self, chunks: List[Dict]) -> List[Dict]:
+        """
+        Embed code chunks with metadata.
+
+        Args:
+            chunks: List of code chunks from parser
+
+        Returns:
+            Chunks with embeddings added
+        """
+        if not chunks:
+            return []
+
+        # Extract texts to embed
+        texts = []
+        for chunk in chunks:
+            # Combine type, name (if exists), and content for embedding
+            text_parts = [chunk['type']]
+
+            if 'name' in chunk:
+                text_parts.append(chunk['name'])
+
+            text_parts.append(chunk['content'])
+
+            texts.append(' '.join(text_parts))
+
+        # Generate embeddings in batch
+        embeddings = self.generate_embeddings_batch(texts)
+
+        # Add embeddings to chunks
+        for chunk, embedding in zip(chunks, embeddings):
+            chunk['embedding'] = embedding
+
+        return chunks

coding_assistant/embeddings/store.py

@@ -0,0 +1,187 @@
+"""Vector store using ChromaDB for semantic code search."""
+from pathlib import Path
+from typing import List, Dict, Optional
+import chromadb
+from chromadb.config import Settings
+import hashlib
+import warnings
+import sys
+import os
+
+
+class VectorStore:
+    """Manages code embeddings in ChromaDB."""
+
+    def __init__(self, persist_dir: Optional[Path] = None):
+        """
+        Initialize ChromaDB vector store.
+
+        Args:
+            persist_dir: Directory to persist the database
+        """
+        if persist_dir is None:
+            persist_dir = Path.home() / ".coding_assistant" / "chroma_db"
+
+        persist_dir.mkdir(parents=True, exist_ok=True)
+
+        # Suppress ChromaDB telemetry warnings during initialization
+        # Save original stderr
+        original_stderr = sys.stderr
+        try:
+            # Redirect stderr to devnull during initialization
+            sys.stderr = open(os.devnull, 'w')
+
+            self.client = chromadb.PersistentClient(
+                path=str(persist_dir),
+                settings=Settings(
+                    anonymized_telemetry=False,
+                    allow_reset=True
+                )
+            )
+
+            # Get or create collection
+            self.collection = self.client.get_or_create_collection(
+                name="codebase",
+                metadata={"hnsw:space": "cosine"}  # Use cosine similarity
+            )
+        finally:
+            # Restore original stderr
+            sys.stderr.close()
+            sys.stderr = original_stderr
+
+        from rich.console import Console
+        console = Console()
+        console.print(f"[bold #10B981]✓[/bold #10B981] [dim]ChromaDB initialized ({self.collection.count()} chunks)[/dim]")
+
+    def add_chunks(self, chunks: List[Dict]):
+        """
+        Add code chunks with embeddings to the store.
+
+        Args:
+            chunks: List of chunks with embeddings
+        """
+        if not chunks:
+            return
+
+        ids = []
+        embeddings = []
+        documents = []
+        metadatas = []
+
+        for chunk in chunks:
+            # Create unique ID from file path + chunk type + position
+            chunk_id = self._generate_chunk_id(chunk)
+            ids.append(chunk_id)
+
+            embeddings.append(chunk['embedding'])
+            documents.append(chunk['content'])
+
+            # Store metadata (without embedding to save space)
+            metadata = {
+                'file_path': chunk['file_path'],
+                'type': chunk['type'],
+                'start_line': chunk['start_line'],
+                'end_line': chunk['end_line']
+            }
+
+            if 'name' in chunk:
+                metadata['name'] = chunk['name']
+
+            metadatas.append(metadata)
+
+        # Add to ChromaDB
+        self.collection.add(
+            ids=ids,
+            embeddings=embeddings,
+            documents=documents,
+            metadatas=metadatas
+        )
+
+        print(f"✓ Added {len(chunks)} chunks to vector store")
+
+    def search(
+        self,
+        query_embedding: List[float],
+        n_results: int = 10,
+        filter_metadata: Optional[Dict] = None
+    ) -> List[Dict]:
+        """
+        Search for similar code chunks.
+
+        Args:
+            query_embedding: Query embedding vector
+            n_results: Number of results to return
+            filter_metadata: Optional metadata filters
+
+        Returns:
+            List of matching chunks with similarity scores
+        """
+        # Suppress ChromaDB warnings during query
+        original_stderr = sys.stderr
+        try:
+            sys.stderr = open(os.devnull, 'w')
+
+            results = self.collection.query(
+                query_embeddings=[query_embedding],
+                n_results=n_results,
+                where=filter_metadata,
+                include=['documents', 'metadatas', 'distances']
+            )
+        finally:
+            sys.stderr.close()
+            sys.stderr = original_stderr
+
+        return self._format_results(results)
+
+    def delete_by_file(self, file_path: str):
+        """
+        Delete all chunks from a specific file.
+
+        Args:
+            file_path: Path to the file
+        """
+        # Get all IDs for this file
+        results = self.collection.get(
+            where={"file_path": file_path}
+        )
+
+        if results['ids']:
+            self.collection.delete(ids=results['ids'])
+            print(f"✓ Deleted {len(results['ids'])} chunks from {file_path}")
+
+    def clear(self):
+        """Clear all chunks from the store."""
+        self.client.delete_collection("codebase")
+        self.collection = self.client.create_collection(
+            name="codebase",
+            metadata={"hnsw:space": "cosine"}
+        )
+        print("✓ Vector store cleared")
+
+    def count(self) -> int:
+        """Get total number of chunks in the store."""
+        return self.collection.count()
+
+    def _generate_chunk_id(self, chunk: Dict) -> str:
+        """Generate unique ID for a chunk."""
+        # Use file path + type + start line
+        id_string = f"{chunk['file_path']}:{chunk['type']}:{chunk['start_line']}"
+        return hashlib.md5(id_string.encode()).hexdigest()
+
+    def _format_results(self, raw_results: Dict) -> List[Dict]:
+        """Format ChromaDB results into usable structure."""
+        formatted = []
+
+        if not raw_results['ids'] or not raw_results['ids'][0]:
+            return formatted
+
+        for i in range(len(raw_results['ids'][0])):
+            formatted.append({
+                'id': raw_results['ids'][0][i],
+                'content': raw_results['documents'][0][i],
+                'metadata': raw_results['metadatas'][0][i],
+                'distance': raw_results['distances'][0][i],
+                'similarity': 1 - raw_results['distances'][0][i]  # Convert distance to similarity
+            })
+
+        return formatted

coding_assistant/exceptions/__init__.py

@@ -0,0 +1,50 @@
+"""Custom exceptions for the AI Coding Assistant."""
+
+from coding_assistant.exceptions.base import AssistantError
+from coding_assistant.exceptions.llm import (
+    LLMError,
+    LLMConnectionError,
+    LLMResponseError,
+    LLMTimeoutError
+)
+from coding_assistant.exceptions.storage import (
+    StorageError,
+    IndexError,
+    IndexNotFoundError,
+    SessionError
+)
+from coding_assistant.exceptions.validation import (
+    ValidationError,
+    FileNotFoundError,
+    InvalidQueryError,
+    InvalidParameterError,
+    ConfigurationError
+)
+from coding_assistant.exceptions.recovery import ErrorRecovery
+
+__all__ = [
+    # Base
+    'AssistantError',
+
+    # LLM
+    'LLMError',
+    'LLMConnectionError',
+    'LLMResponseError',
+    'LLMTimeoutError',
+
+    # Storage
+    'StorageError',
+    'IndexError',
+    'IndexNotFoundError',
+    'SessionError',
+
+    # Validation
+    'ValidationError',
+    'FileNotFoundError',
+    'InvalidQueryError',
+    'InvalidParameterError',
+    'ConfigurationError',
+
+    # Recovery
+    'ErrorRecovery',
+]