From 330e7910685d33bcfe45a35b0a755480d7e28554 Mon Sep 17 00:00:00 2001 From: Greg Richardson Date: Tue, 4 Jun 2024 16:08:51 -0600 Subject: [PATCH] feat(ai-sandbox): multi-message inputs + more tests --- package-lock.json | 15 ++- packages/ai-sandbox/jest.config.js | 3 +- packages/ai-sandbox/package.json | 1 + packages/ai-sandbox/src/llm.test.ts | 187 +++++++++++++++++++++++----- packages/ai-sandbox/src/llm.ts | 187 +++++++++++++--------------- packages/ai-sandbox/test/util.ts | 45 +++++++ 6 files changed, 307 insertions(+), 131 deletions(-) create mode 100644 packages/ai-sandbox/test/util.ts diff --git a/package-lock.json b/package-lock.json index a99f7238255..6ad9feff833 100644 --- a/package-lock.json +++ b/package-lock.json @@ -52212,6 +52212,7 @@ "@types/common-tags": "^1.8.4", "babel-jest": "^29.7.0", "babel-plugin-transform-import-meta": "^2.2.1", + "chalk": "^5.3.0", "common-tags": "^1.8.2", "dotenv": "^16.3.1", "jest": "^29.7.0", @@ -52220,6 +52221,18 @@ "typescript": "^5.4.3" } }, + "packages/ai-sandbox/node_modules/chalk": { + "version": "5.3.0", + "resolved": "https://registry.npmjs.org/chalk/-/chalk-5.3.0.tgz", + "integrity": "sha512-dLitG79d+GV1Nb/VYcCDFivJeK1hiukt9QjRNVOsUtTy1rR1YJsmpGGTZ3qJos+uw7WmWF4wUwBd9jxjocFC2w==", + "dev": true, + "engines": { + "node": "^12.17.0 || ^14.13 || >=16.0.0" + }, + "funding": { + "url": "https://github.com/chalk/chalk?sponsor=1" + } + }, "packages/api-types": { "version": "0.0.0", "license": "MIT", @@ -53686,4 +53699,4 @@ "dev": true } } -} \ No newline at end of file +} diff --git a/packages/ai-sandbox/jest.config.js b/packages/ai-sandbox/jest.config.js index 5a8a440e66a..d4e2ef44f2c 100644 --- a/packages/ai-sandbox/jest.config.js +++ b/packages/ai-sandbox/jest.config.js @@ -7,7 +7,8 @@ module.exports = { '^.+\\.(js|jsx)$': 'babel-jest', }, setupFiles: ['./test/setup.ts'], + testTimeout: 30000, transformIgnorePatterns: [ - 'node_modules/(?!(@jitl/quickjs-wasmfile-.*|mdast-.*|micromark|micromark-.*|unist-.*|decode-named-character-reference|character-entities)/)', + 'node_modules/(?!(@jitl/quickjs-wasmfile-.*|mdast-.*|micromark|micromark-.*|unist-.*|decode-named-character-reference|character-entities|chalk)/)', ], } diff --git a/packages/ai-sandbox/package.json b/packages/ai-sandbox/package.json index 4ffe728b98c..459954c8892 100644 --- a/packages/ai-sandbox/package.json +++ b/packages/ai-sandbox/package.json @@ -22,6 +22,7 @@ "@types/common-tags": "^1.8.4", "babel-jest": "^29.7.0", "babel-plugin-transform-import-meta": "^2.2.1", + "chalk": "^5.3.0", "common-tags": "^1.8.2", "dotenv": "^16.3.1", "jest": "^29.7.0", diff --git a/packages/ai-sandbox/src/llm.test.ts b/packages/ai-sandbox/src/llm.test.ts index 877f0546a44..e51732d95eb 100644 --- a/packages/ai-sandbox/src/llm.test.ts +++ b/packages/ai-sandbox/src/llm.test.ts @@ -1,31 +1,84 @@ import { describe, expect, test } from '@jest/globals' +import { codeBlock } from 'common-tags' import OpenAI from 'openai' -import { asSingleOutput, interpretPrompt } from './llm' +import { ChatCompletionMessageParam } from 'openai/resources' +import { withMetadata } from '../test/util' +import { interpretPrompt } from './llm' import { FunctionUnavailableError } from './runtime' const openAiKey = process.env.OPENAI_API_KEY const openai = new OpenAI({ apiKey: openAiKey }) describe('single stage llm', () => { - test('digits of pi', async () => { - const { exports, source, error } = await interpretPrompt( - openai, - 'list the first 10 digits of pi' - ) + test.concurrent('digits of pi', async () => { + const { exports, source, error } = await interpretPrompt(openai, [ + { + role: 'user', + content: 'list the first 10 digits of pi as an array', + }, + ]) - if (error) { - throw error - } + withMetadata({ source, exports }, () => { + if (error) { + throw error + } - console.log(source) - const output = asSingleOutput(exports) - expect(output).toBe('3141592653') + const { default: output } = exports + expect(output).toStrictEqual([3, 1, 4, 1, 5, 9, 2, 6, 5, 3]) + }) }) - test('with async function binding', async () => { + test.concurrent('complex math using digits of pi', async () => { + const { exports, source, error } = await interpretPrompt(openai, [ + { + role: 'user', + content: 'add every other digit from the first 10 digits of pi', + }, + ]) + + withMetadata({ source, exports }, () => { + if (error) { + throw error + } + + const { default: output } = exports + expect(output).toBe(19) + }) + }) + + test.concurrent('count words', async () => { + const { exports, source, error } = await interpretPrompt(openai, [ + { + role: 'user', + content: codeBlock` + how many words are there in the paragraph below? + + On a sunny day, Maria crafted spicy tacos with fresh salsa. + Friends gathered, tasted, and celebrated, declaring her creation the best in town. + Happiness filled the air. + `, + }, + ]) + + withMetadata({ source, exports }, () => { + if (error) { + throw error + } + + const { default: output } = exports + expect(output).toBe(27) + }) + }) + + test.concurrent('async function binding', async () => { const { exports, source, error } = await interpretPrompt( openai, - 'what is the weather like in singapore?', + [ + { + role: 'user', + content: 'what is the weather like in singapore?', + }, + ], { getWeather: { description: 'Gets the weather in the specified city.', @@ -35,28 +88,104 @@ describe('single stage llm', () => { } ) - if (error) { - throw error - } + withMetadata({ source, exports }, () => { + if (error) { + throw error + } - console.log(source) - const output = asSingleOutput(exports) - expect(output).toBe('hot') + const { default: output } = exports + expect(output).toBe('hot') + }) }) - test('with function unavailable', async () => { + test.concurrent('function unavailable throws', async () => { + const { exports, source, error } = await interpretPrompt(openai, [ + { + role: 'user', + content: 'what is the weather like in singapore?', + }, + ]) + + withMetadata({ source, exports }, () => { + if (!error) { + throw new Error('Expected error to be thrown') + } + + expect(error).toBeInstanceOf(FunctionUnavailableError) + }) + }) + + test.concurrent('function call to another LLM', async () => { const { exports, source, error } = await interpretPrompt( openai, - 'what is the weather like in singapore?' + [ + { + role: 'user', + content: codeBlock` + Extract the nouns from this text and sort them alphabetically (case insensitive): + + On a sunny day, Maria crafted spicy tacos with fresh salsa. + Friends gathered, tasted, and celebrated, declaring her creation the best in town. + Happiness filled the air. + `, + }, + ], + { + getNouns: { + description: 'Returns a list of nouns in a paragraph', + typeDef: '(text: string): Promise', + fn: async (text: string) => { + const chatCompletion = await openai.chat.completions.create({ + model: 'gpt-4o-2024-05-13', + messages: [ + { + role: 'system', + content: codeBlock` + Extract the nouns from this text and return in the following JSON format: + { "nouns": [...] } + `, + }, + { role: 'user', content: text }, + ], + response_format: { + type: 'json_object', + }, + max_tokens: 1024, + temperature: 0, + }) + + const [firstChoice] = chatCompletion.choices + const response = firstChoice.message.content + + if (!response) { + throw new Error('Failed to extract nouns') + } + + const { nouns } = JSON.parse(response) + + return nouns + }, + }, + } ) - if (!error) { - throw new Error('Expected error to be thrown') - } + withMetadata({ source, exports }, () => { + if (error) { + throw error + } - console.dir(error, { depth: null, showHidden: true }) - - console.log(source) - expect(error).toBeInstanceOf(FunctionUnavailableError) + const { default: output } = exports + expect(output).toStrictEqual([ + 'air', + 'creation', + 'day', + 'Friends', + 'Happiness', + 'Maria', + 'salsa', + 'tacos', + 'town', + ]) + }) }) }) diff --git a/packages/ai-sandbox/src/llm.ts b/packages/ai-sandbox/src/llm.ts index 5efb4be3c33..83976a717b8 100644 --- a/packages/ai-sandbox/src/llm.ts +++ b/packages/ai-sandbox/src/llm.ts @@ -2,7 +2,7 @@ import { codeBlock } from 'common-tags' import { fromMarkdown } from 'mdast-util-from-markdown' import { Code } from 'mdast-util-from-markdown/lib' import OpenAI from 'openai' -import { ChatCompletionMessageParam } from 'openai/resources' +import { ChatCompletionMessageParam, ChatCompletionSystemMessageParam } from 'openai/resources' import { ExecuteOptions, executeJS } from './runtime' export type FunctionBinding = { @@ -13,107 +13,106 @@ export type FunctionBinding = { export async function interpretPrompt( openai: OpenAI, - prompt: string, + messages: ChatCompletionMessageParam[], functionBindings: Record = {} ) { - const messages: ChatCompletionMessageParam[] = [ - { - role: 'system', - content: codeBlock` - You are a JavaScript code interpreter. Given the following JavaScript functions - and user prompt, generate the JavaScript code to accomplish that task. + const systemMessage: ChatCompletionSystemMessageParam = { + role: 'system', + content: codeBlock` + You are a JavaScript code interpreter. Given the following JavaScript functions + and user prompt, generate the JavaScript code to accomplish that task. - Here are the requirements: - - Use modern ES2020 syntax - - You are writing code for ES modules, so use imports and exports - - Top level await is supported, so use it if needed - - Output JavaScript and nothing else, wrapped in a markdown code block - - Create your own functions as necessary to keep the code modular - - Prefer creating more generic functions that are reusable than single purpose functions - Example prompt: count to 10 - Bad response: - \`\`\`javascript - function countToTen() { - const output = []; - for (let i = 1; i <= 10; i++) { - output.push(i); - } - return output.join(', '); + Here are the requirements: + - Use modern ES2020 syntax + - You are writing code for ES modules, so use imports and exports + - Top level await is supported, so use it if needed + - Output JavaScript and nothing else, wrapped in a markdown code block + - Use Math.* functions when performing math operations + - Create your own functions as necessary to keep the code modular + - Prefer creating more generic functions that are reusable than single purpose functions + Example prompt: count to 10 + Bad response: + \`\`\`javascript + function countToTen() { + const output = []; + for (let i = 1; i <= 10; i++) { + output.push(i); } + return output.join(', '); + } - export const output = countToTen(); - \`\`\` - Good response: - \`\`\`javascript - function countTo(digits) { - const output = []; - for (let i = 1; i <= n; i++) { - output.push(i); - } - return output.join(', '); + export const output = countToTen(); + \`\`\` + Good response: + \`\`\`javascript + function countTo(digits) { + const output = []; + for (let i = 1; i <= n; i++) { + output.push(i); } + return output.join(', '); + } - export const output = countTo(10); - \`\`\` - - Example prompt: get unfinished todos (with function available: \`function getTodos(isFinished: boolean): Promise\`) - Bad response: - \`\`\`javascript - async function getUnfinishedTodos() { - return getTodos(false); - } + export const output = countTo(10); + \`\`\` + + Example prompt: get unfinished todos (with function available: \`function getTodos(isFinished: boolean): Promise\`) + Bad response: + \`\`\`javascript + async function getUnfinishedTodos() { + return getTodos(false); + } - export const unfinishedTodos = await getUnfinishedTodos(); - \`\`\` - Good response: - \`\`\`javascript - export const unfinishedTodos = await getTodos(false); - \`\`\` - - Use as little functions/variables as possible - (ie. don't create a function or variable that wraps another - just call the function/variable directly) + export const unfinishedTodos = await getUnfinishedTodos(); + \`\`\` + Good response: + \`\`\`javascript + export const unfinishedTodos = await getTodos(false); + \`\`\` + - Use as little functions/variables as possible + (ie. don't create a function or variable that wraps another - just call the function/variable directly) + - If applicable, break the problem into intermediate steps and export each step along the way - Your result should be exported from a \`const\` like so: - \`\`\`javascript - // Here you can: import modules, create functions, perform loops or other logic - // Next export your results using _appropriate variable names_: - export const myResult = ... // call a function, etc - export const myOtherResult = ... // call a function, etc - // ... - \`\`\` + Your results should be exported like so: + \`\`\`javascript + // Here you can: import modules, create functions, perform loops or other logic + // Next export your results using _appropriate variable names_: + export const myDescriptiveStep = ... // call a function, etc (use a better name) + // ... any other intermediate steps _must be exported_ and _given a suitable name_ + export const myOutput = ... // call a function, etc (use a better name) + export default myOutput; // then export your final output as default + \`\`\` - The following functions are available in \`globalThis\`: - ${ - Object.keys(functionBindings).length > 0 - ? codeBlock` - \`\`\`typescript - ${Object.entries(functionBindings) - .map( - ([name, { description, typeDef }]) => codeBlock` - ${formatJsdocComment(description)} - ${formatFunctionDefinition(name, typeDef)} - ` - ) - .join('\n')} - \`\`\` - ` - : 'Only built-in JS functions are available' - } + The following functions are available in \`globalThis\`: + ${ + Object.keys(functionBindings).length > 0 + ? codeBlock` + \`\`\`typescript + ${Object.entries(functionBindings) + .map( + ([name, { description, typeDef }]) => codeBlock` + ${formatJsdocComment(description)} + ${formatFunctionDefinition(name, typeDef)} + ` + ) + .join('\n')} + \`\`\` + ` + : 'Only built-in JS functions are available' + } - If the function you need is not available, throw a - \`new FunctionUnavailableError('')\` - which is available in \`globalThis\`. - _Important:_ \`fetch\` is not available, do not make up functions or imports that are not provided. - `, - }, - { - role: 'user', - content: prompt, - }, - ] + If the function you need is not available, throw a + \`new FunctionUnavailableError('')\` + which is available in \`globalThis\`. + _Important:_ \`fetch\` is not available, do not make up functions or imports that are not provided. + `, + } + + const messagesExcludingSystem = messages.filter((message) => message.role !== 'system') const chatCompletion = await openai.chat.completions.create({ model: 'gpt-4o-2024-05-13', - messages, + messages: [systemMessage, ...messagesExcludingSystem], max_tokens: 1024, temperature: 0, }) @@ -146,18 +145,6 @@ export async function interpretPrompt( } } -export function asSingleOutput(exports: Record) { - const outputs = Object.values(exports) - - if (outputs.length > 1) { - throw new Error('Expected a single export') - } - - const [output] = outputs - - return output -} - function extractMarkdownCodeBlock(markdown: string, language?: string) { const mdTree = fromMarkdown(markdown) diff --git a/packages/ai-sandbox/test/util.ts b/packages/ai-sandbox/test/util.ts new file mode 100644 index 00000000000..913a9dea9ae --- /dev/null +++ b/packages/ai-sandbox/test/util.ts @@ -0,0 +1,45 @@ +import chalk from 'chalk' + +/** + * Prints the provided metadata along with any assertion errors. + * Works both synchronously and asynchronously. + * + * Useful for providing extra context for failed tests. + */ +export function withMetadata>( + metadata: Record, + fn: () => T +): T { + /** + * Prepends metadata to an Error's stack trace. + */ + function modifyError(err: unknown) { + if (err instanceof Error && err.stack) { + const formattedMetadata = Object.entries(metadata).map(([key, value]) => { + const formattedValue = + typeof value === 'string' ? value : JSON.stringify(value, undefined, 2) + return `${chalk.bold.dim(key)}:\n\n${chalk.green.dim(formattedValue)}` + }) + err.stack = `${formattedMetadata.join('\n\n')}\n\n${err.stack}` + } + + return err + } + + // Execute the function and handle both + // synchronous or asynchronous scenarios + try { + const maybePromise = fn() + + if (maybePromise instanceof Promise) { + return maybePromise.catch((err) => { + // Re-throw the error + throw modifyError(err) + }) as T + } + return maybePromise + } catch (err) { + // Re-throw the error + throw modifyError(err) + } +}