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);