1 files changed, 45 insertions, 82 deletions

diff --git a/src/client/views/nodes/ChatBox/MessageComponent.tsx b/src/client/views/nodes/ChatBox/MessageComponent.tsx
index 15c0811fb..1baf6d7d5 100644
--- a/src/client/views/nodes/ChatBox/MessageComponent.tsx
+++ b/src/client/views/nodes/ChatBox/MessageComponent.tsx
@@ -1,110 +1,73 @@
-/* eslint-disable react/require-default-props */
 import React from 'react';
 import { observer } from 'mobx-react';
-import { MathJax, MathJaxContext } from 'better-react-mathjax';
-import ReactMarkdown from 'react-markdown';
-import { AssistantMessage } from './types';
+import { AssistantMessage, CHUNK_TYPE, Citation } from './types';
 import { TbInfoCircleFilled } from 'react-icons/tb';
 
 interface MessageComponentProps {
     message: AssistantMessage;
-    toggleToolLogs: (index: number) => void;
-    expandedLogIndex: number | null;
     index: number;
-    showModal: () => void;
-    goToLinkedDoc: (url: string) => void;
-    setCurrentFile: (file: { url: string }) => void;
-    onFollowUpClick: (question: string) => void; // New prop
-    isCurrent?: boolean;
+    onFollowUpClick: (question: string) => void;
+    onCitationClick: (citation: Citation) => void;
+    updateMessageCitations: (index: number, citations: Citation[]) => void;
 }
 
-const MessageComponent: React.FC<MessageComponentProps> = function ({
-    message,
-    toggleToolLogs,
-    expandedLogIndex,
-    goToLinkedDoc,
-    index,
-    showModal,
-    setCurrentFile,
-    onFollowUpClick, // New prop
-    isCurrent = false,
-}) {
-    const LinkRenderer = ({ href, children }: { href: string; children: React.ReactNode }) => {
-        const regex = /([a-zA-Z0-9_.!-]+)~~~(citation|file_path)/;
-        const matches = href.match(regex);
-        const url = matches ? matches[1] : href;
-        const linkType = matches ? matches[2] : null;
-        if (linkType === 'citation') {
-            children = <TbInfoCircleFilled />;
-        }
-        const style = {
-            color: 'lightblue',
-            verticalAlign: linkType === 'citation' ? 'super' : 'baseline',
-            fontSize: linkType === 'citation' ? 'smaller' : 'inherit',
-        };
+const MessageComponent: React.FC<MessageComponentProps> = function ({ message, index, onFollowUpClick, onCitationClick, updateMessageCitations }) {
+    const LinkRenderer = ({ children }: { children: React.ReactNode }) => {
+        const text = children as string;
+        const citationRegex = /<citation chunk_id="([^"]*)" type="([^"]*)">([^<]*)<\/citation>/g;
+        const parts = [];
+        let lastIndex = 0;
+        let match;
+        const citations: Citation[] = [];
 
-        return (
-            <a
-                href="#"
-                onClick={e => {
-                    e.preventDefault();
-                    if (linkType === 'citation') {
-                        goToLinkedDoc(url);
-                    } else if (linkType === 'file_path') {
-                        showModal();
-                        setCurrentFile({ url });
-                    }
-                }}
-                style={style}>
-                {children}
-            </a>
-        );
-    };
+        while ((match = citationRegex.exec(text)) !== null) {
+            const [fullMatch, chunkId, type, content] = match;
+            const citation: Citation = { chunk_id: chunkId, type: type as CHUNK_TYPE, text: content };
+            citations.push(citation);
+
+            parts.push(text.slice(lastIndex, match.index));
+            parts.push(
+                <a
+                    key={chunkId}
+                    href="#"
+                    onClick={e => {
+                        e.preventDefault();
+                        onCitationClick(citation);
+                    }}
+                    style={{
+                        color: 'lightblue',
+                        verticalAlign: 'super',
+                        fontSize: 'smaller',
+                    }}>
+                    <TbInfoCircleFilled />
+                </a>
+            );
+            lastIndex = match.index + fullMatch.length;
+        }
 
-    const parseMessage = (text: string) => {
-        const answerMatch = text.match(/<answer>([\s\S]*?)<\/answer>/);
-        const followUpMatch = text.match(/<follow_up_question>([\s\S]*?)<\/follow_up_question>/);
+        parts.push(text.slice(lastIndex));
 
-        const answer = answerMatch ? answerMatch[1] : text;
-        const followUpQuestions = followUpMatch
-            ? followUpMatch[1]
-                  .split('\n')
-                  .filter(q => q.trim())
-                  .map(q => q.replace(/^\d+\.\s*/, '').trim())
-            : [];
+        // Update the message's citations in the ChatBox's history
+        updateMessageCitations(index, citations);
 
-        return { answer, followUpQuestions };
+        return <>{parts}</>;
     };
 
-    const { answer, followUpQuestions } = parseMessage(message.text);
-    console.log('Parsed answer:', answer);
-    console.log('Parsed follow-up questions:', followUpQuestions);
     return (
         <div className={`message ${message.role}`}>
-            <ReactMarkdown components={{ a: LinkRenderer }}>{answer}</ReactMarkdown>
-            {message.image && <img src={message.image} alt="" />}
-            {followUpQuestions.length > 0 && (
+            <div>
+                <LinkRenderer>{message.text}</LinkRenderer>
+            </div>
+            {message.follow_up_questions && message.follow_up_questions.length > 0 && (
                 <div className="follow-up-questions">
                     <h4>Follow-up Questions:</h4>
-                    {followUpQuestions.map((question, idx) => (
+                    {message.follow_up_questions.map((question, idx) => (
                         <button key={idx} className="follow-up-button" onClick={() => onFollowUpClick(question)}>
                             {question}
                         </button>
                     ))}
                 </div>
             )}
-            <div className="message-footer">
-                {message.tool_logs && (
-                    <button className="toggle-logs-button" onClick={() => toggleToolLogs(index)}>
-                        {expandedLogIndex === index ? 'Hide Code Interpreter Logs' : 'Show Code Interpreter Logs'}
-                    </button>
-                )}
-                {expandedLogIndex === index && (
-                    <div className="tool-logs">
-                        <pre>{message.tool_logs}</pre>
-                    </div>
-                )}
-            </div>
         </div>
     );
 };