rixty/SillyTavern
1
0
Fork 0
mirror of https://github.com/SillyTavern/SillyTavern.git synced 2025-06-05 21:59:27 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add ability to insert role-typed prompt injections

Browse Source
This commit is contained in:
Cohee
2024-03-23 17:36:43 +02:00
parent 76cde592ad
commit 607df2f555
5 changed files with 201 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files

233
public/script.js
View File

@@ -495,6 +495,9 @@ const durationSaveEdit = 1000;
const saveSettingsDebounced = debounce(() => saveSettings(), durationSaveEdit);
export const saveCharacterDebounced = debounce(() => $('#create_button').trigger('click'), durationSaveEdit);
/**
* @enum {string} System message types
*/
const system_message_types = {
HELP: 'help',
WELCOME: 'welcome',
@@ -511,12 +514,24 @@ const system_message_types = {
MACROS: 'macros',
};
/**
* @enum {number} Extension prompt types
*/
const extension_prompt_types = {
IN_PROMPT: 0,
IN_CHAT: 1,
BEFORE_PROMPT: 2,
};
/**
* @enum {number} Extension prompt roles
*/
export const extension_prompt_roles = {
SYSTEM: 0,
USER: 1,
ASSISTANT: 2,
};
export const MAX_INJECTION_DEPTH = 1000;
let system_messages = {};
@@ -2462,11 +2477,22 @@ function getExtensionPromptByName(moduleName) {
}
}
function getExtensionPrompt(position = 0, depth = undefined, separator = '\n') {
/**
* Returns the extension prompt for the given position, depth, and role.
* If multiple prompts are found, they are joined with a separator.
* @param {number} [position] Position of the prompt
* @param {number} [depth] Depth of the prompt
* @param {string} [separator] Separator for joining multiple prompts
* @param {number} [role] Role of the prompt
* @returns {string} Extension prompt
*/
function getExtensionPrompt(position = extension_prompt_types.IN_PROMPT, depth = undefined, separator = '\n', role = undefined) {
let extension_prompt = Object.keys(extension_prompts)
.sort()
.map((x) => extension_prompts[x])
.filter(x => x.position == position && x.value && (depth === undefined || x.depth == depth))
.filter(x => x.position == position && x.value)
.filter(x => x.depth === undefined || x.depth === depth)
.filter(x => x.role === undefined || x.role === role)
.map(x => x.value.trim())
.join(separator);
if (extension_prompt.length && !extension_prompt.startsWith(separator)) {
@@ -3160,6 +3186,21 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
console.debug('Skipping extension interceptors for dry run');
}
// Adjust token limit for Horde
let adjustedParams;
if (main_api == 'koboldhorde' && (horde_settings.auto_adjust_context_length || horde_settings.auto_adjust_response_length)) {
try {
adjustedParams = await adjustHordeGenerationParams(max_context, amount_gen);
}
catch {
unblockGeneration();
return Promise.resolve();
}
if (horde_settings.auto_adjust_context_length) {
this_max_context = (adjustedParams.maxContextLength - adjustedParams.maxLength);
}
}
console.log(`Core/all messages: ${coreChat.length}/${chat.length}`);
// kingbri MARK: - Make sure the prompt bias isn't the same as the user bias
@@ -3172,6 +3213,32 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
}
//////////////////////////////////
// Extension added strings
// Set non-WI AN
setFloatingPrompt();
// Add WI to prompt (and also inject WI to AN value via hijack)
const chatForWI = coreChat.map(x => `${x.name}: ${x.mes}`).reverse();
let { worldInfoString, worldInfoBefore, worldInfoAfter, worldInfoDepth } = await getWorldInfoPrompt(chatForWI, this_max_context, dryRun);
if (skipWIAN !== true) {
console.log('skipWIAN not active, adding WIAN');
// Add all depth WI entries to prompt
flushWIDepthInjections();
if (Array.isArray(worldInfoDepth)) {
worldInfoDepth.forEach((e) => {
const joinedEntries = e.entries.join('\n');
setExtensionPrompt(`customDepthWI-${e.depth}-${e.role}`, joinedEntries, extension_prompt_types.IN_CHAT, e.depth, false, e.role);
});
}
} else {
console.log('skipping WIAN');
}
// Inject all Depth prompts. Chat Completion does it separately
if (main_api !== 'openai') {
doChatInject(coreChat, isContinue);
}
// Insert character jailbreak as the last user message (if exists, allowed, preferred, and not using Chat Completion)
if (power_user.context.allow_jailbreak && power_user.prefer_character_jailbreak && main_api !== 'openai' && jailbreak) {
@@ -3190,9 +3257,7 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
let chat2 = [];
let continue_mag = '';
for (let i = coreChat.length - 1, j = 0; i >= 0; i--, j++) {
// For OpenAI it's only used in WI
if (main_api == 'openai' && (!world_info || world_info.length === 0)) {
console.debug('No WI, skipping chat2 for OAI');
if (main_api == 'openai') {
break;
}
@@ -3215,49 +3280,12 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
}
}
// Adjust token limit for Horde
let adjustedParams;
if (main_api == 'koboldhorde' && (horde_settings.auto_adjust_context_length || horde_settings.auto_adjust_response_length)) {
try {
adjustedParams = await adjustHordeGenerationParams(max_context, amount_gen);
}
catch {
unblockGeneration();
return Promise.resolve();
}
if (horde_settings.auto_adjust_context_length) {
this_max_context = (adjustedParams.maxContextLength - adjustedParams.maxLength);
}
}
// Extension added strings
// Set non-WI AN
setFloatingPrompt();
// Add WI to prompt (and also inject WI to AN value via hijack)
let { worldInfoString, worldInfoBefore, worldInfoAfter, worldInfoDepth } = await getWorldInfoPrompt(chat2, this_max_context, dryRun);
if (skipWIAN !== true) {
console.log('skipWIAN not active, adding WIAN');
// Add all depth WI entries to prompt
flushWIDepthInjections();
if (Array.isArray(worldInfoDepth)) {
worldInfoDepth.forEach((e) => {
const joinedEntries = e.entries.join('\n');
setExtensionPrompt(`customDepthWI-${e.depth}`, joinedEntries, extension_prompt_types.IN_CHAT, e.depth);
});
}
} else {
console.log('skipping WIAN');
}
// Add persona description to prompt
addPersonaDescriptionExtensionPrompt();
// Call combined AN into Generate
let allAnchors = getAllExtensionPrompts();
const beforeScenarioAnchor = getExtensionPrompt(extension_prompt_types.BEFORE_PROMPT).trimStart();
const afterScenarioAnchor = getExtensionPrompt(extension_prompt_types.IN_PROMPT);
let zeroDepthAnchor = getExtensionPrompt(extension_prompt_types.IN_CHAT, 0, ' ');
const storyStringParams = {
description: description,
@@ -3560,40 +3588,6 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
// Deep clone
let finalMesSend = structuredClone(mesSend);
// TODO: Rewrite getExtensionPrompt to not require multiple for loops
// Set all extension prompts where insertion depth > mesSend length
if (finalMesSend.length) {
for (let upperDepth = MAX_INJECTION_DEPTH; upperDepth >= finalMesSend.length; upperDepth--) {
const upperAnchor = getExtensionPrompt(extension_prompt_types.IN_CHAT, upperDepth);
if (upperAnchor && upperAnchor.length) {
finalMesSend[0].extensionPrompts.push(upperAnchor);
}
}
}
finalMesSend.forEach((mesItem, index) => {
if (index === 0) {
return;
}
const anchorDepth = Math.abs(index - finalMesSend.length);
// NOTE: Depth injected here!
const extensionAnchor = getExtensionPrompt(extension_prompt_types.IN_CHAT, anchorDepth);
if (anchorDepth >= 0 && extensionAnchor && extensionAnchor.length) {
mesItem.extensionPrompts.push(extensionAnchor);
}
});
// TODO: Move zero-depth anchor append to work like CFG and bias appends
if (zeroDepthAnchor?.length && !isContinue) {
console.debug(/\s/.test(finalMesSend[finalMesSend.length - 1].message.slice(-1)));
finalMesSend[finalMesSend.length - 1].message +=
/\s/.test(finalMesSend[finalMesSend.length - 1].message.slice(-1))
? zeroDepthAnchor
: `${zeroDepthAnchor}`;
}
let cfgPrompt = {};
if (cfgGuidanceScale && cfgGuidanceScale?.value !== 1) {
cfgPrompt = getCfgPrompt(cfgGuidanceScale, isNegative);
@@ -3969,6 +3963,57 @@ async function Generate(type, { automatic_trigger, force_name2, quiet_prompt, qu
}
}
/**
* 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 {void}
*/
function doChatInject(messages, isContinue) {
let totalInsertedMessages = 0;
messages.reverse();
for (let i = 0; i <= MAX_INJECTION_DEPTH; i++) {
// Order of priority (most important go lower)
const roles = [extension_prompt_roles.SYSTEM, extension_prompt_roles.USER, extension_prompt_roles.ASSISTANT];
const names = {
[extension_prompt_roles.SYSTEM]: '',
[extension_prompt_roles.USER]: name1,
[extension_prompt_roles.ASSISTANT]: name2,
}
const roleMessages = [];
const separator = '\n';
for (const role of roles) {
// Get extension prompt
const extensionPrompt = String(getExtensionPrompt(extension_prompt_types.IN_CHAT, i, separator, role)).trimStart();
const isNarrator = role === extension_prompt_roles.SYSTEM;
const isUser = role === extension_prompt_roles.USER;
const name = names[role];
if (extensionPrompt) {
roleMessages.push({
name: name,
is_user: isUser,
mes: extensionPrompt,
extra: {
type: isNarrator ? system_message_types.NARRATOR : null,
}
});
}
}
if (roleMessages.length) {
const depth = isContinue && i === 0 ? 1 : i;
const injectIdx = depth + totalInsertedMessages;
messages.splice(injectIdx, 0, ...roleMessages);
totalInsertedMessages += roleMessages.length;
}
}
messages.reverse();
}
function flushWIDepthInjections() {
//prevent custom depth WI entries (which have unique random key names) from duplicating
for (const key of Object.keys(extension_prompts)) {
@@ -4229,25 +4274,6 @@ function addChatsSeparator(mesSendString) {
}
}
// There's a TODO related to zero-depth anchors; not removing this function until that's resolved
// eslint-disable-next-line no-unused-vars
function appendZeroDepthAnchor(force_name2, zeroDepthAnchor, finalPrompt) {
const trimBothEnds = !force_name2;
let trimmedPrompt = (trimBothEnds ? zeroDepthAnchor.trim() : zeroDepthAnchor.trimEnd());
if (trimBothEnds && !finalPrompt.endsWith('\n')) {
finalPrompt += '\n';
}
finalPrompt += trimmedPrompt;
if (force_name2) {
finalPrompt += ' ';
}
return finalPrompt;
}
async function DupeChar() {
if (!this_chid) {
toastr.warning('You must first select a character to duplicate!');
@@ -5715,7 +5741,7 @@ async function uploadUserAvatar(e) {
}
const rawFile = formData.get('avatar');
if (rawFile instanceof File){
if (rawFile instanceof File) {
const convertedFile = await ensureImageFormatSupported(rawFile);
formData.set('avatar', convertedFile);
}
@@ -6656,10 +6682,17 @@ function select_rm_characters() {
* @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} 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.
*/
export function setExtensionPrompt(key, value, position, depth, scan = false) {
extension_prompts[key] = { value: String(value), position: Number(position), depth: Number(depth), scan: !!scan };
export function setExtensionPrompt(key, value, position, depth, scan = false, role = extension_prompt_roles.SYSTEM) {
extension_prompts[key] = {
value: String(value),
position: Number(position),
depth: Number(depth),
scan: !!scan,
role: Number(role ?? extension_prompt_roles.SYSTEM),
};
}
/**
@@ -7223,7 +7256,7 @@ async function createOrEditCharacter(e) {
formData.set('fav', fav_ch_checked);
const rawFile = formData.get('avatar');
if (rawFile instanceof File){
if (rawFile instanceof File) {
const convertedFile = await ensureImageFormatSupported(rawFile);
formData.set('avatar', convertedFile);
}
@@ -8450,7 +8483,7 @@ jQuery(async function () {
$('#advanced_div').click(function () {
if (!is_advanced_char_open) {
is_advanced_char_open = true;
$('#character_popup').css({'display': 'flex', 'opacity': 0.0}).addClass('open');
$('#character_popup').css({ 'display': 'flex', 'opacity': 0.0 }).addClass('open');
$('#character_popup').transition({
opacity: 1.0,
duration: animation_duration,