The Code Interpreter SDK allows you to run AI-generated code in a secure small VM - E2B sandbox - made for AI code execution.
Inside the sandbox is a Jupyter server running that you can control from our SDK through the notebook.execCell() method.Check out the SDK’s repository on GitHub.
Copy
Ask AI
pip install e2b_code_interpreter python-dotenv
Get your E2B API key here and save it to .env in your root directory.
We’ll be using Anthropic’s Claude 3 Opus model but E2B works with any LLM so feel free to pick any you want!
Usually, all you need from the model is just support for tool use. If the LLM doesn’t support tool use, you can ask the LLM to respond with Markdown or XML and parse the LLM’s output on your own. Then just pass the parsed code from code blocks to the code interpreter.
Create the model.py file and paste the following code.
Copy
Ask AI
# model.pyMODEL_NAME = "claude-3-opus-20240229"SYSTEM_PROMPT = """## your job & contextyou are a python data scientist. you are given tasks to complete and you run python code to solve them.- the python code runs in jupyter notebook.- every time you call `execute_python` tool, the python code is executed in a separate cell. it's okay to multiple calls to `execute_python`.- display visualizations using matplotlib or any other visualization library directly in the notebook. don't worry about saving the visualizations to a file.- you have access to the internet and can make api requests.- you also have access to the filesystem and can read/write files.- you can install any pip package (if it exists) if you need to but the usual packages for data analysis are already preinstalled.- you can run any python code you want, everything is running in a secure sandbox environment."""tools = [ { "name": "execute_python", "description": "Execute python code in a Jupyter notebook cell and returns any result, stdout, stderr, display_data, and error.", "input_schema": { "type": "object", "properties": { "code": { "type": "string", "description": "The python code to execute in a single cell." } }, "required": ["code"] } }]
This defines our system prompt and the tools dictionary with available tools for the LLM - namely the "execute_python" tool.
A little bit later, we’ll connect "execute_python" to the E2B’s code interpreter.
We’ll create a new function called code_interpret() in a separate file code_interpreter.py.
Copy
Ask AI
# code_interpreter.pyfrom e2b_code_interpreter import CodeInterpreterdef code_interpret(code_interpreter: CodeInterpreter, code: str): print(f"\n{'='*50}\n> Running following AI-generated code:\n{code}\n{'='*50}") exec = code_interpreter.notebook.exec_cell( code, # You can stream logs from the code interpreter # on_stderr=lambda stderr: print("\n[Code Interpreter stdout]", stderr), # on_stdout=lambda stdout: print("\n[Code Interpreter stderr]", stdout), # # You can also stream additional results like charts, images, etc. # on_result=... ) if exec.error: print("[Code Interpreter error]", exec.error) # Runtime error else: return exec.results, exec.logs
This function takes the CodeInterpreter object from our SDK and code as paramaters.
The code parameter is the code generated by the LLM.Inside the function, we call the code_interpreter.notebook.exec_cell() method. The exec_cell() takes the code argument and executes this code inside E2B sandbox.
Now we’ll put everything together.
Create the main.py file, import dependencies, and create the chat() function that will do the LLM calling and tool parsing.
Now we are ready to run our program. At the end of main.py add the following code prompting the LLM to visualize a distribution of the height of men and print the median.
Copy
Ask AI
# main.pydef main(): user_message = "Visualize a distribution of height of men based on the latest data you know. Also, print the median value." # Create the CodeInterpreter object and save it as code_interpreter with CodeInterpreter() as code_interpreter: code_interpreter_results, code_interpreter_logs = chat( code_interpreter, user_message, ) print(code_interpreter_logs) first_result= code_interpreter_results[0] print(first_result)
After running your code with the following command
Copy
Ask AI
$ python main.py
you should see results similar to this:
Copy
Ask AI
stdout=['The median male height is 175.5 cm\n'] stderr=[]<Figure size 800x400 with 1 Axes>
We got our median in the logs (stdout, and stderr) but we also something intering in first_result.
This looks like a plot. Let’s save it to a file. Add the following to the end of main.py and run the code again with python main.py in your terminal.
Copy
Ask AI
# main.py # If we received a chart in PNG form, we can visualize it if first_result.png: # Decode the base64 encoded PNG data png_data = base64.b64decode(first_result.png) # Generate a unique filename for the PNG filename = f"chart.png" # Save the decoded PNG data to a file with open(filename, "wb") as f: f.write(png_data) print(f"Saved chart to {filename}")
The chart was saved in the chart.png file and it should look similar to this: