From 66dc48c8a04e723fb7142162b5a3072354427172 Mon Sep 17 00:00:00 2001 From: Dustin Healy Date: Thu, 24 Jul 2025 05:53:38 -0700 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20WIP:=20Enhance=20Bedrock=20endpo?= =?UTF-8?q?int=20configuration=20with=20user-provided=20credentials.=20(St?= =?UTF-8?q?ill=20needs=20to=20implement=20user=5Fprovided=20bearer=20token?= =?UTF-8?q?=20support,=20but=20the=20UI=20is=20there=20for=20it)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Added support for user-provided AWS credentials (Access Key ID, Secret Access Key, Session Token, Bearer Token) in the Bedrock endpoint configuration. - Localized new strings for Bedrock configuration in translation files. --- api/server/services/Config/EndpointService.js | 4 +- .../services/Config/getEndpointsConfig.js | 17 +++ .../services/Endpoints/bedrock/options.js | 35 ++++-- .../Chat/Menus/Endpoints/DialogManager.tsx | 20 ++++ .../Input/SetKeyDialog/BedrockConfig.tsx | 111 ++++++++++++++++++ .../Input/SetKeyDialog/SetKeyDialog.tsx | 98 +++++++++++++++- client/src/locales/en/translation.json | 6 + packages/data-provider/src/types.ts | 4 + 8 files changed, 280 insertions(+), 15 deletions(-) create mode 100644 client/src/components/Input/SetKeyDialog/BedrockConfig.tsx diff --git a/api/server/services/Config/EndpointService.js b/api/server/services/Config/EndpointService.js index d8277dd67f..cee17bc446 100644 --- a/api/server/services/Config/EndpointService.js +++ b/api/server/services/Config/EndpointService.js @@ -45,7 +45,9 @@ module.exports = { EModelEndpoint.azureAssistants, ), [EModelEndpoint.bedrock]: generateConfig( - process.env.BEDROCK_AWS_SECRET_ACCESS_KEY ?? process.env.BEDROCK_AWS_DEFAULT_REGION, + process.env.BEDROCK_AWS_SECRET_ACCESS_KEY ?? + process.env.BEDROCK_AWS_BEARER_TOKEN ?? + process.env.BEDROCK_AWS_DEFAULT_REGION, ), /* key will be part of separate config */ [EModelEndpoint.agents]: generateConfig('true', undefined, EModelEndpoint.agents), diff --git a/api/server/services/Config/getEndpointsConfig.js b/api/server/services/Config/getEndpointsConfig.js index 670bc22d11..cfb0e97ba3 100644 --- a/api/server/services/Config/getEndpointsConfig.js +++ b/api/server/services/Config/getEndpointsConfig.js @@ -74,6 +74,23 @@ async function getEndpointsConfig(req) { }; } + // Add individual credential flags for Bedrock + if (mergedConfig[EModelEndpoint.bedrock]) { + const userProvideAccessKeyId = process.env.BEDROCK_AWS_ACCESS_KEY_ID === 'user_provided'; + const userProvideSecretAccessKey = + process.env.BEDROCK_AWS_SECRET_ACCESS_KEY === 'user_provided'; + const userProvideSessionToken = process.env.BEDROCK_AWS_SESSION_TOKEN === 'user_provided'; + const userProvideBearerToken = process.env.BEDROCK_AWS_BEARER_TOKEN === 'user_provided'; + + mergedConfig[EModelEndpoint.bedrock] = { + ...mergedConfig[EModelEndpoint.bedrock], + userProvideAccessKeyId, + userProvideSecretAccessKey, + userProvideSessionToken, + userProvideBearerToken, + }; + } + const endpointsConfig = orderEndpointsConfig(mergedConfig); await cache.set(CacheKeys.ENDPOINT_CONFIG, endpointsConfig); diff --git a/api/server/services/Endpoints/bedrock/options.js b/api/server/services/Endpoints/bedrock/options.js index a31d6e10c4..b4d2663c79 100644 --- a/api/server/services/Endpoints/bedrock/options.js +++ b/api/server/services/Endpoints/bedrock/options.js @@ -8,27 +8,43 @@ const { bedrockOutputParser, removeNullishValues, } = require('librechat-data-provider'); -const { getUserKey, checkUserKeyExpiry } = require('~/server/services/UserService'); +const { getUserKeyValues, checkUserKeyExpiry } = require('~/server/services/UserService'); const getOptions = async ({ req, overrideModel, endpointOption }) => { const { BEDROCK_AWS_SECRET_ACCESS_KEY, BEDROCK_AWS_ACCESS_KEY_ID, BEDROCK_AWS_SESSION_TOKEN, + BEDROCK_AWS_BEARER_TOKEN, BEDROCK_REVERSE_PROXY, BEDROCK_AWS_DEFAULT_REGION, PROXY, } = process.env; const expiresAt = req.body.key; - const isUserProvided = BEDROCK_AWS_SECRET_ACCESS_KEY === AuthType.USER_PROVIDED; + const isUserProvided = + BEDROCK_AWS_SECRET_ACCESS_KEY === AuthType.USER_PROVIDED || + BEDROCK_AWS_BEARER_TOKEN === AuthType.USER_PROVIDED; - let credentials = isUserProvided - ? await getUserKey({ userId: req.user.id, name: EModelEndpoint.bedrock }) - : { - accessKeyId: BEDROCK_AWS_ACCESS_KEY_ID, - secretAccessKey: BEDROCK_AWS_SECRET_ACCESS_KEY, - ...(BEDROCK_AWS_SESSION_TOKEN && { sessionToken: BEDROCK_AWS_SESSION_TOKEN }), - }; + let userValues = null; + if (isUserProvided) { + if (expiresAt) { + checkUserKeyExpiry(expiresAt, EModelEndpoint.bedrock); + } + userValues = await getUserKeyValues({ userId: req.user.id, name: EModelEndpoint.bedrock }); + } + + let credentials; + if (isUserProvided) { + credentials = JSON.parse(userValues.apiKey); + } else if (BEDROCK_AWS_BEARER_TOKEN) { + credentials = { bearerToken: BEDROCK_AWS_BEARER_TOKEN }; + } else { + credentials = { + accessKeyId: BEDROCK_AWS_ACCESS_KEY_ID, + secretAccessKey: BEDROCK_AWS_SECRET_ACCESS_KEY, + ...(BEDROCK_AWS_SESSION_TOKEN && { sessionToken: BEDROCK_AWS_SESSION_TOKEN }), + }; + } if (!credentials) { throw new Error('Bedrock credentials not provided. Please provide them again.'); @@ -36,6 +52,7 @@ const getOptions = async ({ req, overrideModel, endpointOption }) => { if ( !isUserProvided && + !credentials.bearerToken && (credentials.accessKeyId === undefined || credentials.accessKeyId === '') && (credentials.secretAccessKey === undefined || credentials.secretAccessKey === '') ) { diff --git a/client/src/components/Chat/Menus/Endpoints/DialogManager.tsx b/client/src/components/Chat/Menus/Endpoints/DialogManager.tsx index fc6278c20a..be5f3cb299 100644 --- a/client/src/components/Chat/Menus/Endpoints/DialogManager.tsx +++ b/client/src/components/Chat/Menus/Endpoints/DialogManager.tsx @@ -25,6 +25,26 @@ const DialogManager = ({ endpointType={getEndpointField(endpointsConfig, keyDialogEndpoint, 'type')} onOpenChange={onOpenChange} userProvideURL={getEndpointField(endpointsConfig, keyDialogEndpoint, 'userProvideURL')} + userProvideAccessKeyId={getEndpointField( + endpointsConfig, + keyDialogEndpoint, + 'userProvideAccessKeyId', + )} + userProvideSecretAccessKey={getEndpointField( + endpointsConfig, + keyDialogEndpoint, + 'userProvideSecretAccessKey', + )} + userProvideSessionToken={getEndpointField( + endpointsConfig, + keyDialogEndpoint, + 'userProvideSessionToken', + )} + userProvideBearerToken={getEndpointField( + endpointsConfig, + keyDialogEndpoint, + 'userProvideBearerToken', + )} /> )} diff --git a/client/src/components/Input/SetKeyDialog/BedrockConfig.tsx b/client/src/components/Input/SetKeyDialog/BedrockConfig.tsx new file mode 100644 index 0000000000..c36e741bcb --- /dev/null +++ b/client/src/components/Input/SetKeyDialog/BedrockConfig.tsx @@ -0,0 +1,111 @@ +import React from 'react'; +import { EModelEndpoint } from 'librechat-data-provider'; +import { useFormContext, Controller } from 'react-hook-form'; +import { useLocalize } from '~/hooks'; +import InputWithLabel from './InputWithLabel'; + +const BedrockConfig = ({ + userProvideAccessKeyId, + userProvideSecretAccessKey, + userProvideSessionToken, + userProvideBearerToken, +}: { + endpoint: EModelEndpoint | string; + userProvideURL?: boolean | null; + userProvideAccessKeyId?: boolean; + userProvideSecretAccessKey?: boolean; + userProvideSessionToken?: boolean; + userProvideBearerToken?: boolean; +}) => { + const { control } = useFormContext(); + const localize = useLocalize(); + + const renderFields = () => { + const fields: React.ReactNode[] = []; + + if (userProvideAccessKeyId) { + fields.push( + ( + + )} + />, + ); + } + + if (userProvideSecretAccessKey) { + if (fields.length > 0) fields.push(
); + fields.push( + ( + + )} + />, + ); + } + + if (userProvideSessionToken) { + if (fields.length > 0) fields.push(
); + fields.push( + ( + + )} + />, + ); + } + + if (userProvideBearerToken) { + if (fields.length > 0) fields.push(
); + fields.push( + ( + + )} + />, + ); + } + + return <>{fields}; + }; + + return
{renderFields()}
; +}; + +export default BedrockConfig; diff --git a/client/src/components/Input/SetKeyDialog/SetKeyDialog.tsx b/client/src/components/Input/SetKeyDialog/SetKeyDialog.tsx index 0cd8cf3d72..4cb5117fab 100644 --- a/client/src/components/Input/SetKeyDialog/SetKeyDialog.tsx +++ b/client/src/components/Input/SetKeyDialog/SetKeyDialog.tsx @@ -12,6 +12,7 @@ import CustomConfig from './CustomEndpoint'; import GoogleConfig from './GoogleConfig'; import OpenAIConfig from './OpenAIConfig'; import OtherConfig from './OtherConfig'; +import BedrockConfig from './BedrockConfig'; import HelpText from './HelpText'; const endpointComponents = { @@ -22,6 +23,7 @@ const endpointComponents = { [EModelEndpoint.gptPlugins]: OpenAIConfig, [EModelEndpoint.assistants]: OpenAIConfig, [EModelEndpoint.azureAssistants]: OpenAIConfig, + [EModelEndpoint.bedrock]: BedrockConfig, default: OtherConfig, }; @@ -32,6 +34,7 @@ const formSet: Set = new Set([ EModelEndpoint.gptPlugins, EModelEndpoint.assistants, EModelEndpoint.azureAssistants, + EModelEndpoint.bedrock, ]); const EXPIRY = { @@ -50,10 +53,18 @@ const SetKeyDialog = ({ endpoint, endpointType, userProvideURL, + userProvideAccessKeyId, + userProvideSecretAccessKey, + userProvideSessionToken, + userProvideBearerToken, }: Pick & { endpoint: EModelEndpoint | string; endpointType?: EModelEndpoint; userProvideURL?: boolean | null; + userProvideAccessKeyId?: boolean; + userProvideSecretAccessKey?: boolean; + userProvideSessionToken?: boolean; + userProvideBearerToken?: boolean; }) => { const methods = useForm({ defaultValues: { @@ -63,6 +74,10 @@ const SetKeyDialog = ({ azureOpenAIApiInstanceName: '', azureOpenAIApiDeploymentName: '', azureOpenAIApiVersion: '', + bedrockAccessKeyId: '', + bedrockSecretAccessKey: '', + bedrockSessionToken: '', + bedrockBearerToken: '', // TODO: allow endpoint definitions from user // name: '', // TODO: add custom endpoint models defined by user @@ -102,6 +117,7 @@ const SetKeyDialog = ({ // TODO: handle other user provided options besides baseURL and apiKey methods.handleSubmit((data) => { const isAzure = endpoint === EModelEndpoint.azureOpenAI; + const isBedrock = endpoint === EModelEndpoint.bedrock; const isOpenAIBase = isAzure || endpoint === EModelEndpoint.openAI || @@ -115,6 +131,9 @@ const SetKeyDialog = ({ if (!isAzure && key.startsWith('azure')) { return false; } + if (!isBedrock && key.startsWith('bedrock')) { + return false; + } if (isOpenAIBase && key === 'baseURL') { return false; } @@ -124,16 +143,67 @@ const SetKeyDialog = ({ return data[key] === ''; }); - if (emptyValues.length > 0) { + if (isBedrock) { + const missingFields: string[] = []; + let hasValidCredentials = false; + + if (userProvideBearerToken && !data.bedrockBearerToken) { + missingFields.push('AWS Bedrock Bearer Token'); + } else if (userProvideBearerToken && data.bedrockBearerToken) { + hasValidCredentials = true; + } + + if (userProvideAccessKeyId && !data.bedrockAccessKeyId) { + missingFields.push('AWS Access Key ID'); + } + if (userProvideSecretAccessKey && !data.bedrockSecretAccessKey) { + missingFields.push('AWS Secret Access Key'); + } + + if ( + userProvideAccessKeyId && + userProvideSecretAccessKey && + data.bedrockAccessKeyId && + data.bedrockSecretAccessKey + ) { + hasValidCredentials = true; + } + + if (missingFields.length > 0) { + showToast({ + message: `${localize('com_endpoint_config_required_fields')} ${missingFields.join(', ')}`, + status: 'error', + }); + onOpenChange(true); + return; + } + + if (!hasValidCredentials) { + showToast({ + message: localize('com_endpoint_config_bedrock_credentials_required'), + status: 'error', + }); + onOpenChange(true); + return; + } + } else if (emptyValues.length > 0) { showToast({ - message: 'The following fields are required: ' + emptyValues.join(', '), + message: `${localize('com_endpoint_config_required_fields')} ${emptyValues.join(', ')}`, status: 'error', }); onOpenChange(true); return; } - const { apiKey, baseURL, ...azureOptions } = data; + const { + apiKey, + baseURL, + bedrockAccessKeyId, + bedrockSecretAccessKey, + bedrockSessionToken, + bedrockBearerToken, + ...azureOptions + } = data; const userProvidedData = { apiKey, baseURL }; if (isAzure) { userProvidedData.apiKey = JSON.stringify({ @@ -142,6 +212,20 @@ const SetKeyDialog = ({ azureOpenAIApiDeploymentName: azureOptions.azureOpenAIApiDeploymentName, azureOpenAIApiVersion: azureOptions.azureOpenAIApiVersion, }); + } else if (isBedrock) { + // Prioritize bearer token if provided + if (bedrockBearerToken) { + userProvidedData.apiKey = JSON.stringify({ + bearerToken: bedrockBearerToken, + }); + } else { + // Use access keys + userProvidedData.apiKey = JSON.stringify({ + accessKeyId: bedrockAccessKeyId, + secretAccessKey: bedrockSecretAccessKey, + ...(bedrockSessionToken && { sessionToken: bedrockSessionToken }), + }); + } } saveKey(JSON.stringify(userProvidedData)); @@ -171,8 +255,8 @@ const SetKeyDialog = ({ {expiryTime === 'never' ? localize('com_endpoint_config_key_never_expires') : `${localize('com_endpoint_config_key_encryption')} ${new Date( - expiryTime ?? 0, - ).toLocaleString()}`} + expiryTime ?? 0, + ).toLocaleString()}`} diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index 95aa05fee2..0891b5dfb3 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -187,6 +187,12 @@ "com_endpoint_config_key_never_expires": "Your key will never expire", "com_endpoint_config_placeholder": "Set your Key in the Header menu to chat.", "com_endpoint_config_value": "Enter value for", + "com_endpoint_config_bedrock_access_key_id": "AWS Access Key ID", + "com_endpoint_config_bedrock_secret_access_key": "AWS Secret Access Key", + "com_endpoint_config_bedrock_session_token": "AWS Session Token", + "com_endpoint_config_bedrock_bearer_token": "AWS Bedrock Bearer Token", + "com_endpoint_config_bedrock_credentials_required": "Please provide either Access Keys (Access Key ID + Secret Access Key) or Bearer Token", + "com_endpoint_config_required_fields": "The following fields are required:", "com_endpoint_context": "Context", "com_endpoint_context_info": "The maximum number of tokens that can be used for context. Use this for control of how many tokens are sent per request. If unspecified, will use system defaults based on known models' context size. Setting higher values may result in errors and/or higher token cost.", "com_endpoint_context_tokens": "Max Context Tokens", diff --git a/packages/data-provider/src/types.ts b/packages/data-provider/src/types.ts index 5799ca50e8..2fb13f67d2 100644 --- a/packages/data-provider/src/types.ts +++ b/packages/data-provider/src/types.ts @@ -330,6 +330,10 @@ export type TConfig = { modelDisplayLabel?: string; userProvide?: boolean | null; userProvideURL?: boolean | null; + userProvideAccessKeyId?: boolean; + userProvideSecretAccessKey?: boolean; + userProvideSessionToken?: boolean; + userProvideBearerToken?: boolean; disableBuilder?: boolean; retrievalModels?: string[]; capabilities?: string[];