mirror of
https://github.com/danny-avila/LibreChat.git
synced 2025-12-24 04:10:15 +01:00
🧮 feat: Improve structured token spending and testing; fix: Anthropic Cache Spend (#3766)
* chore: update jest and mongodb-memory-server dependencies * fix: Anthropic edge case of increasing balance * refactor: Update token calculations in Transaction model/spec * refactor: `spendTokens` always record transactions, add Tx related tests * feat: Add error handling for CHECK_BALANCE environment variable * feat: Update import path for Balance model in Balance controller * chore: remove logging * refactor: Improve structured token spend logging in spendTokens.js * ci: add unit tests for spend token * ci: Improve structured token spend unit testing * chore: improve logging phrase for no tx spent from balance
This commit is contained in:
Danny Avila
GitHub
parent
ea5140ff0f
commit
d54458b3a6
8 changed files with 850 additions and 82 deletions
|
|
@ -1,5 +1,5 @@
|
|||
const mongoose = require('mongoose');
|
||||
const { isEnabled } = require('../server/utils/handleText');
|
||||
const { isEnabled } = require('~/server/utils/handleText');
|
||||
const transactionSchema = require('./schema/transaction');
|
||||
const { getMultiplier, getCacheMultiplier } = require('./tx');
|
||||
const { logger } = require('~/config');
|
||||
|
|
@ -76,7 +76,7 @@ transactionSchema.statics.createStructured = async function (txData) {
|
|||
await transaction.save();
|
||||
|
||||
if (!isEnabled(process.env.CHECK_BALANCE)) {
|
||||
return transaction;
|
||||
return;
|
||||
}
|
||||
|
||||
let balance = await Balance.findOne({ user: transaction.user }).lean();
|
||||
|
|
@ -122,28 +122,33 @@ transactionSchema.methods.calculateStructuredTokenValue = function () {
|
|||
read: readMultiplier,
|
||||
};
|
||||
|
||||
const totalTokens = (this.inputTokens || 0) + (this.writeTokens || 0) + (this.readTokens || 0);
|
||||
const totalPromptTokens =
|
||||
Math.abs(this.inputTokens || 0) +
|
||||
Math.abs(this.writeTokens || 0) +
|
||||
Math.abs(this.readTokens || 0);
|
||||
|
||||
if (totalTokens > 0) {
|
||||
if (totalPromptTokens > 0) {
|
||||
this.rate =
|
||||
(inputMultiplier * (this.inputTokens || 0) +
|
||||
writeMultiplier * (this.writeTokens || 0) +
|
||||
readMultiplier * (this.readTokens || 0)) /
|
||||
totalTokens;
|
||||
(Math.abs(inputMultiplier * (this.inputTokens || 0)) +
|
||||
Math.abs(writeMultiplier * (this.writeTokens || 0)) +
|
||||
Math.abs(readMultiplier * (this.readTokens || 0))) /
|
||||
totalPromptTokens;
|
||||
} else {
|
||||
this.rate = inputMultiplier; // Default to input rate if no tokens
|
||||
this.rate = Math.abs(inputMultiplier); // Default to input rate if no tokens
|
||||
}
|
||||
|
||||
this.tokenValue =
|
||||
this.inputTokens * inputMultiplier +
|
||||
(this.writeTokens || 0) * writeMultiplier +
|
||||
(this.readTokens || 0) * readMultiplier;
|
||||
} else {
|
||||
const multiplier = Math.abs(
|
||||
getMultiplier({ tokenType: this.tokenType, model, endpointTokenConfig }),
|
||||
this.tokenValue = -(
|
||||
Math.abs(this.inputTokens || 0) * inputMultiplier +
|
||||
Math.abs(this.writeTokens || 0) * writeMultiplier +
|
||||
Math.abs(this.readTokens || 0) * readMultiplier
|
||||
);
|
||||
this.rate = multiplier;
|
||||
this.tokenValue = this.rawAmount * multiplier;
|
||||
|
||||
this.rawAmount = -totalPromptTokens;
|
||||
} else if (this.tokenType === 'completion') {
|
||||
const multiplier = getMultiplier({ tokenType: this.tokenType, model, endpointTokenConfig });
|
||||
this.rate = Math.abs(multiplier);
|
||||
this.tokenValue = -Math.abs(this.rawAmount) * multiplier;
|
||||
this.rawAmount = -Math.abs(this.rawAmount);
|
||||
}
|
||||
|
||||
if (this.context && this.tokenType === 'completion' && this.context === 'incomplete') {
|
||||
|
|
|
|||
348
api/models/Transaction.spec.js
Normal file
348
api/models/Transaction.spec.js
Normal file
|
|
@ -0,0 +1,348 @@
|
|||
const mongoose = require('mongoose');
|
||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||
const Balance = require('./Balance');
|
||||
const { spendTokens, spendStructuredTokens } = require('./spendTokens');
|
||||
const { getMultiplier, getCacheMultiplier } = require('./tx');
|
||||
|
||||
let mongoServer;
|
||||
|
||||
beforeAll(async () => {
|
||||
mongoServer = await MongoMemoryServer.create();
|
||||
const mongoUri = mongoServer.getUri();
|
||||
await mongoose.connect(mongoUri);
|
||||
});
|
||||
|
||||
afterAll(async () => {
|
||||
await mongoose.disconnect();
|
||||
await mongoServer.stop();
|
||||
});
|
||||
|
||||
beforeEach(async () => {
|
||||
await mongoose.connection.dropDatabase();
|
||||
});
|
||||
|
||||
describe('Regular Token Spending Tests', () => {
|
||||
test('Balance should decrease when spending tokens with spendTokens', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000; // $10.00
|
||||
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
|
||||
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 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
|
||||
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: 0,
|
||||
};
|
||||
|
||||
// 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;
|
||||
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 () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000; // $10.00
|
||||
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 = {};
|
||||
|
||||
const result = await spendTokens(txData, tokenUsage);
|
||||
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
|
||||
test('spendTokens should handle only prompt tokens', async () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 10000000; // $10.00
|
||||
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 };
|
||||
|
||||
await spendTokens(txData, tokenUsage);
|
||||
|
||||
const updatedBalance = await Balance.findOne({ user: userId });
|
||||
|
||||
const promptMultiplier = getMultiplier({ model, tokenType: 'prompt' });
|
||||
const expectedCost = 100 * promptMultiplier;
|
||||
expect(updatedBalance.tokenCredits).toBeCloseTo(initialBalance - expectedCost, 0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Structured Token Spending Tests', () => {
|
||||
test('Balance should decrease and rawAmount should be set when spending a large number of structured tokens', async () => {
|
||||
// Arrange
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 17613154.55; // $17.61
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-3-5-sonnet';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'c23a18da-706c-470a-ac28-ec87ed065199',
|
||||
model,
|
||||
context: 'message',
|
||||
endpointTokenConfig: null, // We'll use the default rates
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
promptTokens: {
|
||||
input: 11,
|
||||
write: 140522,
|
||||
read: 0,
|
||||
},
|
||||
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);
|
||||
|
||||
const expectedPromptCost =
|
||||
tokenUsage.promptTokens.input * promptMultiplier +
|
||||
tokenUsage.promptTokens.write * writeMultiplier +
|
||||
tokenUsage.promptTokens.read * readMultiplier;
|
||||
const expectedCompletionCost = tokenUsage.completionTokens * completionMultiplier;
|
||||
const expectedTotalCost = expectedPromptCost + expectedCompletionCost;
|
||||
const expectedBalance = initialBalance - expectedTotalCost;
|
||||
|
||||
console.log('Expected Cost:', expectedTotalCost);
|
||||
console.log('Expected Balance:', expectedBalance);
|
||||
|
||||
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;
|
||||
|
||||
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 () => {
|
||||
const userId = new mongoose.Types.ObjectId();
|
||||
const initialBalance = 17613154.55;
|
||||
await Balance.create({ user: userId, tokenCredits: initialBalance });
|
||||
|
||||
const model = 'claude-3-5-sonnet';
|
||||
const txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-convo',
|
||||
model,
|
||||
context: 'message',
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
promptTokens: {
|
||||
input: 10,
|
||||
write: 100,
|
||||
read: 5,
|
||||
},
|
||||
completionTokens: 0,
|
||||
};
|
||||
|
||||
process.env.CHECK_BALANCE = 'true';
|
||||
const result = await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
expect(result.prompt).toBeDefined();
|
||||
expect(result.completion).toBeUndefined();
|
||||
expect(result.prompt.prompt).toBeLessThan(0);
|
||||
});
|
||||
|
||||
test('should handle only prompt tokens in structured spending', 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 txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-convo',
|
||||
model,
|
||||
context: 'message',
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
promptTokens: {
|
||||
input: 10,
|
||||
write: 100,
|
||||
read: 5,
|
||||
},
|
||||
};
|
||||
|
||||
process.env.CHECK_BALANCE = 'true';
|
||||
const result = await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
expect(result.prompt).toBeDefined();
|
||||
expect(result.completion).toBeUndefined();
|
||||
expect(result.prompt.prompt).toBeLessThan(0);
|
||||
});
|
||||
|
||||
test('should handle undefined token counts in structured spending', 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 txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-convo',
|
||||
model,
|
||||
context: 'message',
|
||||
};
|
||||
|
||||
const tokenUsage = {};
|
||||
|
||||
process.env.CHECK_BALANCE = 'true';
|
||||
const result = await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
expect(result).toEqual({
|
||||
prompt: undefined,
|
||||
completion: undefined,
|
||||
});
|
||||
});
|
||||
|
||||
test('should handle incomplete context for completion tokens', 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 txData = {
|
||||
user: userId,
|
||||
conversationId: 'test-convo',
|
||||
model,
|
||||
context: 'incomplete',
|
||||
};
|
||||
|
||||
const tokenUsage = {
|
||||
promptTokens: {
|
||||
input: 10,
|
||||
write: 100,
|
||||
read: 5,
|
||||
},
|
||||
completionTokens: 50,
|
||||
};
|
||||
|
||||
process.env.CHECK_BALANCE = 'true';
|
||||
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
|
||||
});
|
||||
});
|
||||
|
|
@ -32,35 +32,34 @@ const spendTokens = async (txData, tokenUsage) => {
|
|||
);
|
||||
let prompt, completion;
|
||||
try {
|
||||
if (promptTokens >= 0) {
|
||||
if (promptTokens !== undefined) {
|
||||
prompt = await Transaction.create({
|
||||
...txData,
|
||||
tokenType: 'prompt',
|
||||
rawAmount: -promptTokens,
|
||||
rawAmount: -Math.max(promptTokens, 0),
|
||||
});
|
||||
}
|
||||
|
||||
if (!completionTokens && isNaN(completionTokens)) {
|
||||
logger.debug('[spendTokens] !completionTokens', { prompt, completion });
|
||||
return;
|
||||
if (completionTokens !== undefined) {
|
||||
completion = await Transaction.create({
|
||||
...txData,
|
||||
tokenType: 'completion',
|
||||
rawAmount: -Math.max(completionTokens, 0),
|
||||
});
|
||||
}
|
||||
|
||||
completion = await Transaction.create({
|
||||
...txData,
|
||||
tokenType: 'completion',
|
||||
rawAmount: -completionTokens,
|
||||
});
|
||||
|
||||
prompt &&
|
||||
completion &&
|
||||
if (prompt || completion) {
|
||||
logger.debug('[spendTokens] Transaction data record against balance:', {
|
||||
user: txData.user,
|
||||
prompt: prompt.prompt,
|
||||
promptRate: prompt.rate,
|
||||
completion: completion.completion,
|
||||
completionRate: completion.rate,
|
||||
balance: completion.balance,
|
||||
prompt: prompt?.prompt,
|
||||
promptRate: prompt?.rate,
|
||||
completion: completion?.completion,
|
||||
completionRate: completion?.rate,
|
||||
balance: completion?.balance ?? prompt?.balance,
|
||||
});
|
||||
} else {
|
||||
logger.debug('[spendTokens] No transactions incurred against balance');
|
||||
}
|
||||
} catch (err) {
|
||||
logger.error('[spendTokens]', err);
|
||||
}
|
||||
|
|
@ -102,14 +101,12 @@ const spendStructuredTokens = async (txData, tokenUsage) => {
|
|||
try {
|
||||
if (promptTokens) {
|
||||
const { input = 0, write = 0, read = 0 } = promptTokens;
|
||||
const promptAmount = input + write + read;
|
||||
prompt = await Transaction.createStructured({
|
||||
...txData,
|
||||
tokenType: 'prompt',
|
||||
rawAmount: -promptAmount,
|
||||
inputTokens: input,
|
||||
writeTokens: write,
|
||||
readTokens: read,
|
||||
inputTokens: -input,
|
||||
writeTokens: -write,
|
||||
readTokens: -read,
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -121,19 +118,23 @@ const spendStructuredTokens = async (txData, tokenUsage) => {
|
|||
});
|
||||
}
|
||||
|
||||
prompt &&
|
||||
completion &&
|
||||
if (prompt || completion) {
|
||||
logger.debug('[spendStructuredTokens] Transaction data record against balance:', {
|
||||
user: txData.user,
|
||||
prompt: prompt.tokenValue,
|
||||
promptRate: prompt.rate,
|
||||
completion: completion.tokenValue,
|
||||
completionRate: completion.rate,
|
||||
balance: completion.balance,
|
||||
prompt: prompt?.prompt,
|
||||
promptRate: prompt?.rate,
|
||||
completion: completion?.completion,
|
||||
completionRate: completion?.rate,
|
||||
balance: completion?.balance ?? prompt?.balance,
|
||||
});
|
||||
} else {
|
||||
logger.debug('[spendStructuredTokens] No transactions incurred against balance');
|
||||
}
|
||||
} catch (err) {
|
||||
logger.error('[spendStructuredTokens]', err);
|
||||
}
|
||||
|
||||
return { prompt, completion };
|
||||
};
|
||||
|
||||
module.exports = { spendTokens, spendStructuredTokens };
|
||||
|
|
|
|||
197
api/models/spendTokens.spec.js
Normal file
197
api/models/spendTokens.spec.js
Normal file
|
|
@ -0,0 +1,197 @@
|
|||
const mongoose = require('mongoose');
|
||||
|
||||
jest.mock('./Transaction', () => ({
|
||||
Transaction: {
|
||||
create: jest.fn(),
|
||||
createStructured: jest.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
jest.mock('./Balance', () => ({
|
||||
findOne: jest.fn(),
|
||||
findOneAndUpdate: jest.fn(),
|
||||
}));
|
||||
|
||||
jest.mock('~/config', () => ({
|
||||
logger: {
|
||||
debug: jest.fn(),
|
||||
error: jest.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
// 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';
|
||||
});
|
||||
|
||||
it('should create transactions for both prompt and completion tokens', async () => {
|
||||
const txData = {
|
||||
user: new mongoose.Types.ObjectId(),
|
||||
conversationId: 'test-convo',
|
||||
model: 'gpt-3.5-turbo',
|
||||
context: 'test',
|
||||
};
|
||||
const tokenUsage = {
|
||||
promptTokens: 100,
|
||||
completionTokens: 50,
|
||||
};
|
||||
|
||||
Transaction.create.mockResolvedValueOnce({ tokenType: 'prompt', rawAmount: -100 });
|
||||
Transaction.create.mockResolvedValueOnce({ tokenType: 'completion', rawAmount: -50 });
|
||||
Balance.findOne.mockResolvedValue({ tokenCredits: 10000 });
|
||||
Balance.findOneAndUpdate.mockResolvedValue({ tokenCredits: 9850 });
|
||||
|
||||
await spendTokens(txData, tokenUsage);
|
||||
|
||||
expect(Transaction.create).toHaveBeenCalledTimes(2);
|
||||
expect(Transaction.create).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
tokenType: 'prompt',
|
||||
rawAmount: -100,
|
||||
}),
|
||||
);
|
||||
expect(Transaction.create).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
tokenType: 'completion',
|
||||
rawAmount: -50,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle zero completion tokens', async () => {
|
||||
const txData = {
|
||||
user: new mongoose.Types.ObjectId(),
|
||||
conversationId: 'test-convo',
|
||||
model: 'gpt-3.5-turbo',
|
||||
context: 'test',
|
||||
};
|
||||
const tokenUsage = {
|
||||
promptTokens: 100,
|
||||
completionTokens: 0,
|
||||
};
|
||||
|
||||
Transaction.create.mockResolvedValueOnce({ tokenType: 'prompt', rawAmount: -100 });
|
||||
Transaction.create.mockResolvedValueOnce({ tokenType: 'completion', rawAmount: -0 });
|
||||
Balance.findOne.mockResolvedValue({ tokenCredits: 10000 });
|
||||
Balance.findOneAndUpdate.mockResolvedValue({ tokenCredits: 9850 });
|
||||
|
||||
await spendTokens(txData, tokenUsage);
|
||||
|
||||
expect(Transaction.create).toHaveBeenCalledTimes(2);
|
||||
expect(Transaction.create).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
tokenType: 'prompt',
|
||||
rawAmount: -100,
|
||||
}),
|
||||
);
|
||||
expect(Transaction.create).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
tokenType: 'completion',
|
||||
rawAmount: -0, // Changed from 0 to -0
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle undefined token counts', async () => {
|
||||
const txData = {
|
||||
user: new mongoose.Types.ObjectId(),
|
||||
conversationId: 'test-convo',
|
||||
model: 'gpt-3.5-turbo',
|
||||
context: 'test',
|
||||
};
|
||||
const tokenUsage = {};
|
||||
|
||||
await spendTokens(txData, tokenUsage);
|
||||
|
||||
expect(Transaction.create).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should not update balance when CHECK_BALANCE is false', async () => {
|
||||
process.env.CHECK_BALANCE = 'false';
|
||||
const txData = {
|
||||
user: new mongoose.Types.ObjectId(),
|
||||
conversationId: 'test-convo',
|
||||
model: 'gpt-3.5-turbo',
|
||||
context: 'test',
|
||||
};
|
||||
const tokenUsage = {
|
||||
promptTokens: 100,
|
||||
completionTokens: 50,
|
||||
};
|
||||
|
||||
Transaction.create.mockResolvedValueOnce({ tokenType: 'prompt', rawAmount: -100 });
|
||||
Transaction.create.mockResolvedValueOnce({ tokenType: 'completion', rawAmount: -50 });
|
||||
|
||||
await spendTokens(txData, tokenUsage);
|
||||
|
||||
expect(Transaction.create).toHaveBeenCalledTimes(2);
|
||||
expect(Balance.findOne).not.toHaveBeenCalled();
|
||||
expect(Balance.findOneAndUpdate).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should create structured transactions for both prompt and completion tokens', async () => {
|
||||
const txData = {
|
||||
user: new mongoose.Types.ObjectId(),
|
||||
conversationId: 'test-convo',
|
||||
model: 'claude-3-5-sonnet',
|
||||
context: 'test',
|
||||
};
|
||||
const tokenUsage = {
|
||||
promptTokens: {
|
||||
input: 10,
|
||||
write: 100,
|
||||
read: 5,
|
||||
},
|
||||
completionTokens: 50,
|
||||
};
|
||||
|
||||
Transaction.createStructured.mockResolvedValueOnce({
|
||||
rate: 3.75,
|
||||
user: txData.user.toString(),
|
||||
balance: 9570,
|
||||
prompt: -430,
|
||||
});
|
||||
Transaction.create.mockResolvedValueOnce({
|
||||
rate: 15,
|
||||
user: txData.user.toString(),
|
||||
balance: 8820,
|
||||
completion: -750,
|
||||
});
|
||||
|
||||
const result = await spendStructuredTokens(txData, tokenUsage);
|
||||
|
||||
expect(Transaction.createStructured).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
tokenType: 'prompt',
|
||||
inputTokens: -10,
|
||||
writeTokens: -100,
|
||||
readTokens: -5,
|
||||
}),
|
||||
);
|
||||
expect(Transaction.create).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
tokenType: 'completion',
|
||||
rawAmount: -50,
|
||||
}),
|
||||
);
|
||||
expect(result).toEqual({
|
||||
prompt: expect.objectContaining({
|
||||
rate: 3.75,
|
||||
user: txData.user.toString(),
|
||||
balance: 9570,
|
||||
prompt: -430,
|
||||
}),
|
||||
completion: expect.objectContaining({
|
||||
rate: 15,
|
||||
user: txData.user.toString(),
|
||||
balance: 8820,
|
||||
completion: -750,
|
||||
}),
|
||||
});
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue