mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-02-21 09:54:08 +01:00
7a1d2969b8
- Added support for the new Gemini 3.1 models, including 'gemini-3.1-pro-preview' and 'gemini-3.1-pro-preview-customtools'. - Updated pricing logic to apply standard and premium rates based on token usage thresholds for the new models. - Enhanced tests to validate pricing behavior for both standard and premium scenarios. - Modified configuration files to include Gemini 3.1 models in the default model lists and token value mappings. - Updated environment example file to reflect the new model options.
1191 lines
38 KiB
JavaScript
1191 lines
38 KiB
JavaScript
const mongoose = require('mongoose');
|
|
const { MongoMemoryServer } = require('mongodb-memory-server');
|
|
const { createTransaction, createAutoRefillTransaction } = require('./Transaction');
|
|
const { tokenValues, premiumTokenValues, getCacheMultiplier } = require('./tx');
|
|
const { spendTokens, spendStructuredTokens } = require('./spendTokens');
|
|
|
|
require('~/db/models');
|
|
|
|
jest.mock('~/config', () => ({
|
|
logger: {
|
|
debug: jest.fn(),
|
|
error: jest.fn(),
|
|
},
|
|
}));
|
|
|
|
describe('spendTokens', () => {
|
|
let mongoServer;
|
|
let userId;
|
|
let Transaction;
|
|
let Balance;
|
|
|
|
beforeAll(async () => {
|
|
mongoServer = await MongoMemoryServer.create();
|
|
await mongoose.connect(mongoServer.getUri());
|
|
|
|
Transaction = mongoose.model('Transaction');
|
|
Balance = mongoose.model('Balance');
|
|
});
|
|
|
|
afterAll(async () => {
|
|
await mongoose.disconnect();
|
|
await mongoServer.stop();
|
|
});
|
|
|
|
beforeEach(async () => {
|
|
// Clear collections before each test
|
|
await Transaction.deleteMany({});
|
|
await Balance.deleteMany({});
|
|
|
|
// Create a new user ID for each test
|
|
userId = new mongoose.Types.ObjectId();
|
|
|
|
// Balance config is now passed directly in txData
|
|
});
|
|
|
|
it('should create transactions for both prompt and completion tokens', async () => {
|
|
// Create a balance for the user
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: 10000,
|
|
});
|
|
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-convo',
|
|
model: 'gpt-3.5-turbo',
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
const tokenUsage = {
|
|
promptTokens: 100,
|
|
completionTokens: 50,
|
|
};
|
|
|
|
await spendTokens(txData, tokenUsage);
|
|
|
|
// Verify transactions were created
|
|
const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
|
|
expect(transactions).toHaveLength(2);
|
|
|
|
// Check completion transaction
|
|
expect(transactions[0].tokenType).toBe('completion');
|
|
expect(transactions[0].rawAmount).toBe(-50);
|
|
|
|
// Check prompt transaction
|
|
expect(transactions[1].tokenType).toBe('prompt');
|
|
expect(transactions[1].rawAmount).toBe(-100);
|
|
|
|
// Verify balance was updated
|
|
const balance = await Balance.findOne({ user: userId });
|
|
expect(balance).toBeDefined();
|
|
expect(balance.tokenCredits).toBeLessThan(10000); // Balance should be reduced
|
|
});
|
|
|
|
it('should handle zero completion tokens', async () => {
|
|
// Create a balance for the user
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: 10000,
|
|
});
|
|
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-convo',
|
|
model: 'gpt-3.5-turbo',
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
const tokenUsage = {
|
|
promptTokens: 100,
|
|
completionTokens: 0,
|
|
};
|
|
|
|
await spendTokens(txData, tokenUsage);
|
|
|
|
// Verify transactions were created
|
|
const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
|
|
expect(transactions).toHaveLength(2);
|
|
|
|
// Check completion transaction
|
|
expect(transactions[0].tokenType).toBe('completion');
|
|
// In JavaScript -0 and 0 are different but functionally equivalent
|
|
// Use Math.abs to handle both 0 and -0
|
|
expect(Math.abs(transactions[0].rawAmount)).toBe(0);
|
|
|
|
// Check prompt transaction
|
|
expect(transactions[1].tokenType).toBe('prompt');
|
|
expect(transactions[1].rawAmount).toBe(-100);
|
|
});
|
|
|
|
it('should handle undefined token counts', async () => {
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-convo',
|
|
model: 'gpt-3.5-turbo',
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
const tokenUsage = {};
|
|
|
|
await spendTokens(txData, tokenUsage);
|
|
|
|
// Verify no transactions were created
|
|
const transactions = await Transaction.find({ user: userId });
|
|
expect(transactions).toHaveLength(0);
|
|
});
|
|
|
|
it('should not update balance when the balance feature is disabled', async () => {
|
|
// Balance is now passed directly in txData
|
|
// Create a balance for the user
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: 10000,
|
|
});
|
|
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-convo',
|
|
model: 'gpt-3.5-turbo',
|
|
context: 'test',
|
|
balance: { enabled: false },
|
|
};
|
|
const tokenUsage = {
|
|
promptTokens: 100,
|
|
completionTokens: 50,
|
|
};
|
|
|
|
await spendTokens(txData, tokenUsage);
|
|
|
|
// Verify transactions were created
|
|
const transactions = await Transaction.find({ user: userId });
|
|
expect(transactions).toHaveLength(2);
|
|
|
|
// Verify balance was not updated (should still be 10000)
|
|
const balance = await Balance.findOne({ user: userId });
|
|
expect(balance.tokenCredits).toBe(10000);
|
|
});
|
|
|
|
it('should not allow balance to go below zero when spending tokens', async () => {
|
|
// Create a balance with a low amount
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: 5000,
|
|
});
|
|
|
|
const txData = {
|
|
user: userId,
|
|
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
|
|
const tokenUsage = {
|
|
promptTokens: 1000,
|
|
completionTokens: 500,
|
|
};
|
|
|
|
await spendTokens(txData, tokenUsage);
|
|
|
|
// Verify transactions were created
|
|
const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
|
|
expect(transactions).toHaveLength(2);
|
|
|
|
// Verify balance was reduced to exactly 0, not negative
|
|
const balance = await Balance.findOne({ user: userId });
|
|
expect(balance).toBeDefined();
|
|
expect(balance.tokenCredits).toBe(0);
|
|
|
|
// Check that the transaction records show the adjusted values
|
|
const transactionResults = await Promise.all(
|
|
transactions.map((t) =>
|
|
createTransaction({
|
|
...txData,
|
|
tokenType: t.tokenType,
|
|
rawAmount: t.rawAmount,
|
|
}),
|
|
),
|
|
);
|
|
|
|
// The second transaction should have an adjusted value since balance is already 0
|
|
expect(transactionResults[1]).toEqual(
|
|
expect.objectContaining({
|
|
balance: 0,
|
|
}),
|
|
);
|
|
});
|
|
|
|
it('should handle multiple transactions in sequence with low balance and not increase balance', async () => {
|
|
// This test is specifically checking for the issue reported in production
|
|
// where the balance increases after a transaction when it should remain at 0
|
|
// Create a balance with a very low amount
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: 100,
|
|
});
|
|
|
|
// First transaction - should reduce balance to 0
|
|
const txData1 = {
|
|
user: userId,
|
|
conversationId: 'test-convo-1',
|
|
model: 'gpt-4',
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
const tokenUsage1 = {
|
|
promptTokens: 100,
|
|
completionTokens: 50,
|
|
};
|
|
|
|
await spendTokens(txData1, tokenUsage1);
|
|
|
|
// Check balance after first transaction
|
|
let balance = await Balance.findOne({ user: userId });
|
|
expect(balance.tokenCredits).toBe(0);
|
|
|
|
// Second transaction - should keep balance at 0, not make it negative or increase it
|
|
const txData2 = {
|
|
user: userId,
|
|
conversationId: 'test-convo-2',
|
|
model: 'gpt-4',
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
const tokenUsage2 = {
|
|
promptTokens: 200,
|
|
completionTokens: 100,
|
|
};
|
|
|
|
await spendTokens(txData2, tokenUsage2);
|
|
|
|
// Check balance after second transaction - should still be 0
|
|
balance = await Balance.findOne({ user: userId });
|
|
expect(balance.tokenCredits).toBe(0);
|
|
|
|
// Verify all transactions were created
|
|
const transactions = await Transaction.find({ user: userId });
|
|
expect(transactions).toHaveLength(4); // 2 transactions (prompt+completion) for each call
|
|
|
|
// Let's examine the actual transaction records to see what's happening
|
|
const transactionDetails = await Transaction.find({ user: userId }).sort({ createdAt: 1 });
|
|
|
|
// Log the transaction details for debugging
|
|
console.log('Transaction details:');
|
|
transactionDetails.forEach((tx, i) => {
|
|
console.log(`Transaction ${i + 1}:`, {
|
|
tokenType: tx.tokenType,
|
|
rawAmount: tx.rawAmount,
|
|
tokenValue: tx.tokenValue,
|
|
model: tx.model,
|
|
});
|
|
});
|
|
|
|
// Check the return values from Transaction.create directly
|
|
// This is to verify that the incrementValue is not becoming positive
|
|
const directResult = await createTransaction({
|
|
user: userId,
|
|
conversationId: 'test-convo-3',
|
|
model: 'gpt-4',
|
|
tokenType: 'completion',
|
|
rawAmount: -100,
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
});
|
|
|
|
console.log('Direct Transaction.create result:', directResult);
|
|
|
|
// The completion value should never be positive
|
|
expect(directResult.completion).not.toBeGreaterThan(0);
|
|
});
|
|
|
|
it('should ensure tokenValue is always negative for spending tokens', async () => {
|
|
// Create a balance for the user
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: 10000,
|
|
});
|
|
|
|
// Test with various models to check multiplier calculations
|
|
const models = ['gpt-3.5-turbo', 'gpt-4', 'claude-3-5-sonnet'];
|
|
|
|
for (const model of models) {
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: `test-convo-${model}`,
|
|
model,
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
const tokenUsage = {
|
|
promptTokens: 100,
|
|
completionTokens: 50,
|
|
};
|
|
|
|
await spendTokens(txData, tokenUsage);
|
|
|
|
// Get the transactions for this model
|
|
const transactions = await Transaction.find({
|
|
user: userId,
|
|
model,
|
|
});
|
|
|
|
// Verify tokenValue is negative for all transactions
|
|
transactions.forEach((tx) => {
|
|
console.log(`Model ${model}, Type ${tx.tokenType}: tokenValue = ${tx.tokenValue}`);
|
|
expect(tx.tokenValue).toBeLessThan(0);
|
|
});
|
|
}
|
|
});
|
|
|
|
it('should handle structured transactions in sequence with low balance', async () => {
|
|
// Create a balance with a very low amount
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: 100,
|
|
});
|
|
|
|
// First transaction - should reduce balance to 0
|
|
const txData1 = {
|
|
user: userId,
|
|
conversationId: 'test-convo-1',
|
|
model: 'claude-3-5-sonnet',
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
const tokenUsage1 = {
|
|
promptTokens: {
|
|
input: 10,
|
|
write: 100,
|
|
read: 5,
|
|
},
|
|
completionTokens: 50,
|
|
};
|
|
|
|
await spendStructuredTokens(txData1, tokenUsage1);
|
|
|
|
// Check balance after first transaction
|
|
let balance = await Balance.findOne({ user: userId });
|
|
expect(balance.tokenCredits).toBe(0);
|
|
|
|
// Second transaction - should keep balance at 0, not make it negative or increase it
|
|
const txData2 = {
|
|
user: userId,
|
|
conversationId: 'test-convo-2',
|
|
model: 'claude-3-5-sonnet',
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
const tokenUsage2 = {
|
|
promptTokens: {
|
|
input: 20,
|
|
write: 200,
|
|
read: 10,
|
|
},
|
|
completionTokens: 100,
|
|
};
|
|
|
|
await spendStructuredTokens(txData2, tokenUsage2);
|
|
|
|
// Check balance after second transaction - should still be 0
|
|
balance = await Balance.findOne({ user: userId });
|
|
expect(balance.tokenCredits).toBe(0);
|
|
|
|
// Verify all transactions were created
|
|
const transactions = await Transaction.find({ user: userId });
|
|
expect(transactions).toHaveLength(4); // 2 transactions (prompt+completion) for each call
|
|
|
|
// Let's examine the actual transaction records to see what's happening
|
|
const transactionDetails = await Transaction.find({ user: userId }).sort({ createdAt: 1 });
|
|
|
|
// Log the transaction details for debugging
|
|
console.log('Structured transaction details:');
|
|
transactionDetails.forEach((tx, i) => {
|
|
console.log(`Transaction ${i + 1}:`, {
|
|
tokenType: tx.tokenType,
|
|
rawAmount: tx.rawAmount,
|
|
tokenValue: tx.tokenValue,
|
|
inputTokens: tx.inputTokens,
|
|
writeTokens: tx.writeTokens,
|
|
readTokens: tx.readTokens,
|
|
model: tx.model,
|
|
});
|
|
});
|
|
});
|
|
|
|
it('should not allow balance to go below zero when spending structured tokens', async () => {
|
|
// Create a balance with a low amount
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: 5000,
|
|
});
|
|
|
|
const txData = {
|
|
user: userId,
|
|
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
|
|
const tokenUsage = {
|
|
promptTokens: {
|
|
input: 100,
|
|
write: 1000,
|
|
read: 50,
|
|
},
|
|
completionTokens: 500,
|
|
};
|
|
|
|
const result = await spendStructuredTokens(txData, tokenUsage);
|
|
|
|
// Verify transactions were created
|
|
const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
|
|
expect(transactions).toHaveLength(2);
|
|
|
|
// Verify balance was reduced to exactly 0, not negative
|
|
const balance = await Balance.findOne({ user: userId });
|
|
expect(balance).toBeDefined();
|
|
expect(balance.tokenCredits).toBe(0);
|
|
|
|
// The result should show the adjusted values
|
|
expect(result).toEqual({
|
|
prompt: expect.objectContaining({
|
|
user: userId.toString(),
|
|
balance: expect.any(Number),
|
|
}),
|
|
completion: expect.objectContaining({
|
|
user: userId.toString(),
|
|
balance: 0, // Final balance should be 0
|
|
}),
|
|
});
|
|
});
|
|
|
|
it('should handle multiple concurrent transactions correctly with a high balance', async () => {
|
|
// Create a balance with a high amount
|
|
const initialBalance = 10000000;
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: initialBalance,
|
|
});
|
|
|
|
// Simulate the recordCollectedUsage function from the production code
|
|
const conversationId = 'test-concurrent-convo';
|
|
const context = 'message';
|
|
const model = 'gpt-4';
|
|
|
|
const amount = 50;
|
|
// Create `amount` of usage records to simulate multiple transactions
|
|
const collectedUsage = Array.from({ length: amount }, (_, i) => ({
|
|
model,
|
|
input_tokens: 100 + i * 10, // Increasing input tokens
|
|
output_tokens: 50 + i * 5, // Increasing output tokens
|
|
input_token_details: {
|
|
cache_creation: i % 2 === 0 ? 20 : 0, // Some have cache creation
|
|
cache_read: i % 3 === 0 ? 10 : 0, // Some have cache read
|
|
},
|
|
}));
|
|
|
|
// Process all transactions concurrently to simulate race conditions
|
|
const promises = [];
|
|
let expectedTotalSpend = 0;
|
|
|
|
for (let i = 0; i < collectedUsage.length; i++) {
|
|
const usage = collectedUsage[i];
|
|
if (!usage) {
|
|
continue;
|
|
}
|
|
|
|
const cache_creation = Number(usage.input_token_details?.cache_creation) || 0;
|
|
const cache_read = Number(usage.input_token_details?.cache_read) || 0;
|
|
|
|
const txMetadata = {
|
|
context,
|
|
conversationId,
|
|
user: userId,
|
|
model: usage.model,
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
// Calculate expected spend for this transaction
|
|
const promptTokens = usage.input_tokens;
|
|
const completionTokens = usage.output_tokens;
|
|
|
|
// For regular transactions
|
|
if (cache_creation === 0 && cache_read === 0) {
|
|
// Add to expected spend using the correct multipliers from tx.js
|
|
// For gpt-4, the multipliers are: prompt=30, completion=60
|
|
expectedTotalSpend += promptTokens * 30; // gpt-4 prompt rate is 30
|
|
expectedTotalSpend += completionTokens * 60; // gpt-4 completion rate is 60
|
|
|
|
promises.push(
|
|
spendTokens(txMetadata, {
|
|
promptTokens,
|
|
completionTokens,
|
|
}),
|
|
);
|
|
} else {
|
|
// For structured transactions with cache operations
|
|
// The multipliers for claude models with cache operations are different
|
|
// But since we're using gpt-4 in the test, we need to use appropriate values
|
|
expectedTotalSpend += promptTokens * 30; // Base prompt rate for gpt-4
|
|
// Since gpt-4 doesn't have cache multipliers defined, we'll use the prompt rate
|
|
expectedTotalSpend += cache_creation * 30; // Write rate (using prompt rate as fallback)
|
|
expectedTotalSpend += cache_read * 30; // Read rate (using prompt rate as fallback)
|
|
expectedTotalSpend += completionTokens * 60; // Completion rate for gpt-4
|
|
|
|
promises.push(
|
|
spendStructuredTokens(txMetadata, {
|
|
promptTokens: {
|
|
input: promptTokens,
|
|
write: cache_creation,
|
|
read: cache_read,
|
|
},
|
|
completionTokens,
|
|
}),
|
|
);
|
|
}
|
|
}
|
|
|
|
// Wait for all transactions to complete
|
|
await Promise.all(promises);
|
|
|
|
// Verify final balance
|
|
const finalBalance = await Balance.findOne({ user: userId });
|
|
expect(finalBalance).toBeDefined();
|
|
|
|
// The final balance should be the initial balance minus the expected total spend
|
|
const expectedFinalBalance = initialBalance - expectedTotalSpend;
|
|
|
|
console.log('Initial balance:', initialBalance);
|
|
console.log('Expected total spend:', expectedTotalSpend);
|
|
console.log('Expected final balance:', expectedFinalBalance);
|
|
console.log('Actual final balance:', finalBalance.tokenCredits);
|
|
|
|
// Allow for small rounding differences
|
|
expect(finalBalance.tokenCredits).toBeCloseTo(expectedFinalBalance, 0);
|
|
|
|
// Verify all transactions were created
|
|
const transactions = await Transaction.find({
|
|
user: userId,
|
|
conversationId,
|
|
});
|
|
|
|
// We should have 2 transactions (prompt + completion) for each usage record
|
|
// Some might be structured, some regular
|
|
expect(transactions.length).toBeGreaterThanOrEqual(collectedUsage.length);
|
|
|
|
// Log transaction details for debugging
|
|
console.log('Transaction summary:');
|
|
let totalTokenValue = 0;
|
|
transactions.forEach((tx) => {
|
|
console.log(`${tx.tokenType}: rawAmount=${tx.rawAmount}, tokenValue=${tx.tokenValue}`);
|
|
totalTokenValue += tx.tokenValue;
|
|
});
|
|
console.log('Total token value from transactions:', totalTokenValue);
|
|
|
|
// The difference between expected and actual is significant
|
|
// This is likely due to the multipliers being different in the test environment
|
|
// Let's adjust our expectation based on the actual transactions
|
|
const actualSpend = initialBalance - finalBalance.tokenCredits;
|
|
console.log('Actual spend:', actualSpend);
|
|
|
|
// Instead of checking the exact balance, let's verify that:
|
|
// 1. The balance was reduced (tokens were spent)
|
|
expect(finalBalance.tokenCredits).toBeLessThan(initialBalance);
|
|
// 2. The total token value from transactions matches the actual spend
|
|
expect(Math.abs(totalTokenValue)).toBeCloseTo(actualSpend, -3); // Allow for larger differences
|
|
});
|
|
|
|
// Add this new test case
|
|
it('should handle multiple concurrent balance increases correctly', async () => {
|
|
// Start with zero balance
|
|
const initialBalance = 0;
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: initialBalance,
|
|
});
|
|
|
|
const numberOfRefills = 25;
|
|
const refillAmount = 1000;
|
|
|
|
const promises = [];
|
|
for (let i = 0; i < numberOfRefills; i++) {
|
|
promises.push(
|
|
createAutoRefillTransaction({
|
|
user: userId,
|
|
tokenType: 'credits',
|
|
context: 'concurrent-refill-test',
|
|
rawAmount: refillAmount,
|
|
balance: { enabled: true },
|
|
}),
|
|
);
|
|
}
|
|
|
|
// Wait for all refill transactions to complete
|
|
const results = await Promise.all(promises);
|
|
|
|
// Verify final balance
|
|
const finalBalance = await Balance.findOne({ user: userId });
|
|
expect(finalBalance).toBeDefined();
|
|
|
|
// The final balance should be the initial balance plus the sum of all refills
|
|
const expectedFinalBalance = initialBalance + numberOfRefills * refillAmount;
|
|
|
|
console.log('Initial balance (Increase Test):', initialBalance);
|
|
console.log(`Performed ${numberOfRefills} refills of ${refillAmount} each.`);
|
|
console.log('Expected final balance (Increase Test):', expectedFinalBalance);
|
|
console.log('Actual final balance (Increase Test):', finalBalance.tokenCredits);
|
|
|
|
// Use toBeCloseTo for safety, though toBe should work for integer math
|
|
expect(finalBalance.tokenCredits).toBeCloseTo(expectedFinalBalance, 0);
|
|
|
|
// Verify all transactions were created
|
|
const transactions = await Transaction.find({
|
|
user: userId,
|
|
context: 'concurrent-refill-test',
|
|
});
|
|
|
|
// We should have one transaction for each refill attempt
|
|
expect(transactions.length).toBe(numberOfRefills);
|
|
|
|
// Optional: Verify the sum of increments from the results matches the balance change
|
|
const totalIncrementReported = results.reduce((sum, result) => {
|
|
// Assuming createAutoRefillTransaction returns an object with the increment amount
|
|
// Adjust this based on the actual return structure.
|
|
// Let's assume it returns { balance: newBalance, transaction: { rawAmount: ... } }
|
|
// Or perhaps we check the transaction.rawAmount directly
|
|
return sum + (result?.transaction?.rawAmount || 0);
|
|
}, 0);
|
|
console.log('Total increment reported by results:', totalIncrementReported);
|
|
expect(totalIncrementReported).toBe(expectedFinalBalance - initialBalance);
|
|
|
|
// Optional: Check the sum of tokenValue from saved transactions
|
|
let totalTokenValueFromDb = 0;
|
|
transactions.forEach((tx) => {
|
|
// For refills, rawAmount is positive, and tokenValue might be calculated based on it
|
|
// Let's assume tokenValue directly reflects the increment for simplicity here
|
|
// If calculation is involved, adjust accordingly
|
|
totalTokenValueFromDb += tx.rawAmount; // Or tx.tokenValue if that holds the increment
|
|
});
|
|
console.log('Total rawAmount from DB transactions:', totalTokenValueFromDb);
|
|
expect(totalTokenValueFromDb).toBeCloseTo(expectedFinalBalance - initialBalance, 0);
|
|
});
|
|
|
|
it('should create structured transactions for both prompt and completion tokens', async () => {
|
|
// Create a balance for the user
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: 10000,
|
|
});
|
|
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-convo',
|
|
model: 'claude-3-5-sonnet',
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
const tokenUsage = {
|
|
promptTokens: {
|
|
input: 10,
|
|
write: 100,
|
|
read: 5,
|
|
},
|
|
completionTokens: 50,
|
|
};
|
|
|
|
const result = await spendStructuredTokens(txData, tokenUsage);
|
|
|
|
// Verify transactions were created
|
|
const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
|
|
expect(transactions).toHaveLength(2);
|
|
|
|
// Check completion transaction
|
|
expect(transactions[0].tokenType).toBe('completion');
|
|
expect(transactions[0].rawAmount).toBe(-50);
|
|
|
|
// Check prompt transaction
|
|
expect(transactions[1].tokenType).toBe('prompt');
|
|
expect(transactions[1].inputTokens).toBe(-10);
|
|
expect(transactions[1].writeTokens).toBe(-100);
|
|
expect(transactions[1].readTokens).toBe(-5);
|
|
|
|
// Verify result contains transaction info
|
|
expect(result).toEqual({
|
|
prompt: expect.objectContaining({
|
|
user: userId.toString(),
|
|
prompt: expect.any(Number),
|
|
}),
|
|
completion: expect.objectContaining({
|
|
user: userId.toString(),
|
|
completion: expect.any(Number),
|
|
}),
|
|
});
|
|
|
|
// Verify balance was updated
|
|
const balance = await Balance.findOne({ user: userId });
|
|
expect(balance).toBeDefined();
|
|
expect(balance.tokenCredits).toBeLessThan(10000); // Balance should be reduced
|
|
});
|
|
|
|
describe('premium token pricing', () => {
|
|
it('should charge standard rates for claude-opus-4-6 when prompt tokens are below threshold', async () => {
|
|
const initialBalance = 100000000;
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: initialBalance,
|
|
});
|
|
|
|
const model = 'claude-opus-4-6';
|
|
const promptTokens = 100000;
|
|
const completionTokens = 500;
|
|
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-standard-pricing',
|
|
model,
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
await spendTokens(txData, { promptTokens, completionTokens });
|
|
|
|
const expectedCost =
|
|
promptTokens * tokenValues[model].prompt + completionTokens * tokenValues[model].completion;
|
|
|
|
const balance = await Balance.findOne({ user: userId });
|
|
expect(balance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
|
});
|
|
|
|
it('should charge premium rates for claude-opus-4-6 when prompt tokens exceed threshold', async () => {
|
|
const initialBalance = 100000000;
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: initialBalance,
|
|
});
|
|
|
|
const model = 'claude-opus-4-6';
|
|
const promptTokens = 250000;
|
|
const completionTokens = 500;
|
|
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-premium-pricing',
|
|
model,
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
await spendTokens(txData, { promptTokens, completionTokens });
|
|
|
|
const expectedCost =
|
|
promptTokens * premiumTokenValues[model].prompt +
|
|
completionTokens * premiumTokenValues[model].completion;
|
|
|
|
const balance = await Balance.findOne({ user: userId });
|
|
expect(balance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
|
});
|
|
|
|
it('should charge premium rates for both prompt and completion in structured tokens when above threshold', async () => {
|
|
const initialBalance = 100000000;
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: initialBalance,
|
|
});
|
|
|
|
const model = 'claude-opus-4-6';
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-structured-premium',
|
|
model,
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
const tokenUsage = {
|
|
promptTokens: {
|
|
input: 200000,
|
|
write: 10000,
|
|
read: 5000,
|
|
},
|
|
completionTokens: 1000,
|
|
};
|
|
|
|
const result = await spendStructuredTokens(txData, tokenUsage);
|
|
|
|
const premiumPromptRate = premiumTokenValues[model].prompt;
|
|
const premiumCompletionRate = premiumTokenValues[model].completion;
|
|
const writeRate = getCacheMultiplier({ model, cacheType: 'write' });
|
|
const readRate = getCacheMultiplier({ model, cacheType: 'read' });
|
|
|
|
const expectedPromptCost =
|
|
tokenUsage.promptTokens.input * premiumPromptRate +
|
|
tokenUsage.promptTokens.write * writeRate +
|
|
tokenUsage.promptTokens.read * readRate;
|
|
const expectedCompletionCost = tokenUsage.completionTokens * premiumCompletionRate;
|
|
|
|
expect(result.prompt.prompt).toBeCloseTo(-expectedPromptCost, 0);
|
|
expect(result.completion.completion).toBeCloseTo(-expectedCompletionCost, 0);
|
|
});
|
|
|
|
it('should charge standard rates for structured tokens when below threshold', async () => {
|
|
const initialBalance = 100000000;
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: initialBalance,
|
|
});
|
|
|
|
const model = 'claude-opus-4-6';
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-structured-standard',
|
|
model,
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
const tokenUsage = {
|
|
promptTokens: {
|
|
input: 50000,
|
|
write: 10000,
|
|
read: 5000,
|
|
},
|
|
completionTokens: 1000,
|
|
};
|
|
|
|
const result = await spendStructuredTokens(txData, tokenUsage);
|
|
|
|
const standardPromptRate = tokenValues[model].prompt;
|
|
const standardCompletionRate = tokenValues[model].completion;
|
|
const writeRate = getCacheMultiplier({ model, cacheType: 'write' });
|
|
const readRate = getCacheMultiplier({ model, cacheType: 'read' });
|
|
|
|
const expectedPromptCost =
|
|
tokenUsage.promptTokens.input * standardPromptRate +
|
|
tokenUsage.promptTokens.write * writeRate +
|
|
tokenUsage.promptTokens.read * readRate;
|
|
const expectedCompletionCost = tokenUsage.completionTokens * standardCompletionRate;
|
|
|
|
expect(result.prompt.prompt).toBeCloseTo(-expectedPromptCost, 0);
|
|
expect(result.completion.completion).toBeCloseTo(-expectedCompletionCost, 0);
|
|
});
|
|
|
|
it('should charge standard rates for gemini-3.1-pro-preview when prompt tokens are below threshold', async () => {
|
|
const initialBalance = 100000000;
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: initialBalance,
|
|
});
|
|
|
|
const model = 'gemini-3.1-pro-preview';
|
|
const promptTokens = 100000;
|
|
const completionTokens = 500;
|
|
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-gemini31-standard-pricing',
|
|
model,
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
await spendTokens(txData, { promptTokens, completionTokens });
|
|
|
|
const expectedCost =
|
|
promptTokens * tokenValues['gemini-3.1'].prompt +
|
|
completionTokens * tokenValues['gemini-3.1'].completion;
|
|
|
|
const balance = await Balance.findOne({ user: userId });
|
|
expect(balance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
|
});
|
|
|
|
it('should charge premium rates for gemini-3.1-pro-preview when prompt tokens exceed threshold', async () => {
|
|
const initialBalance = 100000000;
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: initialBalance,
|
|
});
|
|
|
|
const model = 'gemini-3.1-pro-preview';
|
|
const promptTokens = 250000;
|
|
const completionTokens = 500;
|
|
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-gemini31-premium-pricing',
|
|
model,
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
await spendTokens(txData, { promptTokens, completionTokens });
|
|
|
|
const expectedCost =
|
|
promptTokens * premiumTokenValues['gemini-3.1'].prompt +
|
|
completionTokens * premiumTokenValues['gemini-3.1'].completion;
|
|
|
|
const balance = await Balance.findOne({ user: userId });
|
|
expect(balance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
|
});
|
|
|
|
it('should charge premium rates for gemini-3.1-pro-preview-customtools when prompt tokens exceed threshold', async () => {
|
|
const initialBalance = 100000000;
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: initialBalance,
|
|
});
|
|
|
|
const model = 'gemini-3.1-pro-preview-customtools';
|
|
const promptTokens = 250000;
|
|
const completionTokens = 500;
|
|
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-gemini31-customtools-premium',
|
|
model,
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
await spendTokens(txData, { promptTokens, completionTokens });
|
|
|
|
const expectedCost =
|
|
promptTokens * premiumTokenValues['gemini-3.1'].prompt +
|
|
completionTokens * premiumTokenValues['gemini-3.1'].completion;
|
|
|
|
const balance = await Balance.findOne({ user: userId });
|
|
expect(balance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
|
});
|
|
|
|
it('should charge premium rates for structured gemini-3.1 tokens when total input exceeds threshold', async () => {
|
|
const initialBalance = 100000000;
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: initialBalance,
|
|
});
|
|
|
|
const model = 'gemini-3.1-pro-preview';
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-gemini31-structured-premium',
|
|
model,
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
const tokenUsage = {
|
|
promptTokens: {
|
|
input: 200000,
|
|
write: 10000,
|
|
read: 5000,
|
|
},
|
|
completionTokens: 1000,
|
|
};
|
|
|
|
const result = await spendStructuredTokens(txData, tokenUsage);
|
|
|
|
const premiumPromptRate = premiumTokenValues['gemini-3.1'].prompt;
|
|
const premiumCompletionRate = premiumTokenValues['gemini-3.1'].completion;
|
|
const writeRate = getCacheMultiplier({ model, cacheType: 'write' });
|
|
const readRate = getCacheMultiplier({ model, cacheType: 'read' });
|
|
|
|
const expectedPromptCost =
|
|
tokenUsage.promptTokens.input * premiumPromptRate +
|
|
tokenUsage.promptTokens.write * writeRate +
|
|
tokenUsage.promptTokens.read * readRate;
|
|
const expectedCompletionCost = tokenUsage.completionTokens * premiumCompletionRate;
|
|
|
|
expect(result.prompt.prompt).toBeCloseTo(-expectedPromptCost, 0);
|
|
expect(result.completion.completion).toBeCloseTo(-expectedCompletionCost, 0);
|
|
});
|
|
|
|
it('should not apply premium pricing to non-premium models regardless of prompt size', async () => {
|
|
const initialBalance = 100000000;
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: initialBalance,
|
|
});
|
|
|
|
const model = 'claude-opus-4-5';
|
|
const promptTokens = 300000;
|
|
const completionTokens = 500;
|
|
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-no-premium',
|
|
model,
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
await spendTokens(txData, { promptTokens, completionTokens });
|
|
|
|
const expectedCost =
|
|
promptTokens * tokenValues[model].prompt + completionTokens * tokenValues[model].completion;
|
|
|
|
const balance = await Balance.findOne({ user: userId });
|
|
expect(balance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
|
});
|
|
});
|
|
|
|
describe('inputTokenCount Normalization', () => {
|
|
it('should normalize negative promptTokens to zero for inputTokenCount', async () => {
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: 100000000,
|
|
});
|
|
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-negative-prompt',
|
|
model: 'claude-opus-4-6',
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
await spendTokens(txData, { promptTokens: -500, completionTokens: 100 });
|
|
|
|
const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
|
|
|
|
const completionTx = transactions.find((t) => t.tokenType === 'completion');
|
|
const promptTx = transactions.find((t) => t.tokenType === 'prompt');
|
|
|
|
expect(Math.abs(promptTx.rawAmount)).toBe(0);
|
|
expect(completionTx.rawAmount).toBe(-100);
|
|
|
|
const standardCompletionRate = tokenValues['claude-opus-4-6'].completion;
|
|
expect(completionTx.rate).toBe(standardCompletionRate);
|
|
});
|
|
|
|
it('should use normalized inputTokenCount for premium threshold check on completion', async () => {
|
|
const initialBalance = 100000000;
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: initialBalance,
|
|
});
|
|
|
|
const model = 'claude-opus-4-6';
|
|
const promptTokens = 250000;
|
|
const completionTokens = 500;
|
|
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-normalized-premium',
|
|
model,
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
await spendTokens(txData, { promptTokens, completionTokens });
|
|
|
|
const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
|
|
const completionTx = transactions.find((t) => t.tokenType === 'completion');
|
|
const promptTx = transactions.find((t) => t.tokenType === 'prompt');
|
|
|
|
const premiumPromptRate = premiumTokenValues[model].prompt;
|
|
const premiumCompletionRate = premiumTokenValues[model].completion;
|
|
expect(promptTx.rate).toBe(premiumPromptRate);
|
|
expect(completionTx.rate).toBe(premiumCompletionRate);
|
|
});
|
|
|
|
it('should keep inputTokenCount as zero when promptTokens is zero', async () => {
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: 100000000,
|
|
});
|
|
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-zero-prompt',
|
|
model: 'claude-opus-4-6',
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
await spendTokens(txData, { promptTokens: 0, completionTokens: 100 });
|
|
|
|
const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
|
|
const completionTx = transactions.find((t) => t.tokenType === 'completion');
|
|
const promptTx = transactions.find((t) => t.tokenType === 'prompt');
|
|
|
|
expect(Math.abs(promptTx.rawAmount)).toBe(0);
|
|
|
|
const standardCompletionRate = tokenValues['claude-opus-4-6'].completion;
|
|
expect(completionTx.rate).toBe(standardCompletionRate);
|
|
});
|
|
|
|
it('should not trigger premium pricing with negative promptTokens on premium model', async () => {
|
|
const initialBalance = 100000000;
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: initialBalance,
|
|
});
|
|
|
|
const model = 'claude-opus-4-6';
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-negative-no-premium',
|
|
model,
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
await spendTokens(txData, { promptTokens: -300000, completionTokens: 500 });
|
|
|
|
const transactions = await Transaction.find({ user: userId }).sort({ tokenType: 1 });
|
|
const completionTx = transactions.find((t) => t.tokenType === 'completion');
|
|
|
|
const standardCompletionRate = tokenValues[model].completion;
|
|
expect(completionTx.rate).toBe(standardCompletionRate);
|
|
});
|
|
|
|
it('should normalize negative structured token values to zero in spendStructuredTokens', async () => {
|
|
const initialBalance = 100000000;
|
|
await Balance.create({
|
|
user: userId,
|
|
tokenCredits: initialBalance,
|
|
});
|
|
|
|
const model = 'claude-opus-4-6';
|
|
const txData = {
|
|
user: userId,
|
|
conversationId: 'test-negative-structured',
|
|
model,
|
|
context: 'test',
|
|
balance: { enabled: true },
|
|
};
|
|
|
|
const tokenUsage = {
|
|
promptTokens: { input: -100, write: 50, read: -30 },
|
|
completionTokens: -200,
|
|
};
|
|
|
|
await spendStructuredTokens(txData, tokenUsage);
|
|
|
|
const transactions = await Transaction.find({
|
|
user: userId,
|
|
conversationId: 'test-negative-structured',
|
|
}).sort({ tokenType: 1 });
|
|
|
|
const completionTx = transactions.find((t) => t.tokenType === 'completion');
|
|
const promptTx = transactions.find((t) => t.tokenType === 'prompt');
|
|
|
|
expect(Math.abs(promptTx.inputTokens)).toBe(0);
|
|
expect(promptTx.writeTokens).toBe(-50);
|
|
expect(Math.abs(promptTx.readTokens)).toBe(0);
|
|
|
|
expect(Math.abs(completionTx.rawAmount)).toBe(0);
|
|
|
|
const standardRate = tokenValues[model].completion;
|
|
expect(completionTx.rate).toBe(standardRate);
|
|
});
|
|
});
|
|
});
|