mirror of
https://github.com/SillyTavern/SillyTavern.git
synced 2025-02-22 06:57:41 +01:00
Move API tokenization endpoint into /tokenizers
Requires extracting some more functions out of server.js.
This commit is contained in:
valadaptive
parent
8bad059a62
commit
04e92efe29
152
server.js
152
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.
|
||||
|
||||
72
src/additional-headers.js
Normal file
72
src/additional-headers.js
Normal file
@ -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,
|
||||
};
|
||||
@ -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,
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user