Add prompt injection filters

d6f34f7b2c812462030be6af198b4b982808d895

Cohee <18619528+Cohee1207@users.noreply.github.com>

Parents: 9c43999e4b3975455e4c73132412baf7df705eed

public/script.js+75 -33 Viewed

+    }
+}
 -function getAllExtensionPrompts() {
 +/**
 -    const value = Object
 + * Returns all extension prompts combined.
 -        .values(extension_prompts)
 + * @returns {Promise<string>} Combined extension prompts
 -        .filter(x => x.value)
 + */
 -        .map(x => x.value.trim())
 +async function getAllExtensionPrompts() {
 -        .join('\n');
 +    const values = [];
 -    return value.length ? substituteParams(value) : '';
 +    for (const prompt of Object.values(extension_prompts)) {
 +        const value = prompt?.value?.trim();
++
 +        if (!value) {
 +            continue;
+        }
 -// Wrapper to fetch extension prompts by module name
 +        const hasFilter = typeof prompt.filter === 'function';
 -export function getExtensionPromptByName(moduleName) {
 +        if (hasFilter && !await prompt.filter()) {
 -    if (moduleName) {
 +            continue;
 -        return substituteParams(extension_prompts[moduleName]?.value);
 -    } else {
 -        return;
+        }
++
 +        values.push(value);
 +    }
++
 +    return substituteParams(values.join('\n'));
 +}
++
 +/**
 + * Wrapper to fetch extension prompts by module name
 + * @param {string} moduleName Module name
 + * @returns {Promise<string>} Extension prompt
 + */
 +export async function getExtensionPromptByName(moduleName) {
 +    if (!moduleName) {
 +        return '';
 +    }
++
 +    const prompt = extension_prompts[moduleName];
++
 +    if (!prompt) {
 +        return '';
 +    }
++
 +    const hasFilter = typeof prompt.filter === 'function';
++
 +    if (hasFilter && !await prompt.filter()) {
 +        return '';
 +    }
++
 +    return substituteParams(prompt.value);
+}
 /**

  * @param {string} [separator] Separator for joining multiple prompts
  * @param {number} [role] Role of the prompt
  * @param {boolean} [wrap] Wrap start and end with a separator
  * @returns {Promise<string>} Extension prompt
  */
 export async function getExtensionPrompt(position = extension_prompt_types.IN_PROMPT, depth = undefined, separator = '\n', role = undefined, wrap = true) {
 -    let extension_prompt = Object.keys(extension_prompts)
 +    const filterByFunction = async (prompt) => {
 +        const hasFilter = typeof prompt.filter === 'function';
 +        if (hasFilter && !await prompt.filter()) {
 +            return false;
 +        }
 +        return true;
 +    };
 +    const promptPromises = Object.keys(extension_prompts)
         .sort()
         .map((x) => extension_prompts[x])
         .filter(x => x.position == position && x.value)
         .filter(x => depth === undefined || x.depth === undefined || x.depth === depth)
         .filter(x => role === undefined || x.role === undefined || x.role === role)
 -        .map(x => x.value.trim())
 +        .filter(filterByFunction);
 -        .join(separator);
 +    const prompts = await Promise.all(promptPromises);
 -    if (wrap && extension_prompt.length && !extension_prompt.startsWith(separator)) {
++
 -        extension_prompt = separator + extension_prompt;
 +    let values = prompts.map(x => x.value.trim()).join(separator);
 +    if (wrap && values.length && !values.startsWith(separator)) {
 +        values = separator + values;
+    }
     if (wrap && extension_promptvalues.length && !extension_promptvalues.endsWith(separator)) {
         extension_promptvalues = extension_promptvalues + separator;
+    }
     if (extension_promptvalues.length) {
         extension_promptvalues = substituteParams(extension_promptvalues);
+    }
     return extension_promptvalues;
+}
 export function baseChatReplace(value, name1, name2) {

     // Inject all Depth prompts. Chat Completion does it separately
     let injectedIndices = [];
     if (main_api !== 'openai') {
         injectedIndices = await doChatInject(coreChat, isContinue);
+    }
     // Insert character jailbreak as the last user message (if exists, allowed, preferred, and not using Chat Completion)

+    }
     // Call combined AN into Generate
     const beforeScenarioAnchor = (await getExtensionPrompt(extension_prompt_types.BEFORE_PROMPT)).trimStart();
     const afterScenarioAnchor = await getExtensionPrompt(extension_prompt_types.IN_PROMPT);
     const storyStringParams = {
         description: description,

             ...thisPromptBits[currentArrayEntry],
             rawPrompt: generate_data.prompt || generate_data.input,
             mesId: getNextMessageId(type),
             allAnchors: await getAllExtensionPrompts(),
             chatInjects: injectedIndices?.map(index => arrMes[arrMes.length - index - 1])?.join('') || '',
             summarizeString: (extension_prompts['1_memory']?.value || ''),
             authorsNoteString: (extension_prompts['2_floating_prompt']?.value || ''),

  * Injects extension prompts into chat messages.
  * @param {object[]} messages Array of chat messages
  * @param {boolean} isContinue Whether the generation is a continuation. If true, the extension prompts of depth 0 are injected at position 1.
  * @returns {Promise<number[]>} Array of indices where the extension prompts were injected
  */
 async function doChatInject(messages, isContinue) {
     const injectedIndices = [];
     let totalInsertedMessages = 0;
     messages.reverse();

         const wrap = false;
         for (const role of roles) {
             const extensionPrompt = String(await getExtensionPrompt(extension_prompt_types.IN_CHAT, i, separator, role, wrap)).trimStart();
             const isNarrator = role === extension_prompt_roles.SYSTEM;
             const isUser = role === extension_prompt_roles.USER;
             const name = names[role];

  * @param {number} depth Insertion depth. 0 represets the last message in context. Expected values up to MAX_INJECTION_DEPTH.
  * @param {number} role Extension prompt role. Defaults to SYSTEM.
  * @param {boolean} scan Should the prompt be included in the world info scan.
 + * @param {(function(): Promise<boolean>|boolean)} filter Filter function to determine if the prompt should be injected.
  */
 export function setExtensionPrompt(key, value, position, depth, scan = false, role = extension_prompt_roles.SYSTEM, filter = null) {
     extension_prompts[key] = {
         value: String(value),
         position: Number(position),
         depth: Number(depth),
         scan: !!scan,
         role: Number(role ?? extension_prompt_roles.SYSTEM),
 +        filter: filter,
     };
+}

public/scripts/openai.js+10 -6 Viewed

+ *
  * @param {Prompt[]} prompts - Array containing injection prompts.
  * @param {Object[]} messages - Array containing all messages.
 + * @returns {Promise<Object[]>} - Array containing all messages with injections.
  */
 async function populationInjectionPrompts(prompts, messages) {
     let totalInsertedMessages = 0;
     const roleTypes = {

             // Get prompts for current role
             const rolePrompts = depthPrompts.filter(prompt => prompt.role === role).map(x => x.content).join(separator);
             // Get extension prompt
             const extensionPrompt = await getExtensionPrompt(extension_prompt_types.IN_CHAT, i, separator, roleTypes[role], wrap);
             const jointPrompt = [rolePrompts, extensionPrompt].filter(x => x).map(x => x.trim()).join(separator);

+    }
     // Add in-chat injections
     messages = await populationInjectionPrompts(absolutePrompts, messages);
     // Decide whether dialogue examples should always be added
     if (power_user.pin_examples) {

  * @param {string} options.systemPromptOverride
  * @param {string} options.jailbreakPromptOverride
  * @param {string} options.personaDescription
  * @returns {Promise<Object>} prompts - The prepared and merged system and user-defined prompts.
  */
 async function preparePromptsForChatCompletion({ Scenario, charPersonality, name2, worldInfoBefore, worldInfoAfter, charDescription, quietPrompt, bias, extensionPrompts, systemPromptOverride, jailbreakPromptOverride, personaDescription }) {
     const scenarioText = Scenario && oai_settings.scenario_format ? substituteParams(oai_settings.scenario_format) : '';
     const charPersonalityText = charPersonality && oai_settings.personality_format ? substituteParams(oai_settings.personality_format) : '';
     const groupNudge = substituteParams(oai_settings.group_nudge_prompt);

             if (!extensionPrompts[key].value) continue;
             if (![extension_prompt_types.BEFORE_PROMPT, extension_prompt_types.IN_PROMPT].includes(prompt.position)) continue;
 +            const hasFilter = typeof prompt.filter === 'function';
 +            if (hasFilter && !await prompt.filter()) continue;
++
             systemPrompts.push({
                 identifier: key.replace(/\W/g, '_'),
                 position: getPromptPosition(prompt.position),

     try {
         // Merge markers and ordered user prompts with system prompts
         const prompts = await preparePromptsForChatCompletion({
             Scenario,
             charPersonality,
             name2,

public/scripts/world-info.js+1 -1 Viewed

     // Put this code here since otherwise, the chat reference is modified
     for (const key of Object.keys(context.extensionPrompts)) {
         if (context.extensionPrompts[key]?.scan) {
             const prompt = await getExtensionPromptByName(key);
             if (prompt) {
                 buffer.addInject(prompt);
+            }