diff --git a/public/script.js b/public/script.js index 35942ccf7..3b35c3886 100644 --- a/public/script.js +++ b/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'); diff --git a/public/scripts/extensions/infinity-context/index.js b/public/scripts/extensions/infinity-context/index.js index 01c0f438e..3ea77e6ff 100644 --- a/public/scripts/extensions/infinity-context/index.js +++ b/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'); } } @@ -168,7 +235,7 @@ async function onExportClick() { if (exportResult.ok) { const data = await exportResult.json(); - const blob = new Blob([JSON.stringify(data, null, 2)], {type : 'application/json'}); + const blob = new Blob([JSON.stringify(data, null, 2)], { type: 'application/json' }); const href = URL.createObjectURL(blob); const link = document.createElement('a'); link.href = href; @@ -198,12 +265,11 @@ async function onSelectImportFile(e) { const text = await getFileText(file); const imported = JSON.parse(text); - imported.chat_id = currentChatId; - + const url = new URL(getApiUrl()); url.pathname = '/api/chromadb/import'; - + const importResult = await fetch(url, { method: 'POST', headers: postHeaders, @@ -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); }); +