Working now somewhat

1 files changed, 21 insertions, 6 deletions

diff --git a/src/client/views/nodes/ChatBox/ChatBox.tsx b/src/client/views/nodes/ChatBox/ChatBox.tsx
index 3ecb2d340..2ce1ebdd2 100644
--- a/src/client/views/nodes/ChatBox/ChatBox.tsx
+++ b/src/client/views/nodes/ChatBox/ChatBox.tsx
@@ -16,7 +16,7 @@ import { ASSISTANT_ROLE, AssistantMessage, AI_Document, convertToAIDocument, Cit
 import { Vectorstore } from './vectorstore/VectorstoreUpload';
 import { CollectionFreeFormDocumentView } from '../CollectionFreeFormDocumentView';
 import { CollectionFreeFormView } from '../../collections/collectionFreeForm';
-import { ChatBot } from './ChatBot';
+import { Agent } from './Agent';
 import dotenv from 'dotenv';
 dotenv.config();
 
@@ -34,7 +34,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     private documents: AI_Document[] = [];
     private _oldWheel: any;
     private vectorstore: Vectorstore;
-    private chatbot: ChatBot; // Add the ChatBot instance
+    private agent: Agent; // Add the ChatBot instance
 
     public static LayoutString(fieldKey: string) {
         return FieldView.LayoutString(ChatBox, fieldKey);
@@ -48,7 +48,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         this.openai = this.initializeOpenAI();
         this.getOtherDocs();
         this.vectorstore = new Vectorstore();
-        this.chatbot = new ChatBot(this.vectorstore); // Initialize the ChatBot
+        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 })),
@@ -58,6 +58,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         );
     }
 
+    @action
     getOtherDocs = async () => {
         const visible_docs = (CollectionFreeFormDocumentView.from(this._props.DocumentView?.())?._props.parent as CollectionFreeFormView)?.childDocs
             .filter(doc => doc != this.Document)
@@ -76,6 +77,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                 doc['ai_document'] = document_json;
             }
         });
+        this.isInitializing = false;
     };
 
     @action
@@ -120,7 +122,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                     this.history.push({ role: ASSISTANT_ROLE.USER, text: trimmedText });
                 });
                 this.isLoading = true;
-                const response = await this.chatbot.ask(trimmedText); // Use the chatbot to get the response
+                const response = await this.agent.askAgent(trimmedText); // Use the chatbot to get the response
                 runInAction(() => {
                     this.history.push(this.parseAssistantResponse(response));
                 });
@@ -142,8 +144,21 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         const answerElement = xmlDoc.querySelector('answer');
         const followUpQuestionsElement = xmlDoc.querySelector('follow_up_questions');
 
-        const text = answerElement ? answerElement.innerHTML || '' : ''; // Use innerHTML to preserve citation tags
-        const followUpQuestions = followUpQuestionsElement ? Array.from(followUpQuestionsElement.querySelectorAll('question')).map(q => q.textContent || '') : [];
+        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,