From ed74eedc5b866e50a4ffde319dd0aa5c0f27413a Mon Sep 17 00:00:00 2001
From: SillyLossy <titan0599@gmail.com>
Date: Tue, 30 May 2023 23:57:56 +0300
Subject: [PATCH] Display code blocks in translated messages

---
 public/script.js                             | 10 +++++++++-
 public/scripts/extensions/translate/index.js |  8 ++++----
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/public/script.js b/public/script.js
index c9a3741e0..6ef31814d 100644
--- a/public/script.js
+++ b/public/script.js
@@ -1104,6 +1104,14 @@ function getMessageFromTemplate({ mesId, characterName, isUser, avatarImg, bias,
     return mes;
 }
 
+export function updateMessageBlock(messageId, message) {
+    const messageElement = $(`#chat [mesid="${messageId}"]`);
+    const text = message?.extra?.display_text ?? message.mes;
+    messageElement.find('.mes_text').html(messageFormatting(text, message.name, message.is_system, message.is_user));
+    addCopyToCodeBlocks(messageElement)
+    appendImageToMessage(message, messageElement);
+}
+
 export function appendImageToMessage(mes, messageElement) {
     if (mes.extra?.image) {
         const image = messageElement.find('.mes_img');
@@ -1117,7 +1125,7 @@ export function appendImageToMessage(mes, messageElement) {
     }
 }
 
-function addCopyToCodeBlocks(messageElement) {
+export function addCopyToCodeBlocks(messageElement) {
     const codeBlocks = $(messageElement).find("pre code");
     for (let i = 0; i < codeBlocks.length; i++) {
         hljs.highlightElement(codeBlocks.get(i));
diff --git a/public/scripts/extensions/translate/index.js b/public/scripts/extensions/translate/index.js
index 1c12e0f07..3c828ff60 100644
--- a/public/scripts/extensions/translate/index.js
+++ b/public/scripts/extensions/translate/index.js
@@ -3,10 +3,10 @@ import {
     eventSource,
     event_types,
     getRequestHeaders,
-    messageFormatting,
     reloadCurrentChat,
     saveSettingsDebounced,
     substituteParams,
+    updateMessageBlock,
 } from "../../../script.js";
 import { extension_settings, getContext } from "../../extensions.js";
 
@@ -168,7 +168,7 @@ async function translateIncomingMessage(messageId) {
     const translation = await translate(textToTranslate, extension_settings.translate.target_language);
     message.extra.display_text = translation;
 
-    $(`#chat .mes[mesid="${messageId}"] .mes_text`).html(messageFormatting(translation, message.name, message.is_system, message.is_user));
+    updateMessageBlock(messageId, message);
 }
 
 async function translateProviderGoogle(text, lang) {
@@ -211,8 +211,8 @@ async function translateOutgoingMessage(messageId) {
 
     const originalText = message.mes;
     message.extra.display_text = originalText;
-    $(`#chat .mes[mesid="${messageId}"] .mes_text`).html(messageFormatting(originalText, message.name, message.is_system, message.is_user));
     message.mes = await translate(originalText, extension_settings.translate.internal_language);
+    updateMessageBlock(messageId, message);
 
     console.log('translateOutgoingMessage', messageId);
 }
@@ -299,7 +299,7 @@ async function onMessageTranslateClick() {
     // If the message is already translated, revert it back to the original text
     if (message?.extra?.display_text) {
         delete message.extra.display_text;
-        $(`#chat .mes[mesid="${messageId}"] .mes_text`).html(messageFormatting(message.mes, message.name, message.is_system, message.is_user));
+        updateMessageBlock(messageId, message);
     }
     // If the message is not translated, translate it
     else {