✨ feat: add transformToChatMessage & pushChat hook (#157)

ant-design · Apr 22, 2024 · c2e20ce · c2e20ce
1 parent 6196aba
commit c2e20ce
Show file tree

Hide file tree

Showing 5 changed files with 162 additions and 3 deletions.
diff --git a/src/ProChat/demos/sse-trans.tsx b/src/ProChat/demos/sse-trans.tsx
@@ -0,0 +1,122 @@
+/**
+ * compact: true
+ */
+import { ProChat, ProChatInstance } from '@ant-design/pro-chat';
+import { MockSSEResponse } from '../mocks/sseResponse';
+
+import { Avatar, Card, Spin } from 'antd';
+import { useTheme } from 'antd-style';
+import { useRef, useState } from 'react';
+
+const dataArray = [
+  `data: {"id": "chatcmpl-6w****KZb6hx****RzIghUz****Qy", "object": "chat.completion.chunk", "created": 1703582861554, "model": "gpt-3.5-turbo-0301", "choices": [{"delta": {"content": "Searching","type":"function"}, "index": 0, "finish_reason": null}]}`,
+  `data: {"id": "chatcmpl-6w****KZb6hx****RzIghUz****Qy", "object": "chat.completion.chunk", "created": 1703582861554, "model": "gpt-3.5-turbo-0301", "choices": [{"delta": {"content": "苹果"}, "index": 0, "finish_reason": null}]}`,
+  `data: {"id": "chatcmpl-6w****KZb6hx****RzIghUz****Qy", "object": "chat.completion.chunk", "created": 1703582861554, "model": "gpt-3.5-turbo-0301", "choices": [{"delta": {"content": "公司"}, "index": 0, "finish_reason": null}]}`,
+  `data: {"id": "chatcmpl-6w****KZb6hx****RzIghUz****Qy", "object": "chat.completion.chunk", "created": 1703582861554, "model": "gpt-3.5-turbo-0301", "choices": [{"delta": {"content": "是"}, "index": 0, "finish_reason": null}]}`,
+  `data: {"id": "chatcmpl-6w****KZb6hx****RzIghUz****Qy", "object": "chat.completion.chunk", "created": 1703582861554, "model": "gpt-3.5-turbo-0301", "choices": [{"delta": {"content": "一"}, "index": 0, "finish_reason": null}]}`,
+  `data: {"id": "chatcmpl-6w****KZb6hx****RzIghUz****Qy", "object": "chat.completion.chunk", "created": 1703582861554, "model": "gpt-3.5-turbo-0301", "choices": [{"delta": {"content": "家"}, "index": 0, "finish_reason": null}]}`,
+  `data: {"id": "chatcmpl-6w****KZb6hx****RzIghUz****Qy", "object": "chat.completion.chunk", "created": 1703582861554, "model": "gpt-3.5-turbo-0301", "choices": [{"delta": {"content": "科技"}, "index": 0, "finish_reason": null}]}`,
+  `data: {"id": "chatcmpl-6w****KZb6hx****RzIghUz****Qy", "object": "chat.completion.chunk", "created": 1703582861554, "model": "gpt-3.5-turbo-0301", "choices": [{"delta": {"content": "公司"}, "index": 0, "finish_reason": "complete"}]}`,
+];
+
+const LoadingSearch = () => {
+  const [loading, setLoading] = useState(true);
+
+  setTimeout(() => {
+    setLoading(false);
+  }, 2000);
+
+  return (
+    <Card>
+      <Spin spinning={loading}>{loading ? '正在后台帮你执行操作' : '操作完成！'}</Spin>
+    </Card>
+  );
+};
+
+export default () => {
+  const theme = useTheme();
+  const chatRef = useRef<ProChatInstance>();
+  return (
+    <div style={{ background: theme.colorBgLayout }}>
+      <ProChat
+        chatRef={chatRef}
+        style={{ height: 800 }}
+        transformToChatMessage={async (pre) => {
+          try {
+            const preJson = JSON.parse(pre);
+            const { delta } = preJson;
+            const { content, type } = delta;
+            if (content === 'Searching' && type === 'function') {
+              chatRef.current.pushChat({
+                content: 'Doding!',
+                id: 'opDqGn0G',
+                role: 'function',
+              });
+            } else {
+              return content;
+            }
+          } catch (error) {}
+          return '';
+        }}
+        chatItemRenderConfig={{
+          avatarRender: (item) => {
+            if (item?.originData?.role === 'function') {
+              return (
+                <Avatar
+                  size={'large'}
+                  src={'https://gw.alipayobjects.com/zos/rmsportal/KDpgvguMpGfqaHPjicRK.svg'}
+                />
+              );
+            }
+          },
+
+          contentRender: (item) => {
+            if (item?.originData?.role === 'function') {
+              return <LoadingSearch />;
+            }
+          },
+        }}
+        request={async () => {
+          const mockResponse = new MockSSEResponse(dataArray);
+          const response = mockResponse.getResponse();
+
+          // 确保服务器响应是成功的
+          if (!response.ok || !response.body) {
+            throw new Error(`HTTP error! status: ${response.status}`);
+          }
+
+          // 获取 reader
+          const reader = response.body.getReader();
+          const decoder = new TextDecoder('utf-8');
+          const encoder = new TextEncoder();
+
+          const readableStream = new ReadableStream({
+            async start(controller) {
+              function push() {
+                reader
+                  .read()
+                  .then(({ done, value }) => {
+                    if (done) {
+                      controller.close();
+                      return;
+                    }
+                    const chunk = decoder.decode(value, { stream: true });
+                    const message = chunk.replace('data: ', '');
+                    const parsed = JSON.parse(message);
+                    controller.enqueue(encoder.encode(JSON.stringify(parsed.choices[0])));
+                    push();
+                  })
+                  .catch((err) => {
+                    console.error('读取流中的数据时发生错误', err);
+                    controller.error(err);
+                  });
+              }
+              push();
+            },
+          });
+          return new Response(readableStream);
+        }}
+      />
+    </div>
+  );
+};
diff --git a/src/ProChat/hooks/useProChat.ts b/src/ProChat/hooks/useProChat.ts
@@ -14,6 +14,11 @@ export interface ProChatInstance
    * @returns ChatStore['chats']
    */
   getChats: () => ChatStore['chats'];
+  /**
+   * 往数据流中推送消息
+   * @returns void
+   */
+  pushChat: (chats: { id?: string; content: string; role: string }) => void;
   /**
    * 获取当前聊天消息列表
    * @returns ChatMessage[]
@@ -60,6 +65,13 @@ export const useProChat = () => {
   } = storeApi.getState();
 
   const getChats = useRefFunction(() => storeApi.getState().chats);
+  const pushChat = useRefFunction((chat) => {
+    return dispatchMessage({
+      ...chat,
+      type: 'addMessage',
+      message: chat?.content,
+    });
+  });
   const getChatMessages = useRefFunction(() => chatSelectors.currentChats(storeApi.getState()));
 
   const setMessageContent = useRefFunction((id: string, content: string) => {
@@ -75,6 +87,7 @@ export const useProChat = () => {
   return useMemo<ProChatInstance>(() => {
     return {
       getChats,
+      pushChat,
       getChatMessages,
       resendMessage,
       sendMessage,

diff --git a/src/ProChat/index.md b/src/ProChat/index.md
@@ -22,6 +22,12 @@ description: a Chat Solution
 
 <code src="./demos/sse.tsx"></code>
 
+## SSE 的特殊处理
+
+如果你觉得在 Request 里面封装 Reader 不是很方便，要对每次的 Stream 流做处理，我们提供了一个 transformToChatMessage api 来帮你做这个事情
+
+<code src="./demos/sse-trans.tsx"></code>
+
 ## 大数据渲染
 
 <code src="./demos/bigData.tsx"></code>

diff --git a/src/ProChat/store/action.ts b/src/ProChat/store/action.ts
@@ -103,7 +103,14 @@ export interface ChatAction {
    * @returns  消息 id
    */
   getMessageId: (messages: ChatMessage[], parentId: string) => Promise<string>;
-
+  /**
+   * SSE 时候每一条特殊处理转换的方法，处理完后才进行拼接
+   * @returns  string
+   */
+  transformToChatMessage?: (
+    preChatMessage: string,
+    currentContent: string,
+  ) => Promise<string> | string;
   /**
    * 用于更新一条消息的内容（目前仅用于平滑输出时候，其余情况请直接使用 dispatchMessage ）
    */
@@ -167,6 +174,7 @@ export const chatAction: StateCreator<ChatStore, [['zustand/devtools', never]],
       config,
       defaultModelFetcher,
       createSmoothMessage,
+      transformToChatMessage,
       updateMessageContent,
     } = get();
     const abortController = toggleChatLoading(
@@ -235,7 +243,7 @@ export const chatAction: StateCreator<ChatStore, [['zustand/devtools', never]],
         }
         await updateMessageContent(assistantId, content);
       },
-      onMessageHandle: (text) => {
+      onMessageHandle: async (text) => {
         output += text;
 
         if (!isAnimationActive && !isFunctionCall) startAnimation();
@@ -244,7 +252,9 @@ export const chatAction: StateCreator<ChatStore, [['zustand/devtools', never]],
           // aborted 后停止当前输出
           return;
         } else {
-          outputQueue.push(text);
+          outputQueue.push(
+            transformToChatMessage ? await transformToChatMessage(text, output) : text,
+          );
         }
 
         // TODO: need a function call judge callback

diff --git a/src/ProChat/store/initialState.ts b/src/ProChat/store/initialState.ts
@@ -29,6 +29,14 @@ export interface ChatPropsState<T extends Record<string, any> = Record<string, a
   displayMode: 'chat' | 'docs';
   userMeta: MetaData;
   assistantMeta: MetaData;
+  /**
+   * SSE 时候每一条特殊处理转换的方法，处理完后才进行拼接
+   * @returns  string
+   */
+  transformToChatMessage?: (
+    preChatMessage: string,
+    currentContent: string,
+  ) => Promise<string> | string;
   /**
    * 帮助消息
    */