Support uploading file to e2b code interpreter tool

templates/components/engines/python/agent/tools/interpreter.py

@@ -3,7 +3,7 @@
 import base64
 import uuid
 from pydantic import BaseModel
-from typing import List, Tuple, Dict
+from typing import List, Tuple, Dict, Optional
 from llama_index.core.tools import FunctionTool
 from e2b_code_interpreter import CodeInterpreter
 from e2b_code_interpreter.models import Logs
@@ -14,14 +14,15 @@
 
 class InterpreterExtraResult(BaseModel):
     type: str
-    filename: str
-    url: str
+    content: Optional[str] = None
+    filename: Optional[str] = None
+    url: Optional[str] = None
 
 
 class E2BToolOutput(BaseModel):
     is_error: bool
     logs: Logs
-    results: List[InterpreterExtraResult] = []
+    results: List[InterpreterExtraResult | str] = []
 
 
 class E2BCodeInterpreter:
@@ -62,8 +63,9 @@ def get_file_url(self, filename: str) -> str:
 
     def parse_result(self, result) -> List[InterpreterExtraResult]:
         """
-        The result could include multiple formats (e.g. png, svg, etc.) but encoded in base64
-        We save each result to disk and return saved file metadata (extension, filename, url)
+        The result format could be either a base64 string (png, svg, etc.) or a raw text (text, html, markdown,...)
+        If it's base64, we save each result to disk and return saved file metadata (extension, filename, url),
+        otherwise just return the raw text content
         """
         if not result:
             return []
@@ -72,31 +74,46 @@ def parse_result(self, result) -> List[InterpreterExtraResult]:
 
         try:
             formats = result.formats()
-            base64_data_arr = [result[format] for format in formats]
-
-            for ext, base64_data in zip(formats, base64_data_arr):
-                if ext and base64_data:
-                    result = self.save_to_disk(base64_data, ext)
-                    filename = result["filename"]
-                    output.append(
-                        InterpreterExtraResult(
-                            type=ext, filename=filename, url=self.get_file_url(filename)
+            data_list = [result[format] for format in formats]
+
+            for ext, data in zip(formats, data_list):
+                match ext:
+                    case "png" | "jpeg" | "svg":
+                        result = self.save_to_disk(data, ext)
+                        filename = result["filename"]
+                        output.append(
+                            InterpreterExtraResult(
+                                type=ext,
+                                filename=filename,
+                                url=self.get_file_url(filename),
+                            )
                         )
-                    )
+                        break
+                    case "text" | "html" | "markdown":
+                        output.append(InterpreterExtraResult(type=ext, content=data))
         except Exception as error:
             logger.error("Error when saving data to disk", error)
 
         return output
 
-    def interpret(self, code: str) -> E2BToolOutput:
+    def interpret(self, code: str, file_path: Optional[str] = None) -> E2BToolOutput:
         with CodeInterpreter(api_key=self.api_key) as interpreter:
+            # Upload file to E2B sandbox
+            if file_path is not None:
+                with open(file_path, "rb") as f:
+                    remote_path = interpreter.upload_file(f)
+
+            # Execute the code to analyze the file
             logger.info(
                 f"\n{'='*50}\n> Running following AI-generated code:\n{code}\n{'='*50}"
             )
             exec = interpreter.notebook.exec_cell(code)
 
             if exec.error:
-                output = E2BToolOutput(is_error=True, logs=[exec.error])
+                logger.error(
+                    f"Error when executing code in E2B sandbox: {exec.error} {exec.logs}"
+                )
+                output = E2BToolOutput(is_error=True, logs=exec.logs, results=[])
             else:
                 if len(exec.results) == 0:
                     output = E2BToolOutput(is_error=False, logs=exec.logs, results=[])
@@ -108,9 +125,15 @@ def interpret(self, code: str) -> E2BToolOutput:
             return output
 
 
-def code_interpret(code: str) -> Dict:
+def code_interpret(code: str, local_file_path: str) -> Dict:
     """
-    Execute python code in a Jupyter notebook cell and return any result, stdout, stderr, display_data, and error.
+    Use this tool to analyze the provided data in a sandbox environment.
+    The tool will:
+        1. Upload the provided file from local to the sandbox. The uploaded file path will be /home/user/{filename}
+        2. Execute python code in a Jupyter notebook cell to analyze the uploaded file in the sandbox.
+        3. Get the result from the code in stdout, stderr, display_data, and error.
+    You must to provide the code and the provided file path to run this tool.
+    Your code should read the file from the sandbox path /home/user/{filename}.
     """
     api_key = os.getenv("E2B_API_KEY")
     filesever_url_prefix = os.getenv("FILESERVER_URL_PREFIX")
@@ -126,7 +149,7 @@ def code_interpret(code: str) -> Dict:
     interpreter = E2BCodeInterpreter(
         api_key=api_key, filesever_url_prefix=filesever_url_prefix
     )
-    output = interpreter.interpret(code)
+    output = interpreter.interpret(code, local_file_path)
     return output.dict()