diff --git a/js/README.md b/js/README.md index 1ef22070..1dc75b2c 100644 --- a/js/README.md +++ b/js/README.md @@ -18,23 +18,23 @@ npm install @e2b/code-interpreter ### Minimal example with the sharing context ```js -import { CodeInterpreter } from '@e2b/code-interpreter' +import { Sandbox } from '@e2b/code-interpreter' -const sandbox = await CodeInterpreter.create() -await sandbox.notebook.execCell('x = 1') +const sandbox = await Sandbox.create() +await sandbox.runCode('x = 1') -const execution = await sandbox.notebook.execCell('x+=1; x') +const execution = await sandbox.runCode('x+=1; x') console.log(execution.text) // outputs 2 -await sandbox.close() +await sandbox.kill() ``` ### Get charts and any display-able data ```js -import { CodeInterpreter } from '@e2b/code-interpreter' +import { Sandbox } from '@e2b/code-interpreter' -const sandbox = await CodeInterpreter.create() +const sandbox = await Sandbox.create() const code = ` import matplotlib.pyplot as plt @@ -48,20 +48,20 @@ plt.show() ` // you can install dependencies in "jupyter notebook style" -await sandbox.notebook.execCell('!pip install matplotlib') +await sandbox.runCode('!pip install matplotlib') -const execution = await sandbox.notebook.execCell(code) +const execution = await sandbox.runCode(code) // this contains the image data, you can e.g. save it to file or send to frontend execution.results[0].png -await sandbox.close() +await sandbox.kill() ``` ### Streaming code output ```js -import { CodeInterpreter } from '@e2b/code-interpreter' +import { Sandbox } from '@e2b/code-interpreter' const code = ` import time @@ -75,13 +75,13 @@ time.sleep(3) print("world") ` -const sandbox = await CodeInterpreter.create() +const sandbox = await Sandbox.create() -await sandbox.notebook.execCell(code, { +await sandbox.runCode(code, { onStdout: (out) => console.log(out), onStderr: (outErr) => console.error(outErr), onResult: (result) => console.log(result.text), }) -await sandbox.close() +await sandbox.kill() ``` diff --git a/js/example.mts b/js/example.mts index 3fe5d652..217601bd 100644 --- a/js/example.mts +++ b/js/example.mts @@ -1,6 +1,6 @@ import dotenv from 'dotenv' -import { CodeInterpreter } from './dist' +import { Sandbox } from './dist' dotenv.config() @@ -22,10 +22,10 @@ import pandas pandas.DataFrame({"a": [1, 2, 3]}) ` -const sandbox = await CodeInterpreter.create() +const sandbox = await Sandbox.create() console.log(sandbox.sandboxId) -const execution = await sandbox.notebook.execCell(code, { +const execution = await sandbox.runCode(code, { onStdout(msg) { console.log('stdout', msg) }, diff --git a/js/src/consts.ts b/js/src/consts.ts new file mode 100644 index 00000000..535ded60 --- /dev/null +++ b/js/src/consts.ts @@ -0,0 +1,2 @@ +export const DEFAULT_TIMEOUT_MS = 60_000 // 1 minute +export const JUPYTER_PORT = 49999 diff --git a/js/src/index.ts b/js/src/index.ts index 7ac8bd62..80a23823 100644 --- a/js/src/index.ts +++ b/js/src/index.ts @@ -1,7 +1,7 @@ export * from 'e2b' -export { Sandbox, JupyterExtension } from './sandbox' - +export { Sandbox } from './sandbox' +export type { Context } from './sandbox' export type { Logs, ExecutionError, diff --git a/js/src/sandbox.ts b/js/src/sandbox.ts index 2c4bcd7f..51e915b1 100644 --- a/js/src/sandbox.ts +++ b/js/src/sandbox.ts @@ -1,65 +1,19 @@ -import { ConnectionConfig, Sandbox as BaseSandbox, TimeoutError } from 'e2b' +import { Sandbox as BaseSandbox, InvalidArgumentError } from 'e2b' import { Result, Execution, OutputMessage, parseOutput, extractError } from './messaging' - -function formatRequestTimeoutError(error: unknown) { - if (error instanceof Error && error.name === 'AbortError') { - return new TimeoutError('Request timed out — the \'requestTimeoutMs\' option can be used to increase this timeout') - } - - return error -} - -function formatExecutionTimeoutError(error: unknown) { - if (error instanceof Error && error.name === 'AbortError') { - return new TimeoutError('Execution timed out — the \'timeoutMs\' option can be used to increase this timeout') - } - - return error -} - -async function* readLines(stream: ReadableStream) { - const reader = stream.getReader(); - let buffer = '' - - try { - while (true) { - const { done, value } = await reader.read(); - - if (value !== undefined) { - buffer += new TextDecoder().decode(value) - } - - if (done) { - if (buffer.length > 0) { - yield buffer - } - break - } - - let newlineIdx = -1 - - do { - newlineIdx = buffer.indexOf('\n') - if (newlineIdx !== -1) { - yield buffer.slice(0, newlineIdx) - buffer = buffer.slice(newlineIdx + 1) - } - } while (newlineIdx !== -1) - } - } finally { - reader.releaseLock() - } +import { formatExecutionTimeoutError, formatRequestTimeoutError, readLines } from "./utils"; +import { JUPYTER_PORT, DEFAULT_TIMEOUT_MS } from './consts' +export type Context = { + id: string + language: string + cwd: string } /** * Code interpreter module for executing code in a stateful context. */ -export class JupyterExtension { - private static readonly execTimeoutMs = 300_000 - private static readonly defaultKernelID = 'default' - - constructor(private readonly url: string, private readonly connectionConfig: ConnectionConfig) { } +export class Sandbox extends BaseSandbox { + protected static override readonly defaultTemplate: string = 'code-interpreter-beta' /** * Runs the code in the specified context, if not specified, the default context is used. @@ -67,7 +21,8 @@ export class JupyterExtension { * * @param code The code to execute * @param opts Options for executing the code - * @param opts.kernelID The context ID to run the code in + * @param opts.language Based on the value, a default context for the language is used. If not defined and no context is provided, the default Python context is used. + * @param opts.context Concrete context to run the code in. If not specified, the default context for the language is used. It's mutually exclusive with the language. * @param opts.onStdout Callback for handling stdout messages * @param opts.onStderr Callback for handling stderr messages * @param opts.onResult Callback for handling the final result @@ -76,10 +31,11 @@ export class JupyterExtension { * @param opts.requestTimeoutMs Max time to wait for the request to finish * @returns Execution object */ - async execCell( + async runCode( code: string, opts?: { - kernelID?: string, + language?: string, + context?: Context, onStdout?: (output: OutputMessage) => (Promise | any), onStderr?: (output: OutputMessage) => (Promise | any), onResult?: (data: Result) => (Promise | any), @@ -88,6 +44,10 @@ export class JupyterExtension { requestTimeoutMs?: number, }, ): Promise { + if (opts?.context && opts?.language) { + throw new InvalidArgumentError("You can provide context or language, but not both at the same time.") + } + const controller = new AbortController() const requestTimeout = opts?.requestTimeoutMs ?? this.connectionConfig.requestTimeoutMs @@ -98,14 +58,15 @@ export class JupyterExtension { : undefined try { - const res = await fetch(`${this.url}/execute`, { + const res = await fetch(`${this.jupyterUrl}/execute`, { method: 'POST', headers: { 'Content-Type': 'application/json', }, body: JSON.stringify({ code, - context_id: opts?.kernelID, + context_id: opts?.context?.id, + language: opts?.language, env_vars: opts?.envs, }), signal: controller.signal, @@ -123,7 +84,7 @@ export class JupyterExtension { clearTimeout(reqTimer) - const bodyTimeout = opts?.timeoutMs ?? JupyterExtension.execTimeoutMs + const bodyTimeout = opts?.timeoutMs ?? DEFAULT_TIMEOUT_MS const bodyTimer = bodyTimeout ? setTimeout(() => { @@ -154,28 +115,28 @@ export class JupyterExtension { * Creates a new context to run code in. * * @param cwd The working directory for the context - * @param kernelName The name of the context + * @param language The name of the context * @param requestTimeoutMs Max time to wait for the request to finish * @returns The context ID */ - async createKernel({ + async createCodeContext({ cwd, - kernelName, + language, requestTimeoutMs, }: { cwd?: string, - kernelName?: string, + language?: string, requestTimeoutMs?: number, - } = {}): Promise { + } = {}): Promise { try { - const res = await fetch(`${this.url}/contexts`, { + const res = await fetch(`${this.jupyterUrl}/contexts`, { method: 'POST', headers: { 'Content-Type': 'application/json', }, body: JSON.stringify({ - name: kernelName, + language: language, cwd, }), keepalive: true, @@ -187,117 +148,13 @@ export class JupyterExtension { throw error } - const data = await res.json() - return data.id + return await res.json() } catch (error) { throw formatRequestTimeoutError(error) } } - /** - * Restarts the context. - * Restarting will clear all variables, imports, and other settings set during previous executions. - * - * @param kernelID The context ID to restart - * @param requestTimeoutMs Max time to wait for the request to finish - */ - async restartKernel({ - kernelID, - requestTimeoutMs, - }: { - kernelID?: string, - requestTimeoutMs?: number, - } = {}): Promise { - try { - kernelID = kernelID || JupyterExtension.defaultKernelID - const res = await fetch(`${this.url}/contexts/${kernelID}/restart`, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - }, - keepalive: true, - signal: this.connectionConfig.getSignal(requestTimeoutMs), - }) - - const error = await extractError(res) - if (error) { - throw error - } - } catch (error) { - throw formatRequestTimeoutError(error) - } + protected get jupyterUrl(): string { + return `${this.connectionConfig.debug ? 'http' : 'https'}://${this.getHost(JUPYTER_PORT)}` } - - /** - * Shuts down the context. - * - * @param kernelID The context ID to shut down - * @param requestTimeoutMs Max time to wait for the request to finish - */ - async shutdownKernel({ - kernelID, - requestTimeoutMs, - }: { - kernelID?: string, - requestTimeoutMs?: number, - } = {}): Promise { - try { - - kernelID = kernelID || JupyterExtension.defaultKernelID - - const res = await fetch(`${this.url}/contexts/${kernelID}`, { - method: 'DELETE', - keepalive: true, - signal: this.connectionConfig.getSignal(requestTimeoutMs), - }) - - const error = await extractError(res) - if (error) { - throw error - } - } catch (error) { - throw formatRequestTimeoutError(error) - } - } - - /** - * Lists all available contexts. - * - * @param requestTimeoutMs Max time to wait for the request to finish - * @returns List of context IDs and names - */ - async listKernels({ - requestTimeoutMs, - }: { - requestTimeoutMs?: number, - } = {}): Promise<{ kernelID: string, name: string }[]> { - try { - const res = await fetch(`${this.url}/contexts`, { - keepalive: true, - signal: this.connectionConfig.getSignal(requestTimeoutMs), - }) - - const error = await extractError(res) - if (error) { - throw error - } - - return (await res.json()).map((kernel: any) => ({ kernelID: kernel.id, name: kernel.name })) - } catch (error) { - throw formatRequestTimeoutError(error) - } - } -} - -/** - * Code interpreter module for executing code in a stateful context. - */ -export class Sandbox extends BaseSandbox { - protected static override readonly defaultTemplate: string = 'code-interpreter-beta' - protected static readonly jupyterPort = 49999 - - readonly notebook = new JupyterExtension( - `${this.connectionConfig.debug ? 'http' : 'https'}://${this.getHost(Sandbox.jupyterPort)}`, - this.connectionConfig, - ) } diff --git a/js/src/utils.ts b/js/src/utils.ts new file mode 100644 index 00000000..4910d5de --- /dev/null +++ b/js/src/utils.ts @@ -0,0 +1,51 @@ +import { TimeoutError } from 'e2b' + +export function formatRequestTimeoutError(error: unknown) { + if (error instanceof Error && error.name === 'AbortError') { + return new TimeoutError('Request timed out — the \'requestTimeoutMs\' option can be used to increase this timeout') + } + + return error +} + +export function formatExecutionTimeoutError(error: unknown) { + if (error instanceof Error && error.name === 'AbortError') { + return new TimeoutError('Execution timed out — the \'timeoutMs\' option can be used to increase this timeout') + } + + return error +} + +export async function* readLines(stream: ReadableStream) { + const reader = stream.getReader(); + let buffer = '' + + try { + while (true) { + const { done, value } = await reader.read(); + + if (value !== undefined) { + buffer += new TextDecoder().decode(value) + } + + if (done) { + if (buffer.length > 0) { + yield buffer + } + break + } + + let newlineIdx = -1 + + do { + newlineIdx = buffer.indexOf('\n') + if (newlineIdx !== -1) { + yield buffer.slice(0, newlineIdx) + buffer = buffer.slice(newlineIdx + 1) + } + } while (newlineIdx !== -1) + } + } finally { + reader.releaseLock() + } +} diff --git a/js/tests/bash.test.ts b/js/tests/bash.test.ts index ff20756a..7734499f 100644 --- a/js/tests/bash.test.ts +++ b/js/tests/bash.test.ts @@ -4,7 +4,7 @@ import { isDebug, sandboxTest } from './setup' // Skip this test if we are running in debug mode — the pwd and user in the testing docker container are not the same as in the actual sandbox. sandboxTest.skipIf(isDebug)('bash', async ({ sandbox }) => { - const result = await sandbox.notebook.execCell('!pwd') + const result = await sandbox.runCode('!pwd') expect(result.logs.stdout.join().trim()).toEqual('/home/user') }) diff --git a/js/tests/basic.test.ts b/js/tests/basic.test.ts index eb5a8650..c0360d02 100644 --- a/js/tests/basic.test.ts +++ b/js/tests/basic.test.ts @@ -3,7 +3,7 @@ import { expect } from 'vitest' import { sandboxTest } from './setup' sandboxTest('basic', async ({ sandbox }) => { - const result = await sandbox.notebook.execCell('x =1; x') + const result = await sandbox.runCode('x =1; x') expect(result.text).toEqual('1') }) diff --git a/js/tests/benchmarking.js b/js/tests/benchmarking.js index e94792a9..f7ba5bbf 100644 --- a/js/tests/benchmarking.js +++ b/js/tests/benchmarking.js @@ -15,14 +15,14 @@ async function main() { createSandboxTime += new Date() - startTime startTime = new Date() - await sandbox.notebook.execCell('x = 1') + await sandbox.runCode('x = 1') fistExecTime += new Date() - startTime startTime = new Date() - const result = await sandbox.notebook.execCell('x+=1; x') + const result = await sandbox.runCode('x+=1; x') secondExecTime += new Date() - startTime - await sandbox.close() + await sandbox.kill() } console.log('Average create sandbox time:', createSandboxTime / iterations) console.log('Average first exec time:', fistExecTime / iterations) diff --git a/js/tests/data.test.ts b/js/tests/data.test.ts index aede4587..b18eeedc 100644 --- a/js/tests/data.test.ts +++ b/js/tests/data.test.ts @@ -3,7 +3,7 @@ import { expect } from 'vitest' import { sandboxTest } from './setup' sandboxTest('get data', async ({ sandbox }) => { - const execution = await sandbox.notebook.execCell(` + const execution = await sandbox.runCode(` import pandas as pd pd.DataFrame({"a": [1, 2, 3]}) `) diff --git a/js/tests/defaultKernels.test.ts b/js/tests/defaultKernels.test.ts new file mode 100644 index 00000000..4c0779d1 --- /dev/null +++ b/js/tests/defaultKernels.test.ts @@ -0,0 +1,9 @@ +import { expect } from 'vitest' + +import { sandboxTest } from './setup' + +sandboxTest('test js kernel', async ({ sandbox }) => { + const output = await sandbox.runCode('console.log("Hello World!")', { language: 'js' }) + console.log(output) + expect(output.logs.stdout).toEqual(['Hello World!\n']) +}) diff --git a/js/tests/displayData.test.ts b/js/tests/displayData.test.ts index 6c145c0b..5e4684d1 100644 --- a/js/tests/displayData.test.ts +++ b/js/tests/displayData.test.ts @@ -4,7 +4,7 @@ import { sandboxTest } from './setup' sandboxTest('display data', async ({ sandbox }) => { // plot random graph - const result = await sandbox.notebook.execCell(` + const result = await sandbox.runCode(` import matplotlib.pyplot as plt import numpy as np diff --git a/js/tests/envVars.test.ts b/js/tests/envVars.test.ts index 73e3c9e3..ad7d8fa2 100644 --- a/js/tests/envVars.test.ts +++ b/js/tests/envVars.test.ts @@ -8,7 +8,7 @@ sandboxTest.skipIf(isDebug)('env vars', async () => { const sandbox = await Sandbox.create({ envs: { TEST_ENV_VAR: 'supertest' }, }) - const result = await sandbox.notebook.execCell( + const result = await sandbox.runCode( `import os; x = os.getenv('TEST_ENV_VAR'); x` ) @@ -16,7 +16,7 @@ sandboxTest.skipIf(isDebug)('env vars', async () => { }) sandboxTest('env vars on sandbox', async ({ sandbox }) => { - const result = await sandbox.notebook.execCell( + const result = await sandbox.runCode( "import os; os.getenv('FOO')", { envs: { FOO: 'bar' } } ) @@ -28,28 +28,28 @@ sandboxTest('env vars on sandbox override', async () => { const sandbox = await Sandbox.create({ envs: { FOO: 'bar', SBX: 'value' }, }) - await sandbox.notebook.execCell( + await sandbox.runCode( "import os; os.environ['FOO'] = 'bar'; os.environ['RUNTIME_ENV'] = 'js_runtime'" ) - const result = await sandbox.notebook.execCell( + const result = await sandbox.runCode( "import os; os.getenv('FOO')", { envs: { FOO: 'baz' } } ) expect(result.results[0].text.trim()).toEqual('baz') - const result2 = await sandbox.notebook.execCell( + const result2 = await sandbox.runCode( "import os; os.getenv('RUNTIME_ENV')" ) expect(result2.results[0].text.trim()).toEqual('js_runtime') if (!isDebug) { - const result3 = await sandbox.notebook.execCell( + const result3 = await sandbox.runCode( "import os; os.getenv('SBX')" ) expect(result3.results[0].text.trim()).toEqual('value') } - const result4 = await sandbox.notebook.execCell("import os; os.getenv('FOO')") + const result4 = await sandbox.runCode("import os; os.getenv('FOO')") expect(result4.results[0].text.trim()).toEqual('bar') }) diff --git a/js/tests/executionCount.test.ts b/js/tests/executionCount.test.ts index 7baef60e..750aa915 100644 --- a/js/tests/executionCount.test.ts +++ b/js/tests/executionCount.test.ts @@ -4,8 +4,8 @@ import { isDebug, sandboxTest } from './setup' // Skip this test if we are running in debug mode — we don't create new sandbox for each test so the execution number is not reset. sandboxTest.skipIf(isDebug)('execution count', async ({ sandbox }) => { - await sandbox.notebook.execCell('!pwd') - const result = await sandbox.notebook.execCell('!pwd') + await sandbox.runCode('!pwd') + const result = await sandbox.runCode('!pwd') expect(result.executionCount).toEqual(2) }) diff --git a/js/tests/graphs/bar.test.ts b/js/tests/graphs/bar.test.ts index c812906f..d4f56587 100644 --- a/js/tests/graphs/bar.test.ts +++ b/js/tests/graphs/bar.test.ts @@ -22,7 +22,7 @@ plt.title('Book Sales by Authors') plt.tight_layout() plt.show() ` - const result = await sandbox.notebook.execCell(code) + const result = await sandbox.runCode(code) const graph = result.results[0].graph expect(graph).toBeDefined() diff --git a/js/tests/graphs/boxAndWhisker.test.ts b/js/tests/graphs/boxAndWhisker.test.ts index 80398965..802cae49 100644 --- a/js/tests/graphs/boxAndWhisker.test.ts +++ b/js/tests/graphs/boxAndWhisker.test.ts @@ -35,7 +35,7 @@ ax.legend() plt.tight_layout() plt.show() ` - const result = await sandbox.notebook.execCell(code) + const result = await sandbox.runCode(code) const graph = result.results[0].graph expect(graph).toBeDefined() diff --git a/js/tests/graphs/line.test.ts b/js/tests/graphs/line.test.ts index e9ed1714..69b15d6b 100644 --- a/js/tests/graphs/line.test.ts +++ b/js/tests/graphs/line.test.ts @@ -29,7 +29,7 @@ plt.title('Plot of sin(x) and cos(x)') # Display the plot plt.show() ` - const result = await sandbox.notebook.execCell(code) + const result = await sandbox.runCode(code) const graph = result.results[0].graph expect(graph).toBeDefined() diff --git a/js/tests/graphs/log.test.ts b/js/tests/graphs/log.test.ts index 6572dc72..f2a1e33c 100644 --- a/js/tests/graphs/log.test.ts +++ b/js/tests/graphs/log.test.ts @@ -29,7 +29,7 @@ plt.grid(True) plt.show() ` - const result = await sandbox.notebook.execCell(code) + const result = await sandbox.runCode(code) const graph = result.results[0].graph expect(graph).toBeDefined() expect(graph.type).toBe('line') diff --git a/js/tests/graphs/pie.test.ts b/js/tests/graphs/pie.test.ts index 9286a026..72830243 100644 --- a/js/tests/graphs/pie.test.ts +++ b/js/tests/graphs/pie.test.ts @@ -27,7 +27,7 @@ plt.title('Will I wake up early tomorrow?') # Step 5: Show the plot plt.show() ` - const result = await sandbox.notebook.execCell(code) + const result = await sandbox.runCode(code) const graph = result.results[0].graph expect(graph).toBeDefined() diff --git a/js/tests/graphs/scales.test.ts b/js/tests/graphs/scales.test.ts index 8599e34f..dd89fe85 100644 --- a/js/tests/graphs/scales.test.ts +++ b/js/tests/graphs/scales.test.ts @@ -17,7 +17,7 @@ sandboxTest('datetime scale', async ({ sandbox }) => { plt.show() ` - const result = await sandbox.notebook.execCell(code) + const result = await sandbox.runCode(code) const graph = result.results[0].graph expect(graph).toBeDefined() @@ -41,7 +41,7 @@ sandboxTest('categorical scale', async ({ sandbox }) => { plt.show() ` - const result = await sandbox.notebook.execCell(code) + const result = await sandbox.runCode(code) const graph = result.results[0].graph expect(graph).toBeTruthy() diff --git a/js/tests/graphs/scatter.test.ts b/js/tests/graphs/scatter.test.ts index b9b68b18..7708e43b 100644 --- a/js/tests/graphs/scatter.test.ts +++ b/js/tests/graphs/scatter.test.ts @@ -22,7 +22,7 @@ plt.scatter(x2, y2, c='red', label='Dataset 2') plt.show() ` - const result = await sandbox.notebook.execCell(code) + const result = await sandbox.runCode(code) const graph = result.results[0].graph expect(graph).toBeDefined() diff --git a/js/tests/graphs/supergraph.test.ts b/js/tests/graphs/supergraph.test.ts index fbaed47d..74426bb3 100644 --- a/js/tests/graphs/supergraph.test.ts +++ b/js/tests/graphs/supergraph.test.ts @@ -32,7 +32,7 @@ axs[1].grid(True) plt.show() ` - const result = await sandbox.notebook.execCell(code) + const result = await sandbox.runCode(code) const graph = result.results[0].graph expect(graph).toBeDefined() diff --git a/js/tests/graphs/unknown.test.ts b/js/tests/graphs/unknown.test.ts index 70939c8b..842d05ec 100644 --- a/js/tests/graphs/unknown.test.ts +++ b/js/tests/graphs/unknown.test.ts @@ -28,7 +28,7 @@ plt.title('Two Concentric Circles') # Show the plot plt.show() ` - const result = await sandbox.notebook.execCell(code) + const result = await sandbox.runCode(code) const graph = result.results[0].graph expect(graph).toBeDefined() diff --git a/js/tests/kernels.test.ts b/js/tests/kernels.test.ts index 99518130..810a9c23 100644 --- a/js/tests/kernels.test.ts +++ b/js/tests/kernels.test.ts @@ -1,46 +1,20 @@ import { expect } from 'vitest' -import { isDebug, sandboxTest } from './setup' +import { sandboxTest } from './setup' sandboxTest('create new kernel', async ({ sandbox }) => { - await sandbox.notebook.createKernel() + await sandbox.createCodeContext() }) sandboxTest('independence of kernels', async ({ sandbox }) => { - await sandbox.notebook.execCell('x = 1') - const kernelID = await sandbox.notebook.createKernel() - const output = await sandbox.notebook.execCell('x', { kernelID }) + await sandbox.runCode('x = 1') + const context = await sandbox.createCodeContext() + const output = await sandbox.runCode('x', { context }) expect(output.error!.value).toEqual("name 'x' is not defined") }) -sandboxTest('restart kernel', async ({ sandbox }) => { - await sandbox.notebook.execCell('x = 1') - await sandbox.notebook.restartKernel() - - const output = await sandbox.notebook.execCell('x') - - expect(output.error!.value).toEqual("name 'x' is not defined") -}) - -// Skip this test if we are running in debug mode — we don't know how many kernels are in the local debug testing container. -sandboxTest.skipIf(isDebug)('list kernels', async ({ sandbox }) => { - let kernels = await sandbox.notebook.listKernels() - expect(kernels.length).toEqual(1) - - const kernelID = await sandbox.notebook.createKernel() - kernels = await sandbox.notebook.listKernels() - expect(kernels.map((kernel) => kernel.kernelID)).toContain(kernelID) - expect(kernels.length).toEqual(2) -}) - -// Skip this test if we are running in debug mode — we don't know how many kernels are in the local debug testing container. -sandboxTest.skipIf(isDebug)('shutdown kernel', async ({ sandbox }) => { - let kernels = await sandbox.notebook.listKernels() - expect(kernels.length).toEqual(1) - - const kernelID = await sandbox.notebook.shutdownKernel() - kernels = await sandbox.notebook.listKernels() - expect(kernels).not.toContain(kernelID) - expect(kernels.length).toEqual(0) +sandboxTest('pass context and language', async ({ sandbox }) => { + const context = await sandbox.createCodeContext() + await expect(sandbox.runCode({context, language: 'python'})).rejects.toThrowError() }) diff --git a/js/tests/reconnect.test.ts b/js/tests/reconnect.test.ts index 0f97e450..c80e3cf8 100644 --- a/js/tests/reconnect.test.ts +++ b/js/tests/reconnect.test.ts @@ -6,7 +6,7 @@ import { sandboxTest } from './setup' sandboxTest('reconnect', async ({ sandbox }) => { sandbox = await Sandbox.connect(sandbox.sandboxId) - const result = await sandbox.notebook.execCell('x =1; x') + const result = await sandbox.runCode('x =1; x') expect(result.text).toEqual('1') }) diff --git a/js/tests/runtimes/bun/run.test.ts b/js/tests/runtimes/bun/run.test.ts index 2b00c005..1497cdb6 100644 --- a/js/tests/runtimes/bun/run.test.ts +++ b/js/tests/runtimes/bun/run.test.ts @@ -5,7 +5,7 @@ import { Sandbox } from '../../../src' test('Bun test', async () => { const sbx = await Sandbox.create({ timeoutMs: 5_000 }) try { - const result = await sbx.notebook.execCell('print("Hello, World!")') + const result = await sbx.runCode('print("Hello, World!")') expect(result.logs.stdout.join('')).toEqual('Hello, World!\n') } finally { await sbx.kill() diff --git a/js/tests/runtimes/deno/run.test.ts b/js/tests/runtimes/deno/run.test.ts index 19601f57..34343294 100644 --- a/js/tests/runtimes/deno/run.test.ts +++ b/js/tests/runtimes/deno/run.test.ts @@ -8,7 +8,7 @@ import { Sandbox } from '../../../dist/index.mjs' Deno.test('Deno test', async () => { const sbx = await Sandbox.create({ timeoutMs: 5_000 }) try { - const result = await sbx.notebook.execCell('print("Hello, World!")') + const result = await sbx.runCode('print("Hello, World!")') assertEquals(result.logs.stdout.join(''), 'Hello, World!\n') } finally { await sbx.kill() diff --git a/js/tests/statefulness.test.ts b/js/tests/statefulness.test.ts index 8989f22b..fd5a7e4e 100644 --- a/js/tests/statefulness.test.ts +++ b/js/tests/statefulness.test.ts @@ -4,9 +4,9 @@ import { isDebug, sandboxTest } from './setup' // Skip this test if we are running in debug mode — the execution is persisted between all tests so the result is not reset. sandboxTest.skipIf(isDebug)('statefulness', async ({ sandbox }) => { - await sandbox.notebook.execCell('x = 1') + await sandbox.runCode('x = 1') - const result = await sandbox.notebook.execCell('x += 1; x') + const result = await sandbox.runCode('x += 1; x') expect(result.text).toEqual('2') }) diff --git a/js/tests/streaming.test.ts b/js/tests/streaming.test.ts index ac2b8009..bbf76f81 100644 --- a/js/tests/streaming.test.ts +++ b/js/tests/streaming.test.ts @@ -6,7 +6,7 @@ import { sandboxTest } from './setup' sandboxTest('streaming output', async ({ sandbox }) => { const out: OutputMessage[] = [] - await sandbox.notebook.execCell('print(1)', { + await sandbox.runCode('print(1)', { onStdout: (msg) => out.push(msg), }) @@ -16,7 +16,7 @@ sandboxTest('streaming output', async ({ sandbox }) => { sandboxTest('streaming error', async ({ sandbox }) => { const out: OutputMessage[] = [] - await sandbox.notebook.execCell('import sys;print(1, file=sys.stderr)', { + await sandbox.runCode('import sys;print(1, file=sys.stderr)', { onStderr: (msg) => out.push(msg), }) @@ -38,7 +38,7 @@ sandboxTest('streaming result', async ({ sandbox }) => { x ` - await sandbox.notebook.execCell(code, { + await sandbox.runCode(code, { onResult: (result) => out.push(result), }) diff --git a/python/README.md b/python/README.md index 92b44bbc..12a1acd6 100644 --- a/python/README.md +++ b/python/README.md @@ -18,12 +18,12 @@ pip install e2b-code-interpreter ### Minimal example with the sharing context ```python -from e2b_code_interpreter import CodeInterpreter +from e2b_code_interpreter import Sandbox -with CodeInterpreter() as sandbox: - sandbox.notebook.exec_cell("x = 1") +with Sandbox() as sandbox: + sandbox.run_code("x = 1") - execution = sandbox.notebook.exec_cell("x+=1; x") + execution = sandbox.run_code("x+=1; x") print(execution.text) # outputs 2 ``` @@ -36,7 +36,7 @@ import io from matplotlib import image as mpimg, pyplot as plt -from e2b_code_interpreter import CodeInterpreter +from e2b_code_interpreter import Sandbox code = """ import matplotlib.pyplot as plt @@ -49,12 +49,12 @@ plt.plot(x, y) plt.show() """ -with CodeInterpreter() as sandbox: +with Sandbox() as sandbox: # you can install dependencies in "jupyter notebook style" - sandbox.notebook.exec_cell("!pip install matplotlib") + sandbox.run_code("!pip install matplotlib") # plot random graph - execution = sandbox.notebook.exec_cell(code) + execution = sandbox.run_code(code) # there's your image image = execution.results[0].png @@ -71,7 +71,7 @@ plt.show() ### Streaming code output ```python -from e2b_code_interpreter import CodeInterpreter +from e2b_code_interpreter import Sandbox code = """ import time @@ -85,6 +85,6 @@ time.sleep(3) print("world") """ -with CodeInterpreter() as sandbox: - sandbox.notebook.exec_cell(code, on_stdout=print, on_stderr=print, on_result=(lambda result: print(result.text))) +with Sandbox() as sandbox: + sandbox.run_code(code, on_stdout=print, on_stderr=print, on_result=(lambda result: print(result.text))) ``` diff --git a/python/async_example.py b/python/async_example.py index 5d9df50b..02eab3a7 100644 --- a/python/async_example.py +++ b/python/async_example.py @@ -2,7 +2,7 @@ from dotenv import load_dotenv -from e2b_code_interpreter import AsyncCodeInterpreter +from e2b_code_interpreter import AsyncSandbox load_dotenv() @@ -47,16 +47,12 @@ async def create_sbx(sbx, i: int): await asyncio.sleep(i * 0.01) - return ( - await sbx.notebook.exec_cell(f"os.getenv('TEST')", envs={"TEST": str(i)}) - ).text + return (await sbx.run_code(f"os.getenv('TEST')", envs={"TEST": str(i)})).text async def run(): - sbx = await AsyncCodeInterpreter.create( - debug=True, - ) - result = await sbx.notebook.exec_cell(code) + sbx = await AsyncSandbox.create(debug=True) + result = await sbx.run_code(code) print("".join(result.logs.stdout)) print("".join(result.logs.stderr)) diff --git a/python/e2b_code_interpreter/__init__.py b/python/e2b_code_interpreter/__init__.py index b83170d6..7562f51a 100644 --- a/python/e2b_code_interpreter/__init__.py +++ b/python/e2b_code_interpreter/__init__.py @@ -2,6 +2,7 @@ from .code_interpreter_sync import Sandbox from .code_interpreter_async import AsyncSandbox from .models import ( + Context, Execution, ExecutionError, Result, diff --git a/python/e2b_code_interpreter/code_interpreter_async.py b/python/e2b_code_interpreter/code_interpreter_async.py index c2b90eec..b99c265d 100644 --- a/python/e2b_code_interpreter/code_interpreter_async.py +++ b/python/e2b_code_interpreter/code_interpreter_async.py @@ -1,19 +1,23 @@ import logging import httpx -from typing import Optional, List, Dict -from httpx import AsyncHTTPTransport, AsyncClient +from typing import Optional, Dict +from httpx import AsyncClient -from e2b import AsyncSandbox as BaseAsyncSandbox, ConnectionConfig +from e2b import ( + AsyncSandbox as BaseAsyncSandbox, + ConnectionConfig, + InvalidArgumentException, +) from e2b_code_interpreter.constants import ( - DEFAULT_KERNEL_ID, DEFAULT_TEMPLATE, JUPYTER_PORT, + DEFAULT_TIMEOUT, ) from e2b_code_interpreter.models import ( Execution, - Kernel, + Context, Result, aextract_exception, parse_output, @@ -28,31 +32,25 @@ logger = logging.getLogger(__name__) -class JupyterExtension: - """ - Code interpreter module for executing code in a stateful context. - """ +class AsyncSandbox(BaseAsyncSandbox): + default_template = DEFAULT_TEMPLATE + + def __init__(self, sandbox_id: str, connection_config: ConnectionConfig): + super().__init__(sandbox_id=sandbox_id, connection_config=connection_config) - _exec_timeout = 300 + @property + def _jupyter_url(self) -> str: + return f"{'http' if self.connection_config.debug else 'https'}://{self.get_host(JUPYTER_PORT)}" @property def _client(self) -> AsyncClient: return AsyncClient(transport=self._transport) - def __init__( - self, - url: str, - transport: AsyncHTTPTransport, - connection_config: ConnectionConfig, - ): - self._url = url - self._transport = transport - self._connection_config = connection_config - - async def exec_cell( + async def run_code( self, code: str, - kernel_id: Optional[str] = None, + language: Optional[str] = None, + context: Optional[Context] = None, on_stdout: Optional[OutputHandler[OutputMessage]] = None, on_stderr: Optional[OutputHandler[OutputMessage]] = None, on_result: Optional[OutputHandler[Result]] = None, @@ -65,7 +63,8 @@ async def exec_cell( You can reference previously defined variables, imports, and functions in the code. :param code: The code to execute - :param kernel_id: The context id + :param language Based on the value, a default context for the language is used. If not defined and no context is provided, the default Python context is used. + :param context Concrete context to run the code in. If not specified, the default context for the language is used. It's mutually exclusive with the language. :param on_stdout: Callback for stdout messages :param on_stderr: Callback for stderr messages :param on_result: Callback for the `Result` object @@ -76,16 +75,23 @@ async def exec_cell( """ logger.debug(f"Executing code {code}") - timeout = None if timeout == 0 else (timeout or self._exec_timeout) + if context and language: + raise InvalidArgumentException( + "You can provide context or language, but not both at the same time." + ) + + timeout = None if timeout == 0 else (timeout or DEFAULT_TIMEOUT) request_timeout = request_timeout or self._connection_config.request_timeout + context_id = context.id if context else None try: async with self._client.stream( "POST", - f"{self._url}/execute", + f"{self._jupyter_url}/execute", json={ "code": code, - "context_id": kernel_id, + "context_id": context_id, + "language": language, "env_vars": envs, }, timeout=(request_timeout, timeout, request_timeout, request_timeout), @@ -112,27 +118,27 @@ async def exec_cell( except httpx.TimeoutException: raise format_request_timeout_error() - async def create_kernel( + async def create_code_context( self, cwd: Optional[str] = None, - kernel_name: Optional[str] = None, + language: Optional[str] = None, envs: Optional[Dict[str, str]] = None, request_timeout: Optional[float] = None, - ) -> str: + ) -> Context: """ Creates a new context to run code in. :param cwd: Set the current working directory for the context - :param kernel_name: Type of the context + :param language: Language of the context. If not specified, the default Python context is used. :param envs: Environment variables :param request_timeout: Max time to wait for the request to finish :return: Context id """ - logger.debug(f"Creating new kernel {kernel_name}") + logger.debug(f"Creating new {language} context") data = {} - if kernel_name: - data["name"] = kernel_name + if language: + data["language"] = language if cwd: data["cwd"] = cwd if envs: @@ -140,7 +146,7 @@ async def create_kernel( try: response = await self._client.post( - f"{self._url}/contexts", + f"{self._jupyter_url}/contexts", json=data, timeout=request_timeout or self._connection_config.request_timeout, ) @@ -150,110 +156,6 @@ async def create_kernel( raise err data = response.json() - return data["id"] - except httpx.TimeoutException: - raise format_request_timeout_error() - - async def shutdown_kernel( - self, - kernel_id: Optional[str] = None, - request_timeout: Optional[float] = None, - ) -> None: - """ - Shuts down a context. - - :param kernel_id: Context id - :param request_timeout: Max time to wait for the request to finish - """ - kernel_id = kernel_id or DEFAULT_KERNEL_ID - - logger.debug(f"Shutting down a kernel with id {kernel_id}") - - try: - response = await self._client.delete( - url=f"{self._url}/contexts/{kernel_id}", - timeout=request_timeout or self._connection_config.request_timeout, - ) - - err = await aextract_exception(response) - if err: - raise err + return Context.from_json(data) except httpx.TimeoutException: raise format_request_timeout_error() - - async def restart_kernel( - self, - kernel_id: Optional[str] = None, - request_timeout: Optional[float] = None, - ) -> None: - """ - Restarts the context. - Restarting will clear all variables, imports, and other settings set during previous executions. - - :param kernel_id: Context id - :param request_timeout: Max time to wait for the request to finish - """ - kernel_id = kernel_id or DEFAULT_KERNEL_ID - - logger.debug(f"Restarting kernel {kernel_id}") - - try: - response = await self._client.post( - f"{self._url}/contexts/{kernel_id}/restart", - timeout=request_timeout or self._connection_config.request_timeout, - ) - - err = await aextract_exception(response) - if err: - raise err - except httpx.TimeoutException: - raise format_request_timeout_error() - - async def list_kernels( - self, - request_timeout: Optional[float] = None, - ) -> List[Kernel]: - """ - Lists all available contexts. - - :param request_timeout: Max time to wait for the request to finish - :return: List of Kernel objects - """ - logger.debug("Listing kernels") - - try: - response = await self._client.get( - f"{self._url}/contexts", - timeout=request_timeout or self._connection_config.request_timeout, - ) - - err = await aextract_exception(response) - if err: - raise err - - return [Kernel(k["id"], k["name"]) for k in response.json()] - except httpx.TimeoutException: - raise format_request_timeout_error() - - -class AsyncSandbox(BaseAsyncSandbox): - default_template = DEFAULT_TEMPLATE - _jupyter_port = JUPYTER_PORT - - @property - def notebook(self) -> JupyterExtension: - """ - Code interpreter module for executing code in a stateful context. - """ - return self._notebook - - def __init__(self, sandbox_id: str, connection_config: ConnectionConfig): - super().__init__(sandbox_id=sandbox_id, connection_config=connection_config) - - jupyter_url = f"{'http' if self.connection_config.debug else 'https'}://{self.get_host(self._jupyter_port)}" - - self._notebook = JupyterExtension( - jupyter_url, - self._transport, - self.connection_config, - ) diff --git a/python/e2b_code_interpreter/code_interpreter_sync.py b/python/e2b_code_interpreter/code_interpreter_sync.py index 9f4a79e9..2df46bf0 100644 --- a/python/e2b_code_interpreter/code_interpreter_sync.py +++ b/python/e2b_code_interpreter/code_interpreter_sync.py @@ -1,18 +1,18 @@ import logging import httpx -from typing import Optional, Dict, List -from httpx import HTTPTransport, Client -from e2b import Sandbox as BaseSandbox, ConnectionConfig +from typing import Optional, Dict +from httpx import Client +from e2b import Sandbox as BaseSandbox, InvalidArgumentException from e2b_code_interpreter.constants import ( - DEFAULT_KERNEL_ID, DEFAULT_TEMPLATE, JUPYTER_PORT, + DEFAULT_TIMEOUT, ) from e2b_code_interpreter.models import ( Execution, - Kernel, + Context, Result, extract_exception, parse_output, @@ -27,31 +27,22 @@ logger = logging.getLogger(__name__) -class JupyterExtension: - """ - Code interpreter module for executing code in a stateful context. - """ +class Sandbox(BaseSandbox): + default_template = DEFAULT_TEMPLATE - _exec_timeout = 300 + @property + def _jupyter_url(self) -> str: + return f"{'http' if self.connection_config.debug else 'https'}://{self.get_host(JUPYTER_PORT)}" @property def _client(self) -> Client: return Client(transport=self._transport) - def __init__( - self, - url: str, - transport: HTTPTransport, - connection_config: ConnectionConfig, - ): - self._url = url - self._transport = transport - self._connection_config = connection_config - - def exec_cell( + def run_code( self, code: str, - kernel_id: Optional[str] = None, + language: Optional[str] = None, + context: Optional[Context] = None, on_stdout: Optional[OutputHandler[OutputMessage]] = None, on_stderr: Optional[OutputHandler[OutputMessage]] = None, on_result: Optional[OutputHandler[Result]] = None, @@ -64,7 +55,8 @@ def exec_cell( You can reference previously defined variables, imports, and functions in the code. :param code: The code to execute - :param kernel_id: The context id + :param language Based on the value, a default context for the language is used. If not defined and no context is provided, the default Python context is used. + :param context Concrete context to run the code in. If not specified, the default context for the language is used. It's mutually exclusive with the language. :param on_stdout: Callback for stdout messages :param on_stderr: Callback for stderr messages :param on_result: Callback for the `Result` object @@ -75,16 +67,23 @@ def exec_cell( """ logger.debug(f"Executing code {code}") - timeout = None if timeout == 0 else (timeout or self._exec_timeout) + if language and context: + raise InvalidArgumentException( + "You can provide context or language, but not both at the same time." + ) + + timeout = None if timeout == 0 else (timeout or DEFAULT_TIMEOUT) request_timeout = request_timeout or self._connection_config.request_timeout + context_id = context.id if context else None try: with self._client.stream( "POST", - f"{self._url}/execute", + f"{self._jupyter_url}/execute", json={ "code": code, - "context_id": kernel_id, + "context_id": context_id, + "language": language, "env_vars": envs, }, timeout=(request_timeout, timeout, request_timeout, request_timeout), @@ -110,27 +109,27 @@ def exec_cell( except httpx.TimeoutException: raise format_request_timeout_error() - def create_kernel( + def create_code_context( self, cwd: Optional[str] = None, - kernel_name: Optional[str] = None, + language: Optional[str] = None, envs: Optional[Dict[str, str]] = None, request_timeout: Optional[float] = None, - ) -> str: + ) -> Context: """ Creates a new context to run code in. :param cwd: Set the current working directory for the context - :param kernel_name: Type of the context + :param language: Language of the context. If not specified, the default Python context is used. :param envs: Environment variables :param request_timeout: Max time to wait for the request to finish :return: Context id """ - logger.debug(f"Creating new kernel {kernel_name}") + logger.debug(f"Creating new {language} context") data = {} - if kernel_name: - data["name"] = kernel_name + if language: + data["language"] = language if cwd: data["cwd"] = cwd if envs: @@ -138,7 +137,7 @@ def create_kernel( try: response = self._client.post( - f"{self._url}/contexts", + f"{self._jupyter_url}/contexts", json=data, timeout=request_timeout or self._connection_config.request_timeout, ) @@ -148,129 +147,6 @@ def create_kernel( raise err data = response.json() - return data["id"] - except httpx.TimeoutException: - raise format_request_timeout_error() - - def shutdown_kernel( - self, - kernel_id: Optional[str] = None, - request_timeout: Optional[float] = None, - ) -> None: - """ - Shuts down a context. - - :param kernel_id: Context id to shut down - :param request_timeout: Max time to wait for the request to finish - """ - kernel_id = kernel_id or DEFAULT_KERNEL_ID - - logger.debug(f"Shutting down a kernel with id {kernel_id}") - - try: - response = self._client.delete( - url=f"{self._url}/contexts/{kernel_id}", - timeout=request_timeout or self._connection_config.request_timeout, - ) - err = extract_exception(response) - if err: - raise err - except httpx.TimeoutException: - raise format_request_timeout_error() - - def restart_kernel( - self, - kernel_id: Optional[str] = None, - request_timeout: Optional[float] = None, - ) -> None: - """ - Restarts the context. - Restarting will clear all variables, imports, and other settings set during previous executions. - - :param kernel_id: Context id - :param request_timeout: Max time to wait for the request to finish - """ - kernel_id = kernel_id or DEFAULT_KERNEL_ID - - logger.debug(f"Restarting kernel {kernel_id}") - - try: - response = self._client.post( - f"{self._url}/contexts/{kernel_id}/restart", - timeout=request_timeout or self._connection_config.request_timeout, - ) - - err = extract_exception(response) - if err: - raise err + return Context.from_json(data) except httpx.TimeoutException: raise format_request_timeout_error() - - def list_kernels( - self, - request_timeout: Optional[float] = None, - ) -> List[Kernel]: - """ - Lists all available contexts. - - :param request_timeout: Max time to wait for the request to finish - :return: List of Kernel objects - """ - logger.debug("Listing kernels") - - try: - response = self._client.get( - f"{self._url}/contexts", - timeout=request_timeout or self._connection_config.request_timeout, - ) - - err = extract_exception(response) - if err: - raise err - - return [Kernel(kernel_id=k["id"], name=k["name"]) for k in response.json()] - except httpx.TimeoutException: - raise format_request_timeout_error() - - -class Sandbox(BaseSandbox): - default_template = DEFAULT_TEMPLATE - _jupyter_port = JUPYTER_PORT - - @property - def notebook(self) -> JupyterExtension: - """ - Code interpreter module for executing code in a stateful context. - """ - return self._notebook - - def __init__( - self, - template: Optional[str] = None, - timeout: Optional[int] = None, - metadata: Optional[Dict[str, str]] = None, - envs: Optional[Dict[str, str]] = None, - api_key: Optional[str] = None, - domain: Optional[str] = None, - debug: Optional[bool] = None, - sandbox_id: Optional[str] = None, - request_timeout: Optional[float] = None, - ): - super().__init__( - template=template, - timeout=timeout, - metadata=metadata, - envs=envs, - api_key=api_key, - domain=domain, - debug=debug, - sandbox_id=sandbox_id, - request_timeout=request_timeout, - ) - - jupyter_url = f"{'http' if self.connection_config.debug else 'https'}://{self.get_host(self._jupyter_port)}" - self._notebook = JupyterExtension( - jupyter_url, - self._transport, - self.connection_config, - ) diff --git a/python/e2b_code_interpreter/constants.py b/python/e2b_code_interpreter/constants.py index b0f4ec93..642dabc0 100644 --- a/python/e2b_code_interpreter/constants.py +++ b/python/e2b_code_interpreter/constants.py @@ -1,3 +1,3 @@ DEFAULT_TEMPLATE = "code-interpreter-beta" JUPYTER_PORT = 49999 -DEFAULT_KERNEL_ID = "default" +DEFAULT_TIMEOUT = 300 diff --git a/python/e2b_code_interpreter/models.py b/python/e2b_code_interpreter/models.py index 6a25509b..f55cf99d 100644 --- a/python/e2b_code_interpreter/models.py +++ b/python/e2b_code_interpreter/models.py @@ -80,8 +80,6 @@ class Result: 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 display calls don't have to have text representation, for the actual result the representation is always present for the result, the other representations are always optional. - - The class also provides methods to display the data in a Jupyter notebook. """ def __getitem__(self, item): @@ -420,10 +418,20 @@ def parse_output( @dataclass -class Kernel: - kernel_id: str - name: str - - def __init__(self, kernel_id: str, name: str, **kwargs): - self.kernel_id = kernel_id - self.name = name +class Context: + id: str + language: str + cwd: str + + def __init__(self, context_id: str, language: str, cwd: str, **kwargs): + self.id = context_id + self.language = language + self.cwd = cwd + + @classmethod + def from_json(cls, data: Dict[str, str]): + return cls( + context_id=data.get("id"), + language=data.get("language"), + cwd=data.get("cwd"), + ) diff --git a/python/example.py b/python/example.py index b0a89750..cb79406d 100644 --- a/python/example.py +++ b/python/example.py @@ -1,10 +1,8 @@ import asyncio -import time -from time import sleep from dotenv import load_dotenv -from e2b_code_interpreter import AsyncCodeInterpreter +from e2b_code_interpreter import Sandbox load_dotenv() @@ -36,8 +34,8 @@ async def run(): - sbx = await AsyncCodeInterpreter.create(timeout=60) - e = await sbx.notebook.exec_cell(code) + sbx = Sandbox(timeout=60) + e = sbx.run_code(code) print(e.results[0].graph) diff --git a/python/tests/async/test_async_bash.py b/python/tests/async/test_async_bash.py index bcfdb207..a9e2cb0c 100644 --- a/python/tests/async/test_async_bash.py +++ b/python/tests/async/test_async_bash.py @@ -2,5 +2,5 @@ async def test_bash(async_sandbox: AsyncSandbox): - result = await async_sandbox.notebook.exec_cell("!pwd") + result = await async_sandbox.run_code("!pwd") assert "".join(result.logs.stdout).strip() == "/home/user" diff --git a/python/tests/async/test_async_basic.py b/python/tests/async/test_async_basic.py index 84ffd29a..486fd219 100644 --- a/python/tests/async/test_async_basic.py +++ b/python/tests/async/test_async_basic.py @@ -2,5 +2,5 @@ async def test_basic(async_sandbox: AsyncSandbox): - result = await async_sandbox.notebook.exec_cell("x =1; x") + result = await async_sandbox.run_code("x =1; x") assert result.text == "1" diff --git a/python/tests/async/test_async_custom_repr_object.py b/python/tests/async/test_async_custom_repr_object.py index c8deb21a..7b35e925 100644 --- a/python/tests/async/test_async_custom_repr_object.py +++ b/python/tests/async/test_async_custom_repr_object.py @@ -8,5 +8,5 @@ async def test_bash(async_sandbox: AsyncSandbox): - execution = await async_sandbox.notebook.exec_cell(code) + execution = await async_sandbox.run_code(code) assert execution.results[0].formats() == ["latex"] diff --git a/python/tests/async/test_async_data.py b/python/tests/async/test_async_data.py index d0bd1a4c..ad43ffec 100644 --- a/python/tests/async/test_async_data.py +++ b/python/tests/async/test_async_data.py @@ -3,7 +3,7 @@ async def test_data(async_sandbox: AsyncSandbox): # plot random graph - result = await async_sandbox.notebook.exec_cell( + result = await async_sandbox.run_code( """ import pandas as pd pd.DataFrame({"a": [1, 2, 3]}) diff --git a/python/tests/async/test_async_default_kernels.py b/python/tests/async/test_async_default_kernels.py new file mode 100644 index 00000000..a632bda6 --- /dev/null +++ b/python/tests/async/test_async_default_kernels.py @@ -0,0 +1,8 @@ +from e2b_code_interpreter.code_interpreter_async import AsyncSandbox + + +async def test_js_kernel(async_sandbox: AsyncSandbox): + execution = await async_sandbox.run_code( + "console.log('Hello, World!')", language="js" + ) + assert execution.logs.stdout == ["Hello, World!\n"] diff --git a/python/tests/async/test_async_display_data.py b/python/tests/async/test_async_display_data.py index a714a4cc..e3756212 100644 --- a/python/tests/async/test_async_display_data.py +++ b/python/tests/async/test_async_display_data.py @@ -3,7 +3,7 @@ async def test_display_data(async_sandbox: AsyncSandbox): # plot random graph - result = await async_sandbox.notebook.exec_cell( + result = await async_sandbox.run_code( """ import matplotlib.pyplot as plt import numpy as np diff --git a/python/tests/async/test_async_env_vars.py b/python/tests/async/test_async_env_vars.py index cf54661c..f587dcb6 100644 --- a/python/tests/async/test_async_env_vars.py +++ b/python/tests/async/test_async_env_vars.py @@ -1,15 +1,18 @@ +import pytest + from e2b_code_interpreter.code_interpreter_async import AsyncSandbox +@pytest.mark.skip_debug() async def test_env_vars_sandbox(): sbx = await AsyncSandbox.create(envs={"FOO": "bar"}) - result = await sbx.notebook.exec_cell("import os; os.getenv('FOO')") + result = await sbx.run_code("import os; os.getenv('FOO')") assert result.text == "bar" await sbx.kill() -async def test_env_vars_in_exec_cell(async_sandbox: AsyncSandbox): - result = await async_sandbox.notebook.exec_cell( +async def test_env_vars_in_run_code(async_sandbox: AsyncSandbox): + result = await async_sandbox.run_code( "import os; os.getenv('FOO')", envs={"FOO": "bar"} ) assert result.text == "bar" @@ -17,24 +20,22 @@ async def test_env_vars_in_exec_cell(async_sandbox: AsyncSandbox): async def test_env_vars_override(debug: bool): sbx = await AsyncSandbox.create(envs={"FOO": "bar", "SBX": "value"}) - await sbx.notebook.exec_cell( + await sbx.run_code( "import os; os.environ['FOO'] = 'bar'; os.environ['RUNTIME_ENV'] = 'async_python_runtime'" ) - result = await sbx.notebook.exec_cell( - "import os; os.getenv('FOO')", envs={"FOO": "baz"} - ) + result = await sbx.run_code("import os; os.getenv('FOO')", envs={"FOO": "baz"}) assert result.text == "baz" # This can fail if running in debug mode (there's a race condition with the restart kernel test) - result = await sbx.notebook.exec_cell("import os; os.getenv('RUNTIME_ENV')") + result = await sbx.run_code("import os; os.getenv('RUNTIME_ENV')") assert result.text == "async_python_runtime" if not debug: - result = await sbx.notebook.exec_cell("import os; os.getenv('SBX')") + result = await sbx.run_code("import os; os.getenv('SBX')") assert result.text == "value" # This can fail if running in debug mode (there's a race condition with the restart kernel test) - result = await sbx.notebook.exec_cell("import os; os.getenv('FOO')") + result = await sbx.run_code("import os; os.getenv('FOO')") assert result.text == "bar" await sbx.kill() diff --git a/python/tests/async/test_async_execution_count.py b/python/tests/async/test_async_execution_count.py index e163750d..c4955838 100644 --- a/python/tests/async/test_async_execution_count.py +++ b/python/tests/async/test_async_execution_count.py @@ -5,6 +5,6 @@ @pytest.mark.skip_debug() async def test_execution_count(async_sandbox: AsyncSandbox): - await async_sandbox.notebook.exec_cell("echo 'E2B is awesome!'") - result = await async_sandbox.notebook.exec_cell("!pwd") + await async_sandbox.run_code("echo 'E2B is awesome!'") + result = await async_sandbox.run_code("!pwd") assert result.execution_count == 2 diff --git a/python/tests/async/test_async_kernels.py b/python/tests/async/test_async_kernels.py index f5d6b33b..ce88967b 100644 --- a/python/tests/async/test_async_kernels.py +++ b/python/tests/async/test_async_kernels.py @@ -1,50 +1,25 @@ import pytest +from e2b import InvalidArgumentException from e2b_code_interpreter.code_interpreter_async import AsyncSandbox async def test_create_new_kernel(async_sandbox: AsyncSandbox): - await async_sandbox.notebook.create_kernel() + await async_sandbox.create_code_context() async def test_independence_of_kernels(async_sandbox: AsyncSandbox): - kernel_id = await async_sandbox.notebook.create_kernel() - await async_sandbox.notebook.exec_cell("x = 1") + context = await async_sandbox.create_code_context() + await async_sandbox.run_code("x = 1") - r = await async_sandbox.notebook.exec_cell("x", kernel_id=kernel_id) + r = await async_sandbox.run_code("x", context=context) assert r.error is not None assert r.error.value == "name 'x' is not defined" -@pytest.mark.skip_debug() -async def test_restart_kernel(async_sandbox: AsyncSandbox): - await async_sandbox.notebook.exec_cell("x = 1") - await async_sandbox.notebook.restart_kernel() - - r = await async_sandbox.notebook.exec_cell("x") - assert r.error is not None - assert r.error.value == "name 'x' is not defined" - - -@pytest.mark.skip_debug() -async def test_list_kernels(async_sandbox: AsyncSandbox): - kernels = await async_sandbox.notebook.list_kernels() - assert len(kernels) == 1 - - kernel_id = await async_sandbox.notebook.create_kernel() - kernels = await async_sandbox.notebook.list_kernels() - assert kernel_id in [kernel.kernel_id for kernel in kernels] - assert len(kernels) == 2 - - -@pytest.mark.skip_debug() -async def test_shutdown(async_sandbox: AsyncSandbox): - kernel_id = await async_sandbox.notebook.create_kernel() - kernels = await async_sandbox.notebook.list_kernels() - assert kernel_id in [kernel.kernel_id for kernel in kernels] - assert len(kernels) == 2 - - await async_sandbox.notebook.shutdown_kernel(kernel_id) - kernels = await async_sandbox.notebook.list_kernels() - assert kernel_id not in [kernel.kernel_id for kernel in kernels] - assert len(kernels) == 1 +async def test_pass_context_and_language(async_sandbox: AsyncSandbox): + context = await async_sandbox.create_code_context(language="python") + with pytest.raises(InvalidArgumentException): + await async_sandbox.run_code( + "console.log('Hello, World!')", language="js", context=context + ) diff --git a/python/tests/async/test_async_reconnect.py b/python/tests/async/test_async_reconnect.py index 0b91febf..9ed16f4a 100644 --- a/python/tests/async/test_async_reconnect.py +++ b/python/tests/async/test_async_reconnect.py @@ -5,5 +5,5 @@ async def test_reconnect(async_sandbox: AsyncSandbox): sandbox_id = async_sandbox.sandbox_id sandbox2 = await AsyncSandbox.connect(sandbox_id) - result = await sandbox2.notebook.exec_cell("x =1; x") + result = await sandbox2.run_code("x =1; x") assert result.text == "1" diff --git a/python/tests/async/test_async_statefulness.py b/python/tests/async/test_async_statefulness.py index be9612aa..e33dd808 100644 --- a/python/tests/async/test_async_statefulness.py +++ b/python/tests/async/test_async_statefulness.py @@ -2,9 +2,7 @@ async def test_stateful(async_sandbox: AsyncSandbox): - await async_sandbox.notebook.exec_cell("async_test_stateful = 1") + await async_sandbox.run_code("async_test_stateful = 1") - result = await async_sandbox.notebook.exec_cell( - "async_test_stateful+=1; async_test_stateful" - ) + result = await async_sandbox.run_code("async_test_stateful+=1; async_test_stateful") assert result.text == "2" diff --git a/python/tests/async/test_async_streaming.py b/python/tests/async/test_async_streaming.py index 148045de..ceb5c851 100644 --- a/python/tests/async/test_async_streaming.py +++ b/python/tests/async/test_async_streaming.py @@ -8,7 +8,7 @@ def test(line) -> None: out.append(line) return - await async_sandbox.notebook.exec_cell("print(1)", on_stdout=test) + await async_sandbox.run_code("print(1)", on_stdout=test) assert len(out) == 1 assert out[0].line == "1\n" @@ -17,7 +17,7 @@ def test(line) -> None: async def test_streaming_error(async_sandbox: AsyncSandbox): out = [] - await async_sandbox.notebook.exec_cell( + await async_sandbox.run_code( "import sys;print(1, file=sys.stderr)", on_stderr=out.append ) @@ -40,6 +40,6 @@ async def test_streaming_result(async_sandbox: AsyncSandbox): """ out = [] - await async_sandbox.notebook.exec_cell(code, on_result=out.append) + await async_sandbox.run_code(code, on_result=out.append) assert len(out) == 2 diff --git a/python/tests/benchmarking.py b/python/tests/benchmarking.py index 39496deb..c78d674c 100644 --- a/python/tests/benchmarking.py +++ b/python/tests/benchmarking.py @@ -18,14 +18,14 @@ create_sandbox_time += time.time() - start_time start_time = time.time() - sandbox.notebook.exec_cell("x = 1") + sandbox.run_code("x = 1") first_exec_time += time.time() - start_time start_time = time.time() - result = sandbox.notebook.exec_cell("x+=1; x") + result = sandbox.run_code("x+=1; x") second_exec_time += time.time() - start_time - sandbox.close() + sandbox.kill() print(f"Average Create Sandbox Time: {create_sandbox_time / iterations}s") diff --git a/python/tests/graphs/test_bar.py b/python/tests/graphs/test_bar.py index 486c582b..481d4306 100644 --- a/python/tests/graphs/test_bar.py +++ b/python/tests/graphs/test_bar.py @@ -22,7 +22,7 @@ async def test_graph_bar(async_sandbox: AsyncSandbox): - result = await async_sandbox.notebook.exec_cell(code) + result = await async_sandbox.run_code(code) graph = result.results[0].graph assert graph diff --git a/python/tests/graphs/test_box_and_whiskers.py b/python/tests/graphs/test_box_and_whiskers.py index 17073341..fe561d23 100644 --- a/python/tests/graphs/test_box_and_whiskers.py +++ b/python/tests/graphs/test_box_and_whiskers.py @@ -36,7 +36,7 @@ async def test_box_and_whiskers(async_sandbox: AsyncSandbox): - result = await async_sandbox.notebook.exec_cell(code) + result = await async_sandbox.run_code(code) graph = result.results[0].graph assert graph diff --git a/python/tests/graphs/test_line.py b/python/tests/graphs/test_line.py index 7e40ab46..7f88a6f3 100644 --- a/python/tests/graphs/test_line.py +++ b/python/tests/graphs/test_line.py @@ -32,7 +32,7 @@ async def test_line_graph(async_sandbox: AsyncSandbox): - result = await async_sandbox.notebook.exec_cell(code) + result = await async_sandbox.run_code(code) graph = result.results[0].graph assert graph diff --git a/python/tests/graphs/test_log_graph.py b/python/tests/graphs/test_log_graph.py index 900be0f3..31b62d6a 100644 --- a/python/tests/graphs/test_log_graph.py +++ b/python/tests/graphs/test_log_graph.py @@ -30,7 +30,7 @@ async def test_log_graph(async_sandbox: AsyncSandbox): - result = await async_sandbox.notebook.exec_cell(code) + result = await async_sandbox.run_code(code) graph = result.results[0].graph assert graph diff --git a/python/tests/graphs/test_pie.py b/python/tests/graphs/test_pie.py index d3510c29..d78d14d7 100644 --- a/python/tests/graphs/test_pie.py +++ b/python/tests/graphs/test_pie.py @@ -28,7 +28,7 @@ async def test_pie_graph(async_sandbox: AsyncSandbox): - result = await async_sandbox.notebook.exec_cell(code) + result = await async_sandbox.run_code(code) graph = result.results[0].graph assert graph diff --git a/python/tests/graphs/test_scale.py b/python/tests/graphs/test_scale.py index 98d3914f..38d1bf6c 100644 --- a/python/tests/graphs/test_scale.py +++ b/python/tests/graphs/test_scale.py @@ -18,7 +18,7 @@ async def test_datetime_scale(async_sandbox: AsyncSandbox): plt.show() """ - result = await async_sandbox.notebook.exec_cell(code) + result = await async_sandbox.run_code(code) graph = result.results[0].graph assert graph @@ -42,7 +42,7 @@ async def test_categorical_scale(async_sandbox: AsyncSandbox): plt.show() """ - result = await async_sandbox.notebook.exec_cell(code) + result = await async_sandbox.run_code(code) graph = result.results[0].graph assert graph diff --git a/python/tests/graphs/test_scatter.py b/python/tests/graphs/test_scatter.py index 2845e6eb..445daf16 100644 --- a/python/tests/graphs/test_scatter.py +++ b/python/tests/graphs/test_scatter.py @@ -23,7 +23,7 @@ async def test_scatter_graph(async_sandbox: AsyncSandbox): - result = await async_sandbox.notebook.exec_cell(code) + result = await async_sandbox.run_code(code) graph = result.results[0].graph assert graph diff --git a/python/tests/graphs/test_supergraph.py b/python/tests/graphs/test_supergraph.py index c2f1fc9a..0f71315b 100644 --- a/python/tests/graphs/test_supergraph.py +++ b/python/tests/graphs/test_supergraph.py @@ -33,7 +33,7 @@ async def test_super_graph(async_sandbox: AsyncSandbox): - result = await async_sandbox.notebook.exec_cell(code) + result = await async_sandbox.run_code(code) graph = result.results[0].graph assert graph diff --git a/python/tests/graphs/test_unknown.py b/python/tests/graphs/test_unknown.py index 48a39b5f..9d919252 100644 --- a/python/tests/graphs/test_unknown.py +++ b/python/tests/graphs/test_unknown.py @@ -28,7 +28,7 @@ async def test_unknown_graphs(async_sandbox: AsyncSandbox): - result = await async_sandbox.notebook.exec_cell(code) + result = await async_sandbox.run_code(code) graph = result.results[0].graph assert graph diff --git a/python/tests/sync/test_bash.py b/python/tests/sync/test_bash.py index e3d29caa..7de9ccb6 100644 --- a/python/tests/sync/test_bash.py +++ b/python/tests/sync/test_bash.py @@ -2,5 +2,5 @@ def test_bash(sandbox: Sandbox): - result = sandbox.notebook.exec_cell("!pwd") + result = sandbox.run_code("!pwd") assert "".join(result.logs.stdout).strip() == "/home/user" diff --git a/python/tests/sync/test_basic.py b/python/tests/sync/test_basic.py index 4400888a..4bfc503d 100644 --- a/python/tests/sync/test_basic.py +++ b/python/tests/sync/test_basic.py @@ -2,5 +2,5 @@ def test_basic(sandbox: Sandbox): - result = sandbox.notebook.exec_cell("x =1; x") + result = sandbox.run_code("x =1; x") assert result.text == "1" diff --git a/python/tests/sync/test_custom_repr_object.py b/python/tests/sync/test_custom_repr_object.py index 070fafab..8423c950 100644 --- a/python/tests/sync/test_custom_repr_object.py +++ b/python/tests/sync/test_custom_repr_object.py @@ -8,5 +8,5 @@ def test_bash(sandbox: Sandbox): - execution = sandbox.notebook.exec_cell(code) + execution = sandbox.run_code(code) assert execution.results[0].formats() == ["latex"] diff --git a/python/tests/sync/test_data.py b/python/tests/sync/test_data.py index 3c5399e0..eecc11cf 100644 --- a/python/tests/sync/test_data.py +++ b/python/tests/sync/test_data.py @@ -3,7 +3,7 @@ def test_data(sandbox: Sandbox): # plot random graph - result = sandbox.notebook.exec_cell( + result = sandbox.run_code( """ import pandas as pd pd.DataFrame({"a": [1, 2, 3]}) diff --git a/python/tests/sync/test_default_kernels.py b/python/tests/sync/test_default_kernels.py new file mode 100644 index 00000000..cb21cdc5 --- /dev/null +++ b/python/tests/sync/test_default_kernels.py @@ -0,0 +1,6 @@ +from e2b_code_interpreter.code_interpreter_sync import Sandbox + + +def test_js_kernel(sandbox: Sandbox): + execution = sandbox.run_code("console.log('Hello, World!')", language="js") + assert execution.logs.stdout == ["Hello, World!\n"] diff --git a/python/tests/sync/test_display_data.py b/python/tests/sync/test_display_data.py index 50d7a5ae..50443957 100644 --- a/python/tests/sync/test_display_data.py +++ b/python/tests/sync/test_display_data.py @@ -3,7 +3,7 @@ def test_display_data(sandbox: Sandbox): # plot random graph - result = sandbox.notebook.exec_cell( + result = sandbox.run_code( """ import matplotlib.pyplot as plt import numpy as np diff --git a/python/tests/sync/test_env_vars.py b/python/tests/sync/test_env_vars.py index f32eff32..4d56b681 100644 --- a/python/tests/sync/test_env_vars.py +++ b/python/tests/sync/test_env_vars.py @@ -1,38 +1,39 @@ +import pytest + from e2b_code_interpreter.code_interpreter_sync import Sandbox +@pytest.mark.skip_debug() async def test_env_vars_sandbox(): sbx = Sandbox(envs={"FOO": "bar"}) - result = sbx.notebook.exec_cell("import os; os.getenv('FOO')") + result = sbx.run_code("import os; os.getenv('FOO')") assert result.text == "bar" sbx.kill() -async def test_env_vars_in_exec_cell(sandbox: Sandbox): - result = sandbox.notebook.exec_cell( - "import os; os.getenv('FOO')", envs={"FOO": "bar"} - ) +async def test_env_vars_in_run_code(sandbox: Sandbox): + result = sandbox.run_code("import os; os.getenv('FOO')", envs={"FOO": "bar"}) assert result.text == "bar" async def test_env_vars_override(debug: bool): sbx = Sandbox(envs={"FOO": "bar", "SBX": "value"}) - sbx.notebook.exec_cell( + sbx.run_code( "import os; os.environ['FOO'] = 'bar'; os.environ['RUNTIME_ENV'] = 'python_runtime'" ) - result = sbx.notebook.exec_cell("import os; os.getenv('FOO')", envs={"FOO": "baz"}) + result = sbx.run_code("import os; os.getenv('FOO')", envs={"FOO": "baz"}) assert result.text == "baz" # This can fail if running in debug mode (there's a race condition with the restart kernel test) - result = sbx.notebook.exec_cell("import os; os.getenv('RUNTIME_ENV')") + result = sbx.run_code("import os; os.getenv('RUNTIME_ENV')") assert result.text == "python_runtime" if not debug: - result = sbx.notebook.exec_cell("import os; os.getenv('SBX')") + result = sbx.run_code("import os; os.getenv('SBX')") assert result.text == "value" # This can fail if running in debug mode (there's a race condition with the restart kernel test) - result = sbx.notebook.exec_cell("import os; os.getenv('FOO')") + result = sbx.run_code("import os; os.getenv('FOO')") assert result.text == "bar" sbx.kill() diff --git a/python/tests/sync/test_execution_count.py b/python/tests/sync/test_execution_count.py index 4782aafc..98e2039f 100644 --- a/python/tests/sync/test_execution_count.py +++ b/python/tests/sync/test_execution_count.py @@ -5,6 +5,6 @@ @pytest.mark.skip_debug() def test_execution_count(sandbox: Sandbox): - sandbox.notebook.exec_cell("echo 'E2B is awesome!'") - result = sandbox.notebook.exec_cell("!pwd") + sandbox.run_code("echo 'E2B is awesome!'") + result = sandbox.run_code("!pwd") assert result.execution_count == 2 diff --git a/python/tests/sync/test_kernels.py b/python/tests/sync/test_kernels.py index 6be6268c..17d6430f 100644 --- a/python/tests/sync/test_kernels.py +++ b/python/tests/sync/test_kernels.py @@ -1,50 +1,23 @@ import pytest +from e2b import InvalidArgumentException from e2b_code_interpreter.code_interpreter_sync import Sandbox def test_create_new_kernel(sandbox: Sandbox): - sandbox.notebook.create_kernel() + sandbox.create_code_context() def test_independence_of_kernels(sandbox: Sandbox): - kernel_id = sandbox.notebook.create_kernel() - sandbox.notebook.exec_cell("x = 1") + context = sandbox.create_code_context() + sandbox.run_code("x = 1") - r = sandbox.notebook.exec_cell("x", kernel_id=kernel_id) + r = sandbox.run_code("x", context=context) assert r.error is not None assert r.error.value == "name 'x' is not defined" -@pytest.mark.skip_debug() -def test_restart_kernel(sandbox: Sandbox): - sandbox.notebook.exec_cell("x = 1") - sandbox.notebook.restart_kernel() - - r = sandbox.notebook.exec_cell("x") - assert r.error is not None - assert r.error.value == "name 'x' is not defined" - - -@pytest.mark.skip_debug() -def test_list_kernels(sandbox: Sandbox): - kernels = sandbox.notebook.list_kernels() - assert len(kernels) == 1 - - kernel_id = sandbox.notebook.create_kernel() - kernels = sandbox.notebook.list_kernels() - assert kernel_id in [kernel.kernel_id for kernel in kernels] - assert len(kernels) == 2 - - -@pytest.mark.skip_debug() -def test_shutdown(sandbox: Sandbox): - kernel_id = sandbox.notebook.create_kernel() - kernels = sandbox.notebook.list_kernels() - assert kernel_id in [kernel.kernel_id for kernel in kernels] - assert len(kernels) == 2 - - sandbox.notebook.shutdown_kernel(kernel_id) - kernels = sandbox.notebook.list_kernels() - assert kernel_id not in [kernel.kernel_id for kernel in kernels] - assert len(kernels) == 1 +def test_pass_context_and_language(sandbox: Sandbox): + context = sandbox.create_code_context(language="python") + with pytest.raises(InvalidArgumentException): + sandbox.run_code("console.log('Hello, World!')", language="js", context=context) diff --git a/python/tests/sync/test_reconnect.py b/python/tests/sync/test_reconnect.py index 6dbc5557..27800a84 100644 --- a/python/tests/sync/test_reconnect.py +++ b/python/tests/sync/test_reconnect.py @@ -5,5 +5,5 @@ def test_reconnect(sandbox: Sandbox): sandbox_id = sandbox.sandbox_id sandbox2 = Sandbox.connect(sandbox_id) - result = sandbox2.notebook.exec_cell("x =1; x") + result = sandbox2.run_code("x =1; x") assert result.text == "1" diff --git a/python/tests/sync/test_statefulness.py b/python/tests/sync/test_statefulness.py index bcf30586..044478db 100644 --- a/python/tests/sync/test_statefulness.py +++ b/python/tests/sync/test_statefulness.py @@ -2,7 +2,7 @@ def test_stateful(sandbox: Sandbox): - sandbox.notebook.exec_cell("test_stateful = 1") + sandbox.run_code("test_stateful = 1") - result = sandbox.notebook.exec_cell("test_stateful+=1; test_stateful") + result = sandbox.run_code("test_stateful+=1; test_stateful") assert result.text == "2" diff --git a/python/tests/sync/test_streaming.py b/python/tests/sync/test_streaming.py index 9ef912ba..fa7a3c4a 100644 --- a/python/tests/sync/test_streaming.py +++ b/python/tests/sync/test_streaming.py @@ -8,7 +8,7 @@ def test(line) -> None: out.append(line) return - sandbox.notebook.exec_cell("print(1)", on_stdout=test) + sandbox.run_code("print(1)", on_stdout=test) assert len(out) == 1 assert out[0].line == "1\n" @@ -17,9 +17,7 @@ def test(line) -> None: def test_streaming_error(sandbox: Sandbox): out = [] - sandbox.notebook.exec_cell( - "import sys;print(1, file=sys.stderr)", on_stderr=out.append - ) + sandbox.run_code("import sys;print(1, file=sys.stderr)", on_stderr=out.append) assert len(out) == 1 assert out[0].line == "1\n" @@ -40,6 +38,6 @@ def test_streaming_result(sandbox: Sandbox): """ out = [] - sandbox.notebook.exec_cell(code, on_result=out.append) + sandbox.run_code(code, on_result=out.append) assert len(out) == 2 diff --git a/template/README.md b/template/README.md index 8c1aaa34..4d00d588 100644 --- a/template/README.md +++ b/template/README.md @@ -23,21 +23,23 @@ If you want to customize the Code Interprerter sandbox (e.g.: add a preinstalled **Python** ```python - from e2b_code_interpreter import CodeInterpreter - sandbox = CodeInterpreter(template="your-custom-sandbox-name") - execution = sandbox.notebook.exec_cell("print('hello')") - sandbox.close() + from e2b_code_interpreter import Sandbox + sandbox = Sandbox(template="your-custom-sandbox-name") + execution = sandbox.run_code("print('hello')") + sandbox.kill() # Or you can use `with` which handles closing the sandbox for you - with CodeInterpreter(template="your-custom-sandbox-name") as sandbox: - execution = sandbox.notebook.exec_cell("print('hello')") + with Sandbox(template="your-custom-sandbox-name") as sandbox: + execution = sandbox.run_code("print('hello')") ``` **JavaScript/TypeScript** + ```js - import { CodeInterpreter } from '@e2b/code-interpreter' - const sandbox = await CodeInterpreter.create({ template: 'your-custom-sandbox-name' }) - const execution = await sandbox.notebook.execCell('print("hello")') - await sandbox.close() + import {Sandbox} from '@e2b/code-interpreter' + +const sandbox = await Sandbox.create({template: 'your-custom-sandbox-name'}) +const execution = await sandbox.runCode('print("hello")') +await sandbox.kill() ``` diff --git a/template/server/api/models/context.py b/template/server/api/models/context.py index 49ab1fc6..e02484d5 100644 --- a/template/server/api/models/context.py +++ b/template/server/api/models/context.py @@ -4,5 +4,5 @@ class Context(BaseModel): id: StrictStr = Field(description="Context ID") - name: StrictStr = Field(description="Context name") + language: StrictStr = Field(description="Language of the context") cwd: StrictStr = Field(description="Current working directory of the context") diff --git a/template/server/api/models/create_context.py b/template/server/api/models/create_context.py index c24c7fab..1e1fefb7 100644 --- a/template/server/api/models/create_context.py +++ b/template/server/api/models/create_context.py @@ -8,6 +8,6 @@ class CreateContext(BaseModel): default="/home/user", description="Current working directory", ) - name: Optional[StrictStr] = Field( - default="python", description="Name of the kernel" + language: Optional[StrictStr] = Field( + default="python", description="Language of the context" ) diff --git a/template/server/api/models/execution_request.py b/template/server/api/models/execution_request.py index ba49f2a8..bb691e9e 100644 --- a/template/server/api/models/execution_request.py +++ b/template/server/api/models/execution_request.py @@ -7,7 +7,10 @@ class ExecutionRequest(BaseModel): code: StrictStr = Field(description="Code to be executed") - context_id: Optional[StrictStr] = Field(default="default", description="Context ID") + context_id: Optional[StrictStr] = Field(default=None, description="Context ID") + language: Optional[StrictStr] = Field( + default=None, description="Language of the code" + ) env_vars: Optional[EnvVars] = Field( description="Environment variables", default=None ) diff --git a/template/server/api/models/result.py b/template/server/api/models/result.py index 09247117..d1ae16a4 100644 --- a/template/server/api/models/result.py +++ b/template/server/api/models/result.py @@ -18,8 +18,6 @@ class Result(BaseModel): 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 display calls don't have to have text representation, for the actual result the representation is always present for the result, the other representations are always optional. - - The class also provides methods to display the data in a Jupyter notebook. """ type: OutputType = OutputType.RESULT diff --git a/template/server/consts.py b/template/server/consts.py new file mode 100644 index 00000000..e1b92a06 --- /dev/null +++ b/template/server/consts.py @@ -0,0 +1 @@ +JUPYTER_BASE_URL = "http://localhost:8888" diff --git a/template/server/contexts.py b/template/server/contexts.py new file mode 100644 index 00000000..87b8cad5 --- /dev/null +++ b/template/server/contexts.py @@ -0,0 +1,63 @@ +import logging +import uuid +from typing import Optional + +from api.models.context import Context +from fastapi.responses import PlainTextResponse + +from consts import JUPYTER_BASE_URL +from errors import ExecutionError +from messaging import ContextWebSocket + +logger = logging.Logger(__name__) + + +def normalize_language(language: Optional[str]) -> str: + if not language: + return "python" + + language = language.lower().strip() + + if language == "js": + return "javascript" + + return language + + +async def create_context(client, websockets: dict, language: str, cwd: str) -> Context: + data = { + "path": str(uuid.uuid4()), + "kernel": {"name": language}, + "type": "notebook", + "name": str(uuid.uuid4()), + } + logger.debug(f"Creating new {language} context") + + response = await client.post(f"{JUPYTER_BASE_URL}/api/sessions", json=data) + + if not response.is_success: + return PlainTextResponse( + f"Failed to create context: {response.text}", + status_code=500, + ) + + session_data = response.json() + session_id = session_data["id"] + context_id = session_data["kernel"]["id"] + + logger.debug(f"Created context {context_id}") + + ws = ContextWebSocket(context_id, session_id, language, cwd) + await ws.connect() + websockets[context_id] = ws + + logger.info(f"Setting working directory to {cwd}") + try: + await ws.change_current_directory(cwd) + except ExecutionError as e: + return PlainTextResponse( + "Failed to set working directory", + status_code=500, + ) + + return Context(language=language, id=context_id, cwd=cwd) diff --git a/template/server/main.py b/template/server/main.py index 202a6dc6..96e0b856 100644 --- a/template/server/main.py +++ b/template/server/main.py @@ -4,15 +4,16 @@ from typing import Dict, Union, Literal, List -from pydantic import StrictStr from contextlib import asynccontextmanager -from fastapi import FastAPI, HTTPException +from fastapi import FastAPI +from fastapi.responses import PlainTextResponse from api.models.context import Context from api.models.create_context import CreateContext from api.models.execution_request import ExecutionRequest -from errors import ExecutionError -from messaging import JupyterKernelWebSocket +from consts import JUPYTER_BASE_URL +from contexts import create_context, normalize_language +from messaging import ContextWebSocket from stream import StreamingListJsonResponse @@ -22,10 +23,8 @@ http_logger.setLevel(logging.WARNING) -JUPYTER_BASE_URL = "http://localhost:8888" - -websockets: Dict[Union[StrictStr, Literal["default"]], JupyterKernelWebSocket] = {} -global default_kernel_id +websockets: Dict[Union[str, Literal["default"]], ContextWebSocket] = {} +default_websockets: Dict[str, str] = {} global client @@ -34,16 +33,18 @@ async def lifespan(app: FastAPI): global client client = httpx.AsyncClient() - global default_kernel_id with open("/root/.jupyter/kernel_id") as file: - default_kernel_id = file.read().strip() + default_context_id = file.read().strip() - default_ws = JupyterKernelWebSocket( - default_kernel_id, + default_ws = ContextWebSocket( + default_context_id, str(uuid.uuid4()), "python", "/home/user", ) + default_websockets["python"] = default_context_id + websockets["default"] = default_ws + websockets[default_context_id] = default_ws logger.info("Connecting to default runtime") await default_ws.connect() @@ -65,104 +66,83 @@ async def lifespan(app: FastAPI): @app.get("/health") -async def health(): +async def get_health(): return "OK" @app.post("/execute") -async def execute(request: ExecutionRequest): +async def post_execute(request: ExecutionRequest): logger.info(f"Executing code: {request.code}") - if request.context_id: - ws = websockets.get(request.context_id) + if request.context_id and request.language: + return PlainTextResponse( + "Only one of context_id or language can be provided", + status_code=400, + ) + + context_id = None + if request.language: + language = normalize_language(request.language) + context_id = default_websockets.get(language) + + if not context_id: + context = await create_context(client, websockets, language, "/home/user") + context_id = context.id + + elif request.context_id: + context_id = request.context_id - if not ws: - raise HTTPException( - status_code=404, - detail=f"Kernel {request.context_id} not found", - ) + if context_id: + ws = websockets.get(context_id, None) else: ws = websockets["default"] + if not ws: + return PlainTextResponse( + f"Context {request.context_id} not found", + status_code=404, + ) + return StreamingListJsonResponse( ws.execute(request.code, env_vars=request.env_vars) ) @app.post("/contexts") -async def create_context(request: CreateContext) -> Context: - logger.info(f"Creating new kernel") - - data = { - "path": str(uuid.uuid4()), - "kernel": {"name": request.name}, - "type": "notebook", - "name": str(uuid.uuid4()), - } - logger.debug(f"Creating new kernel with data: {data}") - - response = await client.post(f"{JUPYTER_BASE_URL}/api/sessions", json=data) - - if not response.is_success: - raise HTTPException( - status_code=500, - detail=f"Failed to create kernel: {response.text}", - ) - - session_data = response.json() - session_id = session_data["id"] - kernel_id = session_data["kernel"]["id"] - - logger.debug(f"Created kernel {kernel_id}") - - ws = JupyterKernelWebSocket( - kernel_id, - session_id, - request.name, - request.cwd, - ) - await ws.connect() - - websockets[kernel_id] = ws +async def post_contexts(request: CreateContext) -> Context: + logger.info(f"Creating a new context") - if request.cwd: - logger.info(f"Setting working directory to {request.cwd}") - try: - await ws.change_current_directory(request.cwd) - except ExecutionError as e: - raise HTTPException( - status_code=500, - detail="Failed to set working directory", - ) from e + language = normalize_language(request.language) + cwd = request.cwd or "/home/user" - return Context(name=request.name, id=kernel_id, cwd=request.cwd) + return await create_context(client, websockets, language, cwd) @app.get("/contexts") -async def list_contexts() -> List[Context]: - logger.info(f"Listing kernels") +async def get_contexts() -> List[Context]: + logger.info(f"Listing contexts") - kernel_ids = list(websockets.keys()) + context_ids = list(websockets.keys()) return [ Context( - id=websockets[kernel_id].kernel_id, - name=websockets[kernel_id].name, - cwd=websockets[kernel_id].cwd, + id=websockets[context_id].context_id, + language=websockets[context_id].language, + cwd=websockets[context_id].cwd, ) - for kernel_id in kernel_ids + for context_id in context_ids ] @app.post("/contexts/{context_id}/restart") async def restart_context(context_id: str) -> None: - logger.info(f"Restarting kernel {context_id}") + logger.info(f"Restarting context {context_id}") ws = websockets.get(context_id, None) if not ws: - raise HTTPException( + return PlainTextResponse( + f"Context {context_id} not found", status_code=404, - detail=f"Kernel {context_id} not found", ) session_id = ws.session_id @@ -170,18 +150,18 @@ async def restart_context(context_id: str) -> None: await ws.close() response = await client.post( - f"{JUPYTER_BASE_URL}/api/kernels/{ws.kernel_id}/restart" + f"{JUPYTER_BASE_URL}/api/kernels/{ws.context_id}/restart" ) if not response.is_success: - raise HTTPException( + return PlainTextResponse( + f"Failed to restart context {context_id}", status_code=500, - detail=f"Failed to restart kernel {context_id}", ) - ws = JupyterKernelWebSocket( - ws.kernel_id, + ws = ContextWebSocket( + ws.context_id, session_id, - ws.name, + ws.language, ws.cwd, ) @@ -192,13 +172,13 @@ async def restart_context(context_id: str) -> None: @app.delete("/contexts/{context_id}") async def remove_context(context_id: str) -> None: - logger.info(f"Removing kernel {context_id}") + logger.info(f"Removing context {context_id}") ws = websockets.get(context_id, None) if not ws: - raise HTTPException( + return PlainTextResponse( + f"Context {context_id} not found", status_code=404, - detail=f"Kernel {context_id} not found", ) try: @@ -206,11 +186,11 @@ async def remove_context(context_id: str) -> None: except: pass - response = await client.delete(f"{JUPYTER_BASE_URL}/api/kernels/{ws.kernel_id}") + response = await client.delete(f"{JUPYTER_BASE_URL}/api/kernels/{ws.context_id}") if not response.is_success: - raise HTTPException( + return PlainTextResponse( + f"Failed to remove context {context_id}", status_code=500, - detail=f"Failed to remove context {context_id}", ) del websockets[context_id] diff --git a/template/server/messaging.py b/template/server/messaging.py index e5673358..cfa6786c 100644 --- a/template/server/messaging.py +++ b/template/server/messaging.py @@ -44,21 +44,21 @@ def __init__(self, in_background: bool = False): self.in_background = in_background -class JupyterKernelWebSocket: +class ContextWebSocket: _ws: Optional[WebSocketClientProtocol] = None _receive_task: Optional[asyncio.Task] = None def __init__( self, - kernel_id: str, + context_id: str, session_id: str, - name: str, + language: str, cwd: str, ): - self.name = name + self.language = language self.cwd = cwd - self.kernel_id = kernel_id - self.url = f"ws://localhost:8888/api/kernels/{kernel_id}/channels" + self.context_id = context_id + self.url = f"ws://localhost:8888/api/kernels/{context_id}/channels" self.session_id = session_id self._executions: Dict[str, Execution] = {} @@ -206,12 +206,12 @@ async def _process_message(self, data: dict): data["msg_type"] == "status" and data["content"]["execution_state"] == "restarting" ): - logger.error("Kernel is restarting") + logger.error("Context is restarting") for execution in self._executions.values(): await execution.queue.put( Error( - name="KernelRestarting", - value="Kernel was restarted", + name="ContextRestarting", + value="Context was restarted", traceback="", ) ) @@ -315,7 +315,7 @@ async def _process_message(self, data: dict): logger.warning(f"[UNHANDLED MESSAGE TYPE]: {data['msg_type']}") async def close(self): - logger.debug(f"Closing WebSocket {self.kernel_id}") + logger.debug(f"Closing WebSocket {self.context_id}") if self._ws is not None: await self._ws.close() diff --git a/template/test.Dockerfile b/template/test.Dockerfile index 91334f98..97b7113d 100644 --- a/template/test.Dockerfile +++ b/template/test.Dockerfile @@ -18,6 +18,11 @@ ENV PIP_DEFAULT_TIMEOUT=100 \ COPY ./requirements.txt requirements.txt RUN pip install --no-cache-dir -r requirements.txt && ipython kernel install --name "python3" --user +# Javascript Kernel +RUN npm install -g node-gyp +RUN npm install -g --unsafe-perm ijavascript +RUN ijsinstall --install=global + # Create separate virtual environment for server RUN python -m venv $SERVER_PATH/.venv