346e77ec · permissionBRICK/SillyTavern

     resetScrollHeight,
     stringFormat,
 } from './utils.js';
 import { countTokensOpenAIcountTokensOpenAIAsync, getTokenizerModel } from './tokenizers.js';
 import { isMobile } from './RossAscends-mods.js';
 import { saveLogprobsForActiveMessage } from './logprobs.js';
 import { SlashCommandParser } from './slash-commands/SlashCommandParser.js';

     // Reserve budget for new chat message
     const newChat = selected_group ? oai_settings.new_group_chat_prompt : oai_settings.new_chat_prompt;
     const newChatMessage = newawait Message.createAsync('system', substituteParams(newChat), 'newMainChat');
     chatCompletion.reserveBudget(newChatMessage);
     // Reserve budget for group nudge
     let groupNudgeMessage = null;
     const noGroupNudgeTypes = ['impersonate'];
     if (selected_group && prompts.has('groupNudge') && !noGroupNudgeTypes.includes(type)) {
         groupNudgeMessage = await Message.fromPromptfromPromptAsync(prompts.get('groupNudge'));
         chatCompletion.reserveBudget(groupNudgeMessage);
+    }

         };
         const continuePrompt = new Prompt(promptObject);
         const preparedPrompt = promptManager.preparePrompt(continuePrompt);
         continueMessage = await Message.fromPromptfromPromptAsync(preparedPrompt);
         chatCompletion.reserveBudget(continueMessage);
