#569 Use main API as summary source

public/scripts/extensions/memory/index.js

@@ -1,10 +1,9 @@
-import { getStringHash, debounce } from "../../utils.js";
+import { getStringHash, debounce, waitUntilCondition } from "../../utils.js";
-import { getContext, getApiUrl, extension_settings, ModuleWorkerWrapper, doExtrasFetch } from "../../extensions.js";
+import { getContext, getApiUrl, extension_settings, doExtrasFetch, modules } from "../../extensions.js";
-import { extension_prompt_types, is_send_press, saveSettingsDebounced } from "../../../script.js";
+import { eventSource, event_types, extension_prompt_types, generateQuietPrompt, is_send_press, saveSettingsDebounced, substituteParams } from "../../../script.js";
 export { MODULE_NAME };
 
 const MODULE_NAME = '1_memory';
-const UPDATE_INTERVAL = 5000;
 
 let lastCharacterId = null;
 let lastGroupId = null;
@@ -13,9 +12,16 @@ let lastMessageHash = null;
 let lastMessageId = null;
 let inApiCall = false;
 
-const formatMemoryValue = (value) => value ? `Context: ${value.trim()}` : '';
+const formatMemoryValue = (value) => value ? `Summary: ${value.trim()}` : '';
 const saveChatDebounced = debounce(() => getContext().saveChat(), 2000);
 
