✨ feat: support useChat and chatRef

arvinxx · arvinxx · commit cf98b0f4cfea · 2023-11-22T21:25:12.000+08:00
diff --git a/.dumirc.ts b/.dumirc.ts
@@ -5,7 +5,11 @@ export default defineConfig({
   themeConfig: {
     name: '@ant-design/pro-chat',
     github: homepage,
+    siteToken: {
+      demoInheritSiteTheme: true,
+    },
   },
+  mfsu: false,
   outputPath: 'docs-dist',
   html2sketch: {},
   extraBabelPlugins: ['antd-style'],
diff --git a/src/ProChat/container/Provider.tsx b/src/ProChat/container/Provider.tsx
@@ -1,17 +1,45 @@
+import StoreUpdater, { ProChatChatReference } from '@/ProChat/container/StoreUpdater';
 import { memo, ReactNode } from 'react';
 import { DevtoolsOptions } from 'zustand/middleware';
 import { ChatProps, createStore, Provider, useStoreApi } from '../store';
 
 interface ProChatProviderProps extends ChatProps {
   children: ReactNode;
   devtoolOptions?: boolean | DevtoolsOptions;
+  chatRef?: ProChatChatReference;
 }
 
 export const ProChatProvider = memo<ProChatProviderProps>(
-  ({ children, devtoolOptions, ...props }) => {
+  ({
+    children,
+    devtoolOptions,
+    chats,
+    onChatsChange,
+    loading,
+    helloMessage,
+    userMeta,
+    assistantMeta,
+    request,
+    chatRef,
+    ...props
+  }) => {
     let isWrapped = true;
 
-    const Content = <>{children}</>;
+    const Content = (
+      <>
+        {children}
+        <StoreUpdater
+          chatRef={chatRef}
+          init={!loading}
+          helloMessage={helloMessage}
+          chats={chats}
+          userMeta={userMeta}
+          request={request}
+          assistantMeta={assistantMeta}
+          onChatsChange={onChatsChange}
+        />
+      </>
+    );
 
     try {
       useStoreApi();
@@ -23,6 +51,20 @@ export const ProChatProvider = memo<ProChatProviderProps>(
       return Content;
     }
 
-    return <Provider createStore={() => createStore(props, devtoolOptions)}>{Content}</Provider>;
+    return (
+      <Provider createStore={() => createStore(props, devtoolOptions)}>
+        {Content}
+        <StoreUpdater
+          chatRef={chatRef}
+          init={!loading}
+          helloMessage={helloMessage}
+          chats={chats}
+          userMeta={userMeta}
+          request={request}
+          assistantMeta={assistantMeta}
+          onChatsChange={onChatsChange}
+        />
+      </Provider>
+    );
   },
 );
diff --git a/src/ProChat/container/StoreUpdater.tsx b/src/ProChat/container/StoreUpdater.tsx
@@ -1,15 +1,31 @@
-import { memo } from 'react';
+import { memo, MutableRefObject, useImperativeHandle } from 'react';
 import { createStoreUpdater } from 'zustand-utils';
 
+import { ProChatInstance, useProChat } from '../hooks/useProChat';
 import { ChatProps, ChatState, useStoreApi } from '../store';
 
-export type StoreUpdaterProps = Partial<
-  Pick<ChatState, 'chats' | 'config' | 'init' | 'onChatsChange' | 'helloMessage' | 'request'>
-> &
-  Pick<ChatProps, 'userMeta' | 'assistantMeta'>;
+export type ProChatChatReference = MutableRefObject<ProChatInstance | undefined>;
+
+export interface StoreUpdaterProps
+  extends Partial<
+      Pick<ChatState, 'chats' | 'config' | 'init' | 'onChatsChange' | 'helloMessage' | 'request'>
+    >,
+    Pick<ChatProps, 'userMeta' | 'assistantMeta'> {
+  chatRef?: ProChatChatReference;
+}
 
 const StoreUpdater = memo<StoreUpdaterProps>(
-  ({ init, onChatsChange, request, userMeta, assistantMeta, helloMessage, chats, config }) => {
+  ({
+    init,
+    onChatsChange,
+    chatRef,
+    request,
+    userMeta,
+    assistantMeta,
+    helloMessage,
+    chats,
+    config,
+  }) => {
     const storeApi = useStoreApi();
     const useStoreUpdater = createStoreUpdater(storeApi);
 
@@ -25,6 +41,10 @@ const StoreUpdater = memo<StoreUpdaterProps>(
     useStoreUpdater('onChatsChange', onChatsChange);
 
     useStoreUpdater('request', request);
+
+    const instance = useProChat();
+    useImperativeHandle(chatRef, () => instance);
+
     return null;
   },
 );
diff --git a/src/ProChat/container/index.tsx b/src/ProChat/container/index.tsx
@@ -6,61 +6,24 @@ import App from './App';
 import { DevtoolsOptions } from 'zustand/middleware';
 import { ChatProps } from '../store';
 import { ProChatProvider } from './Provider';
-import StoreUpdater from './StoreUpdater';
+import { ProChatChatReference } from './StoreUpdater';
 
 export interface ProChatProps extends ChatProps {
   renderInput?: ReactNode;
   __PRO_CHAT_STORE_DEVTOOLS__?: boolean | DevtoolsOptions;
   showTitle?: boolean;
   style?: CSSProperties;
   className?: string;
+  chatRef?: ProChatChatReference;
 }
 
 export const ProChat = memo<ProChatProps>(
-  ({
-    renderInput,
-    __PRO_CHAT_STORE_DEVTOOLS__,
-    chats,
-    onChatsChange,
-    initialChats,
-    loading,
-    helloMessage,
-    userMeta,
-    assistantMeta,
-    showTitle,
-    request,
-    onResetMessage,
-    style,
-    className,
-    autocompleteRequest,
-    ...props
-  }) => {
+  ({ renderInput, __PRO_CHAT_STORE_DEVTOOLS__, showTitle, style, className, ...props }) => {
     return (
-      <ProChatProvider
-        initialChats={initialChats}
-        chats={chats}
-        loading={loading}
-        helloMessage={helloMessage}
-        userMeta={userMeta}
-        assistantMeta={assistantMeta}
-        request={request}
-        onResetMessage={onResetMessage}
-        autocompleteRequest={autocompleteRequest}
-        {...props}
-        devtoolOptions={__PRO_CHAT_STORE_DEVTOOLS__}
-      >
+      <ProChatProvider {...props} devtoolOptions={__PRO_CHAT_STORE_DEVTOOLS__}>
         <Container>
           <App chatInput={renderInput} showTitle={showTitle} style={style} className={className} />
         </Container>
-        <StoreUpdater
-          init={!loading}
-          helloMessage={helloMessage}
-          chats={chats}
-          userMeta={userMeta}
-          request={request}
-          assistantMeta={assistantMeta}
-          onChatsChange={onChatsChange}
-        />
       </ProChatProvider>
     );
   },
diff --git a/src/ProChat/demos/initialChats.tsx b/src/ProChat/demos/initialChats.tsx
@@ -10,7 +10,7 @@ export default () => {
   const theme = useTheme();
   return (
     <div style={{ background: theme.colorBgLayout }}>
-      <ProChat initialChats={example.chats} config={example.config} />
+      <ProChat initialChats={example.chats} />
     </div>
   );
 };
diff --git a/src/ProChat/demos/meta.tsx b/src/ProChat/demos/meta.tsx
@@ -13,7 +13,7 @@ export default () => {
           avatar: 'https://gw.alipayobjects.com/zos/rmsportal/KDpgvguMpGfqaHPjicRK.svg',
           title: 'Ant Design',
         }}
-        assistantMeta={{ avatar: '🛸', title: '三体世界', backgroundColor: 'blue' }}
+        assistantMeta={{ avatar: '🛸', title: '三体世界', backgroundColor: '#67dedd' }}
         initialChats={chats.chats}
       />
     </div>
diff --git a/src/ProChat/demos/use-pro-chat.tsx b/src/ProChat/demos/use-pro-chat.tsx
@@ -0,0 +1,100 @@
+/**
+ * compact: true
+ */
+import { ProChat, ProChatProvider, useProChat } from '@ant-design/pro-chat';
+import { Button, Divider, Flex, message } from 'antd';
+import { useTheme } from 'antd-style';
+
+import { MockResponse } from '@/ProChat/mocks/streamResponse';
+import { example } from '../mocks/basic';
+
+const Chat = () => {
+  const theme = useTheme();
+  return (
+    <div style={{ background: theme.colorBgLayout }}>
+      <ProChat
+        request={async (messages) => {
+          const mockedData: string = `这是一段模拟的流式字符串数据。本次会话传入了${messages.length}条消息`;
+
+          const mockResponse = new MockResponse(mockedData, 100);
+
+          return mockResponse.getResponse();
+        }}
+      />
+    </div>
+  );
+};
+
+const Control = () => {
+  const proChat = useProChat();
+
+  return (
+    <Flex style={{ padding: 24 }} gap={8} justify={'space-between'}>
+      <Flex gap={8}>
+        <Button
+          type={'primary'}
+          onClick={() => {
+            proChat.sendMessage('这是程序化发送的消息');
+          }}
+        >
+          发送一条消息
+        </Button>
+        <Button
+          onClick={() => {
+            const messages = proChat.getChatMessages();
+
+            const msg = messages.at(-1);
+            if (msg) {
+              message.info(msg.content);
+            } else {
+              message.warning('会话为空');
+            }
+          }}
+        >
+          获取最新会话消息
+        </Button>
+      </Flex>
+
+      <Button
+        onClick={() => {
+          const messages = proChat.getChatMessages();
+          const { id, content } = messages[0] || {};
+
+          if (!id) return;
+          proChat.setMessageContent(id, content + '👋');
+        }}
+      >
+        修改首条消息，添加表情：👋
+      </Button>
+      <Flex gap={8}>
+        <Button
+          danger
+          onClick={() => {
+            const messages = proChat.getChatMessages();
+            proChat.deleteMessage(messages[0].id);
+            message.success('已删除第一条消息');
+          }}
+        >
+          删除第一条消息
+        </Button>
+        <Button
+          type={'primary'}
+          danger
+          onClick={() => {
+            proChat.clearMessage();
+          }}
+        >
+          清空消息
+        </Button>
+      </Flex>
+    </Flex>
+  );
+};
+
+export default () => (
+  <ProChatProvider initialChats={example.chats}>
+    <Control />
+    <Divider>🔼 程序化控制 | 🔽 用户控制</Divider>
+    <Chat />
+  </ProChatProvider>
+);
diff --git a/src/ProChat/demos/use-ref.tsx b/src/ProChat/demos/use-ref.tsx
@@ -0,0 +1,44 @@
+/**
+ * compact: true
+ */
+import { ProChat, ProChatInstance } from '@ant-design/pro-chat';
+import { useTheme } from 'antd-style';
+import { useRef } from 'react';
+
+import { MockResponse } from '@/ProChat/mocks/streamResponse';
+import { Button } from 'antd';
+import { example } from '../mocks/basic';
+
+export default () => {
+  const theme = useTheme();
+  const proChatRef = useRef<ProChatInstance>();
+
+  return (
+    <div style={{ background: theme.colorBgLayout }}>
+      <Button
+        type={'primary'}
+        onClick={() => {
+          if (!proChatRef.current) return;
+          const messages = proChatRef.current.getChatMessages();
+          const { id, content } = messages[0] || {};
+
+          if (!id) return;
+          proChatRef.current.setMessageContent(id, content + '👋');
+        }}
+      >
+        修改首条消息，添加表情：👋
+      </Button>
+      <ProChat
+        initialChats={example.chats}
+        chatRef={proChatRef}
+        request={async (messages) => {
+          const mockedData: string = `这是一段模拟的流式字符串数据。本次会话传入了${messages.length}条消息`;
+
+          const mockResponse = new MockResponse(mockedData, 100);
+
+          return mockResponse.getResponse();
+        }}
+      />
+    </div>
+  );
+};
diff --git a/src/ProChat/hooks/useProChat.ts b/src/ProChat/hooks/useProChat.ts
diff --git a/src/ProChat/index.md b/src/ProChat/index.md
diff --git a/src/ProChat/index.tsx b/src/ProChat/index.tsx
diff --git a/src/ProChat/store/store.ts b/src/ProChat/store/store.ts
