Refactor message handling to real-time WebSocket updates

Author: bogdan-ts

ui/packages/platform/src/pages/Bot/Command/Command.tsx

@@ -85,10 +85,11 @@ export const Command = React.memo((props: Props) => {
     wsLoading,
     loading,
     sendMessage,
-    chatVisibility
+    chatVisibility,
+    isStreamingInProcess
   } = useAiBot();
 
-  const sendDisabled = error !== null || loading || wsLoading || wsReadyState !== ReadyState.OPEN;
+  const sendDisabled = error !== null || loading || wsLoading || wsReadyState !== ReadyState.OPEN || isStreamingInProcess;
 
   // Handle value.
   const [value, setValue] = useState('')

ui/packages/platform/src/pages/Bot/Messages/Message/Message.tsx

@@ -203,6 +203,7 @@ const useStyles = makeStyles(
         marginBlockEnd: '1em',
         marginInlineStart: 0,
         marginInlineEnd: 0,
+        //animation: `$typing 0.5s steps(30, end), $blinkCaret 0.75s step-end infinite`,
       },
       '& .MuiExpansionPanel-root div': {
         marginBlockStart: 0,
@@ -216,6 +217,7 @@ const useStyles = makeStyles(
       marginInlineStart: 0,
       marginInlineEnd: 0,
       fontSize: 14,
+      color: colors.pgaiDarkGray,
       '&:after': {
         overflow: 'hidden',
         display: 'inline-block',
@@ -229,7 +231,16 @@ const useStyles = makeStyles(
       'to': {
         width: '0.9em'
       },
-    }
+    },
+    '@keyframes typing': {
+      from: { width: 0 },
+      to: { width: '100%' },
+    },
+    '@keyframes blinkCaret': {
+      from: { borderRightColor: 'transparent' },
+      to: { borderRightColor: 'transparent' },
+      '50%': { borderRightColor: 'black' },
+    },
   }),
 
 )
@@ -336,7 +347,8 @@ export const Message = React.memo((props: MessageProps) => {
         </div>
         <div>
           {isLoading
-            ? <div className={classes.markdown}>
+            ?
+            <div className={classes.markdown}>
               <div className={classes.loading}>
                 {stateMessage && stateMessage.state ? stateMessage.state : 'Thinking'}
               </div>

ui/packages/platform/src/pages/Bot/Messages/Messages.tsx

@@ -111,23 +111,22 @@ const useStyles = makeStyles(
       padding: 10
     }
   }),
-
 )
 
-type MessagesProps = {
-  messages: BotMessageWithDebugInfo[] | null
-  isLoading: boolean
-  isWaitingForAnswer: boolean
-}
-
 type Time = string
 
 type FormattedTime = {
   [id: string]: Time
 }
 
 export const Messages = React.memo(() => {
-  const { messages, loading: isLoading, wsLoading: isWaitingForAnswer, stateMessage } = useAiBot();
+  const {
+    messages,
+    loading: isLoading,
+    wsLoading: isWaitingForAnswer,
+    stateMessage,
+    currentStreamMessage
+  } = useAiBot();
 
   const rootRef = useRef<HTMLDivElement>(null);
   const wrapperRef = useRef<HTMLDivElement>(null);
@@ -275,6 +274,14 @@ export const Messages = React.memo(() => {
                 />
               )
             })}
+          {
+            currentStreamMessage && <Message
+              id={null}
+              isAi
+              content={currentStreamMessage.content}
+              aiModel={currentStreamMessage.ai_model}
+            />
+          }
           {isWaitingForAnswer &&
             <Message id={null} isLoading isAi={true} stateMessage={stateMessage} />
           }

ui/packages/platform/src/pages/Bot/hooks.tsx

@@ -8,7 +8,7 @@
 import React, { createContext, useCallback, useContext, useEffect, useState } from "react";
 import useWebSocket, {ReadyState} from "react-use-websocket";
 import { useLocation } from "react-router-dom";
