1 files changed, 87 insertions, 64 deletions
diff --git a/src/client/views/nodes/ChatBox/MessageComponent.tsx b/src/client/views/nodes/ChatBox/MessageComponent.tsx
index f27a18891..812e52ee0 100644
--- a/src/client/views/nodes/ChatBox/MessageComponent.tsx
+++ b/src/client/views/nodes/ChatBox/MessageComponent.tsx
@@ -1,82 +1,105 @@
-/* eslint-disable jsx-a11y/control-has-associated-label */
-/* eslint-disable react/require-default-props */
-import { MathJax, MathJaxContext } from 'better-react-mathjax';
+import React, { useState } from 'react';
 import { observer } from 'mobx-react';
-import React from 'react';
-import * as Tb from 'react-icons/tb';
+import { AssistantMessage, Citation, MessageContent, PROCESSING_TYPE, ProcessingInfo, TEXT_TYPE } from './types';
 import ReactMarkdown from 'react-markdown';
-import './MessageComponent.scss';
-import { AssistantMessage } 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 [dropdownOpen, setDropdownOpen] = useState(false);
+
+    const renderContent = (item: MessageContent) => {
+        const i = item.index;
+        //console.log('item', item, 'index', i);
+        if (item.type === TEXT_TYPE.GROUNDED) {
+            const citation_ids = item.citation_ids || [];
+            return (
+                <span key={i} className="grounded-text">
+                    <ReactMarkdown>{item.text}</ReactMarkdown>
+                    {citation_ids.map((id, idx) => {
+                        const citation = message.citations?.find(c => c.citation_id === id);
+                        if (!citation) return null;
+                        return (
+                            <button key={i + idx} className="citation-button" onClick={() => onCitationClick(citation)}>
+                                {i + 1}
+                            </button>
+                        );
+                    })}
+                </span>
+            );
+        } else if (item.type === TEXT_TYPE.NORMAL) {
+            return (
+                <span key={i} className="normal-text">
+                    <ReactMarkdown>{item.text}</ReactMarkdown>
+                </span>
+            );
+        } else if ('query' in item) {
+            return (
+                <span key={i} className="query-text">
+                    <ReactMarkdown>{JSON.stringify(item.query)}</ReactMarkdown>
+                </span>
+            );
+        } else {
+            return (
+                <span key={i}>
+                    <ReactMarkdown>{JSON.stringify(item)}</ReactMarkdown>
+                </span>
+            );
+        }
+    };
+
+    const hasProcessingInfo = message.processing_info && message.processing_info.length > 0;
+
+    const renderProcessingInfo = (info: ProcessingInfo) => {
+        if (info.type === PROCESSING_TYPE.THOUGHT) {
+            return (
+                <div key={info.index} className="dropdown-item">
+                    <strong>Thought:</strong> {info.content}
+                </div>
+            );
+        } else if (info.type === PROCESSING_TYPE.ACTION) {
+            return (
+                <div key={info.index} className="dropdown-item">
+                    <strong>Action:</strong> {info.content}
+                </div>
+            );
+        } else {
+            return null;
+        }
     };
 
-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'}
+            {hasProcessingInfo && (
+                <div className="processing-info">
+                    <button className="toggle-info" onClick={() => setDropdownOpen(!dropdownOpen)}>
+                        {dropdownOpen ? 'Hide Agent Thoughts/Actions' : 'Show Agent Thoughts/Actions'}
                     </button>
-                )}
-                {expandedLogIndex === index && (
-                    <div className="tool-logs">
-                        <pre>{message.tool_logs}</pre>
+                    {dropdownOpen && <div className="info-content">{message.processing_info.map(renderProcessingInfo)}</div>}
+                    <br />
+                </div>
+            )}
+            <div className="message-content">{message.content && message.content.map(messageFragment => <React.Fragment key={messageFragment.index}>{renderContent(messageFragment)}</React.Fragment>)}</div>
+            {message.follow_up_questions && message.follow_up_questions.length > 0 && (
+                <div className="follow-up-questions">
+                    <h4>Follow-up Questions:</h4>
+                    <div className="questions-list">
+                        {message.follow_up_questions.map((question, idx) => (
+                            <button key={idx} className="follow-up-button" onClick={() => onFollowUpClick(question)}>
+                                {question}
+                            </button>
+                        ))}
                     </div>
-                )}
-            </div>
+                </div>
+            )}
         </div>
     );
 };
 
-export default observer(MessageComponent);
+export default observer(MessageComponentBox);