gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Stablize mermaid render in history messages

Browse Source
This commit is contained in:
yangdx
2025-04-26 12:30:29 +08:00
parent e0d99d3c29
commit 9792841a07
2 changed files with 109 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

118
lightrag_webui/src/components/retrieval/ChatMessage.tsx
View File

@@ -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) => <CodeHighlight {...props} isComplete={isComplete} />,
p: ({ children }) => <p className="my-2">{children}</p>,
h1: ({ children }) => <h1 className="text-xl font-bold mt-4 mb-2">{children}</h1>,
h2: ({ children }) => <h2 className="text-lg font-bold mt-4 mb-2">{children}</h2>,
h3: ({ children }) => <h3 className="text-base font-bold mt-3 mb-2">{children}</h3>,
h4: ({ children }) => <h4 className="text-base font-semibold mt-3 mb-2">{children}</h4>,
ul: ({ children }) => <ul className="list-disc pl-5 my-2">{children}</ul>,
ol: ({ children }) => <ol className="list-decimal pl-5 my-2">{children}</ol>,
li: ({ children }) => <li className="my-1">{children}</li>
}}
// 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'
<CodeHighlight
{...props}
renderAsDiagram={message.mermaidRendered ?? false}
/>
),
p: ({ children }: { children?: ReactNode }) => <p className="my-2">{children}</p>,
h1: ({ children }: { children?: ReactNode }) => <h1 className="text-xl font-bold mt-4 mb-2">{children}</h1>,
h2: ({ children }: { children?: ReactNode }) => <h2 className="text-lg font-bold mt-4 mb-2">{children}</h2>,
h3: ({ children }: { children?: ReactNode }) => <h3 className="text-base font-bold mt-3 mb-2">{children}</h3>,
h4: ({ children }: { children?: ReactNode }) => <h4 className="text-base font-semibold mt-3 mb-2">{children}</h4>,
ul: ({ children }: { children?: ReactNode }) => <ul className="list-disc pl-5 my-2">{children}</ul>,
ol: ({ children }: { children?: ReactNode }) => <ol className="list-decimal pl-5 my-2">{children}</ol>,
li: ({ children }: { children?: ReactNode }) => <li className="my-1">{children}</li>
}), [message.mermaidRendered])} // Dependency ensures update if mermaid state changes
>
{message.content}
</ReactMarkdown>
@@ -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 = '<div class="flex justify-center items-center p-4"><svg class="animate-spin h-5 w-5 text-primary" xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="0 0 24 24"><circle class="opacity-25" cx="12" cy="12" r="10" stroke="currentColor" stroke-width="4"></circle><path class="opacity-75" fill="currentColor" d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"></path></svg></div>';
// 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 = '<p class="text-sm text-muted-foreground">Waiting for complete diagram...</p>';
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 += '<p class="text-orange-500 text-xs">Diagram interactions might be limited.</p>';
}
}
} 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 (
<SyntaxHighlighter
style={theme === 'dark' ? oneDark : oneLight}
@@ -273,4 +292,7 @@ const CodeHighlight = ({ className, children, node, isComplete = true, ...props
{children}
</code>
);
};
});
// Assign display name for React DevTools
CodeHighlight.displayName = 'CodeHighlight';

56
lightrag_webui/src/features/RetrievalTesting.tsx
View File

@@ -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<MessageWithError[]>(
() => useSettingsStore.getState().retrievalHistory || []
)
const [messages, setMessages] = useState<MessageWithError[]>(() => {
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')}
</div>
) : (
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 (
<div
key={idx}
key={message.id} // Use stable ID for key
className={`flex ${message.role === 'user' ? 'justify-end' : 'justify-start'}`}
>
{<ChatMessage message={message} isComplete={isMessageComplete} />}
{<ChatMessage message={message} />}
</div>
);
})