Merge pull request #16 from VectorInstitute/e2b_template

Added option to override E2B template.

Author: jacobthebanana

e2b.Dockerfile

@@ -0,0 +1,16 @@
+FROM e2bdev/code-interpreter:latest
+
+# All downloaded files will be available under /data
+WORKDIR /data
+
+# # Example: download given link to publicly-shared Google Drive file and unzip.
+# # Set permission to open for the sandbox user.
+# RUN python3 -m pip install gdown && \
+# python3 -m gdown -O local_sqlite.zip "EXAMPLE_FILE_ID" && \
+# unzip local_sqlite.zip && \
+# chmod -R a+wr /data && \
+# rm -v local_sqlite.zip
+
+# Example: download file from public URL- e.g., HuggingFace, or Github
+RUN wget --content-disposition \
+"https://huggingface.co/datasets/vector-institute/hotpotqa/resolve/main/data/validation-00000-of-00001.parquet"

e2b_template.md

@@ -0,0 +1,64 @@
+# Custom E2B Template
+
+The sandboxes are non-persistent. Each time your agent runs a command, all data files need to be uploaded again. If your data files are large (10MB total or more), it might take minutes before any command starts running. For performance reasons, you should consider bundling these data files into a custom E2B template. Here are the quick steps for setting that up.
+
+Estimated time: 30 minutes.
+
+## Prerequisites
+
+Install Docker Engine ([Docker CE](https://docs.docker.com/engine/install/)).
+
+Install Node Version Manager ([NVM](https://github.com/nvm-sh/nvm)).
+
+Install the latest LTS Node version (22 as of 2025AUG03.)
+
+Install the E2B utils and log into your account:
+
+```bash
+npm i -g @e2b/cli
+e2b auth login
+```
+
+## Steps
+
+Modify `e2b.Dockerfile`. Replace the URL with the link to your data files.
+Important, try not to touch the `/home/user` folder, as that might create permission issues and break the Python installation.
+
+
+## Testing Locally
+
+Try building the image locally.
+
+```bash
+docker build .  -f e2b.Dockerfile  --tag e2b-example
+```
+
+If the build works, run the following to enter the shell environment within the container. Make sure you can find your files under `/data`, as that is where your agent would access those data files.
+
+```bash
+docker run -it --entrypoint /bin/bash e2b-example
+
+# Within the container
+ls -lh /data
+
+# You should see your data file listed.
+```
+
+## Push to E2B
+
+If the local tests looks reasonable, push the image to E2B as a template.
+
+```bash
+# The command "/root/.jupyter/start-up.sh" is from the E2B base image
+# If you are only adding data files, you don't need to modify this line.
+
+e2b template build -c "/root/.jupyter/start-up.sh"
+```
+
+If the build is finished properly, you should see output like the following:
+
+> ✅ Building sandbox template 9p6favrrqijhasgkq1tv finished.
+
+## Modify your "System Prompt"
+
+Previously, the agent assumes that all data files are under the initial working directory. That is not the case here. You should modify the system prompt and instruct the agent to look under `/data` (or whatever folder you specified in the Dockerfile.)

src/utils/tools/code_interpreter.py

@@ -108,6 +108,7 @@ def __init__(
         self,
         local_files: "Sequence[Path | str]| None" = None,
         timeout_seconds: int = 30,
+        template_name: str | None = None,
     ):
         """Configure your Code Interpreter session.
 
@@ -121,14 +122,19 @@ def __init__(
                 to upload to sandbox working directory. Folders will be flattened.
             timeout_seconds : int
                 Limit executions to this duration.
+            template_name : str | None
+                Optionally, override the default e2b template name.
+                See e2b_template.md for details.
         """
         self.timeout_seconds = timeout_seconds
         self.local_files = []
+        self.template_name = template_name
 
         # Recursively find files if the given path is a folder.
         if local_files:
             for _path in local_files:
                 self.local_files.extend(_enumerate_files(_path))
+        self.template_name = template_name
 
     async def run_code(self, code: str) -> str:
         """Run the given Python code in a sandbox environment.
@@ -138,7 +144,9 @@ async def run_code(self, code: str) -> str:
             code : str
                 Python logic to execute.
         """
-        sbx = await AsyncSandbox.create(timeout=self.timeout_seconds)
+        sbx = await AsyncSandbox.create(
+            timeout=self.timeout_seconds, template=self.template_name
+        )
         await _upload_files(sbx, self.local_files)
 
         try: