🔧 fix: Handle Concurrent File Mgmt. For Agents (#5159)

* 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
This commit is contained in:
Thinger Soft 2025-01-02 14:29:07 +01:00 committed by GitHub
parent 6c9a468b8e
commit 65b2d647a1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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');
}
};
/**