Hello World.ts

npm init -y \
&& npm i --save-dev typescript tsx @types/node \
&& npm i @e2b/code-interpreter

E2B_API_KEY="e2b-api-key"

import { Anthropic } from '@anthropic-ai/sdk'

export const MODEL_NAME = 'claude-3-opus-20240229'

export const SYSTEM_PROMPT = `
## your job & context
you are a python data scientist. you are given tasks to complete and you run python code to solve them.
- the python code runs in jupyter notebook.
- every time you call \`execute_python\` tool, the python code is executed in a separate cell. it's okay to multiple calls to \`execute_python\`.
- display visualizations using matplotlib or any other visualization library directly in the notebook. don't worry about saving the visualizations to a file.
- you have access to the internet and can make api requests.
- you also have access to the filesystem and can read/write files.
- you can install any pip package (if it exists) if you need to but the usual packages for data analysis are already preinstalled.
- you can run any python code you want, everything is running in a secure sandbox environment.
`

export const tools: Anthropic.Tool[] = [
  {
    name: 'execute_python',
    description: 'Execute python code in a Jupyter notebook cell and returns any result, stdout, stderr, display_data, and error.',
    input_schema: {
      type: 'object',
      properties: {
        code: {
          type: 'string',
          description: 'The python code to execute in a single cell.'
        }
      },
      required: ['code']
    }
  }
]

import { CodeInterpreter } from '@e2b/code-interpreter'

export async function codeInterpret(codeInterpreter: CodeInterpreter, code: string) {
  console.log(`\n${'='.repeat(50)}\n> Running following AI-generated code:\n${code}\n${'='.repeat(50)}`);

  const exec = await codeInterpreter.notebook.execCell(
    code,
    {
      // You can stream logs from the code interpreter
      // onStderr: (stderr: string) => console.log("\n[Code Interpreter stdout]", stderr),
      // onStdout: (stdout: string) => console.log("\n[Code Interpreter stderr]", stdout),
      //
      // You can also stream additional results like charts, images, etc.
      // onResult: ...
    }
  )

  if (exec.error) {
    console.log('[Code Interpreter error]', exec.error) // Runtime error
    return undefined
  }

  return exec
}

ANTHROPIC_API_KEY="anthropic-api-key"

npm i @anthropic-ai/sdk@0.28.0

import * as fs from 'fs'

import 'dotenv/config'
import { CodeInterpreter, Execution } from '@e2b/code-interpreter'
import Anthropic from '@anthropic-ai/sdk'

import {
  MODEL_NAME,
  SYSTEM_PROMPT,
  tools,
} from './model'
import { codeInterpret } from './codeInterpreter'

const anthropic = new Anthropic()

async function chat(codeInterpreter: CodeInterpreter, userMessage: string): Promise<Execution | undefined> {
  console.log('Waiting for Claude...')

  const msg = await anthropic.messages.create({
    model: MODEL_NAME,
    system: SYSTEM_PROMPT,
    max_tokens: 4096,
    messages: [{role: 'user', content: userMessage}],
    tools,
  })

  console.log(`\n${'='.repeat(50)}\nModel response: ${msg.content}\n${'='.repeat(50)}`)
  console.log(msg)

  if (msg.stop_reason === 'tool_use') {
    const toolBlock = msg.content.find((block) => block.type === 'tool_use');
    const toolName = toolBlock.name
    const toolInput = toolBlock.input

    console.log(`\n${'='.repeat(50)}\nUsing tool: ${toolName}\n${'='.repeat(50)}`);

    if (toolName === 'execute_python') {
      const code = toolInput.code
      return codeInterpret(codeInterpreter, code)
    }
    return undefined
  }
}

async function run() {
  const userMessage = 'Visualize a distribution of height of men based on the latest data you know. Also print the median value.'

  const codeInterpreter = await CodeInterpreter.create()

  const codeOutput = await chat(codeInterpreter, userMessage)
  if (!codeOutput) {
    console.log('No code output')
    return
  }

  const logs = codeOutput.logs
  console.log(logs)

  if (codeOutput.results.length == 0) {
    console.log('No results')
    return
  }

  const firstResult = codeOutput.results[0]
  // Print description of the first rich result
  console.log(firstResult.text)

  await codeInterpreter.close()
}

run()

$ tsx index.ts

async function run() {
  const userMessage = 'Visualize a distribution of height of men based on the latest data you know. Also print the median value.'

  const codeInterpreter = await CodeInterpreter.create()

  const codeOutput = await chat(codeInterpreter, userMessage)
  if (!codeOutput) {
    console.log('No code output')
    return
  }

  const logs = codeOutput.logs
  console.log(logs)

  if (codeOutput.results.length == 0) {
    console.log('No results')
    return
  }

  const firstResult = codeOutput.results[0]
  // Print description of the first rich result
  console.log(firstResult.text)

  // If we received a chart in PNG form, we can visualize it
  if (firstResult.png) {
      // Decode the base64 encoded PNG data
      const pngData = Buffer.from(firstResult.png, 'base64');

      // Generate a unique filename for the PNG
      const filename = 'chart.png';

      // Save the decoded PNG data to a file
      fs.writeFileSync(filename, pngData);

      console.log(`Saved chart to ${filename}`);
  }

  await codeInterpreter.close()
}

run()