feat(ai-sandbox): multi-message inputs + more tests

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
     }
   }
-}
+}

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)/)',
   ],
 }

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",

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<string[]>',
+          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',
+      ])
+    })
   })
 })

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<string, FunctionBinding> = {}
 ) {
-  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<string>\`)
-            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<string>\`)
+          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('<description of what was unavailable>')\`
-        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('<description of what was unavailable>')\`
+      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<string, any>) {
-  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)
 

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<T extends void | Promise<void>>(
+  metadata: Record<string, any>,
+  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)
+  }
+}