e2b-code-interpreter 0.0.4__tar.gz → 0.0.6__tar.gz

{e2b_code_interpreter-0.0.4 → e2b_code_interpreter-0.0.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: e2b-code-interpreter
-Version: 0.0.4
+Version: 0.0.6
 Summary: E2B Code Interpreter - Stateful code execution
 Home-page: https://e2b.dev/
 License: Apache-2.0

{e2b_code_interpreter-0.0.4 → e2b_code_interpreter-0.0.6}/e2b_code_interpreter/main.py

@@ -52,13 +52,13 @@ class CodeInterpreter(Sandbox):
 
 
 class JupyterExtension:
-    _default_kernel_id: Optional[str] = None
-    _connected_kernels: Dict[str, Future[JupyterKernelWebSocket]] = {}
 
     def __init__(self, sandbox: CodeInterpreter, timeout: Optional[float] = TIMEOUT):
         self._sandbox = sandbox
         self._kernel_id_set = Future()
         self._start_connecting_to_default_kernel(timeout=timeout)
+        self._connected_kernels: Dict[str, Future[JupyterKernelWebSocket]] = {}
+        self._default_kernel_id: Optional[str] = None
 
     def exec_cell(
         self,

{e2b_code_interpreter-0.0.4 → e2b_code_interpreter-0.0.6}/e2b_code_interpreter/messaging.py

@@ -43,30 +43,29 @@ class CellExecution:
         self.on_result = on_result
 
 
-class JupyterKernelWebSocket(BaseModel):
-    model_config = ConfigDict(arbitrary_types_allowed=True)
+class JupyterKernelWebSocket:
 
-    url: str
-
-    _cells: Dict[str, CellExecution] = {}
-    _waiting_for_replies: Dict[str, DeferredFuture] = PrivateAttr(default_factory=dict)
-    _queue_in: Queue = PrivateAttr(default_factory=Queue)
-    _queue_out: Queue = PrivateAttr(default_factory=Queue)
-    _process_cleanup: List[Callable[[], Any]] = PrivateAttr(default_factory=list)
-    _closed: bool = PrivateAttr(default=False)
+    def __init__(self, url: str):
+        self.url = url
+        self._cells: Dict[str, CellExecution] = {}
+        self._waiting_for_replies: Dict[str, DeferredFuture] = {}
+        self._queue_in = Queue()
+        self._queue_out = Queue()
+        self._stopped = threading.Event()
 
     def process_messages(self):
-        while True:
-            data = self._queue_out.get()
+        while not self._stopped.is_set():
+            if self._queue_out.empty():
+                time.sleep(0.01)
+                continue
 
+            data = self._queue_out.get()
             logger.debug(f"WebSocket received message: {data}".strip())
             self._receive_message(json.loads(data))
             self._queue_out.task_done()
 
     def connect(self, timeout: float = TIMEOUT):
         started = threading.Event()
-        stopped = threading.Event()
-        self._process_cleanup.append(stopped.set)
 
         threading.Thread(
             target=self.process_messages, daemon=True, name="e2b-process-messages"
@@ -78,7 +77,7 @@ class JupyterKernelWebSocket(BaseModel):
                 queue_in=self._queue_in,
                 queue_out=self._queue_out,
                 started=started,
-                stopped=stopped,
+                stopped=self._stopped
             ).run,
             daemon=True,
             name="e2b-code-interpreter-websocket",
@@ -91,7 +90,7 @@ class JupyterKernelWebSocket(BaseModel):
             while (
                 not started.is_set()
                 and time.time() - start_time < timeout
-                and not self._closed
+                and not self._stopped.is_set()
             ):
                 time.sleep(0.1)
 
@@ -245,12 +244,7 @@ class JupyterKernelWebSocket(BaseModel):
 
     def close(self):
         logger.debug("Closing WebSocket")
-        self._closed = True
-
-        for cancel in self._process_cleanup:
-            cancel()
-
-        self._process_cleanup.clear()
+        self._stopped.set()
 
         for handler in self._waiting_for_replies.values():
             logger.debug(f"Cancelling waiting for execution result for {handler}")

{e2b_code_interpreter-0.0.4 → e2b_code_interpreter-0.0.6}/e2b_code_interpreter/models.py

@@ -1,6 +1,6 @@
 import copy
 from typing import List, Optional, Iterable, Dict
-from pydantic import BaseModel
+from pydantic import BaseModel, field_serializer
 
 
 class Error(BaseModel):
@@ -234,6 +234,25 @@ class Execution(BaseModel):
             if d.is_main_result:
                 return d.text
 
+    def to_json(self) -> str:
+        """
+        Returns the JSON representation of the Execution object.
+        """
+        return self.model_dump_json(exclude_none=True)
+
+    @field_serializer("results", when_used="json")
+    def serialize_results(results: List[Result]) -> List[Dict[str, str]]:
+        """
+        Serializes the results to JSON.
+        This method is used by the Pydantic JSON encoder.
+        """
+        serialized = []
+        for result in results:
+            serialized_dict = {key: result[key] for key in result.formats()}
+            serialized_dict['text'] = result.text
+            serialized.append(serialized_dict)
+        return serialized
+
 
 class KernelException(Exception):
     """

{e2b_code_interpreter-0.0.4 → e2b_code_interpreter-0.0.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "e2b-code-interpreter"
-version = "0.0.4"
+version = "0.0.6"
 description = "E2B Code Interpreter - Stateful code execution"
 authors = ["e2b <hello@e2b.dev>"]
 license = "Apache-2.0"