From ee67c455b155ff244f5698080bb326566a07bc54 Mon Sep 17 00:00:00 2001
From: BetaHuhn <hello@mxis.ch>
Date: Thu, 31 Aug 2023 15:57:12 +0200
Subject: [PATCH] feat: stop generation

---
 src/components/App.tsx        |   2 +-
 src/components/ChatHeader.tsx |  18 +++--
 src/components/Layout.tsx     | 114 ++++----------------------------
 src/routes/ChatRoute.tsx      | 120 ++++++++++++++--------------------
 src/utils/openai.ts           |  11 ++--
 5 files changed, 84 insertions(+), 181 deletions(-)

diff --git a/src/components/App.tsx b/src/components/App.tsx
index c1d6f59..a9bceff 100644
--- a/src/components/App.tsx
+++ b/src/components/App.tsx
@@ -109,7 +109,7 @@ export function App() {
           }}
         >
           <Layout />
-          <Notifications />
+          <Notifications position="bottom-right" style={{ marginBottom: '4rem' }} />
         </MantineProvider>
       </ColorSchemeProvider>
     </Router>
diff --git a/src/components/ChatHeader.tsx b/src/components/ChatHeader.tsx
index 7e1d333..03d5a32 100644
--- a/src/components/ChatHeader.tsx
+++ b/src/components/ChatHeader.tsx
@@ -68,11 +68,19 @@ export function ChatHeader() {
             />
 
             <EditChatModal chat={chat!}>
-            <Tooltip label="Chat Settings">
-                <ActionIcon size="xl">
-                <IconAdjustments size={20} />
-                </ActionIcon>
-            </Tooltip>
+                <Tooltip label="Chat Settings">
+                    <ActionIcon
+                        size="xl"
+                        sx={(theme) => ({
+                            [theme.fn.smallerThan('md')]: {
+                                marginRight: '2.5rem',
+                                paddingBottom: 3,
+                            },
+                        })}
+                    >
+                        <IconAdjustments size={20} />
+                    </ActionIcon>
+                </Tooltip>
             </EditChatModal>
         </Box>
     </Header>
diff --git a/src/components/Layout.tsx b/src/components/Layout.tsx
index 125c258..d8050d3 100644
--- a/src/components/Layout.tsx
+++ b/src/components/Layout.tsx
@@ -153,29 +153,20 @@ export function Layout() {
                           }}
                         />
                       </Link>
-                      <Box
-                        style={{
-                          display: "flex",
-                          alignItems: "center"
-                        }}
-                      >
-                        <SettingsModal>
-                            <Tooltip label="Settings">
-                              <ActionIcon size="xl">
-                                <IconSettings size={20} />
-                              </ActionIcon>
-                            </Tooltip>
-                        </SettingsModal>
-                        <MediaQuery largerThan="md" styles={{ display: "none" }}>
-                          <Burger
-                            opened={opened}
-                            onClick={() => setOpened((o) => !o)}
-                            size="sm"
-                            color={theme.colors.gray[6]}
-                            className="app-region-no-drag"
-                          />
-                        </MediaQuery>
-                      </Box>
+                      <SettingsModal>
+                        <Tooltip label="Settings">
+                          <ActionIcon
+                            size="xl"
+                            sx={(theme) => ({
+                              [theme.fn.smallerThan('md')]: {
+                                marginRight: '2.5rem',
+                              },
+                            })}
+                          >
+                            <IconSettings size={20} />
+                          </ActionIcon>
+                        </Tooltip>
+                      </SettingsModal>
                     </Box>
                   </Navbar.Section>
                   <Navbar.Section
