E2B tool

.github/workflows/run-tests.yml

@@ -1283,6 +1283,7 @@ jobs:
       PERPLEXITYAI_API_KEY: ${{ secrets.PERPLEXITYAI_API_KEY }}
       ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
       GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }}
+      E2B_API_KEY: ${{ secrets.E2B_API_KEY }}
 
     steps:
       - name: Git checkout

client/astra_assistants/astra_assistants_event_handler.py

@@ -40,8 +40,7 @@ def run_tool(self, tool_call):
         if tool_name in self.tools:
             tool = self.tools[tool_name]
             arguments = json.loads(tool_call.function.arguments)
-            query = arguments['query']
-            results = tool.search(query)
+            results = tool.call(arguments)
             return results
         else:
             self.logger.error(f"Tool {tool_name} not found.")

client/astra_assistants/tools/astra_data_api.py

@@ -28,7 +28,8 @@ def __init__(self, db_url: str, collection_name, vectorize: bool = False, openai
         self.openai_client = openai_client
         self.embedding_model = embedding_model
 
-    def search(self, query):
+    def call(self, arguments):
+        query = arguments['arguments']
         try:
             if self.vectorize:
                 results = self.collection.find(

client/astra_assistants/tools/e2b_code_interpreter.py

@@ -0,0 +1,31 @@
+from e2b import Sandbox
+
+from e2b_code_interpreter import CodeInterpreter
+
+from astra_assistants.tools.tool_interface import ToolInterface
+
+
+class E2BCodeInterpreter(ToolInterface):
+
+    def __init__(self):
+        print("initializing code interpreter")
+
+        running_sandboxes = Sandbox.list()
+        # Find the sandbox by metadata
+        for running_sandbox in running_sandboxes:
+            if running_sandbox.metadata.get("user_id", "") == 'uniqueID':
+                sandbox = Sandbox.reconnect(running_sandbox.sandbox_id)
+                sandbox.close()
+        else:
+            # Sandbox not found
+            pass
+        self.code_interpreter = CodeInterpreter()
+
+        print("initialized")
+
+    def call(self, arguments):
+        code = arguments['arguments']
+        exec = self.code_interpreter.notebook.exec_cell(
+            code,
+        )
+        return exec.text

client/astra_assistants/tools/tool_interface.py

@@ -7,7 +7,7 @@ def search(self, query):
         pass
 class ToolInterface(ABC):
     @abstractmethod
-    def search(self, query):
+    def call(self, query):
         pass
 
     def name(self):
@@ -17,7 +17,7 @@ def tool_choice_object(self):
         return {"type": "function", "function": {"name": self.name()}}
 
     def to_function(self):
-        search_sig = inspect.signature(self.search)
+        search_sig = inspect.signature(self.call)
         parameters = {
             name: {
                 "type": "string",  # Assuming all parameters are strings for simplicity

client/pyproject.toml

@@ -20,6 +20,7 @@ pytest = "^8.0.2"
 
 [tool.poetry.group.optional.dependencies]
 astrapy = "^1.2.1"
+e2b-code-interpreter = "^0.0.9"
 
 [build-system]
 requires = ["poetry-core"]