1 files changed, 17 insertions, 3 deletions

diff --git a/src/client/views/nodes/ChatBox/ChatBox.tsx b/src/client/views/nodes/ChatBox/ChatBox.tsx
index 9e604073d..6269b8768 100644
--- a/src/client/views/nodes/ChatBox/ChatBox.tsx
+++ b/src/client/views/nodes/ChatBox/ChatBox.tsx
@@ -54,7 +54,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         } else {
             this.vectorstore_id = StrCast(this.dataDoc.vectorstore_id);
         }
-        this.vectorstore = new Vectorstore(this.vectorstore_id);
+        this.vectorstore = new Vectorstore(this.vectorstore_id, this.retrieveDocIds);
         this.agent = new Agent(this.vectorstore, this.retrieveSummaries, this.retrieveFormattedHistory);
 
         reaction(
@@ -113,7 +113,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
             } catch (err) {
                 console.error('Error:', err);
                 runInAction(() => {
-                    this.history.push({ role: ASSISTANT_ROLE.USER, content: [{ index: 0, type: TEXT_TYPE.NORMAL, text: 'Sorry, I encountered an error while processing your request.', citation_ids: null }] });
+                    this.history.push({ role: ASSISTANT_ROLE.ASSISTANT, content: [{ index: 0, type: TEXT_TYPE.NORMAL, text: 'Sorry, I encountered an error while processing your request.', citation_ids: null }] });
                 });
             } finally {
                 runInAction(() => {
@@ -224,7 +224,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         } else {
             runInAction(() => {
                 this.history.push({
-                    role: ASSISTANT_ROLE.USER,
+                    role: ASSISTANT_ROLE.ASSISTANT,
                     content: [{ index: 0, type: TEXT_TYPE.NORMAL, text: 'Welcome to the Document Analyser Assistant! Link a document or ask questions to get started.', citation_ids: null }],
                 });
             });
@@ -265,6 +265,16 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     }
 
     @computed
+    get docIds() {
+        return LinkManager.Instance.getAllRelatedLinks(this.Document)
+            .map(d => DocCast(LinkManager.getOppositeAnchor(d, this.Document)))
+            .map(d => DocCast(d?.annotationOn, d))
+            .filter(d => d)
+            .filter(d => d.ai_doc_id)
+            .map(d => StrCast(d.ai_doc_id));
+    }
+
+    @computed
     get summaries(): string {
         return (
             LinkManager.Instance.getAllRelatedLinks(this.Document)
@@ -295,6 +305,10 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         return this.formattedHistory;
     };
 
+    retrieveDocIds = () => {
+        return this.docIds;
+    };
+
     @action
     handleFollowUpClick = (question: string) => {
         console.log('Follow-up question clicked:', question);