feat: agent share global methods/controllers

This commit is contained in:
Danny Avila 2024-09-03 18:17:44 -04:00
parent 72dfcb9dc9
commit 3f694f2e25
No known key found for this signature in database
GPG key ID: 2DD9CC89B9B50364
7 changed files with 129 additions and 13 deletions

View file

@ -1,4 +1,11 @@
const mongoose = require('mongoose'); 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 agentSchema = require('./schema/agent');
const Agent = mongoose.model('agent', agentSchema); const Agent = mongoose.model('agent', agentSchema);
@ -48,7 +55,11 @@ const updateAgent = async (searchParameter, updateData, session = null) => {
* @returns {Promise<void>} Resolves when the agent has been successfully deleted. * @returns {Promise<void>} Resolves when the agent has been successfully deleted.
*/ */
const deleteAgent = async (searchParameter) => { 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<Object>} A promise that resolves to an object containing the agents data and pagination info. * @returns {Promise<Object>} A promise that resolves to an object containing the agents data and pagination info.
*/ */
const getListAgents = async (searchParameter) => { 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, id: 1,
name: 1, name: 1,
avatar: 1, avatar: 1,
}).lean(); }).lean();
const hasMore = agents.length > 0; const hasMore = agents.length > 0;
const firstId = agents.length > 0 ? agents[0].id : null; const firstId = agents.length > 0 ? agents[0].id : null;
const lastId = agents.length > 0 ? agents[agents.length - 1].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<MongoAgent>} 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 = { module.exports = {
createAgent, createAgent,
getAgent, getAgent,
updateAgent, updateAgent,
deleteAgent, deleteAgent,
getListAgents, getListAgents,
updateAgentProjects,
}; };

View file

@ -1,9 +1,7 @@
const { model } = require('mongoose'); 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 projectSchema = require('~/models/schema/projectSchema');
const { GLOBAL_PROJECT_NAME } = Constants;
const Project = model('Project', projectSchema); const Project = model('Project', projectSchema);
/** /**
@ -84,10 +82,55 @@ const removeGroupFromAllProjects = async (promptGroupId) => {
await Project.updateMany({}, { $pull: { promptGroupIds: 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<MongoProject>} 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<MongoProject>} 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<void>}
*/
const removeAgentFromAllProjects = async (agentId) => {
await Project.updateMany({}, { $pull: { agentIds: agentId } });
};
module.exports = { module.exports = {
getProjectById, getProjectById,
getProjectByName, getProjectByName,
/* prompts */
addGroupIdsToProject, addGroupIdsToProject,
removeGroupIdsFromProject, removeGroupIdsFromProject,
removeGroupFromAllProjects, removeGroupFromAllProjects,
/* agents */
addAgentIdsToProject,
removeAgentIdsFromProject,
removeAgentFromAllProjects,
}; };

View file

@ -9,8 +9,6 @@ const {
const { Prompt, PromptGroup } = require('./schema/promptSchema'); const { Prompt, PromptGroup } = require('./schema/promptSchema');
const { logger } = require('~/config'); const { logger } = require('~/config');
const { GLOBAL_PROJECT_NAME } = Constants;
/** /**
* Create a pipeline for the aggregation to get prompt groups * Create a pipeline for the aggregation to get prompt groups
* @param {Object} query * @param {Object} query
@ -125,7 +123,7 @@ const getAllPromptGroups = async (req, filter) => {
let combinedQuery = query; let combinedQuery = query;
if (searchShared) { 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) { if (project && project.promptGroupIds.length > 0) {
const projectQuery = { _id: { $in: project.promptGroupIds }, ...query }; const projectQuery = { _id: { $in: project.promptGroupIds }, ...query };
delete projectQuery.author; delete projectQuery.author;
@ -179,7 +177,7 @@ const getPromptGroups = async (req, filter) => {
if (searchShared) { if (searchShared) {
// const projects = req.user.projects || []; // TODO: handle multiple projects // 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) { if (project && project.promptGroupIds.length > 0) {
const projectQuery = { _id: { $in: project.promptGroupIds }, ...query }; const projectQuery = { _id: { $in: project.promptGroupIds }, ...query };
delete projectQuery.author; delete projectQuery.author;

View file

@ -57,6 +57,11 @@ const agentSchema = mongoose.Schema(
ref: 'User', ref: 'User',
required: true, required: true,
}, },
projectIds: {
type: [mongoose.Schema.Types.ObjectId],
ref: 'Project',
index: true,
},
}, },
{ {
timestamps: true, timestamps: true,

View file

@ -21,6 +21,11 @@ const projectSchema = new Schema(
ref: 'PromptGroup', ref: 'PromptGroup',
default: [], default: [],
}, },
agentIds: {
type: [String],
ref: 'Agent',
default: [],
},
}, },
{ {
timestamps: true, timestamps: true,

View file

@ -9,6 +9,7 @@ const {
} = require('~/models/Agent'); } = require('~/models/Agent');
const { getStrategyFunctions } = require('~/server/services/Files/strategies'); const { getStrategyFunctions } = require('~/server/services/Files/strategies');
const { uploadImageBuffer } = require('~/server/services/Files/process'); const { uploadImageBuffer } = require('~/server/services/Files/process');
const { updateAgentProjects } = require('~/models/Agent');
const { deleteFileByFilter } = require('~/models/File'); const { deleteFileByFilter } = require('~/models/File');
const { logger } = require('~/config'); const { logger } = require('~/config');
@ -82,7 +83,14 @@ const getAgentHandler = async (req, res) => {
const updateAgentHandler = async (req, res) => { const updateAgentHandler = async (req, res) => {
try { try {
const id = req.params.id; 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); return res.json(updatedAgent);
} catch (error) { } catch (error) {
logger.error('[/Agents/:id] Error updating Agent', error); logger.error('[/Agents/:id] Error updating Agent', error);

View file

@ -6,8 +6,6 @@ const { isEnabled } = require('~/server/utils');
const { getLogStores } = require('~/cache'); const { getLogStores } = require('~/cache');
const { logger } = require('~/config'); const { logger } = require('~/config');
const { GLOBAL_PROJECT_NAME } = Constants;
const router = express.Router(); const router = express.Router();
const emailLoginEnabled = const emailLoginEnabled =
process.env.ALLOW_EMAIL_LOGIN === undefined || isEnabled(process.env.ALLOW_EMAIL_LOGIN); 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; 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(); const ldap = getLdapConfig();