From 3a3ff89047db3d308c364be091eaa2cf81595c50 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Fri, 8 Sep 2023 15:25:10 +0300
Subject: [PATCH] Add button to vectorize all chat

---
 public/scripts/extensions/vectors/index.js    | 43 +++++++++++++++++++
 .../scripts/extensions/vectors/settings.html  | 13 ++++++
 2 files changed, 56 insertions(+)

diff --git a/public/scripts/extensions/vectors/index.js b/public/scripts/extensions/vectors/index.js
index c04a6780e..7ad0c6fb3 100644
--- a/public/scripts/extensions/vectors/index.js
+++ b/public/scripts/extensions/vectors/index.js
@@ -17,6 +17,48 @@ const settings = {
 
 const moduleWorker = new ModuleWorkerWrapper(synchronizeChat);
 
+async function onVectorizeAllClick() {
+    try {
+        if (!settings.enabled) {
+            return;
+        }
+
+        const chatId = getCurrentChatId();
+        const batchSize = 5;
+        const elapsedLog = [];
+        let finished = false;
+        $('#vectorize_progress').show();
+        $('#vectorize_progress_percent').text('0');
+        $('#vectorize_progress_eta').text('...');
+
+        while (!finished) {
+            const startTime = Date.now();
+            const remaining = await synchronizeChat(batchSize);
+            const elapsed = Date.now() - startTime;
+            elapsedLog.push(elapsed);
+            finished = remaining <= 0;
+
+            const total = getContext().chat.length;
+            const processed = total - remaining;
+            const processedPercent = Math.round((processed / total) * 100); // percentage of the work done
+            const averageElapsed = elapsedLog.slice(-5).reduce((a, b) => a + b, 0) / elapsedLog.length; // average time needed to process one item
+            const pace = averageElapsed / batchSize; // time needed to process one item
+            const remainingTime = Math.round(pace * remaining / 1000);
+
+            $('#vectorize_progress_percent').text(processedPercent);
+            $('#vectorize_progress_eta').text(remainingTime);
+
+            if (chatId !== getCurrentChatId()) {
+                throw new Error('Chat changed');
+            }
+        }
+    } catch (error) {
+        console.error('Vectors: Failed to vectorize all', error);
+    } finally {
+        $('#vectorize_progress').hide();
+    }
+}
+
 async function synchronizeChat(batchSize = 10) {
     try {
         if (!settings.enabled) {
@@ -283,6 +325,7 @@ jQuery(async () => {
         Object.assign(extension_settings.vectors, settings);
         saveSettingsDebounced();
     });
+    $('#vectors_vectorize_all').on('click', onVectorizeAllClick);
 
     eventSource.on(event_types.CHAT_CHANGED, onChatEvent);
     eventSource.on(event_types.MESSAGE_DELETED, onChatEvent);
diff --git a/public/scripts/extensions/vectors/settings.html b/public/scripts/extensions/vectors/settings.html
index a76db1aa8..0d1d7309a 100644
--- a/public/scripts/extensions/vectors/settings.html
+++ b/public/scripts/extensions/vectors/settings.html
@@ -16,6 +16,19 @@
                 <option value="local">Local</option>
                 <option value="openai">OpenAI</option>
             </select>
+            <div>
+                Old messages are vectorized gradually as you chat.
+                To process all previous messages, click the button below.
+            </div>
+            <div id="vectors_vectorize_all" class="menu_button menu_button_icon">
+                Vectorize All
+            </div>
+            <div id="vectorize_progress" style="display: none;">
+                <small>
+                    Processed <span id="vectorize_progress_percent">0</span>% of messages.
+                    ETA: <span id="vectorize_progress_eta">...</span> seconds.
+                </small>
+            </div>
         </div>
     </div>
 </div>