Support images in custom prompt post-processing

This commit is contained in:
Cohee 2024-11-22 17:39:49 +00:00
parent 56137d2b57
commit 6ae120900d

View File

@ -628,11 +628,30 @@ export function convertMistralMessages(messages, charName = '', userName = '') {
export function mergeMessages(messages, charName, userName, strict) {
let mergedMessages = [];
/** @type {Map<string,object>} */
const contentTokens = new Map();
// Remove names from the messages
messages.forEach((message) => {
if (!message.content) {
message.content = '';
}
// Flatten contents and replace image URLs with random tokens
if (Array.isArray(message.content)) {
const text = message.content.map((content) => {
if (content.type === 'text') {
return content.text;
}
// Could be extended with other non-text types
if (content.type === 'image_url') {
const token = crypto.randomBytes(32).toString('base64');
contentTokens.set(token, content);
return token;
}
return '';
}).join('\n\n');
message.content = text;
}
if (message.role === 'system' && message.name === 'example_assistant') {
if (charName && !message.content.startsWith(`${charName}: `)) {
message.content = `${charName}: ${message.content}`;
@ -673,6 +692,32 @@ export function mergeMessages(messages, charName, userName, strict) {
});
}
// Check for content tokens and replace them with the actual content objects
if (contentTokens.size > 0) {
mergedMessages.forEach((message) => {
const hasValidToken = Array.from(contentTokens.keys()).some(token => message.content.includes(token));
if (hasValidToken) {
const splitContent = message.content.split('\n\n');
const mergedContent = [];
splitContent.forEach((content) => {
if (contentTokens.has(content)) {
mergedContent.push(contentTokens.get(content));
} else {
if (mergedContent.length > 0 && mergedContent[mergedContent.length - 1].type === 'text') {
mergedContent[mergedContent.length - 1].text += `\n\n${content}`;
} else {
mergedContent.push({ type: 'text', text: content });
}
}
});
message.content = mergedContent;
}
});
}
if (strict) {
for (let i = 0; i < mergedMessages.length; i++) {
// Force mid-prompt system messages to be user messages