From 9792841a07c124501cad9534180b3a6cd1a3f095 Mon Sep 17 00:00:00 2001 From: yangdx Date: Sat, 26 Apr 2025 12:30:29 +0800 Subject: [PATCH] Stablize mermaid render in history messages --- .../src/components/retrieval/ChatMessage.tsx | 118 +++++++++++------- .../src/features/RetrievalTesting.tsx | 56 ++++++--- 2 files changed, 109 insertions(+), 65 deletions(-) diff --git a/lightrag_webui/src/components/retrieval/ChatMessage.tsx b/lightrag_webui/src/components/retrieval/ChatMessage.tsx index 81888889..1647c2a2 100644 --- a/lightrag_webui/src/components/retrieval/ChatMessage.tsx +++ b/lightrag_webui/src/components/retrieval/ChatMessage.tsx @@ -1,4 +1,4 @@ -import { ReactNode, useCallback, useEffect, useRef } from 'react' +import { ReactNode, useCallback, useEffect, useMemo, useRef, memo, useState } from 'react' // Import useMemo import { Message } from '@/api/lightrag' import useTheme from '@/hooks/useTheme' import Button from '@/components/ui/Button' @@ -19,10 +19,17 @@ import { LoaderIcon, CopyIcon } from 'lucide-react' import { useTranslation } from 'react-i18next' export type MessageWithError = Message & { + id: string // Unique identifier for stable React keys isError?: boolean + /** + * Indicates if the mermaid diagram in this message has been rendered. + * Used to persist the rendering state across updates and prevent flickering. + */ + mermaidRendered?: boolean } -export const ChatMessage = ({ message, isComplete = true }: { message: MessageWithError, isComplete?: boolean }) => { +// Restore original component definition and export +export const ChatMessage = ({ message }: { message: MessageWithError }) => { // Remove isComplete prop const { t } = useTranslation() const handleCopyMarkdown = useCallback(async () => { if (message.content) { @@ -50,17 +57,23 @@ export const ChatMessage = ({ message, isComplete = true }: { message: MessageWi remarkPlugins={[remarkGfm, remarkMath]} rehypePlugins={[rehypeReact]} skipHtml={false} - components={{ - code: (props) => , - p: ({ children }) =>

{children}

, - h1: ({ children }) =>

{children}

, - h2: ({ children }) =>

{children}

, - h3: ({ children }) =>

{children}

, - h4: ({ children }) =>

{children}

, - ul: ({ children }) => , - ol: ({ children }) =>
    {children}
