From c49f883e1ac38ea9b7d229cd97e089e5f233dd23 Mon Sep 17 00:00:00 2001
From: Danny Avila <danny@librechat.ai>
Date: Mon, 14 Apr 2025 22:38:35 -0400
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=81=20refactor:=20Token=20Event=20Hand?=
 =?UTF-8?q?ler=20and=20Standardize=20`maxTokens`=20Key=20(#6886)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* 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                      | 10 ++++++++++
 api/server/services/Endpoints/bedrock/options.js   |  9 ++-------
 api/server/services/Endpoints/custom/initialize.js |  7 +++----
 api/server/services/Endpoints/openAI/initialize.js | 11 +++--------
 api/server/services/Endpoints/openAI/llm.js        |  6 ++++++
 5 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/api/app/clients/generators.js b/api/app/clients/generators.js
index 971ac16da2..9814cac7a5 100644
--- a/api/app/clients/generators.js
+++ b/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,
 };
diff --git a/api/server/services/Endpoints/bedrock/options.js b/api/server/services/Endpoints/bedrock/options.js
index 6740ae882e..1936a8f483 100644
--- a/api/server/services/Endpoints/bedrock/options.js
+++ b/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),
     },
   ];
 
diff --git a/api/server/services/Endpoints/custom/initialize.js b/api/server/services/Endpoints/custom/initialize.js
index c44b9e0e3e..592440db54 100644
--- a/api/server/services/Endpoints/custom/initialize.js
+++ b/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;
diff --git a/api/server/services/Endpoints/openAI/initialize.js b/api/server/services/Endpoints/openAI/initialize.js
index 1661f9e6c6..714ed5a1e6 100644
--- a/api/server/services/Endpoints/openAI/initialize.js
+++ b/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;
diff --git a/api/server/services/Endpoints/openAI/llm.js b/api/server/services/Endpoints/openAI/llm.js
index a8aeeb5b9d..fad9139a16 100644
--- a/api/server/services/Endpoints/openAI/llm.js
+++ b/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,