Upload chart_executor.py

utils/chart_executor.py

@@ -0,0 +1,170 @@
+"""
+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