From ca5a6831b69b06492fa0d8c52e74c8d661215fd4 Mon Sep 17 00:00:00 2001
From: Debanjum <debanjum@gmail.com>
Date: Wed, 30 Oct 2024 04:24:22 -0700
Subject: [PATCH] Add ability to delete messages from the web app

---
 .../components/chatHistory/chatHistory.tsx    | 26 +++++++++++++
 .../components/chatMessage/chatMessage.tsx    | 37 +++++++++++++++++++
 src/interface/web/app/factchecker/page.tsx    | 10 +++++
 src/khoj/database/adapters/__init__.py        |  6 +--
 4 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/src/interface/web/app/components/chatHistory/chatHistory.tsx b/src/interface/web/app/components/chatHistory/chatHistory.tsx
index 9e156142..5075f382 100644
--- a/src/interface/web/app/components/chatHistory/chatHistory.tsx
+++ b/src/interface/web/app/components/chatHistory/chatHistory.tsx
@@ -278,6 +278,16 @@ export default function ChatHistory(props: ChatHistoryProps) {
         return data.agent?.persona;
     }
 
+    const handleDeleteMessage = (turnId?: string) => {
+        setData((prevData) => {
+            if (!prevData || !turnId) return prevData;
+            return {
+                ...prevData,
+                chat: prevData.chat.filter((msg) => msg.turnId !== turnId),
+            };
+        });
+    };
+
     if (!props.conversationId && !props.publicConversationSlug) {
         return null;
     }
