From af22dfb6864907941052813944a36db52a85cde1 Mon Sep 17 00:00:00 2001
From: SuperTurk <alim@refty.co>
Date: Sat, 6 Jan 2024 11:12:25 +0100
Subject: [PATCH] feat: New way to stop the loading task (#631)

---
 .../organisms/chat/inputBox/SubmitButton.tsx  | 49 +++++++++++++++++++
 .../organisms/chat/inputBox/index.tsx         |  2 -
 .../organisms/chat/inputBox/input.tsx         | 16 +-----
 .../components/organisms/chat/stopButton.tsx  | 31 ------------
 4 files changed, 51 insertions(+), 47 deletions(-)
 create mode 100644 frontend/src/components/organisms/chat/inputBox/SubmitButton.tsx
 delete mode 100644 frontend/src/components/organisms/chat/stopButton.tsx

diff --git a/frontend/src/components/organisms/chat/inputBox/SubmitButton.tsx b/frontend/src/components/organisms/chat/inputBox/SubmitButton.tsx
new file mode 100644
index 0000000000..2d3f3b4e44
--- /dev/null
+++ b/frontend/src/components/organisms/chat/inputBox/SubmitButton.tsx
@@ -0,0 +1,49 @@
+import StopCircle from '@mui/icons-material/StopCircle';
+import Telegram from '@mui/icons-material/Telegram';
+import { Tooltip } from '@mui/material';
+import Box from '@mui/material/Box';
+import IconButton from '@mui/material/IconButton';
+import InputAdornment from '@mui/material/InputAdornment';
+
+import { useChatData, useChatInteract } from '@chainlit/react-client';
+
+interface SubmitButtonProps {
+  disabled?: boolean;
+  onSubmit: () => void;
+}
+
+const SubmitButton = ({ disabled, onSubmit }: SubmitButtonProps) => {
+  const { loading } = useChatData();
+  const { stopTask } = useChatInteract();
+
+  const handleClick = () => {
+    stopTask();
+  };
+
+  return (
+    <Box
+      sx={{
+        mr: 1,
+        color: 'text.secondary'
+      }}
+    >
+      {!loading ? (
+        <Tooltip title="Send message">
+          <InputAdornment position="end">
+            <IconButton disabled={disabled} color="inherit" onClick={onSubmit}>
+              <Telegram />
+            </IconButton>
+          </InputAdornment>
+        </Tooltip>
+      ) : (
+        <Tooltip title="Stop task">
+          <IconButton id="stop-button" onClick={handleClick}>
+            <StopCircle />
+          </IconButton>
+        </Tooltip>
+      )}
+    </Box>
+  );
+};
+
+export { SubmitButton };
diff --git a/frontend/src/components/organisms/chat/inputBox/index.tsx b/frontend/src/components/organisms/chat/inputBox/index.tsx
index 7d3239f723..0a55c11703 100644
--- a/frontend/src/components/organisms/chat/inputBox/index.tsx
+++ b/frontend/src/components/organisms/chat/inputBox/index.tsx
@@ -11,7 +11,6 @@ import { IAttachment } from 'state/chat';
 import { IProjectSettings } from 'state/project';
 import { inputHistoryState } from 'state/userInputHistory';
 
-import StopButton from '../stopButton';
 import Input from './input';
 import WaterMark from './waterMark';
 
@@ -106,7 +105,6 @@ const InputBox = memo(
           justifyContent: 'center'
         }}
       >
-        <StopButton />
         <Box>
           <Input
             fileSpec={fileSpec}
diff --git a/frontend/src/components/organisms/chat/inputBox/input.tsx b/frontend/src/components/organisms/chat/inputBox/input.tsx
index 15793ab1b6..13b08c30ee 100644
--- a/frontend/src/components/organisms/chat/inputBox/input.tsx
+++ b/frontend/src/components/organisms/chat/inputBox/input.tsx
@@ -2,7 +2,6 @@ import { memo, useCallback, useEffect, useRef, useState } from 'react';
 import { useRecoilState, useSetRecoilState } from 'recoil';
 import 'regenerator-runtime';
 
-import SendIcon from '@mui/icons-material/Telegram';
 import TuneIcon from '@mui/icons-material/Tune';
 import { Box, IconButton, Stack, TextField } from '@mui/material';
 import InputAdornment from '@mui/material/InputAdornment';
@@ -16,6 +15,7 @@ import { IAttachment, attachmentsState } from 'state/chat';
 import { chatSettingsOpenState, projectSettingsState } from 'state/project';
 import { inputHistoryState } from 'state/userInputHistory';
 
+import { SubmitButton } from './SubmitButton';
 import UploadButton from './UploadButton';
 import SpeechButton from './speechButton';
 
@@ -164,11 +164,6 @@ const Input = memo(
         />
       </>
     );
-    const endAdornment = (
-      <IconButton disabled={disabled} color="inherit" onClick={() => submit()}>
-        <SendIcon />
-      </IconButton>
-    );
 
     return (
       <Stack
@@ -226,14 +221,7 @@ const Input = memo(
                 {startAdornment}
               </InputAdornment>
             ),
-            endAdornment: (
-              <InputAdornment
-                position="end"
-                sx={{ mr: 1, color: 'text.secondary' }}
-              >
-                {endAdornment}
-              </InputAdornment>
-            )
+            endAdornment: <SubmitButton onSubmit={submit} disabled={disabled} />
           }}
         />
       </Stack>
diff --git a/frontend/src/components/organisms/chat/stopButton.tsx b/frontend/src/components/organisms/chat/stopButton.tsx
deleted file mode 100644
index 7403009f7e..0000000000
--- a/frontend/src/components/organisms/chat/stopButton.tsx
+++ /dev/null
@@ -1,31 +0,0 @@
-import CloseIcon from '@mui/icons-material/Close';
-import { Box } from '@mui/material';
-
-import { useChatData, useChatInteract } from '@chainlit/react-client';
-import { GreyButton } from '@chainlit/react-components';
-
-export default function StopButton() {
-  const { loading } = useChatData();
-  const { stopTask } = useChatInteract();
-
-  if (!loading) {
-    return null;
-  }
-
-  const handleClick = () => {
-    stopTask();
-  };
-
-  return (
-    <Box margin="auto">
-      <GreyButton
-        id="stop-button"
-        startIcon={<CloseIcon />}
-        variant="contained"
-        onClick={handleClick}
-      >
-        Stop task
-      </GreyButton>
-    </Box>
-  );
-}