Add graphs as separate result attribute, improve schema, add typing, add supergraph

Author: jakubno

js/README.md

@@ -24,7 +24,7 @@ const sandbox = await CodeInterpreter.create()
 await sandbox.notebook.execCell('x = 1')
 
 const execution = await sandbox.notebook.execCell('x+=1; x')
-console.log(execution.text)  // outputs 2
+console.log(execution.text) // outputs 2
 
 await sandbox.close()
 ```
@@ -48,7 +48,7 @@ plt.show()
 `
 
 // you can install dependencies in "jupyter notebook style"
-await sandbox.notebook.execCell("!pip install matplotlib")
+await sandbox.notebook.execCell('!pip install matplotlib')
 
 const execution = await sandbox.notebook.execCell(code)
 
@@ -80,7 +80,7 @@ const sandbox = await CodeInterpreter.create()
 await sandbox.notebook.execCell(code, {
   onStdout: (out) => console.log(out),
   onStderr: (outErr) => console.error(outErr),
-  onResult: (result) => console.log(result.text)
+  onResult: (result) => console.log(result.text),
 })
 
 await sandbox.close()

js/src/graphs.ts

@@ -0,0 +1,114 @@
+export enum GraphType {
+  LINE = 'line',
+  SCATTER = 'scatter',
+  BAR = 'bar',
+  PIE = 'pie',
+  BOX_AND_WHISKER = 'box_and_whisker',
+  SUPERGRAPH = 'supergraph',
+  UNKNOWN = 'unknown',
+}
+
+export type Graph = {
+  type: GraphType
+  title: string
+  elements: any[]
+}
+
+type Graph2D = Graph & {
+  x_label?: string
+  y_label?: string
+  x_unit?: string
+  y_unit?: string
+}
+
+export type PointData = {
+  label: string
+  points: [number, number][]
+}
+
+type PointGraph = Graph2D & {
+  x_ticks: number[]
+  x_tick_labels: string[]
+  y_ticks: number[]
+  y_tick_labels: string[]
+  elements: PointData[]
+}
+
+export type LineGraph = PointGraph & {
+  type: GraphType.LINE
+}
+
+export type ScatterGraph = PointGraph & {
+  type: GraphType.SCATTER
+}
+
+export type BarData = {
+  label: string
+  value: string
+}
+
+export type BarGraph = Graph2D & {
+  type: GraphType.BAR
+  elements: BarData[]
+}
+
+export type PieData = {
+  label: string
+  angle: number
+  radius: number
+}
+
+export type PieGraph = Graph & {
+  type: GraphType.PIE
+  elements: PieData[]
+}
+
+export type BoxAndWhiskerData = {
+  label: string
+  min: number
+  first_quartile: number
+  median: number
+  third_quartile: number
+  max: number
+}
+
+export type BoxAndWhiskerGraph = Graph2D & {
+  type: GraphType.BOX_AND_WHISKER
+  elements: BoxAndWhiskerData[]
+}
+
+export type SuperGraph = Graph & {
+  type: GraphType.SUPERGRAPH
+  elements: Graph[]
+}
+
+export type GraphTypes =
+  | LineGraph
+  | ScatterGraph
+  | BarGraph
+  | PieGraph
+  | BoxAndWhiskerGraph
+  | SuperGraph
+export function deserializeGraph(data: any): Graph {
+  switch (data.type) {
+    case GraphType.LINE:
+      return { ...data } as LineGraph
+    case GraphType.SCATTER:
+      return { ...data } as ScatterGraph
+    case GraphType.BAR:
+      return { ...data } as BarGraph
+    case GraphType.PIE:
+      return { ...data } as PieGraph
+    case GraphType.BOX_AND_WHISKER:
+      return { ...data } as BoxAndWhiskerGraph
+    case GraphType.SUPERGRAPH:
+      const graphs = data.data.map((g: any) => deserializeGraph(g))
+      delete data.data
+      return {
+        ...data,
+        data: graphs,
+      } as SuperGraph
+    default:
+      return { ...data, type: GraphType.UNKNOWN } as Graph
+  }
+}