@@ -271,83 +262,6 @@ export function Layout() {
                       {tab === "Prompts" && <CreatePromptModal />}
                     </Box>
                   </Navbar.Section>
-                  {/* <Navbar.Section sx={{ borderTop: border }} p="xs">
-                    <Center>
-                      {config.allowDarkModeToggle && (
-                        <Tooltip
-                          label={colorScheme === "dark" ? "Light Mode" : "Dark Mode"}
-                        >
-                          <ActionIcon
-                            sx={{ flex: 1 }}
-                            size="xl"
-                            onClick={() => toggleColorScheme()}
-                          >
-                            {colorScheme === "dark" ? (
-                              <IconSunHigh size={20} />
-                            ) : (
-                              <IconMoonStars size={20} />
-                            )}
-                          </ActionIcon>
-                        </Tooltip>
-                      )}
-                      {config.allowSettingsModal && (
-                        <SettingsModal>
-                          <Tooltip label="Settings">
-                            <ActionIcon sx={{ flex: 1 }} size="xl">
-                              <IconSettings size={20} />
-                            </ActionIcon>
-                          </Tooltip>
-                        </SettingsModal>
-                      )}
-                      {config.githubUrl && (
-                        <Tooltip label="Source Code">
-                          <ActionIcon
-                            component="a"
-                            href={config.githubUrl}
-                            target="_blank"
-                            sx={{ flex: 1 }}
-                            size="xl"
-                          >
-                            <IconBrandGithub size={20} />
-                          </ActionIcon>
-                        </Tooltip>
-                      )}
-                      {config.showTwitterLink && (
-                        <Tooltip label="Follow on Twitter">
-                          <ActionIcon
-                            component="a"
-                            href="https://twitter.com/deiucanta"
-                            target="_blank"
-                            sx={{ flex: 1 }}
-                            size="xl"
-                          >
-                            <IconBrandTwitter size={20} />
-                          </ActionIcon>
-                        </Tooltip>
-                      )}
-                      {config.showFeedbackLink && (
-                        <Tooltip label="Give Feedback">
-                          <ActionIcon
-                            component="a"
-                            href="https://feedback.chatpad.ai"
-                            onClick={(event) => {
-                              if (window.todesktop) {
-                                event.preventDefault();
-                                window.todesktop.contents.openUrlInBrowser(
-                                  "https://feedback.chatpad.ai"
-                                );
-                              }
-                            }}
-                            target="_blank"
-                            sx={{ flex: 1 }}
-                            size="xl"
-                          >
-                            <IconMessage size={20} />
-                          </ActionIcon>
-                        </Tooltip>
-                      )}
-                    </Center>
-                  </Navbar.Section> */}
                 </Navbar>
               }
               header={
diff --git a/src/routes/ChatRoute.tsx b/src/routes/ChatRoute.tsx
index 1c36f5a..3a5d2be 100644
--- a/src/routes/ChatRoute.tsx
+++ b/src/routes/ChatRoute.tsx
@@ -4,7 +4,6 @@ import {
   Card,
   Container,
   Flex,
-  MediaQuery,
   Select,
   Skeleton,
   Stack,
@@ -12,7 +11,7 @@ import {
 } from "@mantine/core";
 import { notifications } from "@mantine/notifications";
 import { KeyboardEvent, useState, type ChangeEvent, useEffect } from "react";
-import { AiOutlineSend } from "react-icons/ai";
+import { IconSend, IconPlayerStop } from "@tabler/icons-react";
 import { MessageItem } from "../components/MessageItem";
 import { Message, Prompt, detaDB, generateKey } from "../db";
 import { useChatId } from "../hooks/useChatId";
@@ -56,6 +55,9 @@ export function ChatRoute() {
   const [content, setContent] = useState("");
   const [contentDraft, setContentDraft] = useState("");
   const [submitting, setSubmitting] = useState(false);
+  const [generating, setGenerating] = useState(false);
+  const [coolDown, setCoolDown] = useState(false);
+  const [chatStream, setChatStream] = useState<any>(null);
   const [promptKey, setPromptKey] = useState<string | null>(null);
   const [newPromptTitle, setNewPromptTitle] = useState<string | null>(null);
   const [selectedModel, setSelectedModel] = useState<string | null>(null);
@@ -69,13 +71,8 @@ export function ChatRoute() {
     }
   }, [chat]);
 
-  // const chat = useLiveQuery(async () => {
-  //   if (!chatId) return null;
-  //   return db.chats.get(chatId);
-  // }, [chatId]);
-
   const submit = async () => {
-    if (submitting) return;
+    if (submitting || coolDown) return;
 
     if (!chatId) {
       notifications.show({
@@ -161,9 +158,10 @@ export function ChatRoute() {
       }, generateKey())
 
       setMessages(current => [...current, systemMessage as unknown as Message])
+      setGenerating(true)
 
       const messageId = systemMessage!.key as string
-      await createStreamChatCompletion(
+      const chatCompletionStream = createStreamChatCompletion(
         { ...settings, openAiModel: model },
         [
           {
@@ -186,9 +184,16 @@ export function ChatRoute() {
       
             return message;
           }));
-        }
+        },
+        () => {
+          setGenerating(false)
+          setChatStream(null)
+          setCoolDown(true);
+          setTimeout(() => setCoolDown(false), 500);
+        },
       );
 
+      setChatStream(chatCompletionStream)
       setSubmitting(false);
 
       if (chat?.description === "New Chat") {
@@ -247,6 +252,7 @@ export function ChatRoute() {
         }
       }
     } catch (error: any) {
+      setGenerating(false)
       if (error.toJSON().message === "Network Error") {
         notifications.show({
           title: "Error",
@@ -267,6 +273,13 @@ export function ChatRoute() {
     }
   };
 
+  const stopGeneration = async () => {
+    if (chatStream) {
+      chatStream.abort()
+      setChatStream(null)
+    }
+  }
+
   const onUserMsgToggle = (event: KeyboardEvent<HTMLTextAreaElement>) => {
     const { selectionStart, selectionEnd } = event.currentTarget;
     if (
@@ -345,7 +358,7 @@ export function ChatRoute() {
               mb="sm"
               style={{
                 display: "flex",
-                justifyContent: "flex-end",
+                justifyContent: "flex-start",
                 alignItems: "center",
                 gap: 10,
               }}
@@ -384,66 +397,18 @@ export function ChatRoute() {
 
               {newPromptTitle && <CreatePromptModal title={newPromptTitle} open={true} />}
             </Box>
-            // <SimpleGrid
-            //   mb="sm"
-            //   spacing="xs"
-            //   breakpoints={[
-            //     { minWidth: "sm", cols: 4 },
-            //     { maxWidth: "sm", cols: 2 },
-            //   ]}
-            // >
-            //   <Select
-            //     value={writingCharacter}
-            //     onChange={setWritingCharacter}
-            //     data={config.writingCharacters}
-            //     placeholder="Character"
-            //     variant="filled"
-            //     searchable
-            //     clearable
-            //     sx={{ flex: 1 }}
-            //   />
-            //   <Select
-            //     value={writingTone}
-            //     onChange={setWritingTone}
-            //     data={config.writingTones}
-            //     placeholder="Tone"
-            //     variant="filled"
-            //     searchable
-            //     clearable
-            //     sx={{ flex: 1 }}
-            //   />
-            //   <Select
-            //     value={writingStyle}
-            //     onChange={setWritingStyle}
-            //     data={config.writingStyles}
-            //     placeholder="Style"
-            //     variant="filled"
-            //     searchable
-            //     clearable
-            //     sx={{ flex: 1 }}
-            //   />
-            //   <Select
-            //     value={writingFormat}
-            //     onChange={setWritingFormat}
-            //     data={config.writingFormats}
-            //     placeholder="Format"
-            //     variant="filled"
-            //     searchable
-            //     clearable
-            //     sx={{ flex: 1 }}
-            //   />
-            // </SimpleGrid>
           )}
-          <Flex gap="sm">
+          <Flex gap="sm" style={{ position: 'relative' }}>
             <Textarea
               key={chatId}
               sx={{ flex: 1 }}
+              styles={{ input: { paddingRight: 60, paddingTop: '0.7rem !important', paddingBottom: '0.7rem !important' } }}
               placeholder="Your message here..."
               autosize
               autoFocus
               disabled={submitting}
               minRows={1}
-              maxRows={5}
+              maxRows={8}
               value={content}
               onChange={onContentChange}
               onKeyDown={async (event) => {
@@ -460,16 +425,29 @@ export function ChatRoute() {
                 }
               }}
             />
-            <MediaQuery largerThan="sm" styles={{ display: "none" }}>
-              <Button
-                h="auto"
-                onClick={() => {
+            <Button
+              w="auto"
+              h="auto"
+              px={8}
+              py={5}
+              onClick={() => {
+                if (generating) {
+                  stopGeneration();
+
+                  notifications.show({
+                    title: "Stopped",
+                    color: "green",
+                    message: "Generation stopped.",
+                  });
+                } else {
                   submit();
-                }}
-              >
-                <AiOutlineSend />
-              </Button>
-            </MediaQuery>
+                }
+              }}
+              style={{ position: 'absolute', right: 5, bottom: 5 }}
+              disabled={!generating && !content}
+            >
+              {generating ? (<IconPlayerStop />) : (<IconSend />)}
+            </Button>
           </Flex>
         </Container>
       </Box>
diff --git a/src/utils/openai.ts b/src/utils/openai.ts
index 6120f73..ab8562b 100644
--- a/src/utils/openai.ts
+++ b/src/utils/openai.ts
@@ -21,12 +21,13 @@ function getClient(
   return new OpenAIApi(configuration);
 }
 
-export async function createStreamChatCompletion(
+export function createStreamChatCompletion(
   settings: Settings,
   messages: ChatCompletionRequestMessage[],
   chatId: string,
   messageId: string,
-  onContent: (content: string, isFinal: boolean) => void
+  onContent: (content: string, isFinal: boolean) => void,
+  onDone: () => void
 ) {
   const model = settings?.openAiModel ?? config.defaultModel;
 
@@ -47,9 +48,11 @@ export async function createStreamChatCompletion(
           }
           onContent(content, isFinal);
         },
-        onDone() {},
+        onDone() {
+          onDone()
+        },
         onError(error) {
-          console.error(error);
+          throw error;
         },
       },
     }