diff --git a/.github/workflows/template.yml b/.github/workflows/template.yml index a18b9fd..80a84f5 100644 --- a/.github/workflows/template.yml +++ b/.github/workflows/template.yml @@ -47,4 +47,4 @@ jobs: - name: Build e2b run: e2b template build env: - E2B_ACCESS_TOKEN: ${{ secrets.E2B_ACCESS_TOKEN }} \ No newline at end of file + E2B_ACCESS_TOKEN: ${{ secrets.E2B_ACCESS_TOKEN }} diff --git a/js/example.mts b/js/example.mts index 4afc07d..3fe5d65 100644 --- a/js/example.mts +++ b/js/example.mts @@ -15,7 +15,6 @@ plt.plot(x, y) plt.show() x = np.linspace(0, 10, 100) - plt.plot(x, y) plt.show() @@ -23,7 +22,7 @@ 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) diff --git a/js/package.json b/js/package.json index efdf53e..7eb8183 100644 --- a/js/package.json +++ b/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", diff --git a/js/src/messaging.ts b/js/src/messaging.ts index 8c22a89..960ef66 100644 --- a/js/src/messaging.ts +++ b/js/src/messaging.ts @@ -59,12 +59,16 @@ export class ExecutionError { */ export type MIMEType = string +type Data = { + data: Record +} + /** - * 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 /** * 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) diff --git a/js/tests/data.test.ts b/js/tests/data.test.ts new file mode 100644 index 0000000..8a638ed --- /dev/null +++ b/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() +}) diff --git a/python/e2b_code_interpreter/models.py b/python/e2b_code_interpreter/models.py index a759b7d..5667a63 100644 --- a/python/e2b_code_interpreter/models.py +++ b/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: diff --git a/python/pyproject.toml b/python/pyproject.toml index a06f9b6..8b070ec 100644 --- a/python/pyproject.toml +++ b/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 "] license = "Apache-2.0" diff --git a/python/tests/async/test_async_data.py b/python/tests/async/test_async_data.py new file mode 100644 index 0000000..f02d9a6 --- /dev/null +++ b/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 diff --git a/python/tests/sync/test_data.py b/python/tests/sync/test_data.py new file mode 100644 index 0000000..0352de4 --- /dev/null +++ b/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 diff --git a/template/e2b.Dockerfile b/template/e2b.Dockerfile index 3856852..0583b94 100644 --- a/template/e2b.Dockerfile +++ b/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 diff --git a/template/ipython_startup_script.py b/template/ipython_startup_script.py new file mode 100644 index 0000000..6565d69 --- /dev/null +++ b/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_) diff --git a/template/server/api/models/result.py b/template/server/api/models/result.py index 6e1f833..178f891 100644 --- a/template/server/api/models/result.py +++ b/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]: