From f31b996cb5fdc7b6228cb0f72c5ecab02e0d6a3e Mon Sep 17 00:00:00 2001
From: kingbri <bdashore3@proton.me>
Date: Fri, 17 Nov 2023 00:32:49 -0500
Subject: [PATCH] Text Generation: Add TabbyAPI support

TabbyAPI is an exllamav2 only API server that aims to provide a simple
experience for loading and chatting with exl2 models.

SillyTavern currently doesn't have the ability to load and unload models,
so only add the OAI compatible completion endpoints.

The repository can be found here:
https://github.com/theroyallab/tabbyAPI

Signed-off-by: kingbri <bdashore3@proton.me>
---
 public/index.html                  | 24 +++++++++++++++++++++++-
 public/script.js                   |  8 ++++++++
 public/scripts/secrets.js          |  2 ++
 public/scripts/textgen-settings.js |  8 ++++++++
 server.js                          | 27 +++++++++++++++++++++------
 src/secrets.js                     |  1 +
 6 files changed, 63 insertions(+), 7 deletions(-)

diff --git a/public/index.html b/public/index.html
index eb136b760..c4ed5bdc0 100644
--- a/public/index.html
+++ b/public/index.html
@@ -1599,6 +1599,7 @@
                                     <option value="ooba">Default (oobabooga)</option>
                                     <option value="mancer">Mancer</option>
                                     <option value="aphrodite">Aphrodite</option>
+                                    <option value="tabby">TabbyAPI</option>
                                 </select>
                             </div>
                             <div data-tg-type="mancer" class="flex-container flexFlowColumn">
@@ -1659,8 +1660,29 @@
                                     <input id="aphrodite_api_url_text" class="text_pole wide100p" maxlength="500" value="" autocomplete="off" data-server-history="aphrodite">
                                 </div>
                             </div>
+                            <div data-tg-type="tabby">
+                                <div class="flex-container flexFlowColumn">
+                                    <a href="https://github.com/theroyallab/tabbyAPI" target="_blank">
+                                        theroyallab/tabbyAPI
+                                    </a>
+                                </div>
+                                <h4 data-i18n="Tabby API key">Tabby API key</h4>
+                                <div class="flex-container">
+                                    <input id="api_key_tabby" name="api_key_tabby" class="text_pole flex1 wide100p" maxlength="500" size="35" type="text" autocomplete="off">
+                                    <div title="Clear your API key" data-i18n="[title]Clear your API key" class="menu_button fa-solid fa-circle-xmark clear-api-key" data-key="api_key_tabby">
+                                    </div>
+                                </div>
+                                <div data-for="api_key_tabby" class="neutral_warning" data-i18n="For privacy reasons, your API key will be hidden after you reload the page.">
+                                    For privacy reasons, your API key will be hidden after you reload the page.
+                                </div>
+                                <div class="flex1">
+                                    <h4 data-i18n="API url">API URL</h4>
+                                    <small data-i18n="Example: http://127.0.0.1:5000">Example: http://127.0.0.1:5000</small>
+                                    <input id="tabby_api_url_text" class="text_pole wide100p" maxlength="500" value="" autocomplete="off" data-server-history="tabby">
+                                </div>
+                            </div>
                             <div class="flex-container">
-                                <div id="api_button_textgenerationwebui" class="api_button menu_button" type="submit" data-i18n="Connect" data-server-connect="ooba_blocking,aphrodite">Connect</div>
+                                <div id="api_button_textgenerationwebui" class="api_button menu_button" type="submit" data-i18n="Connect" data-server-connect="ooba_blocking,aphrodite,tabby">Connect</div>
                                 <div class="api_loading menu_button" data-i18n="Cancel">Cancel</div>
                             </div>
                             <label class="checkbox_label margin-bot-10px" for="legacy_api_textgenerationwebui">
