From 7f1bcb1c9b932b0add34219e5cc2fba86ae4ca49 Mon Sep 17 00:00:00 2001
From: baaalint <balint_kantor@mckinsey.com>
Date: Fri, 6 Sep 2024 13:10:50 +0200
Subject: [PATCH] feat(chat): disable message send upon error

---
 ui/src/atoms/index.ts                         |  1 +
 ui/src/components/feature/ChatSessionList.tsx |  4 ++-
 ui/src/components/shared/Message.tsx          | 29 +++++++++++++++----
 3 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/ui/src/atoms/index.ts b/ui/src/atoms/index.ts
index d5e96e0..8066cc4 100644
--- a/ui/src/atoms/index.ts
+++ b/ui/src/atoms/index.ts
@@ -26,6 +26,7 @@ export const userWithTokenAtom = atomWithStorage<User | null>('user', localStora
 export const publicUserAtom = atom<User>({});
 export const usernameAtom = atom<string>('');
 export const isTypingAtom = atom<boolean>(false);
+export const canSendMessageAtom = atom<boolean>(true);
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 export const selectedRowAtom = atom<any>({});
diff --git a/ui/src/components/feature/ChatSessionList.tsx b/ui/src/components/feature/ChatSessionList.tsx
index 5e9840d..f80c41d 100644
--- a/ui/src/components/feature/ChatSessionList.tsx
+++ b/ui/src/components/feature/ChatSessionList.tsx
@@ -28,7 +28,7 @@ import {
 import Client from '@services/Api'
 import { colors } from '@shared/theme'
 import { Session } from '@shared/types/session'
-import { isTypingAtom, messagesAtom, sessionIdAtom, usernameAtom } from 'atoms'
+import { canSendMessageAtom, isTypingAtom, messagesAtom, sessionIdAtom, usernameAtom } from 'atoms'
 import { useAtom } from 'jotai'
 import { useEffect, useRef, useState } from 'react'
 import { useLocation, useNavigate } from 'react-router-dom'
@@ -51,6 +51,7 @@ const ChatSessionList = (props: Props) => {
   const [isEditing, setIsEditing] = useState(false)
   const [description, setDescription] = useState('')
   const [, fetchSessions] = useAtom(sessionsWithFetchAtom)
+  const [, setCanSendMessage] = useAtom(canSendMessageAtom)
   const inputRef = useRef<HTMLInputElement>(null)
 
   const { colorMode } = useColorMode()
@@ -68,6 +69,7 @@ const ChatSessionList = (props: Props) => {
     props.setNew(false)
     setIsTyping(false)
     setMessages(session.history!)
+    setCanSendMessage(true)
     navigate(`/chat/${session.uid}`)
   }
 
diff --git a/ui/src/components/shared/Message.tsx b/ui/src/components/shared/Message.tsx
index 2de1252..7fb0768 100644
--- a/ui/src/components/shared/Message.tsx
+++ b/ui/src/components/shared/Message.tsx
@@ -13,18 +13,23 @@
 // limitations under the License.
 
 import { ArrowUpIcon, AttachmentIcon } from '@chakra-ui/icons'
-import { Flex, IconButton, Input } from '@chakra-ui/react'
+import { Flex, IconButton, Input, useToast } from '@chakra-ui/react'
 import Client from '@services/Api'
-import { messagesAtom, sessionIdAtom } from 'atoms'
+import { canSendMessageAtom, messagesAtom, sessionIdAtom } from 'atoms'
 import { useAtom } from 'jotai'
 import { useState } from 'react'
 
 const Message = () => {
   const [inputValue, setInputValue] = useState('')
+
   const [sessionId] = useAtom(sessionIdAtom)
   const [, setMessages] = useAtom(messagesAtom)
+  const [canSendMessage, setCanSendMessage] = useAtom(canSendMessageAtom)
+
+  const toast = useToast()
 
   const submitMessage = async () => {
+    setCanSendMessage(false)
     setMessages(prevMessages => {
       const safeMessages = Array.isArray(prevMessages) ? prevMessages : []
       return [...safeMessages, { role: 'Human', content: inputValue, sources: [] }]
@@ -40,10 +45,22 @@ const Message = () => {
       question: inputValue,
       session_id: sessionId,
       data_source: 'default'
+    }).then(res => {
+      if (res.error) {
+        toast({
+          title: 'An unexpected error occured',
+          description: res.error,
+          status: 'error',
+          duration: 5000,
+          isClosable: true
+        })
+        setCanSendMessage(false)
+        return res
+      }
+      setCanSendMessage(true)
+      return res
     })
 
-    console.log('RESULT', result)
-
     setMessages(prevMessages => {
       const safeMessages = Array.isArray(prevMessages) ? prevMessages : []
       return [...safeMessages.slice(0, -1), { role: 'AI', content: result.data.data.answer, sources: result.sources }]
@@ -71,9 +88,9 @@ const Message = () => {
         placeholder="Send message..."
         value={inputValue}
         onChange={e => setInputValue(e.target.value)}
-        onKeyDown={handleKeyPress}
+        onKeyDown={e => canSendMessage && handleKeyPress(e)}
       />
-      <IconButton aria-label="Send" icon={<ArrowUpIcon />} onClick={handleClick} />
+      <IconButton isDisabled={!canSendMessage} aria-label="Send" icon={<ArrowUpIcon />} onClick={handleClick} />
     </Flex>
   )
 }