From 865256f5c089d5e53e586b2b7e232b6f5d15b14a Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Wed, 8 Nov 2023 03:38:04 +0200
Subject: [PATCH] Fix ooba tokenization via API. Fix requiring streaming URL to
 generate

---
 public/index.html            |  2 --
 public/script.js             | 36 ++++++++++++-------------
 public/scripts/tokenizers.js |  8 ++++--
 server.js                    | 52 ++++++++++++++++++++++++++----------
 4 files changed, 61 insertions(+), 37 deletions(-)

diff --git a/public/index.html b/public/index.html
index 5809c61fc..ec8b41d10 100644
--- a/public/index.html
+++ b/public/index.html
@@ -205,7 +205,6 @@
                                     <i data-newbie-hidden data-preset-manager-delete="textgenerationwebui" class="menu_button fa-solid fa-trash-can" title="Delete the preset" data-i18n="[title]Delete the preset"></i>
                                 </div>
                             </div>
-                            <hr>
                         </div>
 
                         <div data-newbie-hidden id="ai_module_block_novel" class="width100p">
@@ -342,7 +341,6 @@
                                     </div>
                                 </div>
                             </div>
-                            <hr>
                             <div id="range_block_novel">
                                 <div class="range-block">
                                     <label class="checkbox_label widthFreeExpand">
diff --git a/public/script.js b/public/script.js
index 0265d37ec..2e8b59a5e 100644
--- a/public/script.js
+++ b/public/script.js
@@ -869,15 +869,7 @@ async function getStatus() {
 
     const url = main_api == "textgenerationwebui" ? '/api/textgenerationwebui/status' : '/getstatus';
 
-    let endpoint = api_server;
-
-    if (main_api == "textgenerationwebui") {
-        endpoint = api_server_textgenerationwebui;
-    }
-
-    if (main_api == "textgenerationwebui" && isMancer()) {
-        endpoint = MANCER_SERVER;
-    }
+    let endpoint = getAPIServerUrl();
 
     if (!endpoint) {
         console.warn("No endpoint for status check");
@@ -946,6 +938,22 @@ export function resultCheckStatus() {
     stopStatusLoading();
 }
 
+export function getAPIServerUrl() {
+    if (main_api == "textgenerationwebui") {
+        if (isMancer()) {
+            return MANCER_SERVER;
+        }
+
+        return api_server_textgenerationwebui;
+    }
+
+    if (main_api == "kobold") {
+        return api_server;
+    }
+
+    return "";
+}
+
 export async function selectCharacterById(id) {
     if (characters[id] == undefined) {
         return;
@@ -2534,16 +2542,6 @@ async function Generate(type, { automatic_trigger, force_name2, resolve, reject,
         return;
     }
 
-    if (
-        main_api == 'textgenerationwebui' &&
-        textgenerationwebui_settings.streaming &&
-        textgenerationwebui_settings.type === textgen_types.OOBA &&
-        !textgenerationwebui_settings.streaming_url) {
-        toastr.error('Streaming URL is not set. Look it up in the console window when starting TextGen Web UI');
-        unblockGeneration();
-        return;
-    }
-
     if (main_api == 'kobold' && kai_settings.streaming_kobold && !kai_flags.can_use_streaming) {
         toastr.error('Streaming is enabled, but the version of Kobold used does not support token streaming.', undefined, { timeOut: 10000, preventDuplicates: true, });
         unblockGeneration();
diff --git a/public/scripts/tokenizers.js b/public/scripts/tokenizers.js
index 34b3a0822..c53312631 100644
--- a/public/scripts/tokenizers.js
+++ b/public/scripts/tokenizers.js
@@ -1,4 +1,4 @@
-import { characters, main_api, nai_settings, online_status, this_chid } from "../script.js";
+import { api_server, api_server_textgenerationwebui, characters, getAPIServerUrl, main_api, nai_settings, online_status, this_chid } from "../script.js";
 import { power_user, registerDebugFunction } from "./power-user.js";
 import { chat_completion_sources, model_list, oai_settings } from "./openai.js";
 import { groups, selected_group } from "./group-chats.js";
@@ -376,7 +376,11 @@ function countTokensRemote(endpoint, str, padding) {
         async: false,
         type: 'POST',
         url: endpoint,
-        data: JSON.stringify({ text: str }),
+        data: JSON.stringify({
+            text: str,
+            api: main_api,
+            url: getAPIServerUrl(),
+        }),
         dataType: "json",
         contentType: "application/json",
         success: function (data) {
diff --git a/server.js b/server.js
index 3fac483c2..8fcd95959 100644
--- a/server.js
+++ b/server.js
@@ -3324,27 +3324,57 @@ app.post("/tokenize_via_api", jsonParser, async function (request, response) {
         return response.sendStatus(400);
     }
     const text = request.body.text || '';
+    const api = request.body.api;
+    const baseUrl = request.body.url;
 
     try {
-        const args = {
-            body: JSON.stringify({ "prompt": text }),
-            headers: { "Content-Type": "application/json" }
-        };
+        if (api == 'textgenerationwebui') {
+            const args = {
+                method: 'POST',
+                body: JSON.stringify({ "prompt": text }),
+                headers: { "Content-Type": "application/json" }
+            };
 
-        if (main_api == 'textgenerationwebui') {
             setAdditionalHeaders(request, args, null);
 
-            const data = await postAsync(api_server + "/v1/token-count", args);
+            const url = new URL(baseUrl);
+            url.pathname = '/api/v1/token-count'
+
+            const result = await fetch(url, args);
+
+            if (!result.ok) {
+                console.log(`API returned error: ${result.status} ${result.statusText}`);
+                return response.send({ error: true });
+            }
+
+            const data = await result.json();
             return response.send({ count: data['results'][0]['tokens'] });
         }
 
-        else if (main_api == 'kobold') {
-            const data = await postAsync(api_server + "/extra/tokencount", args);
+        else if (api == 'kobold') {
+            const args = {
+                method: 'POST',
+                body: JSON.stringify({ "prompt": text }),
+                headers: { "Content-Type": "application/json" }
+            };
+
+            const url = new URL(baseUrl);
+            url.pathname = '/api/extra/tokencount';
+
+            const result = await fetch(url, args);
+
+            if (!result.ok) {
+                console.log(`API returned error: ${result.status} ${result.statusText}`);
+                return response.send({ error: true });
+            }
+
+            const data = await result.json();
             const count = data['value'];
             return response.send({ count: count });
         }
 
         else {
+            console.log('Unknown API', api);
             return response.send({ error: true });
         }
     } catch (error) {
@@ -3371,12 +3401,6 @@ async function fetchJSON(url, args = {}) {
 
     throw response;
 }
-/**
- * Convenience function for fetch requests (default POST with no timeout) returning as JSON.
- * @param {string} url
- * @param {import('node-fetch').RequestInit} args
- */
-async function postAsync(url, args) { return fetchJSON(url, { method: 'POST', timeout: 0, ...args }) }
 
 // ** END **