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

feat: Accurate Token Usage Tracking & Optional Balance (#1018)

Browse source 
* refactor(Chains/llms): allow passing callbacks

* refactor(BaseClient): accurately count completion tokens as generation only

* refactor(OpenAIClient): remove unused getTokenCountForResponse, pass streaming var and callbacks in initializeLLM

* wip: summary prompt tokens

* refactor(summarizeMessages): new cut-off strategy that generates a better summary by adding context from beginning, truncating the middle, and providing the end
wip: draft out relevant providers and variables for token tracing

* refactor(createLLM): make streaming prop false by default

* chore: remove use of getTokenCountForResponse

* refactor(agents): use BufferMemory as ConversationSummaryBufferMemory token usage not easy to trace

* chore: remove passing of streaming prop, also console log useful vars for tracing

* feat: formatFromLangChain helper function to count tokens for ChatModelStart

* refactor(initializeLLM): add role for LLM tracing

* chore(formatFromLangChain): update JSDoc

* feat(formatMessages): formats langChain messages into OpenAI payload format

* chore: install openai-chat-tokens

* refactor(formatMessage): optimize conditional langChain logic
fix(formatFromLangChain): fix destructuring

* feat: accurate prompt tokens for ChatModelStart before generation

* refactor(handleChatModelStart): move to callbacks dir, use factory function

* refactor(initializeLLM): rename 'role' to 'context'

* feat(Balance/Transaction): new schema/models for tracking token spend
refactor(Key): factor out model export to separate file

* refactor(initializeClient): add req,res objects to client options

* feat: add-balance script to add to an existing users' token balance
refactor(Transaction): use multiplier map/function, return balance update

* refactor(Tx): update enum for tokenType, return 1 for multiplier if no map match

* refactor(Tx): add fair fallback value multiplier incase the config result is undefined

* refactor(Balance): rename 'tokens' to 'tokenCredits'

* feat: balance check, add tx.js for new tx-related methods and tests

* chore(summaryPrompts): update prompt token count

* refactor(callbacks): pass req, res
wip: check balance

* refactor(Tx): make convoId a String type, fix(calculateTokenValue)

* refactor(BaseClient): add conversationId as client prop when assigned

* feat(RunManager): track LLM runs with manager, track token spend from LLM,
refactor(OpenAIClient): use RunManager to create callbacks, pass user prop to langchain api calls

* feat(spendTokens): helper to spend prompt/completion tokens

* feat(checkBalance): add helper to check, log, deny request if balance doesn't have enough funds
refactor(Balance): static check method to return object instead of boolean now
wip(OpenAIClient): implement use of checkBalance

* refactor(initializeLLM): add token buffer to assure summary isn't generated when subsequent payload is too large
refactor(OpenAIClient): add checkBalance
refactor(createStartHandler): add checkBalance

* chore: remove prompt and completion token logging from route handler

* chore(spendTokens): add JSDoc

* feat(logTokenCost): record transactions for basic api calls

* chore(ask/edit): invoke getResponseSender only once per API call

* refactor(ask/edit): pass promptTokens to getIds and include in abort data

* refactor(getIds -> getReqData): rename function

* refactor(Tx): increase value if incomplete message

* feat: record tokenUsage when message is aborted

* refactor: subtract tokens when payload includes function_call

* refactor: add namespace for token_balance

* fix(spendTokens): only execute if corresponding token type amounts are defined

* refactor(checkBalance): throws Error if not enough token credits

* refactor(runTitleChain): pass and use signal, spread object props in create helpers, and use 'call' instead of 'run'

* fix(abortMiddleware): circular dependency, and default to empty string for completionTokens

* fix: properly cancel title requests when there isn't enough tokens to generate

* feat(predictNewSummary): custom chain for summaries to allow signal passing
refactor(summaryBuffer): use new custom chain

* feat(RunManager): add getRunByConversationId method, refactor: remove run and throw llm error on handleLLMError

* refactor(createStartHandler): if summary, add error details to runs

* fix(OpenAIClient): support aborting from summarization & showing error to user
refactor(summarizeMessages): remove unnecessary operations counting summaryPromptTokens and note for alternative, pass signal to summaryBuffer

* refactor(logTokenCost -> recordTokenUsage): rename

* refactor(checkBalance): include promptTokens in errorMessage

* refactor(checkBalance/spendTokens): move to models dir

* fix(createLanguageChain): correctly pass config

* refactor(initializeLLM/title): add tokenBuffer of 150 for balance check

* refactor(openAPIPlugin): pass signal and memory, filter functions by the one being called

* refactor(createStartHandler): add error to run if context is plugins as well

* refactor(RunManager/handleLLMError): throw error immediately if plugins, don't remove run

* refactor(PluginsClient): pass memory and signal to tools, cleanup error handling logic

* chore: use absolute equality for addTitle condition

* refactor(checkBalance): move checkBalance to execute after userMessage and tokenCounts are saved, also make conditional

* style: icon changes to match official

* fix(BaseClient): getTokenCountForResponse -> getTokenCount

* fix(formatLangChainMessages): add kwargs as fallback prop from lc_kwargs, update JSDoc

* refactor(Tx.create): does not update balance if CHECK_BALANCE is not enabled

* fix(e2e/cleanUp): cleanup new collections, import all model methods from index

* fix(config/add-balance): add uncaughtException listener

* fix: circular dependency

* refactor(initializeLLM/checkBalance): append new generations to errorMessage if cost exceeds balance

* fix(handleResponseMessage): only record token usage in this method if not error and completion is not skipped

* fix(createStartHandler): correct condition for generations

* chore: bump postcss due to moderate severity vulnerability

* chore: bump zod due to low severity vulnerability

* chore: bump openai & data-provider version

* feat(types): OpenAI Message types

* chore: update bun lockfile

* refactor(CodeBlock): add error block formatting

* refactor(utils/Plugin): factor out formatJSON and cn to separate files (json.ts and cn.ts), add extractJSON

* chore(logViolation): delete user_id after error is logged

* refactor(getMessageError -> Error): change to React.FC, add token_balance handling, use extractJSON to determine JSON instead of regex

* fix(DALL-E): use latest openai SDK

* chore: reorganize imports, fix type issue

* feat(server): add balance route

* fix(api/models): add auth

* feat(data-provider): /api/balance query

* feat: show balance if checking is enabled, refetch on final message or error

* chore: update docs, .env.example with token_usage info, add balance script command

* fix(Balance): fallback to empty obj for balance query

* style: slight adjustment of balance element

* docs(token_usage): add PR notes
This commit is contained in:
Danny Avila 2023-10-05 18:34:10 -04:00 • committed by GitHub
parent be71a1947b
commit 365c39c405
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
81 changed files with 1606 additions and 293 deletions
Download patch file Download diff file Expand all files Collapse all files

31
api/server/routes/ask/anthropic.js
View file

@ -26,18 +26,26 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
console.log('ask log');
console.dir({ text, conversationId, endpointOption }, { depth: null });
let userMessage;
let promptTokens;
let userMessageId;
let responseMessageId;
let lastSavedTimestamp = 0;
let saveDelay = 100;
const sender = getResponseSender(endpointOption);
const user = req.user.id;
const getIds = (data) => {
userMessage = data.userMessage;
userMessageId = data.userMessage.messageId;
responseMessageId = data.responseMessageId;
if (!conversationId) {
conversationId = data.conversationId;
const getReqData = (data = {}) => {
for (let key in data) {
if (key === 'userMessage') {
userMessage = data[key];
userMessageId = data[key].messageId;
} else if (key === 'responseMessageId') {
responseMessageId = data[key];
} else if (key === 'promptTokens') {
promptTokens = data[key];
} else if (!conversationId && key === 'conversationId') {
conversationId = data[key];
}
}
};
@ -49,7 +57,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
lastSavedTimestamp = currentTimestamp;
saveMessage({
messageId: responseMessageId,
sender: getResponseSender(endpointOption),
sender,
conversationId,
parentMessageId: overrideParentMessageId ?? userMessageId,
text: partialText,
@ -69,18 +77,19 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
const getAbortData = () => ({
conversationId,
messageId: responseMessageId,
sender: getResponseSender(endpointOption),
sender,
parentMessageId: overrideParentMessageId ?? userMessageId,
text: getPartialText(),
userMessage,
promptTokens,
});
const { abortController, onStart } = createAbortController(req, res, getAbortData);
const { client } = await initializeClient(req, endpointOption);
const { client } = await initializeClient({ req, res, endpointOption });
let response = await client.sendMessage(text, {
getIds,
getReqData,
// debug: true,
user,
conversationId,
@ -123,7 +132,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
handleAbortError(res, req, error, {
partialText,
conversationId,
sender: getResponseSender(endpointOption),
sender,
messageId: responseMessageId,
parentMessageId: userMessageId ?? parentMessageId,
});

21
api/server/routes/ask/google.js
View file

@ -52,18 +52,25 @@ router.post('/', setHeaders, async (req, res) => {
const ask = async ({ text, endpointOption, parentMessageId = null, conversationId, req, res }) => {
let userMessage;
let userMessageId;
// let promptTokens;
let responseMessageId;
let lastSavedTimestamp = 0;
const { overrideParentMessageId = null } = req.body;
const user = req.user.id;
try {
const getIds = (data) => {
userMessage = data.userMessage;
userMessageId = userMessage.messageId;
responseMessageId = data.responseMessageId;
if (!conversationId) {
conversationId = data.conversationId;
const getReqData = (data = {}) => {
for (let key in data) {
if (key === 'userMessage') {
userMessage = data[key];
userMessageId = data[key].messageId;
} else if (key === 'responseMessageId') {
responseMessageId = data[key];
// } else if (key === 'promptTokens') {
// promptTokens = data[key];
} else if (!conversationId && key === 'conversationId') {
conversationId = data[key];
}
}
sendMessage(res, { message: userMessage, created: true });
@ -121,7 +128,7 @@ const ask = async ({ text, endpointOption, parentMessageId = null, conversationI
const client = new GoogleClient(key, clientOptions);
let response = await client.sendMessage(text, {
getIds,
getReqData,
user,
conversationId,
parentMessageId,

33
api/server/routes/ask/gptPlugins.js
View file

@ -29,22 +29,30 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
console.dir({ text, conversationId, endpointOption }, { depth: null });
let metadata;
let userMessage;
let promptTokens;
let userMessageId;
let responseMessageId;
let lastSavedTimestamp = 0;
let saveDelay = 100;
const sender = getResponseSender(endpointOption);
const newConvo = !conversationId;
const user = req.user.id;
const plugins = [];
const addMetadata = (data) => (metadata = data);
const getIds = (data) => {
userMessage = data.userMessage;
userMessageId = userMessage.messageId;
responseMessageId = data.responseMessageId;
if (!conversationId) {
conversationId = data.conversationId;
const getReqData = (data = {}) => {
for (let key in data) {
if (key === 'userMessage') {
userMessage = data[key];
userMessageId = data[key].messageId;
} else if (key === 'responseMessageId') {
responseMessageId = data[key];
} else if (key === 'promptTokens') {
promptTokens = data[key];
} else if (!conversationId && key === 'conversationId') {
conversationId = data[key];
}
}
};
@ -67,7 +75,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
lastSavedTimestamp = currentTimestamp;
saveMessage({
messageId: responseMessageId,
sender: getResponseSender(endpointOption),
sender,
conversationId,
parentMessageId: overrideParentMessageId || userMessageId,
text: partialText,
@ -135,26 +143,27 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
};
const getAbortData = () => ({
sender: getResponseSender(endpointOption),
sender,
conversationId,
messageId: responseMessageId,
parentMessageId: overrideParentMessageId ?? userMessageId,
text: getPartialText(),
plugins: plugins.map((p) => ({ ...p, loading: false })),
userMessage,
promptTokens,
});
const { abortController, onStart } = createAbortController(req, res, getAbortData);
try {
endpointOption.tools = await validateTools(user, endpointOption.tools);
const { client } = await initializeClient(req, endpointOption);
const { client } = await initializeClient({ req, res, endpointOption });
let response = await client.sendMessage(text, {
user,
conversationId,
parentMessageId,
overrideParentMessageId,
getIds,
getReqData,
onAgentAction,
onChainEnd,
onToolStart,
@ -194,7 +203,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
});
res.end();
if (parentMessageId == '00000000-0000-0000-0000-000000000000' && newConvo) {
if (parentMessageId === '00000000-0000-0000-0000-000000000000' && newConvo) {
addTitle(req, {
text,
response,
@ -206,7 +215,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
handleAbortError(res, req, error, {
partialText,
conversationId,
sender: getResponseSender(endpointOption),
sender,
messageId: responseMessageId,
parentMessageId: userMessageId ?? parentMessageId,
});

38
api/server/routes/ask/openAI.js
View file

@ -27,21 +27,29 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
console.dir({ text, conversationId, endpointOption }, { depth: null });
let metadata;
let userMessage;
let promptTokens;
let userMessageId;
let responseMessageId;
let lastSavedTimestamp = 0;
let saveDelay = 100;
const sender = getResponseSender(endpointOption);
const newConvo = !conversationId;
const user = req.user.id;
const addMetadata = (data) => (metadata = data);
const getIds = (data) => {
userMessage = data.userMessage;
userMessageId = userMessage.messageId;
responseMessageId = data.responseMessageId;
if (!conversationId) {
conversationId = data.conversationId;
const getReqData = (data = {}) => {
for (let key in data) {
if (key === 'userMessage') {
userMessage = data[key];
userMessageId = data[key].messageId;
} else if (key === 'responseMessageId') {
responseMessageId = data[key];
} else if (key === 'promptTokens') {
promptTokens = data[key];
} else if (!conversationId && key === 'conversationId') {
conversationId = data[key];
}
}
};
@ -53,7 +61,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
lastSavedTimestamp = currentTimestamp;
saveMessage({
messageId: responseMessageId,
sender: getResponseSender(endpointOption),
sender,
conversationId,
parentMessageId: overrideParentMessageId ?? userMessageId,
text: partialText,
@ -72,25 +80,26 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
});
const getAbortData = () => ({
sender: getResponseSender(endpointOption),
sender,
conversationId,
messageId: responseMessageId,
parentMessageId: overrideParentMessageId ?? userMessageId,
text: getPartialText(),
userMessage,
promptTokens,
});
const { abortController, onStart } = createAbortController(req, res, getAbortData);
try {
const { client } = await initializeClient(req, endpointOption);
const { client } = await initializeClient({ req, res, endpointOption });
let response = await client.sendMessage(text, {
user,
parentMessageId,
conversationId,
overrideParentMessageId,
getIds,
getReqData,
onStart,
addMetadata,
abortController,
@ -109,11 +118,6 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
response = { ...response, ...metadata };
}
console.log(
'promptTokens, completionTokens:',
response.promptTokens,
response.completionTokens,
);
await saveMessage({ ...response, user });
sendMessage(res, {
@ -125,7 +129,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
});
res.end();
if (parentMessageId == '00000000-0000-0000-0000-000000000000' && newConvo) {
if (parentMessageId === '00000000-0000-0000-0000-000000000000' && newConvo) {
addTitle(req, {
text,
response,
@ -137,7 +141,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
handleAbortError(res, req, error, {
partialText,
conversationId,
sender: getResponseSender(endpointOption),
sender,
messageId: responseMessageId,
parentMessageId: userMessageId ?? parentMessageId,
});

8
api/server/routes/balance.js Normal file
View file

@ -0,0 +1,8 @@
const express = require('express');
const router = express.Router();
const controller = require('../controllers/Balance');
const { requireJwtAuth } = require('../middleware/');
router.get('/', requireJwtAuth, controller);
module.exports = router;

7
api/server/routes/config.js
View file

@ -1,5 +1,6 @@
const express = require('express');
const router = express.Router();
const { isEnabled } = require('../utils');
router.get('/', async function (req, res) {
try {
@ -18,8 +19,9 @@ router.get('/', async function (req, res) {
const discordLoginEnabled =
!!process.env.DISCORD_CLIENT_ID && !!process.env.DISCORD_CLIENT_SECRET;
const serverDomain = process.env.DOMAIN_SERVER || 'http://localhost:3080';
const registrationEnabled = process.env.ALLOW_REGISTRATION?.toLowerCase() === 'true';
const socialLoginEnabled = process.env.ALLOW_SOCIAL_LOGIN?.toLowerCase() === 'true';
const registrationEnabled = isEnabled(process.env.ALLOW_REGISTRATION);
const socialLoginEnabled = isEnabled(process.env.ALLOW_SOCIAL_LOGIN);
const checkBalance = isEnabled(process.env.CHECK_BALANCE);
const emailEnabled =
!!process.env.EMAIL_SERVICE &&
!!process.env.EMAIL_USERNAME &&
@ -39,6 +41,7 @@ router.get('/', async function (req, res) {
registrationEnabled,
socialLoginEnabled,
emailEnabled,
checkBalance,
});
} catch (err) {
console.error(err);

26
api/server/routes/edit/anthropic.js
View file

@ -30,15 +30,24 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
console.dir({ text, generation, isContinued, conversationId, endpointOption }, { depth: null });
let metadata;
let userMessage;
let promptTokens;
let lastSavedTimestamp = 0;
let saveDelay = 100;
const sender = getResponseSender(endpointOption);
const userMessageId = parentMessageId;
const user = req.user.id;
const addMetadata = (data) => (metadata = data);
const getIds = (data) => {
userMessage = data.userMessage;
responseMessageId = data.responseMessageId;
const getReqData = (data = {}) => {
for (let key in data) {
if (key === 'userMessage') {
userMessage = data[key];
} else if (key === 'responseMessageId') {
responseMessageId = data[key];
} else if (key === 'promptTokens') {
promptTokens = data[key];
}
}
};
const { onProgress: progressCallback, getPartialText } = createOnProgress({
@ -49,7 +58,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
lastSavedTimestamp = currentTimestamp;
saveMessage({
messageId: responseMessageId,
sender: getResponseSender(endpointOption),
sender,
conversationId,
parentMessageId: overrideParentMessageId ?? userMessageId,
text: partialText,
@ -70,15 +79,16 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
const getAbortData = () => ({
conversationId,
messageId: responseMessageId,
sender: getResponseSender(endpointOption),
sender,
parentMessageId: overrideParentMessageId ?? userMessageId,
text: getPartialText(),
userMessage,
promptTokens,
});
const { abortController, onStart } = createAbortController(req, res, getAbortData);
const { client } = await initializeClient(req, endpointOption);
const { client } = await initializeClient({ req, res, endpointOption });
let response = await client.sendMessage(text, {
user,
@ -95,7 +105,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
text,
parentMessageId: overrideParentMessageId ?? userMessageId,
}),
getIds,
getReqData,
onStart,
addMetadata,
abortController,
@ -125,7 +135,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
handleAbortError(res, req, error, {
partialText,
conversationId,
sender: getResponseSender(endpointOption),
sender,
messageId: responseMessageId,
parentMessageId: userMessageId ?? parentMessageId,
});

26
api/server/routes/edit/gptPlugins.js
View file

@ -31,8 +31,10 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
console.dir({ text, generation, isContinued, conversationId, endpointOption }, { depth: null });
let metadata;
let userMessage;
let promptTokens;
let lastSavedTimestamp = 0;
let saveDelay = 100;
const sender = getResponseSender(endpointOption);
const userMessageId = parentMessageId;
const user = req.user.id;
@ -44,9 +46,16 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
};
const addMetadata = (data) => (metadata = data);
const getIds = (data) => {
userMessage = data.userMessage;
responseMessageId = data.responseMessageId;
const getReqData = (data = {}) => {
for (let key in data) {
if (key === 'userMessage') {
userMessage = data[key];
} else if (key === 'responseMessageId') {
responseMessageId = data[key];
} else if (key === 'promptTokens') {
promptTokens = data[key];
}
}
};
const {
@ -66,7 +75,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
lastSavedTimestamp = currentTimestamp;
saveMessage({
messageId: responseMessageId,
sender: getResponseSender(endpointOption),
sender,
conversationId,
parentMessageId: overrideParentMessageId || userMessageId,
text: partialText,
@ -106,19 +115,20 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
};
const getAbortData = () => ({
sender: getResponseSender(endpointOption),
sender,
conversationId,
messageId: responseMessageId,
parentMessageId: overrideParentMessageId ?? userMessageId,
text: getPartialText(),
plugin: { ...plugin, loading: false },
userMessage,
promptTokens,
});
const { abortController, onStart } = createAbortController(req, res, getAbortData);
try {
endpointOption.tools = await validateTools(user, endpointOption.tools);
const { client } = await initializeClient(req, endpointOption);
const { client } = await initializeClient({ req, res, endpointOption });
let response = await client.sendMessage(text, {
user,
@ -129,7 +139,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
parentMessageId,
responseMessageId,
overrideParentMessageId,
getIds,
getReqData,
onAgentAction,
onChainEnd,
onStart,
@ -170,7 +180,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
handleAbortError(res, req, error, {
partialText,
conversationId,
sender: getResponseSender(endpointOption),
sender,
messageId: responseMessageId,
parentMessageId: userMessageId ?? parentMessageId,
});

31
api/server/routes/edit/openAI.js
View file

@ -30,15 +30,24 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
console.dir({ text, generation, isContinued, conversationId, endpointOption }, { depth: null });
let metadata;
let userMessage;
let promptTokens;
let lastSavedTimestamp = 0;
let saveDelay = 100;
const sender = getResponseSender(endpointOption);
const userMessageId = parentMessageId;
const user = req.user.id;
const addMetadata = (data) => (metadata = data);
const getIds = (data) => {
userMessage = data.userMessage;
responseMessageId = data.responseMessageId;
const getReqData = (data = {}) => {
for (let key in data) {
if (key === 'userMessage') {
userMessage = data[key];
} else if (key === 'responseMessageId') {
responseMessageId = data[key];
} else if (key === 'promptTokens') {
promptTokens = data[key];
}
}
};
const { onProgress: progressCallback, getPartialText } = createOnProgress({
@ -50,7 +59,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
lastSavedTimestamp = currentTimestamp;
saveMessage({
messageId: responseMessageId,
sender: getResponseSender(endpointOption),
sender,
conversationId,
parentMessageId: overrideParentMessageId || userMessageId,
text: partialText,
@ -70,18 +79,19 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
});
const getAbortData = () => ({
sender: getResponseSender(endpointOption),
sender,
conversationId,
messageId: responseMessageId,
parentMessageId: overrideParentMessageId ?? userMessageId,
text: getPartialText(),
userMessage,
promptTokens,
});
const { abortController, onStart } = createAbortController(req, res, getAbortData);
try {
const { client } = await initializeClient(req, endpointOption);
const { client } = await initializeClient({ req, res, endpointOption });
let response = await client.sendMessage(text, {
user,
@ -92,7 +102,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
parentMessageId,
responseMessageId,
overrideParentMessageId,
getIds,
getReqData,
onStart,
addMetadata,
abortController,
@ -107,11 +117,6 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
response = { ...response, ...metadata };
}
console.log(
'promptTokens, completionTokens:',
response.promptTokens,
response.completionTokens,
);
await saveMessage({ ...response, user });
sendMessage(res, {
@ -127,7 +132,7 @@ router.post('/', validateEndpoint, buildEndpointOption, setHeaders, async (req,
handleAbortError(res, req, error, {
partialText,
conversationId,
sender: getResponseSender(endpointOption),
sender,
messageId: responseMessageId,
parentMessageId: userMessageId ?? parentMessageId,
});

4
api/server/routes/endpoints/anthropic/initializeClient.js
View file

@ -1,7 +1,7 @@
const { AnthropicClient } = require('../../../../app');
const { getUserKey, checkUserKeyExpiry } = require('../../../services/UserService');
const initializeClient = async (req) => {
const initializeClient = async ({ req, res }) => {
const { ANTHROPIC_API_KEY } = process.env;
const { key: expiresAt } = req.body;
@ -16,7 +16,7 @@ const initializeClient = async (req) => {
key = await getUserKey({ userId: req.user.id, name: 'anthropic' });
}
let anthropicApiKey = isUserProvided ? key : ANTHROPIC_API_KEY;
const client = new AnthropicClient(anthropicApiKey);
const client = new AnthropicClient(anthropicApiKey, { req, res });
return {
client,
anthropicApiKey,

4
api/server/routes/endpoints/gptPlugins/initializeClient.js
View file

@ -3,7 +3,7 @@ const { isEnabled } = require('../../../utils');
const { getAzureCredentials } = require('../../../../utils');
const { getUserKey, checkUserKeyExpiry } = require('../../../services/UserService');
const initializeClient = async (req, endpointOption) => {
const initializeClient = async ({ req, res, endpointOption }) => {
const {
PROXY,
OPENAI_API_KEY,
@ -20,6 +20,8 @@ const initializeClient = async (req, endpointOption) => {
debug: isEnabled(DEBUG_PLUGINS),
reverseProxyUrl: OPENAI_REVERSE_PROXY ?? null,
proxy: PROXY ?? null,
req,
res,
...endpointOption,
};

4
api/server/routes/endpoints/openAI/initializeClient.js
View file

@ -3,7 +3,7 @@ const { isEnabled } = require('../../../utils');
const { getAzureCredentials } = require('../../../../utils');
const { getUserKey, checkUserKeyExpiry } = require('../../../services/UserService');
const initializeClient = async (req, endpointOption) => {
const initializeClient = async ({ req, res, endpointOption }) => {
const {
PROXY,
OPENAI_API_KEY,
@ -19,6 +19,8 @@ const initializeClient = async (req, endpointOption) => {
contextStrategy,
reverseProxyUrl: OPENAI_REVERSE_PROXY ?? null,
proxy: PROXY ?? null,
req,
res,
...endpointOption,
};

2
api/server/routes/index.js
View file

@ -10,6 +10,7 @@ const auth = require('./auth');
const keys = require('./keys');
const oauth = require('./oauth');
const endpoints = require('./endpoints');
const balance = require('./balance');
const models = require('./models');
const plugins = require('./plugins');
const user = require('./user');
@ -29,6 +30,7 @@ module.exports = {
user,
tokenizer,
endpoints,
balance,
models,
plugins,
config,

5
api/server/routes/models.js
View file

@ -1,7 +1,8 @@
const express = require('express');
const router = express.Router();
const modelController = require('../controllers/ModelController');
const controller = require('../controllers/ModelController');
const { requireJwtAuth } = require('../middleware/');
router.get('/', modelController);
router.get('/', requireJwtAuth, controller);
module.exports = router;