Feature/scroll (#1694)

---------

Co-authored-by: “Pablo <“pablo@danswer.ai”>

web/src/app/chat/ChatPage.tsx

@@ -26,9 +26,9 @@ import {
   getCitedDocumentsFromMessage,
   getHumanAndAIMessageFromMessageNumber,
   getLastSuccessfulMessageId,
-  handleAutoScroll,
   handleChatFeedback,
   nameChatSession,
+  PacketType,
   personaIncludesRetrieval,
   processRawChatHistory,
   removeMessage,
@@ -37,6 +37,7 @@ import {
   updateModelOverrideForChatSession,
   updateParentChildren,
   uploadFilesForChat,
+  useScrollonStream,
 } from "./lib";
 import { useContext, useEffect, useRef, useState } from "react";
 import { usePopup } from "@/components/admin/connectors/Popup";
@@ -50,7 +51,7 @@ import { DanswerInitializingLoader } from "@/components/DanswerInitializingLoade
 import { FeedbackModal } from "./modal/FeedbackModal";
 import { ShareChatSessionModal } from "./modal/ShareChatSessionModal";
 import { ChatPersonaSelector } from "./ChatPersonaSelector";
-import { FiShare2 } from "react-icons/fi";
+import { FiArrowDown, FiShare2 } from "react-icons/fi";
 import { ChatIntro } from "./ChatIntro";
 import { AIMessage, HumanMessage } from "./message/Messages";
 import { ThreeDots } from "react-loader-spinner";
@@ -77,6 +78,7 @@ import { TbLayoutSidebarRightExpand } from "react-icons/tb";
 import { SIDEBAR_WIDTH_CONST } from "@/lib/constants";
 
 import ResizableSection from "@/components/resizable/ResizableSection";
+import { Button } from "@tremor/react";
 
 const MAX_INPUT_HEIGHT = 200;
 const TEMP_USER_MESSAGE_ID = -1;
@@ -426,49 +428,111 @@ export function ChatPage({
       availableDocumentSets,
     });
 
+  const [currentFeedback, setCurrentFeedback] = useState<
+    [FeedbackType, number] | null
+  >(null);
+
+  const [sharingModalVisible, setSharingModalVisible] =
+    useState<boolean>(false);
+
   // state for cancelling streaming
   const [isCancelled, setIsCancelled] = useState(false);
-  const isCancelledRef = useRef(isCancelled);
+  const [aboveHorizon, setAboveHorizon] = useState(false);
+
+  const scrollableDivRef = useRef<HTMLDivElement>(null);
+  const lastMessageRef = useRef<HTMLDivElement>(null);
+  const inputRef = useRef<HTMLDivElement>(null);
+  const endDivRef = useRef<HTMLDivElement>(null);
+  const endPaddingRef = useRef<HTMLDivElement>(null);
+
+  const previousHeight = useRef<number>(
+    inputRef.current?.getBoundingClientRect().height!
+  );
+  const scrollDist = useRef<number>(0);
+
+  const updateScrollTracking = () => {
+    const scrollDistance =
+      endDivRef?.current?.getBoundingClientRect()?.top! -
+      inputRef?.current?.getBoundingClientRect()?.top!;
+    scrollDist.current = scrollDistance;
+    setAboveHorizon(scrollDist.current > 500);
+  };
+
+  scrollableDivRef?.current?.addEventListener("scroll", updateScrollTracking);
+
+  const handleInputResize = () => {
+    setTimeout(() => {
+      if (inputRef.current && lastMessageRef.current) {
+        let newHeight: number =
+          inputRef.current?.getBoundingClientRect().height!;
+        const heightDifference = newHeight - previousHeight.current;
+        if (
+          previousHeight.current &&
+          heightDifference != 0 &&
+          endPaddingRef.current &&
+          scrollableDivRef &&
+          scrollableDivRef.current
+        ) {
+          endPaddingRef.current.style.transition = "height 0.3s ease-out";
+          endPaddingRef.current.style.height = `${Math.max(newHeight - 50, 0)}px`;
+
+          scrollableDivRef?.current.scrollBy({
+            left: 0,
+            top: Math.max(heightDifference, 0),
+            behavior: "smooth",
+          });
+        }
+        previousHeight.current = newHeight;
+      }
+    }, 100);
+  };
+
+  const clientScrollToBottom = (fast?: boolean) => {
+    setTimeout(
+      () => {
+        endDivRef.current?.scrollIntoView({ behavior: "smooth" });
+        setHasPerformedInitialScroll(true);
+      },
+      fast ? 50 : 500
+    );
+  };
+
+  const isCancelledRef = useRef<boolean>(isCancelled); // scroll is cancelled
   useEffect(() => {
     isCancelledRef.current = isCancelled;
   }, [isCancelled]);
 
-  const [currentFeedback, setCurrentFeedback] = useState<
+  const distance = 500; // distance that should "engage" the scroll
-    [FeedbackType, number] | null
+  const debounce = 100; // time for debouncing
-  >(null);
-  const [sharingModalVisible, setSharingModalVisible] =
-    useState<boolean>(false);
 
-  // auto scroll as message comes out
+  useScrollonStream({
-  const scrollableDivRef = useRef<HTMLDivElement>(null);
+    isStreaming,
-  const endDivRef = useRef<HTMLDivElement>(null);
+    scrollableDivRef,
-  useEffect(() => {
+    scrollDist,
-    if (isStreaming || !message) {
+    endDivRef,
-      handleAutoScroll(endDivRef, scrollableDivRef);
+    distance,
-    }
+    debounce,
   });
 
-  // scroll to bottom initially
   const [hasPerformedInitialScroll, setHasPerformedInitialScroll] =
     useState(false);
+
+  // on new page
   useEffect(() => {
-    endDivRef.current?.scrollIntoView();
+    clientScrollToBottom();
-    setHasPerformedInitialScroll(true);
+  }, [chatSessionId]);
-  }, [isFetchingChatMessages]);
 
   // handle re-sizing of the text area
   const textAreaRef = useRef<HTMLTextAreaElement>(null);
   useEffect(() => {
-    const textarea = textAreaRef.current;
+    handleInputResize();
-    if (textarea) {
-      textarea.style.height = "0px";
-      textarea.style.height = `${Math.min(
-        textarea.scrollHeight,
-        MAX_INPUT_HEIGHT
-      )}px`;
-    }
   }, [message]);
 
+  // tracks scrolling
+  useEffect(() => {
+    updateScrollTracking();
+  }, [messageHistory]);
+
   // used for resizing of the document sidebar
   const masterFlexboxRef = useRef<HTMLDivElement>(null);
   const [maxDocumentSidebarWidth, setMaxDocumentSidebarWidth] = useState<
@@ -502,6 +566,53 @@ export function ChatPage({
     documentSidebarInitialWidth = Math.min(700, maxDocumentSidebarWidth);
   }
 
+  class CurrentMessageFIFO {
+    private stack: PacketType[] = [];
+    isComplete: boolean = false;
+    error: string | null = null;
+
+    push(packetBunch: PacketType) {
+      this.stack.push(packetBunch);
+    }
+
+    nextPacket(): PacketType | undefined {
+      return this.stack.shift();
+    }
+
+    isEmpty(): boolean {
+      return this.stack.length === 0;
+    }
+  }
+  async function updateCurrentMessageFIFO(
+    stack: CurrentMessageFIFO,
+    params: any
+  ) {
+    try {
+      for await (const packetBunch of sendMessage(params)) {
+        for (const packet of packetBunch) {
+          stack.push(packet);
+        }
+
+        if (isCancelledRef.current) {
+          setIsCancelled(false);
+          break;
+        }
+      }
+    } catch (error) {
+      stack.error = String(error);
+    } finally {
+      stack.isComplete = true;
+    }
+  }
+
+  const resetInputBar = () => {
+    setMessage("");
+    setCurrentMessageFiles([]);
+    if (endPaddingRef.current) {
+      endPaddingRef.current.style.height = `95px`;
+    }
+  };
+
   const onSubmit = async ({
     messageIdToResend,
     messageOverride,
@@ -515,6 +626,7 @@ export function ChatPage({
     forceSearch?: boolean;
     isSeededChat?: boolean;
   } = {}) => {
+    clientScrollToBottom();
     let currChatSessionId: number;
     let isNewSession = chatSessionId === null;
     const searchParamBasedChatSessionName =
@@ -596,9 +708,7 @@ export function ChatPage({
     if (!parentMessage && frozenCompleteMessageMap.size === 2) {
       parentMessage = frozenCompleteMessageMap.get(SYSTEM_MESSAGE_ID) || null;
     }
-    setMessage("");
+    resetInputBar();
-    setCurrentMessageFiles([]);
-
     setIsStreaming(true);
     let answer = "";
     let query: string | null = null;
@@ -614,7 +724,9 @@ export function ChatPage({
     try {
       const lastSuccessfulMessageId =
         getLastSuccessfulMessageId(currMessageHistory);
-      for await (const packetBunch of sendMessage({
+
+      const stack = new CurrentMessageFIFO();
+      updateCurrentMessageFIFO(stack, {
         message: currMessage,
         fileDescriptors: currentMessageFiles,
         parentMessageId: lastSuccessfulMessageId,
@@ -647,8 +759,32 @@ export function ChatPage({
         systemPromptOverride:
           searchParams.get(SEARCH_PARAM_NAMES.SYSTEM_PROMPT) || undefined,
         useExistingUserMessage: isSeededChat,
-      })) {
+      });
-        for (const packet of packetBunch) {
+      const updateFn = (messages: Message[]) => {
+        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 delay = (ms: number) => {
+        return new Promise((resolve) => setTimeout(resolve, ms));
+      };
+
+      await delay(50);
+      while (!stack.isComplete || !stack.isEmpty()) {
+        await delay(2);
+
+        if (!stack.isEmpty()) {
+          const packet = stack.nextPacket();
+
+          if (packet) {
             if (Object.hasOwn(packet, "answer_piece")) {
               answer += (packet as AnswerPiecePacket).answer_piece;
             } else if (Object.hasOwn(packet, "top_documents")) {
@@ -660,6 +796,14 @@ export function ChatPage({
                 // we have to use -1)
                 setSelectedMessageForDocDisplay(TEMP_USER_MESSAGE_ID);
               }
+            } else if (Object.hasOwn(packet, "tool_name")) {
+              toolCalls = [
+                {
+                  tool_name: (packet as ToolCallMetadata).tool_name,
+                  tool_args: (packet as ToolCallMetadata).tool_args,
+                  tool_result: (packet as ToolCallMetadata).tool_result,
+                },
+              ];
             } else if (Object.hasOwn(packet, "file_ids")) {
               aiMessageImages = (packet as ImageGenerationDisplay).file_ids.map(
                 (fileId) => {
@@ -682,20 +826,7 @@ export function ChatPage({
             } else if (Object.hasOwn(packet, "message_id")) {
               finalMessage = packet as BackendMessage;
             }
-        }
+
-        const updateFn = (messages: Message[]) => {
-          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 =
@@ -717,18 +848,21 @@ export function ChatPage({
                 type: error ? "error" : "assistant",
                 retrievalType,
                 query: finalMessage?.rephrased_query || query,
-            documents: finalMessage?.context_docs?.top_documents || documents,
+                documents:
+                  finalMessage?.context_docs?.top_documents || documents,
                 citations: finalMessage?.citations || {},
                 files: finalMessage?.files || aiMessageImages || [],
                 toolCalls: finalMessage?.tool_calls || toolCalls,
                 parentMessageId: newUserMessageId,
               },
             ]);
+          }
           if (isCancelledRef.current) {
             setIsCancelled(false);
             break;
           }
         }
+      }
     } catch (e: any) {
       const errorMsg = e.message;
       upsertToCompleteMessageMap({
@@ -817,7 +951,6 @@ export function ChatPage({
     if (persona && persona.id !== livePersona.id) {
       // remove uploaded files
       setCurrentMessageFiles([]);
-
       setSelectedPersona(persona);
       textAreaRef.current?.focus();
       router.push(buildChatUrl(searchParams, null, persona.id));
@@ -1051,7 +1184,7 @@ export function ChatPage({
                               ? completeMessageMap.get(message.parentMessageId)
                               : null;
                             return (
-                              <div key={i}>
+                              <div key={`${i}-${existingChatSessionId}`}>
                                 <HumanMessage
                                   content={message.message}
                                   files={message.files}
@@ -1107,8 +1240,15 @@ export function ChatPage({
                             const previousMessage =
                               i !== 0 ? messageHistory[i - 1] : null;
                             return (
+                              <div
+                                key={`${i}-${existingChatSessionId}`}
+                                ref={
+                                  i == messageHistory.length - 1
+                                    ? lastMessageRef
+                                    : null
+                                }
+                              >
                                 <AIMessage
-                                key={message.messageId}
                                   messageId={message.messageId}
                                   content={message.message}
                                   files={message.files}
@@ -1117,7 +1257,9 @@ export function ChatPage({
                                   citedDocuments={getCitedDocumentsFromMessage(
                                     message
                                   )}
-                                toolCall={message?.toolCalls?.[0]}
+                                  toolCall={
+                                    message.toolCalls && message.toolCalls[0]
+                                  }
                                   isComplete={
                                     i !== messageHistory.length - 1 ||
                                     !isStreaming
@@ -1127,7 +1269,8 @@ export function ChatPage({
                                       message.documents.length > 0) === true
                                   }
                                   handleFeedback={
-                                  i === messageHistory.length - 1 && isStreaming
+                                    i === messageHistory.length - 1 &&
+                                    isStreaming
                                       ? undefined
                                       : (feedbackType) =>
                                           setCurrentFeedback([
@@ -1166,7 +1309,9 @@ export function ChatPage({
                                         }
                                       : undefined
                                   }
-                                isCurrentlyShowingRetrieved={isShowingRetrieved}
+                                  isCurrentlyShowingRetrieved={
+                                    isShowingRetrieved
+                                  }
                                   handleShowRetrieved={(messageNumber) => {
                                     if (isShowingRetrieved) {
                                       setSelectedMessageForDocDisplay(null);
@@ -1200,10 +1345,11 @@ export function ChatPage({
                                   }}
                                   retrievalDisabled={retrievalDisabled}
                                 />
+                              </div>
                             );
                           } else {
                             return (
-                              <div key={i}>
+                              <div key={`${i}-${existingChatSessionId}`}>
                                 <AIMessage
                                   messageId={message.messageId}
                                   personaName={livePersona.name}
@@ -1222,7 +1368,9 @@ export function ChatPage({
                           messageHistory.length > 0 &&
                           messageHistory[messageHistory.length - 1].type ===
                             "user" && (
-                            <div key={messageHistory.length}>
+                            <div
+                              key={`${messageHistory.length}-${existingChatSessionId}`}
+                            >
                               <AIMessage
                                 messageId={null}
                                 personaName={livePersona.name}
@@ -1245,7 +1393,8 @@ export function ChatPage({
                           )}
 
                         {/* Some padding at the bottom so the search bar has space at the bottom to not cover the last message*/}
-                        <div className={`min-h-[100px] w-full`}></div>
+                        <div ref={endPaddingRef} className=" h-[95px]" />
+                        <div ref={endDivRef}></div>
 
                         {livePersona &&
                           livePersona.starter_messages &&
@@ -1270,7 +1419,10 @@ export function ChatPage({
                             >
                               {livePersona.starter_messages.map(
                                 (starterMessage, i) => (
-                                  <div key={i} className="w-full">
+                                  <div
+                                    key={`${i}-${existingChatSessionId}`}
+                                    className="w-full"
+                                  >
                                     <StarterMessage
                                       starterMessage={starterMessage}
                                       onClick={() =>
@@ -1289,8 +1441,21 @@ export function ChatPage({
                       </div>
                     </div>
 
-                    <div className="absolute bottom-0 z-10 w-full">
+                    <div
-                      <div className="w-full pb-4">
+                      ref={inputRef}
+                      className="absolute bottom-0 z-10 w-full"
+                    >
+                      <div className="w-full relative pb-4">
+                        {aboveHorizon && (
+                          <div className="pointer-events-none w-full bg-transparent flex sticky justify-center">
+                            <button
+                              onClick={() => clientScrollToBottom(true)}
+                              className="p-1 pointer-events-auto rounded-2xl bg-background-strong border border-border mb-2 mx-auto "
+                            >
+                              <FiArrowDown size={18} />
+                            </button>
+                          </div>
+                        )}
                         <ChatInputBar
                           message={message}
                           setMessage={setMessage}

web/src/app/chat/lib.tsx

@@ -5,7 +5,14 @@ import {
 } from "@/lib/search/interfaces";
 import { handleStream } from "@/lib/search/streamingUtils";
 import { FeedbackType } from "./types";
-import { RefObject } from "react";
+import {
+  Dispatch,
+  MutableRefObject,
+  RefObject,
+  SetStateAction,
+  useEffect,
+  useRef,
+} from "react";
 import {
   BackendMessage,
   ChatSession,
@@ -65,6 +72,14 @@ export async function createChatSession(
   return chatSessionResponseJson.chat_session_id;
 }
 
+export type PacketType =
+  | ToolCallMetadata
+  | BackendMessage
+  | AnswerPiecePacket
+  | DocumentsResponse
+  | ImageGenerationDisplay
+  | StreamingError;
+
 export async function* sendMessage({
   message,
   fileDescriptors,
@@ -150,14 +165,7 @@ export async function* sendMessage({
     throw Error(`Failed to send message - ${errorMsg}`);
   }
 
-  yield* handleStream<
+  yield* handleStream<PacketType>(sendMessageResponse);
-    | AnswerPiecePacket
-    | DocumentsResponse
-    | BackendMessage
-    | ImageGenerationDisplay
-    | ToolCallMetadata
-    | StreamingError
-  >(sendMessageResponse);
 }
 
 export async function nameChatSession(chatSessionId: number, message: string) {
@@ -250,24 +258,6 @@ export async function* simulateLLMResponse(input: string, delay: number = 30) {
   }
 }
 
-export function handleAutoScroll(
-  endRef: RefObject<any>,
-  scrollableRef: RefObject<any>,
-  buffer: number = 300
-) {
-  // Auto-scrolls if the user is within `buffer` of the bottom of the scrollableRef
-  if (endRef && endRef.current && scrollableRef && scrollableRef.current) {
-    if (
-      scrollableRef.current.scrollHeight -
-        scrollableRef.current.scrollTop -
-        buffer <=
-      scrollableRef.current.clientHeight
-    ) {
-      endRef.current.scrollIntoView({ behavior: "smooth" });
-    }
-  }
-}
-
 export function getHumanAndAIMessageFromMessageNumber(
   messageHistory: Message[],
   messageId: number
@@ -565,3 +555,92 @@ export async function uploadFilesForChat(
 
   return [responseJson.files as FileDescriptor[], null];
 }
+
+export async function useScrollonStream({
+  isStreaming,
+  scrollableDivRef,
+  scrollDist,
+  endDivRef,
+  distance,
+  debounce,
+}: {
+  isStreaming: boolean;
+  scrollableDivRef: RefObject<HTMLDivElement>;
+  scrollDist: MutableRefObject<number>;
+  endDivRef: RefObject<HTMLDivElement>;
+  distance: number;
+  debounce: number;
+}) {
+  const preventScrollInterference = useRef<boolean>(false);
+  const preventScroll = useRef<boolean>(false);
+  const blockActionRef = useRef<boolean>(false);
+  const previousScroll = useRef<number>(0);
+
+  useEffect(() => {
+    if (isStreaming && scrollableDivRef && scrollableDivRef.current) {
+      let newHeight: number = scrollableDivRef.current?.scrollTop!;
+      const heightDifference = newHeight - previousScroll.current;
+      previousScroll.current = newHeight;
+
+      // Prevent streaming scroll
+      if (heightDifference < 0 && !preventScroll.current) {
+        scrollableDivRef.current.style.scrollBehavior = "auto";
+        scrollableDivRef.current.scrollTop = scrollableDivRef.current.scrollTop;
+        scrollableDivRef.current.style.scrollBehavior = "smooth";
+        preventScrollInterference.current = true;
+        preventScroll.current = true;
+
+        setTimeout(() => {
+          preventScrollInterference.current = false;
+        }, 2000);
+        setTimeout(() => {
+          preventScroll.current = false;
+        }, 10000);
+      }
+
+      // Ensure can scroll if scroll down
+      else if (!preventScrollInterference.current) {
+        preventScroll.current = false;
+      }
+      if (
+        scrollDist.current < distance &&
+        !blockActionRef.current &&
+        !blockActionRef.current &&
+        !preventScroll.current &&
+        endDivRef &&
+        endDivRef.current
+      ) {
+        // catch up if necessary!
+        const scrollAmount = scrollDist.current + 10000;
+        if (scrollDist.current > 140) {
+          endDivRef.current.scrollIntoView();
+        } else {
+          blockActionRef.current = true;
+
+          scrollableDivRef?.current?.scrollBy({
+            left: 0,
+            top: Math.max(0, scrollAmount),
+            behavior: "smooth",
+          });
+
+          setTimeout(() => {
+            blockActionRef.current = false;
+          }, debounce);
+        }
+      }
+    }
+  });
+
+  // scroll on end of stream if within distance
+  useEffect(() => {
+    if (scrollableDivRef?.current && !isStreaming) {
+      if (scrollDist.current < distance) {
+        scrollableDivRef?.current?.scrollBy({
+          left: 0,
+          top: Math.max(scrollDist.current + 600, 0),
+          behavior: "smooth",
+        });
+      }
+    }
+  }, [isStreaming]);
+}