From 0728efe71aab4c49618f309ebd1d8b1c21523916 Mon Sep 17 00:00:00 2001 From: Nate Kelley Date: Thu, 6 Mar 2025 09:35:15 -0700 Subject: [PATCH] chat stream message --- .../api/buster_rest/chats/queryRequests.ts | 2 + .../StreamingMessageCode.tsx | 2 +- ...AppendThought.ts => useBlackBoxMessage.ts} | 59 +++---- .../NewChatProvider/useChatStreamMessage.ts | 156 +++++++----------- web/src/context/Chats/interfaces.ts | 2 +- .../ReasoningMessageFile.tsx | 6 +- web/src/lib/chat.ts | 3 +- 7 files changed, 94 insertions(+), 136 deletions(-) rename web/src/context/Chats/NewChatProvider/{useAutoAppendThought.ts => useBlackBoxMessage.ts} (71%) diff --git a/web/src/api/buster_rest/chats/queryRequests.ts b/web/src/api/buster_rest/chats/queryRequests.ts index 759e9151e..927b896ca 100644 --- a/web/src/api/buster_rest/chats/queryRequests.ts +++ b/web/src/api/buster_rest/chats/queryRequests.ts @@ -40,6 +40,7 @@ export const prefetchGetListChats = async ( export const useGetChat = (params: Parameters[0]) => { const queryFn = useMemoizedFn(async () => { return await getChat(params).then((chat) => { + console.log('TODO move this to put message in a better spot'); return updateChatToIChat(chat, true).iChat; }); }); @@ -67,6 +68,7 @@ export const prefetchGetChat = async ( ...queryKeys.chatsGetChat(params.id), queryFn: async () => { return await getChat_server(params).then((chat) => { + console.log('TODO move this to put message in a better spot'); return updateChatToIChat(chat, true).iChat; }); } diff --git a/web/src/components/ui/streaming/StreamingMessageCode/StreamingMessageCode.tsx b/web/src/components/ui/streaming/StreamingMessageCode/StreamingMessageCode.tsx index 3b4cf0cbc..5b81e8969 100644 --- a/web/src/components/ui/streaming/StreamingMessageCode/StreamingMessageCode.tsx +++ b/web/src/components/ui/streaming/StreamingMessageCode/StreamingMessageCode.tsx @@ -51,7 +51,7 @@ export const StreamingMessageCode: React.FC< version_id, buttons }) => { - const showLoader = status === 'loading'; + const showLoader = status === 'loading' && !isCompletedStream; const { text = '', modified } = file; const [lineSegments, setLineSegments] = useState([]); diff --git a/web/src/context/Chats/NewChatProvider/useAutoAppendThought.ts b/web/src/context/Chats/NewChatProvider/useBlackBoxMessage.ts similarity index 71% rename from web/src/context/Chats/NewChatProvider/useAutoAppendThought.ts rename to web/src/context/Chats/NewChatProvider/useBlackBoxMessage.ts index 514e59b92..c16600557 100644 --- a/web/src/context/Chats/NewChatProvider/useAutoAppendThought.ts +++ b/web/src/context/Chats/NewChatProvider/useBlackBoxMessage.ts @@ -5,39 +5,39 @@ import type { import { useMemoizedFn } from 'ahooks'; import sample from 'lodash/sample'; import { useBusterChatContextSelector } from '../ChatProvider'; +import random from 'lodash/random'; +import last from 'lodash/last'; +import { timeout } from '@/lib/timeout'; +import { useState } from 'react'; + +export const useBlackBoxMessage = () => { + const [boxBoxMessages, setBoxBoxMessages] = useState>({}); -export const useAutoAppendThought = () => { const onUpdateChatMessage = useBusterChatContextSelector((x) => x.onUpdateChatMessage); const getChatMessageMemoized = useBusterChatContextSelector((x) => x.getChatMessageMemoized); - const removeAutoThoughts = useMemoizedFn( - (reasoningMessages: BusterChatMessageReasoning[]): BusterChatMessageReasoning[] => { - return reasoningMessages.filter((rm) => rm.id !== AUTO_THOUGHT_ID); - } - ); + const removeAutoThoughts = useMemoizedFn(() => { + // return reasoningMessages.filter((rm) => rm.id !== AUTO_THOUGHT_ID); + }); - const autoAppendThought = useMemoizedFn( - ( - reasoningMessages: BusterChatMessageReasoning[], - chatId: string - ): BusterChatMessageReasoning[] => { - const lastReasoningMessage = reasoningMessages[reasoningMessages.length - 1]; - const lastMessageIsCompleted = - !lastReasoningMessage || lastReasoningMessage?.status === 'completed'; + const autoAppendThought = useMemoizedFn(({ messageId }: { messageId: string }) => { + // + // const lastReasoningMessage = reasoningMessages[reasoningMessages.length - 1]; + // const lastMessageIsCompleted = + // !lastReasoningMessage || lastReasoningMessage?.status === 'completed'; - if (lastMessageIsCompleted) { - _loopAutoThought(chatId); + // if (lastMessageIsCompleted) { + // _loopAutoThought(chatId); - return [...reasoningMessages, createAutoThought()]; - } + // return [...reasoningMessages, createAutoThought()]; + // } - return removeAutoThoughts(reasoningMessages); - } - ); + return removeAutoThoughts(); + }); const _loopAutoThought = useMemoizedFn(async (chatId: string) => { - // const randomDelay = random(3000, 5000); - // await timeout(randomDelay); + const randomDelay = random(3000, 5000); + await timeout(randomDelay); // const chatMessages = getChatMessagesMemoized(chatId); // const lastMessage = last(chatMessages); // const isCompletedStream = !!lastMessage?.isCompletedStream; @@ -56,7 +56,6 @@ export const useAutoAppendThought = () => { // isCompletedStream: false // }); // _loopAutoThought(chatId); - // } }); return { autoAppendThought, removeAutoThoughts }; @@ -69,18 +68,6 @@ const getRandomThought = (currentThought?: string): string => { return sample(thoughts) ?? DEFAULT_THOUGHTS[0]; }; -const AUTO_THOUGHT_ID = 'stub-thought-id'; -const createAutoThought = (currentThought?: string): BusterChatMessageReasoning_pills => { - return { - id: AUTO_THOUGHT_ID, - type: 'pills', - title: getRandomThought(currentThought), - secondary_title: '', - pill_containers: [], - status: 'loading' - }; -}; - const DEFAULT_THOUGHTS = [ 'Thinking through next steps...', 'Looking through context...', diff --git a/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts b/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts index ad8f12f8a..538e1262a 100644 --- a/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts +++ b/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts @@ -4,9 +4,7 @@ import type { BusterChat, BusterChatMessageReasoning_files, BusterChatMessageReasoning_text, - BusterChatMessageReasoning_pills, BusterChatResponseMessage_text, - BusterChatMessageResponse, BusterChatMessageReasoning_file } from '@/api/asset_interfaces'; import type { @@ -15,7 +13,7 @@ import type { ChatEvent_GeneratingTitle } from '@/api/buster_socket/chats'; import { updateChatToIChat } from '@/lib/chat'; -import { useAutoAppendThought } from './useAutoAppendThought'; +import { useBlackBoxMessage } from './useBlackBoxMessage'; import { useAppLayoutContextSelector } from '@/context/BusterAppLayout'; import { BusterRoutes } from '@/routes'; import { useSocketQueryOn } from '@/api/buster_socket_query'; @@ -31,13 +29,16 @@ export const useChatStreamMessage = () => { const onChangePage = useAppLayoutContextSelector((x) => x.onChangePage); const onUpdateChat = useBusterChatContextSelector((x) => x.onUpdateChat); const onUpdateChatMessage = useBusterChatContextSelector((x) => x.onUpdateChatMessage); - const chatRef = useRef>>({}); + const chatRef = useRef>({}); + const chatRefMessages = useRef>({}); const [isPending, startTransition] = useTransition(); + const { autoAppendThought } = useBlackBoxMessage(); + const onUpdateChatMessageTransition = useMemoizedFn( - (chatMessage: Parameters[0], chatId: string) => { - const currentChatMessage = chatRef.current[chatId]?.messages?.[chatMessage.id]; - const iChatMessage = create(currentChatMessage, (draft) => { + (chatMessage: Parameters[0]) => { + const currentChatMessage = chatRefMessages.current[chatMessage.id]; + const iChatMessage: IBusterChatMessage = create(currentChatMessage, (draft) => { Object.assign(draft || {}, chatMessage); })!; @@ -49,32 +50,14 @@ export const useChatStreamMessage = () => { } ); - const { autoAppendThought } = useAutoAppendThought(); - const initializeOrUpdateMessage = useMemoizedFn( - ( - chatId: string, - messageId: string, - updateFn: (draft: Record>) => void - ) => { - chatRef.current = create(chatRef.current, (draft) => { - if (!draft[chatId]) draft[chatId] = {}; - if (!draft[chatId].messages) draft[chatId].messages = {}; - if (!draft[chatId].messages[messageId]) { - draft[chatId].messages[messageId] = { - id: messageId, - request_message: null, - response_message_ids: [], - response_messages: {}, - reasoning_message_ids: [], - reasoning_messages: {}, - created_at: new Date().toISOString(), - final_reasoning_message: null - }; - } - + (messageId: string, updateFn: (draft: IBusterChatMessage) => void) => { + const currentMessage = chatRefMessages.current[messageId]; + const updatedMessage = create(currentMessage || {}, (draft) => { updateFn(draft); }); + chatRefMessages.current[messageId] = updatedMessage; + onUpdateChatMessage(updatedMessage); } ); @@ -84,6 +67,7 @@ export const useChatStreamMessage = () => { const options = queryKeys.chatsMessages(message.id); const queryKey = options.queryKey; queryClient.setQueryData(queryKey, message); + chatRefMessages.current[message.id] = message; } } ); @@ -92,7 +76,6 @@ export const useChatStreamMessage = () => { const { iChat, iChatMessages } = updateChatToIChat(d, false); chatRef.current = create(chatRef.current, (draft) => { draft[iChat.id] = iChat; - draft[iChat.id].messages = iChatMessages; }); normalizeChatMessage(iChatMessages); onUpdateChat(iChat); @@ -109,8 +92,8 @@ export const useChatStreamMessage = () => { const { iChat, iChatMessages } = updateChatToIChat(d, true); chatRef.current = create(chatRef.current, (draft) => { draft[iChat.id] = iChat; - draft[iChat.id].messages = iChatMessages; }); + normalizeChatMessage(iChatMessages); onUpdateChat(iChat); onChangePage({ @@ -141,8 +124,7 @@ export const useChatStreamMessage = () => { const currentTitle = chatRef.current[chat_id]?.title || ''; const newTitle = isCompleted ? title : currentTitle + title_chunk; chatRef.current = create(chatRef.current, (draft) => { - if (!draft[chat_id]) draft[chat_id] = {}; - draft[chat_id].title = newTitle; + if (newTitle) draft[chat_id].title = newTitle; }); onUpdateChat({ id: chat_id, @@ -158,28 +140,29 @@ export const useChatStreamMessage = () => { const responseMessageId = response_message.id; const existingMessage = - chatRef.current[chat_id]?.messages?.[message_id]?.response_messages?.[responseMessageId]; + chatRefMessages.current[message_id]?.response_messages?.[responseMessageId]; const isNewMessage = !existingMessage; if (isNewMessage) { - initializeOrUpdateMessage(chat_id, message_id, (draft) => { - const chat = draft[chat_id]; - if (!chat?.messages?.[message_id]) return; - if (!chat.messages[message_id].response_messages) - chat.messages[message_id].response_messages = {}; - chat.messages[message_id].response_messages[responseMessageId] = response_message; - chat.messages[message_id].response_message_ids.push(responseMessageId); + initializeOrUpdateMessage(message_id, (draft) => { + if (!draft.response_messages) { + draft.response_messages = {}; + } + draft.response_messages[responseMessageId] = response_message; + if (!draft.response_message_ids) { + draft.response_message_ids = []; + } + draft.response_message_ids.push(responseMessageId); }); } if (response_message.type === 'text') { const existingResponseMessageText = existingMessage as BusterChatResponseMessage_text; const isStreaming = - response_message.message_chunk !== undefined || response_message.message_chunk !== null; + response_message.message_chunk !== undefined && response_message.message_chunk !== null; - initializeOrUpdateMessage(chat_id, message_id, (draft) => { - const responseMessage = - draft[chat_id]?.messages?.[message_id]?.response_messages?.[responseMessageId]; + initializeOrUpdateMessage(message_id, (draft) => { + const responseMessage = draft.response_messages?.[responseMessageId]; if (!responseMessage) return; const messageText = responseMessage as BusterChatMessageReasoning_text; Object.assign(messageText, { @@ -193,18 +176,12 @@ export const useChatStreamMessage = () => { }); } - const response_messages = chatRef.current[chat_id]?.messages?.[message_id]?.response_messages; - const response_message_ids = - chatRef.current[chat_id]?.messages?.[message_id]?.response_message_ids; - - onUpdateChatMessageTransition( - { - id: message_id, - response_messages, - response_message_ids - }, - chat_id - ); + const currentMessage = chatRefMessages.current[message_id]; + onUpdateChatMessageTransition({ + id: message_id, + response_messages: currentMessage?.response_messages, + response_message_ids: currentMessage?.response_message_ids + }); } ); @@ -214,17 +191,19 @@ export const useChatStreamMessage = () => { const reasoningMessageId = reasoning.id; const existingMessage = - chatRef.current[chat_id]?.messages?.[message_id]?.reasoning_messages?.[reasoningMessageId]; + chatRefMessages.current[message_id]?.reasoning_messages?.[reasoningMessageId]; const isNewMessage = !existingMessage; if (isNewMessage) { - initializeOrUpdateMessage(chat_id, message_id, (draft) => { - const chat = draft[chat_id]; - if (!chat?.messages?.[message_id]) return; - if (!chat.messages[message_id].reasoning_messages) - chat.messages[message_id].reasoning_messages = {}; - chat.messages[message_id].reasoning_messages[reasoningMessageId] = reasoning; - chat.messages[message_id].reasoning_message_ids.push(reasoningMessageId); + initializeOrUpdateMessage(message_id, (draft) => { + if (!draft.reasoning_messages) { + draft.reasoning_messages = {}; + } + draft.reasoning_messages[reasoningMessageId] = reasoning; + if (!draft.reasoning_message_ids) { + draft.reasoning_message_ids = []; + } + draft.reasoning_message_ids.push(reasoningMessageId); }); } @@ -234,9 +213,8 @@ export const useChatStreamMessage = () => { const isStreaming = reasoning.message_chunk !== null || reasoning.message_chunk !== undefined; - initializeOrUpdateMessage(chat_id, message_id, (draft) => { - const reasoningMessage = - draft[chat_id]?.messages?.[message_id]?.reasoning_messages?.[reasoningMessageId]; + initializeOrUpdateMessage(message_id, (draft) => { + const reasoningMessage = draft.reasoning_messages?.[reasoningMessageId]; if (!reasoningMessage) return; const messageText = reasoningMessage as BusterChatMessageReasoning_text; @@ -254,14 +232,12 @@ export const useChatStreamMessage = () => { case 'files': { const existingReasoningMessageFiles = existingMessage as BusterChatMessageReasoning_files; - initializeOrUpdateMessage(chat_id, message_id, (draft) => { - const chat = draft[chat_id]; - if (!chat?.messages?.[message_id]?.reasoning_messages?.[reasoningMessageId]) return; + initializeOrUpdateMessage(message_id, (draft) => { + const reasoningMessage = draft.reasoning_messages?.[reasoningMessageId]; + if (!reasoningMessage) return; const messageFiles = create( - chat.messages[message_id].reasoning_messages[ - reasoningMessageId - ] as BusterChatMessageReasoning_files, + reasoningMessage as BusterChatMessageReasoning_files, (draft) => { draft.file_ids = existingReasoningMessageFiles?.file_ids || []; @@ -305,15 +281,14 @@ export const useChatStreamMessage = () => { } ); - chat.messages[message_id].reasoning_messages[reasoningMessageId] = messageFiles; + draft.reasoning_messages[reasoningMessageId] = messageFiles; }); break; } case 'pills': { - initializeOrUpdateMessage(chat_id, message_id, (draft) => { - if (!draft[chat_id]?.messages?.[message_id]?.reasoning_messages?.[reasoningMessageId]) - return; - draft[chat_id].messages[message_id].reasoning_messages[reasoningMessageId]! = reasoning; + initializeOrUpdateMessage(message_id, (draft) => { + if (!draft.reasoning_messages?.[reasoningMessageId]) return; + draft.reasoning_messages[reasoningMessageId] = reasoning; }); break; @@ -324,20 +299,13 @@ export const useChatStreamMessage = () => { } } - const reasoning_messages = - chatRef.current[chat_id]?.messages?.[message_id]?.reasoning_messages; - const reasoning_message_ids = - chatRef.current[chat_id]?.messages?.[message_id]?.reasoning_message_ids; - - onUpdateChatMessageTransition( - { - id: message_id, - reasoning_messages, - reasoning_message_ids, - isCompletedStream: false - }, - chat_id - ); + const currentMessage = chatRefMessages.current[message_id]; + onUpdateChatMessageTransition({ + id: message_id, + reasoning_messages: currentMessage?.reasoning_messages, + reasoning_message_ids: currentMessage?.reasoning_message_ids, + isCompletedStream: false + }); } ); diff --git a/web/src/context/Chats/interfaces.ts b/web/src/context/Chats/interfaces.ts index c18050deb..a18fead8d 100644 --- a/web/src/context/Chats/interfaces.ts +++ b/web/src/context/Chats/interfaces.ts @@ -1,6 +1,6 @@ import type { BusterChat, BusterChatMessage } from '@/api/asset_interfaces'; -export interface IBusterChat extends BusterChat { +export interface IBusterChat extends Omit { isNewChat: boolean; } diff --git a/web/src/controllers/ReasoningController/ReasoningMessages/ReasoningMessage_Files/ReasoningMessageFile.tsx b/web/src/controllers/ReasoningController/ReasoningMessages/ReasoningMessage_Files/ReasoningMessageFile.tsx index 293548975..6e2053850 100644 --- a/web/src/controllers/ReasoningController/ReasoningMessages/ReasoningMessage_Files/ReasoningMessageFile.tsx +++ b/web/src/controllers/ReasoningController/ReasoningMessages/ReasoningMessage_Files/ReasoningMessageFile.tsx @@ -52,19 +52,19 @@ const StreamingMessageStatus = React.memo( const content = useMemo(() => { if (status === 'loading') return ( - + Running SQL... ); if (status === 'completed') return ( - + Completed ); if (status === 'failed') return ( - + Failed ); diff --git a/web/src/lib/chat.ts b/web/src/lib/chat.ts index 9721681d1..b9b314123 100644 --- a/web/src/lib/chat.ts +++ b/web/src/lib/chat.ts @@ -1,10 +1,11 @@ import type { BusterChat, BusterChatMessage } from '@/api/asset_interfaces'; import type { IBusterChat, IBusterChatMessage } from '@/context/Chats/interfaces'; import { create } from 'mutative'; +import omit from 'lodash/omit'; const chatUpgrader = (chat: BusterChat, { isNewChat }: { isNewChat: boolean }): IBusterChat => { return { - ...chat, + ...omit(chat, 'messages'), isNewChat }; };