🪙 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

api/models/Balance.js

@@ -10,8 +10,9 @@ balanceSchema.statics.check = async function ({
   valueKey,
   tokenType,
   amount,
+  endpointTokenConfig,
 }) {
-  const multiplier = getMultiplier({ valueKey, tokenType, model, endpoint });
+  const multiplier = getMultiplier({ valueKey, tokenType, model, endpoint, endpointTokenConfig });
   const tokenCost = amount * multiplier;
   const { tokenCredits: balance } = (await this.findOne({ user }, 'tokenCredits').lean()) ?? {};
 
@@ -24,6 +25,7 @@ balanceSchema.statics.check = async function ({
     amount,
     balance,
     multiplier,
+    endpointTokenConfig: !!endpointTokenConfig,
   });
 
   if (!balance) {

api/models/Transaction.js

@@ -10,8 +10,8 @@ transactionSchema.methods.calculateTokenValue = function () {
   if (!this.valueKey || !this.tokenType) {
     this.tokenValue = this.rawAmount;
   }
-  const { valueKey, tokenType, model } = this;
-  const multiplier = getMultiplier({ valueKey, tokenType, model });
+  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') {
@@ -25,6 +25,7 @@ transactionSchema.statics.create = async function (transactionData) {
   const Transaction = this;
 
   const transaction = new Transaction(transactionData);
+  transaction.endpointTokenConfig = transactionData.endpointTokenConfig;
   transaction.calculateTokenValue();
 
   // Save the transaction

api/models/checkBalance.js

@@ -14,6 +14,7 @@ const { logViolation } = require('../cache');
  * @param {('prompt' | 'completion')} params.txData.tokenType - The type of token.
  * @param {number} params.txData.amount - The amount of tokens.
  * @param {string} params.txData.model - The model name or identifier.
+ * @param {string} [params.txData.endpointTokenConfig] - The token configuration for the endpoint.
  * @returns {Promise<boolean>} Returns true if the user can spend the amount, otherwise denies the request.
  * @throws {Error} Throws an error if there's an issue with the balance check.
  */

api/models/spendTokens.js

@@ -11,6 +11,7 @@ const { logger } = require('~/config');
  * @param {String} txData.conversationId - The ID of the conversation.
  * @param {String} txData.model - The model name.
  * @param {String} txData.context - The context in which the transaction is made.
+ * @param {String} [txData.endpointTokenConfig] - The current endpoint token config.
  * @param {String} [txData.valueKey] - The value key (optional).
  * @param {Object} tokenUsage - The number of tokens used.
  * @param {Number} tokenUsage.promptTokens - The number of prompt tokens used.

api/models/tx.js

@@ -57,9 +57,14 @@ const getValueKey = (model, endpoint) => {
  * @param {string} [params.tokenType] - The type of token (e.g., 'prompt' or 'completion').
  * @param {string} [params.model] - The model name to derive the value key from if not provided.
  * @param {string} [params.endpoint] - The endpoint name to derive the value key from if not provided.
+ * @param {EndpointTokenConfig} [params.endpointTokenConfig] - The token configuration for the endpoint.
  * @returns {number} The multiplier for the given parameters, or a default value if not found.
  */
-const getMultiplier = ({ valueKey, tokenType, model, endpoint }) => {
+const getMultiplier = ({ valueKey, tokenType, model, endpoint, endpointTokenConfig }) => {
+  if (endpointTokenConfig) {
+    return endpointTokenConfig?.[model]?.[tokenType] ?? defaultRate;
+  }
+
   if (valueKey && tokenType) {
     return tokenValues[valueKey][tokenType] ?? defaultRate;
   }