fix: use accordion for steps and fix streaming auto scroll

Chainlit · Feb 8, 2025 · 098097e · 098097e
1 parent a552bd2
commit 098097e
Show file tree

Hide file tree

Showing 9 changed files with 114 additions and 94 deletions.
diff --git a/frontend/src/components/CodeSnippet.tsx b/frontend/src/components/CodeSnippet.tsx
@@ -1,3 +1,4 @@
+import { cn } from '@/lib/utils';
 import hljs from 'highlight.js';
 import { useEffect, useRef } from 'react';
 
@@ -64,7 +65,9 @@ export default function CodeSnippet({ ...props }: CodeProps) {
   ) : null;
 
   const nonHighlightedCode = showSyntaxHighlighter ? null : (
-    <div className="p-2 rounded-b-md min-h-20 overflow-x-auto bg-accent">
+    <div
+      className={cn('rounded-b-md overflow-x-auto bg-accent', code && 'p-2')}
+    >
       <code className="whitespace-pre-wrap">{code}</code>
     </div>
   );

diff --git a/frontend/src/components/chat/Footer.tsx b/frontend/src/components/chat/Footer.tsx
@@ -1,4 +1,5 @@
 import { cn, hasMessage } from '@/lib/utils';
+import { MutableRefObject } from 'react';
 
 import { FileSpec, useChatMessages } from '@chainlit/react-client';
 
@@ -11,23 +12,20 @@ interface Props {
   fileSpec: FileSpec;
   onFileUpload: (payload: File[]) => void;
   onFileUploadError: (error: string) => void;
-  setAutoScroll: (autoScroll: boolean) => void;
-  autoScroll: boolean;
+  autoScrollRef: MutableRefObject<boolean>;
   showIfEmptyThread?: boolean;
 }
 
