From c54a57019e39f782b6b8901f3acd3407ad1bf2a4 Mon Sep 17 00:00:00 2001
From: Danny Avila <danny@librechat.ai>
Date: Tue, 15 Oct 2024 19:37:41 -0400
Subject: [PATCH] =?UTF-8?q?=F0=9F=95=92=20feat:=20Add=205-second=20timeout?=
 =?UTF-8?q?=20for=20Fetching=20Model=20Lists=20(#4423)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* refactor: add 5 second timeout for fetching AI provider model lists

* ci: fix test due to recent changes
---
 api/app/clients/OllamaClient.js          | 4 +++-
 api/server/services/ModelService.js      | 7 +++++--
 api/server/services/ModelService.spec.js | 4 +++-
 api/utils/{logAxiosError.js => axios.js} | 2 +-
 api/utils/index.js                       | 8 ++++----
 5 files changed, 16 insertions(+), 9 deletions(-)
 rename api/utils/{logAxiosError.js => axios.js} (97%)

diff --git a/api/app/clients/OllamaClient.js b/api/app/clients/OllamaClient.js
index c88ef72d58..d86e120f43 100644
--- a/api/app/clients/OllamaClient.js
+++ b/api/app/clients/OllamaClient.js
@@ -60,7 +60,9 @@ class OllamaClient {
     try {
       const ollamaEndpoint = deriveBaseURL(baseURL);
       /** @type {Promise<AxiosResponse<OllamaListResponse>>} */
-      const response = await axios.get(`${ollamaEndpoint}/api/tags`);
+      const response = await axios.get(`${ollamaEndpoint}/api/tags`, {
+        timeout: 5000,
+      });
       models = response.data.models.map((tag) => tag.name);
       return models;
     } catch (error) {
diff --git a/api/server/services/ModelService.js b/api/server/services/ModelService.js
index 46fe33c437..0547d03187 100644
--- a/api/server/services/ModelService.js
+++ b/api/server/services/ModelService.js
@@ -1,7 +1,7 @@
 const axios = require('axios');
 const { HttpsProxyAgent } = require('https-proxy-agent');
 const { EModelEndpoint, defaultModels, CacheKeys } = require('librechat-data-provider');
-const { extractBaseURL, inputSchema, processModelData, logAxiosError } = require('~/utils');
+const { inputSchema, logAxiosError, extractBaseURL, processModelData } = require('~/utils');
 const { OllamaClient } = require('~/app/clients/OllamaClient');
 const getLogStores = require('~/cache/getLogStores');
 
@@ -66,6 +66,7 @@ const fetchModels = async ({
       headers: {
         Authorization: `Bearer ${apiKey}`,
       },
+      timeout: 5000,
     };
 
     if (process.env.PROXY) {
@@ -149,6 +150,7 @@ const fetchOpenAIModels = async (opts, _models = []) => {
       baseURL,
       azure: opts.azure,
       user: opts.user,
+      name: baseURL,
     });
   }
 
@@ -175,7 +177,8 @@ const fetchOpenAIModels = async (opts, _models = []) => {
  * @param {object} opts - The options for fetching the models.
  * @param {string} opts.user - The user ID to send to the API.
  * @param {boolean} [opts.azure=false] - Whether to fetch models from Azure.
- * @param {boolean} [opts.plugins=false] - Whether to fetch models from the plugins.
+ * @param {boolean} [opts.plugins=false] - Whether to fetch models for the plugins endpoint.
+ * @param {boolean} [opts.assistants=false] - Whether to fetch models for the Assistants endpoint.
  */
 const getOpenAIModels = async (opts) => {
   let models = defaultModels[EModelEndpoint.openAI];
diff --git a/api/server/services/ModelService.spec.js b/api/server/services/ModelService.spec.js
index 4e4647ee35..a383db1e3c 100644
--- a/api/server/services/ModelService.spec.js
+++ b/api/server/services/ModelService.spec.js
@@ -291,7 +291,9 @@ describe('fetchModels with Ollama specific logic', () => {
     });
 
     expect(models).toEqual(['Ollama-Base', 'Ollama-Advanced']);
-    expect(axios.get).toHaveBeenCalledWith('https://api.ollama.test.com/api/tags'); // Adjusted to expect only one argument if no options are passed
+    expect(axios.get).toHaveBeenCalledWith('https://api.ollama.test.com/api/tags', {
+      timeout: 5000,
+    });
   });
 
   it('should handle errors gracefully when fetching Ollama models fails', async () => {
diff --git a/api/utils/logAxiosError.js b/api/utils/axios.js
similarity index 97%
rename from api/utils/logAxiosError.js
rename to api/utils/axios.js
index 17fac85f47..8b12a5ca99 100644
--- a/api/utils/logAxiosError.js
+++ b/api/utils/axios.js
@@ -42,4 +42,4 @@ const logAxiosError = ({ message, error }) => {
   }
 };
 
-module.exports = logAxiosError;
+module.exports = { logAxiosError };
diff --git a/api/utils/index.js b/api/utils/index.js
index 29357f7adb..62d61586bf 100644
--- a/api/utils/index.js
+++ b/api/utils/index.js
@@ -1,17 +1,17 @@
 const loadYaml = require('./loadYaml');
+const axiosHelpers = require('./axios');
 const tokenHelpers = require('./tokens');
 const azureUtils = require('./azureUtils');
 const deriveBaseURL = require('./deriveBaseURL');
-const logAxiosError = require('./logAxiosError');
 const extractBaseURL = require('./extractBaseURL');
 const findMessageContent = require('./findMessageContent');
 
 module.exports = {
   loadYaml,
-  ...tokenHelpers,
-  ...azureUtils,
   deriveBaseURL,
-  logAxiosError,
   extractBaseURL,
+  ...azureUtils,
+  ...axiosHelpers,
+  ...tokenHelpers,
   findMessageContent,
 };