From 08383a691006cb342a9cf275b7f9a9e212b76f75 Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Wed, 6 Dec 2023 11:16:12 +0700
Subject: [PATCH] fix: delete message break the entire thread

---
 .../conversational-extension/src/index.ts     | 20 ++++++++++++++-----
 web/hooks/useSendChatMessage.ts               | 16 ++++++++++-----
 web/screens/Chat/MessageToolbar/index.tsx     | 11 +++++++++-
 3 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/extensions/conversational-extension/src/index.ts b/extensions/conversational-extension/src/index.ts
index 4461cfaf7d..2c2c7e47f6 100644
--- a/extensions/conversational-extension/src/index.ts
+++ b/extensions/conversational-extension/src/index.ts
@@ -7,7 +7,9 @@ import { join } from 'path'
  * JSONConversationalExtension is a ConversationalExtension implementation that provides
  * functionality for managing threads.
  */
-export default class JSONConversationalExtension implements ConversationalExtension {
+export default class JSONConversationalExtension
+  implements ConversationalExtension
+{
   private static readonly _homeDir = 'threads'
   private static readonly _threadInfoFileName = 'thread.json'
   private static readonly _threadMessagesFileName = 'messages.jsonl'
@@ -67,7 +69,10 @@ export default class JSONConversationalExtension implements ConversationalExtens
    */
   async saveThread(thread: Thread): Promise<void> {
     try {
-      const threadDirPath = join(JSONConversationalExtension._homeDir, thread.id)
+      const threadDirPath = join(
+        JSONConversationalExtension._homeDir,
+        thread.id
+      )
       const threadJsonPath = join(
         threadDirPath,
         JSONConversationalExtension._threadInfoFileName
@@ -119,7 +124,7 @@ export default class JSONConversationalExtension implements ConversationalExtens
       await fs.mkdir(threadDirPath)
       await fs.writeFile(
         threadMessagePath,
-        messages.map((msg) => JSON.stringify(msg)).join('\n')
+        messages.map((msg) => JSON.stringify(msg)).join('\n') + '\n'
       )
       Promise.resolve()
     } catch (err) {
@@ -153,7 +158,10 @@ export default class JSONConversationalExtension implements ConversationalExtens
 
     const threadDirs: string[] = []
     for (let i = 0; i < fileInsideThread.length; i++) {
-      const path = join(JSONConversationalExtension._homeDir, fileInsideThread[i])
+      const path = join(
+        JSONConversationalExtension._homeDir,
+        fileInsideThread[i]
+      )
       const isDirectory = await fs.isDirectory(path)
       if (!isDirectory) {
         console.debug(`Ignore ${path} because it is not a directory`)
@@ -182,7 +190,9 @@ export default class JSONConversationalExtension implements ConversationalExtens
       }
 
       const files: string[] = await fs.listFiles(threadDirPath)
-      if (!files.includes(JSONConversationalExtension._threadMessagesFileName)) {
+      if (
+        !files.includes(JSONConversationalExtension._threadMessagesFileName)
+      ) {
         throw Error(`${threadDirPath} not contains message file`)
       }
 
diff --git a/web/hooks/useSendChatMessage.ts b/web/hooks/useSendChatMessage.ts
index 3e4a8f0261..26861508da 100644
--- a/web/hooks/useSendChatMessage.ts
+++ b/web/hooks/useSendChatMessage.ts
@@ -48,7 +48,7 @@ export default function useSendChatMessage() {
   const { startModel } = useActiveModel()
   const [queuedMessage, setQueuedMessage] = useState(false)
 
-  const resendChatMessage = async () => {
+  const resendChatMessage = async (currentMessage: ThreadMessage) => {
     if (!activeThread) {
       console.error('No active thread')
       return
@@ -68,10 +68,16 @@ export default function useSendChatMessage() {
         return systemMessage
       })
       .concat(
-        currentMessages.map<ChatCompletionMessage>((msg) => ({
-          role: msg.role,
-          content: msg.content[0]?.text.value ?? '',
-        }))
+        currentMessages
+          .filter(
+            (e) =>
+              currentMessage.role === ChatCompletionRole.User ||
+              e.id !== currentMessage.id
+          )
+          .map<ChatCompletionMessage>((msg) => ({
+            role: msg.role,
+            content: msg.content[0]?.text.value ?? '',
+          }))
       )
 
     const messageRequest: MessageRequest = {
diff --git a/web/screens/Chat/MessageToolbar/index.tsx b/web/screens/Chat/MessageToolbar/index.tsx
index e5f3c245ff..f877e1bdb2 100644
--- a/web/screens/Chat/MessageToolbar/index.tsx
+++ b/web/screens/Chat/MessageToolbar/index.tsx
@@ -4,6 +4,7 @@ import {
   ExtensionType,
   ThreadMessage,
   events,
+  ChatCompletionRole,
 } from '@janhq/core'
 import { ConversationalExtension, InferenceExtension } from '@janhq/core'
 import { useAtomValue, useSetAtom } from 'jotai'
@@ -52,6 +53,14 @@ const MessageToolbar = ({ message }: { message: ThreadMessage }) => {
     }
   }
 
+  const onRegenerateClick = async () => {
+    if (message.role !== ChatCompletionRole.User) {
+      // Delete last response before regenerating
+      await onDeleteClick()
+    }
+    resendChatMessage(message)
+  }
+
   return (
     <div className={twMerge('flex flex-row items-center')}>
       <div className="flex overflow-hidden rounded-md border border-border bg-background/20">
@@ -67,7 +76,7 @@ const MessageToolbar = ({ message }: { message: ThreadMessage }) => {
           message.id === messages[messages.length - 1]?.id && (
             <div
               className="cursor-pointer border-r border-border px-2 py-2 hover:bg-background/80"
-              onClick={resendChatMessage}
+              onClick={onRegenerateClick}
             >
               <RefreshCcw size={14} />
             </div>