librechat-balances/services/balance.js

const { ObjectId } = require('mongodb');

async function getStats(db) {
  const balances = db.collection('balances');
  const users = db.collection('users');

  const [totalUsers, balanceAgg] = await Promise.all([
    users.countDocuments(),
    balances.aggregate([
      {
        $group: {
          _id: null,
          totalCredits: { $sum: '$tokenCredits' },
          avgCredits: { $avg: '$tokenCredits' },
          minCredits: { $min: '$tokenCredits' },
          maxCredits: { $max: '$tokenCredits' },
          usersWithBalance: { $sum: 1 },
          zeroBalance: {
            $sum: { $cond: [{ $lte: ['$tokenCredits', 0] }, 1, 0] },
          },
        },
      },
    ]).toArray(),
  ]);

  const stats = balanceAgg[0] || {
    totalCredits: 0,
    avgCredits: 0,
    minCredits: 0,
    maxCredits: 0,
    usersWithBalance: 0,
    zeroBalance: 0,
  };

  return {
    totalUsers,
    usersWithBalance: stats.usersWithBalance,
    usersWithoutBalance: totalUsers - stats.usersWithBalance,
    zeroBalance: stats.zeroBalance,
    totalCredits: stats.totalCredits,
    avgCredits: Math.round(stats.avgCredits || 0),
    minCredits: stats.minCredits || 0,
    maxCredits: stats.maxCredits || 0,
  };
}

async function getAllBalances(db, page, limit) {
  const skip = (page - 1) * limit;

  const [results, total] = await Promise.all([
    db.collection('balances').aggregate([
      {
        $lookup: {
          from: 'users',
          localField: 'user',
          foreignField: '_id',
          as: 'userInfo',
        },
      },
      { $unwind: { path: '$userInfo', preserveNullAndEmptyArrays: true } },
      {
        $project: {
          userId: '$user',
          tokenCredits: 1,
          email: '$userInfo.email',
          name: '$userInfo.name',
          username: '$userInfo.username',
        },
      },
      { $sort: { tokenCredits: -1 } },
      { $skip: skip },
      { $limit: limit },
    ]).toArray(),
    db.collection('balances').countDocuments(),
  ]);

  return {
    users: results,
    total,
    page,
    totalPages: Math.ceil(total / limit),
  };
}

async function getLowBalances(db, threshold) {
  return db.collection('balances').aggregate([
    { $match: { tokenCredits: { $lt: threshold } } },
    {
      $lookup: {
        from: 'users',
        localField: 'user',
        foreignField: '_id',
        as: 'userInfo',
      },
    },
    { $unwind: { path: '$userInfo', preserveNullAndEmptyArrays: true } },
    {
      $project: {
        userId: '$user',
        tokenCredits: 1,
        email: '$userInfo.email',
        name: '$userInfo.name',
        username: '$userInfo.username',
      },
    },
    { $sort: { tokenCredits: 1 } },
  ]).toArray();
}

async function searchUsers(db, query) {
  const regex = new RegExp(query, 'i');

  const users = await db.collection('users').find({
    $or: [
      { email: regex },
      { name: regex },
      { username: regex },
    ],
  }, {
    projection: { _id: 1, email: 1, name: 1, username: 1 },
  }).limit(20).toArray();

  if (users.length === 0) return [];

  const userIds = users.map(u => u._id);
  const balances = await db.collection('balances').find({
    user: { $in: userIds },
  }).toArray();

  const balanceMap = new Map();
  for (const b of balances) {
    balanceMap.set(b.user.toString(), b.tokenCredits);
  }

  return users.map(u => ({
    userId: u._id,
    email: u.email,
    name: u.name,
    username: u.username,
    tokenCredits: balanceMap.get(u._id.toString()) ?? 0,
  }));
}

async function getUserBalance(db, userId) {
  let objectId;
  try {
    objectId = new ObjectId(userId);
  } catch {
    return null;
  }

  const [user, balance] = await Promise.all([
    db.collection('users').findOne({ _id: objectId }, {
      projection: { email: 1, name: 1, username: 1 },
    }),
    db.collection('balances').findOne({ user: objectId }),
  ]);

  if (!user) return null;

  return {
    userId: user._id,
    email: user.email,
    name: user.name,
    username: user.username,
    tokenCredits: balance?.tokenCredits ?? 0,
  };
}

async function setBalance(db, userId, amount) {
  const objectId = new ObjectId(userId);
  await db.collection('balances').updateOne(
    { user: objectId },
    { $set: { tokenCredits: amount } },
    { upsert: true },
  );
  return getUserBalance(db, userId);
}

async function addBalance(db, userId, amount) {
  const objectId = new ObjectId(userId);
  await db.collection('balances').updateOne(
    { user: objectId },
    { $inc: { tokenCredits: amount } },
    { upsert: true },
  );
  return getUserBalance(db, userId);
}

async function addToAll(db, amount) {
  const result = await db.collection('balances').updateMany(
    {},
    { $inc: { tokenCredits: amount } },
  );
  return { modifiedCount: result.modifiedCount };
}

async function setAll(db, amount) {
  const result = await db.collection('balances').updateMany(
    {},
    { $set: { tokenCredits: amount } },
  );
  return { modifiedCount: result.modifiedCount };
}

module.exports = {
  getStats,
  getAllBalances,
  getLowBalances,
  searchUsers,
  getUserBalance,
  setBalance,
  addBalance,
  addToAll,
  setAll,
};