Add WebLLM source to /summarize command

30c19eb8fdaae54fd2d6db31469f91b2a1bbcbcb

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

public/scripts/extensions/memory/index.js+12 -9 Viewed

 /**
  * Count the number of tokens in the provided text.
  * @param {string} text Text to count tokens for
+ * @param {number} padding Number of additional tokens to add to the count
  * @returns {Promise<number>} Number of tokens in the text
  */
 async function countSourceTokens(text, padding = 0) {

 async function onChatEvent() {
     // Module not enabled
-    if (extension_settings.memory.source === summary_sources.extras) {
+    if (extension_settings.memory.source === summary_sources.extras && !modules.includes('summarize')) {
-        if (!modules.includes('summarize')) {
+        return;
-            return;
-        }
+    }
-    if (extension_settings.memory.source === summary_sources.webllm) {
+    // WebLLM is not supported
-        if (!isWebLlmSupported()) {
+    if (extension_settings.memory.source === summary_sources.webllm && !isWebLlmSupported()) {
-            return;
+        return;
-        }
+    }
     const context = getContext();

                 return await callExtrasSummarizeAPI(text);
             case summary_sources.main:
                 return await generateRaw(text, '', false, false, prompt, extension_settings.memory.overrideResponseLength);
+            case summary_sources.webllm: {
+                const messages = [{ role: 'system', content: prompt }, { role: 'user', content: text }].filter(m => m.content);
+                const params = extension_settings.memory.overrideResponseLength > 0 ? { max_tokens: extension_settings.memory.overrideResponseLength } : {};
+                return await generateWebLlmChatPrompt(messages, params);
+            }
             default:
                 toastr.warning('Invalid summarization source specified');
                 return '';

         name: 'summarize',
         callback: summarizeCallback,
         namedArgumentList: [
-            new SlashCommandNamedArgument('source', 'API to use for summarization', [ARGUMENT_TYPE.STRING], false, false, '', ['main', 'extras']),
+            new SlashCommandNamedArgument('source', 'API to use for summarization', [ARGUMENT_TYPE.STRING], false, false, '', Object.values(summary_sources)),
             SlashCommandNamedArgument.fromProps({
                 name: 'prompt',
                 description: 'prompt to use for summarization',