Code Interpreter Agent

import sys
import io
import contextlib
from typing import Optional

from openai import OpenAI
from pydantic import BaseModel, Field
from tensorlake import application, function, Image

# Define the execution environment
# We install pandas and numpy to support data analysis tasks
image = Image(name="code-interpreter").run("pip install openai pandas numpy")

class ExecutionResult(BaseModel):
    stdout: str
    stderr: str
    result: Optional[str] = None

@application()
@function(image=image, secrets=["OPENAI_API_KEY"])
def code_interpreter(prompt: str) -> str:
    """
    Main entry point for the Code Interpreter Agent.
    """
    client = OpenAI()
    
    # Step 1: Generate code based on user prompt
    completion = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "You are a Python code generator. Output only valid Python code to solve the user's problem. Do not include markdown blocks."},
            {"role": "user", "content": prompt}
        ]
    )
    code = completion.choices[0].message.content
    
    # Step 2: Execute the generated code securely
    print(f"Executing generated code:
{code}")
    execution_result = execute_python_code(code)
    
    # Step 3: formulate final answer
    final_response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "Answer the user's question based on the code execution result."},
            {"role": "user", "content": f"Question: {prompt}
Code Output:
{execution_result.stdout}
Errors:
{execution_result.stderr}"}
        ]
    )
    
    return final_response.choices[0].message.content

@function(image=image)
def execute_python_code(code: str) -> ExecutionResult:
    """
    Executes Python code in a secure sandbox and captures output.
    """
    stdout = io.StringIO()
    stderr = io.StringIO()
    
    try:
        # Redirect stdout and stderr to capture print statements
        with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr):
            exec(code, {"__name__": "__main__"})
    except Exception as e:
        print(f"Execution error: {e}", file=stderr)
        
    return ExecutionResult(
        stdout=stdout.getvalue(),
        stderr=stderr.getvalue()
    )

if __name__ == "__main__":
    from tensorlake.applications import run_local_application
    
    # Run the application locally
    result = run_local_application(code_interpreter, "Calculate the sum of the first 50 prime numbers.")
    print(f"Code Interpreter Output:\n{result}")

export OPENAI_API_KEY=your_key_here
python app.py

tl secrets set OPENAI_API_KEY=your_key_here
tl app deploy app.py