Add query editing in Chat UI (#899)

backend/danswer/chat/process_message.py

@@ -164,6 +164,7 @@ def stream_chat_message(
         reference_doc_ids = new_msg_req.search_doc_ids
         retrieval_options = new_msg_req.retrieval_options
         persona = chat_session.persona
+        query_override = new_msg_req.query_override
 
         if reference_doc_ids is None and retrieval_options is None:
             raise RuntimeError(
@@ -259,8 +260,12 @@ def stream_chat_message(
             ]
 
         elif run_search:
-            rephrased_query = history_based_query_rephrase(
-                query_message=final_msg, history=history_msgs, llm=llm
+            rephrased_query = (
+                history_based_query_rephrase(
+                    query_message=final_msg, history=history_msgs, llm=llm
+                )
+                if query_override is None
+                else query_override
             )
 
             (

backend/danswer/server/query_and_chat/models.py

@@ -65,6 +65,9 @@ class CreateChatMessageRequest(BaseModel):
     # If search_doc_ids provided, then retrieval options are unused
     search_doc_ids: list[int] | None
     retrieval_options: RetrievalDetails | None
+    # allows the caller to specify the exact search query they want to use
+    # will disable query re-wording if specified
+    query_override: str | None = None
 
     @root_validator
     def check_search_doc_ids_or_retrieval_options(cls: BaseModel, values: dict) -> dict:

web/src/app/chat/Chat.tsx

@@ -205,7 +205,10 @@ export const Chat = ({
     documentSidebarInitialWidth = Math.min(700, maxDocumentSidebarWidth);
   }
 
-  const onSubmit = async (messageOverride?: string) => {
+  const onSubmit = async ({
+    messageIdToResend,
+    queryOverride,
+  }: { messageIdToResend?: number; queryOverride?: string } = {}) => {
     let currChatSessionId: number;
     let isNewSession = chatSessionId === null;
     if (isNewSession) {
@@ -215,8 +218,26 @@ export const Chat = ({
     }
     setChatSessionId(currChatSessionId);
 
-    const currMessage = messageOverride || message;
-    const currMessageHistory = messageHistory;
+    const messageToResend = messageHistory.find(
+      (message) => message.messageId === messageIdToResend
+    );
+    const messageToResendIndex = messageToResend
+      ? messageHistory.indexOf(messageToResend)
+      : null;
+    if (!messageToResend && messageIdToResend !== undefined) {
+      setPopup({
+        message:
+          "Failed to re-send message - please refresh the page and try again.",
+        type: "error",
+      });
+      return;
+    }
+
+    const currMessage = messageToResend ? messageToResend.message : message;
+    const currMessageHistory =
+      messageToResendIndex !== null
+        ? messageHistory.slice(0, messageToResendIndex)
+        : messageHistory;
     setMessageHistory([
       ...currMessageHistory,
       {
@@ -256,6 +277,7 @@ export const Chat = ({
               document.db_doc_id !== undefined && document.db_doc_id !== null
           )
           .map((document) => document.db_doc_id as number),
+        queryOverride,
       })) {
         for (const packet of packetBunch) {
           if (Object.hasOwn(packet, "answer_piece")) {
@@ -440,6 +462,8 @@ export const Chat = ({
                         selectedMessageForDocDisplay === message.messageId) ||
                       (selectedMessageForDocDisplay === -1 &&
                         i === messageHistory.length - 1);
+                    const previousMessage =
+                      i !== 0 ? messageHistory[i - 1] : null;
                     return (
                       <div key={i}>
                         <AIMessage
@@ -463,6 +487,34 @@ export const Chat = ({
                                     message.messageId as number,
                                   ])
                           }
+                          handleSearchQueryEdit={
+                            i === messageHistory.length - 1 && !isStreaming
+                              ? (newQuery) => {
+                                  if (!previousMessage) {
+                                    setPopup({
+                                      type: "error",
+                                      message:
+                                        "Cannot edit query of first message - please refresh the page and try again.",
+                                    });
+                                    return;
+                                  }
+
+                                  if (previousMessage.messageId === null) {
+                                    setPopup({
+                                      type: "error",
+                                      message:
+                                        "Cannot edit query of a pending message - please wait a few seconds and try again.",
+                                    });
+                                    return;
+                                  }
+                                  onSubmit({
+                                    messageIdToResend:
+                                      previousMessage.messageId,
+                                    queryOverride: newQuery,
+                                  });
+                                }
+                              : undefined
+                          }
                           isCurrentlyShowingRetrieved={isShowingRetrieved}
                           handleShowRetrieved={(messageNumber) => {
                             if (isShowingRetrieved) {
@@ -527,36 +579,6 @@ export const Chat = ({
 
             <div className="absolute bottom-0 z-10 w-full bg-background border-t border-border">
               <div className="w-full pb-4 pt-2">
-                {/* {(isStreaming || messageHistory.length > 0) && (
-              <div className="flex justify-center w-full">
-                <div className="w-[800px] flex">
-                  <div className="cursor-pointer flex w-fit p-2 rounded border border-neutral-400 text-sm hover:bg-neutral-200 ml-auto mr-4">
-                    {isStreaming ? (
-                      <div
-                        onClick={() => setIsCancelled(true)}
-                        className="flex"
-                      >
-                        <FiStopCircle className="my-auto mr-1" />
-                        <div>Stop Generating</div>
-                      </div>
-                    ) : (
-                      <div
-                        className="flex"
-                        onClick={() => {
-                          if (chatSessionId) {
-                            handleRegenerate(chatSessionId);
-                          }
-                        }}
-                      >
-                        <FiRefreshCcw className="my-auto mr-1" />
-                        <div>Regenerate</div>
-                      </div>
-                    )}
-                  </div>
-                </div>
-              </div>
-            )} */}
-
                 <div className="flex">
                   <div className="w-searchbar-xs 2xl:w-searchbar-sm 3xl:w-searchbar mx-auto px-4 pt-1 flex">
                     {selectedDocuments.length > 0 ? (

web/src/app/chat/lib.tsx

@@ -45,6 +45,7 @@ export interface SendMessageRequest {
   promptId: number | null | undefined;
   filters: Filters | null;
   selectedDocumentIds: number[] | null;
+  queryOverride?: string;
 }
 
 export async function* sendMessage({
@@ -54,6 +55,7 @@ export async function* sendMessage({
   promptId,
   filters,
   selectedDocumentIds,
+  queryOverride,
 }: SendMessageRequest) {
   const documentsAreSelected =
     selectedDocumentIds && selectedDocumentIds.length > 0;
@@ -71,11 +73,14 @@ export async function* sendMessage({
       retrieval_options: !documentsAreSelected
         ? {
             run_search:
-              promptId === null || promptId === undefined ? "always" : "auto",
+              promptId === null || promptId === undefined || queryOverride
+                ? "always"
+                : "auto",
             real_time: true,
             filters: filters,
           }
         : null,
+      query_override: queryOverride,
     }),
   });
   if (!sendMessageResponse.ok) {

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

@@ -14,10 +14,10 @@ import { SearchSummary, ShowHideDocsButton } from "./SearchSummary";
 import { SourceIcon } from "@/components/SourceIcon";
 import { ThreeDots } from "react-loader-spinner";
 
-const Hoverable: React.FC<{ children: JSX.Element; onClick?: () => void }> = ({
-  children,
-  onClick,
-}) => {
+export const Hoverable: React.FC<{
+  children: JSX.Element;
+  onClick?: () => void;
+}> = ({ children, onClick }) => {
   return (
     <div
       className="hover:bg-neutral-300 p-2 rounded h-fit cursor-pointer"
@@ -38,6 +38,7 @@ export const AIMessage = ({
   handleFeedback,
   isCurrentlyShowingRetrieved,
   handleShowRetrieved,
+  handleSearchQueryEdit,
 }: {
   messageId: number | null;
   content: string | JSX.Element;
@@ -48,6 +49,7 @@ export const AIMessage = ({
   handleFeedback?: (feedbackType: FeedbackType) => void;
   isCurrentlyShowingRetrieved?: boolean;
   handleShowRetrieved?: (messageNumber: number | null) => void;
+  handleSearchQueryEdit?: (query: string) => void;
 }) => {
   const [copyClicked, setCopyClicked] = useState(false);
   return (
@@ -90,6 +92,7 @@ export const AIMessage = ({
                     messageId={messageId}
                     isCurrentlyShowingRetrieved={isCurrentlyShowingRetrieved}
                     handleShowRetrieved={handleShowRetrieved}
+                    handleSearchQueryEdit={handleSearchQueryEdit}
                   />
                 </div>
               )}

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

@@ -3,8 +3,9 @@ import {
   EmphasizedClickable,
 } from "@/components/BasicClickable";
 import { HoverPopup } from "@/components/HoverPopup";
-import { DanswerDocument } from "@/lib/search/interfaces";
-import { FiBookOpen, FiSearch } from "react-icons/fi";
+import { useEffect, useRef, useState } from "react";
+import { FiCheck, FiEdit2, FiSearch, FiX } from "react-icons/fi";
+import { Hoverable } from "./Messages";
 
 export function ShowHideDocsButton({
   messageId,
@@ -33,53 +34,142 @@ export function ShowHideDocsButton({
   );
 }
 
-function SearchingForDisplay({
-  query,
-  isHoverable,
-}: {
-  query: string;
-  isHoverable?: boolean;
-}) {
-  return (
-    <div className={`flex p-1 rounded ${isHoverable && "cursor-default"}`}>
-      <FiSearch className="mr-2 my-auto" size={14} />
-      <p className="line-clamp-1 break-all xl:w-64 2xl:w-80 3xl:w-96">
-        Searching for: <i>{query}</i>
-      </p>
-    </div>
-  );
-}
-
 export function SearchSummary({
   query,
   hasDocs,
   messageId,
   isCurrentlyShowingRetrieved,
   handleShowRetrieved,
+  handleSearchQueryEdit,
 }: {
   query: string;
   hasDocs: boolean;
   messageId: number | null;
   isCurrentlyShowingRetrieved: boolean;
   handleShowRetrieved: (messageId: number | null) => void;
+  handleSearchQueryEdit?: (query: string) => void;
 }) {
+  const [isEditing, setIsEditing] = useState(false);
+  const [finalQuery, setFinalQuery] = useState(query);
+  const [isOverflowed, setIsOverflowed] = useState(false);
+  const searchingForRef = useRef<HTMLDivElement>(null);
+  const editQueryRef = useRef<HTMLInputElement>(null);
+
+  useEffect(() => {
+    const checkOverflow = () => {
+      const current = searchingForRef.current;
+      if (current) {
+        setIsOverflowed(
+          current.scrollWidth > current.clientWidth ||
+            current.scrollHeight > current.clientHeight
+        );
+      }
+    };
+
+    checkOverflow();
+    window.addEventListener("resize", checkOverflow); // Recheck on window resize
+
+    return () => window.removeEventListener("resize", checkOverflow);
+  }, []);
+
+  useEffect(() => {
+    if (isEditing && editQueryRef.current) {
+      editQueryRef.current.focus();
+    }
+  }, [isEditing]);
+
+  const searchingForDisplay = (
+    <div className={`flex p-1 rounded ${isOverflowed && "cursor-default"}`}>
+      <FiSearch className="mr-2 my-auto" size={14} />
+      <div className="line-clamp-1 break-all px-0.5" ref={searchingForRef}>
+        Searching for: <i>{finalQuery}</i>
+      </div>
+    </div>
+  );
+
+  const editInput = handleSearchQueryEdit ? (
+    <div className="flex w-full mr-3">
+      <div className="my-2 w-full">
+        <input
+          ref={editQueryRef}
+          value={finalQuery}
+          onChange={(e) => setFinalQuery(e.target.value)}
+          onKeyDown={(event) => {
+            if (event.key === "Enter") {
+              setIsEditing(false);
+              if (!finalQuery) {
+                setFinalQuery(query);
+              } else if (finalQuery !== query) {
+                handleSearchQueryEdit(finalQuery);
+              }
+              event.preventDefault();
+            } else if (event.key === "Escape") {
+              setFinalQuery(query);
+              setIsEditing(false);
+              event.preventDefault();
+            }
+          }}
+          className="px-1 py-0.5 h-[28px] text-sm mr-2 w-full rounded-sm border border-border-strong"
+        />
+      </div>
+      <div className="ml-2 my-auto flex">
+        <div
+          onClick={() => {
+            if (!finalQuery) {
+              setFinalQuery(query);
+            } else if (finalQuery !== query) {
+              handleSearchQueryEdit(finalQuery);
+            }
+            setIsEditing(false);
+          }}
+          className={`hover:bg-black/10 p-1 -m-1 rounded`}
+        >
+          <FiCheck size={14} />
+        </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;
+
   return (
     <div className="flex">
-      <div className="text-sm my-2">
-        {query.length >= 40 ? (
-          <HoverPopup
-            mainContent={<SearchingForDisplay query={query} isHoverable />}
-            popupContent={
-              <div className="xl:w-64 2xl:w-80 3xl:w-96">
-                <b>Full query:</b> <div className="mt-1 italic">{query}</div>
-              </div>
-            }
-            direction="top"
-          />
-        ) : (
-          <SearchingForDisplay query={query} />
-        )}
-      </div>
+      {isEditing ? (
+        editInput
+      ) : (
+        <>
+          <div className="text-sm my-2">
+            {isOverflowed ? (
+              <HoverPopup
+                mainContent={searchingForDisplay}
+                popupContent={
+                  <div>
+                    <b>Full query:</b>{" "}
+                    <div className="mt-1 italic">{query}</div>
+                  </div>
+                }
+                direction="top"
+              />
+            ) : (
+              searchingForDisplay
+            )}
+          </div>
+          {handleSearchQueryEdit && (
+            <div className="my-auto">
+              <Hoverable onClick={() => setIsEditing(true)}>
+                <FiEdit2 className="my-auto" size="14" />
+              </Hoverable>
+            </div>
+          )}
+        </>
+      )}
       {hasDocs && (
         <ShowHideDocsButton
           messageId={messageId}