Merge pull request #2988 from QuantumEntangledAndy/feat/cachedVectorSummaries Add client side cacheing of vector summaries

e01a243ce5ee0abd6479d17a2b6e2f58ebf781f5

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

Signed
1 files changed, +54 -43Showing whitespace changes
public/scripts/extensions/vectors/index.js+54 -43
@@ -36,6 +36,7 @@ import { generateWebLlmChatPrompt, isWebLlmSupported } from '../shared.js';
36/**36/**
37 * @typedef {object} HashedMessage37 * @typedef {object} HashedMessage
38 * @property {string} text - The hashed message text38 * @property {string} text - The hashed message text
39 * @property {number} hash - The hash used as the vector key
39 */40 */
4041
41const MODULE_NAME = 'vectors';42const MODULE_NAME = 'vectors';
@@ -96,6 +97,8 @@ const settings = {
9697
97const moduleWorker = new ModuleWorkerWrapper(synchronizeChat);98const moduleWorker = new ModuleWorkerWrapper(synchronizeChat);
9899
100const cachedSummaries = new Map();
101
99/**102/**
100 * Gets the Collection ID for a file embedded in the chat.103 * Gets the Collection ID for a file embedded in the chat.
101 * @param {string} fileUrl URL of the file104 * @param {string} fileUrl URL of the file
@@ -118,6 +121,10 @@ async function onVectorizeAllClick() {
118 return;121 return;
119 }122 }
120123
124 // Clear all cached summaries to ensure that new ones are created
125 // upon request of a full vectorise
126 cachedSummaries.clear();
127
121 const batchSize = 5;128 const batchSize = 5;
122 const elapsedLog = [];129 const elapsedLog = [];
123 let finished = false;130 let finished = false;
@@ -200,11 +207,10 @@ function splitByChunks(items) {
200207
201/**208/**
202 * Summarizes messages using the Extras API method.209 * Summarizes messages using the Extras API method.
203 * @param {HashedMessage[]} hashedMessages Array of hashed messages210 * @param {HashedMessage} element hashed message
204 * @returns {Promise<HashedMessage[]>} Summarized messages211 * @returns {Promise<boolean>} Sucess
205 */212 */
206async function summarizeExtra(hashedMessages) {213async function summarizeExtra(element) {
207 for (const element of hashedMessages) {
208 try {214 try {
209 const url = new URL(getApiUrl());215 const url = new URL(getApiUrl());
210 url.pathname = '/api/summarize';216 url.pathname = '/api/summarize';
@@ -228,42 +234,37 @@ async function summarizeExtra(hashedMessages) {
228 }234 }
229 catch (error) {235 catch (error) {
230 console.log(error);236 console.log(error);
231 }237 return false;
232 }238 }
233239
234 return hashedMessages;240 return true;
235}241}
236242
237/**243/**
238 * Summarizes messages using the main API method.244 * Summarizes messages using the main API method.
239 * @param {HashedMessage[]} hashedMessages Array of hashed messages245 * @param {HashedMessage} element hashed message
240 * @returns {Promise<HashedMessage[]>} Summarized messages246 * @returns {Promise<boolean>} Sucess
241 */247 */
242async function summarizeMain(hashedMessages) {248async function summarizeMain(element) {
243 for (const element of hashedMessages) {
244 element.text = await generateRaw(element.text, '', false, false, settings.summary_prompt);249 element.text = await generateRaw(element.text, '', false, false, settings.summary_prompt);
245 }250 return true;
246
247 return hashedMessages;
248}251}
249252
250/**253/**
251 * Summarizes messages using WebLLM.254 * Summarizes messages using WebLLM.
252 * @param {HashedMessage[]} hashedMessages Array of hashed messages255 * @param {HashedMessage} element hashed message
253 * @returns {Promise<HashedMessage[]>} Summarized messages256 * @returns {Promise<boolean>} Sucess
254 */257 */
255async function summarizeWebLLM(hashedMessages) {258async function summarizeWebLLM(element) {
256 if (!isWebLlmSupported()) {259 if (!isWebLlmSupported()) {
257 console.warn('Vectors: WebLLM is not supported');260 console.warn('Vectors: WebLLM is not supported');
258 return hashedMessages;261 return false;
259 }262 }
260263
261 for (const element of hashedMessages) {
262 const messages = [{ role: 'system', content: settings.summary_prompt }, { role: 'user', content: element.text }];264 const messages = [{ role: 'system', content: settings.summary_prompt }, { role: 'user', content: element.text }];
263 element.text = await generateWebLlmChatPrompt(messages);265 element.text = await generateWebLlmChatPrompt(messages);
264 }
265266
266 return hashedMessages;267 return true;
267}268}
268269
269/**270/**
@@ -273,16 +274,35 @@ async function summarizeWebLLM(hashedMessages) {
273 * @returns {Promise<HashedMessage[]>} Summarized messages274 * @returns {Promise<HashedMessage[]>} Summarized messages
274 */275 */
275async function summarize(hashedMessages, endpoint = 'main') {276async function summarize(hashedMessages, endpoint = 'main') {
277 for (const element of hashedMessages) {
278 const cachedSummary = cachedSummaries.get(element.hash);
279 if (!cachedSummary) {
280 let success = true;
276 switch (endpoint) {281 switch (endpoint) {
277 case 'main':282 case 'main':
278 return await summarizeMain(hashedMessages);283 success = await summarizeMain(element);
284 break;
279 case 'extras':285 case 'extras':
280 return await summarizeExtra(hashedMessages);286 success = await summarizeExtra(element);
287 break;
281 case 'webllm':288 case 'webllm':
282 return await summarizeWebLLM(hashedMessages);289 success = await summarizeWebLLM(element);
290 break;
283 default:291 default:
284 console.error('Unsupported endpoint', endpoint);292 console.error('Unsupported endpoint', endpoint);
293 success = false;
294 break;
295 }
296 if (success) {
297 cachedSummaries.set(element.hash, element.text);
298 } else {
299 break;
300 }
301 } else {
302 element.text = cachedSummary;
303 }
285 }304 }
305 return hashedMessages;
286}306}
287307
288async function synchronizeChat(batchSize = 5) {308async function synchronizeChat(batchSize = 5) {
@@ -307,16 +327,15 @@ async function synchronizeChat(batchSize = 5) {
307 return -1;327 return -1;
308 }328 }
309329
310 let hashedMessages = context.chat.filter(x => !x.is_system).map(x => ({ text: String(substituteParams(x.mes)), hash: getStringHash(substituteParams(x.mes)), index: context.chat.indexOf(x) }));330 const hashedMessages = context.chat.filter(x => !x.is_system).map(x => ({ text: String(substituteParams(x.mes)), hash: getStringHash(substituteParams(x.mes)), index: context.chat.indexOf(x) }));
311 const hashesInCollection = await getSavedHashes(chatId);331 const hashesInCollection = await getSavedHashes(chatId);
312332
313 if (settings.summarize) {333 let newVectorItems = hashedMessages.filter(x => !hashesInCollection.includes(x.hash));
314 hashedMessages = await summarize(hashedMessages, settings.summary_source);
315 }
316
317 const newVectorItems = hashedMessages.filter(x => !hashesInCollection.includes(x.hash));
318 const deletedHashes = hashesInCollection.filter(x => !hashedMessages.some(y => y.hash === x));334 const deletedHashes = hashesInCollection.filter(x => !hashedMessages.some(y => y.hash === x));
319335
336 if (settings.summarize) {
337 newVectorItems = await summarize(newVectorItems, settings.summary_source);
338 }
320339
321 if (newVectorItems.length > 0) {340 if (newVectorItems.length > 0) {
322 const chunkedBatch = splitByChunks(newVectorItems.slice(0, batchSize));341 const chunkedBatch = splitByChunks(newVectorItems.slice(0, batchSize));
@@ -687,25 +706,17 @@ const onChatEvent = debounce(async () => await moduleWorker.update(), debounce_t
687 * @returns {Promise<string>} Text to query706 * @returns {Promise<string>} Text to query
688 */707 */
689async function getQueryText(chat, initiator) {708async function getQueryText(chat, initiator) {
690 let queryText = '';709 let hashedMessages = chat
691 let i = 0;710 .map(x => ({ text: String(substituteParams(x.mes)), hash: getStringHash(substituteParams(x.mes)) }))
692711 .filter(x => x.text)
693 let hashedMessages = chat.map(x => ({ text: String(substituteParams(x.mes)) }));712 .reverse()
713 .slice(0, settings.query);
694714
695 if (initiator === 'chat' && settings.enabled_chats && settings.summarize && settings.summarize_sent) {715 if (initiator === 'chat' && settings.enabled_chats && settings.summarize && settings.summarize_sent) {
696 hashedMessages = await summarize(hashedMessages, settings.summary_source);716 hashedMessages = await summarize(hashedMessages, settings.summary_source);
697 }717 }
698718
699 for (const message of hashedMessages.slice().reverse()) {719 const queryText = hashedMessages.map(x => x.text).join('\n');
700 if (message.text) {
701 queryText += message.text + '\n';
702 i++;
703 }
704
705 if (i === settings.query) {
706 break;
707 }
708 }
709720
710 return collapseNewlines(queryText).trim();721 return collapseNewlines(queryText).trim();
711}722}