1 files changed, 20 insertions, 31 deletions
diff --git a/src/client/views/nodes/ChatBox/ChatBox.tsx b/src/client/views/nodes/ChatBox/ChatBox.tsx
index 73f35f501..3ecb2d340 100644
--- a/src/client/views/nodes/ChatBox/ChatBox.tsx
+++ b/src/client/views/nodes/ChatBox/ChatBox.tsx
@@ -1,33 +1,24 @@
-import { MathJaxContext } from 'better-react-mathjax';
 import { action, makeObservable, observable, observe, reaction, runInAction } from 'mobx';
 import { observer } from 'mobx-react';
 import OpenAI, { ClientOptions } from 'openai';
-import { ImageFile, Message } from 'openai/resources/beta/threads/messages';
-import { RunStep } from 'openai/resources/beta/threads/runs/steps';
 import * as React from 'react';
 import { Doc } from '../../../../fields/Doc';
-import { Id } from '../../../../fields/FieldSymbols';
 import { CsvCast, DocCast, PDFCast, StrCast } from '../../../../fields/Types';
-import { CsvField } from '../../../../fields/URLField';
 import { Networking } from '../../../Network';
-import { DocUtils } from '../../../documents/DocUtils';
 import { DocumentType } from '../../../documents/DocumentTypes';
 import { Docs } from '../../../documents/Documents';
-import { DocumentManager } from '../../../util/DocumentManager';
 import { LinkManager } from '../../../util/LinkManager';
 import { ViewBoxAnnotatableComponent } from '../../DocComponent';
 import { FieldView, FieldViewProps } from '../FieldView';
 import './ChatBox.scss';
 import MessageComponent from './MessageComponent';
-import { ASSISTANT_ROLE, AssistantMessage, AI_Document, convertToAIDocument } from './types';
-import { Annotation } from 'mobx/dist/internal';
-import { FormEvent } from 'react';
-import { url } from 'inspector';
+import { ASSISTANT_ROLE, AssistantMessage, AI_Document, convertToAIDocument, Citation } from './types';
 import { Vectorstore } from './vectorstore/VectorstoreUpload';
-import { DocumentView } from '../DocumentView';
 import { CollectionFreeFormDocumentView } from '../CollectionFreeFormDocumentView';
 import { CollectionFreeFormView } from '../../collections/collectionFreeForm';
 import { ChatBot } from './ChatBot';
+import dotenv from 'dotenv';
+dotenv.config();
 
 @observer
 export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
@@ -151,7 +142,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         const answerElement = xmlDoc.querySelector('answer');
         const followUpQuestionsElement = xmlDoc.querySelector('follow_up_questions');
 
-        const text = answerElement ? answerElement.textContent || '' : '';
+        const text = answerElement ? answerElement.innerHTML || '' : ''; // Use innerHTML to preserve citation tags
         const followUpQuestions = followUpQuestionsElement ? Array.from(followUpQuestionsElement.querySelectorAll('question')).map(q => q.textContent || '') : [];
 
         return {
@@ -161,6 +152,19 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         };
     }
 
+    @action
+    updateMessageCitations = (index: number, citations: Citation[]) => {
+        if (this.history[index]) {
+            this.history[index].citations = citations;
+        }
+    };
+
+    @action
+    handleCitationClick = (citation: Citation) => {
+        console.log('Citation clicked:', citation);
+        // You can implement additional functionality here, such as showing a modal with the full citation content
+    };
+
     // @action
     // uploadLinks = async (linkedDocs: Doc[]) => {
     //     if (this.isInitializing) {
@@ -259,7 +263,6 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     };
     render() {
         return (
-            /** <MathJaxContext config={this.mathJaxConfig}> **/
             <div className="chatBox">
                 {this.isInitializing && <div className="initializing-overlay">Initializing...</div>}
                 <div
@@ -271,29 +274,16 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                     }}>
                     <div className="messages">
                         {this.history.map((message, index) => (
-                            <MessageComponent
-                                key={index}
-                                message={message}
-                                toggleToolLogs={this.toggleToolLogs}
-                                expandedLogIndex={this.expandedScratchpadIndex}
-                                index={index}
-                                showModal={() => {}} // Implement this method if needed
-                                goToLinkedDoc={() => {}} // Implement this method if needed
-                                setCurrentFile={() => {}} // Implement this method if needed
-                                onFollowUpClick={this.handleFollowUpClick}
-                            />
+                            <MessageComponent key={index} message={message} index={index} onFollowUpClick={this.handleFollowUpClick} onCitationClick={this.handleCitationClick} updateMessageCitations={this.updateMessageCitations} />
                         ))}
                         {this.current_message && (
                             <MessageComponent
                                 key={this.history.length}
                                 message={this.current_message}
-                                toggleToolLogs={this.toggleToolLogs}
-                                expandedLogIndex={this.expandedScratchpadIndex}
                                 index={this.history.length}
-                                showModal={() => {}} // Implement this method if needed
-                                goToLinkedDoc={() => {}} // Implement this method if needed
-                                setCurrentFile={() => {}} // Implement this method if needed
                                 onFollowUpClick={this.handleFollowUpClick}
+                                onCitationClick={this.handleCitationClick}
+                                updateMessageCitations={this.updateMessageCitations}
                             />
                         )}
                     </div>
@@ -305,7 +295,6 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                     </button>
                 </form>
             </div>
-            /** </MathJaxContext> **/
         );
     }
 }