1 files changed, 61 insertions, 66 deletions
diff --git a/src/client/views/nodes/ChatBox/MessageComponent.tsx b/src/client/views/nodes/ChatBox/MessageComponent.tsx
index f27a18891..fb4a56bc3 100644
--- a/src/client/views/nodes/ChatBox/MessageComponent.tsx
+++ b/src/client/views/nodes/ChatBox/MessageComponent.tsx
@@ -1,82 +1,77 @@
-/* eslint-disable jsx-a11y/control-has-associated-label */
-/* eslint-disable react/require-default-props */
-import { MathJax, MathJaxContext } from 'better-react-mathjax';
-import { observer } from 'mobx-react';
 import React from 'react';
-import * as Tb from 'react-icons/tb';
+import { observer } from 'mobx-react';
 import ReactMarkdown from 'react-markdown';
-import './MessageComponent.scss';
-import { AssistantMessage } from './types';
+import { AssistantMessage, Citation } from './types';
 
-const TbCircles = [
-    Tb.TbCircleNumber0Filled,
-    Tb.TbCircleNumber1Filled,
-    Tb.TbCircleNumber2Filled,
-    Tb.TbCircleNumber3Filled,
-    Tb.TbCircleNumber4Filled,
-    Tb.TbCircleNumber5Filled,
-    Tb.TbCircleNumber6Filled,
-    Tb.TbCircleNumber7Filled,
-    Tb.TbCircleNumber8Filled,
-    Tb.TbCircleNumber9Filled,
-];
 interface MessageComponentProps {
     message: AssistantMessage;
-    toggleToolLogs: (index: number) => void;
-    expandedLogIndex: number | null;
     index: number;
-    showModal: () => void;
-    goToLinkedDoc: (url: string) => void;
-    setCurrentFile: (file: { url: string }) => void;
-    isCurrent?: boolean;
+    onFollowUpClick: (question: string) => void;
+    onCitationClick: (citation: Citation) => void;
+    updateMessageCitations: (index: number, citations: Citation[]) => void;
 }
 
-const LinkRendererWrapper = (goToLinkedDoc: (url: string) => void, showModal: () => void, setCurrentFile: (file: { url: string }) => void) =>
-    function LinkRenderer({ href, children }: { href?: string; children?: React.ReactNode }) {
-        const Children = TbCircles[Number(children)]; // pascal case variable needed to convert IconType to JSX.Element tag
-        const [, aurl, linkType] = href?.match(/([a-zA-Z0-9_.!-]+)~~~(citation|file_path)/) ?? [undefined, href, null];
-        const renderType = (content: JSX.Element | null, click: (url: string) => void):JSX.Element => (
-            // eslint-disable-next-line jsx-a11y/anchor-is-valid
-            <a className={`MessageComponent-${linkType}`}
-                href="#"
-                onClick={e => {
-                    e.preventDefault();
-                    aurl && click(aurl);
-                }}>
-                {content}
-            </a> 
-        ); // prettier-ignore
-        switch (linkType) {
-            case 'citation':   return renderType(<Children />, (url: string) => goToLinkedDoc(url));
-            case 'file_path':  return renderType(null, (url: string) => { showModal(); setCurrentFile({ url }); });
-            default:           return null;
-        } // prettier-ignore
+const MessageComponentBox: React.FC<MessageComponentProps> = function ({ message, index, onFollowUpClick, onCitationClick, updateMessageCitations }) {
+    const renderContent = (content: string) => {
+        if (!message.citations || message.citations.length === 0) {
+            return <ReactMarkdown>{content}</ReactMarkdown>;
+        }
+
+        const parts = [];
+        let lastIndex = 0;
+
+        message.citations.forEach((citation, idx) => {
+            const location = citation.text_location;
+            const textBefore = content.slice(lastIndex, location);
+            parts.push(<ReactMarkdown key={`md-${idx}`}>{textBefore}</ReactMarkdown>);
+            const citationButton = (
+                <button
+                    key={`citation-${idx}`}
+                    className="citation-button"
+                    onClick={() => onCitationClick(citation)}
+                    style={{
+                        display: 'inline-flex',
+                        alignItems: 'center',
+                        justifyContent: 'center',
+                        width: '20px',
+                        height: '20px',
+                        borderRadius: '50%',
+                        border: 'none',
+                        background: '#ff6347',
+                        color: 'white',
+                        fontSize: '12px',
+                        fontWeight: 'bold',
+                        cursor: 'pointer',
+                        margin: '0 2px',
+                        padding: 0,
+                    }}>
+                    {idx + 1}
+                </button>
+            );
+            parts.push(citationButton);
+            lastIndex = location;
+        });
+
+        parts.push(<ReactMarkdown key="md-last">{content.slice(lastIndex)}</ReactMarkdown>);
+
+        return parts;
     };
 
-const MessageComponent: React.FC<MessageComponentProps> = function ({ message, toggleToolLogs, expandedLogIndex, goToLinkedDoc, index, showModal, setCurrentFile, isCurrent = false }) {
-    // const messageClass = `${message.role} ${isCurrent ? 'current-message' : ''}`;
     return (
         <div className={`message ${message.role}`}>
-            <MathJaxContext>
-                <MathJax dynamic hideUntilTypeset="every">
-                    <ReactMarkdown components={{ a: LinkRendererWrapper(goToLinkedDoc, showModal, setCurrentFile) }}>{message.text}</ReactMarkdown>
-                </MathJax>
-            </MathJaxContext>
-            {message.image && <img src={message.image} alt="" />}
-            <div className="message-footer">
-                {message.tool_logs && (
-                    <button type="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>{renderContent(message.text_content)}</div>
+            {message.follow_up_questions && message.follow_up_questions.length > 0 && (
+                <div className="follow-up-questions">
+                    <h4>Follow-up Questions:</h4>
+                    {message.follow_up_questions.map((question, idx) => (
+                        <button key={idx} className="follow-up-button" onClick={() => onFollowUpClick(question)}>
+                            {question}
+                        </button>
+                    ))}
+                </div>
+            )}
         </div>
     );
 };
 
-export default observer(MessageComponent);
+export default observer(MessageComponentBox);