LibreChat/api/models/Transaction.js
Danny Avila 30e143e96d
🪙 feat: Use OpenRouter Model Data for Token Cost and Context (#1703)
* feat: use openrouter data for model token cost/context

* chore: add ttl for tokenConfig and refetch models if cache expired
2024-02-02 00:42:11 -05:00

46 lines
1.6 KiB
JavaScript

const mongoose = require('mongoose');
const { isEnabled } = require('../server/utils/handleText');
const transactionSchema = require('./schema/transaction');
const { getMultiplier } = require('./tx');
const Balance = require('./Balance');
const cancelRate = 1.15;
// Method to calculate and set the tokenValue for a transaction
transactionSchema.methods.calculateTokenValue = function () {
if (!this.valueKey || !this.tokenType) {
this.tokenValue = this.rawAmount;
}
const { valueKey, tokenType, model, endpointTokenConfig } = this;
const multiplier = getMultiplier({ valueKey, tokenType, model, endpointTokenConfig });
this.rate = multiplier;
this.tokenValue = this.rawAmount * multiplier;
if (this.context && this.tokenType === 'completion' && this.context === 'incomplete') {
this.tokenValue = Math.ceil(this.tokenValue * cancelRate);
this.rate *= cancelRate;
}
};
// Static method to create a transaction and update the balance
transactionSchema.statics.create = async function (transactionData) {
const Transaction = this;
const transaction = new Transaction(transactionData);
transaction.endpointTokenConfig = transactionData.endpointTokenConfig;
transaction.calculateTokenValue();
// Save the transaction
await transaction.save();
if (!isEnabled(process.env.CHECK_BALANCE)) {
return;
}
// Adjust the user's balance
return await Balance.findOneAndUpdate(
{ user: transaction.user },
{ $inc: { tokenCredits: transaction.tokenValue } },
{ upsert: true, new: true },
).lean();
};
module.exports = mongoose.model('Transaction', transactionSchema);