1 files changed, 35 insertions, 60 deletions

diff --git a/src/client/views/nodes/ChatBox/MessageComponent.tsx b/src/client/views/nodes/ChatBox/MessageComponent.tsx
index 38faf7e00..76faff10b 100644
--- a/src/client/views/nodes/ChatBox/MessageComponent.tsx
+++ b/src/client/views/nodes/ChatBox/MessageComponent.tsx
@@ -1,9 +1,6 @@
 import React from 'react';
 import { observer } from 'mobx-react';
-import { AssistantMessage, CHUNK_TYPE, Citation } from './types';
-import { TbInfoCircleFilled } from 'react-icons/tb';
-import { Docs } from '../../../documents/Documents';
-import { DocumentType } from '../../../documents/DocumentTypes';
+import { AssistantMessage, Citation } from './types';
 
 interface MessageComponentProps {
     message: AssistantMessage;
@@ -14,59 +11,53 @@ interface MessageComponentProps {
 }
 
 const MessageComponentBox: React.FC<MessageComponentProps> = function ({ message, index, onFollowUpClick, onCitationClick, updateMessageCitations }) {
-    // public static LayoutString(fieldKey: string) {
-    //     return FieldView.LayoutString(MessageComponentBox, fieldKey);
-    // }
-
-    // the presentation view that renders this slide
-
-    // @computed
-    // get chatBoxView() {
-    //     return this.DocumentView?.().containerViewPath?.().lastElement()?.ComponentView as ChatBox;
-    // }
+    const renderContent = (content: string) => {
+        if (!message.citations || message.citations.length === 0) {
+            return content;
+        }
 
-    const renderContent = (text: string) => {
-        const citationRegex = /<citation chunk_id="([^"]*)" type="([^"]*)">([^<]*)<\/citation>/g;
         const parts = [];
         let lastIndex = 0;
-        let match;
-        const citations: Citation[] = [];
-
-        while ((match = citationRegex.exec(text)) !== null) {
-            const [fullMatch, chunkId, type, content] = match;
-            const citation: Citation = { chunk_id: chunkId, type: type as CHUNK_TYPE, text: content };
-            citations.push(citation);
 
-            parts.push(text.slice(lastIndex, match.index));
-            parts.push(
-                <a
-                    key={chunkId}
-                    href="#"
-                    onClick={e => {
-                        e.preventDefault();
-                        onCitationClick(citation);
-                    }}
+        message.citations.forEach((citation, idx) => {
+            const location = citation.location;
+            const textBefore = content.slice(lastIndex, location);
+            const citationButton = (
+                <button
+                    key={idx}
+                    className="citation-button"
+                    onClick={() => onCitationClick(citation)}
                     style={{
-                        color: 'lightblue',
-                        verticalAlign: 'super',
-                        fontSize: 'smaller',
+                        display: 'inline-flex',
+                        alignItems: 'center',
+                        justifyContent: 'center',
+                        width: '20px',
+                        height: '20px',
+                        borderRadius: '50%',
+                        border: 'none',
+                        background: '#007bff',
+                        color: 'white',
+                        fontSize: '12px',
+                        fontWeight: 'bold',
+                        cursor: 'pointer',
+                        margin: '0 2px',
+                        padding: 0,
                     }}>
-                    <TbInfoCircleFilled />
-                </a>
+                    {idx + 1}
+                </button>
             );
-            lastIndex = match.index + fullMatch.length;
-        }
-
-        parts.push(text.slice(lastIndex));
+            parts.push(textBefore, citationButton);
+            lastIndex = location;
+        });
 
-        updateMessageCitations(index, citations);
+        parts.push(content.slice(lastIndex));
 
-        return <>{parts}</>;
+        return parts;
     };
 
     return (
         <div className={`message ${message.role}`}>
-            <div>{renderContent(message.text)}</div>
+            <div>{renderContent(message.text_content)}</div>
             {message.follow_up_questions && message.follow_up_questions.length > 0 && (
                 <div className="follow-up-questions">
                     <h4>Follow-up Questions:</h4>
@@ -81,20 +72,4 @@ const MessageComponentBox: React.FC<MessageComponentProps> = function ({ message
     );
 };
 
-// Docs.Prototypes.TemplateMap.set(DocumentType.MESSAGE, {
-//     layout: { view: MessageComponentBox, dataField: 'data' },
-//     options: {
-//         acl: '',
-//         _height: 35,
-//         _xMargin: 10,
-//         _yMargin: 10,
-//         _layout_nativeDimEditable: true,
-//         _layout_reflowVertical: true,
-//         _layout_reflowHorizontal: true,
-//         defaultDoubleClick: 'ignore',
-//         systemIcon: 'BsFileEarmarkTextFill',
-//         layout_borderRounding: '10px',
-//     },
-// });
-
 export default observer(MessageComponentBox);