1 files changed, 5 insertions, 6 deletions
diff --git a/src/client/views/nodes/ChatBox/MessageComponent.tsx b/src/client/views/nodes/ChatBox/MessageComponent.tsx
index fb4a56bc3..56fde8bb2 100644
--- a/src/client/views/nodes/ChatBox/MessageComponent.tsx
+++ b/src/client/views/nodes/ChatBox/MessageComponent.tsx
@@ -1,7 +1,7 @@
 import React from 'react';
 import { observer } from 'mobx-react';
-import ReactMarkdown from 'react-markdown';
 import { AssistantMessage, Citation } from './types';
+import Markdown from 'react-markdown';
 
 interface MessageComponentProps {
     message: AssistantMessage;
@@ -14,7 +14,7 @@ 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 <ReactMarkdown>{content}</ReactMarkdown>;
+            return content;
         }
 
         const parts = [];
@@ -23,10 +23,9 @@ const MessageComponentBox: React.FC<MessageComponentProps> = function ({ message
         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}`}
+                    key={idx}
                     className="citation-button"
                     onClick={() => onCitationClick(citation)}
                     style={{
@@ -48,11 +47,11 @@ const MessageComponentBox: React.FC<MessageComponentProps> = function ({ message
                     {idx + 1}
                 </button>
             );
-            parts.push(citationButton);
+            parts.push(textBefore, citationButton);
             lastIndex = location;
         });
 
-        parts.push(<ReactMarkdown key="md-last">{content.slice(lastIndex)}</ReactMarkdown>);
+        parts.push(content.slice(lastIndex));
 
         return parts;
     };