Permanent ChromaDB synchronization state (using IndexedDB object store)

This commit is contained in:
SillyLossy
2023-06-02 23:21:46 +03:00
parent 745ceb0f31
commit 6fec40d543
4 changed files with 127 additions and 43 deletions

View File

@ -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 () {
eventSource.emit(event_types.MESSAGE_SWIPED, (chat.length - 1));
saveChatConditional();
complete: async function () {
await eventSource.emit(event_types.MESSAGE_SWIPED, (chat.length - 1));
await saveChatConditional();
}
});
}
@ -4997,16 +4997,16 @@ const swipe_right = () => {
duration: swipe_duration,
easing: animation_easing,
queue: false,
complete: function () {
eventSource.emit(event_types.MESSAGE_SWIPED, (chat.length - 1));
complete: async function () {
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();
}
}
}

View File

@ -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 = {};
function invalidateMessageSyncState(messageId) {
async function invalidateMessageSyncState(messageId) {
console.log('CHROMADB: invalidating message sync state', messageId);
const state = getChatSyncState();
state[messageId] = false;
const state = await getChatSyncState();
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) {
const syncState = getChatSyncState();
async function filterSyncedMessages(splitMessages) {
const syncState = await getChatSyncState();
const removeIndices = [];
const syncedIndices = [];
for (let i = 0; i < splittedMessages.length; i++) {
const index = splittedMessages[i].index;
for (let i = 0; i < splitMessages.length; i++) {
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));
}
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 })));
return splitMessages.filter((_, i) => !removeIndices.includes(i));
}
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');
}
}

View File

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

View File

@ -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);
};
});
}
}