From d0ab9c7ec4d85be054b8504fabce819c750a0119 Mon Sep 17 00:00:00 2001
From: josc146 <josStorer@outlook.com>
Date: Mon, 25 Mar 2024 16:08:29 +0800
Subject: [PATCH] add system role support for preset

---
 frontend/src/components/MarkdownRender.tsx    |  4 +-
 frontend/src/pages/Chat.tsx                   | 45 ++++++++++++++-----
 .../pages/PresetsManager/MessagesEditor.tsx   |  2 +-
 frontend/src/types/chat.ts                    |  3 +-
 frontend/src/utils/index.tsx                  |  8 ++--
 5 files changed, 44 insertions(+), 18 deletions(-)

diff --git a/frontend/src/components/MarkdownRender.tsx b/frontend/src/components/MarkdownRender.tsx
index 27977a7..6a81136 100644
--- a/frontend/src/components/MarkdownRender.tsx
+++ b/frontend/src/components/MarkdownRender.tsx
@@ -28,10 +28,10 @@ const MarkdownRender: FC<ReactMarkdownOptions & { disabled?: boolean }> = (props
   return (
     <div dir="auto" className="prose markdown-body" style={{ maxWidth: '100%' }}>
       {props.disabled ?
-        <div style={{ whiteSpace: 'pre-wrap' }}>
+        <div style={{ whiteSpace: 'pre-wrap' }} className={props.className}>
           {props.children}
         </div> :
-        <ReactMarkdown
+        <ReactMarkdown className={props.className}
           allowedElements={[
             'div',
             'p',
diff --git a/frontend/src/pages/Chat.tsx b/frontend/src/pages/Chat.tsx
index f6434dd..a50f0bf 100644
--- a/frontend/src/pages/Chat.tsx
+++ b/frontend/src/pages/Chat.tsx
@@ -48,12 +48,13 @@ import {
   toastWithButton
 } from '../utils';
 import { useMediaQuery } from 'usehooks-ts';
-import { botName, ConversationMessage, MessageType, Role, userName, welcomeUuid } from '../types/chat';
+import { botName, ConversationMessage, MessageType, Role, systemName, userName, welcomeUuid } from '../types/chat';
 import { Labeled } from '../components/Labeled';
 import { ValuedSlider } from '../components/ValuedSlider';
 import { PresetsButton } from './PresetsManager/PresetsButton';
 import { webOpenOpenFileDialog } from '../utils/web-file-operations';
 import { defaultPenaltyDecay } from './defaultConfigs';
+import { AvatarProps } from '@fluentui/react-avatar';
 
 let chatSseControllers: {
   [id: string]: AbortController
@@ -92,6 +93,18 @@ const MoreUtilsButton: FC<{
   </Menu>;
 });
 
+const HiddenAvatar: FC<AvatarProps> = ({ children, hidden, ...props }) => {
+  if (hidden) {
+    return <div>{children}</div>;
+  } else {
+    return (
+      <Avatar {...props}>
+        {children}
+      </Avatar>
+    );
+  }
+};
+
 const ChatMessageItem: FC<{
   uuid: string,
   onSubmit: (message: string | null, answerId: string | null,
@@ -129,7 +142,9 @@ const ChatMessageItem: FC<{
   return <div
     className={classnames(
       'flex gap-2 mb-2 overflow-hidden',
-      messageItem.side === 'left' ? 'flex-row' : 'flex-row-reverse'
+      messageItem.side === 'left' && 'flex-row',
+      messageItem.side === 'right' && 'flex-row-reverse',
+      messageItem.side === 'center' && 'flex-row justify-center'
     )}
     onMouseEnter={() => {
       const utils = document.getElementById('utils-' + uuid);
@@ -140,22 +155,26 @@ const ChatMessageItem: FC<{
       if (utils) utils.classList.add('invisible');
     }}
   >
-    <Avatar
+    <HiddenAvatar
       color={messageItem.color}
       name={messageItem.sender}
       image={avatarImg ? { src: avatarImg } : undefined}
+      hidden={messageItem.side === 'center'}
     />
     <div
       className={classnames(
-        'flex p-2 rounded-lg overflow-hidden',
+        'flex rounded-lg overflow-hidden',
         editing ? 'grow' : '',
-        commonStore.settings.darkMode ? 'bg-neutral-800 border-neutral-600 border-[1px]' : (messageItem.side === 'left' ? 'bg-gray-200' : 'bg-blue-500'),
-        commonStore.settings.darkMode ? 'text-white' : (messageItem.side === 'left' ? 'text-gray-600' : 'text-white')
+        messageItem.side === 'center' ? 'p-1' : 'p-2',
+        commonStore.settings.darkMode ? 'bg-neutral-800 border-neutral-600 border-[1px]' : (messageItem.side === 'right' ? 'bg-blue-500' : 'bg-gray-200'),
+        commonStore.settings.darkMode ? 'text-white' : (messageItem.side === 'right' ? 'text-white' : 'text-gray-600'),
+        messageItem.side === 'center' && 'text-opacity-60 bg-opacity-30 border-opacity-30'
       )}
     >
       {!editing ?
         <div className="flex flex-col">
-          <MarkdownRender disabled={!commonStore.chatParams.markdown}>{messageItem.content}</MarkdownRender>
+          <MarkdownRender className={classnames(messageItem.side === 'center' && 'text-xs')}
+            disabled={!commonStore.chatParams.markdown || messageItem.side === 'center'}>{messageItem.content}</MarkdownRender>
           {uuid in commonStore.attachments &&
             <div className="flex grow">
               <div className="grow" />
@@ -375,7 +394,11 @@ const SidePanel: FC = observer(() => {
             return;
           const messageItem = commonStore.conversation[uuid];
           if (messageItem.type !== MessageType.Error) {
-            savedContent += `${messageItem.sender === userName ? user : bot}: ${messageItem.content}\n\n`;
+            if (messageItem.sender === userName) {
+              savedContent += `${user}: ${messageItem.content}\n\n`;
+            } else if (messageItem.sender === botName) {
+              savedContent += `${bot}: ${messageItem.content}\n\n`;
+            }
           }
         });
 
@@ -418,7 +441,7 @@ const ChatPanel: FC = observer(() => {
         [welcomeUuid]: {
           sender: botName,
           type: MessageType.Normal,
-          color: 'colorful',
+          color: 'neutral',
           avatarImg: logo,
           time: new Date().toISOString(),
           content: commonStore.platform === 'web' ? t('Hello, what can I do for you?') : t('Hello! I\'m RWKV, an open-source and commercially usable large language model.'),
@@ -499,6 +522,8 @@ const ChatPanel: FC = observer(() => {
         messages.push({ role: 'user', content: messageItem.content });
       } else if (messageItem.done && messageItem.type === MessageType.Normal && messageItem.sender === botName) {
         messages.push({ role: 'assistant', content: messageItem.content });
+      } else if (messageItem.done && messageItem.type === MessageType.Normal && messageItem.sender === systemName) {
+        messages.push({ role: 'system', content: messageItem.content });
       }
     });
 
@@ -509,7 +534,7 @@ const ChatPanel: FC = observer(() => {
     commonStore.conversation[answerId] = {
       sender: botName,
       type: MessageType.Normal,
-      color: 'colorful',
+      color: 'neutral',
       avatarImg: logo,
       time: new Date().toISOString(),
       content: '',
diff --git a/frontend/src/pages/PresetsManager/MessagesEditor.tsx b/frontend/src/pages/PresetsManager/MessagesEditor.tsx
index 74b886d..aa52c5c 100644
--- a/frontend/src/pages/PresetsManager/MessagesEditor.tsx
+++ b/frontend/src/pages/PresetsManager/MessagesEditor.tsx
@@ -126,7 +126,7 @@ const MessagesEditor: FC = observer(() => {
                             }}>
                             <Option value="user">{t('user')!}</Option>
                             <Option value="assistant">{t('assistant')!}</Option>
-                            {/* TODO <Option value="system">{t('system')!}</Option>*/}
+                            <Option value="system">{t('system')!}</Option>
                           </Dropdown>
                           <Textarea resize="vertical" className="grow" value={item.content}
                             style={{ minWidth: 0, borderRadius: 0 }}
diff --git a/frontend/src/types/chat.ts b/frontend/src/types/chat.ts
index 74e211b..c93eec7 100644
--- a/frontend/src/types/chat.ts
+++ b/frontend/src/types/chat.ts
@@ -2,6 +2,7 @@ import { ApiParameters } from './configs';
 
 export const userName = 'M E';
 export const botName = 'A I';
+export const systemName = 'System';
 export const welcomeUuid = 'welcome';
 
 export enum MessageType {
@@ -9,7 +10,7 @@ export enum MessageType {
   Error
 }
 
-export type Side = 'left' | 'right'
+export type Side = 'left' | 'right' | 'center'
 export type Color = 'neutral' | 'brand' | 'colorful'
 export type MessageItem = {
   sender: string,
diff --git a/frontend/src/utils/index.tsx b/frontend/src/utils/index.tsx
index 5e70b06..bf27546 100644
--- a/frontend/src/utils/index.tsx
+++ b/frontend/src/utils/index.tsx
@@ -26,7 +26,7 @@ import { DataProcessParameters, LoraFinetuneParameters } from '../types/train';
 import { InstrumentTypeNameMap, MidiMessage, tracksMinimalTotalTime } from '../types/composition';
 import logo from '../assets/images/logo.png';
 import { Preset } from '../types/presets';
-import { botName, Conversation, MessageType, Role, userName } from '../types/chat';
+import { botName, Conversation, MessageType, Role, systemName, userName } from '../types/chat';
 import { v4 as uuid } from 'uuid';
 import { findLastIndex } from 'lodash-es';
 
@@ -635,13 +635,13 @@ export function newChatConversation() {
     const newUuid = uuid();
     conversationOrder.push(newUuid);
     conversation[newUuid] = {
-      sender: role === 'user' ? userName : botName,
+      sender: role === 'user' ? userName : role === 'assistant' ? botName : systemName,
       type: MessageType.Normal,
-      color: role === 'user' ? 'brand' : 'colorful',
+      color: role === 'user' ? 'brand' : 'neutral',
       avatarImg: role === 'user' ? undefined : logo,
       time: new Date().toISOString(),
       content: content,
-      side: role === 'user' ? 'right' : 'left',
+      side: role === 'user' ? 'right' : role === 'assistant' ? 'left' : 'center',
       done: true
     };
   };