File size: 5,497 Bytes
5af7d8e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | """
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
|