Skip to content

Commit

Permalink
feat: initial persistence (#3)
Browse files Browse the repository at this point in the history
  • Loading branch information
kirjavascript authored Jul 25, 2024
1 parent 2b1bf0f commit 5db834e
Show file tree
Hide file tree
Showing 9 changed files with 251 additions and 53 deletions.
13 changes: 13 additions & 0 deletions .editorconfig
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
root = true

[*]
indent_style = space
end_of_line = lf
charset = utf-8
trim_trailing_whitespace = true
insert_final_newline = true
max_line_length = 120
indent_size = 2

[*.md]
trim_trailing_whitespace = false
85 changes: 51 additions & 34 deletions packages/bolt/app/components/chat/Chat.client.tsx
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import type { Message } from 'ai';
import { useChat } from 'ai/react';
import { useAnimate } from 'framer-motion';
import { useEffect, useRef, useState } from 'react';
Expand All @@ -8,6 +9,7 @@ import { workbenchStore } from '~/lib/stores/workbench';
import { cubicEasingFn } from '~/utils/easings';
import { createScopedLogger } from '~/utils/logger';
import { BaseChat } from './BaseChat';
import { useChatHistory } from '~/lib/persistence';

const toastAnimation = cssTransition({
enter: 'animated fadeInRight',
Expand All @@ -17,9 +19,25 @@ const toastAnimation = cssTransition({
const logger = createScopedLogger('Chat');

export function Chat() {
const { ready, initialMessages, storeMessageHistory } = useChatHistory();

return (
<>
{ready && <ChatImpl initialMessages={initialMessages} storeMessageHistory={storeMessageHistory} />}
<ToastContainer position="bottom-right" stacked pauseOnFocusLoss transition={toastAnimation} />;
</>
);
}

interface ChatProps {
initialMessages: Message[];
storeMessageHistory: (messages: Message[]) => Promise<void>;
}

export function ChatImpl({ initialMessages, storeMessageHistory }: ChatProps) {
const textareaRef = useRef<HTMLTextAreaElement>(null);

const [chatStarted, setChatStarted] = useState(false);
const [chatStarted, setChatStarted] = useState(initialMessages.length > 0);

const [animationScope, animate] = useAnimate();

Expand All @@ -32,6 +50,7 @@ export function Chat() {
onFinish: () => {
logger.debug('Finished streaming');
},
initialMessages,
});

const { enhancingPrompt, promptEnhanced, enhancePrompt, resetEnhancer } = usePromptEnhancer();
Expand All @@ -41,6 +60,7 @@ export function Chat() {

useEffect(() => {
parseMessages(messages, isLoading);
storeMessageHistory(messages).catch((error) => toast.error(error.message));
}, [messages, isLoading, parseMessages]);

const scrollTextArea = () => {
Expand Down Expand Up @@ -97,38 +117,35 @@ export function Chat() {
const [messageRef, scrollRef] = useSnapScroll();

return (
<>
<BaseChat
ref={animationScope}
textareaRef={textareaRef}
input={input}
chatStarted={chatStarted}
isStreaming={isLoading}
enhancingPrompt={enhancingPrompt}
promptEnhanced={promptEnhanced}
sendMessage={sendMessage}
messageRef={messageRef}
scrollRef={scrollRef}
handleInputChange={handleInputChange}
handleStop={abort}
messages={messages.map((message, i) => {
if (message.role === 'user') {
return message;
}

return {
...message,
content: parsedMessages[i] || '',
};
})}
enhancePrompt={() => {
enhancePrompt(input, (input) => {
setInput(input);
scrollTextArea();
});
}}
/>
<ToastContainer position="bottom-right" stacked={true} pauseOnFocusLoss={true} transition={toastAnimation} />
</>
<BaseChat
ref={animationScope}
textareaRef={textareaRef}
input={input}
chatStarted={chatStarted}
isStreaming={isLoading}
enhancingPrompt={enhancingPrompt}
promptEnhanced={promptEnhanced}
sendMessage={sendMessage}
messageRef={messageRef}
scrollRef={scrollRef}
handleInputChange={handleInputChange}
handleStop={abort}
messages={messages.map((message, i) => {
if (message.role === 'user') {
return message;
}

return {
...message,
content: parsedMessages[i] || '',
};
})}
enhancePrompt={() => {
enhancePrompt(input, (input) => {
setInput(input);
scrollTextArea();
});
}}
/>
);
}
12 changes: 12 additions & 0 deletions packages/bolt/app/lib/.server/login.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,19 @@
import { env } from 'node:process';
import { isAuthenticated } from './sessions';
import { json, redirect, type LoaderFunctionArgs } from '@remix-run/cloudflare';

export function verifyPassword(password: string, cloudflareEnv: Env) {
const loginPassword = env.LOGIN_PASSWORD || cloudflareEnv.LOGIN_PASSWORD;

return password === loginPassword;
}

export async function handleAuthRequest({ request, context }: LoaderFunctionArgs, body: object = {}) {
const authenticated = await isAuthenticated(request, context.cloudflare.env);

if (import.meta.env.DEV || authenticated) {
return json(body);
}

return redirect('/login');
}
16 changes: 7 additions & 9 deletions packages/bolt/app/lib/hooks/useSnapScroll.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,15 +9,13 @@ export function useSnapScroll() {
const messageRef = useCallback((node: HTMLDivElement | null) => {
if (node) {
const observer = new ResizeObserver(() => {
if (autoScrollRef.current) {
if (scrollNodeRef.current) {
const { scrollHeight, clientHeight } = scrollNodeRef.current;
const scrollTarget = scrollHeight - clientHeight;

scrollNodeRef.current.scrollTo({
top: scrollTarget,
});
}
if (autoScrollRef.current && scrollNodeRef.current) {
const { scrollHeight, clientHeight } = scrollNodeRef.current;
const scrollTarget = scrollHeight - clientHeight;

scrollNodeRef.current.scrollTo({
top: scrollTarget,
});
}
});

Expand Down
67 changes: 67 additions & 0 deletions packages/bolt/app/lib/persistence/db.ts
Original file line number Diff line number Diff line change
@@ -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);
});
}
2 changes: 2 additions & 0 deletions packages/bolt/app/lib/persistence/index.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
export * from './db';
export * from './useChatHistory';
86 changes: 86 additions & 0 deletions packages/bolt/app/lib/persistence/useChatHistory.ts
Original file line number Diff line number Diff line change
@@ -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);
}
},
};
}
14 changes: 4 additions & 10 deletions packages/bolt/app/routes/_index.tsx
Original file line number Diff line number Diff line change
@@ -1,22 +1,16 @@
import { json, redirect, type LoaderFunctionArgs, type MetaFunction } from '@remix-run/cloudflare';
import { type LoaderFunctionArgs, type MetaFunction } from '@remix-run/cloudflare';
import { ClientOnly } from 'remix-utils/client-only';
import { BaseChat } from '~/components/chat/BaseChat';
import { Chat } from '~/components/chat/Chat.client';
import { Header } from '~/components/Header';
import { isAuthenticated } from '~/lib/.server/sessions';
import { handleAuthRequest } from '~/lib/.server/login';

export const meta: MetaFunction = () => {
return [{ title: 'Bolt' }, { name: 'description', content: 'Talk with Bolt, an AI assistant from StackBlitz' }];
};

export async function loader({ request, context }: LoaderFunctionArgs) {
const authenticated = await isAuthenticated(request, context.cloudflare.env);

if (import.meta.env.DEV || authenticated) {
return json({});
}

return redirect('/login');
export async function loader(args: LoaderFunctionArgs) {
return handleAuthRequest(args);
}

export default function Index() {
Expand Down
9 changes: 9 additions & 0 deletions packages/bolt/app/routes/chat.$id.tsx
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
import type { LoaderFunctionArgs } from '@remix-run/cloudflare';
import { default as IndexRoute } from './_index';
import { handleAuthRequest } from '~/lib/.server/login';

export async function loader(args: LoaderFunctionArgs) {
return handleAuthRequest(args, { id: args.params.id });
}

export default IndexRoute;

0 comments on commit 5db834e

Please sign in to comment.