From 96caddfd714d7b96fe4c6e38a2e9f14d7aef4622 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Sun, 19 Nov 2023 17:14:53 +0200
Subject: [PATCH] Add koboldcpp as Text Completion source

---
 public/index.html                  | 19 ++++++++++++++++---
 public/script.js                   |  3 +++
 public/scripts/textgen-settings.js | 10 +++++++++-
 public/scripts/tokenizers.js       |  5 +++--
 server.js                          | 19 ++++++++++++++-----
 5 files changed, 45 insertions(+), 11 deletions(-)

diff --git a/public/index.html b/public/index.html
index c00c0da2f..c1e24f2c5 100644
--- a/public/index.html
+++ b/public/index.html
@@ -1467,10 +1467,10 @@
                 <div class="flex-container flexFlowColumn">
                     <div id="main-API-selector-block">
                         <select id="main_api">
-                            <option value="kobold"><span data-i18n="KoboldAI">KoboldAI</span></option>
+                            <option value="kobold"><span data-i18n="KoboldAI">KoboldAI Classic</span></option>
                             <option value="koboldhorde"><span data-i18n="KoboldAI Horde">KoboldAI Horde</span></option>
                             <option value="novel"><span data-i18n="NovelAI">NovelAI</span></option>
-                            <option value="textgenerationwebui"><span data-i18n="Text Completion">Text Completion (ooba, Mancer, Aphrodite)</span></option>
+                            <option value="textgenerationwebui"><span data-i18n="Text Completion">Text Completion (ooba, Mancer, Aphrodite, KoboldCpp)</span></option>
                             <option value="openai"><span data-i18n="Chat Completion (OpenAI, Claude, Window/OpenRouter, Scale, AI21)">Chat Completion (OpenAI, Claude, Window, OpenRouter, Scale, AI21, PaLM)</span></option>
                         </select>
                     </div>
@@ -1600,6 +1600,7 @@
                                     <option value="mancer">Mancer</option>
                                     <option value="aphrodite">Aphrodite</option>
                                     <option value="tabby">TabbyAPI</option>
+                                    <option value="koboldcpp">KoboldCpp</option>
                                 </select>
                             </div>
                             <div data-tg-type="mancer" class="flex-container flexFlowColumn">
@@ -1681,8 +1682,20 @@
                                     <input id="tabby_api_url_text" class="text_pole wide100p" maxlength="500" value="" autocomplete="off" data-server-history="tabby">
                                 </div>
                             </div>
+                            <div data-tg-type="koboldcpp">
+                                <div class="flex-container flexFlowColumn">
+                                    <a href="https://github.com/LostRuins/koboldcpp" target="_blank">
+                                        LostRuins/koboldcpp
+                                    </a>
+                                </div>
+                                <div class="flex1">
+                                    <h4 data-i18n="API url">API URL</h4>
+                                    <small data-i18n="Example: http://127.0.0.1:5001">Example: http://127.0.0.1:5001</small>
+                                    <input id="koboldcpp_api_url_text" class="text_pole wide100p" maxlength="500" value="" autocomplete="off" data-server-history="koboldcpp">
+                                </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,tabby">Connect</div>
+                                <div id="api_button_textgenerationwebui" class="api_button menu_button" type="submit" data-i18n="Connect" data-server-connect="ooba_blocking,aphrodite,tabby,koboldcpp">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 e7aa43689..4c22c2ccf 100644
--- a/public/script.js
+++ b/public/script.js
@@ -24,6 +24,7 @@ import {
     textgenerationwebui_banned_in_macros,
     isOoba,
     MANCER_SERVER,
+    isKoboldCpp,
 } from "./scripts/textgen-settings.js";
 
 import {
@@ -886,6 +887,7 @@ async function getStatus() {
                 use_aphrodite: main_api == "textgenerationwebui" ? isAphrodite() : false,
                 use_ooba: main_api == "textgenerationwebui" ? isOoba() : false,
                 use_tabby: main_api == "textgenerationwebui" ? isTabby() : false,
+                use_koboldcpp: main_api == "textgenerationwebui" ? isKoboldCpp() : false,
                 legacy_api: main_api == "textgenerationwebui" ? textgenerationwebui_settings.legacy_api && !isMancer() : false,
             }),
             signal: abortStatusCheck.signal,
