From 022b586f55d302a81a958f2b6029d1e5e98d5891 Mon Sep 17 00:00:00 2001 From: Danny Banks Date: Thu, 14 Nov 2024 13:29:02 -0800 Subject: [PATCH] chore: cleaning up ai streaming hook (#6086) --- .../context/ResponseComponentsContext.tsx | 7 +- .../ResponseComponentsContext.spec.tsx | 1 - .../views/default/MessageList.tsx | 2 +- .../react-ai/src/hooks/useAIConversation.tsx | 129 +++++++++--------- packages/react-ai/src/types.ts | 6 + 5 files changed, 76 insertions(+), 69 deletions(-) diff --git a/packages/react-ai/src/components/AIConversation/context/ResponseComponentsContext.tsx b/packages/react-ai/src/components/AIConversation/context/ResponseComponentsContext.tsx index 9df54c48555..fc71250db77 100644 --- a/packages/react-ai/src/components/AIConversation/context/ResponseComponentsContext.tsx +++ b/packages/react-ai/src/components/AIConversation/context/ResponseComponentsContext.tsx @@ -51,7 +51,12 @@ export const convertResponseComponentsToToolConfiguration = ( const { props } = responseComponents[toolName]; const requiredProps: string[] = []; Object.keys(props).forEach((propName) => { - if (props[propName].required) requiredProps.push(propName); + 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, diff --git a/packages/react-ai/src/components/AIConversation/context/__tests__/ResponseComponentsContext.spec.tsx b/packages/react-ai/src/components/AIConversation/context/__tests__/ResponseComponentsContext.spec.tsx index 773e5253438..768afebe31a 100644 --- a/packages/react-ai/src/components/AIConversation/context/__tests__/ResponseComponentsContext.spec.tsx +++ b/packages/react-ai/src/components/AIConversation/context/__tests__/ResponseComponentsContext.spec.tsx @@ -41,7 +41,6 @@ describe('convertResponseComponentsToToolConfiguration', () => { type: 'string', description: 'The response you want to render in the component.', - required: true, }, foobar: { type: 'number', diff --git a/packages/react-ai/src/components/AIConversation/views/default/MessageList.tsx b/packages/react-ai/src/components/AIConversation/views/default/MessageList.tsx index 6683fa0ba35..2369b15e190 100644 --- a/packages/react-ai/src/components/AIConversation/views/default/MessageList.tsx +++ b/packages/react-ai/src/components/AIConversation/views/default/MessageList.tsx @@ -29,7 +29,7 @@ const PlaceholderMessage = ({ role }: { role: string }) => { )} > - + diff --git a/packages/react-ai/src/hooks/useAIConversation.tsx b/packages/react-ai/src/hooks/useAIConversation.tsx index 6c44723bb68..f3ece989f09 100644 --- a/packages/react-ai/src/hooks/useAIConversation.tsx +++ b/packages/react-ai/src/hooks/useAIConversation.tsx @@ -33,6 +33,16 @@ interface AIConversationState { conversation?: Conversation; } +// The "states" the hook can be in +// initial: default, nothing happened yet +// loading: the hook has either hit AppSync to create or get a conversation +// initialized: the hook has successfully gotten a conversation and is ready to rock +const INITIALIZE_REF = ['initial', 'initialLoading', 'initialized'] as const; + +function hasStarted(state: (typeof INITIALIZE_REF)[number]) { + return ['initialLoading', 'initialized'].includes(state); +} + export type UseAIConversationHook = ( routeName: T, input?: UseAIConversationInput @@ -63,7 +73,7 @@ export function createUseAIConversation< // 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.useRef(false); + const initRef = React.useRef<(typeof INITIALIZE_REF)[number]>('initial'); const [dataState, setDataState] = React.useState< DataClientState @@ -76,30 +86,11 @@ export function createUseAIConversation< const { id, onInitialize, onMessage } = input; React.useEffect(() => { - // We don't want to run the effect multiple times - // because that could create multiple conversation records - if (initRef.current) return; - initRef.current = true; async function initialize() { - // no client route would mean that the user - // is not using TypeScript and entered 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; - } + // 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 @@ -135,13 +126,33 @@ export function createUseAIConversation< 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(); return () => { contentBlocksRef.current = undefined; + if (hasStarted(initRef.current)) return; setDataState({ ...INITIAL_STATE, data: { messages: [], conversation: undefined }, @@ -166,14 +177,9 @@ export function createUseAIConversation< // this is sent after the last content chunk, verify this matches the // previous contentBlockDeltaIndex contentBlockDoneAtIndex, - // this is the text of the content block - text, - // this is a toolUse block, will always come in a single event - // toolUse, // this is the final event of the conversation turn stopReason, conversationId, - // associatedUserMessageId, id, } = event; @@ -212,43 +218,25 @@ export function createUseAIConversation< } // 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]]; - - setDataState((prev) => { - const message: ConversationMessage = { - id, - conversationId, - // TODO: use better logic here - content: [{ text: text ?? '' }], - createdAt: new Date().toISOString(), - role: 'assistant', - isLoading: true, - }; - return { - ...prev, - data: { - ...prev.data, - messages: [...prev.data.messages.slice(0, -1), message], - }, - }; - }); - return; - } - - // 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), - ]; + // 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) => { @@ -272,9 +260,6 @@ export function createUseAIConversation< }); }, error: (error) => { - error.errors.map((e) => { - return e.message; - }); setDataState((prev) => { return { ...prev, @@ -325,6 +310,18 @@ export function createUseAIConversation< }, })); conversation.sendMessage(input); + } else { + setDataState((prev) => ({ + ...prev, + ...ERROR_STATE, + messages: [ + { + message: 'No conversation found', + errorInfo: null, + errorType: '', + }, + ], + })); } }, [conversation] diff --git a/packages/react-ai/src/types.ts b/packages/react-ai/src/types.ts index c81937ff910..69a58a68a8f 100644 --- a/packages/react-ai/src/types.ts +++ b/packages/react-ai/src/types.ts @@ -5,6 +5,12 @@ export type Conversation = NonNullable< Awaited>['data'] >; +// the JS client looks like this: +// client.conversations.[name].onStreamEvent({ +// next: (event) => {}, +// error: (error) => {} +// }) +// This type gets the 'event' ^ export type ConversationStreamEvent = Parameters< Parameters[0]['next'] >[0];