From 69fdb9090f43816f99d545ce08fad0f4d25f1271 Mon Sep 17 00:00:00 2001
From: Cohee <titan0599@gmail.com>
Date: Tue, 27 Jun 2023 18:39:08 +0300
Subject: [PATCH] Add Chub downloader for characters and lorebooks

---
 public/index.html            |   3 +-
 public/script.js             |  50 +++++++++++++
 public/scripts/world-info.js | 137 ++++++++++++++++++-----------------
 public/style.css             |   3 +-
 server.js                    | 124 +++++++++++++++++++++++++++++--
 5 files changed, 242 insertions(+), 75 deletions(-)

diff --git a/public/index.html b/public/index.html
index 447d3c8fa..f632f15b3 100644
--- a/public/index.html
+++ b/public/index.html
@@ -2267,7 +2267,7 @@
                             <label for="show_card_avatar_urls"><input id="show_card_avatar_urls" type="checkbox" />
                                 <span data-i18n="Show avatar filenames">Show avatar filenames</span>
                             </label>
-                            
+
                             <div class="inline-drawer wide100p flexFlowColumn">
                                 <div class="inline-drawer-toggle inline-drawer-header">
                                     <b>Auto-swipe</b>
@@ -2676,6 +2676,7 @@
                         <form id="form_character_search_form" action="javascript:void(null);">
                             <div id="rm_button_create" title="Create New Character" class="menu_button fa-solid fa-user-plus "></div>
                             <div id="character_import_button" title="Import Character from File" class="menu_button fa-solid fa-file-arrow-up faSmallFontSquareFix"></div>
+                            <div id="external_import_button" title="Import content from external URL" class="menu_button fa-solid fa-cloud-arrow-down faSmallFontSquareFix"></div>
                             <div id="rm_button_group_chats" title="Create New Chat Group" class="menu_button fa-solid fa-users-gear "></div>
                             <input id="character_search_bar" class="text_pole width100p" type="search" placeholder="Search..." maxlength="50" />
                             <select id="character_sort_order" title="Characters sorting order">
