From 53ed7acecaaba6be62b168c561bd22bb9e1e5cb8 Mon Sep 17 00:00:00 2001 From: SillyLossy Date: Tue, 23 May 2023 00:30:52 +0300 Subject: [PATCH] Infinity context: inject and purge chat data --- .../extensions/infinity-context/index.js | 77 +++++++++++++++---- .../extensions/infinity-context/style.css | 4 +- 2 files changed, 66 insertions(+), 15 deletions(-) diff --git a/public/scripts/extensions/infinity-context/index.js b/public/scripts/extensions/infinity-context/index.js index 37c279555..5d1c034d6 100644 --- a/public/scripts/extensions/infinity-context/index.js +++ b/public/scripts/extensions/infinity-context/index.js @@ -1,12 +1,11 @@ -import { saveSettingsDebounced, getCurrentChatId } from "../../../script.js"; +import { saveSettingsDebounced, getCurrentChatId, system_message_types } from "../../../script.js"; +import { humanizedDateTime } from "../../RossAscends-mods.js"; import { getApiUrl, extension_settings } from "../../extensions.js"; -import { getFileText, getStringHash, splitRecursive } from "../../utils.js"; +import { getFileText, onlyUnique, splitRecursive } from "../../utils.js"; export { MODULE_NAME }; const MODULE_NAME = 'chromadb'; -const fileSplitLength = 2048; - const defaultSettings = { keep_context: 10, keep_context_min: 1, @@ -22,6 +21,11 @@ const defaultSettings = { split_length_min: 64, split_length_max: 4096, split_length_step: 64, + + file_split_length: 1024, + file_split_length_min: 512, + file_split_length_max: 4096, + file_split_length_step: 128, }; const postHeaders = { @@ -37,6 +41,7 @@ async function loadSettings() { $('#chromadb_keep_context').val(extension_settings.chromadb.keep_context).trigger('input'); $('#chromadb_n_results').val(extension_settings.chromadb.n_results).trigger('input'); $('#chromadb_split_length').val(extension_settings.chromadb.split_length).trigger('input'); + $('#chromadb_file_split_length').val(extension_settings.chromadb.file_split_length).trigger('input'); } function onKeepContextInput() { @@ -57,6 +62,12 @@ function onSplitLengthInput() { saveSettingsDebounced(); } +function onFileSplitLengthInput() { + extension_settings.chromadb.file_split_length = Number($('#chromadb_file_split_length').val()); + $('#chromadb_file_split_length_value').text(extension_settings.chromadb.file_split_length); + saveSettingsDebounced(); +} + async function addMessages(chat_id, messages) { const url = new URL(getApiUrl()); url.pathname = '/api/chromadb'; @@ -98,6 +109,22 @@ async function addMessages(chat_id, messages) { return { count: 0 }; } +async function onPurgeClick() { + const chat_id = getCurrentChatId(); + const url = new URL(getApiUrl()); + url.pathname = '/api/chromadb/purge'; + + const purgeResult = await fetch(url, { + method: 'POST', + headers: postHeaders, + body: JSON.stringify({ chat_id }), + }); + + if (purgeResult.ok) { + toastr.success('ChromaDB context has been successfully cleared'); + } +} + async function queryMessages(chat_id, query) { const url = new URL(getApiUrl()); url.pathname = '/api/chromadb/query'; @@ -125,17 +152,28 @@ async function onSelectInjectFile(e) { } try { + toastr.info('This may take some time, depending on the file size', 'Processing...'); const currentChatId = getCurrentChatId(); const text = await getFileText(file); - const split = splitRecursive(text, fileSplitLength); + const split = splitRecursive(text, extension_settings.chromadb.file_split_length).filter(onlyUnique); - const messages = split.map(m => ({ - id: `${getStringHash(file.name)}-${getStringHash(m)}`, - role: 'assistant', // probably need a system role? + const messages = split.map(m => ({ + id: `${file.name}-${split.indexOf(m)}`, + role: 'system', content: m, date: Date.now(), - meta: file.name, + meta: JSON.stringify({ + name: file.name, + is_user: false, + is_name: false, + is_system: false, + send_date: humanizedDateTime(), + mes: m, + extra: { + type: system_message_types.NARRATOR, + } + }), })); const url = new URL(getApiUrl()); @@ -150,8 +188,10 @@ async function onSelectInjectFile(e) { if (addMessagesResult.ok) { const addMessagesData = await addMessagesResult.json(); - toastr.info(`Number of chunks: ${addMessagesData.count}`, 'Injected successfully!'); + toastr.success(`Number of chunks: ${addMessagesData.count}`, 'Injected successfully!'); return addMessagesData; + } else { + throw new Error(); } } catch (error) { @@ -204,10 +244,19 @@ jQuery(async () => { -
`; @@ -216,8 +265,10 @@ jQuery(async () => { $('#chromadb_keep_context').on('input', onKeepContextInput); $('#chromadb_n_results').on('input', onNResultsInput); $('#chromadb_split_length').on('input', onSplitLengthInput); + $('#chromadb_file_split_length').on('input', onFileSplitLengthInput); $('#chromadb_inject').on('click', () => $('#chromadb_inject_file').trigger('click')); $('#chromadb_inject_file').on('change', onSelectInjectFile); + $('#chromadb_purge').on('click', onPurgeClick); await loadSettings(); }); diff --git a/public/scripts/extensions/infinity-context/style.css b/public/scripts/extensions/infinity-context/style.css index 92f81b745..44faf33a7 100644 --- a/public/scripts/extensions/infinity-context/style.css +++ b/public/scripts/extensions/infinity-context/style.css @@ -1,7 +1,7 @@ -#chromadb_inject { +.chromadb_settings .menu_button { width: fit-content; display: flex; gap: 10px; - align-items: center; + align-items: baseline; flex-direction: row; }