feat: initial persistence (#3)

This commit is contained in:
Kirjava
2024-07-25 14:03:38 +01:00
committed by GitHub
parent 2b1bf0fdaf
commit 5db834e2f7
9 changed files with 249 additions and 51 deletions

View File

@@ -0,0 +1,67 @@
import type { ChatHistory } from './useChatHistory';
import type { Message } from 'ai';
import { createScopedLogger } from '~/utils/logger';
const logger = createScopedLogger('ChatHistory');
// this is used at the top level and never rejects
export async function openDatabase(): Promise<IDBDatabase | undefined> {
return new Promise((resolve) => {
const request = indexedDB.open('boltHistory', 1);
request.onupgradeneeded = (event: IDBVersionChangeEvent) => {
const db = (event.target as IDBOpenDBRequest).result;
if (!db.objectStoreNames.contains('chats')) {
const store = db.createObjectStore('chats', { keyPath: 'id' });
store.createIndex('id', 'id', { unique: true });
}
};
request.onsuccess = (event: Event) => {
resolve((event.target as IDBOpenDBRequest).result);
};
request.onerror = (event: Event) => {
resolve(undefined);
logger.error((event.target as IDBOpenDBRequest).error);
};
});
}
export async function setMessages(db: IDBDatabase, id: string, messages: Message[]): Promise<void> {
return new Promise((resolve, reject) => {
const transaction = db.transaction('chats', 'readwrite');
const store = transaction.objectStore('chats');
const request = store.put({
id,
messages,
});
request.onsuccess = () => resolve();
request.onerror = () => reject(request.error);
});
}
export async function getMessages(db: IDBDatabase, id: string): Promise<ChatHistory> {
return new Promise((resolve, reject) => {
const transaction = db.transaction('chats', 'readonly');
const store = transaction.objectStore('chats');
const request = store.get(id);
request.onsuccess = () => resolve(request.result as ChatHistory);
request.onerror = () => reject(request.error);
});
}
export async function getNextID(db: IDBDatabase): Promise<string> {
return new Promise((resolve, reject) => {
const transaction = db.transaction('chats', 'readonly');
const store = transaction.objectStore('chats');
const request = store.count();
request.onsuccess = () => resolve(String(request.result));
request.onerror = () => reject(request.error);
});
}

View File

@@ -0,0 +1,2 @@
export * from './db';
export * from './useChatHistory';

View File

@@ -0,0 +1,86 @@
import { useNavigate, useLoaderData } from '@remix-run/react';
import { useState, useEffect } from 'react';
import type { Message } from 'ai';
import { openDatabase, setMessages, getMessages, getNextID } from './db';
import { toast } from 'react-toastify';
export interface ChatHistory {
id: string;
displayName?: string;
messages: Message[];
}
const persistenceEnabled = !import.meta.env.VITE_DISABLE_PERSISTENCE;
const db = persistenceEnabled ? await openDatabase() : undefined;
export function useChatHistory() {
const navigate = useNavigate();
const { id: chatId } = useLoaderData<{ id?: string }>();
const [initialMessages, setInitialMessages] = useState<Message[]>([]);
const [ready, setReady] = useState<boolean>(false);
const [entryId, setEntryId] = useState<string | undefined>();
useEffect(() => {
if (!db) {
setReady(true);
if (persistenceEnabled) {
toast.error(`Chat persistence is unavailable`);
}
return;
}
if (chatId) {
getMessages(db, chatId)
.then((storedMessages) => {
if (storedMessages && storedMessages.messages.length > 0) {
setInitialMessages(storedMessages.messages);
} else {
navigate(`/`, { replace: true });
}
setReady(true);
})
.catch((error) => {
toast.error(error.message);
});
}
}, []);
return {
ready: !chatId || ready,
initialMessages,
storeMessageHistory: async (messages: Message[]) => {
if (!db || messages.length === 0) {
return;
}
if (initialMessages.length === 0) {
if (!entryId) {
const nextId = await getNextID(db);
await setMessages(db, nextId, messages);
setEntryId(nextId);
/**
* FIXME: Using the intended navigate function causes a rerender for <Chat /> that breaks the app.
*
* `navigate(`/chat/${nextId}`, { replace: true });`
*/
const url = new URL(window.location.href);
url.pathname = `/chat/${nextId}`;
window.history.replaceState({}, '', url);
} else {
await setMessages(db, entryId, messages);
}
} else {
await setMessages(db, chatId as string, messages);
}
},
};
}