js/src/index.ts

@@ -2,8 +2,30 @@ export * from 'e2b'
 
 export { CodeInterpreter, JupyterExtension } from './codeInterpreter'
 
-export type { Logs, ExecutionError, Result, Execution, MIMEType, RawData, OutputMessage } from './messaging'
-
+export type {
+  Logs,
+  ExecutionError,
+  Result,
+  Execution,
+  MIMEType,
+  RawData,
+  OutputMessage,
+} from './messaging'
+export type {
+  GraphType,
+  GraphTypes,
+  Graph,
+  BarGraph,
+  BarData,
+  LineGraph,
+  ScatterGraph,
+  BoxAndWhiskerGraph,
+  BoxAndWhiskerData,
+  PieGraph,
+  PieData,
+  SuperGraph,
+  PointData,
+} from './graphs'
 import { CodeInterpreter } from './codeInterpreter'
 
 export default CodeInterpreter

js/src/messaging.ts

@@ -1,4 +1,5 @@
 import { NotFoundError, SandboxError, TimeoutError } from 'e2b'
+import { GraphTypes } from './graphs'
 
 export async function extractError(res: Response) {
   if (res.ok) {
@@ -24,9 +25,8 @@ export class OutputMessage {
      * Unix epoch in nanoseconds
      */
     public readonly timestamp: number,
-    public readonly error: boolean,
-  ) {
-  }
+    public readonly error: boolean
+  ) {}
 
   public toString() {
     return this.line
@@ -50,25 +50,26 @@ export class ExecutionError {
     /**
      * The raw traceback of the error.
      **/
-    public traceback: string,
-  ) { }
+    public traceback: string
+  ) {}
 }
 
 /**
  * Represents a MIME type.
  */
 export type MIMEType = string
 
-type Data = {
+type E2BData = {
   data: Record<string, unknown>
+  graph: GraphTypes
 }
 
 /**
  * Dictionary that maps MIME types to their corresponding representations of the data.
  */
 export type RawData = {
   [key: MIMEType]: string
-} & Data
+} & E2BData
 
 /**
  * Represents the data to be displayed as a result of executing a cell in a Jupyter notebook.
@@ -124,6 +125,10 @@ export class Result {
    * Contains the data from DataFrame.
    */
   readonly data?: Record<string, unknown>
+  /**
+   * Contains the graph data.
+   */
+  readonly graph?: GraphTypes
   /**
    * Extra data that can be included. Not part of the standard types.
    */
@@ -146,10 +151,12 @@ 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
 
+    this.data = data['data']
+    this.graph = data['graph']
+
     this.extra = {}
 
     for (const key of Object.keys(data)) {
@@ -166,7 +173,7 @@ export class Result {
           'json',
           'javascript',
           'data',
-          'extra'
+          'extra',
         ].includes(key)
       ) {
         this.extra[key] = data[key]
@@ -234,7 +241,7 @@ export class Result {
       latex: this.latex,
       json: this.json,
       javascript: this.javascript,
-      ...(Object.keys(this.extra).length > 0 ? { extra: this.extra } : {})
+      ...(Object.keys(this.extra).length > 0 ? { extra: this.extra } : {}),
     }
   }
 }
