From 584d02b92a2949aeb7927cc44c9a82f769dfb2c9 Mon Sep 17 00:00:00 2001
From: ItzCrazyKns <95534749+ItzCrazyKns@users.noreply.github.com>
Date: Thu, 20 Mar 2025 10:56:03 +0530
Subject: [PATCH] feat(app): add thinking model support

---
 src/components/MessageBox.tsx | 38 +++++++++++++++++++++++++++----
 src/components/ThinkBox.tsx   | 43 +++++++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+), 5 deletions(-)
 create mode 100644 src/components/ThinkBox.tsx

diff --git a/src/components/MessageBox.tsx b/src/components/MessageBox.tsx
index 499c3d6..4473c3b 100644
--- a/src/components/MessageBox.tsx
+++ b/src/components/MessageBox.tsx
@@ -12,13 +12,18 @@ import {
   Layers3,
   Plus,
 } from 'lucide-react';
-import Markdown from 'markdown-to-jsx';
+import Markdown, { MarkdownToJSX } from 'markdown-to-jsx';
 import Copy from './MessageActions/Copy';
 import Rewrite from './MessageActions/Rewrite';
 import MessageSources from './MessageSources';
 import SearchImages from './SearchImages';
 import SearchVideos from './SearchVideos';
 import { useSpeech } from 'react-text-to-speech';
+import ThinkBox from './ThinkBox';
+
+const ThinkTagProcessor = ({ children }: { children: React.ReactNode }) => {
+  return <ThinkBox content={children as string} />;
+};
 
 const MessageBox = ({
   message,
@@ -44,27 +49,48 @@ const MessageBox = ({
 
   useEffect(() => {
     const regex = /\[(\d+)\]/g;
+    let processedMessage = message.content;
+
+    if (message.role === 'assistant' && message.content.includes('<think>')) {
+      const openThinkTag = processedMessage.match(/<think>/g)?.length || 0;
+      const closeThinkTag = processedMessage.match(/<\/think>/g)?.length || 0;
+
+      if (openThinkTag > closeThinkTag) {
+        processedMessage += '</think> <a> </a>'; // The extra <a> </a> is to prevent the the think component from looking bad
+      }
+    }
 
     if (
       message.role === 'assistant' &&
       message?.sources &&
       message.sources.length > 0
     ) {
-      return setParsedMessage(
-        message.content.replace(
+      setParsedMessage(
+        processedMessage.replace(
           regex,
           (_, number) =>
-            `<a href="${message.sources?.[number - 1]?.metadata?.url}" target="_blank" className="bg-light-secondary dark:bg-dark-secondary px-1 rounded ml-1 no-underline text-xs text-black/70 dark:text-white/70 relative">${number}</a>`,
+            `<a href="${
+              message.sources?.[number - 1]?.metadata?.url
+            }" target="_blank" className="bg-light-secondary dark:bg-dark-secondary px-1 rounded ml-1 no-underline text-xs text-black/70 dark:text-white/70 relative">${number}</a>`,
         ),
       );
+      return;
     }
 
     setSpeechMessage(message.content.replace(regex, ''));
-    setParsedMessage(message.content);
+    setParsedMessage(processedMessage);
   }, [message.content, message.sources, message.role]);
 
   const { speechStatus, start, stop } = useSpeech({ text: speechMessage });
 
+  const markdownOverrides: MarkdownToJSX.Options = {
+    overrides: {
+      think: {
+        component: ThinkTagProcessor,
+      },
+    },
+  };
+
   return (
     <div>
       {message.role === 'user' && (
@@ -111,11 +137,13 @@ const MessageBox = ({
                   Answer
                 </h3>
               </div>
+
               <Markdown
                 className={cn(
                   '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={markdownOverrides}
               >
                 {parsedMessage}
               </Markdown>
diff --git a/src/components/ThinkBox.tsx b/src/components/ThinkBox.tsx
new file mode 100644
index 0000000..9c6a576
--- /dev/null
+++ b/src/components/ThinkBox.tsx
@@ -0,0 +1,43 @@
+'use client';
+
+import { useState } from 'react';
+import { cn } from '@/lib/utils';
+import { ChevronDown, ChevronUp, BrainCircuit } from 'lucide-react';
+
+interface ThinkBoxProps {
+  content: string;
+}
+
+const ThinkBox = ({ content }: ThinkBoxProps) => {
+  const [isExpanded, setIsExpanded] = useState(false);
+
+  return (
+    <div className="my-4 bg-light-secondary/50 dark:bg-dark-secondary/50 rounded-xl border border-light-200 dark:border-dark-200 overflow-hidden">
+      <button
+        onClick={() => setIsExpanded(!isExpanded)}
+        className="w-full flex items-center justify-between px-4 py-1 text-black/90 dark:text-white/90 hover:bg-light-200 dark:hover:bg-dark-200 transition duration-200"
+      >
+        <div className="flex items-center space-x-2">
+          <BrainCircuit
+            size={20}
+            className="text-[#9C27B0] dark:text-[#CE93D8]"
+          />
+          <p className="font-medium text-sm">Thinking Process</p>
+        </div>
+        {isExpanded ? (
+          <ChevronUp size={18} className="text-black/70 dark:text-white/70" />
+        ) : (
+          <ChevronDown size={18} className="text-black/70 dark:text-white/70" />
+        )}
+      </button>
+
+      {isExpanded && (
+        <div className="px-4 py-3 text-black/80 dark:text-white/80 text-sm border-t border-light-200 dark:border-dark-200 bg-light-100/50 dark:bg-dark-100/50 whitespace-pre-wrap">
+          {content}
+        </div>
+      )}
+    </div>
+  );
+};
+
+export default ThinkBox;