From a0e60b3badcf21ed349ce60fe9b522828e89ea86 Mon Sep 17 00:00:00 2001
From: Alim TUNC <alim@refty.co>
Date: Mon, 28 Aug 2023 15:17:29 +0200
Subject: [PATCH] refactor: Scroll to bottom of chat prompt on submit

---
 .../components/organisms/playground/chat.tsx  | 161 +++++++++---------
 .../components/organisms/playground/index.tsx |  13 +-
 .../organisms/playground/submitButton.tsx     |   4 +-
 3 files changed, 95 insertions(+), 83 deletions(-)

diff --git a/src/chainlit/frontend/src/components/organisms/playground/chat.tsx b/src/chainlit/frontend/src/components/organisms/playground/chat.tsx
index 5cb28031e9..ecbeb335cc 100644
--- a/src/chainlit/frontend/src/components/organisms/playground/chat.tsx
+++ b/src/chainlit/frontend/src/components/organisms/playground/chat.tsx
@@ -1,5 +1,5 @@
 import { EditorState } from 'draft-js';
-import { Fragment } from 'react';
+import { Fragment, forwardRef } from 'react';
 import { useRecoilValue, useSetRecoilState } from 'recoil';
 
 import { Box, Stack, Typography } from '@mui/material';
@@ -17,92 +17,93 @@ interface Props {
   restoredTime: number;
 }
 
-export default function ChatPromptPlayground({
-  hasTemplate,
-  prompt,
-  restoredTime
-}: Props) {
-  const setPlayground = useSetRecoilState(playgroundState);
-  const mode = useRecoilValue(modeState);
+export const ChatPromptPlayground = forwardRef(
+  ({ hasTemplate, prompt, restoredTime }: Props, ref) => {
+    const setPlayground = useSetRecoilState(playgroundState);
+    const mode = useRecoilValue(modeState);
 
-  const messages = prompt.messages;
+    const messages = prompt.messages;
 
-  if (!messages) {
-    return null;
-  }
+    if (!messages) {
+      return null;
+    }
 
-  const onChange = (index: number, nextState: EditorState) => {
-    const text = nextState.getCurrentContent().getPlainText();
-    const key = hasTemplate ? 'template' : 'formatted';
+    const onChange = (index: number, nextState: EditorState) => {
+      const text = nextState.getCurrentContent().getPlainText();
+      const key = hasTemplate ? 'template' : 'formatted';
 
-    setPlayground((old) => ({
-      ...old,
-      prompt: {
-        ...old.prompt!,
-        messages: old.prompt?.messages?.map((message, mIndex) => {
-          if (mIndex === index) {
-            return {
-              ...message,
-              [key]: text
-            };
-          }
-          return message;
-        })
-      }
-    }));
-  };
+      setPlayground((old) => ({
+        ...old,
+        prompt: {
+          ...old.prompt!,
+          messages: old.prompt?.messages?.map((message, mIndex) => {
+            if (mIndex === index) {
+              return {
+                ...message,
+                [key]: text
+              };
+            }
+            return message;
+          })
+        }
+      }));
+    };
 
-  const title =
-    mode === 'Formatted'
-      ? hasTemplate
-        ? 'Formatted messages [Read Only]'
-        : 'Formatted messages'
-      : 'Prompt messages';
+    const title =
+      mode === 'Formatted'
+        ? hasTemplate
+          ? 'Formatted messages [Read Only]'
+          : 'Formatted messages'
+        : 'Prompt messages';
 
-  return (
-    <Stack
-      sx={{
-        width: '100%'
-      }}
-    >
-      <Typography fontSize="14px" fontWeight={700} color={grey[400]}>
-        {title}
-      </Typography>
-      <Box
-        key={restoredTime} // This will re-mount the component with restored messages
+    return (
+      <Stack
         sx={{
-          flex: 1,
-          height: 'auto',
-          overflow: 'scroll',
-          marginTop: 1,
-          gap: 1
+          width: '100%'
         }}
       >
-        {messages.length > 0 ? (
-          <Stack>
-            {messages.map((message, index) => (
-              <Fragment key={`prompt-message-${index}`}>
-                <PromptMessage
-                  message={message}
-                  prompt={prompt}
-                  mode={mode}
-                  index={index}
-                  onChange={onChange}
-                />
-                {index !== messages.length - 1 ? (
-                  <Box
-                    sx={{
-                      border: (theme) => `1px solid ${theme.palette.divider}`,
-                      borderRadius: 1
-                    }}
+        <Typography fontSize="14px" fontWeight={700} color={grey[400]}>
+          {title}
+        </Typography>
+        <Box
+          key={restoredTime} // This will re-mount the component with restored messages
+          ref={ref}
+          sx={{
+            flex: 1,
+            height: 'auto',
+            overflow: 'scroll',
+            marginTop: 1,
+            gap: 1
+          }}
+        >
+          {messages.length > 0 ? (
+            <Stack>
+              {messages.map((message, index) => (
+                <Fragment key={`prompt-message-${index}`}>
+                  <PromptMessage
+                    message={message}
+                    prompt={prompt}
+                    mode={mode}
+                    index={index}
+                    onChange={onChange}
                   />
-                ) : null}
-              </Fragment>
-            ))}
-          </Stack>
-        ) : null}
-        <Completion completion={prompt.completion} chatMode />
-      </Box>
-    </Stack>
-  );
-}
+                  {index !== messages.length - 1 ? (
+                    <Box
+                      sx={{
+                        border: (theme) => `1px solid ${theme.palette.divider}`,
+                        borderRadius: 1
+                      }}
+                    />
+                  ) : null}
+                </Fragment>
+              ))}
+            </Stack>
+          ) : null}
+          <Completion completion={prompt.completion} chatMode />
+        </Box>
+      </Stack>
+    );
+  }
+);
+
+export default ChatPromptPlayground;
diff --git a/src/chainlit/frontend/src/components/organisms/playground/index.tsx b/src/chainlit/frontend/src/components/organisms/playground/index.tsx
index 7be9b4e666..72e0ecf146 100644
--- a/src/chainlit/frontend/src/components/organisms/playground/index.tsx
+++ b/src/chainlit/frontend/src/components/organisms/playground/index.tsx
@@ -1,4 +1,4 @@
-import { useEffect, useState } from 'react';
+import { useEffect, useRef, useState } from 'react';
 import { useRecoilState } from 'recoil';
 import { useToggle } from 'usehooks-ts';
 
