From a19918e33cd689c89eb386415da512b5c2d1e614 Mon Sep 17 00:00:00 2001
From: Shiv Sakhuja <shivsak.s@gmail.com>
Date: Wed, 12 Feb 2025 03:44:04 -0800
Subject: [PATCH] Allow user to pass a template ID for E2B sandbox

---
 athina/steps/code_execution_v2.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/athina/steps/code_execution_v2.py b/athina/steps/code_execution_v2.py
index b68bc2d..fb2a6ed 100644
--- a/athina/steps/code_execution_v2.py
+++ b/athina/steps/code_execution_v2.py
@@ -104,6 +104,7 @@ class CodeExecutionV2(Step):
     DEFAULT_TIMEOUT: ClassVar[int] = 60  # 1 minute default timeout for sandbox
     MAX_TIMEOUT: ClassVar[int] = 300  # 5 minute limit for e2b sandbox execution
     sandbox_timeout: Optional[int] = None
+    template: Optional[str] = None
 
     def __init__(
         self,
@@ -132,6 +133,7 @@ def _create_or_initialize_sandbox(self):
 
             if self._sandbox is None:
                 self._sandbox = Sandbox(
+                    template=self.template,
                     timeout=min(
                         self.sandbox_timeout or self.DEFAULT_TIMEOUT, self.MAX_TIMEOUT
                     ),