From aedfa532351e6ae44b800898301ae9d7e780532e Mon Sep 17 00:00:00 2001
From: arvinxx <arvinx@foxmail.com>
Date: Sun, 23 Jul 2023 22:12:10 +0800
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat:=20=E5=85=BC=E5=AE=B9=E6=8F=92?=
 =?UTF-8?q?=E4=BB=B6=E6=A8=A1=E5=BC=8F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/Chat/types.ts      |  2 +-
 src/ChatItem/index.tsx | 26 +++++++++++++++++---------
 src/ChatList/index.tsx | 19 +++++++++++++++++--
 src/types/llm.ts       |  2 +-
 4 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/src/Chat/types.ts b/src/Chat/types.ts
index 67bfb4a4..ef0abb52 100644
--- a/src/Chat/types.ts
+++ b/src/Chat/types.ts
@@ -1,6 +1,6 @@
 export type InternalChatContext = Omit<ChatContext, 'id'>;
 
-export type MessageRoleType = 'user' | 'system' | 'assistant';
+export type MessageRoleType = 'user' | 'system' | 'assistant' | 'function';
 
 /**
  * 聊天消息错误对象
diff --git a/src/ChatItem/index.tsx b/src/ChatItem/index.tsx
index d42564d1..150e94a8 100644
--- a/src/ChatItem/index.tsx
+++ b/src/ChatItem/index.tsx
@@ -45,13 +45,14 @@ export interface ChatItemProps {
   /**
    * @description The message content of the chat item
    */
-  message?: string;
+  message?: ReactNode;
   messageExtra?: ReactNode;
   /**
    * @description Callback when the message content changes
    * @param value - The new message content
    */
   onChange?: (value: string) => void;
+
   /**
    * @description Callback when the editing mode changes
    * @param editing - The new editing mode
@@ -66,6 +67,7 @@ export interface ChatItemProps {
    * @description Whether the chat item is primary
    */
   primary?: boolean;
+  renderMessage?: (content: ReactNode) => ReactNode;
   /**
    * @description Whether to show the title of the chat item
    */
@@ -100,6 +102,7 @@ const ChatItem = memo<ChatItemProps>(
     onChange,
     onEditingChange,
     messageExtra,
+    renderMessage,
     ...properties
   }) => {
     const { cx, styles } = useStyles({
@@ -111,6 +114,18 @@ const ChatItem = memo<ChatItemProps>(
       type,
     });
 
+    const content = (
+      <EditableMessage
+        classNames={{ input: styles.editingInput }}
+        editButtonSize={'small'}
+        editing={editing}
+        onChange={onChange}
+        onEditingChange={onEditingChange}
+        value={String(message || '...')}
+      />
+    );
+    const messageContent = renderMessage ? renderMessage(content) : content;
+
     return (
       <div className={cx(styles.container, className)} {...properties}>
         <div className={styles.avatarContainer}>
@@ -136,14 +151,7 @@ const ChatItem = memo<ChatItemProps>(
               <Alert className={styles.alert} showIcon {...alert} />
             ) : (
               <Flexbox className={cx(styles.message, editing && styles.editingContainer)}>
-                <EditableMessage
-                  classNames={{ input: styles.editingInput }}
-                  editButtonSize={'small'}
-                  editing={editing}
-                  onChange={onChange}
-                  onEditingChange={onEditingChange}
-                  value={String(message || '...')}
-                />
+                {messageContent}
                 {messageExtra && !editing ? (
                   <div className={styles.messageExtra}>{messageExtra}</div>
                 ) : null}
diff --git a/src/ChatList/index.tsx b/src/ChatList/index.tsx
index 7ee6b3db..946e32f2 100644
--- a/src/ChatList/index.tsx
+++ b/src/ChatList/index.tsx
@@ -1,6 +1,6 @@
 import { App } from 'antd';
 import copy from 'copy-to-clipboard';
-import { ReactNode, memo, useState } from 'react';
+import { ReactNode, memo, useCallback, useState } from 'react';
 
 import ChatItem, { ChatItemProps } from '@/ChatItem';
 import type { DivProps } from '@/types';
@@ -9,7 +9,10 @@ import { ChatMessage } from '@/types/chatMessage';
 import ActionsBar from './ActionsBar';
 import { useStyles } from './style';
 
-type MessageExtra = (props: ChatMessage) => ReactNode;
+export type MessageExtra = (props: ChatMessage) => ReactNode;
+
+export type RenderMessage = (content: ReactNode, message: ChatMessage) => ReactNode;
+
 export interface ChatListProps extends DivProps {
   /**
    * @description Data of chat messages to be displayed
@@ -22,6 +25,7 @@ export interface ChatListProps extends DivProps {
    */
   onActionClick?: (actionKey: string, messageId: string) => void;
   onMessageChange?: (id: string, content: string) => void;
+  renderMessage?: RenderMessage;
   renderMessageExtra?: MessageExtra;
   /**
    * @description Whether to show name of the chat item
@@ -39,6 +43,7 @@ export interface ItemProps extends ChatMessage {
   MessageExtra?: MessageExtra;
   onActionClick?: (actionKey: string, messageId: string) => void;
   onMessageChange?: (id: string, content: string) => void;
+  renderMessage?: RenderMessage;
   showTitle?: boolean;
   type: 'docs' | 'chat';
 }
@@ -49,6 +54,7 @@ const Item = ({
   onActionClick,
   onMessageChange,
   type,
+  renderMessage,
   ...item
 }: ItemProps) => {
   const renderMessageExtra = MessageExtra ? <MessageExtra {...item} /> : undefined;
@@ -56,6 +62,12 @@ const Item = ({
   const [editing, setEditing] = useState(false);
 
   const { message } = App.useApp();
+
+  const innerRenderMessage = useCallback(
+    (content: ReactNode) => renderMessage?.(content, item),
+    [renderMessage],
+  );
+
   return (
     <ChatItem
       actions={
@@ -87,6 +99,7 @@ const Item = ({
       onEditingChange={setEditing}
       placement={type === 'chat' ? (item.role === 'user' ? 'right' : 'left') : 'left'}
       primary={item.role === 'user'}
+      renderMessage={renderMessage ? innerRenderMessage : undefined}
       showTitle={showTitle}
       time={item.updateAt || item.createAt}
       type={type === 'chat' ? 'block' : 'pure'}
@@ -103,6 +116,7 @@ const ChatList = memo<ChatListProps>(
     type = 'chat',
     showTitle,
     onMessageChange,
+    renderMessage,
     ...props
   }) => {
     const { cx, styles } = useStyles();
@@ -115,6 +129,7 @@ const ChatList = memo<ChatListProps>(
             key={item.id}
             onActionClick={onActionClick}
             onMessageChange={onMessageChange}
+            renderMessage={renderMessage}
             showTitle={showTitle}
             type={type}
             {...item}
diff --git a/src/types/llm.ts b/src/types/llm.ts
index 7d7bcc52..09b44093 100644
--- a/src/types/llm.ts
+++ b/src/types/llm.ts
@@ -37,7 +37,7 @@ export interface LMParameters {
   top_p?: number;
 }
 
-export type LLMRoleType = 'user' | 'system' | 'assistant';
+export type LLMRoleType = 'user' | 'system' | 'assistant' | 'function';
 
 export interface LLMMessage {
   content: string;