File size: 2,119 Bytes
7ccd501 | 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 | # --- Data Science Stack ---
from typing import Optional, Tuple
import uuid
import pandas as pd
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
import datetime as dt
import io
import contextlib
import os
import traceback
from dotenv import load_dotenv
load_dotenv()
# Set Matplotlib to non-interactive mode (server backend)
matplotlib.use('Agg')
def execute_python_code(code: str, csv_url: Optional[str] = None) -> Tuple[Optional[bytes], Optional[str], str]:
"""
Executes Python code and captures the Matplotlib plot into memory bytes.
"""
# 1. Define available libraries
local_scope = {
"pd": pd, "np": np, "plt": plt, "sns": sns, "dt": dt,
"uuid": uuid, "os": os, "csv_url": csv_url
}
# 2. [FIX] Load the CSV into 'df' so the AI code can find it
if csv_url:
try:
# Read the CSV from the URL
df = pd.read_csv(csv_url)
# Inject it into the local_scope with the variable name 'df'
local_scope["df"] = df
except Exception as e:
# Return early if we can't even load the data
return None, f"System Error: Failed to load CSV data. {str(e)}", ""
stdout_capture = io.StringIO()
try:
plt.clf()
plt.close('all')
with contextlib.redirect_stdout(stdout_capture):
# 3. Execute the code (now 'df' is defined in local_scope)
exec(code, {}, local_scope)
# Capture the current figure into a BytesIO buffer
buf = io.BytesIO()
fig = plt.gcf()
# Check if the figure actually has content (axes)
if not fig.get_axes():
return None, "No plot was generated by the code.", stdout_capture.getvalue()
fig.savefig(buf, format='png', bbox_inches='tight')
buf.seek(0)
image_bytes = buf.getvalue()
buf.close()
return image_bytes, None, stdout_capture.getvalue()
except Exception:
return None, traceback.format_exc(), stdout_capture.getvalue() |