-import { BotMessage, DebugMessage, AiModel, StateMessage } from "../../types/api/entities/bot";
+import { BotMessage, DebugMessage, AiModel, StateMessage, StreamMessage } from "../../types/api/entities/bot";
 import {getChatsWithWholeThreads} from "../../api/bot/getChatsWithWholeThreads";
 import {getChats} from "api/bot/getChats";
 import {useAlertSnackbar} from "@postgres.ai/shared/components/AlertSnackbar/useAlertSnackbar";
@@ -58,6 +58,8 @@ type UseAiBotReturnType = {
   aiModelsLoading: UseAiModelsList['loading'];
   debugMessagesLoading: boolean;
   stateMessage: StateMessage | null;
+  isStreamingInProcess: boolean
+  currentStreamMessage: StreamMessage | null
 }
 
 type UseAiBotArgs = {
@@ -90,6 +92,8 @@ export const useAiBotProviderValue = (args: UseAiBotArgs): UseAiBotReturnType =>
   const [wsLoading, setWsLoading] = useState<boolean>(false);
   const [chatVisibility, setChatVisibility] = useState<UseAiBotReturnType['chatVisibility']>('public');
   const [stateMessage, setStateMessage] = useState<StateMessage | null>(null)
+  const [currentStreamMessage, setCurrentStreamMessage] = useState<StreamMessage | null>(null)
+  const [isStreamingInProcess, setStreamingInProcess] = useState<boolean>(false)
 
   const [isChangeVisibilityLoading, setIsChangeVisibilityLoading] = useState<boolean>(false);
 
@@ -102,51 +106,35 @@ export const useAiBotProviderValue = (args: UseAiBotArgs): UseAiBotReturnType =>
 
   const onWebSocketMessage = (event: WebSocketEventMap['message']) => {
     if (event.data) {
-      const messageData: BotMessage | DebugMessage | StateMessage = JSON.parse(event.data);
+      const messageData: BotMessage | DebugMessage | StateMessage | StreamMessage = JSON.parse(event.data);
       if (messageData) {
         const isThreadMatching = threadId && threadId === messageData.thread_id;
         const isParentMatching = !threadId && 'parent_id' in messageData && messageData.parent_id && messages;
         const isDebugMessage = messageData.type === 'debug';
         const isStateMessage = messageData.type === 'state';
-        if (isThreadMatching || isParentMatching || isDebugMessage || isStateMessage) {
-          if (isDebugMessage) {
-            let currentDebugMessages = [...(debugMessages || [])];
-            currentDebugMessages.push(messageData)
-            setDebugMessages(currentDebugMessages)
-          } else if (isStateMessage) {
-            if (isThreadMatching || !threadId) {
-              if (messageData.state) {
-                setStateMessage(messageData)
-              } else {
-                setStateMessage(null)
-              }
-            }
-          } else {
-            // Check if the last message needs its data updated
-            let currentMessages = [...(messages || [])];
-            const lastMessage = currentMessages[currentMessages.length - 1];
-            if (lastMessage && !lastMessage.id && messageData.parent_id) {
-              lastMessage.id = messageData.parent_id;
-              lastMessage.created_at = messageData.created_at;
-              lastMessage.is_public = messageData.is_public;
-            }
-
-            currentMessages.push(messageData);
-            setMessages(currentMessages);
-            setWsLoading(false);
-            if (document.visibilityState === "hidden") {
-              if (Notification.permission === "granted") {
-                new Notification("New message", {
-                  body: 'New message from Postgres.AI Bot',
-                  icon: '/images/bot_avatar.png'
-                });
-              }
-            }
+        const isStreamMessage = messageData.type === 'stream';
+
+        if (isThreadMatching || isParentMatching || isDebugMessage || isStateMessage || isStreamMessage) {
+          switch (messageData.type) {
+            case 'debug':
+              handleDebugMessage(messageData)
+              break;
+            case 'state':
+              handleStateMessage(messageData, Boolean(isThreadMatching))
+              break;
+            case 'stream':
+              handleStreamMessage(messageData, Boolean(isThreadMatching))
+              break;
+            case 'message':
+              handleBotMessage(messageData)
+              break;
           }
         } else if (threadId !== messageData.thread_id) {
           const threadInList = chatsList?.find((item) => item.thread_id === messageData.thread_id)
           if (!threadInList) getChatsList()
-          setWsLoading(false);
+          if (currentStreamMessage) setCurrentStreamMessage(null)
+          if (wsLoading) setWsLoading(false);
+          if (isStreamingInProcess) setStreamingInProcess(false)
         }
       } else {
         showMessage('An error occurred. Please try again')
@@ -158,6 +146,56 @@ export const useAiBotProviderValue = (args: UseAiBotArgs): UseAiBotReturnType =>
     setLoading(false);
   }
 
+  const handleDebugMessage = (message: DebugMessage) => {
+    let currentDebugMessages = [...(debugMessages || [])];
+    currentDebugMessages.push(message)
+    setDebugMessages(currentDebugMessages)
+  }
+
+  const handleStateMessage = (message: StateMessage, isThreadMatching?: boolean) => {
+    if (isThreadMatching || !threadId) {
+      if (message.state) {
+        setStateMessage(message)
+      } else {
+        setStateMessage(null)
+      }
+    }
+  }
+
+  const handleStreamMessage = (message: StreamMessage, isThreadMatching?: boolean) => {
+    if (isThreadMatching || !threadId) {
+      if (!isStreamingInProcess) setStreamingInProcess(true)
+      setCurrentStreamMessage(message)
+      setWsLoading(false);
+    }
+  }
+
+  const handleBotMessage = (message: BotMessage) => {
+    if (messages && messages.length > 0) {
+      let currentMessages = [...messages];
+      const lastMessage = currentMessages[currentMessages.length - 1];
+      if (lastMessage && !lastMessage.id && message.parent_id) {
+        lastMessage.id = message.parent_id;
+        lastMessage.created_at = message.created_at;
+        lastMessage.is_public = message.is_public;
+      }
+
+      currentMessages.push(message);
+      if (currentStreamMessage) setCurrentStreamMessage(null)
+      setMessages(currentMessages);
+      setWsLoading(false);
+      setStreamingInProcess(false);
+      if (document.visibilityState === "hidden") {
+        if (Notification.permission === "granted") {
+          new Notification("New message", {
+            body: 'New message from Postgres.AI Bot',
+            icon: '/images/bot_avatar.png'
+          });
+        }
+      }
+    }
+  }
+
   const onWebSocketOpen = () => {
     console.log('WebSocket connection established');
     if (threadId) {
@@ -381,6 +419,8 @@ export const useAiBotProviderValue = (args: UseAiBotArgs): UseAiBotReturnType =>
     debugMessages,
     debugMessagesLoading,
     stateMessage,
+    isStreamingInProcess,
+    currentStreamMessage
   }
 }
 

ui/packages/platform/src/pages/Bot/index.tsx

@@ -156,9 +156,9 @@ export const BotPage = (props: BotPageProps) => {
   }
 
   useEffect(() => {
-    if (!match.params.threadId && !prevThreadId && messages && messages.length > 1 && messages[1].parent_id) {
+    if (!match.params.threadId && !prevThreadId && messages && messages.length > 0 && messages[0].id) {
       // hack that skip additional loading chats_ancestors_and_descendants
-      history.replace(`/${match.params.org}/bot/${messages[1].parent_id}`, { skipReloading: true })
+      history.replace(`/${match.params.org}/bot/${messages[0].id}`, { skipReloading: true })
       getChatsList();
     } else if (prevThreadId && !match.params.threadId) {
       clearChat()

ui/packages/platform/src/types/api/entities/bot.ts

@@ -40,4 +40,11 @@ export type StateMessage = {
   type: 'state'
   state: string | null
   thread_id: string
+}
+
+export type StreamMessage = {
+  type: 'stream'
+  content: string
+  ai_model: string
+  thread_id: string
 }

ui/packages/platform/src/utils/format.ts

@@ -264,6 +264,8 @@ const Format = {
   },
 
   timeAgo: function (date: string | Date): string | null {
+    if (!date) return null
+
     const now = new Date();
     const past = new Date(date);
     const diff = Math.abs(now.getTime() - past.getTime());