Made it so chunk Ids are seperately managed and made sure the doc id is sonsistent and not created in python spawn

1 files changed, 76 insertions, 137 deletions

diff --git a/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts b/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts
index 14cffcb70..c8a6bb16b 100644
--- a/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts
+++ b/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts
@@ -1,4 +1,4 @@
-import { action, makeObservable, observable, ObservableMap, reaction, runInAction } from 'mobx';
+import { action, computed, makeObservable, observable, ObservableMap, reaction, runInAction } from 'mobx';
 import { observer } from 'mobx-react';
 import { v4 as uuidv4 } from 'uuid';
 import { Doc, StrListCast } from '../../../../../fields/Doc';
@@ -31,7 +31,7 @@ export class AgentDocumentManager {
     private chatBox: ChatBox;
     private chatBoxDocument: Doc | null = null;
     private fieldMetadata: Record<string, any> = {};
-    private readonly DOCUMENT_ID_FIELD = '_dash_document_id';
+    @observable private documentIdsFromChunkIds: ObservableMap<string, string>;
 
     /**
      * Creates a new DocumentManager
@@ -40,8 +40,17 @@ export class AgentDocumentManager {
     constructor(chatBox: ChatBox) {
         makeObservable(this);
         const agentDoc = DocCast(chatBox.Document.agentDocument) ?? new Doc();
+        const chunkIds = DocCast(agentDoc.chunkIds) ?? new Doc();
+
         agentDoc.title = chatBox.Document.title + '_agentDocument';
+        chunkIds.title = '_chunkIds';
         chatBox.Document.agentDocument = agentDoc;
+        DocCast(chatBox.Document.agentDocument)!.chunkIds = chunkIds;
+        this.documentIdsFromChunkIds = StrListCast(chunkIds.mapping).reduce((mapping, content) => {
+            const [chunkId, docId] = content.split(':');
+            mapping.set(chunkId, docId);
+            return mapping;
+        }, new ObservableMap<string, string>());
         this.documentsById = StrListCast(agentDoc.mapping).reduce((mapping, content) => {
             const [id, layoutId, docId] = content.split(':');
             const layoutDoc = DocServer.GetCachedRefField(layoutId);
@@ -66,6 +75,19 @@ export class AgentDocumentManager {
             }
             //{ fireImmediately: true }
         );
+        reaction(
+            () => this.documentIdsFromChunkIds.values(),
+            () => {
+                if (this.chatBoxDocument && DocCast(this.chatBoxDocument.agentDocument)) {
+                    // Store the mapping with chunkId:docId format for consistency
+                    const chunkIdsDoc = DocCast(DocCast(this.chatBoxDocument.agentDocument)!.chunkIds);
+                    if (chunkIdsDoc) {
+                        chunkIdsDoc.mapping = new List<string>(Array.from(this.documentIdsFromChunkIds.entries()).map(([chunkId, docId]) => `${chunkId}:${docId}`));
+                    }
+                }
+            }
+            //{ fireImmediately: true }
+        );
         this.processDocument(this.chatBoxDocument);
         this.initializeFieldMetadata();
     }
@@ -120,7 +142,7 @@ export class AgentDocumentManager {
         try {
             // Use the LinkManager approach which is proven to work in ChatBox
             if (this.chatBoxDocument) {
-                console.log('Finding documents linked to ChatBox document with ID:', this.chatBoxDocument.id);
+                console.log('Finding documents linked to ChatBox document with ID:', this.chatBoxDocument[Id]);
 
                 // Get directly linked documents via LinkManager
                 const linkedDocs = LinkManager.Instance.getAllRelatedLinks(this.chatBoxDocument)
@@ -134,57 +156,10 @@ export class AgentDocumentManager {
                 linkedDocs.forEach((doc: Doc | undefined) => {
                     if (doc) {
                         this.processDocument(doc);
-                        console.log('Processed linked document:', doc.id, doc.title, doc.type);
+                        console.log('Processed linked document:', doc[Id], doc.title, doc.type);
                     }
                 });
-
-                // Include the ChatBox document itself
-                this.processDocument(this.chatBoxDocument);
-
-                // If we have access to the Document's parent, try to find sibling documents
-                if (this.chatBoxDocument.parent) {
-                    const parent = this.chatBoxDocument.parent;
-                    console.log('Found parent document, checking for siblings');
-
-                    // Check if parent is a Doc type and has a childDocs function
-                    if (parent && typeof parent === 'object' && 'childDocs' in parent && typeof parent.childDocs === 'function') {
-                        try {
-                            const siblingDocs = parent.childDocs();
-                            if (Array.isArray(siblingDocs)) {
-                                console.log(`Found ${siblingDocs.length} sibling documents via parent.childDocs()`);
-                                siblingDocs.forEach((doc: Doc) => {
-                                    if (doc) {
-                                        this.processDocument(doc);
-                                    }
-                                });
-                            }
-                        } catch (e) {
-                            console.warn('Error accessing parent.childDocs:', e);
-                        }
-                    }
-                }
-            } else if (this.chatBox && this.chatBox.linkedDocs) {
-                // If we have direct access to the linkedDocs computed property from ChatBox
-                console.log('Using ChatBox.linkedDocs directly');
-                const linkedDocs = this.chatBox.linkedDocs;
-                if (Array.isArray(linkedDocs)) {
-                    console.log(`Found ${linkedDocs.length} documents via ChatBox.linkedDocs`);
-                    linkedDocs.forEach((doc: Doc) => {
-                        if (doc) {
-                            this.processDocument(doc);
-                        }
-                    });
-                }
-
-                // Process the ChatBox document if available
-                if (this.chatBox.Document) {
-                    this.processDocument(this.chatBox.Document);
-                }
-            } else {
-                console.warn('No ChatBox document reference available for finding linked documents');
             }
-
-            console.log(`DocumentMetadataTool found ${this.documentsById.size} total documents`);
         } catch (error) {
             console.error('Error finding documents in Freeform view:', error);
         }
@@ -201,6 +176,7 @@ export class AgentDocumentManager {
         // Only add if we haven't already processed this document
         if (!this.documentsById.has(docId)) {
             this.documentsById.set(docId, { layoutDoc: doc, dataDoc: doc[DocData] });
+            console.log('Added document to documentsById:', doc[Id], docId, doc[Id], doc[DocData][Id]);
         }
         return docId;
     }
@@ -213,37 +189,12 @@ export class AgentDocumentManager {
     private ensureDocumentId(doc: Doc): string {
         let docId: string | undefined;
 
-        // First try to get the ID from our custom field
-        if (doc[this.DOCUMENT_ID_FIELD]) {
-            docId = String(doc[this.DOCUMENT_ID_FIELD]);
-            return docId;
-        }
-
-        // Try different ways to get a document ID
-
         // 1. Try the direct id property if it exists
-        if (doc.id && typeof doc.id === 'string') {
-            docId = doc.id;
-        }
-        // 2. Try doc._id if it exists
-        else if (doc._id && typeof doc._id === 'string') {
-            docId = doc._id;
-        }
-        // 3. Try doc.data?.id if it exists
-        else if (doc.data && typeof doc.data === 'object' && 'id' in doc.data && typeof doc.data.id === 'string') {
-            docId = doc.data.id;
-        }
-        // 4. If none of the above work, generate a UUID
-        else {
-            docId = uuidv4();
-            console.log(`Generated new UUID for document with title: ${doc.title || 'Untitled'}`);
-        }
-
-        // Store the ID in the document's metadata so it persists
-        try {
-            doc[this.DOCUMENT_ID_FIELD] = docId;
-        } catch (e) {
-            console.warn(`Could not assign ID to document property`, e);
+        if (doc[Id]) {
+            console.log('Found document ID (normal):', doc[Id]);
+            docId = doc[Id];
+        } else {
+            throw new Error('No document ID found');
         }
 
         return docId;
@@ -256,13 +207,13 @@ export class AgentDocumentManager {
      */
     public extractDocumentMetadata(id: string) {
         if (!id) return null;
-        const doc = this.documentsById.get(id);
-        if (!doc) return null;
-        const layoutDoc = doc.layoutDoc;
-        const dataDoc = doc.dataDoc;
+        const agentDoc = this.documentsById.get(id);
+        if (!agentDoc) return null;
+        const layoutDoc = agentDoc.layoutDoc;
+        const dataDoc = agentDoc.dataDoc;
 
         const metadata: Record<string, any> = {
-            id: layoutDoc.dash_document_id || layoutDoc.id || '',
+            id: layoutDoc[Id] || dataDoc[Id] || '',
             title: layoutDoc.title || '',
             type: layoutDoc.type || '',
             fields: {
@@ -355,7 +306,7 @@ export class AgentDocumentManager {
         if (value instanceof Doc) {
             return {
                 type: 'Doc',
-                id: value.id || this.ensureDocumentId(value),
+                id: value[Id] || this.ensureDocumentId(value),
                 title: value.title || '',
                 docType: value.type || '',
             };
@@ -1011,33 +962,17 @@ export class AgentDocumentManager {
      * Returns a list of all document IDs in the manager.
      * @returns An array of document IDs (strings).
      */
-    public listDocs(): string[] {
-        return Array.from(this.documentsById.keys());
+    @computed
+    public get listDocs(): string[] {
+        console.log(
+            Array.from(this.documentsById.entries()).map(([id, agentDoc]) => JSON.stringify({ id, title: agentDoc.layoutDoc.title, type: agentDoc.layoutDoc.type, summary: agentDoc.layoutDoc.summary || 'No summary available for this document.' }))
+        );
+        return Array.from(this.documentsById.entries()).map(([id, agentDoc]) => JSON.stringify({ id, title: agentDoc.layoutDoc.title, type: agentDoc.layoutDoc.type, summary: agentDoc.layoutDoc.summary || 'No summary available for this document.' }));
     }
 
-    /**
-     * Adds a document with a custom ID to the manager
-     * @param doc The document to add
-     * @param customId The custom ID to assign to the document
-     * @returns The customId that was assigned
-     */
-    @action
-    public addCustomId(doc: Doc, customId: string): string {
-        if (!doc) {
-            console.error('Cannot add null document with custom ID');
-            return '';
-        }
-
-        // Set the custom ID in the document's metadata
-        doc[this.DOCUMENT_ID_FIELD] = customId;
-
-        // Store the document in our map
-        this.documentsById.set(customId, {
-            layoutDoc: doc,
-            dataDoc: doc,
-        });
-
-        return customId;
+    @computed
+    public get docIds(): string[] {
+        return Array.from(this.documentsById.keys());
     }
 
     /**
@@ -1078,11 +1013,8 @@ export class AgentDocumentManager {
             // Ensure each chunk ID can be linked back to its parent document
             // Store a mapping from chunk ID to parent document ID
             // This allows us to easily find a document by any of its chunk IDs
-            if (!this.documentsById.has(chunkId)) {
-                this.documentsById.set(chunkId, {
-                    layoutDoc: doc,
-                    dataDoc: docInfo.dataDoc,
-                });
+            if (!this.documentIdsFromChunkIds.has(chunkId) && doc) {
+                this.documentIdsFromChunkIds.set(chunkId, doc[Id]);
             }
         }
     }
@@ -1092,11 +1024,25 @@ export class AgentDocumentManager {
      * @param chunkId The chunk ID to look up
      * @returns The parent document ID if found
      */
-    public getDocIdByChunkId(chunkId: string): string | undefined {
-        const docInfo = this.documentsById.get(chunkId);
+    public getDocByChunkId(chunkId: string): Doc | undefined {
+        // First, look up the document ID using the chunk ID mapping
+        const docId = this.documentIdsFromChunkIds.get(chunkId);
+        console.log('this.documentIdsFromChunkIds', this.documentIdsFromChunkIds);
+        console.log('docId', docId);
+        if (!docId) {
+            if (this.documentsById.has(chunkId)) {
+                return this.documentsById.get(chunkId)?.layoutDoc;
+            } else {
+                console.error('No document found for chunkId and docId', chunkId);
+                return undefined;
+            }
+        }
+        // Then get the document using the document ID
+        const docInfo = this.documentsById.get(docId);
         if (docInfo) {
-            return docInfo.layoutDoc[this.DOCUMENT_ID_FIELD] as string;
+            return docInfo.layoutDoc;
         }
+        console.error('No document found for docId', docId);
         return undefined;
     }
 
@@ -1157,7 +1103,7 @@ export class AgentDocumentManager {
                 return baseChunk;
             }
         });
-
+        console.log('simplifiedChunks', simplifiedChunks);
         // Update the document with all simplified chunks at once
         doc.chunk_simpl = JSON.stringify({ chunks: simplifiedChunks });
 
@@ -1165,32 +1111,25 @@ export class AgentDocumentManager {
     }
 
     /**
-     * Gets the simplified chunks from a document
-     * @param doc The document to get simplified chunks from
-     * @returns Array of simplified chunks or empty array if none exist
+     * Gets a specific simplified chunk by ID
+     * @param doc The document containing chunks
+     * @param chunkId The ID of the chunk to retrieve
+     * @returns The simplified chunk if found, undefined otherwise
      */
-    public getSimplifiedChunks(doc: Doc): any[] {
+    public getSimplifiedChunkById(doc: Doc, chunkId: string): any | undefined {
+        let chunks: any[] = [];
         if (!doc || !doc.chunk_simpl) {
+            chunks = [];
+            console.warn('No chunk found for chunkId', chunkId, '. Checking if document exists in documentsById.');
             return [];
         }
-
         try {
             const parsed = JSON.parse(StrCast(doc.chunk_simpl));
-            return parsed.chunks || [];
+            chunks = parsed.chunks || [];
         } catch (e) {
             console.error('Error parsing simplified chunks:', e);
             return [];
         }
-    }
-
-    /**
-     * Gets a specific simplified chunk by ID
-     * @param doc The document containing chunks
-     * @param chunkId The ID of the chunk to retrieve
-     * @returns The simplified chunk if found, undefined otherwise
-     */
-    public getSimplifiedChunkById(doc: Doc, chunkId: string): any | undefined {
-        const chunks = this.getSimplifiedChunks(doc);
         return chunks.find(chunk => chunk.chunkId === chunkId);
     }