@aws-amplify/ui-react-ai 0.3.2 → 1.0.0

package/dist/esm/hooks/useAIConversation.mjs

@@ -1,128 +1,268 @@
 import React__default from 'react';
-import { useAIContext } from './AIContextProvider.mjs';
+import { INITIAL_STATE, ERROR_STATE, LOADING_STATE } from './shared.mjs';
+import { isFunction } from '@aws-amplify/ui';
+import { contentFromEvents } from './contentFromEvents.mjs';
+import { exhaustivelyListMessages } from './exhaustivelyListMessages.mjs';
 
-function createNewConversationMessageInRoute({ previousValue, routeName, conversationId, messages, }) {
-    return {
-        ...previousValue,
-        [routeName]: {
-            ...previousValue[routeName],
-            [conversationId]: 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__default.useState(messagesFromAIContext ?? []);
-        const [conversation, setConversation] = React__default.useState(undefined);
-        const [waitingForAIResponse, setWaitingForAIResponse] = React__default.useState(false);
-        const [errorMessage, setErrorMessage] = React__default.useState();
-        const [hasError, setHasError] = React__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__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__default.useRef('initial');
+        const [dataState, setDataState] = React__default.useState(() => ({
+            ...INITIAL_STATE,
+            data: { messages: [], conversation: undefined },
+        }));
+        const { conversation } = dataState.data;
+        const { id, onInitialize, onMessage } = input;
         React__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__default.useEffect(() => {
-            if (!!messagesFromAIContext && messagesFromAIContext !== localMessages)
-                setLocalMessages(messagesFromAIContext);
-        }, [messagesFromAIContext, localMessages]);
-        const sendMessage = React__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__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 (isFunction(onInitialize)) {
+                onInitialize(conversation);
+            }
             return () => {
-                subscription?.unsubscribe();
+                contentBlocksRef.current = undefined;
+                subscription.unsubscribe();
             };
-        }, [conversation, routeName, setRouteToConversationsMap, input]);
-        const getSnapshot = React__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__default.useSyncExternalStore(subscribe, getSnapshot, getSnapshot);
-        return [
-            {
-                data: { messages: messagesFromStore },
-                isLoading: waitingForAIResponse,
-                message: errorMessage,
-                hasError,
-            },
-            sendMessage,
-        ];
+        }, [conversation, onInitialize, onMessage, setDataState]);
+        const handleSendMessage = React__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;
 }

package/dist/esm/hooks/useAIGeneration.mjs

@@ -1,13 +1,6 @@
 import * as React from 'react';
+import { INITIAL_STATE, LOADING_STATE, ERROR_STATE } from './shared.mjs';
 
-// default state
-const INITIAL_STATE = {
-    hasError: false,
-    isLoading: false,
-    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.useState(() => ({

package/dist/esm/index.mjs

@@ -1,4 +1,3 @@
 export { createAIConversation } from './components/AIConversation/createAIConversation.mjs';
 export { AIConversation } from './components/AIConversation/AIConversation.mjs';
-export { AIContextProvider } from './hooks/AIContextProvider.mjs';
 export { createAIHooks } from './hooks/createAIHooks.mjs';

package/dist/esm/version.mjs

@@ -0,0 +1,3 @@
+const VERSION = '1.0.0';
+
+export { VERSION };