Merge pull request #170 from athina-ai/fix/allow-template-sandbox-in-…

…code-execution Allow user to pass a template ID for E2B sandbox
athina-ai · Feb 13, 2025 · 9f35f84 · 9f35f84
2 parents 03d7df1 + a19918e
commit 9f35f84
Showing 1 changed file with 2 additions and 0 deletions.
diff --git 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
                     ),