fix(ai): fix renders for empty toolUse messages (#5799)

* fix(ai): fix renders for empty toolUse messages

* Create ten-beds-jump.md

* memoize response components value

* update component tool prefix

Author: thaddmt

.changeset/ten-beds-jump.md

@@ -0,0 +1,5 @@
+---
+"@aws-amplify/ui-react-ai": patch
+---
+
+fix(ai): fix renders for empty toolUse messages

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

@@ -4,18 +4,36 @@ import { ToolConfiguration } from '../../../types';
 
 type ResponseComponentsContextProps = ResponseComponents | undefined;
 
+export const RESPONSE_COMPONENT_PREFIX = 'AMPLIFY_UI_';
+
 export const ResponseComponentsContext =
   React.createContext<ResponseComponentsContextProps>(undefined);
 
+const prependResponseComponents = (responseComponents?: ResponseComponents) => {
+  if (!responseComponents) return responseComponents;
+  return Object.keys(responseComponents).reduce(
+    (prev, key) => (
+      (prev[`${RESPONSE_COMPONENT_PREFIX}${key}`] = responseComponents[key]),
+      prev
+    ),
+    {} as ResponseComponents
+  );
+};
+
 export const ResponseComponentsProvider = ({
   children,
   responseComponents,
 }: {
   children?: React.ReactNode;
   responseComponents?: ResponseComponents;
 }): JSX.Element => {
+  const _responseComponents = React.useMemo(
+    () => prependResponseComponents(responseComponents),
+    [responseComponents]
+  );
+
   return (
-    <ResponseComponentsContext.Provider value={responseComponents}>
+    <ResponseComponentsContext.Provider value={_responseComponents}>
       {children}
     </ResponseComponentsContext.Provider>
   );

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

@@ -29,7 +29,10 @@ export {
   ControlsProvider,
 } from './ControlsContext';
 export { LoadingContextProvider } from './LoadingContext';
-export { ResponseComponentsProvider } from './ResponseComponentsContext';
+export {
+  ResponseComponentsProvider,
+  RESPONSE_COMPONENT_PREFIX,
+} from './ResponseComponentsContext';
 export { SendMessageContextProvider } from './SendMessageContext';
 
 export * from './elements';

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

@@ -12,7 +12,10 @@ import { convertBufferToBase64 } from '../../utils';
 import { ActionsBarControl } from './ActionsBarControl';
 import { AvatarControl } from './AvatarControl';
 import { ConversationMessage } from '../../../../types';
-import { ResponseComponentsContext } from '../../context/ResponseComponentsContext';
+import {
+  RESPONSE_COMPONENT_PREFIX,
+  ResponseComponentsContext,
+} from '../../context/ResponseComponentsContext';
 import { ControlsContext } from '../../context/ControlsContext';
 
 const { Image, Span, Text, View } = AIConversationElements;
@@ -83,7 +86,11 @@ export const MessageControl: MessageControl = ({ message }) => {
         } else if (content.toolUse) {
           // For now tool use is limited to custom response components
           const { name, input } = content.toolUse;
-          if (!responseComponents || !name) {
+          if (
+            !responseComponents ||
+            !name ||
+            !name.startsWith(RESPONSE_COMPONENT_PREFIX)
+          ) {
             return;
           } else {
             const response = responseComponents[name];
@@ -206,9 +213,19 @@ export const MessagesControl: MessagesControl = ({ renderMessage }) => {
     return <controls.MessageList messages={messages!} />;
   }
 
+  const messagesWithRenderableContent =
+    messages?.filter((message) =>
+      message.content.some(
+        (content) =>
+          content.image ??
+          content.text ??
+          content.toolUse?.name.startsWith(RESPONSE_COMPONENT_PREFIX)
+      )
+    ) ?? [];
+
   return (
     <Layout>
-      {messages?.map((message, index) => {
+      {messagesWithRenderableContent?.map((message, index) => {
         return renderMessage ? (
           renderMessage(message)
         ) : (

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

@@ -54,7 +54,38 @@ const AIResponseComponentMessage: ConversationMessage = {
   content: [
     {
       toolUse: {
-        name: 'annoyingComponent',
+        name: 'AMPLIFY_UI_annoyingComponent',
+        input: { text: 'ahoy matey' },
+        toolUseId: 'tooluseID',
+      },
+    },
+  ],
+  role: 'assistant',
+  createdAt: new Date(2023, 4, 21, 15, 25).toDateString(),
+};
+const ToolUseMessage: ConversationMessage = {
+  conversationId: 'foobar',
+  id: '3',
+  content: [
+    {
+      toolUse: {
+        name: '"generateRecipe"',
+        input: { text: 'ahoy matey' },
+        toolUseId: 'tooluseID',
+      },
+    },
+  ],
+  role: 'assistant',
+  createdAt: new Date(2023, 4, 21, 15, 25).toDateString(),
+};
+const TextAndToolUseMessage: ConversationMessage = {
+  conversationId: 'foobar',
+  id: '3',
+  content: [
+    { text: 'hey what up' },
+    {
+      toolUse: {
+        name: '"generateRecipe"',
         input: { text: 'ahoy matey' },
         toolUseId: 'tooluseID',
       },
@@ -345,4 +376,15 @@ describe('MessageControl', () => {
     const message = screen.getByText('argh matey! ahoy matey');
     expect(message).toBeInTheDocument();
   });
+
+  it('renders text when sent with a tooluse content', () => {
+    render(<MessageControl message={TextAndToolUseMessage} />);
+    const message = screen.getByText('hey what up');
+    expect(message).toBeInTheDocument();
+  });
+
+  it('renders nothing when only a toolUse block is sent', () => {
+    const { container } = render(<MessageControl message={ToolUseMessage} />);
+    expect(container.firstChild).toBeEmptyDOMElement();
+  });
 });

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

@@ -4,6 +4,7 @@ import { MessageControl } from '../Controls/MessagesControl';
 import {
   AvatarsContext,
   MessageVariantContext,
+  RESPONSE_COMPONENT_PREFIX,
   RoleContext,
   useConversationDisplayText,
 } from '../../context';
@@ -99,9 +100,19 @@ export const MessageList: ControlsContextProps['MessageList'] = ({
   messages,
 }) => {
   const isLoading = React.useContext(LoadingContext);
+  const messagesWithRenderableContent =
+    messages?.filter((message) =>
+      message.content.some(
+        (content) =>
+          content.image ??
+          content.text ??
+          content.toolUse?.name.startsWith(RESPONSE_COMPONENT_PREFIX)
+      )
+    ) ?? [];
+
   return (
     <View className={ComponentClassName.AIConversationMessageList}>
-      {messages.map((message, i) => (
+      {messagesWithRenderableContent.map((message, i) => (
         <Message key={`message-${i}`} message={message} />
       ))}
       {isLoading ? <LoadingMessage /> : null}