fix: Math formatting BIG FIX

Author: Veikkosuhonen

src/client/components/ChatV2/Conversation.tsx

@@ -13,7 +13,7 @@ import 'katex/dist/contrib/mhchem'
 import { ArrowRight } from '@mui/icons-material'
 import FormatQuoteIcon from '@mui/icons-material/FormatQuote'
 import { t } from 'i18next'
-import { useId, useMemo } from 'react'
+import React, { useId, useMemo } from 'react'
 import { useTranslation } from 'react-i18next'
 import type { AssistantMessage, ChatMessage, MessageGenerationInfo, ToolCallResultEvent, ToolCallStatusEvent, UserMessage } from '../../../shared/chat'
 import useLocalStorageState from '../../hooks/useLocalStorageState'
@@ -122,7 +122,7 @@ const AssistantMessageInfo = ({ message }: { message: AssistantMessage }) => {
 }
 
 const AssistantMessageItem = ({ message, setActiveToolResult }: { message: AssistantMessage; setActiveToolResult: (data: ToolCallResultEvent) => void }) => {
-  const processedContent = preprocessMath(message.content)
+  const processedContent = React.useMemo(() => preprocessMath(message.content), [message.content])
   const katexOptions = {
     macros: {
       '\\abs': '\\left|#1\\right|',
@@ -158,8 +158,8 @@ const AssistantMessageItem = ({ message, setActiveToolResult }: { message: Assis
     extensions: ['mhchem.js'],
     output: 'htmlAndMathml',
     errorColor: '#cc0000',
-    throwOnError: false,
-    strict: false, // disables logging katex warnings/errors – if debugging, turn these two on
+    throwOnError: true,
+    strict: "error", // disables logging katex warnings/errors – if debugging, turn these two on
   }
   let codeCount = 0
 

src/client/components/ChatV2/util.ts

@@ -1,4 +1,11 @@
 export const preprocessMath = (content: string): string => {
+  // console.time('preprocessMath')
+  const result = _preprocessMath(content)
+  // console.timeEnd('preprocessMath')
+  return result
+}
+
+const _preprocessMath = (content: string): string => {
   // For preprocessing math in assistant messages from LaTex(-ish :D) format to KaTex-recognizable format
   // Consider upgrading to MathJax for more consistent formatting support if problems arise
 
@@ -17,6 +24,7 @@ export const preprocessMath = (content: string): string => {
   })
 
   processedContent = processedContent
+
     // Convert Latex align environments -> Katex aligned display math
     .replace(/\\begin\{(?:align\*?|aligned)\}([\s\S]*?)\\end\{(?:align\*?|aligned)\}/g, (_, innerContent) => {
       const alignedLines = innerContent
@@ -28,35 +36,20 @@ export const preprocessMath = (content: string): string => {
     })
 
     // Convert Latex display math \[...\] -> Katex display math `$$...$$`
-    .replace(/\\\[([\s\S]*?)\\\]/g, (match, innerContent) => {
-      return `$$${innerContent.trim()}$$`
+    .replace(/\\\[([\s\S]*?)\\\]/g, (match, innerContent: string) => {
+      return `\n$$\n${innerContent.replaceAll('\n', ' ').trim()}\n$$\n`
     })
 
     // Convert Latex inline math \(...\) -> Katex display math `$$...$$`
-    .replace(/\\\(([\s\S]*?)\\\)/g, (match, innerContent) => {
-      return `$$${innerContent.trim()}$$`
+    .replace(/\\\(([\s\S]*?)\\\)/g, (match, innerContent: string) => {
+      return `$$ ${innerContent.replaceAll('\n', ' ').trim()} $$`
     })
 
     // Convert text mode parentheses
     .replace(/\\text\{([^}]*\([^}]*\)[^}]*)\}/g, (match, innerContent) => {
       return `\\text{${innerContent.replace(/\(/g, '\\(').replace(/\)/g, '\\)')}}`
     })
-
-    // Support for matrix environments
-    .replace(/\\begin\{(p|b|v|V)?matrix\*?\}([\s\S]*?)\\end\{(p|b|v|V)?matrix\*?\}/g, (match, matrixTypeGroup, innerContent) => {
-      const matrixType = matrixTypeGroup || ''
-
-      const matrixContent = innerContent
-        .split('\n')
-        .map((line) => line.trim())
-        .filter((line) => line.length > 0)
-        .join(' \\\\\n')
-
-      const kaTeXEnv = `\\begin{${matrixType}matrix}`
-      const kaTeXEndEnv = `\\end{${matrixType}matrix}`
-
-      return `$$\n${kaTeXEnv}\n${matrixContent}\n${kaTeXEndEnv}\n$$`
-    })
+      
 
   codeBlocks.forEach((block, index) => {
     processedContent = processedContent.replace(`__CODE_BLOCK_${index}__`, block)

src/server/services/langchain/MockModel.ts

@@ -2,7 +2,7 @@ import type { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager
 import { AIMessage, AIMessageChunk, type BaseMessage, isHumanMessage, isSystemMessage, isToolMessage } from '@langchain/core/messages'
 import type { ChatGenerationChunk, ChatResult } from '@langchain/core/outputs'
 import { FakeStreamingChatModel } from '@langchain/core/utils/testing'
-import { basicTestContent } from './mockContent'
+import { basicTestContent, mathTestContent, mathTestContent2 } from './mockContent'
 import { StructuredTool } from '@langchain/core/tools'
 
 /**
@@ -37,6 +37,9 @@ export class MockModel extends FakeStreamingChatModel {
     } else if (lastHumanMessage.startsWith('rag')) {
       // Do a tool call
       this.chunks = toolCallChunks
+    } else if (lastHumanMessage.startsWith('math')) {
+      // Do a tool call
+      this.chunks = mathChunks
     } else if (lastHumanMessage.startsWith('temperature')) {
       // Echo the temperature
       this.chunks = [new AIMessageChunk(`Temperature: ${this.temperature}`)]
@@ -74,3 +77,9 @@ const toolCallChunks = [
     ],
   }),
 ]
+
+const mathChunks = [
+  new AIMessageChunk({
+    content: mathTestContent,
+  }),
+]

src/server/services/langchain/mockContent.ts

@@ -45,22 +45,22 @@ And for some general variables:
 ### Matrix Examples
 
 A simple 2x2 matrix:
-\\begin{bmatrix}
+\\[\\begin{bmatrix}
   1 & 0 \\
   0 & 1
-\\end{bmatrix}
+\\end{bmatrix}\\]
 
 A determinant:
-\\begin{vmatrix}
+\\[\\begin{vmatrix}
   a & b \\
   c & d
-\\end{vmatrix}
+\\end{vmatrix}\\]
 
 A general matrix with fractions:
-\\begin{pmatrix}
+\\[\\begin{pmatrix}
   \\frac{1}{3} & \\sqrt{3} \\\\
   \\sin(\\theta) & \\cos(\\phi)
-\\end{pmatrix}
+\\end{pmatrix}\\]
 
 ---
 
@@ -102,10 +102,10 @@ Now for an alignment:
 \\end{align*}
 
 And finally, a small matrix:
-\\begin{pmatrix}
+\\[\\begin{pmatrix}
   5 & 6 \\\\
   7 & 8
-\\end{pmatrix}
+\\end{pmatrix}\\]
 
 ---
 
@@ -127,6 +127,89 @@ A simple chemical reaction: \\(\\ch{H2 + Cl2 -> 2HCl}\\).
 An ion with charge: \\(\\ch{Fe^3+}\\) and an isotope \\(\\ch{^238U}\\).
 
 `
+
+export const mathTestContent2 = `
+---
+
+**1. Vector Addition**
+
+\\[
+\\mathbf{a} + \\mathbf{b} = 
+\\begin{pmatrix}
+a_1 \\\\ a_2 \\\\ a_3
+\\end{pmatrix}
++
+\\begin{pmatrix}
+b_1 \\\\ b_2 \\\\ b_3
+\\end{pmatrix}
+=
+\\begin{pmatrix}
+a_1 + b_1 \\\\ a_2 + b_2 \\\\ a_3 + b_3
+\\end{pmatrix}
+\\]
+
+---
+
+**2. Matrix Multiplication**
+
+\\[
+AB = 
+\\begin{pmatrix}
+1 & 2 \\\\
+3 & 4
+\\end{pmatrix}
+\\begin{pmatrix}
+5 & 6 \\\\
+7 & 8
+\\end{pmatrix}
+=
+\\begin{pmatrix}
+1 \\times 5 + 2 \\times 7 & 1 \\times 6 + 2 \\times 8 \\\\
+3 \\times 5 + 4 \\times 7 & 3 \\times 6 + 4 \\times 8
+\\end{pmatrix}
+=
+\\begin{pmatrix}
+19 & 22 \\\\
+43 & 50
+\\end{pmatrix}
+\\]
+
+---
+
+**3. Solving a System of Linear Equations**
+
+\\[
+\\begin{cases}
+2x + 3y = 5 \\\\
+4x - y = 11
+\\end{cases}
+\\]
+
+---
+
+**4. Eigenvalue Equation**
+
+\\[
+A\\mathbf{v} = \\lambda \\mathbf{v}
+\\]
+
+where \\( A \\) is a matrix, \\( \\mathbf{v} \\) is an eigenvector, and \\( \\lambda \\) is an eigenvalue.
+
+---
+
+**5. Determinant of a 2x2 Matrix**
+
+\\[
+\\det
+\\begin{pmatrix}
+a & b \\\\
+c & d
+\\end{pmatrix}
+= ad - bc
+\\]
+
+---`
+
 export const codeTestContent = `
 ### ✅ JavaScript