From ff59a2e41dbe9c49a247699ae82014f48f4a16ee Mon Sep 17 00:00:00 2001
From: Danny Avila <110412045+danny-avila@users.noreply.github.com>
Date: Fri, 15 Dec 2023 02:40:15 -0500
Subject: [PATCH] =?UTF-8?q?fix:=20Avoid=20Throwing=20Errors=20for=20Unsupp?=
 =?UTF-8?q?orted=20Token=20Count=20Endpoints=20=F0=9F=AA=99=20(#1356)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 api/app/clients/BaseClient.js                   | 5 +++--
 api/app/clients/callbacks/createStartHandler.js | 5 +++--
 packages/data-provider/src/schemas.ts           | 6 ++++++
 3 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js
index 3f05a297fc..1ed41b746c 100644
--- a/api/app/clients/BaseClient.js
+++ b/api/app/clients/BaseClient.js
@@ -1,8 +1,9 @@
 const crypto = require('crypto');
-const TextStream = require('./TextStream');
+const { supportsBalanceCheck } = require('librechat-data-provider');
 const { getConvo, getMessages, saveMessage, updateMessage, saveConvo } = require('~/models');
 const { addSpaceIfNeeded, isEnabled } = require('~/server/utils');
 const checkBalance = require('~/models/checkBalance');
+const TextStream = require('./TextStream');
 const { logger } = require('~/config');
 
 class BaseClient {
@@ -424,7 +425,7 @@ class BaseClient {
       await this.saveMessageToDatabase(userMessage, saveOptions, user);
     }
 
-    if (isEnabled(process.env.CHECK_BALANCE)) {
+    if (isEnabled(process.env.CHECK_BALANCE) && supportsBalanceCheck[this.options.endpoint]) {
       await checkBalance({
         req: this.options.req,
         res: this.options.res,
diff --git a/api/app/clients/callbacks/createStartHandler.js b/api/app/clients/callbacks/createStartHandler.js
index e7dfd15569..4bc32bc0c2 100644
--- a/api/app/clients/callbacks/createStartHandler.js
+++ b/api/app/clients/callbacks/createStartHandler.js
@@ -1,5 +1,5 @@
 const { promptTokensEstimate } = require('openai-chat-tokens');
-const { EModelEndpoint } = require('librechat-data-provider');
+const { EModelEndpoint, supportsBalanceCheck } = require('librechat-data-provider');
 const { formatFromLangChain } = require('~/app/clients/prompts');
 const checkBalance = require('~/models/checkBalance');
 const { isEnabled } = require('~/server/utils');
@@ -49,7 +49,8 @@ const createStartHandler = ({
     prelimPromptTokens += tokenBuffer;
 
     try {
-      if (isEnabled(process.env.CHECK_BALANCE)) {
+      // TODO: if plugins extends to non-OpenAI models, this will need to be updated
+      if (isEnabled(process.env.CHECK_BALANCE) && supportsBalanceCheck[EModelEndpoint.openAI]) {
         const generations =
           initialMessageCount && messages.length > initialMessageCount
             ? messages.slice(initialMessageCount)
diff --git a/packages/data-provider/src/schemas.ts b/packages/data-provider/src/schemas.ts
index 8562e19e65..feac357594 100644
--- a/packages/data-provider/src/schemas.ts
+++ b/packages/data-provider/src/schemas.ts
@@ -138,6 +138,12 @@ export const supportsFiles = {
   [EModelEndpoint.assistant]: true,
 };
 
+export const supportsBalanceCheck = {
+  [EModelEndpoint.openAI]: true,
+  [EModelEndpoint.azureOpenAI]: true,
+  [EModelEndpoint.gptPlugins]: true,
+};
+
 export const visionModels = ['gpt-4-vision', 'llava-13b'];
 
 export const eModelEndpointSchema = z.nativeEnum(EModelEndpoint);