utils/chart_executor.py

"""
Code Executor for LLM-Generated Visualization Code

Safely executes matplotlib code generated by LLM to produce charts.
Converts matplotlib figures to Base64 for HTML embedding.
"""

import sys
import io
import base64
import traceback
from typing import Dict, Optional, Any
import warnings


class ChartExecutor:
    """
    Executes LLM-generated matplotlib code and returns Base64 images.
    """
    
    def __init__(self):
        self._setup_matplotlib()
    
    def _setup_matplotlib(self):
        """Setup matplotlib for non-interactive use."""
        try:
            import matplotlib
            matplotlib.use('Agg')
            import matplotlib.pyplot as plt
            plt.rcParams['font.family'] = ['Microsoft YaHei', 'SimHei', 'DejaVu Sans', 'sans-serif']
            plt.rcParams['axes.unicode_minus'] = False
            self.plt = plt
            self.np = __import__('numpy')
            self.available = True
        except ImportError as e:
            warnings.warn(f"Matplotlib not available: {e}")
            self.available = False
    
    def execute_chart_code(self, code: str, chart_name: str = "chart") -> Dict[str, Any]:
        """
        Execute matplotlib code and return Base64 image.
        
        Args:
            code: Python code string containing matplotlib commands
            chart_name: Name for the chart
            
        Returns:
            Dict with 'success', 'image_base64' or 'error'
        """
        if not self.available:
            return {
                "success": False,
                "error": "Matplotlib not available",
                "image_base64": None
            }
        
        try:
            # Create isolated namespace
            namespace = {
                'plt': self.plt,
                'np': self.np,
                'matplotlib': __import__('matplotlib'),
            }
            
            # Redirect savefig to capture output
            buf = io.BytesIO()
            
            # Modify code to save to buffer instead of file
            modified_code = self._modify_code_for_buffer(code)
            
            # Execute code
            exec(modified_code, namespace)
            
            # Get the current figure
            fig = self.plt.gcf()
            
            # Save to buffer
            fig.savefig(buf, format='png', dpi=150, bbox_inches='tight',
                       facecolor='white', edgecolor='none')
            buf.seek(0)
            
            # Convert to Base64
            img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
            
            # Clean up
            self.plt.close('all')
            
            return {
                "success": True,
                "image_base64": f"data:image/png;base64,{img_base64}",
                "chart_name": chart_name
            }
            
        except Exception as e:
            self.plt.close('all')
            return {
                "success": False,
                "error": str(e),
                "traceback": traceback.format_exc(),
                "image_base64": None
            }
    
    def _modify_code_for_buffer(self, code: str) -> str:
        """Remove savefig and show commands from code."""
        lines = code.split('\n')
        modified_lines = []
        
        for line in lines:
            # Skip savefig and show commands
            stripped = line.strip()
            if stripped.startswith('plt.savefig') or stripped.startswith('plt.show'):
                continue
            if 'savefig(' in stripped or '.show(' in stripped:
                continue
            modified_lines.append(line)
        
        return '\n'.join(modified_lines)
    
    def execute_multiple_charts(self, code_dict: Dict[str, str]) -> Dict[str, Dict]:
        """
        Execute multiple chart codes.
        
        Args:
            code_dict: Dict mapping chart names to code strings
            
        Returns:
            Dict mapping chart names to execution results
        """
        results = {}
        
        for name, code in code_dict.items():
            if code and isinstance(code, str) and 'plt' in code:
                results[name] = self.execute_chart_code(code, name)
            else:
                results[name] = {
                    "success": False,
                    "error": "Invalid or empty code",
                    "image_base64": None
                }
        
        return results


def generate_fallback_chart_html(chart_name: str, description: str) -> str:
    """
    Generate a placeholder HTML when chart generation fails.
    """
    return f'''
    <div style="border:2px dashed #ccc; padding:40px; text-align:center; 
                background:#f9f9f9; border-radius:8px; margin:15px 0;">
        <div style="font-size:48px; color:#ccc;">📊</div>
        <div style="font-size:16px; color:#666; margin-top:10px;">
            {chart_name}
        </div>
        <div style="font-size:12px; color:#999; margin-top:5px;">
            {description}
        </div>
    </div>
    '''


# Singleton instance
_executor = None

def get_executor() -> ChartExecutor:
    """Get or create the chart executor singleton."""
    global _executor
    if _executor is None:
        _executor = ChartExecutor()
    return _executor