diff --git a/api/server/services/Config/EndpointService.js b/api/server/services/Config/EndpointService.js index d8277dd67f76..cee17bc446d2 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 670bc22d1196..cfb0e97ba3af 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 a31d6e10c494..b4d2663c7918 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 userValues = null; + if (isUserProvided) { + if (expiresAt) { + checkUserKeyExpiry(expiresAt, EModelEndpoint.bedrock); + } + userValues = await getUserKeyValues({ userId: req.user.id, name: EModelEndpoint.bedrock }); + } - 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 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 fc6278c20a0d..be5f3cb29935 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 000000000000..c36e741bcba6 --- /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 0cd8cf3d72c8..4cb5117fab4d 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 95aa05fee284..0891b5dfb3ec 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 5799ca50e88d..2fb13f67d214 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[];