new attempt with new citation unification

1 files changed, 115 insertions, 53 deletions

diff --git a/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts b/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts
index 4eeac3c6a..c3beebcde 100644
--- a/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts
+++ b/src/client/views/nodes/chatbot/utils/AgentDocumentManager.ts
@@ -165,22 +165,18 @@ export class AgentDocumentManager {
         }
     }
 
-    public addCustomId(doc: Doc, id: string) {
-        doc.id = id;
-        doc.DOCUMENT_ID_FIELD = id;
-    }
-
     /**
      * Process a document by ensuring it has an ID and adding it to the appropriate collections
      * @param doc The document to process
      */
-    public processDocument(doc: Doc) {
+    public processDocument(doc: Doc): string {
         // Ensure document has a persistent ID
         const docId = this.ensureDocumentId(doc);
         // Only add if we haven't already processed this document
         if (!this.documentsById.has(docId)) {
             this.documentsById.set(docId, { layoutDoc: doc, dataDoc: doc[DocData] });
         }
+        return docId;
     }
 
     /**
@@ -232,7 +228,9 @@ export class AgentDocumentManager {
      * @param docId The ID of the document to extract metadata from
      * @returns An object containing the document's metadata
      */
-    public extractDocumentMetadata(doc?: AgentDocument) {
+    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;
@@ -729,16 +727,14 @@ export class AgentDocumentManager {
      */
     public getDocumentMetadata(documentId?: string): any {
         if (documentId) {
-            const doc = this.documentsById.get(documentId);
-            // Get metadata for a specific document
-            return this.extractDocumentMetadata(doc);
+            console.log(`Returning document metadata for docID, ${documentId}:`, this.extractDocumentMetadata(documentId));
+            return this.extractDocumentMetadata(documentId);
         } else {
             // Get metadata for all documents
             const documentsMetadata: Record<string, any> = {};
-            for (const doc of this.documentsById.values()) {
-                documentsMetadata.add(this.extractDocumentMetadata(doc) ?? { documentId: doc.layoutDoc.id, title: doc.layoutDoc.title, type: doc.layoutDoc.type });
+            for (const documentId of this.documentsById.keys()) {
+                documentsMetadata.add(this.extractDocumentMetadata(documentId));
             }
-
             return {
                 documentCount: this.documentsById.size,
                 documents: documentsMetadata,
@@ -845,14 +841,15 @@ export class AgentDocumentManager {
         return Object.values(supportedDocTypes).includes(docType as supportedDocTypes);
     }
     /**
-     * Creates a document in the dashboard.
+     * Creates a document in the dashboard and returns its ID.
+     * This is a public API used by tools like SearchTool.
      *
-     * @param {string} doc_type - The type of document to create.
-     * @param {string} data - The data used to generate the document.
-     * @param {DocumentOptions} options - Configuration options for the document.
-     * @returns {Promise<void>} A promise that resolves once the document is created and displayed.
+     * @param docType The type of document to create
+     * @param data The data for the document
+     * @param options Optional configuration options
+     * @returns The ID of the created document
      */
-    createDocInDash = (docType: string, title: string, data: string) => {
+    public createDocInDash(docType: string, data: string, options?: any): string {
         // Validate doc_type
         if (!this.isValidDocType(docType)) {
             throw new Error(`Invalid document type: ${docType}`);
@@ -862,10 +859,10 @@ export class AgentDocumentManager {
             // Create simple document with just title and data
             const simpleDoc: parsedDoc = {
                 doc_type: docType,
-                title: title,
+                title: options?.title ?? `Untitled Document ${this.documentsById.size + 1}`,
                 data: data,
-                x: 0,
-                y: 0,
+                x: options?.x ?? 0,
+                y: options?.y ?? 0,
                 _width: 300,
                 _height: 300,
                 _layout_fitWidth: false,
@@ -884,46 +881,111 @@ export class AgentDocumentManager {
                 }
             };
             const doc = this.chatBox.whichDoc(simpleDoc, false);
-            if (doc) linkAndShowDoc(doc);
-            return doc;
+            if (doc) {
+                linkAndShowDoc(doc);
+                const id = this.processDocument(doc);
+                return id;
+            } else {
+                throw new Error(`Error creating document. Created document not found.`);
+            }
         } catch (error) {
             throw new Error(`Error creating document: ${error}`);
         }
-    };
+    }
 
     public has(docId: string) {
         return this.documentsById.has(docId);
     }
 
-    public listDocs() {
-        // List all available documents in simple format
-        const docs = Array.from(this.documentsById.entries()).map(([id, doc]) => ({
-            id,
-            title: doc.layoutDoc.title || 'Untitled Document',
-            type: doc.layoutDoc.type || doc.dataDoc.type || 'Unknown Type',
-        }));
-
-        if (docs.length === 0) {
-            return [
-                {
-                    type: 'text',
-                    text: 'No documents found in the current view.',
-                },
-            ];
-        }
-
-        return [
-            {
-                type: 'text',
-                text: `Found ${docs.length} document(s) in the current view:\n${JSON.stringify(docs, null, 2)}`,
-            },
-        ];
+    /**
+     * 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());
+    }
+
+    /**
+     * 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
+     */
+    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;
     }
 
-    public createAgentDoc(doc: Doc) {
-        // Ideally check if Doc is already in there.
-        const agentDoc = { layoutDoc: doc, dataDoc: doc[DocData] };
-        this.documentsById.set(this.ensureDocumentId(doc), agentDoc);
-        return agentDoc;
+    /**
+     * Gets a document by its ID
+     * @param docId The ID of the document to retrieve
+     * @returns The document if found, undefined otherwise
+     */
+    public getDocument(docId: string): Doc | undefined {
+        const docInfo = this.documentsById.get(docId);
+        return docInfo?.layoutDoc;
+    }
+
+    /**
+     * Registers chunk IDs associated with a document in the manager
+     * @param docId The parent document ID
+     * @param chunkIds Array of chunk IDs associated with this document
+     */
+    public registerChunkIds(docId: string, chunkIds: string[]): void {
+        // Get the document if it exists
+        const docInfo = this.documentsById.get(docId);
+        if (!docInfo) {
+            console.warn(`Cannot register chunks for unknown document ID: ${docId}`);
+            return;
+        }
+
+        // Store chunk IDs on the document for future reference
+        const doc = docInfo.layoutDoc;
+        if (!doc.chunk_ids) {
+            doc.chunk_ids = JSON.stringify(chunkIds);
+        } else {
+            // Merge with existing chunk IDs if they exist
+            const existingIds = JSON.parse(doc.chunk_ids as string);
+            const updatedIds = [...new Set([...existingIds, ...chunkIds])]; // Remove duplicates
+            doc.chunk_ids = JSON.stringify(updatedIds);
+        }
+
+        // Ensure each chunk ID can be linked back to its parent document
+        chunkIds.forEach(chunkId => {
+            // 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,
+                });
+            }
+        });
+    }
+
+    /**
+     * Gets a document ID by a chunk ID
+     * @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);
+        if (docInfo) {
+            return docInfo.layoutDoc[this.DOCUMENT_ID_FIELD] as string;
+        }
+        return undefined;
     }
 }