1
import { localforage } from '../lib.js';
2
import { characters, event_types, eventSource, main_api, nai_settings, online_status, this_chid } from '../script.js';
3
import { power_user, registerDebugFunction } from './power-user.js';
4
import { chat_completion_sources, model_list, oai_settings } from './openai.js';
5
import { groups, selected_group } from './group-chats.js';
6
import { getStringHash } from './utils.js';
7
import { kai_flags, kai_settings } from './kai-settings.js';
8
import { textgen_types, textgenerationwebui_settings as textgen_settings, getTextGenServer, getTextGenModel } from './textgen-settings.js';
9
import { getCurrentDreamGenModelTokenizer, getCurrentOpenRouterModelTokenizer, openRouterModels } from './textgen-models.js';
10
export { BYTES_PER_TOKEN as CHARACTERS_PER_TOKEN_RATIO };
11

12
export const BYTES_PER_TOKEN = 3.35;
13
export const TOKENIZER_WARNING_KEY = 'tokenizationWarningShown';
14
export const TOKENIZER_SUPPORTED_KEY = 'tokenizationSupported';
15

16
export const tokenizers = {
17
    NONE: 0,
18
    GPT2: 1,
19
    OPENAI: 2,
20
    LLAMA: 3,
21
    NERD: 4,
22
    NERD2: 5,
23
    API_CURRENT: 6,
24
    MISTRAL: 7,
25
    YI: 8,
26
    API_TEXTGENERATIONWEBUI: 9,
27
    API_KOBOLD: 10,
28
    CLAUDE: 11,
29
    LLAMA3: 12,
30
    GEMMA: 13,
31
    JAMBA: 14,
32
    QWEN2: 15,
33
    COMMAND_R: 16,
34
    NEMO: 17,
35
    DEEPSEEK: 18,
36
    COMMAND_A: 19,
37
    BEST_MATCH: 99,
38
};
39

40
// A list of local tokenizers that support encoding and decoding token ids.
41
export const ENCODE_TOKENIZERS = [
42
    tokenizers.LLAMA,
43
    tokenizers.MISTRAL,
44
    tokenizers.YI,
45
    tokenizers.LLAMA3,
46
    tokenizers.GEMMA,
47
    tokenizers.JAMBA,
48
    tokenizers.QWEN2,
49
    tokenizers.COMMAND_R,
50
    tokenizers.COMMAND_A,
51
    tokenizers.NEMO,
52
    tokenizers.DEEPSEEK,
53
    // uncomment when NovelAI releases Kayra and Clio weights, lol
54
    //tokenizers.NERD,
55
    //tokenizers.NERD2,
56
];
57

58
/**
59
 * A list of Text Completion sources that support remote tokenization.
60
 * Populated in initTokenziers due to circular dependencies.
61
 * @type {string[]}
62
 */
63
export const TEXTGEN_TOKENIZERS = [];
64

65
const TOKENIZER_URLS = {
66
    [tokenizers.GPT2]: {
67
        encode: '/api/tokenizers/gpt2/encode',
68
        decode: '/api/tokenizers/gpt2/decode',
69
        count: '/api/tokenizers/gpt2/encode',
70
    },
71
    [tokenizers.OPENAI]: {
72
        encode: '/api/tokenizers/openai/encode',
73
        decode: '/api/tokenizers/openai/decode',
74
        count: '/api/tokenizers/openai/encode',
75
    },
76
    [tokenizers.LLAMA]: {
77
        encode: '/api/tokenizers/llama/encode',
78
        decode: '/api/tokenizers/llama/decode',
79
        count: '/api/tokenizers/llama/encode',
80
    },
81
    [tokenizers.NERD]: {
82
        encode: '/api/tokenizers/nerdstash/encode',
83
        decode: '/api/tokenizers/nerdstash/decode',
84
        count: '/api/tokenizers/nerdstash/encode',
85
    },
86
    [tokenizers.NERD2]: {
87
        encode: '/api/tokenizers/nerdstash_v2/encode',
88
        decode: '/api/tokenizers/nerdstash_v2/decode',
89
        count: '/api/tokenizers/nerdstash_v2/encode',
90
    },
91
    [tokenizers.API_KOBOLD]: {
92
        count: '/api/tokenizers/remote/kobold/count',
93
        encode: '/api/tokenizers/remote/kobold/count',
94
    },
95
    [tokenizers.MISTRAL]: {
96
        encode: '/api/tokenizers/mistral/encode',
97
        decode: '/api/tokenizers/mistral/decode',
98
        count: '/api/tokenizers/mistral/encode',
99
    },
100
    [tokenizers.YI]: {
101
        encode: '/api/tokenizers/yi/encode',
102
        decode: '/api/tokenizers/yi/decode',
103
        count: '/api/tokenizers/yi/encode',
104
    },
105
    [tokenizers.CLAUDE]: {
106
        encode: '/api/tokenizers/claude/encode',
107
        decode: '/api/tokenizers/claude/decode',
108
        count: '/api/tokenizers/claude/encode',
109
    },
110
    [tokenizers.LLAMA3]: {
111
        encode: '/api/tokenizers/llama3/encode',
112
        decode: '/api/tokenizers/llama3/decode',
113
        count: '/api/tokenizers/llama3/encode',
114
    },
115
    [tokenizers.GEMMA]: {
116
        encode: '/api/tokenizers/gemma/encode',
117
        decode: '/api/tokenizers/gemma/decode',
118
        count: '/api/tokenizers/gemma/encode',
119
    },
120
    [tokenizers.JAMBA]: {
121
        encode: '/api/tokenizers/jamba/encode',
122
        decode: '/api/tokenizers/jamba/decode',
123
        count: '/api/tokenizers/jamba/encode',
124
    },
125
    [tokenizers.QWEN2]: {
126
        encode: '/api/tokenizers/qwen2/encode',
127
        decode: '/api/tokenizers/qwen2/decode',
128
        count: '/api/tokenizers/qwen2/encode',
129
    },
130
    [tokenizers.COMMAND_R]: {
131
        encode: '/api/tokenizers/command-r/encode',
132
        decode: '/api/tokenizers/command-r/decode',
133
        count: '/api/tokenizers/command-r/encode',
134
    },
135
    [tokenizers.COMMAND_A]: {
136
        encode: '/api/tokenizers/command-a/encode',
137
        decode: '/api/tokenizers/command-a/decode',
138
        count: '/api/tokenizers/command-a/encode',
139
    },
140
    [tokenizers.NEMO]: {
141
        encode: '/api/tokenizers/nemo/encode',
142
        decode: '/api/tokenizers/nemo/decode',
143
        count: '/api/tokenizers/nemo/encode',
144
    },
145
    [tokenizers.DEEPSEEK]: {
146
        encode: '/api/tokenizers/deepseek/encode',
147
        decode: '/api/tokenizers/deepseek/decode',
148
        count: '/api/tokenizers/deepseek/encode',
149
    },
150
    [tokenizers.API_TEXTGENERATIONWEBUI]: {
151
        encode: '/api/tokenizers/remote/textgenerationwebui/encode',
152
        count: '/api/tokenizers/remote/textgenerationwebui/encode',
153
    },
154
};
155

