Andreas/LibreChat
1
0
Fork 0
mirror of https://github.com/danny-avila/LibreChat.git synced 2026-01-14 14:38:51 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

refactor: get balance config primarily from appConfig

Browse source
This commit is contained in:
Danny Avila 2025-08-18 03:45:02 -04:00
parent d853c10920
commit 1d2be247cf
No known key found for this signature in database
GPG key ID: BF31EEB2C5CA0956
11 changed files with 89 additions and 79 deletions
Download patch file Download diff file Expand all files Collapse all files

14
api/models/Transaction.js
View file

@ -1,5 +1,4 @@
const { logger } = require('@librechat/data-schemas');
const { getBalanceConfig } = require('~/server/services/Config');
const { getMultiplier, getCacheMultiplier } = require('./tx');
const { Transaction, Balance } = require('~/db/models');
@ -187,9 +186,10 @@ async function createAutoRefillTransaction(txData) {
/**
* Static method to create a transaction and update the balance
* @param {txData} txData - Transaction data.
* @param {txData} _txData - Transaction data.
*/
async function createTransaction(txData) {
async function createTransaction(_txData) {
const { balance, ...txData } = _txData;
if (txData.rawAmount != null && isNaN(txData.rawAmount)) {
return;
}
@ -199,8 +199,6 @@ async function createTransaction(txData) {
calculateTokenValue(transaction);
await transaction.save();
const balance = await getBalanceConfig();
if (!balance?.enabled) {
return;
}
@ -221,9 +219,10 @@ async function createTransaction(txData) {
/**
* Static method to create a structured transaction and update the balance
* @param {txData} txData - Transaction data.
* @param {txData} _txData - Transaction data.
*/
async function createStructuredTransaction(txData) {
async function createStructuredTransaction(_txData) {
const { balance, ...txData } = _txData;
const transaction = new Transaction({
...txData,
endpointTokenConfig: txData.endpointTokenConfig,
@ -233,7 +232,6 @@ async function createStructuredTransaction(txData) {
await transaction.save();
const balance = await getBalanceConfig();
if (!balance?.enabled) {
return;
}

21
api/models/Transaction.spec.js
View file

@ -1,14 +1,11 @@
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 { createTransaction } = require('./Transaction');
const { Balance } = require('~/db/models');
// Mock the custom config module so we can control the balance flag.
jest.mock('~/server/services/Config');
let mongoServer;
beforeAll(async () => {
mongoServer = await MongoMemoryServer.create();
@ -23,8 +20,6 @@ afterAll(async () => {
beforeEach(async () => {
await mongoose.connection.dropDatabase();
// Default: enable balance updates in tests.
getBalanceConfig.mockResolvedValue({ enabled: true });
});
describe('Regular Token Spending Tests', () => {
@ -41,6 +36,7 @@ describe('Regular Token Spending Tests', () => {
model,
context: 'test',
endpointTokenConfig: null,
balance: { enabled: true },
};
const tokenUsage = {
@ -74,6 +70,7 @@ describe('Regular Token Spending Tests', () => {
model,
context: 'test',
endpointTokenConfig: null,
balance: { enabled: true },
};
const tokenUsage = {
@ -104,6 +101,7 @@ describe('Regular Token Spending Tests', () => {
model,
context: 'test',
endpointTokenConfig: null,
balance: { enabled: true },
};
const tokenUsage = {};
@ -128,6 +126,7 @@ describe('Regular Token Spending Tests', () => {
model,
context: 'test',
endpointTokenConfig: null,
balance: { enabled: true },
};
const tokenUsage = { promptTokens: 100 };
@ -143,8 +142,7 @@ describe('Regular Token Spending Tests', () => {
});
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 } });
// Arrange: Balance config is now passed directly in txData
const userId = new mongoose.Types.ObjectId();
const initialBalance = 10000000;
await Balance.create({ user: userId, tokenCredits: initialBalance });
@ -156,6 +154,7 @@ describe('Regular Token Spending Tests', () => {
model,
context: 'test',
endpointTokenConfig: null,
balance: { enabled: false },
};
const tokenUsage = {
@ -186,6 +185,7 @@ describe('Structured Token Spending Tests', () => {
model,
context: 'message',
endpointTokenConfig: null,
balance: { enabled: true },
};
const tokenUsage = {
@ -239,6 +239,7 @@ describe('Structured Token Spending Tests', () => {
conversationId: 'test-convo',
model,
context: 'message',
balance: { enabled: true },
};
const tokenUsage = {
@ -271,6 +272,7 @@ describe('Structured Token Spending Tests', () => {
conversationId: 'test-convo',
model,
context: 'message',
balance: { enabled: true },
};
const tokenUsage = {
@ -302,6 +304,7 @@ describe('Structured Token Spending Tests', () => {
conversationId: 'test-convo',
model,
context: 'message',
balance: { enabled: true },
};
const tokenUsage = {};
@ -328,6 +331,7 @@ describe('Structured Token Spending Tests', () => {
conversationId: 'test-convo',
model,
context: 'incomplete',
balance: { enabled: true },
};
const tokenUsage = {
@ -364,6 +368,7 @@ describe('NaN Handling Tests', () => {
endpointTokenConfig: null,
rawAmount: NaN,
tokenType: 'prompt',
balance: { enabled: true },
};
// Act

16
api/models/spendTokens.js
View file

@ -5,13 +5,7 @@ const { createTransaction, createStructuredTransaction } = require('./Transactio
*
* @function
* @async
* @param {Object} txData - Transaction data.
* @param {mongoose.Schema.Types.ObjectId} txData.user - The user ID.
* @param {String} txData.conversationId - The ID of the conversation.
* @param {String} txData.model - The model name.
* @param {String} txData.context - The context in which the transaction is made.
* @param {EndpointTokenConfig} [txData.endpointTokenConfig] - The current endpoint token config.
* @param {String} [txData.valueKey] - The value key (optional).
* @param {txData} txData - Transaction data.
* @param {Object} tokenUsage - The number of tokens used.
* @param {Number} tokenUsage.promptTokens - The number of prompt tokens used.
* @param {Number} tokenUsage.completionTokens - The number of completion tokens used.
@ -69,13 +63,7 @@ const spendTokens = async (txData, tokenUsage) => {
*
* @function
* @async
* @param {Object} txData - Transaction data.
* @param {mongoose.Schema.Types.ObjectId} txData.user - The user ID.
* @param {String} txData.conversationId - The ID of the conversation.
* @param {String} txData.model - The model name.
* @param {String} txData.context - The context in which the transaction is made.
* @param {EndpointTokenConfig} [txData.endpointTokenConfig] - The current endpoint token config.
* @param {String} [txData.valueKey] - The value key (optional).
* @param {txData} txData - Transaction data.
* @param {Object} tokenUsage - The number of tokens used.
* @param {Object} tokenUsage.promptTokens - The number of prompt tokens used.
* @param {Number} tokenUsage.promptTokens.input - The number of input tokens.

26
api/models/spendTokens.spec.js
View file

@ -5,7 +5,6 @@ const { createTransaction, createAutoRefillTransaction } = require('./Transactio
require('~/db/models');
// Mock the logger to prevent console output during tests
jest.mock('~/config', () => ({
logger: {
debug: jest.fn(),
@ -13,10 +12,6 @@ jest.mock('~/config', () => ({
},
}));
// Mock the Config service
const { getBalanceConfig } = require('~/server/services/Config');
jest.mock('~/server/services/Config');
describe('spendTokens', () => {
let mongoServer;
let userId;
@ -44,8 +39,7 @@ describe('spendTokens', () => {
// Create a new user ID for each test
userId = new mongoose.Types.ObjectId();
// Mock the balance config to be enabled by default
getBalanceConfig.mockResolvedValue({ enabled: true });
// Balance config is now passed directly in txData
});
it('should create transactions for both prompt and completion tokens', async () => {
@ -60,6 +54,7 @@ describe('spendTokens', () => {
conversationId: 'test-convo',
model: 'gpt-3.5-turbo',
context: 'test',
balance: { enabled: true },
};
const tokenUsage = {
promptTokens: 100,
@ -98,6 +93,7 @@ describe('spendTokens', () => {
conversationId: 'test-convo',
model: 'gpt-3.5-turbo',
context: 'test',
balance: { enabled: true },
};
const tokenUsage = {
promptTokens: 100,
@ -127,6 +123,7 @@ describe('spendTokens', () => {
conversationId: 'test-convo',
model: 'gpt-3.5-turbo',
context: 'test',
balance: { enabled: true },
};
const tokenUsage = {};
@ -138,8 +135,7 @@ describe('spendTokens', () => {
});
it('should not update balance when the balance feature is disabled', async () => {
// Override configuration: disable balance updates
getBalanceConfig.mockResolvedValue({ enabled: false });
// Balance is now passed directly in txData
// Create a balance for the user
await Balance.create({
user: userId,
@ -151,6 +147,7 @@ describe('spendTokens', () => {
conversationId: 'test-convo',
model: 'gpt-3.5-turbo',
context: 'test',
balance: { enabled: false },
};
const tokenUsage = {
promptTokens: 100,
@ -180,6 +177,7 @@ describe('spendTokens', () => {
conversationId: 'test-convo',
model: 'gpt-4', // Using a more expensive model
context: 'test',
balance: { enabled: true },
};
// Spending more tokens than the user has balance for
@ -233,6 +231,7 @@ describe('spendTokens', () => {
conversationId: 'test-convo-1',
model: 'gpt-4',
context: 'test',
balance: { enabled: true },
};
const tokenUsage1 = {
@ -252,6 +251,7 @@ describe('spendTokens', () => {
conversationId: 'test-convo-2',
model: 'gpt-4',
context: 'test',
balance: { enabled: true },
};
const tokenUsage2 = {
@ -292,6 +292,7 @@ describe('spendTokens', () => {
tokenType: 'completion',
rawAmount: -100,
context: 'test',
balance: { enabled: true },
});
console.log('Direct Transaction.create result:', directResult);
@ -316,6 +317,7 @@ describe('spendTokens', () => {
conversationId: `test-convo-${model}`,
model,
context: 'test',
balance: { enabled: true },
};
const tokenUsage = {
@ -352,6 +354,7 @@ describe('spendTokens', () => {
conversationId: 'test-convo-1',
model: 'claude-3-5-sonnet',
context: 'test',
balance: { enabled: true },
};
const tokenUsage1 = {
@ -375,6 +378,7 @@ describe('spendTokens', () => {
conversationId: 'test-convo-2',
model: 'claude-3-5-sonnet',
context: 'test',
balance: { enabled: true },
};
const tokenUsage2 = {
@ -426,6 +430,7 @@ describe('spendTokens', () => {
conversationId: 'test-convo',
model: 'claude-3-5-sonnet', // Using a model that supports structured tokens
context: 'test',
balance: { enabled: true },
};
// Spending more tokens than the user has balance for
@ -505,6 +510,7 @@ describe('spendTokens', () => {
conversationId,
user: userId,
model: usage.model,
balance: { enabled: true },
};
// Calculate expected spend for this transaction
@ -617,6 +623,7 @@ describe('spendTokens', () => {
tokenType: 'credits',
context: 'concurrent-refill-test',
rawAmount: refillAmount,
balance: { enabled: true },
}),
);
}
@ -683,6 +690,7 @@ describe('spendTokens', () => {
conversationId: 'test-convo',
model: 'claude-3-5-sonnet',
context: 'test',
balance: { enabled: true },
};
const tokenUsage = {
promptTokens: {