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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: e2b-code-interpreter
-Version: 0.0.6
+Version: 0.0.8
 Summary: E2B Code Interpreter - Stateful code execution
 Home-page: https://e2b.dev/
 License: Apache-2.0
@@ -14,7 +14,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: e2b (>=0.17.0)
+Requires-Dist: e2b (>=0.17.1)
 Requires-Dist: pydantic (>1,<3)
 Requires-Dist: websocket-client (>=1.7.0,<2.0.0)
 Project-URL: Bug Tracker, https://github.com/e2b-dev/code-interpreter/issues

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

@@ -2,6 +2,8 @@ from __future__ import annotations
 
 import logging
 import threading
+import uuid
+
 import requests
 
 from concurrent.futures import Future
@@ -91,16 +93,16 @@ class JupyterExtension:
             ws = ws_future.result(timeout=timeout)
         else:
             logger.debug(f"Creating new websocket connection to kernel {kernel_id}")
-            ws = self._connect_to_kernel_ws(kernel_id, timeout=timeout)
+            ws = self._connect_to_kernel_ws(kernel_id, None, timeout=timeout)
 
-        session_id = ws.send_execution_message(code, on_stdout, on_stderr, on_result)
+        message_id = ws.send_execution_message(code, on_stdout, on_stderr, on_result)
         logger.debug(
-            f"Sent execution message to kernel {kernel_id}, session_id: {session_id}"
+            f"Sent execution message to kernel {kernel_id}, message_id: {message_id}"
         )
 
-        result = ws.get_result(session_id, timeout=timeout)
+        result = ws.get_result(message_id, timeout=timeout)
         logger.debug(
-            f"Received result from kernel {kernel_id}, session_id: {session_id}, result: {result}"
+            f"Received result from kernel {kernel_id}, message_id: {message_id}, result: {result}"
         )
 
         return result
@@ -139,24 +141,35 @@ class JupyterExtension:
         :param timeout: Timeout for the kernel creation request.
         :return: Kernel id of the created kernel
         """
-        data = {"cwd": cwd}
-        if kernel_name:
-            data["kernel_name"] = kernel_name
+        kernel_name = kernel_name or "python3"
+
+        data = {"path": str(uuid.uuid4()), "kernel": {"name": kernel_name}, "type": "notebook", "name": str(uuid.uuid4())}
         logger.debug(f"Creating kernel with data: {data}")
 
         response = requests.post(
-            f"{self._sandbox.get_protocol()}://{self._sandbox.get_hostname(8888)}/api/kernels",
+            f"{self._sandbox.get_protocol()}://{self._sandbox.get_hostname(8888)}/api/sessions",
             json=data,
             timeout=timeout,
         )
         if not response.ok:
             raise KernelException(f"Failed to create kernel: {response.text}")
 
-        kernel_id = response.json()["id"]
+        session_data = response.json()
+        session_id = session_data["id"]
+        kernel_id = session_data["kernel"]["id"]
+
+        response = requests.patch(
+            f"{self._sandbox.get_protocol()}://{self._sandbox.get_hostname(8888)}/api/sessions/{session_id}",
+            json={"path": cwd},
+            timeout=timeout,
+        )
+        if not response.ok:
+            raise KernelException(f"Failed to create kernel: {response.text}")
+
         logger.debug(f"Created kernel {kernel_id}")
 
         threading.Thread(
-            target=self._connect_to_kernel_ws, args=(kernel_id, timeout)
+            target=self._connect_to_kernel_ws, args=(kernel_id, session_id, timeout)
         ).start()
         return kernel_id
 
@@ -186,7 +199,7 @@ class JupyterExtension:
         logger.debug(f"Restarted kernel {kernel_id}")
 
         threading.Thread(
-            target=self._connect_to_kernel_ws, args=(kernel_id, timeout)
+            target=self._connect_to_kernel_ws, args=(kernel_id, None, timeout)
         ).start()
 
     def shutdown_kernel(
@@ -243,7 +256,7 @@ class JupyterExtension:
             ws.result().close()
 
     def _connect_to_kernel_ws(
-        self, kernel_id: str, timeout: Optional[float] = TIMEOUT
+        self, kernel_id: str, session_id: Optional[str], timeout: Optional[float] = TIMEOUT
     ) -> JupyterKernelWebSocket:
         """
         Establishes a WebSocket connection to a specified Jupyter kernel.
@@ -257,9 +270,12 @@ class JupyterExtension:
         future = Future()
         self._connected_kernels[kernel_id] = future
 
+        session_id = session_id or str(uuid.uuid4())
         ws = JupyterKernelWebSocket(
             url=f"{self._sandbox.get_protocol('ws')}://{self._sandbox.get_hostname(8888)}/api/kernels/{kernel_id}/channels",
+            session_id=session_id
         )
