Streamline token counting

By using TokenHandler instead of custom token handling
This commit is contained in:
maver
2023-06-10 18:13:59 +02:00
parent b8d08161ab
commit 4d8d4cd262
2 changed files with 58 additions and 100 deletions

View File

@ -21,6 +21,7 @@ import {
replaceBiasMarkup,
is_send_press,
saveSettings,
Generate,
main_api,
} from "../script.js";
import {groups, selected_group} from "./group-chats.js";
@ -64,7 +65,8 @@ export {
sendOpenAIRequest,
setOpenAIOnlineStatus,
getChatCompletionModel,
countTokens
countTokens,
TokenHandler
}
let openai_msgs = [];
@ -266,7 +268,7 @@ function setOpenAIMessageExamples(mesExamplesArray) {
}
}
function setupOpenAIPromptManager(settings) {
function setupOpenAIPromptManager(openAiSettings) {
promptManager = new PromptManager();
const configuration = {
prefix: 'openai_',
@ -279,7 +281,13 @@ function setupOpenAIPromptManager(settings) {
saveSettingsDebounced();
}
promptManager.init(configuration, settings);
promptManager.tryGenerate = () => {
return Generate('normal', {}, true);
}
promptManager.tokenHandler = tokenHandler;
promptManager.init(configuration, openAiSettings);
promptManager.render();
@ -415,17 +423,16 @@ async function prepareOpenAIMessages({ systemPrompt, name2, storyString, worldIn
// Handle impersonation
if (type === "impersonate") chatCompletion.replace('main', chatCompletion.makeSystemMessage(substituteParams(oai_settings.impersonation_prompt)));
promptManager.updatePrompts(chatCompletion.getPromptsWithTokenCount());
const tokenHandler = promptManager.getTokenHandler();
tokenHandler?.setCounts(
{...tokenHandler.getCounts(), ...chatCompletion.getTokenCounts()}
);
// Save settings with updated token calculation and return context
return promptManager.saveServiceSettings().then(() => {
promptManager.render();
const openai_msgs_tosend = chatCompletion.getChat();
openai_messages_count = openai_msgs_tosend.filter(x => x.role === "user" || x.role === "assistant").length;
console.log(openai_msgs_tosend);
return [openai_msgs_tosend, false];
});
}
@ -944,6 +951,14 @@ class TokenHandler {
};
}
getCounts() {
return this.counts;
}
setCounts(counts) {
this.counts = counts;
}
uncount(value, type) {
this.counts[type] -= value;
}
@ -962,6 +977,8 @@ class TokenHandler {
}
}
const tokenHandler = new TokenHandler(countTokens);
function countTokens(messages, full = false) {
let chatId = 'undefined';