1 files changed, 25 insertions, 8 deletions

diff --git a/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts b/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts
index e9d41efbd..784e90c3c 100644
--- a/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts
+++ b/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts
@@ -173,6 +173,16 @@ export class AgentDocumentManager {
     public processDocument(doc: Doc): string {
         // Ensure document has a persistent ID
         const docId = this.ensureDocumentId(doc);
+        if (doc.chunk_simplified) {
+            const newChunks: SimplifiedChunk[] = [];
+            for (const chunk of JSON.parse(StrCast(doc.chunk_simplified))) {
+                console.log('chunk', chunk);
+                newChunks.push(chunk as SimplifiedChunk);
+            }
+            console.log('Added simplified chunks to simplifiedChunks:', docId, newChunks);
+            this.addSimplifiedChunks(newChunks);
+            //DocCast(DocCast(this.chatBoxDocument!.agentDocument)!.chunk_simpl)!.mapping = new List<string>(Array.from(this.simplifiedChunks.values()).map(chunk => JSON.stringify(chunk)));
+        }
         // Only add if we haven't already processed this document
         if (!this.documentsById.has(docId)) {
             this.documentsById.set(docId, { layoutDoc: doc, dataDoc: doc[DocData] });
@@ -1005,21 +1015,28 @@ export class AgentDocumentManager {
      * @returns The updated document with simplified chunks
      */
     @action
-    public addSimplifiedChunks(chunks: RAGChunk[], docType: string) {
+    public addSimplifiedChunks(simplifiedChunks: SimplifiedChunk[]) {
+        simplifiedChunks.forEach(chunk => {
+            this.simplifiedChunks.set(chunk.chunkId, chunk);
+        });
+    }
+
+    public getSimplifiedChunks(chunks: RAGChunk[], docType: string): SimplifiedChunk[] {
         console.log('chunks', chunks, 'simplifiedChunks', this.simplifiedChunks);
+        const simplifiedChunks: SimplifiedChunk[] = [];
         // Create array of simplified chunks based on document type
         for (const chunk of chunks) {
             // Common properties across all chunk types
             const baseChunk: SimplifiedChunk = {
                 chunkId: chunk.id,
-                text: chunk.metadata.text,
+                //text: chunk.metadata.text,
                 doc_id: chunk.metadata.doc_id,
                 chunkType: chunk.metadata.type || CHUNK_TYPE.TEXT,
             };
 
             // Add type-specific properties
             if (docType === 'video' || docType === 'audio') {
-                this.simplifiedChunks.set(chunk.id, {
+                simplifiedChunks.push({
                     ...baseChunk,
                     start_time: chunk.metadata.start_time,
                     end_time: chunk.metadata.end_time,
@@ -1027,14 +1044,14 @@ export class AgentDocumentManager {
                     chunkType: docType === 'video' ? CHUNK_TYPE.VIDEO : CHUNK_TYPE.AUDIO,
                 } as SimplifiedChunk);
             } else if (docType === 'pdf') {
-                this.simplifiedChunks.set(chunk.id, {
+                simplifiedChunks.push({
                     ...baseChunk,
                     startPage: chunk.metadata.start_page,
                     endPage: chunk.metadata.end_page,
                     location: chunk.metadata.location,
                 } as SimplifiedChunk);
             } else if (docType === 'csv') {
-                this.simplifiedChunks.set(chunk.id, {
+                simplifiedChunks.push({
                     ...baseChunk,
                     rowStart: (chunk.metadata as any).row_start,
                     rowEnd: (chunk.metadata as any).row_end,
@@ -1043,9 +1060,10 @@ export class AgentDocumentManager {
                 } as SimplifiedChunk);
             } else {
                 // Default for other document types
-                this.simplifiedChunks.set(chunk.id, baseChunk as SimplifiedChunk);
+                simplifiedChunks.push(baseChunk as SimplifiedChunk);
             }
         }
+        return simplifiedChunks;
     }
 
     /**
@@ -1054,9 +1072,8 @@ export class AgentDocumentManager {
      * @param chunkId The ID of the chunk to retrieve
      * @returns The simplified chunk if found, undefined otherwise
      */
+    @action
     public getSimplifiedChunkById(chunkId: string): any | undefined {
-        console.log('chunkId', chunkId, 'simplifiedChunks', this.simplifiedChunks);
-        console.log('doc', this.getDocument(this.simplifiedChunks.get(chunkId)?.doc_id || ''));
         return { foundChunk: this.simplifiedChunks.get(chunkId), doc: this.getDocument(this.simplifiedChunks.get(chunkId)?.doc_id || chunkId) };
     }