diff --git a/api/models/Agent.js b/api/models/Agent.js index 1ee783b101..cffced3e95 100644 --- a/api/models/Agent.js +++ b/api/models/Agent.js @@ -1,4 +1,11 @@ const mongoose = require('mongoose'); +const { GLOBAL_PROJECT_NAME } = require('librechat-data-provider').Constants; +const { + getProjectByName, + addAgentIdsToProject, + removeAgentIdsFromProject, + removeAgentFromAllProjects, +} = require('./Project'); const agentSchema = require('./schema/agent'); const Agent = mongoose.model('agent', agentSchema); @@ -48,7 +55,11 @@ const updateAgent = async (searchParameter, updateData, session = null) => { * @returns {Promise} Resolves when the agent has been successfully deleted. */ const deleteAgent = async (searchParameter) => { - return await Agent.findOneAndDelete(searchParameter); + const agent = await Agent.findOneAndDelete(searchParameter); + if (agent) { + await removeAgentFromAllProjects(agent.id); + } + return agent; }; /** @@ -58,11 +69,24 @@ const deleteAgent = async (searchParameter) => { * @returns {Promise} A promise that resolves to an object containing the agents data and pagination info. */ const getListAgents = async (searchParameter) => { - const agents = await Agent.find(searchParameter, { + const { author, ...otherParams } = searchParameter; + + let query = { + $or: [{ author }, { projectIds: { $exists: true, $ne: [], $not: { $size: 0 } } }], + ...otherParams, + }; + + const globalProject = await getProjectByName(GLOBAL_PROJECT_NAME, 'agentIds'); + if (globalProject && globalProject.agentIds.length > 0) { + query.$or.push({ _id: { $in: globalProject.agentIds } }); + } + + const agents = await Agent.find(query, { id: 1, name: 1, avatar: 1, }).lean(); + const hasMore = agents.length > 0; const firstId = agents.length > 0 ? agents[0].id : null; const lastId = agents.length > 0 ? agents[agents.length - 1].id : null; @@ -75,10 +99,45 @@ const getListAgents = async (searchParameter) => { }; }; +/** + * Updates the projects associated with an agent, adding and removing project IDs as specified. + * This function also updates the corresponding projects to include or exclude the agent ID. + * + * @param {string} agentId - The ID of the agent to update. + * @param {string[]} [projectIds] - Array of project IDs to add to the agent. + * @param {string[]} [removeProjectIds] - Array of project IDs to remove from the agent. + * @returns {Promise} The updated agent document. + * @throws {Error} If there's an error updating the agent or projects. + */ +const updateAgentProjects = async (agentId, projectIds, removeProjectIds) => { + const updateOps = {}; + + if (removeProjectIds && removeProjectIds.length > 0) { + for (const projectId of removeProjectIds) { + await removeAgentIdsFromProject(projectId, [agentId]); + } + updateOps.$pull = { projectIds: { $in: removeProjectIds } }; + } + + if (projectIds && projectIds.length > 0) { + for (const projectId of projectIds) { + await addAgentIdsToProject(projectId, [agentId]); + } + updateOps.$addToSet = { projectIds: { $each: projectIds } }; + } + + if (Object.keys(updateOps).length === 0) { + return await Agent.findById(agentId).lean(); + } + + return await Agent.findByIdAndUpdate(agentId, updateOps, { new: true, lean: true }); +}; + module.exports = { createAgent, getAgent, updateAgent, deleteAgent, getListAgents, + updateAgentProjects, }; diff --git a/api/models/Project.js b/api/models/Project.js index 32e44e9e07..17ef3093a5 100644 --- a/api/models/Project.js +++ b/api/models/Project.js @@ -1,9 +1,7 @@ const { model } = require('mongoose'); -const { Constants } = require('librechat-data-provider'); +const { GLOBAL_PROJECT_NAME } = require('librechat-data-provider').Constants; const projectSchema = require('~/models/schema/projectSchema'); -const { GLOBAL_PROJECT_NAME } = Constants; - const Project = model('Project', projectSchema); /** @@ -84,10 +82,55 @@ const removeGroupFromAllProjects = async (promptGroupId) => { await Project.updateMany({}, { $pull: { promptGroupIds: promptGroupId } }); }; +/** + * Add an array of agent IDs to a project's agentIds array, ensuring uniqueness. + * + * @param {string} projectId - The ID of the project to update. + * @param {string[]} agentIds - The array of agent IDs to add to the project. + * @returns {Promise} The updated project document. + */ +const addAgentIdsToProject = async function (projectId, agentIds) { + return await Project.findByIdAndUpdate( + projectId, + { $addToSet: { agentIds: { $each: agentIds } } }, + { new: true }, + ); +}; + +/** + * Remove an array of agent IDs from a project's agentIds array. + * + * @param {string} projectId - The ID of the project to update. + * @param {string[]} agentIds - The array of agent IDs to remove from the project. + * @returns {Promise} The updated project document. + */ +const removeAgentIdsFromProject = async function (projectId, agentIds) { + return await Project.findByIdAndUpdate( + projectId, + { $pull: { agentIds: { $in: agentIds } } }, + { new: true }, + ); +}; + +/** + * Remove an agent ID from all projects. + * + * @param {string} agentId - The ID of the agent to remove from projects. + * @returns {Promise} + */ +const removeAgentFromAllProjects = async (agentId) => { + await Project.updateMany({}, { $pull: { agentIds: agentId } }); +}; + module.exports = { getProjectById, getProjectByName, + /* prompts */ addGroupIdsToProject, removeGroupIdsFromProject, removeGroupFromAllProjects, + /* agents */ + addAgentIdsToProject, + removeAgentIdsFromProject, + removeAgentFromAllProjects, }; diff --git a/api/models/Prompt.js b/api/models/Prompt.js index 8a54cdccc3..548589b4d7 100644 --- a/api/models/Prompt.js +++ b/api/models/Prompt.js @@ -9,8 +9,6 @@ const { const { Prompt, PromptGroup } = require('./schema/promptSchema'); const { logger } = require('~/config'); -const { GLOBAL_PROJECT_NAME } = Constants; - /** * Create a pipeline for the aggregation to get prompt groups * @param {Object} query @@ -125,7 +123,7 @@ const getAllPromptGroups = async (req, filter) => { let combinedQuery = query; if (searchShared) { - const project = await getProjectByName(GLOBAL_PROJECT_NAME, 'promptGroupIds'); + const project = await getProjectByName(Constants.GLOBAL_PROJECT_NAME, 'promptGroupIds'); if (project && project.promptGroupIds.length > 0) { const projectQuery = { _id: { $in: project.promptGroupIds }, ...query }; delete projectQuery.author; @@ -179,7 +177,7 @@ const getPromptGroups = async (req, filter) => { if (searchShared) { // const projects = req.user.projects || []; // TODO: handle multiple projects - const project = await getProjectByName(GLOBAL_PROJECT_NAME, 'promptGroupIds'); + const project = await getProjectByName(Constants.GLOBAL_PROJECT_NAME, 'promptGroupIds'); if (project && project.promptGroupIds.length > 0) { const projectQuery = { _id: { $in: project.promptGroupIds }, ...query }; delete projectQuery.author; diff --git a/api/models/schema/agent.js b/api/models/schema/agent.js index 97f0527916..819398ee7c 100644 --- a/api/models/schema/agent.js +++ b/api/models/schema/agent.js @@ -57,6 +57,11 @@ const agentSchema = mongoose.Schema( ref: 'User', required: true, }, + projectIds: { + type: [mongoose.Schema.Types.ObjectId], + ref: 'Project', + index: true, + }, }, { timestamps: true, diff --git a/api/models/schema/projectSchema.js b/api/models/schema/projectSchema.js index 0e27c6a8f9..dfa68a06c2 100644 --- a/api/models/schema/projectSchema.js +++ b/api/models/schema/projectSchema.js @@ -21,6 +21,11 @@ const projectSchema = new Schema( ref: 'PromptGroup', default: [], }, + agentIds: { + type: [String], + ref: 'Agent', + default: [], + }, }, { timestamps: true, diff --git a/api/server/controllers/agents/v1.js b/api/server/controllers/agents/v1.js index 2a9911c541..d1c2a91241 100644 --- a/api/server/controllers/agents/v1.js +++ b/api/server/controllers/agents/v1.js @@ -9,6 +9,7 @@ const { } = require('~/models/Agent'); const { getStrategyFunctions } = require('~/server/services/Files/strategies'); const { uploadImageBuffer } = require('~/server/services/Files/process'); +const { updateAgentProjects } = require('~/models/Agent'); const { deleteFileByFilter } = require('~/models/File'); const { logger } = require('~/config'); @@ -82,7 +83,14 @@ const getAgentHandler = async (req, res) => { const updateAgentHandler = async (req, res) => { try { const id = req.params.id; - const updatedAgent = await updateAgent({ id, author: req.user.id }, req.body); + const { projectIds, removeProjectIds, ...updateData } = req.body; + + const updatedAgent = await updateAgent({ id, author: req.user.id }, updateData); + + if (projectIds || removeProjectIds) { + await updateAgentProjects(id, projectIds, removeProjectIds); + } + return res.json(updatedAgent); } catch (error) { logger.error('[/Agents/:id] Error updating Agent', error); diff --git a/api/server/routes/config.js b/api/server/routes/config.js index f61c64d95f..f6669169ac 100644 --- a/api/server/routes/config.js +++ b/api/server/routes/config.js @@ -6,8 +6,6 @@ const { isEnabled } = require('~/server/utils'); const { getLogStores } = require('~/cache'); const { logger } = require('~/config'); -const { GLOBAL_PROJECT_NAME } = Constants; - const router = express.Router(); const emailLoginEnabled = process.env.ALLOW_EMAIL_LOGIN === undefined || isEnabled(process.env.ALLOW_EMAIL_LOGIN); @@ -34,7 +32,7 @@ router.get('/', async function (req, res) { return today.getMonth() === 1 && today.getDate() === 11; }; - const instanceProject = await getProjectByName(GLOBAL_PROJECT_NAME, '_id'); + const instanceProject = await getProjectByName(Constants.GLOBAL_PROJECT_NAME, '_id'); const ldap = getLdapConfig();