From 4370db6bdc7f6959caa378c27add3d0c16508eb8 Mon Sep 17 00:00:00 2001 From: Cohee <18619528+Cohee1207@users.noreply.github.com> Date: Tue, 23 Apr 2024 03:09:52 +0300 Subject: [PATCH] Implement World Info activation using Vector Storage --- public/index.html | 11 +- public/script.js | 1 + public/scripts/extensions/vectors/index.js | 135 ++++++++++++++++++ .../scripts/extensions/vectors/settings.html | 40 ++++++ public/scripts/world-info.js | 77 ++++++++-- src/endpoints/characters.js | 1 + 6 files changed, 250 insertions(+), 15 deletions(-) diff --git a/public/index.html b/public/index.html index ab7329d76..75aaf4eac 100644 --- a/public/index.html +++ b/public/index.html @@ -4955,10 +4955,11 @@ - + + + +
@@ -6065,4 +6066,4 @@ - \ No newline at end of file + diff --git a/public/script.js b/public/script.js index 3f28b2e7b..a1df1eeb2 100644 --- a/public/script.js +++ b/public/script.js @@ -449,6 +449,7 @@ export const event_types = { CHARACTER_DUPLICATED: 'character_duplicated', SMOOTH_STREAM_TOKEN_RECEIVED: 'smooth_stream_token_received', FILE_ATTACHMENT_DELETED: 'file_attachment_deleted', + WORLDINFO_FORCE_ACTIVATE: 'worldinfo_force_activate', }; export const eventSource = new EventEmitter(); diff --git a/public/scripts/extensions/vectors/index.js b/public/scripts/extensions/vectors/index.js index 6ff5c6af5..5cdd0830b 100644 --- a/public/scripts/extensions/vectors/index.js +++ b/public/scripts/extensions/vectors/index.js @@ -23,6 +23,7 @@ import { collapseNewlines } from '../../power-user.js'; import { SECRET_KEYS, secret_state, writeSecret } from '../../secrets.js'; import { getDataBankAttachments, getFileAttachment } from '../../chats.js'; import { debounce, getStringHash as calculateHash, waitUntilCondition, onlyUnique, splitRecursive } from '../../utils.js'; +import { getSortedEntries } from '../../world-info.js'; const MODULE_NAME = 'vectors'; @@ -66,6 +67,11 @@ const settings = { file_position_db: extension_prompt_types.IN_PROMPT, file_depth_db: 4, file_depth_role_db: extension_prompt_roles.SYSTEM, + + // For World Info + enabled_world_info: false, + enabled_for_all: false, + max_entries: 5, }; const moduleWorker = new ModuleWorkerWrapper(synchronizeChat); @@ -472,6 +478,10 @@ async function rearrangeChat(chat) { await processFiles(chat); } + if (settings.enabled_world_info) { + await activateWorldInfo(chat); + } + if (!settings.enabled_chats) { return; } @@ -845,6 +855,7 @@ async function purgeVectorIndex(collectionId) { function toggleSettings() { $('#vectors_files_settings').toggle(!!settings.enabled_files); $('#vectors_chats_settings').toggle(!!settings.enabled_chats); + $('#vectors_world_info_settings').toggle(!!settings.enabled_world_info); $('#together_vectorsModel').toggle(settings.source === 'togetherai'); $('#openai_vectorsModel').toggle(settings.source === 'openai'); $('#cohere_vectorsModel').toggle(settings.source === 'cohere'); @@ -934,6 +945,111 @@ async function onPurgeFilesClick() { } } +async function activateWorldInfo(chat) { + if (!settings.enabled_world_info) { + console.debug('Vectors: Disabled for World Info'); + return; + } + + const entries = await getSortedEntries(); + + if (!Array.isArray(entries) || entries.length === 0) { + console.debug('Vectors: No WI entries found'); + return; + } + + // Group entries by "world" field + const groupedEntries = {}; + + for (const entry of entries) { + // Skip orphaned entries. Is it even possible? + if (!entry.world) { + console.debug('Vectors: Skipped orphaned WI entry', entry); + continue; + } + + // Skip disabled entries + if (entry.disable) { + console.debug('Vectors: Skipped disabled WI entry', entry); + continue; + } + + // Skip entries without content + if (!entry.content) { + console.debug('Vectors: Skipped WI entry without content', entry); + continue; + } + + // Skip non-vectorized entries + if (!entry.vectorized && !settings.enabled_for_all) { + console.debug('Vectors: Skipped non-vectorized WI entry', entry); + continue; + } + + if (!Object.hasOwn(groupedEntries, entry.world)) { + groupedEntries[entry.world] = []; + } + + groupedEntries[entry.world].push(entry); + } + + const collectionIds = []; + + if (Object.keys(groupedEntries).length === 0) { + console.debug('Vectors: No WI entries to synchronize'); + return; + } + + // Synchronize collections + for (const world in groupedEntries) { + const collectionId = `world_${getStringHash(world)}`; + const hashesInCollection = await getSavedHashes(collectionId); + const newEntries = groupedEntries[world].filter(x => !hashesInCollection.includes(getStringHash(x.content))); + const deletedHashes = hashesInCollection.filter(x => !groupedEntries[world].some(y => getStringHash(y.content) === x)); + + if (newEntries.length > 0) { + console.log(`Vectors: Found ${newEntries.length} new WI entries for world ${world}`); + await insertVectorItems(collectionId, newEntries.map(x => ({ hash: getStringHash(x.content), text: x.content, index: x.uid }))); + } + + if (deletedHashes.length > 0) { + console.log(`Vectors: Deleted ${deletedHashes.length} old hashes for world ${world}`); + await deleteVectorItems(collectionId, deletedHashes); + } + + collectionIds.push(collectionId); + } + + // Perform a multi-query + const queryText = await getQueryText(chat); + + if (queryText.length === 0) { + console.debug('Vectors: No text to query for WI'); + return; + } + + const queryResults = await queryMultipleCollections(collectionIds, queryText, settings.max_entries); + const activatedHashes = Object.values(queryResults).flatMap(x => x.hashes).filter(onlyUnique); + const activatedEntries = []; + + // Activate entries found in the query results + for (const entry of entries) { + const hash = getStringHash(entry.content); + + if (activatedHashes.includes(hash)) { + activatedEntries.push(entry); + } + } + + if (activatedEntries.length === 0) { + console.debug('Vectors: No activated WI entries found'); + return; + } + + console.log(`Vectors: Activated ${activatedEntries.length} WI entries`, activatedEntries); + await eventSource.emit(event_types.WORLDINFO_FORCE_ACTIVATE, activatedEntries); +} + jQuery(async () => { if (!extension_settings.vectors) { extension_settings.vectors = settings; @@ -1134,6 +1250,25 @@ jQuery(async () => { saveSettingsDebounced(); }); + $('#vectors_enabled_world_info').prop('checked', settings.enabled_world_info).on('input', () => { + settings.enabled_world_info = !!$('#vectors_enabled_world_info').prop('checked'); + Object.assign(extension_settings.vectors, settings); + saveSettingsDebounced(); + toggleSettings(); + }); + + $('#vectors_enabled_for_all').prop('checked', settings.enabled_for_all).on('input', () => { + settings.enabled_for_all = !!$('#vectors_enabled_for_all').prop('checked'); + Object.assign(extension_settings.vectors, settings); + saveSettingsDebounced(); + }); + + $('#vectors_max_entries').val(settings.max_entries).on('input', () => { + settings.max_entries = Number($('#vectors_max_entries').val()); + Object.assign(extension_settings.vectors, settings); + saveSettingsDebounced(); + }); + const validSecret = !!secret_state[SECRET_KEYS.NOMICAI]; const placeholder = validSecret ? '✔️ Key saved' : '❌ Missing key'; $('#api_key_nomicai').attr('placeholder', placeholder); diff --git a/public/scripts/extensions/vectors/settings.html b/public/scripts/extensions/vectors/settings.html index 02499d120..bcaf1c06e 100644 --- a/public/scripts/extensions/vectors/settings.html +++ b/public/scripts/extensions/vectors/settings.html @@ -97,6 +97,46 @@
+

+ World Info settings +

+ + + +
+
+ +
    +
  • + Checked: all entries except ❌ status can be activated. +
    • +
  • + Unchecked: only entries with 🔗 status can be activated. +
    • +
+
+
+
+ +
+
+ + +
+
+ +
+
+
+

File vectorization settings

diff --git a/public/scripts/world-info.js b/public/scripts/world-info.js index 2932e9050..767f923d0 100644 --- a/public/scripts/world-info.js +++ b/public/scripts/world-info.js @@ -82,6 +82,11 @@ class WorldInfoBuffer { /** @typedef {{scanDepth?: number, caseSensitive?: boolean, matchWholeWords?: boolean}} WIScanEntry The entry that triggered the scan */ // End typedef area + /** + * @type {object[]} Array of entries that need to be activated no matter what + */ + static externalActivations = []; + /** * @type {string[]} Array of messages sorted by ascending depth */ @@ -220,6 +225,23 @@ class WorldInfoBuffer { getDepth() { return world_info_depth + this.#skew; } + + /** + * Check if the current entry is externally activated. + * @param {object} entry WI entry to check + * @returns {boolean} True if the entry is forcefully activated + */ + isExternallyActivated(entry) { + // Entries could be copied with structuredClone, so we need to compare them by string representation + return WorldInfoBuffer.externalActivations.some(x => JSON.stringify(x) === JSON.stringify(entry)); + } + + /** + * Clears the force activations buffer. + */ + cleanExternalActivations() { + WorldInfoBuffer.externalActivations.splice(0, WorldInfoBuffer.externalActivations.length); + } } export function getWorldInfoSettings() { @@ -362,6 +384,10 @@ function setWorldInfoSettings(settings, data) { $('.chat_lorebook_button').toggleClass('world_set', hasWorldInfo); }); + eventSource.on(event_types.WORLDINFO_FORCE_ACTIVATE, (entries) => { + WorldInfoBuffer.externalActivations.push(...entries); + }); + // Add slash commands registerWorldInfoSlashCommands(); } @@ -564,6 +590,7 @@ function registerWorldInfoSlashCommands() { return ''; } + registerSlashCommand('world', onWorldInfoChange, [], '[optional state=off|toggle] [optional silent=true] (optional name) – sets active World, or unsets if no args provided, use state=off and state=toggle to deactivate or toggle a World, use silent=true to suppress toast messages', true, true); registerSlashCommand('getchatbook', getChatBookCallback, ['getchatlore', 'getchatwi'], '– get a name of the chat-bound lorebook or create a new one if was unbound, and pass it down the pipe', true, true); registerSlashCommand('findentry', findBookEntryCallback, ['findlore', 'findwi'], '(file=bookName field=field [texts]) – find a UID of the record from the specified book using the fuzzy match of a field value (default: key) and pass it down the pipe, e.g. /findentry file=chatLore field=key Shadowfang', true, true); registerSlashCommand('getentryfield', getEntryFieldCallback, ['getlorefield', 'getwifield'], '(file=bookName field=field [UID]) – get a field value (default: content) of the record with the UID from the specified book and pass it down the pipe, e.g. /getentryfield file=chatLore field=content 123', true, true); @@ -964,6 +991,7 @@ const originalDataKeyMap = { 'caseSensitive': 'extensions.case_sensitive', 'scanDepth': 'extensions.scan_depth', 'automationId': 'extensions.automation_id', + 'vectorized': 'extensions.vectorized', }; function setOriginalDataValue(data, uid, key, value) { @@ -1071,6 +1099,16 @@ function getWorldEntry(name, data, entry) { ); } + // Verify names to exist in the system + if (data.entries[uid]?.characterFilter?.names?.length > 0) { + for (const name of [...data.entries[uid].characterFilter.names]) { + if (!getContext().characters.find(x => x.avatar.replace(/\.[^/.]+$/, '') === name)) { + console.warn(`World Info: Character ${name} not found. Removing from the entry filter.`, entry); + data.entries[uid].characterFilter.names = data.entries[uid].characterFilter.names.filter(x => x !== name); + } + } + } + setOriginalDataValue(data, uid, 'character_filter', data.entries[uid].characterFilter); saveWorldInfo(name, data); }); @@ -1454,22 +1492,37 @@ function getWorldEntry(name, data, entry) { case 'constant': data.entries[uid].constant = true; data.entries[uid].disable = false; + data.entries[uid].vectorized = false; setOriginalDataValue(data, uid, 'enabled', true); setOriginalDataValue(data, uid, 'constant', true); + setOriginalDataValue(data, uid, 'extensions.vectorized', false); template.removeClass('disabledWIEntry'); break; case 'normal': data.entries[uid].constant = false; data.entries[uid].disable = false; + data.entries[uid].vectorized = false; setOriginalDataValue(data, uid, 'enabled', true); setOriginalDataValue(data, uid, 'constant', false); + setOriginalDataValue(data, uid, 'extensions.vectorized', false); + template.removeClass('disabledWIEntry'); + break; + case 'vectorized': + data.entries[uid].constant = false; + data.entries[uid].disable = false; + data.entries[uid].vectorized = true; + setOriginalDataValue(data, uid, 'enabled', true); + setOriginalDataValue(data, uid, 'constant', false); + setOriginalDataValue(data, uid, 'extensions.vectorized', true); template.removeClass('disabledWIEntry'); break; case 'disabled': data.entries[uid].constant = false; data.entries[uid].disable = true; + data.entries[uid].vectorized = false; setOriginalDataValue(data, uid, 'enabled', false); setOriginalDataValue(data, uid, 'constant', false); + setOriginalDataValue(data, uid, 'extensions.vectorized', false); template.addClass('disabledWIEntry'); break; } @@ -1480,6 +1533,8 @@ function getWorldEntry(name, data, entry) { const entryState = function () { if (entry.constant === true) { return 'constant'; + } else if (entry.vectorized === true) { + return 'vectorized'; } else if (entry.disable === true) { return 'disabled'; } else { @@ -1719,6 +1774,7 @@ const newEntryTemplate = { comment: '', content: '', constant: false, + vectorized: false, selective: true, selectiveLogic: world_info_logic.AND_ANY, addMemo: false, @@ -1925,7 +1981,7 @@ async function getCharacterLore() { } const data = await loadWorldInfoData(worldName); - const newEntries = data ? Object.keys(data.entries).map((x) => data.entries[x]) : []; + const newEntries = data ? Object.keys(data.entries).map((x) => data.entries[x]).map(x => ({ ...x, world: worldName })) : []; entries = entries.concat(newEntries); } @@ -1941,7 +1997,7 @@ async function getGlobalLore() { let entries = []; for (const worldName of selected_world_info) { const data = await loadWorldInfoData(worldName); - const newEntries = data ? Object.keys(data.entries).map((x) => data.entries[x]) : []; + const newEntries = data ? Object.keys(data.entries).map((x) => data.entries[x]).map(x => ({ ...x, world: worldName })) : []; entries = entries.concat(newEntries); } @@ -1963,14 +2019,14 @@ async function getChatLore() { } const data = await loadWorldInfoData(chatWorld); - const entries = data ? Object.keys(data.entries).map((x) => data.entries[x]) : []; + const entries = data ? Object.keys(data.entries).map((x) => data.entries[x]).map(x => ({ ...x, world: chatWorld })) : []; console.debug(`Chat lore has ${entries.length} entries`); return entries; } -async function getSortedEntries() { +export async function getSortedEntries() { try { const globalLore = await getGlobalLore(); const characterLore = await getCharacterLore(); @@ -2098,7 +2154,7 @@ async function checkWorldInfo(chat, maxContext) { continue; } - if (entry.constant) { + if (entry.constant || buffer.isExternallyActivated(entry)) { entry.content = substituteParams(entry.content); activatedNow.add(entry); continue; @@ -2295,6 +2351,8 @@ async function checkWorldInfo(chat, maxContext) { context.setExtensionPrompt(NOTE_MODULE_NAME, ANWithWI, chat_metadata[metadata_keys.position], chat_metadata[metadata_keys.depth], extension_settings.note.allowWIScan, chat_metadata[metadata_keys.role]); } + buffer.cleanExternalActivations(); + return { worldInfoBefore, worldInfoAfter, WIDepthEntries, allActivatedEntries }; } @@ -2381,6 +2439,7 @@ function convertAgnaiMemoryBook(inputObj) { content: entry.entry, constant: false, selective: false, + vectorized: false, selectiveLogic: world_info_logic.AND_ANY, order: entry.weight, position: 0, @@ -2415,6 +2474,7 @@ function convertRisuLorebook(inputObj) { content: entry.content, constant: entry.alwaysActive, selective: entry.selective, + vectorized: false, selectiveLogic: world_info_logic.AND_ANY, order: entry.insertorder, position: world_info_position.before, @@ -2454,6 +2514,7 @@ function convertNovelLorebook(inputObj) { content: entry.text, constant: false, selective: false, + vectorized: false, selectiveLogic: world_info_logic.AND_ANY, order: entry.contextConfig?.budgetPriority ?? 0, position: 0, @@ -2510,6 +2571,7 @@ function convertCharacterBook(characterBook) { matchWholeWords: entry.extensions?.match_whole_words ?? null, automationId: entry.extensions?.automation_id ?? '', role: entry.extensions?.role ?? extension_prompt_roles.SYSTEM, + vectorized: entry.extensions?.vectorized ?? false, }; }); @@ -2785,11 +2847,6 @@ function assignLorebookToChat() { jQuery(() => { - $(document).ready(function () { - registerSlashCommand('world', onWorldInfoChange, [], '[optional state=off|toggle] [optional silent=true] (optional name) – sets active World, or unsets if no args provided, use state=off and state=toggle to deactivate or toggle a World, use silent=true to suppress toast messages', true, true); - }); - - $('#world_info').on('mousedown change', async function (e) { // If there's no world names, don't do anything if (world_names.length === 0) { diff --git a/src/endpoints/characters.js b/src/endpoints/characters.js index 556a8fecc..571eec8c0 100644 --- a/src/endpoints/characters.js +++ b/src/endpoints/characters.js @@ -439,6 +439,7 @@ function convertWorldInfoToCharacterBook(name, entries) { case_sensitive: entry.caseSensitive ?? null, automation_id: entry.automationId ?? '', role: entry.role ?? 0, + vectorized: entry.vectorized ?? false, }, };