-export default function ChatFooter({
-  autoScroll,
-  showIfEmptyThread,
-  ...props
-}: Props) {
+export default function ChatFooter({ showIfEmptyThread, ...props }: Props) {
   const { messages } = useChatMessages();
   if (!hasMessage(messages) && !showIfEmptyThread) return null;
 
   return (
     <div className={cn('relative flex flex-col items-center gap-2 w-full')}>
-      {!autoScroll ? (
-        <ScrollDownButton onClick={() => props.setAutoScroll(true)} />
+      {!props.autoScrollRef.current ? (
+        <ScrollDownButton
+          onClick={() => (props.autoScrollRef.current = true)}
+        />
       ) : null}
       <MessageComposer {...props} />
       <WaterMark />

diff --git a/frontend/src/components/chat/MessageComposer/index.tsx b/frontend/src/components/chat/MessageComposer/index.tsx
@@ -1,4 +1,4 @@
-import { useCallback, useRef, useState } from 'react';
+import { MutableRefObject, useCallback, useRef, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { useRecoilState, useSetRecoilState } from 'recoil';
 import { v4 as uuidv4 } from 'uuid';
@@ -29,14 +29,14 @@ interface Props {
   fileSpec: FileSpec;
   onFileUpload: (payload: File[]) => void;
   onFileUploadError: (error: string) => void;
-  setAutoScroll: (autoScroll: boolean) => void;
+  autoScrollRef: MutableRefObject<boolean>;
 }
 
 export default function MessageComposer({
   fileSpec,
   onFileUpload,
   onFileUploadError,
-  setAutoScroll
+  autoScrollRef
 }: Props) {
   const inputRef = useRef<InputMethods>(null);
   const [value, setValue] = useState('');
@@ -88,7 +88,9 @@ export default function MessageComposer({
         ?.filter((a) => !!a.serverId)
         .map((a) => ({ id: a.serverId! }));
 
-      setAutoScroll(true);
+      if (autoScrollRef) {
+        autoScrollRef.current = true;
+      }
       sendMessage(message, fileReferences);
     },
     [user, sendMessage]
@@ -107,7 +109,9 @@ export default function MessageComposer({
       };
 
       replyMessage(message);
-      setAutoScroll(true);
+      if (autoScrollRef) {
+        autoScrollRef.current = true;
+      }
     },
     [user, replyMessage]
   );

diff --git a/frontend/src/components/chat/Messages/Message/Content/index.tsx b/frontend/src/components/chat/Messages/Message/Content/index.tsx
@@ -38,11 +38,9 @@ const MessageContent = memo(
     const isMessage = message.type.includes('message');
 
     const outputMarkdown = (
-      <div className="flex flex-col gap-2">
+      <>
         {!isMessage && displayInput ? (
-          <div className="text-lg font-semibold leading-none tracking-tight">
-            Output
-          </div>
+          <div className="font-medium">Output</div>
         ) : null}
         <Markdown
           allowHtml={allowHtml}
@@ -51,7 +49,7 @@ const MessageContent = memo(
         >
           {output}
         </Markdown>
-      </div>
+      </>
     );
 
     let inputMarkdown;
@@ -73,18 +71,17 @@ const MessageContent = memo(
         });
 
       inputMarkdown = (
-        <div className="flex flex-col gap-2">
-          <div className="text-lg font-semibold leading-none tracking-tight">
-            Input
-          </div>
+        <>
+          <div className="font-medium">Input</div>
+
           <Markdown
             allowHtml={allowHtml}
             latex={latex}
             refElements={inputRefElements}
           >
             {input}
           </Markdown>
-        </div>
+        </>
       );
     }
 

diff --git a/frontend/src/components/chat/Messages/Message/Step.tsx b/frontend/src/components/chat/Messages/Message/Step.tsx
@@ -1,9 +1,14 @@
 import { cn } from '@/lib/utils';
-import { ChevronDown, ChevronUp } from 'lucide-react';
-import { PropsWithChildren, useMemo, useState } from 'react';
+import { PropsWithChildren, useMemo } from 'react';
 
 import type { IStep } from '@chainlit/react-client';
 
+import {
+  Accordion,
+  AccordionContent,
+  AccordionItem,
+  AccordionTrigger
+} from '@/components/ui/accordion';
 import { Translator } from 'components/i18n';
 
 interface Props {
@@ -16,52 +21,71 @@ export default function Step({
   children,
   isRunning
 }: PropsWithChildren<Props>) {
-  const [open, setOpen] = useState(false);
   const using = useMemo(() => {
     return isRunning && step.start && !step.end && !step.isError;
   }, [step, isRunning]);
 
   const hasContent = step.input || step.output || step.steps?.length;
   const isError = step.isError;
-
   const stepName = step.name;
 
-  return (
-    <div className="flex flex-col flex-grow w-0">
-      <p
-        className={cn(
-          'flex items-center gap-1 group/step',
-          isError && 'text-red-500',
-          hasContent && 'cursor-pointer',
-          !using && 'text-muted-foreground hover:text-foreground',
-          using && 'loading-shimmer'
-        )}
-        onClick={() => setOpen(!open)}
-        id={`step-${stepName}`}
-      >
-        {using ? (
-          <>
-            <Translator path="chat.messages.status.using" /> {stepName}
-          </>
-        ) : (
-          <>
-            <Translator path="chat.messages.status.used" /> {stepName}
-          </>
-        )}
-        {hasContent ? (
-          open ? (
-            <ChevronUp className="invisible group-hover/step:visible !size-4" />
+  // If there's no content, just render the status without accordion
+  if (!hasContent) {
+    return (
+      <div className="flex flex-col flex-grow w-0">
+        <p
+          className={cn(
+            'flex items-center gap-1 font-medium',
+            isError && 'text-red-500',
+            !using && 'text-muted-foreground',
+            using && 'loading-shimmer'
+          )}
+          id={`step-${stepName}`}
+        >
+          {using ? (
+            <>
+              <Translator path="chat.messages.status.using" /> {stepName}
+            </>
           ) : (
-            <ChevronDown className="invisible group-hover/step:visible !size-4" />
-          )
-        ) : null}
-      </p>
+            <>
+              <Translator path="chat.messages.status.used" /> {stepName}
+            </>
+          )}
+        </p>
+      </div>
+    );
+  }
 
-      {open && (
-        <div className="flex-grow mt-4 ml-2 pl-4 border-l-2 border-primary">
-          {children}
-        </div>
-      )}
+  return (
+    <div className="flex flex-col flex-grow w-0">
+      <Accordion type="single" collapsible className="w-full">
+        <AccordionItem value={step.id} className="border-none">
+          <AccordionTrigger
+            className={cn(
+              'flex items-center gap-1 justify-start transition-none p-0 hover:no-underline',
+              isError && 'text-red-500',
+              !using && 'text-muted-foreground hover:text-foreground',
+              using && 'loading-shimmer'
+            )}
+            id={`step-${stepName}`}
+          >
+            {using ? (
+              <>
+                <Translator path="chat.messages.status.using" /> {stepName}
+              </>
+            ) : (
+              <>
+                <Translator path="chat.messages.status.used" /> {stepName}
+              </>
+            )}
+          </AccordionTrigger>
+          <AccordionContent>
+            <div className="flex-grow mt-4 ml-1 pl-4 border-l-2 border-primary">
+              {children}
+            </div>
+          </AccordionContent>
+        </AccordionItem>
+      </Accordion>
     </div>
   );
 }
diff --git a/frontend/src/components/chat/ScrollContainer.tsx b/frontend/src/components/chat/ScrollContainer.tsx
@@ -1,42 +1,35 @@
 import { cn } from '@/lib/utils';
-import { useEffect, useRef } from 'react';
+import { MutableRefObject, useEffect, useRef } from 'react';
 
-import { useChatMessages, useChatSession } from '@chainlit/react-client';
+import { useChatMessages } from '@chainlit/react-client';
 
 interface Props {
-  setAutoScroll?: (autoScroll: boolean) => void;
-  autoScroll?: boolean;
+  autoScrollRef?: MutableRefObject<boolean>;
   children: React.ReactNode;
   className?: string;
 }
 
 export default function ScrollContainer({
-  setAutoScroll,
-  autoScroll,
+  autoScrollRef,
   children,
   className
 }: Props) {
   const ref = useRef<HTMLDivElement>(null);
   const { messages } = useChatMessages();
-  const { session } = useChatSession();
 
   useEffect(() => {
-    setAutoScroll?.(true);
-  }, [session?.socket.id]);
-
-  useEffect(() => {
-    if (!ref.current || !autoScroll) {
+    if (!ref.current || !autoScrollRef?.current) {
       return;
     }
     ref.current.scrollTop = ref.current.scrollHeight;
-  }, [messages, autoScroll]);
+  }, [messages]);
 
   const handleScroll = () => {
-    if (!ref.current || !setAutoScroll) return;
+    if (!ref.current || !autoScrollRef) return;
 
     const { scrollTop, scrollHeight, clientHeight } = ref.current;
     const atBottom = scrollTop + clientHeight >= scrollHeight - 10;
-    setAutoScroll(atBottom);
+    autoScrollRef.current = atBottom;
   };
 
   return (

diff --git a/frontend/src/components/chat/WelcomeScreen.tsx b/frontend/src/components/chat/WelcomeScreen.tsx
@@ -1,5 +1,11 @@
 import { cn, hasMessage } from '@/lib/utils';
-import { useContext, useEffect, useMemo, useState } from 'react';
+import {
+  MutableRefObject,
+  useContext,
+  useEffect,
+  useMemo,
+  useState
+} from 'react';
 
 import {
   ChainlitContext,
@@ -19,7 +25,7 @@ interface Props {
   fileSpec: FileSpec;
   onFileUpload: (payload: File[]) => void;
   onFileUploadError: (error: string) => void;
-  setAutoScroll: (autoScroll: boolean) => void;
+  autoScrollRef: MutableRefObject<boolean>;
 }
 
 export default function WelcomeScreen(props: Props) {

diff --git a/frontend/src/components/chat/index.tsx b/frontend/src/components/chat/index.tsx
@@ -1,4 +1,4 @@
-import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
+import { useCallback, useEffect, useMemo, useRef } from 'react';
 import { useNavigate } from 'react-router-dom';
 import { useSetRecoilState } from 'recoil';
 import { toast } from 'sonner';
@@ -35,7 +35,7 @@ const Chat = () => {
   const setAttachments = useSetRecoilState(attachmentsState);
   const setThreads = useSetRecoilState(threadHistoryState);
 
-  const [autoScroll, setAutoScroll] = useState(true);
+  const autoScrollRef = useRef(true);
   const { error, disabled } = useChatData();
   const { uploadFile } = useChatInteract();
   const uploadFileRef = useRef(uploadFile);
@@ -200,7 +200,7 @@ const Chat = () => {
         </div>
       ) : null}
       <ErrorBoundary>
-        <ScrollContainer autoScroll={autoScroll} setAutoScroll={setAutoScroll}>
+        <ScrollContainer autoScrollRef={autoScrollRef}>
           <div
             className="flex flex-col mx-auto w-full flex-grow p-4"
             style={{
@@ -212,7 +212,7 @@ const Chat = () => {
               fileSpec={fileSpec}
               onFileUpload={onFileUpload}
               onFileUploadError={onFileUploadError}
-              setAutoScroll={setAutoScroll}
+              autoScrollRef={autoScrollRef}
             />
             <MessagesContainer navigate={navigate} />
           </div>
@@ -227,8 +227,7 @@ const Chat = () => {
             fileSpec={fileSpec}
             onFileUpload={onFileUpload}
             onFileUploadError={onFileUploadError}
-            setAutoScroll={setAutoScroll}
-            autoScroll={autoScroll}
+            autoScrollRef={autoScrollRef}
           />
         </div>
       </ErrorBoundary>