Merge branch 'lanyi-expanded-agent-paper-main' of https://github.com/brown-dash/Dash-Web into lanyi-expanded-agent-paper-main

1 files changed, 96 insertions, 101 deletions

diff --git a/src/client/views/nodes/chatbot/chatboxcomponents/ChatBox.tsx b/src/client/views/nodes/chatbot/chatboxcomponents/ChatBox.tsx
index 195e85412..c09df166d 100644
--- a/src/client/views/nodes/chatbot/chatboxcomponents/ChatBox.tsx
+++ b/src/client/views/nodes/chatbot/chatboxcomponents/ChatBox.tsx
@@ -15,7 +15,7 @@ import * as React from 'react';
 import { v4 as uuidv4 } from 'uuid';
 import { ClientUtils, OmitKeys } from '../../../../../ClientUtils';
 import { Doc, DocListCast, Opt } from '../../../../../fields/Doc';
-import { DocData, DocLayout, DocViews } from '../../../../../fields/DocSymbols';
+import { DocData, DocViews } from '../../../../../fields/DocSymbols';
 import { Id } from '../../../../../fields/FieldSymbols';
 import { RichTextField } from '../../../../../fields/RichTextField';
 import { ScriptField } from '../../../../../fields/ScriptField';
@@ -44,7 +44,6 @@ import './ChatBox.scss';
 import MessageComponentBox from './MessageComponent';
 import { OpenWhere } from '../../OpenWhere';
 import { Upload } from '../../../../../server/SharedMediaTypes';
-import { DocumentMetadataTool } from '../tools/DocumentMetadataTool';
 import { AgentDocumentManager } from '../utils/AgentDocumentManager';
 import { AiOutlineSend } from 'react-icons/ai';
 import { SnappingManager } from '../../../../util/SnappingManager';
@@ -88,12 +87,12 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
 
     // Private properties for managing OpenAI API, vector store, agent, and UI elements
     private openai!: OpenAI; // Using definite assignment assertion
-    private vectorstore_id: string;
-    private vectorstore: Vectorstore;
-    private agent: Agent;
-    private messagesRef: React.RefObject<HTMLDivElement>;
+    private vectorstore_id: string | undefined;
+    private vectorstore: Vectorstore | undefined;
+    private agent: Agent | undefined;
+    private messagesRef: React.RefObject<HTMLDivElement> = React.createRef();
     private _textInputRef: HTMLInputElement | undefined | null;
