attempt at adding links

4 files changed, 130 insertions, 102 deletions

diff --git a/src/client/views/nodes/ChatBox/ChatBox.scss b/src/client/views/nodes/ChatBox/ChatBox.scss
index f1e3d3d67..75171fe56 100644
--- a/src/client/views/nodes/ChatBox/ChatBox.scss
+++ b/src/client/views/nodes/ChatBox/ChatBox.scss
@@ -5,6 +5,10 @@ $button-color: #007bff;
 $button-hover-color: darken($button-color, 10%);
 $shadow-color: rgba(0, 0, 0, 0.075);
 $border-radius: 8px;
+$citation-color: #ff6347;
+$citation-hover-color: darken($citation-color, 10%);
+$follow-up-bg-color: #e9ecef;
+$follow-up-hover-bg-color: #dee2e6;
 
 .chatBox {
     display: flex;
@@ -40,51 +44,41 @@ $border-radius: 8px;
         .messages {
             display: flex;
             flex-direction: column;
+
             .message {
-                padding: 10px;
+                padding: 10px 15px;
                 margin-bottom: 10px;
                 border-radius: $border-radius;
                 background-color: lighten($background-color, 5%);
                 box-shadow: 0 2px 5px $shadow-color;
                 align-items: flex-start;
                 max-width: 90%;
-                width: 100%;
+                width: fit-content;
                 word-break: break-word;
+                position: relative;
 
-                .message-footer {
-                    width: 100%;
-
-                    .toggle-logs-button {
-                        margin-top: 10px;
-                        width: 95%;
-                        text-align: center;
-                        background-color: $button-color;
-                        color: #fff;
-                        border: none;
-                        border-radius: $border-radius;
-                        cursor: pointer;
-                        box-shadow: 0 2px 4px $shadow-color;
-                        &:hover {
-                            background-color: $button-hover-color;
-                        }
-                    }
-                    .tool-logs {
-                        width: 100%;
-                        background-color: $input-background;
-                        color: $text-color;
-                        margin-top: 5px;
-                        font-family: monospace;
-                        overflow-x: auto;
-                        max-height: 150px;
-                        overflow-y: auto;
+                .citation-button {
+                    background-color: $citation-color;
+                    color: #fff;
+                    border: none;
+                    border-radius: 50%;
+                    cursor: pointer;
+                    width: 20px;
+                    height: 20px;
+                    display: inline-flex;
+                    align-items: center;
+                    justify-content: center;
+                    font-size: 12px;
+                    font-weight: bold;
+                    margin: 0 2px;
+                    padding: 0;
+                    transition: background-color 0.3s;
+
+                    &:hover {
+                        background-color: $citation-hover-color;
                     }
                 }
 
-                .custom-link {
-                    color: lightblue;
-                    text-decoration: underline;
-                    cursor: pointer;
-                }
                 &.user {
                     align-self: flex-end;
                     background-color: $button-color;
@@ -108,6 +102,40 @@ $border-radius: 8px;
                     border-radius: 50%;
                 }
             }
+
+            .follow-up-questions {
+                margin-top: 10px;
+                width: 100%;
+
+                h4 {
+                    margin-bottom: 5px;
+                    font-size: 14px;
+                }
+
+                .follow-up-button {
+                    background-color: $follow-up-bg-color;
+                    border: 1px solid #ddd;
+                    border-radius: 8px;
+                    padding: 8px 10px;
+                    margin: 4px 0;
+                    cursor: pointer;
+                    transition: background-color 0.3s;
+                    display: block;
+                    width: 100%;
+                    text-align: left;
+                    white-space: normal;
+                    word-wrap: break-word;
+                    font-size: 12px;
+                    color: $text-color;
+                    min-height: 40px;
+                    height: auto;
+                    line-height: 1.3;
+
+                    &:hover {
+                        background-color: $follow-up-hover-bg-color;
+                    }
+                }
+            }
         }
     }
 
@@ -212,37 +240,3 @@ $border-radius: 8px;
         }
     }
 }
-
-.follow-up-questions {
-    margin-top: 10px;
-    width: 100%;
-
-    h4 {
-        margin-bottom: 5px;
-        font-size: 14px;
-    }
-
-    .follow-up-button {
-        background-color: #f0f0f0;
-        border: 1px solid #ddd;
-        border-radius: 8px;
-        padding: 8px 10px;
-        margin: 4px 0;
-        cursor: pointer;
-        transition: background-color 0.3s;
-        display: block;
-        width: 100%;
-        text-align: left;
-        white-space: normal;
-        word-wrap: break-word;
-        font-size: 12px;
-        color: $text-color;
-        min-height: 40px;
-        height: auto; // Allow the button to expand as needed
-        line-height: 1.3;
-
-        &:hover {
-            background-color: #e0e0e0;
-        }
-    }
-}
diff --git a/src/client/views/nodes/ChatBox/ChatBox.tsx b/src/client/views/nodes/ChatBox/ChatBox.tsx
index bae6bbaa6..4d7381a57 100644
--- a/src/client/views/nodes/ChatBox/ChatBox.tsx
+++ b/src/client/views/nodes/ChatBox/ChatBox.tsx
@@ -12,15 +12,20 @@ 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 } from './types';
+import { ASSISTANT_ROLE, AssistantMessage, AI_Document, convertToAIDocument, Citation, CHUNK_TYPE, Chunk, getChunkType } from './types';
 import { Vectorstore } from './vectorstore/VectorstoreUpload';
 import { CollectionFreeFormDocumentView } from '../CollectionFreeFormDocumentView';
 import { CollectionFreeFormView } from '../../collections/collectionFreeForm';
 import { Agent } from './Agent';
 import dotenv from 'dotenv';
