feat(ai): add fallback response component (aws-amplify#5968)

Co-authored-by: dindjarinjs <[email protected]>

.changeset/purple-bobcats-refuse.md

@@ -0,0 +1,12 @@
+---
+"@aws-amplify/ui-react-ai": minor
+---
+
+feat(ai): add fallback response component
+
+```tsx
+<AIConversation
+  FallBackResponseComponent={(props) => <>{JSON.stringify(props)}</>}
+  //...
+/>
+```

examples/next/pages/ui/components/ai/ai-conversation-response-components/amplify_outputs.js

@@ -0,0 +1,2 @@
+import amplifyOutputs from '@environments/ai/gen2/amplify_outputs';
+export default amplifyOutputs;

examples/next/pages/ui/components/ai/ai-conversation-response-components/fallback.page.tsx

@@ -0,0 +1,51 @@
+import { AIConversation, ConversationMessage } from '@aws-amplify/ui-react-ai';
+import '@aws-amplify/ui-react/styles.css';
+
+const messages: ConversationMessage[] = [
+  {
+    role: 'user',
+    content: [
+      {
+        text: 'hello',
+      },
+    ],
+    conversationId: '1',
+    id: '2',
+    createdAt: new Date(2023, 4, 21, 15, 24).toDateString(),
+  },
+  {
+    role: 'assistant',
+    content: [
+      {
+        toolUse: {
+          name: 'AMPLIFY_UI_foobar',
+          input: { foo: 'bar' },
+          toolUseId: '1234',
+        },
+      },
+    ],
+    conversationId: '1',
+    id: '2',
+    createdAt: new Date(2023, 4, 21, 15, 24).toDateString(),
+  },
+];
+
+// Note: because response components are sent in the message
+// there could be cases where the AIConversation component
+// gets rendered with an existing conversation and does not
+// have the React component needed to render the response
+// component. For example in the Amplify console, we don't
+// have customers' React code running in our console.
+// Because of this, this example page isn't actually
+// using a live conversation route.
+export default function Example() {
+  return (
+    <AIConversation
+      messages={messages}
+      handleSendMessage={() => {}}
+      FallbackResponseComponent={(props) => {
+        return <div>{JSON.stringify(props)}</div>;
+      }}
+    />
+  );
+}

examples/next/pages/ui/components/ai/ai-conversation-response-components/index.page.tsx

@@ -0,0 +1,58 @@
+import { Amplify } from 'aws-amplify';
+import { createAIHooks, AIConversation } from '@aws-amplify/ui-react-ai';
+import { generateClient } from 'aws-amplify/api';
+import '@aws-amplify/ui-react/styles.css';
+
+import outputs from './amplify_outputs';
+import type { Schema } from '@environments/ai/gen2/amplify/data/resource';
+import { Authenticator, Card } from '@aws-amplify/ui-react';
+
+const client = generateClient<Schema>({ authMode: 'userPool' });
+const { useAIConversation } = createAIHooks(client);
+
+Amplify.configure(outputs);
+
+function Chat() {
+  const [
+    {
+      data: { messages },
+      isLoading,
+    },
+    sendMessage,
+  ] = useAIConversation('pirateChat');
+
+  return (
+    <Card variation="outlined" width="50%" height="300px" margin="0 auto">
+      <AIConversation
+        messages={messages}
+        handleSendMessage={sendMessage}
+        isLoading={isLoading}
+        allowAttachments
+        responseComponents={{
+          WeatherCard: {
+            description:
+              'Used to display the weather of a given city to the user',
+            component: ({ city }) => {
+              return <Card>{city}</Card>;
+            },
+            props: {
+              city: {
+                type: 'string',
+                required: true,
+              },
+            },
+          },
+        }}
+        variant="bubble"
+      />
+    </Card>
+  );
+}
+
+export default function Example() {
+  return (
+    <Authenticator>
+      <Chat />
+    </Authenticator>
+  );
+}

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

@@ -17,6 +17,7 @@ import {
   ResponseComponentsProvider,
   SendMessageContextProvider,
   WelcomeMessageProvider,
+  FallbackComponentProvider,
   MessageRendererProvider,
 } from './context';
 import { AttachmentProvider } from './context/AttachmentContext';
@@ -42,6 +43,7 @@ export const AIConversationProvider = ({
   suggestedPrompts,
   variant,
   welcomeMessage,
+  FallbackResponseComponent,
   messageRenderer,
 }: AIConversationProviderProps): React.JSX.Element => {
   const _displayText = {
@@ -53,33 +55,37 @@ export const AIConversationProvider = ({
       <ControlsProvider controls={controls}>
         <SuggestedPromptProvider suggestedPrompts={suggestedPrompts}>
           <WelcomeMessageProvider welcomeMessage={welcomeMessage}>
-            <MessageRendererProvider {...messageRenderer}>
-              <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>
-            </MessageRendererProvider>
+            <FallbackComponentProvider
+              FallbackComponent={FallbackResponseComponent}
+            >
+              <MessageRendererProvider {...messageRenderer}>
+                <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>
+              </MessageRendererProvider>
+            </FallbackComponentProvider>
           </WelcomeMessageProvider>
         </SuggestedPromptProvider>
       </ControlsProvider>

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

@@ -0,0 +1,20 @@
+import React from 'react';
+import { AIConversationInput } from '../types';
+
+export const FallbackComponentContext = React.createContext<
+  AIConversationInput['FallbackResponseComponent'] | undefined
+>(undefined);
+
+export const FallbackComponentProvider = ({
+  children,
+  FallbackComponent,
+}: {
+  children?: React.ReactNode;
+  FallbackComponent?: AIConversationInput['FallbackResponseComponent'];
+}): JSX.Element => {
+  return (
+    <FallbackComponentContext.Provider value={FallbackComponent}>
+      {children}
+    </FallbackComponentContext.Provider>
+  );
+};

packages/react-ai/src/components/AIConversation/context/index.ts

@@ -44,4 +44,8 @@ export {
   WelcomeMessageContext,
   WelcomeMessageProvider,
 } from './WelcomeMessageContext';
+export {
+  FallbackComponentContext,
+  FallbackComponentProvider,
+} from './FallbackComponentContext';
 export * from './elements';

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

@@ -25,6 +25,7 @@ export function createAIConversation(input: AIConversationInput = {}): {
     displayText,
     allowAttachments,
     messageRenderer,
+    FallbackResponseComponent,
   } = input;
 
   function AIConversation(props: AIConversationProps): JSX.Element {
@@ -43,6 +44,7 @@ export function createAIConversation(input: AIConversationInput = {}): {
       handleSendMessage,
       isLoading,
       messageRenderer,
+      FallbackResponseComponent,
     };
     return (
       <AIConversationProvider {...providerProps}>

packages/react-ai/src/components/AIConversation/types.ts

@@ -34,6 +34,7 @@ export interface AIConversationInput {
   suggestedPrompts?: SuggestedPrompt[];
   actions?: CustomAction[];
   responseComponents?: ResponseComponents;
+  FallbackResponseComponent?: React.ComponentType<any>;
   variant?: MessageVariant;
   controls?: ControlsContextProps;
   allowAttachments?: boolean;

packages/react-ai/src/components/AIConversation/views/Controls/MessagesControl.tsx

@@ -2,6 +2,7 @@ import React from 'react';
 import { withBaseElementProps } from '@aws-amplify/ui-react-core/elements';
 
 import {
+  FallbackComponentContext,
   MessageRendererContext,
   MessagesContext,
   MessageVariantContext,
@@ -67,21 +68,25 @@ const ToolContent = ({
 }: {
   toolUse: NonNullable<ConversationMessage['content'][number]['toolUse']>;
 }) => {
-  const responseComponents = React.useContext(ResponseComponentsContext);
+  const responseComponents = React.useContext(ResponseComponentsContext) ?? {};
+  const FallbackComponent = React.useContext(FallbackComponentContext);
 
   // For now tool use is limited to custom response components
   const { name, input } = toolUse;
 
-  if (
-    !responseComponents ||
-    !name ||
-    !name.startsWith(RESPONSE_COMPONENT_PREFIX)
-  ) {
+  if (!name || !name.startsWith(RESPONSE_COMPONENT_PREFIX)) {
     return;
   } else {
     const response = responseComponents[name];
-    const CustomComponent = response.component;
-    return <CustomComponent {...(input as object)} />;
+    if (response) {
+      const CustomComponent = response.component;
+      return <CustomComponent {...(input as object)} />;
+    }
+    // fallback if there is a UI component message but we don't have
+    // a React component that matches
+    if (FallbackComponent) {
+      return <FallbackComponent {...(input as object)} />;
+    }
   }
 };
 

packages/react-ai/src/components/AIConversation/views/Controls/__tests__/MessagesControl.spec.tsx

@@ -13,7 +13,11 @@ import { MessagesControl, MessageControl } from '../MessagesControl';
 import { convertBufferToBase64 } from '../../../utils';
 import { ConversationMessage } from '../../../../../types';
 import { ResponseComponentsProvider } from '../../../context/ResponseComponentsContext';
-import { MessageRendererProvider } from '../../../context';
+import {
+  FallbackComponentProvider,
+  MessageRendererProvider,
+} from '../../../context';
+import { View } from '@aws-amplify/ui-react';
 
 const AITextMessage: ConversationMessage = {
   conversationId: 'foobar',
@@ -340,6 +344,18 @@ describe('MessageControl', () => {
     expect(message).toBeInTheDocument();
   });
 
+  it('renders fallback response component if no response component is found', async () => {
+    render(
+      <FallbackComponentProvider
+        FallbackComponent={() => <View testId="fallback" />}
+      >
+        <MessageControl message={AIResponseComponentMessage} />
+      </FallbackComponentProvider>
+    );
+    const fallbackComponent = await screen.findByTestId('fallback');
+    expect(fallbackComponent).toBeInTheDocument();
+  });
+
   it('renders text when sent with a tooluse content', () => {
     render(<MessageControl message={TextAndToolUseMessage} />);
     const message = screen.getByText('hey what up');