Add WebLLM extension summarization

8685c2f471bda4784ed52600620bd7019216a293

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

Parents: 77ab694ea0326f381a97057e320723d530aac74e
3 files changed, +145 -23Showing whitespace changes
public/scripts/extensions/memory/index.js+130 -18
@@ -25,6 +25,7 @@ import { SlashCommandParser } from '../../slash-commands/SlashCommandParser.js';
2525import { SlashCommand } from '../../slash-commands/SlashCommand.js';
2626import { ARGUMENT_TYPE, SlashCommandArgument, SlashCommandNamedArgument } from '../../slash-commands/SlashCommandArgument.js';
2727import { MacrosParser } from '../../macros.js';
28+import { countWebLlmTokens, generateWebLlmChatPrompt, getWebLlmContextSize, isWebLlmSupported } from '../shared.js';
2829export { MODULE_NAME };
2930
3031const MODULE_NAME = '1_memory';
@@ -36,6 +37,40 @@ let lastMessageHash = null;
3637let lastMessageId = null;
3738let inApiCall = false;
3839
40+/**
41+ * Count the number of tokens in the provided text.
42+ * @param {string} text Text to count tokens for
43+ * @returns {Promise<number>} Number of tokens in the text
44+ */
45+async function countSourceTokens(text, padding = 0) {
46+ if (extension_settings.memory.source === summary_sources.webllm) {
47+ const count = await countWebLlmTokens(text);
48+ return count + padding;
49+ }
50+
51+ if (extension_settings.memory.source === summary_sources.extras) {
52+ const count = getTextTokens(tokenizers.GPT2, text).length;
53+ return count + padding;
54+ }
55+
56+ return await getTokenCountAsync(text, padding);
57+}
58+
59+async function getSourceContextSize() {
60+ const overrideLength = extension_settings.memory.overrideResponseLength;
61+
62+ if (extension_settings.memory.source === summary_sources.webllm) {
63+ const maxContext = await getWebLlmContextSize();
64+ return overrideLength > 0 ? (maxContext - overrideLength) : Math.round(maxContext * 0.75);
65+ }
66+
67+ if (extension_settings.source === summary_sources.extras) {
68+ return 1024;
69+ }
70+
71+ return getMaxContextSize(overrideLength);
72+}
73+
3974const formatMemoryValue = function (value) {
4075 if (!value) {
4176 return '';
@@ -55,6 +90,7 @@ const saveChatDebounced = debounce(() => getContext().saveChat(), debounce_timeo
5590const summary_sources = {
5691 'extras': 'extras',
5792 'main': 'main',
93+ 'webllm': 'webllm',
5894};
5995
6096const prompt_builders = {
@@ -130,12 +166,12 @@ function loadSettings() {
130166
131167async function onPromptForceWordsAutoClick() {
132168 const context = getContext();
133169 const maxPromptLength = getMaxContextSizeawait getSourceContextSize(extension_settings.memory.overrideResponseLength);
134170 const chat = context.chat;
135171 const allMessages = chat.filter(m => !m.is_system && m.mes).map(m => m.mes);
136172 const messagesWordCount = allMessages.map(m => extractAllWords(m)).flat().length;
137173 const averageMessageWordCount = messagesWordCount / allMessages.length;
138174 const tokensPerWord = await getTokenCountAsynccountSourceTokens(allMessages.join('\n')) / messagesWordCount;
139175 const wordsPerToken = 1 / tokensPerWord;
140176 const maxPromptLengthWords = Math.round(maxPromptLength * wordsPerToken);
141177 // How many words should pass so that messages will start be dropped out of context;
@@ -168,15 +204,15 @@ async function onPromptForceWordsAutoClick() {
168204
169205async function onPromptIntervalAutoClick() {
170206 const context = getContext();
171207 const maxPromptLength = getMaxContextSizeawait getSourceContextSize(extension_settings.memory.overrideResponseLength);
172208 const chat = context.chat;
173209 const allMessages = chat.filter(m => !m.is_system && m.mes).map(m => m.mes);
174210 const messagesWordCount = allMessages.map(m => extractAllWords(m)).flat().length;
175211 const messagesTokenCount = await getTokenCountAsynccountSourceTokens(allMessages.join('\n'));
176212 const tokensPerWord = messagesTokenCount / messagesWordCount;
177213 const averageMessageTokenCount = messagesTokenCount / allMessages.length;
178214 const targetSummaryTokens = Math.round(extension_settings.memory.promptWords * tokensPerWord);
179215 const promptTokens = await getTokenCountAsynccountSourceTokens(extension_settings.memory.prompt);
180216 const promptAllowance = maxPromptLength - promptTokens - targetSummaryTokens;
181217 const maxMessagesPerSummary = extension_settings.memory.maxMessagesPerRequest || 0;
182218 const averageMessagesPerPrompt = Math.floor(promptAllowance / averageMessageTokenCount);
@@ -213,8 +249,8 @@ function onSummarySourceChange(event) {
213249
214250function switchSourceControls(value) {
215251 $('#memory_settings [data-summary-source]').each((_, element) => {
216252 const source = $(element).datadataset.summarySource.split('summary-source,').map(s => s.trim());
217253 $(element).toggle(source === .includes(value));
218254 });
219255}
220256
@@ -359,6 +395,12 @@ async function onChatEvent() {
359395 }
360396 }
361397
398+ if (extension_settings.memory.source === summary_sources.webllm) {
399+ if (!isWebLlmSupported()) {
400+ return;
401+ }
402+ }
403+
362404 const context = getContext();
363405 const chat = context.chat;
364406
@@ -431,8 +473,12 @@ async function forceSummarizeChat() {
431473 return '';
432474 }
433475
434476 const toast = toastr.info('Summarizing chat...', 'Please wait', { timeOut: 0, extendedTimeOut: 0 });
435- const value = await summarizeChatMain(context, true, skipWIAN);
477+ const value = extension_settings.memory.source === summary_sources.main
478+ ? await summarizeChatMain(context, true, skipWIAN)
479+ : await summarizeChatWebLLM(context, true);
480+
481+ toastr.clear(toast);
436482
437483 if (!value) {
438484 toastr.warning('Failed to summarize chat');
@@ -484,16 +530,25 @@ async function summarizeChat(context) {
484530 case summary_sources.main:
485531 await summarizeChatMain(context, false, skipWIAN);
486532 break;
533+ case summary_sources.webllm:
534+ await summarizeChatWebLLM(context, false);
535+ break;
487536 default:
488537 break;
489538 }
490539}
491540
492-async function summarizeChatMain(context, force, skipWIAN) {
541+/**
493-
542+ * Check if the chat should be summarized based on the current conditions.
543+ * Return summary prompt if it should be summarized.
544+ * @param {any} context ST context
545+ * @param {boolean} force Summarize the chat regardless of the conditions
546+ * @returns {Promise<string>} Summary prompt or empty string
547+ */
548+async function getSummaryPromptForNow(context, force) {
494549 if (extension_settings.memory.promptInterval === 0 && !force) {
495550 console.debug('Prompt interval is set to 0, skipping summarization');
496551 return '';
497552 }
498553
499554 try {
@@ -505,17 +560,17 @@ async function summarizeChatMain(context, force, skipWIAN) {
505560 waitUntilCondition(() => is_send_press === false, 30000, 100);
506561 } catch {
507562 console.debug('Timeout waiting for is_send_press');
508563 return '';
509564 }
510565
511566 if (!context.chat.length) {
512567 console.debug('No messages in chat to summarize');
513568 return '';
514569 }
515570
516571 if (context.chat.length < extension_settings.memory.promptInterval && !force) {
517572 console.debug(`Not enough messages in chat to summarize (chat: ${context.chat.length}, interval: ${extension_settings.memory.promptInterval})`);
518573 return '';
519574 }
520575
521576 let messagesSinceLastSummary = 0;
@@ -539,7 +594,7 @@ async function summarizeChatMain(context, force, skipWIAN) {
539594
540595 if (!conditionSatisfied && !force) {
541596 console.debug(`Summary conditions not satisfied (messages: ${messagesSinceLastSummary}, interval: ${extension_settings.memory.promptInterval}, words: ${wordsSinceLastSummary}, force words: ${extension_settings.memory.promptForceWords})`);
542597 return '';
543598 }
544599
545600 console.log('Summarizing chat, messages since last summary: ' + messagesSinceLastSummary, 'words since last summary: ' + wordsSinceLastSummary);
@@ -547,6 +602,63 @@ async function summarizeChatMain(context, force, skipWIAN) {
547602
548603 if (!prompt) {
549604 console.debug('Summarization prompt is empty. Skipping summarization.');
605+ return '';
606+ }
607+
608+ return prompt;
609+}
610+
611+async function summarizeChatWebLLM(context, force) {
612+ if (!isWebLlmSupported()) {
613+ return;
614+ }
615+
616+ const prompt = await getSummaryPromptForNow(context, force);
617+
618+ if (!prompt) {
619+ return;
620+ }
621+
622+ const { rawPrompt, lastUsedIndex } = await getRawSummaryPrompt(context, prompt);
623+
624+ if (lastUsedIndex === null || lastUsedIndex === -1) {
625+ if (force) {
626+ toastr.info('To try again, remove the latest summary.', 'No messages found to summarize');
627+ }
628+
629+ return null;
630+ }
631+
632+ const messages = [
633+ { role: 'system', content: prompt },
634+ { role: 'user', content: rawPrompt },
635+ ];
636+
637+ const params = {};
638+
639+ if (extension_settings.memory.overrideResponseLength > 0) {
640+ params.max_tokens = extension_settings.memory.overrideResponseLength;
641+ }
642+
643+ const summary = await generateWebLlmChatPrompt(messages, params);
644+ const newContext = getContext();
645+
646+ // something changed during summarization request
647+ if (newContext.groupId !== context.groupId ||
648+ newContext.chatId !== context.chatId ||
649+ (!newContext.groupId && (newContext.characterId !== context.characterId))) {
650+ console.log('Context changed, summary discarded');
651+ return;
652+ }
653+
654+ setMemoryContext(summary, true, lastUsedIndex);
655+ return summary;
656+}
657+
658+async function summarizeChatMain(context, force, skipWIAN) {
659+ const prompt = await getSummaryPromptForNow(context, force);
660+
661+ if (!prompt) {
550662 return;
551663 }
552664
@@ -634,7 +746,7 @@ async function getRawSummaryPrompt(context, prompt) {
634746 chat.pop(); // We always exclude the last message from the buffer
635747 const chatBuffer = [];
636748 const PADDING = 64;
637749 const PROMPT_SIZE = getMaxContextSizeawait getSourceContextSize(extension_settings.memory.overrideResponseLength);
638750 let latestUsedMessage = null;
639751
640752 for (let index = latestSummaryIndex + 1; index < chat.length; index++) {
@@ -651,7 +763,7 @@ async function getRawSummaryPrompt(context, prompt) {
651763 const entry = `${message.name}:\n${message.mes}`;
652764 chatBuffer.push(entry);
653765
654766 const tokens = await getTokenCountAsynccountSourceTokens(getMemoryString(true), PADDING);
655767
656768 if (tokens > PROMPT_SIZE) {
657769 chatBuffer.pop();
public/scripts/extensions/memory/settings.html+4 -3
@@ -13,6 +13,7 @@
1313 <select id="summary_source">
1414 <option value="main" data-i18n="ext_sum_main_api">Main API</option>
1515 <option value="extras">Extras API</option>
16+ <option value="webllm" data-i18n="ext_sum_webllm">WebLLM Extension</option>
1617 </select><br>
1718
1819 <div class="flex-container justifyspacebetween alignitemscenter">
@@ -24,7 +25,7 @@
2425
2526 <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>
2627 <div class="memory_contents_controls">
2728 <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">
2829 <i class="fa-solid fa-database"></i>
2930 <span data-i18n="ext_sum_force_text">Summarize now</span>
3031 </div>
@@ -58,7 +59,7 @@
5859 <span data-i18n="ext_sum_prompt_builder_3">Classic, blocking</span>
5960 </label>
6061 </div>
6162 <div data-summary-source="main,webllm">
6263 <label for="memory_prompt" class="title_restorable">
6364 <span data-i18n="Summary Prompt">Summary Prompt</span>
6465 <div id="memory_prompt_restore" data-i18n="[title]ext_sum_restore_default_prompt_tip" title="Restore default prompt" class="right_menu_button">
@@ -74,7 +75,7 @@
7475 </label>
7576 <input id="memory_override_response_length" type="range" value="{{defaultSettings.overrideResponseLength}}" min="{{defaultSettings.overrideResponseLengthMin}}" max="{{defaultSettings.overrideResponseLengthMax}}" step="{{defaultSettings.overrideResponseLengthStep}}" />
7677 <label for="memory_max_messages_per_request">
7778 <span data-i18n="ext_sum_raw_max_msg">[Raw/WebLLM] Max messages per request</span> (<span id="memory_max_messages_per_request_value"></span>)
7879 <small class="memory_disabled_hint" data-i18n="ext_sum_0_unlimited">0 = unlimited</small>
7980 </label>
8081 <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
@@ -183,15 +183,21 @@ function throwIfInvalidModel(useReverseProxy) {
183183 */
184184export function isWebLlmSupported() {
185185 if (!('gpu' in navigator)) {
186+ const warningKey = 'webllm_browser_warning_shown';
187+ if (!sessionStorage.getItem(warningKey)) {
186188 toastr.error('Your browser does not support the WebGPU API. Please use a different browser.', 'WebLLM', {
187189 preventDuplicates: true,
188190 timeOut: 0,
189191 extendedTimeOut: 0,
190192 });
193+ sessionStorage.setItem(warningKey, '1');
194+ }
191195 return false;
192196 }
193197
194198 if (!('llm' in SillyTavern)) {
199+ const warningKey = 'webllm_extension_warning_shown';
200+ if (!sessionStorage.getItem(warningKey)) {
195201 toastr.error('WebLLM extension is not installed. Click here to install it.', 'WebLLM', {
196202 timeOut: 0,
197203 extendedTimeOut: 0,
@@ -209,6 +215,8 @@ export function isWebLlmSupported() {
209215 }
210216 },
211217 });
218+ sessionStorage.setItem(warningKey, '1');
219+ }
212220 return false;
213221 }
214222
@@ -218,15 +226,16 @@ export function isWebLlmSupported() {
218226/**
219227 * Generates text in response to a chat prompt using WebLLM.
220228 * @param {any[]} messages Messages to use for generating
229+ * @param {object} params Additional parameters
221230 * @returns {Promise<string>} Generated response
222231 */
223232export async function generateWebLlmChatPrompt(messages, params = {}) {
224233 if (!isWebLlmSupported()) {
225234 throw new Error('WebLLM extension is not installed.');
226235 }
227236
228237 const engine = SillyTavern.llm;
229238 const response = await engine.generateChatPrompt(messages, params);
230239 return response;
231240}
232241