Add ability to insert role-typed prompt injections

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

View File

@ -10,6 +10,7 @@ import {
characters,
event_types,
eventSource,
extension_prompt_roles,
extension_prompt_types,
Generate,
getExtensionPrompt,
@ -648,6 +649,12 @@ function formatWorldInfo(value) {
function populationInjectionPrompts(prompts, messages) {
let totalInsertedMessages = 0;
const roleTypes = {
'system': extension_prompt_roles.SYSTEM,
'user': extension_prompt_roles.USER,
'assistant': extension_prompt_roles.ASSISTANT,
};
for (let i = 0; i <= MAX_INJECTION_DEPTH; i++) {
// Get prompts for current depth
const depthPrompts = prompts.filter(prompt => prompt.injection_depth === i && prompt.content);
@ -655,14 +662,15 @@ function populationInjectionPrompts(prompts, messages) {
// Order of priority (most important go lower)
const roles = ['system', 'user', 'assistant'];
const roleMessages = [];
const separator = '\n';
for (const role of roles) {
// Get prompts for current role
const rolePrompts = depthPrompts.filter(prompt => prompt.role === role).map(x => x.content).join('\n');
// Get extension prompt (only for system role)
const extensionPrompt = role === 'system' ? getExtensionPrompt(extension_prompt_types.IN_CHAT, i) : '';
const rolePrompts = depthPrompts.filter(prompt => prompt.role === role).map(x => x.content).join(separator);
// Get extension prompt
const extensionPrompt = getExtensionPrompt(extension_prompt_types.IN_CHAT, i, separator, roleTypes[role]);
const jointPrompt = [rolePrompts, extensionPrompt].filter(x => x).map(x => x.trim()).join('\n');
const jointPrompt = [rolePrompts, extensionPrompt].filter(x => x).map(x => x.trim()).join(separator);
if (jointPrompt && jointPrompt.length) {
roleMessages.push({ 'role': role, 'content': jointPrompt });