Implement World Info activation using Vector Storage

This commit is contained in:
Cohee 2024-04-23 03:09:52 +03:00
parent 6d1933c8f3
commit 4370db6bdc
6 changed files with 250 additions and 15 deletions

View File

@ -4955,10 +4955,11 @@
<textarea class="text_pole" rows="1" name="comment" maxlength="5000" data-i18n="[placeholder]Entry Title/Memo" placeholder="Entry Title/Memo"></textarea>
</div>
<!-- <span class="world_entry_form_position_value"></span> -->
<select data-i18n="[title]WI Entry Status:🔵 Constant🟢 Normal❌ Disabled" title="WI Entry Status:&#13;🔵 Constant&#13;🟢 Normal&#13;❌ Disabled" name="entryStateSelector" class="text_pole widthNatural margin0">
<option title="WI Entry Status:&#13;🔵 Constant&#13;🟢 Normal&#13;❌ Disabled" value="constant" data-i18n="[title]WI Entry Status:🔵 Constant🟢 Normal❌ Disabled">🔵</option>
<option title="WI Entry Status:&#13;🔵 Constant&#13;🟢 Normal&#13;❌ Disabled" value="normal" data-i18n="[title]WI Entry Status:🔵 Constant🟢 Normal❌ Disabled">🟢</option>
<option title="WI Entry Status:&#13;🔵 Constant&#13;🟢 Normal&#13;❌ Disabled" value="disabled" data-i18n="[title]WI Entry Status:🔵 Constant🟢 Normal❌ Disabled"></option>
<select data-i18n="[title]WI Entry Status:🔵 Constant🟢 Normal🔗 Vectorized❌ Disabled" title="WI Entry Status:&#13;🔵 Constant&#13;🟢 Normal&#13;🔗 Vectorized&#13;❌ Disabled" name="entryStateSelector" class="text_pole widthNatural margin0">
<option value="constant" title="Constant" data-i18n="[title]Constant">🔵</option>
<option value="normal" title="Normal" data-i18n="[title]Normal">🟢</option>
<option value="vectorized" title="Vectorized" data-i18n="[title]Vectorized">🔗</option>
<option value="disabled" title="Disabled" data-i18n="[title]Disabled"></option>
</select>
</div>
<div class="WIEnteryHeaderControls flex-container">

View File

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

View File

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

View File

@ -97,6 +97,46 @@
<hr>
<h4>
World Info settings
</h4>
<label class="checkbox_label" for="vectors_enabled_world_info" title="Enable activation of World Info entries based on vector similarity.">
<input id="vectors_enabled_world_info" type="checkbox" class="checkbox">
Enabled for World Info
</label>
<div id="vectors_world_info_settings" class="marginTopBot5">
<div class="flex-container">
<label for="vectors_enabled_for_all" class="checkbox_label">
<input id="vectors_enabled_for_all" type="checkbox" />
<span>Enabled for all entries</span>
</label>
<ul class="margin0">
<li>
<small>Checked: all entries except ❌ status can be activated.</small>
</li>
<li>
<small>Unchecked: only entries with 🔗 status can be activated.</small>
</li>
</ul>
</div>
<div class="flex-container">
<div class="flex1">
<!-- Vacant for future use -->
</div>
<div class="flex1" title="Maximum number of entries to be activated">
<label for="vectors_max_entries" >
<small>Max Entries</small>
</label>
<input id="vectors_max_entries" type="number" class="text_pole widthUnset" min="1" max="9999" />
</div>
<div class="flex1">
<!-- Vacant for future use -->
</div>
</div>
</div>
<h4>
File vectorization settings
</h4>

View File

@ -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, [], '<span class="monospace">[optional state=off|toggle] [optional silent=true] (optional name)</span> sets active World, or unsets if no args provided, use <code>state=off</code> and <code>state=toggle</code> to deactivate or toggle a World, use <code>silent=true</code> 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'], '<span class="monospace">(file=bookName field=field [texts])</span> 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. <tt>/findentry file=chatLore field=key Shadowfang</tt>', true, true);
registerSlashCommand('getentryfield', getEntryFieldCallback, ['getlorefield', 'getwifield'], '<span class="monospace">(file=bookName field=field [UID])</span> get a field value (default: content) of the record with the UID from the specified book and pass it down the pipe, e.g. <tt>/getentryfield file=chatLore field=content 123</tt>', 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, [], '<span class="monospace">[optional state=off|toggle] [optional silent=true] (optional name)</span> sets active World, or unsets if no args provided, use <code>state=off</code> and <code>state=toggle</code> to deactivate or toggle a World, use <code>silent=true</code> 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) {

View File

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