From 4b09a065908d572c4527a1962c25b74749f4cc55 Mon Sep 17 00:00:00 2001
From: timothycarambat <rambat1010@gmail.com>
Date: Thu, 5 Dec 2024 11:57:07 -0800
Subject: [PATCH] persist token window for NIM and not only on model change

---
 server/utils/AiProviders/nvidiaNim/index.js | 4 ++--
 server/utils/helpers/updateENV.js           | 3 +++
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/server/utils/AiProviders/nvidiaNim/index.js b/server/utils/AiProviders/nvidiaNim/index.js
index 6deb7b2e4a..554b0eec54 100644
--- a/server/utils/AiProviders/nvidiaNim/index.js
+++ b/server/utils/AiProviders/nvidiaNim/index.js
@@ -45,7 +45,7 @@ class NvidiaNimLLM {
   }
 
   /**
-   * Set the model token limit `NVIDIA_NIM_LLM_TOKEN_LIMIT` for the given model ID
+   * Set the model token limit `NVIDIA_NIM_LLM_MODEL_TOKEN_LIMIT` for the given model ID
    * @param {string} modelId
    * @param {string} basePath
    * @returns {Promise<void>}
@@ -69,7 +69,7 @@ class NvidiaNimLLM {
     if (!model.length) return;
     const modelInfo = model.find((model) => model.id === modelId);
     if (!modelInfo) return;
-    process.env.NVIDIA_NIM_LLM_TOKEN_LIMIT = Number(
+    process.env.NVIDIA_NIM_LLM_MODEL_TOKEN_LIMIT = Number(
       modelInfo.max_model_len || 4096
     );
   }
diff --git a/server/utils/helpers/updateENV.js b/server/utils/helpers/updateENV.js
index 3165dc40a3..2ae930a6c3 100644
--- a/server/utils/helpers/updateENV.js
+++ b/server/utils/helpers/updateENV.js
@@ -965,6 +965,9 @@ function dumpENV() {
     "SIMPLE_SSO_ENABLED",
     // Community Hub
     "COMMUNITY_HUB_BUNDLE_DOWNLOADS_ENABLED",
+
+    // Nvidia NIM Keys that are automatically managed
+    "NVIDIA_NIM_LLM_MODEL_TOKEN_LIMIT",
   ];
 
   // Simple sanitization of each value to prevent ENV injection via newline or quote escaping.