From 9f914a34f9b8fd7509385996c145d56e2ee1def1 Mon Sep 17 00:00:00 2001 From: Nate Kelley Date: Tue, 11 Feb 2025 15:02:02 -0700 Subject: [PATCH] thinking step --- .../chat/chatMessageInterfaces.ts | 1 + .../_components/NewChatModal/NewChatModal.tsx | 2 - .../Streaming/StreamingMessage_Text.tsx | 34 +++++----- .../ReasoningMessageSelector.tsx | 2 +- .../ChatResponseMessageSelector.tsx | 61 +++++++++-------- .../ChatResponseMessages.tsx | 66 +++++++++++-------- .../ChatContent/MessageContainer.tsx | 2 +- .../Chats/ChatProvider/useChatSelectors.ts | 10 +-- .../ChatProvider/useChatSubscriptions.ts | 2 +- .../NewChatProvider/useAutoAppendThought.ts | 10 +-- .../NewChatProvider/useChatUpdateMessage.ts | 27 ++++++-- web/src/utils/chat.ts | 7 +- 12 files changed, 128 insertions(+), 96 deletions(-) diff --git a/web/src/api/asset_interfaces/chat/chatMessageInterfaces.ts b/web/src/api/asset_interfaces/chat/chatMessageInterfaces.ts index b6fe987d2..766d9039d 100644 --- a/web/src/api/asset_interfaces/chat/chatMessageInterfaces.ts +++ b/web/src/api/asset_interfaces/chat/chatMessageInterfaces.ts @@ -22,6 +22,7 @@ export type BusterChatMessage_text = { type: 'text'; message: string; message_chunk?: string; + is_final_message?: boolean; //defaults to false }; export type BusterChatMessage_fileMetadata = { diff --git a/web/src/app/app/_components/NewChatModal/NewChatModal.tsx b/web/src/app/app/_components/NewChatModal/NewChatModal.tsx index 61956e6d8..b79d6fa5b 100644 --- a/web/src/app/app/_components/NewChatModal/NewChatModal.tsx +++ b/web/src/app/app/_components/NewChatModal/NewChatModal.tsx @@ -198,8 +198,6 @@ const NewChatInput: React.FC<{ const onSelectSearchAsset = useBusterNewChatContextSelector((x) => x.onSelectSearchAsset); const [loadingNewChat, setLoadingNewChat] = useState(false); - console.log(selectedChatDataSource); - const onStartNewChatPreflight = useMemoizedFn(async () => { setLoadingNewChat(true); await onStartNewChat({ prompt, datasetId: selectedChatDataSource?.id }); diff --git a/web/src/app/app/_components/Streaming/StreamingMessage_Text.tsx b/web/src/app/app/_components/Streaming/StreamingMessage_Text.tsx index 74999ad05..1acef73a8 100644 --- a/web/src/app/app/_components/Streaming/StreamingMessage_Text.tsx +++ b/web/src/app/app/_components/Streaming/StreamingMessage_Text.tsx @@ -1,28 +1,30 @@ -import React, { useEffect, useState } from 'react'; +import React, { useEffect, useRef, useTransition } from 'react'; import { AnimatePresence, motion } from 'framer-motion'; import { itemAnimationConfig } from './animationConfig'; +import { useMemoizedFn } from 'ahooks'; interface StreamingMessage_TextProps { isCompletedStream: boolean; - message: { - message_chunk?: string; - message?: string; - }; + message: string; } export const StreamingMessage_Text: React.FC = React.memo( - ({ message: messageProp, isCompletedStream }) => { - const { message_chunk, message } = messageProp; + ({ message, isCompletedStream }) => { + const [isPending, startTransition] = useTransition(); + const textChunksRef = useRef([]); - const [textChunks, setTextChunks] = useState([]); + const setTextChunks = useMemoizedFn((updater: (prevChunks: string[]) => string[]) => { + textChunksRef.current = updater(textChunksRef.current); + startTransition(() => { + //just used to trigger UI update + }); + }); useEffect(() => { - if (message_chunk && !message) { - // Handle streaming chunks - setTextChunks((prevChunks) => [...prevChunks, message_chunk || '']); - } else if (message) { + if (message) { // Handle complete message - const currentText = textChunks.join(''); + const currentText = textChunksRef.current.join(''); + if (message.startsWith(currentText)) { const remainingText = message.slice(currentText.length); if (remainingText) { @@ -30,14 +32,14 @@ export const StreamingMessage_Text: React.FC = React } } else { // If there's a mismatch, just use the complete message - setTextChunks([message]); + setTextChunks(() => [message]); } } - }, [message_chunk, message]); + }, [message]); return (
- {textChunks.map((chunk, index) => ( + {textChunksRef.current.map((chunk, index) => ( {chunk} diff --git a/web/src/app/app/_controllers/ReasoningController/ReasoningMessages/ReasoningMessageSelector.tsx b/web/src/app/app/_controllers/ReasoningController/ReasoningMessages/ReasoningMessageSelector.tsx index 7b87aa321..e8ff92a62 100644 --- a/web/src/app/app/_controllers/ReasoningController/ReasoningMessages/ReasoningMessageSelector.tsx +++ b/web/src/app/app/_controllers/ReasoningController/ReasoningMessages/ReasoningMessageSelector.tsx @@ -22,7 +22,7 @@ const ReasoningMessageRecord: Record< text: (props) => ( ), file: ReasoningMessage_File diff --git a/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatResponseMessages/ChatResponseMessageSelector.tsx b/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatResponseMessages/ChatResponseMessageSelector.tsx index 8881ea0cc..7ea9addd7 100644 --- a/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatResponseMessages/ChatResponseMessageSelector.tsx +++ b/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatResponseMessages/ChatResponseMessageSelector.tsx @@ -16,7 +16,10 @@ const ChatResponseMessageRecord: Record< React.FC > = { text: (props) => ( - + ), file: ChatResponseMessage_File }; @@ -27,37 +30,37 @@ export interface ChatResponseMessageSelectorProps { isLastMessageItem: boolean; } -export const ChatResponseMessageSelector: React.FC = ({ - responseMessage, - isCompletedStream, - isLastMessageItem -}) => { - const messageType = responseMessage.type; - const ChatResponseMessage = ChatResponseMessageRecord[messageType]; - const { cx, styles } = useStyles(); +export const ChatResponseMessageSelector: React.FC = React.memo( + ({ responseMessage, isCompletedStream, isLastMessageItem }) => { + const { cx, styles } = useStyles(); + const messageType = responseMessage.type; + const ChatResponseMessage = ChatResponseMessageRecord[messageType]; - const typeClassRecord: Record = useMemo(() => { - return { - text: cx(styles.textCard, 'text-card'), - file: cx(styles.fileCard, 'file-card') - }; - }, []); + const typeClassRecord: Record = useMemo(() => { + return { + text: cx(styles.textCard, 'text-card'), + file: cx(styles.fileCard, 'file-card') + }; + }, []); - const getContainerClass = useMemoizedFn((item: BusterChatMessageResponse) => { - return typeClassRecord[item.type]; - }); + const getContainerClass = useMemoizedFn((item: BusterChatMessageResponse) => { + return typeClassRecord[item.type]; + }); - return ( -
- - -
- ); -}; + return ( +
+ + +
+ ); + } +); + +ChatResponseMessageSelector.displayName = 'ChatResponseMessageSelector'; const VerticalDivider: React.FC<{ className?: string }> = React.memo(({ className }) => { const { cx, styles } = useStyles(); diff --git a/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatResponseMessages/ChatResponseMessages.tsx b/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatResponseMessages/ChatResponseMessages.tsx index 7a8467bfe..3239aeffa 100644 --- a/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatResponseMessages/ChatResponseMessages.tsx +++ b/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatResponseMessages/ChatResponseMessages.tsx @@ -7,6 +7,7 @@ import type { import { MessageContainer } from '../MessageContainer'; import { ChatResponseMessageSelector } from './ChatResponseMessageSelector'; import { ChatResponseReasoning } from './ChatResponseReasoning'; +import { ShimmerText } from '@/components/text'; interface ChatResponseMessagesProps { responseMessages: BusterChatMessageResponse[]; @@ -17,40 +18,41 @@ interface ChatResponseMessagesProps { export const ChatResponseMessages: React.FC = React.memo( ({ responseMessages, reasoningMessages, isCompletedStream, messageId }) => { - const firstResponseMessage = responseMessages[0] as BusterChatMessage_text; - const restResponseMessages = useMemo(() => { - if (!firstResponseMessage) return []; - return responseMessages.slice(1); - }, [firstResponseMessage, responseMessages]); - const lastMessageIndex = responseMessages.length - 1; + const showDefaultMessage = responseMessages.length === 0; + + const reasonginStepIndex = useMemo(() => { + const lastTextMessage = responseMessages.findLast( + (message) => message.type === 'text' && message.is_final_message !== false + ) as BusterChatMessage_text; + + if (!lastTextMessage) return -1; + if (lastTextMessage?.message_chunk) return -1; + + return responseMessages.findIndex((message) => message.id === lastTextMessage.id); + }, [responseMessages]); + return ( - {firstResponseMessage && ( - - )} + {showDefaultMessage && } - {firstResponseMessage && ( - - )} + {responseMessages.map((responseMessage, index) => ( + + - {restResponseMessages.map((responseMessage, index) => ( - + {index === reasonginStepIndex && ( + + )} + ))} ); @@ -58,3 +60,11 @@ export const ChatResponseMessages: React.FC = React.m ); ChatResponseMessages.displayName = 'ChatResponseMessages'; + +const DefaultFirstMessage: React.FC = () => { + return ( +
+ +
+ ); +}; diff --git a/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/MessageContainer.tsx b/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/MessageContainer.tsx index 250536669..a273ec44c 100644 --- a/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/MessageContainer.tsx +++ b/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/MessageContainer.tsx @@ -11,7 +11,7 @@ export const MessageContainer: React.FC<{ }> = React.memo(({ children, senderName, senderId, senderAvatar, className = '' }) => { const { cx } = useStyles(); return ( -
+
{senderName ? ( ) : ( diff --git a/web/src/context/Chats/ChatProvider/useChatSelectors.ts b/web/src/context/Chats/ChatProvider/useChatSelectors.ts index 28a3adebf..afdfc5137 100644 --- a/web/src/context/Chats/ChatProvider/useChatSelectors.ts +++ b/web/src/context/Chats/ChatProvider/useChatSelectors.ts @@ -17,25 +17,25 @@ export const useChatSelectors = ({ const getChatMessages = useCallback( (chatId: string): IBusterChatMessage[] => { - const chatMessageIds = chatsRef.current[chatId].messages || []; - return chatMessageIds.map((messageId) => chatsMessagesRef.current[messageId]); + return getChatMessagesMemoized(chatId); }, [chatsMessagesRef, isPending, chatsRef] ); const getChatMessage = useCallback( (messageId: string): IBusterChatMessage | undefined => { - return chatsMessagesRef.current[messageId]; + return getChatMessageMemoized(messageId); }, [chatsMessagesRef, isPending] ); const getChatMessagesMemoized = useMemoizedFn((chatId: string) => { - return getChatMessages(chatId); + const chatMessageIds = chatsRef.current[chatId].messages || []; + return chatMessageIds.map((messageId) => chatsMessagesRef.current[messageId]); }); const getChatMessageMemoized = useMemoizedFn((messageId: string) => { - return getChatMessage(messageId); + return chatsMessagesRef.current[messageId]; }); return { diff --git a/web/src/context/Chats/ChatProvider/useChatSubscriptions.ts b/web/src/context/Chats/ChatProvider/useChatSubscriptions.ts index 05cc09d8c..b8a22c139 100644 --- a/web/src/context/Chats/ChatProvider/useChatSubscriptions.ts +++ b/web/src/context/Chats/ChatProvider/useChatSubscriptions.ts @@ -17,7 +17,7 @@ export const useChatSubscriptions = ({ const busterSocket = useBusterWebSocket(); const _onGetChat = useMemoizedFn((chat: BusterChat): IBusterChat => { - const { iChat, iChatMessages } = updateChatToIChat(chat); + const { iChat, iChatMessages } = updateChatToIChat(chat, false); chatsRef.current[chat.id] = iChat; chatsMessagesRef.current = { diff --git a/web/src/context/Chats/NewChatProvider/useAutoAppendThought.ts b/web/src/context/Chats/NewChatProvider/useAutoAppendThought.ts index ea4374e07..744102d12 100644 --- a/web/src/context/Chats/NewChatProvider/useAutoAppendThought.ts +++ b/web/src/context/Chats/NewChatProvider/useAutoAppendThought.ts @@ -11,9 +11,7 @@ import random from 'lodash/random'; export const useAutoAppendThought = () => { const onUpdateChatMessage = useBusterChatContextSelector((x) => x.onUpdateChatMessage); - const getChatMemoized = useBusterChatContextSelector((x) => x.getChatMemoized); const getChatMessagesMemoized = useBusterChatContextSelector((x) => x.getChatMessagesMemoized); - const getChatMessageMemoized = useBusterChatContextSelector((x) => x.getChatMessageMemoized); const removeAutoThoughts = useMemoizedFn( (reasoningMessages: BusterChatMessageReasoning[]): BusterChatMessageReasoning[] => { @@ -26,8 +24,9 @@ export const useAutoAppendThought = () => { reasoningMessages: BusterChatMessageReasoning[], chatId: string ): BusterChatMessageReasoning[] => { + const lastReasoningMessage = reasoningMessages[reasoningMessages.length - 1]; const lastMessageIsCompleted = - reasoningMessages[reasoningMessages.length - 1].status === 'completed'; + !lastReasoningMessage || lastReasoningMessage?.status === 'completed'; if (lastMessageIsCompleted) { _loopAutoThought(chatId); @@ -40,13 +39,14 @@ export const useAutoAppendThought = () => { ); const _loopAutoThought = useMemoizedFn(async (chatId: string) => { - const randomDelay = random(3500, 5500); + const randomDelay = random(3000, 5000); await timeout(randomDelay); const chatMessages = getChatMessagesMemoized(chatId); const lastMessage = last(chatMessages); const isCompletedStream = !!lastMessage?.isCompletedStream; const lastReasoningMessage = last(lastMessage?.reasoning); - const lastReasoningMessageIsAutoAppended = lastReasoningMessage?.id === AUTO_THOUGHT_ID; + const lastReasoningMessageIsAutoAppended = + !lastReasoningMessage || lastReasoningMessage?.id === AUTO_THOUGHT_ID; if (!isCompletedStream && lastReasoningMessageIsAutoAppended && lastMessage) { const lastMessageId = lastMessage?.id!; diff --git a/web/src/context/Chats/NewChatProvider/useChatUpdateMessage.ts b/web/src/context/Chats/NewChatProvider/useChatUpdateMessage.ts index ee2e0cde1..e3ceaf053 100644 --- a/web/src/context/Chats/NewChatProvider/useChatUpdateMessage.ts +++ b/web/src/context/Chats/NewChatProvider/useChatUpdateMessage.ts @@ -1,7 +1,7 @@ import { useMemoizedFn } from 'ahooks'; import { useBusterChatContextSelector } from '../ChatProvider'; import { useBusterWebSocket } from '@/context/BusterWebSocket'; -import { BusterChat } from '@/api/asset_interfaces'; +import { BusterChat, BusterChatMessage_text } from '@/api/asset_interfaces'; import { ChatEvent_GeneratingReasoningMessage, ChatEvent_GeneratingResponseMessage, @@ -11,6 +11,7 @@ import { updateChatToIChat } from '@/utils/chat'; import { useAutoAppendThought } from './useAutoAppendThought'; import { useAppLayoutContextSelector } from '@/context/BusterAppLayout'; import { BusterRoutes } from '@/routes'; +import last from 'lodash/last'; export const useChatUpdateMessage = () => { const busterSocket = useBusterWebSocket(); @@ -40,12 +41,27 @@ export const useChatUpdateMessage = () => { (d: ChatEvent_GeneratingResponseMessage) => { const { message_id, response_message } = d; const currentResponseMessages = getChatMessageMemoized(message_id)?.response_messages ?? []; - const isNewMessage = !currentResponseMessages.some(({ id }) => id === message_id); + const responseMessageId = response_message.id; + const foundResponseMessage = currentResponseMessages.find( + ({ id }) => id === responseMessageId + ); + const isNewMessage = !foundResponseMessage; + + if (response_message.type === 'text') { + const existingMessage = (foundResponseMessage as BusterChatMessage_text)?.message || ''; + const isStreaming = !!response_message.message_chunk; + if (isStreaming) { + response_message.message = existingMessage + response_message.message_chunk; + } + } + onUpdateChatMessage({ id: message_id, response_messages: isNewMessage ? [...currentResponseMessages, response_message] - : currentResponseMessages.map((rm) => (rm.id === message_id ? response_message : rm)) + : currentResponseMessages.map((rm) => + rm.id === responseMessageId ? response_message : rm + ) }); } ); @@ -68,7 +84,7 @@ export const useChatUpdateMessage = () => { ); const completeChatCallback = useMemoizedFn((d: BusterChat) => { - const { iChat, iChatMessages } = updateChatToIChat(d); + const { iChat, iChatMessages } = updateChatToIChat(d, false); onBulkSetChatMessages(iChatMessages); onUpdateChat(iChat); }); @@ -81,10 +97,9 @@ export const useChatUpdateMessage = () => { }); const initializeChatCallback = useMemoizedFn((d: BusterChat) => { - const { iChat, iChatMessages } = updateChatToIChat(d); + const { iChat, iChatMessages } = updateChatToIChat(d, true); onBulkSetChatMessages(iChatMessages); onUpdateChat(iChat); - onChangePage({ route: BusterRoutes.APP_CHAT_ID, chatId: iChat.id diff --git a/web/src/utils/chat.ts b/web/src/utils/chat.ts index c93af729e..daf7d4a71 100644 --- a/web/src/utils/chat.ts +++ b/web/src/utils/chat.ts @@ -35,9 +35,12 @@ const chatMessageUpgrader = ( ); }; -export const updateChatToIChat = (chat: BusterChat) => { +export const updateChatToIChat = (chat: BusterChat, isNewChat: boolean) => { const iChat = chatUpgrader(chat); - const iChatMessages = chatMessageUpgrader(chat.messages); + const iChatMessages = chatMessageUpgrader(chat.messages, { + isCompletedStream: !isNewChat, + messageId: chat.messages[0].id + }); return { iChat, iChatMessages