Secure code execution in LlamaIndex with Azure Container Apps dynamic sessions

One of the many amazing feats that LLMs are capable of is generating executable code. This can be used to solve a variety of complex problems that require calculations and fixed logic that traditional computing excels at but LLMs can struggle to perform directly. When building agents to perform complex tasks, equipping your agent with code execution as an available tool can be a powerful strategy.

However, this strategy comes with a major drawback: executable code can be flawed or even dangerous to execute, and detecting whether code will be problematic prior to executing it is arguably an expression of the Halting Problem, making it impossible to guarantee success at detection.

The solution is sandboxing, to isolate potentially problematic code from the host environment. Now, thanks to dynamic sessions in Azure Container Apps, the ability to execute sandboxed code generated by an LLM is simple directly from LlamaIndex. It’s implemented as a tool that can be used by any LlamaIndex agent.

In this blog post we’ll show you exactly how to use the new Azure Code Interpreter tool and walk you through a couple of examples of how to make the most of it. You can see the full code in this notebook and read more in the tool documentation on LlamaHub and on learn.microsoft.com.

Set up Azure Container Apps dynamic sessions

First, install our python packages including the tool:

pip install llama-index
pip install llama-index-llms-azure
pip install llama-index-tools-azure-code-interpreter

In the notebook we’re using GPT 3.5 Turbo hosted on Azure as the LLM, but you can use any LLM capable of tool use:

from llama_index.llms.azure_openai import AzureOpenAI
llm = AzureOpenAI(
    model="gpt-35-turbo",
    deployment_name="gpt-35-deploy",
    api_key=api_key,
    azure_endpoint=azure_endpoint,
    api_version=api_version,
)

Once you’ve got your LLM set up, you’ll need to create a session pool to host your executions. Doing this will give you a pool management endpoint URL that you can provide to LlamaIndex like this:

# Import the AzureCodeInterpreterToolSpec from llama_index
from llama_index.tools.azure_code_interpreter import (
    AzureCodeInterpreterToolSpec,
)

# Create the AzureCodeInterpreterToolSpec with the pool_managment_endpoint set to your session management endpoint
# It is optional to set the local_save_path, but it is recommended to set it to a path where the tool can automatically save any intermediate data generated from Python code's output.
azure_code_interpreter_spec = AzureCodeInterpreterToolSpec(
    pool_managment_endpoint="your-pool-management-endpoint",
    local_save_path="local-file-path-to-save-intermediate-data",
)

This sets up a tool ready to be used with LlamaIndex. You’re now ready to set up your agent:

# Import the ReActAgent
from llama_index.core.agent import ReActAgent

# Create the ReActAgent and inject the tools defined in the AzureDynamicSessionsToolSpec
agent = ReActAgent.from_tools(
    azure_code_interpreter_spec.to_tool_list(), llm=llm, verbose=True
)

In this example we’re providing only a single tool, but you could provide any other tools you like to your ReAct agent. Now you’ve got an agent, you’re ready to ask it to perform tasks!

Dynamic sessions code interpreter in action

In our first example, we’re going to ask the agent the time in Seattle. This is usually a tricky task for LLMs, which don’t know what time it is anywhere!

# Test the agent with simple answers that could leverage Python codes
print(agent.chat("Tell me the current time in Seattle."))

The agent generates python code to determine the time and convert it to the correct time zone. It passes this code to Azure Container Apps dynamic sessions, which execute the code and return the answer:

Thought: To provide the current time in Seattle, I need to calculate it based on the current UTC time and adjust for Seattle's time zone, which is Pacific Daylight Time (PDT) during daylight saving time and Pacific Standard Time (PST) outside of daylight saving time. PDT is UTC-7, and PST is UTC-8. I can use the code interpreter tool to get the current UTC time and adjust it accordingly.
Action: code_interpreter
Action Input: {'python_code': "from datetime import datetime, timedelta; import pytz; utc_now = datetime.now(pytz.utc); seattle_time = utc_now.astimezone(pytz.timezone('America/Los_Angeles')); seattle_time.strftime('%Y-%m-%d %H:%M:%S %Z%z')"}
Observation: {'$id': '1', 'status': 'Success', 'stdout': '', 'stderr': '', 'result': '2024-05-04 13:54:09 PDT-0700', 'executionTimeInMilliseconds': 120}
Thought: I can answer without using any more tools. I'll use the user's language to answer.
Answer: The current time in Seattle is 2024-05-04 13:54:09 PDT.
The current time in Seattle is 2024-05-04 13:54:09 PDT.

You can also use the tool to safely inspect and manipulate data, as in this example where we ask it to open a CSV file and answer questions about it:

