added image citation highlights

1 files changed, 67 insertions, 61 deletions

diff --git a/src/client/views/nodes/ChatBox/ChatBox.tsx b/src/client/views/nodes/ChatBox/ChatBox.tsx
index c7ae9a354..8b4a7bd0a 100644
--- a/src/client/views/nodes/ChatBox/ChatBox.tsx
+++ b/src/client/views/nodes/ChatBox/ChatBox.tsx
@@ -2,7 +2,7 @@ import { action, computed, makeObservable, observable, observe, reaction, runInA
 import { observer } from 'mobx-react';
 import OpenAI, { ClientOptions } from 'openai';
 import * as React from 'react';
-import { Doc } from '../../../../fields/Doc';
+import { Doc, DocListCast } from '../../../../fields/Doc';
 import { CsvCast, DocCast, PDFCast, StrCast } from '../../../../fields/Types';
 import { Networking } from '../../../Network';
 import { DocumentType } from '../../../documents/DocumentTypes';
@@ -12,7 +12,7 @@ import { ViewBoxAnnotatableComponent } from '../../DocComponent';
 import { FieldView, FieldViewProps } from '../FieldView';
 import './ChatBox.scss';
 import MessageComponentBox from './MessageComponent';
-import { ASSISTANT_ROLE, AssistantMessage, AI_Document, convertToAIDocument, Citation, CHUNK_TYPE, Chunk, getChunkType } from './types';
+import { ASSISTANT_ROLE, AssistantMessage, AI_Document, Citation, CHUNK_TYPE, Chunk, getChunkType } from './types';
 import { Vectorstore } from './vectorstore/VectorstoreUpload';
 import { CollectionFreeFormDocumentView } from '../CollectionFreeFormDocumentView';
 import { CollectionFreeFormView } from '../../collections/collectionFreeForm';
@@ -26,6 +26,7 @@ import { UUID } from 'bson';
 import { v4 as uuidv4 } from 'uuid';
 import { aS } from '@fullcalendar/core/internal-common';
 import { computeRect } from '@fullcalendar/core/internal';
+import { DocUtils } from '../../../documents/DocUtils';
 
 dotenv.config();
 
@@ -35,7 +36,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     @observable.deep current_message: AssistantMessage | undefined = undefined;
 
     @observable isLoading: boolean = false;
-    @observable isInitializing: boolean = false;
+    @observable isUploadingDocs: boolean = false;
     @observable expandedScratchpadIndex: number | null = null;
     @observable inputValue: string = '';
     @observable private linked_docs_to_add: ObservableSet<Doc> = observable.set();
@@ -75,9 +76,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
 
     @action
     addDocToVectorstore = async (newLinkedDoc: Doc) => {
-        this.isInitializing = true;
         await this.vectorstore.addAIDoc(newLinkedDoc);
-        this.isInitializing = false;
     };
 
     // @action
@@ -158,75 +157,76 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         const currentLinkedDocs: Doc[] = this.linkedDocs;
         const chunk_id = citation.chunk_id;
         for (let doc of currentLinkedDocs) {
-            //console.log(JSON.parse(StrCast(doc.chunk_simpl)));
-            console.log(JSON.stringify(StrCast(doc.chunk_simpl)));
-            const doc_chunk_simpl = JSON.parse(StrCast(doc.chunk_simpl) as string);
-            console.log(doc_chunk_simpl);
-            const text_chunks = doc_chunk_simpl.text_chunks as [{ chunk_id: string; start_page: number; end_page: number }] | [];
-            const image_chunks = doc_chunk_simpl.image_chunks as [{ chunk_id: string; location: string; page: number; page_width: number; page_height: number }] | [];
-
-            const found_text_chunk = text_chunks.find(chunk => chunk.chunk_id === chunk_id);
-            if (found_text_chunk) {
-                const doc_url = CsvCast(doc.data, PDFCast(doc.data)).url.pathname;
-                console.log('URL: ' + doc_url);
-
-                //const ai_field_id = doc[this.Document[Id] + '_ai_field_id'];
-                DocumentManager.Instance.showDocument(doc, { willZoomCentered: true }, () => {
-                    console.log(doc.data);
-                    //look at context path for each docview and choose the doc view that has as
-                    //its parent the same collection view the chatbox is in
-                    const first_view = Array.from(doc[DocViews])[0];
-                    first_view.ComponentView?.search?.(citation.direct_text);
-                });
-            }
-
-            const found_image_chunk = image_chunks.find(chunk => chunk.chunk_id === chunk_id);
-            if (found_image_chunk) {
-                const location_string: string = found_image_chunk.location;
+            if (doc.chunk_simpl) {
+                //console.log(JSON.parse(StrCast(doc.chunk_simpl)));
+                const doc_chunk_simpl = JSON.parse(StrCast(doc.chunk_simpl));
+                console.log(doc_chunk_simpl);
+                const text_chunks = doc_chunk_simpl.text_chunks as [{ chunk_id: string; start_page: number; end_page: number }] | [];
+                const image_chunks = doc_chunk_simpl.image_chunks as [{ chunk_id: string; location: string; page: number }] | [];
+
+                const found_text_chunk = text_chunks.find(chunk => chunk.chunk_id === chunk_id);
+                if (found_text_chunk) {
+                    const doc_url = CsvCast(doc.data, PDFCast(doc.data)).url.pathname;
+                    console.log('URL: ' + doc_url);
+
+                    //const ai_field_id = doc[this.Document[Id] + '_ai_field_id'];
+                    DocumentManager.Instance.showDocument(doc, { willZoomCentered: true }, () => {
+                        console.log(doc.data);
+                        //look at context path for each docview and choose the doc view that has as
+                        //its parent the same collection view the chatbox is in
+                        const first_view = Array.from(doc[DocViews])[0];
+                        first_view.ComponentView?.search?.(citation.direct_text);
+                    });
+                }
 
-                // Extract variables from location_string
-                const values = location_string.replace(/[\[\]]/g, '').split(',');
+                const found_image_chunk = image_chunks.find(chunk => chunk.chunk_id === chunk_id);
+                if (found_image_chunk) {
+                    const location_string: string = found_image_chunk.location;
 
-                // Ensure we have exactly 4 values
-                if (values.length !== 4) {
-                    console.error('Location string must contain exactly 4 numbers');
-                    return; // or handle this error as appropriate
-                }
+                    // Extract variables from location_string
+                    const values = location_string.replace(/[\[\]]/g, '').split(',');
 
-                const x1 = parseInt(values[0]) * (parseInt(StrCast(doc.width)) / found_image_chunk.page_width);
-                const y1 = parseInt(values[1]) * (parseInt(StrCast(doc.height)) / found_image_chunk.page_height);
-                const x2 = parseInt(values[2]) * (parseInt(StrCast(doc.width)) / found_image_chunk.page_width);
-                const y2 = parseInt(values[3]) * (parseInt(StrCast(doc.height)) / found_image_chunk.page_height);
+                    // Ensure we have exactly 4 values
+                    if (values.length !== 4) {
+                        console.error('Location string must contain exactly 4 numbers');
+                        return; // or handle this error as appropriate
+                    }
 
-                // Parse values to numbers
-                // const [x1, y1, x2, y2] = values.map(Number);
+                    const x1 = parseFloat(values[0]) * Doc.NativeWidth(doc);
+                    const y1 = parseFloat(values[1]) * Doc.NativeHeight(doc);
+                    const x2 = parseFloat(values[2]) * Doc.NativeWidth(doc);
+                    const y2 = parseFloat(values[3]) * Doc.NativeHeight(doc);
 
-                // Check if any parsing resulted in NaN
-                if ([x1, y1, x2, y2].some(isNaN)) {
-                    console.error('All values in location string must be valid numbers');
-                    return; // or handle this error as appropriate
-                }
+                    const annotationKey = Doc.LayoutFieldKey(doc) + '_annotations';
 
-                const highlight_doc = Docs.Create.FreeformDocument([], {
-                    x: x1,
-                    y: y1,
-                    _width: x2 - x1,
-                    _height: y2 - y1,
-                    backgroundColor: 'rgba(255, 255, 0, 0.5)',
-                });
+                    const existingDoc = DocListCast(doc[DocData][annotationKey]).find(d => d.citation_id === citation.citation_id);
+                    const highlight_doc = existingDoc ?? this.createImageCitationHighlight(x1, y1, x2, y2, citation, annotationKey, doc);
 
-                Doc.AddDocToList(doc[DocData], Doc.LayoutFieldKey(doc) + '_annotations', highlight_doc);
-                highlight_doc.annotationOn = doc;
-                Doc.SetContainer(highlight_doc, doc);
-                DocumentManager.Instance.showDocument(highlight_doc, { willZoomCentered: true }, () => {});
+                    DocumentManager.Instance.showDocument(highlight_doc, { willZoomCentered: true }, () => {});
+                }
             }
         }
         // You can implement additional functionality here, such as showing a modal with the full citation content
     };
 
+    createImageCitationHighlight = (x1: number, y1: number, x2: number, y2: number, citation: Citation, annotationKey: string, pdfDoc: Doc): Doc => {
+        const highlight_doc = Docs.Create.FreeformDocument([], {
+            x: x1,
+            y: y1,
+            _width: x2 - x1,
+            _height: y2 - y1,
+            backgroundColor: 'rgba(255, 255, 0, 0.5)',
+        });
+        highlight_doc[DocData].citation_id = citation.citation_id;
+        Doc.AddDocToList(pdfDoc[DocData], annotationKey, highlight_doc);
+        highlight_doc.annotationOn = pdfDoc;
+        Doc.SetContainer(highlight_doc, pdfDoc);
+        return highlight_doc;
+    };
+
     // @action
     // uploadLinks = async (linkedDocs: Doc[]) => {
-    //     if (this.isInitializing) {
+    //     if (this.isUploadingDocs) {
     //         console.log('Initialization in progress, upload aborted.');
     //         return;
     //     }
@@ -293,7 +293,13 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
 
         observe(this.linked_docs_to_add, change => {
             if (change.type === 'add') {
+                runInAction(() => {
+                    this.isUploadingDocs = true;
+                });
                 this.addDocToVectorstore(change.newValue);
+                runInAction(() => {
+                    this.isUploadingDocs = false;
+                });
             } else if (change.type === 'delete') {
                 console.log('Deleted docs: ', change.oldValue);
             }
@@ -358,7 +364,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     render() {
         return (
             <div className="chatBox">
-                {this.isInitializing && <div className="initializing-overlay">Initializing...</div>}
+                {this.isUploadingDocs && <div className="uploading-overlay"></div>}
                 <div
                     className="scroll-box chat-content"
                     ref={r => {