Mintplex-Labs · timothycarambat · Jan 22, 2025 · Jan 22, 2025 · Jan 22, 2025 · Jan 22, 2025
diff --git a/frontend/src/components/WorkspaceChat/ChatContainer/ChatHistory/HistoricalMessage/index.jsx b/frontend/src/components/WorkspaceChat/ChatContainer/ChatHistory/HistoricalMessage/index.jsx
@@ -10,6 +10,12 @@ import createDOMPurify from "dompurify";
 import { EditMessageForm, useEditMessage } from "./Actions/EditMessage";
 import { useWatchDeleteMessage } from "./Actions/DeleteMessage";
 import TTSMessage from "./Actions/TTSButton";
+import {
+  THOUGHT_REGEX_CLOSE,
+  THOUGHT_REGEX_COMPLETE,
+  THOUGHT_REGEX_OPEN,
+  ThoughtChainComponent,
+} from "../ThoughtContainer";
 
 const DOMPurify = createDOMPurify(window);
 const HistoricalMessage = ({
@@ -97,11 +103,10 @@ const HistoricalMessage = ({
             />
           ) : (
             <div className="break-words">
-              <span
-                className="flex flex-col gap-y-1"
-                dangerouslySetInnerHTML={{
-                  __html: DOMPurify.sanitize(renderMarkdown(message)),
-                }}
+              <RenderChatContent
+                role={role}
+                message={message}
+                expanded={isLastMessage}
               />
               <ChatAttachments attachments={attachments} />
             </div>
@@ -179,3 +184,62 @@ function ChatAttachments({ attachments = [] }) {
     </div>
   );
 }
+
+const RenderChatContent = memo(
+  ({ role, message, expanded = false }) => {
+    // If the message is not from the assistant, we can render it directly
+    // as normal since the user cannot think (lol)
+    if (role !== "assistant")
+      return (
+        <span
+          className="flex flex-col gap-y-1"
+          dangerouslySetInnerHTML={{
+            __html: DOMPurify.sanitize(renderMarkdown(message)),
+          }}
+        />
+      );
+    let thoughtChain = null;
+    let msgToRender = message;
+
+    // If the message is a perfect thought chain, we can render it directly
+    // Complete == open and close tags match perfectly.
+    if (message.match(THOUGHT_REGEX_COMPLETE)) {
+      thoughtChain = message.match(THOUGHT_REGEX_COMPLETE)?.[0];
+      msgToRender = message.replace(THOUGHT_REGEX_COMPLETE, "");
+    }
+
+    // If the message is a thought chain but not a complete thought chain (matching opening tags but not closing tags),
+    // we can render it as a thought chain if we can at least find a closing tag
+    // This can occur when the assistant starts with <thinking> and then <response>'s later.
+    if (
+      message.match(THOUGHT_REGEX_OPEN) &&
+      message.match(THOUGHT_REGEX_CLOSE)
+    ) {
+      const closingTag = message.match(THOUGHT_REGEX_CLOSE)?.[0];
+      const splitMessage = message.split(closingTag);
+      thoughtChain = splitMessage[0] + closingTag;
+      msgToRender = splitMessage[1];
+    }
+
+    return (
+      <>
+        {thoughtChain && (
+          <ThoughtChainComponent content={thoughtChain} expanded={expanded} />
+        )}
+        <span
+          className="flex flex-col gap-y-1"
+          dangerouslySetInnerHTML={{
+            __html: DOMPurify.sanitize(renderMarkdown(msgToRender)),
+          }}
+        />
+      </>
+    );
+  },
+  (prevProps, nextProps) => {
+    return (
+      prevProps.role === nextProps.role &&
+      prevProps.message === nextProps.message &&
+      prevProps.expanded === nextProps.expanded
+    );
+  }
+);
diff --git a/frontend/src/components/WorkspaceChat/ChatContainer/ChatHistory/PromptReply/index.jsx b/frontend/src/components/WorkspaceChat/ChatContainer/ChatHistory/PromptReply/index.jsx
@@ -1,8 +1,14 @@
-import { memo } from "react";
+import { memo, useRef, useEffect } from "react";
 import { Warning } from "@phosphor-icons/react";
 import UserIcon from "../../../../UserIcon";
 import renderMarkdown from "@/utils/chat/markdown";
 import Citations from "../Citation";
+import {
+  THOUGHT_REGEX_CLOSE,
+  THOUGHT_REGEX_COMPLETE,
+  THOUGHT_REGEX_OPEN,
+  ThoughtChainComponent,
+} from "../ThoughtContainer";
 
 const PromptReply = ({
   uuid,
@@ -61,9 +67,9 @@ const PromptReply = ({
       <div className="py-8 px-4 w-full flex gap-x-5 md:max-w-[80%] flex-col">
         <div className="flex gap-x-5">
           <WorkspaceProfileImage workspace={workspace} />
-          <span
-            className="break-words"
-            dangerouslySetInnerHTML={{ __html: renderMarkdown(reply) }}
+          <RenderAssistantChatContent
+            key={`${uuid}-prompt-reply-content`}
+            message={reply}
           />
         </div>
         <Citations sources={sources} />
@@ -88,4 +94,51 @@ export function WorkspaceProfileImage({ workspace }) {
   return <UserIcon user={{ uid: workspace.slug }} role="assistant" />;
 }
 
+function RenderAssistantChatContent({ message }) {
+  const contentRef = useRef("");
+  const thoughtChainRef = useRef(null);
+
+  useEffect(() => {
+    const thinking =
+      message.match(THOUGHT_REGEX_OPEN) && !message.match(THOUGHT_REGEX_CLOSE);
+
+    if (thinking && thoughtChainRef.current) {
+      thoughtChainRef.current.updateContent(message);
+      return;
+    }
+
+    const completeThoughtChain = message.match(THOUGHT_REGEX_COMPLETE)?.[0];
+    const msgToRender = message.replace(THOUGHT_REGEX_COMPLETE, "");
+
+    if (completeThoughtChain && thoughtChainRef.current) {
+      thoughtChainRef.current.updateContent(completeThoughtChain);
+    }
+
+    contentRef.current = msgToRender;
+  }, [message]);
+
+  const thinking =
+    message.match(THOUGHT_REGEX_OPEN) && !message.match(THOUGHT_REGEX_CLOSE);
+  if (thinking)
+    return (
+      <ThoughtChainComponent ref={thoughtChainRef} content="" expanded={true} />
+    );
+
+  return (
+    <div className="flex flex-col gap-y-1">
+      {message.match(THOUGHT_REGEX_COMPLETE) && (
+        <ThoughtChainComponent
+          ref={thoughtChainRef}
+          content=""
+          expanded={true}
+        />
+      )}
+      <span
+        className="break-words"
+        dangerouslySetInnerHTML={{ __html: renderMarkdown(contentRef.current) }}
+      />
+    </div>
+  );
+}
+
 export default memo(PromptReply);
diff --git a/frontend/src/components/WorkspaceChat/ChatContainer/ChatHistory/ThoughtContainer/index.jsx b/frontend/src/components/WorkspaceChat/ChatContainer/ChatHistory/ThoughtContainer/index.jsx
@@ -0,0 +1,129 @@
+import { useState, forwardRef, useImperativeHandle } from "react";
+import renderMarkdown from "@/utils/chat/markdown";
+import { Brain, CaretDown } from "@phosphor-icons/react";
+import DOMPurify from "dompurify";
+import truncate from "truncate";
+import { isMobile } from "react-device-detect";
+
+const THOUGHT_KEYWORDS = ["thought", "thinking", "think", "thought_chain"];
+const CLOSING_TAGS = [...THOUGHT_KEYWORDS, "response", "answer"];
+export const THOUGHT_REGEX_OPEN = new RegExp(
+  THOUGHT_KEYWORDS.map((keyword) => `<${keyword}\\s*(?:[^>]*?)?\\s*>`).join("|")
+);
+export const THOUGHT_REGEX_CLOSE = new RegExp(
+  CLOSING_TAGS.map((keyword) => `</${keyword}\\s*(?:[^>]*?)?>`).join("|")
+);
+export const THOUGHT_REGEX_COMPLETE = new RegExp(
+  THOUGHT_KEYWORDS.map(
+    (keyword) =>
+      `<${keyword}\\s*(?:[^>]*?)?\\s*>[\\s\\S]*?<\\/${keyword}\\s*(?:[^>]*?)?>`
+  ).join("|")
+);
+const THOUGHT_PREVIEW_LENGTH = isMobile ? 25 : 50;
+
+/**
+ * Component to render a thought chain.
+ * @param {string} content - The content of the thought chain.
+ * @param {boolean} expanded - Whether the thought chain is expanded.
+ * @returns {JSX.Element}
+ */
+export const ThoughtChainComponent = forwardRef(
+  ({ content: initialContent, expanded }, ref) => {
+    const [content, setContent] = useState(initialContent);
+    const [isExpanded, setIsExpanded] = useState(expanded);
+    useImperativeHandle(ref, () => ({
+      updateContent: (newContent) => {
+        setContent(newContent);
+      },
+    }));
+
+    const isThinking =
+      content.match(THOUGHT_REGEX_OPEN) && !content.match(THOUGHT_REGEX_CLOSE);
+    const isComplete =
+      content.match(THOUGHT_REGEX_COMPLETE) ||
+      content.match(THOUGHT_REGEX_CLOSE);
+    const tagStrippedContent = content
+      .replace(THOUGHT_REGEX_OPEN, "")
+      .replace(THOUGHT_REGEX_CLOSE, "");
+    const autoExpand =
+      isThinking && tagStrippedContent.length > THOUGHT_PREVIEW_LENGTH;
+    const canExpand = tagStrippedContent.length > THOUGHT_PREVIEW_LENGTH;
+    if (!content || !content.length) return null;
+
+    function handleExpandClick() {
+      if (!canExpand) return;
+      setIsExpanded(!isExpanded);
+    }
+
+    return (
+      <div className="flex justify-start items-end transition-all duration-200 w-full md:max-w-[800px]">
+        <div className="pb-2 w-full flex gap-x-5 flex-col relative">
+          <div
+            style={{
+              transition: "all 0.1s ease-in-out",
+              borderRadius: isExpanded || autoExpand ? "6px" : "24px",
+            }}
+            className={`${isExpanded || autoExpand ? "" : `${canExpand ? "hover:bg-theme-sidebar-item-hover" : ""}`} items-start bg-theme-bg-chat-input py-2 px-4 flex gap-x-2 border border-theme-sidebar-border`}
+          >
+            {isThinking || isComplete ? (
+              <Brain
+                data-tooltip-id="cot-thinking"
+                data-tooltip-content={
+                  isThinking
+                    ? "Model is thinking..."
+                    : "Model has finished thinking"
+                }
+                className={`w-4 h-4 mt-1 ${isThinking ? "text-blue-500 animate-pulse" : "text-green-400"}`}
+                aria-label={
+                  isThinking
+                    ? "Model is thinking..."
+                    : "Model has finished thinking"
+                }
+              />
+            ) : null}
+            <div className="flex-1 overflow-hidden">
+              {!isExpanded && !autoExpand ? (
+                <span
+                  className="text-xs text-theme-text-secondary font-mono inline-block w-full"
+                  dangerouslySetInnerHTML={{
+                    __html: DOMPurify.sanitize(
+                      truncate(tagStrippedContent, THOUGHT_PREVIEW_LENGTH)
+                    ),
+                  }}
+                />
+              ) : (
+                <span
+                  className="text-xs text-theme-text-secondary font-mono inline-block w-full"
+                  dangerouslySetInnerHTML={{
+                    __html: DOMPurify.sanitize(
+                      renderMarkdown(tagStrippedContent)
+                    ),
+                  }}
+                />
+              )}
+            </div>
+            <div className="flex items-center gap-x-2">
+              {!autoExpand && canExpand ? (
+                <button
+                  onClick={handleExpandClick}
+                  data-tooltip-id="expand-cot"
+                  data-tooltip-content={
+                    isExpanded ? "Hide thought chain" : "Show thought chain"
+                  }
+                  className="border-none text-theme-text-secondary hover:text-theme-text-primary transition-colors p-1 rounded-full hover:bg-theme-sidebar-item-hover"
+                  aria-label={
+                    isExpanded ? "Hide thought chain" : "Show thought chain"
+                  }
+                >
+                  <CaretDown
+                    className={`w-4 h-4 transform transition-transform duration-200 ${isExpanded ? "rotate-180" : ""}`}
+                  />
+                </button>
+              ) : null}
+            </div>
+          </div>
+        </div>
+      </div>
+    );
+  }
+);
diff --git a/frontend/src/components/WorkspaceChat/ChatContainer/ChatTooltips/index.jsx b/frontend/src/components/WorkspaceChat/ChatContainer/ChatTooltips/index.jsx
@@ -73,6 +73,12 @@ export function ChatTooltips() {
         delayShow={300}
         className="tooltip !text-xs"
       />
+      <Tooltip
+        id="cot-thinking"
+        place="bottom"
+        delayShow={500}
+        className="tooltip !text-xs"
+      />
     </>
   );
 }