diff --git a/public/script.js b/public/script.js
index 9ff299d9c..79550ceaa 100644
--- a/public/script.js
+++ b/public/script.js
@@ -19,6 +19,7 @@ import {
     getTextGenUrlSourceId,
     isMancer,
     isAphrodite,
+    isTabby,
     textgen_types,
     textgenerationwebui_banned_in_macros,
     isOoba,
@@ -882,6 +883,7 @@ async function getStatus() {
                 use_mancer: main_api == "textgenerationwebui" ? isMancer() : false,
                 use_aphrodite: main_api == "textgenerationwebui" ? isAphrodite() : false,
                 use_ooba: main_api == "textgenerationwebui" ? isOoba() : false,
+                use_tabby: main_api == "textgenerationwebui" ? isTabby() : false,
                 legacy_api: main_api == "textgenerationwebui" ? textgenerationwebui_settings.legacy_api && !isMancer() : false,
             }),
             signal: abortStatusCheck.signal,
@@ -5411,6 +5413,7 @@ async function getSettings() {
         api_server_textgenerationwebui = settings.api_server_textgenerationwebui;
         $("#textgenerationwebui_api_url_text").val(api_server_textgenerationwebui);
         $("#aphrodite_api_url_text").val(api_server_textgenerationwebui);
+        $("#tabby_api_url_text").val(api_server_textgenerationwebui);
 
         selected_button = settings.selected_button;
 
@@ -8008,6 +8011,11 @@ jQuery(async function () {
             await writeSecret(SECRET_KEYS.APHRODITE, aphroditeKey);
         }
 
+        const tabbyKey = String($("#api_key_tabby").val()).trim();
+        if (tabbyKey.length) {
+            await writeSecret(SECRET_KEYS.TABBY, tabbyKey)
+        }
+
         const urlSourceId = getTextGenUrlSourceId();
 
         if (urlSourceId && $(urlSourceId).val() !== "") {
diff --git a/public/scripts/secrets.js b/public/scripts/secrets.js
index 1e0133c6e..e699c403e 100644
--- a/public/scripts/secrets.js
+++ b/public/scripts/secrets.js
@@ -4,6 +4,7 @@ export const SECRET_KEYS = {
     HORDE: 'api_key_horde',
     MANCER: 'api_key_mancer',
     APHRODITE: 'api_key_aphrodite',
+    TABBY: 'api_key_tabby',
     OPENAI: 'api_key_openai',
     NOVEL: 'api_key_novel',
     CLAUDE: 'api_key_claude',
@@ -27,6 +28,7 @@ const INPUT_MAP = {
     [SECRET_KEYS.SCALE_COOKIE]: '#scale_cookie',
     [SECRET_KEYS.PALM]: '#api_key_palm',
     [SECRET_KEYS.APHRODITE]: '#api_key_aphrodite',
+    [SECRET_KEYS.TABBY]: '#api_key_tabby'
 }
 
 async function clearSecret() {
diff --git a/public/scripts/textgen-settings.js b/public/scripts/textgen-settings.js
index b397c0a4f..7891ed5f6 100644
--- a/public/scripts/textgen-settings.js
+++ b/public/scripts/textgen-settings.js
@@ -28,6 +28,7 @@ export const textgen_types = {
     OOBA: 'ooba',
     MANCER: 'mancer',
     APHRODITE: 'aphrodite',
+    TABBY: 'tabby'
 };
 
 // Maybe let it be configurable in the future?
@@ -283,6 +284,10 @@ export function isAphrodite() {
     return textgenerationwebui_settings.type === textgen_types.APHRODITE;
 }
 
+export function isTabby() {
+    return textgenerationwebui_settings.type === textgen_types.TABBY;
+}
+
 export function isOoba() {
     return textgenerationwebui_settings.type === textgen_types.OOBA;
 }
@@ -293,6 +298,8 @@ export function getTextGenUrlSourceId() {
             return "#textgenerationwebui_api_url_text";
         case textgen_types.APHRODITE:
             return "#aphrodite_api_url_text";
+        case textgen_types.TABBY:
+            return "#tabby_api_url_text";
     }
 }
 
@@ -551,6 +558,7 @@ export function getTextGenGenerationData(finalPrompt, this_amount_gen, isImperso
         'custom_token_bans': isAphrodite() ? toIntArray(getCustomTokenBans()) : getCustomTokenBans(),
         'use_mancer': isMancer(),
         'use_aphrodite': isAphrodite(),
+        'use_tabby': isTabby(),
         'use_ooba': isOoba(),
         'api_server': isMancer() ? MANCER_SERVER : api_server_textgenerationwebui,
         'legacy_api': textgenerationwebui_settings.legacy_api && !isMancer(),
diff --git a/server.js b/server.js
index d4708854b..08ce7e641 100644
--- a/server.js
+++ b/server.js
@@ -164,6 +164,15 @@ function getAphroditeHeaders() {
     }) : {};
 }
 
+function getTabbyHeaders() {
+    const apiKey = readSecret(SECRET_KEYS.TABBY)
+
+    return apiKey ? ({
+        "x-api-key": apiKey,
+        "Authorization": `Bearer ${apiKey}`,
+    }) : {};
+}
+
 function getOverrideHeaders(urlHost) {
     const overrideHeaders = config.requestOverrides?.find((e) => e.hosts?.includes(urlHost))?.headers;
     if (overrideHeaders && urlHost) {
@@ -186,6 +195,8 @@ function setAdditionalHeaders(request, args, server) {
         headers = getMancerHeaders();
     } else if (request.body.use_aphrodite) {
         headers = getAphroditeHeaders();
+    } else if (request.body.use_tabby) {
+        headers = getTabbyHeaders();
     } else {
         headers = server ? getOverrideHeaders((new URL(server))?.host) : {};
     }
@@ -520,6 +531,9 @@ app.post("/api/textgenerationwebui/status", jsonParser, async function (request,
         else if (request.body.use_mancer) {
             url += "/oai/v1/models";
         }
+        else if (request.body.use_tabby) {
+            url += "/v1/model/list"
+        }
 
         const modelsReply = await fetch(url, args);
 
@@ -546,20 +560,21 @@ app.post("/api/textgenerationwebui/status", jsonParser, async function (request,
         // Set result to the first model ID
         result = modelIds[0] || 'Valid';
 
-        if (request.body.use_ooba) {
+        if (request.body.use_ooba || request.body.use_tabby) {
             try {
-                const modelInfoUrl = baseUrl + '/v1/internal/model/info';
+                const modelInfoPath = request.body.use_ooba ? "/v1/internal/model/info" : "/v1/model"
+                const modelInfoUrl = baseUrl + modelInfoPath;
                 const modelInfoReply = await fetch(modelInfoUrl, args);
 
                 if (modelInfoReply.ok) {
                     const modelInfo = await modelInfoReply.json();
-                    console.log('Ooba model info:', modelInfo);
+                    console.log(`${request.body.use_ooba ? "Ooba" : "Tabby"} model info: ${modelInfo}`);
 
-                    const modelName = modelInfo?.model_name;
+                    const modelName = request.body.use_ooba ? modelInfo?.model_name : modelInfo?.id;
                     result = modelName || result;
                 }
             } catch (error) {
-                console.error('Failed to get Ooba model info:', error);
+                console.error(`Failed to get ${request.body.use_ooba ? "Ooba" : "Tabby"} model info: ${error}`);
             }
         }
 
@@ -593,7 +608,7 @@ app.post("/api/textgenerationwebui/generate", jsonParser, async function (reques
         if (request.body.legacy_api) {
             url += "/v1/generate";
         }
-        else if (request.body.use_aphrodite || request.body.use_ooba) {
+        else if (request.body.use_aphrodite || request.body.use_ooba || request.body.use_tabby) {
             url += "/v1/completions";
         }
         else if (request.body.use_mancer) {
diff --git a/src/secrets.js b/src/secrets.js
index 6da764529..1018911df 100644
--- a/src/secrets.js
+++ b/src/secrets.js
@@ -8,6 +8,7 @@ const SECRET_KEYS = {
     HORDE: 'api_key_horde',
     MANCER: 'api_key_mancer',
     APHRODITE: 'api_key_aphrodite',
+    TABBY: 'api_key_tabby',
     OPENAI: 'api_key_openai',
     NOVEL: 'api_key_novel',
     CLAUDE: 'api_key_claude',