E2B upgrade to 1.0.1 (#78)

* bump e2b and update tool * extras in tests * litellm version * lock
datastax · Oct 30, 2024 · 013e4fb · 013e4fb
1 parent 8acee8b
commit 013e4fb
Show file tree

Hide file tree

Showing 4 changed files with 159 additions and 281 deletions.
diff --git a/.github/workflows/run-tests.yml b/.github/workflows/run-tests.yml
@@ -1438,7 +1438,7 @@ jobs:
       - name: get client dependencies
         working-directory: ./client
         run: |
-          poetry install --with optional
+          poetry install --extras tools
       - name: get server dependencies
         run: |
           poetry install
@@ -1494,7 +1494,7 @@ jobs:
       - name: get client dependencies
         working-directory: ./client
         run: |
-          poetry install --with optional
+          poetry install --extras tools
       - name: get server dependencies
         run: |
           poetry install
@@ -1550,7 +1550,7 @@ jobs:
       - name: get client dependencies
         working-directory: ./client
         run: |
-          poetry install --with optional
+          poetry install --extras tools
       - name: get server dependencies
         run: |
           poetry install
@@ -1606,7 +1606,7 @@ jobs:
       - name: get client dependencies
         working-directory: ./client
         run: |
-          poetry install --with optional
+          poetry install --extras tools
       - name: get server dependencies
         run: |
           poetry install

diff --git a/client/astra_assistants/tools/e2b_code_interpreter.py b/client/astra_assistants/tools/e2b_code_interpreter.py
@@ -1,6 +1,4 @@
-from e2b import Sandbox
-
-from e2b_code_interpreter import CodeInterpreter
+from e2b_code_interpreter import Sandbox
 
 from astra_assistants.tools.tool_interface import ToolInterface
 
@@ -13,18 +11,16 @@ def __init__(self):
         running_sandboxes = Sandbox.list()
         # Find the sandbox by metadata
         for running_sandbox in running_sandboxes:
-            sandbox = Sandbox.reconnect(running_sandbox.sandbox_id)
-            sandbox.close()
+            sandbox = Sandbox.connect(running_sandbox.sandbox_id)
+            sandbox.kill()
         else:
             # Sandbox not found
             pass
-        self.code_interpreter = CodeInterpreter()
+        self.code_interpreter = Sandbox()
 
         print("initialized")
 
     def call(self, arguments):
         code = arguments['arguments']
-        exec = self.code_interpreter.notebook.exec_cell(
-            code,
-        )
-        return exec.text
+        execution = self.code_interpreter.run_code(code)
+        return execution.text