156
const textEncoder = new TextEncoder();
157
const objectStore = localforage.createInstance({ name: 'SillyTavern_ChatCompletions' });
158

159
let tokenCache = {};
160

161
/**
162
 * Guesstimates the token count for a string.
163
 * @param {string} str String to tokenize.
164
 * @returns {number} Token count.
165
 */
166
export function guesstimate(str) {
167
    const byteLength = textEncoder.encode(str).length;
168
    return Math.ceil(byteLength / BYTES_PER_TOKEN);
169
}
170

171
async function loadTokenCache() {
172
    try {
173
        console.debug('Chat Completions: loading token cache');
174
        tokenCache = await objectStore.getItem('tokenCache') || {};
175
    } catch (e) {
176
        console.log('Chat Completions: unable to load token cache, using default value', e);
177
        tokenCache = {};
178
    }
179
}
180

181
export async function saveTokenCache() {
182
    try {
183
        console.debug('Chat Completions: saving token cache');
184
        await objectStore.setItem('tokenCache', tokenCache);
185
    } catch (e) {
186
        console.log('Chat Completions: unable to save token cache', e);
187
    }
188
}
189

190
async function resetTokenCache() {
191
    try {
192
        console.debug('Chat Completions: resetting token cache');
193
        Object.keys(tokenCache).forEach(key => delete tokenCache[key]);
194
        await objectStore.removeItem('tokenCache');
195
        toastr.success('Token cache cleared. Please reload the chat to re-tokenize it.');
196
    } catch (e) {
197
        console.log('Chat Completions: unable to reset token cache', e);
198
    }
199
}
200

201
/**
202
 * @typedef {object} Tokenizer
203
 * @property {number} tokenizerId - The id of the tokenizer option
204
 * @property {string} tokenizerKey - Internal name/key of the tokenizer
205
 * @property {string} tokenizerName - Human-readable detailed name of the tokenizer (as displayed in the UI)
206
 */
207

208
/**
209
 * Gets all tokenizers available to the user.
210
 * @returns {Tokenizer[]} Tokenizer info.
211
 */
212
export function getAvailableTokenizers() {
213
    const tokenizerOptions = $('#tokenizer').find('option').toArray();
214
    return tokenizerOptions.map(tokenizerOption => ({
215
        tokenizerId: Number(tokenizerOption.value),
216
        tokenizerKey: Object.entries(tokenizers).find(([_, value]) => value === Number(tokenizerOption.value))[0].toLocaleLowerCase(),
217
        tokenizerName: tokenizerOption.text,
218
    }));
219
}
220

221
/**
222
 * Selects tokenizer if not already selected.
223
 * @param {number} tokenizerId Tokenizer ID.
224
 */
225
export function selectTokenizer(tokenizerId) {
226
    if (tokenizerId !== power_user.tokenizer) {
227
        const tokenizer = getAvailableTokenizers().find(tokenizer => tokenizer.tokenizerId === tokenizerId);
228
        if (!tokenizer) {
229
            console.warn('Failed to find tokenizer with id', tokenizerId);
230
            return;
231
        }
232
        $('#tokenizer').val(tokenizer.tokenizerId).trigger('change');
233
        toastr.info(`Tokenizer: "${tokenizer.tokenizerName}" selected`);
234
    }
235
}
236

237
/**
238
 * Gets the friendly name of the current tokenizer.
239
 * @param {string} forApi API to get the tokenizer for. Defaults to the main API.
240
 * @returns {Tokenizer} Tokenizer info
241
 */
242
export function getFriendlyTokenizerName(forApi) {
243
    if (!forApi) {
244
        forApi = main_api;
245
    }
246

247
    const tokenizerOption = $('#tokenizer').find(':selected');
248
    let tokenizerId = Number(tokenizerOption.val());
249
    let tokenizerName = tokenizerOption.text();
250

251
    if (forApi !== 'openai' && tokenizerId === tokenizers.BEST_MATCH) {
252
        tokenizerId = getTokenizerBestMatch(forApi);
253

254
        switch (tokenizerId) {
255
            case tokenizers.API_KOBOLD:
256
                tokenizerName = 'API (KoboldAI Classic)';
257
                break;
258
            case tokenizers.API_TEXTGENERATIONWEBUI:
259
                tokenizerName = 'API (Text Completion)';
260
                break;
261
            default:
262
                tokenizerName = $(`#tokenizer option[value="${tokenizerId}"]`).text();
263
                break;
264
        }
265
    }
266

267
    tokenizerName = forApi == 'openai'
268
        ? getTokenizerModel()
269
        : tokenizerName;
270

271
    tokenizerId = forApi == 'openai'
272
        ? tokenizers.OPENAI
273
        : tokenizerId;
274

275
    const tokenizerKey = Object.entries(tokenizers).find(([_, value]) => value === tokenizerId)[0].toLocaleLowerCase();
276

277
    return { tokenizerName, tokenizerKey, tokenizerId };
278
}
279

280
/**
281
 * Gets the best tokenizer for the current API.
282
 * @param {string} forApi API to get the tokenizer for. Defaults to the main API.
283
 * @returns {number} Tokenizer type.
284
 */
