From dec7879cc13bb8b23e6da5d51160af42a650e7c8 Mon Sep 17 00:00:00 2001
From: Danny Avila <messagedaniel@protonmail.com>
Date: Fri, 29 Mar 2024 11:49:38 -0400
Subject: [PATCH] =?UTF-8?q?=E2=9A=93=20refactor(loadConfigModels):=20Stric?=
 =?UTF-8?q?ter=20Default=20Model=20Fallback=20(#2239)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* chore: add TEndpoint type/typedef

* refactor(loadConfigModels.spec): stricter default model matching (fails with current impl.)

* refactor(loadConfigModels): return default models on endpoint basis and not fetch basis

* refactor: rename `uniqueKeyToNameMap` to `uniqueKeyToEndpointsMap` for clarity
---
 .../services/Config/loadConfigModels.js       | 26 +++++++++++++------
 .../services/Config/loadConfigModels.spec.js  | 18 +++++++------
 api/typedefs.js                               |  6 +++++
 packages/data-provider/src/config.ts          |  2 ++
 4 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/api/server/services/Config/loadConfigModels.js b/api/server/services/Config/loadConfigModels.js
index bfdb80b4ac..b3997a2ada 100644
--- a/api/server/services/Config/loadConfigModels.js
+++ b/api/server/services/Config/loadConfigModels.js
@@ -46,12 +46,23 @@ async function loadConfigModels(req) {
       (endpoint.models.fetch || endpoint.models.default),
   );
 
-  const fetchPromisesMap = {}; // Map for promises keyed by unique combination of baseURL and apiKey
-  const uniqueKeyToNameMap = {}; // Map to associate unique keys with endpoint names
+  /**
+   * @type {Record<string, string[]>}
+   * Map for promises keyed by unique combination of baseURL and apiKey */
+  const fetchPromisesMap = {};
+  /**
+   * @type {Record<string, string[]>}
+   * Map to associate unique keys with endpoint names; note: one key may can correspond to multiple endpoints */
+  const uniqueKeyToEndpointsMap = {};
+  /**
+   * @type {Record<string, Partial<TEndpoint>>}
+   * Map to associate endpoint names to their configurations */
+  const endpointsMap = {};
 
   for (let i = 0; i < customEndpoints.length; i++) {
     const endpoint = customEndpoints[i];
     const { models, name, baseURL, apiKey } = endpoint;
+    endpointsMap[name] = endpoint;
 
     const API_KEY = extractEnvVariable(apiKey);
     const BASE_URL = extractEnvVariable(baseURL);
@@ -69,11 +80,9 @@ async function loadConfigModels(req) {
           apiKey: API_KEY,
           name,
           userIdQuery: models.userIdQuery,
-        }).then((result) => {
-          return !result?.length ? models.default ?? [] : result;
         });
-      uniqueKeyToNameMap[uniqueKey] = uniqueKeyToNameMap[uniqueKey] || [];
-      uniqueKeyToNameMap[uniqueKey].push(name);
+      uniqueKeyToEndpointsMap[uniqueKey] = uniqueKeyToEndpointsMap[uniqueKey] || [];
+      uniqueKeyToEndpointsMap[uniqueKey].push(name);
       continue;
     }
 
@@ -88,10 +97,11 @@ async function loadConfigModels(req) {
   for (let i = 0; i < fetchedData.length; i++) {
     const currentKey = uniqueKeys[i];
     const modelData = fetchedData[i];
-    const associatedNames = uniqueKeyToNameMap[currentKey];
+    const associatedNames = uniqueKeyToEndpointsMap[currentKey];
 
     for (const name of associatedNames) {
-      modelsConfig[name] = modelData;
+      const endpoint = endpointsMap[name];
+      modelsConfig[name] = !modelData?.length ? endpoint.models.default ?? [] : modelData;
     }
   }
 
diff --git a/api/server/services/Config/loadConfigModels.spec.js b/api/server/services/Config/loadConfigModels.spec.js
index 695a59b03f..1b7dec5fd7 100644
--- a/api/server/services/Config/loadConfigModels.spec.js
+++ b/api/server/services/Config/loadConfigModels.spec.js
@@ -267,6 +267,15 @@ describe('loadConfigModels', () => {
     getCustomConfig.mockResolvedValue({
       endpoints: {
         custom: [
+          {
+            name: 'EndpointWithSameFetchKey',
+            apiKey: 'API_KEY',
+            baseURL: 'http://example.com',
+            models: {
+              fetch: true,
+              default: ['defaultModel1'],
+            },
+          },
           {
             name: 'EmptyFetchModel',
             apiKey: 'API_KEY',
@@ -283,14 +292,7 @@ describe('loadConfigModels', () => {
     fetchModels.mockResolvedValue([]);
 
     const result = await loadConfigModels(mockRequest);
-
-    expect(fetchModels).toHaveBeenCalledWith(
-      expect.objectContaining({
-        name: 'EmptyFetchModel',
-        apiKey: 'API_KEY',
-      }),
-    );
-
+    expect(fetchModels).toHaveBeenCalledTimes(1);
     expect(result.EmptyFetchModel).toEqual(['defaultModel1', 'defaultModel2']);
   });
 
diff --git a/api/typedefs.js b/api/typedefs.js
index def59971dc..24870b63ba 100644
--- a/api/typedefs.js
+++ b/api/typedefs.js
@@ -277,6 +277,12 @@
  * @memberof typedefs
  */
 
+/**
+ * @exports TEndpoint
+ * @typedef {import('librechat-data-provider').TEndpoint} TEndpoint
+ * @memberof typedefs
+ */
+
 /**
  * @exports TEndpointsConfig
  * @typedef {import('librechat-data-provider').TEndpointsConfig} TEndpointsConfig
diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts
index e1b9271e9b..56faa115df 100644
--- a/packages/data-provider/src/config.ts
+++ b/packages/data-provider/src/config.ts
@@ -150,6 +150,8 @@ export const endpointSchema = z.object({
   customOrder: z.number().optional(),
 });
 
+export type TEndpoint = z.infer<typeof endpointSchema>;
+
 export const azureEndpointSchema = z
   .object({
     groups: azureGroupConfigsSchema,