🐞 fix: Balance and Token Usage Improvements (#2350)

* fix(processModelData): handle `openrouter/auto` edge case

* fix(Tx.create): prevent negative multiplier edge case and prevent balance from becoming negative

* fix(NavLinks): render 0 balance properly

* refactor(NavLinks): show only up to 2 decimal places for balance

* fix(OpenAIClient/titleConvo): fix cohere condition and record token usage for `this.options.titleMethod === 'completion'`

api/app/clients/OpenAIClient.js

@@ -792,14 +792,18 @@ ${convo}
         },
       ];
 
+      const promptTokens = this.getTokenCountForMessage(instructionsPayload[0]);
+
       try {
         let useChatCompletion = true;
-        if (CohereConstants.API_URL) {
+        if (this.options.reverseProxyUrl === CohereConstants.API_URL) {
           useChatCompletion = false;
         }
         title = (
           await this.sendPayload(instructionsPayload, { modelOptions, useChatCompletion })
         ).replaceAll('"', '');
+        const completionTokens = this.getTokenCount(title);
+        this.recordTokenUsage({ promptTokens, completionTokens, context: 'title' });
       } catch (e) {
         logger.error(
           '[OpenAIClient] There was an issue generating the title with the completion method',
@@ -951,12 +955,12 @@ ${convo}
     }
   }
 
-  async recordTokenUsage({ promptTokens, completionTokens }) {
+  async recordTokenUsage({ promptTokens, completionTokens, context = 'message' }) {
     await spendTokens(
       {
+        context,
         user: this.user,
         model: this.modelOptions.model,
-        context: 'message',
         conversationId: this.conversationId,
         endpointTokenConfig: this.options.endpointTokenConfig,
       },

api/models/Transaction.js

@@ -12,7 +12,7 @@ transactionSchema.methods.calculateTokenValue = function () {
     this.tokenValue = this.rawAmount;
   }
   const { valueKey, tokenType, model, endpointTokenConfig } = this;
-  const multiplier = getMultiplier({ valueKey, tokenType, model, endpointTokenConfig });
+  const multiplier = Math.abs(getMultiplier({ valueKey, tokenType, model, endpointTokenConfig }));
   this.rate = multiplier;
   this.tokenValue = this.rawAmount * multiplier;
   if (this.context && this.tokenType === 'completion' && this.context === 'incomplete') {
@@ -36,18 +36,24 @@ transactionSchema.statics.create = async function (transactionData) {
     return;
   }
 
-  // Adjust the user's balance
+  let balance = await Balance.findOne({ user: transaction.user }).lean();
-  const updatedBalance = await Balance.findOneAndUpdate(
+  let incrementValue = transaction.tokenValue;
+
+  if (balance && balance?.tokenCredits + incrementValue < 0) {
+    incrementValue = -balance.tokenCredits;
+  }
+
+  balance = await Balance.findOneAndUpdate(
     { user: transaction.user },
-    { $inc: { tokenCredits: transaction.tokenValue } },
+    { $inc: { tokenCredits: incrementValue } },
     { upsert: true, new: true },
   ).lean();
 
   return {
     rate: transaction.rate,
     user: transaction.user.toString(),
-    balance: updatedBalance.tokenCredits,
+    balance: balance.tokenCredits,
-    [transaction.tokenType]: transaction.tokenValue,
+    [transaction.tokenType]: incrementValue,
   };
 };
 

api/utils/tokens.js

@@ -215,6 +215,12 @@ function processModelData(input) {
 
   for (const model of data) {
     const modelKey = model.id;
+    if (modelKey === 'openrouter/auto') {
+      model.pricing = {
+        prompt: '0.00001',
+        completion: '0.00003',
+      };
+    }
     const prompt = parseFloat(model.pricing.prompt) * 1000000;
     const completion = parseFloat(model.pricing.completion) * 1000000;
 

client/src/components/Nav/NavLinks.tsx

@@ -59,9 +59,11 @@ function NavLinks() {
       <Menu as="div" className="group relative">
         {({ open }) => (
           <>
-            {startupConfig?.checkBalance && balanceQuery.data && (
+            {startupConfig?.checkBalance &&
+              balanceQuery.data &&
+              !isNaN(parseFloat(balanceQuery.data)) && (
               <div className="m-1 ml-3 whitespace-nowrap text-left text-sm text-black dark:text-gray-200">
-                {`Balance: ${balanceQuery.data}`}
+                {`Balance: ${parseFloat(balanceQuery.data).toFixed(2)}`}
               </div>
             )}
             <Menu.Button