1 files changed, 27 insertions, 2 deletions
diff --git a/src/client/views/nodes/chatbot/chatboxcomponents/ChatBox.tsx b/src/client/views/nodes/chatbot/chatboxcomponents/ChatBox.tsx
index e09b4313f..43765c1ce 100644
--- a/src/client/views/nodes/chatbot/chatboxcomponents/ChatBox.tsx
+++ b/src/client/views/nodes/chatbot/chatboxcomponents/ChatBox.tsx
@@ -44,6 +44,7 @@ import { ProgressBar } from './ProgressBar';
 import { OpenWhere } from '../../OpenWhere';
 import { Upload } from '../../../../../server/SharedMediaTypes';
 import { DocumentMetadataTool } from '../tools/DocumentMetadataTool';
+import { AgentDocumentManager } from '../utils/AgentDocumentManager';
 
 dotenv.config();
 
@@ -76,6 +77,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     private agent: Agent;
     private messagesRef: React.RefObject<HTMLDivElement>;
     private _textInputRef: HTMLInputElement | undefined | null;
+    private docManager: AgentDocumentManager;
 
     /**
      * Static method that returns the layout string for the field.
@@ -107,7 +109,8 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
             this.vectorstore_id = StrCast(this.dataDoc.vectorstore_id);
         }
         this.vectorstore = new Vectorstore(this.vectorstore_id, this.retrieveDocIds);
-        this.agent = new Agent(this.vectorstore, this.retrieveSummaries, this.retrieveFormattedHistory, this.retrieveCSVData, this.addLinkedUrlDoc, this.createImageInDash, this.createCSVInDash, this);
+        this.docManager = new AgentDocumentManager(this);
+        this.agent = new Agent(this.vectorstore, this.retrieveSummaries, this.retrieveFormattedHistory, this.retrieveCSVData, this.addLinkedUrlDoc, this.getLinkedUrlDocIds, this.createImageInDash, this.createCSVInDash, this, this.docManager);
 
         // Reinitialize the DocumentMetadataTool with a direct reference to this ChatBox instance
         // This ensures the tool can properly access documents in the same Freeform view
@@ -380,7 +383,7 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     @action
     addLinkedUrlDoc = async (url: string, id: string) => {
         const doc = Docs.Create.WebDocument(url, { data_useCors: true });
-
+        this.docManager.addCustomId(doc, id);
         const linkDoc = Docs.Create.LinkDocument(this.Document, doc);
         LinkManager.Instance.addLink(linkDoc);
 
@@ -391,6 +394,28 @@ export class ChatBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         };
 
         doc.chunk_simpl = JSON.stringify({ chunks: [chunkToAdd] });
+        this.docManager.processDocument(doc);
+    };
+
+    /**
+     * Retrieves the IDs of linked url documents.
+     * @returns An array of document IDs.
+     */
+    @action
+    getLinkedUrlDocIds = () => {
+        const linkedDocs: Doc[] = this.linkedDocs;
+        const linkedUrlDocIds: string[] = [];
+
+        for (const doc of linkedDocs) {
+            if (doc.chunk_simpl) {
+                const docChunkSimpl = JSON.parse(StrCast(doc.chunk_simpl)) as { chunks: SimplifiedChunk[] };
+                const foundChunk = docChunkSimpl.chunks.find(chunk => chunk.chunkType === CHUNK_TYPE.URL);
+                if (foundChunk) {
+                    linkedUrlDocIds.push(foundChunk.chunkId);
+                }
+            }
+        }
+        return linkedUrlDocIds;
     };
 
     /**