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

Merge branch 'add-mermaid'

Browse Source
This commit is contained in:
yangdx
2025-04-22 08:52:33 +08:00
parent e7063b5f1e 333c43fbdf
commit 29d01dc10d
2 changed files with 159 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
lightrag_webui/src/api/lightrag.ts
View File

@@ -333,7 +333,6 @@ export const queryTextStream = async (
try {
const parsed = JSON.parse(line);
if (parsed.response) {
console.log('Received chunk:', parsed.response); // Log for debugging
onChunk(parsed.response);
} else if (parsed.error && onError) {
onError(parsed.error);

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

@@ -1,4 +1,4 @@
import { ReactNode, useCallback } from 'react'
import { ReactNode, useCallback, useEffect, useRef } from 'react'
import { Message } from '@/api/lightrag'
import useTheme from '@/hooks/useTheme'
import Button from '@/components/ui/Button'
@@ -8,6 +8,7 @@ import ReactMarkdown from 'react-markdown'
import remarkGfm from 'remark-gfm'
import rehypeReact from 'rehype-react'
import remarkMath from 'remark-math'
import mermaid from 'mermaid'
import type { Element } from 'hast'
@@ -23,6 +24,8 @@ export type MessageWithError = Message & {
export const ChatMessage = ({ message }: { message: MessageWithError }) => {
const { t } = useTranslation()
// Remove extra spaces around bold text
message.content = message.content.replace(/\*\ {3}/g, '').replace(/\ {4}\*\*/g, '**')
const handleCopyMarkdown = useCallback(async () => {
if (message.content) {
@@ -32,7 +35,7 @@ export const ChatMessage = ({ message }: { message: MessageWithError }) => {
console.error(t('chat.copyError'), err)
}
}
}, [message])
}, [message, t]) // Added t to dependency array
return (
<div
@@ -56,7 +59,7 @@ export const ChatMessage = ({ message }: { message: MessageWithError }) => {
>
{message.content}
</ReactMarkdown>
{message.role === 'assistant' && message.content.length > 0 && (
{message.role === 'assistant' && message.content && message.content.length > 0 && ( // Added check for message.content existence
<Button
onClick={handleCopyMarkdown}
className="absolute right-0 bottom-0 size-6 rounded-md opacity-20 transition-opacity hover:opacity-100"
@@ -64,11 +67,11 @@ export const ChatMessage = ({ message }: { message: MessageWithError }) => {
variant="default"
size="icon"
>
<CopyIcon />
<CopyIcon className="size-4" /> {/* Explicit size */}
</Button>
)}
</pre>
{message.content.length === 0 && <LoaderIcon className="animate-spin duration-2000" />}
{message.content === '' && <LoaderIcon className="animate-spin duration-2000" />} {/* Check for empty string specifically */}
</div>
)
}
@@ -77,39 +80,176 @@ interface CodeHighlightProps {
inline?: boolean
className?: string
children?: ReactNode
node?: Element
node?: Element // Keep node for inline check
}
const isInlineCode = (node: Element): boolean => {
const textContent = (node.children || [])
// Helper function remains the same
const isInlineCode = (node?: Element): boolean => {
if (!node || !node.children) return false;
const textContent = node.children
.filter((child) => child.type === 'text')
.map((child) => (child as any).value)
.join('')
.join('');
// Consider inline if it doesn't contain newline or is very short
return !textContent.includes('\n') || textContent.length < 40;
};
return !textContent.includes('\n')
}
const CodeHighlight = ({ className, children, node, ...props }: CodeHighlightProps) => {
const { theme } = useTheme()
const match = className?.match(/language-(\w+)/)
const language = match ? match[1] : undefined
const inline = node ? isInlineCode(node) : false
const { theme } = useTheme();
const match = className?.match(/language-(\w+)/);
const language = match ? match[1] : undefined;
const inline = isInlineCode(node); // Use the helper function
const mermaidRef = useRef<HTMLDivElement>(null);
const debounceTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null); // Use ReturnType for better typing
// Handle Mermaid rendering with debounce
useEffect(() => {
// Clear any existing timer when dependencies change
if (debounceTimerRef.current) {
clearTimeout(debounceTimerRef.current);
}
if (language === 'mermaid' && mermaidRef.current) {
const container = mermaidRef.current; // Capture ref value for use inside timeout/callbacks
// Set a new timer to render after a short delay
debounceTimerRef.current = setTimeout(() => {
// Ensure container still exists when timer fires
if (!container) return;
try {
// Initialize mermaid config (safe to call multiple times)
mermaid.initialize({
startOnLoad: false,
theme: theme === 'dark' ? 'dark' : 'default',
securityLevel: 'loose',
});
// Show loading indicator while processing
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
// 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') ||
rawContent.startsWith('classDiagram') ||
rawContent.startsWith('stateDiagram') ||
rawContent.startsWith('gantt') ||
rawContent.startsWith('pie') ||
rawContent.startsWith('flowchart') ||
rawContent.startsWith('erDiagram')
);
if (!looksPotentiallyComplete) {
console.log("Mermaid content might be incomplete, skipping render attempt:", rawContent);
// 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
}
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) {
const title = parts.slice(1).join(' ').replace(/["']/g, '');
return `subgraph "${title}"`;
}
}
return trimmedLine;
})
.filter(line => !line.trim().startsWith('linkStyle')) // Keep filtering 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) {
container.innerHTML = svg;
if (bindFunctions) {
try { // Add try-catch around bindFunctions as it can also throw
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 {
console.log("Mermaid container changed before rendering completed.");
}
})
.catch(error => {
console.error('Mermaid rendering promise error (debounced):', error);
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.appendChild(errorPre);
}
});
} catch (error) {
console.error('Mermaid synchronous error (debounced):', error);
console.error('Failed content (debounced):', String(children));
if (mermaidRef.current === container) {
const errorMessage = error instanceof Error ? error.message : String(error);
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.appendChild(errorPre);
}
}
}, 300); // 300ms debounce delay
}
// Cleanup function to clear the timer
return () => {
if (debounceTimerRef.current) {
clearTimeout(debounceTimerRef.current);
}
};
}, [language, children, theme]); // Dependencies
// Render based on language type
if (language === 'mermaid') {
// Container for Mermaid diagram
return <div className="mermaid-diagram-container my-4 overflow-x-auto" ref={mermaidRef}></div>;
}
// Handle non-Mermaid code blocks
return !inline ? (
<SyntaxHighlighter
style={theme === 'dark' ? oneDark : oneLight}
PreTag="div"
PreTag="div" // Use div for block code
language={language}
{...props}
>
{String(children).replace(/\n$/, '')}
</SyntaxHighlighter>
) : (
// Handle inline code
<code
className={cn(className, 'mx-1 rounded-xs bg-black/10 px-1 dark:bg-gray-100/20')}
className={cn(className, 'mx-1 rounded-sm bg-muted px-1 py-0.5 text-sm')} // Adjusted styling for inline code
{...props}
>
{children}
</code>
)
}
);
};