# Upload a sample temperature file of a day in Redmond Washington and ask a question about it
res = azure_code_interpreter_spec.upload_file(
    local_file_path="./TemperatureData.csv"
)
if len(res) != 0:
    print(
        agent.chat("Find the highest temperature in the file that I uploaded.")
    )

It doesn’t just read data from the CSV, it performs math on it to determine the highest temperature:

Thought: I need to use the list_files tool to get the metadata for the uploaded file, and then use python to read the file and find the highest temperature.
Action: list_files
Action Input: {}
Observation: [RemoteFileMetadata(filename='TemperatureData.csv', size_in_bytes=514, file_full_path='/mnt/data/TemperatureData.csv')]
Thought: I have the metadata for the file. I need to use python to read the file and find the highest temperature.
Action: code_interpreter
Action Input: {'python_code': "import csv\n\nwith open('/mnt/data/TemperatureData.csv', 'r') as f:\n    reader = csv.reader(f)\n    next(reader)\n    highest_temp = float('-inf')\n    for row in reader:\n        temp = float(row[1])\n        if temp > highest_temp:\n            highest_temp = temp\nprint(highest_temp)"}
Observation: {'$id': '1', 'status': 'Success', 'stdout': '12.4\n', 'stderr': '', 'result': '', 'executionTimeInMilliseconds': 26}
Thought: I have the highest temperature. I can answer the question.
Answer: The highest temperature in the file is 12.4 degrees.
The highest temperature in the file is 12.4 degrees.

The tool can also safely modify your uploaded files, as in this example where we get it to sort the data in the CSV:

# Ask the LLM to perform modifications on the sample temperature file
print(
    agent.chat(
        "Rearrange the temperature data in a descending order and save it back to the original csv file."
    )
)

In this example, you’ll see that the LLM initially produces code with a syntax error, but the agent is smart enough to figure out what went wrong and correct itself:

Thought: To proceed, I need to first list the files in the session pool to identify the uploaded temperature data file.
Action: list_files
Action Input: {}
Observation: [RemoteFileMetadata(filename='TemperatureCurve.png', size_in_bytes=35512, file_full_path='/mnt/data/TemperatureCurve.png'), RemoteFileMetadata(filename='TemperatureData.csv', size_in_bytes=514, file_full_path='/mnt/data/TemperatureData.csv')]
Thought: The temperature data file is identified as 'TemperatureData.csv'. I need to read its contents, rearrange the temperature data in descending order, and then save the changes back to the original CSV file.
Action: code_interpreter
Action Input: {'python_code': "import pandas as pd\n\ndf = pd.read_csv('/mnt/data/TemperatureData.csv')\ndf_sorted = df.sort_values(by='Temperature', ascending=False)\ndf_sorted.to_csv('/mnt/data/TemperatureData.csv', index=False)\n'filename': 'TemperatureData.csv'"}
Observation: {'$id': '1', 'status': 'Failure', 'stdout': '', 'stderr': "illegal target for annotation (3407605186.py, line 6)  Cell In[2363], line 6\n    'filename': 'TemperatureData.csv'\n    ^\nSyntaxError: illegal target for annotation\n\n", 'result': None, 'executionTimeInMilliseconds': 9}
Thought: There was a syntax error in the Python code due to an incorrect line. I need to correct the code and execute it again to sort the temperature data and save it back to the CSV file.
Action: code_interpreter
Action Input: {'python_code': "import pandas as pd\n\ndf = pd.read_csv('/mnt/data/TemperatureData.csv')\ndf_sorted = df.sort_values(by='Temperature', ascending=False)\ndf_sorted.to_csv('/mnt/data/TemperatureData.csv', index=False)"}
Observation: {'$id': '1', 'status': 'Success', 'stdout': '', 'stderr': '', 'result': '', 'executionTimeInMilliseconds': 91}
Thought: I can answer without using any more tools. I'll use the user's language to answer.
Answer: The temperature data has been successfully rearranged in descending order and saved back to the original CSV file, 'TemperatureData.csv'.
The temperature data has been successfully rearranged in descending order and saved back to the original CSV file, 'TemperatureData.csv'.

Modifying files would not be useful if you couldn’t retrieve them after modification, which is done like this:

# Download the modified file
azure_code_interpreter_spec.download_file_to_local(
    remote_file_path="TemperatureData.csv",
    local_file_path="/.../SortedTemperatureData.csv",
)

Endless possibilities

The scope of tasks that you can achieve with sandboxed code execution is as broad as programming itself, and having safe execution guaranteed allows you to confidently hand agents tasks that previously you might have been hesitant about. We think this is an amazing addition to our LLM agent capabilities and we’re excited to see what you build with it.