1 files changed, 41 insertions, 43 deletions
diff --git a/src/client/views/nodes/ChatBox/MessageComponent.tsx b/src/client/views/nodes/ChatBox/MessageComponent.tsx
index 56fde8bb2..fd7c445c5 100644
--- a/src/client/views/nodes/ChatBox/MessageComponent.tsx
+++ b/src/client/views/nodes/ChatBox/MessageComponent.tsx
@@ -1,6 +1,6 @@
 import React from 'react';
 import { observer } from 'mobx-react';
-import { AssistantMessage, Citation } from './types';
+import { AssistantMessage, Citation, MessageContent, TEXT_TYPE } from './types';
 import Markdown from 'react-markdown';
 
 interface MessageComponentProps {
@@ -12,53 +12,51 @@ interface MessageComponentProps {
 }
 
 const MessageComponentBox: React.FC<MessageComponentProps> = function ({ message, index, onFollowUpClick, onCitationClick, updateMessageCitations }) {
-    const renderContent = (content: string) => {
-        if (!message.citations || message.citations.length === 0) {
-            return content;
-        }
-
-        const parts = [];
-        let lastIndex = 0;
-
-        message.citations.forEach((citation, idx) => {
-            const location = citation.text_location;
-            const textBefore = content.slice(lastIndex, location);
-            const citationButton = (
-                <button
-                    key={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>
+    const renderContent = (item: MessageContent) => {
+        const i = item.index;
+        if (item.type === TEXT_TYPE.GROUNDED) {
+            const citation_ids = item.citation_ids || [];
+            return (
+                <span key={i} className="grounded-text">
+                    {item.text}
+                    {citation_ids.map((id, idx) => {
+                        const citation = message.citations?.find(c => c.citation_id === id);
+                        if (!citation) return null;
+                        return (
+                            <button
+                                key={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>
+                        );
+                    })}
+                </span>
             );
-            parts.push(textBefore, citationButton);
-            lastIndex = location;
-        });
-
-        parts.push(content.slice(lastIndex));
-
-        return parts;
+        } else {
+            return <span key={i}>{item.text}</span>;
+        }
     };
 
     return (
         <div className={`message ${message.role}`}>
-            <div>{renderContent(message.text_content)}</div>
+            <div>{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>