Merge pull request #33 from e2b-dev/beta-df

Extract data from DataFrames

.github/workflows/template.yml

@@ -47,4 +47,4 @@ jobs:
       - name: Build e2b
         run: e2b template build
         env:
-          E2B_ACCESS_TOKEN: ${{ secrets.E2B_ACCESS_TOKEN }}
+          E2B_ACCESS_TOKEN: ${{ secrets.E2B_ACCESS_TOKEN }}

js/example.mts

@@ -15,15 +15,14 @@ plt.plot(x, y)
 plt.show()
 
 x = np.linspace(0, 10, 100)
-
 plt.plot(x, y)
 plt.show()
 
 import pandas
 pandas.DataFrame({"a": [1, 2, 3]})
 `
 
-const sandbox = await CodeInterpreter.connect("", { debug: true })
+const sandbox = await CodeInterpreter.create()
 console.log(sandbox.sandboxId)
 
 const execution = await sandbox.notebook.execCell(code, {
@@ -35,4 +34,5 @@ const execution = await sandbox.notebook.execCell(code, {
   },
 })
 console.log(execution.results[0].formats())
+console.log(execution.results[0].data)
 console.log(execution.results.length)

js/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@e2b/code-interpreter",
-  "version": "0.0.9-beta.39",
+  "version": "0.0.9-beta-df.5",
   "description": "E2B Code Interpreter - Stateful code execution",
   "homepage": "https://e2b.dev",
   "license": "MIT",

js/src/messaging.ts

@@ -59,12 +59,16 @@ export class ExecutionError {
  */
 export type MIMEType = string
 
+type Data = {
+  data: Record<string, unknown>
+}
+
 /**
- * Dictionary that maps MIME types to their corresponding string representations of the data.
+ * Dictionary that maps MIME types to their corresponding representations of the data.
  */
 export type RawData = {
   [key: MIMEType]: string
-}
+} & Data
 
 /**
  * Represents the data to be displayed as a result of executing a cell in a Jupyter notebook.
@@ -116,6 +120,10 @@ export class Result {
    * JavaScript representation of the data.
    */
   readonly javascript?: string
+  /**
+   * Contains the data from DataFrame.
+   */
+  readonly data?: Record<string, unknown>
   /**
    * Extra data that can be included. Not part of the standard types.
    */
@@ -138,6 +146,7 @@ export class Result {
     this.latex = data['latex']
     this.json = data['json']
     this.javascript = data['javascript']
+    this.data = data['data']
     this.isMainResult = isMainResult
     this.raw = data
 
@@ -155,7 +164,9 @@ export class Result {
           'pdf',
           'latex',
           'json',
-          'javascript'
+          'javascript',
+          'data',
+          'extra'
         ].includes(key)
       ) {
         this.extra[key] = data[key]
@@ -197,6 +208,9 @@ export class Result {
     if (this.javascript) {
       formats.push('javascript')
     }
+    if (this.data) {
+      formats.push('data')
+    }
 
     for (const key of Object.keys(this.extra)) {
       formats.push(key)

js/tests/data.test.ts

@@ -0,0 +1,10 @@
+import { expect } from 'vitest'
+
+import { sandboxTest } from './setup'
+
+sandboxTest('get data', async ({ sandbox }) => {
+  const execution = await sandbox.notebook.execCell('pd.DataFrame({"a": [1, 2, 3]})')
+
+  const result = execution.results[0]
+  expect(result.data).toBeDefined()
+})

python/e2b_code_interpreter/models.py

@@ -88,6 +88,7 @@ def __getitem__(self, item):
     latex: Optional[str] = None
     json: Optional[dict] = None
     javascript: Optional[str] = None
+    data: Optional[dict] = None
     is_main_result: bool = False
     """Whether this data is the result of the cell. Data can be produced by display calls of which can be multiple in a cell."""
     extra: Optional[dict] = None
@@ -120,6 +121,8 @@ def formats(self) -> Iterable[str]:
             formats.append("json")
         if self.javascript:
             formats.append("javascript")
+        if self.data:
+            formats.append("data")
 
         if self.extra:
             for key in self.extra:

python/pyproject.toml

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

python/tests/async/test_async_data.py

@@ -0,0 +1,16 @@
+from e2b_code_interpreter.code_interpreter_async import AsyncCodeInterpreter
+
+
+async def test_data(async_sandbox: AsyncCodeInterpreter):
+    # plot random graph
+    result = await async_sandbox.notebook.exec_cell(
+        """
+        pd.DataFrame({"a": [1, 2, 3]})
+        """
+    )
+
+    # there's your image
+    data = result.results[0]
+    assert data.data
+    assert "a" in data.data
+    assert len(data.data["a"]) == 3

python/tests/sync/test_data.py

@@ -0,0 +1,16 @@
+from e2b_code_interpreter.code_interpreter_sync import CodeInterpreter
+
+
+def test_data(sandbox: CodeInterpreter):
+    # plot random graph
+    result = sandbox.notebook.exec_cell(
+        """
+        pd.DataFrame({"a": [1, 2, 3]})
+        """
+    )
+
+    # there's your image
+    data = result.results[0]
+    assert data.data
+    assert "a" in data.data
+    assert len(data.data["a"]) == 3

template/e2b.Dockerfile

@@ -52,5 +52,8 @@ COPY ./jupyter_server_config.py $JUPYTER_CONFIG_PATH/
 RUN mkdir -p $IPYTHON_CONFIG_PATH/profile_default
 COPY ipython_kernel_config.py $IPYTHON_CONFIG_PATH/profile_default/
 
+RUN mkdir -p $IPYTHON_CONFIG_PATH/profile_default/startup
+COPY ./ipython_startup_script.py $IPYTHON_CONFIG_PATH/profile_default/startup/00-startup.py
+
 # Setup entrypoint for local development
 ENTRYPOINT $JUPYTER_CONFIG_PATH/start-up.sh

template/ipython_startup_script.py

@@ -0,0 +1,18 @@
+import pandas as pd
+
+
+def _repr_mimebundle_(self, include=None, exclude=None):
+    data = {
+        'text/html': self.to_html(),
+        'text/plain': self.to_string(),
+        'e2b/df': self.to_dict(orient='list'),
+    }
+
+    if include:
+        data = {k: v for (k, v) in data.items() if k in include}
+    if exclude:
+        data = {k: v for (k, v) in data.items() if k not in exclude}
+    return data
+
+
+setattr(pd.DataFrame, '_repr_mimebundle_', _repr_mimebundle_)

template/server/api/models/result.py

@@ -35,6 +35,7 @@ class Result(BaseModel):
     latex: Optional[str] = None
     json: Optional[dict] = None
     javascript: Optional[str] = None
+    data: Optional[dict] = None
     extra: Optional[dict] = None
     "Extra data that can be included. Not part of the standard types."
 
@@ -58,6 +59,7 @@ def __init__(self, is_main_result: bool, data: [str, str]):
         self.latex = data.pop("text/latex", None)
         self.json = data.pop("application/json", None)
         self.javascript = data.pop("application/javascript", None)
+        self.data = data.pop("e2b/df", None)
         self.extra = data
 
     def formats(self) -> Iterable[str]: