From 9f0530f4223f18cf62cfe22b1de5158e3bcc0541 Mon Sep 17 00:00:00 2001
From: Cohee <titan0599@gmail.com>
Date: Fri, 23 Jun 2023 22:50:16 +0300
Subject: [PATCH] Export characters with world info embeds.

---
 public/index.html            |  3 +-
 public/script.js             |  3 +-
 public/scripts/world-info.js | 64 ++++++++++++++++++++++++++++++++++--
 server.js                    | 49 +++++++++++++++++++++++++--
 4 files changed, 112 insertions(+), 7 deletions(-)

diff --git a/public/index.html b/public/index.html
index f8f7d8721..343455b3d 100644
--- a/public/index.html
+++ b/public/index.html
@@ -2887,8 +2887,7 @@
             <div class="range-block-counter justifyLeft flex-container flexFlowColumn margin-bot-10px">
                 A selected World Info / Lorebook will be bound to this character.
                 When generating an AI reply, it will be combined with the entries from a global World Info / Lorebook selector.
-                <!-- Telling lies? Uncomment when actually implemented. -->
-                <!-- Exporting a character would also export the selected World Info file embedded in the JSON data. -->
+                Exporting a character would also export the selected World Info file embedded in the JSON data.
             </div>
             <div class="range-block-range wide100p">
                 <select class="character_world_info_selector">
diff --git a/public/script.js b/public/script.js
index 88897114a..44b6ed717 100644
--- a/public/script.js
+++ b/public/script.js
@@ -4891,7 +4891,8 @@ function checkEmbeddedWorld(chid) {
 
         // Only show the alert once per character
         const checkKey = `AlertWI_${characters[chid].avatar}`;
-        if (!localStorage.getItem(checkKey) && !(characters[chid]?.data?.extensions?.world)) {
+        const worldName = characters[chid]?.data?.extensions?.world;
+        if (!localStorage.getItem(checkKey) && (!worldName || !world_names.includes(worldName))) {
             toastr.info(
                 'To import and use it, select "Import Embedded World Info" in the Options menu.',
                 `${characters[chid].name} has an embedded World/Lorebook`,
diff --git a/public/scripts/world-info.js b/public/scripts/world-info.js
index cb43d319b..e7af45a15 100644
--- a/public/scripts/world-info.js
+++ b/public/scripts/world-info.js
@@ -250,6 +250,41 @@ function displayWorldEntries(name, data) {
     $("#world_popup_entries_list").disableSelection();
 }
 
+function setOriginalDataValue(data, uid, key, value) {
+    if (data.originalData && Array.isArray(data.originalData.entries)) {
+        let originalEntry = data.originalData.entries.find(x => x.uid === uid);
+
+        if (!originalEntry) {
+            return;
+        }
+
+        const keyParts = key.split('.');
+        let currentObject = originalEntry;
+
+        for (let i = 0; i < keyParts.length - 1; i++) {
+            const part = keyParts[i];
+
+            if (!currentObject.hasOwnProperty(part)) {
+                currentObject[part] = {};
+            }
+
+            currentObject = currentObject[part];
+        }
+
+        currentObject[keyParts[keyParts.length - 1]] = value;
+    }
+}
+
+function deleteOriginalDataValue(data, uid) {
+    if (data.originalData && Array.isArray(data.originalData.entries)) {
+        const originalIndex = data.originalData.entries.findIndex(x => x.uid === uid);
+
+        if (originalIndex >= 0) {
+            data.originalData.entries.splice(originalIndex, 1);
+        }
+    }
+}
+
 function appendWorldEntry(name, data, entry) {
     const template = $("#entry_edit_template .world_entry").clone();
     template.data("uid", entry.uid);
@@ -270,6 +305,8 @@ function appendWorldEntry(name, data, entry) {
             .split(",")
             .map((x) => x.trim())
             .filter((x) => x);
+
+        setOriginalDataValue(data, uid, "keys", data.entries[uid].key);
         saveWorldInfo(name, data);
     });
     keyInput.val(entry.key.join(",")).trigger("input");
@@ -286,6 +323,8 @@ function appendWorldEntry(name, data, entry) {
             .split(",")
             .map((x) => x.trim())
             .filter((x) => x);
+
+        setOriginalDataValue(data, uid, "secondary_keys", data.entries[uid].keysecondary);
         saveWorldInfo(name, data);
     });
 
@@ -300,6 +339,8 @@ function appendWorldEntry(name, data, entry) {
         const uid = $(this).data("uid");
         const value = $(this).val();
         data.entries[uid].comment = value;
+
+        setOriginalDataValue(data, uid, "comment", data.entries[uid].comment);
         saveWorldInfo(name, data);
     });
     commentToggle.data("uid", entry.uid);
@@ -333,6 +374,8 @@ function appendWorldEntry(name, data, entry) {
         const uid = $(this).data("uid");
         const value = $(this).val();
         data.entries[uid].content = value;
+
+        setOriginalDataValue(data, uid, "content", data.entries[uid].content);
         saveWorldInfo(name, data);
 
         // count tokens
@@ -348,6 +391,8 @@ function appendWorldEntry(name, data, entry) {
         const uid = $(this).data("uid");
         const value = $(this).prop("checked");
         data.entries[uid].selective = value;
+
+        setOriginalDataValue(data, uid, "selective", data.entries[uid].selective);
         saveWorldInfo(name, data);
 
         const keysecondary = $(this)
@@ -378,6 +423,7 @@ function appendWorldEntry(name, data, entry) {
         const uid = $(this).data("uid");
         const value = $(this).prop("checked");
         data.entries[uid].constant = value;
+        setOriginalDataValue(data, uid, "constant", data.entries[uid].constant);
         saveWorldInfo(name, data);
     });
     constantInput.prop("checked", entry.constant).trigger("input");
@@ -390,6 +436,7 @@ function appendWorldEntry(name, data, entry) {
         const value = Number($(this).val());
 
         data.entries[uid].order = !isNaN(value) ? value : 0;
+        setOriginalDataValue(data, uid, "insertion_order", data.entries[uid].order);
         saveWorldInfo(name, data);
     });
     orderInput.val(entry.order).trigger("input");
@@ -405,6 +452,10 @@ function appendWorldEntry(name, data, entry) {
         const uid = $(this).data("uid");
         const value = Number($(this).val());
         data.entries[uid].position = !isNaN(value) ? value : 0;
+        // Spec v2 only supports before_char and after_char
+        setOriginalDataValue(data, uid, "position", data.entries[uid].position == 0 ? 'before_char' : 'after_char');
+        // Write the original value as extensions field
+        setOriginalDataValue(data, uid, "extensions.position", data.entries[uid].position);
         saveWorldInfo(name, data);
     });
     template
@@ -422,6 +473,7 @@ function appendWorldEntry(name, data, entry) {
         const uid = $(this).data("uid");
         const value = $(this).prop("checked");
         data.entries[uid].disable = value;
+        setOriginalDataValue(data, uid, "enabled", !data.entries[uid].disable);
         saveWorldInfo(name, data);
     });
     disableInput.prop("checked", entry.disable).trigger("input");