, - li: ({ children }) =>
  • {children}
    • - }} + // Memoize the components object to prevent unnecessary re-renders of ReactMarkdown children + components={useMemo(() => ({ + code: (props: any) => ( // Add type annotation if needed, e.g., props: CodeProps from 'react-markdown/lib/ast-to-react' + + ), + p: ({ children }: { children?: ReactNode }) =>

    {children}

    , + h1: ({ children }: { children?: ReactNode }) =>

    {children}

    , + h2: ({ children }: { children?: ReactNode }) =>

    {children}

    , + h3: ({ children }: { children?: ReactNode }) =>

    {children}

    , + h4: ({ children }: { children?: ReactNode }) =>

    {children}

    , + ul: ({ children }: { children?: ReactNode }) =>
      {children}
    , + ol: ({ children }: { children?: ReactNode }) =>
      {children}
    , + li: ({ children }: { children?: ReactNode }) =>
  • {children}
    • + }), [message.mermaidRendered])} // Dependency ensures update if mermaid state changes > {message.content} @@ -81,12 +94,14 @@ export const ChatMessage = ({ message, isComplete = true }: { message: MessageWi ) } +// Remove the incorrect memo export line + interface CodeHighlightProps { inline?: boolean className?: string children?: ReactNode node?: Element // Keep node for inline check - isComplete?: boolean // Flag to indicate if the message is complete + renderAsDiagram?: boolean // Flag to indicate if rendering as diagram should be attempted } // Helper function remains the same @@ -101,8 +116,10 @@ const isInlineCode = (node?: Element): boolean => { }; -const CodeHighlight = ({ className, children, node, isComplete = true, ...props }: CodeHighlightProps) => { +// Memoize the CodeHighlight component +const CodeHighlight = memo(({ className, children, node, renderAsDiagram = false, ...props }: CodeHighlightProps) => { const { theme } = useTheme(); + const [hasRendered, setHasRendered] = useState(false); // State to track successful render const match = className?.match(/language-(\w+)/); const language = match ? match[1] : undefined; const inline = isInlineCode(node); // Use the helper function @@ -111,35 +128,37 @@ const CodeHighlight = ({ className, children, node, isComplete = true, ...props // Handle Mermaid rendering with debounce useEffect(() => { - // Clear any existing timer when dependencies change - if (debounceTimerRef.current) { - clearTimeout(debounceTimerRef.current); - } + // Effect should run when renderAsDiagram becomes true or hasRendered changes. + // The actual rendering logic inside checks language and hasRendered state. + if (renderAsDiagram && !hasRendered && language === 'mermaid' && mermaidRef.current) { + const container = mermaidRef.current; // Capture ref value - if (language === 'mermaid' && mermaidRef.current) { - const container = mermaidRef.current; // Capture ref value for use inside timeout/callbacks + // Clear previous timer if dependencies change before timeout (e.g., renderAsDiagram flips quickly) + if (debounceTimerRef.current) { + clearTimeout(debounceTimerRef.current); + } - // Set a new timer to render after a short delay debounceTimerRef.current = setTimeout(() => { - // Ensure container still exists when timer fires - if (!container) return; + if (!container) return; // Container might have unmounted + + // Double check hasRendered state inside timeout, in case it changed rapidly + if (hasRendered) return; try { - // Initialize mermaid config (safe to call multiple times) + // Initialize mermaid config mermaid.initialize({ startOnLoad: false, theme: theme === 'dark' ? 'dark' : 'default', securityLevel: 'loose', }); - // Show loading indicator while processing + // Show loading indicator container.innerHTML = '
    '; // Preprocess mermaid content - const rawContent = String(children).replace(/\n$/, '').trim(); // Trim whitespace as well + const rawContent = String(children).replace(/\n$/, '').trim(); // Heuristic check for potentially complete graph definition - // Looks for graph type declaration and some content beyond it. const looksPotentiallyComplete = rawContent.length > 10 && ( rawContent.startsWith('graph') || rawContent.startsWith('sequenceDiagram') || @@ -151,19 +170,17 @@ const CodeHighlight = ({ className, children, node, isComplete = true, ...props rawContent.startsWith('erDiagram') ); - if (!looksPotentiallyComplete) { console.log('Mermaid content might be incomplete, skipping render attempt:', rawContent); - // Keep loading indicator or show a message + // Optionally keep loading indicator or show a message // container.innerHTML = '

    Waiting for complete diagram...

    '; - return; // Don't attempt to render potentially incomplete content + return; } const processedContent = rawContent .split('\n') .map(line => { const trimmedLine = line.trim(); - // Keep subgraph processing if (trimmedLine.startsWith('subgraph')) { const parts = trimmedLine.split(' '); if (parts.length > 1) { @@ -173,26 +190,25 @@ const CodeHighlight = ({ className, children, node, isComplete = true, ...props } return trimmedLine; }) - .filter(line => !line.trim().startsWith('linkStyle')) // Keep filtering linkStyle + .filter(line => !line.trim().startsWith('linkStyle')) .join('\n'); const mermaidId = `mermaid-${Date.now()}`; mermaid.render(mermaidId, processedContent) .then(({ svg, bindFunctions }) => { - // Check ref again inside async callback - // Ensure the container is still the one we intended to update - if (mermaidRef.current === container) { + // Check ref and hasRendered state again inside async callback + if (mermaidRef.current === container && !hasRendered) { container.innerHTML = svg; + setHasRendered(true); // Mark as rendered successfully if (bindFunctions) { - try { // Add try-catch around bindFunctions as it can also throw + try { bindFunctions(container); } catch (bindError) { console.error('Mermaid bindFunctions error:', bindError); - // Optionally display a message in the container container.innerHTML += '

    Diagram interactions might be limited.

    '; } } - } else { + } else if (mermaidRef.current !== container) { console.log('Mermaid container changed before rendering completed.'); } }) @@ -201,11 +217,10 @@ const CodeHighlight = ({ className, children, node, isComplete = true, ...props console.error('Failed content (debounced):', processedContent); if (mermaidRef.current === container) { const errorMessage = error instanceof Error ? error.message : String(error); - // Make error display more robust const errorPre = document.createElement('pre'); errorPre.className = 'text-red-500 text-xs whitespace-pre-wrap break-words'; errorPre.textContent = `Mermaid diagram error: ${errorMessage}\n\nContent:\n${processedContent}`; - container.innerHTML = ''; // Clear previous content + container.innerHTML = ''; container.appendChild(errorPre); } }); @@ -218,24 +233,28 @@ const CodeHighlight = ({ className, children, node, isComplete = true, ...props const errorPre = document.createElement('pre'); errorPre.className = 'text-red-500 text-xs whitespace-pre-wrap break-words'; errorPre.textContent = `Mermaid diagram setup error: ${errorMessage}`; - container.innerHTML = ''; // Clear previous content + container.innerHTML = ''; container.appendChild(errorPre); } } - }, 300); // 300ms debounce delay + }, 300); // Debounce delay } - // Cleanup function to clear the timer + // Cleanup function to clear the timer on unmount or before re-running effect return () => { if (debounceTimerRef.current) { clearTimeout(debounceTimerRef.current); } }; - }, [language, children, theme]); // Dependencies + // Dependencies: renderAsDiagram ensures effect runs when diagram should be shown. + // children, language, theme trigger re-render if code/context changes. + // Dependencies are minimal: only run when the intent to render changes or the rendered state changes. + // Access children, theme, language inside the effect when needed. + }, [renderAsDiagram, hasRendered, language]); // Keep language to ensure it IS mermaid // Render based on language type - // If it's a mermaid language block and the message is not complete, display as plain text - if (language === 'mermaid' && !isComplete) { + // If it's a mermaid language block and rendering as diagram is not requested (e.g., incomplete stream), display as plain text + if (language === 'mermaid' && !renderAsDiagram) { return ( ); -}; +}); + +// Assign display name for React DevTools +CodeHighlight.displayName = 'CodeHighlight'; diff --git a/lightrag_webui/src/features/RetrievalTesting.tsx b/lightrag_webui/src/features/RetrievalTesting.tsx index a2e5f569..6e7bdf7e 100644 --- a/lightrag_webui/src/features/RetrievalTesting.tsx +++ b/lightrag_webui/src/features/RetrievalTesting.tsx @@ -2,7 +2,7 @@ import Input from '@/components/ui/Input' import Button from '@/components/ui/Button' import { useCallback, useEffect, useRef, useState } from 'react' import { throttle } from '@/lib/utils' -import { queryText, queryTextStream, Message } from '@/api/lightrag' +import { queryText, queryTextStream } from '@/api/lightrag' import { errorMessage } from '@/lib/utils' import { useSettingsStore } from '@/stores/settings' import { useDebounce } from '@/hooks/useDebounce' @@ -14,9 +14,18 @@ import type { QueryMode } from '@/api/lightrag' export default function RetrievalTesting() { const { t } = useTranslation() - const [messages, setMessages] = useState( - () => useSettingsStore.getState().retrievalHistory || [] - ) + const [messages, setMessages] = useState(() => { + const history = useSettingsStore.getState().retrievalHistory || [] + // Ensure each message from history has a unique ID and mermaidRendered status + return history.map((msg, index) => { + const msgWithError = msg as MessageWithError // Cast to access potential properties + return { + ...msg, + id: msgWithError.id || `hist-${Date.now()}-${index}`, // Add ID if missing + mermaidRendered: msgWithError.mermaidRendered ?? true // Assume historical mermaid is rendered + } + }) + }) const [inputValue, setInputValue] = useState('') const [isLoading, setIsLoading] = useState(false) const [inputError, setInputError] = useState('') // Error message for input @@ -81,14 +90,17 @@ export default function RetrievalTesting() { // Create messages // Save the original input (with prefix if any) in userMessage.content for display - const userMessage: Message = { + const userMessage: MessageWithError = { + id: crypto.randomUUID(), // Add unique ID content: inputValue, role: 'user' } - const assistantMessage: Message = { + const assistantMessage: MessageWithError = { + id: crypto.randomUUID(), // Add unique ID content: '', - role: 'assistant' + role: 'assistant', + mermaidRendered: false } const prevMessages = [...messages] @@ -113,12 +125,28 @@ export default function RetrievalTesting() { // Create a function to update the assistant's message const updateAssistantMessage = (chunk: string, isError?: boolean) => { assistantMessage.content += chunk + + // Detect if the assistant message contains a complete mermaid code block + // Simple heuristic: look for ```mermaid ... ``` + const mermaidBlockRegex = /```mermaid\s+([\s\S]+?)```/g + let mermaidRendered = false + let match + while ((match = mermaidBlockRegex.exec(assistantMessage.content)) !== null) { + // If the block is not too short, consider it complete + if (match[1] && match[1].trim().length > 10) { + mermaidRendered = true + break + } + } + assistantMessage.mermaidRendered = mermaidRendered + setMessages((prev) => { const newMessages = [...prev] const lastMessage = newMessages[newMessages.length - 1] if (lastMessage.role === 'assistant') { lastMessage.content = assistantMessage.content lastMessage.isError = isError + lastMessage.mermaidRendered = assistantMessage.mermaidRendered } return newMessages }) @@ -279,20 +307,14 @@ export default function RetrievalTesting() { {t('retrievePanel.retrieval.startPrompt')} ) : ( - messages.map((message, idx) => { - // Determine if this message is complete: - // 1. If it's not the last message, it's complete - // 2. If it's the last message but we're not receiving a streaming response, it's complete - // 3. If it's the last message and we're receiving a streaming response, it's not complete - const isLastMessage = idx === messages.length - 1; - const isMessageComplete = !isLastMessage || !isReceivingResponseRef.current; - + messages.map((message) => { // Remove unused idx + // isComplete logic is now handled internally based on message.mermaidRendered return (
    - {} + {}
    ); })