e2b-code-interpreter 0.0.1a0__tar.gz

e2b_code_interpreter-0.0.1a0/PKG-INFO

@@ -0,0 +1,151 @@
+Metadata-Version: 2.1
+Name: e2b-code-interpreter
+Version: 0.0.1a0
+Summary: E2B Code Interpreter - Stateful code execution
+Home-page: https://e2b.dev/
+License: Apache-2.0
+Author: e2b
+Author-email: hello@e2b.dev
+Requires-Python: >=3.8,<4.0
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Requires-Dist: e2b (>=0.14.11)
+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
+Project-URL: Repository, https://github.com/e2b-dev/e2b-code-interpreter/tree/python
+Description-Content-Type: text/markdown
+
+# Code interpreter extension for Python
+
+The repository contains a template and modules for the code interpreter sandbox. It is based on the Jupyter server and implements the Jupyter Kernel messaging protocol. This allows for sharing context between code executions and improves support for plotting charts and other display-able data.
+
+## Key Features
+
+- **Stateful Execution**: Unlike traditional sandboxes that treat each code execution independently, this package maintains context across executions.
+- **Displaying Graph & Data**: Implements parts of the [Jupyter Kernel messaging protocol](https://jupyter-client.readthedocs.io/en/latest/messaging.html), which support for interactive features like plotting charts, rendering DataFrames, etc.
+
+## Installation
+
+```sh
+pip install e2b-code-interpreter
+```
+
+## Examples
+
+### Minimal example with the sharing context
+
+```python
+from e2b_code_interpreter import CodeInterpreter
+
+with CodeInterpreter() as sandbox:
+    sandbox.exec_cell("x = 1")
+
+    result = sandbox.exec_cell("x += 1; x")
+    print(result.text)  # outputs 2
+
+```
+
+### Get charts and any display-able data
+
+```python
+import base64
+import io
+
+from matplotlib import image as mpimg, pyplot as plt
+
+from e2b_code_interpreter import CodeInterpreter
+
+code = """
+import matplotlib.pyplot as plt
+import numpy as np
+
+x = np.linspace(0, 20, 100)
+y = np.sin(x)
+
+plt.plot(x, y)
+plt.show()
+"""
+
+with CodeInterpreter() as sandbox:
+    # you can install dependencies in "jupyter notebook style"
+    sandbox.exec_cell("!pip install matplotlib")
+
+    # plot random graph
+    result = sandbox.exec_cell(code)
+
+    # there's your image
+    image = result.display_data[0]["image/png"]
+
+    # example how to show the image / prove it works
+    i = base64.b64decode(image)
+    i = io.BytesIO(i)
+    i = mpimg.imread(i, format='PNG')
+
+    plt.imshow(i, interpolation='nearest')
+    plt.show()
+```
+
+### Streaming code output
+
+```python
+from e2b_code_interpreter import CodeInterpreter
+
+code = """
+import time
+
+print("hello")
+time.sleep(5)
+print("world")
+"""
+
+with CodeInterpreter() as sandbox:
+    sandbox.exec_cell(code, on_stdout=print, on_stderr=print)
+```
+
+### Pre-installed Python packages inside the sandbox
+
+The full and always up-to-date list can be found in the [`requirements.txt`](https://github.com/e2b-dev/E2B/blob/stateful-code-interpreter/sandboxes/code-interpreter-stateful/requirements.txt) file.
+
+```text
+# Jupyter server requirements
+jupyter-server==2.13.0
+ipykernel==6.29.3
+ipython==8.22.2
+
+# Other packages
+aiohttp==3.9.3
+beautifulsoup4==4.12.3
+bokeh==3.3.4
+gensim==4.3.2
+imageio==2.34.0
+joblib==1.3.2
+librosa==0.10.1
+matplotlib==3.8.3
+nltk==3.8.1
+numpy==1.26.4
+opencv-python==4.9.0.80
+openpyxl==3.1.2
+pandas==1.5.3
+plotly==5.19.0
+pytest==8.1.0
+python-docx==1.1.0
+pytz==2024.1
+requests==2.26.0
+scikit-image==0.22.0
+scikit-learn==1.4.1.post1
+scipy==1.12.0
+seaborn==0.13.2
+soundfile==0.12.1
+spacy==3.7.4
+textblob==0.18.0
+tornado==6.4
+urllib3==1.26.7
+xarray==2024.2.0
+xlrd==2.0.1
+```
+

e2b_code_interpreter-0.0.1a0/README.md

@@ -0,0 +1,128 @@
+# Code interpreter extension for Python
+
+The repository contains a template and modules for the code interpreter sandbox. It is based on the Jupyter server and implements the Jupyter Kernel messaging protocol. This allows for sharing context between code executions and improves support for plotting charts and other display-able data.
+
+## Key Features
+
+- **Stateful Execution**: Unlike traditional sandboxes that treat each code execution independently, this package maintains context across executions.
+- **Displaying Graph & Data**: Implements parts of the [Jupyter Kernel messaging protocol](https://jupyter-client.readthedocs.io/en/latest/messaging.html), which support for interactive features like plotting charts, rendering DataFrames, etc.
+
+## Installation
+
+```sh
+pip install e2b-code-interpreter
+```
+
+## Examples
+
+### Minimal example with the sharing context
+
+```python
+from e2b_code_interpreter import CodeInterpreter
+
+with CodeInterpreter() as sandbox:
+    sandbox.exec_cell("x = 1")
+
+    result = sandbox.exec_cell("x += 1; x")
+    print(result.text)  # outputs 2
+
+```
+
+### Get charts and any display-able data
+
+```python
+import base64
+import io
+
+from matplotlib import image as mpimg, pyplot as plt
+
+from e2b_code_interpreter import CodeInterpreter
+
+code = """
+import matplotlib.pyplot as plt
+import numpy as np
+
+x = np.linspace(0, 20, 100)
+y = np.sin(x)
+
+plt.plot(x, y)
+plt.show()
+"""
+
+with CodeInterpreter() as sandbox:
+    # you can install dependencies in "jupyter notebook style"
+    sandbox.exec_cell("!pip install matplotlib")
+
+    # plot random graph
+    result = sandbox.exec_cell(code)
+
+    # there's your image
+    image = result.display_data[0]["image/png"]
+
+    # example how to show the image / prove it works
+    i = base64.b64decode(image)
+    i = io.BytesIO(i)
+    i = mpimg.imread(i, format='PNG')
+
+    plt.imshow(i, interpolation='nearest')
+    plt.show()
+```
+
+### Streaming code output
+
+```python
+from e2b_code_interpreter import CodeInterpreter
+
+code = """
+import time
+
+print("hello")
+time.sleep(5)
+print("world")
+"""
+
+with CodeInterpreter() as sandbox:
+    sandbox.exec_cell(code, on_stdout=print, on_stderr=print)
+```
+
+### Pre-installed Python packages inside the sandbox
+
+The full and always up-to-date list can be found in the [`requirements.txt`](https://github.com/e2b-dev/E2B/blob/stateful-code-interpreter/sandboxes/code-interpreter-stateful/requirements.txt) file.
+
+```text
+# Jupyter server requirements
+jupyter-server==2.13.0
+ipykernel==6.29.3
+ipython==8.22.2
+
+# Other packages
+aiohttp==3.9.3
+beautifulsoup4==4.12.3
+bokeh==3.3.4
+gensim==4.3.2
+imageio==2.34.0
+joblib==1.3.2
+librosa==0.10.1
+matplotlib==3.8.3
+nltk==3.8.1
+numpy==1.26.4
+opencv-python==4.9.0.80
+openpyxl==3.1.2
+pandas==1.5.3
+plotly==5.19.0
+pytest==8.1.0
+python-docx==1.1.0
+pytz==2024.1
+requests==2.26.0
+scikit-image==0.22.0
+scikit-learn==1.4.1.post1
+scipy==1.12.0
+seaborn==0.13.2
+soundfile==0.12.1
+spacy==3.7.4
+textblob==0.18.0
+tornado==6.4
+urllib3==1.26.7
+xarray==2024.2.0
+xlrd==2.0.1
+```

e2b_code_interpreter-0.0.1a0/e2b_code_interpreter/__init__.py

@@ -0,0 +1,2 @@
+from .main import CodeInterpreter, JupyterExtension
+from .models import Result, Error, Data, KernelException

e2b_code_interpreter-0.0.1a0/e2b_code_interpreter/main.py

@@ -0,0 +1,289 @@
+from __future__ import annotations
+
+import logging
+import threading
+from concurrent.futures import Future
+from typing import Any, Callable, List, Optional, Dict
+
+import requests
+from e2b import EnvVars, ProcessMessage, Sandbox
+from e2b.constants import TIMEOUT
+
+from e2b_code_interpreter.messaging import JupyterKernelWebSocket
+from e2b_code_interpreter.models import KernelException, Result
+
+
+logger = logging.getLogger(__name__)
+
+
+class CodeInterpreter(Sandbox):
+    """
+    E2B code interpreter sandbox extension.
+    """
+
+    template = "code-interpreter-stateful"
+
+    def __init__(
+        self,
+        template: Optional[str] = None,
+        api_key: Optional[str] = None,
+        cwd: Optional[str] = None,
+        env_vars: Optional[EnvVars] = None,
+        timeout: Optional[float] = TIMEOUT,
+        on_stdout: Optional[Callable[[ProcessMessage], Any]] = None,
+        on_stderr: Optional[Callable[[ProcessMessage], Any]] = None,
+        on_exit: Optional[Callable[[int], Any]] = None,
+        **kwargs,
+    ):
+        super().__init__(
+            template=template or self.template,
+            api_key=api_key,
+            cwd=cwd,
+            env_vars=env_vars,
+            timeout=timeout,
+            on_stdout=on_stdout,
+            on_stderr=on_stderr,
+            on_exit=on_exit,
+            **kwargs,
+        )
+        self.notebook = JupyterExtension(self, timeout=timeout)
+        # Close all the websocket connections when the interpreter is closed
+        self._process_cleanup.append(self.notebook.close)
+
+
+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)
+
+    def exec_cell(
+        self,
+        code: str,
+        kernel_id: Optional[str] = None,
+        on_stdout: Optional[Callable[[ProcessMessage], Any]] = None,
+        on_stderr: Optional[Callable[[ProcessMessage], Any]] = None,
+        on_display_data: Optional[Callable[[Dict[str, Any]], Any]] = None,
+        timeout: Optional[float] = TIMEOUT,
+    ) -> Result:
+        """
+        Execute code in a notebook cell.
+
+        :param code: Code to execute
+        :param kernel_id: The ID of the kernel to execute the code on. If not provided, the default kernel is used.
+        :param on_stdout: A callback function to handle standard output messages from the code execution.
+        :param on_stderr: A callback function to handle standard error messages from the code execution.
+        :param on_display_data: A callback function to handle display data messages from the code execution.
+        :param timeout: Timeout for the call
+
+        :return: Result of the execution
+        """
+        kernel_id = kernel_id or self.default_kernel_id
+        ws_future = self._connected_kernels.get(kernel_id)
+
+        logger.debug(f"Executing code in kernel {kernel_id}")
+
+        if ws_future:
+            logger.debug(f"Using existing websocket connection to kernel {kernel_id}")
+            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)
+
+        session_id = ws.send_execution_message(code, on_stdout, on_stderr, on_display_data)
+        logger.debug(
+            f"Sent execution message to kernel {kernel_id}, session_id: {session_id}"
+        )
+
+        result = ws.get_result(session_id, timeout=timeout)
+        logger.debug(
+            f"Received result from kernel {kernel_id}, session_id: {session_id}, result: {result}"
+        )
+
+        return result
+
+    @property
+    def default_kernel_id(self) -> str:
+        """
+        Get the default kernel id
+
+        :return: Default kernel id
+        """
+        if not self._default_kernel_id:
+            logger.debug("Waiting for default kernel id")
+            self._default_kernel_id = self._kernel_id_set.result()
+
+        return self._default_kernel_id
+
+    def create_kernel(
+        self,
+        cwd: str = "/home/user",
+        kernel_name: Optional[str] = None,
+        timeout: Optional[float] = TIMEOUT,
+    ) -> str:
+        """
+        Creates a new kernel, this can be useful if you want to have multiple independent code execution environments.
+
+        The kernel can be optionally configured to start in a specific working directory and/or
+        with a specific kernel name. If no kernel name is provided, the default kernel will be used.
+        Once the kernel is created, this method establishes a WebSocket connection to the new kernel for
+        real-time communication.
+
+        :param cwd: Sets the current working directory for the kernel. Defaults to "/home/user".
+        :param kernel_name:
+            Specifies which kernel should be used, useful if you have multiple kernel types.
+            If not provided, the default kernel will be used.
+        :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
+        logger.debug(f"Creating kernel with data: {data}")
+
+        response = requests.post(
+            f"{self._sandbox.get_protocol()}://{self._sandbox.get_hostname(8888)}/api/kernels",
+            json=data,
+            timeout=timeout,
+        )
+        if not response.ok:
+            raise KernelException(f"Failed to create kernel: {response.text}")
+
+        kernel_id = response.json()["id"]
+        logger.debug(f"Created kernel {kernel_id}")
+
+        threading.Thread(
+            target=self._connect_to_kernel_ws, args=(kernel_id, timeout)
+        ).start()
+        return kernel_id
+
+    def restart_kernel(
+        self, kernel_id: Optional[str] = None, timeout: Optional[float] = TIMEOUT
+    ) -> None:
+        """
+        Restarts an existing Jupyter kernel. This can be useful to reset the kernel's state or to recover from errors.
+
+        :param kernel_id: The unique identifier of the kernel to restart. If not provided, the default kernel is restarted.
+        :param timeout: The timeout in milliseconds for the kernel restart request.
+        """
+        kernel_id = kernel_id or self.default_kernel_id
+        logger.debug(f"Restarting kernel {kernel_id}")
+
+        self._connected_kernels[kernel_id].result().close()
+        del self._connected_kernels[kernel_id]
+        logger.debug(f"Closed websocket connection to kernel {kernel_id}")
+
+        response = requests.post(
+            f"{self._sandbox.get_protocol()}://{self._sandbox.get_hostname(8888)}/api/kernels/{kernel_id}/restart",
+            timeout=timeout,
+        )
+        if not response.ok:
+            raise KernelException(f"Failed to restart kernel {kernel_id}")
+
+        logger.debug(f"Restarted kernel {kernel_id}")
+
+        threading.Thread(
+            target=self._connect_to_kernel_ws, args=(kernel_id, timeout)
+        ).start()
+
+    def shutdown_kernel(
+        self, kernel_id: Optional[str] = None, timeout: Optional[float] = TIMEOUT
+    ) -> None:
+        """
+        Shuts down an existing Jupyter kernel. This method is used to gracefully terminate a kernel's process.
+
+        :param kernel_id: The unique identifier of the kernel to shutdown. If not provided, the default kernel is shutdown.
+        :param timeout: The timeout for the kernel shutdown request.
+        """
+        kernel_id = kernel_id or self.default_kernel_id
+        logger.debug(f"Shutting down kernel {kernel_id}")
+
+        self._connected_kernels[kernel_id].result().close()
+        del self._connected_kernels[kernel_id]
+        logger.debug(f"Closed websocket connection to kernel {kernel_id}")
+
+        response = requests.delete(
+            f"{self._sandbox.get_protocol()}://{self._sandbox.get_hostname(8888)}/api/kernels/{kernel_id}",
+            timeout=timeout,
+        )
+        if not response.ok:
+            raise KernelException(f"Failed to shutdown kernel {kernel_id}")
+
+        logger.debug(f"Shutdown kernel {kernel_id}")
+
+    def list_kernels(self, timeout: Optional[float] = TIMEOUT) -> List[str]:
+        """
+        Lists all available Jupyter kernels.
+
+        This method fetches a list of all currently available Jupyter kernels from the server. It can be used
+        to retrieve the IDs of all kernels that are currently running or available for connection.
+
+        :param timeout: The timeout for the kernel list request.
+        :return: List of kernel ids
+        """
+        response = requests.get(
+            f"{self._sandbox.get_protocol()}://{self._sandbox.get_hostname(8888)}/api/kernels",
+            timeout=timeout,
+        )
+
+        if not response.ok:
+            raise KernelException(f"Failed to list kernels: {response.text}")
+
+        return [kernel["id"] for kernel in response.json()]
+
+    def close(self):
+        """
+        Close all the websocket connections to the kernels. It doesn't shutdown the kernels.
+        """
+        logger.debug("Closing all websocket connections")
+        for ws in self._connected_kernels.values():
+            ws.result().close()
+
+    def _connect_to_kernel_ws(
+        self, kernel_id: str, timeout: Optional[float] = TIMEOUT
+    ) -> JupyterKernelWebSocket:
+        """
+        Establishes a WebSocket connection to a specified Jupyter kernel.
+
+        :param kernel_id: Kernel id
+        :param timeout: The timeout for the kernel connection request.
+
+        :return: Websocket connection
+        """
+        logger.debug(f"Connecting to kernel's ({kernel_id}) websocket")
+        future = Future()
+        self._connected_kernels[kernel_id] = future
+
+        ws = JupyterKernelWebSocket(
+            url=f"{self._sandbox.get_protocol('ws')}://{self._sandbox.get_hostname(8888)}/api/kernels/{kernel_id}/channels",
+        )
+        ws.connect(timeout=timeout)
+        logger.debug(f"Connected to kernel's ({kernel_id}) websocket.")
+
+        future.set_result(ws)
+        return ws
+
+    def _start_connecting_to_default_kernel(
+        self, timeout: Optional[float] = TIMEOUT
+    ) -> None:
+        """
+        Start connecting to the default kernel in a separate thread to avoid blocking the main thread.
+        :param timeout: Timeout for the call
+        """
+        logger.debug("Starting to connect to the default kernel")
+
+        def setup_default_kernel():
+            kernel_id = self._sandbox.filesystem.read(
+                "/root/.jupyter/kernel_id", timeout=timeout
+            )
+            if kernel_id is None and not self._sandbox.is_open:
+                return
+
+            logger.debug(f"Default kernel id: {kernel_id}")
+            self._connect_to_kernel_ws(kernel_id, timeout=timeout)
+            self._kernel_id_set.set_result(kernel_id)
+
+        threading.Thread(target=setup_default_kernel).start()

e2b_code_interpreter-0.0.1a0/e2b_code_interpreter/messaging.py

@@ -0,0 +1,253 @@
+import json
+import logging
+import threading
+import time
+import uuid
+from concurrent.futures import Future
+from queue import Queue
+from typing import Callable, Dict, List, 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 Result, Data, Error, MIMEType
+
+logger = logging.getLogger(__name__)
+
+
+class CellExecution:
+    """
+    Represents the execution of a cell in the Jupyter kernel.
+    It's an internal class used by JupyterKernelWebSocket.
+    """
+
+    input_accepted: bool = False
+    on_stdout: Optional[Callable[[ProcessMessage], None]] = None
+    on_stderr: Optional[Callable[[ProcessMessage], None]] = None
+    on_display_data: Optional[Callable[[Dict[MIMEType, str]], None]] = None
+
+    def __init__(
+        self,
+        on_stdout: Optional[Callable[[ProcessMessage], None]] = None,
+        on_stderr: Optional[Callable[[ProcessMessage], None]] = None,
+        on_display_data: Optional[Callable[[Dict[MIMEType, str]], None]] = None,
+    ):
+        self.partial_result = Result()
+        self.result = Future()
+        self.on_stdout = on_stdout
+        self.on_stderr = on_stderr
+        self.on_display_data = on_display_data
+
+
+class JupyterKernelWebSocket(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    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 process_messages(self):
+        while True:
+            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"
+        ).start()
+
+        threading.Thread(
+            target=WebSocket(
+                url=self.url,
+                queue_in=self._queue_in,
+                queue_out=self._queue_out,
+                started=started,
+                stopped=stopped,
+            ).run,
+            daemon=True,
+            name="e2b-code-interpreter-websocket",
+        ).start()
+
+        logger.debug("WebSocket waiting to start")
+
+        try:
+            start_time = time.time()
+            while (
+                not started.is_set()
+                and time.time() - start_time < timeout
+                and not self._closed
+            ):
+                time.sleep(0.1)
+
+            if not started.is_set():
+                raise TimeoutException("WebSocket failed to start")
+        except BaseException as e:
+            self.close()
+            raise Exception(f"WebSocket failed to start: {e}") from e
+
+        logger.debug("WebSocket started")
+
+    @staticmethod
+    def _get_execute_request(msg_id: str, code: str) -> str:
+        return json.dumps(
+            {
+                "header": {
+                    "msg_id": msg_id,
+                    "username": "e2b",
+                    "session": str(uuid.uuid4()),
+                    "msg_type": "execute_request",
+                    "version": "5.3",
+                },
+                "parent_header": {},
+                "metadata": {},
+                "content": {
+                    "code": code,
+                    "silent": False,
+                    "store_history": False,
+                    "user_expressions": {},
+                    "allow_stdin": False,
+                },
+            }
+        )
+
+    def send_execution_message(
+        self,
+        code: str,
+        on_stdout: Optional[Callable[[ProcessMessage], None]] = None,
+        on_stderr: Optional[Callable[[ProcessMessage], None]] = None,
+        on_display_data: Optional[Callable[[Dict[MIMEType, str]], None]] = None,
+    ) -> str:
+        message_id = str(uuid.uuid4())
+        logger.debug(f"Sending execution message: {message_id}")
+
+        self._cells[message_id] = CellExecution(
+            on_stdout=on_stdout,
+            on_stderr=on_stderr,
+            on_display_data=on_display_data,
+        )
+        request = self._get_execute_request(message_id, code)
+        self._queue_in.put(request)
+        return message_id
+
+    def get_result(self, message_id: str, timeout: Optional[float] = TIMEOUT) -> Result:
+        result = self._cells[message_id].result.result(timeout=timeout)
+        logger.debug(f"Got result for message: {message_id}")
+        del self._cells[message_id]
+        return result
+
+    def _receive_message(self, data: dict):
+        """
+        Process messages from the WebSocket
+
+        Message types:
+        https://jupyter-client.readthedocs.io/en/stable/messaging.html
+
+        :param data: The message data
+        """
+        parent_msg_ig = data["parent_header"]["msg_id"]
+        logger.debug(f"Received message {data['msg_type']} for {parent_msg_ig}")
+
+        cell = self._cells.get(parent_msg_ig)
+        if not cell:
+            return
+
+        result = cell.partial_result
+
+        if data["msg_type"] == "error":
+            logger.debug(f"Cell {parent_msg_ig} finished execution with error")
+            result.error = Error(
+                name=data["content"]["ename"],
+                value=data["content"]["evalue"],
+                traceback_raw=data["content"]["traceback"],
+            )
+
+        elif data["msg_type"] == "stream":
+            if data["content"]["name"] == "stdout":
+                result.logs.stdout.append(data["content"]["text"])
+                if cell.on_stdout:
+                    cell.on_stdout(
+                        ProcessMessage(
+                            line=data["content"]["text"],
+                            timestamp=time.time_ns(),
+                        )
+                    )
+
+            elif data["content"]["name"] == "stderr":
+                result.logs.stderr.append(data["content"]["text"])
+                if cell.on_stderr:
+                    cell.on_stderr(
+                        ProcessMessage(
+                            line=data["content"]["text"],
+                            error=True,
+                            timestamp=time.time_ns(),
+                        )
+                    )
+
+        elif data["msg_type"] in "display_data":
+            result.data.append(Data(is_main_result=False, data=data["content"]["data"]))
+            if cell.on_display_data:
+                cell.on_display_data(data["content"]["data"])
+        elif data["msg_type"] == "execute_result":
+            result.data.append(Data(is_main_result=True, data=data["content"]["data"]))
+        elif data["msg_type"] == "status":
+            if data["content"]["execution_state"] == "idle":
+                if cell.input_accepted:
+                    logger.debug(f"Cell {parent_msg_ig} finished execution")
+                    cell.result.set_result(result)
+
+            elif data["content"]["execution_state"] == "error":
+                logger.debug(f"Cell {parent_msg_ig} finished execution with error")
+                result.error = Error(
+                    name=data["content"]["ename"],
+                    value=data["content"]["evalue"],
+                    traceback_raw=data["content"]["traceback"],
+                )
+                cell.result.set_result(result)
+
+        elif data["msg_type"] == "execute_reply":
+            if data["content"]["status"] == "error":
+                logger.debug(f"Cell {parent_msg_ig} finished execution with error")
+                result.error = Error(
+                    name=data["content"]["ename"],
+                    value=data["content"]["evalue"],
+                    traceback_raw=data["content"]["traceback"],
+                )
+            elif data["content"]["status"] == "ok":
+                pass
+
+        elif data["msg_type"] == "execute_input":
+            logger.debug(f"Input accepted for {parent_msg_ig}")
+            cell.input_accepted = True
+        else:
+            logger.error(f"[UNHANDLED MESSAGE TYPE]: {data['msg_type']}")
+            print("[UNHANDLED MESSAGE TYPE]:", data["msg_type"])
+
+    def close(self):
+        logger.debug("Closing WebSocket")
+        self._closed = True
+
+        for cancel in self._process_cleanup:
+            cancel()
+
+        self._process_cleanup.clear()
+
+        for handler in self._waiting_for_replies.values():
+            logger.debug(f"Cancelling waiting for execution result for {handler}")
+            handler.cancel()
+            del handler

e2b_code_interpreter-0.0.1a0/e2b_code_interpreter/models.py

@@ -0,0 +1,215 @@
+import copy
+from typing import List, Optional, Iterable, Dict
+from pydantic import BaseModel
+
+
+class Error(BaseModel):
+    """
+    Represents an error that occurred during the execution of a cell.
+    The error contains the name of the error, the value of the error, and the traceback.
+    """
+
+    name: str
+    "Name of the exception."
+    value: str
+    "Value of the exception."
+    traceback_raw: List[str]
+    "List of strings representing the traceback."
+
+    @property
+    def traceback(self) -> str:
+        """
+        Returns the traceback as a single string.
+
+        :return: The traceback as a single string.
+        """
+        return "\n".join(self.traceback_raw)
+
+
+class MIMEType(str):
+    """
+    Represents a MIME type.
+    """
+
+
+class Data:
+    """
+    Represents the data to be displayed as a result of executing a cell in a Jupyter notebook.
+    This is result returned by ipython kernel: https://ipython.readthedocs.io/en/stable/development/execution.html#execution-semantics
+
+
+    The result can contain multiple types of data, such as text, images, plots, etc. Each type of data is represented
+    as a string, and the result can contain multiple types of data. The text representation is always present, and
+    the other representations are optional.
+
+    The class also provides methods to display the data in a Jupyter notebook.
+    """
+    text: str
+    "Text representation of the data. Always present."
+    html: Optional[str] = None
+    markdown: Optional[str] = None
+    svg: Optional[str] = None
+    png: Optional[str] = None
+    jpeg: Optional[str] = None
+    pdf: Optional[str] = None
+    latex: Optional[str] = None
+    json: Optional[dict] = None
+    javascript: Optional[str] = None
+    extra: Optional[dict] = None
+    "Extra data that can be included. Not part of the standard types."
+
+    is_main_result: bool
+    "Whether this data is the result of the cell. Data can be produced by display calls of which can be multiple in a cell."
+
+    raw: Dict[MIMEType, str]
+    "Dictionary that maps MIME types to their corresponding string representations of the data."
+
+    def __init__(self, is_main_result: bool, data: [MIMEType, str]):
+        self.is_main_result = is_main_result
+        self.raw = copy.deepcopy(data)
+
+        self.text = data.pop("text/plain")
+        self.html = data.pop("text/html", None)
+        self.markdown = data.pop("text/markdown", None)
+        self.svg = data.pop("image/svg+xml", None)
+        self.png = data.pop("image/png", None)
+        self.jpeg = data.pop("image/jpeg", None)
+        self.pdf = data.pop("application/pdf", None)
+        self.latex = data.pop("text/latex", None)
+        self.json = data.pop("application/json", None)
+        self.javascript = data.pop("application/javascript", None)
+        self.extra = data
+
+    def keys(self) -> Iterable[str]:
+        """
+        Returns the MIME types of the data.
+
+        :return: The MIME types of the data.
+        """
+        return self.raw.keys()
+
+    def __str__(self) -> str:
+        """
+        Returns the text representation of the data.
+
+        :return: The text representation of the data.
+        """
+        return self.text
+
+    def _repr_html_(self) -> str:
+        """
+        Returns the HTML representation of the data.
+
+        :return: The HTML representation of the data.
+        """
+        return self.html
+
+    def _repr_markdown_(self) -> str:
+        """
+        Returns the Markdown representation of the data.
+
+        :return: The Markdown representation of the data.
+        """
+        return self.markdown
+
+    def _repr_svg_(self) -> str:
+        """
+        Returns the SVG representation of the data.
+
+        :return: The SVG representation of the data.
+        """
+        return self.svg
+
+    def _repr_png_(self) -> str:
+        """
+        Returns the base64 representation of the PNG data.
+
+        :return: The base64 representation of the PNG data.
+        """
+        return self.png
+
+    def _repr_jpeg_(self) -> str:
+        """
+        Returns the base64 representation of the JPEG data.
+
+        :return: The base64 representation of the JPEG data.
+        """
+        return self.jpeg
+
+    def _repr_pdf_(self) -> str:
+        """
+        Returns the PDF representation of the data.
+
+        :return: The PDF representation of the data.
+        """
+        return self.pdf
+
+    def _repr_latex_(self) -> str:
+        """
+        Returns the LaTeX representation of the data.
+
+        :return: The LaTeX representation of the data.
+        """
+        return self.latex
+
+    def _repr_json_(self) -> dict:
+        """
+        Returns the JSON representation of the data.
+
+        :return: The JSON representation of the data.
+        """
+        return self.json
+
+    def _repr_javascript_(self) -> str:
+        """
+        Returns the JavaScript representation of the data.
+
+        :return: The JavaScript representation of the data.
+        """
+        return self.javascript
+
+
+class Logs(BaseModel):
+    """
+    Data printed to stdout and stderr during execution, usually by print statements, logs, warnings, subprocesses, etc.
+    """
+
+    stdout: List[str] = []
+    "List of strings printed to stdout by prints, subprocesses, etc."
+    stderr: List[str] = []
+    "List of strings printed to stderr by prints, subprocesses, etc."
+
+
+class Result(BaseModel):
+    """
+    Represents the result of a cell execution.
+    """
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    data: List[Data] = []
+    "List of result of the cell (interactively interpreted last line), display calls, e.g. matplotlib plots."
+    logs: Logs = Logs()
+    "Logs printed to stdout and stderr during execution."
+    error: Optional[Error] = None
+    "Error object if an error occurred, None otherwise."
+
+    @property
+    def text(self) -> Optional[str]:
+        """
+        Returns the text representation of the result.
+
+        :return: The text representation of the result.
+        """
+        for d in self.data:
+            if d.is_main_result:
+                return d.text
+
+
+class KernelException(Exception):
+    """
+    Exception raised when a kernel operation fails.
+    """
+
+    pass

e2b_code_interpreter-0.0.1a0/pyproject.toml

@@ -0,0 +1,31 @@
+[tool.poetry]
+name = "e2b-code-interpreter"
+version = "0.0.1a0"
+description = "E2B Code Interpreter - Stateful code execution"
+authors = ["e2b <hello@e2b.dev>"]
+license = "Apache-2.0"
+readme = "README.md"
+homepage = "https://e2b.dev/"
+repository = "https://github.com/e2b-dev/e2b-code-interpreter/tree/python"
+packages = [{ include = "e2b_code_interpreter" }]
+
+[tool.poetry.dependencies]
+python = "^3.8"
+
+e2b = ">=0.14.11"
+pydantic = ">1, <3"
+websocket-client = "^1.7.0"
+
+[tool.poetry.group.dev.dependencies]
+black = "^24.3.0"
+pytest = "^7.4.0"
+python-dotenv = "^1.0.0"
+pytest-dotenv = "^0.5.2"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry.urls]
+"Bug Tracker" = "https://github.com/e2b-dev/code-interpreter/issues"