CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutSign UpSign In
huggingface

Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.

GitHub Repository: huggingface/notebooks
 Path: blob/main/smolagents_doc/en/pytorch/secure_code_execution.ipynb
Views: 2935
Kernel: Unknown Kernel
# Installation
! pip install smolagents
# To install from source instead of the last release, comment the command above and uncomment the following one.
# ! pip install git+https://github.com/huggingface/smolagents.git

Secure code execution

[!TIP] If you're new to building agents, make sure to first read the intro to agents and the guided tour of smolagents.

Code agents

Multiple research papers have shown that having the LLM write its actions (the tool calls) in code is much better than the current standard format for tool calling, which is across the industry different shades of "writing actions as a JSON of tools names and arguments to use".

Why is code better? Well, because we crafted our code languages specifically to be great at expressing actions performed by a computer. If JSON snippets were a better way, this package would have been written in JSON snippets and the devil would be laughing at us.

Code is just a better way to express actions on a computer. It has better:

  • Composability: could you nest JSON actions within each other, or define a set of JSON actions to re-use later, the same way you could just define a python function?

  • Object management: how do you store the output of an action like generate_image in JSON?

  • Generality: code is built to express simply anything you can have a computer do.

  • Representation in LLM training corpus: why not leverage this benediction of the sky that plenty of quality actions have already been included in LLM training corpus?

This is illustrated on the figure below, taken from Executable Code Actions Elicit Better LLM Agents.

This is why we put emphasis on proposing code agents, in this case python agents, which meant putting higher effort on building secure python interpreters.

Local code execution??

By default, the CodeAgent runs LLM-generated code in your environment.

This is inherently risky, LLM-generated code could be harmful to your environment. One could argue that on the spectrum of agency, code agents give much higher agency to the LLM on your system than other less agentic setups: this goes hand-in-hand with higher risk.

So you need to be mindful of security.

To add a first layer of security, code execution in smolagents is not performed by the vanilla Python interpreter. We have re-built a more secure LocalPythonExecutor from the ground up.

To be precise, this interpreter works by loading the Abstract Syntax Tree (AST) from your Code and executes it operation by operation, making sure to always follow certain rules:

  • By default, imports are disallowed unless they have been explicitly added to an authorization list by the user.

    • Even so, because some innocuous packages like re can give access to potentially harmful packages as in re.subprocess, subpackages that match a list of dangerous patterns are not imported.

  • The total count of elementary operations processed is capped to prevent infinite loops and resource bloating.

  • Any operation that has not been explicitly defined in our custom interpreter will raise an error.

As a result, this interpreter is safer. We have used it on a diversity of use cases, without ever observing any damage to the environment.

However, this solution is certainly not watertight, as no local python sandbox can really be: one could imagine occasions where LLMs fine-tuned for malignant actions could still hurt your environment.

For instance, if you have allowed an innocuous package like Pillow to process images, the LLM could generate thousands of image saves to bloat your hard drive. Other examples of attacks can be found here.

Running these targeted malicious code snippet require a supply chain attack, meaning the LLM you use has been intoxicated.

The likelihood of this happening is low when using well-known LLMs from trusted inference providers, but it is still non-zero.

[!WARNING] The only way to run LLM-generated code securely is to isolate the execution from your local environment.

So if you want to exercise caution, you should use a remote execution sandbox.

Here are examples of how to do it.

Sandbox setup for secure code execution

When working with AI agents that execute code, security is paramount. This guide describes how to set up and use secure sandboxes for your agent applications using either E2B cloud sandboxes or local Docker containers.

E2B setup

Installation

  1. Create an E2B account at e2b.dev

  2. Install the required packages:

pip install 'smolagents[e2b]'

Running your agent in E2B: mono agents

We provide a simple way to use an E2B Sandbox: simply add executor_type="e2b" to the agent initialization, like:

from smolagents import HfApiModel, CodeAgent

agent = CodeAgent(model=HfApiModel(), tools=[], executor_type="e2b")

agent.run("Can you give me the 100th Fibonacci number?")

However, this does not work (yet) with more complicated multi-agent setups.

Running your agent in E2B: multi-agents

To use multi-agents in an E2B sandbox, you need to run your agents completely from within E2B.

Here is how to do it:

from e2b_code_interpreter import Sandbox
import os

# Create the sandbox
sandbox = Sandbox()

# Install required packages
sandbox.commands.run("pip install smolagents")

