From 6ae120900ddabe7accf2313d1c8e52a1a86cf18f Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Fri, 22 Nov 2024 17:39:49 +0000
Subject: [PATCH] Support images in custom prompt post-processing

---
 src/prompt-converters.js | 45 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)

diff --git a/src/prompt-converters.js b/src/prompt-converters.js
index ccea89139..344a889b0 100644
--- a/src/prompt-converters.js
+++ b/src/prompt-converters.js
@@ -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