From 65179e8b0519aa4ccf28afc4c429262ecf7a62f3 Mon Sep 17 00:00:00 2001
From: "A.J. Shulman" <Shulman.aj@gmail.com>
Date: Tue, 16 Jul 2024 11:11:53 -0400
Subject: attempting annotations for image chunks

---
 src/client/views/nodes/ChatBox/ChatBox.tsx | 111 ++++++++++++++++++-----------
 1 file changed, 71 insertions(+), 40 deletions(-)

(limited to 'src/client/views/nodes/ChatBox/ChatBox.tsx')

diff --git a/src/client/views/nodes/ChatBox/ChatBox.tsx b/src/client/views/nodes/ChatBox/ChatBox.tsx
index 4d1cd38a0..a47e9a95b 100644
--- a/src/client/views/nodes/ChatBox/ChatBox.tsx
+++ b/src/client/views/nodes/ChatBox/ChatBox.tsx
@@ -1,4 +1,4 @@
-import { action, computed, makeObservable, observable, observe, reaction, runInAction } from 'mobx';
+import { action, computed, makeObservable, observable, observe, reaction, runInAction, ObservableSet } from 'mobx';
 import { observer } from 'mobx-react';
 import OpenAI, { ClientOptions } from 'openai';
 import * as React from 'react';
@@ -25,6 +25,7 @@ import { DocumentManager } from '../../../util/DocumentManager';
 import { UUID } from 'bson';
 import { v4 as uuidv4 } from 'uuid';
 import { aS } from '@fullcalendar/core/internal-common';
+import { computeRect } from '@fullcalendar/core/internal';
 
 dotenv.config();
 
@@ -34,10 +35,10 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     @observable.deep current_message: AssistantMessage | undefined = undefined;
 
     @observable isLoading: boolean = false;
-    @observable isInitializing: boolean = true;
+    @observable isInitializing: boolean = false;
     @observable expandedScratchpadIndex: number | null = null;
     @observable inputValue: string = '';
-    @observable private linked_docs_to_add: Doc[] = [];
+    @observable private linked_docs_to_add: ObservableSet<Doc> = observable.set();
     private openai: OpenAI;
     private vectorstore_id: string;
     private documents: AI_Document[] = [];
@@ -55,13 +56,19 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         this.history = [{ role: ASSISTANT_ROLE.ASSISTANT, text_content: 'Welcome to the Document Analyser Assistant! Link a document or ask questions to get started.' }];
         this.openai = this.initializeOpenAI();
         if (StrCast(this.dataDoc.vectorstore_id) == '') {
+            console.log('new_id');
             this.vectorstore_id = uuidv4();
             this.dataDoc.vectorstore_id = this.vectorstore_id;
         } else {
             this.vectorstore_id = StrCast(this.dataDoc.vectorstore_id);
         }
         this.vectorstore = new Vectorstore(this.vectorstore_id);
-        this.agent = new Agent(this.vectorstore); // Initialize the Agent
+        this.agent = new Agent(
+            this.vectorstore,
+            () => this.summaries,
+            () => this.formattedHistory
+        );
+
         reaction(
             () => this.history.map((msg: AssistantMessage) => ({ role: msg.role, text_content: msg.text_content, follow_up_questions: msg.follow_up_questions, citations: msg.citations })),
             serializableHistory => {
@@ -70,8 +77,11 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         );
     }
 
