Persist token cache to indexedDb

This commit is contained in:
Cohee 2023-08-13 22:45:27 +03:00
parent 23a4705d36
commit e6bf870c55
2 changed files with 41 additions and 3 deletions

View File

@ -45,6 +45,7 @@ import {
} from "./secrets.js"; } from "./secrets.js";
import { import {
IndexedDBStore,
delay, delay,
download, download,
getFileText, getFileText,
@ -120,7 +121,39 @@ const openrouter_website_model = 'OR_Website';
let biasCache = undefined; let biasCache = undefined;
let model_list = []; let model_list = [];
const tokenCache = {}; const objectStore = new IndexedDBStore('SillyTavern', 'chat_completions');
const tokenCache = await loadTokenCache();
async function loadTokenCache() {
try {
console.debug('Chat Completions: loading token cache from IndexedDB')
return await objectStore.get('tokenCache') || {};
} catch (e) {
console.log('Chat Completions: unable to load token cache from IndexedDB, using default value', e);
return {};
}
}
async function saveTokenCache() {
try {
console.debug('Chat Completions: saving token cache to IndexedDB')
await objectStore.put('tokenCache', tokenCache);
} catch (e) {
console.log('Chat Completions: unable to save token cache to IndexedDB', e);
}
}
async function resetTokenCache() {
try {
console.debug('Chat Completions: resetting token cache in IndexedDB');
Object.keys(tokenCache).forEach(key => delete tokenCache[key]);
await objectStore.delete('tokenCache');
} catch (e) {
console.log('Chat Completions: unable to reset token cache in IndexedDB', e);
}
}
window['resetTokenCache'] = resetTokenCache;
export const chat_completion_sources = { export const chat_completion_sources = {
OPENAI: 'openai', OPENAI: 'openai',
@ -795,6 +828,8 @@ function prepareOpenAIMessages({
const chat = chatCompletion.getChat(); const chat = chatCompletion.getChat();
openai_messages_count = chat.filter(x => x?.role === "user" || x?.role === "assistant")?.length || 0; openai_messages_count = chat.filter(x => x?.role === "user" || x?.role === "assistant")?.length || 0;
// Save token cache to IndexedDB storage (async, no need to await)
saveTokenCache();
return [chat, promptManager.tokenHandler.counts]; return [chat, promptManager.tokenHandler.counts];
} }

View File

@ -375,15 +375,18 @@ export class IndexedDBStore {
this.dbName = dbName; this.dbName = dbName;
this.storeName = storeName; this.storeName = storeName;
this.db = null; this.db = null;
this.version = Date.now();
} }
async open() { async open() {
return new Promise((resolve, reject) => { return new Promise((resolve, reject) => {
const request = indexedDB.open(this.dbName); const request = indexedDB.open(this.dbName, this.version);
request.onupgradeneeded = (event) => { request.onupgradeneeded = (event) => {
const db = event.target.result; const db = event.target.result;
db.createObjectStore(this.storeName, { keyPath: null, autoIncrement: false }); if (!db.objectStoreNames.contains(this.storeName)) {
db.createObjectStore(this.storeName, { keyPath: null, autoIncrement: false });
}
}; };
request.onsuccess = (event) => { request.onsuccess = (event) => {