Merge pull request #2988 from QuantumEntangledAndy/feat/cachedVectorSummaries Add client side cacheing of vector summaries
Signed| @@ -36,6 +36,7 @@ import { generateWebLlmChatPrompt, isWebLlmSupported } from '../shared.js'; | |||
| 36 | /** | 36 | /** |
| 37 | * @typedef {object} HashedMessage | 37 | * @typedef {object} HashedMessage |
| 38 | * @property {string} text - The hashed message text | 38 | * @property {string} text - The hashed message text |
| 39 | * @property {number} hash - The hash used as the vector key | ||
| 39 | */ | 40 | */ |
| 40 | 41 | ||
| 41 | const MODULE_NAME = 'vectors'; | 42 | const MODULE_NAME = 'vectors'; |
| @@ -96,6 +97,8 @@ const settings = { | |||
| 96 | 97 | ||
| 97 | const moduleWorker = new ModuleWorkerWrapper(synchronizeChat); | 98 | const moduleWorker = new ModuleWorkerWrapper(synchronizeChat); |
| 98 | 99 | ||
| 100 | const 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 file | 104 | * @param {string} fileUrl URL of the file |
| @@ -118,6 +121,10 @@ async function onVectorizeAllClick() { | |||
| 118 | return; | 121 | return; |
| 119 | } | 122 | } |
| 120 | 123 | ||
| 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) { | |||
| 200 | 207 | ||
| 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 messages | 210 | * @param {HashedMessage} element hashed message |
| 204 | * @returns {Promise<HashedMessage[]>} Summarized messages | 211 | * @returns {Promise<boolean>} Sucess |
| 205 | */ | 212 | */ |
| 206 | async function summarizeExtra(hashedMessages) { | 213 | async 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 | } |
| 233 | 239 | ||
| 234 | return hashedMessages; | 240 | return true; |
| 235 | } | 241 | } |
| 236 | 242 | ||
| 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 messages | 245 | * @param {HashedMessage} element hashed message |
| 240 | * @returns {Promise<HashedMessage[]>} Summarized messages | 246 | * @returns {Promise<boolean>} Sucess |
| 241 | */ | 247 | */ |
| 242 | async function summarizeMain(hashedMessages) { | 248 | async 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 | } |
| 249 | 252 | ||
| 250 | /** | 253 | /** |
| 251 | * Summarizes messages using WebLLM. | 254 | * Summarizes messages using WebLLM. |
| 252 | * @param {HashedMessage[]} hashedMessages Array of hashed messages | 255 | * @param {HashedMessage} element hashed message |
| 253 | * @returns {Promise<HashedMessage[]>} Summarized messages | 256 | * @returns {Promise<boolean>} Sucess |
| 254 | */ | 257 | */ |
| 255 | async function summarizeWebLLM(hashedMessages) { | 258 | async 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 | } |
| 260 | 263 | ||
| 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 | } | ||
| 265 | 266 | ||
| 266 | return hashedMessages; | 267 | return true; |
| 267 | } | 268 | } |
| 268 | 269 | ||
| 269 | /** | 270 | /** |
| @@ -273,16 +274,35 @@ async function summarizeWebLLM(hashedMessages) { | |||
| 273 | * @returns {Promise<HashedMessage[]>} Summarized messages | 274 | * @returns {Promise<HashedMessage[]>} Summarized messages |
| 274 | */ | 275 | */ |
| 275 | async function summarize(hashedMessages, endpoint = 'main') { | 276 | async 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 | } |
| 287 | 307 | ||
| 288 | async function synchronizeChat(batchSize = 5) { | 308 | async function synchronizeChat(batchSize = 5) { |
| @@ -307,16 +327,15 @@ async function synchronizeChat(batchSize = 5) { | |||
| 307 | return -1; | 327 | return -1; |
| 308 | } | 328 | } |
| 309 | 329 | ||
| 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); |
| 312 | 332 | ||
| 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)); |
| 319 | 335 | ||
| 336 | if (settings.summarize) { | ||
| 337 | newVectorItems = await summarize(newVectorItems, settings.summary_source); | ||
| 338 | } | ||
| 320 | 339 | ||
| 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 query | 706 | * @returns {Promise<string>} Text to query |
| 688 | */ | 707 | */ |
| 689 | async function getQueryText(chat, initiator) { | 708 | async 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)) })) |
| 692 | 711 | .filter(x => x.text) | |
| 693 | let hashedMessages = chat.map(x => ({ text: String(substituteParams(x.mes)) })); | 712 | .reverse() |
| 713 | .slice(0, settings.query); | ||
| 694 | 714 | ||
| 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 | } |
| 698 | 718 | ||
| 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 | } | ||
| 709 | 720 | ||
| 710 | return collapseNewlines(queryText).trim(); | 721 | return collapseNewlines(queryText).trim(); |
| 711 | } | 722 | } |