-    addDocsToVectorstore = async (linkedDocs: Doc[]) => {
-        await this.vectorstore.addAIDocs(linkedDocs);
+    @action
+    addDocToVectorstore = async (newLinkedDoc: Doc) => {
+        this.isInitializing = true;
+        await this.vectorstore.addAIDoc(newLinkedDoc);
+        this.isInitializing = false;
     };
 
     // @action
@@ -152,20 +162,18 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         const currentLinkedDocs: Doc[] = this.linkedDocs;
         const chunk_id = citation.chunk_id;
         for (let doc of currentLinkedDocs) {
-            const doc_chunks: Chunk[] = JSON.parse(StrCast(doc.ai_document)).chunks;
-            const chunk_file_name = doc_chunks.find(chunk => chunk.id === chunk_id)?.metadata.file_path;
+            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 + ' Citation URL: ' + chunk_file_name);
+            console.log('URL: ' + doc_url);
             //const ai_field_id = doc[this.Document[Id] + '_ai_field_id'];
-            if (chunk_file_name == doc_url) {
-                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);
-                });
-            }
+            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);
+            });
         }
         // You can implement additional functionality here, such as showing a modal with the full citation content
     };
@@ -234,32 +242,29 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                 return linkedDocs;
             },
 
-            linked => this.linked_docs_to_add.push(...linked.filter(linkedDoc => !this.linked_docs_to_add.includes(linkedDoc)))
+            linked => linked.forEach(doc => this.linked_docs_to_add.add(doc))
         );
 
-        observe(
-            // right now this skips during initialization which is necessary because it would be blank
-            // However, it will upload the same link twice when it is
-            this.linked_docs_to_add,
-            change => {
-                // observe pushes/splices on a user link DB 'data' field (should only happen for local changes)
-                switch (change.type as any) {
-                    case 'splice':
-                        if ((change as any).addedCount > 0) {
-                            // maybe check here if its already in the urls datadoc array so doesn't add twice
-                            console.log((change as any).added as Doc[]);
-                            this.addDocsToVectorstore((change as any).added as Doc[]);
-                        }
-                        // (change as any).removed.forEach((link: any) => remLinkFromDoc(toRealField(link)));
-                        break;
-                    case 'update': // let oldValue = change.oldValue;
-                    default:
-                }
-            },
-            true
-        );
+        observe(this.linked_docs_to_add, change => {
+            if (change.type === 'add') {
+                this.addDocToVectorstore(change.newValue);
+            } else if (change.type === 'delete') {
+                console.log('Deleted docs: ', change.oldValue);
+            }
+        });
     }
 
+    // case 'splice':
+    //                     if ((change as any).addedCount > 0) {
+    //                         // maybe check here if its already in the urls datadoc array so doesn't add twice
+    //                         console.log((change as any).added as Doc[]);
+    //                         this.addDocsToVectorstore((change as any).added as Doc[]);
+    //                     }
+    //                     // (change as any).removed.forEach((link: any) => remLinkFromDoc(toRealField(link)));
+    //                     break;
+    //                 case 'update': // let oldValue = change.oldValue;
+    //                 default:
+
     @computed
     get linkedDocs() {
         //return (CollectionFreeFormDocumentView.from(this._props.DocumentView?.())?._props.parent as CollectionFreeFormView)?.childDocs.filter(doc => doc != this.Document) ?? [];
@@ -269,6 +274,32 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
             .filter(d => d);
     }
 
+    @computed
+    get summaries(): string {
+        return (
+            LinkManager.Instance.getAllRelatedLinks(this.Document)
+                .map(d => DocCast(LinkManager.getOppositeAnchor(d, this.Document)))
+                .map(d => DocCast(d?.annotationOn, d))
+                .filter(d => d)
+                .map((doc, index) => `${index + 1}) ${doc.summary}`)
+                .join('\n') + '\n'
+        );
+    }
+
+    @computed
+    get formattedHistory(): string {
+        let history = '<chat_history>\n';
+        for (const message of this.history) {
+            history += `<${message.role}>${message.text_content}</${message.role}>\n`;
+        }
+        history += '</chat_history>';
+        return history;
+    }
+
+    retrieveSummaries(): string {
+        return this.summaries;
+    }
+
     @action
     handleFollowUpClick = (question: string) => {
         console.log('Follow-up question clicked:', question);
-- 
cgit v1.2.3-70-g09d2

