Implement World Info activation using Vector Storage

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