make response messages a helper function

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);
+  });
+});

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;
+};

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
       });
     }
   );