Vector Storage: summarize with WebLLM extension

public/scripts/extensions/vectors/index.js

@@ -30,6 +30,12 @@ import { textgen_types, textgenerationwebui_settings } from '../../textgen-setti
 import { SlashCommandParser } from '../../slash-commands/SlashCommandParser.js';
 import { SlashCommand } from '../../slash-commands/SlashCommand.js';
 import { ARGUMENT_TYPE, SlashCommandArgument, SlashCommandNamedArgument } from '../../slash-commands/SlashCommandArgument.js';
+import { generateWebLlmChatPrompt, isWebLlmSupported } from '../shared.js';
+
+/**
+ * @typedef {object} HashedMessage
+ * @property {string} text - The hashed message text
+ */
 
 const MODULE_NAME = 'vectors';
 
@@ -191,6 +197,11 @@ function splitByChunks(items) {
     return chunkedItems;
 }
 
+/**
+ * Summarizes messages using the Extras API method.
+ * @param {HashedMessage[]} hashedMessages Array of hashed messages
+ * @returns {Promise<HashedMessage[]>} Summarized messages
+ */
 async function summarizeExtra(hashedMessages) {
     for (const element of hashedMessages) {
         try {
@@ -222,6 +233,11 @@ async function summarizeExtra(hashedMessages) {
     return hashedMessages;
 }
 
+/**
+ * Summarizes messages using the main API method.
+ * @param {HashedMessage[]} hashedMessages Array of hashed messages
+ * @returns {Promise<HashedMessage[]>} Summarized messages
+ */
 async function summarizeMain(hashedMessages) {
     for (const element of hashedMessages) {
         element.text = await generateRaw(element.text, '', false, false, settings.summary_prompt);
@@ -230,12 +246,39 @@ async function summarizeMain(hashedMessages) {
     return hashedMessages;
 }
 
+/**
+ * Summarizes messages using WebLLM.
+ * @param {HashedMessage[]} hashedMessages Array of hashed messages
+ * @returns {Promise<HashedMessage[]>} Summarized messages
+ */
+async function summarizeWebLLM(hashedMessages) {
+    if (!isWebLlmSupported()) {
+        console.warn('Vectors: WebLLM is not supported');
+        return hashedMessages;
+    }
+
+    for (const element of hashedMessages) {
+        const messages = [{ role:'system', content: settings.summary_prompt }, { role:'user', content: element.text }];
+        element.text = await generateWebLlmChatPrompt(messages);
+    }
+
+    return hashedMessages;
+}
+
+/**
+ * Summarizes messages using the chosen method.
+ * @param {HashedMessage[]} hashedMessages Array of hashed messages
+ * @param {string} endpoint Type of endpoint to use
+ * @returns {Promise<HashedMessage[]>} Summarized messages
+ */
 async function summarize(hashedMessages, endpoint = 'main') {
     switch (endpoint) {
         case 'main':
             return await summarizeMain(hashedMessages);
         case 'extras':
             return await summarizeExtra(hashedMessages);
+        case 'webllm':
+            return await summarizeWebLLM(hashedMessages);
         default:
             console.error('Unsupported endpoint', endpoint);
     }

public/scripts/extensions/vectors/settings.html

@@ -378,10 +378,11 @@
                         <select id="vectors_summary_source" class="text_pole">
                             <option value="main" data-i18n="Main API">Main API</option>
                             <option value="extras" data-i18n="Extras API">Extras API</option>
+                            <option value="webllm" data-i18n="WebLLM Extension">WebLLM Extension</option>
                         </select>
 
                         <label for="vectors_summary_prompt" title="Summary Prompt:">Summary Prompt:</label>
-                        <small data-i18n="Only used when Main API is selected.">Only used when Main API is selected.</small>
+                        <small data-i18n="Only used when Main API or WebLLM Extension is selected.">Only used when Main API or WebLLM Extension is selected.</small>
                         <textarea id="vectors_summary_prompt" class="text_pole textarea_compact" rows="6" placeholder="This prompt will be sent to AI to request the summary generation."></textarea>
                     </div>
                 </div>