@aws-amplify/ui-react-ai 0.4.0 → 1.1.0

package/dist/index.js

@@ -5,9 +5,9 @@ Object.defineProperty(exports, '__esModule', { value: true });
 var React = require('react');
 var elements = require('@aws-amplify/ui-react-core/elements');
 var uiReactCore = require('@aws-amplify/ui-react-core');
+var ui = require('@aws-amplify/ui');
 var uiReact = require('@aws-amplify/ui-react');
 var internal = require('@aws-amplify/ui-react/internal');
-var ui = require('@aws-amplify/ui');
 
 function _interopNamespace(e) {
     if (e && e.__esModule) return e;
@@ -29,6 +29,11 @@ function _interopNamespace(e) {
 
 var React__namespace = /*#__PURE__*/_interopNamespace(React);
 
+const AIContextContext = React__namespace["default"].createContext(undefined);
+const AIContextProvider = ({ children, aiContext, }) => {
+    return (React__namespace["default"].createElement(AIContextContext.Provider, { value: aiContext }, children));
+};
+
 const ActionsContext = React__namespace["default"].createContext(undefined);
 const ActionsProvider = ({ children, actions, }) => {
     return (React__namespace["default"].createElement(ActionsContext.Provider, { value: actions }, children));
@@ -42,7 +47,8 @@ const AvatarsProvider = ({ children, avatars, }) => {
 const ConversationInputContext = React__namespace["default"].createContext({});
 const ConversationInputContextProvider = ({ children, }) => {
     const [input, setInput] = React__namespace["default"].useState();
-    const providerValue = React__namespace["default"].useMemo(() => ({ input, setInput }), [input, setInput]);
+    const [error, setError] = React__namespace["default"].useState();
+    const providerValue = React__namespace["default"].useMemo(() => ({ input, setInput, error, setError }), [input, setInput, error, setError]);
     return (React__namespace["default"].createElement(ConversationInputContext.Provider, { value: providerValue }, children));
 };
 
@@ -77,32 +83,67 @@ function formatDate(date) {
     return `${dateString} at ${timeString}`;
 }
 function arrayBufferToBase64(buffer) {
-    let binary = '';
-    const bytes = new Uint8Array(buffer);
-    const len = bytes.byteLength;
-    for (let i = 0; i < len; i++) {
-        binary += String.fromCharCode(bytes[i]);
-    }
-    return window.btoa(binary);
-}
-function convertBufferToBase64(buffer, format) {
-    let base64string = '';
     // Use node-based buffer if available
     // fall back on browser if not
     if (typeof Buffer !== 'undefined') {
-        base64string = Buffer.from(new Uint8Array(buffer)).toString('base64');
+        return Buffer.from(new Uint8Array(buffer)).toString('base64');
     }
     else {
-        base64string = arrayBufferToBase64(buffer);
+        let binary = '';
+        const bytes = new Uint8Array(buffer);
+        const len = bytes.byteLength;
+        for (let i = 0; i < len; i++) {
+            binary += String.fromCharCode(bytes[i]);
+        }
+        return window.btoa(binary);
     }
+}
+function convertBufferToBase64(buffer, format) {
+    const base64string = arrayBufferToBase64(buffer);
     return `data:image/${format};base64,${base64string}`;
 }
 function getImageTypeFromMimeType(mimeType) {
     return mimeType.split('/')[1];
 }
+async function attachmentsValidator({ files, maxAttachments, maxAttachmentSize, }) {
+    const acceptedFiles = [];
+    const rejectedFiles = [];
+    let hasMaxSizeError = false;
+    for (const file of files) {
+        const arrayBuffer = await file.arrayBuffer();
+        const base64 = arrayBufferToBase64(arrayBuffer);
+        if (base64.length < maxAttachmentSize) {
+            acceptedFiles.push(file);
+        }
+        else {
+            rejectedFiles.push(file);
+            hasMaxSizeError = true;
+        }
+    }
+    if (acceptedFiles.length > maxAttachments) {
+        return {
+            acceptedFiles: acceptedFiles.slice(0, maxAttachments),
+            rejectedFiles: [...acceptedFiles.slice(maxAttachments), ...rejectedFiles],
+            hasMaxAttachmentsError: true,
+            hasMaxAttachmentSizeError: hasMaxSizeError,
+        };
+    }
+    return {
+        acceptedFiles,
+        rejectedFiles,
+        hasMaxAttachmentsError: false,
+        hasMaxAttachmentSizeError: hasMaxSizeError,
+    };
+}
 
 const defaultAIConversationDisplayTextEn = {
     getMessageTimestampText: (date) => formatDate(date),
+    getMaxAttachmentErrorText(count) {
+        return `Cannot choose more than ${count} ${count === 1 ? 'file' : 'files'}. `;
+    },
+    getAttachmentSizeErrorText(sizeText) {
+        return `File size must be below ${sizeText}.`;
+    },
 };
 
 const { ConversationDisplayTextContext, ConversationDisplayTextProvider, useConversationDisplayText, } = uiReactCore.createContextUtilities({
@@ -141,8 +182,12 @@ const convertResponseComponentsToToolConfiguration = (responseComponents) => {
         const { props } = responseComponents[toolName];
         const requiredProps = [];
         Object.keys(props).forEach((propName) => {
-            if (props[propName].required)
+            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,
@@ -171,9 +216,18 @@ const { MessageRendererContext, MessageRendererProvider, useMessageRenderer, } =
     errorMessage: '`useMessageRenderer` must be used with an AIConversation component',
 });
 
-const AttachmentContext = React__namespace.createContext(false);
-const AttachmentProvider = ({ children, allowAttachments, }) => {
-    return (React__namespace.createElement(AttachmentContext.Provider, { value: allowAttachments ?? false }, children));
+const AttachmentContext = React__namespace.createContext({
+    allowAttachments: false,
+    // We save attachments as base64 strings into dynamodb for conversation history
+    // DynamoDB has a max size of 400kb for records
+    // This can be overridden so cutsomers could provide a lower number
+    // or a higher number if in the future we support larger sizes.
+    maxAttachmentSize: 400000,
+    maxAttachments: 20,
+});
+const AttachmentProvider = ({ children, allowAttachments = false, maxAttachmentSize = 400000, maxAttachments = 20, }) => {
+    const providerValue = React__namespace.useMemo(() => ({ maxAttachmentSize, maxAttachments, allowAttachments }), [maxAttachmentSize, maxAttachments, allowAttachments]);
+    return (React__namespace.createElement(AttachmentContext.Provider, { value: providerValue }, children));
 };
 
 const WelcomeMessageContext = React__namespace.createContext(undefined);
@@ -181,6 +235,11 @@ const WelcomeMessageProvider = ({ children, welcomeMessage, }) => {
     return (React__namespace.createElement(WelcomeMessageContext.Provider, { value: welcomeMessage }, children));
 };
 
+const FallbackComponentContext = React__namespace["default"].createContext(undefined);
+const FallbackComponentProvider = ({ children, FallbackComponent, }) => {
+    return (React__namespace["default"].createElement(FallbackComponentContext.Provider, { value: FallbackComponent }, children));
+};
+
 const DEFAULT_ICON_PATHS = {
     attach: 'M720-330q0 104-73 177T470-80q-104 0-177-73t-73-177v-370q0-75 52.5-127.5T400-880q75 0 127.5 52.5T580-700v350q0 46-32 78t-78 32q-46 0-78-32t-32-78v-370h80v370q0 13 8.5 21.5T470-320q13 0 21.5-8.5T500-350v-350q-1-42-29.5-71T400-800q-42 0-71 29t-29 71v370q-1 71 49 120.5T470-160q70 0 119-49.5T640-330v-390h80v390Z',
     close: 'm256-200-56-56 224-224-224-224 56-56 224 224 224-224 56 56-224 224 224 224-56 56-224-224-224 224Z',
@@ -197,7 +256,7 @@ const DEFAULT_ICON_ATTRIBUTES = {
     fill: 'none',
     xmlns: 'http://www.w3.org/2000/svg',
 };
-const BaseIconElement = elements.defineBaseElement({
+const BaseIconElement = elements.defineBaseElementWithRef({
     type: 'svg',
     displayName: 'Icon',
 });
@@ -213,44 +272,44 @@ const getIconProps = ({ variant, ...props }) => {
 };
 const IconElement = elements.withBaseElementProps(BaseIconElement, getIconProps);
 
-const LabelElement$1 = elements.defineBaseElement({
+const LabelElement$1 = elements.defineBaseElementWithRef({
     type: 'label',
     displayName: 'Label',
 });
-const TextElement = elements.defineBaseElement({
+const TextElement = elements.defineBaseElementWithRef({
     type: 'p',
     displayName: 'Text',
 });
-const UnorderedListElement = elements.defineBaseElement({
+const UnorderedListElement = elements.defineBaseElementWithRef({
     type: 'ul',
     displayName: 'UnorderedList',
 });
-const ListItemElement = elements.defineBaseElement({
+const ListItemElement = elements.defineBaseElementWithRef({
     type: 'li',
     displayName: 'ListItem',
 });
-const HeadingElement = elements.defineBaseElement({
+const HeadingElement = elements.defineBaseElementWithRef({
     type: 'h2',
     displayName: 'Title',
 });
-const ImageElement = elements.defineBaseElement({
+const ImageElement = elements.defineBaseElementWithRef({
     type: 'img',
     displayName: 'Image',
 });
-const InputElement = elements.defineBaseElement({
+const InputElement = elements.defineBaseElementWithRef({
     type: 'input',
     displayName: 'Input',
 });
-const ButtonElement = elements.defineBaseElement({ type: 'button', displayName: 'Button' });
-const ViewElement = elements.defineBaseElement({
+const ButtonElement = elements.defineBaseElementWithRef({ type: 'button', displayName: 'Button' });
+const ViewElement = elements.defineBaseElementWithRef({
     type: 'div',
     displayName: 'View',
 });
-const SpanElement = elements.defineBaseElement({
+const SpanElement = elements.defineBaseElementWithRef({
     type: 'span',
     displayName: 'Span',
 });
-const TextAreaElement = elements.defineBaseElement({
+const TextAreaElement = elements.defineBaseElementWithRef({
     type: 'textarea',
     displayName: 'TextArea',
 });
@@ -269,9 +328,9 @@ const AIConversationElements = {
     View: ViewElement,
 };
 
-const { Button: Button$4, Span: Span$3, View: View$6 } = AIConversationElements;
+const { Button: Button$4, Span: Span$2, View: View$6 } = AIConversationElements;
 const ACTIONS_BAR_BLOCK = 'ai-actions-bar';
-const ActionIcon = elements.withBaseElementProps(Span$3, {
+const ActionIcon = elements.withBaseElementProps(Span$2, {
     'aria-hidden': 'true',
     className: `${ACTIONS_BAR_BLOCK}__icon`,
 });
@@ -286,14 +345,13 @@ const Container$3 = elements.withBaseElementProps(View$6, {
 });
 const ActionsBarControl = ({ message, focusable, }) => {
     const actions = React__namespace["default"].useContext(ActionsContext);
-    return (React__namespace["default"].createElement(Container$3, null, actions?.map((action, index) => (React__namespace["default"].createElement(ActionButton, { "aria-label": action.displayName, key: index, onClick: () => action.handler(message), tabIndex: focusable ? 0 : -1 },
-        React__namespace["default"].createElement(ActionIcon, { "data-testid": `action-icon-${action.displayName}` }, action.icon))))));
+    return (React__namespace["default"].createElement(Container$3, null, actions?.map((action, index) => (React__namespace["default"].createElement(ActionButton, { key: index, onClick: () => action.handler(message), tabIndex: focusable ? 0 : -1 }, action.component)))));
 };
 ActionsBarControl.Button = ActionButton;
 ActionsBarControl.Container = Container$3;
 ActionsBarControl.Icon = ActionIcon;
 
-const { Icon: Icon$3, Span: Span$2, Text: Text$2, View: View$5 } = AIConversationElements;
+const { Icon: Icon$3, Span: Span$1, Text: Text$2, View: View$5 } = AIConversationElements;
 const AVATAR_BLOCK = 'ai-avatar';
 const DEFAULT_USER_ICON = elements.withBaseElementProps(Icon$3, {
     variant: 'user-avatar',
@@ -304,7 +362,7 @@ const DEFAULT_AI_ICON = () => (React__namespace["default"].createElement("svg",
 const AvatarDisplayName = elements.withBaseElementProps(Text$2, {
     className: `${AVATAR_BLOCK}__display-name`,
 });
-const AvatarIcon = elements.withBaseElementProps(Span$2, {
+const AvatarIcon = elements.withBaseElementProps(Span$1, {
     'aria-hidden': true,
     className: `${AVATAR_BLOCK}__icon`,
 });
@@ -385,7 +443,7 @@ AttachFileControl.Icon = AttachFileIcon;
 AttachFileControl.Button = AttachFileButton;
 AttachFileControl.Container = AttachFileContainer;
 
-const { Button: Button$2, Icon: Icon$1, ListItem, UnorderedList: ListElement, Span: Span$1, Text: Text$1, View: View$3, } = AIConversationElements;
+const { Button: Button$2, Icon: Icon$1, ListItem, UnorderedList: ListElement, Span, Text: Text$1, View: View$3, } = AIConversationElements;
 const IMAGE_LIST_BLOCK = 'ai-attachment-list';
 const IMAGE_ITEM_BLOCK = 'ai-attachment';
 const REMOVE_IMAGE_BLOCK = 'ai-remove-attachment';
@@ -416,7 +474,7 @@ const FileNameText = elements.withBaseElementProps(Text$1, {
 const FileSizeText = elements.withBaseElementProps(Text$1, {
     className: `${IMAGE_TEXT_BLOCK}__file-size`,
 });
-const Separator$1 = elements.withBaseElementProps(Span$1, {
+const Separator = elements.withBaseElementProps(Span, {
     'aria-hidden': true,
     className: `${IMAGE_TEXT_BLOCK}__separator`,
     children: '|',
@@ -427,13 +485,13 @@ const TextContainer = elements.withBaseElementProps(View$3, {
 const TextControl = ({ fileName, fileSize }) => {
     return (React__namespace["default"].createElement(TextContainer, null,
         React__namespace["default"].createElement(FileNameText, null, fileName),
-        React__namespace["default"].createElement(Separator$1, null),
+        React__namespace["default"].createElement(Separator, null),
         React__namespace["default"].createElement(FileSizeText, null, fileSize)));
 };
 TextControl.Container = TextContainer;
 TextControl.FileName = FileNameText;
 TextControl.FileSize = FileSizeText;
-TextControl.Separator = Separator$1;
+TextControl.Separator = Separator;
 const Container$1 = elements.withBaseElementProps(ListItem, {
     className: `${IMAGE_ITEM_BLOCK}__list-item`,
 });
@@ -548,12 +606,16 @@ const InputContainer = elements.withBaseElementProps(View$2, {
     className: `${FIELD_BLOCK}__input-container`,
 });
 const FormControl = () => {
-    const { input, setInput } = React__namespace["default"].useContext(ConversationInputContext);
+    const { input, setInput, error, setError } = React__namespace["default"].useContext(ConversationInputContext);
     const handleSendMessage = React__namespace["default"].useContext(SendMessageContext);
-    const allowAttachments = React__namespace["default"].useContext(AttachmentContext);
-    const ref = React__namespace["default"].useRef(null);
+    const { allowAttachments, maxAttachmentSize, maxAttachments } = React__namespace["default"].useContext(AttachmentContext);
+    const displayText = useConversationDisplayText();
     const responseComponents = React__namespace["default"].useContext(ResponseComponentsContext);
+    const isLoading = React__namespace["default"].useContext(LoadingContext);
+    const aiContext = React__namespace["default"].useContext(AIContextContext);
+    const ref = React__namespace["default"].useRef(null);
     const controls = React__namespace["default"].useContext(ControlsContext);
+    const [composing, setComposing] = React__namespace["default"].useState(false);
     const submitMessage = async () => {
         ref.current?.reset();
         const submittedContent = [];
@@ -578,6 +640,7 @@ const FormControl = () => {
         if (handleSendMessage) {
             handleSendMessage({
                 content: submittedContent,
+                aiContext: ui.isFunction(aiContext) ? aiContext() : undefined,
                 toolConfiguration: convertResponseComponentsToToolConfiguration(responseComponents),
             });
         }
@@ -590,7 +653,7 @@ const FormControl = () => {
     };
     const handleOnKeyDown = (event) => {
         const { key, shiftKey } = event;
-        if (key === 'Enter' && !shiftKey) {
+        if (key === 'Enter' && !shiftKey && !composing) {
             event.preventDefault();
             const hasInput = !!input?.text || (input?.files?.length && input?.files?.length > 0);
             if (hasInput) {
@@ -598,15 +661,43 @@ const FormControl = () => {
             }
         }
     };
+    const onValidate = React__namespace["default"].useCallback(async (files) => {
+        const previousFiles = input?.files ?? [];
+        const { acceptedFiles, hasMaxAttachmentsError, hasMaxAttachmentSizeError, } = await attachmentsValidator({
+            files: [...files, ...previousFiles],
+            maxAttachments,
+            maxAttachmentSize,
+        });
+        if (hasMaxAttachmentsError || hasMaxAttachmentSizeError) {
+            const errors = [];
+            if (hasMaxAttachmentsError) {
+                errors.push(displayText.getMaxAttachmentErrorText(maxAttachments));
+            }
+            if (hasMaxAttachmentSizeError) {
+                errors.push(displayText.getAttachmentSizeErrorText(
+                // base64 size is about 137% that of the file size
+                // https://en.wikipedia.org/wiki/Base64#MIME
+                ui.humanFileSize((maxAttachmentSize - 814) / 1.37, true)));
+            }
+            setError?.(errors.join(' '));
+        }
+        else {
+            setError?.(undefined);
+        }
+        setInput?.((prevValue) => ({
+            ...prevValue,
+            files: acceptedFiles,
+        }));
+    }, [setInput, input, displayText, maxAttachmentSize, maxAttachments, setError]);
     if (controls?.Form) {
-        return (React__namespace["default"].createElement(controls.Form, { handleSubmit: handleSubmit, input: input, setInput: setInput, allowAttachments: allowAttachments }));
+        return (React__namespace["default"].createElement(controls.Form, { handleSubmit: handleSubmit, input: input, setInput: setInput, onValidate: onValidate, allowAttachments: allowAttachments, isLoading: isLoading, error: error, setError: setError }));
     }
     return (React__namespace["default"].createElement("form", { className: `${FIELD_BLOCK}__form`, onSubmit: handleSubmit, method: "post", ref: ref },
         allowAttachments ? React__namespace["default"].createElement(AttachFileControl, null) : null,
         React__namespace["default"].createElement(InputContainer, null,
             React__namespace["default"].createElement(VisuallyHidden, null,
                 React__namespace["default"].createElement(Label, null)),
-            React__namespace["default"].createElement(TextInput, { onKeyDown: handleOnKeyDown }),
+            React__namespace["default"].createElement(TextInput, { onKeyDown: handleOnKeyDown, onCompositionStart: () => setComposing(true), onCompositionEnd: () => setComposing(false) }),
             React__namespace["default"].createElement(AttachmentListControl, null)),
         React__namespace["default"].createElement(SendButton, null,
             React__namespace["default"].createElement(SendIcon, null))));
@@ -618,61 +709,52 @@ FormControl.TextInput = TextInput;
 FormControl.SendButton = SendButton;
 FormControl.SendIcon = SendIcon;
 
-const { Image, Span, Text, View: View$1 } = AIConversationElements;
-const MESSAGES_BLOCK = 'ai-messages';
-const MESSAGE_BLOCK = 'ai-message';
-const MediaContentBase = elements.withBaseElementProps(Image, {
-    alt: 'Image attachment',
-});
-const MediaContent = React__namespace["default"].forwardRef(function MediaContent(props, ref) {
+const { Text, View: View$1 } = AIConversationElements;
+const MESSAGES_BLOCK = 'amplify-ai-conversation__message__list';
+const MESSAGE_BLOCK = 'amplify-ai-conversation__message';
+const MediaContent = (props) => {
     const variant = React__namespace["default"].useContext(MessageVariantContext);
     const role = React__namespace["default"].useContext(RoleContext);
-    return (React__namespace["default"].createElement(MediaContentBase, { ref: ref, className: `${MESSAGE_BLOCK}__image ${MESSAGE_BLOCK}__image--${variant} ${MESSAGE_BLOCK}__image--${role}`, ...props }));
-});
+    return (React__namespace["default"].createElement(uiReact.Image, { className: ui.classNames(`${MESSAGE_BLOCK}__image`, variant && `${MESSAGE_BLOCK}__image--${variant}`, `${MESSAGE_BLOCK}__image--${role}`), ...props }));
+};
 const TextContent = React__namespace["default"].forwardRef(function TextContent(props, ref) {
     return React__namespace["default"].createElement(Text, { ref: ref, className: `${MESSAGE_BLOCK}__text`, ...props });
 });
-const ContentContainer = React__namespace["default"].forwardRef(function ContentContainer(props, ref) {
-    const variant = React__namespace["default"].useContext(MessageVariantContext);
-    return (React__namespace["default"].createElement(View$1, { "data-testid": 'content', className: `${MESSAGE_BLOCK}__content ${MESSAGE_BLOCK}__content--${variant}`, ref: ref, ...props }));
-});
 const ToolContent = ({ toolUse, }) => {
-    const responseComponents = React__namespace["default"].useContext(ResponseComponentsContext);
+    const responseComponents = React__namespace["default"].useContext(ResponseComponentsContext) ?? {};
+    const FallbackComponent = React__namespace["default"].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 React__namespace["default"].createElement(CustomComponent, { ...input });
+        if (response) {
+            const CustomComponent = response.component;
+            return React__namespace["default"].createElement(CustomComponent, { ...input });
+        }
+        // fallback if there is a UI component message but we don't have
+        // a React component that matches
+        if (FallbackComponent) {
+            return React__namespace["default"].createElement(FallbackComponent, { ...input });
+        }
     }
 };
 const MessageControl = ({ message }) => {
     const messageRenderer = React__namespace["default"].useContext(MessageRendererContext);
-    return (React__namespace["default"].createElement(ContentContainer, null, message.content.map((content, index) => {
+    return (React__namespace["default"].createElement(React__namespace["default"].Fragment, null, message.content.map((content, index) => {
         if (content.text) {
             return messageRenderer?.text ? (React__namespace["default"].createElement(React__namespace["default"].Fragment, { key: index }, messageRenderer.text({ text: content.text }))) : (React__namespace["default"].createElement(TextContent, { "data-testid": 'text-content', key: index }, content.text));
         }
         else if (content.image) {
-            return messageRenderer?.image ? (React__namespace["default"].createElement(React__namespace["default"].Fragment, { key: index }, messageRenderer?.image({ image: content.image }))) : (React__namespace["default"].createElement(MediaContent, { "data-testid": 'image-content', key: index, src: convertBufferToBase64(content.image?.source.bytes, content.image?.format) }));
+            return messageRenderer?.image ? (React__namespace["default"].createElement(React__namespace["default"].Fragment, { key: index }, messageRenderer?.image({ image: content.image }))) : (React__namespace["default"].createElement(MediaContent, { "data-testid": 'image-content', key: index, alt: "", src: convertBufferToBase64(content.image?.source.bytes, content.image?.format) }));
         }
         else if (content.toolUse) {
             return React__namespace["default"].createElement(ToolContent, { toolUse: content.toolUse, key: index });
         }
     })));
 };
-MessageControl.Container = ContentContainer;
-MessageControl.MediaContent = MediaContent;
-MessageControl.TextContent = TextContent;
-const Separator = elements.withBaseElementProps(Span, {
-    'aria-hidden': true,
-    children: '|',
-    className: `${MESSAGE_BLOCK}__separator`,
-});
 const Timestamp = elements.withBaseElementProps(Text, {
     className: `${MESSAGE_BLOCK}__timestamp`,
 });
@@ -733,7 +815,6 @@ const MessagesControl = () => {
             React__namespace["default"].createElement(MessageContainer, { "data-testid": `message`, key: `message-${index}`, tabIndex: focusedItemIndex === index ? 0 : -1, onFocus: () => handleFocus(index), onKeyDown: (event) => onKeyDown(index, event), ref: (el) => (messagesRef.current[index] = el) },
                 React__namespace["default"].createElement(HeaderContainer, null,
                     React__namespace["default"].createElement(AvatarControl, null),
-                    React__namespace["default"].createElement(Separator, null),
                     React__namespace["default"].createElement(Timestamp, null, getMessageTimestampText(new Date(message.createdAt)))),
                 React__namespace["default"].createElement(MessageControl, { message: message }),
                 message.role === 'assistant' ? (React__namespace["default"].createElement(ActionsBarControl, { message: message, focusable: focusedItemIndex === index })) : null)));
@@ -745,7 +826,6 @@ MessagesControl.Container = MessageContainer;
 MessagesControl.HeaderContainer = HeaderContainer;
 MessagesControl.Layout = Layout;
 MessagesControl.Message = MessageControl;
-MessagesControl.Separator = Separator;
 
 const { View, Button } = AIConversationElements;
 const PROMPT_BLOCK = 'ai-prompts';
@@ -789,25 +869,27 @@ PromptControl.Container = Container;
 PromptControl.PromptGroup = PromptGroup;
 PromptControl.PromptCard = PromptCard;
 
-const AIConversationProvider = ({ actions, allowAttachments, avatars, children, controls, displayText, elements: elements$1, handleSendMessage, isLoading, messages, responseComponents, suggestedPrompts, variant, welcomeMessage, }) => {
+const AIConversationProvider = ({ aiContext, actions, allowAttachments, avatars, children, controls, displayText, handleSendMessage, isLoading, maxAttachmentSize, maxAttachments, messages, messageRenderer, responseComponents, suggestedPrompts, variant, welcomeMessage, FallbackResponseComponent, }) => {
     const _displayText = {
         ...defaultAIConversationDisplayTextEn,
         ...displayText,
     };
-    return (React__namespace["default"].createElement(elements.ElementsProvider, { elements: elements$1 },
-        React__namespace["default"].createElement(ControlsProvider, { controls: controls },
-            React__namespace["default"].createElement(SuggestedPromptProvider, { suggestedPrompts: suggestedPrompts },
-                React__namespace["default"].createElement(WelcomeMessageProvider, { welcomeMessage: welcomeMessage },
-                    React__namespace["default"].createElement(ResponseComponentsProvider, { responseComponents: responseComponents },
-                        React__namespace["default"].createElement(AttachmentProvider, { allowAttachments: allowAttachments },
-                            React__namespace["default"].createElement(ConversationDisplayTextProvider, { ..._displayText },
-                                React__namespace["default"].createElement(ConversationInputContextProvider, null,
-                                    React__namespace["default"].createElement(SendMessageContextProvider, { handleSendMessage: handleSendMessage },
-                                        React__namespace["default"].createElement(AvatarsProvider, { avatars: avatars },
-                                            React__namespace["default"].createElement(ActionsProvider, { actions: actions },
-                                                React__namespace["default"].createElement(MessageVariantProvider, { variant: variant },
-                                                    React__namespace["default"].createElement(MessagesProvider, { messages: messages },
-                                                        React__namespace["default"].createElement(LoadingContextProvider, { isLoading: isLoading }, children)))))))))))))));
+    return (React__namespace["default"].createElement(ControlsProvider, { controls: controls },
+        React__namespace["default"].createElement(SuggestedPromptProvider, { suggestedPrompts: suggestedPrompts },
+            React__namespace["default"].createElement(WelcomeMessageProvider, { welcomeMessage: welcomeMessage },
+                React__namespace["default"].createElement(FallbackComponentProvider, { FallbackComponent: FallbackResponseComponent },
+                    React__namespace["default"].createElement(MessageRendererProvider, { ...messageRenderer },
+                        React__namespace["default"].createElement(ResponseComponentsProvider, { responseComponents: responseComponents },
+                            React__namespace["default"].createElement(AttachmentProvider, { allowAttachments: allowAttachments, maxAttachmentSize: maxAttachmentSize, maxAttachments: maxAttachments },
+                                React__namespace["default"].createElement(ConversationDisplayTextProvider, { ..._displayText },
+                                    React__namespace["default"].createElement(ConversationInputContextProvider, null,
+                                        React__namespace["default"].createElement(SendMessageContextProvider, { handleSendMessage: handleSendMessage },
+                                            React__namespace["default"].createElement(AvatarsProvider, { avatars: avatars },
+                                                React__namespace["default"].createElement(ActionsProvider, { actions: actions },
+                                                    React__namespace["default"].createElement(MessageVariantProvider, { variant: variant },
+                                                        React__namespace["default"].createElement(MessagesProvider, { messages: messages },
+                                                            React__namespace["default"].createElement(AIContextProvider, { aiContext: aiContext },
+                                                                React__namespace["default"].createElement(LoadingContextProvider, { isLoading: isLoading }, children)))))))))))))))));
 };
 
 const DefaultMessageControl = () => {
@@ -820,15 +902,11 @@ const DefaultMessageControl = () => {
     }
 };
 
-/**
- * @experimental
- */
 function createAIConversation(input = {}) {
-    const { elements, suggestedPrompts, actions, responseComponents, variant, controls, displayText, allowAttachments, messageRenderer, } = input;
+    const { suggestedPrompts, actions, responseComponents, variant, controls, displayText, allowAttachments, messageRenderer, FallbackResponseComponent, } = input;
     function AIConversation(props) {
         const { messages, avatars, handleSendMessage, isLoading } = props;
         const providerProps = {
-            elements,
             actions,
             suggestedPrompts,
             responseComponents,
@@ -841,6 +919,7 @@ function createAIConversation(input = {}) {
             handleSendMessage,
             isLoading,
             messageRenderer,
+            FallbackResponseComponent,
         };
         return (React__namespace["default"].createElement(AIConversationProvider, { ...providerProps },
             React__namespace["default"].createElement(ViewElement, null,
@@ -857,6 +936,16 @@ function createAIConversation(input = {}) {
     return { AIConversation };
 }
 
+const PlaceholderMessage = ({ role }) => {
+    const variant = React__namespace.useContext(MessageVariantContext);
+    return (React__namespace.createElement(uiReact.View, { className: ui.classNames(ui.ComponentClassName.AIConversationMessage, ui.classNameModifier(ui.ComponentClassName.AIConversationMessage, variant), ui.classNameModifier(ui.ComponentClassName.AIConversationMessage, role)) },
+        React__namespace.createElement(uiReact.View, { className: ui.ComponentClassName.AIConversationMessageAvatar },
+            React__namespace.createElement(uiReact.Avatar, null)),
+        React__namespace.createElement(uiReact.View, { className: ui.ComponentClassName.AIConversationMessageBody },
+            React__namespace.createElement(uiReact.Placeholder, { width: "25%" }),
+            React__namespace.createElement(uiReact.Placeholder, { width: "50%" }),
+            React__namespace.createElement(uiReact.Placeholder, { width: "25%" }))));
+};
 const MessageMeta = ({ message }) => {
     // need to pass this in as props in order for it to be overridable
     const avatars = React__namespace.useContext(AvatarsContext);
@@ -868,29 +957,30 @@ const MessageMeta = ({ message }) => {
         React__namespace.createElement(uiReact.Text, { className: ui.ComponentClassName.AIConversationMessageSenderUsername }, avatar?.username),
         React__namespace.createElement(uiReact.Text, { className: ui.ComponentClassName.AIConversationMessageSenderTimestamp }, getMessageTimestampText(new Date(message.createdAt)))));
 };
-const LoadingMessage = () => {
-    const avatars = React__namespace.useContext(AvatarsContext);
-    const variant = React__namespace.useContext(MessageVariantContext);
-    const avatar = avatars?.ai;
-    return (React__namespace.createElement(uiReact.View, { className: ui.classNames(ui.ComponentClassName.AIConversationMessage, ui.classNameModifier(ui.ComponentClassName.AIConversationMessage, variant), ui.classNameModifier(ui.ComponentClassName.AIConversationMessage, 'assistant')) },
-        React__namespace.createElement(uiReact.View, { className: ui.ComponentClassName.AIConversationMessageAvatar },
-            React__namespace.createElement(uiReact.Avatar, { isLoading: true }, avatar?.avatar)),
-        React__namespace.createElement(uiReact.View, { className: ui.ComponentClassName.AIConversationMessageBody },
-            React__namespace.createElement(uiReact.View, { className: ui.ComponentClassName.AIConversationMessageSender },
-                React__namespace.createElement(uiReact.Text, { className: ui.ComponentClassName.AIConversationMessageSenderUsername }, avatar?.username)))));
+const MessageActions = ({ message }) => {
+    const actions = React__namespace.useContext(ActionsContext);
+    if (!actions)
+        return null;
+    return (React__namespace.createElement(uiReact.View, { className: ui.ComponentClassName.AIConversationMessageActions }, actions.map((action, i) => {
+        return (React__namespace.createElement(uiReact.Button, { key: i, size: "small", onClick: () => {
+                action.handler(message);
+            } }, action.component));
+    })));
 };
 const Message = ({ message }) => {
     const avatars = React__namespace.useContext(AvatarsContext);
     const variant = React__namespace.useContext(MessageVariantContext);
+    const { isLoading } = message;
     const avatar = message.role === 'assistant' ? avatars?.ai : avatars?.user;
     return (React__namespace.createElement(RoleContext.Provider, { value: message.role },
         React__namespace.createElement(uiReact.View, { className: ui.classNames(ui.ComponentClassName.AIConversationMessage, ui.classNameModifier(ui.ComponentClassName.AIConversationMessage, variant), ui.classNameModifier(ui.ComponentClassName.AIConversationMessage, message.role)) },
             React__namespace.createElement(uiReact.View, { className: ui.ComponentClassName.AIConversationMessageAvatar },
-                React__namespace.createElement(uiReact.Avatar, null, avatar?.avatar)),
+                React__namespace.createElement(uiReact.Avatar, { isLoading: isLoading }, avatar?.avatar)),
             React__namespace.createElement(uiReact.View, { className: ui.ComponentClassName.AIConversationMessageBody },
                 React__namespace.createElement(MessageMeta, { message: message }),
                 React__namespace.createElement(uiReact.View, { className: ui.ComponentClassName.AIConversationMessageContent },
-                    React__namespace.createElement(MessageControl, { message: message }))))));
+                    React__namespace.createElement(MessageControl, { message: message })),
+                message.role === 'assistant' ? (React__namespace.createElement(MessageActions, { message: message })) : null))));
 };
 const MessageList = ({ messages, }) => {
     const isLoading = React__namespace.useContext(LoadingContext);
@@ -898,8 +988,10 @@ const MessageList = ({ messages, }) => {
         content.text ??
         content.toolUse?.name.startsWith(RESPONSE_COMPONENT_PREFIX))) ?? [];
     return (React__namespace.createElement(uiReact.View, { className: ui.ComponentClassName.AIConversationMessageList },
-        messagesWithRenderableContent.map((message, i) => (React__namespace.createElement(Message, { key: `message-${i}`, message: message }))),
-        isLoading ? React__namespace.createElement(LoadingMessage, null) : null));
+        isLoading ? (React__namespace.createElement(React__namespace.Fragment, null,
+            React__namespace.createElement(PlaceholderMessage, { role: "user" }),
+            React__namespace.createElement(PlaceholderMessage, { role: "assistant" }))) : null,
+        messagesWithRenderableContent.map((message, i) => (React__namespace.createElement(Message, { key: `message-${i}`, message: message })))));
 };
 
 const Attachment = ({ file, handleRemove, }) => {
@@ -932,27 +1024,24 @@ function isHTMLFormElement(target) {
  * Will conditionally render the DropZone if allowAttachments
  * is true
  */
-const FormWrapper = ({ children, allowAttachments, setInput, }) => {
+const FormWrapper = ({ children, allowAttachments, onValidate, }) => {
     if (allowAttachments) {
         return (React__namespace.createElement(uiReact.DropZone, { className: ui.ComponentClassName.AIConversationFormDropzone, onDropComplete: ({ acceptedFiles }) => {
-                setInput?.((prevInput) => ({
-                    ...prevInput,
-                    files: [...(prevInput?.files ?? []), ...acceptedFiles],
-                }));
+                onValidate(acceptedFiles);
             } }, children));
     }
     else {
         return children;
     }
 };
-const Form = ({ setInput, input, handleSubmit, allowAttachments, }) => {
+const Form = ({ setInput, input, handleSubmit, allowAttachments, onValidate, isLoading, error, }) => {
     const icons = internal.useIcons('aiConversation');
     const sendIcon = icons?.send ?? React__namespace.createElement(internal.IconSend, null);
     const attachIcon = icons?.attach ?? React__namespace.createElement(internal.IconAttach, null);
     const hiddenInput = React__namespace.useRef(null);
-    const isLoading = React__namespace.useContext(LoadingContext);
+    const [composing, setComposing] = React__namespace.useState(false);
     const isInputEmpty = !input?.text?.length && !input?.files?.length;
-    return (React__namespace.createElement(FormWrapper, { allowAttachments: allowAttachments, setInput: setInput },
+    return (React__namespace.createElement(FormWrapper, { onValidate: onValidate, allowAttachments: allowAttachments },
         React__namespace.createElement(uiReact.View, { as: "form", className: ui.ComponentClassName.AIConversationForm, onSubmit: handleSubmit },
             allowAttachments ? (React__namespace.createElement(uiReact.Button, { className: ui.ComponentClassName.AIConversationFormAttach, onClick: () => {
                     hiddenInput?.current?.click();
@@ -963,24 +1052,20 @@ const Form = ({ setInput, input, handleSubmit, allowAttachments, }) => {
                 React__namespace.createElement("span", null, attachIcon),
                 React__namespace.createElement(uiReact.VisuallyHidden, null,
                     React__namespace.createElement("input", { type: "file", tabIndex: -1, ref: hiddenInput, onChange: (e) => {
-                            const { files } = e.target;
-                            if (!files || files.length === 0) {
+                            if (!e.target.files || e.target.files.length === 0) {
                                 return;
                             }
-                            setInput((prevValue) => ({
-                                ...prevValue,
-                                files: [...(prevValue?.files ?? []), ...Array.from(files)],
-                            }));
-                        }, multiple: true, accept: "*", "data-testid": "hidden-file-input" })))) : null,
-            React__namespace.createElement(uiReact.TextAreaField, { className: ui.ComponentClassName.AIConversationFormField, label: "input", labelHidden: true, autoResize: true, flex: "1", rows: 1, value: input?.text ?? '', testId: "text-input", onKeyDown: (e) => {
+                            onValidate(Array.from(e.target.files));
+                        }, multiple: true, accept: ".jpeg,.png,.webp,.gif", "data-testid": "hidden-file-input" })))) : null,
+            React__namespace.createElement(uiReact.TextAreaField, { className: ui.ComponentClassName.AIConversationFormField, label: "input", labelHidden: true, autoResize: true, flex: "1", rows: 1, value: input?.text ?? '', testId: "text-input", onCompositionStart: () => setComposing(true), onCompositionEnd: () => setComposing(false), onKeyDown: (e) => {
                     // Submit on enter key if shift is not pressed also
-                    const shouldSubmit = !e.shiftKey && e.key === 'Enter';
+                    const shouldSubmit = !e.shiftKey && e.key === 'Enter' && !composing;
                     if (shouldSubmit && isHTMLFormElement(e.target)) {
                         e.target.form.requestSubmit();
                         e.preventDefault();
                     }
                 }, onChange: (e) => {
-                    setInput((prevValue) => ({
+                    setInput?.((prevValue) => ({
                         ...prevValue,
                         text: e.target.value,
                     }));
@@ -990,6 +1075,7 @@ const Form = ({ setInput, input, handleSubmit, allowAttachments, }) => {
                 // eslint-disable-next-line @typescript-eslint/prefer-nullish-coalescing
                 isDisabled: isLoading || isInputEmpty },
                 React__namespace.createElement("span", null, sendIcon))),
+        error ? (React__namespace.createElement(uiReact.Message, { className: ui.ComponentClassName.AIConversationFormError, variation: "plain", colorTheme: "warning" }, error)) : null,
         React__namespace.createElement(Attachments, { setInput: setInput, files: input?.files })));
 };
 
@@ -1004,9 +1090,9 @@ const PromptList = ({ setInput, suggestedPrompts = [], }) => {
     })));
 };
 
-const VERSION = '0.4.0';
+const VERSION = '1.1.0';
 
-function Provider({ actions, avatars, controls, handleSendMessage, messages, responseComponents, suggestedPrompts, variant, isLoading, displayText, allowAttachments, messageRenderer, children, }) {
+function AIConversationBase({ avatars, controls, ...rest }) {
     uiReactCore.useSetUserAgent({
         componentName: 'AIConversation',
         packageName: 'react-ai',
@@ -1016,83 +1102,53 @@ function Provider({ actions, avatars, controls, handleSendMessage, messages, res
     const defaultAvatars = {
         ai: {
             username: 'Assistant',
-            avatar: icons?.assistant ?? React__namespace.createElement(internal.IconAssistant, null),
+            avatar: icons?.assistant ?? React__namespace.createElement(internal.IconAssistant, { testId: "icon-assistant" }),
         },
         user: {
             username: 'User',
-            avatar: icons?.user ?? React__namespace.createElement(internal.IconUser, null),
+            avatar: icons?.user ?? React__namespace.createElement(internal.IconUser, { testId: "icon-user" }),
         },
     };
     const providerProps = {
-        messages,
-        handleSendMessage,
+        ...rest,
         avatars: {
             ...defaultAvatars,
             ...avatars,
         },
-        isLoading,
-        elements: {
-            Text: uiReact.Text,
-        },
-        actions,
-        suggestedPrompts,
-        responseComponents,
-        variant,
         controls: {
             MessageList,
             PromptList,
             Form,
             ...controls,
         },
-        displayText,
-        allowAttachments,
-        messageRenderer,
     };
-    return (React__namespace.createElement(AIConversationProvider, { ...providerProps }, children));
-}
-function AIConversationBase(props) {
-    return (React__namespace.createElement(Provider, { ...props },
-        React__namespace.createElement(uiReact.Flex, { className: ui.ComponentClassName.AIConversation },
+    return (React__namespace.createElement(AIConversationProvider, { ...providerProps },
+        React__namespace.createElement(uiReact.Flex, { className: ui.ComponentClassName.AIConversation, testId: "ai-conversation" },
             React__namespace.createElement(uiReact.ScrollView, { autoScroll: "smooth", flex: "1" },
                 React__namespace.createElement(DefaultMessageControl, null),
                 React__namespace.createElement(MessagesControl, null)),
             React__namespace.createElement(FormControl, null))));
 }
-/**
- * @experimental
- */
 const AIConversation = Object.assign(AIConversationBase, {
-    Provider,
+    Provider: AIConversationProvider,
     DefaultMessage: DefaultMessageControl,
     Messages: MessagesControl,
     Form: FormControl,
 });
 
-const AIContext = React__namespace["default"].createContext(undefined);
-const useAIContext = () => {
-    const context = React__namespace["default"].useContext(AIContext);
-    const [routeToConversationsMap, setRouteToConversationsMap] = React__namespace["default"].useState({});
-    if (context) {
-        return context;
-    }
-    return { routeToConversationsMap, setRouteToConversationsMap };
-};
-/**
- * @experimental
- */
-const AIContextProvider = ({ children, }) => {
-    const context = useAIContext();
-    return React__namespace["default"].createElement(AIContext.Provider, { value: context }, children);
-};
-
 // default state
 const INITIAL_STATE = {
     hasError: false,
     isLoading: false,
     messages: undefined,
 };
-const LOADING_STATE = { hasError: false, isLoading: true, messages: undefined };
+const LOADING_STATE = {
+    hasError: false,
+    isLoading: true,
+    messages: undefined,
+};
 const ERROR_STATE = { hasError: true, isLoading: false };
+
 function createUseAIGeneration(client) {
     const useAIGeneration = (routeName) => {
         const [dataState, setDataState] = React__namespace.useState(() => ({
@@ -1119,142 +1175,314 @@ function createUseAIGeneration(client) {
     return useAIGeneration;
 }
 
-function createNewConversationMessageInRoute({ previousValue, routeName, conversationId, messages, }) {
+const contentFromEvents = (contentBlocks) => {
+    if (!contentBlocks)
+        return [];
+    return contentBlocks.map((contentBlock) => {
+        const isTextBlock = contentBlock.some((event) => event.text);
+        if (isTextBlock) {
+            return {
+                text: contentBlock
+                    .map((event) => {
+                    return event.text;
+                })
+                    .join(''),
+            };
+        }
+        // tool use is never chunked
+        if (contentBlock[0].toolUse) {
+            return { toolUse: contentBlock[0].toolUse };
+        }
+    });
+};
+
+async function exhaustivelyListMessages({ conversation, messages = [], nextToken, }) {
+    const result = await conversation.listMessages({ nextToken });
+    if (result.data) {
+        messages?.push(...result.data);
+    }
+    if (result.nextToken) {
+        return exhaustivelyListMessages({
+            conversation,
+            messages,
+            nextToken: result.nextToken,
+        });
+    }
     return {
-        ...previousValue,
-        [routeName]: {
-            ...previousValue[routeName],
-            [conversationId]: messages,
-        },
+        ...result,
+        data: messages,
     };
 }
+
+function hasStarted(state) {
+    return ['initialLoading', 'initialized'].includes(state);
+}
 function createUseAIConversation(client) {
+    // This is a bit complicated so buckle up.
+    // The way the data client works is conversation.get() or conversation.create()
+    // is an async function because it makes a graphql call to appsync
+    // then it returns a conversation object, which is like a normal
+    // data client record, except that it also has functions on it,
+    // like sendMessage and onStreamEvent. onStreamEvent sets up a
+    // subscription using a websocket connection, which ideally we only want to
+    // do once per conversation. Because we can only subscribe AFTER the
+    // async call to get/create the conversation is made, the cleanup
+    // function in the effect will won't actually unsubscribe
     const useAIConversation = (routeName, input = {}) => {
         const clientRoute = client.conversations[routeName];
-        const { routeToConversationsMap, setRouteToConversationsMap } = useAIContext();
-        const messagesFromAIContext = input.id
-            ? routeToConversationsMap[routeName]?.[input.id]
-            : undefined;
-        const [localMessages, setLocalMessages] = React__namespace["default"].useState(messagesFromAIContext ?? []);
-        const [conversation, setConversation] = React__namespace["default"].useState(undefined);
-        const [waitingForAIResponse, setWaitingForAIResponse] = React__namespace["default"].useState(false);
-        const [errorMessage, setErrorMessage] = React__namespace["default"].useState();
-        const [hasError, setHasError] = React__namespace["default"].useState(false);
-        // On hook initialization get conversation and load all messages
+        // We need to keep track of the stream events as the come in
+        // for an assistant message, but don't need to keep them in state
+        const contentBlocksRef = React__namespace["default"].useRef();
+        // 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__namespace["default"].useRef('initial');
+        const [dataState, setDataState] = React__namespace["default"].useState(() => ({
+            ...INITIAL_STATE,
+            data: { messages: [], conversation: undefined },
+        }));
+        const { conversation } = dataState.data;
+        const { id, onInitialize, onMessage } = input;
         React__namespace["default"].useEffect(() => {
             async function initialize() {
-                const { data: conversation } = input.id
-                    ? await clientRoute.get({ id: input.id })
-                    : await clientRoute.create();
-                if (!conversation) {
-                    const errorString = 'No conversation found';
-                    setHasError(true);
-                    setErrorMessage(errorString);
-                    throw new Error(errorString);
+                // 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
+                if (id) {
+                    setDataState({
+                        ...LOADING_STATE,
+                        data: { messages: [], conversation: undefined },
+                    });
                 }
-                const { data: messages } = await conversation.listMessages();
-                setLocalMessages(messages);
-                setConversation(conversation);
-                setRouteToConversationsMap((previousValue) => {
-                    return createNewConversationMessageInRoute({
-                        previousValue,
-                        routeName: routeName,
-                        conversationId: conversation.id,
-                        messages,
+                const { data: conversation, errors } = id
+                    ? await clientRoute.get({ id })
+                    : await clientRoute.create();
+                if (errors ?? !conversation) {
+                    setDataState({
+                        ...ERROR_STATE,
+                        data: { messages: [] },
+                        messages: errors,
                     });
+                }
+                else {
+                    if (id) {
+                        const { data: messages } = await exhaustivelyListMessages({
+                            conversation,
+                        });
+                        setDataState({
+                            ...INITIAL_STATE,
+                            data: { messages, conversation },
+                        });
+                    }
+                    else {
+                        setDataState({
+                            ...INITIAL_STATE,
+                            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();
-        }, [clientRoute, input.id, routeName, setRouteToConversationsMap]);
-        // Update messages to match what is in AIContext if they aren't equal
+            return () => {
+                contentBlocksRef.current = undefined;
+                if (hasStarted(initRef.current))
+                    return;
+                setDataState({
+                    ...INITIAL_STATE,
+                    data: { messages: [], conversation: undefined },
+                });
+            };
+        }, [clientRoute, id, setDataState]);
+        // Run a separate effect that is triggered by the conversation state
+        // so that we know we have a conversation object to set up the subscription
+        // and also unsubscribe on cleanup
         React__namespace["default"].useEffect(() => {
-            if (!!messagesFromAIContext && messagesFromAIContext !== localMessages)
-                setLocalMessages(messagesFromAIContext);
-        }, [messagesFromAIContext, localMessages]);
-        const sendMessage = React__namespace["default"].useCallback((input) => {
-            const { content, aiContext, toolConfiguration } = input;
-            conversation
-                ?.sendMessage({ content, aiContext, toolConfiguration })
-                .then((value) => {
-                const { data: sentMessage } = value;
-                if (sentMessage) {
-                    setWaitingForAIResponse(true);
-                    setLocalMessages((previousLocalMessages) => [
-                        ...previousLocalMessages,
-                        sentMessage,
-                    ]);
-                    setRouteToConversationsMap((previousValue) => {
-                        return createNewConversationMessageInRoute({
-                            previousValue,
-                            routeName: routeName,
-                            conversationId: conversation.id,
-                            messages: [
-                                ...previousValue[routeName][conversation.id],
-                                sentMessage,
-                            ],
+            if (!conversation)
+                return;
+            const subscription = conversation.onStreamEvent({
+                next: (event) => {
+                    const { 
+                    // messages have a content block array,
+                    // this is the index of the content block that was updated
+                    contentBlockIndex, 
+                    // this is the index of the content chunk, ensure these are in order!
+                    contentBlockDeltaIndex, 
+                    // this is sent after the last content chunk, verify this matches the
+                    // previous contentBlockDeltaIndex
+                    contentBlockDoneAtIndex, 
+                    // this is the final event of the conversation turn
+                    stopReason, conversationId, id, } = event;
+                    // return early for content blocks being done
+                    // or conversation turn being over
+                    if (contentBlockDoneAtIndex) {
+                        return;
+                    }
+                    // stop reason will signify end of conversation turn
+                    if (stopReason) {
+                        // remove loading state from streamed message
+                        setDataState((prev) => {
+                            return {
+                                ...prev,
+                                data: {
+                                    ...prev.data,
+                                    messages: prev.data.messages.map((message) => ({
+                                        ...message,
+                                        isLoading: false,
+                                    })),
+                                },
+                            };
                         });
-                    });
-                }
-            })
-                .catch((reason) => {
-                setHasError(true);
-                setErrorMessage(`error sending message ${reason}`);
-            });
-        }, [conversation, routeName, setRouteToConversationsMap]);
-        const subscribe = React__namespace["default"].useCallback((handleStoreChange) => {
-            const subscription = conversation &&
-                conversation.onMessage((message) => {
-                    if (input.onResponse)
-                        input.onResponse(message);
-                    setWaitingForAIResponse(false);
-                    setLocalMessages((previousLocalMessages) => [
-                        ...previousLocalMessages,
-                        message,
-                    ]);
-                    setRouteToConversationsMap((previousValue) => {
-                        return createNewConversationMessageInRoute({
-                            previousValue,
-                            routeName: routeName,
-                            conversationId: conversation.id,
-                            messages: [
-                                ...previousValue[routeName][conversation.id],
-                                message,
-                            ],
+                        onMessage?.({
+                            id,
+                            conversationId,
+                            content: contentFromEvents(contentBlocksRef.current),
+                            createdAt: new Date().toISOString(),
+                            role: 'assistant',
+                            isLoading: true,
                         });
+                        // clear out the stream cache
+                        contentBlocksRef.current = undefined;
+                        return;
+                    }
+                    // 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]];
+                    }
+                    else {
+                        // 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) => {
+                        const message = {
+                            id,
+                            conversationId,
+                            content: contentFromEvents(contentBlocksRef.current),
+                            createdAt: new Date().toISOString(),
+                            role: 'assistant',
+                            isLoading: true,
+                        };
+                        return {
+                            ...prev,
+                            data: {
+                                ...prev.data,
+                                // TODO: we are assuming we only update the last
+                                // message, but maybe we should match it by message ID?
+                                messages: [...prev.data.messages.slice(0, -1), message],
+                            },
+                        };
                     });
-                    handleStoreChange(); // should cause a re-render
-                });
+                },
+                error: (error) => {
+                    setDataState((prev) => {
+                        return {
+                            ...prev,
+                            ...ERROR_STATE,
+                            messages: error.errors,
+                        };
+                    });
+                },
+            });
+            if (ui.isFunction(onInitialize)) {
+                onInitialize(conversation);
+            }
             return () => {
-                subscription?.unsubscribe();
+                contentBlocksRef.current = undefined;
+                subscription.unsubscribe();
             };
-        }, [conversation, routeName, setRouteToConversationsMap, input]);
-        const getSnapshot = React__namespace["default"].useCallback(() => localMessages, [localMessages]);
-        // Using useSyncExternalStore to subscribe to external data updates
-        // Have to provide third optional argument in next - https://github.com/vercel/next.js/issues/54685
-        const messagesFromStore = React__namespace["default"].useSyncExternalStore(subscribe, getSnapshot, getSnapshot);
-        return [
-            {
-                data: { messages: messagesFromStore },
-                isLoading: waitingForAIResponse,
-                message: errorMessage,
-                hasError,
-            },
-            sendMessage,
-        ];
+        }, [conversation, onInitialize, onMessage, setDataState]);
+        const handleSendMessage = React__namespace["default"].useCallback((input) => {
+            const { content } = input;
+            if (conversation) {
+                setDataState((prevState) => ({
+                    ...prevState,
+                    data: {
+                        ...prevState.data,
+                        // optimistically add user and assistant messages
+                        messages: [
+                            ...prevState.data.messages,
+                            {
+                                content,
+                                role: 'user',
+                                createdAt: new Date().toISOString(),
+                                id: 'temp-id',
+                                conversationId: conversation.id ?? '',
+                            },
+                            {
+                                content: [{ text: ' ' }],
+                                role: 'assistant',
+                                createdAt: new Date().toISOString(),
+                                id: 'temp-id-2',
+                                conversationId: conversation.id ?? '',
+                                isLoading: true,
+                            },
+                        ],
+                    },
+                }));
+                conversation.sendMessage(input);
+            }
+            else {
+                setDataState((prev) => ({
+                    ...prev,
+                    ...ERROR_STATE,
+                    messages: [
+                        {
+                            message: 'No conversation found',
+                            errorInfo: null,
+                            errorType: '',
+                        },
+                    ],
+                }));
+            }
+        }, [conversation]);
+        return [dataState, handleSendMessage];
     };
     return useAIConversation;
 }
 
-/**
- * @experimental
- */
 function createAIHooks(_client) {
     const useAIConversation = createUseAIConversation(_client);
     const useAIGeneration = createUseAIGeneration(_client);
     return { useAIConversation, useAIGeneration };
 }
 
-exports.AIContextProvider = AIContextProvider;
 exports.AIConversation = AIConversation;
 exports.createAIConversation = createAIConversation;
 exports.createAIHooks = createAIHooks;