feat: add DOT diagram generation and management

src/core/diagram.py

@@ -0,0 +1,470 @@
+"""
+Diagram Generator Module
+
+Handles Graphviz rendering, sanitization, and diagram management.
+"""
+
+import os
+import re
+import json
+import logging
+import tempfile
+import hashlib
+from datetime import datetime
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional, List, Tuple, Dict
+
+from ..config import get_config, DIAGRAMS_DIR
+
+logger = logging.getLogger("codeatlas.diagram")
+
+# Try to import graphviz
+try:
+    import graphviz
+    GRAPHVIZ_AVAILABLE = True
+except ImportError:
+    GRAPHVIZ_AVAILABLE = False
+    logger.warning("Graphviz not available")
+
+
+@dataclass
+class LayoutOptions:
+    """Options for diagram layout."""
+    direction: str = "TB"  # TB, LR, BT, RL
+    splines: str = "polyline"  # polyline, ortho, spline, line
+    nodesep: float = 0.5
+    ranksep: float = 0.75
+    zoom: float = 1.0
+    
+    DIRECTION_MAP = {
+        "Top → Down": "TB",
+        "Left → Right": "LR", 
+        "Bottom → Up": "BT",
+        "Right → Left": "RL",
+    }
+    
+    @classmethod
+    def from_ui(cls, direction: str, splines: str, nodesep: float, ranksep: float, zoom: float = 1.0):
+        """Create from UI values."""
+        return cls(
+            direction=cls.DIRECTION_MAP.get(direction, "TB"),
+            splines=splines,
+            nodesep=nodesep,
+            ranksep=ranksep,
+            zoom=zoom,
+        )
+
+
+@dataclass
+class DiagramInfo:
+    """Information about a saved diagram."""
+    filename: str
+    repo_name: str
+    timestamp: str
+    formatted_timestamp: str
+    file_path: Path
+    # Metadata fields (loaded from JSON if available)
+    model_name: str = ""
+    files_processed: int = 0
+    total_characters: int = 0
+    node_count: int = 0
+    edge_count: int = 0
+
+
+class DiagramGenerator:
+    """Generates and manages architecture diagrams."""
+    
+    def __init__(self):
+        self.config = get_config()
+        self.diagrams_dir = self.config.diagrams_dir
+    
+    def render(
+        self, 
+        dot_source: str, 
+        layout: Optional[LayoutOptions] = None,
+        repo_name: str = "",
+        save_to_history: bool = False,
+        metadata: Optional[Dict] = None
+    ) -> str:
+        """Render DOT source to HTML with embedded SVG.
+        
+        Args:
+            dot_source: Graphviz DOT source code
+            layout: Layout options
+            repo_name: Repository name for saving
+            save_to_history: Whether to save to history
+            metadata: Optional metadata dict (model, files, chars) to save with diagram
+            
+        Returns:
+            HTML string with SVG diagram
+        """
+        if dot_source.startswith("Error:"):
+            return self._error_html(dot_source)
+        
+        if not GRAPHVIZ_AVAILABLE:
+            return self._fallback_html(dot_source)
+        
+        if layout is None:
+            layout = LayoutOptions()
+        
+        try:
+            # Clean and prepare DOT source
+            dot_source = self._prepare_dot(dot_source)
+            
+            # Save raw diagram if requested
+            if save_to_history:
+                self._save_diagram(dot_source, "raw", repo_name, metadata)
+            
+            # Sanitize DOT code
+            dot_source = self._sanitize_dot(dot_source)
+            
+            # Apply layout settings
+            dot_source = self._apply_layout(dot_source, layout)
+            
+            # Render to SVG
+            svg_content = self._render_svg(dot_source)
+            
+            # Wrap in HTML
+            return self._wrap_svg(svg_content, layout.zoom)
+            
+        except Exception as e:
+            logger.exception("Rendering failed")
+            return self._fallback_html(dot_source, error=str(e))
+    
+    def _prepare_dot(self, dot_source: str) -> str:
+        """Prepare DOT source by removing markdown and ensuring structure."""
+        dot_source = dot_source.strip()
+        
+        # Remove markdown code fences
+        if dot_source.startswith("```"):
+            lines = dot_source.split("\n")
+            if lines[0].startswith("```"):
+                lines = lines[1:]
+            if lines and lines[-1].strip() == "```":
+                lines = lines[:-1]
+            dot_source = "\n".join(lines)
+        
+        # Ensure digraph wrapper
+        if "digraph" not in dot_source and "graph" not in dot_source:
+            dot_source = f"digraph G {{\n{dot_source}\n}}"
+        
+        return dot_source
+    
+    def _sanitize_dot(self, dot_source: str) -> str:
+        """Sanitize DOT source to fix common LLM output issues."""
+        # Check for error responses
+        if "Error" in dot_source and any(x in dot_source for x in ["429", "RESOURCE_EXHAUSTED", "quota"]):
+            raise ValueError("Rate limited - received error instead of diagram")
+        
+        lines = dot_source.split("\n")
+        sanitized = []
+        brace_count = 0
+        
+        for i, line in enumerate(lines):
+            stripped = line.strip()
+            if not stripped or stripped.startswith("//") or stripped.startswith("#"):
+                sanitized.append(line)
+                continue
+            
+            brace_count += line.count("{") - line.count("}")
+            
+            # Remove HTML-like tags that Graphviz doesn't support
+            line = re.sub(r"<[^>]+>", "", line)
+            
+            # Fix incomplete edges at end of file
+            is_last_few = i >= len(lines) - 3
+            if is_last_few:
+                if stripped.endswith("->") or re.match(r".*->\s*$", stripped):
+                    continue
+                if "->" in stripped and not stripped.endswith(";") and "[" not in stripped:
+                    parts = stripped.split("->")
+                    if len(parts) == 2 and not parts[1].strip():
+                        continue
+                # Fix unclosed quotes
+                if stripped.count('"') % 2 == 1:
+                    line = line.rstrip() + '"];'
+            
+            sanitized.append(line)
+        
+        result = "\n".join(sanitized)
+        
+        # Balance braces
+        if brace_count > 0:
+            result += "\n" + "}" * brace_count
+        
+        return result
+    
+    def _apply_layout(self, dot_source: str, layout: LayoutOptions) -> str:
+        """Apply layout settings to DOT source."""
+        # Remove existing layout attributes
+        dot_source = re.sub(r"rankdir\s*=\s*\w+\s*;?", "", dot_source)
+        dot_source = re.sub(r"splines\s*=\s*\w+\s*;?", "", dot_source)
+        dot_source = re.sub(r"nodesep\s*=\s*[\d.]+\s*;?", "", dot_source)
+        dot_source = re.sub(r"ranksep\s*=\s*[\d.]+\s*;?", "", dot_source)
+        
+        # Add new layout settings
+        layout_settings = f"""
+    rankdir={layout.direction};
+    splines={layout.splines};
+    nodesep={layout.nodesep};
+    ranksep={layout.ranksep};
+    pad=0.5;
+    node [shape=box, style="rounded,filled", fontname="Helvetica", fontsize=12, width=2.0, height=0.6, margin="0.2,0.1"];
+    edge [fontname="Helvetica", fontsize=10, arrowsize=0.8];
+"""
+        dot_source = dot_source.replace("{", "{" + layout_settings, 1)
+        
+        # For ortho splines, convert label to xlabel
+        if layout.splines == "ortho":
+            lines = dot_source.split("\n")
+            converted = []
+            for line in lines:
+                if "->" in line and "label=" in line and "xlabel=" not in line:
+                    line = line.replace("label=", "xlabel=")
+                converted.append(line)
+            dot_source = "\n".join(converted)
+        
+        return dot_source
+    
+    def _render_svg(self, dot_source: str) -> str:
+        """Render DOT to SVG using Graphviz."""
+        graph = graphviz.Source(dot_source)
+        
+        with tempfile.TemporaryDirectory() as tmpdir:
+            output_path = os.path.join(tmpdir, "diagram")
+            graph.render(output_path, format="svg", cleanup=True)
+            
+            with open(output_path + ".svg", "r", encoding="utf-8") as f:
+                return f.read()
+    
+    def _wrap_svg(self, svg_content: str, zoom: float = 1.0) -> str:
+        """Wrap SVG in responsive HTML container."""
+        # Make SVG responsive
+        svg_content = re.sub(r'<svg([^>]*?)width="[^"]*"', r'<svg\1width="100%"', svg_content)
+        svg_content = re.sub(r'height="[^"]*"', 'height="auto"', svg_content, count=1)
+        
+        # Apply zoom
+        transform = f'style="transform: scale({zoom}); transform-origin: top left;"' if zoom != 1.0 else ""
+        
+        return f'''<div class="diagram-box">
+            <div class="diagram-inner" {transform}>
+                {svg_content}
+            </div>
+        </div>'''
+    
+    def _save_diagram(
+        self, 
+        dot_content: str, 
+        prefix: str, 
+        repo_name: str = "",
+        metadata: Optional[Dict] = None
+    ) -> Path:
+        """Save DOT content to history with optional metadata."""
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        
+        # Sanitize repo name
+        safe_repo = repo_name.replace("/", "_").replace(" ", "_") if repo_name else ""
+        safe_repo = "".join(c for c in safe_repo if c.isalnum() or c in "_-")[:50]
+        
+        filename = f"{prefix}_{safe_repo}_{timestamp}.dot" if safe_repo else f"{prefix}_{timestamp}.dot"
+        filepath = self.diagrams_dir / filename
+        
+        try:
+            with open(filepath, "w", encoding="utf-8") as f:
+                f.write(dot_content)
+            logger.info(f"Saved diagram: {filepath}")
+            
+            # Save metadata as JSON if provided
+            if metadata:
+                # Add node/edge counts from DOT content
+                node_count, edge_count = self._count_nodes_edges(dot_content)
+                metadata["node_count"] = node_count
+                metadata["edge_count"] = edge_count
+                metadata["repo_name"] = repo_name
+                metadata["timestamp"] = timestamp
+                
+                meta_filepath = filepath.with_suffix(".json")
+                with open(meta_filepath, "w", encoding="utf-8") as f:
+                    json.dump(metadata, f, indent=2)
+                logger.info(f"Saved metadata: {meta_filepath}")
+            
+            return filepath
+        except Exception as e:
+            logger.warning(f"Failed to save diagram: {e}")
+            return None
+    
+    def _count_nodes_edges(self, dot_content: str) -> Tuple[int, int]:
+        """Count nodes and edges in DOT source."""
+        # Count edges (lines with ->), handling quoted node names
+        edge_pattern = r'(?:"[^"]+"|[\w]+)\s*->\s*(?:"[^"]+"|[\w]+)'
+        edges = len(re.findall(edge_pattern, dot_content))
+        
+        # Count unique node names (excluding keywords)
+        keywords = {'digraph', 'graph', 'subgraph', 'node', 'edge', 'rankdir', 'splines', 'nodesep', 'ranksep', 'label', 'style', 'shape', 'color', 'fillcolor', 'fontname', 'fontsize', 'margin', 'pad', 'width', 'height', 'arrowsize', 'cluster', 'tb', 'lr', 'bt', 'rl'}
+        
+        nodes = set()
+        
+        # Match quoted node names in definitions: "Node Name" [...]
+        quoted_defs = re.findall(r'^\s*"([^"]+)"\s*\[', dot_content, re.MULTILINE)
+        for node in quoted_defs:
+            nodes.add(node)
+        
+        # Match unquoted node definitions: NodeName [...]  
+        unquoted_defs = re.findall(r'^\s*(\w+)\s*\[', dot_content, re.MULTILINE)
+        for node in unquoted_defs:
+            if node.lower() not in keywords and not node.isdigit():
+                nodes.add(node)
+        
+        # Match nodes in edges (both quoted and unquoted)
+        edge_nodes = re.findall(r'(?:"([^"]+)"|(\w+))\s*->\s*(?:"([^"]+)"|(\w+))', dot_content)
+        for match in edge_nodes:
+            # Each match is (quoted_src, unquoted_src, quoted_dst, unquoted_dst)
+            src = match[0] or match[1]
+            dst = match[2] or match[3]
+            if src and src.lower() not in keywords and not src.isdigit():
+                nodes.add(src)
+            if dst and dst.lower() not in keywords and not dst.isdigit():
+                nodes.add(dst)
+        
+        return len(nodes), edges
+    
+    def get_history(self, limit: int = 50) -> List[DiagramInfo]:
+        """Get list of saved diagrams.
+        
+        Args:
+            limit: Maximum number of diagrams to return
+            
+        Returns:
+            List of DiagramInfo objects
+        """
+        if not self.diagrams_dir.exists():
+            return []
+        
+        files = [f for f in self.diagrams_dir.iterdir() if f.name.startswith("raw_") and f.suffix == ".dot"]
+        
+        # Sort by timestamp (newest first)
+        def extract_timestamp(path: Path) -> str:
+            name = path.stem.replace("raw_", "")
+            parts = name.split("_")
+            if len(parts) >= 2:
+                return parts[-2] + parts[-1]
+            return "0"
+        
+        files.sort(key=extract_timestamp, reverse=True)
+        
+        # Build DiagramInfo list (no deduplication - show all history)
+        diagrams = []
+        for f in files[:limit]:
+            name = f.stem.replace("raw_", "")
+            parts = name.split("_")
+            
+            if len(parts) >= 2 and len(parts[-2]) == 8 and len(parts[-1]) == 6:
+                repo_parts = parts[:-2]
+                repo_name = repo_parts[-1] if repo_parts else "local"
+                date_part = parts[-2]
+                time_part = parts[-1]
+                
+                try:
+                    formatted_ts = f"{date_part[:4]}-{date_part[4:6]}-{date_part[6:8]} {time_part[:2]}:{time_part[2:4]}"
+                except:
+                    formatted_ts = f"{date_part}_{time_part}"
+                
+                # Load metadata if available
+                metadata = self._load_metadata(f)
+                
+                diagrams.append(DiagramInfo(
+                    filename=f.name,
+                    repo_name=metadata.get("repo_name", repo_name) if metadata else repo_name,
+                    timestamp=f"{date_part}_{time_part}",
+                    formatted_timestamp=formatted_ts,
+                    file_path=f,
+                    model_name=metadata.get("model_name", "") if metadata else "",
+                    files_processed=metadata.get("files_processed", 0) if metadata else 0,
+                    total_characters=metadata.get("total_characters", 0) if metadata else 0,
+                    node_count=metadata.get("node_count", 0) if metadata else 0,
+                    edge_count=metadata.get("edge_count", 0) if metadata else 0,
+                ))
+        
+        return diagrams
+    
+    def _load_metadata(self, dot_filepath: Path) -> Optional[Dict]:
+        """Load metadata JSON for a diagram file."""
+        meta_filepath = dot_filepath.with_suffix(".json")
+        if not meta_filepath.exists():
+            return None
+        try:
+            with open(meta_filepath, "r", encoding="utf-8") as f:
+                return json.load(f)
+        except Exception as e:
+            logger.warning(f"Failed to load metadata: {e}")
+            return None
+    
+    def get_history_choices(self) -> List[Tuple[str, str]]:
+        """Get history as choices for Gradio dropdown."""
+        diagrams = self.get_history()
+        return [(f"{d.repo_name} — {d.formatted_timestamp}", d.filename) for d in diagrams]
+    
+    def load_from_history(self, filename: str) -> Optional[str]:
+        """Load a diagram from history.
+        
+        Args:
+            filename: Name of the diagram file
+            
+        Returns:
+            DOT source or None if not found
+        """
+        filepath = self.diagrams_dir / filename
+        if not filepath.exists():
+            return None
+        
+        try:
+            return filepath.read_text(encoding="utf-8")
+        except Exception as e:
+            logger.warning(f"Failed to load diagram: {e}")
+            return None
+    
+    def load_from_history_with_metadata(self, filename: str) -> Tuple[Optional[str], Optional[Dict]]:
+        """Load a diagram and its metadata from history.
+        
+        Args:
+            filename: Name of the diagram file
+            
+        Returns:
+            Tuple of (DOT source, metadata dict) or (None, None) if not found
+        """
+        filepath = self.diagrams_dir / filename
+        if not filepath.exists():
+            return None, None
+        
+        try:
+            dot_source = filepath.read_text(encoding="utf-8")
+            metadata = self._load_metadata(filepath)
+            
+            # If no metadata file, compute node/edge counts from DOT
+            if metadata is None:
+                node_count, edge_count = self._count_nodes_edges(dot_source)
+                metadata = {
+                    "node_count": node_count,
+                    "edge_count": edge_count,
+                }
+            
+            return dot_source, metadata
+        except Exception as e:
+            logger.warning(f"Failed to load diagram: {e}")
+            return None, None
+    
+    def _error_html(self, message: str) -> str:
+        """Generate error display HTML."""
+        return f'''<div style="color:#dc2626; padding:20px; text-align:center;">
+            <strong>⚠️ {message}</strong>
+        </div>'''
+    
+    def _fallback_html(self, content: str, error: str = None) -> str:
+        """Generate fallback HTML when Graphviz is unavailable."""
+        error_msg = f"<p style='color: #dc2626;'>Rendering error: {error}</p>" if error else ""
+        return f'''<div style="background: #f9fafb; padding: 1.5rem; border-radius: 12px;">
+            <strong>📊 Architecture (Text View)</strong>
+            {error_msg}
+            <pre style="background: #fff; padding: 1rem; border-radius: 8px; overflow-x: auto;">{content[:2000]}</pre>
+        </div>'''