From c2799c9494b2306aaea656c019fd02b8cfff718a Mon Sep 17 00:00:00 2001
From: josc146 <josStorer@outlook.com>
Date: Tue, 26 Mar 2024 15:02:27 +0800
Subject: [PATCH] add additional finish conditions to provide better ollama
 support

---
 frontend/src/pages/Chat.tsx        | 23 ++++++++++++++++-------
 frontend/src/pages/Completion.tsx  | 13 +++++++++++--
 frontend/src/pages/Composition.tsx | 15 ++++++++++++---
 3 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/frontend/src/pages/Chat.tsx b/frontend/src/pages/Chat.tsx
index 731c223..03ff1db 100644
--- a/frontend/src/pages/Chat.tsx
+++ b/frontend/src/pages/Chat.tsx
@@ -545,6 +545,16 @@ const ChatPanel: FC = observer(() => {
     commonStore.setConversationOrder(commonStore.conversationOrder);
     setTimeout(scrollToBottom);
     let answer = '';
+    let finished = false;
+    const finish = () => {
+      finished = true;
+      if (answerId! in chatSseControllers)
+        delete chatSseControllers[answerId!];
+      commonStore.conversation[answerId!].done = true;
+      commonStore.conversation[answerId!].content = commonStore.conversation[answerId!].content.trim();
+      commonStore.setConversation(commonStore.conversation);
+      commonStore.setConversationOrder([...commonStore.conversationOrder]);
+    };
     const chatSseController = new AbortController();
     chatSseControllers[answerId] = chatSseController;
     fetchEventSource( // https://api.openai.com/v1/chat/completions || http://127.0.0.1:${port}/v1/chat/completions
@@ -572,13 +582,8 @@ const ChatPanel: FC = observer(() => {
         signal: chatSseController?.signal,
         onmessage(e) {
           scrollToBottom();
-          if (e.data.trim() === '[DONE]') {
-            if (answerId! in chatSseControllers)
-              delete chatSseControllers[answerId!];
-            commonStore.conversation[answerId!].done = true;
-            commonStore.conversation[answerId!].content = commonStore.conversation[answerId!].content.trim();
-            commonStore.setConversation(commonStore.conversation);
-            commonStore.setConversationOrder([...commonStore.conversationOrder]);
+          if (!finished && e.data.trim() === '[DONE]') {
+            finish();
             return;
           }
           let data;
@@ -591,6 +596,10 @@ const ChatPanel: FC = observer(() => {
           if (data.model)
             commonStore.setLastModelName(data.model);
           if (data.choices && Array.isArray(data.choices) && data.choices.length > 0) {
+            if (!finished && data.choices[0]?.finish_reason) {
+              finish();
+              return;
+            }
             answer += data.choices[0]?.delta?.content || '';
             commonStore.conversation[answerId!].content = answer;
             commonStore.setConversation(commonStore.conversation);
diff --git a/frontend/src/pages/Completion.tsx b/frontend/src/pages/Completion.tsx
index 5b614c2..a6484f2 100644
--- a/frontend/src/pages/Completion.tsx
+++ b/frontend/src/pages/Completion.tsx
@@ -80,6 +80,11 @@ const CompletionPanel: FC = observer(() => {
     prompt += params.injectStart.replaceAll('\\n', '\n');
 
     let answer = '';
+    let finished = false;
+    const finish = () => {
+      finished = true;
+      commonStore.setCompletionGenerating(false);
+    };
     completionSseController = new AbortController();
     fetchEventSource( // https://api.openai.com/v1/completions || http://127.0.0.1:${port}/v1/completions
       getServerRoot(port, true) + '/v1/completions',
@@ -103,8 +108,8 @@ const CompletionPanel: FC = observer(() => {
         signal: completionSseController?.signal,
         onmessage(e) {
           scrollToBottom();
-          if (e.data.trim() === '[DONE]') {
-            commonStore.setCompletionGenerating(false);
+          if (!finished && e.data.trim() === '[DONE]') {
+            finish();
             return;
           }
           let data;
@@ -117,6 +122,10 @@ const CompletionPanel: FC = observer(() => {
           if (data.model)
             commonStore.setLastModelName(data.model);
           if (data.choices && Array.isArray(data.choices) && data.choices.length > 0) {
+            if (!finished && data.choices[0]?.finish_reason) {
+              finish();
+              return;
+            }
             answer += data.choices[0]?.text || data.choices[0]?.delta?.content || '';
             setPrompt(prompt + answer.replace(/\s+$/, '') + params.injectEnd.replaceAll('\\n', '\n'));
           }
diff --git a/frontend/src/pages/Composition.tsx b/frontend/src/pages/Composition.tsx
index 43d872a..2981c17 100644
--- a/frontend/src/pages/Composition.tsx
+++ b/frontend/src/pages/Composition.tsx
@@ -185,6 +185,12 @@ const CompositionPanel: FC = observer(() => {
     }
 
     let answer = '';
+    let finished = false;
+    const finish = () => {
+      finished = true;
+      commonStore.setCompositionGenerating(false);
+      generateNs(commonStore.compositionParams.autoPlay);
+    };
     compositionSseController = new AbortController();
     fetchEventSource( // https://api.openai.com/v1/completions || http://127.0.0.1:${port}/v1/completions
       getServerRoot(port, true) + '/v1/completions',
@@ -205,9 +211,8 @@ const CompositionPanel: FC = observer(() => {
         signal: compositionSseController?.signal,
         onmessage(e) {
           scrollToBottom();
-          if (e.data.trim() === '[DONE]') {
-            commonStore.setCompositionGenerating(false);
-            generateNs(commonStore.compositionParams.autoPlay);
+          if (!finished && e.data.trim() === '[DONE]') {
+            finish();
             return;
           }
           let data;
@@ -220,6 +225,10 @@ const CompositionPanel: FC = observer(() => {
           if (data.model)
             commonStore.setLastModelName(data.model);
           if (data.choices && Array.isArray(data.choices) && data.choices.length > 0) {
+            if (!finished && data.choices[0]?.finish_reason) {
+              finish();
+              return;
+            }
             answer += data.choices[0]?.text || data.choices[0]?.delta?.content || '';
             setPrompt(prompt + answer.replace(/\s+$/, ''));
           }