diff --git a/public/script.js b/public/script.js
index 7f621320a..6a612bf0b 100644
--- a/public/script.js
+++ b/public/script.js
@@ -32,6 +32,7 @@ import {
     importEmbeddedWorldInfo,
     checkEmbeddedWorld,
     setWorldInfoButtonClass,
+    importWorldInfo,
 } from "./scripts/world-info.js";
 
 import {
@@ -7924,6 +7925,55 @@ $(document).ready(function () {
         restoreCaretPosition($(this).get(0), caretPosition);
     });
 
+    $('#external_import_button').on('click', async () => {
+        const html = `<h3>Enter the URL of the content to import</h3>
+        Supported sources:<br>
+        <ul class="justifyLeft">
+            <li>Chub characters (direct link or id)<br>Example: <tt>lorebooks/bartleby/example-lorebook</tt></li>
+            <li>Chub lorebooks (direct link or id)<br>Example: <tt>Anonymous/example-character</tt></li>
+            <li>More coming soon...</li>
+        <ul>`
+        const input = await callPopup(html, 'input');
+
+        if (!input) {
+            console.debug('Custom content import cancelled');
+            return;
+        }
+
+        const url = input.trim();
+        console.debug('Custom content import started', url);
+
+        const request = await fetch('/import_custom', {
+            method: 'POST',
+            headers: getRequestHeaders(),
+            body: JSON.stringify({ url }),
+        });
+
+        if (!request.ok) {
+            toastr.info(request.statusText, 'Custom content import failed');
+            console.error('Custom content import failed', request.status, request.statusText);
+            return;
+        }
+
+        const data = await request.blob();
+        const customContentType = request.headers.get('X-Custom-Content-Type');
+        const fileName = request.headers.get('Content-Disposition').split('filename=')[1].replace(/"/g, '');
+        const file = new File([data], fileName, { type: data.type });
+
+        switch (customContentType) {
+            case 'character':
+                processDroppedFiles([file]);
+                break;
+            case 'lorebook':
+                await importWorldInfo(file);
+                break;
+            default:
+                toastr.warn('Unknown content type');
+                console.error('Unknown content type', customContentType);
+                break;
+        }
+    });
+
     const $dropzone = $(document.body);
 
     $dropzone.on('dragover', (event) => {
diff --git a/public/scripts/world-info.js b/public/scripts/world-info.js
index 8942040df..81df9197f 100644
--- a/public/scripts/world-info.js
+++ b/public/scripts/world-info.js
@@ -1183,6 +1183,76 @@ function onWorldInfoChange(_, text) {
     saveSettingsDebounced();
 }
 
+export async function importWorldInfo(file) {
+    if (!file) {
+        return;
+    }
+
+    const formData = new FormData();
+    formData.append('avatar', file);
+
+    try {
+        let jsonData;
+
+        if (file.name.endsWith('.png')) {
+            const buffer = new Uint8Array(await getFileBuffer(file));
+            jsonData = extractDataFromPng(buffer, 'naidata');
+        } else {
+            // File should be a JSON file
+            jsonData = await parseJsonFile(file);
+        }
+
+        if (jsonData === undefined || jsonData === null) {
+            toastr.error(`File is not valid: ${file.name}`);
+            return;
+        }
+
+        // Convert Novel Lorebook
+        if (jsonData.lorebookVersion !== undefined) {
+            console.log('Converting Novel Lorebook');
+            formData.append('convertedData', JSON.stringify(convertNovelLorebook(jsonData)));
+        }
+
+        // Convert Agnai Memory Book
+        if (jsonData.kind === 'memory') {
+            console.log('Converting Agnai Memory Book');
+            formData.append('convertedData', JSON.stringify(convertAgnaiMemoryBook(jsonData)));
+        }
+
+        // Convert Risu Lorebook
+        if (jsonData.type === 'risu') {
+            console.log('Converting Risu Lorebook');
+            formData.append('convertedData', JSON.stringify(convertRisuLorebook(jsonData)));
+        }
+    } catch (error) {
+        toastr.error(`Error parsing file: ${error}`);
+        return;
+    }
+
+    jQuery.ajax({
+        type: "POST",
+        url: "/importworldinfo",
+        data: formData,
+        beforeSend: () => { },
+        cache: false,
+        contentType: false,
+        processData: false,
+        success: async function (data) {
+            if (data.name) {
+                await updateWorldInfoList();
+
+                const newIndex = world_names.indexOf(data.name);
+                if (newIndex >= 0) {
+                    $("#world_editor_select").val(newIndex).trigger('change');
+                }
+
+                toastr.info(`World Info "${data.name}" imported successfully!`);
+            }
+        },
+        error: (jqXHR, exception) => { },
+    });
+}
+
 jQuery(() => {
 
     $(document).ready(function () {
@@ -1219,72 +1289,7 @@ jQuery(() => {
     $("#world_import_file").on("change", async function (e) {
         const file = e.target.files[0];
 
-        if (!file) {
-            return;
-        }
-
-        const formData = new FormData($("#form_world_import").get(0));
-
-        try {
-            let jsonData;
-
-            if (file.name.endsWith('.png')) {
-                const buffer = new Uint8Array(await getFileBuffer(file));
-                jsonData = extractDataFromPng(buffer, 'naidata');
-            } else {
-                // File should be a JSON file
-                jsonData = await parseJsonFile(file);
-            }
-
-            if (jsonData === undefined || jsonData === null) {
-                toastr.error(`File is not valid: ${file.name}`);
-                return;
-            }
-
-            // Convert Novel Lorebook
-            if (jsonData.lorebookVersion !== undefined) {
-                console.log('Converting Novel Lorebook');
-                formData.append('convertedData', JSON.stringify(convertNovelLorebook(jsonData)));
-            }
-
-            // Convert Agnai Memory Book
-            if (jsonData.kind === 'memory') {
-                console.log('Converting Agnai Memory Book');
-                formData.append('convertedData', JSON.stringify(convertAgnaiMemoryBook(jsonData)));
-            }
-
-            // Convert Risu Lorebook
-            if (jsonData.type === 'risu') {
-                console.log('Converting Risu Lorebook');
-                formData.append('convertedData', JSON.stringify(convertRisuLorebook(jsonData)));
-            }
-        } catch (error) {
-            toastr.error(`Error parsing file: ${error}`);
-            return;
-        }
-
-        jQuery.ajax({
-            type: "POST",
-            url: "/importworldinfo",
-            data: formData,
-            beforeSend: () => { },
-            cache: false,
-            contentType: false,
-            processData: false,
-            success: async function (data) {
-                if (data.name) {
-                    await updateWorldInfoList();
-
-                    const newIndex = world_names.indexOf(data.name);
-                    if (newIndex >= 0) {
-                        $("#world_editor_select").val(newIndex).trigger('change');
-                    }
-
-                    toastr.info(`World Info "${data.name}" imported successfully!`);
-                }
-            },
-            error: (jqXHR, exception) => { },
-        });
+        await importWorldInfo(file);
 
         // Will allow to select the same file twice in a row
         $("#form_world_import").trigger("reset");
diff --git a/public/style.css b/public/style.css
index a78fb5b08..5dd93fe1a 100644
--- a/public/style.css
+++ b/public/style.css
@@ -1311,13 +1311,14 @@ select option:not(:checked) {
     margin: 0;
     flex: 1;
     border-radius: 7px;
+    height: auto;
 }
 
 #character_search_bar {
     margin: 0;
     flex: 1;
     /* padding-left: 0.75em; */
-    height: fit-content;
+    height: auto;
 }
 
 input[type=search]::-webkit-search-cancel-button {
diff --git a/server.js b/server.js
index 1be7e7dba..97ca1ff78 100644
--- a/server.js
+++ b/server.js
@@ -1085,19 +1085,19 @@ async function charaWrite(img_url, data, target_img, response = undefined, mes =
 async function tryReadImage(img_url, crop) {
     try {
         let rawImg = await jimp.read(img_url);
-		let final_width = rawImg.bitmap.width, final_height = rawImg.bitmap.height
+        let final_width = rawImg.bitmap.width, final_height = rawImg.bitmap.height
 
         // Apply crop if defined
         if (typeof crop == 'object' && [crop.x, crop.y, crop.width, crop.height].every(x => typeof x === 'number')) {
             rawImg = rawImg.crop(crop.x, crop.y, crop.width, crop.height);
-			// Apply standard resize if requested
-			if (crop.want_resize) {
-				final_width = AVATAR_WIDTH
-				final_height = AVATAR_HEIGHT
-			}
+            // Apply standard resize if requested
+            if (crop.want_resize) {
+                final_width = AVATAR_WIDTH
+                final_height = AVATAR_HEIGHT
+            }
         }
 
-		const image = await rawImg.cover(final_width, final_height).getBufferAsync(jimp.MIME_PNG);
+        const image = await rawImg.cover(final_width, final_height).getBufferAsync(jimp.MIME_PNG);
         return image;
     }
     // If it's an unsupported type of image (APNG) - just read the file as buffer
@@ -3857,6 +3857,116 @@ app.post('/upload_sprite', urlencodedParser, async (request, response) => {
     }
 });
 
+app.post('/import_custom', jsonParser, async (request, response) => {
+    if (!request.body.url) {
+        return response.sendStatus(400);
+    }
+
+    try {
+        const url = request.body.url;
+        let result;
+
+        const chubParsed = parseChubUrl(url);
+
+        if (chubParsed?.type === 'character') {
+            console.log('Downloading chub character:', chubParsed.id);
+            result = await downloadChubCharacter(chubParsed.id);
+        }
+        else if (chubParsed?.type === 'lorebook') {
+            console.log('Downloading chub lorebook:', chubParsed.id);
+            result = await downloadChubLorebook(chubParsed.id);
+        }
+        else {
+            return response.sendStatus(404);
+        }
+
+        response.set('Content-Type', result.fileType);
+        response.set('Content-Disposition', `attachment; filename="${result.fileName}"`);
+        response.set('X-Custom-Content-Type', chubParsed?.type);
+        return response.send(result.buffer);
+    } catch (error) {
+        console.log('Importing custom content failed', error);
+        return response.sendStatus(500);
+    }
+});
+
+async function downloadChubLorebook(id) {
+    const fetch = require('node-fetch').default;
+
+    const result = await fetch('https://api.chub.ai/api/lorebooks/download', {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify({
+            "fullPath": id,
+            "format": "SILLYTAVERN",
+        }),
+    });
+
+    if (!result.ok) {
+        throw new Error('Failed to download lorebook');
+    }
+
+    const name = id.split('/').pop();
+    const buffer = await result.buffer();
+    const fileName = `${sanitize(name)}.json`;
+    const fileType = result.headers.get('content-type');
+
+    return { buffer, fileName, fileType };
+}
+
+async function downloadChubCharacter(id) {
+    const fetch = require('node-fetch').default;
+
+    const result = await fetch('https://api.chub.ai/api/characters/download', {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify({
+            "format": "tavern",
+            "fullPath": id,
+        })
+    });
+
+    if (!result.ok) {
+        throw new Error('Failed to download character');
+    }
+
+    const buffer = await result.buffer();
+    const fileName = result.headers.get('content-disposition').split('filename=')[1];
+    const fileType = result.headers.get('content-type');
+
+    return { buffer, fileName, fileType };
+}
+
+function parseChubUrl(str) {
+    const splitStr = str.split('/');
+    const length = splitStr.length;
+
+    if (length < 2) {
+        return null;
+    }
+
+    const domainIndex = splitStr.indexOf('chub.ai');
+    const lastTwo = domainIndex !== -1 ? splitStr.slice(domainIndex + 1) : splitStr;
+
+    const firstPart = lastTwo[0].toLowerCase();
+
+    if (firstPart === 'characters' || firstPart === 'lorebooks') {
+        const type = firstPart === 'characters' ? 'character' : 'lorebook';
+        const id = type === 'character' ? lastTwo.slice(1).join('/') : lastTwo.join('/');
+        return {
+            id: id,
+            type: type
+        };
+    } else if (length === 2) {
+        return {
+            id: lastTwo.join('/'),
+            type: 'character'
+        };
+    }
+
+    return null;
+}
+
 function importRisuSprites(data) {
     try {
         const name = data?.data?.name;