From 7ae6e35746649b37fb974e33e7018ec23faf346a Mon Sep 17 00:00:00 2001
From: Faisal Amir <urmauur@gmail.com>
Date: Mon, 12 Feb 2024 13:00:24 +0700
Subject: [PATCH] feat: add edit messages users (#1974)

* feat: add edit message user

* fix: delete messages should delete the rest below

---------

Co-authored-by: Louis <louis@jan.ai>
---
 web/containers/Providers/Jotai.tsx           |   1 +
 web/helpers/atoms/ChatMessage.atom.ts        |   2 +
 web/hooks/useSendChatMessage.ts              |   8 +-
 web/screens/Chat/EditChatInput/index.tsx     | 159 +++++++++++++++++++
 web/screens/Chat/MessageToolbar/index.tsx    |  30 +++-
 web/screens/Chat/SimpleTextMessage/index.tsx |  33 ++--
 6 files changed, 220 insertions(+), 13 deletions(-)
 create mode 100644 web/screens/Chat/EditChatInput/index.tsx

diff --git a/web/containers/Providers/Jotai.tsx b/web/containers/Providers/Jotai.tsx
index 103f0d9ee..5907ac746 100644
--- a/web/containers/Providers/Jotai.tsx
+++ b/web/containers/Providers/Jotai.tsx
@@ -8,6 +8,7 @@ type Props = {
   children: ReactNode
 }
 
+export const editPromptAtom = atom<string>('')
 export const currentPromptAtom = atom<string>('')
 export const fileUploadAtom = atom<FileInfo[]>([])
 export const appDownloadProgress = atom<number>(-1)
diff --git a/web/helpers/atoms/ChatMessage.atom.ts b/web/helpers/atoms/ChatMessage.atom.ts
index 45cc773e6..904ad344d 100644
--- a/web/helpers/atoms/ChatMessage.atom.ts
+++ b/web/helpers/atoms/ChatMessage.atom.ts
@@ -109,6 +109,8 @@ export const deleteMessageAtom = atom(null, (get, set, id: string) => {
   }
 })
 
