Skip to content

Commit d54458b

Browse files
authored
đź§® feat: Improve structured token spending and testing; fix: Anthropic Cache Spend (danny-avila#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
1 parent ea5140f commit d54458b

File tree

8 files changed

+850
-82
lines changed

8 files changed

+850
-82
lines changed

‎api/models/Transaction.js‎

Lines changed: 23 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
const mongoose = require('mongoose');
2-
const { isEnabled } = require('../server/utils/handleText');
2+
const { isEnabled } = require('~/server/utils/handleText');
33
const transactionSchema = require('./schema/transaction');
44
const { getMultiplier, getCacheMultiplier } = require('./tx');
55
const { logger } = require('~/config');
@@ -76,7 +76,7 @@ transactionSchema.statics.createStructured = async function (txData) {
7676
await transaction.save();
7777

7878
if (!isEnabled(process.env.CHECK_BALANCE)) {
79-
return transaction;
79+
return;
8080
}
8181

8282
let balance = await Balance.findOne({ user: transaction.user }).lean();
@@ -122,28 +122,33 @@ transactionSchema.methods.calculateStructuredTokenValue = function () {
122122
read: readMultiplier,
123123
};
124124

125-
const totalTokens = (this.inputTokens || 0) + (this.writeTokens || 0) + (this.readTokens || 0);
125+
const totalPromptTokens =
126+
Math.abs(this.inputTokens || 0) +
127+
Math.abs(this.writeTokens || 0) +
128+
Math.abs(this.readTokens || 0);
126129

127-
if (totalTokens > 0) {
130+
if (totalPromptTokens > 0) {
128131
this.rate =
129-
(inputMultiplier * (this.inputTokens || 0) +
130-
writeMultiplier * (this.writeTokens || 0) +
131-
readMultiplier * (this.readTokens || 0)) /
132-
totalTokens;
132+
(Math.abs(inputMultiplier * (this.inputTokens || 0)) +
133+
Math.abs(writeMultiplier * (this.writeTokens || 0)) +
134+
Math.abs(readMultiplier * (this.readTokens || 0))) /
135+
totalPromptTokens;
133136
} else {
134-
this.rate = inputMultiplier; // Default to input rate if no tokens
137+
this.rate = Math.abs(inputMultiplier); // Default to input rate if no tokens
135138
}
136139

137-
this.tokenValue =
138-
this.inputTokens * inputMultiplier +
139-
(this.writeTokens || 0) * writeMultiplier +
140-
(this.readTokens || 0) * readMultiplier;
141-
} else {
142-
const multiplier = Math.abs(
143-
getMultiplier({ tokenType: this.tokenType, model, endpointTokenConfig }),
140+
this.tokenValue = -(
141+
Math.abs(this.inputTokens || 0) * inputMultiplier +
142+
Math.abs(this.writeTokens || 0) * writeMultiplier +
143+
Math.abs(this.readTokens || 0) * readMultiplier
144144
);
145-
this.rate = multiplier;
146-
this.tokenValue = this.rawAmount * multiplier;
145+
146+
this.rawAmount = -totalPromptTokens;
147+
} else if (this.tokenType === 'completion') {
148+
const multiplier = getMultiplier({ tokenType: this.tokenType, model, endpointTokenConfig });
149+
this.rate = Math.abs(multiplier);
150+
this.tokenValue = -Math.abs(this.rawAmount) * multiplier;
151+
this.rawAmount = -Math.abs(this.rawAmount);
147152
}
148153

149154
if (this.context && this.tokenType === 'completion' && this.context === 'incomplete') {

0 commit comments

Comments
 (0)