OPENAI_API_KEY=your_openai_api_key
E2B_API_KEY=your_e2b_api_key

%pip install openai=1.45.1 e2b_code_interpreter=0.0.10 python-dotenv==1.0.1

import os
from dotenv import load_dotenv
load_dotenv()

# Get your API keys
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
E2B_API_KEY = os.getenv("E2B_API_KEY")

O1_PROMPT = """You're a data scientist analyzing survival data from the Titanic Disaster. You are given tasks to complete and you run Python code to solve them.

Information about the Titanic dataset:
- It's in the `/home/user/train.csv` and `/home/user/test.csv` files
- The CSV files use `,` as the delimiter
- They have the following columns:
  - PassengerId: Unique passenger ID
  - Pclass: 1st, 2nd, 3rd (Ticket class)
  - Name: Passenger name
  - Sex: Gender
  - Age: Age in years
  - SibSp: Number of siblings/spouses aboard
  - Parch: Number of parents/children aboard
  - Ticket: Ticket number
  - Fare: Passenger fare
  - Cabin: Cabin number
  - Embarked: Port of Embarkation (C = Cherbourg, Q = Queenstown, S = Southampton)

Generally, you follow these rules:
- **Always format your response in Markdown**
- **Always respond only with code in code blocks like this:**
```python
{code}
```
- The Python code runs in a Jupyter notebook.
- Every time you generate Python code, it is executed in a separate cell.
- Display visualizations using `matplotlib` or any other visualization library directly in the notebook.
- You have access to the internet and can make API requests.
- You also have access to the filesystem and can read/write files.
- Install all packages before using them by running `!pip install {package}`.
- You can run any Python code you want; everything is running in a secure sandbox environment.
  """

GPT_4O_PROMPT = """You are an expert software engineer that receives an execution plan, and then creates a single Python script that does everything in the plan. It will be executed in a single Python notebook cell."""

import re

def match_code_blocks(llm_response):
	  pattern = re.compile(r'```python\n(.*?)\n```', re.DOTALL)  # Match everything between ```python and ```
    matches = pattern.findall(llm_response)
    if matches:
        code = "\n".join(matches)
        print("> LLM-generated code:")
        print(code)
        return code
    return ""

from openai import OpenAI

def chat(e2b_code_interpreter, user_message):
    print(f"\n{'='*50}\nUser message: {user_message}\n{'='*50}")
    
    client = OpenAI(api_key=OPENAI_API_KEY)
    
    # Get the plan from o1-mini
    response_o1 = client.chat.completions.create(
        model="o1-mini",
        messages=[
            {"role": "user", "content": O1_PROMPT},
            {"role": "user", "content": user_message}
        ]
    )
    content_o1 = response_o1.choices[0].message.content
    
    # Use gpt-4o-mini to extract the final code
    response_4o = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": GPT_4O_PROMPT},
            {"role": "user", "content": f"This is the plan I received, please write fully functional code that I can run in one notebook cell, and list all its dependencies: {content_o1}"}
        ]
    )
    content_4o = response_4o.choices[0].message.content
    python_code = match_code_blocks(content_4o)
    
    if python_code != "":
        code_interpreter_results = code_interpret(e2b_code_interpreter, python_code)
        return code_interpreter_results
    else:
        print(f"Failed to match any Python code in model's response: {content_4o}")
        return []

def code_interpret(e2b_code_interpreter, code):
    print("Running code interpreter...")
    exec = e2b_code_interpreter.notebook.exec_cell(
        code,
        on_stderr=lambda stderr: print("[Code Interpreter]", stderr),
        on_stdout=lambda stdout: print("[Code Interpreter]", stdout),
    )
    
    if exec.error:
        print("[Code Interpreter ERROR]", exec.error)
    else:
        return exec.results

def upload_dataset(code_interpreter):
    print("Uploading testing and training datasets to Code Interpreter sandbox...")
    with open("./test.csv", "rb") as f:
        remote_path = code_interpreter.upload_file(f)
        print("Uploaded test.csv at", remote_path)
    with open("./train.csv", "rb") as f:
        remote_path = code_interpreter.upload_file(f)
        print("Uploaded train.csv at", remote_path)

from e2b_code_interpreter import CodeInterpreter

with CodeInterpreter(api_key=E2B_API_KEY) as code_interpreter:
    # Upload the dataset to the code interpreter sandbox
    upload_dataset(code_interpreter)
    
    code_results = chat(
        code_interpreter,
        "Clean the data, train a decision tree to predict the survival of passengers, and visualize the learning curve. Then run the model on the test dataset and print the results."
    )
    if code_results:
        first_result = code_results[0]
    else:
        raise Exception("No code interpreter results")

Uploading testing and training datasets to Code Interpreter sandbox...
Uploaded at /home/user/train.csv

==================================================
User message: Clean the data, train a decision tree to predict the survival of passengers, and visualize the learning curve.
==================================================
> LLM-generated code:
# Install necessary packages
!pip install pandas numpy matplotlib seaborn scikit-learn

# Import libraries
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import train_test_split, learning_curve
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import accuracy_score

# Load the data (Make sure to replace these paths with the correct ones)
train = pd.read_csv('/home/user/train.csv')
test = pd.read_csv('/home/user/test.csv')

# Display first few rows
print(train.head())

# Data Cleaning
# Combine train and test for consistent preprocessing
combined = pd.concat([train, test], sort=False)
# ... [code continues] ...