From 710f4c3f653f5ecd7e4d8ba199d2ea9ba9fd0138 Mon Sep 17 00:00:00 2001
From: Jing Hua <tohjinghua123@gmail.com>
Date: Sat, 4 Mar 2023 10:46:41 +0800
Subject: [PATCH] feat: regenerate response button

---
 .../Chat/ChatContent/ChatContent.tsx          | 81 ++--------------
 .../Chat/ChatContent/Message/Message.tsx      |  1 +
 .../ChatContent/Message/MessageContent.tsx    | 26 +++++
 src/hooks/useSubmit.ts                        | 96 +++++++++++++++++++
 4 files changed, 130 insertions(+), 74 deletions(-)
 create mode 100644 src/hooks/useSubmit.ts

diff --git a/src/components/Chat/ChatContent/ChatContent.tsx b/src/components/Chat/ChatContent/ChatContent.tsx
index a72b5c3..35360f1 100644
--- a/src/components/Chat/ChatContent/ChatContent.tsx
+++ b/src/components/Chat/ChatContent/ChatContent.tsx
@@ -1,4 +1,4 @@
-import React, { useState } from 'react';
+import React, { createRef, useState } from 'react';
 import ScrollToBottom from 'react-scroll-to-bottom';
 import useStore from '@store/store';
 
@@ -13,83 +13,14 @@ import { parseEventSource } from '@api/helper';
 
 import RefreshIcon from '@icon/RefreshIcon';
 import { MessageInterface } from '@type/chat';
