citation parsing working much better

1 files changed, 14 insertions, 38 deletions

diff --git a/src/client/views/nodes/ChatBox/ChatBox.tsx b/src/client/views/nodes/ChatBox/ChatBox.tsx
index 9b2a92564..bae6bbaa6 100644
--- a/src/client/views/nodes/ChatBox/ChatBox.tsx
+++ b/src/client/views/nodes/ChatBox/ChatBox.tsx
@@ -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 } from './types';
+import { ASSISTANT_ROLE, AssistantMessage, AI_Document, convertToAIDocument, Citation, CHUNK_TYPE } from './types';
 import { Vectorstore } from './vectorstore/VectorstoreUpload';
 import { CollectionFreeFormDocumentView } from '../CollectionFreeFormDocumentView';
 import { CollectionFreeFormView } from '../../collections/collectionFreeForm';
@@ -20,6 +20,7 @@ import { Agent } from './Agent';
 import dotenv from 'dotenv';
 import { DocData } from '../../../../fields/DocSymbols';
 import { DocumentView } from '../DocumentView';
+import { AnswerParser } from './AnswerParser';
 dotenv.config();
 
 @observer
@@ -47,14 +48,13 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     constructor(props: FieldViewProps) {
         super(props);
         makeObservable(this);
-        this.openai = this.initializeOpenAI();
-        this.history = [{ role: ASSISTANT_ROLE.ASSISTANT, text: 'Welcome to the Document Analyser Assistant! Link a document or ask questions to get started.' }];
+        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();
         this.agent = new Agent(this.vectorstore); // Initialize the Agent
 
         reaction(
-            () => this.history.map((msg: AssistantMessage) => ({ role: msg.role, text: msg.text, follow_up_questions: msg.follow_up_questions, citations: msg.citations })),
+            () => this.history.map((msg: AssistantMessage) => ({ role: msg.role, text_content: msg.text_content, follow_up_questions: msg.follow_up_questions, citations: msg.citations })),
             serializableHistory => {
                 this.dataDoc.data = JSON.stringify(serializableHistory);
             }
@@ -110,19 +110,19 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
             try {
                 textInput.value = '';
                 runInAction(() => {
-                    this.history.push({ role: ASSISTANT_ROLE.USER, text: trimmedText });
+                    this.history.push({ role: ASSISTANT_ROLE.USER, text_content: trimmedText });
                     this.isLoading = true;
                 });
 
                 const response = await this.agent.askAgent(trimmedText); // Use the chatbot to get the response
                 runInAction(() => {
-                    this.history.push(this.parseAssistantResponse(response));
+                    this.history.push(AnswerParser.parse(response));
                 });
                 this.dataDoc.data = JSON.stringify(this.history);
             } catch (err) {
                 console.error('Error:', err);
                 runInAction(() => {
-                    this.history.push({ role: ASSISTANT_ROLE.ASSISTANT, text: 'Sorry, I encountered an error while processing your request.' });
+                    this.history.push({ role: ASSISTANT_ROLE.ASSISTANT, text_content: 'Sorry, I encountered an error while processing your request.' });
                 });
             } finally {
                 runInAction(() => {
@@ -132,35 +132,6 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         }
     };
 
-    parseAssistantResponse(response: string): AssistantMessage {
-        const parser = new DOMParser();
-        const xmlDoc = parser.parseFromString(response, 'text/xml');
-        const answerElement = xmlDoc.querySelector('answer');
-        const followUpQuestionsElement = xmlDoc.querySelector('follow_up_questions');
-
-        let text = '';
-        let followUpQuestions: string[] = [];
-
-        if (answerElement) {
-            // Remove the follow_up_questions element from the answer
-            const followUpElement = answerElement.querySelector('follow_up_questions');
-            if (followUpElement) {
-                followUpElement.remove();
-            }
-            text = answerElement.innerHTML.trim();
-        }
-
-        if (followUpQuestionsElement) {
-            followUpQuestions = Array.from(followUpQuestionsElement.querySelectorAll('question')).map(q => q.textContent || '');
-        }
-
-        return {
-            role: ASSISTANT_ROLE.ASSISTANT,
-            text,
-            follow_up_questions: followUpQuestions,
-        };
-    }
-
     @action
     updateMessageCitations = (index: number, citations: Citation[]) => {
         if (this.history[index]) {
@@ -219,7 +190,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                     this.history.push(
                         ...storedHistory.map((msg: AssistantMessage) => ({
                             role: msg.role,
-                            text: msg.text,
+                            text_content: msg.text_content,
                             follow_up_questions: msg.follow_up_questions,
                             citations: msg.citations,
                         }))
@@ -246,6 +217,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                         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[]);
+                            console.log('here!');
                             this.addDocsToVectorstore((change as any).added as Doc[]);
                         }
                         // (change as any).removed.forEach((link: any) => remLinkFromDoc(toRealField(link)));
@@ -265,7 +237,11 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
 
     @computed
     get visibleDocs() {
-        return (CollectionFreeFormDocumentView.from(this._props.DocumentView?.())?._props.parent as CollectionFreeFormView)?.childDocs.filter(doc => doc != this.Document) ?? [];
+        //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)))
+            .map(d => DocCast(d?.annotationOn, d))
+            .filter(d => d);
     }
 
     @action