LibreChat/api/app/clients/prompts/formatMessages.spec.js
Danny Avila 365c39c405
feat: Accurate Token Usage Tracking & Optional Balance (#1018)
* 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
2023-10-05 18:34:10 -04:00

247 lines
6.2 KiB
JavaScript

const { formatMessage, formatLangChainMessages, formatFromLangChain } = require('./formatMessages');
const { HumanMessage, AIMessage, SystemMessage } = require('langchain/schema');
describe('formatMessage', () => {
it('formats user message', () => {
const input = {
message: {
sender: 'user',
text: 'Hello',
},
userName: 'John',
};
const result = formatMessage(input);
expect(result).toEqual({
role: 'user',
content: 'Hello',
name: 'John',
});
});
it('formats a realistic user message', () => {
const input = {
message: {
_id: '6512cdfb92cbf69fea615331',
messageId: 'b620bf73-c5c3-4a38-b724-76886aac24c4',
__v: 0,
cancelled: false,
conversationId: '5c23d24f-941f-4aab-85df-127b596c8aa5',
createdAt: Date.now(),
error: false,
finish_reason: null,
isCreatedByUser: true,
isEdited: false,
model: null,
parentMessageId: '00000000-0000-0000-0000-000000000000',
sender: 'User',
text: 'hi',
tokenCount: 5,
unfinished: false,
updatedAt: Date.now(),
user: '6512cdf475f05c86d44c31d2',
},
userName: 'John',
};
const result = formatMessage(input);
expect(result).toEqual({
role: 'user',
content: 'hi',
name: 'John',
});
});
it('formats assistant message', () => {
const input = {
message: {
sender: 'assistant',
text: 'Hi there',
},
assistantName: 'Assistant',
};
const result = formatMessage(input);
expect(result).toEqual({
role: 'assistant',
content: 'Hi there',
name: 'Assistant',
});
});
it('formats system message', () => {
const input = {
message: {
role: 'system',
text: 'Hi there',
},
};
const result = formatMessage(input);
expect(result).toEqual({
role: 'system',
content: 'Hi there',
});
});
it('formats user message with langChain', () => {
const input = {
message: {
sender: 'user',
text: 'Hello',
},
userName: 'John',
langChain: true,
};
const result = formatMessage(input);
expect(result).toBeInstanceOf(HumanMessage);
expect(result.lc_kwargs.content).toEqual(input.message.text);
expect(result.lc_kwargs.name).toEqual(input.userName);
});
it('formats assistant message with langChain', () => {
const input = {
message: {
sender: 'assistant',
text: 'Hi there',
},
assistantName: 'Assistant',
langChain: true,
};
const result = formatMessage(input);
expect(result).toBeInstanceOf(AIMessage);
expect(result.lc_kwargs.content).toEqual(input.message.text);
expect(result.lc_kwargs.name).toEqual(input.assistantName);
});
it('formats system message with langChain', () => {
const input = {
message: {
role: 'system',
text: 'This is a system message.',
},
langChain: true,
};
const result = formatMessage(input);
expect(result).toBeInstanceOf(SystemMessage);
expect(result.lc_kwargs.content).toEqual(input.message.text);
});
it('formats langChain messages into OpenAI payload format', () => {
const human = {
message: new HumanMessage({
content: 'Hello',
}),
};
const system = {
message: new SystemMessage({
content: 'Hello',
}),
};
const ai = {
message: new AIMessage({
content: 'Hello',
}),
};
const humanResult = formatMessage(human);
const systemResult = formatMessage(system);
const aiResult = formatMessage(ai);
expect(humanResult).toEqual({
role: 'user',
content: 'Hello',
});
expect(systemResult).toEqual({
role: 'system',
content: 'Hello',
});
expect(aiResult).toEqual({
role: 'assistant',
content: 'Hello',
});
});
});
describe('formatLangChainMessages', () => {
it('formats an array of messages for LangChain', () => {
const messages = [
{
role: 'system',
content: 'This is a system message',
},
{
sender: 'user',
text: 'Hello',
},
{
sender: 'assistant',
text: 'Hi there',
},
];
const formatOptions = {
userName: 'John',
assistantName: 'Assistant',
};
const result = formatLangChainMessages(messages, formatOptions);
expect(result).toHaveLength(3);
expect(result[0]).toBeInstanceOf(SystemMessage);
expect(result[1]).toBeInstanceOf(HumanMessage);
expect(result[2]).toBeInstanceOf(AIMessage);
expect(result[0].lc_kwargs.content).toEqual(messages[0].content);
expect(result[1].lc_kwargs.content).toEqual(messages[1].text);
expect(result[2].lc_kwargs.content).toEqual(messages[2].text);
expect(result[1].lc_kwargs.name).toEqual(formatOptions.userName);
expect(result[2].lc_kwargs.name).toEqual(formatOptions.assistantName);
});
describe('formatFromLangChain', () => {
it('should merge kwargs and additional_kwargs', () => {
const message = {
kwargs: {
content: 'some content',
name: 'dan',
additional_kwargs: {
function_call: {
name: 'dall-e',
arguments: '{\n "input": "Subject: hedgehog, Style: cute"\n}',
},
},
},
};
const expected = {
content: 'some content',
name: 'dan',
function_call: {
name: 'dall-e',
arguments: '{\n "input": "Subject: hedgehog, Style: cute"\n}',
},
};
expect(formatFromLangChain(message)).toEqual(expected);
});
it('should handle messages without additional_kwargs', () => {
const message = {
kwargs: {
content: 'some content',
name: 'dan',
},
};
const expected = {
content: 'some content',
name: 'dan',
};
expect(formatFromLangChain(message)).toEqual(expected);
});
it('should handle empty messages', () => {
const message = {
kwargs: {},
};
const expected = {};
expect(formatFromLangChain(message)).toEqual(expected);
});
});
});