@@ -432,6 +484,7 @@ function appendWorldEntry(name, data, entry) {
         const uid = $(this).data("uid");
         const value = $(this).prop("checked");
         data.entries[uid].excludeRecursion = value;
+        setOriginalDataValue(data, uid, "extensions.exclude_recursion", data.entries[uid].excludeRecursion);
         saveWorldInfo(name, data);
     });
     excludeRecursionInput.prop("checked", entry.excludeRecursion).trigger("input");
@@ -442,6 +495,7 @@ function appendWorldEntry(name, data, entry) {
     deleteButton.on("click", function () {
         const uid = $(this).data("uid");
         deleteWorldInfoEntry(data, uid);
+        deleteOriginalDataValue(data, uid);
         $(this).closest(".world_entry").remove();
         saveWorldInfo(name, data);
     });
@@ -877,9 +931,14 @@ function convertNovelLorebook(inputObj) {
 }
 
 function convertCharacterBook(characterBook) {
-    const result = { entries: {} };
+    const result = { entries: {}, originalData: characterBook };
 
     characterBook.entries.forEach((entry, index) => {
+        // Not in the spec, but this is needed to find the entry in the original data
+        if (entry.id === undefined) {
+            entry.id = index;
+        }
+
         result.entries[index] = {
             uid: entry.id || index,
             key: entry.keys,
@@ -889,7 +948,8 @@ function convertCharacterBook(characterBook) {
             constant: entry.constant || false,
             selective: entry.selective || false,
             order: entry.insertion_order,
-            position: entry.position === "before_char" ? world_info_position.before : world_info_position.after,
+            position: entry.extensions?.position ?? (entry.position === "before_char" ? world_info_position.before : world_info_position.after),
+            excludeRecursion: entry.extensions?.exclude_recursion ?? false,
             disable: !entry.enabled,
             addMemo: entry.comment ? true : false,
         };
diff --git a/server.js b/server.js
index 597839522..a1ae35fdb 100644
--- a/server.js
+++ b/server.js
@@ -835,8 +835,24 @@ function charaFormatData(data) {
     //_.set(char, 'data.extensions.avatar', 'none');
     //_.set(char, 'data.extensions.chat', data.ch_name + ' - ' + humanizedISO8601DateTime());
 
-    // TODO: Character book
-    _//.set(char, 'data.character_book', undefined);
+    if (data.world) {
+        try {
+            const file = readWorldInfoFile(data.world);
+
+            // File was imported - save it to the character book
+            if (file && file.originalData) {
+                _.set(char, 'data.character_book', file.originalData);
+            }
+
+            // File was not imported - convert the world info to the character book
+            if (file && file.entries) {
+                _.set(char, 'data.character_book', convertWorldInfoToCharacterBook(data.world, file.entries));
+            }
+
+        } catch {
+            console.debug(`Failed to read world info file: ${data.world}. Character book will not be available.`);
+        }
+    }
 
     return char;
 }
@@ -1437,6 +1453,35 @@ app.post('/savetheme', jsonParser, (request, response) => {
     return response.sendStatus(200);
 });
 
+function convertWorldInfoToCharacterBook(name, entries) {
+    const result = { entries: [], name };
+
+    for (const index in entries) {
+        const entry = entries[index];
+
+        const originalEntry = {
+            id: entry.uid,
+            keys: entry.key,
+            secondary_keys: entry.keysecondary,
+            comment: entry.comment,
+            content: entry.content,
+            constant: entry.constant,
+            selective: entry.selective,
+            insertion_order: entry.order,
+            enabled: !entry.disable,
+            position: entry.position == 0 ? 'before_char' : 'after_char',
+            extensions: {
+                position: entry.position,
+                exclude_recursion: entry.excludeRecursion,
+            }
+        };
+
+        result.entries.push(originalEntry);
+    }
+
+    return result;
+}
+
 function readWorldInfoFile(worldInfoName) {
     if (!worldInfoName) {
         return { entries: {} };