diff --git a/api/server/routes/assistants/chat.js b/api/server/routes/assistants/chat.js index 095446988c..8796bbcf30 100644 --- a/api/server/routes/assistants/chat.js +++ b/api/server/routes/assistants/chat.js @@ -97,11 +97,16 @@ router.post('/', validateModel, buildEndpointOption, setHeaders, async (req, res const cache = getLogStores(CacheKeys.ABORT_KEYS); const cacheKey = `${req.user.id}:${conversationId}`; + /** @type {Run | undefined} - The completed run, undefined if incomplete */ + let completedRun; + const handleError = async (error) => { if (error.message === 'Run cancelled') { return res.end(); } - if (error.message === 'Request closed') { + if (error.message === 'Request closed' && completedRun) { + return; + } else if (error.message === 'Request closed') { logger.debug('[/assistants/chat/] Request aborted on close'); } @@ -161,7 +166,9 @@ router.post('/', validateModel, buildEndpointOption, setHeaders, async (req, res try { res.on('close', async () => { - await handleError(new Error('Request closed')); + if (!completedRun) { + await handleError(new Error('Request closed')); + } }); if (convoId && !_thread_id) { @@ -322,6 +329,8 @@ router.post('/', validateModel, buildEndpointOption, setHeaders, async (req, res }); } + completedRun = response.run; + /** @type {ResponseMessage} */ const responseMessage = { ...openai.responseMessage, @@ -367,7 +376,7 @@ router.post('/', validateModel, buildEndpointOption, setHeaders, async (req, res if (!response.run.usage) { await sleep(3000); - const completedRun = await openai.beta.threads.runs.retrieve(thread_id, run.id); + completedRun = await openai.beta.threads.runs.retrieve(thread_id, run.id); if (completedRun.usage) { await recordUsage({ ...completedRun.usage, diff --git a/api/server/services/Config/EndpointService.js b/api/server/services/Config/EndpointService.js index eeab6a4c7d..8bfc2f6695 100644 --- a/api/server/services/Config/EndpointService.js +++ b/api/server/services/Config/EndpointService.js @@ -1,4 +1,5 @@ const { EModelEndpoint } = require('librechat-data-provider'); +const { isUserProvided, generateConfig } = require('~/server/utils'); const { OPENAI_API_KEY: openAIApiKey, @@ -9,17 +10,16 @@ const { BINGAI_TOKEN: bingToken, PLUGINS_USE_AZURE, GOOGLE_KEY: googleKey, + OPENAI_REVERSE_PROXY, + AZURE_OPENAI_BASEURL, + ASSISTANTS_BASE_URL, } = process.env ?? {}; const useAzurePlugins = !!PLUGINS_USE_AZURE; const userProvidedOpenAI = useAzurePlugins - ? azureOpenAIApiKey === 'user_provided' - : openAIApiKey === 'user_provided'; - -function isUserProvided(key) { - return key ? { userProvide: key === 'user_provided' } : false; -} + ? isUserProvided(azureOpenAIApiKey) + : isUserProvided(openAIApiKey); module.exports = { config: { @@ -28,11 +28,11 @@ module.exports = { useAzurePlugins, userProvidedOpenAI, googleKey, - [EModelEndpoint.openAI]: isUserProvided(openAIApiKey), - [EModelEndpoint.assistants]: isUserProvided(assistantsApiKey), - [EModelEndpoint.azureOpenAI]: isUserProvided(azureOpenAIApiKey), - [EModelEndpoint.chatGPTBrowser]: isUserProvided(chatGPTToken), - [EModelEndpoint.anthropic]: isUserProvided(anthropicApiKey), - [EModelEndpoint.bingAI]: isUserProvided(bingToken), + [EModelEndpoint.openAI]: generateConfig(openAIApiKey, OPENAI_REVERSE_PROXY), + [EModelEndpoint.assistants]: generateConfig(assistantsApiKey, ASSISTANTS_BASE_URL), + [EModelEndpoint.azureOpenAI]: generateConfig(azureOpenAIApiKey, AZURE_OPENAI_BASEURL), + [EModelEndpoint.chatGPTBrowser]: generateConfig(chatGPTToken), + [EModelEndpoint.anthropic]: generateConfig(anthropicApiKey), + [EModelEndpoint.bingAI]: generateConfig(bingToken), }, }; diff --git a/api/server/services/Config/loadAsyncEndpoints.js b/api/server/services/Config/loadAsyncEndpoints.js index dcc3f62413..409b9485de 100644 --- a/api/server/services/Config/loadAsyncEndpoints.js +++ b/api/server/services/Config/loadAsyncEndpoints.js @@ -1,8 +1,10 @@ const { EModelEndpoint } = require('librechat-data-provider'); const { addOpenAPISpecs } = require('~/app/clients/tools/util/addOpenAPISpecs'); const { availableTools } = require('~/app/clients/tools'); -const { openAIApiKey, azureOpenAIApiKey, useAzurePlugins, userProvidedOpenAI, googleKey } = - require('./EndpointService').config; +const { isUserProvided } = require('~/server/utils'); +const { config } = require('./EndpointService'); + +const { openAIApiKey, azureOpenAIApiKey, useAzurePlugins, userProvidedOpenAI, googleKey } = config; /** * Load async endpoints and return a configuration object @@ -19,7 +21,7 @@ async function loadAsyncEndpoints(req) { } } - if (googleKey === 'user_provided') { + if (isUserProvided(googleKey)) { googleUserProvides = true; if (i <= 1) { i++; @@ -44,6 +46,10 @@ async function loadAsyncEndpoints(req) { plugins, availableAgents: ['classic', 'functions'], userProvide: useAzure ? false : userProvidedOpenAI, + userProvideURL: useAzure + ? false + : config[EModelEndpoint.openAI]?.userProvideURL || + config[EModelEndpoint.azureOpenAI]?.userProvideURL, azure: useAzurePlugins || useAzure, } : false; diff --git a/api/server/services/Endpoints/assistant/initializeClient.js b/api/server/services/Endpoints/assistant/initializeClient.js index 3ffba002cc..c6013b32a5 100644 --- a/api/server/services/Endpoints/assistant/initializeClient.js +++ b/api/server/services/Endpoints/assistant/initializeClient.js @@ -7,12 +7,42 @@ const { checkUserKeyExpiry, } = require('~/server/services/UserService'); const OpenAIClient = require('~/app/clients/OpenAIClient'); +const { isUserProvided } = require('~/server/utils'); const initializeClient = async ({ req, res, endpointOption, initAppClient = false }) => { const { PROXY, OPENAI_ORGANIZATION, ASSISTANTS_API_KEY, ASSISTANTS_BASE_URL } = process.env; + const userProvidesKey = isUserProvided(ASSISTANTS_API_KEY); + const userProvidesURL = isUserProvided(ASSISTANTS_BASE_URL); + + let userValues = null; + if (userProvidesKey || userProvidesURL) { + const expiresAt = await getUserKeyExpiry({ + userId: req.user.id, + name: EModelEndpoint.assistants, + }); + checkUserKeyExpiry( + expiresAt, + 'Your Assistants API key has expired. Please provide your API key again.', + ); + userValues = await getUserKey({ userId: req.user.id, name: EModelEndpoint.assistants }); + try { + userValues = JSON.parse(userValues); + } catch (e) { + throw new Error( + 'Invalid JSON provided for Assistants API user values. Please provide them again.', + ); + } + } + + let apiKey = userProvidesKey ? userValues.apiKey : ASSISTANTS_API_KEY; + let baseURL = userProvidesURL ? userValues.baseURL : ASSISTANTS_BASE_URL; + + if (!apiKey) { + throw new Error('Assistants API key not provided. Please provide it again.'); + } + const opts = {}; - const baseURL = ASSISTANTS_BASE_URL ?? null; if (baseURL) { opts.baseURL = baseURL; @@ -26,29 +56,6 @@ const initializeClient = async ({ req, res, endpointOption, initAppClient = fals opts.organization = OPENAI_ORGANIZATION; } - const credentials = ASSISTANTS_API_KEY; - - const isUserProvided = credentials === 'user_provided'; - - let userKey = null; - if (isUserProvided) { - const expiresAt = await getUserKeyExpiry({ - userId: req.user.id, - name: EModelEndpoint.assistants, - }); - checkUserKeyExpiry( - expiresAt, - 'Your Assistants API key has expired. Please provide your API key again.', - ); - userKey = await getUserKey({ userId: req.user.id, name: EModelEndpoint.assistants }); - } - - let apiKey = isUserProvided ? userKey : credentials; - - if (!apiKey) { - throw new Error(`${EModelEndpoint.assistants} API key not provided.`); - } - /** @type {OpenAIClient} */ const openai = new OpenAI({ apiKey, diff --git a/api/server/services/Endpoints/assistant/initializeClient.spec.js b/api/server/services/Endpoints/assistant/initializeClient.spec.js new file mode 100644 index 0000000000..05851f97e2 --- /dev/null +++ b/api/server/services/Endpoints/assistant/initializeClient.spec.js @@ -0,0 +1,99 @@ +// const OpenAI = require('openai'); +const { HttpsProxyAgent } = require('https-proxy-agent'); +const { getUserKey, getUserKeyExpiry } = require('~/server/services/UserService'); +const initializeClient = require('./initializeClient'); +// const { OpenAIClient } = require('~/app'); + +jest.mock('~/server/services/UserService', () => ({ + getUserKey: jest.fn(), + getUserKeyExpiry: jest.fn(), + checkUserKeyExpiry: jest.requireActual('~/server/services/UserService').checkUserKeyExpiry, +})); + +const today = new Date(); +const tenDaysFromToday = new Date(today.setDate(today.getDate() + 10)); +const isoString = tenDaysFromToday.toISOString(); + +describe('initializeClient', () => { + // Set up environment variables + const originalEnvironment = process.env; + const app = { + locals: {}, + }; + + beforeEach(() => { + jest.resetModules(); // Clears the cache + process.env = { ...originalEnvironment }; // Make a copy + }); + + afterAll(() => { + process.env = originalEnvironment; // Restore original env vars + }); + + test('initializes OpenAI client with default API key and URL', async () => { + process.env.ASSISTANTS_API_KEY = 'default-api-key'; + process.env.ASSISTANTS_BASE_URL = 'https://default.api.url'; + + // Assuming 'isUserProvided' to return false for this test case + jest.mock('~/server/utils', () => ({ + isUserProvided: jest.fn().mockReturnValueOnce(false), + })); + + const req = { user: { id: 'user123' }, app }; + const res = {}; + + const { openai, openAIApiKey } = await initializeClient({ req, res }); + expect(openai.apiKey).toBe('default-api-key'); + expect(openAIApiKey).toBe('default-api-key'); + expect(openai.baseURL).toBe('https://default.api.url'); + }); + + test('initializes OpenAI client with user-provided API key and URL', async () => { + process.env.ASSISTANTS_API_KEY = 'user_provided'; + process.env.ASSISTANTS_BASE_URL = 'user_provided'; + + getUserKey.mockResolvedValue( + JSON.stringify({ apiKey: 'user-api-key', baseURL: 'https://user.api.url' }), + ); + getUserKeyExpiry.mockResolvedValue(isoString); + + const req = { user: { id: 'user123' } }; + const res = {}; + + const { openai, openAIApiKey } = await initializeClient({ req, res }); + expect(openAIApiKey).toBe('user-api-key'); + expect(openai.apiKey).toBe('user-api-key'); + expect(openai.baseURL).toBe('https://user.api.url'); + }); + + test('throws error for invalid JSON in user-provided values', async () => { + process.env.ASSISTANTS_API_KEY = 'user_provided'; + getUserKey.mockResolvedValue('invalid-json'); + getUserKeyExpiry.mockResolvedValue(isoString); + + const req = { user: { id: 'user123' } }; + const res = {}; + + await expect(initializeClient({ req, res })).rejects.toThrow(/Invalid JSON/); + }); + + test('throws error if API key is not provided', async () => { + delete process.env.ASSISTANTS_API_KEY; // Simulate missing API key + + const req = { user: { id: 'user123' } }; + const res = {}; + + await expect(initializeClient({ req, res })).rejects.toThrow(/Assistants API key not/); + }); + + test('initializes OpenAI client with proxy configuration', async () => { + process.env.ASSISTANTS_API_KEY = 'test-key'; + process.env.PROXY = 'http://proxy.server'; + + const req = { user: { id: 'user123' }, app }; + const res = {}; + + const { openai } = await initializeClient({ req, res }); + expect(openai.httpAgent).toBeInstanceOf(HttpsProxyAgent); + }); +}); diff --git a/api/server/services/Endpoints/custom/initializeClient.js b/api/server/services/Endpoints/custom/initializeClient.js index 9e0b0f666f..18ca8c4455 100644 --- a/api/server/services/Endpoints/custom/initializeClient.js +++ b/api/server/services/Endpoints/custom/initializeClient.js @@ -1,8 +1,9 @@ const { - EModelEndpoint, CacheKeys, - extractEnvVariable, envVarRegex, + EModelEndpoint, + FetchTokenConfig, + extractEnvVariable, } = require('librechat-data-provider'); const { getUserKey, checkUserKeyExpiry } = require('~/server/services/UserService'); const getCustomConfig = require('~/server/services/Config/getCustomConfig'); @@ -42,11 +43,53 @@ const initializeClient = async ({ req, res, endpointOption }) => { throw new Error(`Missing Base URL for ${endpoint}.`); } + const userProvidesKey = isUserProvided(CUSTOM_API_KEY); + const userProvidesURL = isUserProvided(CUSTOM_BASE_URL); + + let userValues = null; + if (expiresAt && (userProvidesKey || userProvidesURL)) { + checkUserKeyExpiry( + expiresAt, + `Your API values for ${endpoint} have expired. Please configure them again.`, + ); + userValues = await getUserKey({ userId: req.user.id, name: endpoint }); + try { + userValues = JSON.parse(userValues); + } catch (e) { + throw new Error(`Invalid JSON provided for ${endpoint} user values.`); + } + } + + let apiKey = userProvidesKey ? userValues.apiKey : CUSTOM_API_KEY; + let baseURL = userProvidesURL ? userValues.baseURL : CUSTOM_BASE_URL; + + if (!apiKey) { + throw new Error(`${endpoint} API key not provided.`); + } + + if (!baseURL) { + throw new Error(`${endpoint} Base URL not provided.`); + } + const cache = getLogStores(CacheKeys.TOKEN_CONFIG); - let endpointTokenConfig = await cache.get(endpoint); - if (endpointConfig && endpointConfig.models.fetch && !endpointTokenConfig) { - await fetchModels({ apiKey: CUSTOM_API_KEY, baseURL: CUSTOM_BASE_URL, name: endpoint }); - endpointTokenConfig = await cache.get(endpoint); + const tokenKey = + !endpointConfig.tokenConfig && (userProvidesKey || userProvidesURL) + ? `${endpoint}:${req.user.id}` + : endpoint; + + let endpointTokenConfig = + !endpointConfig.tokenConfig && + FetchTokenConfig[endpoint.toLowerCase()] && + (await cache.get(tokenKey)); + + if ( + FetchTokenConfig[endpoint.toLowerCase()] && + endpointConfig && + endpointConfig.models.fetch && + !endpointTokenConfig + ) { + await fetchModels({ apiKey, baseURL, name: endpoint, user: req.user.id, tokenKey }); + endpointTokenConfig = await cache.get(tokenKey); } const customOptions = { @@ -63,34 +106,6 @@ const initializeClient = async ({ req, res, endpointOption }) => { endpointTokenConfig, }; - const useUserKey = isUserProvided(CUSTOM_API_KEY); - const useUserURL = isUserProvided(CUSTOM_BASE_URL); - - let userValues = null; - if (expiresAt && (useUserKey || useUserURL)) { - checkUserKeyExpiry( - expiresAt, - `Your API values for ${endpoint} have expired. Please configure them again.`, - ); - userValues = await getUserKey({ userId: req.user.id, name: endpoint }); - try { - userValues = JSON.parse(userValues); - } catch (e) { - throw new Error(`Invalid JSON provided for ${endpoint} user values.`); - } - } - - let apiKey = useUserKey ? userValues.apiKey : CUSTOM_API_KEY; - let baseURL = useUserURL ? userValues.baseURL : CUSTOM_BASE_URL; - - if (!apiKey) { - throw new Error(`${endpoint} API key not provided.`); - } - - if (!baseURL) { - throw new Error(`${endpoint} Base URL not provided.`); - } - const clientOptions = { reverseProxyUrl: baseURL ?? null, proxy: PROXY ?? null, diff --git a/api/server/services/Endpoints/gptPlugins/initializeClient.js b/api/server/services/Endpoints/gptPlugins/initializeClient.js index 5dd5d1e181..a596c7820b 100644 --- a/api/server/services/Endpoints/gptPlugins/initializeClient.js +++ b/api/server/services/Endpoints/gptPlugins/initializeClient.js @@ -4,8 +4,8 @@ const { resolveHeaders, } = require('librechat-data-provider'); const { getUserKey, checkUserKeyExpiry } = require('~/server/services/UserService'); +const { isEnabled, isUserProvided } = require('~/server/utils'); const { getAzureCredentials } = require('~/utils'); -const { isEnabled } = require('~/server/utils'); const { PluginsClient } = require('~/app'); const initializeClient = async ({ req, res, endpointOption }) => { @@ -34,43 +34,48 @@ const initializeClient = async ({ req, res, endpointOption }) => { endpoint = EModelEndpoint.azureOpenAI; } + const credentials = { + [EModelEndpoint.openAI]: OPENAI_API_KEY, + [EModelEndpoint.azureOpenAI]: AZURE_API_KEY, + }; + const baseURLOptions = { [EModelEndpoint.openAI]: OPENAI_REVERSE_PROXY, [EModelEndpoint.azureOpenAI]: AZURE_OPENAI_BASEURL, }; - const reverseProxyUrl = baseURLOptions[endpoint] ?? null; + const userProvidesKey = isUserProvided(credentials[endpoint]); + const userProvidesURL = isUserProvided(baseURLOptions[endpoint]); + + let userValues = null; + if (expiresAt && (userProvidesKey || userProvidesURL)) { + checkUserKeyExpiry( + expiresAt, + 'Your OpenAI API values have expired. Please provide them again.', + ); + userValues = await getUserKey({ userId: req.user.id, name: endpoint }); + try { + userValues = JSON.parse(userValues); + } catch (e) { + throw new Error( + `Invalid JSON provided for ${endpoint} user values. Please provide them again.`, + ); + } + } + + let apiKey = userProvidesKey ? userValues.apiKey : credentials[endpoint]; + let baseURL = userProvidesURL ? userValues.baseURL : baseURLOptions[endpoint]; const clientOptions = { contextStrategy, debug: isEnabled(DEBUG_PLUGINS), - reverseProxyUrl, + reverseProxyUrl: baseURL ? baseURL : null, proxy: PROXY ?? null, req, res, ...endpointOption, }; - const credentials = { - [EModelEndpoint.openAI]: OPENAI_API_KEY, - [EModelEndpoint.azureOpenAI]: AZURE_API_KEY, - }; - - const isUserProvided = credentials[endpoint] === 'user_provided'; - - let userKey = null; - if (expiresAt && isUserProvided) { - checkUserKeyExpiry( - expiresAt, - 'Your OpenAI API key has expired. Please provide your API key again.', - ); - userKey = await getUserKey({ - userId: req.user.id, - name: endpoint, - }); - } - - let apiKey = isUserProvided ? userKey : credentials[endpoint]; if (useAzure && azureConfig) { const { modelGroupMap, groupMap } = azureConfig; const { @@ -99,12 +104,12 @@ const initializeClient = async ({ req, res, endpointOption }) => { apiKey = azureOptions.azureOpenAIApiKey; clientOptions.azure = !serverless && azureOptions; } else if (useAzure || (apiKey && apiKey.includes('{"azure') && !clientOptions.azure)) { - clientOptions.azure = isUserProvided ? JSON.parse(userKey) : getAzureCredentials(); + clientOptions.azure = userProvidesKey ? JSON.parse(userValues.apiKey) : getAzureCredentials(); apiKey = clientOptions.azure.azureOpenAIApiKey; } if (!apiKey) { - throw new Error(`${endpoint} API key not provided.`); + throw new Error(`${endpoint} API key not provided. Please provide it again.`); } const client = new PluginsClient(apiKey, clientOptions); diff --git a/api/server/services/Endpoints/gptPlugins/initializeClient.spec.js b/api/server/services/Endpoints/gptPlugins/initializeClient.spec.js index f4539462f3..01718c31d2 100644 --- a/api/server/services/Endpoints/gptPlugins/initializeClient.spec.js +++ b/api/server/services/Endpoints/gptPlugins/initializeClient.spec.js @@ -1,5 +1,5 @@ // gptPlugins/initializeClient.spec.js -const { EModelEndpoint } = require('librechat-data-provider'); +const { EModelEndpoint, validateAzureGroups } = require('librechat-data-provider'); const { getUserKey } = require('~/server/services/UserService'); const initializeClient = require('./initializeClient'); const { PluginsClient } = require('~/app'); @@ -17,6 +17,69 @@ describe('gptPlugins/initializeClient', () => { locals: {}, }; + const validAzureConfigs = [ + { + group: 'librechat-westus', + apiKey: 'WESTUS_API_KEY', + instanceName: 'librechat-westus', + version: '2023-12-01-preview', + models: { + 'gpt-4-vision-preview': { + deploymentName: 'gpt-4-vision-preview', + version: '2024-02-15-preview', + }, + 'gpt-3.5-turbo': { + deploymentName: 'gpt-35-turbo', + }, + 'gpt-3.5-turbo-1106': { + deploymentName: 'gpt-35-turbo-1106', + }, + 'gpt-4': { + deploymentName: 'gpt-4', + }, + 'gpt-4-1106-preview': { + deploymentName: 'gpt-4-1106-preview', + }, + }, + }, + { + group: 'librechat-eastus', + apiKey: 'EASTUS_API_KEY', + instanceName: 'librechat-eastus', + deploymentName: 'gpt-4-turbo', + version: '2024-02-15-preview', + models: { + 'gpt-4-turbo': true, + }, + baseURL: 'https://eastus.example.com', + additionalHeaders: { + 'x-api-key': 'x-api-key-value', + }, + }, + { + group: 'mistral-inference', + apiKey: 'AZURE_MISTRAL_API_KEY', + baseURL: + 'https://Mistral-large-vnpet-serverless.region.inference.ai.azure.com/v1/chat/completions', + serverless: true, + models: { + 'mistral-large': true, + }, + }, + { + group: 'llama-70b-chat', + apiKey: 'AZURE_LLAMA2_70B_API_KEY', + baseURL: + 'https://Llama-2-70b-chat-qmvyb-serverless.region.inference.ai.azure.com/v1/chat/completions', + serverless: true, + models: { + 'llama-70b-chat': true, + }, + }, + ]; + + const { modelNames, modelGroupMap, groupMap } = validateAzureGroups(validAzureConfigs); + beforeEach(() => { jest.resetModules(); // Clears the cache process.env = { ...originalEnvironment }; // Make a copy @@ -142,7 +205,7 @@ describe('gptPlugins/initializeClient', () => { const res = {}; const endpointOption = { modelOptions: { model: 'default-model' } }; - getUserKey.mockResolvedValue('test-user-provided-openai-api-key'); + getUserKey.mockResolvedValue(JSON.stringify({ apiKey: 'test-user-provided-openai-api-key' })); const { openAIApiKey } = await initializeClient({ req, res, endpointOption }); @@ -164,8 +227,10 @@ describe('gptPlugins/initializeClient', () => { getUserKey.mockResolvedValue( JSON.stringify({ - azureOpenAIApiKey: 'test-user-provided-azure-api-key', - azureOpenAIApiDeploymentName: 'test-deployment', + apiKey: JSON.stringify({ + azureOpenAIApiKey: 'test-user-provided-azure-api-key', + azureOpenAIApiDeploymentName: 'test-deployment', + }), }), ); @@ -186,9 +251,7 @@ describe('gptPlugins/initializeClient', () => { const res = {}; const endpointOption = { modelOptions: { model: 'default-model' } }; - await expect(initializeClient({ req, res, endpointOption })).rejects.toThrow( - /Your OpenAI API key has expired/, - ); + await expect(initializeClient({ req, res, endpointOption })).rejects.toThrow(/Your OpenAI API/); }); test('should throw an error if the user-provided Azure key is invalid JSON', async () => { @@ -207,7 +270,7 @@ describe('gptPlugins/initializeClient', () => { getUserKey.mockResolvedValue('invalid-json'); await expect(initializeClient({ req, res, endpointOption })).rejects.toThrow( - /Unexpected token/, + /Invalid JSON provided/, ); }); @@ -229,4 +292,92 @@ describe('gptPlugins/initializeClient', () => { expect(client.options.reverseProxyUrl).toBe('http://reverse.proxy'); expect(client.options.proxy).toBe('http://proxy'); }); + + test('should throw an error when user-provided values are not valid JSON', async () => { + process.env.OPENAI_API_KEY = 'user_provided'; + const req = { + body: { key: new Date(Date.now() + 10000).toISOString(), endpoint: 'openAI' }, + user: { id: '123' }, + app, + }; + const res = {}; + const endpointOption = {}; + + // Mock getUserKey to return a non-JSON string + getUserKey.mockResolvedValue('not-a-json'); + + await expect(initializeClient({ req, res, endpointOption })).rejects.toThrow( + /Invalid JSON provided for openAI user values/, + ); + }); + + test('should initialize client correctly for Azure OpenAI with valid configuration', async () => { + const req = { + body: { + key: null, + endpoint: EModelEndpoint.gptPlugins, + model: modelNames[0], + }, + user: { id: '123' }, + app: { + locals: { + [EModelEndpoint.azureOpenAI]: { + plugins: true, + modelNames, + modelGroupMap, + groupMap, + }, + }, + }, + }; + const res = {}; + const endpointOption = {}; + + const client = await initializeClient({ req, res, endpointOption }); + expect(client.client.options.azure).toBeDefined(); + }); + + test('should initialize client with default options when certain env vars are not set', async () => { + delete process.env.DEBUG_OPENAI; + delete process.env.OPENAI_SUMMARIZE; + + const req = { + body: { key: null, endpoint: 'openAI' }, + user: { id: '123' }, + app, + }; + const res = {}; + const endpointOption = {}; + + const client = await initializeClient({ req, res, endpointOption }); + + expect(client.client.options.debug).toBe(false); + expect(client.client.options.contextStrategy).toBe(null); + }); + + test('should correctly use user-provided apiKey and baseURL when provided', async () => { + process.env.OPENAI_API_KEY = 'user_provided'; + process.env.OPENAI_REVERSE_PROXY = 'user_provided'; + const req = { + body: { + key: new Date(Date.now() + 10000).toISOString(), + endpoint: 'openAI', + }, + user: { + id: '123', + }, + app, + }; + const res = {}; + const endpointOption = {}; + + getUserKey.mockResolvedValue( + JSON.stringify({ apiKey: 'test', baseURL: 'https://user-provided-url.com' }), + ); + + const result = await initializeClient({ req, res, endpointOption }); + + expect(result.openAIApiKey).toBe('test'); + expect(result.client.options.reverseProxyUrl).toBe('https://user-provided-url.com'); + }); }); diff --git a/api/server/services/Endpoints/openAI/initializeClient.js b/api/server/services/Endpoints/openAI/initializeClient.js index 85dfe9bf0f..06dad5963a 100644 --- a/api/server/services/Endpoints/openAI/initializeClient.js +++ b/api/server/services/Endpoints/openAI/initializeClient.js @@ -4,8 +4,8 @@ const { resolveHeaders, } = require('librechat-data-provider'); const { getUserKey, checkUserKeyExpiry } = require('~/server/services/UserService'); +const { isEnabled, isUserProvided } = require('~/server/utils'); const { getAzureCredentials } = require('~/utils'); -const { isEnabled } = require('~/server/utils'); const { OpenAIClient } = require('~/app'); const initializeClient = async ({ req, res, endpointOption }) => { @@ -21,40 +21,48 @@ const initializeClient = async ({ req, res, endpointOption }) => { const { key: expiresAt, endpoint, model: modelName } = req.body; const contextStrategy = isEnabled(OPENAI_SUMMARIZE) ? 'summarize' : null; + const credentials = { + [EModelEndpoint.openAI]: OPENAI_API_KEY, + [EModelEndpoint.azureOpenAI]: AZURE_API_KEY, + }; + const baseURLOptions = { [EModelEndpoint.openAI]: OPENAI_REVERSE_PROXY, [EModelEndpoint.azureOpenAI]: AZURE_OPENAI_BASEURL, }; - const reverseProxyUrl = baseURLOptions[endpoint] ?? null; + const userProvidesKey = isUserProvided(credentials[endpoint]); + const userProvidesURL = isUserProvided(baseURLOptions[endpoint]); + + let userValues = null; + if (expiresAt && (userProvidesKey || userProvidesURL)) { + checkUserKeyExpiry( + expiresAt, + 'Your OpenAI API values have expired. Please provide them again.', + ); + userValues = await getUserKey({ userId: req.user.id, name: endpoint }); + try { + userValues = JSON.parse(userValues); + } catch (e) { + throw new Error( + `Invalid JSON provided for ${endpoint} user values. Please provide them again.`, + ); + } + } + + let apiKey = userProvidesKey ? userValues.apiKey : credentials[endpoint]; + let baseURL = userProvidesURL ? userValues.baseURL : baseURLOptions[endpoint]; const clientOptions = { debug: isEnabled(DEBUG_OPENAI), contextStrategy, - reverseProxyUrl, + reverseProxyUrl: baseURL ? baseURL : null, proxy: PROXY ?? null, req, res, ...endpointOption, }; - const credentials = { - [EModelEndpoint.openAI]: OPENAI_API_KEY, - [EModelEndpoint.azureOpenAI]: AZURE_API_KEY, - }; - - const isUserProvided = credentials[endpoint] === 'user_provided'; - - let userKey = null; - if (expiresAt && isUserProvided) { - checkUserKeyExpiry( - expiresAt, - 'Your OpenAI API key has expired. Please provide your API key again.', - ); - userKey = await getUserKey({ userId: req.user.id, name: endpoint }); - } - - let apiKey = isUserProvided ? userKey : credentials[endpoint]; const isAzureOpenAI = endpoint === EModelEndpoint.azureOpenAI; /** @type {false | TAzureConfig} */ const azureConfig = isAzureOpenAI && req.app.locals[EModelEndpoint.azureOpenAI]; @@ -87,12 +95,12 @@ const initializeClient = async ({ req, res, endpointOption }) => { apiKey = azureOptions.azureOpenAIApiKey; clientOptions.azure = !serverless && azureOptions; } else if (isAzureOpenAI) { - clientOptions.azure = isUserProvided ? JSON.parse(userKey) : getAzureCredentials(); + clientOptions.azure = userProvidesKey ? JSON.parse(userValues.apiKey) : getAzureCredentials(); apiKey = clientOptions.azure.azureOpenAIApiKey; } if (!apiKey) { - throw new Error(`${endpoint} API key not provided.`); + throw new Error(`${endpoint} API key not provided. Please provide it again.`); } const client = new OpenAIClient(apiKey, clientOptions); diff --git a/api/server/services/Endpoints/openAI/initializeClient.spec.js b/api/server/services/Endpoints/openAI/initializeClient.spec.js index 9be110c40b..bf31d2d9bf 100644 --- a/api/server/services/Endpoints/openAI/initializeClient.spec.js +++ b/api/server/services/Endpoints/openAI/initializeClient.spec.js @@ -1,4 +1,4 @@ -const { EModelEndpoint } = require('librechat-data-provider'); +const { EModelEndpoint, validateAzureGroups } = require('librechat-data-provider'); const { getUserKey } = require('~/server/services/UserService'); const initializeClient = require('./initializeClient'); const { OpenAIClient } = require('~/app'); @@ -16,6 +16,69 @@ describe('initializeClient', () => { locals: {}, }; + const validAzureConfigs = [ + { + group: 'librechat-westus', + apiKey: 'WESTUS_API_KEY', + instanceName: 'librechat-westus', + version: '2023-12-01-preview', + models: { + 'gpt-4-vision-preview': { + deploymentName: 'gpt-4-vision-preview', + version: '2024-02-15-preview', + }, + 'gpt-3.5-turbo': { + deploymentName: 'gpt-35-turbo', + }, + 'gpt-3.5-turbo-1106': { + deploymentName: 'gpt-35-turbo-1106', + }, + 'gpt-4': { + deploymentName: 'gpt-4', + }, + 'gpt-4-1106-preview': { + deploymentName: 'gpt-4-1106-preview', + }, + }, + }, + { + group: 'librechat-eastus', + apiKey: 'EASTUS_API_KEY', + instanceName: 'librechat-eastus', + deploymentName: 'gpt-4-turbo', + version: '2024-02-15-preview', + models: { + 'gpt-4-turbo': true, + }, + baseURL: 'https://eastus.example.com', + additionalHeaders: { + 'x-api-key': 'x-api-key-value', + }, + }, + { + group: 'mistral-inference', + apiKey: 'AZURE_MISTRAL_API_KEY', + baseURL: + 'https://Mistral-large-vnpet-serverless.region.inference.ai.azure.com/v1/chat/completions', + serverless: true, + models: { + 'mistral-large': true, + }, + }, + { + group: 'llama-70b-chat', + apiKey: 'AZURE_LLAMA2_70B_API_KEY', + baseURL: + 'https://Llama-2-70b-chat-qmvyb-serverless.region.inference.ai.azure.com/v1/chat/completions', + serverless: true, + models: { + 'llama-70b-chat': true, + }, + }, + ]; + + const { modelNames, modelGroupMap, groupMap } = validateAzureGroups(validAzureConfigs); + beforeEach(() => { jest.resetModules(); // Clears the cache process.env = { ...originalEnvironment }; // Make a copy @@ -38,10 +101,10 @@ describe('initializeClient', () => { const res = {}; const endpointOption = {}; - const client = await initializeClient({ req, res, endpointOption }); + const result = await initializeClient({ req, res, endpointOption }); - expect(client.openAIApiKey).toBe('test-openai-api-key'); - expect(client.client).toBeInstanceOf(OpenAIClient); + expect(result.openAIApiKey).toBe('test-openai-api-key'); + expect(result.client).toBeInstanceOf(OpenAIClient); }); test('should initialize client with Azure credentials when endpoint is azureOpenAI', async () => { @@ -137,9 +200,7 @@ describe('initializeClient', () => { const res = {}; const endpointOption = {}; - await expect(initializeClient({ req, res, endpointOption })).rejects.toThrow( - 'Your OpenAI API key has expired. Please provide your API key again.', - ); + await expect(initializeClient({ req, res, endpointOption })).rejects.toThrow(/Your OpenAI API/); }); test('should throw an error if no API keys are provided in the environment', async () => { @@ -180,7 +241,7 @@ describe('initializeClient', () => { process.env.OPENAI_API_KEY = 'user_provided'; // Mock getUserKey to return the expected key - getUserKey.mockResolvedValue('test-user-provided-openai-api-key'); + getUserKey.mockResolvedValue(JSON.stringify({ apiKey: 'test-user-provided-openai-api-key' })); // Call the initializeClient function const result = await initializeClient({ req, res, endpointOption }); @@ -205,8 +266,93 @@ describe('initializeClient', () => { // Mock getUserKey to return an invalid key getUserKey.mockResolvedValue(invalidKey); + await expect(initializeClient({ req, res, endpointOption })).rejects.toThrow(/Your OpenAI API/); + }); + + test('should throw an error when user-provided values are not valid JSON', async () => { + process.env.OPENAI_API_KEY = 'user_provided'; + const req = { + body: { key: new Date(Date.now() + 10000).toISOString(), endpoint: 'openAI' }, + user: { id: '123' }, + app, + }; + const res = {}; + const endpointOption = {}; + + // Mock getUserKey to return a non-JSON string + getUserKey.mockResolvedValue('not-a-json'); + await expect(initializeClient({ req, res, endpointOption })).rejects.toThrow( - /Your OpenAI API key has expired/, + /Invalid JSON provided for openAI user values/, ); }); + + test('should initialize client correctly for Azure OpenAI with valid configuration', async () => { + const req = { + body: { + key: null, + endpoint: EModelEndpoint.azureOpenAI, + model: modelNames[0], + }, + user: { id: '123' }, + app: { + locals: { + [EModelEndpoint.azureOpenAI]: { + modelNames, + modelGroupMap, + groupMap, + }, + }, + }, + }; + const res = {}; + const endpointOption = {}; + + const client = await initializeClient({ req, res, endpointOption }); + expect(client.client.options.azure).toBeDefined(); + }); + + test('should initialize client with default options when certain env vars are not set', async () => { + delete process.env.DEBUG_OPENAI; + delete process.env.OPENAI_SUMMARIZE; + + const req = { + body: { key: null, endpoint: 'openAI' }, + user: { id: '123' }, + app, + }; + const res = {}; + const endpointOption = {}; + + const client = await initializeClient({ req, res, endpointOption }); + + expect(client.client.options.debug).toBe(false); + expect(client.client.options.contextStrategy).toBe(null); + }); + + test('should correctly use user-provided apiKey and baseURL when provided', async () => { + process.env.OPENAI_API_KEY = 'user_provided'; + process.env.OPENAI_REVERSE_PROXY = 'user_provided'; + const req = { + body: { + key: new Date(Date.now() + 10000).toISOString(), + endpoint: 'openAI', + }, + user: { + id: '123', + }, + app, + }; + const res = {}; + const endpointOption = {}; + + getUserKey.mockResolvedValue( + JSON.stringify({ apiKey: 'test', baseURL: 'https://user-provided-url.com' }), + ); + + const result = await initializeClient({ req, res, endpointOption }); + + expect(result.openAIApiKey).toBe('test'); + expect(result.client.options.reverseProxyUrl).toBe('https://user-provided-url.com'); + }); }); diff --git a/api/server/services/ModelService.js b/api/server/services/ModelService.js index 96091e7b54..7ffa33b418 100644 --- a/api/server/services/ModelService.js +++ b/api/server/services/ModelService.js @@ -20,6 +20,7 @@ const { openAIApiKey, userProvidedOpenAI } = require('./Config/EndpointService') * @param {boolean} [params.azure=false] - Whether to fetch models from Azure. * @param {boolean} [params.userIdQuery=false] - Whether to send the user ID as a query parameter. * @param {boolean} [params.createTokenConfig=true] - Whether to create a token configuration from the API response. + * @param {string} [params.tokenKey] - The cache key to save the token configuration. Uses `name` if omitted. * @returns {Promise} A promise that resolves to an array of model identifiers. * @async */ @@ -31,6 +32,7 @@ const fetchModels = async ({ azure = false, userIdQuery = false, createTokenConfig = true, + tokenKey, }) => { let models = []; @@ -70,7 +72,7 @@ const fetchModels = async ({ if (validationResult.success && createTokenConfig) { const endpointTokenConfig = processModelData(input); const cache = getLogStores(CacheKeys.TOKEN_CONFIG); - await cache.set(name, endpointTokenConfig); + await cache.set(tokenKey ?? name, endpointTokenConfig); } models = input.data.map((item) => item.id); } catch (error) { diff --git a/api/server/utils/handleText.js b/api/server/utils/handleText.js index 8607d71519..9049bcff44 100644 --- a/api/server/utils/handleText.js +++ b/api/server/utils/handleText.js @@ -172,6 +172,27 @@ function isEnabled(value) { */ const isUserProvided = (value) => value === 'user_provided'; +/** + * Generate the configuration for a given key and base URL. + * @param {string} key + * @param {string} baseURL + * @returns {boolean | { userProvide: boolean, userProvideURL?: boolean }} + */ +function generateConfig(key, baseURL) { + if (!key) { + return false; + } + + /** @type {{ userProvide: boolean, userProvideURL?: boolean }} */ + const config = { userProvide: isUserProvided(key) }; + + if (baseURL) { + config.userProvideURL = isUserProvided(baseURL); + } + + return config; +} + module.exports = { createOnProgress, isEnabled, @@ -180,4 +201,5 @@ module.exports = { formatAction, addSpaceIfNeeded, isUserProvided, + generateConfig, }; diff --git a/client/src/components/Input/SetKeyDialog/OpenAIConfig.tsx b/client/src/components/Input/SetKeyDialog/OpenAIConfig.tsx index a72d88dbd9..db014ddd78 100644 --- a/client/src/components/Input/SetKeyDialog/OpenAIConfig.tsx +++ b/client/src/components/Input/SetKeyDialog/OpenAIConfig.tsx @@ -1,80 +1,101 @@ -import { useEffect, useState } from 'react'; import { EModelEndpoint } from 'librechat-data-provider'; -import { useMultipleKeys } from '~/hooks/Input'; +import { useFormContext, Controller } from 'react-hook-form'; import InputWithLabel from './InputWithLabel'; -import type { TConfigProps } from '~/common'; -import { isJson } from '~/utils/json'; - -const OpenAIConfig = ({ userKey, setUserKey, endpoint }: TConfigProps) => { - const [showPanel, setShowPanel] = useState(endpoint === EModelEndpoint.azureOpenAI); - const { getMultiKey: getAzure, setMultiKey: setAzure } = useMultipleKeys(setUserKey); - - useEffect(() => { - if (isJson(userKey)) { - setShowPanel(true); - } - setUserKey(''); - // eslint-disable-next-line react-hooks/exhaustive-deps - }, []); - - useEffect(() => { - if (!showPanel && isJson(userKey)) { - setUserKey(''); - } - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [showPanel]); +const OpenAIConfig = ({ + endpoint, + userProvideURL, +}: { + endpoint: EModelEndpoint | string; + userProvideURL?: boolean | null; +}) => { + const { control } = useFormContext(); + const isAzure = endpoint === EModelEndpoint.azureOpenAI; return ( - <> - {!showPanel ? ( +
+ {!isAzure && ( + ( + + )} + /> + )} + {isAzure && ( <> - setUserKey(e.target.value ?? '')} - label={'OpenAI API Key'} + ( + + )} /> - - ) : ( - <> - - setAzure('azureOpenAIApiInstanceName', e.target.value ?? '', userKey) - } - label={'Azure OpenAI Instance Name'} + ( + + )} /> - - - setAzure('azureOpenAIApiDeploymentName', e.target.value ?? '', userKey) - } - label={'Azure OpenAI Deployment Name'} + ( + + )} /> - - - setAzure('azureOpenAIApiVersion', e.target.value ?? '', userKey) - } - label={'Azure OpenAI API Version'} - /> - - - setAzure('azureOpenAIApiKey', e.target.value ?? '', userKey) - } - label={'Azure OpenAI API Key'} + ( + + )} /> )} - + {userProvideURL && ( + ( + + )} + /> + )} + ); }; diff --git a/client/src/components/Input/SetKeyDialog/SetKeyDialog.tsx b/client/src/components/Input/SetKeyDialog/SetKeyDialog.tsx index 90013b735a..0f9fc7f2e0 100644 --- a/client/src/components/Input/SetKeyDialog/SetKeyDialog.tsx +++ b/client/src/components/Input/SetKeyDialog/SetKeyDialog.tsx @@ -20,9 +20,18 @@ const endpointComponents = { [EModelEndpoint.custom]: CustomConfig, [EModelEndpoint.azureOpenAI]: OpenAIConfig, [EModelEndpoint.gptPlugins]: OpenAIConfig, + [EModelEndpoint.assistants]: OpenAIConfig, default: OtherConfig, }; +const formSet: Set = new Set([ + EModelEndpoint.openAI, + EModelEndpoint.custom, + EModelEndpoint.azureOpenAI, + EModelEndpoint.gptPlugins, + EModelEndpoint.assistants, +]); + const EXPIRY = { THIRTY_MINUTES: { display: 'in 30 minutes', value: 30 * 60 * 1000 }, TWO_HOURS: { display: 'in 2 hours', value: 2 * 60 * 60 * 1000 }, @@ -47,6 +56,10 @@ const SetKeyDialog = ({ defaultValues: { apiKey: '', baseURL: '', + azureOpenAIApiKey: '', + azureOpenAIApiInstanceName: '', + azureOpenAIApiDeploymentName: '', + azureOpenAIApiVersion: '', // TODO: allow endpoint definitions from user // name: '', // TODO: add custom endpoint models defined by user @@ -76,10 +89,26 @@ const SetKeyDialog = ({ onOpenChange(false); }; - if (endpoint === EModelEndpoint.custom || endpointType === EModelEndpoint.custom) { + if (formSet.has(endpoint) || formSet.has(endpointType ?? '')) { // TODO: handle other user provided options besides baseURL and apiKey methods.handleSubmit((data) => { + const isAzure = endpoint === EModelEndpoint.azureOpenAI; + const isOpenAIBase = + isAzure || + endpoint === EModelEndpoint.openAI || + endpoint === EModelEndpoint.gptPlugins || + endpoint === EModelEndpoint.assistants; + if (isAzure) { + data.apiKey = 'n/a'; + } + const emptyValues = Object.keys(data).filter((key) => { + if (!isAzure && key.startsWith('azure')) { + return false; + } + if (isOpenAIBase && key === 'baseURL') { + return false; + } if (key === 'baseURL' && !userProvideURL) { return false; } @@ -92,10 +121,22 @@ const SetKeyDialog = ({ status: 'error', }); onOpenChange(true); - } else { - saveKey(JSON.stringify(data)); - methods.reset(); + return; } + + const { apiKey, baseURL, ...azureOptions } = data; + const userProvidedData = { apiKey, baseURL }; + if (isAzure) { + userProvidedData.apiKey = JSON.stringify({ + azureOpenAIApiKey: azureOptions.azureOpenAIApiKey, + azureOpenAIApiInstanceName: azureOptions.azureOpenAIApiInstanceName, + azureOpenAIApiDeploymentName: azureOptions.azureOpenAIApiDeploymentName, + azureOpenAIApiVersion: azureOptions.azureOpenAIApiVersion, + }); + } + + saveKey(JSON.stringify(userProvidedData)); + methods.reset(); })(); return; } diff --git a/package-lock.json b/package-lock.json index 7f02e47514..638cc7dca9 100644 --- a/package-lock.json +++ b/package-lock.json @@ -11288,9 +11288,9 @@ } }, "node_modules/caniuse-lite": { - "version": "1.0.30001584", - "resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001584.tgz", - "integrity": "sha512-LOz7CCQ9M1G7OjJOF9/mzmqmj3jE/7VOmrfw6Mgs0E8cjOsbRXQJHsPBfmBOXDskXKrHLyyW3n7kpDW/4BsfpQ==", + "version": "1.0.30001591", + "resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001591.tgz", + "integrity": "sha512-PCzRMei/vXjJyL5mJtzNiUCKP59dm8Apqc3PH8gJkMnMXZGox93RbE76jHsmLwmIo6/3nsYIpJtx0O7u5PqFuQ==", "dev": true, "funding": [ { @@ -27993,7 +27993,7 @@ }, "packages/data-provider": { "name": "librechat-data-provider", - "version": "0.4.4", + "version": "0.4.5", "license": "ISC", "dependencies": { "@types/js-yaml": "^4.0.9", diff --git a/packages/data-provider/package.json b/packages/data-provider/package.json index a33a5fbed7..943833e47c 100644 --- a/packages/data-provider/package.json +++ b/packages/data-provider/package.json @@ -1,6 +1,6 @@ { "name": "librechat-data-provider", - "version": "0.4.5", + "version": "0.4.6", "description": "data services for librechat apps", "main": "dist/index.js", "module": "dist/index.es.js", diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index 1e71f140f1..76b6b02d02 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -171,16 +171,20 @@ export const configSchema = z.object({ export type TCustomConfig = z.infer; -export const KnownEndpoints = { - mistral: 'mistral', - openrouter: 'openrouter', - groq: 'groq', - anyscale: 'anyscale', - fireworks: 'fireworks', - ollama: 'ollama', - perplexity: 'perplexity', - 'together.ai': 'together.ai', -} as const; +export enum KnownEndpoints { + mistral = 'mistral', + openrouter = 'openrouter', + groq = 'groq', + anyscale = 'anyscale', + fireworks = 'fireworks', + ollama = 'ollama', + perplexity = 'perplexity', + 'together.ai' = 'together.ai', +} + +export enum FetchTokenConfig { + openrouter = KnownEndpoints.openrouter, +} export const defaultEndpoints: EModelEndpoint[] = [ EModelEndpoint.openAI,