1 files changed, 34 insertions, 9 deletions

diff --git a/src/client/views/nodes/chatbot/chatboxcomponents/MessageComponent.tsx b/src/client/views/nodes/chatbot/chatboxcomponents/MessageComponent.tsx
index 44dedca78..120e20001 100644
--- a/src/client/views/nodes/chatbot/chatboxcomponents/MessageComponent.tsx
+++ b/src/client/views/nodes/chatbot/chatboxcomponents/MessageComponent.tsx
@@ -34,40 +34,57 @@ interface MessageComponentProps {
  * @param {MessageComponentProps} props - The props for the component.
  */
 const MessageComponentBox: React.FC<MessageComponentProps> = ({ message, onFollowUpClick, onCitationClick }) => {
+    // State for managing whether the dropdown is open or closed for processing info
     const [dropdownOpen, setDropdownOpen] = useState(false);
 
+    /**
+     * Renders the content of the message based on the type (e.g., grounded text, normal text).
+     * @param {MessageContent} item - The content item to render.
+     * @returns {JSX.Element} JSX element rendering the content.
+     */
     const renderContent = (item: MessageContent) => {
         const i = item.index;
 
+        // Handle grounded text with citations
         if (item.type === TEXT_TYPE.GROUNDED) {
             const citation_ids = item.citation_ids || [];
             return (
                 <span key={i} className="grounded-text">
-                    <ReactMarkdown>{item.text}</ReactMarkdown>
+                    {item.text}
                     {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)}>
-                                [{idx + 1}]
+                                {i + idx + 1}
                             </button>
                         );
                     })}
+                    <br />
                 </span>
             );
-        } else if (item.type === TEXT_TYPE.NORMAL) {
+        }
+
+        // Handle normal text
+        else if (item.type === TEXT_TYPE.NORMAL) {
             return (
                 <span key={i} className="normal-text">
                     <ReactMarkdown>{item.text}</ReactMarkdown>
                 </span>
             );
-        } else if ('query' in item) {
+        }
+
+        // Handle query type content
+        else if ('query' in item) {
             return (
                 <span key={i} className="query-text">
                     <ReactMarkdown>{JSON.stringify(item.query)}</ReactMarkdown>
                 </span>
             );
-        } else {
+        }
+
+        // Fallback for any other content type
+        else {
             return (
                 <span key={i}>
                     <ReactMarkdown>{JSON.stringify(item)}</ReactMarkdown>
@@ -76,18 +93,24 @@ const MessageComponentBox: React.FC<MessageComponentProps> = ({ message, onFollo
         }
     };
 
+    // Check if the message contains processing information (thoughts/actions)
     const hasProcessingInfo = message.processing_info && message.processing_info.length > 0;
 
+    /**
+     * Renders processing information such as thoughts or actions during message handling.
+     * @param {ProcessingInfo} info - The processing information to render.
+     * @returns {JSX.Element | null} JSX element rendering the processing info or null.
+     */
     const renderProcessingInfo = (info: ProcessingInfo) => {
         if (info.type === PROCESSING_TYPE.THOUGHT) {
             return (
-                <div key={info.index} className="processing-item">
+                <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="processing-item">
+                <div key={info.index} className="dropdown-item">
                     <strong>Action:</strong> {info.content}
                 </div>
             );
@@ -97,8 +120,6 @@ const MessageComponentBox: React.FC<MessageComponentProps> = ({ message, onFollo
 
     return (
         <div className={`message ${message.role}`}>
-            <div className="message-content">{message.content && message.content.map(messageFragment => <React.Fragment key={messageFragment.index}>{renderContent(messageFragment)}</React.Fragment>)}</div>
-
             {/* Processing Information Dropdown */}
             {hasProcessingInfo && (
                 <div className="processing-info">
@@ -106,9 +127,13 @@ const MessageComponentBox: React.FC<MessageComponentProps> = ({ message, onFollo
                         {dropdownOpen ? 'Hide Agent Thoughts/Actions' : 'Show Agent Thoughts/Actions'}
                     </button>
                     {dropdownOpen && <div className="info-content">{message.processing_info.map(renderProcessingInfo)}</div>}
+                    <br />
                 </div>
             )}
 
+            {/* Message Content */}
+            <div className="message-content">{message.content && message.content.map(messageFragment => <React.Fragment key={messageFragment.index}>{renderContent(messageFragment)}</React.Fragment>)}</div>
+
             {/* Follow-up Questions Section */}
             {message.follow_up_questions && message.follow_up_questions.length > 0 && (
                 <div className="follow-up-questions">