+
         ws.connect(timeout=timeout)
         logger.debug(f"Connected to kernel's ({kernel_id}) websocket.")
 
@@ -279,12 +295,14 @@ class JupyterExtension:
             kernel_id = self._sandbox.filesystem.read(
                 "/root/.jupyter/kernel_id", timeout=timeout
             )
+
             if kernel_id is None and not self._sandbox.is_open:
                 return
 
             kernel_id = kernel_id.strip()
+
             logger.debug(f"Default kernel id: {kernel_id}")
-            self._connect_to_kernel_ws(kernel_id, timeout=timeout)
+            self._connect_to_kernel_ws(kernel_id, None, timeout=timeout)
             self._kernel_id_set.set_result(kernel_id)
 
         threading.Thread(target=setup_default_kernel).start()

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

@@ -5,14 +5,13 @@ import time
 import uuid
 from concurrent.futures import Future
 from queue import Queue
-from typing import Callable, Dict, List, Any, Optional
+from typing import Callable, Dict, Any, Optional
 
 from e2b import ProcessMessage
 from e2b.constants import TIMEOUT
 from e2b.sandbox import TimeoutException
 from e2b.sandbox.websocket_client import WebSocket
 from e2b.utils.future import DeferredFuture
-from pydantic import ConfigDict, PrivateAttr, BaseModel
 
 from e2b_code_interpreter.models import Execution, Result, Error
 
@@ -26,6 +25,7 @@ class CellExecution:
     """
 
     input_accepted: bool = False
+
     on_stdout: Optional[Callable[[ProcessMessage], Any]] = None
     on_stderr: Optional[Callable[[ProcessMessage], Any]] = None
     on_result: Optional[Callable[[Result], Any]] = None
@@ -45,8 +45,9 @@ class CellExecution:
 
 class JupyterKernelWebSocket:
 
-    def __init__(self, url: str):
+    def __init__(self, url: str, session_id: str):
         self.url = url
+        self.session_id = session_id
         self._cells: Dict[str, CellExecution] = {}
         self._waiting_for_replies: Dict[str, DeferredFuture] = {}
         self._queue_in = Queue()
@@ -102,14 +103,13 @@ class JupyterKernelWebSocket:
 
         logger.debug("WebSocket started")
 
-    @staticmethod
-    def _get_execute_request(msg_id: str, code: str) -> str:
+    def _get_execute_request(self, msg_id: str, code: str) -> str:
         return json.dumps(
             {
                 "header": {
                     "msg_id": msg_id,
                     "username": "e2b",
-                    "session": str(uuid.uuid4()),
+                    "session": self.session_id,
                     "msg_type": "execute_request",
                     "version": "5.3",
                 },
@@ -118,7 +118,7 @@ class JupyterKernelWebSocket:
                 "content": {
                     "code": code,
                     "silent": False,
-                    "store_history": False,
+                    "store_history": True,
                     "user_expressions": {},
                     "allow_stdin": False,
                 },
@@ -238,6 +238,7 @@ class JupyterKernelWebSocket:
 
         elif data["msg_type"] == "execute_input":
             logger.debug(f"Input accepted for {parent_msg_ig}")
+            cell.partial_result.execution_count = data["content"]["execution_count"]
             cell.input_accepted = True
         else:
             logger.warning(f"[UNHANDLED MESSAGE TYPE]: {data['msg_type']}")

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

@@ -124,6 +124,9 @@ class Result:
         """
         return self.text
 
+    def __repr__(self) -> str:
+        return f"Result({self.text})"
+
     def _repr_html_(self) -> Optional[str]:
         """
         Returns the HTML representation of the data.
@@ -222,6 +225,8 @@ class Execution(BaseModel):
     "Logs printed to stdout and stderr during execution."
     error: Optional[Error] = None
     "Error object if an error occurred, None otherwise."
+    execution_count: Optional[int] = None
+    "Execution count of the cell."
 
     @property
     def text(self) -> Optional[str]:
@@ -249,7 +254,7 @@ class Execution(BaseModel):
         serialized = []
         for result in results:
             serialized_dict = {key: result[key] for key in result.formats()}
-            serialized_dict['text'] = result.text
+            serialized_dict["text"] = result.text
             serialized.append(serialized_dict)
         return serialized
 

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

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "e2b-code-interpreter"
-version = "0.0.6"
+version = "0.0.8"
 description = "E2B Code Interpreter - Stateful code execution"
 authors = ["e2b <hello@e2b.dev>"]
 license = "Apache-2.0"
@@ -14,7 +14,7 @@ python = "^3.8"
 
 pydantic = ">1, <3"
 websocket-client = "^1.7.0"
-e2b = ">=0.17.0"
+e2b = ">=0.17.1"
 
 [tool.poetry.group.dev.dependencies]
 black = "^24.3.0"