Add WebLLM extension summarization

8685c2f471bda4784ed52600620bd7019216a293

Cohee <18619528+Cohee1207@users.noreply.github.com>

Parents: 77ab694ea0326f381a97057e320723d530aac74e

public/scripts/extensions/memory/index.js+130 -18 Viewed

 import { SlashCommand } from '../../slash-commands/SlashCommand.js';
 import { ARGUMENT_TYPE, SlashCommandArgument, SlashCommandNamedArgument } from '../../slash-commands/SlashCommandArgument.js';
 import { MacrosParser } from '../../macros.js';
+import { countWebLlmTokens, generateWebLlmChatPrompt, getWebLlmContextSize, isWebLlmSupported } from '../shared.js';
 export { MODULE_NAME };
 const MODULE_NAME = '1_memory';

 let lastMessageId = null;
 let inApiCall = false;
+/**
+ * Count the number of tokens in the provided text.
+ * @param {string} text Text to count tokens for
+ * @returns {Promise<number>} Number of tokens in the text
+ */
+async function countSourceTokens(text, padding = 0) {
+    if (extension_settings.memory.source === summary_sources.webllm) {
+        const count = await countWebLlmTokens(text);
+        return count + padding;
+    }
+    if (extension_settings.memory.source === summary_sources.extras) {
+        const count = getTextTokens(tokenizers.GPT2, text).length;
+        return count + padding;
+    }
+    return await getTokenCountAsync(text, padding);
+}
+async function getSourceContextSize() {
+    const overrideLength = extension_settings.memory.overrideResponseLength;
+    if (extension_settings.memory.source === summary_sources.webllm) {
+        const maxContext = await getWebLlmContextSize();
+        return overrideLength > 0 ? (maxContext - overrideLength) : Math.round(maxContext * 0.75);
+    }
+    if (extension_settings.source === summary_sources.extras) {
+        return 1024;
+    }
+    return getMaxContextSize(overrideLength);
+}
 const formatMemoryValue = function (value) {
     if (!value) {
         return '';

 const summary_sources = {
     'extras': 'extras',
     'main': 'main',
+    'webllm': 'webllm',
 };
 const prompt_builders = {

 async function onPromptForceWordsAutoClick() {
     const context = getContext();
-    const maxPromptLength = getMaxContextSize(extension_settings.memory.overrideResponseLength);
+    const maxPromptLength = await getSourceContextSize();
     const chat = context.chat;
     const allMessages = chat.filter(m => !m.is_system && m.mes).map(m => m.mes);
     const messagesWordCount = allMessages.map(m => extractAllWords(m)).flat().length;
     const averageMessageWordCount = messagesWordCount / allMessages.length;
-    const tokensPerWord = await getTokenCountAsync(allMessages.join('\n')) / messagesWordCount;
+    const tokensPerWord = await countSourceTokens(allMessages.join('\n')) / messagesWordCount;
     const wordsPerToken = 1 / tokensPerWord;
     const maxPromptLengthWords = Math.round(maxPromptLength * wordsPerToken);
     // How many words should pass so that messages will start be dropped out of context;

 async function onPromptIntervalAutoClick() {
     const context = getContext();
-    const maxPromptLength = getMaxContextSize(extension_settings.memory.overrideResponseLength);
+    const maxPromptLength = await getSourceContextSize();
     const chat = context.chat;
     const allMessages = chat.filter(m => !m.is_system && m.mes).map(m => m.mes);
     const messagesWordCount = allMessages.map(m => extractAllWords(m)).flat().length;
-    const messagesTokenCount = await getTokenCountAsync(allMessages.join('\n'));
+    const messagesTokenCount = await countSourceTokens(allMessages.join('\n'));
     const tokensPerWord = messagesTokenCount / messagesWordCount;
     const averageMessageTokenCount = messagesTokenCount / allMessages.length;
     const targetSummaryTokens = Math.round(extension_settings.memory.promptWords * tokensPerWord);
-    const promptTokens = await getTokenCountAsync(extension_settings.memory.prompt);
+    const promptTokens = await countSourceTokens(extension_settings.memory.prompt);
     const promptAllowance = maxPromptLength - promptTokens - targetSummaryTokens;
     const maxMessagesPerSummary = extension_settings.memory.maxMessagesPerRequest || 0;
     const averageMessagesPerPrompt = Math.floor(promptAllowance / averageMessageTokenCount);

 function switchSourceControls(value) {
     $('#memory_settings [data-summary-source]').each((_, element) => {
-        const source = $(element).data('summary-source');
+        const source = element.dataset.summarySource.split(',').map(s => s.trim());
-        $(element).toggle(source === value);
+        $(element).toggle(source.includes(value));
     });
+}

@@ -359,6 +395,12 @@ async function onChatEvent() {
359	}	395	}
360	}	396	}
361		397
		398	if (extension_settings.memory.source === summary_sources.webllm) {
		399	if (!isWebLlmSupported()) {
		400	return;
		401	}
		402	}
		403
362	const context = getContext();	404	const context = getContext();
363	const chat = context.chat;	405	const chat = context.chat;
364		406

         return '';
+    }
-    toastr.info('Summarizing chat...', 'Please wait');
+    const toast = toastr.info('Summarizing chat...', 'Please wait', { timeOut: 0, extendedTimeOut: 0 });
-    const value = await summarizeChatMain(context, true, skipWIAN);
+    const value = extension_settings.memory.source === summary_sources.main
+        ? await summarizeChatMain(context, true, skipWIAN)
+        : await summarizeChatWebLLM(context, true);
+    toastr.clear(toast);
     if (!value) {
         toastr.warning('Failed to summarize chat');

         case summary_sources.main:
             await summarizeChatMain(context, false, skipWIAN);
             break;
+        case summary_sources.webllm:
+            await summarizeChatWebLLM(context, false);
+            break;
         default:
             break;
+    }
+}
-async function summarizeChatMain(context, force, skipWIAN) {
+/**
+ * Check if the chat should be summarized based on the current conditions.
+ * Return summary prompt if it should be summarized.
+ * @param {any} context ST context
+ * @param {boolean} force Summarize the chat regardless of the conditions
+ * @returns {Promise<string>} Summary prompt or empty string
+ */
+async function getSummaryPromptForNow(context, force) {
     if (extension_settings.memory.promptInterval === 0 && !force) {
         console.debug('Prompt interval is set to 0, skipping summarization');
-        return;
+        return '';
+    }
     try {

@@ -505,17 +560,17 @@ async function summarizeChatMain(context, force, skipWIAN) {
505	waitUntilCondition(() => is_send_press === false, 30000, 100);	560	waitUntilCondition(() => is_send_press === false, 30000, 100);
506	} catch {	561	} catch {
507	console.debug('Timeout waiting for is_send_press');	562	console.debug('Timeout waiting for is_send_press');
508	return;	563	return '';
509	}	564	}
510		565
511	if (!context.chat.length) {	566	if (!context.chat.length) {
512	console.debug('No messages in chat to summarize');	567	console.debug('No messages in chat to summarize');
513	return;	568	return '';
514	}	569	}
515		570
516	if (context.chat.length < extension_settings.memory.promptInterval && !force) {	571	if (context.chat.length < extension_settings.memory.promptInterval && !force) {
517	console.debug(`Not enough messages in chat to summarize (chat: ${context.chat.length}, interval: ${extension_settings.memory.promptInterval})`);	572	console.debug(`Not enough messages in chat to summarize (chat: ${context.chat.length}, interval: ${extension_settings.memory.promptInterval})`);
518	return;	573	return '';
519	}	574	}
520		575
521	let messagesSinceLastSummary = 0;	576	let messagesSinceLastSummary = 0;

     if (!conditionSatisfied && !force) {
         console.debug(`Summary conditions not satisfied (messages: ${messagesSinceLastSummary}, interval: ${extension_settings.memory.promptInterval}, words: ${wordsSinceLastSummary}, force words: ${extension_settings.memory.promptForceWords})`);
-        return;
+        return '';
+    }
     console.log('Summarizing chat, messages since last summary: ' + messagesSinceLastSummary, 'words since last summary: ' + wordsSinceLastSummary);

     if (!prompt) {
         console.debug('Summarization prompt is empty. Skipping summarization.');
+        return '';
+    }
+    return prompt;
+}
+async function summarizeChatWebLLM(context, force) {
+    if (!isWebLlmSupported()) {
+        return;
+    }
+    const prompt = await getSummaryPromptForNow(context, force);
+    if (!prompt) {
+        return;
+    }
+    const { rawPrompt, lastUsedIndex } = await getRawSummaryPrompt(context, prompt);
+    if (lastUsedIndex === null || lastUsedIndex === -1) {
+        if (force) {
+            toastr.info('To try again, remove the latest summary.', 'No messages found to summarize');
+        }
+        return null;
+    }
+    const messages = [
+        { role: 'system', content: prompt },
+        { role: 'user', content: rawPrompt },
+    ];
+    const params = {};
+    if (extension_settings.memory.overrideResponseLength > 0) {
+        params.max_tokens = extension_settings.memory.overrideResponseLength;
+    }
+    const summary = await generateWebLlmChatPrompt(messages, params);
+    const newContext = getContext();
+    // something changed during summarization request
+    if (newContext.groupId !== context.groupId ||
+        newContext.chatId !== context.chatId ||
+        (!newContext.groupId && (newContext.characterId !== context.characterId))) {
+        console.log('Context changed, summary discarded');
+        return;
+    }
+    setMemoryContext(summary, true, lastUsedIndex);
+    return summary;
+}
+async function summarizeChatMain(context, force, skipWIAN) {
+    const prompt = await getSummaryPromptForNow(context, force);
+    if (!prompt) {
         return;
+    }

     chat.pop(); // We always exclude the last message from the buffer
     const chatBuffer = [];
     const PADDING = 64;
-    const PROMPT_SIZE = getMaxContextSize(extension_settings.memory.overrideResponseLength);
+    const PROMPT_SIZE = await getSourceContextSize();
     let latestUsedMessage = null;
     for (let index = latestSummaryIndex + 1; index < chat.length; index++) {

         const entry = `${message.name}:\n${message.mes}`;
         chatBuffer.push(entry);
-        const tokens = await getTokenCountAsync(getMemoryString(true), PADDING);
+        const tokens = await countSourceTokens(getMemoryString(true), PADDING);
         if (tokens > PROMPT_SIZE) {
             chatBuffer.pop();

public/scripts/extensions/memory/settings.html+4 -3 Viewed

                 <select id="summary_source">
                     <option value="main" data-i18n="ext_sum_main_api">Main API</option>
                     <option value="extras">Extras API</option>
+                    <option value="webllm" data-i18n="ext_sum_webllm">WebLLM Extension</option>
                 </select><br>
                 <div class="flex-container justifyspacebetween alignitemscenter">

                 <textarea id="memory_contents" class="text_pole textarea_compact" rows="6" data-i18n="[placeholder]ext_sum_memory_placeholder" placeholder="Summary will be generated here..."></textarea>
                 <div class="memory_contents_controls">
-                    <div id="memory_force_summarize" data-summary-source="main" class="menu_button menu_button_icon" title="Trigger a summary update right now." data-i18n="[title]ext_sum_force_tip">
+                    <div id="memory_force_summarize" data-summary-source="main,webllm" class="menu_button menu_button_icon" title="Trigger a summary update right now." data-i18n="[title]ext_sum_force_tip">
                         <i class="fa-solid fa-database"></i>
                         <span data-i18n="ext_sum_force_text">Summarize now</span>
                     </div>

                             <span data-i18n="ext_sum_prompt_builder_3">Classic, blocking</span>
                         </label>
                     </div>
-                    <div data-summary-source="main">
+                    <div data-summary-source="main,webllm">
                         <label for="memory_prompt" class="title_restorable">
                             <span data-i18n="Summary Prompt">Summary Prompt</span>
                             <div id="memory_prompt_restore" data-i18n="[title]ext_sum_restore_default_prompt_tip" title="Restore default prompt" class="right_menu_button">

                         </label>
                         <input id="memory_override_response_length" type="range" value="{{defaultSettings.overrideResponseLength}}" min="{{defaultSettings.overrideResponseLengthMin}}" max="{{defaultSettings.overrideResponseLengthMax}}" step="{{defaultSettings.overrideResponseLengthStep}}" />
                         <label for="memory_max_messages_per_request">
-                            <span data-i18n="ext_sum_raw_max_msg">[Raw] Max messages per request</span> (<span id="memory_max_messages_per_request_value"></span>)
+                            <span data-i18n="ext_sum_raw_max_msg">[Raw/WebLLM] Max messages per request</span> (<span id="memory_max_messages_per_request_value"></span>)
                             <small class="memory_disabled_hint" data-i18n="ext_sum_0_unlimited">0 = unlimited</small>
                         </label>
                         <input id="memory_max_messages_per_request" type="range" value="{{defaultSettings.maxMessagesPerRequest}}" min="{{defaultSettings.maxMessagesPerRequestMin}}" max="{{defaultSettings.maxMessagesPerRequestMax}}" step="{{defaultSettings.maxMessagesPerRequestStep}}" />

public/scripts/extensions/shared.js+11 -2 Viewed

  */
 export function isWebLlmSupported() {
     if (!('gpu' in navigator)) {
+        const warningKey = 'webllm_browser_warning_shown';
+        if (!sessionStorage.getItem(warningKey)) {
             toastr.error('Your browser does not support the WebGPU API. Please use a different browser.', 'WebLLM', {
                 preventDuplicates: true,
                 timeOut: 0,
                 extendedTimeOut: 0,
             });
+            sessionStorage.setItem(warningKey, '1');
+        }
         return false;
+    }
     if (!('llm' in SillyTavern)) {
+        const warningKey = 'webllm_extension_warning_shown';
+        if (!sessionStorage.getItem(warningKey)) {
             toastr.error('WebLLM extension is not installed. Click here to install it.', 'WebLLM', {
                 timeOut: 0,
                 extendedTimeOut: 0,

@@ -209,6 +215,8 @@ export function isWebLlmSupported() {
209	}	215	}
210	},	216	},
211	});	217	});
		218	sessionStorage.setItem(warningKey, '1');
		219	}
212	return false;	220	return false;
213	}	221	}
214		222

 /**
  * Generates text in response to a chat prompt using WebLLM.
  * @param {any[]} messages Messages to use for generating
+ * @param {object} params Additional parameters
  * @returns {Promise<string>} Generated response
  */
-export async function generateWebLlmChatPrompt(messages) {
+export async function generateWebLlmChatPrompt(messages, params = {}) {
     if (!isWebLlmSupported()) {
         throw new Error('WebLLM extension is not installed.');
+    }
     const engine = SillyTavern.llm;
-    const response = await engine.generateChatPrompt(messages);
+    const response = await engine.generateChatPrompt(messages, params);
     return response;
+}