Spaces:
Running
Running
File size: 4,041 Bytes
09801ca | 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 | import os
import uuid
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
from typing import Optional
from core.llm import chat
import logging
import ast
logger = logging.getLogger(__name__)
# Ensure non-interactive backend for matplotlib
matplotlib.use('Agg')
STATIC_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'static')
CHARTS_DIR = os.path.join(STATIC_DIR, 'charts')
os.makedirs(CHARTS_DIR, exist_ok=True)
def generate_image_dashboard(query: str, df: pd.DataFrame) -> Optional[str]:
"""
Generates a literal .png dashboard image using matplotlib and seaborn,
driven by LLM code generation. Returns the markdown image link.
"""
try:
# Limit rows to 100 for code execution context to keep it fast
sample_df = df.head(100)
# Save temp CSV for the generated code to read safely
temp_csv_path = os.path.join(CHARTS_DIR, f"temp_{uuid.uuid4().hex}.csv")
# Replace backslashes for python code block
temp_csv_path_safe = temp_csv_path.replace("\\", "/")
sample_df.to_csv(temp_csv_path, index=False)
# Output PNG path
img_id = uuid.uuid4().hex
img_filename = f"dashboard_{img_id}.png"
img_filepath = os.path.join(CHARTS_DIR, img_filename)
img_filepath_safe = img_filepath.replace("\\", "/")
# Prompt LLM to write ONLY pure Python code
prompt = f"""You are a Silicon Valley Data Engineer.
The user wants an image dashboard for this query: "{query}"
Write Python code using `matplotlib.pyplot` and `seaborn` to generate a beautiful, dark-mode data visualization.
The data is available at this local path: `{temp_csv_path_safe}`
The output image MUST be saved to: `{img_filepath_safe}`
RULES:
1. ONLY output valid Python code. NO markdown formatting, NO explanations.
2. Use dark mode aesthetics: `plt.style.use('dark_background')` and sleek colors (cyan, purple, etc.).
3. Handle missing values natively.
4. Call `plt.savefig('{img_filepath_safe}', bbox_inches='tight', dpi=150, facecolor='#111827')` at the end.
5. Call `plt.close('all')` at the very end.
6. Make it look like a premium dashboard.
Data Summary:
Columns: {list(df.columns)}
Shape: {df.shape}
Sample Data:
{df.head(3).to_string()}
"""
# Generate code
response = chat(prompt, temperature=0.1, max_tokens=2000)
# Clean response
code = response.strip()
if '```python' in code:
code = code.split('```python')[1].split('```')[0]
elif '```' in code:
code = code.split('```')[1].split('```')[0]
code = code.strip()
# Execute the code safely in this sandbox context
local_vars = {}
global_vars = {
'pd': pd,
'plt': plt,
'sns': sns,
'__builtins__': __builtins__
}
try:
exec(code, global_vars, local_vars)
except Exception as exec_err:
logger.error(f"Image Agent code execution failed: {exec_err}")
return f"\n\n> [!WARNING]\n> **Image Generation Failed**\n> The AI made a syntax error while drawing the chart: `{exec_err}`\n\n"
# Check if the image was actually created
if os.path.exists(img_filepath):
# Clean up temp csv
if os.path.exists(temp_csv_path):
os.remove(temp_csv_path)
return f"\n\n\n\n"
else:
logger.error("Image Agent code executed but no image was saved.")
return f"\n\n> [!WARNING]\n> **Image Generation Failed**\n> The AI executed successfully but failed to save the image file to `{img_filepath}`.\n\n"
except Exception as e:
logger.error(f"Image Agent failed: {e}")
return f"\n\n> [!WARNING]\n> **Image Generation Failed**\n> An unexpected error occurred: `{e}`\n\n"
|