From e0b5df97c485c13282b7479e5f1003b145f43094 Mon Sep 17 00:00:00 2001
From: city-unit <140349364+city-unit@users.noreply.github.com>
Date: Wed, 20 Sep 2023 22:48:05 -0400
Subject: [PATCH] Add branching as distinct from bookmarking

---
 public/index.html           |  3 ++-
 public/script.js            | 17 ++++++++++++++++-
 public/scripts/bookmarks.js | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/public/index.html b/public/index.html
index 12029608e..1e44acd0e 100644
--- a/public/index.html
+++ b/public/index.html
@@ -4123,7 +4123,8 @@
                             <div title="Generate Image" class="sd_message_gen fa-solid fa-paintbrush" data-i18n="[title]Generate Image"></div>
                             <div title="Narrate" class="mes_narrate fa-solid fa-bullhorn" data-i18n="[title]Narrate"></div>
                             <div title="Prompt" class="mes_prompt fa-solid fa-square-poll-horizontal " data-i18n="[title]Prompt"></div>
-                            <div title="Create bookmark" class="mes_create_bookmark fa-regular fa-code-branch" data-i18n="[title]Create Bookmark"></div>
+                            <div title="Create bookmark" class="mes_create_bookmark fa-regular fa-solid fa-book-bookmark" data-i18n="[title]Create Bookmark"></div>
+                            <div title="Create branch" class="mes_create_branch fa-regular fa-code-branch" data-i18n="[title]Create Branch"></div>
                             <div title="Copy" class="mes_copy fa-solid fa-copy " data-i18n="[title]Copy"></div>
                         </div>
                         <div title="Open bookmark chat" class="mes_bookmark fa-solid fa-bookmark" data-i18n="[title]Open bookmark chat"></div>
diff --git a/public/script.js b/public/script.js
index 78ca282f8..9db5878f3 100644
--- a/public/script.js
+++ b/public/script.js
@@ -103,7 +103,8 @@ import {
 
 import {
     createNewBookmark,
-    showBookmarksButtons
+    showBookmarksButtons,
+    createBranch,
 } from "./scripts/bookmarks.js";
 
 import {
@@ -6525,6 +6526,11 @@ function swipe_left() {      // when we swipe left..but no generation.
     }
 }
 
+async function branchChat(mesID) {
+    let name = await createBranch(mesID);
+    await openCharacterChat(name);
+}
+
 // when we click swipe right button
 const swipe_right = () => {
     if (chat.length - 1 === Number(this_edit_mes_id)) {
@@ -6705,6 +6711,8 @@ const swipe_right = () => {
     }
 }
 
+
+
 function displayOverrideWarnings() {
     if (!this_chid || !selected_group) {
         $('.prompt_overridden').hide();
@@ -8426,6 +8434,13 @@ jQuery(async function () {
         }
     });
 
+    $(document).on("click", ".mes_create_branch", async function () {
+        var selected_mes_id = $(this).closest(".mes").attr("mesid");
+        if (selected_mes_id !== undefined) {
+            branchChat(selected_mes_id);
+        }
+    });
+
     $(document).on("click", ".mes_stop", function () {
         if (streamingProcessor) {
             streamingProcessor.abortController.abort();
diff --git a/public/scripts/bookmarks.js b/public/scripts/bookmarks.js
index e22b135fd..29d86dd18 100644
--- a/public/scripts/bookmarks.js
+++ b/public/scripts/bookmarks.js
@@ -133,6 +133,39 @@ async function saveBookmarkMenu() {
     return createNewBookmark(chat.length - 1);
 }
 
+export async function createBranch(mesId) {
+    if (!chat.length) {
+        toastr.warning('The chat is empty.', 'Bookmark creation failed');
+        return;
+    }
+
+    if (mesId < 0 || mesId >= chat.length) {
+        toastr.warning('Invalid message ID.', 'Bookmark creation failed');
+        return;
+    }
+
+    const lastMes = chat[mesId];
+    const mainChat = selected_group ? groups?.find(x => x.id == selected_group)?.chat_id : characters[this_chid].chat;
+    const newMetadata = { main_chat: mainChat };
+    let name = `Branch #${mesId} - ${humanizedDateTime()}`
+
+    if (selected_group) {
+        await saveGroupBookmarkChat(selected_group, name, newMetadata, mesId);
+    } else {
+        await saveChat(name, newMetadata, mesId);
+    }
+    // append to branches list if it exists
+    // otherwise create it
+    if (typeof lastMes.extra !== 'object') {
+        lastMes.extra = {};
+    }
+    if (typeof lastMes.extra['branches'] !== 'object') {
+        lastMes.extra['branches'] = [];
+    }
+    lastMes.extra['branches'].push(name);
+    return name;
+}
+
 async function createNewBookmark(mesId) {
     if (!chat.length) {
         toastr.warning('The chat is empty.', 'Bookmark creation failed');