1 files changed, 81 insertions, 18 deletions

diff --git a/examples/server/webui/src/utils/app.context.tsx b/examples/server/webui/src/utils/app.context.tsx
index ce81fccd..509b4ee0 100644
--- a/examples/server/webui/src/utils/app.context.tsx
+++ b/examples/server/webui/src/utils/app.context.tsx
@@ -15,7 +15,7 @@ import {
 } from './misc';
 import { BASE_URL, CONFIG_DEFAULT, isDev } from '../Config';
 import { matchPath, useLocation, useNavigate } from 'react-router';
-
+import toast from 'react-hot-toast';
 class Timer {
 	static timercount = 1;
 }
@@ -39,7 +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;
   setCanvasData: (data: CanvasData | null) => void;
@@ -136,7 +141,8 @@ export const AppContextProvider = ({
   const generateMessage = async (
     convId: string,
     leafNodeId: Message['id'],
-    onChunk: CallbackGeneratedChunk
+    onChunk: CallbackGeneratedChunk,
+    isContinuation: boolean = false
   ) => {
     if (isGenerating(convId)) return;
 
@@ -160,17 +166,36 @@ export const AppContextProvider = ({
 
     const pendingId = Date.now() + Timer.timercount + 1;
 	Timer.timercount=Timer.timercount+2;
-    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
@@ -254,7 +279,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) {
@@ -271,11 +296,16 @@ export const AppContextProvider = ({
     }
 	finally {
 		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); // trigger scroll to bottom and switch to the last node
   };
 
   const sendMessage = async (
@@ -317,7 +347,7 @@ export const AppContextProvider = ({
     onChunk(currMsgId);
 
     try {
-      await generateMessage(convId, currMsgId, onChunk);
+      await generateMessage(convId, currMsgId, onChunk, false);
       return true;
     } catch (_) {
       // TODO: rollback
@@ -364,6 +394,38 @@ 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,
+    };
+      //children: [], // Clear existing children to start a new branch of generation
+
+    await StorageUtils.updateMessage(updatedAssistantMessage);
+    onChunk;
+  };
+
+ 
   const saveConfig = (config: typeof CONFIG_DEFAULT) => {
     StorageUtils.setConfig(config);
     setConfig(config);
@@ -378,6 +440,7 @@ export const AppContextProvider = ({
         sendMessage,
         stopGenerating,
         replaceMessageAndGenerate,
+        continueMessageAndGenerate,
         canvasData,
         setCanvasData,
         config,