From ae64e46e6f9924de88cde57b9924922df893b3ec Mon Sep 17 00:00:00 2001
From: pablodanswer <pablo@danswer.ai>
Date: Tue, 17 Dec 2024 17:06:25 -0800
Subject: [PATCH] k

---
 web/package-lock.json                 | 11 +++++++----
 web/package.json                      |  2 +-
 web/src/app/chat/message/Messages.tsx | 13 +++++++++----
 web/src/app/chat/message/codeUtils.ts | 14 ++++++++++++++
 4 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/web/package-lock.json b/web/package-lock.json
index 48c45ba0b53..d32345fe64f 100644
--- a/web/package-lock.json
+++ b/web/package-lock.json
@@ -40,7 +40,7 @@
         "favicon-fetch": "^1.0.0",
         "formik": "^2.2.9",
         "js-cookie": "^3.0.5",
-        "katex": "^0.16.15",
+        "katex": "^0.16.17",
         "lodash": "^4.17.21",
         "lucide-react": "^0.454.0",
         "mdast-util-find-and-replace": "^3.0.1",
@@ -9326,9 +9326,9 @@
       }
     },
     "node_modules/katex": {
-      "version": "0.16.15",
-      "resolved": "https://registry.npmjs.org/katex/-/katex-0.16.15.tgz",
-      "integrity": "sha512-yE9YJIEAk2aZ+FL/G8r+UGw0CTUzEA8ZFy6E+8tc3spHUKq3qBnzCkI1CQwGoI9atJhVyFPEypQsTY7mJ1Pi9w==",
+      "version": "0.16.17",
+      "resolved": "https://registry.npmjs.org/katex/-/katex-0.16.17.tgz",
+      "integrity": "sha512-OyzSrXBllz+Jdc9Auiw0kt21gbZ4hkz8Q5srVAb2U9INcYIfGKbxe+bvNvEz1bQ/NrDeRRho5eLCyk/L03maAw==",
       "funding": [
         "https://opencollective.com/katex",
         "https://github.com/sponsors/katex"
@@ -14149,6 +14149,7 @@
       "version": "9.0.1",
       "resolved": "https://registry.npmjs.org/react-markdown/-/react-markdown-9.0.1.tgz",
       "integrity": "sha512-186Gw/vF1uRkydbsOIkcGXw7aHq0sZOCRFFjGrr7b9+nVZg4UfA4enXCaxm4fUzecU38sWfrNDitGhshuU7rdg==",
+      "license": "MIT",
       "dependencies": {
         "@types/hast": "^3.0.0",
         "devlop": "^1.0.0",
@@ -14534,6 +14535,7 @@
       "version": "2.0.0",
       "resolved": "https://registry.npmjs.org/rehype-prism-plus/-/rehype-prism-plus-2.0.0.tgz",
       "integrity": "sha512-FeM/9V2N7EvDZVdR2dqhAzlw5YI49m9Tgn7ZrYJeYHIahM6gcXpH0K1y2gNnKanZCydOMluJvX2cB9z3lhY8XQ==",
+      "license": "MIT",
       "dependencies": {
         "hast-util-to-string": "^3.0.0",
         "parse-numeric-range": "^1.3.0",
@@ -14557,6 +14559,7 @@
       "version": "4.0.0",
       "resolved": "https://registry.npmjs.org/remark-gfm/-/remark-gfm-4.0.0.tgz",
       "integrity": "sha512-U92vJgBPkbw4Zfu/IiW2oTZLSL3Zpv+uI7My2eq8JxKgqraFdU8YUGicEJCEgSbeaG+QDFqIcwwfMTOEelPxuA==",
+      "license": "MIT",
       "dependencies": {
         "@types/mdast": "^4.0.0",
         "mdast-util-gfm": "^3.0.0",
diff --git a/web/package.json b/web/package.json
index e09e648afa1..4a52f11c9bc 100644
--- a/web/package.json
+++ b/web/package.json
@@ -42,7 +42,7 @@
     "favicon-fetch": "^1.0.0",
     "formik": "^2.2.9",
     "js-cookie": "^3.0.5",
-    "katex": "^0.16.15",
+    "katex": "^0.16.17",
     "lodash": "^4.17.21",
     "lucide-react": "^0.454.0",
     "mdast-util-find-and-replace": "^3.0.1",
diff --git a/web/src/app/chat/message/Messages.tsx b/web/src/app/chat/message/Messages.tsx
index f7c5d8a1a44..1a7f885cbd1 100644
--- a/web/src/app/chat/message/Messages.tsx
+++ b/web/src/app/chat/message/Messages.tsx
@@ -66,7 +66,7 @@ import RegenerateOption from "../RegenerateOption";
 import { LlmOverride } from "@/lib/hooks";
 import { ContinueGenerating } from "./ContinueMessage";
 import { MemoizedAnchor, MemoizedParagraph } from "./MemoizedTextComponents";
-import { extractCodeText } from "./codeUtils";
+import { extractCodeText, preprocessLaTeX } from "./codeUtils";
 import ToolResult from "../../../components/tools/ToolResult";
 import CsvContent from "../../../components/tools/CSVContent";
 import SourceCard, {
@@ -75,7 +75,7 @@ import SourceCard, {
 
 import remarkMath from "remark-math";
 import rehypeKatex from "rehype-katex";
-import "katex/dist/katex.min.css"; // Add this import for KaTeX styles
+import "katex/dist/katex.min.css";
 
 const TOOLS_WITH_CUSTOM_HANDLING = [
   SEARCH_TOOL_NAME,
@@ -228,6 +228,7 @@ export const AIMessage = ({
   setPresentingDocument?: (document: OnyxDocument) => void;
 }) => {
   const toolCallGenerating = toolCall && !toolCall.tool_result;
+
   const processContent = (content: string | JSX.Element) => {
     if (typeof content !== "string") {
       return content;
@@ -246,12 +247,16 @@ export const AIMessage = ({
 
       const lastMatch = matches[matches.length - 1];
       if (!lastMatch.endsWith("```")) {
-        return content;
+        return preprocessLaTeX(content);
       }
     }
 
-    return content + (!isComplete && !toolCallGenerating ? " [*]() " : "");
+    return (
+      preprocessLaTeX(content) +
+      (!isComplete && !toolCallGenerating ? " [*]() " : "")
+    );
   };
+
   const finalContent = processContent(content as string);
 
   const [isRegenerateHovered, setIsRegenerateHovered] = useState(false);
diff --git a/web/src/app/chat/message/codeUtils.ts b/web/src/app/chat/message/codeUtils.ts
index a9cd13944f0..3f5fad5dc6b 100644
--- a/web/src/app/chat/message/codeUtils.ts
+++ b/web/src/app/chat/message/codeUtils.ts
@@ -59,3 +59,17 @@ export function extractCodeText(
 
   return codeText || "";
 }
+
+export const preprocessLaTeX = (content: string) => {
+  // Replace block-level LaTeX delimiters \[ \] with $$ $$
+  const blockProcessedContent = content.replace(
+    /\\\[([\s\S]*?)\\\]/g,
+    (_, equation) => `$$${equation}$$`
+  );
+  // Replace inline LaTeX delimiters \( \) with $ $
+  const inlineProcessedContent = blockProcessedContent.replace(
+    /\\\(([\s\S]*?)\\\)/g,
+    (_, equation) => `$${equation}$`
+  );
+  return inlineProcessedContent;
+};