e2b-ai-function.ts

e2b-ai-function.ts

import { CodeInterpreter, type ProcessMessage } from '@e2b/code-interpreter'
import { z } from 'zod'

import { createAIFunction } from '../create-ai-function.js'
import { getEnv } from '../utils.js'

export const e2b = createAIFunction(
  {
    name: 'execute_python',
    description: `
Execute python code in a Jupyter notebook cell and returns any result, stdout, stderr, display_data, and error.

- code has access to the internet and can make api requests
- code has access to the filesystem and can read/write files
- coce can install any pip package (if it exists) if you need to, but the usual packages for data analysis are already preinstalled
- code uses python3
- code is executed in a secure sandbox environment, so you don't need to worry about safety
      `.trim(),
    inputSchema: z.object({
      code: z
        .string()
        .describe('Python code to execute in a single notebook cell.')
    })
  },
  async ({ code }) => {
    const sandbox = await CodeInterpreter.create({
      apiKey: getEnv('E2B_API_KEY')
    })

    try {
      const exec = await sandbox.notebook.execCell(code, {
        onStderr: (msg: ProcessMessage) => {
          console.warn('[Code Interpreter stderr]', msg)
        },

        onStdout: (stdout: ProcessMessage) => {
          console.log('[Code Interpreter stdout]', stdout)
        }
      })

      if (exec.error) {
        console.error('[Code Interpreter error]', exec.error)
        throw new Error(exec.error.value)
      }

      return exec.results.map((result) => result.toJSON())
    } finally {
      await sandbox.close()
    }
  }
)