diff --git a/src/chainlit/frontend/src/components/organisms/playground/chat.tsx b/src/chainlit/frontend/src/components/organisms/playground/chat.tsx
index ecbeb335cc..9289b3fa70 100644
--- a/src/chainlit/frontend/src/components/organisms/playground/chat.tsx
+++ b/src/chainlit/frontend/src/components/organisms/playground/chat.tsx
@@ -57,11 +57,7 @@ export const ChatPromptPlayground = forwardRef(
         : 'Prompt messages';
 
     return (
-      <Stack
-        sx={{
-          width: '100%'
-        }}
-      >
+      <Stack sx={{ width: '100%' }}>
         <Typography fontSize="14px" fontWeight={700} color={grey[400]}>
           {title}
         </Typography>
@@ -87,14 +83,6 @@ export const ChatPromptPlayground = forwardRef(
                     index={index}
                     onChange={onChange}
                   />
-                  {index !== messages.length - 1 ? (
-                    <Box
-                      sx={{
-                        border: (theme) => `1px solid ${theme.palette.divider}`,
-                        borderRadius: 1
-                      }}
-                    />
-                  ) : null}
                 </Fragment>
               ))}
             </Stack>
diff --git a/src/chainlit/frontend/src/components/organisms/playground/editor/MessageWrapper.tsx b/src/chainlit/frontend/src/components/organisms/playground/editor/MessageWrapper.tsx
new file mode 100644
index 0000000000..df1db8db08
--- /dev/null
+++ b/src/chainlit/frontend/src/components/organisms/playground/editor/MessageWrapper.tsx
@@ -0,0 +1,125 @@
+import React from 'react';
+import { useSetRecoilState } from 'recoil';
+import { useToggle } from 'usehooks-ts';
+
+import {
+  Box,
+  SelectChangeEvent,
+  Stack,
+  Theme,
+  Typography
+} from '@mui/material';
+
+import SelectInput from 'components/organisms/inputs/selectInput';
+
+import { IPromptMessage, PromptMessageRole } from 'state/chat';
+import { playgroundState } from 'state/playground';
+
+const roles = ['Assistant', 'System', 'User'];
+
+interface MessageWrapperProps {
+  canSelectRole?: boolean;
+  children: React.ReactElement;
+  index?: number;
+  message?: IPromptMessage;
+  role?: string;
+}
+
+const MessageWrapper = ({
+  canSelectRole,
+  children,
+  index,
+  message,
+  role
+}: MessageWrapperProps): JSX.Element => {
+  const setPlayground = useSetRecoilState(playgroundState);
+  const [isSelectRoleOpen, toggleSelectRole] = useToggle(false);
+
+  const onRoleSelected = (event: SelectChangeEvent) => {
+    const role = event.target.value as PromptMessageRole;
+
+    if (role) {
+      setPlayground((old) => ({
+        ...old,
+        prompt: {
+          ...old.prompt!,
+          messages: old.prompt?.messages?.map((message, mIndex) => ({
+            ...message,
+            ...(mIndex === index ? { role } : {}) // Update role if it's the selected message
+          }))
+        }
+      }));
+    }
+
+    toggleSelectRole();
+  };
+
+  return (
+    <Box key={index}>
+      <Box
+        sx={{
+          border: (theme) => `1px solid ${theme.palette.divider}`,
+          borderRadius: 1
+        }}
+      />
+      <Stack
+        direction="row"
+        width="100%"
+        sx={{
+          gap: 2,
+          flex: 10,
+          paddingY: 1,
+          '&:hover': {
+            background: (theme: Theme) => theme.palette.background.paper
+          }
+        }}
+      >
+        <Box sx={{ flex: 1, paddingLeft: 2 }}>
+          {!isSelectRoleOpen ? (
+            <Typography
+              onClick={() => canSelectRole && toggleSelectRole()}
+              color="text.primary"
+              sx={{
+                p: 1,
+                borderRadius: 0.5,
+                marginTop: 1,
+                cursor: canSelectRole ? 'pointer' : 'default',
+                fontSize: '12px',
+                fontWeight: 700,
+                width: 'fit-content',
+                ...(canSelectRole && {
+                  '&:hover': {
+                    backgroundColor: (theme) => theme.palette.divider
+                  }
+                })
+              }}
+            >
+              {role}
+            </Typography>
+          ) : (
+            <SelectInput
+              defaultOpen
+              items={roles.map((role) => ({
+                label: role,
+                value: role.toLowerCase()
+              }))}
+              id="role-select"
+              value={message?.role}
+              onChange={onRoleSelected}
+              sx={{ width: 'fit-content' }}
+              iconSx={{
+                px: 0,
+                marginRight: '2px !important'
+              }}
+            />
+          )}
+        </Box>
+        <Box width="100%" flex={8}>
+          {children}
+        </Box>
+      </Stack>
+    </Box>
+  );
+};
+
+export default MessageWrapper;
diff --git a/src/chainlit/frontend/src/components/organisms/playground/editor/completion.tsx b/src/chainlit/frontend/src/components/organisms/playground/editor/completion.tsx
index 56f32e43ff..2e32225bc0 100644
--- a/src/chainlit/frontend/src/components/organisms/playground/editor/completion.tsx
+++ b/src/chainlit/frontend/src/components/organisms/playground/editor/completion.tsx
@@ -11,6 +11,8 @@ import EditorWrapper from 'components/organisms/playground/editor/wrapper';
 
 import 'draft-js/dist/Draft.css';
 
