From f73385aa87a91f9152db05800b167a9187e7036f Mon Sep 17 00:00:00 2001 From: Nate Kelley Date: Thu, 6 Mar 2025 10:05:40 -0700 Subject: [PATCH] make response messages a helper function --- .../chatStreamMessageHelper.test.ts | 171 +++++++++++++++++- .../chatStreamMessageHelper.ts | 58 +++++- .../NewChatProvider/useChatStreamMessage.ts | 59 ++---- 3 files changed, 240 insertions(+), 48 deletions(-) diff --git a/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.test.ts b/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.test.ts index ab69b9d67..4cb6bb24c 100644 --- a/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.test.ts +++ b/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.test.ts @@ -1,5 +1,14 @@ -import { initializeOrUpdateMessage, updateChatTitle } from './chatStreamMessageHelper'; +import { + initializeOrUpdateMessage, + updateChatTitle, + updateResponseMessage +} from './chatStreamMessageHelper'; import { IBusterChatMessage, IBusterChat } from '../interfaces'; +import { ChatEvent_GeneratingResponseMessage } from '@/api/buster_socket/chats'; +import { + BusterChatResponseMessage_file, + BusterChatResponseMessage_text +} from '@/api/asset_interfaces'; describe('initializeOrUpdateMessage', () => { it('should initialize a new message when currentMessage is undefined', () => { @@ -119,3 +128,163 @@ describe('updateChatTitle', () => { expect(result.title).toBe('Initial Title'); }); }); + +describe('updateResponseMessage', () => { + it('should create a new message when currentMessage is undefined', () => { + const mockEvent: ChatEvent_GeneratingResponseMessage = { + chat_id: 'test-chat-id', + message_id: 'test-message-id', + response_message: { + id: 'response-1', + type: 'text', + message: '', + message_chunk: 'Hello' + } as BusterChatResponseMessage_text, + progress: 'in_progress' as const + }; + const result = updateResponseMessage('test-message-id', undefined, mockEvent); + + expect(result.id).toBe('test-message-id'); + expect(result.response_message_ids).toContain('response-1'); + expect( + (result.response_messages['response-1'] as BusterChatResponseMessage_text).message + ).toEqual('Hello'); + }); + + it('should update existing message with new response', () => { + const currentMessage: IBusterChatMessage = { + id: 'test-message-id', + isCompletedStream: false, + request_message: { + request: 'test request', + sender_id: 'user1', + sender_name: 'Test User', + sender_avatar: null + }, + response_message_ids: ['response-1'], + reasoning_message_ids: [], + response_messages: { + 'response-1': { + id: 'response-1', + type: 'text', + message: 'Hello', + message_chunk: undefined + } + }, + reasoning_messages: {}, + created_at: new Date().toISOString(), + final_reasoning_message: null + }; + + const mockEvent: ChatEvent_GeneratingResponseMessage = { + chat_id: 'test-chat-id', + message_id: 'test-message-id', + response_message: { + id: 'response-1', + type: 'text', + message: '', + message_chunk: ' World' + } as BusterChatResponseMessage_text, + progress: 'in_progress' as const + }; + + const result = updateResponseMessage('test-message-id', currentMessage, mockEvent); + + expect(result.response_message_ids).toContain('response-1'); + expect( + (result.response_messages['response-1'] as BusterChatResponseMessage_text).message + ).toEqual('Hello World'); + }); + + it('should handle large message chunks correctly', () => { + const currentMessage: IBusterChatMessage = { + id: 'test-message-id', + isCompletedStream: false, + request_message: { + request: 'test request', + sender_id: 'user1', + sender_name: 'Test User', + sender_avatar: null + }, + response_message_ids: ['response-1'], + reasoning_message_ids: [], + response_messages: { + 'response-1': { + id: 'response-1', + type: 'text', + message: 'Initial message', + message_chunk: undefined + } + }, + reasoning_messages: {}, + created_at: new Date().toISOString(), + final_reasoning_message: null + }; + + const largeChunk = + ' '.repeat(1000) + + 'This is a very large message chunk that should be appended correctly. '.repeat(10); + + const mockEvent: ChatEvent_GeneratingResponseMessage = { + chat_id: 'test-chat-id', + message_id: 'test-message-id', + response_message: { + id: 'response-1', + type: 'text', + message: '', + message_chunk: largeChunk + } as BusterChatResponseMessage_text, + progress: 'in_progress' as const + }; + + const result = updateResponseMessage('test-message-id', currentMessage, mockEvent); + + expect(result.response_message_ids).toContain('response-1'); + expect( + (result.response_messages['response-1'] as BusterChatResponseMessage_text).message + ).toEqual('Initial message' + largeChunk); + }); + + it('should handle file type response messages', () => { + const responseMessageFile: BusterChatResponseMessage_file = { + id: 'response-1', + type: 'file', + file_name: 'initial.txt', + file_type: 'metric', + version_number: 1, + version_id: '1', + filter_version_id: '1' + }; + + const currentMessage: IBusterChatMessage = { + id: 'test-message-id', + isCompletedStream: false, + request_message: { + request: 'test request', + sender_id: 'user1', + sender_name: 'Test User', + sender_avatar: null + }, + response_message_ids: ['response-1'], + reasoning_message_ids: [], + response_messages: { + 'response-1': responseMessageFile + }, + reasoning_messages: {}, + created_at: new Date().toISOString(), + final_reasoning_message: null + }; + + const mockEvent: ChatEvent_GeneratingResponseMessage = { + chat_id: 'test-chat-id', + message_id: 'test-message-id', + response_message: { ...responseMessageFile, id: 'response-2' }, + progress: 'in_progress' as const + }; + + const result = updateResponseMessage('test-message-id', currentMessage, mockEvent); + + expect(result.response_message_ids).toContain('response-2'); + expect(result.response_messages['response-2']).toEqual(mockEvent.response_message); + }); +}); diff --git a/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.ts b/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.ts index 1e588f42d..502dd7572 100644 --- a/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.ts +++ b/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.ts @@ -1,6 +1,10 @@ import { create } from 'mutative'; import { IBusterChat, IBusterChatMessage } from '../interfaces'; -import { ChatEvent_GeneratingTitle } from '@/api/buster_socket/chats'; +import { + ChatEvent_GeneratingTitle, + ChatEvent_GeneratingResponseMessage +} from '@/api/buster_socket/chats'; +import { BusterChatResponseMessage_text } from '@/api/asset_interfaces'; const createInitialMessage = (messageId: string): IBusterChatMessage => ({ id: messageId, @@ -41,3 +45,55 @@ export const updateChatTitle = ( if (newTitle) draft.title = newTitle; }); }; + +export const updateResponseMessage = ( + messageId: string, + currentMessage: IBusterChatMessage | undefined, + event: ChatEvent_GeneratingResponseMessage +): IBusterChatMessage => { + const { response_message } = event; + + if (!response_message?.id) { + return currentMessage || createInitialMessage(messageId); + } + + const responseMessageId = response_message.id; + const existingResponseMessage = currentMessage?.response_messages?.[responseMessageId]; + const isNewResponseMessage = !existingResponseMessage; + + let updatedMessage = currentMessage || createInitialMessage(messageId); + + if (isNewResponseMessage) { + updatedMessage = initializeOrUpdateMessage(messageId, updatedMessage, (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 = existingResponseMessage as BusterChatResponseMessage_text; + const isStreaming = + response_message.message_chunk !== undefined && response_message.message_chunk !== null; + + updatedMessage = initializeOrUpdateMessage(messageId, updatedMessage, (draft) => { + const responseMessage = draft.response_messages?.[responseMessageId]; + if (!responseMessage) return; + const messageText = responseMessage as BusterChatResponseMessage_text; + Object.assign(messageText, { + ...existingResponseMessageText, + ...response_message, + message: isStreaming + ? (existingResponseMessageText?.message || '') + (response_message.message_chunk || '') + : response_message.message + }); + }); + } + + return updatedMessage; +}; diff --git a/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts b/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts index 5c2c367c4..dd0ba0bd3 100644 --- a/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts +++ b/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts @@ -22,7 +22,11 @@ import { IBusterChat, IBusterChatMessage } from '../interfaces'; import { queryKeys } from '@/api/query_keys'; import { useQueryClient } from '@tanstack/react-query'; import { create } from 'mutative'; -import { initializeOrUpdateMessage, updateChatTitle } from './chatStreamMessageHelper'; +import { + initializeOrUpdateMessage, + updateChatTitle, + updateResponseMessage +} from './chatStreamMessageHelper'; export const useChatStreamMessage = () => { const queryClient = useQueryClient(); @@ -116,55 +120,18 @@ export const useChatStreamMessage = () => { const _generatingResponseMessageCallback = useMemoizedFn( (_: null, d: ChatEvent_GeneratingResponseMessage) => { - const { message_id, response_message } = d; + const { message_id } = d; - if (!response_message?.id) return; - - const responseMessageId = response_message.id; - const existingResponseMessage = - chatRefMessages.current[message_id]?.response_messages?.[responseMessageId]; - const isNewResponseMessage = !existingResponseMessage; - - let currentMessage = chatRefMessages.current[message_id]; - - if (isNewResponseMessage) { - currentMessage = initializeOrUpdateMessage(message_id, currentMessage, (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 = - existingResponseMessage as BusterChatResponseMessage_text; - const isStreaming = - response_message.message_chunk !== undefined && response_message.message_chunk !== null; - - currentMessage = initializeOrUpdateMessage(message_id, currentMessage, (draft) => { - const responseMessage = draft.response_messages?.[responseMessageId]; - if (!responseMessage) return; - const messageText = responseMessage as BusterChatMessageReasoning_text; - Object.assign(messageText, { - ...existingResponseMessageText, - ...response_message, - message: isStreaming - ? (existingResponseMessageText?.message || '') + - (response_message.message_chunk || '') - : response_message.message - }); - }); - } + const updatedMessage = updateResponseMessage( + message_id, + chatRefMessages.current[message_id], + d + ); onUpdateChatMessageTransition({ id: message_id, - response_messages: currentMessage?.response_messages, - response_message_ids: currentMessage?.response_message_ids + response_messages: updatedMessage?.response_messages, + response_message_ids: updatedMessage?.response_message_ids }); } );