From 7a349da981ff93056b4dfdd2d9bc66cdcbdd0b54 Mon Sep 17 00:00:00 2001 From: Nate Kelley Date: Mon, 10 Feb 2025 16:43:38 -0700 Subject: [PATCH] added functions for streaming content --- .../chat/chatMessageInterfaces.ts | 1 + .../chat/chatProgressInterfaces.ts | 10 +- .../api/buster_socket/chats/chatRequests.ts | 5 +- .../api/buster_socket/chats/chatResponses.ts | 33 ++++- .../buster_socket/chats/eventInterfaces.ts | 43 +++--- web/src/api/buster_socket/chats/index.ts | 1 + .../ReasoningFileButtons.tsx | 4 +- .../ReasoningMessage_File.tsx | 2 + .../ChatContent/ChatInput/useChatInputFlow.ts | 6 +- .../ChatLayout/hooks/useAutoSetLayout.tsx | 1 - .../Chats/ChatProvider/useChatSelectors.ts | 21 ++- .../ChatProvider/useChatSubscriptions.ts | 124 +--------------- .../Chats/NewChatProvider/NewChatProvider.tsx | 59 ++++++-- .../NewChatProvider/useChatUpdateMessage.ts | 133 ++++++++++++++++++ 14 files changed, 268 insertions(+), 175 deletions(-) create mode 100644 web/src/context/Chats/NewChatProvider/useChatUpdateMessage.ts diff --git a/web/src/api/asset_interfaces/chat/chatMessageInterfaces.ts b/web/src/api/asset_interfaces/chat/chatMessageInterfaces.ts index ec116aafb..edbadeaea 100644 --- a/web/src/api/asset_interfaces/chat/chatMessageInterfaces.ts +++ b/web/src/api/asset_interfaces/chat/chatMessageInterfaces.ts @@ -80,6 +80,7 @@ export type BusterChatMessageReasoning_file = { version_number: number; version_id: string; status: 'loading' | 'completed' | 'failed'; + //when we are streaming, the whole file will always be streamed back, not chunks file?: { text: string; line_number: number; diff --git a/web/src/api/asset_interfaces/chat/chatProgressInterfaces.ts b/web/src/api/asset_interfaces/chat/chatProgressInterfaces.ts index d14f1f1a5..723e866c9 100644 --- a/web/src/api/asset_interfaces/chat/chatProgressInterfaces.ts +++ b/web/src/api/asset_interfaces/chat/chatProgressInterfaces.ts @@ -1,4 +1,8 @@ -import type { BusterChatMessage, BusterChatMessageResponse } from './chatMessageInterfaces'; +import type { + BusterChatMessage, + BusterChatMessageReasoning, + BusterChatMessageResponse +} from './chatMessageInterfaces'; enum BusterChatStepProgress { IN_PROGRESS = 'in_progress', @@ -20,6 +24,10 @@ export type ChatPost_generatingMessage = { resposnse_message: BusterChatMessageResponse; } & BusterChatStepBase; +export type ChatPost_generatingReasoning = { + reasoning: BusterChatMessageReasoning; +} & BusterChatStepBase; + export type ChatPost_complete = { message: BusterChatMessage; } & BusterChatStepBase; diff --git a/web/src/api/buster_socket/chats/chatRequests.ts b/web/src/api/buster_socket/chats/chatRequests.ts index 8b8492d07..65568077d 100644 --- a/web/src/api/buster_socket/chats/chatRequests.ts +++ b/web/src/api/buster_socket/chats/chatRequests.ts @@ -9,7 +9,7 @@ export type ChatCreateNewChat = BusterSocketRequestBase< '/chats/post', { /** The ID of the dataset to associate with the chat. Null if no dataset is associated */ - dataset_id: string | null; + dataset_id?: string | null; /** The initial message or prompt to start the chat conversation */ prompt: string; /** Optional ID of an existing chat for follow-up messages. Null for new chats */ @@ -154,4 +154,5 @@ export type ChatEmits = | ChatDeleteChat | ChatUpdateChat | ChatsSearch - | ChatsDuplicateChat; + | ChatsDuplicateChat + | ChatStopChat; diff --git a/web/src/api/buster_socket/chats/chatResponses.ts b/web/src/api/buster_socket/chats/chatResponses.ts index 287a0b2c1..32390c731 100644 --- a/web/src/api/buster_socket/chats/chatResponses.ts +++ b/web/src/api/buster_socket/chats/chatResponses.ts @@ -1,6 +1,10 @@ import type { RustApiError } from '../../buster_rest/errors'; import type { BusterChat, BusterChatListItem } from '../../asset_interfaces/chat'; -import { ChatEvent_GeneratingMessage, ChatEvent_GeneratingTitle } from './eventInterfaces'; +import { + ChatEvent_GeneratingReasoningMessage, + ChatEvent_GeneratingResponseMessage, + ChatEvent_GeneratingTitle +} from './eventInterfaces'; export enum ChatsResponses { '/chats/list:getChatsList' = '/chats/list:getChatsList', @@ -8,7 +12,10 @@ export enum ChatsResponses { '/chats/get:getChat' = '/chats/get:getChat', '/chats/get:getChatAsset' = '/chats/get:getChatAsset', '/chats/post:initializeChat' = '/chats/post:initializeChat', - '/chats/post:generatingTitle' = '/chats/post:generatingTitle' + '/chats/post:generatingTitle' = '/chats/post:generatingTitle', + '/chats/post:generatingResponseMessage' = '/chats/post:generatingResponseMessage', + '/chats/post:generatingReasoningMessage' = '/chats/post:generatingReasoningMessage', + '/chats/post:complete' = '/chats/post:complete' } export type ChatList_getChatsList = { @@ -56,9 +63,21 @@ export type ChatPost_generatingTitle = { onError?: (d: unknown | RustApiError) => void; }; -export type ChatPost_generatingMessage = { - route: '/chats/post:generatingMessage'; - callback: (d: ChatEvent_GeneratingMessage) => void; +export type ChatPost_generatingResponseMessage = { + route: '/chats/post:generatingResponseMessage'; + callback: (d: ChatEvent_GeneratingResponseMessage) => void; + onError?: (d: unknown | RustApiError) => void; +}; + +export type ChatPost_generatingReasoningMessage = { + route: '/chats/post:generatingReasoningMessage'; + callback: (d: ChatEvent_GeneratingReasoningMessage) => void; + onError?: (d: unknown | RustApiError) => void; +}; + +export type ChatPost_complete = { + route: '/chats/post:complete'; + callback: (d: BusterChat) => void; onError?: (d: unknown | RustApiError) => void; }; @@ -71,4 +90,6 @@ export type ChatResponseTypes = | Chat_getChatAsset | ChatPost_initializeChat | ChatPost_generatingTitle - | ChatPost_generatingMessage; + | ChatPost_generatingResponseMessage + | ChatPost_generatingReasoningMessage + | ChatPost_complete; diff --git a/web/src/api/buster_socket/chats/eventInterfaces.ts b/web/src/api/buster_socket/chats/eventInterfaces.ts index 0b730bd9d..77d974799 100644 --- a/web/src/api/buster_socket/chats/eventInterfaces.ts +++ b/web/src/api/buster_socket/chats/eventInterfaces.ts @@ -1,37 +1,34 @@ -import type { BusterChatMessageResponse } from '@/api/asset_interfaces/chat'; +import type { + BusterChatMessageReasoning, + BusterChatMessageResponse +} from '@/api/asset_interfaces/chat'; import type { EventBase } from '../base_interfaces'; -/** - * Chat event interface for title generation process. - * - * @remarks - * This interface extends EventBase to include properties specific to - * the title generation process in chat events. - * - * @public - */ export type ChatEvent_GeneratingTitle = { /** The complete generated title when available */ title?: string; /** A partial chunk of the title during the generation process */ title_chunk?: string; + /** The ID of the chat that the title belongs to */ + chat_id: string; } & EventBase; -/** - * Chat event interface for message generation process. - * - * @remarks - * This interface extends EventBase and handles the message generation process. - * When new messages are received, they are appended to the response_messages array. - * If a message with a matching ID already exists, it will be updated instead of - * creating a duplicate entry. - * - * @public - */ -export type ChatEvent_GeneratingMessage = { +export type ChatEvent_GeneratingResponseMessage = { // We will append each incoming message to the response_messages array // If the message id is already found in the array, we will update the message with the new data // This will happen when we need to "hide" a message /** The chat message response containing the generated content */ - message: BusterChatMessageResponse; + response_message: BusterChatMessageResponse; + /** The ID of the chat that the response message belongs to */ + chat_id: string; + /** The ID of the message that the response message belongs to */ + message_id: string; +} & EventBase; + +export type ChatEvent_GeneratingReasoningMessage = { + reasoning: BusterChatMessageReasoning; + /** The ID of the chat that the reasoning message belongs to */ + chat_id: string; + /** The ID of the message that the reasoning message belongs to */ + message_id: string; } & EventBase; diff --git a/web/src/api/buster_socket/chats/index.ts b/web/src/api/buster_socket/chats/index.ts index 1f6ffab23..5f502ad72 100644 --- a/web/src/api/buster_socket/chats/index.ts +++ b/web/src/api/buster_socket/chats/index.ts @@ -1,2 +1,3 @@ export * from './chatRequests'; export * from './chatResponses'; +export * from './eventInterfaces'; diff --git a/web/src/app/app/_controllers/ReasoningController/ReasoningMessages/ReasoningMessage_File/ReasoningFileButtons.tsx b/web/src/app/app/_controllers/ReasoningController/ReasoningMessages/ReasoningMessage_File/ReasoningFileButtons.tsx index 32a9d5a17..1b3acef77 100644 --- a/web/src/app/app/_controllers/ReasoningController/ReasoningMessages/ReasoningMessage_File/ReasoningFileButtons.tsx +++ b/web/src/app/app/_controllers/ReasoningController/ReasoningMessages/ReasoningMessage_File/ReasoningFileButtons.tsx @@ -18,8 +18,6 @@ export const ReasoningFileButtons = React.memo( chatId: string; isCompletedStream: boolean; }) => { - if (!isCompletedStream) return null; - const onSetSelectedFile = useChatLayoutContextSelector((state) => state.onSetSelectedFile); const link = createChatAssetRoute({ @@ -37,6 +35,8 @@ export const ReasoningFileButtons = React.memo( }); }); + if (!isCompletedStream) return null; + return (
diff --git a/web/src/app/app/_controllers/ReasoningController/ReasoningMessages/ReasoningMessage_File/ReasoningMessage_File.tsx b/web/src/app/app/_controllers/ReasoningController/ReasoningMessages/ReasoningMessage_File/ReasoningMessage_File.tsx index 7889c5226..ea6a7d2a3 100644 --- a/web/src/app/app/_controllers/ReasoningController/ReasoningMessages/ReasoningMessage_File/ReasoningMessage_File.tsx +++ b/web/src/app/app/_controllers/ReasoningController/ReasoningMessages/ReasoningMessage_File/ReasoningMessage_File.tsx @@ -125,3 +125,5 @@ const MemoizedSyntaxHighlighter = React.memo( ); } ); + +MemoizedSyntaxHighlighter.displayName = 'MemoizedSyntaxHighlighter'; diff --git a/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatInput/useChatInputFlow.ts b/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatInput/useChatInputFlow.ts index f79b7b678..7d8544936 100644 --- a/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatInput/useChatInputFlow.ts +++ b/web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatInput/useChatInputFlow.ts @@ -37,16 +37,16 @@ export const useChatInputFlow = ({ }, [hasChat, selectedFileType, selectedFileId]); const onSubmitPreflight = useMemoizedFn(async () => { - if (disableSendButton) return; + if (disableSendButton || !chatId || !currentMessageId) return; if (loading) { - onStopChat({ chatId: chatId! }); + onStopChat({ chatId: chatId!, messageId: currentMessageId }); return; } switch (flow) { case 'followup-chat': - await onFollowUpChat({ prompt: inputValue, messageId: currentMessageId! }); + await onFollowUpChat({ prompt: inputValue, chatId: chatId }); break; case 'followup-metric': diff --git a/web/src/app/app/_layouts/ChatLayout/hooks/useAutoSetLayout.tsx b/web/src/app/app/_layouts/ChatLayout/hooks/useAutoSetLayout.tsx index 913973431..30ac8f32b 100644 --- a/web/src/app/app/_layouts/ChatLayout/hooks/useAutoSetLayout.tsx +++ b/web/src/app/app/_layouts/ChatLayout/hooks/useAutoSetLayout.tsx @@ -16,7 +16,6 @@ export const useAutoSetLayout = ({ }, [defaultSelectedLayout]); const collapseDirection: 'left' | 'right' = useMemo(() => { - console.log(defaultSelectedLayout); return defaultSelectedLayout === 'file' ? 'left' : 'right'; }, [defaultSelectedLayout]); diff --git a/web/src/context/Chats/ChatProvider/useChatSelectors.ts b/web/src/context/Chats/ChatProvider/useChatSelectors.ts index 489578dd1..03120e7dc 100644 --- a/web/src/context/Chats/ChatProvider/useChatSelectors.ts +++ b/web/src/context/Chats/ChatProvider/useChatSelectors.ts @@ -1,5 +1,6 @@ import { MutableRefObject, useCallback } from 'react'; import { IBusterChat, IBusterChatMessage } from '../interfaces'; +import { useMemoizedFn } from 'ahooks'; export const useChatSelectors = ({ isPending, @@ -10,6 +11,10 @@ export const useChatSelectors = ({ chatsRef: MutableRefObject>; chatsMessagesRef: MutableRefObject>; }) => { + const getChatMemoized = useMemoizedFn((chatId: string) => { + return chatsRef.current[chatId]; + }); + const getChatMessages = useCallback( (chatId: string): IBusterChatMessage[] => { const chatMessageIds = chatsRef.current[chatId].messages || []; @@ -25,5 +30,19 @@ export const useChatSelectors = ({ [chatsMessagesRef, isPending] ); - return { getChatMessages, getChatMessage }; + const getChatMessagesMemoized = useMemoizedFn((chatId: string) => { + return getChatMessages(chatId); + }); + + const getChatMessageMemoized = useMemoizedFn((messageId: string) => { + return getChatMessage(messageId); + }); + + return { + getChatMemoized, + getChatMessages, + getChatMessage, + getChatMessagesMemoized, + getChatMessageMemoized + }; }; diff --git a/web/src/context/Chats/ChatProvider/useChatSubscriptions.ts b/web/src/context/Chats/ChatProvider/useChatSubscriptions.ts index 029c44444..06307bee8 100644 --- a/web/src/context/Chats/ChatProvider/useChatSubscriptions.ts +++ b/web/src/context/Chats/ChatProvider/useChatSubscriptions.ts @@ -1,23 +1,10 @@ import { MutableRefObject } from 'react'; import { useBusterWebSocket } from '../../BusterWebSocket'; import { useMemoizedFn } from 'ahooks'; -import { - BusterChat, - BusterChatMessage, - BusterChatMessageReasoning_thought, - BusterChatMessageReasoning_file -} from '@/api/asset_interfaces'; +import type { BusterChat } from '@/api/asset_interfaces'; import { IBusterChat, IBusterChatMessage } from '../interfaces'; import { chatMessageUpgrader, chatUpgrader } from './helpers'; -import { - createMockResponseMessageFile, - createMockResponseMessageText, - createMockResponseMessageThought, - createMockReasoningMessageFile, - MOCK_CHAT -} from './MOCK_CHAT'; -import { useHotkeys } from 'react-hotkeys-hook'; -import { faker } from '@faker-js/faker'; +import { MOCK_CHAT } from './MOCK_CHAT'; export const useChatSubscriptions = ({ chatsRef, @@ -67,113 +54,6 @@ export const useChatSubscriptions = ({ // }); }); - useHotkeys('f', () => { - // Find the last chat message - const lastChatId = Object.keys(chatsRef.current)[Object.keys(chatsRef.current).length - 1]; - const lastChat = chatsRef.current[lastChatId]; - - if (!lastChat?.messages?.length) return; - - const lastMessageId = lastChat.messages[lastChat.messages.length - 1]; - const lastMessage = chatsMessagesRef.current[lastMessageId]; - - if (!lastMessage?.reasoning?.length) return; - - // Find the last reasoning file message - const lastReasoningFile = lastMessage.reasoning - .filter((r: { type: string }) => r.type === 'file') - .pop() as BusterChatMessageReasoning_file | undefined; - - if (!lastReasoningFile) return; - - // Create new file chunk - const newChunk = { - text: faker.lorem.sentence(), - line_number: (lastReasoningFile.file?.length || 0) + 1, - modified: true - }; - - // Create new reasoning file with appended chunk - const updatedReasoningFile = { - ...lastReasoningFile, - file: [...(lastReasoningFile.file || []), newChunk] - }; - - // Create new message with updated reasoning array - const updatedMessage = { - ...lastMessage, - reasoning: lastMessage.reasoning.map((r) => { - if (r.type === 'file' && r.id === lastReasoningFile.id) { - return updatedReasoningFile; - } - return r; - }) - }; - - // Update the refs with new object references - chatsMessagesRef.current = { - ...chatsMessagesRef.current, - [lastMessageId]: updatedMessage - }; - - chatsRef.current = { - ...chatsRef.current, - [lastChatId]: { - ...lastChat, - messages: [...lastChat.messages] - } - }; - - startTransition(() => { - // Force a re-render - chatsRef.current = { ...chatsRef.current }; - chatsMessagesRef.current = { ...chatsMessagesRef.current }; - }); - }); - - useHotkeys('y', () => { - // Find the last chat message - const lastChatId = Object.keys(chatsRef.current)[Object.keys(chatsRef.current).length - 1]; - const lastChat = chatsRef.current[lastChatId]; - - if (!lastChat?.messages?.length) return; - - const lastMessageId = lastChat.messages[lastChat.messages.length - 1]; - const lastMessage = chatsMessagesRef.current[lastMessageId]; - - if (!lastMessage) return; - lastMessage.isCompletedStream = false; - - // Create a new reasoning file message - const newReasoningFile = createMockReasoningMessageFile(); - - // Add the new reasoning file to the reasoning array - const updatedMessage = { - ...lastMessage, - reasoning: [...(lastMessage.reasoning || []), newReasoningFile] - }; - - // Update the refs with new object references - chatsMessagesRef.current = { - ...chatsMessagesRef.current, - [lastMessageId]: updatedMessage - }; - - chatsRef.current = { - ...chatsRef.current, - [lastChatId]: { - ...lastChat, - messages: [...lastChat.messages] - } - }; - - startTransition(() => { - // Force a re-render - chatsRef.current = { ...chatsRef.current }; - chatsMessagesRef.current = { ...chatsMessagesRef.current }; - }); - }); - return { unsubscribeFromChat, subscribeToChat diff --git a/web/src/context/Chats/NewChatProvider/NewChatProvider.tsx b/web/src/context/Chats/NewChatProvider/NewChatProvider.tsx index a49ff7c80..89bf8a014 100644 --- a/web/src/context/Chats/NewChatProvider/NewChatProvider.tsx +++ b/web/src/context/Chats/NewChatProvider/NewChatProvider.tsx @@ -6,8 +6,19 @@ import { } from '@fluentui/react-context-selector'; import { useMemoizedFn } from 'ahooks'; import type { BusterDatasetListItem, BusterSearchResult, FileType } from '@/api/asset_interfaces'; +import { useBusterWebSocket } from '@/context/BusterWebSocket'; +import { useChatUpdateMessage } from './useChatUpdateMessage'; export const useBusterNewChat = () => { + const busterSocket = useBusterWebSocket(); + + const { + completeChatCallback, + startListeningForChatProgress, + stopListeningForChatProgress, + stopChatCallback + } = useChatUpdateMessage(); + const onSelectSearchAsset = useMemoizedFn(async (asset: BusterSearchResult) => { console.log('select search asset'); await new Promise((resolve) => setTimeout(resolve, 1000)); @@ -25,13 +36,6 @@ export const useBusterNewChat = () => { } ); - const onFollowUpChat = useMemoizedFn( - async ({ prompt, messageId }: { prompt: string; messageId: string }) => { - console.log('follow up chat'); - await new Promise((resolve) => setTimeout(resolve, 1000)); - } - ); - const onReplaceMessageInChat = useMemoizedFn( async ({ prompt, messageId }: { prompt: string; messageId: string }) => { console.log('replace message in chat'); @@ -39,18 +43,45 @@ export const useBusterNewChat = () => { } ); - const onStopChat = useMemoizedFn(({ chatId }: { chatId: string }) => { - console.log('stop current chat'); - }); + const onFollowUpChat = useMemoizedFn( + async ({ prompt, chatId }: { prompt: string; chatId: string }) => { + startListeningForChatProgress(); + const result = await busterSocket.emitAndOnce({ + emitEvent: { + route: '/chats/post', + payload: { + dataset_id: null, + prompt, + chat_id: chatId + } + }, + responseEvent: { + route: '/chats/post:complete', + callback: completeChatCallback + } + }); - const onSetSelectedChatDataSource = useMemoizedFn((dataSource: BusterDatasetListItem | null) => { - // - }); + stopListeningForChatProgress(); + } + ); + + const onStopChat = useMemoizedFn( + ({ chatId, messageId }: { chatId: string; messageId: string }) => { + busterSocket.emit({ + route: '/chats/stop', + payload: { + id: chatId, + message_id: messageId + } + }); + stopListeningForChatProgress(); + stopChatCallback(chatId); + } + ); return { onStartNewChat, onSelectSearchAsset, - onSetSelectedChatDataSource, onFollowUpChat, onStartChatFromFile, onReplaceMessageInChat, diff --git a/web/src/context/Chats/NewChatProvider/useChatUpdateMessage.ts b/web/src/context/Chats/NewChatProvider/useChatUpdateMessage.ts new file mode 100644 index 000000000..5c2de4998 --- /dev/null +++ b/web/src/context/Chats/NewChatProvider/useChatUpdateMessage.ts @@ -0,0 +1,133 @@ +import { useMemoizedFn } from 'ahooks'; +import { useBusterChatContextSelector } from '../ChatProvider'; +import { useBusterWebSocket } from '@/context/BusterWebSocket'; +import { BusterChat } from '@/api/asset_interfaces'; +import { + ChatEvent_GeneratingReasoningMessage, + ChatEvent_GeneratingResponseMessage, + ChatEvent_GeneratingTitle +} from '@/api/buster_socket/chats'; + +export const useChatUpdateMessage = () => { + const busterSocket = useBusterWebSocket(); + const onUpdateChat = useBusterChatContextSelector((x) => x.onUpdateChat); + const getChatMemoized = useBusterChatContextSelector((x) => x.getChatMemoized); + const onUpdateChatMessage = useBusterChatContextSelector((x) => x.onUpdateChatMessage); + const getChatMessageMemoized = useBusterChatContextSelector((x) => x.getChatMessageMemoized); + + const _generatingTitleCallback = useMemoizedFn((d: ChatEvent_GeneratingTitle) => { + const { chat_id, title, title_chunk } = d; + const isCompleted = d.progress === 'completed'; + const currentTitle = getChatMemoized(chat_id)?.title; + const newTitle = isCompleted ? title : currentTitle + title_chunk; + onUpdateChat({ + ...getChatMemoized(chat_id), + title: newTitle + }); + }); + + const _generatingResponseMessageCallback = useMemoizedFn( + (d: ChatEvent_GeneratingResponseMessage) => { + const { message_id, response_message } = d; + const currentResponseMessages = getChatMessageMemoized(message_id)?.response_messages ?? []; + const isNewMessage = !currentResponseMessages.some(({ id }) => id === message_id); + onUpdateChatMessage({ + id: message_id, + response_messages: isNewMessage + ? [...currentResponseMessages, response_message] + : currentResponseMessages.map((rm) => (rm.id === message_id ? response_message : rm)) + }); + } + ); + + const _generatingReasoningMessageCallback = useMemoizedFn( + (d: ChatEvent_GeneratingReasoningMessage) => { + const { message_id, reasoning } = d; + const currentReasoning = getChatMessageMemoized(message_id)?.reasoning; + const isNewMessage = !currentReasoning?.some(({ id }) => id === message_id); + const updatedReasoning = isNewMessage + ? [...currentReasoning, reasoning] + : currentReasoning.map((rm) => (rm.id === message_id ? reasoning : rm)); + + onUpdateChatMessage({ + id: message_id, + reasoning: updatedReasoning + }); + } + ); + + const completeChatCallback = useMemoizedFn((d: BusterChat) => { + onUpdateChatMessage({ + ...d, + isCompletedStream: true + }); + }); + + const stopChatCallback = useMemoizedFn((chatId: string) => { + onUpdateChatMessage({ + id: chatId, + isCompletedStream: true + }); + }); + + const listenForGeneratingTitle = useMemoizedFn(() => { + busterSocket.on({ + route: '/chats/post:generatingTitle', + callback: _generatingTitleCallback + }); + }); + + const stopListeningForGeneratingTitle = useMemoizedFn(() => { + busterSocket.off({ + route: '/chats/post:generatingTitle', + callback: _generatingTitleCallback + }); + }); + + const listenForGeneratingResponseMessage = useMemoizedFn(() => { + busterSocket.on({ + route: '/chats/post:generatingResponseMessage', + callback: _generatingResponseMessageCallback + }); + }); + + const stopListeningForGeneratingResponseMessage = useMemoizedFn(() => { + busterSocket.off({ + route: '/chats/post:generatingResponseMessage', + callback: _generatingResponseMessageCallback + }); + }); + + const listenForGeneratingReasoningMessage = useMemoizedFn(() => { + busterSocket.on({ + route: '/chats/post:generatingReasoningMessage', + callback: _generatingReasoningMessageCallback + }); + }); + + const stopListeningForGeneratingReasoningMessage = useMemoizedFn(() => { + busterSocket.off({ + route: '/chats/post:generatingReasoningMessage', + callback: _generatingReasoningMessageCallback + }); + }); + + const startListeningForChatProgress = useMemoizedFn(() => { + listenForGeneratingTitle(); + listenForGeneratingResponseMessage(); + listenForGeneratingReasoningMessage(); + }); + + const stopListeningForChatProgress = useMemoizedFn(() => { + stopListeningForGeneratingTitle(); + stopListeningForGeneratingResponseMessage(); + stopListeningForGeneratingReasoningMessage(); + }); + + return { + completeChatCallback, + startListeningForChatProgress, + stopListeningForChatProgress, + stopChatCallback + }; +};