fix(ai): no more unnecessary re-renders in AIConversation (#5883)

.changeset/metal-pumas-repeat.md

@@ -0,0 +1,5 @@
+---
+"@aws-amplify/ui-react-ai": patch
+---
+
+fix(ai): no more unnecessary re-renders in AIConversation

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

@@ -13,7 +13,7 @@ import { Form } from './views/default/Form';
 import { PromptList } from './views/default/PromptList';
 import { AutoHidablePromptControl } from './views/Controls';
 import { ComponentClassName } from '@aws-amplify/ui';
-import createProvider from './createProvider';
+import { AIConversationProvider } from './AIConversationProvider';
 
 interface AIConversationBaseProps
   extends AIConversationProps,
@@ -44,7 +44,14 @@ function AIConversationBase({
     },
   };
 
-  const Provider = createProvider({
+  const providerProps = {
+    messages,
+    handleSendMessage,
+    avatars: {
+      ...defaultAvatars,
+      ...avatars,
+    },
+    isLoading,
     elements: {
       Text: React.forwardRef<HTMLParagraphElement, TextProps>(
         function _Text(props, ref) {
@@ -64,28 +71,18 @@ function AIConversationBase({
     },
     displayText,
     allowAttachments,
-  });
-
-  const providerProps = {
-    messages,
-    handleSendMessage,
-    avatars: {
-      ...defaultAvatars,
-      ...avatars,
-    },
-    isLoading,
   };
 
   return (
-    <Provider {...providerProps}>
+    <AIConversationProvider {...providerProps}>
       <Flex className={ComponentClassName.AIConversation}>
         <ScrollView autoScroll="smooth" flex="1">
           <AutoHidablePromptControl />
           <MessagesControl />
         </ScrollView>
         <FieldControl />
       </Flex>
-    </Provider>
+    </AIConversationProvider>
   );
 }
 

packages/react-ai/src/components/AIConversation/AIConversationProvider.tsx

@@ -0,0 +1,78 @@
+import React from 'react';
+
+import { ElementsProvider } from '@aws-amplify/ui-react-core/elements';
+
+import { AIConversationInput, AIConversationProps } from './types';
+import { defaultAIConversationDisplayTextEn } from './displayText';
+import {
+  ConversationDisplayTextProvider,
+  SuggestedPromptProvider,
+  ConversationInputContextProvider,
+  AvatarsProvider,
+  ActionsProvider,
+  MessageVariantProvider,
+  MessagesProvider,
+  ControlsProvider,
+  LoadingContextProvider,
+  ResponseComponentsProvider,
+  SendMessageContextProvider,
+} from './context';
+import { AttachmentProvider } from './context/AttachmentContext';
+
+interface AIConversationProviderProps
+  extends AIConversationInput,
+    AIConversationProps {
+  children?: React.ReactNode;
+}
+
+export const AIConversationProvider = ({
+  elements,
+  actions,
+  suggestedPrompts,
+  responseComponents,
+  variant,
+  controls,
+  displayText,
+  allowAttachments,
+  messages,
+  handleSendMessage,
+  avatars,
+  isLoading,
+  children,
+}: AIConversationProviderProps): React.JSX.Element => {
+  const _displayText = {
+    ...defaultAIConversationDisplayTextEn,
+    ...displayText,
+  };
+  return (
+    <ElementsProvider elements={elements}>
+      <ControlsProvider controls={controls}>
+        <SuggestedPromptProvider suggestedPrompts={suggestedPrompts}>
+          <ResponseComponentsProvider responseComponents={responseComponents}>
+            <AttachmentProvider allowAttachments={allowAttachments}>
+              <ConversationDisplayTextProvider {..._displayText}>
+                <ConversationInputContextProvider>
+                  <SendMessageContextProvider
+                    handleSendMessage={handleSendMessage}
+                  >
+                    <AvatarsProvider avatars={avatars}>
+                      <ActionsProvider actions={actions}>
+                        <MessageVariantProvider variant={variant}>
+                          <MessagesProvider messages={messages}>
+                            <LoadingContextProvider isLoading={isLoading}>
+                              {children}
+                            </LoadingContextProvider>
+                          </MessagesProvider>
+                        </MessageVariantProvider>
+                      </ActionsProvider>
+                    </AvatarsProvider>
+                  </SendMessageContextProvider>
+                </ConversationInputContextProvider>
+              </ConversationDisplayTextProvider>
+            </AttachmentProvider>
+          </ResponseComponentsProvider>
+        </SuggestedPromptProvider>
+      </ControlsProvider>
+    </ElementsProvider>
+  );
+};

packages/react-ai/src/components/AIConversation/createAIConversation.tsx

@@ -14,7 +14,7 @@ import {
   MessagesControl,
   PromptControl,
 } from './views';
-import createProvider from './createProvider';
+import { AIConversationProvider } from './AIConversationProvider';
 
 /**
  * @experimental
@@ -33,28 +33,26 @@ export function createAIConversation(input: AIConversationInput = {}): {
     allowAttachments,
   } = input;
 
-  const Provider = createProvider({
-    elements,
-    actions,
-    suggestedPrompts,
-    responseComponents,
-    variant,
-    controls,
-    displayText,
-    allowAttachments,
-  });
-
   function AIConversation(props: AIConversationProps): JSX.Element {
     const { messages, avatars, handleSendMessage, isLoading } = props;
+    const providerProps = {
+      elements,
+      actions,
+      suggestedPrompts,
+      responseComponents,
+      variant,
+      controls,
+      displayText,
+      allowAttachments,
+      messages,
+      avatars,
+      handleSendMessage,
+      isLoading,
+    };
     return (
-      <Provider
-        messages={messages}
-        avatars={avatars}
-        handleSendMessage={handleSendMessage}
-        isLoading={isLoading}
-      >
+      <AIConversationProvider {...providerProps}>
         <Conversation />
-      </Provider>
+      </AIConversationProvider>
     );
   }
 
@@ -67,7 +65,7 @@ export function createAIConversation(input: AIConversationInput = {}): {
     SuggestedPrompts: PromptControl,
   };
 
-  AIConversation.Provider = Provider;
+  AIConversation.Provider = AIConversationProvider;
   AIConversation.Conversation = Conversation;
   AIConversation.Controls = Controls;