+const summary_sources = {
+    'extras': 'extras',
+    'main': 'main',
+};
+
+const defaultPrompt = '[Pause your roleplay. Summarize the most important facts and events that have happened in the chat so far. If a summary already exists in your memory, use that as a base and expand with new facts. Limit the summary to {{words}} words or less. Your response should include nothing but the summary.]';
+
 const defaultSettings = {
     minLongMemory: 16,
     maxLongMemory: 1024,
@@ -38,6 +44,16 @@ const defaultSettings = {
     maxLengthPenalty: 4,
     lengthPenaltyStep: 0.1,
     memoryFrozen: false,
+    source: summary_sources.extras,
+    prompt: defaultPrompt,
+    promptWords: 200,
+    promptMinWords: 25,
+    promptMaxWords: 1000,
+    promptWordsStep: 25,
+    promptInterval: 10,
+    promptMinInterval: 1,
+    promptMaxInterval: 100,
+    promptIntervalStep: 1,
 };
 
 function loadSettings() {
@@ -45,12 +61,42 @@ function loadSettings() {
         Object.assign(extension_settings.memory, defaultSettings);
     }
 
+    if (extension_settings.memory.source === undefined) {
+        extension_settings.memory.source = defaultSettings.source;
+    }
+
+    if (extension_settings.memory.prompt === undefined) {
+        extension_settings.memory.prompt = defaultSettings.prompt;
+    }
+
+    if (extension_settings.memory.promptWords === undefined) {
+        extension_settings.memory.promptWords = defaultSettings.promptWords;
+    }
+
+    if (extension_settings.memory.promptInterval === undefined) {
+        extension_settings.memory.promptInterval = defaultSettings.promptInterval;
+    }
+
+    $('#summary_source').val(extension_settings.memory.source).trigger('change');
     $('#memory_long_length').val(extension_settings.memory.longMemoryLength).trigger('input');
     $('#memory_short_length').val(extension_settings.memory.shortMemoryLength).trigger('input');
     $('#memory_repetition_penalty').val(extension_settings.memory.repetitionPenalty).trigger('input');
     $('#memory_temperature').val(extension_settings.memory.temperature).trigger('input');
     $('#memory_length_penalty').val(extension_settings.memory.lengthPenalty).trigger('input');
     $('#memory_frozen').prop('checked', extension_settings.memory.memoryFrozen).trigger('input');
+    $('#memory_prompt').val(extension_settings.memory.prompt).trigger('input');
+    $('#memory_prompt_words').val(extension_settings.memory.promptWords).trigger('input');
+    $('#memory_prompt_interval').val(extension_settings.memory.promptInterval).trigger('input');
+}
+
+function onSummarySourceChange(event) {
+    const value = event.target.value;
+    extension_settings.memory.source = value;
+    $('#memory_settings [data-source]').each((_, element) => {
+        const source = $(element).data('source');
+        $(element).toggle(source === value);
+    });
+    saveSettingsDebounced();
 }
 
 function onMemoryShortInput() {
@@ -104,6 +150,26 @@ function onMemoryFrozenInput() {
     saveSettingsDebounced();
 }
 
+function onMemoryPromptWordsInput() {
+    const value = $(this).val();
+    extension_settings.memory.promptWords = Number(value);
+    $('#memory_prompt_words_value').text(extension_settings.memory.promptWords);
+    saveSettingsDebounced();
+}
+
+function onMemoryPromptIntervalInput() {
+    const value = $(this).val();
+    extension_settings.memory.promptInterval = Number(value);
+    $('#memory_prompt_interval_value').text(extension_settings.memory.promptInterval);
+    saveSettingsDebounced();
+}
+
+function onMemoryPromptInput() {
+    const value = $(this).val();
+    extension_settings.memory.prompt = value;
+    saveSettingsDebounced();
+}
+
 function saveLastValues() {
     const context = getContext();
     lastGroupId = context.groupId;
@@ -129,7 +195,14 @@ function getLatestMemoryFromChat(chat) {
     return '';
 }
 
-async function moduleWorker() {
+async function onChatEvent() {
+    // Module not enabled
+    if (extension_settings.memory.source === summary_sources.extras) {
+        if (!modules.includes('summarize')) {
+            return;
+        }
+    }
+
     const context = getContext();
     const chat = context.chat;
 
@@ -187,7 +260,93 @@ async function moduleWorker() {
     }
 }
 
+async function forceSummarizeChat() {
+    const context = getContext();
+
+    if (!context.chatId) {
+        toastr.warning('No chat selected');
+        return;
+    }
+
+    toastr.info('Summarizing chat...', 'Please wait');
+    const value = await summarizeChatMain(context, true);
+
+    if (!value) {
+        toastr.warning('Failed to summarize chat');
+        return;
+    }
+}
+
 async function summarizeChat(context) {
+    switch (extension_settings.memory.source) {
+        case summary_sources.extras:
+            await summarizeChatExtras(context);
+            break;
+        case summary_sources.main:
+            await summarizeChatMain(context, false);
+            break;
+        default:
+            break;
+    }
+}
+
+async function summarizeChatMain(context, force) {
+    try {
+        // Wait for the send button to be released
+        waitUntilCondition(() => is_send_press === false, 10000, 100);
+    } catch {
+        console.debug('Timeout waiting for is_send_press');
+        return;
+    }
+
+    if (!context.chat.length) {
+        console.debug('No messages in chat to summarize');
+        return;
+    }
+
+    if (context.chat.length < extension_settings.memory.promptInterval && !force) {
+        console.debug(`Not enough messages in chat to summarize (chat: ${context.chat.length}, interval: ${extension_settings.memory.promptInterval})`);
+        return;
+    }
+
+    let messagesSinceLastSummary = 0;
+    for (let i = context.chat.length - 1; i >= 0; i--) {
+        if (context.chat[i].extra && context.chat[i].extra.memory) {
+            break;
+        }
+        messagesSinceLastSummary++;
+    }
+
+    if (messagesSinceLastSummary < extension_settings.memory.promptInterval && !force) {
+        console.debug(`Not enough messages since last summary (messages: ${messagesSinceLastSummary}, interval: ${extension_settings.memory.promptInterval}`);
+        return;
+    }
+
+    console.log('Summarizing chat, messages since last summary: ' + messagesSinceLastSummary);
+    const prompt = substituteParams(extension_settings.memory.prompt)
+        .replace(/{{words}}/gi, extension_settings.memory.promptWords);
+
+    if (!prompt) {
+        console.debug('Summarization prompt is empty. Skipping summarization.');
+        return;
+    }
+
+    const summary = await generateQuietPrompt(prompt);
+    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);
+    return summary;
+}
+
+async function summarizeChatExtras(context) {
     function getMemoryString() {
         return (longMemory + '\n\n' + memoryBuffer.slice().reverse().join('\n\n')).trim();
     }
@@ -301,6 +460,7 @@ function setMemoryContext(value, saveToMessage) {
     const context = getContext();
     context.setExtensionPrompt(MODULE_NAME, formatMemoryValue(value), extension_prompt_types.AFTER_SCENARIO);
     $('#memory_contents').val(value);
+    console.log('Memory set to: ' + value);
 
     if (saveToMessage && context.chat.length) {
         const idx = context.chat.length - 2;
@@ -315,7 +475,7 @@ function setMemoryContext(value, saveToMessage) {
     }
 }
 
-$(document).ready(function () {
+jQuery(function () {
     function addExtensionControls() {
         const settingsHtml = `
         <div id="memory_settings">
@@ -325,20 +485,34 @@ $(document).ready(function () {
                     <div class="inline-drawer-icon fa-solid fa-circle-chevron-down down"></div>
                 </div>
                 <div class="inline-drawer-content">
+                    <label for="summary_source">Summarization source:</label>
+                    <select id="summary_source">
+                        <option value="main">Main API</option>
+                        <option value="extras">Extras API</option>
+                    </select>
                     <label for="memory_contents">Current summary: </label>
-                <textarea id="memory_contents" class="text_pole" rows="8" placeholder="Context will be generated here..."></textarea>
+                    <textarea id="memory_contents" class="text_pole textarea_compact" rows="6" placeholder="Summary will be generated here..."></textarea>
                     <div class="memory_contents_controls">
-                    <input id="memory_restore" class="menu_button" type="submit" value="Restore previous state" />
+                        <input id="memory_restore" class="menu_button" type="button" value="Restore previous state" />
-                    <label for="memory_frozen"><input id="memory_frozen" type="checkbox" />Stop summarization updates</label>
+                        <label for="memory_frozen"><input id="memory_frozen" type="checkbox" />Pause summarization</label>
                     </div>
-            <!--</div>
+                    <div data-source="main" class="memory_contents_controls">
                     </div>
-            <div class="inline-drawer">
+                    <div data-source="main">
-                <div class="inline-drawer-toggle inline-drawer-header">
+                        <label for="memory_prompt" class="title_restorable">
-                    <b>Summarization parameters</b>
+                            Summarization Prompt
-                    <div class="inline-drawer-icon fa-solid fa-circle-chevron-down down"></div>
+                            <div id="memory_force_summarize" class="menu_button menu_button_icon">
+                                <i class="fa-solid fa-database"></i>
+                                <span>Generate now</span>
                             </div>
-                <div class="inline-drawer-content">-->
+                        </label>
+                        <textarea id="memory_prompt" class="text_pole textarea_compact" rows="6" placeholder="This prompt will be used in summary generation. Insert {{words}} macro to use the "Number of words" parameter."></textarea>
+                        <label for="memory_prompt_words">Number of words in the summary (<span id="memory_prompt_words_value"></span> words)</label>
+                        <input id="memory_prompt_words" type="range" value="${defaultSettings.promptWords}" min="${defaultSettings.promptMinWords}" max="${defaultSettings.promptMaxWords}" step="${defaultSettings.promptWordsStep}" />
+                        <label for="memory_prompt_interval">Update interval (<span id="memory_prompt_interval_value"></span> messages)</label>
+                        <input id="memory_prompt_interval" type="range" value="${defaultSettings.promptInterval}" min="${defaultSettings.promptMinInterval}" max="${defaultSettings.promptMaxInterval}" step="${defaultSettings.promptIntervalStep}" />
+                    </div>
+                    <div data-source="extras">
                         <label for="memory_short_length">Chat to Summarize buffer length (<span id="memory_short_length_tokens"></span> tokens)</label>
                         <input id="memory_short_length" type="range" value="${defaultSettings.shortMemoryLength}" min="${defaultSettings.minShortMemory}" max="${defaultSettings.maxShortMemory}" step="${defaultSettings.shortMemoryStep}" />
                         <label for="memory_long_length">Summary output length (<span id="memory_long_length_tokens"></span> tokens)</label>
@@ -352,6 +526,7 @@ $(document).ready(function () {
                     </div>
                 </div>
             </div>
+        </div>
         `;
         $('#extensions_settings2').append(settingsHtml);
         $('#memory_restore').on('click', onMemoryRestoreClick);
@@ -362,10 +537,18 @@ $(document).ready(function () {
         $('#memory_temperature').on('input', onMemoryTemperatureInput);
         $('#memory_length_penalty').on('input', onMemoryLengthPenaltyInput);
         $('#memory_frozen').on('input', onMemoryFrozenInput);
+        $('#summary_source').on('change', onSummarySourceChange);
+        $('#memory_prompt_words').on('input', onMemoryPromptWordsInput);
+        $('#memory_prompt_interval').on('input', onMemoryPromptIntervalInput);
+        $('#memory_prompt').on('input', onMemoryPromptInput);
+        $('#memory_force_summarize').on('click', forceSummarizeChat);
     }
 
     addExtensionControls();
     loadSettings();
-    const wrapper = new ModuleWorkerWrapper(moduleWorker);
+    eventSource.on(event_types.MESSAGE_RECEIVED, onChatEvent);
-    setInterval(wrapper.update.bind(wrapper), UPDATE_INTERVAL);
+    eventSource.on(event_types.MESSAGE_DELETED, onChatEvent);
+    eventSource.on(event_types.MESSAGE_EDITED, onChatEvent);
+    eventSource.on(event_types.MESSAGE_SWIPED, onChatEvent);
+    eventSource.on(event_types.CHAT_CHANGED, onChatEvent);
 });

public/scripts/extensions/memory/manifest.json

@@ -1,10 +1,10 @@
 {
     "display_name": "Memory",
     "loading_order": 9,
-    "requires": [
+    "requires": [],
+    "optional": [
         "summarize"
     ],
-    "optional": [],
     "js": "index.js",
     "css": "style.css",
     "author": "Cohee#1207",

public/scripts/utils.js

@@ -651,3 +651,20 @@ export function createThumbnail(dataUrl, maxWidth, maxHeight) {
         };
     });
 }
+
+export async function waitUntilCondition(condition, timeout = 1000, interval = 100) {
+    return new Promise((resolve, reject) => {
+        const timeoutId = setTimeout(() => {
+            clearInterval(intervalId);
+            reject(new Error('Timed out waiting for condition to be true'));
+        }, timeout);
+
+        const intervalId = setInterval(() => {
+            if (condition()) {
+                clearTimeout(timeoutId);
+                clearInterval(intervalId);
+                resolve();
+            }
+        }, interval);
+    });
+}