Permanent ChromaDB synchronization state (using IndexedDB object store)

public/script.js

@@ -1018,11 +1018,11 @@ function clearChat() {
     $("#chat").children().remove();
 }
 
-function deleteLastMessage() {
+async function deleteLastMessage() {
     count_view_mes--;
     chat.length = chat.length - 1;
     $('#chat').children('.mes').last().remove();
-    eventSource.emit(event_types.MESSAGE_DELETED, chat.length);
+    await eventSource.emit(event_types.MESSAGE_DELETED, chat.length);
 }
 
 export async function reloadCurrentChat() {
@@ -1839,7 +1839,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);
+                await eventSource.emit(event_types.MESSAGE_DELETED, chat.length);
             }
         }
 
@@ -2417,7 +2417,7 @@ async function Generate(type, { automatic_trigger, force_name2, resolve, reject,
                         if (isImpersonate) {
                             $('#send_textarea').val(getMessage).trigger('input');
                             generatedPromtCache = "";
-                            eventSource.emit(event_types.IMPERSONATE_READY, getMessage);
+                            await eventSource.emit(event_types.IMPERSONATE_READY, getMessage);
                         }
                         else if (type == 'quiet') {
                             resolve(getMessage);
@@ -4838,9 +4838,9 @@ function swipe_left() {      // when we swipe left..but no generation.
                             duration: swipe_duration,
                             easing: animation_easing,
                             queue: false,
-                            complete: function () {
+                            complete: async function () {
-                                eventSource.emit(event_types.MESSAGE_SWIPED, (chat.length - 1));
+                                await eventSource.emit(event_types.MESSAGE_SWIPED, (chat.length - 1));
-                                saveChatConditional();
+                                await saveChatConditional();
                             }
                         });
                     }
@@ -4997,16 +4997,16 @@ const swipe_right = () => {
                             duration: swipe_duration,
                             easing: animation_easing,
                             queue: false,
-                            complete: function () {
+                            complete: async function () {
-                                eventSource.emit(event_types.MESSAGE_SWIPED, (chat.length - 1));
+                                await 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.debug('caught here 2');
                                     is_send_press = true;
                                     $('.mes_buttons:last').hide();
-                                    Generate('swipe');
+                                    await Generate('swipe');
                                 } else {
                                     if (parseInt(chat[chat.length - 1]['swipe_id']) !== chat[chat.length - 1]['swipes'].length) {
-                                        saveChatConditional();
+                                        await saveChatConditional();
                                     }
                                 }
                             }

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

@@ -1,10 +1,11 @@
 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";
+import { getFileText, onlyUnique, splitRecursive, IndexedDBStore } from "../../utils.js";
 export { MODULE_NAME };
 
 const MODULE_NAME = 'chromadb';
+const dbStore = new IndexedDBStore('SillyTavern', MODULE_NAME);
 
 const defaultSettings = {
     strategy: 'original',
@@ -35,22 +36,21 @@ const postHeaders = {
     'Bypass-Tunnel-Reminder': 'bypass',
 };
 
-const chatStateFlags = {};
+async function invalidateMessageSyncState(messageId) {
-
-function invalidateMessageSyncState(messageId) {
     console.log('CHROMADB: invalidating message sync state', messageId);
-    const state = getChatSyncState();
+    const state = await getChatSyncState();
-    state[messageId] = false;
+    state[messageId] = 0;
+    await dbStore.put(getCurrentChatId(), state);
 }
 
-function getChatSyncState() {
+async function getChatSyncState() {
     const currentChatId = getCurrentChatId();
     if (!checkChatId(currentChatId)) {
         return;
     }
 
     const context = getContext();
-    const chatState = chatStateFlags[currentChatId] || [];
+    const chatState = (await dbStore.get(currentChatId)) || [];
 
     // if the chat length has decreased, it means that some messages were deleted
     if (chatState.length > context.chat.length) {
@@ -70,10 +70,10 @@ function getChatSyncState() {
     chatState.length = context.chat.length;
     for (let i = 0; i < chatState.length; i++) {
         if (chatState[i] === undefined) {
-            chatState[i] = false;
+            chatState[i] = 0;
         }
     }
-    chatStateFlags[currentChatId] = chatState;
+    await dbStore.put(currentChatId, chatState);
 
     return chatState;
 }
@@ -144,12 +144,12 @@ async function addMessages(chat_id, messages) {
     url.pathname = '/api/chromadb';
 
     const messagesDeepCopy = JSON.parse(JSON.stringify(messages));
-    let splittedMessages = [];
+    let splitMessages = [];
 
     let id = 0;
     messagesDeepCopy.forEach((m, index) => {
         const split = splitRecursive(m.mes, extension_settings.chromadb.split_length);
-        splittedMessages.push(...split.map(text => ({
+        splitMessages.push(...split.map(text => ({
             ...m,
             mes: text,
             send_date: id,
@@ -159,14 +159,14 @@ async function addMessages(chat_id, messages) {
         })));
     });
 
-    splittedMessages = filterSyncedMessages(splittedMessages);
+    splitMessages = await filterSyncedMessages(splitMessages);
 
     // no messages to add
-    if (splittedMessages.length === 0) {
+    if (splitMessages.length === 0) {
         return { count: 0 };
     }
 
-    const transformedMessages = splittedMessages.map((m) => ({
+    const transformedMessages = splitMessages.map((m) => ({
         id: m.id,
         role: m.is_user ? 'user' : 'assistant',
         content: m.mes,
@@ -182,19 +182,18 @@ async function addMessages(chat_id, messages) {
 
     if (addMessagesResult.ok) {
         const addMessagesData = await addMessagesResult.json();
-
         return addMessagesData; // { count: 1 }
     }
 
     return { count: 0 };
 }
 
-function filterSyncedMessages(splittedMessages) {
+async function filterSyncedMessages(splitMessages) {
-    const syncState = getChatSyncState();
+    const syncState = await getChatSyncState();
     const removeIndices = [];
     const syncedIndices = [];
-    for (let i = 0; i < splittedMessages.length; i++) {
+    for (let i = 0; i < splitMessages.length; i++) {
-        const index = splittedMessages[i].index;
+        const index = splitMessages[i].index;
 
         if (syncState[index]) {
             removeIndices.push(i);
@@ -205,19 +204,14 @@ function filterSyncedMessages(splittedMessages) {
     }
 
     for (const index of syncedIndices) {
-        syncState[index] = true;
+        syncState[index] = 1;
     }
 
-    logSyncState(syncState);
+    console.debug('CHROMADB: sync state', syncState.map((v, i) => ({ id: i, synced: v })));
+    await dbStore.put(getCurrentChatId(), syncState);
 
     // remove messages that are already synced
-    return splittedMessages.filter((_, i) => !removeIndices.includes(i));
+    return splitMessages.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() {
@@ -235,7 +229,7 @@ async function onPurgeClick() {
     });
 
     if (purgeResult.ok) {
-        delete chatStateFlags[chat_id];
+        await dbStore.delete(chat_id);
         toastr.success('ChromaDB context has been successfully cleared');
     }
 }

public/scripts/group-chats.js

@@ -120,7 +120,7 @@ async function regenerateGroup() {
             break;
         }
 
-        deleteLastMessage();
+        await deleteLastMessage();
     }
 
     generateGroupWrapper();

public/scripts/utils.js

@@ -256,7 +256,7 @@ export function countOccurrences(string, character) {
 }
 
 export function isOdd(number) {
-  return number % 2 !== 0;
+    return number % 2 !== 0;
 }
 
 export function timestampToMoment(timestamp) {
@@ -314,3 +314,93 @@ export function splitRecursive(input, length, delimitiers = ['\n\n', '\n', ' ',
     }
     return result;
 }
+
+export class IndexedDBStore {
+    constructor(dbName, storeName) {
+        this.dbName = dbName;
+        this.storeName = storeName;
+        this.db = null;
+    }
+
+    async open() {
+        return new Promise((resolve, reject) => {
+            const request = indexedDB.open(this.dbName);
+
+            request.onupgradeneeded = (event) => {
+                const db = event.target.result;
+                db.createObjectStore(this.storeName, { keyPath: null, autoIncrement: false });
+            };
+
+            request.onsuccess = (event) => {
+                console.debug(`IndexedDBStore.open(${this.dbName})`);
+                this.db = event.target.result;
+                resolve(this.db);
+            };
+
+            request.onerror = (event) => {
+                console.error(`IndexedDBStore.open(${this.dbName})`);
+                reject(event.target.error);
+            };
+        });
+    }
+
+    async get(key) {
+        if (!this.db) await this.open();
+
+        return new Promise((resolve, reject) => {
+            const transaction = this.db.transaction(this.storeName, "readonly");
+            const objectStore = transaction.objectStore(this.storeName);
+            const request = objectStore.get(key);
+
+            request.onsuccess = (event) => {
+                console.debug(`IndexedDBStore.get(${key})`);
+                resolve(event.target.result);
+            };
+
+            request.onerror = (event) => {
+                console.error(`IndexedDBStore.get(${key})`);
+                reject(event.target.error);
+            };
+        });
+    }
+
+    async put(key, object) {
+        if (!this.db) await this.open();
+
+        return new Promise((resolve, reject) => {
+            const transaction = this.db.transaction(this.storeName, "readwrite");
+            const objectStore = transaction.objectStore(this.storeName);
+            const request = objectStore.put(object, key);
+
+            request.onsuccess = (event) => {
+                console.debug(`IndexedDBStore.put(${key})`);
+                resolve(event.target.result);
+            };
+
+            request.onerror = (event) => {
+                console.error(`IndexedDBStore.put(${key})`);
+                reject(event.target.error);
+            };
+        });
+    }
+
+    async delete(key) {
+        if (!this.db) await this.open();
+
+        return new Promise((resolve, reject) => {
+            const transaction = this.db.transaction(this.storeName, "readwrite");
+            const objectStore = transaction.objectStore(this.storeName);
+            const request = objectStore.delete(key);
+
+            request.onsuccess = (event) => {
+                console.debug(`IndexedDBStore.delete(${key})`);
+                resolve(event.target.result);
+            };
+
+            request.onerror = (event) => {
+                console.error(`IndexedDBStore.delete(${key})`);
+                reject(event.target.error);
+            };
+        });
+    }
+}