-    private docManager: AgentDocumentManager;
+    private docManager: AgentDocumentManager | undefined;
 
     /**
      * Static method that returns the layout string for the field.
@@ -109,8 +108,8 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
 
     @action
     toggleCanvasMode = () => {
-        const newMode = !this.docManager.getCanvasMode();
-        this.docManager.setCanvasMode(newMode);
+        const newMode = !this.docManager?.getCanvasMode();
+        this.docManager?.setCanvasMode(newMode);
     };
 
     /**
@@ -121,69 +120,6 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     constructor(props: FieldViewProps) {
         super(props);
         makeObservable(this);
-
-        // At mount time, find the DocumentView whose .Document is the collection container.
-        const parentView = DocumentView.Selected().lastElement();
-        if (!parentView) {
-            console.warn('GPT ChatBox not inside a DocumentView – cannot sort.');
-        }
-
-        this.messagesRef = React.createRef();
-        this.docManager = new AgentDocumentManager(this, parentView);
-
-        // Initialize OpenAI client
-        this.initializeOpenAI();
-
-        // Create a unique vectorstore ID for this ChatBox
-        this.vectorstore_id = uuidv4();
-
-        // Initialize vectorstore with the document manager
-        this.vectorstore = new Vectorstore(this.vectorstore_id, this.docManager);
-
-        // Create an agent with the vectorstore
-        this.agent = new Agent(this.vectorstore, this.retrieveFormattedHistory.bind(this), this.retrieveCSVData.bind(this), this.createImageInDash.bind(this), this.createCSVInDash.bind(this), this.docManager);
-
-        // Set up the tool created callback
-        this.agent.setToolCreatedCallback(this.handleToolCreated);
-
-        // Add event listeners
-        this.addScrollListener();
-
-        // Reaction to update dataDoc when chat history changes
-        reaction(
-            () =>
-                this._history.map((msg: AssistantMessage) => ({
-                    role: msg.role,
-                    content: msg.content,
-                    follow_up_questions: msg.follow_up_questions,
-                    citations: msg.citations,
-                })),
-            serializableHistory => {
-                this.dataDoc.data = JSON.stringify(serializableHistory);
-            }
-        );
-
-        /*
-                reaction(
-                    () => ({ selDoc: DocumentView.Selected().lastElement(), visible: SnappingManager.ChatVisible }),
-                    ({ selDoc, visible }) => {
-                        const hasChildDocs = visible && selDoc?.ComponentView?.hasChildDocs;
-                        if (hasChildDocs) {
-                            this._textToDocMap.clear();
-                            this.setCollectionContext(selDoc.Document);
-                            this.onGptResponse = (sortResult: string, questionType: GPTDocCommand) => this.processGptResponse(selDoc, this._textToDocMap, sortResult, questionType);
-                            this.onQuizRandom = () => this.randomlyChooseDoc(selDoc.Document, hasChildDocs());
-                            this._documentDescriptions = Promise.all(hasChildDocs().map(doc =>
-                                    Doc.getDescription(doc).then(text => this._textToDocMap.set(text.replace(/\n/g, ' ').trim(), doc) && `${DescriptionSeperator}${text}${DescriptionSeperator}`)
-                                )).then(docDescriptions => docDescriptions.join()); // prettier-ignore
-                        }
-                    },
-                    { fireImmediately: true }
-                );
-            }*/
-
-        // Initialize font size from saved preference
-        this.initFontSize();
     }
 
     /**
@@ -203,6 +139,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                 this._currentStep = isAudioOrVideo ? 'Preparing media file...' : 'Processing document...';
             });
 
+            if (!this.docManager || !this.vectorstore) throw new Error('Document manager or vectorstore is not initialized');
             // Process the document first to ensure it has a valid ID
             await this.docManager.processDocument(newLinkedDoc);
 
@@ -337,7 +274,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
      * Adds a scroll event listener to detect user scrolling and handle passive wheel events.
      */
     addScrollListener = () => {
-        if (this.messagesRef.current) {
+        if (this.messagesRef?.current) {
             this.messagesRef.current.addEventListener('wheel', this.onPassiveWheel, { passive: false });
         }
     };
@@ -346,7 +283,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
      * Removes the scroll event listener from the chat messages container.
      */
     removeScrollListener = () => {
-        if (this.messagesRef.current) {
+        if (this.messagesRef?.current) {
             this.messagesRef.current.removeEventListener('wheel', this.onPassiveWheel);
         }
     };
@@ -428,6 +365,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                     });
                 };
 
+                if (!this.agent) throw new Error('Agent is not initialized');
                 // Send the user's question to the assistant and get the final message
                 const finalMessage = await this.agent.askAgent(trimmedText, onProcessingUpdate, onAnswerUpdate);
 
@@ -537,7 +475,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                 case supportedDocTypes.image:      return Docs.Create.ImageDocument(data as string, options);
                 case supportedDocTypes.equation:   return Docs.Create.EquationDocument(data as string, options);
                 case supportedDocTypes.notetaking: return Docs.Create.NoteTakingDocument([], options);
-                case supportedDocTypes.web:        
+                case supportedDocTypes.web:        {
                     // Create web document with enhanced safety options
                     const webOptions = { 
                         ...options, 
@@ -546,10 +484,11 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                     
                     // If iframe_sandbox was passed from AgentDocumentManager, add it to the options
                     if ('_iframe_sandbox' in options) {
-                        (webOptions as any)._iframe_sandbox = options._iframe_sandbox;
+                        webOptions._iframe_sandbox = options._iframe_sandbox;
                     }
                     
                     return Docs.Create.WebDocument(data as string, webOptions);
+                }
                 case supportedDocTypes.dataviz: case supportedDocTypes.table: return Docs.Create.DataVizDocument('/Users/ajshul/Dash-Web/src/server/public/files/csv/0d237e7c-98c9-44d0-aa61-5285fdbcf96c-random_sample.csv.csv', options);
                 case supportedDocTypes.pdf:        return Docs.Create.PdfDocument(data as string, options);
                 case supportedDocTypes.video:      return Docs.Create.VideoDocument(data as string, options);
@@ -688,6 +627,8 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                 citation: JSON.stringify(citation, null, 2),
             });
 
