Add vector storage to WI scanning

This commit is contained in:
Cohee 2023-12-11 22:47:26 +02:00
parent afe3e824b1
commit 1b11ddc26a
5 changed files with 27 additions and 21 deletions

View File

@ -2447,7 +2447,7 @@ function addPersonaDescriptionExtensionPrompt() {
? `${power_user.persona_description}\n${originalAN}` ? `${power_user.persona_description}\n${originalAN}`
: `${originalAN}\n${power_user.persona_description}`; : `${originalAN}\n${power_user.persona_description}`;
setExtensionPrompt(NOTE_MODULE_NAME, ANWithDesc, chat_metadata[metadata_keys.position], chat_metadata[metadata_keys.depth]); setExtensionPrompt(NOTE_MODULE_NAME, ANWithDesc, chat_metadata[metadata_keys.position], chat_metadata[metadata_keys.depth], extension_settings.note.allowWIScan);
} }
} }
@ -3064,12 +3064,12 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
if (selected_group && Array.isArray(groupDepthPrompts) && groupDepthPrompts.length > 0) { if (selected_group && Array.isArray(groupDepthPrompts) && groupDepthPrompts.length > 0) {
groupDepthPrompts.forEach((value, index) => { groupDepthPrompts.forEach((value, index) => {
setExtensionPrompt('DEPTH_PROMPT_' + index, value.text, extension_prompt_types.IN_CHAT, value.depth); setExtensionPrompt('DEPTH_PROMPT_' + index, value.text, extension_prompt_types.IN_CHAT, value.depth, extension_settings.note.allowWIScan);
}); });
} else { } else {
const depthPromptText = baseChatReplace(characters[this_chid].data?.extensions?.depth_prompt?.prompt?.trim(), name1, name2) || ''; const depthPromptText = baseChatReplace(characters[this_chid].data?.extensions?.depth_prompt?.prompt?.trim(), name1, name2) || '';
const depthPromptDepth = characters[this_chid].data?.extensions?.depth_prompt?.depth ?? depth_prompt_depth_default; const depthPromptDepth = characters[this_chid].data?.extensions?.depth_prompt?.depth ?? depth_prompt_depth_default;
setExtensionPrompt('DEPTH_PROMPT', depthPromptText, extension_prompt_types.IN_CHAT, depthPromptDepth); setExtensionPrompt('DEPTH_PROMPT', depthPromptText, extension_prompt_types.IN_CHAT, depthPromptDepth, extension_settings.note.allowWIScan);
} }
// Parse example messages // Parse example messages
@ -6348,9 +6348,10 @@ function select_rm_characters() {
* @param {string} value Prompt injection value. * @param {string} value Prompt injection value.
* @param {number} position Insertion position. 0 is after story string, 1 is in-chat with custom depth. * @param {number} position Insertion position. 0 is after story string, 1 is in-chat with custom depth.
* @param {number} depth Insertion depth. 0 represets the last message in context. Expected values up to MAX_INJECTION_DEPTH. * @param {number} depth Insertion depth. 0 represets the last message in context. Expected values up to MAX_INJECTION_DEPTH.
* @param {boolean} scan Should the prompt be included in the world info scan.
*/ */
export function setExtensionPrompt(key, value, position, depth) { export function setExtensionPrompt(key, value, position, depth, scan = false) {
extension_prompts[key] = { value: String(value), position: Number(position), depth: Number(depth) }; extension_prompts[key] = { value: String(value), position: Number(position), depth: Number(depth), scan: !!scan };
} }
/** /**

View File

@ -313,7 +313,7 @@ export function setFloatingPrompt() {
} }
} }
} }
context.setExtensionPrompt(MODULE_NAME, prompt, chat_metadata[metadata_keys.position], chat_metadata[metadata_keys.depth]); context.setExtensionPrompt(MODULE_NAME, prompt, chat_metadata[metadata_keys.position], chat_metadata[metadata_keys.depth], extension_settings.note.allowWIScan);
$('#extension_floating_counter').text(shouldAddPrompt ? '0' : messagesTillInsertion); $('#extension_floating_counter').text(shouldAddPrompt ? '0' : messagesTillInsertion);
} }

View File

@ -11,6 +11,7 @@ export const EXTENSION_PROMPT_TAG = '3_vectors';
const settings = { const settings = {
// For both // For both
source: 'transformers', source: 'transformers',
include_wi: false,
// For chats // For chats
enabled_chats: false, enabled_chats: false,
@ -254,7 +255,7 @@ async function vectorizeFile(fileText, fileName, collectionId) {
async function rearrangeChat(chat) { async function rearrangeChat(chat) {
try { try {
// Clear the extension prompt // Clear the extension prompt
setExtensionPrompt(EXTENSION_PROMPT_TAG, '', extension_prompt_types.IN_PROMPT, 0); setExtensionPrompt(EXTENSION_PROMPT_TAG, '', extension_prompt_types.IN_PROMPT, 0, settings.include_wi);
if (settings.enabled_files) { if (settings.enabled_files) {
await processFiles(chat); await processFiles(chat);
@ -319,7 +320,7 @@ async function rearrangeChat(chat) {
// Format queried messages into a single string // Format queried messages into a single string
const insertedText = getPromptText(queriedMessages); const insertedText = getPromptText(queriedMessages);
setExtensionPrompt(EXTENSION_PROMPT_TAG, insertedText, settings.position, settings.depth); setExtensionPrompt(EXTENSION_PROMPT_TAG, insertedText, settings.position, settings.depth, settings.include_wi);
} catch (error) { } catch (error) {
console.error('Vectors: Failed to rearrange chat', error); console.error('Vectors: Failed to rearrange chat', error);
} }
@ -574,6 +575,12 @@ jQuery(async () => {
saveSettingsDebounced(); saveSettingsDebounced();
}); });
$('#vectors_include_wi').prop('checked', settings.include_wi).on('input', () => {
settings.include_wi = !!$('#vectors_include_wi').prop('checked');
Object.assign(extension_settings.vectors, settings);
saveSettingsDebounced();
});
toggleSettings(); toggleSettings();
eventSource.on(event_types.MESSAGE_DELETED, onChatEvent); eventSource.on(event_types.MESSAGE_DELETED, onChatEvent);
eventSource.on(event_types.MESSAGE_EDITED, onChatEvent); eventSource.on(event_types.MESSAGE_EDITED, onChatEvent);

View File

@ -23,6 +23,11 @@
<input type="number" id="vectors_query" class="text_pole widthUnset" min="1" max="99" /> <input type="number" id="vectors_query" class="text_pole widthUnset" min="1" max="99" />
</div> </div>
<label class="checkbox_label" for="vectors_include_wi" title="Query results can activate World Info entries.">
<input id="vectors_include_wi" type="checkbox" class="checkbox">
Include in World Info Scanning
</label>
<hr> <hr>
<h4> <h4>

View File

@ -1684,20 +1684,13 @@ async function checkWorldInfo(chat, maxContext) {
// Add the depth or AN if enabled // Add the depth or AN if enabled
// Put this code here since otherwise, the chat reference is modified // Put this code here since otherwise, the chat reference is modified
if (extension_settings.note.allowWIScan) { for (const key of Object.keys(context.extensionPrompts)) {
for (const key of Object.keys(context.extensionPrompts)) { if (context.extensionPrompts[key]?.scan) {
if (key.startsWith('DEPTH_PROMPT')) { const prompt = getExtensionPromptByName(key);
const depthPrompt = getExtensionPromptByName(key); if (prompt) {
if (depthPrompt) { textToScan = `${prompt}\n${textToScan}`;
textToScan = `${depthPrompt}\n${textToScan}`;
}
} }
} }
const anPrompt = getExtensionPromptByName(NOTE_MODULE_NAME);
if (anPrompt) {
textToScan = `${anPrompt}\n${textToScan}`;
}
} }
// Transform the resulting string // Transform the resulting string
@ -1948,7 +1941,7 @@ async function checkWorldInfo(chat, maxContext) {
if (shouldWIAddPrompt) { if (shouldWIAddPrompt) {
const originalAN = context.extensionPrompts[NOTE_MODULE_NAME].value; const originalAN = context.extensionPrompts[NOTE_MODULE_NAME].value;
const ANWithWI = `${ANTopEntries.join('\n')}\n${originalAN}\n${ANBottomEntries.join('\n')}`; const ANWithWI = `${ANTopEntries.join('\n')}\n${originalAN}\n${ANBottomEntries.join('\n')}`;
context.setExtensionPrompt(NOTE_MODULE_NAME, ANWithWI, chat_metadata[metadata_keys.position], chat_metadata[metadata_keys.depth]); context.setExtensionPrompt(NOTE_MODULE_NAME, ANWithWI, chat_metadata[metadata_keys.position], chat_metadata[metadata_keys.depth], extension_settings.note.allowWIScan);
} }
return { worldInfoBefore, worldInfoAfter, WIDepthEntries }; return { worldInfoBefore, worldInfoAfter, WIDepthEntries };