webui: Add editing assistant messages (#11849)

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

@@ -27,13 +27,15 @@ export default function ChatMessage({
   onEditMessage,
   onChangeSibling,
   isPending,
+  onContinueMessage,
 }: {
   msg: Message | PendingMessage;
   siblingLeafNodeIds: Message['id'][];
   siblingCurrIdx: number;
   id?: string;
   onRegenerateMessage(msg: Message): void;
   onEditMessage(msg: Message, content: string): void;
+  onContinueMessage(msg: Message, content: string): void;
   onChangeSibling(sibling: Message['id']): void;
   isPending?: boolean;
 }) {
@@ -123,7 +125,11 @@ export default function ChatMessage({
                 onClick={() => {
                   if (msg.content !== null) {
                     setEditingContent(null);
-                    onEditMessage(msg as Message, editingContent);
+                    if (msg.role === 'user') {
+                      onEditMessage(msg as Message, editingContent);
+                    } else {
+                      onContinueMessage(msg as Message, editingContent);
+                    }
                   }
                 }}
               >
@@ -248,12 +254,22 @@ export default function ChatMessage({
                   <ArrowPathIcon className="h-4 w-4" />
                 </BtnWithTooltips>
               )}
+              {!isPending && (
+                <BtnWithTooltips
+                  className="btn-mini show-on-hover w-8 h-8"
+                  onClick={() => setEditingContent(msg.content)}
+                  disabled={msg.content === null}
+                  tooltipsContent="Edit message"
+                >
+                  <PencilSquareIcon className="h-4 w-4" />
+                </BtnWithTooltips>
+              )}
+              <CopyButton
+                className="btn-mini show-on-hover w-8 h-8"
+                content={msg.content}
+              />
             </>
           )}
-          <CopyButton
-            className="btn-mini show-on-hover w-8 h-8"
-            content={msg.content}
-          />
         </div>
       )}
     </div>

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

@@ -94,6 +94,7 @@ export default function ChatScreen() {
     pendingMessages,
     canvasData,
     replaceMessageAndGenerate,
+    continueMessageAndGenerate,
   } = useAppContext();
 
   const textarea: ChatTextareaApi = useChatTextarea(prefilledMsg.content());
@@ -188,6 +189,20 @@ export default function ChatScreen() {
     scrollToBottom(false);
   };
 