def run_code_raise_errors(sandbox, code: str, verbose: bool = False) -> str:
    execution = sandbox.run_code(
        code,
        envs={'HF_TOKEN': os.getenv('HF_TOKEN')}
    )
    if execution.error:
        execution_logs = "\n".join([str(log) for log in execution.logs.stdout])
        logs = execution_logs
        logs += execution.error.traceback
        raise ValueError(logs)
    return "\n".join([str(log) for log in execution.logs.stdout])

# Define your agent application
agent_code = """
import os
from smolagents import CodeAgent, HfApiModel

# Initialize the agents
agent = CodeAgent(
    model=HfApiModel(token=os.getenv("HF_TOKEN"), provider="together"),
    tools=[],
    name="coder_agent",
    description="This agent takes care of your difficult algorithmic problems using code."
)

manager_agent = CodeAgent(
    model=HfApiModel(token=os.getenv("HF_TOKEN"), provider="together"),
    tools=[],
    managed_agents=[agent],
)

# Run the agent
response = manager_agent.run("What's the 20th Fibonacci number?")
print(response)
"""

# Run the agent code in the sandbox
execution_logs = run_code_raise_errors(sandbox, agent_code)
print(execution_logs)

Docker setup

Installation

  1. Install Docker on your system

  2. Install the required packages:

pip install 'smolagents[docker]'

Setting up the docker sandbox

Create a Dockerfile for your agent environment:

FROM python:3.10-bullseye

# Install build dependencies
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        build-essential \
        python3-dev && \
    pip install --no-cache-dir --upgrade pip && \
    pip install --no-cache-dir smolagents && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

# Set working directory
WORKDIR /app

# Run with limited privileges
USER nobody

# Default command
CMD ["python", "-c", "print('Container ready')"]

Create a sandbox manager to run code:

import docker
import os
from typing import Optional

class DockerSandbox:
    def __init__(self):
        self.client = docker.from_env()
        self.container = None

    def create_container(self):
        try:
            image, build_logs = self.client.images.build(
                path=".",
                tag="agent-sandbox",
                rm=True,
                forcerm=True,
                buildargs={},
                # decode=True
            )
        except docker.errors.BuildError as e:
            print("Build error logs:")
            for log in e.build_log:
                if 'stream' in log:
                    print(log['stream'].strip())
            raise

        # Create container with security constraints and proper logging
        self.container = self.client.containers.run(
            "agent-sandbox",
            command="tail -f /dev/null",  # Keep container running
            detach=True,
            tty=True,
            mem_limit="512m",
            cpu_quota=50000,
            pids_limit=100,
            security_opt=["no-new-privileges"],
            cap_drop=["ALL"],
            environment={
                "HF_TOKEN": os.getenv("HF_TOKEN")
            },
        )

    def run_code(self, code: str) -> Optional[str]:
        if not self.container:
            self.create_container()

        # Execute code in container
        exec_result = self.container.exec_run(
            cmd=["python", "-c", code],
            user="nobody"
        )

        # Collect all output
        return exec_result.output.decode() if exec_result.output else None


    def cleanup(self):
        if self.container:
            try:
                self.container.stop()
            except docker.errors.NotFound:
                # Container already removed, this is expected
                pass
            except Exception as e:
                print(f"Error during cleanup: {e}")
            finally:
                self.container = None  # Clear the reference

# Example usage:
sandbox = DockerSandbox()

try:
    # Define your agent code
    agent_code = """
import os
from smolagents import CodeAgent, HfApiModel

# Initialize the agent
agent = CodeAgent(
    model=HfApiModel(token=os.getenv("HF_TOKEN"), provider="together"),
    tools=[]
)

# Run the agent
response = agent.run("What's the 20th Fibonacci number?")
print(response)
"""

    # Run the code in the sandbox
    output = sandbox.run_code(agent_code)
    print(output)

finally:
    sandbox.cleanup()

Best practices for sandboxes

These key practices apply to both E2B and Docker sandboxes:

  • Resource management

    • Set memory and CPU limits

    • Implement execution timeouts

    • Monitor resource usage

  • Security

    • Run with minimal privileges

    • Disable unnecessary network access

    • Use environment variables for secrets

  • Environment

    • Keep dependencies minimal

    • Use fixed package versions

    • If you use base images, update them regularly

  • Cleanup

    • Always ensure proper cleanup of resources, especially for Docker containers, to avoid having dangling containers eating up resources.

✨ By following these practices and implementing proper cleanup procedures, you can ensure your agent runs safely and efficiently in a sandboxed environment.