-import { DocData } from '../../../../fields/DocSymbols';
+import { DocData, DocViews } from '../../../../fields/DocSymbols';
 import { DocumentView } from '../DocumentView';
 import { AnswerParser } from './AnswerParser';
+import { DocumentManager } from '../../../util/DocumentManager';
+import { UUID } from 'bson';
+import { v4 as uuidv4 } from 'uuid';
+import { aS } from '@fullcalendar/core/internal-common';
+
 dotenv.config();
 
 @observer
@@ -31,11 +36,10 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     @observable isLoading: boolean = false;
     @observable isInitializing: boolean = true;
     @observable expandedScratchpadIndex: number | null = null;
-    @observable linked_docs_to_add: Doc[] = [];
     @observable inputValue: string = '';
-    @observable private _visibleDocs: Doc[] = [];
+    @observable private currently_linked: Doc[] = [];
     private openai: OpenAI;
-    // private vectorstore_id: string;
+    private vectorstore_id: string;
     private documents: AI_Document[] = [];
     private _oldWheel: any;
     private vectorstore: Vectorstore;
@@ -50,9 +54,14 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         makeObservable(this);
         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();
-        this.vectorstore = new Vectorstore();
+        if (StrCast(this.dataDoc.vectorstore_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
-
         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 => {
@@ -61,19 +70,17 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         );
     }
 
-    @action
-    addDocsToVectorstore = async (visible_docs: Doc[]) => {
-        await this.vectorstore.addAIDocs(visible_docs);
-        this.isInitializing = false;
+    addDocsToVectorstore = async (linkedDocs: Doc[]) => {
+        await this.vectorstore.addAIDocs(linkedDocs);
     };
 
-    @action
-    uploadNewDocument = async (newDoc: Doc) => {
-        const local_file_path: string = CsvCast(newDoc.data, PDFCast(newDoc.data)).url.pathname;
-        const { document_json } = await Networking.PostToServer('/createDocument', { file_path: local_file_path });
-        this.documents.push(...document_json.map(convertToAIDocument));
-        newDoc['ai_document'] = document_json;
-    };
+    // @action
+    // uploadNewDocument = async (newDoc: Doc) => {
+    //     const local_file_path: string = CsvCast(newDoc.data, PDFCast(newDoc.data)).url.pathname;
+    //     const { document_json } = await Networking.PostToServer('/createDocument', { file_path: local_file_path });
+    //     this.documents.push(...document_json.map(convertToAIDocument));
+    //     //newDoc['ai_document'] = document_json;
+    // };
 
     @action
     toggleToolLogs = (index: number) => {
@@ -142,6 +149,24 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     @action
     handleCitationClick = (citation: Citation) => {
         console.log('Citation clicked:', citation);
+        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_url = CsvCast(doc.data, PDFCast(doc.data)).url.pathname;
+            console.log('URL: ' + doc_url + ' Citation URL: ' + chunk_file_name);
+            //const ai_field_id = doc[this.Document[Id] + '_ai_field_id'];
+            if (chunk_file_name == doc_url) {
+                DocumentManager.Instance.showDocument(doc, {}, () => {
+                    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
     };
 
@@ -183,6 +208,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
 
     componentDidMount() {
         this._props.setContentViewBox?.(this);
+        this.currently_linked = this.linkedDocs;
         if (this.dataDoc.data) {
             try {
                 const storedHistory = JSON.parse(StrCast(this.dataDoc.data));
@@ -201,15 +227,15 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
             }
         }
         reaction(
-            () => this.visibleDocs,
-            visibleDocs => {
-                this._visibleDocs.push(...visibleDocs.filter(visibleDoc => !this._visibleDocs.includes(visibleDoc)));
+            () => this.linkedDocs,
+            linkedDocs => {
+                this.currently_linked.push(...linkedDocs.filter(linkedDoc => !this.currently_linked.includes(linkedDoc)));
             }
         );
         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._visibleDocs,
+            this.currently_linked,
             change => {
                 // observe pushes/splices on a user link DB 'data' field (should only happen for local changes)
                 switch (change.type as any) {
@@ -228,15 +254,13 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
             },
             true
         );
-        runInAction(() => {
-            if (!this._visibleDocs.length) {
-                this.isInitializing = false;
-            }
-        });
+        if (this.isInitializing) {
+            this.isInitializing = false;
+        }
     }
 
     @computed
-    get visibleDocs() {
+    get linkedDocs() {
         //return (CollectionFreeFormDocumentView.from(this._props.DocumentView?.())?._props.parent as CollectionFreeFormView)?.childDocs.filter(doc => doc != this.Document) ?? [];
         return LinkManager.Instance.getAllRelatedLinks(this.Document)
             .map(d => DocCast(LinkManager.getOppositeAnchor(d, this.Document)))
diff --git a/src/client/views/nodes/ChatBox/MessageComponent.tsx b/src/client/views/nodes/ChatBox/MessageComponent.tsx
index 76faff10b..e18224405 100644
--- a/src/client/views/nodes/ChatBox/MessageComponent.tsx
+++ b/src/client/views/nodes/ChatBox/MessageComponent.tsx
@@ -35,7 +35,7 @@ const MessageComponentBox: React.FC<MessageComponentProps> = function ({ message
                         height: '20px',
                         borderRadius: '50%',
                         border: 'none',
-                        background: '#007bff',
+                        background: '#ff6347',
                         color: 'white',
                         fontSize: '12px',
                         fontWeight: 'bold',
diff --git a/src/client/views/nodes/ChatBox/vectorstore/VectorstoreUpload.ts b/src/client/views/nodes/ChatBox/vectorstore/VectorstoreUpload.ts
index d3b1cb4e7..3a889bff2 100644
--- a/src/client/views/nodes/ChatBox/vectorstore/VectorstoreUpload.ts
+++ b/src/client/views/nodes/ChatBox/vectorstore/VectorstoreUpload.ts
@@ -16,9 +16,10 @@ export class Vectorstore {
     private index!: Index;
     private cohere: CohereClient;
     private indexName: string = 'pdf-chatbot';
+    private id: string;
     documents: AI_Document[] = [];
 
-    constructor() {
+    constructor(id: string) {
         const pineconeApiKey = process.env.PINECONE_API_KEY;
         if (!pineconeApiKey) {
             throw new Error('PINECONE_API_KEY is not defined.');
@@ -30,6 +31,7 @@ export class Vectorstore {
         this.cohere = new CohereClient({
             token: process.env.COHERE_API_KEY,
         });
+        this.id = id;
         this.initializeIndex();
     }
 
@@ -62,23 +64,28 @@ export class Vectorstore {
     }
 
     async addAIDoc(doc: Doc) {
-        if (doc[DocData]?.ai_document) {
-            this.documents.push(convertToAIDocument(JSON.parse(StrCast(doc[DocData].ai_document))));
+        if (doc.ai_document) {
+            if (doc.ai_document === 'IN PROGRESS') {
+                console.log('Already in progress.');
+                return;
+            }
+            this.documents.push(convertToAIDocument(JSON.parse(StrCast(doc.ai_document))));
             console.log(`Document already added: ${doc.file_name}`);
         } else {
+            doc.ai_document = 'IN PROGRESS';
             console.log(doc);
             console.log(PDFCast(doc.data)?.url?.pathname);
             console.log(CsvCast(doc.data)?.url?.pathname);
             const local_file_path: string = CsvCast(doc.data)?.url?.pathname ?? PDFCast(doc.data)?.url?.pathname;
             console.log('Local File Path:', local_file_path);
-            if (local_file_path) {
+            if (local_file_path !== undefined || local_file_path !== null || local_file_path !== '') {
                 const { document_json } = await Networking.PostToServer('/createDocument', { file_path: local_file_path });
                 console.log('Document JSON:', document_json);
                 const ai_document: AI_Document = convertToAIDocument(document_json);
                 this.documents.push(ai_document);
                 await this.indexDocument(ai_document);
                 console.log(`Document added: ${ai_document.file_name}`);
-                doc[DocData].ai_document = JSON.stringify(document_json);
+                doc.ai_document = JSON.stringify(document_json);
             }
         }
     }
@@ -94,7 +101,7 @@ export class Vectorstore {
                 ({
                     id: chunk.id,
                     values: chunk.values,
-                    metadata: chunk.metadata as RecordMetadata,
+                    metadata: { ...chunk.metadata, vectorestore_id: this.id } as RecordMetadata,
                 }) as PineconeRecord
         );
         await this.index.upsert(pineconeRecords);
@@ -125,6 +132,9 @@ export class Vectorstore {
 
             const queryResponse: QueryResponse<RecordMetadata> = await this.index.query({
                 vector: queryEmbedding,
+                filter: {
+                    vectorstore_id: this.id,
+                },
                 topK,
                 includeValues: true,
                 includeMetadata: true,