🪙 feat: Assistants Token Balance & other improvements (#2114)

* chore: add assistants to supportsBalanceCheck

* feat(Transaction): getTransactions and refactor export of model

* refactor: use enum: ViolationTypes.TOKEN_BALANCE

* feat(assistants): check balance

* refactor(assistants): only add promptBuffer if new convo (for title), and remove endpoint definition

* refactor(assistants): Count tokens up to the current context window

* fix(Switcher): make Select list explicitly controlled

* feat(assistants): use assistant's default model when no model is specified instead of the last selected assistant, prevent assistant_id from being recorded in non-assistant endpoints

* chore(assistants/chat): import order

* chore: bump librechat-data-provider due to changes

api/cache/getLogStores.js

@@ -47,7 +47,7 @@ const namespaces = {
   concurrent: createViolationInstance('concurrent'),
   non_browser: createViolationInstance('non_browser'),
   message_limit: createViolationInstance('message_limit'),
-  token_balance: createViolationInstance('token_balance'),
+  token_balance: createViolationInstance(ViolationTypes.TOKEN_BALANCE),
   registrations: createViolationInstance('registrations'),
   [ViolationTypes.FILE_UPLOAD_LIMIT]: createViolationInstance(ViolationTypes.FILE_UPLOAD_LIMIT),
   [ViolationTypes.ILLEGAL_MODEL_REQUEST]: createViolationInstance(

api/models/Transaction.js

@@ -50,4 +50,23 @@ transactionSchema.statics.create = async function (transactionData) {
   };
 };
 
-module.exports = mongoose.model('Transaction', transactionSchema);
+const Transaction = mongoose.model('Transaction', transactionSchema);
+
+/**
+ * Queries and retrieves transactions based on a given filter.
+ * @async
+ * @function getTransactions
+ * @param {Object} filter - MongoDB filter object to apply when querying transactions.
+ * @returns {Promise<Array>} A promise that resolves to an array of matched transactions.
+ * @throws {Error} Throws an error if querying the database fails.
+ */
+async function getTransactions(filter) {
+  try {
+    return await Transaction.find(filter).lean();
+  } catch (error) {
+    console.error('Error querying transactions:', error);
+    throw error;
+  }
+}
+
+module.exports = { Transaction, getTransactions };

api/models/checkBalance.js

@@ -1,5 +1,6 @@
+const { ViolationTypes } = require('librechat-data-provider');
+const { logViolation } = require('~/cache');
 const Balance = require('./Balance');
-const { logViolation } = require('../cache');
 /**
  * Checks the balance for a user and determines if they can spend a certain amount.
  * If the user cannot spend the amount, it logs a violation and denies the request.
@@ -25,7 +26,7 @@ const checkBalance = async ({ req, res, txData }) => {
     return true;
   }
 
-  const type = 'token_balance';
+  const type = ViolationTypes.TOKEN_BALANCE;
   const errorMessage = {
     type,
     balance,

api/models/index.js

@@ -22,14 +22,12 @@ const Key = require('./Key');
 const User = require('./User');
 const Session = require('./Session');
 const Balance = require('./Balance');
-const Transaction = require('./Transaction');
 
 module.exports = {
   User,
   Key,
   Session,
   Balance,
-  Transaction,
 
   hashPassword,
   updateUser,

api/models/spendTokens.js

@@ -1,4 +1,4 @@
-const Transaction = require('./Transaction');
+const { Transaction } = require('./Transaction');
 const { logger } = require('~/config');
 
 /**

api/server/routes/assistants/chat.js

@@ -1,6 +1,12 @@
 const { v4 } = require('uuid');
 const express = require('express');
-const { EModelEndpoint, Constants, RunStatus, CacheKeys } = require('librechat-data-provider');
+const {
+  Constants,
+  RunStatus,
+  CacheKeys,
+  EModelEndpoint,
+  ViolationTypes,
+} = require('librechat-data-provider');
 const {
   initThread,
   recordUsage,
@@ -11,10 +17,13 @@ const {
 } = require('~/server/services/Threads');
 const { runAssistant, createOnTextProgress } = require('~/server/services/AssistantService');
 const { addTitle, initializeClient } = require('~/server/services/Endpoints/assistants');
-const { sendResponse, sendMessage, sleep } = require('~/server/utils');
+const { sendResponse, sendMessage, sleep, isEnabled, countTokens } = require('~/server/utils');
+const { getTransactions } = require('~/models/Transaction');
 const { createRun } = require('~/server/services/Runs');
+const checkBalance = require('~/models/checkBalance');
 const { getConvo } = require('~/models/Conversation');
 const getLogStores = require('~/cache/getLogStores');
+const { getModelMaxTokens } = require('~/utils');
 const { logger } = require('~/config');
 
 const router = express.Router();
@@ -128,6 +137,8 @@ router.post('/', validateModel, buildEndpointOption, setHeaders, async (req, res
           : ''
       }`;
       return sendResponse(res, messageData, errorMessage);
+    } else if (error?.message?.includes(ViolationTypes.TOKEN_BALANCE)) {
+      return sendResponse(res, messageData, error.message);
     } else {
       logger.error('[/assistants/chat/]', error);
     }
@@ -207,6 +218,38 @@ router.post('/', validateModel, buildEndpointOption, setHeaders, async (req, res
       throw new Error('Missing assistant_id');
     }
 
+    if (isEnabled(process.env.CHECK_BALANCE)) {
+      const transactions =
+        (await getTransactions({
+          user: req.user.id,
+          context: 'message',
+          conversationId,
+        })) ?? [];
+
+      const totalPreviousTokens = Math.abs(
+        transactions.reduce((acc, curr) => acc + curr.rawAmount, 0),
+      );
+
+      // TODO: make promptBuffer a config option; buffer for titles, needs buffer for system instructions
+      const promptBuffer = parentMessageId === Constants.NO_PARENT && !_thread_id ? 200 : 0;
+      // 5 is added for labels
+      let promptTokens = (await countTokens(text + (promptPrefix ?? ''))) + 5;
+      promptTokens += totalPreviousTokens + promptBuffer;
+      // Count tokens up to the current context window
+      promptTokens = Math.min(promptTokens, getModelMaxTokens(model));
+
+      await checkBalance({
+        req,
+        res,
+        txData: {
+          model,
+          user: req.user.id,
+          tokenType: 'prompt',
+          amount: promptTokens,
+        },
+      });
+    }
+
     /** @type {{ openai: OpenAIClient }} */
     const { openai: _openai, client } = await initializeClient({
       req,