Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
155 changes: 155 additions & 0 deletions migrations/mongo/agentLimitMigration.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,155 @@
const { MongoClient } = require("mongodb");
const dotenv = require("dotenv");
dotenv.config();

/**
* Migration script for agent_limit field reorganization
*
* This script moves the following fields from root level into a nested agent_limit object:
* - bridge_limit → agent_limit.limit
* - bridge_usage → agent_limit.usage
* - bridge_limit_reset_period → agent_limit.reset_period
* - bridge_limit_start_date → agent_limit.start_date
*
*/

async function migrateAgentLimitFields() {
const client = new MongoClient(process.env.MONGODB_CONNECTION_URI);

try {
console.log("Starting agent_limit field migration...");

await client.connect();
console.log("Connected to MongoDB");

const db = client.db();
const collectionName = "configurations";

console.log(`\nProcessing ${collectionName} collection...`);

const collection = db.collection(collectionName);

// Find ALL documents in the collection to ensure everyone gets agent_limit field
const documents = await collection.find({}).toArray();
console.log(`Found ${documents.length} total documents in ${collectionName}`);

if (documents.length === 0) {
console.log(`No documents found in ${collectionName}`);
return;
}

// Process in batches
const batchSize = 100;
let processedCount = 0;

for (let i = 0; i < documents.length; i += batchSize) {
const batch = documents.slice(i, i + batchSize);
const bulkOps = [];

for (const doc of batch) {
// Check if document already has agent_limit field
if (doc.agent_limit) {
console.log(`Document ${doc._id} already has agent_limit field, skipping...`);
continue;
}

// Build the agent_limit object with new field names
// Use existing bridge_* values if they exist, otherwise use defaults
const agent_limit = {
limit: doc.bridge_limit || 0,
usage: doc.bridge_usage || 0,
reset_period: doc.bridge_limit_reset_period || "monthly",
start_date: doc.bridge_limit_start_date || new Date()
};

// Create update operation
const updateOp = {
updateOne: {
filter: { _id: doc._id },
update: {
$set: {
agent_limit: agent_limit
}
}
}
};

// Only unset old fields if they exist
if (
doc.bridge_limit !== undefined ||
doc.bridge_usage !== undefined ||
doc.bridge_limit_reset_period !== undefined ||
doc.bridge_limit_start_date !== undefined
) {
updateOp.updateOne.update.$unset = {
bridge_limit: 1,
bridge_usage: 1,
bridge_limit_reset_period: 1,
bridge_limit_start_date: 1
};
}

bulkOps.push(updateOp);
}

// Execute batch update
if (bulkOps.length > 0) {
const result = await collection.bulkWrite(bulkOps);
processedCount += result.modifiedCount;
console.log(
`Processed batch ${Math.floor(i / batchSize) + 1}/${Math.ceil(documents.length / batchSize)}: ${result.modifiedCount} documents updated`
);
}
}

console.log(`Completed ${collectionName}: ${processedCount} documents updated`);

// Verify migration - check for documents missing agent_limit field
const missingAgentLimitDocs = await collection.countDocuments({ agent_limit: { $exists: false } });
if (missingAgentLimitDocs === 0) {
console.log(`Verification passed: All documents now have agent_limit field`);
} else {
console.log(` Verification warning: ${missingAgentLimitDocs} documents still missing agent_limit field`);
}

// Check if any old bridge_* fields remain
const verifyQuery = {
$or: [
{ bridge_limit: { $exists: true } },
{ bridge_usage: { $exists: true } },
{ bridge_limit_reset_period: { $exists: true } },
{ bridge_limit_start_date: { $exists: true } }
]
};

const remainingDocs = await collection.countDocuments(verifyQuery);
if (remainingDocs === 0) {
console.log(`Verification passed: All old bridge_* fields removed`);
} else {
console.log(` Verification warning: ${remainingDocs} documents still have old bridge_* fields`);
}

// Check agent_limit fields
const agentLimitDocs = await collection.countDocuments({ agent_limit: { $exists: true } });
console.log(` ${collectionName} now has ${agentLimitDocs} documents with agent_limit field`);

console.log("\nAgent_limit field migration completed successfully!");
} catch (error) {
console.error("Migration failed:", error);
throw error;
} finally {
await client.close();
console.log("MongoDB connection closed");
}
}

// Run the migration
migrateAgentLimitFields()
.then(() => {
console.log("Migration script completed");
process.exit(0);
})
.catch((error) => {
console.error("Migration script failed:", error);
process.exit(1);
});
31 changes: 16 additions & 15 deletions src/controllers/agentConfig.controller.js
Original file line number Diff line number Diff line change
Expand Up @@ -192,11 +192,6 @@ const createAgentController = async (req, res, next) => {
}
}

const agent_limit = agents.bridge_limit;
const agent_usage = agents.bridge_usage;
const agent_limit_reset_period = agents.bridge_limit_reset_period;
const agent_limit_start_date = agents.bridge_limit_start_date;

const useAiData = purpose && Object.keys(agent_data).length > 0;
const aiVal = (aiField, fallback) => (useAiData ? (aiField ?? fallback) : fallback);

Expand All @@ -212,10 +207,7 @@ const createAgentController = async (req, res, next) => {
folder_id: folder_id,
user_id: user_id,
fall_back: aiVal(agent_data?.fall_back, fall_back),
bridge_limit: agent_limit,
bridge_usage: agent_usage,
bridge_limit_reset_period: agent_limit_reset_period,
bridge_limit_start_date: agent_limit_start_date,
agent_limit: agents.agent_limit,
bridge_status: 1,
createdAt: new Date(),
updatedAt: new Date(),
Expand Down Expand Up @@ -342,11 +334,20 @@ const updateAgentController = async (req, res, next) => {
}
}

