csv_chart_service.py

# --- 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()