1 files changed, 67 insertions, 19 deletions

diff --git a/src/client/views/nodes/ChatBox/ChatBox.tsx b/src/client/views/nodes/ChatBox/ChatBox.tsx
index a47e9a95b..5d0a16b4f 100644
--- a/src/client/views/nodes/ChatBox/ChatBox.tsx
+++ b/src/client/views/nodes/ChatBox/ChatBox.tsx
@@ -63,11 +63,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
             this.vectorstore_id = StrCast(this.dataDoc.vectorstore_id);
         }
         this.vectorstore = new Vectorstore(this.vectorstore_id);
-        this.agent = new Agent(
-            this.vectorstore,
-            () => this.summaries,
-            () => this.formattedHistory
-        );
+        this.agent = new Agent(this.vectorstore, this.retrieveSummaries, this.retrieveFormattedHistory);
 
         reaction(
             () => this.history.map((msg: AssistantMessage) => ({ role: msg.role, text_content: msg.text_content, follow_up_questions: msg.follow_up_questions, citations: msg.citations })),
@@ -162,18 +158,66 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         const currentLinkedDocs: Doc[] = this.linkedDocs;
         const chunk_id = citation.chunk_id;
         for (let doc of currentLinkedDocs) {
-            const doc_chunk_ids: string[] = JSON.parse(StrCast(doc.chunk_ids));
-            if (!doc_chunk_ids.includes(chunk_id)) continue;
-            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);
-            });
+            console.log(JSON.parse(StrCast(doc.chunk_simpl)));
+            const doc_chunk_simpl = JSON.parse(StrCast(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;
+
+                // Extract variables from location_string
+                const values = location_string.replace(/[\[\]]/g, '').split(',');
+
+                // 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
+                }
+
+                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);
+
+                // Parse values to numbers
+                // const [x1, y1, x2, y2] = values.map(Number);
+
+                // 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 highlight_doc = Docs.Create.FreeformDocument([], {
+                    x: x1,
+                    y: y1,
+                    _width: x2 - x1,
+                    _height: y2 - y1,
+                    backgroundColor: 'rgba(255, 255, 0, 0.5)',
+                });
+
+                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 }, () => {});
+            }
         }
         // You can implement additional functionality here, such as showing a modal with the full citation content
     };
@@ -296,9 +340,13 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         return history;
     }
 
-    retrieveSummaries(): string {
+    retrieveSummaries = () => {
         return this.summaries;
-    }
+    };
+
+    retrieveFormattedHistory = () => {
+        return this.formattedHistory;
+    };
 
     @action
     handleFollowUpClick = (question: string) => {