From 689af3151a9a50b7b0770e55958ccab2be884bae Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Thu, 28 Mar 2024 02:59:52 +0200
Subject: [PATCH] Pre-populate chat history with injections

---
 public/script.js | 64 ++++++++++++++++++++++++++++++++++++------------
 1 file changed, 49 insertions(+), 15 deletions(-)

diff --git a/public/script.js b/public/script.js
index 634c3b5f4..cfbfc7700 100644
--- a/public/script.js
+++ b/public/script.js
@@ -3384,21 +3384,19 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
     }
 
     // Collect enough messages to fill the context
-    let arrMes = [];
+    let arrMes = new Array(chat2.length);
     let tokenCount = getMessagesTokenCount();
     let lastAddedIndex = -1;
-    for (let i = 0; i < chat2.length; i++) {
-        // not needed for OAI prompting
-        if (main_api == 'openai') {
-            break;
-        }
 
-        const item = chat2[i];
+    // Pre-allocate all injections first.
+    // If it doesn't fit - user shot himself in the foot
+    for (const index of injectedIndices) {
+        const item = chat2[index];
         tokenCount += getTokenCount(item.replace(/\r/gm, ''));
         chatString = item + chatString;
         if (tokenCount < this_max_context) {
-            arrMes[arrMes.length] = item;
-            lastAddedIndex = i;
+            arrMes[index] = item;
+            lastAddedIndex = Math.max(lastAddedIndex, index);
         } else {
             break;
         }
@@ -3407,18 +3405,54 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
         await delay(1);
     }
 
+    for (let i = 0; i < chat2.length; i++) {
+        // not needed for OAI prompting
+        if (main_api == 'openai') {
+            break;
+        }
+
+        // Skip already injected messages
+        if (arrMes[i] !== undefined) {
+            continue;
+        }
+
+        const item = chat2[i];
+        tokenCount += getTokenCount(item.replace(/\r/gm, ''));
+        chatString = item + chatString;
+        if (tokenCount < this_max_context) {
+            arrMes[i] = item;
+            lastAddedIndex = Math.max(lastAddedIndex, i);
+        } else {
+            break;
+        }
+
+        // Prevent UI thread lock on tokenization
+        await delay(1);
+    }
+
+    // Add user alignment message if last message is not a user message
     const stoppedAtUser = userMessageIndices.includes(lastAddedIndex);
     if (addUserAlignment && !stoppedAtUser) {
         tokenCount += getTokenCount(userAlignmentMessage.replace(/\r/gm, ''));
         chatString = userAlignmentMessage + chatString;
-        arrMes[arrMes.length] = userAlignmentMessage;
-        // Injected indices shift by 1 for user alignment message at the beginning
-        injectedIndices.forEach((value, index) => (injectedIndices[index] = value + 1));
-        injectedIndices.push(0);
+        arrMes.push(userAlignmentMessage);
+        injectedIndices.push(arrMes.length - 1);
     }
 
-    // Filter injections which don't fit in the context
-    injectedIndices = injectedIndices.filter(value => value < arrMes.length);
+    // Unsparse the array. Adjust injected indices
+    const newArrMes = [];
+    const newInjectedIndices = [];
+    for (let i = 0; i < arrMes.length; i++) {
+        if (arrMes[i] !== undefined) {
+            newArrMes.push(arrMes[i]);
+            if (injectedIndices.includes(i)) {
+                newInjectedIndices.push(newArrMes.length - 1);
+            }
+        }
+    }
+
+    arrMes = newArrMes;
+    injectedIndices = newInjectedIndices;
 
     if (main_api !== 'openai') {
         setInContextMessages(arrMes.length - injectedIndices.length, type);