+  const handleContinueMessage = async (msg: Message, content: string) => {
+    if (!viewingChat || !continueMessageAndGenerate) return;
+    setCurrNodeId(msg.id);
+    scrollToBottom(false);
+    await continueMessageAndGenerate(
+      viewingChat.conv.id,
+      msg.id,
+      content,
+      onChunk
+    );
+    setCurrNodeId(-1);
+    scrollToBottom(false);
+  };
+
   const hasCanvas = !!canvasData;
 
   useEffect(() => {
@@ -205,7 +220,7 @@ export default function ChatScreen() {
 
   // due to some timing issues of StorageUtils.appendMsg(), we need to make sure the pendingMsg is not duplicated upon rendering (i.e. appears once in the saved conversation and once in the pendingMsg)
   const pendingMsgDisplay: MessageDisplay[] =
-    pendingMsg && messages.at(-1)?.msg.id !== pendingMsg.id
+    pendingMsg && !messages.some((m) => m.msg.id === pendingMsg.id) // Only show if pendingMsg is not an existing message being continued
       ? [
           {
             msg: pendingMsg,
@@ -244,18 +259,35 @@ export default function ChatScreen() {
               </>
             )}
           </div>
-          {[...messages, ...pendingMsgDisplay].map((msg) => (
-            <ChatMessage
-              key={msg.msg.id}
-              msg={msg.msg}
-              siblingLeafNodeIds={msg.siblingLeafNodeIds}
-              siblingCurrIdx={msg.siblingCurrIdx}
-              onRegenerateMessage={handleRegenerateMessage}
-              onEditMessage={handleEditMessage}
-              onChangeSibling={setCurrNodeId}
-              isPending={msg.isPending}
-            />
-          ))}
+          {[...messages, ...pendingMsgDisplay].map((msgDisplay) => {
+            const actualMsgObject = msgDisplay.msg;
+            // Check if the current message from the list is the one actively being generated/continued
+            const isThisMessageTheActivePendingOne =
+              pendingMsg?.id === actualMsgObject.id;
+
+            return (
+              <ChatMessage
+                key={actualMsgObject.id}
+                // If this message is the active pending one, use the live object from pendingMsg state (which has streamed content).
+                // Otherwise, use the version from the messages array (from storage).
+                msg={
+                  isThisMessageTheActivePendingOne
+                    ? pendingMsg
+                    : actualMsgObject
+                }
+                siblingLeafNodeIds={msgDisplay.siblingLeafNodeIds}
+                siblingCurrIdx={msgDisplay.siblingCurrIdx}
+                onRegenerateMessage={handleRegenerateMessage}
+                onEditMessage={handleEditMessage}
+                onChangeSibling={setCurrNodeId}
+                // A message is pending if it's the actively streaming one OR if it came from pendingMsgDisplay (for new messages)
+                isPending={
+                  isThisMessageTheActivePendingOne || msgDisplay.isPending
+                }
+                onContinueMessage={handleContinueMessage}
+              />
+            );
+          })}
         </div>
 
         {/* chat input */}

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

@@ -39,6 +39,12 @@ interface AppContextValue {
     extra: Message['extra'],
     onChunk: CallbackGeneratedChunk
   ) => Promise<void>;
+  continueMessageAndGenerate: (
+    convId: string,
+    messageIdToContinue: Message['id'],
+    newContent: string,
+    onChunk: CallbackGeneratedChunk
+  ) => Promise<void>;
 
   // canvas
   canvasData: CanvasData | null;
@@ -156,7 +162,8 @@ export const AppContextProvider = ({
   const generateMessage = async (
     convId: string,
     leafNodeId: Message['id'],
-    onChunk: CallbackGeneratedChunk
+    onChunk: CallbackGeneratedChunk,
+    isContinuation: boolean = false
   ) => {
     if (isGenerating(convId)) return;
 
@@ -179,17 +186,36 @@ export const AppContextProvider = ({
     }
 
     const pendingId = Date.now() + 1;
-    let pendingMsg: PendingMessage = {
-      id: pendingId,
-      convId,
-      type: 'text',
-      timestamp: pendingId,
-      role: 'assistant',
-      content: null,
-      parent: leafNodeId,
-      children: [],
-    };
-    setPending(convId, pendingMsg);
+    let pendingMsg: Message | PendingMessage;
+
+    if (isContinuation) {
+      const existingAsstMsg = await StorageUtils.getMessage(convId, leafNodeId);
+      if (!existingAsstMsg || existingAsstMsg.role !== 'assistant') {
+        toast.error(
+          'Cannot continue: target message not found or not an assistant message.'
+        );
+        throw new Error(
+          'Cannot continue: target message not found or not an assistant message.'
+        );
+      }
+      pendingMsg = {
+        ...existingAsstMsg,
+        content: existingAsstMsg.content || '',
+      };
+      setPending(convId, pendingMsg as PendingMessage);
+    } else {
+      pendingMsg = {
+        id: pendingId,
+        convId,
+        type: 'text',
+        timestamp: pendingId,
+        role: 'assistant',
+        content: null,
+        parent: leafNodeId,
+        children: [],
+      };
+      setPending(convId, pendingMsg as PendingMessage);
+    }
 
     try {
       // prepare messages for API
@@ -272,7 +298,7 @@ export const AppContextProvider = ({
             predicted_ms: timings.predicted_ms,
           };
         }
-        setPending(convId, pendingMsg);
+        setPending(convId, pendingMsg as PendingMessage);
         onChunk(); // don't need to switch node for pending message
       }
     } catch (err) {
@@ -289,10 +315,15 @@ export const AppContextProvider = ({
     }
 
     if (pendingMsg.content !== null) {
-      await StorageUtils.appendMsg(pendingMsg as Message, leafNodeId);
+      if (isContinuation) {
+        await StorageUtils.updateMessage(pendingMsg as Message);
+      } else if (pendingMsg.content.trim().length > 0) {
+        await StorageUtils.appendMsg(pendingMsg as Message, leafNodeId);
+      }
     }
     setPending(convId, null);
-    onChunk(pendingId); // trigger scroll to bottom and switch to the last node
+    const finalNodeId = (pendingMsg as Message).id;
+    onChunk(finalNodeId);
   };
 
   const sendMessage = async (
@@ -333,7 +364,7 @@ export const AppContextProvider = ({
     onChunk(currMsgId);
 
     try {
-      await generateMessage(convId, currMsgId, onChunk);
+      await generateMessage(convId, currMsgId, onChunk, false);
       return true;
     } catch (_) {
       // TODO: rollback
@@ -380,6 +411,47 @@ export const AppContextProvider = ({
     await generateMessage(convId, parentNodeId, onChunk);
   };
 
+  const continueMessageAndGenerate = async (
+    convId: string,
+    messageIdToContinue: Message['id'],
+    newContent: string,
+    onChunk: CallbackGeneratedChunk
+  ) => {
+    if (isGenerating(convId)) return;
+
+    const existingMessage = await StorageUtils.getMessage(
+      convId,
+      messageIdToContinue
+    );
+    if (!existingMessage || existingMessage.role !== 'assistant') {
+      console.error(
+        'Cannot continue non-assistant message or message not found'
+      );
+      toast.error(
+        'Failed to continue message: Not an assistant message or not found.'
+      );
+      return;
+    }
+
+    const updatedAssistantMessage: Message = {
+      ...existingMessage,
+      content: newContent,
+      timestamp: Date.now(),
+      children: [], // Clear existing children to start a new branch of generation
+      extra: existingMessage.extra, // Preserve existing extra data
+    };
+
+    await StorageUtils.updateMessage(updatedAssistantMessage);
+
+    onChunk(messageIdToContinue);
+
+    try {
+      await generateMessage(convId, messageIdToContinue, onChunk, true);
+    } catch (err) {
+      console.error('Error continuing message');
+    }
+  };
+
   const saveConfig = (config: typeof CONFIG_DEFAULT) => {
     StorageUtils.setConfig(config);
     setConfig(config);
@@ -394,6 +466,7 @@ export const AppContextProvider = ({
         sendMessage,
         stopGenerating,
         replaceMessageAndGenerate,
+        continueMessageAndGenerate,
         canvasData,
         setCanvasData,
         config,

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

@@ -213,6 +213,22 @@ const StorageUtils = {
       localStorage.setItem('theme', theme);
     }
   },
+  async getMessage(
+    convId: string,
+    messageId: Message['id']
+  ): Promise<Message | undefined> {
+    return await db.messages.where({ convId, id: messageId }).first();
+  },
+  async updateMessage(updatedMessage: Message): Promise<void> {
+    await db.transaction('rw', db.conversations, db.messages, async () => {
+      await db.messages.put(updatedMessage);
+      await db.conversations.update(updatedMessage.convId, {
+        lastModified: Date.now(),
+        currNode: updatedMessage.id,
+      });
+    });
+    dispatchConversationChange(updatedMessage.convId);
+  },
 };
 
 export default StorageUtils;