From fb910d0a52b459b1b2cd383cdfb060ab7d466928 Mon Sep 17 00:00:00 2001
From: Jing Hua <tohjinghua123@gmail.com>
Date: Sat, 4 Mar 2023 13:43:20 +0800
Subject: [PATCH] improve useSubmitHook

---
 .../ChatContent/Message/MessageContent.tsx    | 78 +++++++++++++------
 src/hooks/useSubmit.ts                        | 35 +++------
 2 files changed, 66 insertions(+), 47 deletions(-)

diff --git a/src/components/Chat/ChatContent/Message/MessageContent.tsx b/src/components/Chat/ChatContent/Message/MessageContent.tsx
index bf8bc4d..67c2d81 100644
--- a/src/components/Chat/ChatContent/Message/MessageContent.tsx
+++ b/src/components/Chat/ChatContent/Message/MessageContent.tsx
@@ -12,10 +12,10 @@ 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';
 
+import { MessageInterface } from '@type/chat';
+
 const MessageContent = ({
   role,
   content,
@@ -91,6 +91,13 @@ const ContentView = ({
     setMessages(updatedMessages);
   };
 
+  const handleRefresh = () => {
+    const updatedMessages = JSON.parse(JSON.stringify(messages));
+    updatedMessages.splice(updatedMessages.length - 1, 1);
+    setMessages(updatedMessages);
+    handleSubmit();
+  };
+
   return (
     <>
       <div className='markdown prose w-full break-words dark:prose-invert dark'>
@@ -161,11 +168,7 @@ const ContentView = ({
         {isDelete || (
           <>
             {role === 'assistant' && messageIndex === messages?.length - 1 && (
-              <RefreshButton
-                onClick={() => {
-                  handleSubmit(true);
-                }}
-              />
+              <RefreshButton onClick={handleRefresh} />
             )}
             {messageIndex !== 0 && (
               <UpButton onClick={() => handleMove('up')} />
@@ -289,6 +292,34 @@ const EditView = ({
     }
   };
 
+  const handleSave = () => {
+    if (_content === '') return;
+    const updatedMessages: MessageInterface[] = JSON.parse(
+      JSON.stringify(messages)
+    );
+    if (sticky) {
+      updatedMessages.push({ role: inputRole, content: _content });
+      _setContent('');
+    } else {
+      updatedMessages[messageIndex].content = _content;
+      setIsEdit(false);
+    }
+    setMessages(updatedMessages);
+  };
+
+  // only for sticky forms
+  const { handleSubmit } = useSubmit();
+  const handleSaveAndSubmit = () => {
+    if (_content == '') return;
+    const updatedMessages: MessageInterface[] = JSON.parse(
+      JSON.stringify(messages)
+    );
+    updatedMessages.push({ role: inputRole, content: _content });
+    _setContent('');
+    setMessages(updatedMessages);
+    handleSubmit();
+  };
+
   useEffect(() => {
     if (textareaRef.current) {
       textareaRef.current.style.height = 'auto';
@@ -308,28 +339,29 @@ const EditView = ({
         }}
         value={_content}
         onInput={handleInput}
-        rows={sticky ? 10 : 1}
+        rows={1}
       ></textarea>
       <div className='text-center mt-2 flex justify-center'>
+        {sticky && (
+          <button
+            className='btn relative btn-primary mr-2'
+            onClick={handleSaveAndSubmit}
+          >
+            <div className='flex items-center justify-center gap-2'>
+              Save and Submit
+            </div>
+          </button>
+        )}
+
         <button
-          className='btn relative btn-primary mr-2'
-          onClick={() => {
-            if (_content === '') return;
-            const updatedMessages: MessageInterface[] = JSON.parse(
-              JSON.stringify(messages)
-            );
-            if (sticky) {
-              updatedMessages.push({ role: inputRole, content: _content });
-              _setContent('');
-            } else {
-              updatedMessages[messageIndex].content = _content;
-              setIsEdit(false);
-            }
-            setMessages(updatedMessages);
-          }}
+          className={`btn relative mr-2 ${
+            sticky ? 'btn-neutral' : 'btn-primary'
+          }`}
+          onClick={handleSave}
         >
           <div className='flex items-center justify-center gap-2'>Save</div>
         </button>
+
         {sticky || (
           <button
             className='btn relative btn-neutral'
diff --git a/src/hooks/useSubmit.ts b/src/hooks/useSubmit.ts
index 7b886ad..517668d 100644
--- a/src/hooks/useSubmit.ts
+++ b/src/hooks/useSubmit.ts
@@ -7,48 +7,35 @@ import { parseEventSource } from '@api/helper';
 
 const useSubmit = () => {
   const [error, setError] = useState<string>('');
-  const [messages, apiFree, apiKey, setMessages, setGenerating, generating] =
-    useStore((state) => [
-      state.messages,
+  const [apiFree, apiKey, setMessages, setGenerating, generating] = useStore(
+    (state) => [
       state.apiFree,
       state.apiKey,
       state.setMessages,
       state.setGenerating,
       state.generating,
-    ]);
+    ]
+  );
 
-  const handleSubmit = async (refresh?: boolean) => {
+  const handleSubmit = async () => {
     if (generating) return;
+    const messages = useStore.getState().messages;
 
     const updatedMessages: MessageInterface[] = JSON.parse(
       JSON.stringify(messages)
     );
-    if (refresh) {
-      updatedMessages[updatedMessages.length - 1] = {
-        role: 'assistant',
-        content: '',
-      };
-    } else {
-      updatedMessages.push({ role: 'assistant', content: '' });
-    }
+
+    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);
+        stream = await getChatCompletionStreamFree(messages);
       } else if (apiKey) {
-        if (refresh)
-          stream = await getChatCompletionStreamCustom(
-            apiKey,
-            updatedMessages.slice(0, updatedMessages.length - 1)
-          );
-        else stream = await getChatCompletionStreamFree(messages);
+        stream = await getChatCompletionStreamFree(messages);
       }
 
       if (stream) {