+            if (!this.docManager) throw new Error('Document manager is not initialized');
+
             // Get the simplified chunk using the document manager
             const { foundChunk, doc, dataDoc } = this.docManager.getSimplifiedChunkById(chunkId);
             console.log('doc: ', doc);
@@ -722,7 +663,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                 }
                 this.handleOtherChunkTypes(foundChunk, citation, doc, dataDoc);
                 // Show the chunk text in citation popup
-                let chunkText = citation.direct_text || 'Text content not available';
+                const chunkText = citation.direct_text || 'Text content not available';
                 this.showCitationPopup(chunkText);
 
                 // Also navigate to the document
@@ -741,7 +682,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
      * @returns The starting timestamp of the matching segment, or -1 if not found
      */
     getDirectMatchingSegmentStart = (doc: Doc, citationText: string, indexesOfSegments: string[]): number => {
-        if (!doc || !citationText) return -1;
+        if (!doc || !citationText || !this.docManager) return -1;
 
         // Get original segments using document manager
         const original_segments = this.docManager.getOriginalSegments(doc);
@@ -754,7 +695,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
 
         // If specific indexes are provided, filter segments by those indexes
         if (indexesOfSegments && indexesOfSegments.length > 0) {
-            segments = original_segments.filter((segment: any) => indexesOfSegments.includes(segment.index));
+            segments = original_segments.filter(segment => indexesOfSegments.includes(segment.index));
         }
 
         // If no segments match the indexes, use all segments
@@ -763,7 +704,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         }
 
         // First try to find an exact match
-        const exactMatch = segments.find((segment: any) => segment.text && segment.text.includes(citationText));
+        const exactMatch = segments.find(segment => segment.text && segment.text.includes(citationText));
 
         if (exactMatch) {
             return exactMatch.start;
@@ -880,14 +821,16 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                 }
                 break;
             case CHUNK_TYPE.TEXT:
-                this._citationPopup = { text: citation.direct_text ?? 'No text available', visible: true };
-                this.startCitationPopupTimer();
+                {
+                    this._citationPopup = { text: citation.direct_text ?? 'No text available', visible: true };
+                    this.startCitationPopupTimer();
 
-                // Check if the document is a PDF (has a PDF viewer component)
-                const isPDF = PDFCast(dataDoc!.data) !== null || dataDoc!.type === DocumentType.PDF;
+                    // Check if the document is a PDF (has a PDF viewer component)
+                    const isPDF = PDFCast(dataDoc!.data) !== null || dataDoc!.type === DocumentType.PDF;
 
-                // First ensure document is fully visible before trying to access its views
-                this.ensureDocumentIsVisible(dataDoc!, isPDF, citation, foundChunk, doc);
+                    // First ensure document is fully visible before trying to access its views
+                    this.ensureDocumentIsVisible(dataDoc!, isPDF, citation, foundChunk, doc);
+                }
                 break;
             case CHUNK_TYPE.CSV:
             case CHUNK_TYPE.URL:
