fix: remove files and memory when user clean thread (janhq#2524)

Signed-off-by: James <[email protected]> Co-authored-by: James <[email protected]>
feng-huang · Mar 28, 2024 · 1f8dc89 · 1f8dc89
1 parent 54ba410
commit 1f8dc89
Showing 1 changed file with 53 additions and 24 deletions.
diff --git a/web/hooks/useDeleteThread.ts b/web/hooks/useDeleteThread.ts
@@ -1,7 +1,11 @@
+import { useCallback } from 'react'
+
 import {
   ChatCompletionRole,
   ExtensionTypeEnum,
   ConversationalExtension,
+  fs,
+  joinPath,
 } from '@janhq/core'
 
 import { useAtom, useAtomValue, useSetAtom } from 'jotai'
@@ -12,10 +16,11 @@ import { toaster } from '@/containers/Toast'
 
 import { extensionManager } from '@/extension/ExtensionManager'
 
+import { janDataFolderPathAtom } from '@/helpers/atoms/AppConfig.atom'
 import {
+  chatMessages,
   cleanChatMessageAtom as cleanChatMessagesAtom,
   deleteChatMessageAtom as deleteChatMessagesAtom,
-  getCurrentChatMessagesAtom,
 } from '@/helpers/atoms/ChatMessage.atom'
 import {
   threadsAtom,
@@ -26,39 +31,63 @@ import {
 
 export default function useDeleteThread() {
   const [threads, setThreads] = useAtom(threadsAtom)
-  const setCurrentPrompt = useSetAtom(currentPromptAtom)
-  const messages = useAtomValue(getCurrentChatMessagesAtom)
+  const messages = useAtomValue(chatMessages)
+  const janDataFolderPath = useAtomValue(janDataFolderPathAtom)
 
+  const setCurrentPrompt = useSetAtom(currentPromptAtom)
   const setActiveThreadId = useSetAtom(setActiveThreadIdAtom)
   const deleteMessages = useSetAtom(deleteChatMessagesAtom)
   const cleanMessages = useSetAtom(cleanChatMessagesAtom)
+
   const deleteThreadState = useSetAtom(deleteThreadStateAtom)
   const updateThreadLastMessage = useSetAtom(updateThreadStateLastMessageAtom)
 
-  const cleanThread = async (threadId: string) => {
-    if (threadId) {
-      const thread = threads.filter((c) => c.id === threadId)[0]
+  const cleanThread = useCallback(
+    async (threadId: string) => {
       cleanMessages(threadId)
+      const thread = threads.find((c) => c.id === threadId)
+      if (!thread) return
 
-      if (thread) {
-        await extensionManager
-          .get<ConversationalExtension>(ExtensionTypeEnum.Conversational)
-          ?.writeMessages(
-            threadId,
-            messages.filter((msg) => msg.role === ChatCompletionRole.System)
-          )
-
-        thread.metadata = {
-          ...thread.metadata,
-          lastMessage: undefined,
-        }
-        await extensionManager
-          .get<ConversationalExtension>(ExtensionTypeEnum.Conversational)
-          ?.saveThread(thread)
-        updateThreadLastMessage(threadId, undefined)
+      const updatedMessages = (messages[threadId] ?? []).filter(
+        (msg) => msg.role === ChatCompletionRole.System
+      )
+
+      // remove files
+      try {
+        const threadFolderPath = await joinPath([
+          janDataFolderPath,
+          'threads',
+          threadId,
+        ])
+        const threadFilesPath = await joinPath([threadFolderPath, 'files'])
+        const threadMemoryPath = await joinPath([threadFolderPath, 'memory'])
+        await fs.rm(threadFilesPath)
+        await fs.rm(threadMemoryPath)
+      } catch (err) {
+        console.warn('Error deleting thread files', err)
       }
-    }
-  }
+
+      await extensionManager
+        .get<ConversationalExtension>(ExtensionTypeEnum.Conversational)
+        ?.writeMessages(threadId, updatedMessages)
+
+      thread.metadata = {
+        ...thread.metadata,
+        lastMessage: undefined,
+      }
+      await extensionManager
+        .get<ConversationalExtension>(ExtensionTypeEnum.Conversational)
+        ?.saveThread(thread)
+      updateThreadLastMessage(threadId, undefined)
+    },
+    [
+      janDataFolderPath,
+      threads,
+      messages,
+      cleanMessages,
+      updateThreadLastMessage,
+    ]
+  )
 
   const deleteThread = async (threadId: string) => {
     if (!threadId) {