Add edit functionality to PlaygroundMessage component

src/hooks/useMessageOption.tsx

@@ -15,10 +15,12 @@ import {
 } from "@langchain/core/messages"
 import { useStoreMessageOption } from "~store/option"
 import {
+  deleteChatForEdit,
   getPromptById,
   removeMessageUsingHistoryId,
   saveHistory,
-  saveMessage
+  saveMessage,
+  updateMessageByIndex
 } from "~libs/db"
 import { useNavigate } from "react-router-dom"
 import { notification } from "antd"
@@ -114,6 +116,8 @@ export const useMessageOption = () => {
     setSelectedSystemPrompt
   } = useStoreMessageOption()
 
+  // const { notification } = App.useApp()
+
   const navigate = useNavigate()
   const textareaRef = React.useRef<HTMLTextAreaElement>(null)
 
@@ -134,7 +138,9 @@ export const useMessageOption = () => {
   const searchChatMode = async (
     message: string,
     image: string,
-    isRegenerate: boolean
+    isRegenerate: boolean,
+    messages: Message[],
+    history: ChatHistory
   ) => {
     const url = await getOllamaURL()
 
@@ -387,7 +393,9 @@ export const useMessageOption = () => {
   const normalChatMode = async (
     message: string,
     image: string,
-    isRegenerate: boolean
+    isRegenerate: boolean,
+    messages: Message[],
+    history: ChatHistory
   ) => {
     const url = await getOllamaURL()
 
@@ -625,21 +633,42 @@ export const useMessageOption = () => {
   const onSubmit = async ({
     message,
     image,
-    isRegenerate = false
+    isRegenerate = false,
+    messages: chatHistory,
+    memory
   }: {
     message: string
     image: string
     isRegenerate?: boolean
+    messages?: Message[]
+    memory?: ChatHistory
   }) => {
     setStreaming(true)
     if (webSearch) {
-      await searchChatMode(message, image, isRegenerate)
+      await searchChatMode(
+        message,
+        image,
+        isRegenerate,
+        chatHistory || messages,
+        memory || history
+      )
     } else {
-      await normalChatMode(message, image, isRegenerate)
+      await normalChatMode(
+        message,
+        image,
+        isRegenerate,
+        chatHistory || messages,
+        memory || history
+      )
     }
   }
 
   const regenerateLastMessage = async () => {
+    const isOk = validateBeforeSubmit()
+
+    if (!isOk) {
+      return
+    }
     if (history.length > 0) {
       const lastMessage = history[history.length - 2]
       let newHistory = history
@@ -653,7 +682,8 @@ export const useMessageOption = () => {
         await onSubmit({
           message: lastMessage.content,
           image: lastMessage.image || "",
-          isRegenerate: true
+          isRegenerate: true,
+          memory: newHistory
         })
       }
     }
@@ -666,7 +696,61 @@ export const useMessageOption = () => {
     }
   }
 
+  const validateBeforeSubmit = () => {
+    if (!selectedModel || selectedModel?.trim()?.length === 0) {
+      notification.error({
+        message: "Error",
+        description: "Please select a model to continue"
+      })
+      return false
+    }
+
+    return true
+  }
+
+  const editMessage = async (
+    index: number,
+    message: string,
+    isHuman: boolean
+  ) => {
+    // update message and history by index
+    let newMessages = messages
+    let newHistory = history
+
+    if (isHuman) {
+      const isOk = validateBeforeSubmit()
+
+      if (!isOk) {
+        return
+      }
+
+      const currentHumanMessage = newMessages[index]
+      newMessages[index].message = message
+      newHistory[index].content = message
+      const previousMessages = newMessages.slice(0, index + 1)
+      setMessages(previousMessages)
+      const previousHistory = newHistory.slice(0, index + 1)
+      setHistory(previousHistory)
+      await updateMessageByIndex(historyId, index, message)
+      await deleteChatForEdit(historyId, index)
+      await onSubmit({
+        message: message,
+        image: currentHumanMessage.images[0] || "",
+        isRegenerate: true,
+        messages: previousMessages,
+        memory: previousHistory
+      })
+    } else {
+      newMessages[index].message = message
+      setMessages(newMessages)
+      newHistory[index].content = message
+      setHistory(newHistory)
+      await updateMessageByIndex(historyId, index, message)
+    }
+  }
+
   return {
+    editMessage,
     messages,
     setMessages,
     onSubmit,