@@ -37,6 +37,7 @@ export default function PromptPlayground() {
 
   const [restoredTime, setRestoredTime] = useState(0);
   const [isDrawerOpen, toggleDrawer] = useToggle(false);
+  const chatPromptScrollRef = useRef<HTMLDivElement | null>(null);
 
   const theme = useTheme();
   const isSmallScreen = useMediaQuery<Theme>((theme) =>
@@ -117,6 +118,7 @@ export default function PromptPlayground() {
                 prompt={playground.prompt}
               />
               <ChatPromptPlayground
+                ref={chatPromptScrollRef}
                 restoredTime={restoredTime}
                 hasTemplate={hasTemplate}
                 prompt={playground.prompt}
@@ -137,7 +139,14 @@ export default function PromptPlayground() {
             <RestoreIcon />
           </IconButton>
         </Tooltip>
-        <SubmitButton />
+        <SubmitButton
+          onSubmit={() => {
+            if (!chatPromptScrollRef?.current) return;
+
+            chatPromptScrollRef.current.scrollTop =
+              chatPromptScrollRef.current.scrollHeight;
+          }}
+        />
       </ActionBar>
     </Dialog>
   );
diff --git a/src/chainlit/frontend/src/components/organisms/playground/submitButton.tsx b/src/chainlit/frontend/src/components/organisms/playground/submitButton.tsx
index db59945a80..a8fb473b2d 100644
--- a/src/chainlit/frontend/src/components/organisms/playground/submitButton.tsx
+++ b/src/chainlit/frontend/src/components/organisms/playground/submitButton.tsx
@@ -11,7 +11,7 @@ import { accessTokenState, userEnvState } from 'state/user';
 
 import { getProviders } from './helpers';
 
-export default function SubmitButton() {
+export default function SubmitButton({ onSubmit }: { onSubmit: () => void }) {
   const [completionController, setCompletionController] = useState<
     AbortController | undefined
   >();
@@ -46,6 +46,8 @@ export default function SubmitButton() {
         controller,
         accessToken,
         (done, token) => {
+          onSubmit && onSubmit();
+
           if (done) {
             setCompletionController(undefined);
             return;