+export const editMessageAtom = atom('')
+
 export const updateMessageAtom = atom(
   null,
   (
diff --git a/web/hooks/useSendChatMessage.ts b/web/hooks/useSendChatMessage.ts
index d7c2d10fd..e1c91cca2 100644
--- a/web/hooks/useSendChatMessage.ts
+++ b/web/hooks/useSendChatMessage.ts
@@ -24,7 +24,11 @@ import { atom, useAtom, useAtomValue, useSetAtom } from 'jotai'
 import { ulid } from 'ulid'
 
 import { selectedModelAtom } from '@/containers/DropdownListSidebar'
-import { currentPromptAtom, fileUploadAtom } from '@/containers/Providers/Jotai'
+import {
+  currentPromptAtom,
+  editPromptAtom,
+  fileUploadAtom,
+} from '@/containers/Providers/Jotai'
 
 import { getBase64 } from '@/utils/base64'
 import { toRuntimeParams, toSettingParams } from '@/utils/modelParam'
@@ -54,6 +58,7 @@ export default function useSendChatMessage() {
   const updateThread = useSetAtom(updateThreadAtom)
   const updateThreadWaiting = useSetAtom(updateThreadWaitingForResponseAtom)
   const setCurrentPrompt = useSetAtom(currentPromptAtom)
+  const setEditPrompt = useSetAtom(editPromptAtom)
 
   const currentMessages = useAtomValue(getCurrentChatMessagesAtom)
   const { activeModel } = useActiveModel()
@@ -147,6 +152,7 @@ export default function useSendChatMessage() {
     updateThreadWaiting(activeThread.id, true)
     const prompt = message.trim()
     setCurrentPrompt('')
+    setEditPrompt('')
 
     const base64Blob = fileUpload[0]
       ? await getBase64(fileUpload[0].file).then()
diff --git a/web/screens/Chat/EditChatInput/index.tsx b/web/screens/Chat/EditChatInput/index.tsx
new file mode 100644
index 000000000..29b81118b
--- /dev/null
+++ b/web/screens/Chat/EditChatInput/index.tsx
@@ -0,0 +1,159 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+import { useEffect, useRef } from 'react'
+
+import {
+  ConversationalExtension,
+  ExtensionTypeEnum,
+  InferenceEvent,
+  MessageStatus,
+  ThreadMessage,
+  events,
+} from '@janhq/core'
+
+import { Textarea, Button } from '@janhq/uikit'
+import { useAtom, useAtomValue, useSetAtom } from 'jotai'
+
+import { twMerge } from 'tailwind-merge'
+
+import { editPromptAtom } from '@/containers/Providers/Jotai'
+
+import { useActiveModel } from '@/hooks/useActiveModel'
+
+import useSendChatMessage from '@/hooks/useSendChatMessage'
+
+import { extensionManager } from '@/extension'
+
+import {
+  editMessageAtom,
+  getCurrentChatMessagesAtom,
+  setConvoMessagesAtom,
+} from '@/helpers/atoms/ChatMessage.atom'
+import {
+  activeThreadAtom,
+  getActiveThreadIdAtom,
+  waitingToSendMessage,
+} from '@/helpers/atoms/Thread.atom'
+
+type Props = {
+  message: ThreadMessage
+}
+
+const EditChatInput: React.FC<Props> = ({ message }) => {
+  const activeThread = useAtomValue(activeThreadAtom)
+  const { stateModel } = useActiveModel()
+  const messages = useAtomValue(getCurrentChatMessagesAtom)
+
+  const [editPrompt, setEditPrompt] = useAtom(editPromptAtom)
+  const { sendChatMessage } = useSendChatMessage()
+  const setMessages = useSetAtom(setConvoMessagesAtom)
+  const activeThreadId = useAtomValue(getActiveThreadIdAtom)
+
+  const [isWaitingToSend, setIsWaitingToSend] = useAtom(waitingToSendMessage)
+  const textareaRef = useRef<HTMLTextAreaElement>(null)
+  const setEditMessage = useSetAtom(editMessageAtom)
+
+  const onPromptChange = (e: React.ChangeEvent<HTMLTextAreaElement>) => {
+    setEditPrompt(e.target.value)
+  }
+
+  useEffect(() => {
+    if (isWaitingToSend && activeThreadId) {
+      setIsWaitingToSend(false)
+      sendChatMessage(editPrompt)
+    }
+  }, [
+    activeThreadId,
+    isWaitingToSend,
+    editPrompt,
+    setIsWaitingToSend,
+    sendChatMessage,
+  ])
+
+  useEffect(() => {
+    if (textareaRef.current) {
+      textareaRef.current.focus()
+    }
+  }, [activeThreadId])
+
+  useEffect(() => {
+    if (textareaRef.current) {
+      textareaRef.current.style.height = '40px'
+      textareaRef.current.style.height = textareaRef.current.scrollHeight + 'px'
+    }
+  }, [editPrompt])
+
+  useEffect(() => {
+    setEditPrompt(message.content[0]?.text?.value)
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [])
+
+  const sendEditMessage = async () => {
+    setEditMessage('')
+    const messageIdx = messages.findIndex((msg) => msg.id === message.id)
+    const newMessages = messages.slice(0, messageIdx)
+    if (activeThread) {
+      setMessages(activeThread.id, newMessages)
+      await extensionManager
+        .get<ConversationalExtension>(ExtensionTypeEnum.Conversational)
+        ?.writeMessages(
+          activeThread.id,
+          // Remove all of the messages below this
+          newMessages
+        )
+        .then(() => {
+          sendChatMessage(editPrompt)
+        })
+    }
+  }
+
+  const onKeyDown = async (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
+    if (e.key === 'Enter' && !e.shiftKey) {
+      e.preventDefault()
+      if (messages[messages.length - 1]?.status !== MessageStatus.Pending)
+        sendEditMessage()
+      else onStopInferenceClick()
+    }
+  }
+
+  const onStopInferenceClick = async () => {
+    events.emit(InferenceEvent.OnInferenceStopped, {})
+  }
+
+  return (
+    <div className="mx-auto flex w-full flex-shrink-0 items-end justify-center space-x-4 pb-0 pt-1">
+      <div className="relative flex w-full flex-col">
+        <Textarea
+          className={twMerge(
+            'max-h-[400px] resize-none overflow-y-hidden pr-20'
+          )}
+          style={{ height: '40px' }}
+          ref={textareaRef}
+          onKeyDown={onKeyDown}
+          placeholder="Enter your message..."
+          disabled={stateModel.loading || !activeThread}
+          value={editPrompt}
+          onChange={onPromptChange}
+        />
+      </div>
+
+      <div className="flex items-center space-x-2">
+        <Button
+          disabled={
+            stateModel.loading ||
+            !activeThread ||
+            editPrompt.trim().length === 0
+          }
+          themes="primary"
+          onClick={sendEditMessage}
+        >
+          Submit
+        </Button>
+        <Button themes="outline" onClick={() => setEditMessage('')}>
+          Cancel
+        </Button>
+      </div>
+    </div>
+  )
+}
+
+export default EditChatInput
diff --git a/web/screens/Chat/MessageToolbar/index.tsx b/web/screens/Chat/MessageToolbar/index.tsx
index 070022122..ed73d5ecb 100644
--- a/web/screens/Chat/MessageToolbar/index.tsx
+++ b/web/screens/Chat/MessageToolbar/index.tsx
@@ -6,8 +6,14 @@ import {
   ConversationalExtension,
   ContentType,
 } from '@janhq/core'
-import { useAtomValue, useSetAtom } from 'jotai'
-import { RefreshCcw, CopyIcon, Trash2Icon, CheckIcon } from 'lucide-react'
+import { useAtom, useAtomValue, useSetAtom } from 'jotai'
+import {
+  RefreshCcw,
+  CopyIcon,
+  Trash2Icon,
+  CheckIcon,
+  PencilIcon,
+} from 'lucide-react'
 
 import { twMerge } from 'tailwind-merge'
 
@@ -17,12 +23,14 @@ import useSendChatMessage from '@/hooks/useSendChatMessage'
 import { extensionManager } from '@/extension'
 import {
   deleteMessageAtom,
+  editMessageAtom,
   getCurrentChatMessagesAtom,
 } from '@/helpers/atoms/ChatMessage.atom'
 import { activeThreadAtom } from '@/helpers/atoms/Thread.atom'
 
 const MessageToolbar = ({ message }: { message: ThreadMessage }) => {
   const deleteMessage = useSetAtom(deleteMessageAtom)
+  const [editMessage, setEditMessage] = useAtom(editMessageAtom)
   const thread = useAtomValue(activeThreadAtom)
   const messages = useAtomValue(getCurrentChatMessagesAtom)
   const { resendChatMessage } = useSendChatMessage()
@@ -40,6 +48,14 @@ const MessageToolbar = ({ message }: { message: ThreadMessage }) => {
     }
   }
 
+  const onEditClick = async () => {
+    if (!editMessage.length) {
+      setEditMessage(message.id ?? '')
+    } else {
+      setEditMessage('')
+    }
+  }
+
   const onRegenerateClick = async () => {
     if (message.role !== ChatCompletionRole.User) {
       // Delete last response before regenerating
@@ -53,6 +69,16 @@ const MessageToolbar = ({ message }: { message: ThreadMessage }) => {
   return (
     <div className={twMerge('flex flex-row items-center')}>
       <div className="flex overflow-hidden rounded-md border border-border bg-background/20">
+        {message.role === ChatCompletionRole.User &&
+          message.content[0]?.type === ContentType.Text && (
+            <div
+              className="cursor-pointer border-r border-border px-2 py-2 hover:bg-background/80"
+              onClick={onEditClick}
+            >
+              <PencilIcon size={14} />
+            </div>
+          )}
+
         {message.id === messages[messages.length - 1]?.id &&
           messages[messages.length - 1].status !== MessageStatus.Error &&
           messages[messages.length - 1].content[0]?.type !==
diff --git a/web/screens/Chat/SimpleTextMessage/index.tsx b/web/screens/Chat/SimpleTextMessage/index.tsx
index 9be45e7e6..f064132b8 100644
--- a/web/screens/Chat/SimpleTextMessage/index.tsx
+++ b/web/screens/Chat/SimpleTextMessage/index.tsx
@@ -32,10 +32,14 @@ import { usePath } from '@/hooks/usePath'
 import { toGibibytes } from '@/utils/converter'
 import { displayDate } from '@/utils/datetime'
 
+import EditChatInput from '../EditChatInput'
 import Icon from '../FileUploadPreview/Icon'
 import MessageToolbar from '../MessageToolbar'
 
-import { getCurrentChatMessagesAtom } from '@/helpers/atoms/ChatMessage.atom'
+import {
+  editMessageAtom,
+  getCurrentChatMessagesAtom,
+} from '@/helpers/atoms/ChatMessage.atom'
 
 function isMarkdownValue(value: string): boolean {
   const tokenTypes: string[] = []
@@ -54,6 +58,7 @@ const SimpleTextMessage: React.FC<ThreadMessage> = (props) => {
   let text = ''
   const isUser = props.role === ChatCompletionRole.User
   const isSystem = props.role === ChatCompletionRole.System
+  const editMessage = useAtomValue(editMessageAtom)
 
   if (props.content && props.content.length > 0) {
     text = props.content[0]?.text?.value ?? ''
@@ -276,16 +281,24 @@ const SimpleTextMessage: React.FC<ThreadMessage> = (props) => {
           )}
 
           {isUser && !isMarkdownValue(text) ? (
-            <div
-              className={twMerge(
-                'message flex flex-grow flex-col gap-y-2 text-[15px] font-normal leading-relaxed',
-                isUser
-                  ? 'whitespace-pre-wrap break-words'
-                  : 'rounded-xl bg-secondary p-4'
+            <>
+              {editMessage === props.id ? (
+                <div>
+                  <EditChatInput message={props} />
+                </div>
+              ) : (
+                <div
+                  className={twMerge(
+                    'message flex flex-grow flex-col gap-y-2 text-[15px] font-normal leading-relaxed',
+                    isUser
+                      ? 'whitespace-pre-wrap break-words'
+                      : 'rounded-xl bg-secondary p-4'
+                  )}
+                >
+                  {text}
+                </div>
               )}
-            >
-              {text}
-            </div>
+            </>
           ) : (
             <div
               className={twMerge(