if (body.bridge_limit !== undefined) update_fields.bridge_limit = body.bridge_limit;
if (body.bridge_usage !== undefined) update_fields.bridge_usage = body.bridge_usage;
if (body.bridge_limit_reset_period !== undefined) {
update_fields.bridge_limit_reset_period = body.bridge_limit_reset_period;
update_fields.bridge_limit_start_date = new Date();
// Handle agent_limit updates
if (body.agent_limit !== undefined) {
update_fields.agent_limit = {};

if (body.agent_limit.limit !== undefined) {
update_fields.agent_limit.limit = body.agent_limit.limit;
}
if (body.agent_limit.usage !== undefined) {
update_fields.agent_limit.usage = body.agent_limit.usage;
}
if (body.agent_limit.reset_period !== undefined) {
update_fields.agent_limit.reset_period = body.agent_limit.reset_period;
update_fields.agent_limit.start_date = new Date();
}
}

if (page_config) update_fields.page_config = page_config;
Expand Down Expand Up @@ -531,7 +532,7 @@ const updateAgentController = async (req, res, next) => {
await addBulkUserEntries(user_history);

try {
await purgeRelatedBridgeCaches(agent_id, body.bridge_usage !== undefined ? body.bridge_usage : -1);
await purgeRelatedBridgeCaches(agent_id, body.agent_limit?.usage !== undefined ? body.agent_limit.usage : -1);
} catch (e) {
console.error(`Failed clearing agent related cache on update: ${e}`);
}
Expand Down
5 changes: 1 addition & 4 deletions src/controllers/template.controller.js
Original file line number Diff line number Diff line change
Expand Up @@ -29,10 +29,7 @@ const FILTER_BRIDGE_EXCLUDE_KEYS = new Set([
"total_tokens",
"prompt_total_tokens",
"prompt_enhancer_percentage",
"bridge_usage",
"bridge_limit",
"bridge_limit_reset_period",
"bridge_limit_start_date",
"agent_limit",
"last_used",
"responseIds",
"__v",
Expand Down
5 changes: 1 addition & 4 deletions src/db_services/configuration.service.js
Original file line number Diff line number Diff line change
Expand Up @@ -1146,10 +1146,7 @@ const getAllAgentsInOrg = async (org_id, folder_id, user_id, isEmbedUser) => {
connected_agent_details: 1,
bridge_summary: 1,
deletedAt: 1,
bridge_limit: 1,
bridge_usage: 1,
bridge_limit_reset_period: 1,
bridge_limit_start_date: 1,
agent_limit: 1,
Comment thread
Anushtha-Rathore marked this conversation as resolved.
last_used: 1,
variables_path: 1,
users: 1,
Expand Down
34 changes: 18 additions & 16 deletions src/mongoModel/Configuration.model.js
Original file line number Diff line number Diff line change
Expand Up @@ -255,22 +255,24 @@ const configuration = new mongoose.Schema({
type: Date,
default: null
},
bridge_limit: {
type: Number,
default: 0
},
bridge_usage: {
type: Number,
default: 0
},
bridge_limit_reset_period: {
type: String,
enum: ["monthly", "weekly", "daily"],
default: "monthly"
},
bridge_limit_start_date: {
type: Date,
default: Date.now
agent_limit: {
limit: {
type: Number,
default: 0
},
usage: {
type: Number,
default: 0
},
reset_period: {
type: String,
enum: ["monthly", "weekly", "daily"],
default: "monthly"
},
start_date: {
type: Date,
default: Date.now
}
},
last_used: {
type: Date,
Expand Down
20 changes: 12 additions & 8 deletions src/validation/joi_validation/agentConfig.validation.js
Original file line number Diff line number Diff line change
Expand Up @@ -9,10 +9,12 @@ const createBridgeSchema = Joi.object({
.pattern(/^[0-9a-fA-F]{24}$/)
.optional(),
bridgeType: Joi.string().valid("api", "chatbot").optional().default("api"),
bridge_limit: Joi.number().min(0).optional(),
bridge_usage: Joi.number().min(0).optional(),
bridge_limit_reset_period: Joi.string().valid("monthly", "weekly", "daily").optional(),
bridge_limit_start_date: Joi.date().optional()
agent_limit: Joi.object({
limit: Joi.number().min(0).optional(),
usage: Joi.number().min(0).optional(),
reset_period: Joi.string().valid("monthly", "weekly", "daily").optional(),
start_date: Joi.date().optional()
}).optional()
}).unknown(true); // Allow additional fields that might be added dynamically

const updateBridgeSchema = Joi.object({
Expand Down Expand Up @@ -92,10 +94,12 @@ const updateBridgeSchema = Joi.object({
guardrails: Joi.object().optional(),
web_search_filters: Joi.alternatives().try(Joi.array().items(Joi.string()), Joi.object()).optional(),
gtwy_web_search_filters: Joi.alternatives().try(Joi.array().items(Joi.string()), Joi.object()).optional(),
bridge_limit: Joi.number().min(0).optional(),
bridge_usage: Joi.number().min(0).optional(),
bridge_limit_reset_period: Joi.string().valid("monthly", "weekly", "daily").optional(),
bridge_limit_start_date: Joi.date().optional(),
agent_limit: Joi.object({
limit: Joi.number().min(0).optional(),
usage: Joi.number().min(0).optional(),
reset_period: Joi.string().valid("monthly", "weekly", "daily").optional(),
start_date: Joi.date().optional()
}).optional(),
page_config: Joi.object().optional(),
variables_path: Joi.object().optional(),
built_in_tools_data: Joi.object({
Expand Down