+import useSubmit from '@hooks/useSubmit';
 
 const ChatContent = () => {
-  const [
-    messages,
-    inputRole,
-    apiFree,
-    apiKey,
-    setMessages,
-    setGenerating,
-    generating,
-  ] = useStore((state) => [
+  const [messages, inputRole] = useStore((state) => [
     state.messages,
     state.inputRole,
-    state.apiFree,
-    state.apiKey,
-    state.setMessages,
-    state.setGenerating,
-    state.generating,
   ]);
-  const [error, setError] = useState<string>('');
-
-  const handleSubmit = async () => {
-    if (generating) return;
-
-    const updatedMessages: MessageInterface[] = JSON.parse(
-      JSON.stringify(messages)
-    );
-    updatedMessages.push({ role: 'assistant', content: '' });
-    setMessages(updatedMessages);
-    setGenerating(true);
-    let stream;
-
-    try {
-      if (apiFree) {
-        stream = await getChatCompletionStreamFree(messages);
-      } else if (apiKey) {
-        stream = await getChatCompletionStreamCustom(apiKey, messages);
-      }
-
-      if (stream) {
-        const reader = stream.getReader();
-        let reading = true;
-        while (reading) {
-          const { done, value } = await reader.read();
-
-          const result = parseEventSource(new TextDecoder().decode(value));
-
-          if (result === '[DONE]' || done) {
-            reading = false;
-          } else {
-            const resultString = result.reduce((output: string, curr) => {
-              if (typeof curr === 'string') return output;
-              else {
-                const content = curr.choices[0].delta.content;
-                if (content) output += content;
-                return output;
-              }
-            }, '');
-
-            const updatedMessages: MessageInterface[] = JSON.parse(
-              JSON.stringify(useStore.getState().messages)
-            );
-            updatedMessages[updatedMessages.length - 1].content += resultString;
-            setMessages(updatedMessages);
-          }
-        }
-      }
-    } catch (e: unknown) {
-      const err = (e as Error).message;
-      console.log(err);
-      setError(err);
-      setTimeout(() => {
-        setError(''), 10000;
-      });
-    }
-    setGenerating(false);
-  };
+  const { handleSubmit, error } = useSubmit();
 
   return (
     <div className='flex-1 overflow-hidden'>
@@ -122,7 +53,9 @@ const ChatContent = () => {
           <div className='text-center mt-6 flex justify-center gap-2'>
             <button
               className='btn relative btn-primary mt-2'
-              onClick={handleSubmit}
+              onClick={() => {
+                handleSubmit();
+              }}
             >
               Submit
             </button>
diff --git a/src/components/Chat/ChatContent/Message/Message.tsx b/src/components/Chat/ChatContent/Message/Message.tsx
index 68f67b0..2b7c68b 100644
--- a/src/components/Chat/ChatContent/Message/Message.tsx
+++ b/src/components/Chat/ChatContent/Message/Message.tsx
@@ -41,6 +41,7 @@ const Message = ({
             sticky={sticky}
           />
           <MessageContent
+            role={role}
             content={content}
             messageIndex={messageIndex}
             sticky={sticky}
diff --git a/src/components/Chat/ChatContent/Message/MessageContent.tsx b/src/components/Chat/ChatContent/Message/MessageContent.tsx
index 71cfc43..ba708bf 100644
--- a/src/components/Chat/ChatContent/Message/MessageContent.tsx
+++ b/src/components/Chat/ChatContent/Message/MessageContent.tsx
@@ -9,15 +9,20 @@ import EditIcon2 from '@icon/EditIcon2';
 import DeleteIcon from '@icon/DeleteIcon';
 import TickIcon from '@icon/TickIcon';
 import CrossIcon from '@icon/CrossIcon';
+import RefreshIcon from '@icon/RefreshIcon';
 import DownChevronArrow from '@icon/DownChevronArrow';
 
 import { MessageInterface } from '@type/chat';
 
+import useSubmit from '@hooks/useSubmit';
+
 const MessageContent = ({
+  role,
   content,
   messageIndex,
   sticky = false,
 }: {
+  role: string;
   content: string;
   messageIndex: number;
   sticky?: boolean;
@@ -36,6 +41,7 @@ const MessageContent = ({
         />
       ) : (
         <ContentView
+          role={role}
           content={content}
           setIsEdit={setIsEdit}
           messageIndex={messageIndex}
@@ -46,14 +52,18 @@ const MessageContent = ({
 };
 
 const ContentView = ({
+  role,
   content,
   setIsEdit,
   messageIndex,
 }: {
+  role: string;
   content: string;
   setIsEdit: React.Dispatch<React.SetStateAction<boolean>>;
   messageIndex: number;
 }) => {
+  const { handleSubmit, error } = useSubmit();
+
   const [isDelete, setIsDelete] = useState<boolean>(false);
   const [copied, setCopied] = useState<boolean>(false);
 
@@ -150,12 +160,20 @@ const ContentView = ({
       <div className='flex justify-end gap-2 w-full mt-2'>
         {isDelete || (
           <>
+            {role === 'assistant' && messageIndex === messages?.length - 1 && (
+              <RefreshButton
+                onClick={() => {
+                  handleSubmit(true);
+                }}
+              />
+            )}
             {messageIndex !== 0 && (
               <UpButton onClick={() => handleMove('up')} />
             )}
             {messageIndex !== messages?.length - 1 && (
               <DownButton onClick={() => handleMove('down')} />
             )}
+
             <EditButton setIsEdit={setIsEdit} />
             <DeleteButton setIsDelete={setIsDelete} />
           </>
@@ -236,6 +254,14 @@ const UpButton = ({
   );
 };
 
+const RefreshButton = ({
+  onClick,
+}: {
+  onClick: React.MouseEventHandler<HTMLButtonElement>;
+}) => {
+  return <MessageButton icon={<RefreshIcon />} onClick={onClick} />;
+};
+
 const EditView = ({
   content,
   setIsEdit,
diff --git a/src/hooks/useSubmit.ts b/src/hooks/useSubmit.ts
new file mode 100644
index 0000000..7b886ad
--- /dev/null
+++ b/src/hooks/useSubmit.ts
@@ -0,0 +1,96 @@
+import React, { useState } from 'react';
+import useStore from '@store/store';
+import { MessageInterface } from '@type/chat';
+import { getChatCompletionStream as getChatCompletionStreamFree } from '@api/freeApi';
+import { getChatCompletionStream as getChatCompletionStreamCustom } from '@api/customApi';
+import { parseEventSource } from '@api/helper';
+
+const useSubmit = () => {
+  const [error, setError] = useState<string>('');
+  const [messages, apiFree, apiKey, setMessages, setGenerating, generating] =
+    useStore((state) => [
+      state.messages,
+      state.apiFree,
+      state.apiKey,
+      state.setMessages,
+      state.setGenerating,
+      state.generating,
+    ]);
+
+  const handleSubmit = async (refresh?: boolean) => {
+    if (generating) return;
+
+    const updatedMessages: MessageInterface[] = JSON.parse(
+      JSON.stringify(messages)
+    );
+    if (refresh) {
+      updatedMessages[updatedMessages.length - 1] = {
+        role: 'assistant',
+        content: '',
+      };
+    } else {
+      updatedMessages.push({ role: 'assistant', content: '' });
+    }
+    setMessages(updatedMessages);
+    setGenerating(true);
+    let stream;
+
+    try {
+      if (apiFree) {
+        if (refresh)
+          stream = await getChatCompletionStreamFree(
+            updatedMessages.slice(0, updatedMessages.length - 1)
+          );
+        else stream = await getChatCompletionStreamFree(messages);
+      } else if (apiKey) {
+        if (refresh)
+          stream = await getChatCompletionStreamCustom(
+            apiKey,
+            updatedMessages.slice(0, updatedMessages.length - 1)
+          );
+        else stream = await getChatCompletionStreamFree(messages);
+      }
+
+      if (stream) {
+        const reader = stream.getReader();
+        let reading = true;
+        while (reading) {
+          const { done, value } = await reader.read();
+
+          const result = parseEventSource(new TextDecoder().decode(value));
+
+          if (result === '[DONE]' || done) {
+            reading = false;
+          } else {
+            const resultString = result.reduce((output: string, curr) => {
+              if (typeof curr === 'string') return output;
+              else {
+                const content = curr.choices[0].delta.content;
+                if (content) output += content;
+                return output;
+              }
+            }, '');
+
+            const updatedMessages: MessageInterface[] = JSON.parse(
+              JSON.stringify(useStore.getState().messages)
+            );
+            updatedMessages[updatedMessages.length - 1].content += resultString;
+            setMessages(updatedMessages);
+          }
+        }
+      }
+    } catch (e: unknown) {
+      const err = (e as Error).message;
+      console.log(err);
+      setError(err);
+      setTimeout(() => {
+        setError(''), 10000;
+      });
+    }
+    setGenerating(false);
+  };
+
+  return { handleSubmit, error };
+};
+
+export default useSubmit;