app.ts

// Code interpreter with IBM WatsonX AI Graphite
// This example uses the [E2B SDK](https://github.com/e2b-dev/code-interpreter) as a code interpreter for [IBM Granite Code](https://www.ibm.com/granite/docs/models/code/) on [IBM WatsonX](https://www.ibm.com/watsonx).
// The code generated by the LLM runs in a [E2B secure cloud sandbox](https://e2b.dev/docs/sandbox/overview).

// 1. Imports and API keys
// You need to get your API keys and save them in .env file. You can copy and rename the .env.template file. We import all necessary libraries.

import { config } from "dotenv";
import { WatsonXAI } from "@ibm-cloud/watsonx-ai";
import { IamAuthenticator } from "ibm-cloud-sdk-core";
import { Sandbox } from "@e2b/code-interpreter";
import fs from "fs";

// Load environment variables from .env file
config({ path: "./.env" });

// 2. Initialize WatsonX AI Client
// We initialize the WatsonX AI client using the credentials stored in the environment variables. This client will be used to interact with the IBM WatsonX Granite model.

// Environment variables
const WATSONX_PROJECT_ID = process.env.WATSONX_PROJECT_ID;
const WATSONX_URL = process.env.WATSONX_URL || "https://us-south.ml.cloud.ibm.com";
const WATSONX_API_KEY = process.env.WATSONX_API_KEY;

// Initialize WatsonX client
const watsonxAIService = WatsonXAI.newInstance({
version: "2024-05-31",
serviceUrl: WATSONX_URL,
authenticator: new IamAuthenticator({ apikey: WATSONX_API_KEY || "" }),
});

// Parameters for the text generation request
const params = {
modelId: "ibm/granite-34b-code-instruct",
projectId: WATSONX_PROJECT_ID!,
parameters: {
    max_new_tokens: 1024,
},
};

console.log("WatsonX service initialized");

// 3. Configure System Prompt and Generate a Response
// We define a system prompt that instructs the model on how to behave. Then, we send a user query to the model and retrieve a response.

// System prompt configuration
// Because this is Markdown, the LLM will respond in Markdown
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.
- 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 but the usual packages for data analysis are already preinstalled.
`;

// Specify the prompt
const userMessage = "Plot a 3D chart of sin x cos y.";
console.log(`User Message: ${userMessage}`);

const prompt = SYSTEM_PROMPT + userMessage + "\n";

// Generate response using the Watson SDK
console.log("Generating response from WatsonX...");
const response = await watsonxAIService.generateText({ input: prompt, ...params });
const content = response.result.results[0].generated_text;

console.log(`Model response:`);
console.log(`${"=".repeat(50)}\n${content}\n${"=".repeat(50)}`);

// 4. Extract and Display AI-Generated Code
// The AI-generated response may include Markdown formatting. We extract the code block from the response using regular expressions.

// Extract code blocks
const pattern = /```(?:python)?[\n\r](.*?)```/gs;
const results = Array.from(content.matchAll(pattern));
const codeBlock = results.length ? results.map(match => match[1]).join("\n") : content;

// Print the code block
console.log(`AI-generated code:`);
console.log(`${"=".repeat(50)}\n${codeBlock}\n${"=".repeat(50)}`);

// 5. Execute Code in an E2B Code Interpreter
// We run the AI-generated code in a secure E2B sandbox environment, capturing its output, errors, and runtime exceptions.

// Run the code in an E2B code interpreter
const codeInterpreter = await Sandbox.create();
const execution = await codeInterpreter.runCode(codeBlock);

console.log("Stdout:", execution.logs.stdout);
console.log("Stderr:", execution.logs.stderr);
console.log("Python runtime error:", execution.error ?? "None");

// 6. Display Generated Figures
// If the AI-generated code produces visualizations, we retrieve and display them.

// Print the first figure
const figures = execution.results.map(result => {
// In a real script, you'd want to save this to a file or display it another way
const binaryData = Buffer.from(result.png, 'base64');
return binaryData; // In a script, you'd likely save this to a file
});

console.log(`Generated ${figures.length} figures`);
fs.writeFileSync('figure.png', figures[0]);

// 7. Terminate the Code Interpreter
// Finally, we safely terminate the E2B code interpreter to free up resources.

// Kill the code interpreter
await codeInterpreter.kill();