From 2fdec7eb03d1ea9240b63cd437f40a8e41cb7e46 Mon Sep 17 00:00:00 2001 From: 50h100a Date: Wed, 2 Aug 2023 23:25:24 -0400 Subject: [PATCH] Added authentication variant to WebUI API. --- public/script.js | 41 +++++++++++++++++++++++++++++++++-------- server.js | 29 ++++++++++++++++++++++++++--- 2 files changed, 59 insertions(+), 11 deletions(-) diff --git a/public/script.js b/public/script.js index 157d79a99..09c090477 100644 --- a/public/script.js +++ b/public/script.js @@ -688,6 +688,7 @@ let is_get_status = false; let is_get_status_novel = false; let is_api_button_press = false; let is_api_button_press_novel = false; +let api_use_mancer_webui = false; let is_send_press = false; //Send generation let add_mes_without_animation = false; @@ -810,9 +811,9 @@ async function getStatus() { type: "POST", // url: "/getstatus", // data: JSON.stringify({ - api_server: - main_api == "kobold" ? api_server : api_server_textgenerationwebui, + api_server: main_api == "kobold" ? api_server : api_server_textgenerationwebui, main_api: main_api, + use_mancer: main_api == "textgenerationwebui" ? api_use_mancer_webui : false, }), beforeSend: function () { }, cache: false, @@ -2672,6 +2673,7 @@ async function Generate(type, { automatic_trigger, force_name2, resolve, reject, } else if (main_api == 'textgenerationwebui') { generate_data = getTextGenGenerationData(finalPromt, this_amount_gen, isImpersonate); + generate_data.use_mancer = api_use_mancer_webui; } else if (main_api == 'novel') { const this_settings = novelai_settings[novelai_setting_names[nai_settings.preset_settings_novel]]; @@ -2933,6 +2935,14 @@ async function Generate(type, { automatic_trigger, force_name2, resolve, reject, activateSendButtons(); //console.log('runGenerate calling showSwipeBtns'); showSwipeButtons(); + + if (main_api == 'textgenerationwebui' && api_use_mancer_webui) { + const errorText = `

Inferencer endpoint is unhappy!

+ Check you have credits available on your + Mancer account.
+ If you have sufficient credits, please try again later.`; + callPopup(errorText, 'text'); + } } console.debug('/savechat called by /Generate'); @@ -5015,11 +5025,13 @@ async function getSettings(type) { setWorldInfoSettings(settings, data); - api_server_textgenerationwebui = - settings.api_server_textgenerationwebui; + api_server_textgenerationwebui = settings.api_server_textgenerationwebui; $("#textgenerationwebui_api_url_text").val( api_server_textgenerationwebui ); + api_use_mancer_webui = settings.api_use_mancer_webui + $('#use-mancer-api-checkbox').prop("checked", api_use_mancer_webui); + $('#use-mancer-api-checkbox').trigger("change"); selected_button = settings.selected_button; @@ -5051,6 +5063,7 @@ async function saveSettings(type) { username: name1, api_server: api_server, api_server_textgenerationwebui: api_server_textgenerationwebui, + api_use_mancer_webui: api_use_mancer_webui, preset_settings: preset_settings, user_avatar: user_avatar, amount_gen: amount_gen, @@ -7611,16 +7624,28 @@ $(document).ready(function () { } }); + $("#use-mancer-api-checkbox").on("change", function (e) { + const enabled = $("#use-mancer-api-checkbox").prop("checked"); + $("#mancer-api-ui").css("display", enabled?"block":"none"); + api_use_mancer_webui = enabled; + saveSettingsDebounced(); + getStatus(); + }); + $("#api_button_textgenerationwebui").click(function (e) { e.stopPropagation(); if ($("#textgenerationwebui_api_url_text").val() != "") { - let value = formatKoboldUrl($("#textgenerationwebui_api_url_text").val().trim()); - - if (!value) { - callPopup('Please enter a valid URL.', 'text'); + let value = $("#textgenerationwebui_api_url_text").val().trim(); + if (!value || !value.endsWith('/api')) { + callPopup('Please enter a valid URL.
WebUI URLs should end with /api', 'text'); return; } + const mancer_key = $("#api_key_mancer").val().trim(); + if (mancer_key.length) { + writeSecret(SECRET_KEYS.MANCER, mancer_key); + } + $("#textgenerationwebui_api_url_text").val(value); $("#api_loading_textgenerationwebui").css("display", "inline-block"); $("#api_button_textgenerationwebui").css("display", "none"); diff --git a/server.js b/server.js index 5497e3a5d..fc1913e26 100644 --- a/server.js +++ b/server.js @@ -146,6 +146,14 @@ let response_getstatus; let first_run = true; + +function get_mancer_headers() { + const api_key_mancer = readSecret(SECRET_KEYS.MANCER); + return api_key_mancer ? { "X-API-KEY": api_key_mancer} : {}; +} + + + //RossAscends: Added function to format dates used in files and chat timestamps to a humanized format. //Mostly I wanted this to be for file names, but couldn't figure out exactly where the filename save code was as everything seemed to be connected. //During testing, this performs the same as previous date.now() structure. @@ -621,17 +629,21 @@ app.post("/generate_textgenerationwebui", jsonParser, async function (request, r else { const args = { body: JSON.stringify(request.body), - headers: { "Content-Type": "application/json" }, + headers: { "Content-Type": "application/json"}, signal: controller.signal, }; + if (request.body.use_mancer) { + args.headers = Object.assign(args.headers, get_mancer_headers()); + } + try { const data = await postAsync(api_server + "/v1/generate", args); console.log(data); return response_generate.send(data); } catch (error) { - console.log(error); - return response_generate.send({ error: true }); + console.log("Error: %o", error); + return response_generate.send({ error: true}); } } }); @@ -685,6 +697,7 @@ app.post("/getchat", jsonParser, function (request, response) { } }); + app.post("/getstatus", jsonParser, async function (request, response_getstatus = response) { if (!request.body) return response_getstatus.sendStatus(400); api_server = request.body.api_server; @@ -695,6 +708,11 @@ app.post("/getstatus", jsonParser, async function (request, response_getstatus = var args = { headers: { "Content-Type": "application/json" } }; + + if (main_api == 'textgenerationwebui' && request.body.use_mancer) { + args.headers = Object.assign(args.headers, get_mancer_headers()); + } + var url = api_server + "/v1/model"; let version = ''; let koboldVersion = {}; @@ -3403,6 +3421,10 @@ app.post("/tokenize_via_api", jsonParser, async function (request, response) { body: JSON.stringify({ "prompt": text }), headers: { "Content-Type": "application/json" } }; + + if (main_api == 'textgenerationwebui' && request.body.use_mancer) { + args.headers = Object.assign(args.headers, get_mancer_headers()); + } const data = await postAsync(api_server + "/v1/token-count", args); console.log(data); @@ -3609,6 +3631,7 @@ const SECRETS_FILE = './secrets.json'; const SETTINGS_FILE = './public/settings.json'; const SECRET_KEYS = { HORDE: 'api_key_horde', + MANCER: 'api_key_mancer', OPENAI: 'api_key_openai', NOVEL: 'api_key_novel', CLAUDE: 'api_key_claude',