🧮 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 2024-08-24 04:36:08 -04:00 committed by GitHub
parent ea5140ff0f
commit d54458b3a6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 850 additions and 82 deletions

View file

@ -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') {