From 901c27a3f141507d1c1ffa149d320b503e08af3e Mon Sep 17 00:00:00 2001
From: minicx <minicx@disroot.org>
Date: Fri, 13 Dec 2024 01:56:31 +0300
Subject: [PATCH] feat(highlight): add syntax highlighting

- Add syntax highlighting with [highlight.js](https://highlightjs.org/)

Signed-off-by: minicx <minicx@disroot.org>
---
 ui/components/CodeBlock.tsx  | 27 +++++++++++++++++++++++++++
 ui/components/MessageBox.tsx |  6 ++++++
 ui/package.json              |  1 +
 ui/yarn.lock                 |  5 +++++
 4 files changed, 39 insertions(+)
 create mode 100644 ui/components/CodeBlock.tsx

diff --git a/ui/components/CodeBlock.tsx b/ui/components/CodeBlock.tsx
new file mode 100644
index 0000000..66f1942
--- /dev/null
+++ b/ui/components/CodeBlock.tsx
@@ -0,0 +1,27 @@
+import React, { useEffect, useRef } from 'react';
+import hljs from 'highlight.js';
+
+const CodeBlock = ({
+  children,
+  className,
+}: {
+  children: React.ReactNode;
+  className?: string;
+}) => {
+  const ref = useRef<HTMLElement | null>(null);
+
+  useEffect(() => {
+    if (ref.current && className?.includes('lang-')) {
+      hljs.highlightElement(ref.current);
+      // hljs won't reprocess the element unless this attribute is removed
+      ref.current.removeAttribute('data-highlighted');
+    }
+  }, [children, className]);
+  return (
+    <code className={className} ref={ref}>
+      {children}
+    </code>
+  );
+};
+
+export default CodeBlock;
diff --git a/ui/components/MessageBox.tsx b/ui/components/MessageBox.tsx
index f23127c..856daa5 100644
--- a/ui/components/MessageBox.tsx
+++ b/ui/components/MessageBox.tsx
@@ -19,6 +19,7 @@ import MessageSources from './MessageSources';
 import SearchImages from './SearchImages';
 import SearchVideos from './SearchVideos';
 import { useSpeech } from 'react-text-to-speech';
+import CodeBlock from './CodeBlock';
 
 const MessageBox = ({
   message,
@@ -110,6 +111,11 @@ const MessageBox = ({
                   'prose prose-h1:mb-3 prose-h2:mb-2 prose-h2:mt-6 prose-h2:font-[800] prose-h3:mt-4 prose-h3:mb-1.5 prose-h3:font-[600] dark:prose-invert prose-p:leading-relaxed prose-pre:p-0 font-[400]',
                   'max-w-none break-words text-black dark:text-white',
                 )}
+                options={{
+                  overrides: {
+                    code: CodeBlock,
+                  },
+                }}
               >
                 {parsedMessage}
               </Markdown>
diff --git a/ui/package.json b/ui/package.json
index bdda684..4bd3dee 100644
--- a/ui/package.json
+++ b/ui/package.json
@@ -16,6 +16,7 @@
     "@langchain/openai": "^0.0.25",
     "@tailwindcss/typography": "^0.5.12",
     "clsx": "^2.1.0",
+    "highlight.js": "^11.10.0",
     "langchain": "^0.1.30",
     "lucide-react": "^0.363.0",
     "markdown-to-jsx": "^7.6.2",
diff --git a/ui/yarn.lock b/ui/yarn.lock
index 0798ca9..ae63208 100644
--- a/ui/yarn.lock
+++ b/ui/yarn.lock
@@ -1753,6 +1753,11 @@ hasown@^2.0.0, hasown@^2.0.1, hasown@^2.0.2:
   dependencies:
     function-bind "^1.1.2"
 
+highlight.js@^11.10.0:
+  version "11.10.0"
+  resolved "https://registry.yarnpkg.com/highlight.js/-/highlight.js-11.10.0.tgz#6e3600dc4b33d6dc23d5bd94fbf72405f5892b92"
+  integrity sha512-SYVnVFswQER+zu1laSya563s+F8VDGt7o35d4utbamowvUNLLMovFqwCLSocpZTz3MgaSRA1IbqRWZv97dtErQ==
+
 humanize-ms@^1.2.1:
   version "1.2.1"
   resolved "https://registry.yarnpkg.com/humanize-ms/-/humanize-ms-1.2.1.tgz#c46e3159a293f6b896da29316d8b6fe8bb79bbed"