@@ -312,6 +322,8 @@ export default function ChatHistory(props: ChatHistoryProps) {
                                     customClassName="fullHistory"
                                     borderLeftColor={`${data?.agent?.color}-500`}
                                     isLastMessage={index === data.chat.length - 1}
+                                    onDeleteMessage={handleDeleteMessage}
+                                    conversationId={props.conversationId}
                                 />
                                 {chatMessage.trainOfThought && chatMessage.by === "khoj" && (
                                     <TrainOfThoughtComponent
@@ -342,9 +354,14 @@ export default function ChatHistory(props: ChatHistoryProps) {
                                             by: "you",
                                             automationId: "",
                                             images: message.images,
+                                            conversationId: props.conversationId,
+                                            turnId: message.turnId,
+                                            onDeleteMessage: handleDeleteMessage,
                                         }}
                                         customClassName="fullHistory"
                                         borderLeftColor={`${data?.agent?.color}-500`}
+                                        onDeleteMessage={handleDeleteMessage}
+                                        conversationId={props.conversationId}
                                     />
                                     {message.trainOfThought && (
                                         <TrainOfThoughtComponent
@@ -373,7 +390,12 @@ export default function ChatHistory(props: ChatHistoryProps) {
                                                 "memory-type": "",
                                                 "inferred-queries": message.inferredQueries || [],
                                             },
+                                            conversationId: props.conversationId,
+                                            turnId: message.turnId,
+                                            onDeleteMessage: handleDeleteMessage,
                                         }}
+                                        conversationId={props.conversationId}
+                                        onDeleteMessage={handleDeleteMessage}
                                         customClassName="fullHistory"
                                         borderLeftColor={`${data?.agent?.color}-500`}
                                         isLastMessage={true}
@@ -393,7 +415,11 @@ export default function ChatHistory(props: ChatHistoryProps) {
                                 created: new Date().getTime().toString(),
                                 by: "you",
                                 automationId: "",
+                                conversationId: props.conversationId,
+                                onDeleteMessage: handleDeleteMessage,
                             }}
+                            conversationId={props.conversationId}
+                            onDeleteMessage={handleDeleteMessage}
                             customClassName="fullHistory"
                             borderLeftColor={`${data?.agent?.color}-500`}
                             isLastMessage={true}
diff --git a/src/interface/web/app/components/chatMessage/chatMessage.tsx b/src/interface/web/app/components/chatMessage/chatMessage.tsx
index 6d82fe1d..8f94db25 100644
--- a/src/interface/web/app/components/chatMessage/chatMessage.tsx
+++ b/src/interface/web/app/components/chatMessage/chatMessage.tsx
@@ -29,6 +29,7 @@ import {
     Check,
     Code,
     Shapes,
+    Trash,
 } from "@phosphor-icons/react";
 
 import DOMPurify from "dompurify";
@@ -146,6 +147,9 @@ export interface SingleChatMessage {
     intent?: Intent;
     agent?: AgentData;
     images?: string[];
+    conversationId: string;
+    turnId?: string;
+    onDeleteMessage: (turnId: string) => void;
 }
 
 export interface StreamMessage {
@@ -161,6 +165,7 @@ export interface StreamMessage {
     images?: string[];
     intentType?: string;
     inferredQueries?: string[];
+    turnId?: string;
 }
 
 export interface ChatHistoryData {
@@ -242,6 +247,8 @@ interface ChatMessageProps {
     borderLeftColor?: string;
     isLastMessage?: boolean;
     agent?: AgentData;
+    onDeleteMessage: (turnId?: string) => void;
+    conversationId: string;
 }
 
 interface TrainOfThoughtProps {
@@ -654,6 +661,26 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
         });
     }
 
+    const deleteMessage = async (message: SingleChatMessage) => {
+        const response = await fetch("/api/chat/conversation/message", {
+            method: "DELETE",
+            headers: {
+                "Content-Type": "application/json",
+            },
+            body: JSON.stringify({
+                conversation_id: props.conversationId,
+                turn_id: message.turnId,
+            }),
+        });
+
+        if (response.ok) {
+            // Update the UI after successful deletion
+            props.onDeleteMessage(message.turnId);
+        } else {
+            console.error("Failed to delete message");
+        }
+    };
+
     const allReferences = constructAllReferences(
         props.chatMessage.context,
         props.chatMessage.onlineContext,
@@ -737,6 +764,16 @@ const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>((props, ref) =>
                                     />
                                 )}
                             </button>
+                            <button
+                                title="Delete"
+                                className={`${styles.deleteButton}`}
+                                onClick={() => deleteMessage(props.chatMessage)}
+                            >
+                                <Trash
+                                    alt="Delete Message"
+                                    className="hsl(var(--muted-foreground)) hover:text-red-500"
+                                />
+                            </button>
                             {props.chatMessage.by === "khoj" &&
                                 (props.chatMessage.intent ? (
                                     <FeedbackButtons
diff --git a/src/interface/web/app/factchecker/page.tsx b/src/interface/web/app/factchecker/page.tsx
index 8f5538f6..3caed53f 100644
--- a/src/interface/web/app/factchecker/page.tsx
+++ b/src/interface/web/app/factchecker/page.tsx
@@ -195,8 +195,13 @@ function ReferenceVerification(props: ReferenceVerificationProps) {
                     created: new Date().toISOString(),
                     onlineContext: {},
                     codeContext: {},
+                    conversationId: props.conversationId,
+                    turnId: "",
+                    onDeleteMessage: (turnId?: string) => {},
                 }}
                 isMobileWidth={isMobileWidth}
+                onDeleteMessage={(turnId?: string) => {}}
+                conversationId={props.conversationId}
             />
         </div>
     );
@@ -626,7 +631,12 @@ export default function FactChecker() {
                                         created: new Date().toISOString(),
                                         onlineContext: {},
                                         codeContext: {},
+                                        conversationId: conversationID,
+                                        turnId: "",
+                                        onDeleteMessage: (turnId?: string) => {},
                                     }}
+                                    conversationId={conversationID}
+                                    onDeleteMessage={(turnId?: string) => {}}
                                     isMobileWidth={isMobileWidth}
                                 />
                             </div>
diff --git a/src/khoj/database/adapters/__init__.py b/src/khoj/database/adapters/__init__.py
index bfabc17a..d8d8e644 100644
--- a/src/khoj/database/adapters/__init__.py
+++ b/src/khoj/database/adapters/__init__.py
@@ -1351,11 +1351,11 @@ class ConversationAdapters:
     @staticmethod
     def delete_message_by_turn_id(user: KhojUser, conversation_id: str, turn_id: str):
         conversation = ConversationAdapters.get_conversation_by_user(user, conversation_id=conversation_id)
-        if not conversation:
+        if not conversation or not conversation.conversation_log or not conversation.conversation_log.get("chat"):
             return False
         conversation_log = conversation.conversation_log
-        updated_log = [msg for msg in conversation_log if msg.get("turnId") != turn_id]
-        conversation.conversation_log = updated_log
+        updated_log = [msg for msg in conversation_log["chat"] if msg.get("turnId") != turn_id]
+        conversation.conversation_log["chat"] = updated_log
         conversation.save()
         return True