285
export function getTokenizerBestMatch(forApi) {
286
    if (!forApi) {
287
        forApi = main_api;
288
    }
289

290
    if (forApi === 'novel') {
291
        if (nai_settings.model_novel.includes('clio')) {
292
            return tokenizers.NERD;
293
        }
294
        if (nai_settings.model_novel.includes('kayra')) {
295
            return tokenizers.NERD2;
296
        }
297
        if (nai_settings.model_novel.includes('erato')) {
298
            return tokenizers.LLAMA3;
299
        }
300
    }
301
    if (forApi === 'kobold' || forApi === 'textgenerationwebui' || forApi === 'koboldhorde') {
302
        // Try to use the API tokenizer if possible:
303
        // - API must be connected
304
        // - Kobold must pass a version check
305
        // - Tokenizer haven't reported an error previously
306
        const hasTokenizerError = sessionStorage.getItem(TOKENIZER_WARNING_KEY);
307
        const hasValidEndpoint = sessionStorage.getItem(TOKENIZER_SUPPORTED_KEY);
308
        const isConnected = online_status !== 'no_connection';
309
        const isTokenizerSupported = TEXTGEN_TOKENIZERS.includes(textgen_settings.type) && (textgen_settings.type !== textgen_types.OOBA || hasValidEndpoint);
310

311
        if (!hasTokenizerError && isConnected) {
312
            if (forApi === 'kobold' && kai_flags.can_use_tokenization) {
313
                return tokenizers.API_KOBOLD;
314
            }
315

316
            if (forApi === 'textgenerationwebui' && isTokenizerSupported) {
317
                return tokenizers.API_TEXTGENERATIONWEBUI;
318
            }
319
            if (forApi === 'textgenerationwebui' && textgen_settings.type === textgen_types.OPENROUTER) {
320
                return getCurrentOpenRouterModelTokenizer();
321
            }
322
            if (forApi === 'textgenerationwebui' && textgen_settings.type === textgen_types.DREAMGEN) {
323
                return getCurrentDreamGenModelTokenizer();
324
            }
325
        }
326

327
        if (forApi === 'textgenerationwebui') {
328
            const model = String(getTextGenModel() || online_status).toLowerCase();
329
            if (model.includes('llama3') || model.includes('llama-3')) {
330
                return tokenizers.LLAMA3;
331
            }
332
            if (model.includes('mistral') || model.includes('mixtral')) {
333
                return tokenizers.MISTRAL;
334
            }
335
            if (model.includes('gemma')) {
336
                return tokenizers.GEMMA;
337
            }
338
            if (model.includes('nemo') || model.includes('pixtral')) {
339
                return tokenizers.NEMO;
340
            }
341
            if (model.includes('deepseek')) {
342
                return tokenizers.DEEPSEEK;
343
            }
344
            if (model.includes('yi')) {
345
                return tokenizers.YI;
346
            }
347
            if (model.includes('jamba')) {
348
                return tokenizers.JAMBA;
349
            }
350
            if (model.includes('command-r')) {
351
                return tokenizers.COMMAND_R;
352
            }
353
            if (model.includes('command-a')) {
354
                return tokenizers.COMMAND_A;
355
            }
356
            if (model.includes('qwen2')) {
357
                return tokenizers.QWEN2;
358
            }
359
        }
360

361
        return tokenizers.LLAMA;
362
    }
363

364
    return tokenizers.NONE;
365
}
366

367
// Get the current remote tokenizer API based on the current text generation API.
368
function currentRemoteTokenizerAPI() {
369
    switch (main_api) {
370
        case 'kobold':
371
            return tokenizers.API_KOBOLD;
372
        case 'textgenerationwebui':
373
            return tokenizers.API_TEXTGENERATIONWEBUI;
374
        default:
375
            return tokenizers.NONE;
376
    }
377
}
378

379
/**
380
 * Calls the underlying tokenizer model to the token count for a string.
381
 * @param {number} type Tokenizer type.
382
 * @param {string} str String to tokenize.
383
 * @returns {number} Token count.
384
 */
385
function callTokenizer(type, str) {
386
    if (type === tokenizers.NONE) return guesstimate(str);
387

388
    switch (type) {
389
        case tokenizers.API_CURRENT:
390
            return callTokenizer(currentRemoteTokenizerAPI(), str);
391
        case tokenizers.API_KOBOLD:
392
            return countTokensFromKoboldAPI(str);
393
        case tokenizers.API_TEXTGENERATIONWEBUI:
394
            return countTokensFromTextgenAPI(str);
395
        default: {
396
            const endpointUrl = TOKENIZER_URLS[type]?.count;
397
            if (!endpointUrl) {
398
                console.warn('Unknown tokenizer type', type);
399
                return apiFailureTokenCount(str);
400
            }
401
            return countTokensFromServer(endpointUrl, str);
402
        }
403
    }
404
}
405

406
/**
407
 * Calls the underlying tokenizer model to the token count for a string.
408
 * @param {number} type Tokenizer type.
409
 * @param {string} str String to tokenize.
410
 * @returns {Promise<number>} Token count.
411
 */
412
function callTokenizerAsync(type, str) {
413
    return new Promise(resolve => {
414
        if (type === tokenizers.NONE) {
415
            return resolve(guesstimate(str));
416
        }
417

418
        switch (type) {
419
            case tokenizers.API_CURRENT:
420
                return callTokenizerAsync(currentRemoteTokenizerAPI(), str).then(resolve);
421
            case tokenizers.API_KOBOLD:
422
                return countTokensFromKoboldAPI(str, resolve);
423
            case tokenizers.API_TEXTGENERATIONWEBUI:
424
                return countTokensFromTextgenAPI(str, resolve);
425
            default: {
426
                const endpointUrl = TOKENIZER_URLS[type]?.count;
427
                if (!endpointUrl) {
428
                    console.warn('Unknown tokenizer type', type);
429
                    return resolve(apiFailureTokenCount(str));
430
                }
431
                return countTokensFromServer(endpointUrl, str, resolve);
432
            }
433
        }
434
    });
435
}
436

437
/**
438
 * Gets the token count for a string using the current model tokenizer.
439
 * @param {string} str String to tokenize
440
 * @param {number | undefined} padding Optional padding tokens. Defaults to 0.
441
 * @returns {Promise<number>} Token count.
442
 */
