1 files changed, 23 insertions, 11 deletions
diff --git a/src/client/views/nodes/chatbot/chatboxcomponents/MessageComponent.tsx b/src/client/views/nodes/chatbot/chatboxcomponents/MessageComponent.tsx
index e463d15bf..1a3d4dbc6 100644
--- a/src/client/views/nodes/chatbot/chatboxcomponents/MessageComponent.tsx
+++ b/src/client/views/nodes/chatbot/chatboxcomponents/MessageComponent.tsx
@@ -11,6 +11,7 @@ import React, { useState } from 'react';
 import { observer } from 'mobx-react';
 import { AssistantMessage, Citation, MessageContent, PROCESSING_TYPE, ProcessingInfo, TEXT_TYPE } from '../types/types';
 import ReactMarkdown from 'react-markdown';
+import remarkGfm from 'remark-gfm';
 
 /**
  * Props for the MessageComponentBox.
@@ -50,16 +51,27 @@ const MessageComponentBox: React.FC<MessageComponentProps> = ({ message, onFollo
             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>
-                        );
-                    })}
+                    <ReactMarkdown
+                        remarkPlugins={[remarkGfm]}
+                        components={{
+                            p: ({ node, children }) => (
+                                <span className="grounded-text">
+                                    {children}
+                                    {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)} style={{ display: 'inline-flex', alignItems: 'center', marginLeft: '4px' }}>
+                                                {i + idx + 1}
+                                            </button>
+                                        );
+                                    })}
+                                    <br />
+                                </span>
+                            ),
+                        }}>
+                        {item.text}
+                    </ReactMarkdown>
                 </span>
             );
         }
@@ -68,7 +80,7 @@ const MessageComponentBox: React.FC<MessageComponentProps> = ({ message, onFollo
         else if (item.type === TEXT_TYPE.NORMAL) {
             return (
                 <span key={i} className="normal-text">
-                    <ReactMarkdown>{item.text}</ReactMarkdown>
+                    <ReactMarkdown remarkPlugins={[remarkGfm]}>{item.text}</ReactMarkdown>
                 </span>
             );
         }