From 4e553cf6abc57d45c707105458f146646d03b2f5 Mon Sep 17 00:00:00 2001 From: kingbri Date: Thu, 31 Aug 2023 00:09:31 -0400 Subject: [PATCH 1/2] Server: Allow appending of additional headers for local backends This is a useful feature for those who want to utilize APIs with proxy middleware for adding extra features or security. For cloud API safety and abiding by rate limits, this feature only applies to local backends such as ooba or kobold. Signed-off-by: kingbri --- default/config.conf | 4 ++++ server.js | 29 ++++++++++++++++++++++++++--- 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/default/config.conf b/default/config.conf index 92c99ba02..8bec948a5 100644 --- a/default/config.conf +++ b/default/config.conf @@ -15,6 +15,9 @@ const skipContentCheck = false; // If true, no new default content will be deliv // Change this setting only on "trusted networks". Do not change this value unless you are aware of the issues that can arise from changing this setting and configuring a insecure setting. const securityOverride = false; +// Request overrides for additional headers +const requestOverrides = []; + module.exports = { port, whitelist, @@ -28,4 +31,5 @@ module.exports = { allowKeysExposure, securityOverride, skipContentCheck, + requestOverrides, }; diff --git a/server.js b/server.js index fa312caa0..04c18797e 100644 --- a/server.js +++ b/server.js @@ -190,7 +190,14 @@ function get_mancer_headers() { return api_key_mancer ? { "X-API-KEY": api_key_mancer } : {}; } - +function getOverrideHeaders(urlHost) { + const overrideHeaders = config.requestOverrides?.find((e) => e.hosts?.includes(urlHost))?.headers; + if (overrideHeaders && urlHost) { + return overrideHeaders; + } else { + return {}; + } +} //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. @@ -543,7 +550,10 @@ app.post("/generate", jsonParser, async function (request, response_generate) { console.log(this_settings); const args = { body: JSON.stringify(this_settings), - headers: { "Content-Type": "application/json" }, + headers: Object.assign( + { "Content-Type": "application/json" }, + getOverrideHeaders((new URL(api_server))?.host) + ), signal: controller.signal, }; @@ -633,11 +643,19 @@ app.post("/generate_textgenerationwebui", jsonParser, async function (request, r }); async function* readWebsocket() { + const streamingUrlString = request.header('X-Streaming-URL').replace("localhost", "127.0.0.1"); + const streamingUrl = new URL(streamingUrlString); const websocket = new WebSocket(streamingUrl); websocket.on('open', async function () { console.log('WebSocket opened'); - const combined_args = Object.assign(request.body.use_mancer ? get_mancer_headers() : {}, request.body); + const combined_args = Object.assign( + {}, + request.body.use_mancer ? get_mancer_headers() : getOverrideHeaders(streamingUrl?.host), + request.body + ); + console.log(combined_args); + websocket.send(JSON.stringify(combined_args)); }); @@ -719,6 +737,8 @@ app.post("/generate_textgenerationwebui", jsonParser, async function (request, r if (request.body.use_mancer) { args.headers = Object.assign(args.headers, get_mancer_headers()); + } else { + args.headers = Object.assign(args.headers, getOverrideHeaders((new URL(api_server))?.host)); } try { @@ -786,6 +806,7 @@ app.post("/getchat", jsonParser, function (request, response) { } }); +// Only called for kobold and ooba/mancer app.post("/getstatus", jsonParser, async function (request, response) { if (!request.body) return response.sendStatus(400); api_server = request.body.api_server; @@ -800,6 +821,8 @@ app.post("/getstatus", jsonParser, async function (request, response) { if (main_api == 'textgenerationwebui' && request.body.use_mancer) { args.headers = Object.assign(args.headers, get_mancer_headers()); + } else { + args.headers = Object.assign(args.headers, getOverrideHeaders((new URL(api_server))?.host)); } const url = api_server + "/v1/model"; From fce57b41dd919ace27fa33d3f75efa93149d1440 Mon Sep 17 00:00:00 2001 From: kingbri Date: Thu, 31 Aug 2023 00:43:45 -0400 Subject: [PATCH 2/2] Config: Indent by 4 spaces 2 spaces is too small for a config file like this. Signed-off-by: kingbri --- default/config.conf | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/default/config.conf b/default/config.conf index 8bec948a5..962c6ce5b 100644 --- a/default/config.conf +++ b/default/config.conf @@ -19,17 +19,17 @@ const securityOverride = false; const requestOverrides = []; module.exports = { - port, - whitelist, - whitelistMode, - basicAuthMode, - basicAuthUser, - autorun, - enableExtensions, - listen, - disableThumbnails, - allowKeysExposure, - securityOverride, - skipContentCheck, - requestOverrides, + port, + whitelist, + whitelistMode, + basicAuthMode, + basicAuthUser, + autorun, + enableExtensions, + listen, + disableThumbnails, + allowKeysExposure, + securityOverride, + skipContentCheck, + requestOverrides, };