From b6a18323249545e27d23b4200f0359b42ab23269 Mon Sep 17 00:00:00 2001
From: yangdx <gzdaniel@me.com>
Date: Sun, 28 Sep 2025 20:59:56 +0800
Subject: [PATCH] feat: prevent LaTeX parsing errors during streaming by
 implementing completeness detection

- Add latexRendered flag to control KaTeX plugin application
- Implement detectLatexCompleteness function to check formula closure
- Maintain backward compatibility for historical messages
---
 .../src/components/retrieval/ChatMessage.tsx  |  9 +++++--
 .../src/features/RetrievalTesting.tsx         | 25 ++++++++++++++++++-
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/lightrag_webui/src/components/retrieval/ChatMessage.tsx b/lightrag_webui/src/components/retrieval/ChatMessage.tsx
index a3cba698..ffdc41d0 100644
--- a/lightrag_webui/src/components/retrieval/ChatMessage.tsx
+++ b/lightrag_webui/src/components/retrieval/ChatMessage.tsx
@@ -27,6 +27,11 @@ export type MessageWithError = Message & {
    * Used to persist the rendering state across updates and prevent flickering.
    */
   mermaidRendered?: boolean
+  /**
+   * Indicates if the LaTeX formulas in this message are complete and ready for rendering.
+   * Used to prevent red error text during streaming of incomplete LaTeX formulas.
+   */
+  latexRendered?: boolean
 }
 
 // Restore original component definition and export
@@ -138,7 +143,7 @@ export const ChatMessage = ({ message }: { message: MessageWithError }) => { //
                 remarkPlugins={[remarkGfm, remarkFootnotes, remarkMath]}
                 rehypePlugins={[
                   rehypeRaw,
-                  ...(katexPlugin ? [[katexPlugin, {
+                  ...((katexPlugin && (message.latexRendered ?? true)) ? [[katexPlugin, {
                     errorColor: theme === 'dark' ? '#ef4444' : '#dc2626',
                     throwOnError: false,
                     displayMode: false,
@@ -168,7 +173,7 @@ export const ChatMessage = ({ message }: { message: MessageWithError }) => { //
             remarkPlugins={[remarkGfm, remarkFootnotes, remarkMath]}
             rehypePlugins={[
               rehypeRaw,
-              ...(katexPlugin ? [[
+              ...((katexPlugin && (message.latexRendered ?? true)) ? [[
                 katexPlugin,
                 {
                   errorColor: theme === 'dark' ? '#ef4444' : '#dc2626',
diff --git a/lightrag_webui/src/features/RetrievalTesting.tsx b/lightrag_webui/src/features/RetrievalTesting.tsx
index c96541aa..0f3e0869 100644
--- a/lightrag_webui/src/features/RetrievalTesting.tsx
+++ b/lightrag_webui/src/features/RetrievalTesting.tsx
@@ -25,6 +25,22 @@ const generateUniqueId = () => {
   return `id-${Date.now()}-${Math.random().toString(36).substring(2, 9)}`;
 };
 
+// LaTeX completeness detection function
+const detectLatexCompleteness = (content: string): boolean => {
+  // Check for unclosed block-level LaTeX formulas ($$...$$)
+  const blockLatexMatches = content.match(/\$\$/g) || []
+  const hasUnclosedBlock = blockLatexMatches.length % 2 !== 0
+
+  // Check for unclosed inline LaTeX formulas ($...$, but not $$)
+  // Remove all block formulas first to avoid interference
+  const contentWithoutBlocks = content.replace(/\$\$[\s\S]*?\$\$/g, '')
+  const inlineLatexMatches = contentWithoutBlocks.match(/(?<!\$)\$(?!\$)/g) || []
+  const hasUnclosedInline = inlineLatexMatches.length % 2 !== 0
+
+  // LaTeX is complete if there are no unclosed formulas
+  return !hasUnclosedBlock && !hasUnclosedInline
+}
+
 // Robust COT parsing function to handle multiple think blocks and edge cases
 const parseCOTContent = (content: string) => {
   const thinkStartTag = '<think>'
@@ -97,7 +113,8 @@ export default function RetrievalTesting() {
           return {
             ...msg,
             id: msgWithError.id || `hist-${Date.now()}-${index}`, // Add ID if missing
-            mermaidRendered: msgWithError.mermaidRendered ?? true // Assume historical mermaid is rendered
+            mermaidRendered: msgWithError.mermaidRendered ?? true, // Assume historical mermaid is rendered
+            latexRendered: msgWithError.latexRendered ?? true // Assume historical LaTeX is rendered
           }
         } catch (error) {
           console.error('Error processing message:', error)
@@ -203,6 +220,7 @@ export default function RetrievalTesting() {
         content: '',
         role: 'assistant',
         mermaidRendered: false,
+        latexRendered: false,      // Explicitly initialize to false
         thinkingTime: null,        // Explicitly initialize to null
         thinkingContent: undefined, // Explicitly initialize to undefined
         displayContent: undefined,  // Explicitly initialize to undefined
@@ -282,6 +300,10 @@ export default function RetrievalTesting() {
         }
         assistantMessage.mermaidRendered = mermaidRendered
 
+        // Detect if the assistant message contains complete LaTeX formulas
+        const latexRendered = detectLatexCompleteness(assistantMessage.content)
+        assistantMessage.latexRendered = latexRendered
+
         // Single unified update to avoid race conditions
         setMessages((prev) => {
           const newMessages = [...prev]
@@ -295,6 +317,7 @@ export default function RetrievalTesting() {
               isThinking: assistantMessage.isThinking,
               isError: isError,
               mermaidRendered: assistantMessage.mermaidRendered,
+              latexRendered: assistantMessage.latexRendered,
               thinkingTime: assistantMessage.thinkingTime
             })
           }