1 files changed, 13 insertions, 22 deletions
diff --git a/src/client/views/nodes/ChatBox/vectorstore/Vectorstore.ts b/src/client/views/nodes/ChatBox/vectorstore/Vectorstore.ts
index 8e7be6eec..4383bb72d 100644
--- a/src/client/views/nodes/ChatBox/vectorstore/Vectorstore.ts
+++ b/src/client/views/nodes/ChatBox/vectorstore/Vectorstore.ts
@@ -3,6 +3,7 @@ import { CohereClient } from 'cohere-ai';
 import { EmbedResponse } from 'cohere-ai/api';
 import dotenv from 'dotenv';
 import axios from 'axios';
+import { SimplifiedChunk } from '../types';
 
 import { RAGChunk, AI_Document, CHUNK_TYPE } from '../types';
 import { Doc } from '../../../../../fields/Doc';
@@ -19,6 +20,7 @@ export class Vectorstore {
     private indexName: string = 'pdf-chatbot';
     private _id: string;
     private _doc_ids: string[] = [];
+
     documents: AI_Document[] = [];
 
     constructor(id: string, doc_ids: () => string[]) {
@@ -91,31 +93,20 @@ export class Vectorstore {
                     doc.vectorstore_id = JSON.stringify(JSON.parse(StrCast(doc.vectorstore_id)).concat([this._id]));
                 }
                 if (doc.chunk_simpl === undefined || doc.chunk_simpl === null || doc.chunk_simpl === '' || doc.chunk_simpl === '[]') {
-                    doc.chunk_simpl = JSON.stringify({ text_chunks: [], image_chunks: [] });
+                    doc.chunk_simpl = JSON.stringify({ chunks: [] });
                 }
-                let new_chunk_simpl: { text_chunks: { chunk_id: string; start_page: number; end_page: number }[]; image_chunks: { chunk_id: string; location: string; page: number }[] } = {
-                    text_chunks: [],
-                    image_chunks: [],
-                };
 
                 document_json.chunks.forEach((chunk: RAGChunk) => {
-                    let chunk_to_add: { chunk_id: string; start_page: number; end_page: number }[] | { chunk_id: string; location: string; page: number }[];
-                    switch (chunk.metadata.type) {
-                        case CHUNK_TYPE.TEXT:
-                            chunk_to_add = [{ chunk_id: chunk.id, start_page: chunk.metadata.start_page, end_page: chunk.metadata.end_page }];
-                            new_chunk_simpl = JSON.parse(StrCast(doc.chunk_simpl));
-                            new_chunk_simpl.text_chunks = new_chunk_simpl.text_chunks.concat(chunk_to_add);
-                            doc.chunk_simpl = JSON.stringify(new_chunk_simpl);
-                            break;
-                        case CHUNK_TYPE.IMAGE:
-                        case CHUNK_TYPE.TABLE:
-                            console.log('Location:', chunk.metadata.location);
-                            chunk_to_add = [{ chunk_id: chunk.id, location: chunk.metadata.location, page: chunk.metadata.start_page }];
-                            new_chunk_simpl = JSON.parse(StrCast(doc.chunk_simpl));
-                            new_chunk_simpl.image_chunks = new_chunk_simpl.image_chunks.concat(chunk_to_add);
-                            doc.chunk_simpl = JSON.stringify(new_chunk_simpl);
-                            break;
-                    }
+                    const chunkToAdd = {
+                        chunkId: chunk.id,
+                        startPage: chunk.metadata.start_page,
+                        endPage: chunk.metadata.end_page,
+                        location: chunk.metadata.location,
+                        chunkType: chunk.metadata.type as CHUNK_TYPE,
+                    };
+                    const new_chunk_simpl = JSON.parse(StrCast(doc.chunk_simpl));
+                    new_chunk_simpl.chunks = new_chunk_simpl.chunks.concat(chunkToAdd);
+                    doc.chunk_simpl = JSON.stringify(new_chunk_simpl);
                 });
 
                 doc.ai_document_status = 'COMPLETED';