mirror of
https://github.com/danny-avila/LibreChat.git
synced 2025-12-30 15:18:50 +01:00
💵 feat: Add Automatic Balance Refill (#6452)
* 🚀 feat: Add automatic refill settings to balance schema * 🚀 feat: Refactor balance feature to use global interface configuration * 🚀 feat: Implement auto-refill functionality for balance management * 🚀 feat: Enhance auto-refill logic and configuration for balance management * 🚀 chore: Bump version to 0.7.74 in package.json and package-lock.json * 🚀 chore: Bump version to 0.0.5 in package.json and package-lock.json * 🚀 docs: Update comment for balance settings in librechat.example.yaml * chore: space in `.env.example` * 🚀 feat: Implement balance configuration loading and refactor related components * 🚀 test: Refactor tests to use custom config for balance feature * 🚀 fix: Update balance response handling in Transaction.js to use Balance model * 🚀 test: Update AppService tests to include balance configuration in mock setup * 🚀 test: Enhance AppService tests with complete balance configuration scenarios * 🚀 refactor: Rename balanceConfig to balance and update related tests for clarity * 🚀 refactor: Remove loadDefaultBalance and update balance handling in AppService * 🚀 test: Update AppService tests to reflect new balance structure and defaults * 🚀 test: Mock getCustomConfig in BaseClient tests to control balance configuration * 🚀 test: Add get method to mockCache in OpenAIClient tests for improved cache handling * 🚀 test: Mock getCustomConfig in OpenAIClient tests to control balance configuration * 🚀 test: Remove mock for getCustomConfig in OpenAIClient tests to streamline configuration handling * 🚀 fix: Update balance configuration reference in config.js for consistency * refactor: Add getBalanceConfig function to retrieve balance configuration * chore: Comment out example balance settings in librechat.example.yaml * refactor: Replace getCustomConfig with getBalanceConfig for balance handling * fix: tests * refactor: Replace getBalanceConfig call with balance from request locals * refactor: Update balance handling to use environment variables for configuration * refactor: Replace getBalanceConfig calls with balance from request locals * refactor: Simplify balance configuration logic in getBalanceConfig --------- Co-authored-by: Danny Avila <danny@librechat.ai>
This commit is contained in:
parent
cbba914290
commit
3a62a2633d
24 changed files with 334 additions and 147 deletions
|
|
@ -3,6 +3,40 @@ const { balanceSchema } = require('@librechat/data-schemas');
|
|||
const { getMultiplier } = require('./tx');
|
||||
const { logger } = require('~/config');
|
||||
|
||||
/**
|
||||
* Adds a time interval to a given date.
|
||||
* @param {Date} date - The starting date.
|
||||
* @param {number} value - The numeric value of the interval.
|
||||
* @param {'seconds'|'minutes'|'hours'|'days'|'weeks'|'months'} unit - The unit of time.
|
||||
* @returns {Date} A new Date representing the starting date plus the interval.
|
||||
*/
|
||||
const addIntervalToDate = (date, value, unit) => {
|
||||
const result = new Date(date);
|
||||
switch (unit) {
|
||||
case 'seconds':
|
||||
result.setSeconds(result.getSeconds() + value);
|
||||
break;
|
||||
case 'minutes':
|
||||
result.setMinutes(result.getMinutes() + value);
|
||||
break;
|
||||
case 'hours':
|
||||
result.setHours(result.getHours() + value);
|
||||
break;
|
||||
case 'days':
|
||||
result.setDate(result.getDate() + value);
|
||||
break;
|
||||
case 'weeks':
|
||||
result.setDate(result.getDate() + value * 7);
|
||||
break;
|
||||
case 'months':
|
||||
result.setMonth(result.getMonth() + value);
|
||||
break;
|
||||
default:
|
||||
break;
|
||||
}
|
||||
return result;
|
||||
};
|
||||
|
||||
balanceSchema.statics.check = async function ({
|
||||
user,
|
||||
model,
|
||||
|
|
@ -14,9 +48,20 @@ balanceSchema.statics.check = async function ({
|
|||
}) {
|
||||
const multiplier = getMultiplier({ valueKey, tokenType, model, endpoint, endpointTokenConfig });
|
||||
const tokenCost = amount * multiplier;
|
||||
const { tokenCredits: balance } = (await this.findOne({ user }, 'tokenCredits').lean()) ?? {};
|
||||
|
||||
logger.debug('[Balance.check]', {
|
||||
// Retrieve the complete balance record
|
||||
let record = await this.findOne({ user }).lean();
|
||||
if (!record) {
|
||||
logger.debug('[Balance.check] No balance record found for user', { user });
|
||||
return {
|
||||
canSpend: false,
|
||||
balance: 0,
|
||||
tokenCost,
|
||||
};
|
||||
}
|
||||
let balance = record.tokenCredits;
|
||||
|
||||
logger.debug('[Balance.check] Initial state', {
|
||||
user,
|
||||
model,
|
||||
endpoint,
|
||||
|
|
@ -28,15 +73,31 @@ balanceSchema.statics.check = async function ({
|
|||
endpointTokenConfig: !!endpointTokenConfig,
|
||||
});
|
||||
|
||||
if (!balance) {
|
||||
return {
|
||||
canSpend: false,
|
||||
balance: 0,
|
||||
tokenCost,
|
||||
};
|
||||
// Only perform auto-refill if spending would bring the balance to 0 or below
|
||||
if (balance - tokenCost <= 0 && record.autoRefillEnabled && record.refillAmount > 0) {
|
||||
const lastRefillDate = new Date(record.lastRefill);
|
||||
const nextRefillDate = addIntervalToDate(
|
||||
lastRefillDate,
|
||||
record.refillIntervalValue,
|
||||
record.refillIntervalUnit,
|
||||
);
|
||||
const now = new Date();
|
||||
|
||||
if (now >= nextRefillDate) {
|
||||
record = await this.findOneAndUpdate(
|
||||
{ user },
|
||||
{
|
||||
$inc: { tokenCredits: record.refillAmount },
|
||||
$set: { lastRefill: new Date() },
|
||||
},
|
||||
{ new: true },
|
||||
).lean();
|
||||
balance = record.tokenCredits;
|
||||
logger.debug('[Balance.check] Auto-refill performed', { balance });
|
||||
}
|
||||
}
|
||||
|
||||
logger.debug('[Balance.check]', { tokenCost });
|
||||
logger.debug('[Balance.check] Token cost', { tokenCost });
|
||||
|
||||
return { canSpend: balance >= tokenCost, balance, tokenCost };
|
||||
};
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
const mongoose = require('mongoose');
|
||||
const { isEnabled } = require('~/server/utils/handleText');
|
||||
const { transactionSchema } = require('@librechat/data-schemas');
|
||||
const { getBalanceConfig } = require('~/server/services/Config');
|
||||
const { getMultiplier, getCacheMultiplier } = require('./tx');
|
||||
const { logger } = require('~/config');
|
||||
const Balance = require('./Balance');
|
||||
|
|
@ -37,18 +37,19 @@ transactionSchema.statics.create = async function (txData) {
|
|||
|
||||
await transaction.save();
|
||||
|
||||
if (!isEnabled(process.env.CHECK_BALANCE)) {
|
||||
const balance = await getBalanceConfig();
|
||||
if (!balance?.enabled) {
|
||||
return;
|
||||
}
|
||||
|
||||
let balance = await Balance.findOne({ user: transaction.user }).lean();
|
||||
let balanceResponse = await Balance.findOne({ user: transaction.user }).lean();
|
||||
let incrementValue = transaction.tokenValue;
|
||||
|
||||
if (balance && balance?.tokenCredits + incrementValue < 0) {
|
||||
incrementValue = -balance.tokenCredits;
|
||||
if (balanceResponse && balanceResponse.tokenCredits + incrementValue < 0) {
|
||||
incrementValue = -balanceResponse.tokenCredits;
|
||||
}
|
||||
|
||||
balance = await Balance.findOneAndUpdate(
|
||||
balanceResponse = await Balance.findOneAndUpdate(
|
||||
{ user: transaction.user },
|
||||
{ $inc: { tokenCredits: incrementValue } },
|
||||
{ upsert: true, new: true },
|
||||
|
|
@ -57,7 +58,7 @@ transactionSchema.statics.create = async function (txData) {
|
|||
return {
|
||||
rate: transaction.rate,
|
||||
user: transaction.user.toString(),
|
||||
balance: balance.tokenCredits,
|
||||
balance: balanceResponse.tokenCredits,
|
||||
[transaction.tokenType]: incrementValue,
|
||||
};
|
||||
};
|
||||
|
|
@ -78,18 +79,19 @@ transactionSchema.statics.createStructured = async function (txData) {
|
|||
|
||||
await transaction.save();
|
||||
|
||||
if (!isEnabled(process.env.CHECK_BALANCE)) {
|
||||
const balance = await getBalanceConfig();
|
||||
if (!balance?.enabled) {
|
||||
return;
|
||||
}
|
||||
|
||||
let balance = await Balance.findOne({ user: transaction.user }).lean();
|
||||
let balanceResponse = await Balance.findOne({ user: transaction.user }).lean();
|
||||
let incrementValue = transaction.tokenValue;
|
||||
|
||||
if (balance && balance?.tokenCredits + incrementValue < 0) {
|
||||
incrementValue = -balance.tokenCredits;
|
||||
if (balanceResponse && balanceResponse.tokenCredits + incrementValue < 0) {
|
||||
incrementValue = -balanceResponse.tokenCredits;
|
||||
}
|
||||
|
||||
balance = await Balance.findOneAndUpdate(
|
||||
balanceResponse = await Balance.findOneAndUpdate(
|
||||
{ user: transaction.user },
|
||||
{ $inc: { tokenCredits: incrementValue } },
|
||||
{ upsert: true, new: true },
|
||||
|
|
@ -98,7 +100,7 @@ transactionSchema.statics.createStructured = async function (txData) {
|
|||
return {
|
||||
rate: transaction.rate,
|
||||
user: transaction.user.toString(),
|
||||
balance: balance.tokenCredits,
|
||||
balance: balanceResponse.tokenCredits,
|
||||
[transaction.tokenType]: incrementValue,
|
||||
};
|
||||
};
|
||||
|
|
|
|||
|
|
@ -1,9 +1,13 @@
|
|||
const mongoose = require('mongoose');
|
||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||
const { spendTokens, spendStructuredTokens } = require('./spendTokens');
|
||||
const { getBalanceConfig } = require('~/server/services/Config');
|
||||
const { getMultiplier, getCacheMultiplier } = require('./tx');
|
||||
const { Transaction } = require('./Transaction');
|
||||
const Balance = require('./Balance');
|
||||
const { spendTokens, spendStructuredTokens } = require('./spendTokens');
|
||||
const { getMultiplier, getCacheMultiplier } = require('./tx');
|
||||
|
||||
// Mock the custom config module so we can control the balance flag.
|
||||
jest.mock('~/server/services/Config');
|
||||
|
||||
let mongoServer;
|
||||
|
||||
|
|
@ -20,6 +24,8 @@ afterAll(async () => {
|
|||
|
||||
beforeEach(async () => {
|
||||
await mongoose.connection.dropDatabase();
|
||||
// Default: enable balance updates in tests.
|
||||
getBalanceConfig.mockResolvedValue({ enabled: true });
|
||||
});
|
||||
|
||||
describe('Regular Token Spending Tests', () => {
|
||||
|
|
@ -44,34 +50,22 @@ describe('Regular Token Spending Tests', () => {
|
|||
};
|
||||
|
||||
// Act
|
||||
process.env.CHECK_BALANCE = 'true';
|
||||
await spendTokens(txData, tokenUsage);
|
||||
|
||||
// Assert
|
||||
console.log('Initial Balance:', initialBalance);
|
||||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
console.log('Updated Balance:', updatedBalance.tokenCredits);
|
||||
|
||||
const promptMultiplier = getMultiplier({ model, tokenType: 'prompt' });
|
||||
const completionMultiplier = getMultiplier({ model, tokenType: 'completion' });
|
||||
|
||||
const expectedPromptCost = tokenUsage.promptTokens * promptMultiplier;
|
||||
const expectedCompletionCost = tokenUsage.completionTokens * completionMultiplier;
|
||||
const expectedTotalCost = expectedPromptCost + expectedCompletionCost;
|
||||
const expectedTotalCost = 100 * promptMultiplier + 50 * completionMultiplier;
|
||||
const expectedBalance = initialBalance - expectedTotalCost;
|
||||
|
||||
expect(updatedBalance.tokenCredits).toBeLessThan(initialBalance);
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(expectedBalance, 0);
|
||||
|
||||
console.log('Expected Total Cost:', expectedTotalCost);
|
||||
console.log('Actual Balance Decrease:', initialBalance - updatedBalance.tokenCredits);
|
||||
});
|
||||
|
||||
test('spendTokens should handle zero completion tokens', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000; // $10.00
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'gpt-3.5-turbo';
|
||||
|
|
@ -89,24 +83,19 @@ describe('Regular Token Spending Tests', () => {
|
|||
};
|
||||
|
||||
// Act
|
||||
process.env.CHECK_BALANCE = 'true';
|
||||
await spendTokens(txData, tokenUsage);
|
||||
|
||||
// Assert
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
|
||||
const promptMultiplier = getMultiplier({ model, tokenType: 'prompt' });
|
||||
const expectedCost = tokenUsage.promptTokens * promptMultiplier;
|
||||
const expectedCost = 100 * promptMultiplier;
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
|
||||
console.log('Initial Balance:', initialBalance);
|
||||
console.log('Updated Balance:', updatedBalance.tokenCredits);
|
||||
console.log('Expected Cost:', expectedCost);
|
||||
});
|
||||
|
||||
test('spendTokens should handle undefined token counts', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000; // $10.00
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'gpt-3.5-turbo';
|
||||
|
|
@ -120,14 +109,17 @@ describe('Regular Token Spending Tests', () => {
|
|||
|
||||
const tokenUsage = {};
|
||||
|
||||
// Act
|
||||
const result = await spendTokens(txData, tokenUsage);
|
||||
|
||||
// Assert: No transaction should be created
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
|
||||
test('spendTokens should handle only prompt tokens', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000; // $10.00
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'gpt-3.5-turbo';
|
||||
|
|
@ -141,14 +133,44 @@ describe('Regular Token Spending Tests', () => {
|
|||
|
||||
const tokenUsage = { promptTokens: 100 };
|
||||
|
||||
// Act
|
||||
await spendTokens(txData, tokenUsage);
|
||||
|
||||
// Assert
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
|
||||
const promptMultiplier = getMultiplier({ model, tokenType: 'prompt' });
|
||||
const expectedCost = 100 * promptMultiplier;
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
});
|
||||
|
||||
test('spendTokens should not update balance when balance feature is disabled', async () => {
|
||||
// Arrange: Override the config to disable balance updates.
|
||||
getBalanceConfig.mockResolvedValue({ balance: { enabled: false } });
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'gpt-3.5-turbo';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-conversation-id',
|
||||
model,
|
||||
context: 'test',
|
||||
endpointTokenConfig: null,
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
promptTokens: 100,
|
||||
completionTokens: 50,
|
||||
};
|
||||
|
||||
// Act
|
||||
await spendTokens(txData, tokenUsage);
|
||||
|
||||
// Assert: Balance should remain unchanged.
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
expect(updatedBalance.tokenCredits).toBe(initialBalance);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Structured Token Spending Tests', () => {
|
||||
|
|
@ -164,7 +186,7 @@ describe('Structured Token Spending Tests', () => {
|
|||
conversationId: 'c23a18da-706c-470a-ac28-ec87ed065199',
|
||||
model,
|
||||
context: 'message',
|
||||
endpointTokenConfig: null, // We'll use the default rates
|
||||
endpointTokenConfig: null,
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
|
|
@ -176,28 +198,15 @@ describe('Structured Token Spending Tests', () => {
|
|||
completionTokens: 5,
|
||||
};
|
||||
|
||||
// Get the actual multipliers
|
||||
const promptMultiplier = getMultiplier({ model, tokenType: 'prompt' });
|
||||
const completionMultiplier = getMultiplier({ model, tokenType: 'completion' });
|
||||
const writeMultiplier = getCacheMultiplier({ model, cacheType: 'write' });
|
||||
const readMultiplier = getCacheMultiplier({ model, cacheType: 'read' });
|
||||
|
||||
console.log('Multipliers:', {
|
||||
promptMultiplier,
|
||||
completionMultiplier,
|
||||
writeMultiplier,
|
||||
readMultiplier,
|
||||
});
|
||||
|
||||
// Act
|
||||
process.env.CHECK_BALANCE = 'true';
|
||||
const result = await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
// Assert
|
||||
console.log('Initial Balance:', initialBalance);
|
||||
console.log('Updated Balance:', result.completion.balance);
|
||||
console.log('Transaction Result:', result);
|
||||
|
||||
// Calculate expected costs.
|
||||
const expectedPromptCost =
|
||||
tokenUsage.promptTokens.input * promptMultiplier +
|
||||
tokenUsage.promptTokens.write * writeMultiplier +
|
||||
|
|
@ -206,37 +215,21 @@ describe('Structured Token Spending Tests', () => {
|
|||
const expectedTotalCost = expectedPromptCost + expectedCompletionCost;
|
||||
const expectedBalance = initialBalance - expectedTotalCost;
|
||||
|
||||
console.log('Expected Cost:', expectedTotalCost);
|
||||
console.log('Expected Balance:', expectedBalance);
|
||||
|
||||
// Assert
|
||||
expect(result.completion.balance).toBeLessThan(initialBalance);
|
||||
|
||||
// Allow for a small difference (e.g., 100 token credits, which is $0.0001)
|
||||
const allowedDifference = 100;
|
||||
expect(Math.abs(result.completion.balance - expectedBalance)).toBeLessThan(allowedDifference);
|
||||
|
||||
// Check if the decrease is approximately as expected
|
||||
const balanceDecrease = initialBalance - result.completion.balance;
|
||||
expect(balanceDecrease).toBeCloseTo(expectedTotalCost, 0);
|
||||
|
||||
// Check token values
|
||||
const expectedPromptTokenValue = -(
|
||||
tokenUsage.promptTokens.input * promptMultiplier +
|
||||
tokenUsage.promptTokens.write * writeMultiplier +
|
||||
tokenUsage.promptTokens.read * readMultiplier
|
||||
);
|
||||
const expectedCompletionTokenValue = -tokenUsage.completionTokens * completionMultiplier;
|
||||
|
||||
const expectedPromptTokenValue = -expectedPromptCost;
|
||||
const expectedCompletionTokenValue = -expectedCompletionCost;
|
||||
expect(result.prompt.prompt).toBeCloseTo(expectedPromptTokenValue, 1);
|
||||
expect(result.completion.completion).toBe(expectedCompletionTokenValue);
|
||||
|
||||
console.log('Expected prompt tokenValue:', expectedPromptTokenValue);
|
||||
console.log('Actual prompt tokenValue:', result.prompt.prompt);
|
||||
console.log('Expected completion tokenValue:', expectedCompletionTokenValue);
|
||||
console.log('Actual completion tokenValue:', result.completion.completion);
|
||||
});
|
||||
|
||||
test('should handle zero completion tokens in structured spending', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 17613154.55;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
|
@ -258,15 +251,17 @@ describe('Structured Token Spending Tests', () => {
|
|||
completionTokens: 0,
|
||||
};
|
||||
|
||||
process.env.CHECK_BALANCE = 'true';
|
||||
// Act
|
||||
const result = await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
// Assert
|
||||
expect(result.prompt).toBeDefined();
|
||||
expect(result.completion).toBeUndefined();
|
||||
expect(result.prompt.prompt).toBeLessThan(0);
|
||||
});
|
||||
|
||||
test('should handle only prompt tokens in structured spending', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 17613154.55;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
|
@ -287,15 +282,17 @@ describe('Structured Token Spending Tests', () => {
|
|||
},
|
||||
};
|
||||
|
||||
process.env.CHECK_BALANCE = 'true';
|
||||
// Act
|
||||
const result = await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
// Assert
|
||||
expect(result.prompt).toBeDefined();
|
||||
expect(result.completion).toBeUndefined();
|
||||
expect(result.prompt.prompt).toBeLessThan(0);
|
||||
});
|
||||
|
||||
test('should handle undefined token counts in structured spending', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 17613154.55;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
|
@ -310,9 +307,10 @@ describe('Structured Token Spending Tests', () => {
|
|||
|
||||
const tokenUsage = {};
|
||||
|
||||
process.env.CHECK_BALANCE = 'true';
|
||||
// Act
|
||||
const result = await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
// Assert
|
||||
expect(result).toEqual({
|
||||
prompt: undefined,
|
||||
completion: undefined,
|
||||
|
|
@ -320,6 +318,7 @@ describe('Structured Token Spending Tests', () => {
|
|||
});
|
||||
|
||||
test('should handle incomplete context for completion tokens', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 17613154.55;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
|
@ -341,15 +340,18 @@ describe('Structured Token Spending Tests', () => {
|
|||
completionTokens: 50,
|
||||
};
|
||||
|
||||
process.env.CHECK_BALANCE = 'true';
|
||||
// Act
|
||||
const result = await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
expect(result.completion.completion).toBeCloseTo(-50 * 15 * 1.15, 0); // Assuming multiplier is 15 and cancelRate is 1.15
|
||||
// Assert:
|
||||
// (Assuming a multiplier for completion of 15 and a cancel rate of 1.15 as noted in the original test.)
|
||||
expect(result.completion.completion).toBeCloseTo(-50 * 15 * 1.15, 0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('NaN Handling Tests', () => {
|
||||
test('should skip transaction creation when rawAmount is NaN', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
|
@ -365,9 +367,11 @@ describe('NaN Handling Tests', () => {
|
|||
tokenType: 'prompt',
|
||||
};
|
||||
|
||||
// Act
|
||||
const result = await Transaction.create(txData);
|
||||
expect(result).toBeUndefined();
|
||||
|
||||
// Assert: No transaction should be created and balance remains unchanged.
|
||||
expect(result).toBeUndefined();
|
||||
const balance = await Balance.findOne({ user: userId });
|
||||
expect(balance.tokenCredits).toBe(initialBalance);
|
||||
});
|
||||
|
|
|
|||
|
|
@ -19,14 +19,19 @@ jest.mock('~/config', () => ({
|
|||
},
|
||||
}));
|
||||
|
||||
// New config module
|
||||
const { getBalanceConfig } = require('~/server/services/Config');
|
||||
jest.mock('~/server/services/Config');
|
||||
|
||||
// Import after mocking
|
||||
const { spendTokens, spendStructuredTokens } = require('./spendTokens');
|
||||
const { Transaction } = require('./Transaction');
|
||||
const Balance = require('./Balance');
|
||||
|
||||
describe('spendTokens', () => {
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
process.env.CHECK_BALANCE = 'true';
|
||||
getBalanceConfig.mockResolvedValue({ enabled: true });
|
||||
});
|
||||
|
||||
it('should create transactions for both prompt and completion tokens', async () => {
|
||||
|
|
@ -92,7 +97,7 @@ describe('spendTokens', () => {
|
|||
expect(Transaction.create).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
tokenType: 'completion',
|
||||
rawAmount: -0, // Changed from 0 to -0
|
||||
rawAmount: -0,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
|
@ -111,8 +116,9 @@ describe('spendTokens', () => {
|
|||
expect(Transaction.create).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should not update balance when CHECK_BALANCE is false', async () => {
|
||||
process.env.CHECK_BALANCE = 'false';
|
||||
it('should not update balance when the balance feature is disabled', async () => {
|
||||
// Override configuration: disable balance updates.
|
||||
getBalanceConfig.mockResolvedValue({ enabled: false });
|
||||
const txData = {
|
||||
user: new mongoose.Types.ObjectId(),
|
||||
conversationId: 'test-convo',
|
||||
|
|
@ -130,6 +136,7 @@ describe('spendTokens', () => {
|
|||
await spendTokens(txData, tokenUsage);
|
||||
|
||||
expect(Transaction.create).toHaveBeenCalledTimes(2);
|
||||
// When balance updates are disabled, Balance methods should not be called.
|
||||
expect(Balance.findOne).not.toHaveBeenCalled();
|
||||
expect(Balance.findOneAndUpdate).not.toHaveBeenCalled();
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
const bcrypt = require('bcryptjs');
|
||||
const { getBalanceConfig } = require('~/server/services/Config');
|
||||
const signPayload = require('~/server/services/signPayload');
|
||||
const { isEnabled } = require('~/server/utils/handleText');
|
||||
const Balance = require('./Balance');
|
||||
const User = require('./User');
|
||||
|
||||
|
|
@ -13,11 +13,9 @@ const User = require('./User');
|
|||
*/
|
||||
const getUserById = async function (userId, fieldsToSelect = null) {
|
||||
const query = User.findById(userId);
|
||||
|
||||
if (fieldsToSelect) {
|
||||
query.select(fieldsToSelect);
|
||||
}
|
||||
|
||||
return await query.lean();
|
||||
};
|
||||
|
||||
|
|
@ -32,7 +30,6 @@ const findUser = async function (searchCriteria, fieldsToSelect = null) {
|
|||
if (fieldsToSelect) {
|
||||
query.select(fieldsToSelect);
|
||||
}
|
||||
|
||||
return await query.lean();
|
||||
};
|
||||
|
||||
|
|
@ -58,11 +55,12 @@ const updateUser = async function (userId, updateData) {
|
|||
* Creates a new user, optionally with a TTL of 1 week.
|
||||
* @param {MongoUser} data - The user data to be created, must contain user_id.
|
||||
* @param {boolean} [disableTTL=true] - Whether to disable the TTL. Defaults to `true`.
|
||||
* @param {boolean} [returnUser=false] - Whether to disable the TTL. Defaults to `true`.
|
||||
* @returns {Promise<ObjectId>} A promise that resolves to the created user document ID.
|
||||
* @param {boolean} [returnUser=false] - Whether to return the created user object.
|
||||
* @returns {Promise<ObjectId|MongoUser>} A promise that resolves to the created user document ID or user object.
|
||||
* @throws {Error} If a user with the same user_id already exists.
|
||||
*/
|
||||
const createUser = async (data, disableTTL = true, returnUser = false) => {
|
||||
const balance = await getBalanceConfig();
|
||||
const userData = {
|
||||
...data,
|
||||
expiresAt: disableTTL ? null : new Date(Date.now() + 604800 * 1000), // 1 week in milliseconds
|
||||
|
|
@ -74,13 +72,27 @@ const createUser = async (data, disableTTL = true, returnUser = false) => {
|
|||
|
||||
const user = await User.create(userData);
|
||||
|
||||
if (isEnabled(process.env.CHECK_BALANCE) && process.env.START_BALANCE) {
|
||||
let incrementValue = parseInt(process.env.START_BALANCE);
|
||||
await Balance.findOneAndUpdate(
|
||||
{ user: user._id },
|
||||
{ $inc: { tokenCredits: incrementValue } },
|
||||
{ upsert: true, new: true },
|
||||
).lean();
|
||||
// If balance is enabled, create or update a balance record for the user using global.interfaceConfig.balance
|
||||
if (balance?.enabled && balance?.startBalance) {
|
||||
const update = {
|
||||
$inc: { tokenCredits: balance.startBalance },
|
||||
};
|
||||
|
||||
if (
|
||||
balance.autoRefillEnabled &&
|
||||
balance.refillIntervalValue != null &&
|
||||
balance.refillIntervalUnit != null &&
|
||||
balance.refillAmount != null
|
||||
) {
|
||||
update.$set = {
|
||||
autoRefillEnabled: true,
|
||||
refillIntervalValue: balance.refillIntervalValue,
|
||||
refillIntervalUnit: balance.refillIntervalUnit,
|
||||
refillAmount: balance.refillAmount,
|
||||
};
|
||||
}
|
||||
|
||||
await Balance.findOneAndUpdate({ user: user._id }, update, { upsert: true, new: true }).lean();
|
||||
}
|
||||
|
||||
if (returnUser) {
|
||||
|
|
@ -123,7 +135,7 @@ const expires = eval(SESSION_EXPIRY) ?? 1000 * 60 * 15;
|
|||
/**
|
||||
* Generates a JWT token for a given user.
|
||||
*
|
||||
* @param {MongoUser} user - ID of the user for whom the token is being generated.
|
||||
* @param {MongoUser} user - The user for whom the token is being generated.
|
||||
* @returns {Promise<string>} A promise that resolves to a JWT token.
|
||||
*/
|
||||
const generateToken = async (user) => {
|
||||
|
|
@ -146,7 +158,7 @@ const generateToken = async (user) => {
|
|||
/**
|
||||
* Compares the provided password with the user's password.
|
||||
*
|
||||
* @param {MongoUser} user - the user to compare password for.
|
||||
* @param {MongoUser} user - The user to compare the password for.
|
||||
* @param {string} candidatePassword - The password to test against the user's password.
|
||||
* @returns {Promise<boolean>} A promise that resolves to a boolean indicating if the password matches.
|
||||
*/
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue