e2b-dev / code-interpreter

Python & JS/TS SDK for running AI-generated code/code interpreting in your AI app
https://e2b.dev
Apache License 2.0
1.28k stars 96 forks source link

Extract data from DataFrames #33

Closed jakubno closed 3 months ago

jakubno commented 4 months ago

Description

MVP for easy way to extract data from pandas DataFrames.

Steps to try it out (Python only)

  1. Install E2B e2b-code-interpreter==0.0.11b2
  2. You'll be able to access the df variable in the code through the .data property on the result of the execution. As long as the last line is a dataframe, you'll be able to access it

Here's a code snippet to try this

import asyncio

from e2b_code_interpreter import AsyncCodeInterpreter
from dotenv import load_dotenv

load_dotenv()
CODE = """
import pandas
data = {  
  "food": ["bacon", "pulled pork", "oats"],
  "type": ["meat", "meat", "vegetarian"],
  "calories": [420, 380, 390],
}
df = pd.DataFrame(data)

df
"""

CODE2 = """
df2 = df.groupby(["type"]).agg({'calories': 'mean'})
df2
"""

async def main():
    sandbox = await AsyncCodeInterpreter.create(timeout=180)
    execution = await sandbox.notebook.exec_cell(CODE)
    print(execution.results[0].data)
    # {'food': ['bacon', 'pulled pork', 'oats'], 'type': ['meat', 'meat', 'vegetarian'], 'calories': [420, 380, 390]} 
    execution = await sandbox.notebook.exec_cell(CODE2)
    print(execution.results[0].data)
    # {'calories': [400.0, 390.0]}

    await sandbox.kill()

asyncio.run(main())