443
export async function getTokenCountAsync(str, padding = undefined) {
444
    if (typeof str !== 'string' || !str?.length) {
445
        return 0;
446
    }
447

448
    let tokenizerType = power_user.tokenizer;
449
    let modelHash = '';
450

451
    if (main_api === 'openai') {
452
        if (padding === power_user.token_padding) {
453
            // For main "shadow" prompt building
454
            tokenizerType = tokenizers.NONE;
455
        } else {
456
            // For extensions and WI
457
            return counterWrapperOpenAIAsync(str);
458
        }
459
    }
460

461
    if (tokenizerType === tokenizers.BEST_MATCH) {
462
        tokenizerType = getTokenizerBestMatch(main_api);
463
    }
464

465
    if (tokenizerType === tokenizers.API_TEXTGENERATIONWEBUI) {
466
        modelHash = getStringHash(getTextGenModel() || online_status).toString();
467
    }
468

469
    if (padding === undefined) {
470
        padding = 0;
471
    }
472

473
    const cacheObject = getTokenCacheObject();
474
    const hash = getStringHash(str);
475
    const cacheKey = `${tokenizerType}-${hash}${modelHash}+${padding}`;
476

477
    if (typeof cacheObject[cacheKey] === 'number') {
478
        return cacheObject[cacheKey];
479
    }
480

481
    const result = (await callTokenizerAsync(tokenizerType, str)) + padding;
482

483
    if (isNaN(result)) {
484
        console.warn('Token count calculation returned NaN');
485
        return 0;
486
    }
487

488
    cacheObject[cacheKey] = result;
489
    return result;
490
}
491

492
/**
493
 * Gets the token count for a string using the current model tokenizer.
494
 * @param {string} str String to tokenize
495
 * @param {number | undefined} padding Optional padding tokens. Defaults to 0.
496
 * @returns {number} Token count.
497
 * @deprecated Use getTokenCountAsync instead.
498
 */
499
export function getTokenCount(str, padding = undefined) {
500
    if (typeof str !== 'string' || !str?.length) {
501
        return 0;
502
    }
503

504
    let tokenizerType = power_user.tokenizer;
505
    let modelHash = '';
506

507
    if (main_api === 'openai') {
508
        if (padding === power_user.token_padding) {
509
            // For main "shadow" prompt building
510
            tokenizerType = tokenizers.NONE;
511
        } else {
512
            // For extensions and WI
513
            return counterWrapperOpenAI(str);
514
        }
515
    }
516

517
    if (tokenizerType === tokenizers.BEST_MATCH) {
518
        tokenizerType = getTokenizerBestMatch(main_api);
519
    }
520

521
    if (tokenizerType === tokenizers.API_TEXTGENERATIONWEBUI) {
522
        modelHash = getStringHash(getTextGenModel() || online_status).toString();
523
    }
524

525
    if (padding === undefined) {
526
        padding = 0;
527
    }
528

529
    const cacheObject = getTokenCacheObject();
530
    const hash = getStringHash(str);
531
    const cacheKey = `${tokenizerType}-${hash}${modelHash}+${padding}`;
532

533
    if (typeof cacheObject[cacheKey] === 'number') {
534
        return cacheObject[cacheKey];
535
    }
536

537
    const result = callTokenizer(tokenizerType, str) + padding;
538

539
    if (isNaN(result)) {
540
        console.warn('Token count calculation returned NaN');
541
        return 0;
542
    }
543

544
    cacheObject[cacheKey] = result;
545
    return result;
546
}
547

548
/**
549
 * Gets the token count for a string using the OpenAI tokenizer.
550
 * @param {string} text Text to tokenize.
551
 * @returns {number} Token count.
552
 * @deprecated Use counterWrapperOpenAIAsync instead.
553
 */
554
function counterWrapperOpenAI(text) {
555
    const message = { role: 'system', content: text };
556
    return countTokensOpenAI(message, true);
557
}
558

559
/**
560
 * Gets the token count for a string using the OpenAI tokenizer.
561
 * @param {string} text Text to tokenize.
562
 * @returns {Promise<number>} Token count.
563
 */
564
function counterWrapperOpenAIAsync(text) {
565
    const message = { role: 'system', content: text };
566
    return countTokensOpenAIAsync(message, true);
567
}
568

