Implemented autoScroll on generation for React

nlkitai · Apr 24, 2024 · e74684d · e74684d
1 parent a1ce50e
commit e74684d
Show file tree

Hide file tree

Showing 14 changed files with 222 additions and 51 deletions.
diff --git a/packages/react/core/src/exports/AiChat.tsx b/packages/react/core/src/exports/AiChat.tsx
@@ -12,16 +12,24 @@ import {adapterParamToUsableAdapter} from '../utils/adapterParamToUsableAdapter'
 import {chatItemsToChatSegment} from '../utils/chatItemsToChatSegment';
 import {reactPropsToCoreProps} from '../utils/reactPropsToCoreProps';
 import {useAiChatStyle} from './hooks/useAiChatStyle';
+import {useAutoScrollHandler} from './hooks/useAutoScrollHandler';
 import {useSubmitPromptHandler} from './hooks/useSubmitPromptHandler';
 import {AiChatComponentProps} from './props';
 
+const defaultAutoScrollOption = true;
+
 export const AiChat: <AiMsg>(
     props: AiChatComponentProps<AiMsg>,
 ) => ReactElement = function <AiMsg>(
     props: AiChatComponentProps<AiMsg>,
 ): ReactElement {
     const conversationRef = useRef<ImperativeConversationCompProps>(null);
     const exceptionBoxRef = useRef<HTMLDivElement>(null);
+    const conversationContainerRef = useRef<HTMLDivElement>(null);
+    const autoScrollHandler = useAutoScrollHandler(
+        conversationContainerRef,
+        props.conversationOptions?.autoScroll,
+    );
 
     const exceptionBoxController = useMemo(() => {
         return exceptionBoxRef.current ? createExceptionsBoxController(exceptionBoxRef.current) : undefined;
@@ -44,6 +52,7 @@ export const AiChat: <AiMsg>(
         adapterToUse ? {aiChatProps: reactPropsToCoreProps<AiMsg>(props, adapterToUse)} : undefined
     ), [props, adapterToUse]);
 
+    const lastActiveSegmentIdRef = useRef<string | undefined>(undefined);
     const hasValidInput = useMemo(() => prompt.length > 0, [prompt]);
     const handlePromptChange = useCallback((value: string) => setPrompt(value), [setPrompt]);
     const handleSubmitPrompt = useSubmitPromptHandler({
@@ -62,6 +71,21 @@ export const AiChat: <AiMsg>(
         props.initialConversation ? chatItemsToChatSegment(props.initialConversation) : undefined,
     ), [props.initialConversation]);
 
+    const handleLastActiveSegmentChange = useCallback((data: {uid: string; div: HTMLDivElement} | undefined) => {
+        if (!autoScrollHandler) {
+            return;
+        }
+
+        if (data) {
+            lastActiveSegmentIdRef.current = data.uid;
+            autoScrollHandler.handleNewChatSegmentAdded(data.uid, data.div);
+        } else {
+            if (lastActiveSegmentIdRef.current) {
+                autoScrollHandler.handleChatSegmentRemoved(lastActiveSegmentIdRef.current);
+            }
+        }
+    }, [autoScrollHandler]);
+
     const segments = useMemo(() => (
         initialSegment ? [initialSegment, ...chatSegments] : chatSegments
     ), [initialSegment, chatSegments]);
@@ -72,20 +96,23 @@ export const AiChat: <AiMsg>(
         themeId: props.themeId,
     }).join(' ');
 
-    const ForwardConversationComp = forwardRef(ConversationComp<AiMsg>);
+    const ForwardConversationComp = useMemo(() => forwardRef(
+        ConversationComp<AiMsg>,
+    ), []);
 
     return (
         <div className={rootClassNames} style={rootStyle}>
             <div className={compExceptionsBoxClassName} ref={exceptionBoxRef}/>
             <div className="nlux-chtRm-cntr">
-                <div className="nlux-chtRm-cnv-cntr">
+                <div className="nlux-chtRm-cnv-cntr" ref={conversationContainerRef}>
                     <ForwardConversationComp
                         ref={conversationRef}
                         segments={segments}
                         conversationOptions={props.conversationOptions}
                         personaOptions={props.personaOptions}
                         customRenderer={props.aiMessageComponent}
                         syntaxHighlighter={props.syntaxHighlighter}
+                        onLastActiveSegmentChange={handleLastActiveSegmentChange}
                     />
                 </div>
                 <div className="nlux-chtRm-prmptBox-cntr">

diff --git a/packages/react/core/src/exports/hooks/useAutoScrollHandler.ts b/packages/react/core/src/exports/hooks/useAutoScrollHandler.ts
@@ -0,0 +1,56 @@
+import {MutableRefObject, useEffect, useRef, useState} from 'react';
+import {createAutoScrollHandler} from '../../../../../shared/src/interactions/autoScroll/autoScrollHandler';
+import {AutoScrollHandler} from '../../../../../shared/src/interactions/autoScroll/type';
+
+const defaultAutoScrollOption = true;
+
+export const useAutoScrollHandler = (
+    conversationContainerRef: MutableRefObject<HTMLDivElement | null>,
+    autoScroll?: boolean,
+) => {
+    const [autoScrollHandler, setAutoScrollHandler] = useState<AutoScrollHandler>();
+    const [conversationContainer, setConversationContainer] = useState<HTMLDivElement>();
+
+    const autoScrollHandlerRef = useRef(autoScrollHandler);
+    const autoScrollPropRef = useRef(autoScroll);
+
+    useEffect(() => {
+        const currentConversationContainer = conversationContainerRef.current || undefined;
+        if (currentConversationContainer !== conversationContainer) {
+            setConversationContainer(currentConversationContainer);
+        }
+    }); // No dependencies - If statement inside the effect will handle the update
+
+    useEffect(() => {
+        if (!conversationContainer) {
+            if (autoScrollHandlerRef.current) {
+                autoScrollHandlerRef.current.destroy();
+                setAutoScrollHandler(undefined);
+                autoScrollHandlerRef.current = undefined;
+            }
+
+            return;
+        }
+
+        if (autoScrollHandlerRef.current) {
+            autoScrollHandlerRef.current.updateConversationContainer(conversationContainer);
+        } else {
+            autoScrollHandlerRef.current = createAutoScrollHandler(
+                conversationContainer,
+                autoScrollPropRef.current ?? defaultAutoScrollOption,
+            );
+            setAutoScrollHandler(autoScrollHandlerRef.current);
+        }
+    }, [conversationContainer]);
+
+    useEffect(() => {
+        autoScrollPropRef.current = autoScroll;
+        if (autoScrollHandlerRef.current) {
+            autoScrollHandlerRef.current.updateProps({
+                autoScroll: autoScroll ?? defaultAutoScrollOption,
+            });
+        }
+    }, [autoScroll]);
+
+    return autoScrollHandler;
+};
diff --git a/packages/react/core/src/exports/hooks/useSubmitPromptHandler.ts b/packages/react/core/src/exports/hooks/useSubmitPromptHandler.ts
@@ -33,6 +33,8 @@ export const useSubmitPromptHandler = <AiMsg>(props: SubmitPromptHandlerProps<Ai
     } = props;
 
     const hasValidInput = useMemo(() => prompt.length > 0, [prompt]);
+
+    // React functions and state that can be accessed by non-React DOM update code
     const domToReactRef = useRef({
         chatSegments,
         setChatSegments,

diff --git a/packages/react/core/src/logic/ChatSegment/ChatSegmentComp.tsx b/packages/react/core/src/logic/ChatSegment/ChatSegmentComp.tsx
@@ -17,7 +17,7 @@ export const ChatSegmentComp: <AiMsg>(
     props: ChatSegmentProps<AiMsg>,
     ref: Ref<ChatSegmentImperativeProps<AiMsg>>,
 ): ReactNode {
-    const {chatSegment} = props;
+    const {chatSegment, containerRef} = props;
     const chatItemsRef = useMemo(
         () => new Map<string, RefObject<ChatItemImperativeProps>>(), [],
     );
@@ -52,9 +52,6 @@ export const ChatSegmentComp: <AiMsg>(
     const rootClassName = useMemo(() => getChatSegmentClassName(chatSegment.status), [chatSegment.status]);
 
     useImperativeHandle(ref, () => ({
-        scrollToBottom: () => {
-            // TODO - Implement scroll to bottom
-        },
         streamChunk: (messageId: string, chunk: string) => {
             const messageCompRef = chatItemsRef.get(messageId);
             if (messageCompRef?.current) {
@@ -63,24 +60,24 @@ export const ChatSegmentComp: <AiMsg>(
         },
     }), []);
 
+    const ForwardRefChatItemComp = useMemo(() => forwardRef(
+        ChatItemComp<AiMsg>,
+    ), []);
+
     const chatItems = chatSegment.items;
     if (chatItems.length === 0) {
         return null;
     }
 
     return (
-        <div className={rootClassName}>
+        <div className={rootClassName} ref={containerRef}>
             {chatItems.map((chatItem, index) => {
                 let ref: RefObject<ChatItemImperativeProps> | undefined = chatItemsRef.get(chatItem.uid);
                 if (!ref) {
                     ref = createRef<ChatItemImperativeProps>();
                     chatItemsRef.set(chatItem.uid, ref);
                 }
 
-                const ForwardRefChatItemComp = forwardRef(
-                    ChatItemComp<AiMsg>,
-                );
-
                 if (chatItem.participantRole === 'user') {
                     //
                     // User chat item — That should always be in complete state.

diff --git a/packages/react/core/src/logic/ChatSegment/props.ts b/packages/react/core/src/logic/ChatSegment/props.ts
@@ -1,5 +1,5 @@
 import {HighlighterExtension} from '@nlux/core';
-import {FunctionComponent, ReactElement} from 'react';
+import {FunctionComponent, ReactElement, RefObject} from 'react';
 import {ChatSegment} from '../../../../../shared/src/types/chatSegment/chatSegment';
 import {PersonaOptions} from '../../exports/personaOptions';
 
@@ -9,6 +9,7 @@ export type ChatSegmentProps<AiMsg> = {
     loader?: ReactElement;
     personaOptions?: PersonaOptions;
     syntaxHighlighter?: HighlighterExtension;
+    containerRef?: RefObject<HTMLDivElement>;
 };
 
 export type ChatSegmentImperativeProps<AiMsg> = {

diff --git a/packages/react/core/src/logic/Conversation/ConversationComp.tsx b/packages/react/core/src/logic/Conversation/ConversationComp.tsx
@@ -1,4 +1,4 @@
-import {createRef, forwardRef, ReactNode, Ref, RefObject, useEffect, useImperativeHandle, useMemo} from 'react';
+import {createRef, forwardRef, ReactNode, Ref, RefObject, useEffect, useImperativeHandle, useMemo, useRef} from 'react';
 import {WelcomeMessageComp} from '../../ui/WelcomeMessage/WelcomeMessageComp';
 import {ChatSegmentComp} from '../ChatSegment/ChatSegmentComp';
 import {ChatSegmentImperativeProps} from '../ChatSegment/props';
@@ -13,11 +13,19 @@ export const ConversationComp: ConversationCompType = function <AiMsg>(
     props: ConversationCompProps<AiMsg>,
     ref: Ref<ImperativeConversationCompProps>,
 ): ReactNode {
-    const {segments, personaOptions} = props;
+    const {
+        segments,
+        personaOptions,
+        onLastActiveSegmentChange,
+    } = props;
+
     const hasMessages = useMemo(() => segments.some((segment) => segment.items.length > 0), [segments]);
     const hasAiPersona = personaOptions?.bot?.name && personaOptions.bot.picture;
     const showWelcomeMessage = hasAiPersona && !hasMessages;
 
+    const lastSegmentContainerRef = createRef<HTMLDivElement>();
+    const lastActiveSegmentDataReportedRef = useRef<{uid: string; div: HTMLDivElement} | undefined>(undefined);
+
     const chatSegmentsRef = useMemo(
         () => new Map<string, RefObject<ChatSegmentImperativeProps<any>>>(), [],
     );
@@ -37,10 +45,42 @@ export const ConversationComp: ConversationCompType = function <AiMsg>(
         }
     }, [props.segments]);
 
+    const lastActiveSegmentId = useMemo(() => {
+        const lastSegment = segments.length > 0 ? segments[segments.length - 1] : undefined;
+        return lastSegment?.status === 'active' ? lastSegment.uid : undefined;
+    }, [segments]);
+
+    const ForwardRefChatSegmentComp = useMemo(() => forwardRef(
+        ChatSegmentComp<AiMsg>,
+    ), []);
+
+    // Whenever the last active segment div+id changes, call the onLastActiveSegmentChange callback
+    useEffect(() => {
+        if (!onLastActiveSegmentChange) {
+            return;
+        }
+
+        const lastReportedData = lastActiveSegmentDataReportedRef.current;
+        if (lastActiveSegmentId === lastReportedData?.uid
+            && lastSegmentContainerRef.current === lastReportedData?.div) {
+            return;
+        }
+
+        const data = (lastActiveSegmentId && lastSegmentContainerRef.current) ? {
+            uid: lastActiveSegmentId,
+            div: lastSegmentContainerRef.current,
+        } : undefined;
+
+        if (!data && !lastActiveSegmentDataReportedRef.current) {
+            return;
+        }
+
+        onLastActiveSegmentChange(data);
+        lastActiveSegmentDataReportedRef.current = data;
+    }); // No dependencies on purpose — we want to run this effect on every render cycle
+    // 'if' statements inside the effect will prevent unnecessary calls to the callback
+
     useImperativeHandle(ref, () => ({
-        scrollToBottom: () => {
-            // TODO - Implement scroll to bottom
-        },
         streamChunk: (segmentId: string, messageId: string, chunk: string) => {
             const chatSegmentRef = chatSegmentsRef.get(segmentId);
             if (chatSegmentRef?.current) {
@@ -59,26 +99,24 @@ export const ConversationComp: ConversationCompType = function <AiMsg>(
                 />
             )}
             <div className="nlux-chtRm-cnv-sgmts-cntr">
-                {segments.map((segment) => {
+                {segments.map((segment, index) => {
+                    const isLastSegment = index === segments.length - 1;
                     let ref: RefObject<ChatSegmentImperativeProps<any>> | undefined = chatSegmentsRef.get(segment.uid);
                     if (!ref) {
                         ref = createRef();
                         chatSegmentsRef.set(segment.uid, ref);
                     }
 
-                    const ForwardRefChatItemComp = forwardRef(
-                        ChatSegmentComp<AiMsg>,
-                    );
-
                     return (
-                        <ForwardRefChatItemComp
+                        <ForwardRefChatSegmentComp
                             ref={ref}
                             key={segment.uid}
                             chatSegment={segment}
                             personaOptions={personaOptions}
                             loader={props.loader}
                             customRenderer={props.customRenderer}
                             syntaxHighlighter={props.syntaxHighlighter}
+                            containerRef={isLastSegment ? lastSegmentContainerRef : undefined}
                         />
                     );
                 })}

diff --git a/packages/react/core/src/logic/Conversation/props.ts b/packages/react/core/src/logic/Conversation/props.ts
@@ -10,9 +10,12 @@ export type ConversationCompProps<AiMsg> = {
     customRenderer?: FunctionComponent<{message: AiMsg}>;
     syntaxHighlighter?: HighlighterExtension;
     loader?: ReactElement;
+    onLastActiveSegmentChange?: (data: {
+        uid: string;
+        div: HTMLDivElement;
+    } | undefined) => void;
 };
 
 export type ImperativeConversationCompProps = {
-    scrollToBottom: () => void;
     streamChunk: (segmentId: string, messageId: string, chunk: string) => void;
 };
diff --git a/packages/react/core/src/ui/ChatItem/ChatItemComp.tsx b/packages/react/core/src/ui/ChatItem/ChatItemComp.tsx
@@ -59,9 +59,9 @@ export const ChatItemComp: <AiMsg>(
         return () => <>{props.message !== undefined ? props.message : ''}</>;
     }, [props.customRenderer, props.message]);
 
-    const ForwardRefStreamContainerComp = forwardRef(
+    const ForwardRefStreamContainerComp = useMemo(() => forwardRef(
         StreamContainerComp,
-    );
+    ), []);
 
     return (
         <div className={className}>