+    }
     const lastChatPrompt = messages[messages.length - 1];
     const message = newawait Message.createAsync('user', oai_settings.send_if_empty, 'emptyUserMessageReplacement');
     if (lastChatPrompt && lastChatPrompt.role === 'assistant' && oai_settings.send_if_empty && chatCompletion.canAfford(message)) {
         chatCompletion.insert(message, 'chatHistory');
+    }

         // We do not want to mutate the prompt
         const prompt = new Prompt(chatPrompt);
         prompt.identifier = `chatHistory-${messages.length - index}`;
         const chatMessage = await Message.fromPromptfromPromptAsync(promptManager.preparePrompt(prompt));
         if (promptManager.serviceSettings.names_behavior === character_names_behavior.COMPLETION && prompt.name) {
             const messageName = promptManager.isValidName(prompt.name) ? prompt.name : promptManager.sanitizeName(prompt.name);
             await chatMessage.setName(messageName);
+        }
         if (imageInlining && chatPrompt.image) {

         if (canUseTools && Array.isArray(chatPrompt.invocations)) {
             /** @type {import('./tool-calling.js').ToolInvocation[]} */
             const invocations = chatPrompt.invocations;
             const toolCallMessage = newawait Message.createAsync(chatMessage.role, undefined, 'toolCall-' + chatMessage.identifier);
             const toolResultMessages = await Promise.all(invocations.slice().reverse().map((invocation) => new Message.createAsync('tool', invocation.result || '[No content]', invocation.id)));
             await toolCallMessage.setToolCalls(invocations);
             if (chatCompletion.canAffordAll([toolCallMessage, ...toolResultMessages])) {
                 for (const resultMessage of toolResultMessages) {
                     chatCompletion.insertAtStart(resultMessage, 'chatHistory');

             if (type === 'continue' && oai_settings.continue_prefill && chatPrompt === firstNonInjected) {
                 // in case we are using continue_prefill and the latest message is an assistant message, we want to prepend the users assistant prefill on the message
                 if (chatPrompt.role === 'assistant') {
                     const collectioncontinueMessage = new MessageCollection('continuePrefill', newawait Message.createAsync(chatMessage.role, substituteParams(oai_settings.assistant_prefill + '\n\n') + chatMessage.content, chatMessage.identifier));
 +                    const collection = new MessageCollection('continuePrefill', continueMessage);
                     chatCompletion.add(collection, -1);
                     continue;
+                }

  * @param {ChatCompletion} chatCompletion - An instance of ChatCompletion class that will be populated with the prompts.
  * @param {Object[]} messageExamples - Array containing all message examples.
  */
 async function populateDialogueExamples(prompts, chatCompletion, messageExamples) {
     if (!prompts.has('dialogueExamples')) {
         return;
+    }
     chatCompletion.add(new MessageCollection('dialogueExamples'), prompts.index('dialogueExamples'));
     if (Array.isArray(messageExamples) && messageExamples.length) {
         const newExampleChat = newawait Message.createAsync('system', substituteParams(oai_settings.new_example_chat_prompt), 'newChat');
         for (const dialogue of [...messageExamples].forEach((dialogue, dialogueIndex) => {
             letconst examplesAddeddialogueIndex = 0messageExamples.indexOf(dialogue);
+-
 +            const chatMessages = [];
 -            if (chatCompletion.canAfford(newExampleChat)) chatCompletion.insert(newExampleChat, 'dialogueExamples');
             for (let promptIndex = 0; promptIndex < dialogue.length; promptIndex++) {
                 const prompt = dialogue[promptIndex];

@@ -806,19 +806,20 @@ function populateDialogueExamples(prompts, chatCompletion, messageExamples) {
806	806	const content = prompt.content \|\| '';
807	807	const identifier = `dialogueExamples ${dialogueIndex}-${promptIndex}`;
808	808
809	809	const chatMessage = newawait Message.createAsync(role, content, identifier);
810	810	await chatMessage.setName(prompt.name);
811		- if (!chatCompletion.canAfford(chatMessage)) {
	811	+ chatMessages.push(chatMessage);
	812	+ }
	813	+
	814	+ if (!chatCompletion.canAffordAll([newExampleChat, ...chatMessages])) {
812	815	break;
813	816	}
	817	+
	818	+ chatCompletion.insert(newExampleChat, 'dialogueExamples');
	819	+ for (const chatMessage of chatMessages) {
814	820	chatCompletion.insert(chatMessage, 'dialogueExamples');
815		- examplesAdded++;
816	821	}
817		-
818		- if (0 === examplesAdded) {
819		- chatCompletion.removeLastFrom('dialogueExamples');
820	822	}
821		- });
822	823	}
823	824	}
824	825

  */
 async function populateChatCompletion(prompts, chatCompletion, { bias, quietPrompt, quietImage, type, cyclePrompt, messages, messageExamples }) {
     // Helper function for preparing a prompt, that already exists within the prompt collection, for completion
     const addToChatCompletion = async (source, target = null) => {
         // We need the prompts array to determine a position for the source.
         if (false === prompts.has(source)) return;

         const index = target ? prompts.index(target) : prompts.index(source);
         const collection = new MessageCollection(source);
 -        collection.add(Message.fromPrompt(prompt));
 +        const message = await Message.fromPromptAsync(prompt);
 +        collection.add(message);
         chatCompletion.add(collection, index);
     };
     chatCompletion.reserveBudget(3); // every reply is primed with <|start|>assistant<|message|>
     // Character and world information
     await addToChatCompletion('worldInfoBefore');
     await addToChatCompletion('main');
     await addToChatCompletion('worldInfoAfter');
     await addToChatCompletion('charDescription');
     await addToChatCompletion('charPersonality');
     await addToChatCompletion('scenario');
     await addToChatCompletion('personaDescription');
     // Collection of control prompts that will always be positioned last
     chatCompletion.setOverriddenPrompts(prompts.overriddenPrompts);
     const controlPrompts = new MessageCollection('controlPrompts');
     const impersonateMessage = await Message.fromPromptfromPromptAsync(prompts.get('impersonate')) ?? null;
     if (type === 'impersonate') controlPrompts.add(impersonateMessage);
     // Add quiet prompt to control prompts
     // This should always be last, even in control prompts. Add all further control prompts BEFORE this prompt
     const quietPromptMessage = await Message.fromPromptfromPromptAsync(prompts.get('quietPrompt')) ?? null;
     if (quietPromptMessage && quietPromptMessage.content) {
         if (isImageInliningSupported() && quietImage) {
             await quietPromptMessage.addImage(quietImage);

             return acc;
         }, []);
     for (const identifier of [...systemPrompts, ...userRelativePrompts].forEach(identifier => addToChatCompletion(identifier)); {
 +        await addToChatCompletion(identifier);
 +    }
     // Add enhance definition instruction
     if (prompts.has('enhanceDefinitions')) await addToChatCompletion('enhanceDefinitions');
     // Bias
     if (bias && bias.trim().length) await addToChatCompletion('bias');
     // Tavern Extras - Summary
     if (prompts.has('summary')) {
         const summary = prompts.get('summary');
         if (summary.position) {
 -            chatCompletion.insert(Message.fromPrompt(summary), 'main', summary.position);
 +            const message = await Message.fromPromptAsync(summary);
 +            chatCompletion.insert(message, 'main', summary.position);
+        }
+    }

@@ -962,7 +967,8 @@ async function populateChatCompletion(prompts, chatCompletion, { bias, quietProm
962	967	const authorsNote = prompts.get('authorsNote');
963	968
964	969	if (authorsNote.position) {
965		- chatCompletion.insert(Message.fromPrompt(authorsNote), 'main', authorsNote.position);
	970	+ const message = await Message.fromPromptAsync(authorsNote);
	971	+ chatCompletion.insert(message, 'main', authorsNote.position);
966	972	}
967	973	}
968	974

@@ -971,7 +977,8 @@ async function populateChatCompletion(prompts, chatCompletion, { bias, quietProm
971	977	const vectorsMemory = prompts.get('vectorsMemory');
972	978
973	979	if (vectorsMemory.position) {
974		- chatCompletion.insert(Message.fromPrompt(vectorsMemory), 'main', vectorsMemory.position);
	980	+ const message = await Message.fromPromptAsync(vectorsMemory);
	981	+ chatCompletion.insert(message, 'main', vectorsMemory.position);
975	982	}
976	983	}
977	984

@@ -980,7 +987,8 @@ async function populateChatCompletion(prompts, chatCompletion, { bias, quietProm
980	987	const vectorsDataBank = prompts.get('vectorsDataBank');
981	988
982	989	if (vectorsDataBank.position) {
983		- chatCompletion.insert(Message.fromPrompt(vectorsDataBank), 'main', vectorsDataBank.position);
	990	+ const message = await Message.fromPromptAsync(vectorsDataBank);
	991	+ chatCompletion.insert(message, 'main', vectorsDataBank.position);
984	992	}
985	993	}
986	994

@@ -989,13 +997,15 @@ async function populateChatCompletion(prompts, chatCompletion, { bias, quietProm
989	997	const smartContext = prompts.get('smartContext');
990	998
991	999	if (smartContext.position) {
992		- chatCompletion.insert(Message.fromPrompt(smartContext), 'main', smartContext.position);
	1000	+ const message = await Message.fromPromptAsync(smartContext);
	1001	+ chatCompletion.insert(message, 'main', smartContext.position);
993	1002	}
994	1003	}
995	1004
996	1005	// Other relative extension prompts
997	1006	for (const prompt of prompts.collection.filter(p => p.extension && p.position)) {
998		- chatCompletion.insert(Message.fromPrompt(prompt), 'main', prompt.position);
	1007	+ const message = await Message.fromPromptAsync(prompt);
	1008	+ chatCompletion.insert(message, 'main', prompt.position);
999	1009	}
1000	1010
1001	1011	// Pre-allocation of tokens for tool data

         const toolData = {};
         await ToolManager.registerFunctionToolsOpenAI(toolData);
         const toolMessage = [{ role: 'user', content: JSON.stringify(toolData) }];
         const toolTokens = await tokenHandler.countcountAsync(toolMessage);
         chatCompletion.reserveBudget(toolTokens);
+    }

     // Decide whether dialogue examples should always be added
     if (power_user.pin_examples) {
         await populateDialogueExamples(prompts, chatCompletion, messageExamples);
         await populateChatHistory(messages, prompts, chatCompletion, type, cyclePrompt);
     } else {
         await populateChatHistory(messages, prompts, chatCompletion, type, cyclePrompt);
         await populateDialogueExamples(prompts, chatCompletion, messageExamples);
+    }
     chatCompletion.freeBudget(controlPrompts);

         promptManager.setChatCompletion(chatCompletion);
         if (oai_settings.squash_system_messages && dryRun == false) {
             await chatCompletion.squashSystemMessages();
+        }
         // All information is up-to-date, render.

+}
 class TokenHandler {
 -    constructor(countTokenFn) {
 +    /**
 -        this.countTokenFn = countTokenFn;
 +     * @param {(messages: object[] | object, full?: boolean) => Promise<number>} countTokenAsyncFn Function to count tokens
 +     */
 +    constructor(countTokenAsyncFn) {
 +        this.countTokenAsyncFn = countTokenAsyncFn;
         this.counts = {
             'start_chat': 0,
             'prompt': 0,

         this.counts[type] -= value;
+    }
 -    count(messages, full, type) {
 +    /**
 -        const token_count = this.countTokenFn(messages, full);
 +     * Count tokens for a message or messages.
 +     * @param {object|any[]} messages Messages to count tokens for
 +     * @param {boolean} [full] Count full tokens
 +     * @param {string} [type] Identifier for the token count
 +     * @returns {Promise<number>} The token count
 +     */
 +    async countAsync(messages, full, type) {
 +        const token_count = await this.countTokenAsyncFn(messages, full);
         this.counts[type] += token_count;
         return token_count;

@@ -2178,7 +2198,7 @@ class TokenHandler {
2178	2198	}
2179	2199
2180	2200
2181	2201	const tokenHandler = new TokenHandler(countTokensOpenAIcountTokensOpenAIAsync);
2182	2202
2183	2203	// Thrown by ChatCompletion when a requested prompt couldn't be found.
2184	2204	class IdentifierNotFoundError extends Error {

      * @param {string} role - The role of the entity creating the message.
      * @param {string} content - The actual content of the message.
      * @param {string} identifier - A unique identifier for the message.
 +     * @private Don't use this constructor directly. Use createAsync instead.
      */
     constructor(role, content, identifier) {
         this.identifier = identifier;

             this.role = 'system';
+        }
 -        if (typeof this.content === 'string' && this.content.length > 0) {
 -            this.tokens = tokenHandler.count({ role: this.role, content: this.content });
 -        } else {
         this.tokens = 0;
+    }
++
 +    /**
 +     * Create a new Message instance.
 +     * @param {string} role
 +     * @param {string} content
 +     * @param {string} identifier
 +     * @returns {Promise<Message>} Message instance
 +     */
 +    static async createAsync(role, content, identifier) {
 +        const message = new Message(role, content, identifier);
++
 +        if (typeof message.content === 'string' && message.content.length > 0) {
 +            message.tokens = await tokenHandler.countAsync({ role: message.role, content: message.content });
 +        }
++
 +        return message;
+    }
     /**
      * Reconstruct the message from a tool invocation.
      * @param {import('./tool-calling.js').ToolInvocation[]} invocations - The tool invocations to reconstruct the message from.
 +     * @returns {Promise<void>}
      */
     async setToolCalls(invocations) {
         this.tool_calls = invocations.map(i => ({
             id: i.id,
             type: 'function',

                 name: i.name,
             },
         }));
         this.tokens = await tokenHandler.countcountAsync({ role: this.role, tool_calls: JSON.stringify(this.tool_calls) });
+    }
 -    setName(name) {
 +    /**
 +     * Add a name to the message.
 +     * @param {string} name Name to set for the message.
 +     * @returns {Promise<void>}
 +     */
 +    async setName(name) {
         this.name = name;
         this.tokens = await tokenHandler.countcountAsync({ role: this.role, content: this.content, name: this.name });
+    }
 +    /**
 +     * Adds an image to the message.
 +     * @param {string} image Image URL or Data URL.
 +     * @returns {Promise<void>}
 +     */
     async addImage(image) {
         const textContent = this.content;
         const isDataUrl = isDataURL(image);

+    }
     /**
      * Create a new Message instance from a prompt asynchronously.
      * @static
      * @param {Object} prompt - The prompt object.
      * @returns {Promise<Message>} A new instance of Message.
      */
     static fromPromptfromPromptAsync(prompt) {
         return new Message.createAsync(prompt.role, prompt.content, prompt.identifier);
+    }
     /**

     /**
      * Combines consecutive system messages into one if they have no name attached.
 +     * @returns {Promise<void>}
      */
     async squashSystemMessages() {
         const excludeList = ['newMainChat', 'newChat', 'groupNudge'];
         this.messages.collection = this.messages.flatten();

             if (shouldSquash(message)) {
                 if (lastMessage && shouldSquash(lastMessage)) {
                     lastMessage.content += '\n' + message.content;
                     lastMessage.tokens = await tokenHandler.countcountAsync({ role: lastMessage.role, content: lastMessage.content });
+                }
                 else {
                     squashedMessages.push(message);

      */
     init(moduleConfiguration, serviceSettings) {
         this.configuration = Object.assign(this.configuration, moduleConfiguration);
         this.tokenHandler = this.tokenHandler || new TokenHandler(() => { throw new Error('Token handler not set'); });
         this.serviceSettings = serviceSettings;
         this.containerElement = document.getElementById(this.configuration.containerIdentifier);

Merge pull request #2976 from SillyTavern/cc-async-count Chat Completion: switch to async token handling