diff --git a/api/app/clients/OpenAIClient.js b/api/app/clients/OpenAIClient.js index 9be04105cb..f66afda4ab 100644 --- a/api/app/clients/OpenAIClient.js +++ b/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, }, diff --git a/api/models/Transaction.js b/api/models/Transaction.js index 76ea4ffdef..0d11ab5374 100644 --- a/api/models/Transaction.js +++ b/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 - const updatedBalance = await Balance.findOneAndUpdate( + let balance = await Balance.findOne({ user: transaction.user }).lean(); + 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, - [transaction.tokenType]: transaction.tokenValue, + balance: balance.tokenCredits, + [transaction.tokenType]: incrementValue, }; }; diff --git a/api/utils/tokens.js b/api/utils/tokens.js index c35674a5c5..02e80d20be 100644 --- a/api/utils/tokens.js +++ b/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; diff --git a/client/src/components/Nav/NavLinks.tsx b/client/src/components/Nav/NavLinks.tsx index b7643dbe63..54e70a9370 100644 --- a/client/src/components/Nav/NavLinks.tsx +++ b/client/src/components/Nav/NavLinks.tsx @@ -59,9 +59,11 @@ function NavLinks() { {({ open }) => ( <> - {startupConfig?.checkBalance && balanceQuery.data && ( + {startupConfig?.checkBalance && + balanceQuery.data && + !isNaN(parseFloat(balanceQuery.data)) && (
- {`Balance: ${balanceQuery.data}`} + {`Balance: ${parseFloat(balanceQuery.data).toFixed(2)}`}
)}