@@ -5505,6 +5507,7 @@ async function getSettings() {
         $("#textgenerationwebui_api_url_text").val(api_server_textgenerationwebui);
         $("#aphrodite_api_url_text").val(api_server_textgenerationwebui);
         $("#tabby_api_url_text").val(api_server_textgenerationwebui);
+        $('#koboldcpp_api_url_text').val(api_server_textgenerationwebui);
 
         selected_button = settings.selected_button;
 
diff --git a/public/scripts/textgen-settings.js b/public/scripts/textgen-settings.js
index 7891ed5f6..6cc76db50 100644
--- a/public/scripts/textgen-settings.js
+++ b/public/scripts/textgen-settings.js
@@ -28,7 +28,8 @@ export const textgen_types = {
     OOBA: 'ooba',
     MANCER: 'mancer',
     APHRODITE: 'aphrodite',
-    TABBY: 'tabby'
+    TABBY: 'tabby',
+    KOBOLDCPP: 'koboldcpp',
 };
 
 // Maybe let it be configurable in the future?
@@ -292,6 +293,10 @@ export function isOoba() {
     return textgenerationwebui_settings.type === textgen_types.OOBA;
 }
 
+export function isKoboldCpp() {
+    return textgenerationwebui_settings.type === textgen_types.KOBOLDCPP;
+}
+
 export function getTextGenUrlSourceId() {
     switch (textgenerationwebui_settings.type) {
         case textgen_types.OOBA:
@@ -300,6 +305,8 @@ export function getTextGenUrlSourceId() {
             return "#aphrodite_api_url_text";
         case textgen_types.TABBY:
             return "#tabby_api_url_text";
+        case textgen_types.KOBOLDCPP:
+            return "#koboldcpp_api_url_text";
     }
 }
 
@@ -559,6 +566,7 @@ export function getTextGenGenerationData(finalPrompt, this_amount_gen, isImperso
         'use_mancer': isMancer(),
         'use_aphrodite': isAphrodite(),
         'use_tabby': isTabby(),
+        'use_koboldcpp': isKoboldCpp(),
         'use_ooba': isOoba(),
         'api_server': isMancer() ? MANCER_SERVER : api_server_textgenerationwebui,
         'legacy_api': textgenerationwebui_settings.legacy_api && !isMancer(),
diff --git a/public/scripts/tokenizers.js b/public/scripts/tokenizers.js
index 61d4d0041..cbff17f6d 100644
--- a/public/scripts/tokenizers.js
+++ b/public/scripts/tokenizers.js
@@ -4,7 +4,7 @@ import { chat_completion_sources, model_list, oai_settings } from "./openai.js";
 import { groups, selected_group } from "./group-chats.js";
 import { getStringHash } from "./utils.js";
 import { kai_flags } from "./kai-settings.js";
-import { isMancer, isTabby, textgenerationwebui_settings } from "./textgen-settings.js";
+import { isKoboldCpp, isMancer, isTabby, textgenerationwebui_settings } from "./textgen-settings.js";
 
 export const CHARACTERS_PER_TOKEN_RATIO = 3.35;
 const TOKENIZER_WARNING_KEY = 'tokenizationWarningShown';
@@ -369,7 +369,8 @@ function getRemoteTokenizationParams(str) {
         api: main_api,
         url: getAPIServerUrl(),
         legacy_api: main_api === 'textgenerationwebui' && textgenerationwebui_settings.legacy_api && !isMancer(),
-        use_tabby: isTabby()
+        use_tabby: main_api === 'textgenerationwebui' && isTabby(),
+        use_koboldcpp: main_api === 'textgenerationwebui' && isKoboldCpp(),
     };
 }
 
diff --git a/server.js b/server.js
index 914a7ade0..400045f16 100644
--- a/server.js
+++ b/server.js
@@ -539,6 +539,9 @@ app.post("/api/textgenerationwebui/status", jsonParser, async function (request,
         else if (request.body.use_tabby) {
             url += "/v1/model/list"
         }
+        else if (request.body.use_koboldcpp) {
+            url += "/v1/models";
+        }
 
         const modelsReply = await fetch(url, args);
 
@@ -629,7 +632,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 || request.body.use_tabby) {
+        else if (request.body.use_aphrodite || request.body.use_ooba || request.body.use_tabby || request.body.use_koboldcpp) {
             url += "/v1/completions";
         }
         else if (request.body.use_mancer) {
@@ -3445,10 +3448,16 @@ app.post("/tokenize_via_api", jsonParser, async function (request, response) {
             if (legacyApi) {
                 url += '/v1/token-count';
                 args.body = JSON.stringify({ "prompt": text });
-            } else if (request.body.use_tabby) {
+            }
+            else if (request.body.use_tabby) {
                 url += '/v1/token/encode';
                 args.body = JSON.stringify({ "text": text });
-            } else {
+            }
+            else if (request.body.use_koboldcpp) {
+                url += '/api/extra/tokencount';
+                args.body = JSON.stringify({ "prompt": text });
+            }
+            else {
                 url += '/v1/internal/encode';
                 args.body = JSON.stringify({ "text": text });
             }
@@ -3461,8 +3470,8 @@ app.post("/tokenize_via_api", jsonParser, async function (request, response) {
             }
 
             const data = await result.json();
-            const count = legacyApi ? data?.results[0]?.tokens : data?.length;
-            const ids = legacyApi ? [] : data?.tokens;
+            const count = legacyApi ? data?.results[0]?.tokens : (data?.length ?? data?.value);
+            const ids = legacyApi ? [] : (data?.tokens ?? []);
 
             return response.send({ count, ids });
         }