From ef0d1a986bf1b904ad85526636eb9bcf2cd79e5d Mon Sep 17 00:00:00 2001
From: SillyLossy <titan0599@gmail.com>
Date: Mon, 8 May 2023 18:55:01 +0300
Subject: [PATCH] #260 Adjust message clean up. New setting to allow
 impersonation.

---
 public/index.html            |  6 +++++-
 public/script.js             | 23 +++++++++++++++++++----
 public/scripts/power-user.js |  8 ++++++++
 3 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/public/index.html b/public/index.html
index 9b9c30ca2..3570e20e0 100644
--- a/public/index.html
+++ b/public/index.html
@@ -1552,7 +1552,11 @@
                             </label>
                             <label class="checkbox_label" for="allow_name2_display">
                                 <input id="allow_name2_display" type="checkbox" />
-                                Allow {{char}}: display in bot messages
+                                Allow {{char}}: in bot messages
+                            </label>
+                            <label class="checkbox_label" for="allow_name1_display">
+                                <input id="allow_name1_display" type="checkbox" />
+                                Allow {{user}}: in bot messages
                             </label>
                             <label for="auto_scroll_chat_to_bottom">
                                 <input id="auto_scroll_chat_to_bottom" type="checkbox" />
diff --git a/public/script.js b/public/script.js
index 0cedf8504..4e35c4b2b 100644
--- a/public/script.js
+++ b/public/script.js
@@ -1838,13 +1838,16 @@ async function Generate(type, { automatic_trigger, force_name2, resolve, reject,
                 arrMes[arrMes.length] = item;
 
             } else {
-                setInContextMessages(arrMes.length, type);
                 break;
 
             }
             await delay(1); //For disable slow down (encode gpt-2 need fix)
         }
 
+        if (main_api !== 'openai') {
+            setInContextMessages(arrMes.length, type);
+        }
+
         // Estimate how many unpinned example messages fit in the context
         let count_exm_add = 0;
         if (!power_user.pin_examples) {
@@ -2459,7 +2462,6 @@ function extractMessageFromData(data) {
 }
 
 function cleanUpMessage(getMessage, isImpersonate) {
-    const nameToTrim = isImpersonate ? name2 : name1;
     if (power_user.collapse_newlines) {
         getMessage = collapseNewlines(getMessage);
     }
@@ -2474,8 +2476,21 @@ function cleanUpMessage(getMessage, isImpersonate) {
         getMessage = getMessage.replace(/<BOT>/g, name2);
         getMessage = getMessage.replace(/You:/g, name1 + ':');
     }
-    if (getMessage.indexOf(nameToTrim + ":") != -1) {
-        getMessage = getMessage.substr(0, getMessage.indexOf(nameToTrim + ":"));
+
+    let nameToTrim = isImpersonate ?  name2 : name1;
+
+    if (isImpersonate) {
+        nameToTrim = power_user.allow_name2_display ? '' : name2;
+    }
+    else {
+        nameToTrim = power_user.allow_name1_display ? '' : name1;
+    }
+
+    if (nameToTrim && getMessage.indexOf(`${nameToTrim}:`) == 0) {
+        getMessage = getMessage.substr(0, getMessage.indexOf(`${nameToTrim}:`));
+    }
+    if (nameToTrim && getMessage.indexOf(`\n${nameToTrim}:`) > 0) {
+        getMessage = getMessage.substr(0, getMessage.indexOf(`\n${nameToTrim}:`));
     }
     if (getMessage.indexOf('<|endoftext|>') != -1) {
         getMessage = getMessage.substr(0, getMessage.indexOf('<|endoftext|>'));
diff --git a/public/scripts/power-user.js b/public/scripts/power-user.js
index f4ce852c7..6ac313808 100644
--- a/public/scripts/power-user.js
+++ b/public/scripts/power-user.js
@@ -105,6 +105,7 @@ let power_user = {
     auto_fix_generated_markdown: true,
     send_on_enter: send_on_enter_options.AUTO,
     render_formulas: false,
+    allow_name1_display: false,
     allow_name2_display: false,
     hotswap_enabled: true,
     timer_enabled: true,
@@ -477,6 +478,7 @@ function loadPowerUserSettings(settings, data) {
     $("#play_message_sound").prop("checked", power_user.play_message_sound);
     $("#play_sound_unfocused").prop("checked", power_user.play_sound_unfocused);
     $("#auto_save_msg_edits").prop("checked", power_user.auto_save_msg_edits);
+    $("#allow_name1_display").prop("checked", power_user.allow_name1_display);
     $("#allow_name2_display").prop("checked", power_user.allow_name2_display);
     $("#hotswapEnabled").prop("checked", power_user.hotswap_enabled);
     $("#messageTimerEnabled").prop("checked", power_user.timer_enabled);
@@ -898,6 +900,12 @@ $(document).ready(() => {
         reloadMarkdownProcessor(power_user.render_formulas);
         reloadCurrentChat();
         saveSettingsDebounced();
+    });
+
+    $("#allow_name1_display").on("input", function () {
+        power_user.allow_name1_display = !!$(this).prop('checked');
+        reloadCurrentChat();
+        saveSettingsDebounced();
     })
 
     $("#allow_name2_display").on("input", function () {