Merge pull request #2 from unstubbable/structured-input

Model image selection as assistant state

Author: unstubbable

src/app/ai-state.ts

@@ -0,0 +1,35 @@
+export type AIStateItem =
+  | {
+      readonly role: 'user' | 'assistant' | 'system';
+      readonly content: string;
+    }
+  | {
+      readonly role: 'function';
+      readonly content: string;
+      readonly name: string;
+    };
+
+export type UserInput =
+  | {
+      readonly action: 'message';
+      readonly content: string;
+    }
+  | {
+      readonly action: 'select-image';
+      readonly url: string;
+    };
+
+export function fromUserInput(userInput: UserInput): AIStateItem {
+  switch (userInput.action) {
+    case `select-image`:
+      return {
+        role: `assistant`,
+        content: `[user wants to know more about the image ${userInput.url}. keep it short.]`,
+      };
+    case `message`:
+      return {
+        role: `user`,
+        content: userInput.content,
+      };
+  }
+}

src/app/ai.tsx

@@ -1,18 +1,8 @@
 import {createAI} from 'ai/rsc';
 import type * as React from 'react';
+import type {AIStateItem} from './ai-state.js';
 import {submitUserMessage} from './submit-user-message.js';
 
-export type AIStateItem =
-  | {
-      readonly role: 'user' | 'assistant' | 'system';
-      readonly content: string;
-    }
-  | {
-      readonly role: 'function';
-      readonly content: string;
-      readonly name: string;
-    };
-
 export interface UIStateItem {
   readonly id: number;
   readonly role: 'user' | 'assistant' | 'error';

src/app/chat.tsx

@@ -57,7 +57,10 @@ export function Chat({children}: React.PropsWithChildren): React.ReactNode {
       ]);
 
       try {
-        const message = await submitUserMessage(userInput);
+        const message = await submitUserMessage({
+          action: `message`,
+          content: userInput,
+        });
 
         setMessages((prevMessages) => [...prevMessages, message]);
       } catch (error) {

src/app/image-selector.tsx

@@ -27,9 +27,7 @@ export function ImageSelector({
     document.body.scrollIntoView({block: `end`, behavior: `smooth`});
 
     try {
-      const message = await submitUserMessage(
-        `Tell me more about image ${url}, keep it short.`,
-      );
+      const message = await submitUserMessage({action: `select-image`, url});
 
       setMessages((prevMessages) => [
         ...prevMessages.filter(({id}) => id !== optimisticMessageId),

src/app/submit-user-message.tsx

@@ -4,6 +4,7 @@ import {getMutableAIState, render} from 'ai/rsc';
 import {OpenAI} from 'openai';
 import * as React from 'react';
 import {z} from 'zod';
+import {type UserInput, fromUserInput} from './ai-state.js';
 import type {AI, UIStateItem} from './ai.js';
 import {imageSearchParams, searchImages} from './google-image-search.js';
 import {ImageSelector} from './image-selector.js';
@@ -16,11 +17,11 @@ const openai = new OpenAI({apiKey: process.env.OPENAI_API_KEY});
 
 // eslint-disable-next-line @typescript-eslint/require-await
 export async function submitUserMessage(
-  userInput: string,
+  userInput: UserInput,
 ): Promise<UIStateItem> {
   const aiState = getMutableAIState<typeof AI>();
 
-  aiState.update([...aiState.get(), {role: `user`, content: userInput}]);
+  aiState.update([...aiState.get(), fromUserInput(userInput)]);
 
   let lastTextContent: string | undefined;
 
@@ -41,6 +42,8 @@ export async function submitUserMessage(
         Never ask the user whether they want to see images of the discussed subject, always show them unprompted.
 
         Before showing images of a certain artist it might make sense to introduce them to the user first, with a couple of words.
+
+        The user can also select an image if they want to know more about it.
         `,
       },
       ...aiState.get(),
@@ -62,7 +65,7 @@ export async function submitUserMessage(
 
         If the images to search for are distinctly different from each other (e.g. two different, titled paintings by the same artist), split it the search up into multiple search parameter sets.
 
-        Select three images overall, unless the user asks for more or less.
+        Select three images overall, unless the user asks for more or less. When searching for a specific artwork, select only one image.
 
         Never call the function more than once in a message, instead use multiple searches within the same function call.
         `,