+import MessageWrapper from './MessageWrapper';
+
 const styleMap = {
   COMPLETION: {
     backgroundColor: '#d2f4d3',
@@ -118,39 +120,6 @@ export default function Completion({ completion, chatMode }: Props) {
       </Box>
     </Box>
   ) : (
-    <Box>
-      <Box
-        sx={{
-          border: (theme) => `1px solid ${theme.palette.divider}`,
-          borderRadius: 1
-        }}
-      />
-      <Stack
-        direction="row"
-        width="100%"
-        sx={{
-          paddingY: 1,
-          paddingRight: 0,
-          '&:hover': {
-            background: (theme) => theme.palette.background.paper
-          }
-        }}
-      >
-        <Typography
-          color="text.primary"
-          sx={{
-            marginTop: 2,
-            cursor: 'pointer',
-            fontSize: '12px',
-            fontWeight: 700,
-            paddingLeft: 3,
-            paddingRight: 5
-          }}
-        >
-          {'ASSISTANT'}
-        </Typography>
-        <Box width="100%">{renderEditor()}</Box>
-      </Stack>
-    </Box>
+    <MessageWrapper role="ASSISTANT">{renderEditor()}</MessageWrapper>
   );
 }
diff --git a/src/chainlit/frontend/src/components/organisms/playground/editor/promptMessage.tsx b/src/chainlit/frontend/src/components/organisms/playground/editor/promptMessage.tsx
index d5e0e014b3..e86d63b0f9 100644
--- a/src/chainlit/frontend/src/components/organisms/playground/editor/promptMessage.tsx
+++ b/src/chainlit/frontend/src/components/organisms/playground/editor/promptMessage.tsx
@@ -1,20 +1,14 @@
 import { EditorState } from 'draft-js';
-import { useSetRecoilState } from 'recoil';
-import { useToggle } from 'usehooks-ts';
 
-import { Alert, Box, Stack, Typography, useTheme } from '@mui/material';
-import { SelectChangeEvent } from '@mui/material/Select';
+import { Alert, useTheme } from '@mui/material';
 
-import SelectInput from 'components/organisms/inputs/selectInput';
-
-import { IPrompt, IPromptMessage, PromptMessageRole } from 'state/chat';
-import { PromptMode, playgroundState } from 'state/playground';
+import { IPrompt, IPromptMessage } from 'state/chat';
+import { PromptMode } from 'state/playground';
 
+import MessageWrapper from './MessageWrapper';
 import FormattedEditor from './formatted';
 import TemplateEditor from './template';
 
-const roles = ['Assistant', 'System', 'User'];
-
 interface Props {
   message: IPromptMessage;
   prompt: IPrompt;
@@ -30,8 +24,6 @@ export default function PromptMessage({
   index,
   onChange
 }: Props) {
-  const setPlayground = useSetRecoilState(playgroundState);
-  const [isSelectRoleOpen, toggleSelectRole] = useToggle();
   const theme = useTheme();
 
   const templateProps = {
@@ -84,87 +76,17 @@ export default function PromptMessage({
     );
   };
 
-  const onRoleSelected = (event: SelectChangeEvent) => {
-    const role = event.target.value as PromptMessageRole;
-
-    if (role) {
-      setPlayground((old) => ({
-        ...old,
-        prompt: {
-          ...old.prompt!,
-          messages: old.prompt?.messages?.map((message, mIndex) => ({
-            ...message,
-            ...(mIndex === index ? { role } : {}) // Update role if it's the selected message
-          }))
-        }
-      }));
-    }
-
-    toggleSelectRole();
-  };
-
   return (
-    <Stack
-      className="prompt-message"
-      key={index}
-      direction="row"
-      sx={{
-        padding: (theme) => theme.spacing(1, 2),
-        paddingRight: 0,
-        '&:hover': {
-          background: (theme) => theme.palette.background.paper
-        }
-      }}
+    <MessageWrapper
+      canSelectRole
+      index={index}
+      message={message}
+      role={message?.role?.toUpperCase()}
     >
-      <Stack
-        sx={{
-          fontSize: '12px',
-          fontWeight: 700,
-          paddingRight: 2,
-          maxWidth: '100px',
-          width: '100%'
-        }}
-      >
-        {isSelectRoleOpen ? (
-          <SelectInput
-            defaultOpen
-            items={roles.map((role) => ({
-              label: role,
-              value: role.toLowerCase()
-            }))}
-            id="role-select"
-            value={message.role}
-            onChange={onRoleSelected}
-            iconSx={{
-              px: 0,
-              marginRight: '2px !important'
-            }}
-          />
-        ) : (
-          <Typography
-            onClick={toggleSelectRole}
-            color="text.primary"
-            sx={{
-              marginTop: 2,
-              cursor: 'pointer',
-              fontSize: '12px',
-              fontWeight: 700,
-              width: 'fit-content',
-              padding: (theme) => theme.spacing(0.5, 1),
-              '&:hover': {
-                backgroundColor: (theme) => theme.palette.divider
-              }
-            }}
-          >
-            {message?.role?.toUpperCase()}
-          </Typography>
-        )}
-      </Stack>
-
-      <Box sx={{ width: '100%' }}>
+      <>
         {mode === 'Template' ? renderTemplate() : null}
         {mode === 'Formatted' ? renderFormatted() : null}
-      </Box>
-    </Stack>
+      </>
+    </MessageWrapper>
   );
 }