569
export function getTokenizerModel() {
570
    // OpenAI models always provide their own tokenizer
571
    if (oai_settings.chat_completion_source == chat_completion_sources.OPENAI) {
572
        return oai_settings.openai_model;
573
    }
574

575
    const turboTokenizer = 'gpt-3.5-turbo';
576
    const gpt4Tokenizer = 'gpt-4';
577
    const gpt4oTokenizer = 'gpt-4o';
578
    const gpt2Tokenizer = 'gpt2';
579
    const claudeTokenizer = 'claude';
580
    const llamaTokenizer = 'llama';
581
    const llama3Tokenizer = 'llama3';
582
    const mistralTokenizer = 'mistral';
583
    const yiTokenizer = 'yi';
584
    const gemmaTokenizer = 'gemma';
585
    const jambaTokenizer = 'jamba';
586
    const qwen2Tokenizer = 'qwen2';
587
    const commandRTokenizer = 'command-r';
588
    const commandATokenizer = 'command-a';
589
    const nemoTokenizer = 'nemo';
590
    const deepseekTokenizer = 'deepseek';
591

592
    if (oai_settings.chat_completion_source == chat_completion_sources.AZURE_OPENAI) {
593
        return oai_settings.azure_openai_model || turboTokenizer;
594
    }
595

596
    if (oai_settings.chat_completion_source == chat_completion_sources.DEEPSEEK) {
597
        return deepseekTokenizer;
598
    }
599

600
    // And for OpenRouter (if not a site model, then it's impossible to determine the tokenizer)
601
    if (main_api == 'openai' && oai_settings.chat_completion_source == chat_completion_sources.OPENROUTER && oai_settings.openrouter_model ||
602
        main_api == 'textgenerationwebui' && textgen_settings.type === textgen_types.OPENROUTER && textgen_settings.openrouter_model) {
603
        const model = main_api == 'openai'
604
            ? model_list.find(x => x.id === oai_settings.openrouter_model)
605
            : openRouterModels.find(x => x.id === textgen_settings.openrouter_model);
606

607
        if (model?.architecture?.tokenizer === 'Llama2') {
608
            return llamaTokenizer;
609
        } else if (model?.architecture?.tokenizer === 'Llama3') {
610
            return llama3Tokenizer;
611
        } else if (model?.architecture?.tokenizer === 'Mistral') {
612
            return mistralTokenizer;
613
        } else if (model?.architecture?.tokenizer === 'Yi') {
614
            return yiTokenizer;
615
        } else if (model?.architecture?.tokenizer === 'Gemini') {
616
            return gemmaTokenizer;
617
        } else if (model?.architecture?.tokenizer === 'Qwen') {
618
            return qwen2Tokenizer;
619
        } else if (model?.architecture?.tokenizer === 'Cohere') {
620
            if (model?.id && model?.id.includes('command-a')) {
621
                return commandATokenizer;
622
            }
623
            return commandRTokenizer;
624
        } else if (oai_settings.openrouter_model.includes('gpt-4o')) {
625
            return gpt4oTokenizer;
626
        } else if (oai_settings.openrouter_model.includes('gpt-4')) {
627
            return gpt4Tokenizer;
628
        } else if (oai_settings.openrouter_model.includes('gpt-3.5-turbo')) {
629
            return turboTokenizer;
630
        } else if (oai_settings.openrouter_model.includes('claude')) {
631
            return claudeTokenizer;
632
        } else if (oai_settings.openrouter_model.includes('GPT-NeoXT')) {
633
            return gpt2Tokenizer;
634
        } else if (oai_settings.openrouter_model.includes('jamba')) {
635
            return jambaTokenizer;
636
        } else if (oai_settings.openrouter_model.includes('deepseek')) {
637
            return deepseekTokenizer;
638
        }
639
    }
640

641
    if (oai_settings.chat_completion_source == chat_completion_sources.ELECTRONHUB && oai_settings.electronhub_model) {
642
        if (oai_settings.electronhub_model.includes('gpt-4o') || oai_settings.electronhub_model.includes('gpt-5')) {
643
            return gpt4oTokenizer;
644
        } else if (oai_settings.electronhub_model.includes('gpt-4.1') || oai_settings.electronhub_model.includes('gpt-4.5')) {
645
            return gpt4oTokenizer;
646
        } else if (oai_settings.electronhub_model.includes('gpt-4')) {
647
            return gpt4Tokenizer;
648
        } else if (oai_settings.electronhub_model.includes('gpt-3.5-turbo')) {
649
            return turboTokenizer;
650
        } else if (oai_settings.electronhub_model.includes('claude')) {
651
            return claudeTokenizer;
652
        } else if (oai_settings.electronhub_model.includes('jamba')) {
653
            return jambaTokenizer;
654
        } else if (oai_settings.electronhub_model.includes('deepseek') || oai_settings.electronhub_model.includes('sonar-reasoning') || oai_settings.electronhub_model.includes('r1')) {
655
            return deepseekTokenizer;
656
        } else if (oai_settings.electronhub_model.includes('qwen')) {
657
            return qwen2Tokenizer;
658
        } else if (oai_settings.electronhub_model.includes('gemma')) {
659
            return gemmaTokenizer;
660
        } else if (oai_settings.electronhub_model.includes('mistral')) {
661
            return mistralTokenizer;
662
        } else if (oai_settings.electronhub_model.includes('yi')) {
663
            return yiTokenizer;
664
        } else if (oai_settings.electronhub_model.includes('llama3') || oai_settings.electronhub_model.includes('llama-3') || oai_settings.electronhub_model.startsWith('l3')) {
665
            return llama3Tokenizer;
666
        } else if (oai_settings.electronhub_model.includes('llama')) {
667
            return llamaTokenizer;
668
        } else if (oai_settings.electronhub_model.includes('command-a')) {
669
            return commandATokenizer;
670
        } else if (oai_settings.electronhub_model.includes('command-r')) {
671
            return commandRTokenizer;
672
        } else if (oai_settings.electronhub_model.includes('nemo')) {
673
            return nemoTokenizer;
674
        }
675
    }
676

677
    if (oai_settings.chat_completion_source == chat_completion_sources.CHUTES && oai_settings.chutes_model) {
678
        const model = oai_settings.chutes_model.toLowerCase();
679

680
        if (model.includes('deepseek') || model.includes('mai-ds')) {
681
            return deepseekTokenizer;
682
        } else if (model.includes('qwen') || model.includes('qwq') || model.includes('tongyi') || model.includes('kimi')) {
683
            return qwen2Tokenizer;
684
        } else if (model.includes('llama') || model.includes('longcat') || model.includes('hermes')) {
685
            return llama3Tokenizer;
686
        } else if (model.includes('gemma')) {
687
            return gemmaTokenizer;
688
        } else if (model.includes('nemo')) {
689
            return nemoTokenizer;
690
        } else if (model.includes('mistral')) {
691
            return mistralTokenizer;
692
        } else if (model.includes('gpt-oss')) {
693
            return gpt4oTokenizer;
694
        }
695
    }
696

697
    if (oai_settings.chat_completion_source == chat_completion_sources.MINIMAX) {
698
        // MiniMax uses a proprietary tokenizer; fall back to a coarse OpenAI estimation.
699
        return 'gpt-3.5-turbo';
700
    }
701

702
    if (oai_settings.chat_completion_source == chat_completion_sources.WORKERS_AI && oai_settings.workers_ai_model) {
703
        const model = oai_settings.workers_ai_model.toLowerCase();
704

705
        if (model.includes('deepseek')) {
706
            return deepseekTokenizer;
707
        } else if (model.includes('qwen') || model.includes('qwq') || model.includes('kimi')) {
708
            return qwen2Tokenizer;
709
        } else if (model.includes('llama-3') || model.includes('llama-4')) {
710
            return llama3Tokenizer;
711
        } else if (model.includes('llama')) {
712
            return llamaTokenizer;
713
        } else if (model.includes('gemma')) {
714
            return gemmaTokenizer;
715
        } else if (model.includes('mistral')) {
716
            return mistralTokenizer;
717
        } else if (model.includes('phi')) {
718
            return turboTokenizer;
719
        } else if (model.includes('gpt-oss')) {
720
            return gpt4oTokenizer;
721
        }
722
    }
723

724
    if (oai_settings.chat_completion_source == chat_completion_sources.COHERE) {
725
        if (oai_settings.cohere_model.includes('command-a')) {
726
            return commandATokenizer;
727
        }
728
        return commandRTokenizer;
729
    }
730

731
    if (oai_settings.chat_completion_source == chat_completion_sources.MAKERSUITE) {
732
        return gemmaTokenizer;
733
    }
734

735
    if (oai_settings.chat_completion_source == chat_completion_sources.VERTEXAI) {
736
        return gemmaTokenizer;
737
    }
738

739
    if (oai_settings.chat_completion_source == chat_completion_sources.AI21) {
740
        return jambaTokenizer;
741
    }
742

743
    if (oai_settings.chat_completion_source == chat_completion_sources.CLAUDE) {
744
        return claudeTokenizer;
745
    }
746

747
    if (oai_settings.chat_completion_source == chat_completion_sources.MISTRALAI) {
748
        if (oai_settings.mistralai_model.includes('nemo') || oai_settings.mistralai_model.includes('pixtral')) {
749
            return nemoTokenizer;
750
        }
751
        return mistralTokenizer;
752
    }
753

754
    if (oai_settings.chat_completion_source == chat_completion_sources.CUSTOM) {
755
        return oai_settings.custom_model;
756
    }
757

758
    if (oai_settings.chat_completion_source === chat_completion_sources.PERPLEXITY) {
759
        if (oai_settings.perplexity_model.includes('sonar-reasoning') || oai_settings.perplexity_model.includes('r1-1776')) {
760
            return deepseekTokenizer;
761
        }
762
        if (oai_settings.perplexity_model.includes('llama-3') || oai_settings.perplexity_model.includes('llama3')) {
763
            return llama3Tokenizer;
764
        }
765
        if (oai_settings.perplexity_model.includes('llama')) {
766
            return llamaTokenizer;
767
        }
768
        if (oai_settings.perplexity_model.includes('mistral') || oai_settings.perplexity_model.includes('mixtral')) {
769
            return mistralTokenizer;
770
        }
771
    }
772

773
    if (oai_settings.chat_completion_source === chat_completion_sources.GROQ) {
774
        if (oai_settings.groq_model.includes('qwen')) {
775
            return qwen2Tokenizer;
776
        }
777
        if (oai_settings.groq_model.includes('llama-3') || oai_settings.groq_model.includes('llama3')) {
778
            return llama3Tokenizer;
779
        }
780
        if (oai_settings.groq_model.includes('mistral') || oai_settings.groq_model.includes('mixtral')) {
781
            return mistralTokenizer;
782
        }
783
        if (oai_settings.groq_model.includes('gemma')) {
784
            return gemmaTokenizer;
785
        }
786
    }
787

788
    // Default to Turbo 3.5
789
    return turboTokenizer;
790
}
791

