Build a Code Execution Agent with E2B

import warnings
warnings.filterwarnings('ignore')

import os
import base64
import pandas as pd
import time

from dotenv import load_dotenv
load_dotenv()

from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.tools import tool
from langchain.agents import create_agent
from langchain.messages import HumanMessage
from langgraph.checkpoint.memory import InMemorySaver

from e2b_code_interpreter import Sandbox

model = ChatGoogleGenerativeAI(model='gemini-3-flash-preview')
checkpointer = InMemorySaver()

sbx = Sandbox.create(timeout=40*60)

response = sbx.run_code('print(2+2)')
response = sbx.run_code('print(2/0)')
response

Execution(Results: [], Logs: Logs(stdout: [], stderr: []),
          Error: ExecutionError(name='ZeroDivisionError', value='division by zero', ...))

code = """import pandas as pd
import matplotlib.pyplot as plt

data = {
    'product': ['Laptop', 'Mouse', 'Keyboard', 'Monitor', 'Headphones',
                'Laptop', 'Mouse', 'Keyboard', 'Monitor', 'Laptop'],
    'quantity': [5, 15, 10, 3, 8, 4, 20, 12, 2, 6],
    'price': [1200, 25, 75, 350, 150, 1200, 25, 75, 350, 1200]
}

df = pd.DataFrame(data)
df['revenue'] = df['quantity'] * df['price']
product_revenue = df.groupby('product')['revenue'].sum().sort_values(ascending=False)

print("Total Revenue by Product:")
print(product_revenue)

plt.figure(figsize=(10, 6))
product_revenue.plot(kind='bar', color='steelblue')
plt.title('Total Revenue by Product', fontsize=16, fontweight='bold')
plt.xlabel('Product'); plt.ylabel('Revenue ($)')
plt.xticks(rotation=45)
plt.tight_layout()"""

response = sbx.run_code(code)
response

Execution(Results: [Result(<Figure size 1000x600 with 1 Axes>)],
          Logs: Logs(stdout: ['Total Revenue by Product:\nproduct\nLaptop  18000\nMonitor  1750\n...'], stderr: []),
          Error: None)

def get_dataset_info(file_path):
    if file_path.endswith('.csv'):
        df = pd.read_csv(file_path, nrows=3)
    else:
        df = pd.read_excel(file_path, nrows=3)
    return f"Columns: {list(df.columns)}\n Sample data:\n{df.to_string()}"

@tool
def run_python_code(code: str):
    """Execute Python code in the E2B sandbox.

    Args:
        code: Valid executable Python code only.

    Returns:
        Execution result, plus paths to any saved charts.
    """
    print('Running code in sandbox....')
    execution = sbx.run_code(code)
    print('Code execution is done!')

    if execution.error:
        return f"Error: {execution.error.name}\nValue: {execution.error.value}"

    os.makedirs('images', exist_ok=True)
    output = []
    timestamp = int(time.time())
    output.append(str(execution))

    for idx, result in enumerate(execution.results):
        if result.png:
            filename = f'images/{timestamp}_chart_{idx}.png'
            with open(filename, 'wb') as f:
                f.write(base64.b64decode(result.png))
            output.append(f"Chart saved to {filename}")

    return "\n".join(output) if output else "Code executed but no output was returned"

CODE_EXECUTION_PROMPT = """You are a data analysis assistant. You MUST use the available tools to complete tasks.

WORKFLOW - Follow these steps in order:
1. Search for data files using glob_search (for LOCAL file discovery only)
2. Upload file using upload_file (transfers from local to sandbox)
3. ANALYZE THE DATASET FIRST - Use run_python_code to check shape, columns,
   data types, first rows, null values, and df.describe()
4. PERFORM ANALYSIS - Use run_python_code multiple times to clean, aggregate, and answer
5. CREATE VISUALIZATIONS (if requested) - matplotlib with titles/labels, use plt.show()

CRITICAL RULES:
- You MUST call the appropriate tool for each step - do not just think, ACT by calling tools
- NEVER skip the dataset exploration step
- Use run_python_code multiple times rather than one large code block
- All file paths in code must use '/home/user/data/' prefix"""

from scripts import prompts
from langchain.agents.middleware import FilesystemFileSearchMiddleware

agent = create_agent(
    model=model,
    tools=[upload_file, run_python_code],
    system_prompt=prompts.CODE_EXECUTION_PROMPT,
    checkpointer=checkpointer,
    middleware=[
        FilesystemFileSearchMiddleware(
            root_path="./data",
            use_ripgrep=True,
            max_file_size_mb=1000)
    ]
)