rixty
/
SillyTavern
mirror of https://github.com/SillyTavern/SillyTavern.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add client side cacheing of vector summaries

This commit is contained in:
QuantumEntangledAndy 2024-10-15 10:34:03 +07:00
parent bb9700b478
commit bb7e7b645d
No known key found for this signature in database
GPG Key ID: 3EB4B66F30C609B6
1 changed files with 77 additions and 58 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

135
public/scripts/extensions/vectors/index.js
View File

@ -36,6 +36,7 @@ import { generateWebLlmChatPrompt, isWebLlmSupported } from '../shared.js';
/** /**
* @typedef {object} HashedMessage * @typedef {object} HashedMessage
* @property {string} text - The hashed message text * @property {string} text - The hashed message text
* @property {number} hash - The hash used as the vector key
*/ */
const MODULE_NAME = 'vectors'; const MODULE_NAME = 'vectors';
@ -96,6 +97,8 @@ const settings = {
const moduleWorker = new ModuleWorkerWrapper(synchronizeChat); const moduleWorker = new ModuleWorkerWrapper(synchronizeChat);
const cachedSummaries = new Map();
/** /**
* Gets the Collection ID for a file embedded in the chat. * Gets the Collection ID for a file embedded in the chat.
* @param {string} fileUrl URL of the file * @param {string} fileUrl URL of the file
@ -118,6 +121,10 @@ async function onVectorizeAllClick() {
return; return;
} }
// Clear all cached summaries to ensure that new ones are created
// upon request of a full vectorise
cachedSummaries.clear();
const batchSize = 5; const batchSize = 5;
const elapsedLog = []; const elapsedLog = [];
let finished = false; let finished = false;
@ -200,70 +207,64 @@ function splitByChunks(items) {
/** /**
* Summarizes messages using the Extras API method. * Summarizes messages using the Extras API method.
* @param {HashedMessage[]} hashedMessages Array of hashed messages * @param {HashedMessage} element hashed message
* @returns {Promise<HashedMessage[]>} Summarized messages * @returns {Promise<boolean>} Sucess
*/ */
async function summarizeExtra(hashedMessages) { async function summarizeExtra(element) {
for (const element of hashedMessages) { try {
try { const url = new URL(getApiUrl());
const url = new URL(getApiUrl()); url.pathname = '/api/summarize';
url.pathname = '/api/summarize';
const apiResult = await doExtrasFetch(url, { const apiResult = await doExtrasFetch(url, {
method: 'POST', method: 'POST',
headers: { headers: {
'Content-Type': 'application/json', 'Content-Type': 'application/json',
'Bypass-Tunnel-Reminder': 'bypass', 'Bypass-Tunnel-Reminder': 'bypass',
}, },
body: JSON.stringify({ body: JSON.stringify({
text: element.text, text: element.text,
params: {}, params: {},
}), }),
}); });
if (apiResult.ok) { if (apiResult.ok) {
const data = await apiResult.json(); const data = await apiResult.json();
element.text = data.summary; element.text = data.summary;
}
}
catch (error) {
console.log(error);
} }
} }
catch (error) {
console.log(error);
return false;
}
return hashedMessages; return true;
} }
/** /**
* Summarizes messages using the main API method. * Summarizes messages using the main API method.
* @param {HashedMessage[]} hashedMessages Array of hashed messages * @param {HashedMessage} element hashed message
* @returns {Promise<HashedMessage[]>} Summarized messages * @returns {Promise<boolean>} Sucess
*/ */
async function summarizeMain(hashedMessages) { async function summarizeMain(element) {
for (const element of hashedMessages) { element.text = await generateRaw(element.text, '', false, false, settings.summary_prompt);
element.text = await generateRaw(element.text, '', false, false, settings.summary_prompt); return true;
}
return hashedMessages;
} }
/** /**
* Summarizes messages using WebLLM. * Summarizes messages using WebLLM.
* @param {HashedMessage[]} hashedMessages Array of hashed messages * @param {HashedMessage} element hashed message
* @returns {Promise<HashedMessage[]>} Summarized messages * @returns {Promise<boolean>} Sucess
*/ */
async function summarizeWebLLM(hashedMessages) { async function summarizeWebLLM(element) {
if (!isWebLlmSupported()) { if (!isWebLlmSupported()) {
console.warn('Vectors: WebLLM is not supported'); console.warn('Vectors: WebLLM is not supported');
return hashedMessages; return false;
} }
for (const element of hashedMessages) { const messages = [{ role:'system', content: settings.summary_prompt }, { role:'user', content: element.text }];
const messages = [{ role:'system', content: settings.summary_prompt }, { role:'user', content: element.text }]; element.text = await generateWebLlmChatPrompt(messages);
element.text = await generateWebLlmChatPrompt(messages);
}
return hashedMessages; return true;
} }
/** /**
@ -273,16 +274,35 @@ async function summarizeWebLLM(hashedMessages) {
* @returns {Promise<HashedMessage[]>} Summarized messages * @returns {Promise<HashedMessage[]>} Summarized messages
*/ */
async function summarize(hashedMessages, endpoint = 'main') { async function summarize(hashedMessages, endpoint = 'main') {
switch (endpoint) { for (const element of hashedMessages) {
case 'main': const cachedSummary = cachedSummaries.get(element.hash)
return await summarizeMain(hashedMessages); if (!cachedSummary) {
case 'extras': let sucess = true;
return await summarizeExtra(hashedMessages); switch (endpoint) {
case 'webllm': case 'main':
return await summarizeWebLLM(hashedMessages); sucess = await summarizeMain(element);
default: break;
console.error('Unsupported endpoint', endpoint); case 'extras':
sucess = await summarizeExtra(element);
break;
case 'webllm':
sucess = await summarizeWebLLM(element);
break;
default:
console.error('Unsupported endpoint', endpoint);
sucess = false;
break;
}
if (sucess) {
cachedSummaries.set(element.hash, element.text);
} else {
break;
}
} else {
element.text = cachedSummary;
}
} }
return hashedMessages;
} }
async function synchronizeChat(batchSize = 5) { async function synchronizeChat(batchSize = 5) {
@ -307,16 +327,15 @@ async function synchronizeChat(batchSize = 5) {
return -1; return -1;
} }
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) })); 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) }));
const hashesInCollection = await getSavedHashes(chatId); const hashesInCollection = await getSavedHashes(chatId);
if (settings.summarize) { let newVectorItems = hashedMessages.filter(x => !hashesInCollection.includes(x.hash));
hashedMessages = await summarize(hashedMessages, settings.summary_source);
}
const newVectorItems = hashedMessages.filter(x => !hashesInCollection.includes(x.hash));
const deletedHashes = hashesInCollection.filter(x => !hashedMessages.some(y => y.hash === x)); const deletedHashes = hashesInCollection.filter(x => !hashedMessages.some(y => y.hash === x));
if (settings.summarize) {
newVectorItems = await summarize(newVectorItems, settings.summary_source);
}
if (newVectorItems.length > 0) { if (newVectorItems.length > 0) {
const chunkedBatch = splitByChunks(newVectorItems.slice(0, batchSize)); const chunkedBatch = splitByChunks(newVectorItems.slice(0, batchSize));