792
/**
793
 * @param {any[] | Object} messages
794
 * @deprecated Use countTokensOpenAIAsync instead.
795
 */
796
export function countTokensOpenAI(messages, full = false) {
797
    const tokenizerEndpoint = `/api/tokenizers/openai/count?model=${getTokenizerModel()}`;
798
    const cacheObject = getTokenCacheObject();
799

800
    if (!Array.isArray(messages)) {
801
        messages = [messages];
802
    }
803

804
    let token_count = -1;
805

806
    for (const message of messages) {
807
        const model = getTokenizerModel();
808

809
        if (model === 'claude') {
810
            full = true;
811
        }
812

813
        const hash = getStringHash(JSON.stringify(message));
814
        const cacheKey = `${model}-${hash}`;
815
        const cachedCount = cacheObject[cacheKey];
816

817
        if (typeof cachedCount === 'number') {
818
            token_count += cachedCount;
819
        } else {
820
            jQuery.ajax({
821
                async: false,
822
                type: 'POST', //
823
                url: tokenizerEndpoint,
824
                data: JSON.stringify([message]),
825
                dataType: 'json',
826
                contentType: 'application/json',
827
                success: function (data) {
828
                    token_count += Number(data.token_count);
829
                    cacheObject[cacheKey] = Number(data.token_count);
830
                },
831
            });
832
        }
833
    }
834

835
    if (!full) token_count -= 2;
836

837
    return token_count;
838
}
839

840
/**
841
 * Returns the token count for a message using the OpenAI tokenizer.
842
 * @param {object[]|object} messages
843
 * @param {boolean} full
844
 * @returns {Promise<number>} Token count.
845
 */
846
export async function countTokensOpenAIAsync(messages, full = false) {
847
    const tokenizerEndpoint = `/api/tokenizers/openai/count?model=${getTokenizerModel()}`;
848
    const cacheObject = getTokenCacheObject();
849

850
    if (!Array.isArray(messages)) {
851
        messages = [messages];
852
    }
853

854
    let token_count = -1;
855

856
    for (const message of messages) {
857
        const model = getTokenizerModel();
858

859
        if (model === 'claude') {
860
            full = true;
861
        }
862

863
        const hash = getStringHash(JSON.stringify(message));
864
        const cacheKey = `${model}-${hash}`;
865
        const cachedCount = cacheObject[cacheKey];
866

867
        if (typeof cachedCount === 'number') {
868
            token_count += cachedCount;
869
        } else {
870
            const data = await jQuery.ajax({
871
                async: true,
872
                type: 'POST', //
873
                url: tokenizerEndpoint,
874
                data: JSON.stringify([message]),
875
                dataType: 'json',
876
                contentType: 'application/json',
877
            });
878

879
            token_count += Number(data.token_count);
880
            cacheObject[cacheKey] = Number(data.token_count);
881
        }
882
    }
883

884
    if (!full) token_count -= 2;
885

886
    return token_count;
887
}
888

889
/**
890
 * Gets the token cache object for the current chat.
891
 * @returns {Object} Token cache object for the current chat.
892
 */
893
function getTokenCacheObject() {
894
    let chatId = 'undefined';
895

896
    try {
897
        if (selected_group) {
898
            chatId = groups.find(x => x.id == selected_group)?.chat_id;
899
        } else if (this_chid !== undefined) {
900
            chatId = characters[this_chid].chat;
901
        }
902
    } catch {
903
        console.log('No character / group selected. Using default cache item');
904
    }
905

906
    if (typeof tokenCache[chatId] !== 'object') {
907
        tokenCache[chatId] = {};
908
    }
909

910
    return tokenCache[String(chatId)];
911
}
912

