chore: cleaning up ai streaming hook (#6086)

packages/react-ai/src/components/AIConversation/context/ResponseComponentsContext.tsx

@@ -51,7 +51,12 @@ export const convertResponseComponentsToToolConfiguration = (
     const { props } = responseComponents[toolName];
     const requiredProps: string[] = [];
     Object.keys(props).forEach((propName) => {
-      if (props[propName].required) requiredProps.push(propName);
+      if (props[propName].required) {
+        requiredProps.push(propName);
+        // The inputSchema for a tool needs to not
+        // have `required` in the properties
+        props[propName].required = undefined;
+      }
     });
     tools[toolName] = {
       description: responseComponents[toolName].description,

packages/react-ai/src/components/AIConversation/context/__tests__/ResponseComponentsContext.spec.tsx

@@ -41,7 +41,6 @@ describe('convertResponseComponentsToToolConfiguration', () => {
                   type: 'string',
                   description:
                     'The response you want to render in the component.',
-                  required: true,
                 },
                 foobar: {
                   type: 'number',

packages/react-ai/src/components/AIConversation/views/default/MessageList.tsx

@@ -29,7 +29,7 @@ const PlaceholderMessage = ({ role }: { role: string }) => {
       )}
     >
       <View className={ComponentClassName.AIConversationMessageAvatar}>
-        <Avatar> </Avatar>
+        <Avatar />
       </View>
       <View className={ComponentClassName.AIConversationMessageBody}>
         <Placeholder width="25%" />

packages/react-ai/src/hooks/useAIConversation.tsx

@@ -33,6 +33,16 @@ interface AIConversationState {
   conversation?: Conversation;
 }
 
+// The "states" the hook can be in
+// initial: default, nothing happened yet
+// loading: the hook has either hit AppSync to create or get a conversation
+// initialized: the hook has successfully gotten a conversation and is ready to rock
+const INITIALIZE_REF = ['initial', 'initialLoading', 'initialized'] as const;
+
+function hasStarted(state: (typeof INITIALIZE_REF)[number]) {
+  return ['initialLoading', 'initialized'].includes(state);
+}
+
 export type UseAIConversationHook<T extends string> = (
   routeName: T,
   input?: UseAIConversationInput
@@ -63,7 +73,7 @@ export function createUseAIConversation<
     // Using this hook without an existing conversation id means
     // it will create a new conversation when it is executed
     // we don't want to create 2 conversations
-    const initRef = React.useRef(false);
+    const initRef = React.useRef<(typeof INITIALIZE_REF)[number]>('initial');
 
     const [dataState, setDataState] = React.useState<
       DataClientState<AIConversationState>
@@ -76,30 +86,11 @@ export function createUseAIConversation<
     const { id, onInitialize, onMessage } = input;
 
     React.useEffect(() => {
-      // We don't want to run the effect multiple times
-      // because that could create multiple conversation records
-      if (initRef.current) return;
-      initRef.current = true;
       async function initialize() {
-        // no client route would mean that the user
-        // is not using TypeScript and entered the
-        // route name wrong, or there is a mismatch
-        // between the gen2 schema definition and
-        // whats in amplify_outputs
-        if (!clientRoute) {
-          setDataState({
-            ...ERROR_STATE,
-            data: { messages: [] },
-            messages: [
-              {
-                message: 'Conversation route does not exist',
-                errorInfo: null,
-                errorType: '',
-              },
-            ],
-          });
-          return;
-        }
+        // We don't want to run the effect multiple times
+        // because that could create multiple conversation records
+        if (hasStarted(initRef.current)) return;
+        initRef.current = 'initialLoading';
 
         // Only show component loading state if we are
         // actually loading messages
@@ -135,13 +126,33 @@ export function createUseAIConversation<
               data: { conversation, messages: [] },
             });
           }
+          initRef.current = 'initialized';
         }
       }
 
+      // this is a runtime guard to make catch an error if
+      // the route name wrong, or there is a mismatch
+      // between the gen2 schema definition and
+      // whats in amplify_outputs
+      if (!clientRoute) {
+        setDataState({
+          ...ERROR_STATE,
+          data: { messages: [] },
+          messages: [
+            {
+              message: 'Conversation route does not exist',
+              errorInfo: null,
+              errorType: '',
+            },
+          ],
+        });
+        return;
+      }
       initialize();
 
       return () => {
         contentBlocksRef.current = undefined;
+        if (hasStarted(initRef.current)) return;
         setDataState({
           ...INITIAL_STATE,
           data: { messages: [], conversation: undefined },
@@ -166,14 +177,9 @@ export function createUseAIConversation<
             // this is sent after the last content chunk, verify this matches the
             // previous contentBlockDeltaIndex
             contentBlockDoneAtIndex,
-            // this is the text of the content block
-            text,
-            // this is a toolUse block, will always come in a single event
-            // toolUse,
             // this is the final event of the conversation turn
             stopReason,
             conversationId,
-            // associatedUserMessageId,
             id,
           } = event;
 
@@ -212,43 +218,25 @@ export function createUseAIConversation<
           }
 
           // no ref means its the first event for the message stream
+          // so lets create the contentBlocks ref or else we will
+          // add the incoming event to the right content content block
           if (!contentBlocksRef.current) {
             contentBlocksRef.current = [[event]];
-
-            setDataState((prev) => {
-              const message: ConversationMessage = {
-                id,
-                conversationId,
-                // TODO: use better logic here
-                content: [{ text: text ?? '' }],
-                createdAt: new Date().toISOString(),
-                role: 'assistant',
-                isLoading: true,
-              };
-              return {
-                ...prev,
-                data: {
-                  ...prev.data,
-                  messages: [...prev.data.messages.slice(0, -1), message],
-                },
-              };
-            });
-            return;
-          }
-
-          // place the incoming event in the right content block
-          // and order. message content is an array so a single message
-          // can have multiple content blocks, and each content block
-          // can have multiple events/chunks
-          const currentBlock = contentBlocksRef.current[contentBlockIndex];
-          if (!currentBlock) {
-            contentBlocksRef.current[contentBlockIndex] = [event];
           } else {
-            contentBlocksRef.current[contentBlockIndex] = [
-              ...currentBlock.slice(0, contentBlockDeltaIndex),
-              event,
-              ...currentBlock.slice(contentBlockDeltaIndex),
-            ];
+            // place the incoming event in the right content block
+            // and order. message content is an array so a single message
+            // can have multiple content blocks, and each content block
+            // can have multiple events/chunks
+            const currentBlock = contentBlocksRef.current[contentBlockIndex];
+            if (!currentBlock) {
+              contentBlocksRef.current[contentBlockIndex] = [event];
+            } else {
+              contentBlocksRef.current[contentBlockIndex] = [
+                ...currentBlock.slice(0, contentBlockDeltaIndex),
+                event,
+                ...currentBlock.slice(contentBlockDeltaIndex),
+              ];
+            }
           }
 
           setDataState((prev) => {
@@ -272,9 +260,6 @@ export function createUseAIConversation<
           });
         },
         error: (error) => {
-          error.errors.map((e) => {
-            return e.message;
-          });
           setDataState((prev) => {
             return {
               ...prev,
@@ -325,6 +310,18 @@ export function createUseAIConversation<
             },
           }));
           conversation.sendMessage(input);
+        } else {
+          setDataState((prev) => ({
+            ...prev,
+            ...ERROR_STATE,
+            messages: [
+              {
+                message: 'No conversation found',
+                errorInfo: null,
+                errorType: '',
+              },
+            ],
+          }));
         }
       },
       [conversation]

packages/react-ai/src/types.ts

@@ -5,6 +5,12 @@ export type Conversation = NonNullable<
   Awaited<ReturnType<ConversationRoute['create']>>['data']
 >;
 
+// the JS client looks like this:
+// client.conversations.[name].onStreamEvent({
+//   next: (event) => {},
+//   error: (error) => {}
+// })
+// This type gets the 'event' ^
 export type ConversationStreamEvent = Parameters<
   Parameters<Conversation['onStreamEvent']>[0]['next']
 >[0];