@@ -1079,6 +1022,65 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
      * Initializes scroll listeners, sets up document reactions, and loads chat history from dataDoc if available.
      */
     componentDidMount() {
+        // At mount time, find the DocumentView whose .Document is the collection container.
+        const parentView = this.DocumentView?.()._props.containerViewPath?.().lastElement();
+        if (!parentView) {
+            console.warn('GPT ChatBox not inside a DocumentView – cannot perform operations on Documents.');
+            return;
+        }
+
+        this.docManager = new AgentDocumentManager(this, parentView);
+
+        // Initialize OpenAI client
+        this.initializeOpenAI();
+
+        // Create a unique vectorstore ID for this ChatBox
+        this.vectorstore_id = uuidv4();
+
+        // Initialize vectorstore with the document manager
+        this.vectorstore = new Vectorstore(this.vectorstore_id, this.docManager);
+
+        /*
+                reaction(
+                    () => ({ selDoc: DocumentView.Selected().lastElement(), visible: SnappingManager.ChatVisible }),
+                    ({ selDoc, visible }) => {
+                        const hasChildDocs = visible && selDoc?.ComponentView?.hasChildDocs;
+                        if (hasChildDocs) {
+                            this._textToDocMap.clear();
+                            this.setCollectionContext(selDoc.Document);
+                            this.onGptResponse = (sortResult: string, questionType: GPTDocCommand) => this.processGptResponse(selDoc, this._textToDocMap, sortResult, questionType);
+                            this.onQuizRandom = () => this.randomlyChooseDoc(selDoc.Document, hasChildDocs());
+                            this._documentDescriptions = Promise.all(hasChildDocs().map(doc =>
+                                    Doc.getDescription(doc).then(text => this._textToDocMap.set(text.replace(/\n/g, ' ').trim(), doc) && `${DescriptionSeperator}${text}${DescriptionSeperator}`)
+                                )).then(docDescriptions => docDescriptions.join()); // prettier-ignore
+                        }
+                    },
+                    { fireImmediately: true }
+                );
+            }*/
+
+        // Initialize font size from saved preference
+        this.initFontSize();
+        // Create an agent with the vectorstore
+        this.agent = new Agent(this.vectorstore, this.retrieveFormattedHistory.bind(this), this.retrieveCSVData.bind(this), this.createImageInDash.bind(this), this.createCSVInDash.bind(this), this.docManager);
+
+        // Set up the tool created callback
+        this.agent.setToolCreatedCallback(this.handleToolCreated);
+
+        // Reaction to update dataDoc when chat history changes
+        reaction(
+            () =>
+                this._history.map((msg: AssistantMessage) => ({
+                    role: msg.role,
+                    content: msg.content,
+                    follow_up_questions: msg.follow_up_questions,
+                    citations: msg.citations,
+                })),
+            serializableHistory => {
+                this.dataDoc.data = JSON.stringify(serializableHistory);
+            }
+        );
+
         this._props.setContentViewBox?.(this);
         if (this.dataDoc.data) {
             try {
@@ -1146,9 +1148,12 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         // If there are stored doc IDs in our list of docs to add, process them
         if (this._linked_docs_to_add.size > 0) {
             this._linked_docs_to_add.forEach(async doc => {
-                await this.docManager.processDocument(doc);
+                await this.docManager!.processDocument(doc);
             });
         }
+
+        // Add event listeners
+        this.addScrollListener();
     }
 
     /**
@@ -1235,7 +1240,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
 
         try {
             // Perform the deferred tool save operation
-            const saveSuccess = await this.agent.performDeferredToolSave();
+            const saveSuccess = await this.agent?.performDeferredToolSave();
 
             if (saveSuccess) {
                 console.log('Tool saved successfully, proceeding with reload...');
@@ -1447,19 +1452,9 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                 <div className="chat-header">
                     <h2>{this.userName()}&apos;s AI Assistant</h2>
                     <div className="header-controls">
-                        <Tooltip title={
-                            <div className="dash-tooltip">
-                                {this.docManager.getCanvasMode() 
-                                    ? "Click to limit scope to linked documents"
-                                    : "Click to expand scope to all documents on canvas"
-                                }
-                            </div>
-                        } placement="bottom">
-                            <div 
-                                className={`canvas-mode-toggle ${this.docManager.getCanvasMode() ? 'canvas-active' : ''}`}
-                                onClick={this.toggleCanvasMode}
-                            >
-                                {this.docManager.getCanvasMode() ? <MdViewModule /> : <MdLink />}
+                        <Tooltip title={<div className="dash-tooltip">{this.docManager?.getCanvasMode() ? 'Click to limit scope to linked documents' : 'Click to expand scope to all documents on canvas'}</div>} placement="bottom">
+                            <div className={`canvas-mode-toggle ${this.docManager?.getCanvasMode() ? 'canvas-active' : ''}`} onClick={this.toggleCanvasMode}>
+                                {this.docManager?.getCanvasMode() ? <MdViewModule /> : <MdLink />}
                             </div>
                         </Tooltip>
                         <div className="font-size-control" onClick={this.toggleFontSizeModal}>