From 65b2d647a1b4340584b9e7eb9083e502f8ca536d Mon Sep 17 00:00:00 2001
From: Thinger Soft <thingersoft@gmail.com>
Date: Thu, 2 Jan 2025 14:29:07 +0100
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20fix:=20Handle=20Concurrent=20Fil?=
 =?UTF-8?q?e=20Mgmt.=20For=20Agents=20(#5159)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* fix: handle concurrent file upload for agents rag

Closes #4746:

* fix: handle concurrent file deletions for agents rag

Closes #5160:

* refactor: remove useless promise wrapping
---
 api/models/Agent.js | 81 +++++++++++++++++++++++++--------------------
 1 file changed, 45 insertions(+), 36 deletions(-)

diff --git a/api/models/Agent.js b/api/models/Agent.js
index 077ac9d57a..6fa00f56bc 100644
--- a/api/models/Agent.js
+++ b/api/models/Agent.js
@@ -82,7 +82,7 @@ const loadAgent = async ({ req, agent_id }) => {
  */
 const updateAgent = async (searchParameter, updateData) => {
   const options = { new: true, upsert: false };
-  return await Agent.findOneAndUpdate(searchParameter, updateData, options).lean();
+  return Agent.findOneAndUpdate(searchParameter, updateData, options).lean();
 };
 
 /**
@@ -96,25 +96,18 @@ const updateAgent = async (searchParameter, updateData) => {
  */
 const addAgentResourceFile = async ({ agent_id, tool_resource, file_id }) => {
   const searchParameter = { id: agent_id };
-  const agent = await getAgent(searchParameter);
 
-  if (!agent) {
+  // build the update to push or create the file ids set
+  const fileIdsPath = `tool_resources.${tool_resource}.file_ids`;
+  const updateData = { $addToSet: { [fileIdsPath]: file_id } };
+
+  // return the updated agent or throw if no agent matches
+  const updatedAgent = await updateAgent(searchParameter, updateData);
+  if (updatedAgent) {
+    return updatedAgent;
+  } else {
     throw new Error('Agent not found for adding resource file');
   }
-
-  const tool_resources = agent.tool_resources || {};
-
-  if (!tool_resources[tool_resource]) {
-    tool_resources[tool_resource] = { file_ids: [] };
-  }
-
-  if (!tool_resources[tool_resource].file_ids.includes(file_id)) {
-    tool_resources[tool_resource].file_ids.push(file_id);
-  }
-
-  const updateData = { tool_resources };
-
-  return await updateAgent(searchParameter, updateData);
 };
 
 /**
@@ -126,36 +119,52 @@ const addAgentResourceFile = async ({ agent_id, tool_resource, file_id }) => {
  */
 const removeAgentResourceFiles = async ({ agent_id, files }) => {
   const searchParameter = { id: agent_id };
-  const agent = await getAgent(searchParameter);
-
-  if (!agent) {
-    throw new Error('Agent not found for removing resource files');
-  }
-
-  const tool_resources = { ...agent.tool_resources } || {};
 
+  // associate each tool resource with the respective file ids array
   const filesByResource = files.reduce((acc, { tool_resource, file_id }) => {
     if (!acc[tool_resource]) {
-      acc[tool_resource] = new Set();
+      acc[tool_resource] = [];
     }
-    acc[tool_resource].add(file_id);
+    acc[tool_resource].push(file_id);
     return acc;
   }, {});
 
+  // build the update aggregation pipeline wich removes file ids from tool resources array
+  // and eventually deletes empty tool resources
+  const updateData = [];
   Object.entries(filesByResource).forEach(([resource, fileIds]) => {
-    if (tool_resources[resource] && tool_resources[resource].file_ids) {
-      tool_resources[resource].file_ids = tool_resources[resource].file_ids.filter(
-        (id) => !fileIds.has(id),
-      );
+    const toolResourcePath = `tool_resources.${resource}`;
+    const fileIdsPath = `${toolResourcePath}.file_ids`;
 
-      if (tool_resources[resource].file_ids.length === 0) {
-        delete tool_resources[resource];
-      }
-    }
+    // file ids removal stage
+    updateData.push({
+      $set: {
+        [fileIdsPath]: {
+          $filter: {
+            input: `$${fileIdsPath}`,
+            cond: { $not: [{ $in: ['$$this', fileIds] }] },
+          },
+        },
+      },
+    });
+
+    // empty tool resource deletion stage
+    updateData.push({
+      $set: {
+        [toolResourcePath]: {
+          $cond: [{ $eq: [`$${fileIdsPath}`, []] }, '$$REMOVE', `$${toolResourcePath}`],
+        },
+      },
+    });
   });
 
-  const updateData = { tool_resources };
-  return await updateAgent(searchParameter, updateData);
+  // return the updated agent or throw if no agent matches
+  const updatedAgent = await updateAgent(searchParameter, updateData);
+  if (updatedAgent) {
+    return updatedAgent;
+  } else {
+    throw new Error('Agent not found for removing resource files');
+  }
 };
 
 /**