improved consistency across doc types and parsing

1 files changed, 98 insertions, 78 deletions

diff --git a/src/client/views/nodes/chatbot/chatboxcomponents/ChatBox.tsx b/src/client/views/nodes/chatbot/chatboxcomponents/ChatBox.tsx
index b11bf7405..ba30cb42b 100644
--- a/src/client/views/nodes/chatbot/chatboxcomponents/ChatBox.tsx
+++ b/src/client/views/nodes/chatbot/chatboxcomponents/ChatBox.tsx
@@ -18,7 +18,7 @@ import { Doc, DocListCast, Opt } from '../../../../../fields/Doc';
 import { DocData, DocViews } from '../../../../../fields/DocSymbols';
 import { RichTextField } from '../../../../../fields/RichTextField';
 import { ScriptField } from '../../../../../fields/ScriptField';
-import { CsvCast, DocCast, NumCast, PDFCast, RTFCast, StrCast } from '../../../../../fields/Types';
+import { CsvCast, DocCast, NumCast, PDFCast, RTFCast, StrCast, VideoCast, AudioCast } from '../../../../../fields/Types';
 import { DocUtils } from '../../../../documents/DocUtils';
 import { CollectionViewType, DocumentType } from '../../../../documents/DocumentTypes';
 import { Docs, DocumentOptions } from '../../../../documents/Documents';
@@ -48,7 +48,14 @@ import { AgentDocumentManager } from '../utils/AgentDocumentManager';
 
 dotenv.config();
 
