Handle various error scenarios during token counting

public/css/promptmanager.css

@@ -112,6 +112,13 @@
     margin-left: 0.5em;
 }
 
+#openai_prompt_manager .openai_prompt_manager_error {
+    padding: 1em;
+    border: 3px solid var(--fullred);
+    margin-top: 1em;
+    margin-bottom: 0.5em;
+}
+
 #openai_prompt_manager .openai_prompt_manager_header {
     display: flex;
     flex-direction: row;

public/scripts/PromptManager.js

@@ -37,7 +37,7 @@ class PromptCollection {
 
     get(identifier) {
         const index = this.index(identifier);
-        if (0 > index) throw new IdentifierNotFoundError(identifier);
+        if (0 > index) return null;
         return this.collection[index];
     }
 
@@ -63,8 +63,8 @@ function PromptManagerModule() {
     this.containerElement = null;
     this.listElement = null;
     this.activeCharacter = null;
-
     this.tokenHandler = null;
+    this.error = null;
 
     this.handleToggle = () => { };
     this.handleEdit = () => { };
@@ -95,7 +95,9 @@ PromptManagerModule.prototype.init = function (moduleConfiguration, serviceSetti
     this.handleToggle = (event) => {
         const promptID = event.target.closest('.' + this.configuration.prefix + 'prompt_manager_prompt').dataset.pmIdentifier;
         const promptListEntry = this.getPromptListEntry(this.activeCharacter, promptID);
+        const counts = this.tokenHandler.getCounts();
 
+        counts[promptID] = null;
         promptListEntry.enabled = !promptListEntry.enabled;
         this.saveServiceSettings().then(() => this.render());
     };
@@ -191,17 +193,25 @@ PromptManagerModule.prototype.init = function (moduleConfiguration, serviceSetti
         this.saveServiceSettings().then(() => this.render());
     });
 
+    document.getElementById('openai_max_context').addEventListener('change', (event) => {
+        if (this.activeCharacter) this.render();
+    });
+
+    document.getElementById('openai_max_tokens').addEventListener('change', (event) => {
+        if (this.activeCharacter) this.render();
+    });
+
     // Prepare prompt edit form save and close button.
     document.getElementById(this.configuration.prefix + 'prompt_manager_popup_entry_form_save').addEventListener('click', this.handleSavePrompt);
     document.getElementById(this.configuration.prefix + 'prompt_manager_popup_entry_form_close').addEventListener('click', () => {
         this.hideEditForm();
         this.clearEditForm();
     });
-
 };
 
 PromptManagerModule.prototype.render = function () {
     if (null === this.activeCharacter) return;
+    this.error = null;
     this.tryGenerate().then(() => {
         this.renderPromptManager();
         this.renderPromptManagerListItems()
@@ -464,6 +474,8 @@ PromptManagerModule.prototype.getPromptIndexById = function (identifier) {
 PromptManagerModule.prototype.preparePrompt = function (prompt) {
     const groupMembers = this.getActiveGroupCharacters();
     if (0 < groupMembers.length) return {role: prompt.role || 'system', content: substituteParams(prompt.content ?? '', null, null, groupMembers.join(', '))}
+    prompt.content = substituteParams(prompt.content);
+
     return new Prompt(prompt);
 }
 
@@ -519,6 +531,13 @@ PromptManagerModule.prototype.getPromptCollection = function () {
     return promptCollection
 }
 
+PromptManagerModule.prototype.populateTokenHandler = function(messageCollection) {
+    const counts = this.tokenHandler.getCounts();
+    messageCollection.getCollection().forEach((message) => {
+        counts[message.identifier] = message.getTokens();
+    });
+}
+
 // Empties, then re-assembles the container containing the prompt list.
 PromptManagerModule.prototype.renderPromptManager = function () {
     const promptManagerDiv = this.containerElement;
@@ -526,6 +545,13 @@ PromptManagerModule.prototype.renderPromptManager = function () {
 
     const showAdvancedSettings = this.serviceSettings.prompt_manager_settings.showAdvancedSettings;
     const checkSpanClass = showAdvancedSettings ? 'fa-solid fa-toggle-on' : 'fa-solid fa-toggle-off';
+
+    const errorDiv = `
+            <div class="${this.configuration.prefix}prompt_manager_error">
+                <span class="fa-solid tooltip fa-triangle-exclamation text_danger"></span> ${this.error}
+            </div>
+    `;
+    const activeTokenInfo = `<span class="tooltip fa-solid fa-info-circle" title="Including tokens from hidden prompts"></span>`;
     const totalActiveTokens = this.tokenHandler?.getTotal();
 
     promptManagerDiv.insertAdjacentHTML('beforeend', `
@@ -536,12 +562,13 @@ PromptManagerModule.prototype.renderPromptManager = function () {
             </a>
         </div>
         <div class="range-block">
+            ${this.error ? errorDiv : ''}
             <div class="${this.configuration.prefix}prompt_manager_header">
                 <div class="${this.configuration.prefix}prompt_manager_header_advanced">
                     <span class="${checkSpanClass}"></span>
                     <span class="checkbox_label" data-i18n="Show advanced options">Show advanced options</span>
                 </div>
-                <div>Total Tokens: ${totalActiveTokens}</div>
+                <div>Total Tokens: ${totalActiveTokens} ${ showAdvancedSettings ? '' : activeTokenInfo} </div>
             </div>
             <ul id="${this.configuration.prefix}prompt_manager_list" class="text_pole"></ul>
         </div>
@@ -614,7 +641,21 @@ PromptManagerModule.prototype.renderPromptManagerListItems = function () {
         const enabledClass = listEntry.enabled ? '' : `${prefix}prompt_manager_prompt_disabled`;
         const draggableClass = draggableEnabled ? 'draggable' : prompt.marker ? 'droppable' : '';
         const markerClass = prompt.marker ? `${prefix}prompt_manager_marker` : '';
-        const calculatedTokens = this.tokenHandler?.getCounts()[prompt.identifier] ?? '';
+        const tokens = this.tokenHandler?.getCounts()[prompt.identifier] ?? 0;
+
+        let warningClass = '';
+        let warningTitle = '';
+        if ('chatHistory' === prompt.identifier) {
+            if (500 >= tokens) {
+                warningClass = 'fa-solid tooltip fa-triangle-exclamation text_danger';
+                warningTitle = 'Very little of your chat history is being sent, consider deactivating some other prompts.';
+            } else if (1000 >= tokens) {
+                warningClass = 'fa-solid tooltip fa-triangle-exclamation text_warning';
+                warningTitle = 'Only a few messages worth chat history are being sent.';
+            }
+        }
+
+        const calculatedTokens = tokens ? tokens : '-';
 
         let detachSpanHtml = '';
         if (this.isPromptDeletionAllowed(prompt)) {
@@ -638,7 +679,7 @@ PromptManagerModule.prototype.renderPromptManagerListItems = function () {
                         </span>
                     </span>
                 `}
-                <span class="prompt_manager_prompt_tokens" data-pm-tokens="${calculatedTokens}">${calculatedTokens}</span>
+                <span class="prompt_manager_prompt_tokens" data-pm-tokens="${calculatedTokens}"><span class="${warningClass}" title="${warningTitle}"> </span>${calculatedTokens}</span>
             </li>
         `;
     });
@@ -751,18 +792,6 @@ const openAiDefaultPrompts = {
             "system_prompt": true,
             "marker": true,
         },
-        {
-            "identifier": "newMainChat",
-            "name": "Start Chat",
-            "system_prompt": true,
-            "marker": true,
-        },
-        {
-            "identifier": "newExampleChat",
-            "name": "Start Chat",
-            "system_prompt": true,
-            "marker": true,
-        },
         {
             "identifier": "worldInfoAfter",
             "name": "World Info (after)",
@@ -843,18 +872,10 @@ const openAiDefaultPromptList = [
         "identifier": "worldInfoAfter",
         "enabled": true
     },
-    {
-        "identifier": "newExampleChat",
-        "enabled": true
-    },
     {
         "identifier": "dialogueExamples",
         "enabled": true
     },
-    {
-        "identifier": "newMainChat",
-        "enabled": true
-    },
     {
         "identifier": "chatHistory",
         "enabled": true

public/scripts/openai.js

@@ -367,6 +367,108 @@ function formatWorldInfo(value) {
     return stringFormat(oai_settings.wi_format, value);
 }
 
+/**
+ * Populate a chat conversation by adding prompts to the conversation and managing system and user prompts.
+ *
+ * @param {Map} prompts - Map object containing all prompts where the key is the prompt identifier and the value is the prompt object.
+ * @param {ChatCompletion} chatCompletion - An instance of ChatCompletion class that will be populated with the prompts.
+ * @param {Object} options - An object with optional settings.
+ * @param {string} options.bias - A bias to be added in the conversation.
+ * @param {string} options.quietPrompt - A quiet prompt to be used in the conversation.
+ * @param {string} options.type - The type of the chat, can be 'impersonate'.
+ */
+function populateChatCompletion (prompts, chatCompletion, {bias, quietPrompt, type} = {}) {
+    const addToChatCompletion = (source, target = null) => {
+        if (false === prompts.has(source)) return;
+
+        const prompt = prompts.get(source);
+        const index = target ? prompts.index(target) : prompts.index(source);
+        const collection = new MessageCollection(source);
+        collection.addItem(Message.fromPrompt(prompt));
+        chatCompletion.add(collection, index);
+    };
+
+    addToChatCompletion('worldInfoBefore');
+    addToChatCompletion('worldInfoAfter');
+    addToChatCompletion('charDescription');
+    addToChatCompletion('charPersonality');
+    addToChatCompletion('scenario');
+
+    // Add main prompt
+    if (type === "impersonate") addToChatCompletion('impersonate', 'main');
+    else addToChatCompletion('main');
+
+    // Add managed system and user prompts
+    const systemPrompts = ['nsfw', 'jailbreak'];
+    const userPrompts = prompts.collection
+        .filter((prompt) => false === prompt.system_prompt)
+        .reduce((acc, prompt) => {
+            acc.push(prompt.identifier)
+            return acc;
+        }, []);
+
+    [...systemPrompts, ...userPrompts].forEach(identifier => addToChatCompletion(identifier));
+
+    // Add enhance definition instruction
+    if (prompts.has('scenario')) addToChatCompletion('enhanceDefinitions');
+
+    // Insert nsfw avoidance prompt into main, if no nsfw prompt is present
+    if (false === chatCompletion.has('nsfw') && oai_settings.nsfw_avoidance_prompt)
+        if (prompts.has('nsfwAvoidance')) chatCompletion.insert(prompts.get('nsfwAvoidance'), 'main');
+
+    // Insert quiet prompt into main
+    if (quietPrompt) {
+        const quietPromptMessage = Message.fromPrompt(prompts.get('quietPrompt'));
+        chatCompletion.insert(quietPromptMessage, 'main');
+    }
+
+    if (bias && bias.trim().length) addToChatCompletion('bias');
+
+    if (prompts.has('summary')) chatCompletion.insert(Message.fromPrompt(prompts.get('summary')), 'main');
+
+    if (prompts.has('authorsNote')) {
+        const authorsNote = Message.fromPrompt(prompts.get('authorsNote'));
+        if (extension_prompt_types.AFTER_SCENARIO) chatCompletion.insert(authorsNote, 'scenario');
+        else chatCompletion.insert(authorsNote, 'main')
+    }
+
+    // Chat History
+    chatCompletion.add( new MessageCollection('chatHistory'), prompts.index('chatHistory'));
+    const mainChat = selected_group ? '[Start a new group chat. Group members: ${names}]' : '[Start a new Chat]';
+    const mainChatMessage = new Message('system', mainChat, 'newMainChat');
+    // Insert chat messages
+    if (chatCompletion.canAfford(mainChatMessage)) {
+        chatCompletion.insert(mainChatMessage, 'chatHistory');
+
+        [...openai_msgs].forEach((prompt, index) => {
+            const chatMessage = new Message(prompt.role, prompt.content, 'chatHistory-' + index);
+            if (chatCompletion.canAfford(chatMessage)) {
+                chatCompletion.insert(chatMessage, 'chatHistory');
+            }
+        });
+    }
+
+    chatCompletion.add( new MessageCollection('dialogueExamples'), prompts.index('dialogueExamples'));
+    if (openai_msgs_example.length) {
+        // Insert chat message examples if there's enough budget if there is enough budget left for at least one example.
+        const dialogueExampleChat = new Message('system', '[Start a new Chat]', 'newChat');
+        const prompt = openai_msgs_example[0];
+        const dialogueExample = new Message(prompt[0].role || '', prompt[0].content || '', 'dialogueExampleTest');
+
+        if (chatCompletion.canAfford(dialogueExampleChat) &&
+            chatCompletion.canAfford(dialogueExample)) {
+            chatCompletion.insert(dialogueExampleChat, 'dialogueExamples');
+
+            [...openai_msgs_example].forEach((prompt, index) => {
+                const chatMessage = new Message(prompt[0].role || '', prompt[0].content || '', 'dialogueExamples-' + index);
+                if (chatCompletion.canAfford(chatMessage)) {
+                    chatCompletion.insert(chatMessage, 'dialogueExamples');
+                }
+            });
+        }
+    }
+}
+
 /**
  * Take a configuration object and prepares messages for a chat with OpenAI's chat completion API.
  * Handles prompts, prepares chat history, manages token budget, and processes various user settings.
@@ -384,8 +486,6 @@ function formatWorldInfo(value) {
  * @param {string} options.quietPrompt - The quiet prompt to be used in the conversation.
  * @param {Array} options.extensionPrompts - An array of additional prompts.
  * @returns {Promise<Array>} An array where the first element is the prepared chat and the second element is a boolean flag.
- * @throws {TokenBudgetExceededError} If the token budget is exceeded.
- * @throws {IdentifierNotFoundError} If a specific identifier is not found in the message collection.
  */
 async function prepareOpenAIMessages({
                                          name2,
@@ -400,118 +500,61 @@ async function prepareOpenAIMessages({
                                          extensionPrompts,
                                          cyclePrompt
                                      } = {}) {
-
     const prompts = promptManager.getPromptCollection();
     const chatCompletion = new ChatCompletion();
-
+    chatCompletion.setTokenBudget(oai_settings.openai_max_context - oai_settings.openai_max_tokens);
-    chatCompletion.tokenBudget = promptManager.serviceSettings.openai_max_context - promptManager.serviceSettings.amount_gen;
 
     if (power_user.console_log_prompts) chatCompletion.enableLogging();
 
-    // Helper functions
+    // Merge items to send that are managed by the prompt manager with items from other places in silly tavern
-    const addToChatCompletion = (role, content, identifier = null) => {
+    // While the position in this array matters for positioning items inside the chat completion, further elements
-        const collection = new MessageCollection(identifier)
+    // may be added for later reference, as long as the initial order is not altered.
-        if (role && content) collection.addItem(new Message(role, content, identifier));
+    const mappedPrompts = [
-        const index = identifier ? prompts.index(identifier) : null;
+        // Ordered prompts for which a marker should exist
-        chatCompletion.add(collection, index);
+        {role: 'system', content: formatWorldInfo(worldInfoBefore), identifier: 'worldInfoBefore'},
-    };
+        {role: 'system', content: formatWorldInfo(worldInfoAfter), identifier: 'worldInfoAfter'},
+        {role: 'system', content: charDescription, identifier: 'charDescription'},
+        {role: 'system', content: `${name2}'s personality: ${charPersonality}`, identifier: 'charPersonality'},
+        {role: 'system', content: `Circumstances and context of the dialogue: ${Scenario}`, identifier: 'scenario'},
+        // Unordered prompts without marker
+        {role: 'system', content: oai_settings.nsfw_avoidance_prompt, identifier: 'nsfwAvoidance'},
+        {role: 'system', content: oai_settings.impersonation_prompt, identifier: 'impersonate'},
+        {role: 'system', content: quietPrompt, identifier: 'quietPrompt'},
+        {role: 'system', content: bias, identifier: 'bias'}
+    ];
 
-    addToChatCompletion('system', formatWorldInfo(worldInfoBefore), 'worldInfoBefore');
+    // Tavern Extras - Summary
-    addToChatCompletion('system', formatWorldInfo(worldInfoAfter), 'worldInfoAfter');
+    const summary = extensionPrompts['1_memory'];
-    addToChatCompletion('system', substituteParams(charDescription), 'charDescription');
+    if (summary) mappedPrompts.push({role: 'system', content: summary.content, identifier: 'summary'});
-    addToChatCompletion('system', `${name2}'s personality: ${substituteParams(charPersonality)}`, 'charPersonality');
-    addToChatCompletion('system', `Circumstances and context of the dialogue: ${substituteParams(Scenario)}`, 'scenario');
 
-    // Add main prompt
+    // Authors Note
-    if (type === "impersonate") {
+    const authorsNote = extensionPrompts['2_floating_prompt'];
-        const impersonate = substituteParams(oai_settings.impersonation_prompt);
+    if (authorsNote) mappedPrompts.push({role: 'system', content: authorsNote.content, identifier: 'authorsNote'});
-        addToChatCompletion('system', impersonate, 'main');
+
+    mappedPrompts.forEach((prompt) => {
+        const newPrompt = promptManager.preparePrompt(prompt);
+        const markerIndex = prompts.index(prompt.identifier);
+
+        if (-1 !== markerIndex) prompts.collection[markerIndex] = newPrompt;
+        else prompts.add(newPrompt);
+    });
+
+    // Allow subscribers to manipulate the prompts object
+    await eventSource.emit(event_types.OAI_BEFORE_CHATCOMPLETION, prompts);
+
+    try {
+        populateChatCompletion(prompts, chatCompletion, {bias, quietPrompt, type});
+    } catch (error) {
+        if (error instanceof TokenBudgetExceededError) {
+            chatCompletion.log('Token budget exceeded.');
+            promptManager.error = 'Not enough free tokens for mandatory prompts. Raise your token Limit or disable custom prompts.';
         } else {
-        addToChatCompletion('system', prompts.get('main').content, 'main');
+            chatCompletion.log('Unexpected error:');
-    }
+            chatCompletion.log(error);
-
-    // Add managed system and user prompts
-    const systemPrompts = ['nsfw', 'jailbreak'];
-    const userPrompts = prompts.collection
-        .filter((prompt) => false === prompt.system_prompt)
-        .reduce((acc, prompt) => {
-            acc.push(prompt.identifier)
-            return acc;
-        }, []);
-
-    [...systemPrompts, ...userPrompts].forEach(identifier => {
-        if (prompts.has(identifier)) {
-            const prompt = prompts.get(identifier);
-            addToChatCompletion(prompt.role, prompt.content, identifier);
-        }
-    });
-
-    // Add enhance definition instruction
-    if (prompts.has('enhanceDefinitions')) {
-        const prompt = prompts.get('enhanceDefinitions');
-        addToChatCompletion(prompt.role, prompt.content, identifier);
-     }
-
-    // Insert nsfw avoidance prompt into main, if no nsfw prompt is present
-    if (false === chatCompletion.has('nsfw') && oai_settings.nsfw_avoidance_prompt) {
-        const nsfwAvoidanceMessage = new Message('system', oai_settings.nsfw_avoidance_prompt, 'nsfwAvoidance');
-        chatCompletion.insert(nsfwAvoidanceMessage, 'main');
-    }
-
-    // Insert quiet prompt into main
-    if (quietPrompt) {
-        const quietPromptMessage = new Message('system', quietPrompt, 'quietPrompt');
-        chatCompletion.insert(quietPromptMessage, 'main');
-    }
-
-    if (bias && bias.trim().length) {
-        addToChatCompletion('system', bias, 'main');
-    }
-
-    // Add extension prompts
-    if (0 < extensionPrompts.length) {
-        const summary = extensionPrompts['1_memory'] ?? null;
-        if (summary) {
-            const summaryMessage = new Message('system', summary.content, 'authorsNote');
-            chatCompletion.insert(summaryMessage, 'main')
-        }
-
-        const authorsNote = extensionPrompts['2_floating_prompt'] ?? null;
-        if (authorsNote) {
-            const authorsNoteMessage = new Message('system', authorsNote.content, 'authorsNote');
-            if (extension_prompt_types.AFTER_SCENARIO) chatCompletion.insert(authorsNoteMessage, 'scenario')
-            else chatCompletion.insert(authorsNoteMessage, 'main')
         }
     }
 
-    // Chat History
+    promptManager.populateTokenHandler(chatCompletion.getMessages());
-    addToChatCompletion(null, null, 'chatHistory');
-    const mainChat = selected_group ? '[Start a new group chat. Group members: ${names}]' : '[Start a new Chat]';
-    const mainChatMessage = new Message('system', mainChat, 'newMainChat');
-
-    // Insert chat messages
-    if (chatCompletion.canAfford(mainChatMessage)) {
-        chatCompletion.insert(mainChatMessage, 'chatHistory');
-        [...openai_msgs].reverse().forEach((prompt, index) => {
-            const chatMessage = new Message(prompt.role, prompt.content, 'chatHistory-' + index);
-            if (chatCompletion.canAfford(chatMessage)) {
-                chatCompletion.insert(chatMessage, 'chatHistory');
-            }
-        });
-    }
-
-    // Insert chat message examples if there's enough budget
-    addToChatCompletion(null, null, 'dialogueExamples');
-    if (chatCompletion.canAfford(mainChatMessage)) {
-        // Insert dialogue examples messages
-        chatCompletion.insert(mainChatMessage, 'dialogueExamples');
-        [...openai_msgs_example].forEach((prompt, index) => {
-            const chatMessage = new Message(prompt[0].role || '', prompt[0].content || '', 'dialogueExamples-' + index);
-            if (chatCompletion.canAfford(chatMessage)) {
-                chatCompletion.insert(chatMessage, 'dialogueExamples');
-            }
-        });
-    }
 
     const chat = chatCompletion.getChat();
     openai_messages_count = chat.filter(x => x.role === "user" || x.role === "assistant").length;
@@ -996,7 +1039,7 @@ class TokenHandler {
     }
 
     getTotal() {
-        return Object.values(this.counts).reduce((a, b) => a + b);
+        return Object.values(this.counts).reduce((a, b) => a + (isNaN(b) ? 0 : b), 0);
     }
 
     log() {
@@ -1084,6 +1127,10 @@ class Message {
         this.tokens = tokenHandler.count(this);
     }
 
+    static fromPrompt(prompt) {
+        return new Message(prompt.role, prompt.content, prompt.identifier);
+    }
+
     getTokens() {return this.tokens};
 }
 
@@ -1101,6 +1148,10 @@ class MessageCollection  {
         this.identifier = identifier;
     }
 
+    getCollection() {
+        return this.collection;
+    }
+
     addItem(item) {
         this.collection.push(item);
     }
@@ -1123,6 +1174,15 @@ class ChatCompletion {
         this.loggingEnabled = false;
     }
 
+    getMessages() {
+        return this.messages;
+    }
+
+    setTokenBudget(tokenBudget) {
+        this.tokenBudget = tokenBudget;
+        console.log(`Token budget: ${this.tokenBudget}`);
+    }
+
     add(collection, position = null) {
         this.validateMessageCollection(collection);
         this.checkTokenBudget(collection, collection.identifier);
@@ -1147,7 +1207,7 @@ class ChatCompletion {
         if (message.content) {
             this.messages.collection[index].collection.push(message);
             this.decreaseTokenBudgetBy(message.getTokens());
-            this.log(`Added ${message.identifier}. Remaining tokens: ${this.tokenBudget}`);
+            this.log(`Inserted ${message.identifier} into ${identifier}. Remaining tokens: ${this.tokenBudget}`);
         }
     }