From 1ef08110c382bd68b9d2ab20a148dc9c674a4c1d Mon Sep 17 00:00:00 2001
From: Webifi <john@webifi.com>
Date: Tue, 25 Jul 2023 14:53:29 -0500
Subject: [PATCH] Add aggressive stop setting

---
 src/lib/ChatRequestPetals.svelte | 46 ++++++++++++++++++++++----------
 src/lib/Settings.svelte          |  8 ++++++
 src/lib/Types.svelte             |  1 +
 3 files changed, 41 insertions(+), 14 deletions(-)

diff --git a/src/lib/ChatRequestPetals.svelte b/src/lib/ChatRequestPetals.svelte
index f7e7833..96e3667 100644
--- a/src/lib/ChatRequestPetals.svelte
+++ b/src/lib/ChatRequestPetals.svelte
@@ -109,20 +109,38 @@ export const runPetalsCompletionRequest = async (
               chatResponse.updateFromError(err.message)
               throw err
             }
-            window.setTimeout(() => {
-              chatResponse.updateFromAsyncResponse(
-                      {
-                        model,
-                        choices: [{
-                          delta: {
-                            content: response.outputs,
-                            role: 'assistant'
-                          },
-                          finish_reason: (response.stop ? 'stop' : null)
-                        }]
-                      } as any
-              )
-            }, 1)
+            chatResponse.updateFromAsyncResponse(
+                {
+                  model,
+                  choices: [{
+                    delta: {
+                      content: response.outputs,
+                      role: 'assistant'
+                    },
+                    finish_reason: (response.stop ? 'stop' : null)
+                  }]
+                } as any
+            )
+            if (chat.settings.aggressiveStop && !response.stop) {
+              // check if we should've stopped
+              const message = chatResponse.getMessages()[0]
+              const pad = 10 // look back 10 characters + stop sequence
+              if (message) {
+                const mc = (message.content).trim()
+                for (let i = 0, l = stopSequences.length; i < l; i++) {
+                  const ss = stopSequences[i].trim()
+                  const ind = mc.slice(0 - (ss.length + pad)).indexOf(ss)
+                  if (ind > -1) {
+                    const offset = (ss.length + pad) - ind
+                    message.content = mc.slice(0, mc.length - offset)
+                    response.stop = true
+                    updateMessages(chat.id)
+                    chatResponse.finish()
+                    ws.close()
+                  }
+                }
+              }
+            }
           }
         }
         ws.onclose = () => {
diff --git a/src/lib/Settings.svelte b/src/lib/Settings.svelte
index f33324a..8ec32fb 100644
--- a/src/lib/Settings.svelte
+++ b/src/lib/Settings.svelte
@@ -110,6 +110,7 @@ const defaults:ChatSettings = {
   hppWithSummaryPrompt: false,
   imageGenerationSize: '',
   stopSequence: '',
+  aggressiveStop: false,
   userMessageStart: '',
   assistantMessageStart: '',
   systemMessageStart: '',
@@ -524,6 +525,13 @@ const chatSettingsList: ChatSetting[] = [
         },
         hide: isNotPetals
       },
+      {
+        key: 'aggressiveStop',
+        name: 'Use aggressive stop',
+        title: 'Sometimes generation con continue even after a stop sequence. This will stop generation client side if generation continues after stop sequence.',
+        type: 'boolean',
+        hide: isNotPetals
+      },
       {
         key: 'userMessageStart',
         name: 'User Message Start Sequence',
diff --git a/src/lib/Types.svelte b/src/lib/Types.svelte
index 7064115..ae45919 100644
--- a/src/lib/Types.svelte
+++ b/src/lib/Types.svelte
@@ -114,6 +114,7 @@ export type ChatSettings = {
     useResponseAlteration?: boolean;
     responseAlterations?: ResponseAlteration[];
     stopSequence: string;
+    aggressiveStop: boolean;
     userMessageStart: string;
     assistantMessageStart: string;
     systemMessageStart: string;