Andreas/LibreChat
1
0
Fork 0
mirror of https://github.com/danny-avila/LibreChat.git synced 2025-12-17 17:00:15 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

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

Browse source 
* 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
This commit is contained in:
Danny Avila 2024-03-15 19:48:42 -04:00 • committed by GitHub
parent f848d752e0
commit a9d2d3fe40
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 128 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

47
api/server/routes/assistants/chat.js
View file

@ -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,