@@ -274,7 +281,7 @@ export class Execution {
      * Execution count of the cell.
      */
     public executionCount?: number
-  ) { }
+  ) {}
 
   /**
    * Returns the text representation of the main result of the cell.
@@ -294,23 +301,26 @@ export class Execution {
     return {
       results: this.results,
       logs: this.logs,
-      error: this.error
+      error: this.error,
     }
   }
 }
 
 export async function parseOutput(
   execution: Execution,
   line: string,
-  onStdout?: (output: OutputMessage) => (Promise<any> | any),
-  onStderr?: (output: OutputMessage) => (Promise<any> | any),
-  onResult?: (data: Result) => (Promise<any> | any),
+  onStdout?: (output: OutputMessage) => Promise<any> | any,
+  onStderr?: (output: OutputMessage) => Promise<any> | any,
+  onResult?: (data: Result) => Promise<any> | any
 ) {
   const msg = JSON.parse(line)
 
   switch (msg.type) {
     case 'result':
-      const result = new Result({ ...msg, type: undefined, is_main_result: undefined }, msg.is_main_result)
+      const result = new Result(
+        { ...msg, type: undefined, is_main_result: undefined },
+        msg.is_main_result
+      )
       execution.results.push(result)
       if (onResult) {
         await onResult(result)

js/tests/displayData.test.ts

@@ -3,7 +3,6 @@ import { expect } from 'vitest'
 import { sandboxTest } from './setup'
 
 sandboxTest('display data', async ({ sandbox }) => {
-
   // plot random graph
   const result = await sandbox.notebook.execCell(`
         import matplotlib.pyplot as plt

js/tests/envVars.test.ts

@@ -1,30 +1,46 @@
 import { expect } from 'vitest'
 
 import { sandboxTest } from './setup'
-import { CodeInterpreter } from "../src";
+import { CodeInterpreter } from '../src'
 
 // 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('env vars', async () => {
-  const sandbox = await CodeInterpreter.create({ envs: { TEST_ENV_VAR: "supertest" } })
-  const result = await sandbox.notebook.execCell(`import os; x = os.getenv('TEST_ENV_VAR'); x`)
+  const sandbox = await CodeInterpreter.create({
+    envs: { TEST_ENV_VAR: 'supertest' },
+  })
+  const result = await sandbox.notebook.execCell(
+    `import os; x = os.getenv('TEST_ENV_VAR'); x`
+  )
 
   expect(result.results[0].text.trim()).toEqual('supertest')
 })
 
-sandboxTest('env vars on sandbox', async ({sandbox}) => {
-  const result = await sandbox.notebook.execCell("import os; os.getenv('FOO')", {envs: {FOO: "bar"}})
+sandboxTest('env vars on sandbox', async ({ sandbox }) => {
+  const result = await sandbox.notebook.execCell(
+    "import os; os.getenv('FOO')",
+    { envs: { FOO: 'bar' } }
+  )
 
   expect(result.results[0].text.trim()).toEqual('bar')
 })
 
 sandboxTest('env vars on sandbox override', async () => {
-  const sandbox = await CodeInterpreter.create({ envs: { FOO: "bar", SBX: "value" } })
-  await sandbox.notebook.execCell("import os; os.environ['FOO'] = 'bar'; os.environ['RUNTIME_ENV'] = 'value'")
-  const result = await sandbox.notebook.execCell("import os; os.getenv('FOO')", {envs: {FOO: "baz"}})
+  const sandbox = await CodeInterpreter.create({
+    envs: { FOO: 'bar', SBX: 'value' },
+  })
+  await sandbox.notebook.execCell(
+    "import os; os.environ['FOO'] = 'bar'; os.environ['RUNTIME_ENV'] = 'value'"
+  )
+  const result = await sandbox.notebook.execCell(
+    "import os; os.getenv('FOO')",
+    { envs: { FOO: 'baz' } }
+  )
 
   expect(result.results[0].text.trim()).toEqual('baz')
 
-  const result2 = await sandbox.notebook.execCell("import os; os.getenv('RUNTIME_ENV')")
+  const result2 = await sandbox.notebook.execCell(
+    "import os; os.getenv('RUNTIME_ENV')"
+  )
   expect(result2.results[0].text.trim()).toEqual('value')
 
   const result3 = await sandbox.notebook.execCell("import os; os.getenv('SBX')")