Exploring techniques for untrusted Python code execution in agentic workflows

If you ask me, one of the most compelling use cases for LLMs is the ability to let them invoke tools.

In LLM parlance, tools are functions that the model can invoke at will.

Type of tools

I like to distinguish between (for now) two type of tools the LLM can invoke:

• safe(ish) tools
• unsafe tools

Safe tools are functions written by the agent developer. These functions are (hopefully) well tested and devoid of security risks.

Unsafe tools on the other hand are function completely generated by the LLM, and then invoked by the very same. As we can easily imagine, this poses a number of security implications and risks.

Generating and executing code

Large Language Models (LLMs) have shown promise in automated code generation but typically excel only in simpler tasks such as generating standalone code units. That is, for simpler tasks, such as creating a doc or odt files, LLM are able to generate proper code (although not without guidance).

As of today, proper prompting, or the use of few-shot examples is necessary to have the model generate working code through a trial and error process.

Risks

LLMs can generate and execute their own code. At what cost? Here's a non-exhaustive list of the risks:

• escaping
• data loss
• data leak
• resource exhaustion

Data leaks and resource exhaustion are particulary terrifying.

Resource exhaustion can manifest itself in a number of ways since the model might generate:

• code that gets stuck in infinite loops
• regular expressions vulnerable to catastrophic backtracking

Mitigations

• monitoring
• human supervision
• limiting access to data
• data anonymization
• preemptive testing
• sandboxing

Human supervision is particularly challenging. An agent interacting with non technical people cannot be easily scrutinized for insecure code.

Preemptive testing is a technique which can potentially limit damage. The idea is to test the code generated by the LLM before executing it.

Considering the stakes, above all the mitigations proposed, sandboxing of the generated code is the most promising technique. However, sandboxing in Python is notoriously hard.

Executing code, the naive way

It's easy to fall in the trap to directly execute the code generated by an LLM. Think of this example:

@tool("execute_python_code")
def execute_python_code(code_string):
    """Execute the given Python code string and return the local variables."""

    try:
        global_vars = {}
        local_vars = {}
        exec(code_string, global_vars, local_vars)
        
    except Exception as e:
        print(f"Error executing code: {str(e)}")
        return None

This is highly insecure and risky.

Executing code, better ways

There are a number of ways to execute untrusted (Python in this case) code. While sandboxing Python is pratically impossible (see resources), there are a number of ways to control the potential damage.

HuggingFace Python interpreter tool

This Python interpreter from HuggingFace uses ast under the hood. It has the ability to scan the code string with a great degree of accuracy. It prevents infinite loops and has the ability to restrict imports. Here's an example usage:

from transformers.agents import PythonInterpreterTool

python_repl = PythonInterpreterTool(authorized_imports=["odf", "pathlib"])

def generate_odt(code_string: str):
    """Execute the given Python code to save an odt file by using a Python REPL."""
    return python_repl.forward(code_string)

Azure Container Apps dynamic sessions

One interesting technique for running untrusted Python code is offered by Azure Container Apps dynamic sessions.

This service provides strong isolated sandboxing environments, ideal for running LLM-generated code. Here's an example usage in Langchain:

import os

from langchain_azure_dynamic_sessions import SessionsPythonREPLTool
from langchain_core.tools import tool

from dotenv import load_dotenv

load_dotenv()

os.environ["AZURE_TENANT_ID"] = os.getenv("AZURE_TENANT_ID")
os.environ["AZURE_CLIENT_ID"] = os.getenv("AZURE_CLIENT_ID")
os.environ["AZURE_CLIENT_SECRET"] = os.getenv("AZURE_CLIENT_SECRET")

repl = SessionsPythonREPLTool(
    pool_management_endpoint=os.getenv("POOL_MANAGEMENT_ENDPOINT")
)

@tool("generate-odt")
def generate_odt(code_string: str):
    """Execute the given Python code to save an odt file in a Python REPL."""

    output_dir = Path("files")
    properties = repl.execute(code_string)

    if "status" in properties and properties["status"] == "Success":
        filename = properties["result"]
        try:
            output_dir.mkdir(exist_ok=True)

            repl.download_file(
                remote_file_path=filename, local_file_path=output_dir / filename
            )
            return output_dir / filename
        except HTTPError as e:
            return f"Error downloading file: {str(e)}"

    return properties

In this example we use the Azure Container Apps dynamic sessions REPL to run code generated by the LLM. The tool is able to execute the provided code in a sandboxed environment and then download the resulting file.

RestrictedPython

RestrictedPython helps reducing the risk surface when running untrusted code, but it's not enough alone. For example, AgentRun use a combination of Docker and Restricted Python to execute untrusted code.

Resources

LLM and tools

• Function calling
• Anthropic's comprehensive tool use tutorial

Code execution and sandboxing in Python

• The failure of pysandbox
• PyPy's sandboxing features
• RestrictedPython
• Simon Willison's has a lot of great resources about code execution in the context of LLM systems