913
/**
914
 * Count tokens using the server API.
915
 * @param {string} endpoint API endpoint.
916
 * @param {string} str String to tokenize.
917
 * @param {function} [resolve] Promise resolve function.s
918
 * @returns {number} Token count.
919
 */
920
function countTokensFromServer(endpoint, str, resolve) {
921
    const isAsync = typeof resolve === 'function';
922
    let tokenCount = 0;
923

924
    jQuery.ajax({
925
        async: isAsync,
926
        type: 'POST',
927
        url: endpoint,
928
        data: JSON.stringify({ text: str }),
929
        dataType: 'json',
930
        contentType: 'application/json',
931
        success: function (data) {
932
            if (typeof data.count === 'number') {
933
                tokenCount = data.count;
934
            } else {
935
                tokenCount = apiFailureTokenCount(str);
936
            }
937

938
            isAsync && resolve(tokenCount);
939
        },
940
    });
941

942
    return tokenCount;
943
}
944

945
/**
946
 * Count tokens using the AI provider's API.
947
 * @param {string} str String to tokenize.
948
 * @param {function} [resolve] Promise resolve function.
949
 * @returns {number} Token count.
950
 */
951
function countTokensFromKoboldAPI(str, resolve) {
952
    const isAsync = typeof resolve === 'function';
953
    let tokenCount = 0;
954

955
    jQuery.ajax({
956
        async: isAsync,
957
        type: 'POST',
958
        url: TOKENIZER_URLS[tokenizers.API_KOBOLD].count,
959
        data: JSON.stringify({
960
            text: str,
961
            url: kai_settings.api_server,
962
        }),
963
        dataType: 'json',
964
        contentType: 'application/json',
965
        success: function (data) {
966
            if (typeof data.count === 'number') {
967
                tokenCount = data.count;
968
            } else {
969
                tokenCount = apiFailureTokenCount(str);
970
            }
971

972
            isAsync && resolve(tokenCount);
973
        },
974
    });
975

976
    return tokenCount;
977
}
978

979
function getTextgenAPITokenizationParams(str) {
980
    return {
981
        text: str,
982
        api_type: textgen_settings.type,
983
        url: getTextGenServer(),
984
        model: getTextGenModel(),
985
    };
986
}
987

988
/**
989
 * Count tokens using the AI provider's API.
990
 * @param {string} str String to tokenize.
991
 * @param {function} [resolve] Promise resolve function.
992
 * @returns {number} Token count.
993
 */
994
function countTokensFromTextgenAPI(str, resolve) {
995
    const isAsync = typeof resolve === 'function';
996
    let tokenCount = 0;
997

998
    jQuery.ajax({
999
        async: isAsync,
1000
        type: 'POST',
1001
        url: TOKENIZER_URLS[tokenizers.API_TEXTGENERATIONWEBUI].count,
1002
        data: JSON.stringify(getTextgenAPITokenizationParams(str)),
1003
        dataType: 'json',
1004
        contentType: 'application/json',
1005
        success: function (data) {
1006
            if (typeof data.count === 'number') {
1007
                tokenCount = data.count;
1008
            } else {
1009
                tokenCount = apiFailureTokenCount(str);
1010
            }
1011

1012
            isAsync && resolve(tokenCount);
1013
        },
1014
    });
1015

1016
    return tokenCount;
1017
}
1018

1019
function apiFailureTokenCount(str) {
1020
    console.error('Error counting tokens');
1021
    let shouldTryAgain = false;
1022

1023
    if (!sessionStorage.getItem(TOKENIZER_WARNING_KEY)) {
1024
        const bestMatchBefore = getTokenizerBestMatch(main_api);
1025
        sessionStorage.setItem(TOKENIZER_WARNING_KEY, String(true));
1026
        const bestMatchAfter = getTokenizerBestMatch(main_api);
1027
        if ([tokenizers.API_TEXTGENERATIONWEBUI, tokenizers.API_KOBOLD].includes(bestMatchBefore) && bestMatchBefore !== bestMatchAfter) {
1028
            shouldTryAgain = true;
1029
        }
1030
    }
1031

1032
    // Only try again if we guarantee not to be looped by the same error
1033
    if (shouldTryAgain && power_user.tokenizer === tokenizers.BEST_MATCH) {
1034
        return getTokenCount(str);
1035
    }
1036

1037
    return guesstimate(str);
1038
}
1039

1040
/**
1041
 * Calls the underlying tokenizer model to encode a string to tokens.
1042
 * @param {string} endpoint API endpoint.
1043
 * @param {string} str String to tokenize.
1044
 * @param {function} [resolve] Promise resolve function.
1045
 * @returns {number[]} Array of token ids.
1046
 */
1047
function getTextTokensFromServer(endpoint, str, resolve) {
1048
    const isAsync = typeof resolve === 'function';
1049
    let ids = [];
1050
    jQuery.ajax({
1051
        async: isAsync,
1052
        type: 'POST',
1053
        url: endpoint,
1054
        data: JSON.stringify({ text: str }),
1055
        dataType: 'json',
1056
        contentType: 'application/json',
1057
        success: function (data) {
1058
            ids = data.ids;
1059

1060
            // Don't want to break reverse compatibility, so sprinkle in some of the JS magic
1061
            if (Array.isArray(data.chunks)) {
1062
                Object.defineProperty(ids, 'chunks', { value: data.chunks });
1063
            }
1064

1065
            isAsync && resolve(ids);
1066
        },
1067
    });
1068
    return ids;
1069
}
1070

1071
/**
1072
 * Calls the AI provider's tokenize API to encode a string to tokens.
1073
 * @param {string} str String to tokenize.
1074
 * @param {function} [resolve] Promise resolve function.
1075
 * @returns {number[]} Array of token ids.
1076
 */
1077
function getTextTokensFromTextgenAPI(str, resolve) {
1078
    const isAsync = typeof resolve === 'function';
1079
    let ids = [];
1080
    jQuery.ajax({
1081
        async: isAsync,
1082
        type: 'POST',
1083
        url: TOKENIZER_URLS[tokenizers.API_TEXTGENERATIONWEBUI].encode,
1084
        data: JSON.stringify(getTextgenAPITokenizationParams(str)),
1085
        dataType: 'json',
1086
        contentType: 'application/json',
1087
        success: function (data) {
1088
            ids = data.ids;
1089
            isAsync && resolve(ids);
1090
        },
1091
    });
1092
    return ids;
1093
}
1094

