feat(ai-conversation): add allowAttachments prop (#5802)

Co-authored-by: Scott Rees <[email protected]>
aws-amplify · Sep 30, 2024 · 8b4a28b · 8b4a28b
1 parent 22e285f
commit 8b4a28b
Show file tree

Hide file tree

Showing 26 changed files with 217 additions and 99 deletions.
diff --git a/.changeset/seven-bats-travel.md b/.changeset/seven-bats-travel.md
@@ -0,0 +1,15 @@
+---
+"@aws-amplify/ui-react-ai": minor
+---
+
+feat(ai-conversation): add allowAttachments prop
+
+BREAKING - This is a breaking change to an experimental API. Previously, the AIConversation component always allowed attachments. Now you will need to provide the `allowAttachments` prop to get the same behavior. The reason for this change is that attachments can quickly cost a lot based on the token use and we didn't want the default behavior to have that. 
+
+```jsx
+<AIConversation
+  allowAttachments
+  messages={messages}
+  handleSendMessage={handleSendMessage}
+/>
+```
diff --git a/examples/next/pages/ui/components/ai/ai-conversation/index.page.tsx b/examples/next/pages/ui/components/ai/ai-conversation/index.page.tsx
@@ -36,6 +36,7 @@ function Chat() {
         messages={messages}
         handleSendMessage={sendMessage}
         isLoading={isLoading}
+        allowAttachments
         suggestedPrompts={[
           {
             inputText: 'hello',

diff --git a/packages/react-ai/src/components/AIConversation/AIConversation.tsx b/packages/react-ai/src/components/AIConversation/AIConversation.tsx
@@ -30,6 +30,7 @@ function AIConversationBase({
   variant,
   isLoading,
   displayText,
+  allowAttachments,
 }: AIConversationBaseProps): JSX.Element {
   const icons = useIcons('aiConversation');
   const defaultAvatars: Avatars = {
@@ -72,6 +73,7 @@ function AIConversationBase({
       ...avatars,
     },
     isLoading,
+    allowAttachments,
   };
 
   return (

diff --git a/packages/react-ai/src/components/AIConversation/context/AttachmentContext.tsx b/packages/react-ai/src/components/AIConversation/context/AttachmentContext.tsx
@@ -0,0 +1,17 @@
+import * as React from 'react';
+
+export const AttachmentContext = React.createContext(false);
+
+export const AttachmentProvider = ({
+  children,
+  allowAttachments,
+}: {
+  children?: React.ReactNode;
+  allowAttachments?: boolean;
+}): JSX.Element => {
+  return (
+    <AttachmentContext.Provider value={allowAttachments ?? false}>
+      {children}
+    </AttachmentContext.Provider>
+  );
+};
diff --git a/packages/react-ai/src/components/AIConversation/context/ControlsContext.tsx b/packages/react-ai/src/components/AIConversation/context/ControlsContext.tsx
@@ -7,6 +7,7 @@ export interface ControlsContextProps {
   Form?: React.ComponentType<
     {
       handleSubmit: (e: React.FormEvent<HTMLFormElement>) => void;
+      allowAttachments?: boolean;
     } & Required<ConversationInputContext>
   >;
   MessageList?: React.ComponentType<{ messages: ConversationMessage[] }>;

diff --git a/packages/react-ai/src/components/AIConversation/createAIConversation.tsx b/packages/react-ai/src/components/AIConversation/createAIConversation.tsx
@@ -30,6 +30,7 @@ export function createAIConversation(input: AIConversationInput = {}): {
     variant,
     controls,
     displayText,
+    allowAttachments,
   } = input;
 
   const Provider = createProvider({
@@ -40,6 +41,7 @@ export function createAIConversation(input: AIConversationInput = {}): {
     variant,
     controls,
     displayText,
+    allowAttachments,
   });
 
   function AIConversation(props: AIConversationProps): JSX.Element {

diff --git a/packages/react-ai/src/components/AIConversation/createProvider.tsx b/packages/react-ai/src/components/AIConversation/createProvider.tsx
@@ -17,6 +17,7 @@ import {
   ResponseComponentsProvider,
   SendMessageContextProvider,
 } from './context';
+import { AttachmentProvider } from './context/AttachmentContext';
 
 export default function createProvider({
   elements,
@@ -26,6 +27,7 @@ export default function createProvider({
   variant,
   controls,
   displayText,
+  allowAttachments,
 }: Pick<
   AIConversationInput,
   | 'elements'
@@ -35,6 +37,7 @@ export default function createProvider({
   | 'variant'
   | 'controls'
   | 'displayText'
+  | 'allowAttachments'
 >) {
   return function Provider({
     children,
@@ -57,25 +60,27 @@ export default function createProvider({
         <ControlsProvider controls={controls}>
           <SuggestedPromptProvider suggestedPrompts={suggestedPrompts}>
             <ResponseComponentsProvider responseComponents={responseComponents}>
-              <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 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>

diff --git a/packages/react-ai/src/components/AIConversation/types.ts b/packages/react-ai/src/components/AIConversation/types.ts
@@ -31,6 +31,7 @@ export interface AIConversationInput {
   responseComponents?: ResponseComponents;
   variant?: MessageVariant;
   controls?: ControlsContextProps;
+  allowAttachments?: boolean;
 }
 
 export interface AIConversationProps {

diff --git a/packages/react-ai/src/components/AIConversation/views/Controls/AttachFileControl.tsx b/packages/react-ai/src/components/AIConversation/views/Controls/AttachFileControl.tsx
@@ -3,6 +3,7 @@ import React from 'react';
 import { withBaseElementProps } from '@aws-amplify/ui-react-core/elements';
 import { ConversationInputContext } from '../../context';
 import { AIConversationElements } from '../../context/elements';
+import { useDropZone } from '@aws-amplify/ui-react-core';
 
 const { Button, Icon, View } = AIConversationElements;
 
@@ -32,6 +33,17 @@ const AttachFileButton = withBaseElementProps(Button, {
 export const AttachFileControl: AttachFileControl = () => {
   const hiddenInput = React.useRef<HTMLInputElement>(null);
   const { setInput } = React.useContext(ConversationInputContext);
+  const { dragState, ...dropHandlers } = useDropZone({
+    acceptedFileTypes: ['.jpeg'],
+    onDropComplete: ({ acceptedFiles }) => {
+      if (acceptedFiles && acceptedFiles?.length > 0 && setInput) {
+        setInput((prevInput) => ({
+          ...prevInput,
+          files: [...(prevInput?.files ?? []), ...acceptedFiles],
+        }));
+      }
+    },
+  });
 
   function handleButtonClick() {
     if (hiddenInput.current) {
@@ -53,7 +65,7 @@ export const AttachFileControl: AttachFileControl = () => {
   }
 
   return (
-    <AttachFileContainer>
+    <AttachFileContainer {...dropHandlers}>
       <AttachFileButton onClick={handleButtonClick}>
         <AttachFileIcon />
       </AttachFileButton>

diff --git a/packages/react-ai/src/components/AIConversation/views/Controls/FieldControl.tsx b/packages/react-ai/src/components/AIConversation/views/Controls/FieldControl.tsx
@@ -15,6 +15,7 @@ import {
 import { ControlsContext } from '../../context/ControlsContext';
 import { getImageTypeFromMimeType } from '../../utils';
 import { LoadingContext } from '../../context/LoadingContext';
+import { AttachmentContext } from '../../context/AttachmentContext';
 
 const {
   Button,
@@ -148,6 +149,7 @@ const InputContainer = withBaseElementProps(View, {
 export const FieldControl: FieldControl = () => {
   const { input, setInput } = React.useContext(ConversationInputContext);
   const handleSendMessage = React.useContext(SendMessageContext);
+  const allowAttachments = React.useContext(AttachmentContext);
   const ref = React.useRef<HTMLFormElement | null>(null);
   const responseComponents = React.useContext(ResponseComponentsContext);
   const controls = React.useContext(ControlsContext);
@@ -212,6 +214,7 @@ export const FieldControl: FieldControl = () => {
         handleSubmit={handleSubmit}
         input={input!}
         setInput={setInput!}
+        allowAttachments={allowAttachments}
       />
     );
   }
@@ -223,7 +226,7 @@ export const FieldControl: FieldControl = () => {
       method="post"
       ref={ref}
     >
-      <AttachFileControl />
+      {allowAttachments ? <AttachFileControl /> : null}
       <InputContainer>
         <VisuallyHidden>
           <Label />

diff --git a/...ges/react-ai/src/components/AIConversation/views/Controls/__tests__/FieldControl.spec.tsx b/...ges/react-ai/src/components/AIConversation/views/Controls/__tests__/FieldControl.spec.tsx
@@ -5,10 +5,15 @@ import { FieldControl } from '../FieldControl';
 import { ConversationInputContextProvider } from '../../../context/ConversationInputContext';
 import userEvent from '@testing-library/user-event';
 import { SendMessageContextProvider } from '../../../context/SendMessageContext';
+import { AttachmentProvider } from '../../../context/AttachmentContext';
 
 describe('FieldControl', () => {
   it('renders a FieldControl component with the correct elements', () => {
-    const result = render(<FieldControl />);
+    const result = render(
+      <AttachmentProvider allowAttachments>
+        <FieldControl />
+      </AttachmentProvider>
+    );
     expect(result.container).toBeDefined();
 
     const form = screen.findByRole('form');
@@ -23,7 +28,11 @@ describe('FieldControl', () => {
   });
 
   it('renders FieldControl with the correct accessibility roles', () => {
-    render(<FieldControl />);
+    render(
+      <AttachmentProvider allowAttachments>
+        <FieldControl />
+      </AttachmentProvider>
+    );
 
     const actionButtons = screen.getAllByRole('button');
     const sendButton = actionButtons[1];