server : (webui) introduce conversation branching + idb storage

examples/server/webui/package-lock.json

@@ -13,6 +13,7 @@
         "@vscode/markdown-it-katex": "^1.1.1",
         "autoprefixer": "^10.4.20",
         "daisyui": "^4.12.14",
+        "dexie": "^4.0.11",
         "highlight.js": "^11.10.0",
         "katex": "^0.16.15",
         "postcss": "^8.4.49",
@@ -2338,6 +2339,12 @@
         "url": "https://github.com/sponsors/wooorm"
       }
     },
+    "node_modules/dexie": {
+      "version": "4.0.11",
+      "resolved": "https://registry.npmjs.org/dexie/-/dexie-4.0.11.tgz",
+      "integrity": "sha512-SOKO002EqlvBYYKQSew3iymBoN2EQ4BDw/3yprjh7kAfFzjBYkaMNa/pZvcA7HSWlcKSQb9XhPe3wKyQ0x4A8A==",
+      "license": "Apache-2.0"
+    },
     "node_modules/didyoumean": {
       "version": "1.2.2",
       "resolved": "https://registry.npmjs.org/didyoumean/-/didyoumean-1.2.2.tgz",

examples/server/webui/package.json

@@ -16,6 +16,7 @@
     "@vscode/markdown-it-katex": "^1.1.1",
     "autoprefixer": "^10.4.20",
     "daisyui": "^4.12.14",
+    "dexie": "^4.0.11",
     "highlight.js": "^11.10.0",
     "katex": "^0.16.15",
     "postcss": "^8.4.49",

examples/server/webui/src/App.tsx

@@ -1,4 +1,4 @@
-import { HashRouter, Outlet, Route, Routes } from 'react-router';
+import { HashRouter, Outlet, Route, Routes, Navigate } from 'react-router';
 import Header from './components/Header';
 import Sidebar from './components/Sidebar';
 import { AppContextProvider, useAppContext } from './utils/app.context';

examples/server/webui/src/components/ChatMessage.tsx

@@ -3,6 +3,7 @@ import { useAppContext } from '../utils/app.context';
 import { Message, PendingMessage } from '../utils/types';
 import { classNames } from '../utils/misc';
 import MarkdownDisplay, { CopyButton } from './MarkdownDisplay';
+import { ChevronLeftIcon, ChevronRightIcon } from '@heroicons/react/24/outline';
 
 interface SplitMessage {
   content: PendingMessage['content'];
@@ -12,17 +13,24 @@ interface SplitMessage {
 
 export default function ChatMessage({
   msg,
+  siblingLastNodeIds,
+  siblingCurrIdx,
   id,
-  scrollToBottom,
+  onRegenerateMessage,
+  onEditMessage,
+  onChangeSibling,
   isPending,
 }: {
   msg: Message | PendingMessage;
+  siblingLastNodeIds: Message['id'][];
+  siblingCurrIdx: number;
   id?: string;
-  scrollToBottom: (requiresNearBottom: boolean) => void;
+  onRegenerateMessage(msg: Message): void;
+  onEditMessage(msg: Message, content: string): void;
+  onChangeSibling(sibling: Message['id']): void;
   isPending?: boolean;
 }) {
-  const { viewingConversation, replaceMessageAndGenerate, config } =
+  const { viewingChat, config } = useAppContext();
-    useAppContext();
   const [editingContent, setEditingContent] = useState<string | null>(null);
   const timings = useMemo(
     () =>
@@ -37,6 +45,8 @@ export default function ChatMessage({
         : null,
     [msg.timings]
   );
+  const nextSibling = siblingLastNodeIds[siblingCurrIdx + 1];
+  const prevSibling = siblingLastNodeIds[siblingCurrIdx - 1];
 
   // for reasoning model, we split the message into content and thought
   // TODO: implement this as remark/rehype plugin in the future
@@ -64,13 +74,7 @@ export default function ChatMessage({
     return { content: actualContent, thought, isThinking };
   }, [msg]);
 
-  if (!viewingConversation) return null;
+  if (!viewingChat) return null;
-
-  const regenerate = async () => {
-    replaceMessageAndGenerate(viewingConversation.id, msg.id, undefined, () =>
-      scrollToBottom(true)
-    );
-  };
 
   return (
     <div className="group" id={id}>
@@ -105,13 +109,12 @@ export default function ChatMessage({
               </button>
               <button
                 className="btn mt-2"
-                onClick={() =>
+                onClick={() => {
-                  replaceMessageAndGenerate(
+                  if (msg.content !== null) {
-                    viewingConversation.id,
+                    setEditingContent(null);
-                    msg.id,
+                    onEditMessage(msg as Message, editingContent);
-                    editingContent
-                  )
                   }
+                }}
               >
                 Submit
               </button>
@@ -196,10 +199,35 @@ export default function ChatMessage({
       {msg.content !== null && (
         <div
           className={classNames({
-            'mx-4 mt-2 mb-2': true,
+            'flex items-center gap-2 mx-4 mt-2 mb-2': true,
-            'text-right': msg.role === 'user',
+            'flex-row-reverse': msg.role === 'user',
           })}
         >
+          {siblingLastNodeIds && siblingLastNodeIds.length > 1 && (
+            <div className="flex gap-1 items-center opacity-60 text-sm">
+              <button
+                className={classNames({
+                  'btn btn-sm btn-ghost p-1': true,
+                  'opacity-20': !prevSibling,
+                })}
+                onClick={() => prevSibling && onChangeSibling(prevSibling)}
+              >
+                <ChevronLeftIcon className="h-4 w-4" />
+              </button>
+              <span>
+                {siblingCurrIdx + 1} / {siblingLastNodeIds.length}
+              </span>
+              <button
+                className={classNames({
+                  'btn btn-sm btn-ghost p-1': true,
+                  'opacity-20': !nextSibling,
+                })}
+                onClick={() => nextSibling && onChangeSibling(nextSibling)}
+              >
+                <ChevronRightIcon className="h-4 w-4" />
+              </button>
+            </div>
+          )}
           {/* user message */}
           {msg.role === 'user' && (
             <button
@@ -216,7 +244,11 @@ export default function ChatMessage({
               {!isPending && (
                 <button
                   className="badge btn-mini show-on-hover mr-2"
-                  onClick={regenerate}
+                  onClick={() => {
+                    if (msg.content !== null) {
+                      onRegenerateMessage(msg as Message);
+                    }
+                  }}
                   disabled={msg.content === null}
                 >
                   🔄 Regenerate

examples/server/webui/src/components/ChatScreen.tsx

@@ -1,28 +1,54 @@
-import { useEffect, useState } from 'react';
+import { useEffect, useMemo, useState } from 'react';
-import { useAppContext } from '../utils/app.context';
+import { CallbackGeneratedChunk, useAppContext } from '../utils/app.context';
-import StorageUtils from '../utils/storage';
-import { useNavigate } from 'react-router';
 import ChatMessage from './ChatMessage';
-import { CanvasType, PendingMessage } from '../utils/types';
+import { CanvasType, Message, PendingMessage } from '../utils/types';
-import { classNames } from '../utils/misc';
+import { classNames, throttle } from '../utils/misc';
 import CanvasPyInterpreter from './CanvasPyInterpreter';
+import StorageUtils from '../utils/storage';
 
-export default function ChatScreen() {
+export interface MessageDisplay {
-  const {
+  msg: Message | PendingMessage;
-    viewingConversation,
+  siblingLastNodeIds: Message['id'][];
-    sendMessage,
+  siblingCurrIdx: number;
-    isGenerating,
+  isPending?: boolean;
-    stopGenerating,
+}
-    pendingMessages,
-    canvasData,
-  } = useAppContext();
-  const [inputMsg, setInputMsg] = useState('');
-  const navigate = useNavigate();
 
-  const currConvId = viewingConversation?.id ?? '';
+function getListMessageDisplay(
-  const pendingMsg: PendingMessage | undefined = pendingMessages[currConvId];
+  msgs: Readonly<Message[]>,
+  lastNodeId: Message['id']
+): MessageDisplay[] {
+  const currNodes = StorageUtils.filterByLastNodeId(msgs, lastNodeId, true);
+  const res: MessageDisplay[] = [];
+  const nodeMap = new Map<Message['id'], Message>();
+  for (const msg of msgs) {
+    nodeMap.set(msg.id, msg);
+  }
+  const findLastNode = (msgId: Message['id']): Message['id'] => {
+    let currNode: Message | undefined = nodeMap.get(msgId);
+    while (currNode) {
+      if (currNode.children.length === 0) break;
+      currNode = nodeMap.get(currNode.children.at(-1) ?? -1);
+    }
+    return currNode?.id ?? -1;
+  };
+  // traverse the current nodes
+  for (const msg of currNodes) {
+    const parentNode = nodeMap.get(msg.parent ?? -1);
+    if (!parentNode) continue;
+    const siblings = parentNode.children;
+    if (msg.type !== 'root') {
+      res.push({
+        msg,
+        siblingLastNodeIds: siblings.map(findLastNode),
+        siblingCurrIdx: siblings.indexOf(msg.id),
+      });
+    }
+  }
+  return res;
+}
 
-  const scrollToBottom = (requiresNearBottom: boolean) => {
+const scrollToBottom = throttle(
+  (requiresNearBottom: boolean, delay: number = 80) => {
     const mainScrollElem = document.getElementById('main-scroll');
     if (!mainScrollElem) return;
     const spaceToBottom =
@@ -32,35 +58,100 @@ export default function ChatScreen() {
     if (!requiresNearBottom || spaceToBottom < 50) {
       setTimeout(
         () => mainScrollElem.scrollTo({ top: mainScrollElem.scrollHeight }),
-        1
+        delay
       );
     }
+  },
+  80
+);
+
+export default function ChatScreen() {
+  const {
+    viewingChat,
+    sendMessage,
+    isGenerating,
+    stopGenerating,
+    pendingMessages,
+    canvasData,
+    replaceMessageAndGenerate,
+  } = useAppContext();
+  const [inputMsg, setInputMsg] = useState('');
+
+  const [currNodeId, setCurrNodeId] = useState<number>(-1);
+  const messages: MessageDisplay[] = useMemo(() => {
+    if (!viewingChat) return [];
+    else return getListMessageDisplay(viewingChat.messages, currNodeId);
+  }, [currNodeId, viewingChat?.messages]);
+
+  const currConvId = viewingChat?.conv.id ?? null;
+  const pendingMsg: PendingMessage | undefined =
+    pendingMessages[currConvId ?? ''];
+
+  useEffect(() => {
+    // reset to latest node when conversation changes
+    setCurrNodeId(-1);
+    // scroll to bottom when conversation changes
+    scrollToBottom(false, 1);
+  }, [currConvId]);
+
+  const onChunk: CallbackGeneratedChunk = () => {
+    scrollToBottom(true);
   };
 
-  // scroll to bottom when conversation changes
-  useEffect(() => {
-    scrollToBottom(false);
-  }, [viewingConversation?.id]);
-
   const sendNewMessage = async () => {
-    if (inputMsg.trim().length === 0 || isGenerating(currConvId)) return;
+    if (inputMsg.trim().length === 0 || isGenerating(currConvId ?? '')) return;
-    const convId = viewingConversation?.id ?? StorageUtils.getNewConvId();
     const lastInpMsg = inputMsg;
     setInputMsg('');
-    if (!viewingConversation) {
-      // if user is creating a new conversation, redirect to the new conversation
-      navigate(`/chat/${convId}`);
-    }
     scrollToBottom(false);
-    // auto scroll as message is being generated
+    setCurrNodeId(-1);
-    const onChunk = () => scrollToBottom(true);
+    // get the last message node
-    if (!(await sendMessage(convId, inputMsg, onChunk))) {
+    const lastMsgNodeId = messages.at(-1)?.msg.id ?? null;
+    if (!(await sendMessage(currConvId, lastMsgNodeId, inputMsg, onChunk))) {
       // restore the input message if failed
       setInputMsg(lastInpMsg);
     }
   };
 
+  const handleEditMessage = async (msg: Message, content: string) => {
+    if (!viewingChat) return;
+    setCurrNodeId(msg.id);
+    scrollToBottom(false);
+    await replaceMessageAndGenerate(
+      viewingChat.conv.id,
+      msg.parent,
+      content,
+      onChunk
+    );
+    setCurrNodeId(-1);
+    scrollToBottom(false);
+  };
+
+  const handleRegenerateMessage = async (msg: Message) => {
+    if (!viewingChat) return;
+    setCurrNodeId(msg.parent);
+    scrollToBottom(false);
+    await replaceMessageAndGenerate(
+      viewingChat.conv.id,
+      msg.parent,
+      null,
+      onChunk
+    );
+    setCurrNodeId(-1);
+    scrollToBottom(false);
+  };
+
   const hasCanvas = !!canvasData;
+  const pendingMsgDisplay: MessageDisplay[] =
+    pendingMsg && messages.at(-1)?.msg.id !== pendingMsg.id
+      ? [
+          {
+            msg: pendingMsg,
+            siblingLastNodeIds: [],
+            siblingCurrIdx: 0,
+            isPending: true,
+          },
+        ]
+      : [];
 
   return (
     <div
@@ -81,24 +172,19 @@ export default function ChatScreen() {
         <div id="messages-list" className="grow">
           <div className="mt-auto flex justify-center">
             {/* placeholder to shift the message to the bottom */}
-            {viewingConversation ? '' : 'Send a message to start'}
+            {viewingChat ? '' : 'Send a message to start'}
           </div>
-          {viewingConversation?.messages.map((msg) => (
+          {[...messages, ...pendingMsgDisplay].map((msg) => (
             <ChatMessage
-              key={msg.id}
+              key={msg.msg.id}
-              msg={msg}
+              msg={msg.msg}
-              scrollToBottom={scrollToBottom}
+              siblingLastNodeIds={msg.siblingLastNodeIds}
+              siblingCurrIdx={msg.siblingCurrIdx}
+              onRegenerateMessage={handleRegenerateMessage}
+              onEditMessage={handleEditMessage}
+              onChangeSibling={setCurrNodeId}
             />
           ))}
-
-          {pendingMsg && (
-            <ChatMessage
-              msg={pendingMsg}
-              scrollToBottom={scrollToBottom}
-              isPending
-              id="pending-msg"
-            />
-          )}
         </div>
 
         {/* chat input */}
@@ -118,10 +204,10 @@ export default function ChatScreen() {
             id="msg-input"
             dir="auto"
           ></textarea>
-          {isGenerating(currConvId) ? (
+          {isGenerating(currConvId ?? '') ? (
             <button
               className="btn btn-neutral ml-2"
-              onClick={() => stopGenerating(currConvId)}
+              onClick={() => stopGenerating(currConvId ?? '')}
             >
               Stop
             </button>

examples/server/webui/src/components/Header.tsx

@@ -25,12 +25,12 @@ export default function Header() {
     );
   }, [selectedTheme]);
 
-  const { isGenerating, viewingConversation } = useAppContext();
+  const { isGenerating, viewingChat } = useAppContext();
-  const isCurrConvGenerating = isGenerating(viewingConversation?.id ?? '');
+  const isCurrConvGenerating = isGenerating(viewingChat?.conv.id ?? '');
 
   const removeConversation = () => {
-    if (isCurrConvGenerating || !viewingConversation) return;
+    if (isCurrConvGenerating || !viewingChat) return;
-    const convId = viewingConversation.id;
+    const convId = viewingChat?.conv.id;
     if (window.confirm('Are you sure to delete this conversation?')) {
       StorageUtils.remove(convId);
       navigate('/');
@@ -38,9 +38,9 @@ export default function Header() {
   };
 
   const downloadConversation = () => {
-    if (isCurrConvGenerating || !viewingConversation) return;
+    if (isCurrConvGenerating || !viewingChat) return;
-    const convId = viewingConversation.id;
+    const convId = viewingChat?.conv.id;
-    const conversationJson = JSON.stringify(viewingConversation, null, 2);
+    const conversationJson = JSON.stringify(viewingChat, null, 2);
     const blob = new Blob([conversationJson], { type: 'application/json' });
     const url = URL.createObjectURL(blob);
     const a = document.createElement('a');
@@ -75,7 +75,8 @@ export default function Header() {
 
       {/* action buttons (top right) */}
       <div className="flex items-center">
-        <div v-if="messages.length > 0" className="dropdown dropdown-end">
+        {viewingChat && (
+          <div className="dropdown dropdown-end">
             {/* "..." button */}
             <button
               tabIndex={0}
@@ -107,6 +108,8 @@ export default function Header() {
               </li>
             </ul>
           </div>
+        )}
+
         <div className="tooltip tooltip-bottom" data-tip="Settings">
           <button className="btn" onClick={() => setShowSettings(true)}>
             {/* settings button */}

examples/server/webui/src/components/Sidebar.tsx

@@ -1,4 +1,4 @@
-import { useEffect, useMemo, useState } from 'react';
+import { useEffect, useState } from 'react';
 import { classNames } from '../utils/misc';
 import { Conversation } from '../utils/types';
 import StorageUtils from '../utils/storage';
@@ -7,16 +7,17 @@ import { useNavigate, useParams } from 'react-router';
 export default function Sidebar() {
   const params = useParams();
   const navigate = useNavigate();
-  const currConv = useMemo(
-    () => StorageUtils.getOneConversation(params.convId ?? ''),
-    [params.convId]
-  );
 
   const [conversations, setConversations] = useState<Conversation[]>([]);
+  const [currConv, setCurrConv] = useState<Conversation | null>(null);
 
   useEffect(() => {
-    const handleConversationChange = () => {
+    StorageUtils.getOneConversation(params.convId ?? '').then(setCurrConv);
-      setConversations(StorageUtils.getAllConversations());
+  }, [params.convId]);
+
+  useEffect(() => {
+    const handleConversationChange = async () => {
+      setConversations(await StorageUtils.getAllConversations());
     };
     StorageUtils.onConversationChanged(handleConversationChange);
     handleConversationChange();
@@ -82,11 +83,11 @@ export default function Sidebar() {
               onClick={() => navigate(`/chat/${conv.id}`)}
               dir="auto"
             >
-              <span className="truncate">{conv.messages[0].content}</span>
+              <span className="truncate">{conv.name}</span>
             </div>
           ))}
           <div className="text-center text-xs opacity-40 mt-auto mx-4">
-            Conversations are saved to browser's localStorage
+            Conversations are saved to browser's IndexedDB
           </div>
         </div>
       </div>

examples/server/webui/src/utils/app.context.tsx

@@ -5,6 +5,7 @@ import {
   Conversation,
   Message,
   PendingMessage,
+  ViewingChat,
 } from './types';
 import StorageUtils from './storage';
 import {
@@ -13,24 +14,25 @@ import {
   getSSEStreamAsync,
 } from './misc';
 import { BASE_URL, CONFIG_DEFAULT, isDev } from '../Config';
-import { matchPath, useLocation } from 'react-router';
+import { matchPath, useLocation, useNavigate } from 'react-router';
 
 interface AppContextValue {
   // conversations and messages
-  viewingConversation: Conversation | null;
+  viewingChat: ViewingChat | null;
   pendingMessages: Record<Conversation['id'], PendingMessage>;
   isGenerating: (convId: string) => boolean;
   sendMessage: (
-    convId: string,
+    convId: string | null,
+    lastNodeId: Message['id'] | null,
     content: string,
-    onChunk?: CallbackGeneratedChunk
+    onChunk: CallbackGeneratedChunk
   ) => Promise<boolean>;
   stopGenerating: (convId: string) => void;
   replaceMessageAndGenerate: (
     convId: string,
-    origMsgId: Message['id'],
+    parentNodeId: Message['id'], // the parent node of the message to be replaced
-    content?: string,
+    content: string | null,
-    onChunk?: CallbackGeneratedChunk
+    onChunk: CallbackGeneratedChunk
   ) => Promise<void>;
 
   // canvas
@@ -44,23 +46,33 @@ interface AppContextValue {
   setShowSettings: (show: boolean) => void;
 }
 
-// for now, this callback is only used for scrolling to the bottom of the chat
+// this callback is used for scrolling to the bottom of the chat and switching to the last node
-type CallbackGeneratedChunk = () => void;
+export type CallbackGeneratedChunk = () => void;
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 const AppContext = createContext<AppContextValue>({} as any);
 
+const getViewingChat = async (convId: string): Promise<ViewingChat | null> => {
+  const conv = await StorageUtils.getOneConversation(convId);
+  if (!conv) return null;
+  return {
+    conv: conv,
+    // all messages from all branches, not filtered by last node
+    messages: await StorageUtils.getMessages(convId),
+  };
+};
+
 export const AppContextProvider = ({
   children,
 }: {
   children: React.ReactElement;
 }) => {
   const { pathname } = useLocation();
+  const navigate = useNavigate();
   const params = matchPath('/chat/:convId', pathname);
   const convId = params?.params?.convId;
 
-  const [viewingConversation, setViewingConversation] =
+  const [viewingChat, setViewingChat] = useState<ViewingChat | null>(null);
-    useState<Conversation | null>(null);
   const [pendingMessages, setPendingMessages] = useState<
     Record<Conversation['id'], PendingMessage>
   >({});
@@ -75,12 +87,12 @@ export const AppContextProvider = ({
   useEffect(() => {
     // also reset the canvas data
     setCanvasData(null);
-    const handleConversationChange = (changedConvId: string) => {
+    const handleConversationChange = async (changedConvId: string) => {
       if (changedConvId !== convId) return;
-      setViewingConversation(StorageUtils.getOneConversation(convId));
+      setViewingChat(await getViewingChat(changedConvId));
     };
     StorageUtils.onConversationChanged(handleConversationChange);
-    setViewingConversation(StorageUtils.getOneConversation(convId ?? ''));
+    getViewingChat(convId ?? '').then(setViewingChat);
     return () => {
       StorageUtils.offConversationChanged(handleConversationChange);
     };
@@ -118,23 +130,39 @@ export const AppContextProvider = ({
 
   const generateMessage = async (
     convId: string,
-    onChunk?: CallbackGeneratedChunk
+    lastNodeId: Message['id'],
+    onChunk: CallbackGeneratedChunk
   ) => {
     if (isGenerating(convId)) return;
 
     const config = StorageUtils.getConfig();
-    const currConversation = StorageUtils.getOneConversation(convId);
+    const currConversation = await StorageUtils.getOneConversation(convId);
     if (!currConversation) {
       throw new Error('Current conversation is not found');
     }
 
+    const currMessages = StorageUtils.filterByLastNodeId(
+      await StorageUtils.getMessages(convId),
+      lastNodeId,
+      false
+    );
     const abortController = new AbortController();
     setAbort(convId, abortController);
 
+    if (!currMessages) {
+      throw new Error('Current messages are not found');
+    }
+
+    const pendingId = Date.now() + 1;
     let pendingMsg: PendingMessage = {
-      id: Date.now() + 1,
+      id: pendingId,
+      convId,
+      type: 'text',
+      timestamp: pendingId,
       role: 'assistant',
       content: null,
+      parent: lastNodeId,
+      children: [],
     };
     setPending(convId, pendingMsg);
 
@@ -144,7 +172,7 @@ export const AppContextProvider = ({
         ...(config.systemMessage.length === 0
           ? []
           : [{ role: 'system', content: config.systemMessage } as APIMessage]),
-        ...normalizeMsgsForAPI(currConversation?.messages ?? []),
+        ...normalizeMsgsForAPI(currMessages),
       ];
       if (config.excludeThoughtOnReq) {
         messages = filterThoughtFromMsgs(messages);
@@ -205,8 +233,7 @@ export const AppContextProvider = ({
         const lastContent = pendingMsg.content || '';
         if (addedContent) {
           pendingMsg = {
-            id: pendingMsg.id,
+            ...pendingMsg,
-            role: 'assistant',
             content: lastContent + addedContent,
           };
         }
@@ -221,7 +248,7 @@ export const AppContextProvider = ({
           };
         }
         setPending(convId, pendingMsg);
-        onChunk?.();
+        onChunk(); // don't need to switch node for pending message
       }
     } catch (err) {
       setPending(convId, null);
@@ -236,37 +263,53 @@ export const AppContextProvider = ({
       }
     }
 
-    if (pendingMsg.content) {
+    if (pendingMsg.content !== null) {
-      StorageUtils.appendMsg(currConversation.id, {
+      await StorageUtils.appendMsg(pendingMsg as Message, lastNodeId);
-        id: pendingMsg.id,
-        content: pendingMsg.content,
-        role: pendingMsg.role,
-        timings: pendingMsg.timings,
-      });
     }
     setPending(convId, null);
-    onChunk?.(); // trigger scroll to bottom
+    onChunk(); // trigger scroll to bottom and switch to the last node
   };
 
   const sendMessage = async (
-    convId: string,
+    convId: string | null,
+    lastNodeId: Message['id'] | null,
     content: string,
-    onChunk?: CallbackGeneratedChunk
+    onChunk: CallbackGeneratedChunk
   ): Promise<boolean> => {
-    if (isGenerating(convId) || content.trim().length === 0) return false;
+    if (isGenerating(convId ?? '') || content.trim().length === 0) return false;
 
-    StorageUtils.appendMsg(convId, {
+    if (convId === null || convId.length === 0 || lastNodeId === null) {
-      id: Date.now(),
+      const conv = await StorageUtils.createConversation(
+        content.substring(0, 256)
+      );
+      convId = conv.id;
+      lastNodeId = conv.currNode;
+      // if user is creating a new conversation, redirect to the new conversation
+      navigate(`/chat/${convId}`);
+    }
+
+    const now = Date.now();
+    const currMsgId = now;
+    StorageUtils.appendMsg(
+      {
+        id: currMsgId,
+        timestamp: now,
+        type: 'text',
+        convId,
         role: 'user',
         content,
-    });
+        parent: lastNodeId,
+        children: [],
+      },
+      lastNodeId
+    );
+    onChunk();
 
     try {
-      await generateMessage(convId, onChunk);
+      await generateMessage(convId, currMsgId, onChunk);
       return true;
     } catch (_) {
-      // rollback
+      // TODO: rollback
-      StorageUtils.popMsg(convId);
     }
     return false;
   };
@@ -279,22 +322,33 @@ export const AppContextProvider = ({
   // if content is undefined, we remove last assistant message
   const replaceMessageAndGenerate = async (
     convId: string,
-    origMsgId: Message['id'],
+    parentNodeId: Message['id'], // the parent node of the message to be replaced
-    content?: string,
+    content: string | null,
-    onChunk?: CallbackGeneratedChunk
+    onChunk: CallbackGeneratedChunk
   ) => {
     if (isGenerating(convId)) return;
 
-    StorageUtils.filterAndKeepMsgs(convId, (msg) => msg.id < origMsgId);
+    if (content !== null) {
-    if (content) {
+      const now = Date.now();
-      StorageUtils.appendMsg(convId, {
+      const currMsgId = now;
-        id: Date.now(),
+      StorageUtils.appendMsg(
+        {
+          id: currMsgId,
+          timestamp: now,
+          type: 'text',
+          convId,
           role: 'user',
           content,
-      });
+          parent: parentNodeId,
+          children: [],
+        },
+        parentNodeId
+      );
+      parentNodeId = currMsgId;
     }
 
-    await generateMessage(convId, onChunk);
+    onChunk();
+    await generateMessage(convId, parentNodeId, onChunk);
   };
 
   const saveConfig = (config: typeof CONFIG_DEFAULT) => {
@@ -306,7 +360,7 @@ export const AppContextProvider = ({
     <AppContext.Provider
       value={{
         isGenerating,
-        viewingConversation,
+        viewingChat,
         pendingMessages,
         sendMessage,
         stopGenerating,

examples/server/webui/src/utils/misc.ts

@@ -4,7 +4,6 @@ import { APIMessage, Message } from './types';
 
 // ponyfill for missing ReadableStream asyncIterator on Safari
 import { asyncIterator } from '@sec-ant/readable-stream/ponyfill/asyncIterator';
-import { isDev } from '../Config';
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 export const isString = (x: any) => !!x.toLowerCase;
@@ -23,7 +22,7 @@ export async function* getSSEStreamAsync(fetchResponse: Response) {
     .pipeThrough(new TextLineStream());
   // @ts-expect-error asyncIterator complains about type, but it should work
   for await (const line of asyncIterator(lines)) {
-    if (isDev) console.log({ line });
+    //if (isDev) console.log({ line });
     if (line.startsWith('data:') && !line.endsWith('[DONE]')) {
       const data = JSON.parse(line.slice(5));
       yield data;
@@ -55,7 +54,7 @@ export const copyStr = (textToCopy: string) => {
 /**
  * filter out redundant fields upon sending to API
  */
-export function normalizeMsgsForAPI(messages: Message[]) {
+export function normalizeMsgsForAPI(messages: Readonly<Message[]>) {
   return messages.map((msg) => {
     return {
       role: msg.role,
@@ -88,3 +87,23 @@ export function classNames(classes: Record<string, boolean>): string {
 
 export const delay = (ms: number) =>
   new Promise((resolve) => setTimeout(resolve, ms));
+
+export const throttle = <T extends unknown[]>(
+  callback: (...args: T) => void,
+  delay: number
+) => {
+  let isWaiting = false;
+
+  return (...args: T) => {
+    if (isWaiting) {
+      return;
+    }
+
+    callback(...args);
+    isWaiting = true;
+
+    setTimeout(() => {
+      isWaiting = false;
+    }, delay);
+  };
+};

examples/server/webui/src/utils/storage.ts

@@ -2,7 +2,8 @@
 // format: { [convId]: { id: string, lastModified: number, messages: [...] } }
 
 import { CONFIG_DEFAULT } from '../Config';
-import { Conversation, Message } from './types';
+import { Conversation, Message, TimingReport } from './types';
+import Dexie, { Table } from 'dexie';
 
 const event = new EventTarget();
 
@@ -17,85 +18,151 @@ const dispatchConversationChange = (convId: string) => {
   );
 };
 
+const db = new Dexie('LlamacppWebui') as Dexie & {
+  conversations: Table<Conversation>;
+  messages: Table<Message>;
+};
+
+// https://dexie.org/docs/Version/Version.stores()
+db.version(1).stores({
+  // Unlike SQL, you don’t need to specify all properties but only the one you wish to index.
+  conversations: '&id, lastModified',
+  messages: '&id, convId, [convId+id], timestamp',
+});
+
 // convId is a string prefixed with 'conv-'
 const StorageUtils = {
   /**
    * manage conversations
    */
-  getAllConversations(): Conversation[] {
+  async getAllConversations(): Promise<Conversation[]> {
-    const res = [];
+    return (await db.conversations.toArray()).sort(
-    for (const key in localStorage) {
+      (a, b) => b.lastModified - a.lastModified
-      if (key.startsWith('conv-')) {
+    );
-        res.push(JSON.parse(localStorage.getItem(key) ?? '{}'));
-      }
-    }
-    res.sort((a, b) => b.lastModified - a.lastModified);
-    return res;
   },
   /**
    * can return null if convId does not exist
    */
-  getOneConversation(convId: string): Conversation | null {
+  async getOneConversation(convId: string): Promise<Conversation | null> {
-    return JSON.parse(localStorage.getItem(convId) || 'null');
+    return (await db.conversations.where('id').equals(convId).first()) ?? null;
   },
   /**
-   * if convId does not exist, create one
+   * get messages by convId and timeline
    */
-  appendMsg(convId: string, msg: Message): void {
+  async getMessages(convId: string): Promise<Message[]> {
-    if (msg.content === null) return;
+    return await db.messages.where({ convId }).toArray();
-    const conv = StorageUtils.getOneConversation(convId) || {
+  },
-      id: convId,
+  /**
-      lastModified: Date.now(),
+   * use in conjunction with getMessages to filter messages by lastNodeId
-      messages: [],
+   */
+  filterByLastNodeId(
+    msgs: Readonly<Message[]>,
+    lastNodeId: Message['id'],
+    includeRoot: boolean
+  ): Readonly<Message[]> {
+    const res: Message[] = [];
+    const nodeMap = new Map<Message['id'], Message>();
+    for (const msg of msgs) {
+      nodeMap.set(msg.id, msg);
+    }
+    let startNode: Message | undefined = nodeMap.get(lastNodeId);
+    if (!startNode) {
+      // if not found, we return the path with the latest timestamp
+      let latestTime = -1;
+      for (const msg of msgs) {
+        if (msg.timestamp > latestTime) {
+          startNode = msg;
+          latestTime = msg.timestamp;
+        }
+      }
+    }
+    // traverse the path from lastNodeId to root
+    // startNode can never be undefined here
+    let currNode: Message | undefined = startNode;
+    while (currNode) {
+      if (currNode.type !== 'root' || (currNode.type === 'root' && includeRoot))
+        res.push(currNode);
+      currNode = nodeMap.get(currNode.parent ?? -1);
+    }
+    res.sort((a, b) => a.timestamp - b.timestamp);
+    return res;
+  },
+  /**
+   * create a new conversation with a default timeline number 0
+   */
+  async createConversation(name: string): Promise<Conversation> {
+    const now = Date.now();
+    const msgId = now;
+    const conv: Conversation = {
+      id: `conv-${now}`,
+      lastModified: now,
+      currNode: msgId,
+      name,
     };
-    conv.messages.push(msg);
+    await db.conversations.add(conv);
-    conv.lastModified = Date.now();
+    // create a root node
-    localStorage.setItem(convId, JSON.stringify(conv));
+    await db.messages.add({
-    dispatchConversationChange(convId);
+      id: msgId,
+      convId: conv.id,
+      type: 'root',
+      timestamp: now,
+      role: 'system',
+      content: '',
+      parent: -1,
+      children: [],
+    });
+    return conv;
   },
   /**
-   * Get new conversation id
+   * if convId does not exist, throw an error
    */
-  getNewConvId(): string {
+  async appendMsg(
-    return `conv-${Date.now()}`;
+    msg: Exclude<Message, 'parent' | 'children'>,
+    parentNodeId: Message['id']
+  ): Promise<void> {
+    if (msg.content === null) return;
+    const { convId } = msg;
+    await db.transaction('rw', db.conversations, db.messages, async () => {
+      const conv = await StorageUtils.getOneConversation(convId);
+      const parentMsg = await db.messages
+        .where({ convId, id: parentNodeId })
+        .first();
+      // update the currNode of conversation
+      if (!conv) {
+        throw new Error(`Conversation ${convId} does not exist`);
+      }
+      if (!parentMsg) {
+        throw new Error(
+          `Parent message ID ${parentNodeId} does not exist in conversation ${convId}`
+        );
+      }
+      await db.conversations.update(convId, {
+        lastModified: Date.now(),
+        currNode: msg.id,
+      });
+      // update parent
+      await db.messages.update(parentNodeId, {
+        children: [...parentMsg.children, msg.id],
+      });
+      // create message
+      await db.messages.add({
+        ...msg,
+        parent: parentNodeId,
+        children: [],
+      });
+    });
+    dispatchConversationChange(convId);
   },
   /**
    * remove conversation by id
    */
-  remove(convId: string): void {
+  async remove(convId: string): Promise<void> {
-    localStorage.removeItem(convId);
+    await db.transaction('rw', db.conversations, db.messages, async () => {
+      await db.conversations.delete(convId);
+      await db.messages.where({ convId }).delete();
+    });
     dispatchConversationChange(convId);
   },
-  /**
-   * remove all conversations
-   */
-  filterAndKeepMsgs(
-    convId: string,
-    predicate: (msg: Message) => boolean
-  ): void {
-    const conv = StorageUtils.getOneConversation(convId);
-    if (!conv) return;
-    conv.messages = conv.messages.filter(predicate);
-    conv.lastModified = Date.now();
-    localStorage.setItem(convId, JSON.stringify(conv));
-    dispatchConversationChange(convId);
-  },
-  /**
-   * remove last message from conversation
-   */
-  popMsg(convId: string): Message | undefined {
-    const conv = StorageUtils.getOneConversation(convId);
-    if (!conv) return;
-    const msg = conv.messages.pop();
-    conv.lastModified = Date.now();
-    if (conv.messages.length === 0) {
-      StorageUtils.remove(convId);
-    } else {
-      localStorage.setItem(convId, JSON.stringify(conv));
-    }
-    dispatchConversationChange(convId);
-    return msg;
-  },
 
   // event listeners
   onConversationChanged(callback: CallbackConversationChanged) {
@@ -136,3 +203,78 @@ const StorageUtils = {
 };
 
 export default StorageUtils;
+
+// Migration from localStorage to IndexedDB
+
+// these are old types, LS prefix stands for LocalStorage
+interface LSConversation {
+  id: string; // format: `conv-{timestamp}`
+  lastModified: number; // timestamp from Date.now()
+  messages: LSMessage[];
+}
+interface LSMessage {
+  id: number;
+  role: 'user' | 'assistant' | 'system';
+  content: string;
+  timings?: TimingReport;
+}
+async function migrationLStoIDB() {
+  const res: LSConversation[] = [];
+  for (const key in localStorage) {
+    if (key.startsWith('conv-')) {
+      res.push(JSON.parse(localStorage.getItem(key) ?? '{}'));
+    }
+  }
+  if (res.length === 0) return;
+  await db.transaction('rw', db.conversations, db.messages, async () => {
+    for (const conv of res) {
+      const { id: convId, lastModified, messages } = conv;
+      const firstMsg = messages[0];
+      const lastMsg = messages.at(-1);
+      if (messages.length < 2 || !firstMsg || !lastMsg) {
+        console.log(
+          `Skipping conversation ${convId} with ${messages.length} messages`
+        );
+        continue;
+      }
+      const name = firstMsg.content ?? '(no messages)';
+      await db.conversations.add({
+        id: convId,
+        lastModified,
+        currNode: lastMsg.id,
+        name,
+      });
+      const rootId = Date.now();
+      await db.messages.add({
+        id: rootId,
+        convId: convId,
+        type: 'root',
+        timestamp: lastModified,
+        role: 'system',
+        content: '',
+        parent: -1,
+        children: [firstMsg.id],
+      });
+      for (let i = 0; i < messages.length; i++) {
+        const msg = messages[i];
+        await db.messages.add({
+          ...msg,
+          type: 'text',
+          convId: convId,
+          timestamp: msg.id,
+          parent: i === 0 ? rootId : messages[i - 1].id,
+          children: i === messages.length - 1 ? [] : [messages[i + 1].id],
+        });
+      }
+      localStorage.removeItem(convId);
+      console.log(
+        `Migrated conversation ${convId} with ${messages.length} messages`
+      );
+    }
+    console.log(
+      `Migrated ${res.length} conversations from localStorage to IndexedDB`
+    );
+  });
+}
+
+migrationLStoIDB();

examples/server/webui/src/utils/types.ts

@@ -5,11 +5,24 @@ export interface TimingReport {
   predicted_ms: number;
 }
 
+/**
+ * What is conversation "branching"? It is a feature that allows the user to edit an old message in the history, while still keeping the conversation flow.
+ * Inspired by ChatGPT UI where you edit a message, a new branch of the conversation is created, and the old message is still visible.
+ *
+ * We use the same node based structure as ChatGPT, where each message has a parent and children. A "root" message is the first message in a conversation, which will not be displayed in the UI.
+ */
+
 export interface Message {
   id: number;
+  convId: string;
+  type: 'text' | 'root';
+  timestamp: number; // timestamp from Date.now()
   role: 'user' | 'assistant' | 'system';
   content: string;
   timings?: TimingReport;
+  // node based system for branching
+  parent: Message['id'];
+  children: Message['id'][];
 }
 
 export type APIMessage = Pick<Message, 'role' | 'content'>;
@@ -17,7 +30,13 @@ export type APIMessage = Pick<Message, 'role' | 'content'>;
 export interface Conversation {
   id: string; // format: `conv-{timestamp}`
   lastModified: number; // timestamp from Date.now()
-  messages: Message[];
+  currNode: Message['id']; // the current message node being viewed
+  name: string;
+}
+
+export interface ViewingChat {
+  conv: Readonly<Conversation>;
+  messages: Readonly<Message[]>;
 }
 
 export type PendingMessage = Omit<Message, 'content'> & {