feat: implement useSmoothStream hook for dynamic text rendering (#8070)

* feat: implement useSmoothStream hook for dynamic text rendering

- Added a new custom hook `useSmoothStream` to manage smooth text streaming with adjustable delays.
- Integrated `useSmoothStream` into the Markdown component to enhance content display during streaming.
- Implemented logic to handle content updates and reset conditions based on block status.

* feat: enhance chunk processing in useSmoothStream hook

- Updated the `addChunk` function to split text into characters for Chinese and words for English using a regular expression.
- Improved text rendering logic to support mixed language content more effectively.

* refactor: improve regular expression for chunk processing in useSmoothStream hook

- Updated the regular expression to enhance text splitting capabilities, allowing for better handling of mixed language content, including Chinese characters and alphanumeric sequences.

* refactor: simplify character processing in useSmoothStream hook

- Replaced the regular expression for character extraction with a direct conversion of the chunk to an array of characters.
- This change streamlines the chunk processing logic, enhancing performance and readability.

* feat: add post-processing capability to Markdown component

- Enhanced the Markdown component to accept an optional post-processing function for text manipulation during streaming.
- Updated MainTextBlock to utilize the new postProcess prop, allowing for dynamic content adjustments based on citation references.

* fix: update mock Markdown component to support post-processing

- Modified the mock implementation of the Markdown component to accept and apply an optional postProcess function, ensuring that the test environment accurately reflects the updated functionality introduced in the previous commit.

src/renderer/src/hooks/useSmoothStream.ts

@@ -0,0 +1,107 @@
+import { useCallback, useEffect, useRef, useState } from 'react'
+
+interface UseSmoothStreamOptions {
+  onUpdate: (text: string) => void
+  streamDone: boolean
+  // 我们不再需要固定的interval，但可以保留一个最小延迟以保证动画感
+  minDelay?: number
+  initialText?: string
+}
+// 如果不行还可以使用Array.from(chunk)分割
+// const reg = /[\u4E00-\u9FFF]|[a-zA-Z0-9]+|\s+|[^\s\w]/g
+
+export const useSmoothStream = ({ onUpdate, streamDone, minDelay = 10, initialText = '' }: UseSmoothStreamOptions) => {
+  const [chunkQueue, setChunkQueue] = useState<string[]>([])
+  const animationFrameRef = useRef<number | null>(null)
+  const displayedTextRef = useRef<string>(initialText)
+  const lastUpdateTimeRef = useRef<number>(0)
+
+  const addChunk = useCallback((chunk: string) => {
+    // 英文按照word拆分, 中文按照字拆分,使用正则表达式
+    // const words = chunk.match(/[\w\d]+/g)
+    const chars = Array.from(chunk)
+    setChunkQueue((prev) => [...prev, ...(chars || [])])
+  }, [])
+
+  const reset = useCallback(
+    (newText = '') => {
+      if (animationFrameRef.current) {
+        cancelAnimationFrame(animationFrameRef.current)
+      }
+      setChunkQueue([])
+      displayedTextRef.current = newText
+      onUpdate(newText)
+    },
+    [onUpdate]
+  )
+
+  const renderLoop = useCallback(
+    (currentTime: number) => {
+      // 1. 如果队列为空，等待下一帧
+      if (chunkQueue.length === 0) {
+        // 如果流还没结束但队列空了，就等待下一帧
+        if (!streamDone) {
+          animationFrameRef.current = requestAnimationFrame(renderLoop)
+        }
+        return
+      }
+
+      // 2. 时间控制，确保最小延迟
+      if (currentTime - lastUpdateTimeRef.current < minDelay) {
+        animationFrameRef.current = requestAnimationFrame(renderLoop)
+        return
+      }
+      lastUpdateTimeRef.current = currentTime
+
+      setChunkQueue((prevQueue) => {
+        // 3. 动态计算本次渲染的字符数
+        // 如果队列积压严重，就一次性渲染更多字符来"追赶"
+        const charsToRenderCount = Math.max(1, Math.floor(prevQueue.length / 5)) // 每次至少渲染1个，最多渲染队列的1/5
+
+        const charsToRender = prevQueue.slice(0, charsToRenderCount)
+        displayedTextRef.current += charsToRender.join('')
+
+        // 4. 立即更新UI
+        onUpdate(displayedTextRef.current)
+
+        // 返回新的队列
+        return prevQueue.slice(charsToRenderCount)
+      })
+
+      // 5. 请求下一帧动画
+      animationFrameRef.current = requestAnimationFrame(renderLoop)
+    },
+    [chunkQueue, streamDone, onUpdate, minDelay]
+  )
+
+  useEffect(() => {
+    // 启动渲染循环
+    animationFrameRef.current = requestAnimationFrame(renderLoop)
+
+    // 组件卸载时清理
+    return () => {
+      if (animationFrameRef.current) {
+        cancelAnimationFrame(animationFrameRef.current)
+      }
+    }
+  }, [renderLoop]) // 依赖 renderLoop
+
+  // 当外部流结束，且队列即将变空时，进行最后一次"瞬移"渲染
+  useEffect(() => {
+    if (streamDone && chunkQueue.length > 0) {
+      const remainingText = chunkQueue.join('')
+      const finalText = displayedTextRef.current + remainingText
+
+      // 取消正在进行的动画循环
+      if (animationFrameRef.current) {
+        cancelAnimationFrame(animationFrameRef.current)
+      }
+
+      // 直接更新到最终状态
+      onUpdate(finalText)
+      setChunkQueue([]) // 清空队列
+    }
+  }, [streamDone, chunkQueue, onUpdate])
+
+  return { addChunk, reset }
+}

src/renderer/src/pages/home/Markdown/Markdown.tsx

@@ -5,13 +5,15 @@ import 'katex/dist/contrib/mhchem'
 import ImageViewer from '@renderer/components/ImageViewer'
 import MarkdownShadowDOMRenderer from '@renderer/components/MarkdownShadowDOMRenderer'
 import { useSettings } from '@renderer/hooks/useSettings'
+import { useSmoothStream } from '@renderer/hooks/useSmoothStream'
 import { EVENT_NAMES, EventEmitter } from '@renderer/services/EventService'
 import type { MainTextMessageBlock, ThinkingMessageBlock, TranslationMessageBlock } from '@renderer/types/newMessage'
 import { parseJSON } from '@renderer/utils'
 import { removeSvgEmptyLines } from '@renderer/utils/formats'
 import { findCitationInChildren, getCodeBlockId, processLatexBrackets } from '@renderer/utils/markdown'
 import { isEmpty } from 'lodash'
-import { type FC, memo, useCallback, useMemo } from 'react'
+import { type FC, memo, useCallback, useEffect, useMemo, useState } from 'react'
+import { useRef } from 'react'
 import { useTranslation } from 'react-i18next'
 import ReactMarkdown, { type Components, defaultUrlTransform } from 'react-markdown'
 import rehypeKatex from 'rehype-katex'
@@ -35,12 +37,56 @@ const DISALLOWED_ELEMENTS = ['iframe']
 interface Props {
   // message: Message & { content: string }
   block: MainTextMessageBlock | TranslationMessageBlock | ThinkingMessageBlock
+  // 可选的后处理函数，用于在流式渲染过程中处理文本（如引用标签转换）
+  postProcess?: (text: string) => string
 }
 
-const Markdown: FC<Props> = ({ block }) => {
+const Markdown: FC<Props> = ({ block, postProcess }) => {
   const { t } = useTranslation()
   const { mathEngine } = useSettings()
 
+  const isTrulyDone = 'status' in block && block.status === 'success'
+  const [displayedContent, setDisplayedContent] = useState(block.content)
+  const [isStreamDone, setIsStreamDone] = useState(isTrulyDone)
+
+  const prevContentRef = useRef(block.content)
+  const prevBlockIdRef = useRef(block.id)
+
+  const { addChunk, reset } = useSmoothStream({
+    onUpdate: (rawText) => {
+      // 如果提供了后处理函数就调用，否则直接使用原始文本
+      const finalText = postProcess ? postProcess(rawText) : rawText
+      setDisplayedContent(finalText)
+    },
+    streamDone: isStreamDone,
+    initialText: block.content
+  })
+
+  useEffect(() => {
+    const newContent = block.content || ''
+    const oldContent = prevContentRef.current || ''
+
+    const isDifferentBlock = block.id !== prevBlockIdRef.current
+
+    const isContentReset = oldContent && newContent && !newContent.startsWith(oldContent)
+
+    if (isDifferentBlock || isContentReset) {
+      reset(newContent)
+    } else {
+      const delta = newContent.substring(oldContent.length)
+      if (delta) {
+        addChunk(delta)
+      }
+    }
+
+    prevContentRef.current = newContent
+    prevBlockIdRef.current = block.id
+
+    // 更新 stream 状态
+    const isStreaming = 'status' in block && block.status === 'streaming'
+    setIsStreamDone(!isStreaming)
+  }, [block.content, block.id, block.status, addChunk, reset])
+
   const remarkPlugins = useMemo(() => {
     const plugins = [
       [remarkGfm, { singleTilde: false }] as Pluggable,
@@ -54,11 +100,11 @@ const Markdown: FC<Props> = ({ block }) => {
   }, [mathEngine])
 
   const messageContent = useMemo(() => {
-    const empty = isEmpty(block.content)
-    const paused = block.status === 'paused'
-    const content = empty && paused ? t('message.chat.completion.paused') : block.content
-    return removeSvgEmptyLines(processLatexBrackets(content))
-  }, [block, t])
+    if ('status' in block && block.status === 'paused' && isEmpty(block.content)) {
+      return t('message.chat.completion.paused')
+    }
+    return removeSvgEmptyLines(processLatexBrackets(displayedContent))
+  }, [block, displayedContent, t])
 
   const rehypePlugins = useMemo(() => {
     const plugins: any[] = []

src/renderer/src/pages/home/Messages/Blocks/MainTextBlock.tsx

@@ -6,7 +6,7 @@ import { type Model } from '@renderer/types'
 import type { MainTextMessageBlock, Message } from '@renderer/types/newMessage'
 import { determineCitationSource, withCitationTags } from '@renderer/utils/citation'
 import { Flex } from 'antd'
-import React, { useMemo } from 'react'
+import React, { useCallback } from 'react'
 import { useSelector } from 'react-redux'
 import styled from 'styled-components'
 
@@ -25,16 +25,20 @@ const MainTextBlock: React.FC<Props> = ({ block, citationBlockId, role, mentions
 
   const rawCitations = useSelector((state: RootState) => selectFormattedCitationsByBlockId(state, citationBlockId))
 
-  const processedContent = useMemo(() => {
-    if (!block.citationReferences?.length || !citationBlockId || rawCitations.length === 0) {
-      return block.content
-    }
+  // 创建引用处理函数，传递给 Markdown 组件在流式渲染中使用
+  const processContent = useCallback(
+    (rawText: string) => {
+      if (!block.citationReferences?.length || !citationBlockId || rawCitations.length === 0) {
+        return rawText
+      }
 
-    // 确定最适合的 source
-    const sourceType = determineCitationSource(block.citationReferences)
+      // 确定最适合的 source
+      const sourceType = determineCitationSource(block.citationReferences)
 
-    return withCitationTags(block.content, rawCitations, sourceType)
-  }, [block.content, block.citationReferences, citationBlockId, rawCitations])
+      return withCitationTags(rawText, rawCitations, sourceType)
+    },
+    [block.citationReferences, citationBlockId, rawCitations]
+  )
 
   return (
     <>
@@ -51,7 +55,7 @@ const MainTextBlock: React.FC<Props> = ({ block, citationBlockId, role, mentions
           {block.content}
         </p>
       ) : (
-        <Markdown block={{ ...block, content: processedContent }} />
+        <Markdown block={block} postProcess={processContent} />
       )}
     </>
   )

src/renderer/src/pages/home/Messages/Blocks/__tests__/MainTextBlock.test.tsx

@@ -78,11 +78,14 @@ vi.mock('@renderer/services/ModelService', () => ({
 // Mock Markdown component
 vi.mock('@renderer/pages/home/Markdown/Markdown', () => ({
   __esModule: true,
-  default: ({ block }: any) => (
-    <div data-testid="mock-markdown" data-content={block.content}>
-      Markdown: {block.content}
-    </div>
-  )
+  default: ({ block, postProcess }: any) => {
+    const content = postProcess ? postProcess(block.content) : block.content
+    return (
+      <div data-testid="mock-markdown" data-content={content}>
+        Markdown: {content}
+      </div>
+    )
+  }
 }))
 
 describe('MainTextBlock', () => {