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

api/app/clients/PluginsClient.js

@@ -1,9 +1,11 @@
 const OpenAIClient = require('./OpenAIClient');
 const { CallbackManager } = require('langchain/callbacks');
+const { BufferMemory, ChatMessageHistory } = require('langchain/memory');
 const { initializeCustomAgent, initializeFunctionsAgent } = require('./agents');
 const { addImages, buildErrorInput, buildPromptPrefix } = require('./output_parsers');
-// const { createSummaryBufferMemory } = require('./memory');
+const checkBalance = require('../../models/checkBalance');
 const { formatLangChainMessages } = require('./prompts');
+const { isEnabled } = require('../../server/utils');
 const { SelfReflectionTool } = require('./tools');
 const { loadTools } = require('./tools/util');
 
@@ -73,7 +75,11 @@ class PluginsClient extends OpenAIClient {
       temperature: this.agentOptions.temperature,
     };
 
-    const model = this.initializeLLM(modelOptions);
+    const model = this.initializeLLM({
+      ...modelOptions,
+      context: 'plugins',
+      initialMessageCount: this.currentMessages.length + 1,
+    });
 
     if (this.options.debug) {
       console.debug(
@@ -87,8 +93,11 @@ class PluginsClient extends OpenAIClient {
     });
     this.options.debug && console.debug('pastMessages: ', pastMessages);
 
-    // TODO: implement new token efficient way of processing openAPI plugins so they can "share" memory with agent
-    // const memory = createSummaryBufferMemory({ llm: this.initializeLLM(modelOptions), messages: pastMessages });
+    // TODO: use readOnly memory, TokenBufferMemory? (both unavailable in LangChainJS)
+    const memory = new BufferMemory({
+      llm: model,
+      chatHistory: new ChatMessageHistory(pastMessages),
+    });
 
     this.tools = await loadTools({
       user,
@@ -96,7 +105,8 @@ class PluginsClient extends OpenAIClient {
       tools: this.options.tools,
       functions: this.functionsAgent,
       options: {
-        // memory,
+        memory,
+        signal: this.abortController.signal,
         openAIApiKey: this.openAIApiKey,
         conversationId: this.conversationId,
         debug: this.options?.debug,
@@ -198,16 +208,12 @@ class PluginsClient extends OpenAIClient {
         break; // Exit the loop if the function call is successful
       } catch (err) {
         console.error(err);
-        errorMessage = err.message;
-        let content = '';
-        if (content) {
-          errorMessage = content;
-          break;
-        }
         if (attempts === maxAttempts) {
-          this.result.output = `Encountered an error while attempting to respond. Error: ${err.message}`;
+          const { run } = this.runManager.getRunByConversationId(this.conversationId);
+          const defaultOutput = `Encountered an error while attempting to respond. Error: ${err.message}`;
+          this.result.output = run && run.error ? run.error : defaultOutput;
+          this.result.errorMessage = run && run.error ? run.error : err.message;
           this.result.intermediateSteps = this.actions;
-          this.result.errorMessage = errorMessage;
           break;
         }
       }
@@ -215,11 +221,21 @@ class PluginsClient extends OpenAIClient {
   }
 
   async handleResponseMessage(responseMessage, saveOptions, user) {
-    responseMessage.tokenCount = this.getTokenCountForResponse(responseMessage);
-    responseMessage.completionTokens = responseMessage.tokenCount;
+    const { output, errorMessage, ...result } = this.result;
+    this.options.debug &&
+      console.debug('[handleResponseMessage] Output:', { output, errorMessage, ...result });
+    const { error } = responseMessage;
+    if (!error) {
+      responseMessage.tokenCount = this.getTokenCount(responseMessage.text);
+      responseMessage.completionTokens = responseMessage.tokenCount;
+    }
+
+    if (!this.agentOptions.skipCompletion && !error) {
+      await this.recordTokenUsage(responseMessage);
+    }
     await this.saveMessageToDatabase(responseMessage, saveOptions, user);
     delete responseMessage.tokenCount;
-    return { ...responseMessage, ...this.result };
+    return { ...responseMessage, ...result };
   }
 
   async sendMessage(message, opts = {}) {
@@ -229,9 +245,7 @@ class PluginsClient extends OpenAIClient {
       this.setOptions(opts);
       return super.sendMessage(message, opts);
     }
-    if (this.options.debug) {
-      console.log('Plugins sendMessage', message, opts);
-    }
+    this.options.debug && console.log('Plugins sendMessage', message, opts);
     const {
       user,
       isEdited,
@@ -245,7 +259,6 @@ class PluginsClient extends OpenAIClient {
       onToolEnd,
     } = await this.handleStartMethods(message, opts);
 
-    this.conversationId = conversationId;
     this.currentMessages.push(userMessage);
 
     let {
@@ -275,6 +288,21 @@ class PluginsClient extends OpenAIClient {
       this.currentMessages = payload;
     }
     await this.saveMessageToDatabase(userMessage, saveOptions, user);
+
+    if (isEnabled(process.env.CHECK_BALANCE)) {
+      await checkBalance({
+        req: this.options.req,
+        res: this.options.res,
+        txData: {
+          user: this.user,
+          tokenType: 'prompt',
+          amount: promptTokens,
+          debug: this.options.debug,
+          model: this.modelOptions.model,
+        },
+      });
+    }
+
     const responseMessage = {
       messageId: responseMessageId,
       conversationId,
@@ -311,6 +339,13 @@ class PluginsClient extends OpenAIClient {
       return await this.handleResponseMessage(responseMessage, saveOptions, user);
     }
 
+    // If error occurred during generation (likely token_balance)
+    if (this.result?.errorMessage?.length > 0) {
+      responseMessage.error = true;
+      responseMessage.text = this.result.output;
+      return await this.handleResponseMessage(responseMessage, saveOptions, user);
+    }
+
     if (this.agentOptions.skipCompletion && this.result.output && this.functionsAgent) {
       const partialText = opts.getPartialText();
       const trimmedPartial = opts.getPartialText().replaceAll(':::plugin:::\n', '');