mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-03-09 09:32:36 +01:00
🪢 chore: Consolidate Pricing and Tx Imports After tx.js Module Removal (#12086)
* 🧹 chore: resolve imports due to rebase
* chore: Update model mocks in unit tests for consistency
- Consolidated model mock implementations across various test files to streamline setup and reduce redundancy.
- Removed duplicate mock definitions for `getMultiplier` and `getCacheMultiplier`, ensuring a unified approach in `recordCollectedUsage.spec.js`, `openai.spec.js`, `responses.unit.spec.js`, and `abortMiddleware.spec.js`.
- Enhanced clarity and maintainability of test files by aligning mock structures with the latest model updates.
* fix: Safeguard token credit checks in transaction tests
- Updated assertions in `transaction.spec.ts` to handle potential null values for `updatedBalance` by using optional chaining.
- Enhanced robustness of tests related to token credit calculations, ensuring they correctly account for scenarios where the balance may not be found.
* chore: transaction methods with bulk insert functionality
- Introduced `bulkInsertTransactions` method in `transaction.ts` to facilitate batch insertion of transaction documents.
- Updated test file `transactions.bulk-parity.spec.ts` to utilize new pricing function assignments and handle potential null values in calculations, improving test robustness.
- Refactored pricing function initialization for clarity and consistency.
* refactor: Enhance type definitions and introduce new utility functions for model matching
- Added `findMatchingPattern` and `matchModelName` utility functions to improve model name matching logic in transaction methods.
- Updated type definitions for `findMatchingPattern` to accept a more specific tokensMap structure, enhancing type safety.
- Refactored `dbMethods` initialization in `transactions.bulk-parity.spec.ts` to include the new utility functions, improving test clarity and functionality.
* refactor: Update database method imports and enhance transaction handling
- Refactored `abortMiddleware.js` to utilize centralized database methods for message handling and conversation retrieval, improving code consistency.
- Enhanced `bulkInsertTransactions` in `transaction.ts` to handle empty document arrays gracefully and added error logging for better debugging.
- Updated type definitions in `transactions.ts` to enforce stricter typing for token types, enhancing type safety across transaction methods.
- Improved test setup in `transactions.bulk-parity.spec.ts` by refining pricing function assignments and ensuring robust handling of potential null values.
* refactor: Update database method references and improve transaction multiplier handling
- Refactored `client.js` to update database method references for `bulkInsertTransactions` and `updateBalance`, ensuring consistency in method usage.
- Enhanced transaction multiplier calculations in `transaction.spec.ts` to provide fallback values for write and read multipliers, improving robustness in cost calculations across structured token spending tests.
This commit is contained in:
Danny Avila
GitHub
parent
8fafda47c2
commit
d24fe17a4b
17 changed files with 123 additions and 427 deletions
|
|
@ -85,7 +85,10 @@ export interface CreateMethodsDeps {
|
|||
/** Matches a model name to a canonical key. From @librechat/api. */
|
||||
matchModelName?: (model: string, endpoint?: string) => string | undefined;
|
||||
/** Finds the first key in values whose key is a substring of model. From @librechat/api. */
|
||||
findMatchingPattern?: (model: string, values: Record<string, unknown>) => string | undefined;
|
||||
findMatchingPattern?: (
|
||||
model: string,
|
||||
values: Record<string, number | Record<string, number>>,
|
||||
) => string | undefined;
|
||||
/** Removes all ACL permissions for a resource. From PermissionService. */
|
||||
removeAllPermissions?: (params: { resourceType: string; resourceId: unknown }) => Promise<void>;
|
||||
/** Returns a cache store for the given key. From getLogStores. */
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@
|
|||
*/
|
||||
export function findMatchingPattern(
|
||||
modelName: string,
|
||||
tokensMap: Record<string, unknown>,
|
||||
tokensMap: Record<string, number | Record<string, number>>,
|
||||
): string | undefined {
|
||||
const keys = Object.keys(tokensMap);
|
||||
const lowerModelName = modelName.toLowerCase();
|
||||
|
|
|
|||
|
|
@ -247,8 +247,8 @@ describe('Structured Token Spending Tests', () => {
|
|||
|
||||
const promptMultiplier = getMultiplier({ model, tokenType: 'prompt' });
|
||||
const completionMultiplier = getMultiplier({ model, tokenType: 'completion' });
|
||||
const writeMultiplier = getCacheMultiplier({ model, cacheType: 'write' });
|
||||
const readMultiplier = getCacheMultiplier({ model, cacheType: 'read' });
|
||||
const writeMultiplier = getCacheMultiplier({ model, cacheType: 'write' }) ?? promptMultiplier;
|
||||
const readMultiplier = getCacheMultiplier({ model, cacheType: 'read' }) ?? promptMultiplier;
|
||||
|
||||
// Act
|
||||
const result = await spendStructuredTokens(txData, tokenUsage);
|
||||
|
|
@ -256,8 +256,8 @@ describe('Structured Token Spending Tests', () => {
|
|||
// Calculate expected costs.
|
||||
const expectedPromptCost =
|
||||
tokenUsage.promptTokens.input * promptMultiplier +
|
||||
tokenUsage.promptTokens.write * (writeMultiplier ?? 0) +
|
||||
tokenUsage.promptTokens.read * (readMultiplier ?? 0);
|
||||
tokenUsage.promptTokens.write * writeMultiplier +
|
||||
tokenUsage.promptTokens.read * readMultiplier;
|
||||
const expectedCompletionCost = tokenUsage.completionTokens * completionMultiplier;
|
||||
const expectedTotalCost = expectedPromptCost + expectedCompletionCost;
|
||||
const expectedBalance = initialBalance - expectedTotalCost;
|
||||
|
|
@ -813,13 +813,18 @@ describe('Premium Token Pricing Integration Tests', () => {
|
|||
|
||||
const premiumPromptRate = premiumTokenValues[model].prompt;
|
||||
const premiumCompletionRate = premiumTokenValues[model].completion;
|
||||
const writeMultiplier = getCacheMultiplier({ model, cacheType: 'write' });
|
||||
const readMultiplier = getCacheMultiplier({ model, cacheType: 'read' });
|
||||
const promptMultiplier = getMultiplier({
|
||||
model,
|
||||
tokenType: 'prompt',
|
||||
inputTokenCount: totalInput,
|
||||
});
|
||||
const writeMultiplier = getCacheMultiplier({ model, cacheType: 'write' }) ?? promptMultiplier;
|
||||
const readMultiplier = getCacheMultiplier({ model, cacheType: 'read' }) ?? promptMultiplier;
|
||||
|
||||
const expectedPromptCost =
|
||||
tokenUsage.promptTokens.input * premiumPromptRate +
|
||||
tokenUsage.promptTokens.write * (writeMultiplier ?? 0) +
|
||||
tokenUsage.promptTokens.read * (readMultiplier ?? 0);
|
||||
tokenUsage.promptTokens.write * writeMultiplier +
|
||||
tokenUsage.promptTokens.read * readMultiplier;
|
||||
const expectedCompletionCost = tokenUsage.completionTokens * premiumCompletionRate;
|
||||
const expectedTotalCost = expectedPromptCost + expectedCompletionCost;
|
||||
|
||||
|
|
@ -859,13 +864,18 @@ describe('Premium Token Pricing Integration Tests', () => {
|
|||
|
||||
const standardPromptRate = tokenValues[model].prompt;
|
||||
const standardCompletionRate = tokenValues[model].completion;
|
||||
const writeMultiplier = getCacheMultiplier({ model, cacheType: 'write' });
|
||||
const readMultiplier = getCacheMultiplier({ model, cacheType: 'read' });
|
||||
const promptMultiplier = getMultiplier({
|
||||
model,
|
||||
tokenType: 'prompt',
|
||||
inputTokenCount: totalInput,
|
||||
});
|
||||
const writeMultiplier = getCacheMultiplier({ model, cacheType: 'write' }) ?? promptMultiplier;
|
||||
const readMultiplier = getCacheMultiplier({ model, cacheType: 'read' }) ?? promptMultiplier;
|
||||
|
||||
const expectedPromptCost =
|
||||
tokenUsage.promptTokens.input * standardPromptRate +
|
||||
tokenUsage.promptTokens.write * (writeMultiplier ?? 0) +
|
||||
tokenUsage.promptTokens.read * (readMultiplier ?? 0);
|
||||
tokenUsage.promptTokens.write * writeMultiplier +
|
||||
tokenUsage.promptTokens.read * readMultiplier;
|
||||
const expectedCompletionCost = tokenUsage.completionTokens * standardCompletionRate;
|
||||
const expectedTotalCost = expectedPromptCost + expectedCompletionCost;
|
||||
|
||||
|
|
@ -900,7 +910,7 @@ describe('Premium Token Pricing Integration Tests', () => {
|
|||
promptTokens * standardPromptRate + completionTokens * standardCompletionRate;
|
||||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
expect(updatedBalance?.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
});
|
||||
|
||||
test('spendTokens should apply premium pricing for gemini-3.1-pro-preview above threshold', async () => {
|
||||
|
|
@ -929,7 +939,7 @@ describe('Premium Token Pricing Integration Tests', () => {
|
|||
promptTokens * premiumPromptRate + completionTokens * premiumCompletionRate;
|
||||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
expect(updatedBalance?.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
});
|
||||
|
||||
test('spendTokens should apply standard pricing for gemini-3.1-pro-preview at exactly the threshold', async () => {
|
||||
|
|
@ -958,7 +968,7 @@ describe('Premium Token Pricing Integration Tests', () => {
|
|||
promptTokens * standardPromptRate + completionTokens * standardCompletionRate;
|
||||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
expect(updatedBalance?.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
});
|
||||
|
||||
test('spendStructuredTokens should apply premium pricing for gemini-3.1 when total input exceeds threshold', async () => {
|
||||
|
|
@ -992,8 +1002,13 @@ describe('Premium Token Pricing Integration Tests', () => {
|
|||
|
||||
const premiumPromptRate = premiumTokenValues['gemini-3.1'].prompt;
|
||||
const premiumCompletionRate = premiumTokenValues['gemini-3.1'].completion;
|
||||
const writeMultiplier = getCacheMultiplier({ model, cacheType: 'write' });
|
||||
const readMultiplier = getCacheMultiplier({ model, cacheType: 'read' });
|
||||
const promptMultiplier = getMultiplier({
|
||||
model,
|
||||
tokenType: 'prompt',
|
||||
inputTokenCount: totalInput,
|
||||
});
|
||||
const writeMultiplier = getCacheMultiplier({ model, cacheType: 'write' }) ?? promptMultiplier;
|
||||
const readMultiplier = getCacheMultiplier({ model, cacheType: 'read' }) ?? promptMultiplier;
|
||||
|
||||
const expectedPromptCost =
|
||||
tokenUsage.promptTokens.input * premiumPromptRate +
|
||||
|
|
@ -1004,7 +1019,7 @@ describe('Premium Token Pricing Integration Tests', () => {
|
|||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
expect(totalInput).toBeGreaterThan(premiumTokenValues['gemini-3.1'].threshold);
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedTotalCost, 0);
|
||||
expect(updatedBalance?.tokenCredits).toBeCloseTo(initialBalance - expectedTotalCost, 0);
|
||||
});
|
||||
|
||||
test('non-premium models should not be affected by inputTokenCount regardless of prompt size', async () => {
|
||||
|
|
@ -1036,339 +1051,3 @@ describe('Premium Token Pricing Integration Tests', () => {
|
|||
expect(updatedBalance?.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Bulk path parity', () => {
|
||||
/**
|
||||
* Each test here mirrors an existing legacy test above, replacing spendTokens/
|
||||
* spendStructuredTokens with recordCollectedUsage + bulk deps.
|
||||
* The balance deduction and transaction document fields must be numerically identical.
|
||||
*/
|
||||
let bulkDeps;
|
||||
let methods;
|
||||
|
||||
beforeEach(() => {
|
||||
methods = createMethods(mongoose);
|
||||
bulkDeps = {
|
||||
spendTokens: () => Promise.resolve(),
|
||||
spendStructuredTokens: () => Promise.resolve(),
|
||||
pricing: { getMultiplier, getCacheMultiplier },
|
||||
bulkWriteOps: {
|
||||
insertMany: methods.bulkInsertTransactions,
|
||||
updateBalance: methods.updateBalance,
|
||||
},
|
||||
};
|
||||
});
|
||||
|
||||
test('balance should decrease when spending tokens via bulk path', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'gpt-3.5-turbo';
|
||||
const promptTokens = 100;
|
||||
const completionTokens = 50;
|
||||
|
||||
await recordCollectedUsage(bulkDeps, {
|
||||
user: userId.toString(),
|
||||
conversationId: 'test-conversation-id',
|
||||
model,
|
||||
context: 'test',
|
||||
balance: { enabled: true },
|
||||
transactions: { enabled: true },
|
||||
collectedUsage: [{ input_tokens: promptTokens, output_tokens: completionTokens, model }],
|
||||
});
|
||||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
const promptMultiplier = getMultiplier({
|
||||
model,
|
||||
tokenType: 'prompt',
|
||||
inputTokenCount: promptTokens,
|
||||
});
|
||||
const completionMultiplier = getMultiplier({
|
||||
model,
|
||||
tokenType: 'completion',
|
||||
inputTokenCount: promptTokens,
|
||||
});
|
||||
const expectedTotalCost =
|
||||
promptTokens * promptMultiplier + completionTokens * completionMultiplier;
|
||||
const expectedBalance = initialBalance - expectedTotalCost;
|
||||
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(expectedBalance, 0);
|
||||
|
||||
const txns = await Transaction.find({ user: userId }).lean();
|
||||
expect(txns).toHaveLength(2);
|
||||
});
|
||||
|
||||
test('bulk path should not update balance when balance.enabled is false', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'gpt-3.5-turbo';
|
||||
|
||||
await recordCollectedUsage(bulkDeps, {
|
||||
user: userId.toString(),
|
||||
conversationId: 'test-conversation-id',
|
||||
model,
|
||||
context: 'test',
|
||||
balance: { enabled: false },
|
||||
transactions: { enabled: true },
|
||||
collectedUsage: [{ input_tokens: 100, output_tokens: 50, model }],
|
||||
});
|
||||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
expect(updatedBalance.tokenCredits).toBe(initialBalance);
|
||||
const txns = await Transaction.find({ user: userId }).lean();
|
||||
expect(txns).toHaveLength(2); // transactions still recorded
|
||||
});
|
||||
|
||||
test('bulk path should not insert when transactions.enabled is false', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
await recordCollectedUsage(bulkDeps, {
|
||||
user: userId.toString(),
|
||||
conversationId: 'test-conversation-id',
|
||||
model: 'gpt-3.5-turbo',
|
||||
context: 'test',
|
||||
balance: { enabled: true },
|
||||
transactions: { enabled: false },
|
||||
collectedUsage: [{ input_tokens: 100, output_tokens: 50, model: 'gpt-3.5-turbo' }],
|
||||
});
|
||||
|
||||
const txns = await Transaction.find({ user: userId }).lean();
|
||||
expect(txns).toHaveLength(0);
|
||||
const balance = await Balance.findOne({ user: userId });
|
||||
expect(balance.tokenCredits).toBe(initialBalance);
|
||||
});
|
||||
|
||||
test('bulk path handles incomplete context for completion tokens — same CANCEL_RATE as legacy', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 17613154.55;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-3-5-sonnet';
|
||||
const promptTokens = 10;
|
||||
const completionTokens = 50;
|
||||
|
||||
await recordCollectedUsage(bulkDeps, {
|
||||
user: userId.toString(),
|
||||
conversationId: 'test-convo',
|
||||
model,
|
||||
context: 'incomplete',
|
||||
balance: { enabled: true },
|
||||
transactions: { enabled: true },
|
||||
collectedUsage: [{ input_tokens: promptTokens, output_tokens: completionTokens, model }],
|
||||
});
|
||||
|
||||
const txns = await Transaction.find({ user: userId }).lean();
|
||||
const completionTx = txns.find((t) => t.tokenType === 'completion');
|
||||
const completionMultiplier = getMultiplier({
|
||||
model,
|
||||
tokenType: 'completion',
|
||||
inputTokenCount: promptTokens,
|
||||
});
|
||||
expect(completionTx.tokenValue).toBeCloseTo(-completionTokens * completionMultiplier * 1.15, 0);
|
||||
});
|
||||
|
||||
test('bulk path structured tokens — balance deduction matches legacy spendStructuredTokens', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 17613154.55;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-3-5-sonnet';
|
||||
const promptInput = 11;
|
||||
const promptWrite = 140522;
|
||||
const promptRead = 0;
|
||||
const completionTokens = 5;
|
||||
const totalInput = promptInput + promptWrite + promptRead;
|
||||
|
||||
await recordCollectedUsage(bulkDeps, {
|
||||
user: userId.toString(),
|
||||
conversationId: 'test-convo',
|
||||
model,
|
||||
context: 'message',
|
||||
balance: { enabled: true },
|
||||
transactions: { enabled: true },
|
||||
collectedUsage: [
|
||||
{
|
||||
input_tokens: promptInput,
|
||||
output_tokens: completionTokens,
|
||||
model,
|
||||
input_token_details: { cache_creation: promptWrite, cache_read: promptRead },
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
const promptMultiplier = getMultiplier({
|
||||
model,
|
||||
tokenType: 'prompt',
|
||||
inputTokenCount: totalInput,
|
||||
});
|
||||
const completionMultiplier = getMultiplier({
|
||||
model,
|
||||
tokenType: 'completion',
|
||||
inputTokenCount: totalInput,
|
||||
});
|
||||
const writeMultiplier = getCacheMultiplier({ model, cacheType: 'write' }) ?? promptMultiplier;
|
||||
const readMultiplier = getCacheMultiplier({ model, cacheType: 'read' }) ?? promptMultiplier;
|
||||
|
||||
const expectedPromptCost =
|
||||
promptInput * promptMultiplier + promptWrite * writeMultiplier + promptRead * readMultiplier;
|
||||
const expectedCompletionCost = completionTokens * completionMultiplier;
|
||||
const expectedTotalCost = expectedPromptCost + expectedCompletionCost;
|
||||
const expectedBalance = initialBalance - expectedTotalCost;
|
||||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
expect(Math.abs(updatedBalance.tokenCredits - expectedBalance)).toBeLessThan(100);
|
||||
});
|
||||
|
||||
test('premium pricing above threshold via bulk path — same balance as legacy', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-opus-4-6';
|
||||
const promptTokens = 250000;
|
||||
const completionTokens = 500;
|
||||
|
||||
await recordCollectedUsage(bulkDeps, {
|
||||
user: userId.toString(),
|
||||
conversationId: 'test-premium',
|
||||
model,
|
||||
context: 'test',
|
||||
balance: { enabled: true },
|
||||
transactions: { enabled: true },
|
||||
collectedUsage: [{ input_tokens: promptTokens, output_tokens: completionTokens, model }],
|
||||
});
|
||||
|
||||
const premiumPromptRate = premiumTokenValues[model].prompt;
|
||||
const premiumCompletionRate = premiumTokenValues[model].completion;
|
||||
const expectedCost =
|
||||
promptTokens * premiumPromptRate + completionTokens * premiumCompletionRate;
|
||||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
});
|
||||
|
||||
test('real-world multi-entry batch: 5 sequential tool calls — same total deduction as 5 legacy spendTokens calls', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 100000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-opus-4-5-20251101';
|
||||
const calls = [
|
||||
{ input_tokens: 31596, output_tokens: 151 },
|
||||
{ input_tokens: 35368, output_tokens: 150 },
|
||||
{ input_tokens: 58362, output_tokens: 295 },
|
||||
{ input_tokens: 112604, output_tokens: 193 },
|
||||
{ input_tokens: 257440, output_tokens: 2217 },
|
||||
];
|
||||
|
||||
let expectedTotalCost = 0;
|
||||
for (const { input_tokens, output_tokens } of calls) {
|
||||
const pm = getMultiplier({ model, tokenType: 'prompt', inputTokenCount: input_tokens });
|
||||
const cm = getMultiplier({ model, tokenType: 'completion', inputTokenCount: input_tokens });
|
||||
expectedTotalCost += input_tokens * pm + output_tokens * cm;
|
||||
}
|
||||
|
||||
await recordCollectedUsage(bulkDeps, {
|
||||
user: userId.toString(),
|
||||
conversationId: 'test-sequential',
|
||||
model,
|
||||
context: 'message',
|
||||
balance: { enabled: true },
|
||||
transactions: { enabled: true },
|
||||
collectedUsage: calls.map((c) => ({ ...c, model })),
|
||||
});
|
||||
|
||||
const txns = await Transaction.find({ user: userId }).lean();
|
||||
expect(txns).toHaveLength(10); // 5 calls × 2 docs (prompt + completion)
|
||||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedTotalCost, 0);
|
||||
});
|
||||
|
||||
test('bulk path should save transaction but not update balance when balance disabled, transactions enabled', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
await recordCollectedUsage(bulkDeps, {
|
||||
user: userId.toString(),
|
||||
conversationId: 'test-conversation-id',
|
||||
model: 'gpt-3.5-turbo',
|
||||
context: 'test',
|
||||
balance: { enabled: false },
|
||||
transactions: { enabled: true },
|
||||
collectedUsage: [{ input_tokens: 100, output_tokens: 50, model: 'gpt-3.5-turbo' }],
|
||||
});
|
||||
|
||||
const txns = await Transaction.find({ user: userId }).lean();
|
||||
expect(txns).toHaveLength(2);
|
||||
expect(txns[0].rawAmount).toBeDefined();
|
||||
const balance = await Balance.findOne({ user: userId });
|
||||
expect(balance.tokenCredits).toBe(initialBalance);
|
||||
});
|
||||
|
||||
test('bulk path structured tokens should not save when transactions.enabled is false', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
await recordCollectedUsage(bulkDeps, {
|
||||
user: userId.toString(),
|
||||
conversationId: 'test-conversation-id',
|
||||
model: 'claude-3-5-sonnet',
|
||||
context: 'message',
|
||||
balance: { enabled: true },
|
||||
transactions: { enabled: false },
|
||||
collectedUsage: [
|
||||
{
|
||||
input_tokens: 10,
|
||||
output_tokens: 5,
|
||||
model: 'claude-3-5-sonnet',
|
||||
input_token_details: { cache_creation: 100, cache_read: 5 },
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
const txns = await Transaction.find({ user: userId }).lean();
|
||||
expect(txns).toHaveLength(0);
|
||||
const balance = await Balance.findOne({ user: userId });
|
||||
expect(balance.tokenCredits).toBe(initialBalance);
|
||||
});
|
||||
|
||||
test('bulk path structured tokens should save but not update balance when balance disabled', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
await recordCollectedUsage(bulkDeps, {
|
||||
user: userId.toString(),
|
||||
conversationId: 'test-conversation-id',
|
||||
model: 'claude-3-5-sonnet',
|
||||
context: 'message',
|
||||
balance: { enabled: false },
|
||||
transactions: { enabled: true },
|
||||
collectedUsage: [
|
||||
{
|
||||
input_tokens: 10,
|
||||
output_tokens: 5,
|
||||
model: 'claude-3-5-sonnet',
|
||||
input_token_details: { cache_creation: 100, cache_read: 5 },
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
const txns = await Transaction.find({ user: userId }).lean();
|
||||
expect(txns).toHaveLength(2);
|
||||
const promptTx = txns.find((t) => t.tokenType === 'prompt');
|
||||
expect(promptTx.inputTokens).toBe(-10);
|
||||
expect(promptTx.writeTokens).toBe(-100);
|
||||
expect(promptTx.readTokens).toBe(-5);
|
||||
const balance = await Balance.findOne({ user: userId });
|
||||
expect(balance.tokenCredits).toBe(initialBalance);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import logger from '~/config/winston';
|
||||
import type { FilterQuery, Model, Types } from 'mongoose';
|
||||
import type { IBalance, IBalanceUpdate, TransactionData } from '~/types';
|
||||
import type { ITransaction } from '~/schema/transaction';
|
||||
import type { IBalance, IBalanceUpdate } from '~/types';
|
||||
|
||||
const cancelRate = 1.15;
|
||||
|
||||
|
|
@ -408,7 +408,22 @@ export function createTransactionMethods(
|
|||
return Balance.deleteMany(filter);
|
||||
}
|
||||
|
||||
async function bulkInsertTransactions(docs: TransactionData[]): Promise<void> {
|
||||
if (!docs.length) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
const Transaction = mongoose.models.Transaction;
|
||||
await Transaction.insertMany(docs);
|
||||
} catch (error) {
|
||||
logger.error('[bulkInsertTransactions] Error inserting transaction docs:', error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
updateBalance,
|
||||
bulkInsertTransactions,
|
||||
findBalanceByUser,
|
||||
upsertBalanceFields,
|
||||
getTransactions,
|
||||
|
|
|
|||
|
|
@ -20,7 +20,10 @@ export interface TxDeps {
|
|||
/** From @librechat/api — matches a model name to a canonical key. */
|
||||
matchModelName: (model: string, endpoint?: string) => string | undefined;
|
||||
/** From @librechat/api — finds the longest key in `values` whose key is a substring of `model`. */
|
||||
findMatchingPattern: (model: string, values: Record<string, unknown>) => string | undefined;
|
||||
findMatchingPattern: (
|
||||
model: string,
|
||||
values: Record<string, number | Record<string, number>>,
|
||||
) => string | undefined;
|
||||
}
|
||||
|
||||
export const defaultRate = 6;
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue