🔁 refactor: Token Event Handler and Standardize maxTokens Key (#6886)

* refactor: agent token handling to use createHandleLLMNewToken for improved closure

* refactor: update llmConfig to use maxTokens instead of max_tokens for consistency

api/app/clients/generators.js

@@ -1,6 +1,7 @@
 const fetch = require('node-fetch');
 const { GraphEvents } = require('@librechat/agents');
 const { logger, sendEvent } = require('~/config');
+const { sleep } = require('~/server/utils');
 
 /**
  * Makes a function to make HTTP request and logs the process.
@@ -55,7 +56,16 @@ function createStreamEventHandlers(res) {
   };
 }
 
+function createHandleLLMNewToken(streamRate) {
+  return async () => {
+    if (streamRate) {
+      await sleep(streamRate);
+    }
+  };
+}
+
 module.exports = {
   createFetch,
+  createHandleLLMNewToken,
   createStreamEventHandlers,
 };

api/server/services/Endpoints/bedrock/options.js

@@ -8,7 +8,7 @@ const {
   removeNullishValues,
 } = require('librechat-data-provider');
 const { getUserKey, checkUserKeyExpiry } = require('~/server/services/UserService');
-const { sleep } = require('~/server/utils');
+const { createHandleLLMNewToken } = require('~/app/clients/generators');
 
 const getOptions = async ({ req, overrideModel, endpointOption }) => {
   const {
@@ -90,12 +90,7 @@ const getOptions = async ({ req, overrideModel, endpointOption }) => {
 
   llmConfig.callbacks = [
     {
-      handleLLMNewToken: async () => {
-        if (!streamRate) {
-          return;
-        }
-        await sleep(streamRate);
-      },
+      handleLLMNewToken: createHandleLLMNewToken(streamRate),
     },
   ];
 

api/server/services/Endpoints/custom/initialize.js

@@ -9,9 +9,10 @@ const { Providers } = require('@librechat/agents');
 const { getUserKeyValues, checkUserKeyExpiry } = require('~/server/services/UserService');
 const { getLLMConfig } = require('~/server/services/Endpoints/openAI/llm');
 const { getCustomEndpointConfig } = require('~/server/services/Config');
+const { createHandleLLMNewToken } = require('~/app/clients/generators');
 const { fetchModels } = require('~/server/services/ModelService');
-const { isUserProvided, sleep } = require('~/server/utils');
 const OpenAIClient = require('~/app/clients/OpenAIClient');
+const { isUserProvided } = require('~/server/utils');
 const getLogStores = require('~/cache/getLogStores');
 
 const { PROXY } = process.env;
@@ -148,9 +149,7 @@ const initializeClient = async ({ req, res, endpointOption, optionsOnly, overrid
       }
       options.llmConfig.callbacks = [
         {
-          handleLLMNewToken: async () => {
-            await sleep(customOptions.streamRate);
-          },
+          handleLLMNewToken: createHandleLLMNewToken(clientOptions.streamRate),
         },
       ];
       return options;

api/server/services/Endpoints/openAI/initialize.js

@@ -6,16 +6,11 @@ const {
 } = require('librechat-data-provider');
 const { getUserKeyValues, checkUserKeyExpiry } = require('~/server/services/UserService');
 const { getLLMConfig } = require('~/server/services/Endpoints/openAI/llm');
-const { isEnabled, isUserProvided, sleep } = require('~/server/utils');
+const { createHandleLLMNewToken } = require('~/app/clients/generators');
+const { isEnabled, isUserProvided } = require('~/server/utils');
 const OpenAIClient = require('~/app/clients/OpenAIClient');
 const { getAzureCredentials } = require('~/utils');
 
-function createHandleNewToken(streamRate) {
-  async () => {
-    await sleep(streamRate);
-  };
-}
-
 const initializeClient = async ({
   req,
   res,
@@ -152,7 +147,7 @@ const initializeClient = async ({
     }
     options.llmConfig.callbacks = [
       {
-        handleLLMNewToken: createHandleNewToken(streamRate),
+        handleLLMNewToken: createHandleLLMNewToken(streamRate),
       },
     ];
     return options;

api/server/services/Endpoints/openAI/llm.js

@@ -153,6 +153,12 @@ function getLLMConfig(apiKey, options = {}, endpoint = null) {
     delete llmConfig.reasoning_effort;
   }
 
+  if (llmConfig?.['max_tokens'] != null) {
+    /** @type {number} */
+    llmConfig.maxTokens = llmConfig['max_tokens'];
+    delete llmConfig['max_tokens'];
+  }
+
   return {
     /** @type {OpenAIClientOptions} */
     llmConfig,