rixty/SillyTavern
1
0
Fork 0
mirror of https://github.com/SillyTavern/SillyTavern.git synced 2025-06-05 21:59:27 +02:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
SillyLossy
2023-05-30 00:42:54 +03:00
parent 76eca6abc2
commit f11645e66e
2 changed files with 88 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
public/script.js
View File

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

89
public/scripts/extensions/infinity-context/index.js
View File

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