-export type parsedDocData = { doc_type: string; data: unknown };
+export type parsedDocData = {
+    doc_type: string;
+    data: unknown;
+    _disable_resource_loading?: boolean;
+    _sandbox_iframe?: boolean;
+    _iframe_sandbox?: string;
+    data_useCors?: boolean;
+};
 export type parsedDoc = DocumentOptions & parsedDocData;
 /**
  * ChatBox is the main class responsible for managing the interaction between the user and the assistant,
@@ -150,7 +157,14 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     @action
     addDocToVectorstore = async (newLinkedDoc: Doc) => {
         try {
-            this._isUploadingDocs = true;
+            const isAudioOrVideo = VideoCast(newLinkedDoc.data)?.url?.pathname || AudioCast(newLinkedDoc.data)?.url?.pathname;
+
+            // Set UI state to show the processing overlay
+            runInAction(() => {
+                this._isUploadingDocs = true;
+                this._uploadProgress = 0;
+                this._currentStep = isAudioOrVideo ? 'Preparing media file...' : 'Processing document...';
+            });
 
             // Process the document first to ensure it has a valid ID
             this.docManager.processDocument(newLinkedDoc);
@@ -158,15 +172,36 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
             // Add the document to the vectorstore which will also register chunks
             await this.vectorstore.addAIDoc(newLinkedDoc, this.updateProgress);
 
-            // No longer needed as documents are tracked by the AgentDocumentManager
-            // this._linked_docs_to_add.add(newLinkedDoc);
+            // Give a slight delay to show the completion message
+            if (this._uploadProgress === 100) {
+                await new Promise(resolve => setTimeout(resolve, 1000));
+            }
 
-            this._isUploadingDocs = false;
+            // Reset UI state
+            runInAction(() => {
+                this._isUploadingDocs = false;
+                this._uploadProgress = 0;
+                this._currentStep = '';
+            });
 
             return true;
         } catch (err) {
             console.error('Error adding document to vectorstore:', err);
-            this._isUploadingDocs = false;
+
+            // Show error in UI
+            runInAction(() => {
+                this._currentStep = `Error: ${err instanceof Error ? err.message : 'Failed to process document'}`;
+            });
+
+            await new Promise(resolve => setTimeout(resolve, 2000));
+
+            // Reset UI state
+            runInAction(() => {
+                this._isUploadingDocs = false;
+                this._uploadProgress = 0;
+                this._currentStep = '';
+            });
+
             return false;
         }
     };
@@ -178,8 +213,15 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
      */
     @action
     updateProgress = (progress: number, step: string) => {
-        this._uploadProgress = progress;
+        // Ensure progress is within expected bounds
+        const validProgress = Math.min(Math.max(0, progress), 100);
+        this._uploadProgress = validProgress;
         this._currentStep = step;
+
+        // Force UI update
+        if (process.env.NODE_ENV !== 'production') {
+            console.log(`Progress: ${validProgress}%, Step: ${step}`);
+        }
     };
 
     /**
@@ -453,7 +495,19 @@ 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:        return Docs.Create.WebDocument(data as string, { ...options, data_useCors: true });
+                case supportedDocTypes.web:        
+                    // Create web document with enhanced safety options
+                    const webOptions = { 
+                        ...options, 
+                        data_useCors: true
+                    };
+                    
+                    // 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;
+                    }
+                    
+                    return Docs.Create.WebDocument(data as string, webOptions);
                 case supportedDocTypes.dataviz:    return Docs.Create.DataVizDocument('/users/rz/Downloads/addresses.csv', options);
                 case supportedDocTypes.pdf:        return Docs.Create.PdfDocument(data as string, options);
                 case supportedDocTypes.video:      return Docs.Create.VideoDocument(data as string, options);
@@ -607,65 +661,36 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                 return;
             }
 
-            // Process the chunk data
-            let docChunkSimpl: { chunks: SimplifiedChunk[] } = { chunks: [] };
-            try {
-                docChunkSimpl = JSON.parse(StrCast(doc.chunk_simpl) || '{"chunks":[]}');
-            } catch (e) {
-                console.error(`Error parsing chunk_simpl for the found document:`, e);
+            // Get the simplified chunk using the document manager
+            const foundChunk = this.docManager.getSimplifiedChunkById(doc, chunkId);
+            if (!foundChunk) {
+                console.warn(`Chunk not found in document for chunk ID: ${chunkId}`);
+                DocumentManager.Instance.showDocument(doc, { willZoomCentered: true }, () => {});
                 return;
             }
 
-            const foundChunk = docChunkSimpl.chunks.find((chunk: SimplifiedChunk) => chunk.chunkId === chunkId);
+            console.log(`Found chunk in document:`, foundChunk);
 
             // Handle different chunk types
-            if (foundChunk) {
-                console.log(`Found chunk in document:`, foundChunk);
-                if (foundChunk.chunkType === CHUNK_TYPE.AUDIO || foundChunk.chunkType === CHUNK_TYPE.VIDEO) {
-                    const directMatchSegmentStart = this.getDirectMatchingSegmentStart(doc, citation.direct_text || '', foundChunk.indexes || []);
-                    if (directMatchSegmentStart) {
-                        await this.goToMediaTimestamp(doc, directMatchSegmentStart, foundChunk.chunkType);
-                    } else {
-                        console.error('No direct matching segment found for the citation.');
-                    }
-                } else if (foundChunk.chunkType === CHUNK_TYPE.TABLE || foundChunk.chunkType === CHUNK_TYPE.IMAGE) {
-                    this.handleOtherChunkTypes(foundChunk, citation, doc);
-                } else if (foundChunk.chunkType === CHUNK_TYPE.TEXT) {
-                    // Find text from the document's chunks metadata
-                    let chunkText = '';
-
-                    try {
-                        // We already parsed the chunks earlier, so use that
-                        const matchingChunk = docChunkSimpl.chunks.find(c => c.chunkId === foundChunk.chunkId);
-                        if (matchingChunk && 'text' in matchingChunk) {
-                            // If the text property exists on the chunk (even though it's not in the type)
-                            chunkText = String(matchingChunk['text'] || '');
-                        }
-                    } catch (e) {
-                        console.error('Error getting chunk text:', e);
-                    }
-
-                    // Default text if none found
-                    if (!chunkText) {
-                        chunkText = 'Text content not available';
-                    }
-
-                    this._citationPopup = {
-                        text: chunkText,
-                        visible: true,
-                    };
-                }
-                // Handle URL chunks
-                else if (foundChunk.chunkType === CHUNK_TYPE.URL) {
-                    if (foundChunk.url) {
-                        DocumentManager.Instance.showDocument(doc, { willZoomCentered: true }, () => {});
-                        console.log(`Navigated to web document with URL: ${foundChunk.url}`);
-                    } else {
-                        console.warn('URL chunk missing URL:', foundChunk);
-                    }
+            if (foundChunk.chunkType === CHUNK_TYPE.AUDIO || foundChunk.chunkType === CHUNK_TYPE.VIDEO) {
+                const directMatchSegmentStart = this.getDirectMatchingSegmentStart(doc, citation.direct_text || '', foundChunk.indexes || []);
+                if (directMatchSegmentStart) {
+                    await this.goToMediaTimestamp(doc, directMatchSegmentStart, foundChunk.chunkType);
+                } else {
+                    console.error('No direct matching segment found for the citation.');
                 }
+            } else if (foundChunk.chunkType === CHUNK_TYPE.TABLE || foundChunk.chunkType === CHUNK_TYPE.IMAGE) {
+                this.handleOtherChunkTypes(foundChunk, citation, doc);
             } else {
-                console.warn('Navigating to doc. Unable to find chunk or segments for citation', citation);
+                // Show the chunk text in citation popup
+                let chunkText = foundChunk.text || 'Text content not available';
+
+                this._citationPopup = {
+                    text: chunkText,
+                    visible: true,
+                };
+
+                // Also navigate to the document
                 DocumentManager.Instance.showDocument(doc, { willZoomCentered: true }, () => {});
             }
         } catch (error) {
@@ -683,8 +708,8 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     getDirectMatchingSegmentStart = (doc: Doc, citationText: string, indexesOfSegments: string[]): number => {
         if (!doc || !citationText) return -1;
 
-        // Get original segments from the document
-        const original_segments = doc.original_segments ? JSON.parse(StrCast(doc.original_segments)) : [];
+        // Get original segments using document manager
+        const original_segments = this.docManager.getOriginalSegments(doc);
 
         if (!original_segments || !Array.isArray(original_segments) || original_segments.length === 0) {
             return -1;
@@ -993,18 +1018,8 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
      */
     @computed
     get summaries(): string {
-        const linkedDocs = Array.from(this.docManager.listDocs())
-            .map(id => {
-                const doc = this.docManager.extractDocumentMetadata(id);
-                if (doc && doc.fields && (doc.fields.layout.summary || doc.fields.data.summary)) {
-                    return doc.fields.layout.summary || doc.fields.data.summary;
-                }
-                return null;
-            })
-            .filter(Boolean)
-            .join('\n\n');
-
-        return linkedDocs;
+        // Use the document manager to get all summaries
+        return this.docManager.getAllDocumentSummaries();
     }
 
     /**
@@ -1033,7 +1048,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     // Other helper methods for retrieving document data and processing
 
     retrieveSummaries = (): string => {
-        return this.summaries;
+        return this.docManager.getAllDocumentSummaries();
     };
 
     retrieveCSVData = () => {
@@ -1068,8 +1083,13 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                 {this._isUploadingDocs && (
                     <div className="uploading-overlay">
                         <div className="progress-container">
-                            <ProgressBar />
-                            <div className="step-name">{this._currentStep}</div>
+                            <div className="progress-bar-wrapper">
+                                <div className="progress-bar" style={{ width: `${this._uploadProgress}%` }} />
+                            </div>
+                            <div className="progress-details">
+                                <div className="progress-percentage">{Math.round(this._uploadProgress)}%</div>
+                                <div className="step-name">{this._currentStep}</div>
+                            </div>
                         </div>
                     </div>
                 )}