Experimental ChromaDB performance improvement. Keep chat sync flags and only push new messages.

public/script.js

@@ -418,6 +418,7 @@ export const event_types = {
     MESSAGE_SENT: 'message_sent',
     MESSAGE_RECEIVED: 'message_received',
     MESSAGE_EDITED: 'message_edited',
+    MESSAGE_DELETED: 'message_deleted',
     IMPERSONATE_READY: 'impersonate_ready',
 }
 
@@ -1013,6 +1014,7 @@ function deleteLastMessage() {
     count_view_mes--;
     chat.length = chat.length - 1;
     $('#chat').children('.mes').last().remove();
+    eventSource.emit(event_types.MESSAGE_DELETED, chat.length);
 }
 
 export async function reloadCurrentChat() {
@@ -1819,6 +1821,7 @@ async function Generate(type, { automatic_trigger, force_name2, resolve, reject,
                 $('#chat').children().last().hide(500, function () {
                     $(this).remove();
                 });
+                eventSource.emit(event_types.MESSAGE_DELETED, chat.length);
             }
         }
 
@@ -4846,8 +4849,8 @@ function swipe_left() {      // when we swipe left..but no generation.
                             easing: animation_easing,
                             queue: false,
                             complete: function () {
-                                saveChatConditional();
                                 eventSource.emit(event_types.MESSAGE_SWIPED, (chat.length - 1));
+                                saveChatConditional();
                             }
                         });
                     }
@@ -5005,6 +5008,7 @@ const swipe_right = () => {
                             easing: animation_easing,
                             queue: false,
                             complete: function () {
+                                eventSource.emit(event_types.MESSAGE_SWIPED, (chat.length - 1));
                                 if (run_generate && !is_send_press && parseInt(chat[chat.length - 1]['swipe_id']) === chat[chat.length - 1]['swipes'].length) {
                                     console.log('caught here 2');
                                     is_send_press = true;
@@ -5015,7 +5019,6 @@ const swipe_right = () => {
                                         saveChatConditional();
                                     }
                                 }
-                                eventSource.emit(event_types.MESSAGE_SWIPED, (chat.length - 1));
                             }
                         });
                     }
@@ -6113,6 +6116,7 @@ $(document).ready(function () {
             saveChatConditional();
             var $textchat = $("#chat");
             $textchat.scrollTop($textchat[0].scrollHeight);
+            eventSource.emit(event_types.MESSAGE_DELETED, chat.length);
         }
         this_del_mes = 0;
         $('#chat .mes').last().addClass('last_mes');

public/scripts/extensions/infinity-context/index.js

@@ -1,4 +1,4 @@
-import { saveSettingsDebounced, getCurrentChatId, system_message_types } from "../../../script.js";
+import { saveSettingsDebounced, getCurrentChatId, system_message_types, eventSource, event_types } from "../../../script.js";
 import { humanizedDateTime } from "../../RossAscends-mods.js";
 import { getApiUrl, extension_settings, getContext } from "../../extensions.js";
 import { getFileText, onlyUnique, splitRecursive } from "../../utils.js";
@@ -35,6 +35,33 @@ const postHeaders = {
     'Bypass-Tunnel-Reminder': 'bypass',
 };
 
+const chatStateFlags = {};
+
+function invalidateMessageSyncState(messageId) {
+    console.log('CHROMADB: invalidating message sync state', messageId);
+    const state = getChatSyncState();
+    state[messageId] = false;
+}
+
+function getChatSyncState() {
+    const currentChatId = getCurrentChatId();
+    if (!checkChatId(currentChatId)) {
+        return;
+    }
+
+    const context = getContext();
+    const chatState = chatStateFlags[currentChatId] || [];
+    chatState.length = context.chat.length;
+    for (let i = 0; i < chatState.length; i++) {
+        if (chatState[i] === undefined) {
+            chatState[i] = false;
+        }
+    }
+    chatStateFlags[currentChatId] = chatState;
+
+    return chatState;
+}
+
 async function loadSettings() {
     if (Object.keys(extension_settings.chromadb).length === 0) {
         Object.assign(extension_settings.chromadb, defaultSettings);
@@ -96,20 +123,28 @@ async function addMessages(chat_id, messages) {
     url.pathname = '/api/chromadb';
 
     const messagesDeepCopy = JSON.parse(JSON.stringify(messages));
-    const splittedMessages = [];
+    let splittedMessages = [];
 
     let id = 0;
-    messagesDeepCopy.forEach(m => {
+    messagesDeepCopy.forEach((m, index) => {
         const split = splitRecursive(m.mes, extension_settings.chromadb.split_length);
         splittedMessages.push(...split.map(text => ({
             ...m,
             mes: text,
             send_date: id,
             id: `msg-${id++}`,
+            index: index,
             extra: undefined,
         })));
     });
 
+    splittedMessages = filterSyncedMessages(splittedMessages);
+
+    // no messages to add
+    if (splittedMessages.length === 0) {
+        return { count: 0 };
+    }
+
     const transformedMessages = splittedMessages.map((m) => ({
         id: m.id,
         role: m.is_user ? 'user' : 'assistant',
@@ -133,6 +168,37 @@ async function addMessages(chat_id, messages) {
     return { count: 0 };
 }
 
+function filterSyncedMessages(splittedMessages) {
+    const syncState = getChatSyncState();
+    const removeIndices = [];
+    const syncedIndices = [];
+    for (let i = 0; i < splittedMessages.length; i++) {
+        const index = splittedMessages[i].index;
+
+        if (syncState[index]) {
+            removeIndices.push(i);
+            continue;
+        }
+
+        syncedIndices.push(index);
+    }
+
+    for (const index of syncedIndices) {
+        syncState[index] = true;
+    }
+
+    logSyncState(syncState);
+
+    // remove messages that are already synced
+    return splittedMessages.filter((_, i) => !removeIndices.includes(i));
+}
+
+function logSyncState(syncState) {
+    const chat = getContext().chat;
+    console.log('CHROMADB: sync state');
+    console.table(syncState.map((v, i) => ({ synced: v, name: chat[i].name, message: chat[i].mes })));
+}
+
 async function onPurgeClick() {
     const chat_id = getCurrentChatId();
     if (!checkChatId(chat_id)) {
@@ -148,6 +214,7 @@ async function onPurgeClick() {
     });
 
     if (purgeResult.ok) {
+        delete chatStateFlags[chat_id];
         toastr.success('ChromaDB context has been successfully cleared');
     }
 }
@@ -198,7 +265,6 @@ async function onSelectImportFile(e) {
         const text = await getFileText(file);
         const imported = JSON.parse(text);
 
-
         imported.chat_id = currentChatId;
 
         const url = new URL(getApiUrl());
@@ -422,4 +488,13 @@ jQuery(async () => {
     $('#chromadb_purge').on('click', onPurgeClick);
     $('#chromadb_export').on('click', onExportClick);
     await loadSettings();
+
+    // Not sure if this is needed, but it's here just in case
+    eventSource.on(event_types.MESSAGE_DELETED, getChatSyncState);
+    eventSource.on(event_types.MESSAGE_RECEIVED, getChatSyncState);
+    eventSource.on(event_types.MESSAGE_SENT, getChatSyncState);
+    // Will make the sync state update when a message is edited or swiped
+    eventSource.on(event_types.MESSAGE_EDITED, invalidateMessageSyncState);
+    eventSource.on(event_types.MESSAGE_SWIPED, invalidateMessageSyncState);
 });
+