added functions for streaming content

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;

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;

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;

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;

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;

web/src/api/buster_socket/chats/index.ts

@@ -1,2 +1,3 @@
 export * from './chatRequests';
 export * from './chatResponses';
+export * from './eventInterfaces';

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 (
       <div>
         <AppTooltip title="Open file">

web/src/app/app/_controllers/ReasoningController/ReasoningMessages/ReasoningMessage_File/ReasoningMessage_File.tsx

@@ -125,3 +125,5 @@ const MemoizedSyntaxHighlighter = React.memo(
     );
   }
 );
+
+MemoizedSyntaxHighlighter.displayName = 'MemoizedSyntaxHighlighter';

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':

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

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<Record<string, IBusterChat>>;
   chatsMessagesRef: MutableRefObject<Record<string, IBusterChatMessage>>;
 }) => {
+  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
+  };
 };

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

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,

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