diff --git a/api/models/Group.js b/api/models/Group.js new file mode 100644 index 0000000000..69a8a53a6e --- /dev/null +++ b/api/models/Group.js @@ -0,0 +1,6 @@ +const mongoose = require('mongoose'); +const groupSchema = require('~/models/schema/groupSchema'); + +const Group = mongoose.model('Group', groupSchema); + +module.exports = Group; diff --git a/api/models/index.js b/api/models/index.js index 73cfa1c96c..89f6ffb8f2 100644 --- a/api/models/index.js +++ b/api/models/index.js @@ -40,6 +40,7 @@ const { getPreset, getPresets, savePreset, deletePresets } = require('./Preset') const { createToken, findToken, updateToken, deleteTokens } = require('./Token'); const Balance = require('./Balance'); const User = require('./User'); +const Group = require('./Group'); const Key = require('./Key'); module.exports = { @@ -92,6 +93,7 @@ module.exports = { countActiveSessions, User, + Group, Key, Balance, }; diff --git a/api/models/schema/groupSchema.js b/api/models/schema/groupSchema.js new file mode 100644 index 0000000000..b2dc988c7e --- /dev/null +++ b/api/models/schema/groupSchema.js @@ -0,0 +1,29 @@ +const mongoose = require('mongoose'); + +const groupSchema = new mongoose.Schema( + { + name: { + type: String, + required: true, + unique: true, + }, + description: { + type: String, + }, + allowedEndpoints: [ + { + type: String, + required: true, + + }, + ], + allowedModels: [ + { + type: String, + }, + ], + }, + { timestamps: true }, +); + +module.exports = groupSchema; \ No newline at end of file diff --git a/api/models/schema/userSchema.js b/api/models/schema/userSchema.js index bebc7fea1e..32e8a62314 100644 --- a/api/models/schema/userSchema.js +++ b/api/models/schema/userSchema.js @@ -143,6 +143,12 @@ const userSchema = mongoose.Schema( type: Boolean, default: false, }, + groups: [ + { + type: mongoose.Schema.Types.ObjectId, + ref: 'Group', + }, + ], }, { timestamps: true }, diff --git a/config/assign-group.js b/config/assign-group.js new file mode 100644 index 0000000000..d6f407f944 --- /dev/null +++ b/config/assign-group.js @@ -0,0 +1,62 @@ +const path = require('path'); +require('module-alias')({ base: path.resolve(__dirname, '..', 'api') }); +const { askQuestion, silentExit } = require('./helpers'); +const User = require('~/models/User'); +const Group = require('~/models/Group'); +const connect = require('./connect'); + +(async () => { + await connect(); + + console.purple('---------------------------------------'); + console.purple('Assign a Group to a User'); + console.purple('---------------------------------------'); + + // Read arguments from CLI or prompt the user + let userEmail = process.argv[2] || (await askQuestion('User email: ')); + let groupName = process.argv[3] || (await askQuestion('Group name to assign: ')); + + // Validate email format + if (!userEmail.includes('@')) { + console.red('Error: Invalid email address!'); + silentExit(1); + } + + // Find the group by name + const group = await Group.findOne({ name: groupName }); + if (!group) { + console.red('Error: No group with that name was found!'); + silentExit(1); + } + + // Find the user by email + const user = await User.findOne({ email: userEmail }); + if (!user) { + console.red('Error: No user with that email was found!'); + silentExit(1); + } + console.purple(`Found user: ${user.email}`); + + // Assign the group to the user if not already assigned + try { + if (!user.groups) { + user.groups = []; + } + if (!user.groups.includes(group._id)) { + user.groups.push(group._id); + await user.save(); + } + } catch (error) { + console.red('Error assigning group to user: ' + error.message); + silentExit(1); + } + + console.green(`User ${user.email} successfully assigned to group ${group.name}!`); + silentExit(0); +})(); + +process.on('uncaughtException', (err) => { + console.error('There was an uncaught error:'); + console.error(err); + process.exit(1); +}); \ No newline at end of file diff --git a/config/create-group.js b/config/create-group.js new file mode 100644 index 0000000000..5a70fb08bc --- /dev/null +++ b/config/create-group.js @@ -0,0 +1,60 @@ +const path = require('path'); +require('module-alias')({ base: path.resolve(__dirname, '..', 'api') }); +const { askQuestion, silentExit } = require('./helpers'); +const Group = require('~/models/Group'); +const connect = require('./connect'); + +(async () => { + await connect(); + + console.purple('---------------------------------------'); + console.purple('Create a New Group'); + console.purple('---------------------------------------'); + + // Read arguments from CLI or prompt the user + let groupName = process.argv[2] || (await askQuestion('Group name: ')); + let groupDescription = + process.argv[3] || (await askQuestion('Group description (optional): ')); + let allowedEndpointsInput = + process.argv[4] || + (await askQuestion( + 'Allowed endpoints (comma separated, e.g., "assistants,agents", or enter "*" for all): ', + )); + let allowedModelsInput = + process.argv[5] || + (await askQuestion( + 'Allowed models (comma separated, e.g., "gpt-4,chatgpt-4o-latest", or enter "*" for all): ', + )); + + // Process the comma-separated inputs into arrays + const allowedEndpoints = allowedEndpointsInput + .split(',') + .map((s) => s.trim()) + .filter((s) => s); + const allowedModels = allowedModelsInput + .split(',') + .map((s) => s.trim()) + .filter((s) => s); + + // Create the group document + let group; + try { + group = await Group.create({ + name: groupName, + description: groupDescription, + allowedEndpoints, + allowedModels, + }); + } catch (error) { + console.red('Error creating group: ' + error.message); + silentExit(1); + } + console.green(`Group created successfully with id: ${group._id}`); + silentExit(0); +})(); + +process.on('uncaughtException', (err) => { + console.error('There was an uncaught error:'); + console.error(err); + process.exit(1); +}); \ No newline at end of file diff --git a/package.json b/package.json index 99f9fd908e..419401aa95 100644 --- a/package.json +++ b/package.json @@ -32,6 +32,8 @@ "reset-password": "node config/reset-password.js", "ban-user": "node config/ban-user.js", "delete-user": "node config/delete-user.js", + "create-group": "node config/create-group.js", + "assign-group": "node config/assign-group.js", "update-banner": "node config/update-banner.js", "delete-banner": "node config/delete-banner.js", "backend": "cross-env NODE_ENV=production node api/server/index.js",