Chat history editing

backend/danswer/tools/utils.py

@@ -7,7 +7,7 @@ from danswer.llm.utils import get_default_llm_tokenizer
 from danswer.tools.tool import Tool
 
 
-OPEN_AI_TOOL_CALLING_MODELS = {"gpt-3.5-turbo", "gpt-4-turbo", "gpt-4", "gpt-4o"}
+OPEN_AI_TOOL_CALLING_MODELS = {"gpt-3.5-turbo", "gpt-4-turbo", "gpt-4"}
 
 
 def explicit_tool_calling_supported(model_provider: str, model_name: str) -> bool:

web/src/app/chat/ChatPage.tsx

@@ -23,6 +23,7 @@ import { InstantSSRAutoRefresh } from "@/components/SSRAutoRefresh";
 import { Settings } from "../admin/settings/interfaces";
 import {
   buildChatUrl,
+  buildLatestMessageChain,
   createChatSession,
   getCitedDocumentsFromMessage,
   getHumanAndAIMessageFromMessageNumber,
@@ -32,7 +33,10 @@ import {
   nameChatSession,
   personaIncludesRetrieval,
   processRawChatHistory,
+  removeMessage,
   sendMessage,
+  setMessageAsLatest,
+  updateParentChildren,
   uploadFilesForChat,
 } from "./lib";
 import { useContext, useEffect, useRef, useState } from "react";
@@ -67,6 +71,9 @@ import { InputBarPreviewImage } from "./images/InputBarPreviewImage";
 import { Folder } from "./folders/interfaces";
 
 const MAX_INPUT_HEIGHT = 200;
+const TEMP_USER_MESSAGE_ID = -1;
+const TEMP_ASSISTANT_MESSAGE_ID = -2;
+const SYSTEM_MESSAGE_ID = -3;
 
 export function ChatPage({
   user,
@@ -160,7 +167,7 @@ export function ChatPage({
         } else {
           setSelectedPersona(undefined);
         }
-        setMessageHistory([]);
+        setCompleteMessageMap(new Map());
         setChatSessionSharedStatus(ChatSessionSharedStatus.Private);
 
         // if we're supposed to submit on initial load, then do that here
@@ -186,10 +193,11 @@ export function ChatPage({
         )
       );
 
-      const newMessageHistory = processRawChatHistory(chatSession.messages);
+      const newCompleteMessageMap = processRawChatHistory(chatSession.messages);
+      const newMessageHistory = buildLatestMessageChain(newCompleteMessageMap);
       // if the last message is an error, don't overwrite it
       if (messageHistory[messageHistory.length - 1]?.type !== "error") {
-        setMessageHistory(newMessageHistory);
+        setCompleteMessageMap(newCompleteMessageMap);
 
         const latestMessageId =
           newMessageHistory[newMessageHistory.length - 1]?.messageId;
@@ -231,7 +239,77 @@ export function ChatPage({
   const [message, setMessage] = useState(
     searchParams.get(SEARCH_PARAM_NAMES.USER_MESSAGE) || ""
   );
-  const [messageHistory, setMessageHistory] = useState<Message[]>([]);
+  const [completeMessageMap, setCompleteMessageMap] = useState<
+    Map<number, Message>
+  >(new Map());
+  const upsertToCompleteMessageMap = ({
+    messages,
+    completeMessageMapOverride,
+    replacementsMap = null,
+    makeLatestChildMessage = false,
+  }: {
+    messages: Message[];
+    // if calling this function repeatedly with short delay, stay may not update in time
+    // and result in weird behavipr
+    completeMessageMapOverride?: Map<number, Message> | null;
+    replacementsMap?: Map<number, number> | null;
+    makeLatestChildMessage?: boolean;
+  }) => {
+    // deep copy
+    const frozenCompleteMessageMap =
+      completeMessageMapOverride || completeMessageMap;
+    const newCompleteMessageMap = structuredClone(frozenCompleteMessageMap);
+    if (newCompleteMessageMap.size === 0) {
+      const systemMessageId = messages[0].parentMessageId || SYSTEM_MESSAGE_ID;
+      const firstMessageId = messages[0].messageId;
+      const dummySystemMessage: Message = {
+        messageId: systemMessageId,
+        message: "",
+        type: "system",
+        files: [],
+        parentMessageId: null,
+        childrenMessageIds: [firstMessageId],
+        latestChildMessageId: firstMessageId,
+      };
+      newCompleteMessageMap.set(
+        dummySystemMessage.messageId,
+        dummySystemMessage
+      );
+      messages[0].parentMessageId = systemMessageId;
+    }
+
+    messages.forEach((message) => {
+      const idToReplace = replacementsMap?.get(message.messageId);
+      if (idToReplace) {
+        removeMessage(idToReplace, newCompleteMessageMap);
+      }
+
+      // update childrenMessageIds for the parent
+      if (
+        !newCompleteMessageMap.has(message.messageId) &&
+        message.parentMessageId !== null
+      ) {
+        updateParentChildren(message, newCompleteMessageMap, true);
+      }
+      newCompleteMessageMap.set(message.messageId, message);
+    });
+
+    // if specified, make these new message the latest of the current message chain
+    if (makeLatestChildMessage) {
+      const currentMessageChain = buildLatestMessageChain(
+        frozenCompleteMessageMap
+      );
+      const latestMessage = currentMessageChain[currentMessageChain.length - 1];
+      if (latestMessage) {
+        newCompleteMessageMap.get(
+          latestMessage.messageId
+        )!.latestChildMessageId = messages[0].messageId;
+      }
+    }
+    setCompleteMessageMap(newCompleteMessageMap);
+    return newCompleteMessageMap;
+  };
+  const messageHistory = buildLatestMessageChain(completeMessageMap);
   const [currentTool, setCurrentTool] = useState<string | null>(null);
   const [isStreaming, setIsStreaming] = useState(false);
 
@@ -415,6 +493,11 @@ export function ChatPage({
     const messageToResend = messageHistory.find(
       (message) => message.messageId === messageIdToResend
     );
+    const messageToResendParent =
+      messageToResend?.parentMessageId !== null &&
+      messageToResend?.parentMessageId !== undefined
+        ? completeMessageMap.get(messageToResend.parentMessageId)
+        : null;
     const messageToResendIndex = messageToResend
       ? messageHistory.indexOf(messageToResend)
       : null;
@@ -435,19 +518,39 @@ export function ChatPage({
       messageToResendIndex !== null
         ? messageHistory.slice(0, messageToResendIndex)
         : messageHistory;
+    let parentMessage =
+      messageToResendParent ||
+      (currMessageHistory.length > 0
+        ? currMessageHistory[currMessageHistory.length - 1]
+        : null);
     const currFiles = currentMessageFileIds.map((id) => ({
       id,
       type: "image",
     })) as FileDescriptor[];
-    setMessageHistory([
+
-      ...currMessageHistory,
+    // if we're resending, set the parent's child to null
+    // we will use tempMessages until the regenerated message is complete
+    const messageUpdates: Message[] = [
       {
-        messageId: 0,
+        messageId: TEMP_USER_MESSAGE_ID,
         message: currMessage,
         type: "user",
         files: currFiles,
+        parentMessageId: parentMessage?.messageId || null,
       },
-    ]);
+    ];
+    if (parentMessage) {
+      messageUpdates.push({
+        ...parentMessage,
+        childrenMessageIds: (parentMessage.childrenMessageIds || []).concat([
+          TEMP_USER_MESSAGE_ID,
+        ]),
+        latestChildMessageId: TEMP_USER_MESSAGE_ID,
+      });
+    }
+    const frozenCompleteMessageMap = upsertToCompleteMessageMap({
+      messages: messageUpdates,
+    });
     setMessage("");
     setCurrentMessageFileIds([]);
 
@@ -504,7 +607,7 @@ export function ChatPage({
             if (documents && documents.length > 0) {
               // point to the latest message (we don't know the messageId yet, which is why
               // we have to use -1)
-              setSelectedMessageForDocDisplay(-1);
+              setSelectedMessageForDocDisplay(TEMP_USER_MESSAGE_ID);
             }
           } else if (Object.hasOwn(packet, "file_ids")) {
             aiMessageImages = (packet as ImageGenerationDisplay).file_ids.map(
@@ -523,16 +626,35 @@ export function ChatPage({
             finalMessage = packet as BackendMessage;
           }
         }
-        setMessageHistory([
+        const updateFn = (messages: Message[]) => {
-          ...currMessageHistory,
+          const replacementsMap = finalMessage
+            ? new Map([
+                [messages[0].messageId, TEMP_USER_MESSAGE_ID],
+                [messages[1].messageId, TEMP_ASSISTANT_MESSAGE_ID],
+              ] as [number, number][])
+            : null;
+          upsertToCompleteMessageMap({
+            messages: messages,
+            replacementsMap: replacementsMap,
+            completeMessageMapOverride: frozenCompleteMessageMap,
+          });
+        };
+        const newUserMessageId =
+          finalMessage?.parent_message || TEMP_USER_MESSAGE_ID;
+        const newAssistantMessageId =
+          finalMessage?.message_id || TEMP_ASSISTANT_MESSAGE_ID;
+        updateFn([
           {
-            messageId: finalMessage?.parent_message || null,
+            messageId: newUserMessageId,
             message: currMessage,
             type: "user",
             files: currFiles,
+            parentMessageId: parentMessage?.messageId || null,
+            childrenMessageIds: [newAssistantMessageId],
+            latestChildMessageId: newAssistantMessageId,
           },
           {
-            messageId: finalMessage?.message_id || null,
+            messageId: newAssistantMessageId,
             message: error || answer,
             type: error ? "error" : "assistant",
             retrievalType,
@@ -540,6 +662,7 @@ export function ChatPage({
             documents: finalMessage?.context_docs?.top_documents || documents,
             citations: finalMessage?.citations || {},
             files: finalMessage?.files || aiMessageImages || [],
+            parentMessageId: newUserMessageId,
           },
         ]);
         if (isCancelledRef.current) {
@@ -549,21 +672,25 @@ export function ChatPage({
       }
     } catch (e: any) {
       const errorMsg = e.message;
-      setMessageHistory([
+      upsertToCompleteMessageMap({
-        ...currMessageHistory,
+        messages: [
           {
-          messageId: null,
+            messageId: TEMP_USER_MESSAGE_ID,
             message: currMessage,
             type: "user",
             files: currFiles,
+            parentMessageId: null,
           },
           {
-          messageId: null,
+            messageId: TEMP_ASSISTANT_MESSAGE_ID,
             message: errorMsg,
             type: "error",
             files: aiMessageImages || [],
+            parentMessageId: null,
           },
-      ]);
+        ],
+        completeMessageMapOverride: frozenCompleteMessageMap,
+      });
     }
     setIsStreaming(false);
     if (isNewSession) {
@@ -787,11 +914,53 @@ export function ChatPage({
                       >
                         {messageHistory.map((message, i) => {
                           if (message.type === "user") {
+                            const parentMessage = message.parentMessageId
+                              ? completeMessageMap.get(message.parentMessageId)
+                              : null;
                             return (
                               <div key={i}>
                                 <HumanMessage
                                   content={message.message}
                                   files={message.files}
+                                  messageId={message.messageId}
+                                  otherMessagesCanSwitchTo={
+                                    parentMessage?.childrenMessageIds || []
+                                  }
+                                  onEdit={(editedContent) => {
+                                    const parentMessageId =
+                                      message.parentMessageId!;
+                                    const parentMessage =
+                                      completeMessageMap.get(parentMessageId)!;
+                                    upsertToCompleteMessageMap({
+                                      messages: [
+                                        {
+                                          ...parentMessage,
+                                          latestChildMessageId: null,
+                                        },
+                                      ],
+                                    });
+                                    onSubmit({
+                                      messageIdToResend:
+                                        message.messageId || undefined,
+                                      messageOverride: editedContent,
+                                    });
+                                  }}
+                                  onMessageSelection={(messageId) => {
+                                    const newCompleteMessageMap = new Map(
+                                      completeMessageMap
+                                    );
+                                    newCompleteMessageMap.get(
+                                      message.parentMessageId!
+                                    )!.latestChildMessageId = messageId;
+                                    setCompleteMessageMap(
+                                      newCompleteMessageMap
+                                    );
+                                    setSelectedMessageForDocDisplay(messageId);
+
+                                    // set message as latest so we can edit this message
+                                    // and so it sticks around on page reload
+                                    setMessageAsLatest(messageId);
+                                  }}
                                 />
                               </div>
                             );
@@ -800,7 +969,8 @@ export function ChatPage({
                               (selectedMessageForDocDisplay !== null &&
                                 selectedMessageForDocDisplay ===
                                   message.messageId) ||
-                              (selectedMessageForDocDisplay === -1 &&
+                              (selectedMessageForDocDisplay ===
+                                TEMP_USER_MESSAGE_ID &&
                                 i === messageHistory.length - 1);
                             const previousMessage =
                               i !== 0 ? messageHistory[i - 1] : null;
@@ -921,7 +1091,7 @@ export function ChatPage({
                         })}
 
                         {isStreaming &&
-                          messageHistory.length &&
+                          messageHistory.length > 0 &&
                           messageHistory[messageHistory.length - 1].type ===
                             "user" && (
                             <div key={messageHistory.length}>

web/src/app/chat/interfaces.ts

@@ -35,14 +35,18 @@ export interface ChatSession {
 }
 
 export interface Message {
-  messageId: number | null;
+  messageId: number;
   message: string;
-  type: "user" | "assistant" | "error";
+  type: "user" | "assistant" | "system" | "error";
   retrievalType?: RetrievalType;
   query?: string | null;
   documents?: DanswerDocument[] | null;
   citations?: CitationMap;
   files: FileDescriptor[];
+  // for rebuilding the message tree
+  parentMessageId: number | null;
+  childrenMessageIds?: number[];
+  latestChildMessageId?: number | null;
 }
 
 export interface BackendChatSession {

web/src/app/chat/lib.tsx

@@ -154,6 +154,19 @@ export async function nameChatSession(chatSessionId: number, message: string) {
   return response;
 }
 
+export async function setMessageAsLatest(messageId: number) {
+  const response = await fetch("/api/chat/set-message-as-latest", {
+    method: "PUT",
+    headers: {
+      "Content-Type": "application/json",
+    },
+    body: JSON.stringify({
+      message_id: messageId,
+    }),
+  });
+  return response;
+}
+
 export async function handleChatFeedback(
   messageId: number,
   feedback: FeedbackType,
@@ -332,32 +345,13 @@ export function getLastSuccessfulMessageId(messageHistory: Message[]) {
   return lastSuccessfulMessage ? lastSuccessfulMessage?.messageId : null;
 }
 
-export function processRawChatHistory(rawMessages: BackendMessage[]) {
+export function processRawChatHistory(
-  const messageMap: Map<number, BackendMessage> = new Map(
+  rawMessages: BackendMessage[]
-    rawMessages.map((message) => [message.message_id, message])
+): Map<number, Message> {
-  );
+  const messages: Map<number, Message> = new Map();
+  const parentMessageChildrenMap: Map<number, number[]> = new Map();
 
-  const rootMessage = rawMessages.find(
+  rawMessages.forEach((messageInfo) => {
-    (message) => message.parent_message === null
-  );
-
-  const finalMessageList: BackendMessage[] = [];
-  if (rootMessage) {
-    let currMessage: BackendMessage | null = rootMessage;
-    while (currMessage) {
-      finalMessageList.push(currMessage);
-      const childMessageNumber = currMessage.latest_child_message;
-      if (childMessageNumber && messageMap.has(childMessageNumber)) {
-        currMessage = messageMap.get(childMessageNumber) as BackendMessage;
-      } else {
-        currMessage = null;
-      }
-    }
-  }
-
-  const messages: Message[] = finalMessageList
-    .filter((messageInfo) => messageInfo.message_type !== "system")
-    .map((messageInfo) => {
     const hasContextDocs =
       (messageInfo?.context_docs?.top_documents || []).length > 0;
     let retrievalType;
@@ -371,7 +365,7 @@ export function processRawChatHistory(rawMessages: BackendMessage[]) {
       retrievalType = RetrievalType.None;
     }
 
-      return {
+    const message: Message = {
       messageId: messageInfo.message_id,
       message: messageInfo.message,
       type: messageInfo.message_type as "user" | "assistant",
@@ -386,12 +380,110 @@ export function processRawChatHistory(rawMessages: BackendMessage[]) {
             citations: messageInfo?.citations || {},
           }
         : {}),
+      parentMessageId: messageInfo.parent_message,
+      childrenMessageIds: [],
+      latestChildMessageId: messageInfo.latest_child_message,
     };
+
+    messages.set(messageInfo.message_id, message);
+
+    if (messageInfo.parent_message !== null) {
+      if (!parentMessageChildrenMap.has(messageInfo.parent_message)) {
+        parentMessageChildrenMap.set(messageInfo.parent_message, []);
+      }
+      parentMessageChildrenMap
+        .get(messageInfo.parent_message)!
+        .push(messageInfo.message_id);
+    }
+  });
+
+  // Populate childrenMessageIds for each message
+  parentMessageChildrenMap.forEach((childrenIds, parentId) => {
+    childrenIds.sort((a, b) => a - b);
+    const parentMesage = messages.get(parentId);
+    if (parentMesage) {
+      parentMesage.childrenMessageIds = childrenIds;
+    }
   });
 
   return messages;
 }
 
+export function buildLatestMessageChain(
+  messageMap: Map<number, Message>,
+  additionalMessagesOnMainline: Message[] = []
+): Message[] {
+  const rootMessage = Array.from(messageMap.values()).find(
+    (message) => message.parentMessageId === null
+  );
+
+  let finalMessageList: Message[] = [];
+  if (rootMessage) {
+    let currMessage: Message | null = rootMessage;
+    while (currMessage) {
+      finalMessageList.push(currMessage);
+      const childMessageNumber = currMessage.latestChildMessageId;
+      if (childMessageNumber && messageMap.has(childMessageNumber)) {
+        currMessage = messageMap.get(childMessageNumber) as Message;
+      } else {
+        currMessage = null;
+      }
+    }
+  }
+
+  // remove system message
+  if (finalMessageList.length > 0 && finalMessageList[0].type === "system") {
+    finalMessageList = finalMessageList.slice(1);
+  }
+  return finalMessageList.concat(additionalMessagesOnMainline);
+}
+
+export function updateParentChildren(
+  message: Message,
+  completeMessageMap: Map<number, Message>,
+  setAsLatestChild: boolean = false
+) {
+  // NOTE: updates the `completeMessageMap` in place
+  const parentMessage = message.parentMessageId
+    ? completeMessageMap.get(message.parentMessageId)
+    : null;
+  if (parentMessage) {
+    if (setAsLatestChild) {
+      parentMessage.latestChildMessageId = message.messageId;
+    }
+
+    const parentChildMessages = parentMessage.childrenMessageIds || [];
+    if (!parentChildMessages.includes(message.messageId)) {
+      parentChildMessages.push(message.messageId);
+    }
+    parentMessage.childrenMessageIds = parentChildMessages;
+  }
+}
+
+export function removeMessage(
+  messageId: number,
+  completeMessageMap: Map<number, Message>
+) {
+  const messageToRemove = completeMessageMap.get(messageId);
+  if (!messageToRemove) {
+    return;
+  }
+
+  const parentMessage = messageToRemove.parentMessageId
+    ? completeMessageMap.get(messageToRemove.parentMessageId)
+    : null;
+  if (parentMessage) {
+    if (parentMessage.latestChildMessageId === messageId) {
+      parentMessage.latestChildMessageId = null;
+    }
+    const currChildMessage = parentMessage.childrenMessageIds || [];
+    const newChildMessage = currChildMessage.filter((id) => id !== messageId);
+    parentMessage.childrenMessageIds = newChildMessage;
+  }
+
+  completeMessageMap.delete(messageId);
+}
+
 export function personaIncludesRetrieval(selectedPersona: Persona) {
   return selectedPersona.num_chunks !== 0;
 }

web/src/app/chat/message/Messages.tsx

@@ -1,15 +1,15 @@
 import {
-  FiCheck,
-  FiCopy,
   FiCpu,
   FiImage,
   FiThumbsDown,
   FiThumbsUp,
-  FiTool,
   FiUser,
+  FiEdit2,
+  FiChevronRight,
+  FiChevronLeft,
 } from "react-icons/fi";
 import { FeedbackType } from "../types";
-import { useState } from "react";
+import { useEffect, useRef, useState } from "react";
 import ReactMarkdown from "react-markdown";
 import { DanswerDocument } from "@/lib/search/interfaces";
 import { SearchSummary, ShowHideDocsButton } from "./SearchSummary";
@@ -20,25 +20,11 @@ import remarkGfm from "remark-gfm";
 import { CopyButton } from "@/components/CopyButton";
 import { FileDescriptor } from "../interfaces";
 import { InMessageImage } from "../images/InMessageImage";
-import {
+import { IMAGE_GENERATION_TOOL_NAME } from "../tools/constants";
-  IMAGE_GENERATION_TOOL_NAME,
-  SEARCH_TOOL_NAME,
-} from "../tools/constants";
 import { ToolRunningAnimation } from "../tools/ToolRunningAnimation";
+import { Hoverable } from "@/components/Hoverable";
 
-export const Hoverable: React.FC<{
+const ICON_SIZE = 15;
-  children: JSX.Element;
-  onClick?: () => void;
-}> = ({ children, onClick }) => {
-  return (
-    <div
-      className="hover:bg-hover p-2 rounded h-fit cursor-pointer"
-      onClick={onClick}
-    >
-      {children}
-    </div>
-  );
-};
 
 export const AIMessage = ({
   messageId,
@@ -236,14 +222,16 @@ export const AIMessage = ({
             )}
           </div>
           {handleFeedback && (
-            <div className="flex flex-col md:flex-row gap-x-0.5 ml-8 mt-1">
+            <div className="flex flex-col md:flex-row gap-x-0.5 ml-8 mt-1.5">
               <CopyButton content={content.toString()} />
-              <Hoverable onClick={() => handleFeedback("like")}>
+              <Hoverable
-                <FiThumbsUp />
+                icon={FiThumbsUp}
-              </Hoverable>
+                onClick={() => handleFeedback("like")}
-              <Hoverable>
+              />
-                <FiThumbsDown onClick={() => handleFeedback("dislike")} />
+              <Hoverable
-              </Hoverable>
+                icon={FiThumbsDown}
+                onClick={() => handleFeedback("dislike")}
+              />
             </div>
           )}
         </div>
@@ -252,15 +240,88 @@ export const AIMessage = ({
   );
 };
 
+function MessageSwitcher({
+  currentPage,
+  totalPages,
+  handlePrevious,
+  handleNext,
+}: {
+  currentPage: number;
+  totalPages: number;
+  handlePrevious: () => void;
+  handleNext: () => void;
+}) {
+  return (
+    <div className="flex items-center text-sm space-x-0.5">
+      <Hoverable
+        icon={FiChevronLeft}
+        onClick={currentPage === 1 ? undefined : handlePrevious}
+      />
+      <span className="text-emphasis text-medium select-none">
+        {currentPage} / {totalPages}
+      </span>
+      <Hoverable
+        icon={FiChevronRight}
+        onClick={currentPage === totalPages ? undefined : handleNext}
+      />
+    </div>
+  );
+}
+
 export const HumanMessage = ({
   content,
   files,
+  messageId,
+  otherMessagesCanSwitchTo,
+  onEdit,
+  onMessageSelection,
 }: {
-  content: string | JSX.Element;
+  content: string;
   files?: FileDescriptor[];
+  messageId?: number | null;
+  otherMessagesCanSwitchTo?: number[];
+  onEdit?: (editedContent: string) => void;
+  onMessageSelection?: (messageId: number) => void;
 }) => {
+  const textareaRef = useRef<HTMLTextAreaElement>(null);
+
+  const [isHovered, setIsHovered] = useState(false);
+  const [isEditing, setIsEditing] = useState(false);
+  const [editedContent, setEditedContent] = useState(content);
+
+  useEffect(() => {
+    if (!isEditing) {
+      setEditedContent(content);
+    }
+  }, [content]);
+
+  useEffect(() => {
+    if (textareaRef.current) {
+      // Focus the textarea
+      textareaRef.current.focus();
+      // Move the cursor to the end of the text
+      textareaRef.current.selectionStart = textareaRef.current.value.length;
+      textareaRef.current.selectionEnd = textareaRef.current.value.length;
+    }
+  }, [isEditing]);
+
+  const handleEditSubmit = () => {
+    if (editedContent.trim() !== content.trim()) {
+      onEdit?.(editedContent);
+    }
+    setIsEditing(false);
+  };
+
+  const currentMessageInd = messageId
+    ? otherMessagesCanSwitchTo?.indexOf(messageId)
+    : undefined;
+
   return (
-    <div className="py-5 px-5 flex -mr-6 w-full">
+    <div
+      className="pt-5 pb-1 px-5 flex -mr-6 w-full relative"
+      onMouseEnter={() => setIsHovered(true)}
+      onMouseLeave={() => setIsHovered(false)}
+    >
       <div className="mx-auto w-searchbar-xs 2xl:w-searchbar-sm 3xl:w-searchbar">
         <div className="ml-8">
           <div className="flex">
@@ -284,7 +345,102 @@ export const HumanMessage = ({
                 </div>
               )}
 
-              {typeof content === "string" ? (
+              {isEditing ? (
+                <div>
+                  <div
+                    className={`
+                      opacity-100
+                      w-full
+                      flex
+                      flex-col
+                      border 
+                      border-border 
+                      rounded-lg 
+                      bg-background-emphasis 
+                      pb-2
+                      [&:has(textarea:focus)]::ring-1
+                      [&:has(textarea:focus)]::ring-black
+                    `}
+                  >
+                    <textarea
+                      ref={textareaRef}
+                      className={`
+                      m-0 
+                      w-full 
+                      h-auto
+                      shrink
+                      border-0
+                      rounded-lg 
+                      overflow-y-hidden
+                      bg-background-emphasis 
+                      whitespace-normal 
+                      break-word
+                      overscroll-contain
+                      outline-none 
+                      placeholder-gray-400 
+                      resize-none
+                      pl-4
+                      pr-12 
+                      py-4`}
+                      aria-multiline
+                      role="textarea"
+                      value={editedContent}
+                      style={{ scrollbarWidth: "thin" }}
+                      onChange={(e) => {
+                        setEditedContent(e.target.value);
+                        e.target.style.height = "auto";
+                        e.target.style.height = `${e.target.scrollHeight}px`;
+                      }}
+                      onKeyDown={(e) => {
+                        if (e.key === "Escape") {
+                          e.preventDefault();
+                          setEditedContent(content);
+                          setIsEditing(false);
+                        }
+                      }}
+                      // ref={(textarea) => {
+                      //   if (textarea) {
+                      //     textarea.selectionStart = textarea.selectionEnd =
+                      //       textarea.value.length;
+                      //   }
+                      // }}
+                    />
+                    <div className="flex justify-end mt-2 gap-2 pr-4">
+                      <button
+                        className={`
+                          w-fit 
+                          p-1 
+                          bg-accent 
+                          text-inverted 
+                          text-sm
+                          rounded-lg 
+                          hover:bg-accent-hover
+                        `}
+                        onClick={handleEditSubmit}
+                      >
+                        Submit
+                      </button>
+                      <button
+                        className={`
+                          w-fit 
+                          p-1 
+                          bg-hover
+                          bg-background-strong 
+                          text-sm
+                          rounded-lg
+                          hover:bg-hover-emphasis
+                        `}
+                        onClick={() => {
+                          setEditedContent(content);
+                          setIsEditing(false);
+                        }}
+                      >
+                        Cancel
+                      </button>
+                    </div>
+                  </div>
+                </div>
+              ) : typeof content === "string" ? (
                 <ReactMarkdown
                   className="prose max-w-full"
                   components={{
@@ -306,6 +462,43 @@ export const HumanMessage = ({
               )}
             </div>
           </div>
+          <div className="flex flex-col md:flex-row gap-x-0.5 ml-8 mt-1">
+            {currentMessageInd !== undefined &&
+              onMessageSelection &&
+              otherMessagesCanSwitchTo &&
+              otherMessagesCanSwitchTo.length > 1 && (
+                <div className="mr-2">
+                  <MessageSwitcher
+                    currentPage={currentMessageInd + 1}
+                    totalPages={otherMessagesCanSwitchTo.length}
+                    handlePrevious={() =>
+                      onMessageSelection(
+                        otherMessagesCanSwitchTo[currentMessageInd - 1]
+                      )
+                    }
+                    handleNext={() =>
+                      onMessageSelection(
+                        otherMessagesCanSwitchTo[currentMessageInd + 1]
+                      )
+                    }
+                  />
+                </div>
+              )}
+            {onEdit &&
+            isHovered &&
+            !isEditing &&
+            (!files || files.length === 0) ? (
+              <Hoverable
+                icon={FiEdit2}
+                onClick={() => {
+                  setIsEditing(true);
+                  setIsHovered(false);
+                }}
+              />
+            ) : (
+              <div className="h-[27px]" />
+            )}
+          </div>
         </div>
       </div>
     </div>

web/src/app/chat/message/SearchSummary.tsx

@@ -3,9 +3,9 @@ import {
   EmphasizedClickable,
 } from "@/components/BasicClickable";
 import { HoverPopup } from "@/components/HoverPopup";
+import { Hoverable } from "@/components/Hoverable";
 import { useEffect, useRef, useState } from "react";
 import { FiCheck, FiEdit2, FiSearch, FiX } from "react-icons/fi";
-import { Hoverable } from "./Messages";
 
 export function ShowHideDocsButton({
   messageId,
@@ -78,6 +78,12 @@ export function SearchSummary({
     }
   }, [isEditing]);
 
+  useEffect(() => {
+    if (!isEditing) {
+      setFinalQuery(query);
+    }
+  }, [query]);
+
   const searchingForDisplay = (
     <div className={`flex p-1 rounded ${isOverflowed && "cursor-default"}`}>
       <FiSearch className="mr-2 my-auto" size={14} />
@@ -113,7 +119,8 @@ export function SearchSummary({
         />
       </div>
       <div className="ml-2 my-auto flex">
-        <div
+        <Hoverable
+          icon={FiCheck}
           onClick={() => {
             if (!finalQuery) {
               setFinalQuery(query);
@@ -122,19 +129,14 @@ export function SearchSummary({
             }
             setIsEditing(false);
           }}
-          className={`hover:bg-black/10 p-1 -m-1 rounded`}
+        />
-        >
+        <Hoverable
-          <FiCheck size={14} />
+          icon={FiX}
-        </div>
-        <div
           onClick={() => {
             setFinalQuery(query);
             setIsEditing(false);
           }}
-          className={`hover:bg-black/10 p-1 -m-1 rounded ml-2`}
+        />
-        >
-          <FiX size={14} />
-        </div>
       </div>
     </div>
   ) : null;
@@ -163,9 +165,7 @@ export function SearchSummary({
           </div>
           {handleSearchQueryEdit && (
             <div className="my-auto">
-              <Hoverable onClick={() => setIsEditing(true)}>
+              <Hoverable icon={FiEdit2} onClick={() => setIsEditing(true)} />
-                <FiEdit2 className="my-auto" size="14" />
-              </Hoverable>
             </div>
           )}
         </>

web/src/app/chat/modal/ShareChatSessionModal.tsx

@@ -4,7 +4,6 @@ import { Button, Callout, Divider, Text } from "@tremor/react";
 import { Spinner } from "@/components/Spinner";
 import { ChatSessionSharedStatus } from "../interfaces";
 import { FiCopy, FiX } from "react-icons/fi";
-import { Hoverable } from "../message/Messages";
 import { CopyButton } from "@/components/CopyButton";
 
 function buildShareLink(chatSessionId: number) {

web/src/app/chat/shared/[chatId]/SharedChatDisplay.tsx

@@ -2,7 +2,11 @@
 
 import { humanReadableFormat } from "@/lib/time";
 import { BackendChatSession } from "../../interfaces";
-import { getCitedDocumentsFromMessage, processRawChatHistory } from "../../lib";
+import {
+  buildLatestMessageChain,
+  getCitedDocumentsFromMessage,
+  processRawChatHistory,
+} from "../../lib";
 import { AIMessage, HumanMessage } from "../../message/Messages";
 import { Button, Callout, Divider } from "@tremor/react";
 import { useRouter } from "next/navigation";
@@ -40,7 +44,9 @@ export function SharedChatDisplay({
     );
   }
 
-  const messages = processRawChatHistory(chatSession.messages);
+  const messages = buildLatestMessageChain(
+    processRawChatHistory(chatSession.messages)
+  );
 
   return (
     <div className="w-full overflow-hidden">

web/src/components/CopyButton.tsx

@@ -1,6 +1,6 @@
-import { Hoverable } from "@/app/chat/message/Messages";
 import { useState } from "react";
 import { FiCheck, FiCopy } from "react-icons/fi";
+import { Hoverable } from "./Hoverable";
 
 export function CopyButton({
   content,
@@ -13,6 +13,7 @@ export function CopyButton({
 
   return (
     <Hoverable
+      icon={copyClicked ? FiCheck : FiCopy}
       onClick={() => {
         if (content) {
           navigator.clipboard.writeText(content.toString());
@@ -22,8 +23,6 @@ export function CopyButton({
         setCopyClicked(true);
         setTimeout(() => setCopyClicked(false), 3000);
       }}
-    >
+    />
-      {copyClicked ? <FiCheck /> : <FiCopy />}
-    </Hoverable>
   );
 }

web/src/components/Hoverable.tsx

@@ -0,0 +1,18 @@
+import { IconType } from "react-icons";
+
+const ICON_SIZE = 15;
+
+export const Hoverable: React.FC<{
+  icon: IconType;
+  onClick?: () => void;
+  size?: number;
+}> = ({ icon, onClick, size = ICON_SIZE }) => {
+  return (
+    <div
+      className="hover:bg-hover p-1.5 rounded h-fit cursor-pointer"
+      onClick={onClick}
+    >
+      {icon({ size: size, className: "my-auto" })}
+    </div>
+  );
+};

web/tailwind.config.js

@@ -51,9 +51,10 @@ module.exports = {
         "border-strong": "#9ca3af", // gray-400
         "hover-light": "#f3f4f6", // gray-100
         hover: "#e5e7eb", // gray-200
+        "hover-emphasis": "#d1d5db", // gray-300
         popup: "#ffffff", // white
         accent: "#6671d0",
-        "accent-hover": "#6671d0",
+        "accent-hover": "#5964c2",
         highlight: {
           text: "#fef9c3", // yellow-100
         },