👤 feat: User ID in Model Query; chore: cleanup ModelService (danny-av…

…ila#1753) * feat: send the LibreChat user ID as a query param when fetching the list of models * chore: update bun * chore: change bun command for building data-provider * refactor: prefer use of `getCustomConfig` to access custom config, also move to `server/services/Config` * refactor: make endpoints/custom option for the config optional, add userIdQuery, and use modelQueries log store in ModelService * refactor(ModelService): use env variables at runtime, use default models from data-provider, and add tests * docs: add `userIdQuery` * fix(ci): import changed
e-gineering · Feb 8, 2024 · ff05715 · ff05715
1 parent d06e5d2
commit ff05715
Show file tree

Hide file tree

Showing 17 changed files with 340 additions and 84 deletions.
diff --git a/api/cache/getLogStores.js b/api/cache/getLogStores.js
@@ -31,6 +31,10 @@ const genTitle = isEnabled(USE_REDIS) // ttl: 2 minutes
   ? new Keyv({ store: keyvRedis, ttl: 120000 })
   : new Keyv({ namespace: CacheKeys.GEN_TITLE, ttl: 120000 });
 
+const modelQueries = isEnabled(process.env.USE_REDIS)
+  ? new Keyv({ store: keyvRedis })
+  : new Keyv({ namespace: 'models' });
+
 const namespaces = {
   [CacheKeys.CONFIG_STORE]: config,
   pending_req,
@@ -44,6 +48,7 @@ const namespaces = {
   logins: createViolationInstance('logins'),
   [CacheKeys.TOKEN_CONFIG]: tokenConfig,
   [CacheKeys.GEN_TITLE]: genTitle,
+  [CacheKeys.MODEL_QUERIES]: modelQueries,
 };
 
 /**

diff --git a/api/server/controllers/ModelController.js b/api/server/controllers/ModelController.js
@@ -9,8 +9,8 @@ async function modelController(req, res) {
     res.send(cachedModelsConfig);
     return;
   }
-  const defaultModelsConfig = await loadDefaultModels();
-  const customModelsConfig = await loadConfigModels();
+  const defaultModelsConfig = await loadDefaultModels(req);
+  const customModelsConfig = await loadConfigModels(req);
 
   const modelConfig = { ...defaultModelsConfig, ...customModelsConfig };
 

diff --git a/api/server/services/AuthService.js b/api/server/services/AuthService.js
@@ -1,7 +1,7 @@
 const crypto = require('crypto');
 const bcrypt = require('bcryptjs');
 const { registerSchema, errorsToString } = require('~/strategies/validators');
-const getCustomConfig = require('~/cache/getCustomConfig');
+const getCustomConfig = require('~/server/services/Config/getCustomConfig');
 const Token = require('~/models/schema/tokenSchema');
 const { sendEmail } = require('~/server/utils');
 const Session = require('~/models/Session');

diff --git a/api/server/services/AuthService.spec.js b/api/server/services/AuthService.spec.js
@@ -1,7 +1,7 @@
-const getCustomConfig = require('~/cache/getCustomConfig');
+const getCustomConfig = require('~/server/services/Config/getCustomConfig');
 const { isDomainAllowed } = require('./AuthService');
 
-jest.mock('~/cache/getCustomConfig', () => jest.fn());
+jest.mock('~/server/services/Config/getCustomConfig', () => jest.fn());
 
 describe('isDomainAllowed', () => {
   it('should allow domain when customConfig is not available', async () => {

diff --git a/api/cache/getCustomConfig.js → ...server/services/Config/getCustomConfig.js b/api/cache/getCustomConfig.js → ...server/services/Config/getCustomConfig.js
@@ -1,6 +1,6 @@
 const { CacheKeys } = require('librechat-data-provider');
-const loadCustomConfig = require('~/server/services/Config/loadCustomConfig');
-const getLogStores = require('./getLogStores');
+const loadCustomConfig = require('./loadCustomConfig');
+const getLogStores = require('~/cache/getLogStores');
 
 /**
  * Retrieves the configuration object

diff --git a/api/server/services/Config/index.js b/api/server/services/Config/index.js
@@ -1,4 +1,5 @@
 const { config } = require('./EndpointService');
+const getCustomConfig = require('./getCustomConfig');
 const loadCustomConfig = require('./loadCustomConfig');
 const loadConfigModels = require('./loadConfigModels');
 const loadDefaultModels = require('./loadDefaultModels');
@@ -9,6 +10,7 @@ const loadDefaultEndpointsConfig = require('./loadDefaultEConfig');
 
 module.exports = {
   config,
+  getCustomConfig,
   loadCustomConfig,
   loadConfigModels,
   loadDefaultModels,

diff --git a/api/server/services/Config/loadConfigEndpoints.js b/api/server/services/Config/loadConfigEndpoints.js
@@ -1,18 +1,12 @@
-const { CacheKeys, EModelEndpoint } = require('librechat-data-provider');
+const { EModelEndpoint } = require('librechat-data-provider');
 const { isUserProvided, extractEnvVariable } = require('~/server/utils');
-const loadCustomConfig = require('./loadCustomConfig');
-const { getLogStores } = require('~/cache');
+const getCustomConfig = require('./getCustomConfig');
 
 /**
  * Load config endpoints from the cached configuration object
  * @function loadConfigEndpoints */
 async function loadConfigEndpoints() {
-  const cache = getLogStores(CacheKeys.CONFIG_STORE);
-  let customConfig = await cache.get(CacheKeys.CUSTOM_CONFIG);
-
-  if (!customConfig) {
-    customConfig = await loadCustomConfig();
-  }
+  const customConfig = await getCustomConfig();
 
   if (!customConfig) {
     return {};

diff --git a/api/server/services/Config/loadConfigModels.js b/api/server/services/Config/loadConfigModels.js
@@ -1,19 +1,15 @@
-const { CacheKeys, EModelEndpoint } = require('librechat-data-provider');
+const { EModelEndpoint } = require('librechat-data-provider');
 const { isUserProvided, extractEnvVariable } = require('~/server/utils');
 const { fetchModels } = require('~/server/services/ModelService');
-const loadCustomConfig = require('./loadCustomConfig');
-const { getLogStores } = require('~/cache');
+const getCustomConfig = require('./getCustomConfig');
 
 /**
  * Load config endpoints from the cached configuration object
- * @function loadConfigModels */
-async function loadConfigModels() {
-  const cache = getLogStores(CacheKeys.CONFIG_STORE);
-  let customConfig = await cache.get(CacheKeys.CUSTOM_CONFIG);
-
-  if (!customConfig) {
-    customConfig = await loadCustomConfig();
-  }
+ * @function loadConfigModels
+ * @param {Express.Request} req - The Express request object.
+ */
+async function loadConfigModels(req) {
+  const customConfig = await getCustomConfig();
 
   if (!customConfig) {
     return {};
@@ -49,7 +45,14 @@ async function loadConfigModels() {
 
     if (models.fetch && !isUserProvided(API_KEY) && !isUserProvided(BASE_URL)) {
       fetchPromisesMap[BASE_URL] =
-        fetchPromisesMap[BASE_URL] || fetchModels({ baseURL: BASE_URL, apiKey: API_KEY, name });
+        fetchPromisesMap[BASE_URL] ||
+        fetchModels({
+          user: req.user.id,
+          baseURL: BASE_URL,
+          apiKey: API_KEY,
+          name,
+          userIdQuery: models.userIdQuery,
+        });
       baseUrlToNameMap[BASE_URL] = baseUrlToNameMap[BASE_URL] || [];
       baseUrlToNameMap[BASE_URL].push(name);
       continue;

diff --git a/api/server/services/Config/loadCustomConfig.js b/api/server/services/Config/loadCustomConfig.js
@@ -17,6 +17,7 @@ const configPath = path.resolve(projectRoot, 'librechat.yaml');
 async function loadCustomConfig() {
   const customConfig = loadYaml(configPath);
   if (!customConfig) {
+    logger.info('Custom config file missing or YAML format invalid.');
     return null;
   }
 
@@ -25,7 +26,7 @@ async function loadCustomConfig() {
     logger.error(`Invalid custom config file at ${configPath}`, result.error);
     return null;
   } else {
-    logger.info('Loaded custom config file:');
+    logger.info('Custom config file loaded:');
     logger.info(JSON.stringify(customConfig, null, 2));
   }
 

diff --git a/api/server/services/Config/loadDefaultModels.js b/api/server/services/Config/loadDefaultModels.js
@@ -11,13 +11,23 @@ const fitlerAssistantModels = (str) => {
   return /gpt-4|gpt-3\\.5/i.test(str) && !/vision|instruct/i.test(str);
 };
 
-async function loadDefaultModels() {
+/**
+ * Loads the default models for the application.
+ * @async
+ * @function
+ * @param {Express.Request} req - The Express request object.
+ */
+async function loadDefaultModels(req) {
   const google = getGoogleModels();
-  const openAI = await getOpenAIModels();
+  const openAI = await getOpenAIModels({ user: req.user.id });
   const anthropic = getAnthropicModels();
   const chatGPTBrowser = getChatGPTBrowserModels();
-  const azureOpenAI = await getOpenAIModels({ azure: true });
-  const gptPlugins = await getOpenAIModels({ azure: useAzurePlugins, plugins: true });
+  const azureOpenAI = await getOpenAIModels({ user: req.user.id, azure: true });
+  const gptPlugins = await getOpenAIModels({
+    user: req.user.id,
+    azure: useAzurePlugins,
+    plugins: true,
+  });
 
   return {
     [EModelEndpoint.openAI]: openAI,

diff --git a/api/server/services/Endpoints/custom/initializeClient.js b/api/server/services/Endpoints/custom/initializeClient.js
@@ -1,8 +1,8 @@
 const { EModelEndpoint, CacheKeys } = require('librechat-data-provider');
 const { getUserKey, checkUserKeyExpiry } = require('~/server/services/UserService');
+const getCustomConfig = require('~/server/services/Config/getCustomConfig');
 const { isUserProvided, extractEnvVariable } = require('~/server/utils');
 const { fetchModels } = require('~/server/services/ModelService');
-const getCustomConfig = require('~/cache/getCustomConfig');
 const getLogStores = require('~/cache/getLogStores');
 const { OpenAIClient } = require('~/app');