✨ feat: support setMessageItem (#17)

ant-design · Dec 6, 2023 · 30b516d · 30b516d
1 parent c78f532
commit 30b516d
Show file tree

Hide file tree

Showing 6 changed files with 64 additions and 26 deletions.
diff --git a/src/ProChat/container/Provider.tsx b/src/ProChat/container/Provider.tsx
@@ -3,13 +3,14 @@ import { memo, ReactNode } from 'react';
 import { DevtoolsOptions } from 'zustand/middleware';
 import { ChatProps, createStore, Provider, useStoreApi } from '../store';
 
-interface ProChatProviderProps extends ChatProps {
+interface ProChatProviderProps<T extends Record<string, any> = Record<string, any>>
+  extends ChatProps<T> {
   children: ReactNode;
   devtoolOptions?: boolean | DevtoolsOptions;
   chatRef?: ProChatChatReference;
 }
 
-export const ProChatProvider = memo<ProChatProviderProps>(
+export const ProChatProvider = memo<ProChatProviderProps<any>>(
   ({
     children,
     devtoolOptions,

diff --git a/src/ProChat/container/index.tsx b/src/ProChat/container/index.tsx
@@ -1,5 +1,5 @@
 import { App as Container } from 'antd';
-import { CSSProperties, ReactNode, memo } from 'react';
+import { CSSProperties, ReactNode } from 'react';
 
 import App from './App';
 
@@ -8,7 +8,7 @@ import { ChatProps } from '../store';
 import { ProChatProvider } from './Provider';
 import { ProChatChatReference } from './StoreUpdater';
 
-export interface ProChatProps extends ChatProps {
+export interface ProChatProps<T extends Record<string, any>> extends ChatProps<T> {
   renderInput?: ReactNode;
   __PRO_CHAT_STORE_DEVTOOLS__?: boolean | DevtoolsOptions;
   showTitle?: boolean;
@@ -17,14 +17,19 @@ export interface ProChatProps extends ChatProps {
   chatRef?: ProChatChatReference;
 }
 
-export const ProChat = memo<ProChatProps>(
-  ({ renderInput, __PRO_CHAT_STORE_DEVTOOLS__, showTitle, style, className, ...props }) => {
-    return (
-      <ProChatProvider {...props} devtoolOptions={__PRO_CHAT_STORE_DEVTOOLS__}>
-        <Container>
-          <App chatInput={renderInput} showTitle={showTitle} style={style} className={className} />
-        </Container>
-      </ProChatProvider>
-    );
-  },
-);
+export function ProChat<T extends Record<string, any> = Record<string, any>>({
+  renderInput,
+  __PRO_CHAT_STORE_DEVTOOLS__,
+  showTitle,
+  style,
+  className,
+  ...props
+}: ProChatProps<T>) {
+  return (
+    <ProChatProvider {...props} devtoolOptions={__PRO_CHAT_STORE_DEVTOOLS__}>
+      <Container>
+        <App chatInput={renderInput} showTitle={showTitle} style={style} className={className} />
+      </Container>
+    </ProChatProvider>
+  );
+}
diff --git a/src/ProChat/hooks/useProChat.ts b/src/ProChat/hooks/useProChat.ts
@@ -9,9 +9,31 @@ export interface ProChatInstance
     ChatStore,
     'resendMessage' | 'stopGenerateMessage' | 'sendMessage' | 'deleteMessage' | 'clearMessage'
   > {
+  /**
+   * 获取当前聊天列表对象
+   * @returns ChatStore['chats']
+   */
   getChats: () => ChatStore['chats'];
+  /**
+   * 获取当前聊天消息列表
+   * @returns ChatMessage[]
+   */
   getChatMessages: () => ChatMessage[];
+  /**
+   * 设置消息内容
+   * @param id
+   * @param content
+   * @returns  void
+   */
   setMessageContent: (id: string, content: string) => void;
+  /**
+   * 修改消息的某个属性
+   * @param id
+   * @param key
+   * @param value
+   * @returns  void
+   */
+  setMessageValue: (id: string, key: keyof ChatMessage<Record<string, any>>, value: any) => void;
 }
 
 export const useProChat = () => {
@@ -33,6 +55,12 @@ export const useProChat = () => {
     dispatchMessage({ type: 'updateMessage', id, key: 'content', value: content });
   });
 
+  const setMessageValue = useMemoizedFn(
+    (id: string, key: keyof ChatMessage<Record<string, any>>, value) => {
+      dispatchMessage({ type: 'updateMessage', id, key, value });
+    },
+  );
+
   return useMemo<ProChatInstance>(() => {
     return {
       getChats,
@@ -43,6 +71,7 @@ export const useProChat = () => {
       deleteMessage,
       clearMessage,
       setMessageContent,
+      setMessageValue,
     };
   }, []);
 };
diff --git a/src/ProChat/store/initialState.ts b/src/ProChat/store/initialState.ts
@@ -7,16 +7,16 @@ import { FlexBasicProps } from 'react-layout-kit/lib/FlexBasic';
 
 export type ChatRequest = (messages: ChatMessage[], config: ModelConfig) => Promise<Response>;
 
-export interface ChatPropsState {
+export interface ChatPropsState<T extends Record<string, any> = Record<string, any>> {
   /**
    * 语言模型角色设定
    */
   config: ModelConfig;
   /**
    * 聊天记录
    */
-  chats: ChatMessageMap;
-  onChatsChange?: (chats: ChatMessageMap) => void;
+  chats: ChatMessageMap<T>;
+  onChatsChange?: (chats: ChatMessageMap<T>) => void;
   displayMode: 'chat' | 'docs';
   userMeta: MetaData;
   assistantMeta: MetaData;
@@ -31,7 +31,7 @@ export interface ChatPropsState {
    * @param message
    * @returns message id
    */
-  genMessageId?: (message: ChatMessage[], parentId: string) => Promise<string>;
+  genMessageId?: (message: ChatMessage<T>[], parentId: string) => Promise<string>;
 
   /**
    * 重置消息
@@ -64,7 +64,7 @@ export interface ChatPropsState {
   /**
    * 信息框额外渲染
    */
-  messageItemExtraRender?: (message: ChatMessage, type: 'assistant' | 'user') => React.ReactNode;
+  messageItemExtraRender?: (message: ChatMessage<T>, type: 'assistant' | 'user') => React.ReactNode;
 
   /**
    * 信息框顶部的操作列表

diff --git a/src/ProChat/store/store.ts b/src/ProChat/store/store.ts
@@ -9,10 +9,11 @@ import { createWithEqualityFn } from 'zustand/traditional';
 import { ChatAction, chatAction } from './action';
 import { ChatPropsState, ChatState, initialState } from './initialState';
 
-export interface ChatProps extends Partial<ChatPropsState> {
+export interface ChatProps<T extends Record<string, any> = Record<string, any>>
+  extends Partial<ChatPropsState<T>> {
   // init
   loading?: boolean;
-  initialChats?: ChatPropsState['chats'];
+  initialChats?: ChatPropsState<T>['chats'];
   userMeta?: MetaData;
   assistantMeta?: MetaData;
 }

diff --git a/src/types/message.ts b/src/types/message.ts
@@ -9,7 +9,7 @@ export interface ChatMessageError {
   type: string | number;
 }
 
-export interface ChatMessage {
+export interface ChatMessage<T extends Record<string, any> = Record<string, any>> {
   /**
    * @title 内容
    * @description 消息内容
@@ -27,11 +27,13 @@ export interface ChatMessage {
   createAt: number;
   id: string;
   updateAt: number;
-  extra?: Record<string, any>;
+  extra?: T;
 }
 
-export type ChatMessageMap = Record<string, ChatMessage>;
-
+export type ChatMessageMap<T extends Record<string, any> = Record<string, any>> = Record<
+  string,
+  ChatMessage<T>
+>;
 export interface OpenAIFunctionCall {
   arguments?: string;
   name: string;