feat(chat): support rehype plugins and regenerate provider config

Author: mortalYoung

src/chat/button/index.tsx

@@ -17,6 +17,8 @@ export default function Button({ type = 'default', className, children, ...rest
     useEffect(() => {
         if (!document.querySelector(`.${GLOBAL_GRADIENT_CLASSNAME}`)) {
             const div = document.createElement('div');
+            div.style.setProperty('width', '0');
+            div.style.setProperty('height', '0');
             div.className = GLOBAL_GRADIENT_CLASSNAME;
             div.innerHTML = renderToString(
                 <svg width={0} height={0}>

src/chat/content/index.tsx

@@ -26,7 +26,7 @@ const Content = forwardRef<IContentRef, IContentProps>(function (
     { data, placeholder, robotIcon = true, scrollable = true, onRegenerate, onStop },
     forwardedRef
 ) {
-    const { maxRegenerateCount, copy } = useContext();
+    const { maxRegenerateCount, copy, regenerate } = useContext();
     const containerRef = useRef<HTMLDivElement>(null);
 
     const [isStickyAtBottom, setIsStickyAtBottom] = useState<boolean>(true);
@@ -105,14 +105,18 @@ const Content = forwardRef<IContentRef, IContentProps>(function (
             {dataValid ? (
                 <div className="dtc__aigc__content__inner__holder">
                     {data.map((row, idx) => {
+                        const defaultRegenerate =
+                            idx === data.length - 1 && row.messages.length < maxRegenerateCount;
                         return (
                             <React.Fragment key={row.id}>
                                 <Prompt data={row} />
                                 <Message
+                                    prompt={row}
                                     data={row.messages}
                                     regenerate={
-                                        idx === data.length - 1 &&
-                                        row.messages.length < maxRegenerateCount
+                                        typeof regenerate === 'function'
+                                            ? regenerate(row, idx, data)
+                                            : regenerate ?? defaultRegenerate
                                     }
                                     copy={copy}
                                     onRegenerate={(message) => onRegenerate?.(message, row)}

src/chat/index.tsx

@@ -25,12 +25,26 @@ function Chat({
     chat,
     components,
     maxRegenerateCount = DEFAULT_MAX_REGENERATE_COUNT,
+    regenerate,
     copy,
     messageIcons,
+    rehypePlugins,
+    remarkPlugins,
     children,
 }: PropsWithChildren<ChatProviderConfig>) {
     return (
-        <context.Provider value={{ chat, components, maxRegenerateCount, copy, messageIcons }}>
+        <context.Provider
+            value={{
+                chat,
+                components,
+                maxRegenerateCount,
+                copy,
+                messageIcons,
+                regenerate,
+                rehypePlugins,
+                remarkPlugins,
+            }}
+        >
             {children}
         </context.Provider>
     );

src/chat/message/index.tsx

@@ -39,7 +39,7 @@ export default function Message({
     onLazyRendered,
 }: IMessageProps) {
     const divRef = useRef<HTMLDivElement>(null);
-    const { components = {}, messageIcons } = useContext();
+    const { components = {}, messageIcons, rehypePlugins, remarkPlugins } = useContext();
 
     // 当前 Message 的懒加载，是否已经加载过
     const [lazyRendered, setLazyRendered] = useState(false);
@@ -144,6 +144,8 @@ export default function Message({
                             <Markdown
                                 typing={typing}
                                 components={composedComponents}
+                                rehypePlugins={rehypePlugins}
+                                remarkPlugins={remarkPlugins}
                                 onMount={() => {
                                     mountCallback.current();
                                 }}

src/chat/useChat.ts

@@ -169,17 +169,20 @@ export default function useChat<
     function _updateMessage(
         promptId: Id,
         messageId: Id,
-        predicate: (message: Message) => Message
+        predicate: (message: Message) => Message,
+        triggerRerender?: boolean
     ): void;
     function _updateMessage(
         promptId: Id,
         messageId: Id,
-        data: Partial<Omit<MessageProperties, 'id'>>
+        data: Partial<Omit<MessageProperties, 'id'>>,
+        triggerRerender?: boolean
     ): void;
     function _updateMessage(
         promptId: Id,
         messageId: Id,
-        dataOrPredicate: Partial<Omit<MessageProperties, 'id'>> | ((message: Message) => Message)
+        dataOrPredicate: Partial<Omit<MessageProperties, 'id'>> | ((message: Message) => Message),
+        triggerRerender?: boolean
     ) {
         if (!state.current) return;
         state.current = produce(state.current, (draft) => {
@@ -192,7 +195,9 @@ export default function useChat<
                 Object.assign(message, dataOrPredicate);
             }
         });
-        update();
+        if (triggerRerender !== false) {
+            update();
+        }
     }
 
     function _getMessage(promptId: Id, messageId: Id) {

src/chat/useContext.ts

@@ -1,5 +1,6 @@
 import React from 'react';
 import { type Components } from 'react-markdown';
+import { type ReactMarkdownOptions } from 'react-markdown/lib/react-markdown';
 
 import { type ICopyProps } from '../copy';
 import type { Message, Prompt } from './entity';
@@ -22,8 +23,14 @@ export interface IChatContext {
      * 重新回答的最大次数
      */
     maxRegenerateCount: number;
+    /**
+     * 是否支持重新生成
+     */
+    regenerate?: boolean | ((prompt: Prompt, index: number, array: Prompt[]) => boolean);
     copy?: boolean | CopyOptions;
     messageIcons?: React.ReactNode | ((record: Message, prompt: Prompt) => React.ReactNode);
+    rehypePlugins?: ReactMarkdownOptions['rehypePlugins'];
+    remarkPlugins?: ReactMarkdownOptions['remarkPlugins'];
 }
 
 export const context = React.createContext<IChatContext>({