From e19123a6479916cca4619c19c6ef476ea55ad9c6 Mon Sep 17 00:00:00 2001 From: KALNBERZINS Normunds Date: Wed, 3 Jul 2024 12:45:24 +0200 Subject: [PATCH 1/7] feat: issue-1028, fetch models when user enters api key --- api/server/controllers/ModelController.js | 12 ++----- .../services/Config/loadConfigModels.js | 16 ++++++++-- api/server/services/UserService.js | 32 ++++++++++++++++++- .../src/react-query/react-query-service.ts | 1 + 4 files changed, 49 insertions(+), 12 deletions(-) diff --git a/api/server/controllers/ModelController.js b/api/server/controllers/ModelController.js index 022ece4c103..ebe8032316d 100644 --- a/api/server/controllers/ModelController.js +++ b/api/server/controllers/ModelController.js @@ -18,16 +18,10 @@ const getModelsConfig = async (req) => { * @returns {Promise} The models config. */ async function loadModels(req) { - const cache = getLogStores(CacheKeys.CONFIG_STORE); - const cachedModelsConfig = await cache.get(CacheKeys.MODELS_CONFIG); - if (cachedModelsConfig) { - return cachedModelsConfig; - } - const defaultModelsConfig = await loadDefaultModels(req); - const customModelsConfig = await loadConfigModels(req); - - const modelConfig = { ...defaultModelsConfig, ...customModelsConfig }; + const modelConfig = { ...(await loadDefaultModels(req)), ...(await loadConfigModels(req)) }; + // caching for other services + const cache = getLogStores(CacheKeys.CONFIG_STORE); await cache.set(CacheKeys.MODELS_CONFIG, modelConfig); return modelConfig; } diff --git a/api/server/services/Config/loadConfigModels.js b/api/server/services/Config/loadConfigModels.js index cb0b800d740..30c8bb80d8c 100644 --- a/api/server/services/Config/loadConfigModels.js +++ b/api/server/services/Config/loadConfigModels.js @@ -1,5 +1,6 @@ const { EModelEndpoint, extractEnvVariable } = require('librechat-data-provider'); const { fetchModels } = require('~/server/services/ModelService'); +const { getUserKeyWithExpiry } = require('../UserService'); const { isUserProvided } = require('~/server/utils'); const getCustomConfig = require('./getCustomConfig'); @@ -64,13 +65,24 @@ async function loadConfigModels(req) { const { models, name, baseURL, apiKey } = endpoint; endpointsMap[name] = endpoint; - const API_KEY = extractEnvVariable(apiKey); + let API_KEY = extractEnvVariable(apiKey); const BASE_URL = extractEnvVariable(baseURL); const uniqueKey = `${BASE_URL}__${API_KEY}`; modelsConfig[name] = []; - + /** if key user provided and not expired use it instead of user_defined */ + if (models.fetch && isUserProvided(API_KEY)) { + try { + const userKey = await getUserKeyWithExpiry({ userId: req.user.id, name }); + if (userKey.expiresAt && new Date(userKey.expiresAt).getTime() > Date.now()) { + // in case key is valid replace the default key with the user provided key + API_KEY = userKey.apiKey || API_KEY; + } + } catch (e) { + // ignore if key is missing or invalid + } + } if (models.fetch && !isUserProvided(API_KEY) && !isUserProvided(BASE_URL)) { fetchPromisesMap[uniqueKey] = fetchPromisesMap[uniqueKey] || diff --git a/api/server/services/UserService.js b/api/server/services/UserService.js index 6c736e43666..85ee964380a 100644 --- a/api/server/services/UserService.js +++ b/api/server/services/UserService.js @@ -52,7 +52,36 @@ const getUserKey = async ({ userId, name }) => { } return decrypt(keyValue.value); }; - +/** + * Retrieves and decrypts the key object including expiry date for a given user identified by userId and identifier name. + * @param {Object} params - The parameters object. + * @param {string} params.userId - The unique identifier for the user. + * @param {string} params.name - The name associated with the key. + * @returns {Promise>} The decrypted key object. + * @throws {Error} Throws an error if the key is not found, there is a problem during key retrieval, parsing or decryption + * @description This function searches for a user's key in the database using their userId and name. + * If found, it decrypts the value of the key and returns it as object including expiry date. + * If no key is found, it throws an error indicating that there is no user key available. + */ +const getUserKeyWithExpiry = async ({ userId, name }) => { + const keyValue = await Key.findOne({ userId, name }).lean(); + if (!keyValue) { + throw new Error( + JSON.stringify({ + type: ErrorTypes.NO_USER_KEY, + }), + ); + } + try { + return { ...JSON.parse(decrypt(keyValue.value)), expiresAt: keyValue.expiresAt }; + } catch (e) { + throw new Error( + JSON.stringify({ + type: ErrorTypes.INVALID_USER_KEY, + }), + ); + } +}; /** * Retrieves, decrypts, and parses the key values for a given user identified by userId and name. * @param {Object} params - The parameters object. @@ -169,6 +198,7 @@ module.exports = { deleteUserKey, getUserKeyValues, getUserKeyExpiry, + getUserKeyWithExpiry, checkUserKeyExpiry, updateUserPluginsService, }; diff --git a/packages/data-provider/src/react-query/react-query-service.ts b/packages/data-provider/src/react-query/react-query-service.ts index 50632a3f1f1..ba7b21fee90 100644 --- a/packages/data-provider/src/react-query/react-query-service.ts +++ b/packages/data-provider/src/react-query/react-query-service.ts @@ -134,6 +134,7 @@ export const useUpdateUserKeysMutation = (): UseMutationResult< return useMutation((payload: t.TUpdateUserKeyRequest) => dataService.updateUserKey(payload), { onSuccess: (data, variables) => { queryClient.invalidateQueries([QueryKeys.name, variables.name]); + queryClient.invalidateQueries([QueryKeys.models]); // force models reload after user key update }, }); }; From 3ba470bf5b0454d4cace904c777afd69ddc62137 Mon Sep 17 00:00:00 2001 From: KALNBERZINS Normunds Date: Wed, 3 Jul 2024 13:41:16 +0200 Subject: [PATCH 2/7] fix: user_provided key - allow several custom configurations with the same base url --- api/server/services/Config/loadConfigModels.js | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/api/server/services/Config/loadConfigModels.js b/api/server/services/Config/loadConfigModels.js index 30c8bb80d8c..da103faf5d6 100644 --- a/api/server/services/Config/loadConfigModels.js +++ b/api/server/services/Config/loadConfigModels.js @@ -68,8 +68,6 @@ async function loadConfigModels(req) { let API_KEY = extractEnvVariable(apiKey); const BASE_URL = extractEnvVariable(baseURL); - const uniqueKey = `${BASE_URL}__${API_KEY}`; - modelsConfig[name] = []; /** if key user provided and not expired use it instead of user_defined */ if (models.fetch && isUserProvided(API_KEY)) { @@ -83,6 +81,9 @@ async function loadConfigModels(req) { // ignore if key is missing or invalid } } + + const uniqueKey = `${BASE_URL}__${API_KEY}`; + if (models.fetch && !isUserProvided(API_KEY) && !isUserProvided(BASE_URL)) { fetchPromisesMap[uniqueKey] = fetchPromisesMap[uniqueKey] || From e7ee26ad138e0d1ae5a7e0fd3af349f0a5c5df72 Mon Sep 17 00:00:00 2001 From: KALNBERZINS Normunds Date: Mon, 8 Jul 2024 12:04:14 +0200 Subject: [PATCH 3/7] fix: implemented model caching per endpoint+key --- api/server/controllers/ModelController.js | 15 +------ .../services/Config/loadConfigModels.js | 40 +++++++++++++------ 2 files changed, 30 insertions(+), 25 deletions(-) diff --git a/api/server/controllers/ModelController.js b/api/server/controllers/ModelController.js index ebe8032316d..43e6cb58127 100644 --- a/api/server/controllers/ModelController.js +++ b/api/server/controllers/ModelController.js @@ -1,15 +1,8 @@ -const { CacheKeys } = require('librechat-data-provider'); const { loadDefaultModels, loadConfigModels } = require('~/server/services/Config'); -const { getLogStores } = require('~/cache'); const getModelsConfig = async (req) => { - const cache = getLogStores(CacheKeys.CONFIG_STORE); - let modelsConfig = await cache.get(CacheKeys.MODELS_CONFIG); - if (!modelsConfig) { - modelsConfig = await loadModels(req); - } - - return modelsConfig; + // never caching all models (caching per endpoint+key) + return await loadModels(req); }; /** @@ -19,10 +12,6 @@ const getModelsConfig = async (req) => { */ async function loadModels(req) { const modelConfig = { ...(await loadDefaultModels(req)), ...(await loadConfigModels(req)) }; - - // caching for other services - const cache = getLogStores(CacheKeys.CONFIG_STORE); - await cache.set(CacheKeys.MODELS_CONFIG, modelConfig); return modelConfig; } diff --git a/api/server/services/Config/loadConfigModels.js b/api/server/services/Config/loadConfigModels.js index da103faf5d6..0ea00fc85bf 100644 --- a/api/server/services/Config/loadConfigModels.js +++ b/api/server/services/Config/loadConfigModels.js @@ -1,8 +1,9 @@ -const { EModelEndpoint, extractEnvVariable } = require('librechat-data-provider'); +const { EModelEndpoint, extractEnvVariable, CacheKeys } = require('librechat-data-provider'); const { fetchModels } = require('~/server/services/ModelService'); const { getUserKeyWithExpiry } = require('../UserService'); const { isUserProvided } = require('~/server/utils'); const getCustomConfig = require('./getCustomConfig'); +const getLogStores = require('~/cache/getLogStores'); /** * Load config endpoints from the cached configuration object @@ -76,24 +77,36 @@ async function loadConfigModels(req) { if (userKey.expiresAt && new Date(userKey.expiresAt).getTime() > Date.now()) { // in case key is valid replace the default key with the user provided key API_KEY = userKey.apiKey || API_KEY; + } else { + // if key is expired remove it from the cache + await keyRemoveFromCache(getUniqueKey(BASE_URL, userKey.apiKey)); } } catch (e) { // ignore if key is missing or invalid } } - const uniqueKey = `${BASE_URL}__${API_KEY}`; + const uniqueKey = getUniqueKey(BASE_URL, API_KEY); if (models.fetch && !isUserProvided(API_KEY) && !isUserProvided(BASE_URL)) { - fetchPromisesMap[uniqueKey] = - fetchPromisesMap[uniqueKey] || - fetchModels({ - user: req.user.id, - baseURL: BASE_URL, - apiKey: API_KEY, - name, - userIdQuery: models.userIdQuery, - }); + const modelsCache = getLogStores(CacheKeys.MODEL_QUERIES); + const cachedModels = await modelsCache.get(uniqueKey); + if (cachedModels) { + fetchPromisesMap[uniqueKey] = Promise.resolve(cachedModels); + } else { + fetchPromisesMap[uniqueKey] = + fetchPromisesMap[uniqueKey] || + fetchModels({ + user: req.user.id, + baseURL: BASE_URL, + apiKey: API_KEY, + name, + userIdQuery: models.userIdQuery, + }).then((models) => { + // add models to cache + return modelsCache.set(uniqueKey, models).then(() => models); + }); + } uniqueKeyToEndpointsMap[uniqueKey] = uniqueKeyToEndpointsMap[uniqueKey] || []; uniqueKeyToEndpointsMap[uniqueKey].push(name); continue; @@ -120,5 +133,8 @@ async function loadConfigModels(req) { return modelsConfig; } - +const keyRemoveFromCache = async (key) => { + await getLogStores(CacheKeys.MODEL_QUERIES).delete(key); +}; +const getUniqueKey = (url, key) => `${url}__${key}`; module.exports = loadConfigModels; From ca84534a6446e3a8fe64998f9c11da86b84bcdc8 Mon Sep 17 00:00:00 2001 From: KALNBERZINS Normunds Date: Thu, 12 Sep 2024 18:12:44 +0200 Subject: [PATCH 4/7] fix: align decryption with changes in server/utils --- api/server/services/UserService.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/server/services/UserService.js b/api/server/services/UserService.js index b4592d20eb1..9e6ef7b8746 100644 --- a/api/server/services/UserService.js +++ b/api/server/services/UserService.js @@ -73,7 +73,7 @@ const getUserKeyWithExpiry = async ({ userId, name }) => { ); } try { - return { ...JSON.parse(decrypt(keyValue.value)), expiresAt: keyValue.expiresAt }; + return { ...JSON.parse(await decrypt(keyValue.value)), expiresAt: keyValue.expiresAt }; } catch (e) { throw new Error( JSON.stringify({ From fb2b41f214a7b3d4ef3d89e2c062929a74ae0230 Mon Sep 17 00:00:00 2001 From: KALNBERZINS Normunds Date: Tue, 1 Oct 2024 09:26:13 +0200 Subject: [PATCH 5/7] fix: include fix for cannot change expiry to "never" (#4293) --- api/server/services/UserService.js | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/api/server/services/UserService.js b/api/server/services/UserService.js index 9e6ef7b8746..fb36ac1cb8a 100644 --- a/api/server/services/UserService.js +++ b/api/server/services/UserService.js @@ -144,13 +144,15 @@ const updateUserKey = async ({ userId, name, value, expiresAt = null }) => { name, value: encryptedValue, }; - - // Only add expiresAt to the update object if it's not null + const updateQuery = { $set: updateObject }; + // add expiresAt to the update object if it's not null if (expiresAt) { updateObject.expiresAt = new Date(expiresAt); + } else { + // else unset if already present + updateQuery.$unset = { expiresAt }; } - - return await Key.findOneAndUpdate({ userId, name }, updateObject, { + return await Key.findOneAndUpdate({ userId, name }, updateQuery, { upsert: true, new: true, }).lean(); From a5b956c2b964ee8d54fa0c5adc0afc9d142a115f Mon Sep 17 00:00:00 2001 From: KALNBERZINS Normunds Date: Tue, 1 Oct 2024 10:12:36 +0200 Subject: [PATCH 6/7] fix: accommodate never expiring keys (#3252) --- api/server/services/Config/loadConfigModels.js | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/api/server/services/Config/loadConfigModels.js b/api/server/services/Config/loadConfigModels.js index 0ea00fc85bf..05159d9b0c3 100644 --- a/api/server/services/Config/loadConfigModels.js +++ b/api/server/services/Config/loadConfigModels.js @@ -74,8 +74,8 @@ async function loadConfigModels(req) { if (models.fetch && isUserProvided(API_KEY)) { try { const userKey = await getUserKeyWithExpiry({ userId: req.user.id, name }); - if (userKey.expiresAt && new Date(userKey.expiresAt).getTime() > Date.now()) { - // in case key is valid replace the default key with the user provided key + if (!userKey.expiresAt || new Date(userKey.expiresAt).getTime() > Date.now()) { + // in case the key is not expired (expires never if expiresAt is missing) replace the default key with the user provided key API_KEY = userKey.apiKey || API_KEY; } else { // if key is expired remove it from the cache From 8aecbf23ca20662173fcc18e42e06d246bdd2ac7 Mon Sep 17 00:00:00 2001 From: KALNBERZINS Normunds Date: Mon, 11 Nov 2024 16:29:13 +0100 Subject: [PATCH 7/7] fix: adapt to changed getCustomConfig exports --- api/server/services/Config/loadConfigModels.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/server/services/Config/loadConfigModels.js b/api/server/services/Config/loadConfigModels.js index 1eed3605574..95dd29efc0b 100644 --- a/api/server/services/Config/loadConfigModels.js +++ b/api/server/services/Config/loadConfigModels.js @@ -2,7 +2,7 @@ const { EModelEndpoint, extractEnvVariable, CacheKeys } = require('librechat-dat const { fetchModels } = require('~/server/services/ModelService'); const { getUserKeyWithExpiry } = require('../UserService'); const { isUserProvided, normalizeEndpointName } = require('~/server/utils'); -const getCustomConfig = require('./getCustomConfig'); +const { getCustomConfig } = require('./getCustomConfig'); const getLogStores = require('~/cache/getLogStores'); /**