Add on stdout and on stderr for code execution streaming, fix websocket

packages/python-sdk/e2b/templates/stateful_code_interpreter.py

@@ -1,6 +1,7 @@
 import binascii
 import json
 import os
+import time
 import uuid
 from time import sleep
 from typing import Optional, Callable, Any, List
@@ -13,11 +14,17 @@
 from e2b.constants import TIMEOUT
 
 
+class Error(BaseModel):
+    name: str
+    value: str
+    traceback: List[str]
+
+
 class Result(BaseModel):
     output: Optional[str] = None
     stdout: List[str] = []
     stderr: List[str] = []
-    error: Optional[str] = None
+    error: Optional[Error] = None
     # TODO: This will be changed in the future, it's just to enable the use of display_data
     display_data: List[dict] = []
 
@@ -79,7 +86,7 @@ def _connect_kernel(self):
         )
 
     @staticmethod
-    def _send_execute_request(code: str):
+    def _send_execute_request(code: str) -> dict:
         msg_id = str(uuid.uuid4())
         session = str(uuid.uuid4())
 
@@ -103,42 +110,84 @@ def _send_execute_request(code: str):
         }
 
     @staticmethod
-    def _wait_for_result(ws) -> Result:
+    def _wait_for_result(
+        ws,
+        on_stdout: Optional[Callable[[ProcessMessage], Any]],
+        on_stderr: Optional[Callable[[ProcessMessage], Any]],
+    ) -> Result:
         result = Result()
-        was_busy = False
+        input_accepted = False
 
         while True:
             response = json.loads(ws.recv())
             if response["msg_type"] == "error":
-                result.error = response["content"]["traceback"]
+                result.error = Error(
+                    name=response["content"]["ename"],
+                    value=response["content"]["evalue"],
+                    traceback=response["content"]["traceback"],
+                )
+
             elif response["msg_type"] == "stream":
                 if response["content"]["name"] == "stdout":
                     result.stdout.append(response["content"]["text"])
+                    if on_stdout:
+                        on_stdout(
+                            ProcessMessage(
+                                line=response["content"]["text"],
+                                timestamp=time.time_ns(),
+                            )
+                        )
+
                 elif response["content"]["name"] == "stderr":
                     result.stderr.append(response["content"]["text"])
+                    if on_stderr:
+                        on_stderr(
+                            ProcessMessage(
+                                line=response["content"]["text"],
+                                error=True,
+                                timestamp=time.time_ns(),
+                            )
+                        )
 
             elif response["msg_type"] == "display_data":
                 result.display_data.append(response["content"]["data"])
+
             elif response["msg_type"] == "execute_result":
                 result.output = response["content"]["data"]["text/plain"]
 
             elif response["msg_type"] == "status":
                 if response["content"]["execution_state"] == "idle":
-                    if was_busy:
+                    if input_accepted:
                         break
                 elif response["content"]["execution_state"] == "error":
-                    result.error = "An error occurred while executing the code"
+                    result.error = "Kernel execution error"
                     break
-                elif response["content"]["execution_state"] == "busy":
-                    if not was_busy:
-                        was_busy = True
 
+            elif response["msg_type"] == "execute_reply":
+                if response["content"]["status"] == "error":
+                    result.error = Error(
+                        name=response["content"]["ename"],
+                        value=response["content"]["evalue"],
+                        traceback=response["content"]["traceback"],
+                    )
+                elif response["content"]["status"] == "ok":
+                    pass
+
+            elif response["msg_type"] == "execute_input":
+                input_accepted = True
+            else:
+                print("[UNHANDLED MESSAGE TYPE]:", response["msg_type"])
         return result
 
-    def exec_python(self, code: str) -> Result:
+    def exec_python(
+        self,
+        code: str,
+        on_stdout: Optional[Callable[[ProcessMessage], Any]] = None,
+        on_stderr: Optional[Callable[[ProcessMessage], Any]] = None,
+    ) -> Result:
         ws = self._connect_kernel()
         ws.send(json.dumps(self._send_execute_request(code)))
-        result = self._wait_for_result(ws)
+        result = self._wait_for_result(ws, on_stdout, on_stderr)
 
         ws.close()
 

packages/python-sdk/pyproject.toml

@@ -21,6 +21,7 @@ websockets = ">=11.0.3"
 jsonrpcclient = ">=4.0.3"
 requests = ">=2.31.0"
 typing-extensions = ">=4.8.0"
+websocket-client = "^1.7.0"
 
 [tool.poetry.group.dev.dependencies]
 black = "^23.7.0"