From 04e92efe298353df7b139a2a61e403f8e49aba23 Mon Sep 17 00:00:00 2001 From: valadaptive Date: Sat, 9 Dec 2023 19:35:49 -0500 Subject: [PATCH] Move API tokenization endpoint into /tokenizers Requires extracting some more functions out of server.js. --- server.js | 152 +----------------------------------- src/additional-headers.js | 72 +++++++++++++++++ src/endpoints/tokenizers.js | 89 +++++++++++++++++++++ 3 files changed, 162 insertions(+), 151 deletions(-) create mode 100644 src/additional-headers.js diff --git a/server.js b/server.js index 6b4d7199d..b44533d3e 100644 --- a/server.js +++ b/server.js @@ -49,6 +49,7 @@ const { delay, getVersion, getConfigValue, color, uuidv4, tryParse, clientRelati const { ensureThumbnailCache } = require('./src/endpoints/thumbnails'); const { getTokenizerModel, getTiktokenTokenizer, loadTokenizers, TEXT_COMPLETION_MODELS, getSentencepiceTokenizer, sentencepieceTokenizers } = require('./src/endpoints/tokenizers'); const { convertClaudePrompt } = require('./src/chat-completion'); +const { getOverrideHeaders, setAdditionalHeaders } = require('./src/additional-headers'); // Work around a node v20.0.0, v20.1.0, and v20.2.0 bug. The issue was fixed in v20.3.0. // https://github.com/nodejs/node/issues/47822#issuecomment-1564708870 @@ -119,70 +120,6 @@ const listen = getConfigValue('listen', false); const API_OPENAI = 'https://api.openai.com/v1'; const API_CLAUDE = 'https://api.anthropic.com/v1'; -function getMancerHeaders() { - const apiKey = readSecret(SECRET_KEYS.MANCER); - - return apiKey ? ({ - 'X-API-KEY': apiKey, - 'Authorization': `Bearer ${apiKey}`, - }) : {}; -} - -function getAphroditeHeaders() { - const apiKey = readSecret(SECRET_KEYS.APHRODITE); - - return apiKey ? ({ - 'X-API-KEY': apiKey, - 'Authorization': `Bearer ${apiKey}`, - }) : {}; -} - -function getTabbyHeaders() { - const apiKey = readSecret(SECRET_KEYS.TABBY); - - return apiKey ? ({ - 'x-api-key': apiKey, - 'Authorization': `Bearer ${apiKey}`, - }) : {}; -} - -function getOverrideHeaders(urlHost) { - const requestOverrides = getConfigValue('requestOverrides', []); - const overrideHeaders = requestOverrides?.find((e) => e.hosts?.includes(urlHost))?.headers; - if (overrideHeaders && urlHost) { - return overrideHeaders; - } else { - return {}; - } -} - -/** - * Sets additional headers for the request. - * @param {object} request Original request body - * @param {object} args New request arguments - * @param {string|null} server API server for new request - */ -function setAdditionalHeaders(request, args, server) { - let headers; - - switch (request.body.api_type) { - case TEXTGEN_TYPES.MANCER: - headers = getMancerHeaders(); - break; - case TEXTGEN_TYPES.APHRODITE: - headers = getAphroditeHeaders(); - break; - case TEXTGEN_TYPES.TABBY: - headers = getTabbyHeaders(); - break; - default: - headers = server ? getOverrideHeaders((new URL(server))?.host) : {}; - break; - } - - Object.assign(args.headers, headers); -} - const SETTINGS_FILE = './public/settings.json'; const { DIRECTORIES, UPLOADS_PATH, PALM_SAFETY, TEXTGEN_TYPES, CHAT_COMPLETION_SOURCES, AVATAR_WIDTH, AVATAR_HEIGHT } = require('./src/constants'); @@ -1774,93 +1711,6 @@ async function sendAI21Request(request, response) { } -app.post('/api/tokenizers/remote/encode', jsonParser, async function (request, response) { - if (!request.body) { - return response.sendStatus(400); - } - const text = String(request.body.text) || ''; - const api = String(request.body.main_api); - const baseUrl = String(request.body.url); - const legacyApi = Boolean(request.body.legacy_api); - - try { - if (api == 'textgenerationwebui') { - const args = { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - }; - - setAdditionalHeaders(request, args, null); - - // Convert to string + remove trailing slash + /v1 suffix - let url = String(baseUrl).replace(/\/$/, '').replace(/\/v1$/, ''); - - if (legacyApi) { - url += '/v1/token-count'; - args.body = JSON.stringify({ 'prompt': text }); - } else { - switch (request.body.api_type) { - case TEXTGEN_TYPES.TABBY: - url += '/v1/token/encode'; - args.body = JSON.stringify({ 'text': text }); - break; - case TEXTGEN_TYPES.KOBOLDCPP: - url += '/api/extra/tokencount'; - args.body = JSON.stringify({ 'prompt': text }); - break; - default: - url += '/v1/internal/encode'; - args.body = JSON.stringify({ 'text': text }); - break; - } - } - - 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 = legacyApi ? data?.results[0]?.tokens : (data?.length ?? data?.value); - const ids = legacyApi ? [] : (data?.tokens ?? []); - - return response.send({ count, ids }); - } - - else if (api == 'kobold') { - const args = { - method: 'POST', - body: JSON.stringify({ 'prompt': text }), - headers: { 'Content-Type': 'application/json' }, - }; - - let url = String(baseUrl).replace(/\/$/, ''); - url += '/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, ids: [] }); - } - - else { - console.log('Unknown API', api); - return response.send({ error: true }); - } - } catch (error) { - console.log(error); - return response.send({ error: true }); - } -}); - /** * Redirect a deprecated API endpoint URL to its replacement. Because fetch, form submissions, and $.ajax follow * redirects, this is transparent to client-side code. diff --git a/src/additional-headers.js b/src/additional-headers.js new file mode 100644 index 000000000..61ea1790d --- /dev/null +++ b/src/additional-headers.js @@ -0,0 +1,72 @@ +const { TEXTGEN_TYPES } = require('./constants'); +const { SECRET_KEYS, readSecret } = require('./endpoints/secrets'); +const { getConfigValue } = require('./util'); + +function getMancerHeaders() { + const apiKey = readSecret(SECRET_KEYS.MANCER); + + return apiKey ? ({ + 'X-API-KEY': apiKey, + 'Authorization': `Bearer ${apiKey}`, + }) : {}; +} + +function getAphroditeHeaders() { + const apiKey = readSecret(SECRET_KEYS.APHRODITE); + + return apiKey ? ({ + 'X-API-KEY': apiKey, + 'Authorization': `Bearer ${apiKey}`, + }) : {}; +} + +function getTabbyHeaders() { + const apiKey = readSecret(SECRET_KEYS.TABBY); + + return apiKey ? ({ + 'x-api-key': apiKey, + 'Authorization': `Bearer ${apiKey}`, + }) : {}; +} + +function getOverrideHeaders(urlHost) { + const requestOverrides = getConfigValue('requestOverrides', []); + const overrideHeaders = requestOverrides?.find((e) => e.hosts?.includes(urlHost))?.headers; + if (overrideHeaders && urlHost) { + return overrideHeaders; + } else { + return {}; + } +} + +/** + * Sets additional headers for the request. + * @param {object} request Original request body + * @param {object} args New request arguments + * @param {string|null} server API server for new request + */ +function setAdditionalHeaders(request, args, server) { + let headers; + + switch (request.body.api_type) { + case TEXTGEN_TYPES.MANCER: + headers = getMancerHeaders(); + break; + case TEXTGEN_TYPES.APHRODITE: + headers = getAphroditeHeaders(); + break; + case TEXTGEN_TYPES.TABBY: + headers = getTabbyHeaders(); + break; + default: + headers = server ? getOverrideHeaders((new URL(server))?.host) : {}; + break; + } + + Object.assign(args.headers, headers); +} + +module.exports = { + getOverrideHeaders, + setAdditionalHeaders, +}; diff --git a/src/endpoints/tokenizers.js b/src/endpoints/tokenizers.js index 57abc6b8f..8a4db7728 100644 --- a/src/endpoints/tokenizers.js +++ b/src/endpoints/tokenizers.js @@ -6,7 +6,9 @@ const tiktoken = require('@dqbd/tiktoken'); const { Tokenizer } = require('@agnai/web-tokenizers'); const { convertClaudePrompt } = require('../chat-completion'); const { readSecret, SECRET_KEYS } = require('./secrets'); +const { TEXTGEN_TYPES } = require('../constants'); const { jsonParser } = require('../express-common'); +const { setAdditionalHeaders } = require('../additional-headers'); /** * @type {{[key: string]: import("@dqbd/tiktoken").Tiktoken}} Tokenizers cache @@ -534,6 +536,93 @@ router.post('/openai/count', jsonParser, async function (req, res) { } }); +router.post('/remote/encode', jsonParser, async function (request, response) { + if (!request.body) { + return response.sendStatus(400); + } + const text = String(request.body.text) || ''; + const api = String(request.body.main_api); + const baseUrl = String(request.body.url); + const legacyApi = Boolean(request.body.legacy_api); + + try { + if (api == 'textgenerationwebui') { + const args = { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + }; + + setAdditionalHeaders(request, args, null); + + // Convert to string + remove trailing slash + /v1 suffix + let url = String(baseUrl).replace(/\/$/, '').replace(/\/v1$/, ''); + + if (legacyApi) { + url += '/v1/token-count'; + args.body = JSON.stringify({ 'prompt': text }); + } else { + switch (request.body.api_type) { + case TEXTGEN_TYPES.TABBY: + url += '/v1/token/encode'; + args.body = JSON.stringify({ 'text': text }); + break; + case TEXTGEN_TYPES.KOBOLDCPP: + url += '/api/extra/tokencount'; + args.body = JSON.stringify({ 'prompt': text }); + break; + default: + url += '/v1/internal/encode'; + args.body = JSON.stringify({ 'text': text }); + break; + } + } + + 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 = legacyApi ? data?.results[0]?.tokens : (data?.length ?? data?.value); + const ids = legacyApi ? [] : (data?.tokens ?? []); + + return response.send({ count, ids }); + } + + else if (api == 'kobold') { + const args = { + method: 'POST', + body: JSON.stringify({ 'prompt': text }), + headers: { 'Content-Type': 'application/json' }, + }; + + let url = String(baseUrl).replace(/\/$/, ''); + url += '/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, ids: [] }); + } + + else { + console.log('Unknown API', api); + return response.send({ error: true }); + } + } catch (error) { + console.log(error); + return response.send({ error: true }); + } +}); + module.exports = { TEXT_COMPLETION_MODELS, getTokenizerModel,