refactor: Remove old message handling

Chainlit · Oct 24, 2023 · 776a1d0 · 776a1d0
1 parent aebd100
commit 776a1d0
Show file tree

Hide file tree

Showing 11 changed files with 56 additions and 136 deletions.
diff --git a/libs/components/hooks/useChat/state.ts b/libs/components/hooks/useChat/state.ts
@@ -7,7 +7,6 @@ import {
   IAvatarElement,
   IMessage,
   IMessageElement,
-  INestedMessage,
   ITasklistElement
 } from 'src/types';
 import { v4 as uuidv4 } from 'uuid';
@@ -46,12 +45,6 @@ export const messagesState = atom<IMessage[]>({
   default: []
 });
 
-export const nestedMessagesState = atom<INestedMessage[]>({
-  key: 'NestedMessages',
-  dangerouslyAllowMutability: true,
-  default: []
-});
-
 export const tokenCountState = atom<number>({
   key: 'TokenCount',
   default: 0

diff --git a/libs/components/hooks/useChat/useChatInteract.ts b/libs/components/hooks/useChat/useChatInteract.ts
@@ -1,7 +1,7 @@
 import { useCallback } from 'react';
 import { useRecoilValue, useResetRecoilState, useSetRecoilState } from 'recoil';
 import { IAction, IFileElement, IMessage } from 'src/types';
-import { addNestedMessage } from 'utils/message';
+import { addMessages } from 'utils/message';
 
 import {
   actionState,
@@ -13,7 +13,6 @@ import {
   firstUserMessageState,
   loadingState,
   messagesState,
-  nestedMessagesState,
   sessionIdState,
   sessionState,
   tasklistState,
@@ -31,7 +30,6 @@ const useChatInteract = () => {
   const setFirstUserMessage = useSetRecoilState(firstUserMessageState);
   const setLoading = useSetRecoilState(loadingState);
   const setMessages = useSetRecoilState(messagesState);
-  const setNestedMessages = useSetRecoilState(nestedMessagesState);
   const setElements = useSetRecoilState(elementState);
   const setAvatars = useSetRecoilState(avatarState);
   const setTasklists = useSetRecoilState(tasklistState);
@@ -44,7 +42,6 @@ const useChatInteract = () => {
     resetSessionId();
     setFirstUserMessage(undefined);
     setMessages([]);
-    setNestedMessages([]);
     setElements([]);
     setAvatars([]);
     setTasklists([]);
@@ -56,10 +53,7 @@ const useChatInteract = () => {
 
   const sendMessage = useCallback(
     (message: IMessage, files?: IFileElement[]) => {
-      setMessages((oldMessages) => [...oldMessages, { ...message }]);
-      setNestedMessages((oldNestedMessages) =>
-        addNestedMessage(oldNestedMessages, message)
-      );
+      setMessages((oldMessages) => addMessages(oldMessages, message));
 
       session?.socket.emit('ui_message', { message, files });
     },
@@ -69,10 +63,7 @@ const useChatInteract = () => {
   const replyMessage = useCallback(
     (message: IMessage) => {
       if (askUser) {
-        setMessages((oldMessages) => [...oldMessages, { ...message }]);
-        setNestedMessages((oldNestedMessages) =>
-          addNestedMessage(oldNestedMessages, message)
-        );
+        setMessages((oldMessages) => addMessages(oldMessages, message));
         askUser.callback(message);
       }
     },

diff --git a/libs/components/hooks/useChat/useChatMessages.ts b/libs/components/hooks/useChat/useChatMessages.ts
@@ -1,9 +1,9 @@
 import { useRecoilValue } from 'recoil';
 
-import { firstUserMessageState, nestedMessagesState } from './state';
+import { firstUserMessageState, messagesState } from './state';
 
 const useChatMessages = () => {
-  const messages = useRecoilValue(nestedMessagesState);
+  const messages = useRecoilValue(messagesState);
   const firstUserMessage = useRecoilValue(firstUserMessageState);
 
   return {

diff --git a/libs/components/hooks/useChat/useChatSession.ts b/libs/components/hooks/useChat/useChatSession.ts
@@ -1,5 +1,4 @@
 import debounce from 'lodash/debounce';
-import isEqual from 'lodash/isEqual';
 import { useCallback } from 'react';
 import {
   useRecoilState,
@@ -11,7 +10,7 @@ import io from 'socket.io-client';
 import { TFormInput } from 'src/inputs';
 import { IAction, IElement, IMessage } from 'src/types';
 import {
-  addNestedMessage,
+  addMessages,
   deleteMessageById,
   updateMessageById,
   updateMessageContentById
@@ -27,7 +26,6 @@ import {
   firstUserMessageState,
   loadingState,
   messagesState,
-  nestedMessagesState,
   sessionIdState,
   sessionState,
   tasklistState,
@@ -44,7 +42,6 @@ const useChatSession = () => {
   const setFirstUserMessage = useSetRecoilState(firstUserMessageState);
   const setLoading = useSetRecoilState(loadingState);
   const setMessages = useSetRecoilState(messagesState);
-  const setNestedMessages = useSetRecoilState(nestedMessagesState);
   const setAskUser = useSetRecoilState(askUserState);
   const setElements = useSetRecoilState(elementState);
   const setAvatars = useSetRecoilState(avatarState);
@@ -105,93 +102,38 @@ const useChatSession = () => {
       });
 
       socket.on('new_message', (message: IMessage) => {
-        setMessages((oldMessages) => {
-          const index = oldMessages.findIndex((m) => isEqual(m.id, message.id));
-          if (index === -1) {
-            return [...oldMessages, message];
-          } else {
-            return [
-              ...oldMessages.slice(0, index),
-              message,
-              ...oldMessages.slice(index + 1)
-            ];
-          }
-        });
-
-        setNestedMessages((oldNestedMessages) =>
-          addNestedMessage(oldNestedMessages, message)
-        );
+        setMessages((oldMessages) => addMessages(oldMessages, message));
       });
 
       socket.on('init_conversation', (message: IMessage) => {
         setFirstUserMessage(message);
       });
 
       socket.on('update_message', (message: IMessageUpdate) => {
-        setNestedMessages((oldNestedMessages) =>
-          updateMessageById(oldNestedMessages, message.id, message)
+        setMessages((oldMessages) =>
+          updateMessageById(oldMessages, message.id, message)
         );
       });
 
       socket.on('delete_message', (message: IMessage) => {
-        setMessages((oldMessages) => {
-          const index = oldMessages.findIndex((m) => isEqual(m.id, message.id));
-
-          if (index === -1) return oldMessages;
-          return [
-            ...oldMessages.slice(0, index),
-            ...oldMessages.slice(index + 1)
-          ];
-        });
-
-        setNestedMessages((oldNestedMessages) =>
-          deleteMessageById(oldNestedMessages, message.id)
+        setMessages((oldMessages) =>
+          deleteMessageById(oldMessages, message.id)
         );
       });
 
       socket.on('stream_start', (message: IMessage) => {
-        setMessages((oldMessages) => {
-          const index = oldMessages.findIndex((m) => m.id === message.id);
-          if (index === -1) {
-            return [...oldMessages, message];
-          }
-          return oldMessages;
-        });
-
-        setNestedMessages((oldNestedMessages) =>
-          addNestedMessage(oldNestedMessages, message)
-        );
+        setMessages((oldMessages) => addMessages(oldMessages, message));
       });
 
       socket.on('stream_token', ({ id, token, isSequence }: IToken) => {
-        setMessages((oldMessages) => {
-          const index = oldMessages.findIndex((m) => m.id === id);
-          if (index === -1) return oldMessages;
-          const oldMessage = oldMessages[index];
-          const newMessage = { ...oldMessage };
-          if (isSequence) {
-            newMessage.content = token;
-          } else {
-            newMessage.content += token;
-          }
-          return [
-            ...oldMessages.slice(0, index),
-            newMessage,
-            ...oldMessages.slice(index + 1)
-          ];
-        });
-
-        setNestedMessages((oldNestedMessages) =>
-          updateMessageContentById(oldNestedMessages, id, token, isSequence)
+        setMessages((oldMessages) =>
+          updateMessageContentById(oldMessages, id, token, isSequence)
         );
       });
 
       socket.on('ask', ({ msg, spec }, callback) => {
         setAskUser({ spec, callback });
-        setMessages((oldMessages) => [...oldMessages, { ...msg }]);
-        setNestedMessages((oldNestedMessages) =>
-          addNestedMessage(oldNestedMessages, msg)
-        );
+        setMessages((oldMessages) => addMessages(oldMessages, msg));
 
         setLoading(false);
       });

diff --git a/libs/components/src/messages/Message.tsx b/libs/components/src/messages/Message.tsx
@@ -13,12 +13,12 @@ import { MessageContent } from './components/MessageContent';
 
 import { IAction } from 'src/types/action';
 import { IMessageElement } from 'src/types/element';
-import { INestedMessage } from 'src/types/message';
+import { IMessage } from 'src/types/message';
 
 import { Messages } from './Messages';
 
 interface Props {
-  message: INestedMessage;
+  message: IMessage;
   elements: IMessageElement[];
   actions: IAction[];
   indent: number;

diff --git a/libs/components/src/messages/MessageContainer.tsx b/libs/components/src/messages/MessageContainer.tsx
@@ -5,7 +5,7 @@ import Box from '@mui/material/Box';
 
 import { IAction } from 'src/types/action';
 import { IMessageElement } from 'src/types/element';
-import { INestedMessage } from 'src/types/message';
+import { IMessage } from 'src/types/message';
 import { IMessageContext } from 'src/types/messageContext';
 
 import { Messages } from './Messages';
@@ -15,7 +15,7 @@ interface Props {
   autoScroll?: boolean;
   context: IMessageContext;
   elements: IMessageElement[];
-  messages: INestedMessage[];
+  messages: IMessage[];
   setAutoScroll?: (autoScroll: boolean) => void;
 }
 

diff --git a/libs/components/src/messages/Messages.tsx b/libs/components/src/messages/Messages.tsx
@@ -4,12 +4,12 @@ import { isLastMessage } from 'utils/message';
 
 import { IAction } from 'src/types/action';
 import { IMessageElement } from 'src/types/element';
-import { INestedMessage } from 'src/types/message';
+import { IMessage } from 'src/types/message';
 
 import { Message } from './Message';
 
 interface Props {
-  messages: INestedMessage[];
+  messages: IMessage[];
   elements: IMessageElement[];
   actions: IAction[];
   indent: number;

diff --git a/libs/components/src/messages/components/DetailsButton.tsx b/libs/components/src/messages/components/DetailsButton.tsx
@@ -6,10 +6,10 @@ import ExpandLess from '@mui/icons-material/ExpandLess';
 import ExpandMore from '@mui/icons-material/ExpandMore';
 import CircularProgress from '@mui/material/CircularProgress';
 
-import { INestedMessage } from 'src/types/message';
+import { IMessage } from 'src/types/message';
 
 interface Props {
-  message: INestedMessage;
+  message: IMessage;
   opened: boolean;
   loading?: boolean;
   onClick: () => void;

diff --git a/libs/components/src/messages/components/FeedbackButtons.tsx b/libs/components/src/messages/components/FeedbackButtons.tsx
@@ -17,7 +17,7 @@ import ButtonGroup from '@mui/material/ButtonGroup';
 import Stack from '@mui/material/Stack';
 import Tooltip from '@mui/material/Tooltip';
 
-import { nestedMessagesState } from 'hooks/useChat/state';
+import { messagesState } from 'hooks/useChat/state';
 
 import { IMessage } from 'src/types/message';
 
@@ -35,15 +35,15 @@ const FeedbackButtons = ({ message }: Props) => {
   const UpIcon = feedback === 1 ? ThumbUpAlt : ThumbUpAltOutlined;
   const [showFeedbackDialog, setShowFeedbackDialog] = useState<number>();
   const [commentInput, setCommentInput] = useState<string>();
-  const setNestedMessages = useSetRecoilState(nestedMessagesState);
+  const setMessages = useSetRecoilState(messagesState);
 
   const handleFeedbackChanged = (feedback: number, comment?: string) => {
     onFeedbackUpdated &&
       onFeedbackUpdated(
         message.id,
         feedback,
         () =>
-          setNestedMessages((prev) =>
+          setMessages((prev) =>
             updateMessageById(prev, message.id, {
               ...message,
               humanFeedback: feedback,

diff --git a/libs/components/src/types/message.ts b/libs/components/src/types/message.ts
@@ -40,9 +40,6 @@ export interface IMessage {
   prompt?: IPrompt;
   streaming?: boolean;
   waitForAnswer?: boolean;
-}
-
-export interface INestedMessage extends IMessage {
   subMessages?: IMessage[];
 }