1095
/**
1096
 * Calls the AI provider's tokenize API to encode a string to tokens.
1097
 * @param {string} str String to tokenize.
1098
 * @param {function} [resolve] Promise resolve function.
1099
 * @returns {number[]} Array of token ids.
1100
 */
1101
function getTextTokensFromKoboldAPI(str, resolve) {
1102
    const isAsync = typeof resolve === 'function';
1103
    let ids = [];
1104

1105
    jQuery.ajax({
1106
        async: isAsync,
1107
        type: 'POST',
1108
        url: TOKENIZER_URLS[tokenizers.API_KOBOLD].encode,
1109
        data: JSON.stringify({
1110
            text: str,
1111
            url: kai_settings.api_server,
1112
        }),
1113
        dataType: 'json',
1114
        contentType: 'application/json',
1115
        success: function (data) {
1116
            ids = data.ids;
1117
            isAsync && resolve(ids);
1118
        },
1119
    });
1120

1121
    return ids;
1122
}
1123

1124
/**
1125
 * Calls the underlying tokenizer model to decode token ids to text.
1126
 * @param {string} endpoint API endpoint.
1127
 * @param {number[]} ids Array of token ids
1128
 * @param {function} [resolve] Promise resolve function.
1129
 * @returns {({ text: string, chunks?: string[] })} Decoded token text as a single string and individual chunks (if available).
1130
 */
1131
function decodeTextTokensFromServer(endpoint, ids, resolve) {
1132
    const isAsync = typeof resolve === 'function';
1133
    let text = '';
1134
    let chunks = [];
1135
    jQuery.ajax({
1136
        async: isAsync,
1137
        type: 'POST',
1138
        url: endpoint,
1139
        data: JSON.stringify({ ids: ids }),
1140
        dataType: 'json',
1141
        contentType: 'application/json',
1142
        success: function (data) {
1143
            text = data.text;
1144
            chunks = data.chunks;
1145
            isAsync && resolve({ text, chunks });
1146
        },
1147
    });
1148
    return { text, chunks };
1149
}
1150

1151
/**
1152
 * Encodes a string to tokens using the server API.
1153
 * @param {number} tokenizerType Tokenizer type.
1154
 * @param {string} str String to tokenize.
1155
 * @returns {number[]} Array of token ids.
1156
 */
1157
export function getTextTokens(tokenizerType, str) {
1158
    switch (tokenizerType) {
1159
        case tokenizers.API_CURRENT:
1160
            return getTextTokens(currentRemoteTokenizerAPI(), str);
1161
        case tokenizers.API_TEXTGENERATIONWEBUI:
1162
            return getTextTokensFromTextgenAPI(str);
1163
        case tokenizers.API_KOBOLD:
1164
            return getTextTokensFromKoboldAPI(str);
1165
        default: {
1166
            const tokenizerEndpoints = TOKENIZER_URLS[tokenizerType];
1167
            if (!tokenizerEndpoints) {
1168
                apiFailureTokenCount(str);
1169
                console.warn('Unknown tokenizer type', tokenizerType);
1170
                return [];
1171
            }
1172
            let endpointUrl = tokenizerEndpoints.encode;
1173
            if (!endpointUrl) {
1174
                apiFailureTokenCount(str);
1175
                console.warn('This tokenizer type does not support encoding', tokenizerType);
1176
                return [];
1177
            }
1178
            if (tokenizerType === tokenizers.OPENAI) {
1179
                endpointUrl += `?model=${getTokenizerModel()}`;
1180
            }
1181
            return getTextTokensFromServer(endpointUrl, str);
1182
        }
1183
    }
1184
}
1185

1186
/**
1187
 * Decodes token ids to text using the server API.
1188
 * @param {number} tokenizerType Tokenizer type.
1189
 * @param {number[]} ids Array of token ids
1190
 * @returns {({ text: string, chunks?: string[] })} Decoded token text as a single string and individual chunks (if available).
1191
 */
1192
export function decodeTextTokens(tokenizerType, ids) {
1193
    // Currently, neither remote API can decode, but this may change in the future. Put this guard here to be safe
1194
    if (tokenizerType === tokenizers.API_CURRENT) {
1195
        return decodeTextTokens(tokenizers.NONE, ids);
1196
    }
1197
    const tokenizerEndpoints = TOKENIZER_URLS[tokenizerType];
1198
    if (!tokenizerEndpoints) {
1199
        console.warn('Unknown tokenizer type', tokenizerType);
1200
        return { text: '', chunks: [] };
1201
    }
1202
    let endpointUrl = tokenizerEndpoints.decode;
1203
    if (!endpointUrl) {
1204
        console.warn('This tokenizer type does not support decoding', tokenizerType);
1205
        return { text: '', chunks: [] };
1206
    }
1207
    if (tokenizerType === tokenizers.OPENAI) {
1208
        endpointUrl += `?model=${getTokenizerModel()}`;
1209
    }
1210
    return decodeTextTokensFromServer(endpointUrl, ids);
1211
}
1212

1213
export async function initTokenizers() {
1214
    TEXTGEN_TOKENIZERS.push(
1215
        textgen_types.OOBA,
1216
        textgen_types.TABBY,
1217
        textgen_types.KOBOLDCPP,
1218
        textgen_types.LLAMACPP,
1219
        textgen_types.VLLM,
1220
        textgen_types.APHRODITE,
1221
    );
1222
    eventSource.on(event_types.ONLINE_STATUS_CHANGED, async () => {
1223
        // Clear tokenizer warning when (re)connecting to an LLM backend that supports tokenization
1224
        if (main_api === 'textgenerationwebui' && TEXTGEN_TOKENIZERS.includes(textgen_settings.type)) {
1225
            sessionStorage.removeItem(TOKENIZER_WARNING_KEY);
1226
        }
1227
    });
1228
    await loadTokenCache();
1229
    registerDebugFunction('resetTokenCache', 'Reset token cache', 'Purges the calculated token counts. Use this if you want to force a full re-tokenization of all chats or suspect the token counts are wrong.', resetTokenCache);
1230
}
1231