even better

1 files changed, 5 insertions, 1 deletions

diff --git a/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts b/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts
index 784e90c3c..33eec5972 100644
--- a/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts
+++ b/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts
@@ -1007,6 +1007,10 @@ export class AgentDocumentManager {
         return docInfo?.layoutDoc;
     }
 
+    public getDataDocument(docId: string): Doc | undefined {
+        const docInfo = this.documentsById.get(docId);
+        return docInfo?.dataDoc;
+    }
     /**
      * Adds simplified chunks to a document for citation handling
      * @param doc The document to add simplified chunks to
@@ -1074,7 +1078,7 @@ export class AgentDocumentManager {
      */
     @action
     public getSimplifiedChunkById(chunkId: string): any | undefined {
-        return { foundChunk: this.simplifiedChunks.get(chunkId), doc: this.getDocument(this.simplifiedChunks.get(chunkId)?.doc_id || chunkId) };
+        return { foundChunk: this.simplifiedChunks.get(chunkId), doc: this.getDocument(this.simplifiedChunks.get(chunkId)?.doc_id || chunkId), dataDoc: this.getDataDocument(this.simplifiedChunks.get(chunkId)?.doc_id || chunkId) };
     }
 
     /**