From db90ed13f4de49d7cedfde12991e4b2a2833eea3 Mon Sep 17 00:00:00 2001 From: Nate Kelley Date: Thu, 6 Mar 2025 09:51:12 -0700 Subject: [PATCH] move some functionality to a helper for testing --- .../chatStreamMessageHelper.test.ts | 121 ++++++++++++++++++ .../chatStreamMessageHelper.ts | 43 +++++++ .../NewChatProvider/useChatStreamMessage.ts | 68 ++++------ 3 files changed, 191 insertions(+), 41 deletions(-) create mode 100644 web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.test.ts create mode 100644 web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.ts diff --git a/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.test.ts b/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.test.ts new file mode 100644 index 000000000..ab69b9d67 --- /dev/null +++ b/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.test.ts @@ -0,0 +1,121 @@ +import { initializeOrUpdateMessage, updateChatTitle } from './chatStreamMessageHelper'; +import { IBusterChatMessage, IBusterChat } from '../interfaces'; + +describe('initializeOrUpdateMessage', () => { + it('should initialize a new message when currentMessage is undefined', () => { + const messageId = 'test-id'; + const updateFn = (draft: IBusterChatMessage) => { + if (draft.request_message) { + draft.request_message.request = 'test request'; + } + }; + + const result = initializeOrUpdateMessage(messageId, undefined, updateFn); + + expect(result.id).toBe(messageId); + expect(result.isCompletedStream).toBe(false); + expect(result.request_message?.request).toBe('test request'); + expect(result.created_at).toBeDefined(); + expect(result.final_reasoning_message).toBeNull(); + }); + + it('should update an existing message', () => { + const messageId = 'test-id'; + const currentMessage: IBusterChatMessage = { + id: messageId, + isCompletedStream: false, + request_message: { + request: 'original request', + sender_id: 'user1', + sender_name: 'Test User', + sender_avatar: null + }, + response_message_ids: [], + reasoning_message_ids: [], + response_messages: {}, + reasoning_messages: {}, + created_at: new Date().toISOString(), + final_reasoning_message: null + }; + + const updateFn = (draft: IBusterChatMessage) => { + if (draft.request_message) { + draft.request_message.request = 'updated request'; + } + }; + + const result = initializeOrUpdateMessage(messageId, currentMessage, updateFn); + + expect(result.id).toBe(messageId); + expect(result.request_message?.request).toBe('updated request'); + }); +}); + +const mockChat: IBusterChat = { + id: 'test-chat-id', + title: 'Initial Title', + isNewChat: false, + is_favorited: false, + message_ids: [], + created_at: new Date().toISOString(), + updated_at: new Date().toISOString(), + created_by: 'test-user', + created_by_id: 'test-user-id', + created_by_name: 'Test User', + created_by_avatar: null +}; + +describe('updateChatTitle', () => { + it('should update title with a chunk when progress is not completed', () => { + const event = { + chat_id: 'test-chat-id', + title: 'Final Title', + title_chunk: ' New', + progress: 'in_progress' as const + }; + + const result = updateChatTitle(mockChat, event); + expect(result.title).toBe('Initial Title New'); + }); + + it('should set the final title when progress is completed', () => { + const event = { + chat_id: 'test-chat-id', + title: 'Final Title', + title_chunk: '', + progress: 'completed' as const + }; + + const result = updateChatTitle(mockChat, event); + expect(result.title).toBe('Final Title'); + }); + + it('should handle chat with empty initial title', () => { + const chatWithoutTitle: IBusterChat = { + ...mockChat, + title: '' + }; + + const event = { + chat_id: 'test-chat-id', + title: 'Final Title', + title_chunk: ' New', + progress: 'in_progress' as const + }; + + const result = updateChatTitle(chatWithoutTitle, event); + expect(result.title).toBe(' New'); + }); + + it('should not modify title when title and title_chunk are empty', () => { + const event = { + chat_id: 'test-chat-id', + title: '', + title_chunk: '', + progress: 'completed' as const + }; + + const result = updateChatTitle(mockChat, event); + expect(result.title).toBe('Initial Title'); + }); +}); diff --git a/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.ts b/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.ts new file mode 100644 index 000000000..1e588f42d --- /dev/null +++ b/web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.ts @@ -0,0 +1,43 @@ +import { create } from 'mutative'; +import { IBusterChat, IBusterChatMessage } from '../interfaces'; +import { ChatEvent_GeneratingTitle } from '@/api/buster_socket/chats'; + +const createInitialMessage = (messageId: string): IBusterChatMessage => ({ + id: messageId, + isCompletedStream: false, + request_message: { + request: '', + sender_id: '', + sender_name: '', + sender_avatar: null + }, + response_message_ids: [], + reasoning_message_ids: [], + response_messages: {}, + reasoning_messages: {}, + created_at: new Date().toISOString(), + final_reasoning_message: null +}); + +export const initializeOrUpdateMessage = ( + messageId: string, + currentMessage: IBusterChatMessage | undefined, + updateFn: (draft: IBusterChatMessage) => void +) => { + return create(currentMessage || createInitialMessage(messageId), (draft) => { + updateFn(draft); + }); +}; + +export const updateChatTitle = ( + currentChat: IBusterChat, + event: ChatEvent_GeneratingTitle +): IBusterChat => { + const { chat_id, title, title_chunk, progress } = event; + const isCompleted = progress === 'completed'; + const currentTitle = currentChat.title || ''; + const newTitle = isCompleted ? title : currentTitle + title_chunk; + return create(currentChat, (draft) => { + if (newTitle) draft.title = newTitle; + }); +}; diff --git a/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts b/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts index 538e1262a..5c2c367c4 100644 --- a/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts +++ b/web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts @@ -22,6 +22,7 @@ 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'; export const useChatStreamMessage = () => { const queryClient = useQueryClient(); @@ -50,17 +51,6 @@ export const useChatStreamMessage = () => { } ); - const initializeOrUpdateMessage = useMemoizedFn( - (messageId: string, updateFn: (draft: IBusterChatMessage) => void) => { - const currentMessage = chatRefMessages.current[messageId]; - const updatedMessage = create(currentMessage || {}, (draft) => { - updateFn(draft); - }); - chatRefMessages.current[messageId] = updatedMessage; - onUpdateChatMessage(updatedMessage); - } - ); - const normalizeChatMessage = useMemoizedFn( (iChatMessages: Record) => { for (const message of Object.values(iChatMessages)) { @@ -93,7 +83,6 @@ export const useChatStreamMessage = () => { chatRef.current = create(chatRef.current, (draft) => { draft[iChat.id] = iChat; }); - normalizeChatMessage(iChatMessages); onUpdateChat(iChat); onChangePage({ @@ -119,32 +108,27 @@ export const useChatStreamMessage = () => { ); const _generatingTitleCallback = useMemoizedFn((_: null, newData: ChatEvent_GeneratingTitle) => { - const { chat_id, title, title_chunk, progress } = newData; - const isCompleted = progress === 'completed'; - const currentTitle = chatRef.current[chat_id]?.title || ''; - const newTitle = isCompleted ? title : currentTitle + title_chunk; - chatRef.current = create(chatRef.current, (draft) => { - if (newTitle) draft[chat_id].title = newTitle; - }); - onUpdateChat({ - id: chat_id, - title: newTitle - }); + const { chat_id } = newData; + const updatedChat = updateChatTitle(chatRef.current[chat_id], newData); + chatRef.current[chat_id] = updatedChat; + onUpdateChat(updatedChat); }); const _generatingResponseMessageCallback = useMemoizedFn( (_: null, d: ChatEvent_GeneratingResponseMessage) => { - const { message_id, response_message, chat_id } = d; + const { message_id, response_message } = d; if (!response_message?.id) return; const responseMessageId = response_message.id; - const existingMessage = + const existingResponseMessage = chatRefMessages.current[message_id]?.response_messages?.[responseMessageId]; - const isNewMessage = !existingMessage; + const isNewResponseMessage = !existingResponseMessage; - if (isNewMessage) { - initializeOrUpdateMessage(message_id, (draft) => { + let currentMessage = chatRefMessages.current[message_id]; + + if (isNewResponseMessage) { + currentMessage = initializeOrUpdateMessage(message_id, currentMessage, (draft) => { if (!draft.response_messages) { draft.response_messages = {}; } @@ -157,11 +141,12 @@ export const useChatStreamMessage = () => { } if (response_message.type === 'text') { - const existingResponseMessageText = existingMessage as BusterChatResponseMessage_text; + const existingResponseMessageText = + existingResponseMessage as BusterChatResponseMessage_text; const isStreaming = response_message.message_chunk !== undefined && response_message.message_chunk !== null; - initializeOrUpdateMessage(message_id, (draft) => { + currentMessage = initializeOrUpdateMessage(message_id, currentMessage, (draft) => { const responseMessage = draft.response_messages?.[responseMessageId]; if (!responseMessage) return; const messageText = responseMessage as BusterChatMessageReasoning_text; @@ -176,7 +161,6 @@ export const useChatStreamMessage = () => { }); } - const currentMessage = chatRefMessages.current[message_id]; onUpdateChatMessageTransition({ id: message_id, response_messages: currentMessage?.response_messages, @@ -190,12 +174,13 @@ export const useChatStreamMessage = () => { const { message_id, reasoning, chat_id } = d; const reasoningMessageId = reasoning.id; - const existingMessage = + const existingReasoningMessage = chatRefMessages.current[message_id]?.reasoning_messages?.[reasoningMessageId]; - const isNewMessage = !existingMessage; + const isNewReasoningMessage = !existingReasoningMessage; + let currentMessage = chatRefMessages.current[message_id]; - if (isNewMessage) { - initializeOrUpdateMessage(message_id, (draft) => { + if (isNewReasoningMessage) { + currentMessage = initializeOrUpdateMessage(message_id, currentMessage, (draft) => { if (!draft.reasoning_messages) { draft.reasoning_messages = {}; } @@ -209,11 +194,12 @@ export const useChatStreamMessage = () => { switch (reasoning.type) { case 'text': { - const existingReasoningMessageText = existingMessage as BusterChatMessageReasoning_text; + const existingReasoningMessageText = + existingReasoningMessage as BusterChatMessageReasoning_text; const isStreaming = reasoning.message_chunk !== null || reasoning.message_chunk !== undefined; - initializeOrUpdateMessage(message_id, (draft) => { + currentMessage = initializeOrUpdateMessage(message_id, currentMessage, (draft) => { const reasoningMessage = draft.reasoning_messages?.[reasoningMessageId]; if (!reasoningMessage) return; const messageText = reasoningMessage as BusterChatMessageReasoning_text; @@ -230,9 +216,10 @@ export const useChatStreamMessage = () => { break; } case 'files': { - const existingReasoningMessageFiles = existingMessage as BusterChatMessageReasoning_files; + const existingReasoningMessageFiles = + existingReasoningMessage as BusterChatMessageReasoning_files; - initializeOrUpdateMessage(message_id, (draft) => { + currentMessage = initializeOrUpdateMessage(message_id, currentMessage, (draft) => { const reasoningMessage = draft.reasoning_messages?.[reasoningMessageId]; if (!reasoningMessage) return; @@ -286,7 +273,7 @@ export const useChatStreamMessage = () => { break; } case 'pills': { - initializeOrUpdateMessage(message_id, (draft) => { + currentMessage = initializeOrUpdateMessage(message_id, currentMessage, (draft) => { if (!draft.reasoning_messages?.[reasoningMessageId]) return; draft.reasoning_messages[reasoningMessageId] = reasoning; }); @@ -299,7 +286,6 @@ export const useChatStreamMessage = () => { } } - const currentMessage = chatRefMessages.current[message_id]; onUpdateChatMessageTransition({ id: message_id, reasoning_messages: currentMessage?.reasoning_messages,