Bot UI: Display state messages

Author: Bogdan Tsechoev

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

@@ -3,14 +3,14 @@ import cn from "classnames";
 import ReactMarkdown, { Components } from "react-markdown";
 import rehypeRaw from "rehype-raw";
 import remarkGfm from "remark-gfm";
-import { Badge, makeStyles } from "@material-ui/core";
+import { makeStyles } from "@material-ui/core";
 import { colors } from "@postgres.ai/shared/styles/colors";
 import { icons } from "@postgres.ai/shared/styles/icons";
 import { DebugDialog } from "../../DebugDialog/DebugDialog";
 import { CodeBlock } from "./CodeBlock";
 import { disallowedHtmlTagsForMarkdown, permalinkLinkBuilder } from "../../utils";
-import { BotMessage, DebugMessage } from "../../../../types/api/entities/bot";
-import { getDebugMessages } from "../../../../api/bot/getDebugMessages";
+import { StateMessage } from "../../../../types/api/entities/bot";
+
 
 type BaseMessageProps = {
   id: string | null;
@@ -20,6 +20,7 @@ type BaseMessageProps = {
   isLoading?: boolean;
   formattedTime?: string;
   aiModel?: string
+  stateMessage?: StateMessage | null
 }
 
 type AiMessageProps = BaseMessageProps & {
@@ -38,6 +39,7 @@ type LoadingMessageProps = BaseMessageProps & {
   isLoading: true;
   isAi: true;
   content?: undefined
+  stateMessage: StateMessage | null
 }
 
 type MessageProps = AiMessageProps | HumanMessageProps | LoadingMessageProps;
@@ -241,7 +243,8 @@ export const Message = React.memo((props: MessageProps) => {
     name,
     created_at,
     isLoading,
-    aiModel
+    aiModel,
+    stateMessage
   } = props;
 
   const [isDebugVisible, setDebugVisible] = useState(false);
@@ -335,8 +338,8 @@ export const Message = React.memo((props: MessageProps) => {
           {isLoading
             ? <div className={classes.markdown}>
               <div className={classes.loading}>
-                  Thinking
-                </div>
+                {stateMessage && stateMessage.state ? stateMessage.state : 'Thinking'}
+              </div>
             </div>
             : <ReactMarkdown
                 className={classes.markdown}

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

@@ -127,7 +127,7 @@ type FormattedTime = {
 }
 
 export const Messages = React.memo(() => {
-  const { messages, loading: isLoading, wsLoading: isWaitingForAnswer } = useAiBot();
+  const { messages, loading: isLoading, wsLoading: isWaitingForAnswer, stateMessage } = useAiBot();
 
   const rootRef = useRef<HTMLDivElement>(null);
   const wrapperRef = useRef<HTMLDivElement>(null);
@@ -276,7 +276,7 @@ export const Messages = React.memo(() => {
               )
             })}
           {isWaitingForAnswer &&
-            <Message id={null} isLoading isAi={true} />
+            <Message id={null} isLoading isAi={true} stateMessage={stateMessage} />
           }
         </div>
       </div>

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 } from "../../types/api/entities/bot";
+import { BotMessage, DebugMessage, AiModel, StateMessage } 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";
@@ -45,18 +45,19 @@ type UseAiBotReturnType = {
   wsReadyState: ReadyState;
   changeChatVisibility: (threadId: string, isPublic: boolean) => void;
   isChangeVisibilityLoading: boolean;
-  unsubscribe: (threadId: string) => void
-  chatVisibility: 'public' | 'private'
-  debugMessages: DebugMessage[] | null
-  getDebugMessagesForWholeThread: () => void
+  unsubscribe: (threadId: string) => void;
+  chatVisibility: 'public' | 'private';
+  debugMessages: DebugMessage[] | null;
+  getDebugMessagesForWholeThread: () => void;
   chatsList: UseBotChatsListHook['chatsList'];
   chatsListLoading: UseBotChatsListHook['loading'];
-  getChatsList: UseBotChatsListHook['getChatsList']
-  aiModel: UseAiModelsList['aiModel'],
-  setAiModel: UseAiModelsList['setAiModel']
-  aiModels: UseAiModelsList['aiModels']
-  aiModelsLoading: UseAiModelsList['loading']
+  getChatsList: UseBotChatsListHook['getChatsList'];
+  aiModel: UseAiModelsList['aiModel'];
+  setAiModel: UseAiModelsList['setAiModel'];
+  aiModels: UseAiModelsList['aiModels'];
+  aiModelsLoading: UseAiModelsList['loading'];
   debugMessagesLoading: boolean;
+  stateMessage: StateMessage | null;
 }
 
 type UseAiBotArgs = {
@@ -88,6 +89,7 @@ export const useAiBotProviderValue = (args: UseAiBotArgs): UseAiBotReturnType =>
   const [error, setError] = useState<ErrorType | null>(null);
   const [wsLoading, setWsLoading] = useState<boolean>(false);
   const [chatVisibility, setChatVisibility] = useState<UseAiBotReturnType['chatVisibility']>('public');
+  const [stateMessage, setStateMessage] = useState<StateMessage | null>(null)
 
   const [isChangeVisibilityLoading, setIsChangeVisibilityLoading] = useState<boolean>(false);
 
@@ -100,16 +102,25 @@ export const useAiBotProviderValue = (args: UseAiBotArgs): UseAiBotReturnType =>
 
   const onWebSocketMessage = (event: WebSocketEventMap['message']) => {
     if (event.data) {
-      const messageData: BotMessage | DebugMessage = JSON.parse(event.data);
+      const messageData: BotMessage | DebugMessage | StateMessage = 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';
-        if (isThreadMatching || isParentMatching || isDebugMessage) {
+        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 || [])];
@@ -368,7 +379,8 @@ export const useAiBotProviderValue = (args: UseAiBotArgs): UseAiBotReturnType =>
     aiModelsLoading,
     chatVisibility,
     debugMessages,
-    debugMessagesLoading
+    debugMessagesLoading,
+    stateMessage,
   }
 }
 

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

@@ -34,4 +34,10 @@ export type AiModel = {
   vendor: string;
   isThirdParty: boolean;
   freeUseAvailable: boolean;
-};
+};
+
+export type StateMessage = {
+  type: 'state'
+  state: string | null
+  thread_id: string
+}