add Custom Role

This commit is contained in:
被遗忘的记忆
2024-11-30 00:54:49 +08:00
committed by GitHub
parent d022e5db79
commit 0d9f42a701
+72 -25
View File
@@ -6,7 +6,9 @@ export function formatMessages(messages, proxyModel, randomFileName) {
const USE_BACKSPACE_PREFIX = process.env.USE_BACKSPACE_PREFIX === 'true'; const USE_BACKSPACE_PREFIX = process.env.USE_BACKSPACE_PREFIX === 'true';
// 定义角色映射 // 定义角色映射
const roleFeatures = getRoleFeatures(isClaudeModel, USE_BACKSPACE_PREFIX); const roleFeatures = getRoleFeatures(messages, isClaudeModel, USE_BACKSPACE_PREFIX);
messages = removeCustomRoleDefinitions(messages);
messages = convertRoles(messages, roleFeatures); messages = convertRoles(messages, roleFeatures);
@@ -131,11 +133,39 @@ function escapeRegExp(string) {
} }
// 获取角色特征 // 获取角色特征
function getRoleFeatures(isClaudeModel, useBackspacePrefix) { function getRoleFeatures(messages, isClaudeModel, useBackspacePrefix) {
const prefix = useBackspacePrefix ? '\u0008' : ''; let prefix = useBackspacePrefix ? '\u0008' : '';
const systemRole = `${prefix}${isClaudeModel ? 'System' : 'system'}`; let systemRole = `${prefix}${isClaudeModel ? 'System' : 'system'}`;
const userRole = `${prefix}${isClaudeModel ? 'Human' : 'user'}`; let userRole = `${prefix}${isClaudeModel ? 'Human' : 'user'}`;
const assistantRole = `${prefix}${isClaudeModel ? 'Assistant' : 'assistant'}`; let assistantRole = `${prefix}${isClaudeModel ? 'Assistant' : 'assistant'}`;
// 匹配自定义角色
const rolePattern = /\[\|(\w+)::(.*?)\|\]/g;
let customRoles = {};
messages.forEach(message => {
let content = message.content;
let match;
while ((match = rolePattern.exec(content)) !== null) {
const roleKey = match[1]; // 'system', 'user', 'assistant'
customRoles[roleKey.toLowerCase()] = match[2];
}
});
if (Object.keys(customRoles).length > 0) {
prefix = '';
if (customRoles['system']) {
systemRole = customRoles['system'];
}
if (customRoles['user']) {
userRole = customRoles['user'];
}
if (customRoles['assistant']) {
assistantRole = customRoles['assistant'];
}
}
return { return {
systemRole, systemRole,
@@ -145,6 +175,19 @@ function getRoleFeatures(isClaudeModel, useBackspacePrefix) {
}; };
} }
// 移除 messages 自定义角色格式
function removeCustomRoleDefinitions(messages) {
const rolePattern = /\[\|\w+::.*?\|]/g;
return messages.map(message => {
let newContent = message.content.replace(rolePattern, '');
return {
...message,
content: newContent
};
});
}
// 转换角色 // 转换角色
function convertRoles(messages, roleFeatures) { function convertRoles(messages, roleFeatures) {
const { systemRole, userRole, assistantRole } = roleFeatures; const { systemRole, userRole, assistantRole } = roleFeatures;
@@ -171,33 +214,37 @@ function convertRoles(messages, roleFeatures) {
// 替换 content 中的角色定义 // 替换 content 中的角色定义
function replaceRolesInContent(messages, roleFeatures) { function replaceRolesInContent(messages, roleFeatures) {
// 角色名称映射 // 避免重复添加
const roleMap = { const roleMap = {
'System:': roleFeatures.systemRole + ':', 'System:': roleFeatures.systemRole.replace(roleFeatures.prefix, '') + ':',
'system:': roleFeatures.systemRole + ':', 'system:': roleFeatures.systemRole.replace(roleFeatures.prefix, '') + ':',
'Human:': roleFeatures.userRole + ':', 'Human:': roleFeatures.userRole.replace(roleFeatures.prefix, '') + ':',
'human:': roleFeatures.userRole + ':', 'human:': roleFeatures.userRole.replace(roleFeatures.prefix, '') + ':',
'user:': roleFeatures.userRole + ':', 'user:': roleFeatures.userRole.replace(roleFeatures.prefix, '') + ':',
'Assistant:': roleFeatures.assistantRole + ':', 'Assistant:': roleFeatures.assistantRole.replace(roleFeatures.prefix, '') + ':',
'assistant:': roleFeatures.assistantRole + ':', 'assistant:': roleFeatures.assistantRole.replace(roleFeatures.prefix, '') + ':',
}; };
// 避免重复添加前缀 // 构建角色正则
Object.values(roleMap).forEach(label => {
roleMap[label] = label;
});
// 仅匹配角色标签
const escapedLabels = Object.keys(roleMap).map(label => escapeRegExp(label)); const escapedLabels = Object.keys(roleMap).map(label => escapeRegExp(label));
const roleNamesPattern = new RegExp(`(${escapedLabels.join('|')})`, 'g'); const prefixPattern = roleFeatures.prefix ? escapeRegExp(roleFeatures.prefix) : '';
const roleNamesPattern = new RegExp(`(\\n\\n)(${prefixPattern})?(${escapedLabels.join('|')})`, 'g');
return messages.map(message => { return messages.map(message => {
let newContent = message.content; let newContent = message.content;
// 替换内容中角色 if (typeof newContent === 'string') {
newContent = newContent.replace(roleNamesPattern, (match) => { // 仅替换段落开头角色
return roleMap[match] || match; newContent = newContent.replace(roleNamesPattern, (match, p1, p2, p3) => {
}); const newRoleLabel = roleMap[p3] || p3;
const prefixToUse = p2 !== undefined ? p2 : roleFeatures.prefix;
return p1 + (prefixToUse || '') + newRoleLabel;
});
} else {
console.warn('message.content is not a string:', newContent);
newContent = '';
}
return { return {
...message, ...message,