From 53e7c41033d8f7738733c69b6f23d9b7292b4db9 Mon Sep 17 00:00:00 2001 From: Dustin Healy <54083382+dustinhealy@users.noreply.github.com> Date: Sat, 23 May 2026 09:41:38 -0700 Subject: [PATCH] =?UTF-8?q?=F0=9F=AA=99=20feat:=20Add=20AWS=20Bedrock=20AP?= =?UTF-8?q?I=20key=20support=20(#8690)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: Add Bedrock API key support * fix: Respect Bedrock credential mode * fix: Support mixed Bedrock credential forms --------- Co-authored-by: Danny Avila --- .env.example | 3 + api/server/services/Config/EndpointService.js | 15 ++- .../Config/__tests__/EndpointService.spec.js | 36 ++++++ .../Chat/Menus/Endpoints/DialogManager.tsx | 20 ++++ .../Input/SetKeyDialog/BedrockConfig.tsx | 111 +++++++++++++++++ .../Input/SetKeyDialog/SetKeyDialog.tsx | 110 +++++++++++++++-- client/src/locales/en/translation.json | 6 + .../src/endpoints/bedrock/initialize.spec.ts | 77 ++++++++++++ .../api/src/endpoints/bedrock/initialize.ts | 113 ++++++++++++++---- .../src/endpoints/config/endpoints.spec.ts | 50 +++++++- .../api/src/endpoints/config/endpoints.ts | 14 ++- packages/api/src/types/bedrock.ts | 7 ++ packages/data-provider/src/types.ts | 4 + 13 files changed, 527 insertions(+), 39 deletions(-) create mode 100644 client/src/components/Input/SetKeyDialog/BedrockConfig.tsx diff --git a/.env.example b/.env.example index 2020b38242..645a44b833 100644 --- a/.env.example +++ b/.env.example @@ -233,6 +233,9 @@ ANTHROPIC_API_KEY=user_provided # BEDROCK_AWS_SECRET_ACCESS_KEY=someSecretAccessKey # BEDROCK_AWS_SESSION_TOKEN=someSessionToken +# Bedrock API key +# BEDROCK_AWS_BEARER_TOKEN=yourBedrockApiKey + # Note: This example list is not meant to be exhaustive. If omitted, all known, supported model IDs will be included for you. # BEDROCK_AWS_MODELS=anthropic.claude-opus-4-7,anthropic.claude-sonnet-4-6,anthropic.claude-opus-4-6-v1,anthropic.claude-3-5-sonnet-20240620-v1:0,meta.llama3-1-8b-instruct-v1:0 # Cross-region inference model IDs: us.anthropic.claude-opus-4-7,us.anthropic.claude-sonnet-4-6,us.anthropic.claude-opus-4-6-v1,global.anthropic.claude-opus-4-6-v1 diff --git a/api/server/services/Config/EndpointService.js b/api/server/services/Config/EndpointService.js index 058341ca2c..56fa84f2a2 100644 --- a/api/server/services/Config/EndpointService.js +++ b/api/server/services/Config/EndpointService.js @@ -17,6 +17,13 @@ const { const userProvidedOpenAI = isUserProvided(openAIApiKey); const anthropicUsesVertex = isEnabled(process.env.ANTHROPIC_USE_VERTEX); +const firstNonEmpty = (...values) => values.find((value) => value != null && value !== ''); +const bedrockUserProvidedCredential = [ + process.env.BEDROCK_AWS_BEARER_TOKEN, + process.env.BEDROCK_AWS_ACCESS_KEY_ID, + process.env.BEDROCK_AWS_SECRET_ACCESS_KEY, + process.env.BEDROCK_AWS_SESSION_TOKEN, +].find(isUserProvided); module.exports = { config: { @@ -38,7 +45,13 @@ module.exports = { EModelEndpoint.azureAssistants, ), [EModelEndpoint.bedrock]: generateConfig( - process.env.BEDROCK_AWS_SECRET_ACCESS_KEY ?? process.env.BEDROCK_AWS_DEFAULT_REGION, + bedrockUserProvidedCredential ?? + firstNonEmpty( + process.env.BEDROCK_AWS_BEARER_TOKEN, + process.env.BEDROCK_AWS_SECRET_ACCESS_KEY, + process.env.BEDROCK_AWS_PROFILE, + 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/__tests__/EndpointService.spec.js b/api/server/services/Config/__tests__/EndpointService.spec.js index 82f7175d23..0943daca47 100644 --- a/api/server/services/Config/__tests__/EndpointService.spec.js +++ b/api/server/services/Config/__tests__/EndpointService.spec.js @@ -60,6 +60,12 @@ describe('EndpointService', () => { process.env = { ...originalEnv }; delete process.env.ANTHROPIC_API_KEY; delete process.env.ANTHROPIC_USE_VERTEX; + delete process.env.BEDROCK_AWS_ACCESS_KEY_ID; + delete process.env.BEDROCK_AWS_SECRET_ACCESS_KEY; + delete process.env.BEDROCK_AWS_SESSION_TOKEN; + delete process.env.BEDROCK_AWS_BEARER_TOKEN; + delete process.env.BEDROCK_AWS_PROFILE; + delete process.env.BEDROCK_AWS_DEFAULT_REGION; Object.assign(process.env, env); return require('../EndpointService').config; } @@ -80,4 +86,34 @@ describe('EndpointService', () => { expect(config[EModelEndpoint.anthropic]).toEqual({ userProvide: true }); }); + + it('requires a user Bedrock key when bearer token user_provided is set with a legacy static secret', () => { + const config = loadConfig({ + BEDROCK_AWS_SECRET_ACCESS_KEY: 'legacy-secret', + BEDROCK_AWS_BEARER_TOKEN: 'user_provided', + BEDROCK_AWS_DEFAULT_REGION: 'us-east-1', + }); + + expect(config[EModelEndpoint.bedrock]).toEqual({ userProvide: true }); + }); + + it('enables Bedrock with static bearer token before static secret credentials', () => { + const config = loadConfig({ + BEDROCK_AWS_SECRET_ACCESS_KEY: 'legacy-secret', + BEDROCK_AWS_BEARER_TOKEN: 'bedrock-api-key', + BEDROCK_AWS_DEFAULT_REGION: 'us-east-1', + }); + + expect(config[EModelEndpoint.bedrock]).toEqual({ userProvide: false }); + }); + + it('skips blank optional Bedrock env vars before falling back to region', () => { + const config = loadConfig({ + BEDROCK_AWS_BEARER_TOKEN: '', + BEDROCK_AWS_PROFILE: '', + BEDROCK_AWS_DEFAULT_REGION: 'us-east-1', + }); + + expect(config[EModelEndpoint.bedrock]).toEqual({ userProvide: false }); + }); }); diff --git a/client/src/components/Chat/Menus/Endpoints/DialogManager.tsx b/client/src/components/Chat/Menus/Endpoints/DialogManager.tsx index 75eed646b0..1596512294 100644 --- a/client/src/components/Chat/Menus/Endpoints/DialogManager.tsx +++ b/client/src/components/Chat/Menus/Endpoints/DialogManager.tsx @@ -24,6 +24,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 7fec25e4a5..8735333da6 100644 --- a/client/src/components/Input/SetKeyDialog/SetKeyDialog.tsx +++ b/client/src/components/Input/SetKeyDialog/SetKeyDialog.tsx @@ -25,6 +25,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'; import { logger } from '~/utils'; @@ -35,6 +36,7 @@ const endpointComponents = { [EModelEndpoint.azureOpenAI]: OpenAIConfig, [EModelEndpoint.assistants]: OpenAIConfig, [EModelEndpoint.azureAssistants]: OpenAIConfig, + [EModelEndpoint.bedrock]: BedrockConfig, default: OtherConfig, }; @@ -44,6 +46,7 @@ const formSet: Set = new Set([ EModelEndpoint.azureOpenAI, EModelEndpoint.assistants, EModelEndpoint.azureAssistants, + EModelEndpoint.bedrock, ]); const EXPIRY = { @@ -150,10 +153,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: { @@ -163,6 +174,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 @@ -177,6 +192,7 @@ const SetKeyDialog = ({ const localize = useLocalize(); const expirationOptions = Object.values(EXPIRY); + const configuredEndpoint = endpointType ?? endpoint; const handleExpirationChange = (label: string) => { setExpiresAtLabel(label); @@ -212,9 +228,12 @@ const SetKeyDialog = ({ 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 isAzure = configuredEndpoint === EModelEndpoint.azureOpenAI; + const isBedrock = configuredEndpoint === EModelEndpoint.bedrock; const isOpenAIBase = - isAzure || endpoint === EModelEndpoint.openAI || isAssistantsEndpoint(endpoint); + isAzure || + configuredEndpoint === EModelEndpoint.openAI || + isAssistantsEndpoint(configuredEndpoint); if (isAzure) { data.apiKey = 'n/a'; } @@ -223,6 +242,9 @@ const SetKeyDialog = ({ if (!isAzure && key.startsWith('azure')) { return false; } + if (!isBedrock && key.startsWith('bedrock')) { + return false; + } if (isOpenAIBase && key === 'baseURL') { return false; } @@ -232,16 +254,68 @@ const SetKeyDialog = ({ return data[key] === ''; }); - if (emptyValues.length > 0) { + if (isBedrock) { + const bearerToken = userProvideBearerToken ? data.bedrockBearerToken?.trim() : ''; + const accessKeyId = userProvideAccessKeyId ? data.bedrockAccessKeyId?.trim() : ''; + const secretAccessKey = userProvideSecretAccessKey + ? data.bedrockSecretAccessKey?.trim() + : ''; + const sessionToken = userProvideSessionToken ? data.bedrockSessionToken?.trim() : ''; + const accessKeyIdLabel = localize('com_endpoint_config_bedrock_access_key_id'); + const secretAccessKeyLabel = localize('com_endpoint_config_bedrock_secret_access_key'); + const sessionTokenLabel = localize('com_endpoint_config_bedrock_session_token'); + const bearerTokenLabel = localize('com_endpoint_config_bedrock_bearer_token'); + const canSubmitBearerToken = !!bearerToken; + const hasUserProvidedAccessKeyAuth = + !!userProvideAccessKeyId || !!userProvideSecretAccessKey || !!userProvideSessionToken; + const missingFields = [ + !canSubmitBearerToken && !hasUserProvidedAccessKeyAuth && userProvideBearerToken + ? bearerTokenLabel + : '', + !canSubmitBearerToken && userProvideAccessKeyId && !accessKeyId ? accessKeyIdLabel : '', + !canSubmitBearerToken && userProvideSecretAccessKey && !secretAccessKey + ? secretAccessKeyLabel + : '', + !canSubmitBearerToken && userProvideSessionToken && !sessionToken + ? sessionTokenLabel + : '', + ].filter(Boolean); + + if (!canSubmitBearerToken && missingFields.length > 0) { + showToast({ + message: `${localize('com_endpoint_config_required_fields')} ${missingFields.join(', ')}`, + status: NotificationSeverity.ERROR, + }); + onOpenChange(true); + return; + } + + if (!canSubmitBearerToken && !hasUserProvidedAccessKeyAuth) { + showToast({ + message: localize('com_endpoint_config_bedrock_credentials_required'), + status: NotificationSeverity.ERROR, + }); + onOpenChange(true); + return; + } + } else if (emptyValues.length > 0) { showToast({ - message: 'The following fields are required: ' + emptyValues.join(', '), - status: 'error', + message: `${localize('com_endpoint_config_required_fields')} ${emptyValues.join(', ')}`, + status: NotificationSeverity.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({ @@ -250,6 +324,23 @@ const SetKeyDialog = ({ azureOpenAIApiDeploymentName: azureOptions.azureOpenAIApiDeploymentName, azureOpenAIApiVersion: azureOptions.azureOpenAIApiVersion, }); + } else if (isBedrock) { + const bearerToken = userProvideBearerToken ? bedrockBearerToken.trim() : ''; + const accessKeyId = userProvideAccessKeyId ? bedrockAccessKeyId.trim() : ''; + const secretAccessKey = userProvideSecretAccessKey ? bedrockSecretAccessKey.trim() : ''; + const sessionToken = userProvideSessionToken ? bedrockSessionToken.trim() : ''; + + if (bearerToken) { + userProvidedData.apiKey = JSON.stringify({ + bearerToken, + }); + } else { + userProvidedData.apiKey = JSON.stringify({ + ...(accessKeyId && { accessKeyId }), + ...(secretAccessKey && { secretAccessKey }), + ...(sessionToken && { sessionToken }), + }); + } } saveKey(JSON.stringify(userProvidedData)); @@ -270,8 +361,7 @@ const SetKeyDialog = ({ setUserKey(''); }; - const EndpointComponent = - endpointComponents[endpointType ?? endpoint] ?? endpointComponents['default']; + const EndpointComponent = endpointComponents[configuredEndpoint] ?? endpointComponents['default']; const expiryTime = getExpiry(); return ( @@ -305,6 +395,10 @@ const SetKeyDialog = ({ endpoint={endpoint} setUserKey={setUserKey} userProvideURL={userProvideURL} + userProvideAccessKeyId={userProvideAccessKeyId} + userProvideSecretAccessKey={userProvideSecretAccessKey} + userProvideSessionToken={userProvideSessionToken} + userProvideBearerToken={userProvideBearerToken} /> diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index 0e6dfa39f7..2c89ed8e09 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -259,6 +259,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 API Key", + "com_endpoint_config_bedrock_credentials_required": "Please provide either an AWS Bedrock API key or AWS access keys (Access Key ID + Secret Access Key)", + "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/api/src/endpoints/bedrock/initialize.spec.ts b/packages/api/src/endpoints/bedrock/initialize.spec.ts index 666291e7b7..65be1c65b8 100644 --- a/packages/api/src/endpoints/bedrock/initialize.spec.ts +++ b/packages/api/src/endpoints/bedrock/initialize.spec.ts @@ -75,6 +75,11 @@ describe('initializeBedrock', () => { beforeEach(() => { jest.clearAllMocks(); process.env = { ...originalEnv }; + delete process.env.BEDROCK_AWS_BEARER_TOKEN; + delete process.env.BEDROCK_AWS_PROFILE; + delete process.env.BEDROCK_AWS_SESSION_TOKEN; + delete process.env.BEDROCK_REVERSE_PROXY; + delete process.env.PROXY; process.env.BEDROCK_AWS_ACCESS_KEY_ID = 'test-access-key'; process.env.BEDROCK_AWS_SECRET_ACCESS_KEY = 'test-secret-key'; process.env.BEDROCK_AWS_DEFAULT_REGION = 'us-east-1'; @@ -131,6 +136,21 @@ describe('initializeBedrock', () => { }); }); + it('should pass BEDROCK_AWS_BEARER_TOKEN as a BedrockRuntimeClient token', async () => { + process.env.BEDROCK_AWS_BEARER_TOKEN = 'test-bedrock-api-key'; + const params = createMockParams(); + const result = (await initializeBedrock(params)) as BedrockLLMConfigResult; + + expect(result.llmConfig).toHaveProperty('client'); + expect(result.llmConfig.client).toHaveProperty('_isBedrockClient', true); + expect(result.llmConfig.client).toHaveProperty('token', { + token: 'test-bedrock-api-key', + }); + expect(result.llmConfig.client).toHaveProperty('authSchemePreference', ['httpBearerAuth']); + expect(result.llmConfig).not.toHaveProperty('credentials'); + expect(result.llmConfig).not.toHaveProperty('profile'); + }); + it('should pass AWS profile to ChatBedrockConverse when static credentials are unset', async () => { delete process.env.BEDROCK_AWS_ACCESS_KEY_ID; delete process.env.BEDROCK_AWS_SECRET_ACCESS_KEY; @@ -408,6 +428,7 @@ describe('initializeBedrock', () => { describe('User-Provided Credentials', () => { it('should fetch credentials from database when user-provided', async () => { + process.env.BEDROCK_AWS_ACCESS_KEY_ID = AuthType.USER_PROVIDED; process.env.BEDROCK_AWS_SECRET_ACCESS_KEY = AuthType.USER_PROVIDED; const params = createMockParams({ body: { key: '2024-12-31T23:59:59Z' }, @@ -436,6 +457,62 @@ describe('initializeBedrock', () => { expect(mockedCheckUserKeyExpiry).toHaveBeenCalledWith(expiresAt, EModelEndpoint.bedrock); }); + + it('should fetch a user-provided Bedrock API key from database', async () => { + process.env.BEDROCK_AWS_BEARER_TOKEN = AuthType.USER_PROVIDED; + const params = createMockParams(); + (params.db.getUserKey as jest.Mock).mockResolvedValue( + JSON.stringify({ + apiKey: JSON.stringify({ + bearerToken: 'user-bedrock-api-key', + }), + }), + ); + + const result = (await initializeBedrock(params)) as BedrockLLMConfigResult; + + expect(params.db.getUserKey).toHaveBeenCalledWith({ + userId: 'test-user-id', + name: EModelEndpoint.bedrock, + }); + expect(result.llmConfig).toHaveProperty('client'); + expect(result.llmConfig.client).toHaveProperty('token', { + token: 'user-bedrock-api-key', + }); + expect(result.llmConfig.client).toHaveProperty('authSchemePreference', ['httpBearerAuth']); + expect(result.llmConfig).not.toHaveProperty('credentials'); + }); + + it('should not use stored access keys when only bearer token mode is configured', async () => { + process.env.BEDROCK_AWS_BEARER_TOKEN = AuthType.USER_PROVIDED; + const params = createMockParams(); + + await expect(initializeBedrock(params)).rejects.toThrow( + 'Bedrock credentials not provided. Please provide them again.', + ); + }); + + it('should merge user-provided access key ID with static secret access key', async () => { + process.env.BEDROCK_AWS_ACCESS_KEY_ID = AuthType.USER_PROVIDED; + process.env.BEDROCK_AWS_SECRET_ACCESS_KEY = 'static-secret-key'; + const params = createMockParams(); + (params.db.getUserKey as jest.Mock).mockResolvedValue( + JSON.stringify({ + apiKey: JSON.stringify({ + accessKeyId: 'user-access-key', + bearerToken: 'ignored-bedrock-api-key', + }), + }), + ); + + const result = (await initializeBedrock(params)) as BedrockLLMConfigResult; + + expect(result.llmConfig.credentials).toEqual({ + accessKeyId: 'user-access-key', + secretAccessKey: 'static-secret-key', + }); + expect(result.llmConfig).not.toHaveProperty('client'); + }); }); describe('Credentials Edge Cases', () => { diff --git a/packages/api/src/endpoints/bedrock/initialize.ts b/packages/api/src/endpoints/bedrock/initialize.ts index 42b57fffd7..bd19da1acb 100644 --- a/packages/api/src/endpoints/bedrock/initialize.ts +++ b/packages/api/src/endpoints/bedrock/initialize.ts @@ -2,6 +2,7 @@ import { HttpsProxyAgent } from 'https-proxy-agent'; import { NodeHttpHandler } from '@smithy/node-http-handler'; import { fromNodeProviderChain } from '@aws-sdk/credential-providers'; import { BedrockRuntimeClient } from '@aws-sdk/client-bedrock-runtime'; +import type { BedrockRuntimeClientConfig } from '@aws-sdk/client-bedrock-runtime'; import { AuthType, EModelEndpoint, @@ -14,6 +15,7 @@ import type { BaseInitializeParams, InitializeResultBase, BedrockCredentials, + BedrockUserCredentials, GuardrailConfiguration, InferenceProfileConfig, } from '~/types'; @@ -63,19 +65,34 @@ export async function initializeBedrock({ BEDROCK_AWS_ACCESS_KEY_ID, BEDROCK_AWS_SESSION_TOKEN, BEDROCK_AWS_PROFILE, + BEDROCK_AWS_BEARER_TOKEN, BEDROCK_REVERSE_PROXY, BEDROCK_AWS_DEFAULT_REGION, PROXY, } = process.env; const { key: expiresAt } = req.body; - const isUserProvided = BEDROCK_AWS_SECRET_ACCESS_KEY === AuthType.USER_PROVIDED; + const userProvidesAccessKeyId = BEDROCK_AWS_ACCESS_KEY_ID === AuthType.USER_PROVIDED; + const userProvidesSecretAccessKey = BEDROCK_AWS_SECRET_ACCESS_KEY === AuthType.USER_PROVIDED; + const userProvidesSessionToken = BEDROCK_AWS_SESSION_TOKEN === AuthType.USER_PROVIDED; + const userProvidesBearerToken = BEDROCK_AWS_BEARER_TOKEN === AuthType.USER_PROVIDED; + const isUserProvided = + userProvidesAccessKeyId || + userProvidesSecretAccessKey || + userProvidesSessionToken || + userProvidesBearerToken; + const staticAccessKeyId = userProvidesAccessKeyId ? undefined : BEDROCK_AWS_ACCESS_KEY_ID; + const staticSecretAccessKey = userProvidesSecretAccessKey + ? undefined + : BEDROCK_AWS_SECRET_ACCESS_KEY; + const staticSessionToken = userProvidesSessionToken ? undefined : BEDROCK_AWS_SESSION_TOKEN; + const staticBearerToken = userProvidesBearerToken ? undefined : BEDROCK_AWS_BEARER_TOKEN; - const hasAccessKey = BEDROCK_AWS_ACCESS_KEY_ID != null && BEDROCK_AWS_ACCESS_KEY_ID !== ''; - const hasSecretKey = - BEDROCK_AWS_SECRET_ACCESS_KEY != null && BEDROCK_AWS_SECRET_ACCESS_KEY !== ''; + const hasAccessKey = staticAccessKeyId != null && staticAccessKeyId !== ''; + const hasSecretKey = staticSecretAccessKey != null && staticSecretAccessKey !== ''; let credentials: BedrockCredentials | undefined; + let bearerToken: string | undefined; if (isUserProvided) { const userKey = await db.getUserKey({ @@ -87,20 +104,55 @@ export async function initializeBedrock({ throw new Error('Bedrock credentials not provided. Please provide them again.'); } - credentials = JSON.parse(userKey) as BedrockCredentials; + let userCredentials: BedrockUserCredentials; + try { + const storedCredentials = JSON.parse(userKey) as BedrockUserCredentials & { apiKey?: string }; + userCredentials = + typeof storedCredentials.apiKey === 'string' + ? (JSON.parse(storedCredentials.apiKey) as BedrockUserCredentials) + : storedCredentials; + } catch { + throw new Error('Bedrock credentials not provided. Please provide them again.'); + } + + if (userProvidesBearerToken && userCredentials.bearerToken) { + bearerToken = userCredentials.bearerToken; + } else { + const canUseAccessKeys = + userProvidesAccessKeyId || userProvidesSecretAccessKey || userProvidesSessionToken; + const accessKeyId = userProvidesAccessKeyId ? userCredentials.accessKeyId : staticAccessKeyId; + const secretAccessKey = userProvidesSecretAccessKey + ? userCredentials.secretAccessKey + : staticSecretAccessKey; + const sessionToken = userProvidesSessionToken + ? userCredentials.sessionToken + : staticSessionToken; + + if (!canUseAccessKeys || !accessKeyId || !secretAccessKey) { + throw new Error('Bedrock credentials not provided. Please provide them again.'); + } + + credentials = { + accessKeyId, + secretAccessKey, + ...(sessionToken && { sessionToken }), + }; + } if (expiresAt) { checkUserKeyExpiry(expiresAt, EModelEndpoint.bedrock); } + } else if (staticBearerToken) { + bearerToken = staticBearerToken; } else if (hasAccessKey !== hasSecretKey) { throw new Error( 'Both BEDROCK_AWS_ACCESS_KEY_ID and BEDROCK_AWS_SECRET_ACCESS_KEY must be provided together.', ); } else if (hasAccessKey && hasSecretKey) { credentials = { - accessKeyId: BEDROCK_AWS_ACCESS_KEY_ID, - secretAccessKey: BEDROCK_AWS_SECRET_ACCESS_KEY, - ...(BEDROCK_AWS_SESSION_TOKEN && { sessionToken: BEDROCK_AWS_SESSION_TOKEN }), + accessKeyId: staticAccessKeyId, + secretAccessKey: staticSecretAccessKey, + ...(staticSessionToken && { sessionToken: staticSessionToken }), }; } @@ -150,33 +202,48 @@ export async function initializeBedrock({ credentials.accessKeyId !== '' && typeof credentials.secretAccessKey === 'string' && credentials.secretAccessKey !== ''; + const hasBearerToken = typeof bearerToken === 'string' && bearerToken !== ''; - if (PROXY) { - const proxyAgent = new HttpsProxyAgent(PROXY); + if (PROXY || hasBearerToken) { + const proxyAgent = PROXY ? new HttpsProxyAgent(PROXY) : undefined; const credentialProvider = - !hasCompleteCredentials && BEDROCK_AWS_PROFILE + !hasCompleteCredentials && !hasBearerToken && BEDROCK_AWS_PROFILE ? fromNodeProviderChain({ profile: BEDROCK_AWS_PROFILE }) : undefined; - // Create a custom BedrockRuntimeClient with proxy-enabled request handler. + // Create a custom BedrockRuntimeClient for proxy routing or Bedrock API keys. // ChatBedrockConverse will use this pre-configured client directly instead of // creating its own. Credentials are only set if explicitly provided; otherwise // the AWS SDK's default credential provider chain is used (instance profiles, // AWS profiles, environment variables, etc.) - const customClient = new BedrockRuntimeClient({ + const customClientConfig: BedrockRuntimeClientConfig = { region: (llmConfig.region as string) ?? BEDROCK_AWS_DEFAULT_REGION, - ...(hasCompleteCredentials && { - credentials: credentials as { accessKeyId: string; secretAccessKey: string }, - }), - ...(!hasCompleteCredentials && credentialProvider && { credentials: credentialProvider }), - requestHandler: new NodeHttpHandler({ + }; + + if (hasBearerToken && bearerToken) { + customClientConfig.token = { token: bearerToken }; + customClientConfig.authSchemePreference = ['httpBearerAuth']; + } else if (hasCompleteCredentials) { + customClientConfig.credentials = credentials as { + accessKeyId: string; + secretAccessKey: string; + }; + } else if (credentialProvider) { + customClientConfig.credentials = credentialProvider; + } + + if (proxyAgent) { + customClientConfig.requestHandler = new NodeHttpHandler({ httpAgent: proxyAgent, httpsAgent: proxyAgent, - }), - ...(BEDROCK_REVERSE_PROXY && { - endpoint: `https://${BEDROCK_REVERSE_PROXY}`, - }), - }); + }); + } + + if (BEDROCK_REVERSE_PROXY) { + customClientConfig.endpoint = `https://${BEDROCK_REVERSE_PROXY}`; + } + + const customClient = new BedrockRuntimeClient(customClientConfig); llmConfig.client = customClient; } else { diff --git a/packages/api/src/endpoints/config/endpoints.spec.ts b/packages/api/src/endpoints/config/endpoints.spec.ts index 8d77c82960..10feda4034 100644 --- a/packages/api/src/endpoints/config/endpoints.spec.ts +++ b/packages/api/src/endpoints/config/endpoints.spec.ts @@ -1,5 +1,6 @@ import { Types } from 'mongoose'; import { + AuthType, AgentCapabilities, EModelEndpoint, PrincipalType, @@ -37,9 +38,7 @@ function createAppConfigCache() { }; } -type ConfigPrincipals = NonNullable< - Parameters[0] ->; +type ConfigPrincipals = NonNullable[0]>; function createMockDeps(overrides: Partial = {}): EndpointsConfigDeps { return { @@ -198,6 +197,47 @@ describe('createEndpointsConfigService', () => { ]); }); + it('exposes Bedrock user-provided credential options', async () => { + const previousEnv = { + BEDROCK_AWS_ACCESS_KEY_ID: process.env.BEDROCK_AWS_ACCESS_KEY_ID, + BEDROCK_AWS_SECRET_ACCESS_KEY: process.env.BEDROCK_AWS_SECRET_ACCESS_KEY, + BEDROCK_AWS_SESSION_TOKEN: process.env.BEDROCK_AWS_SESSION_TOKEN, + BEDROCK_AWS_BEARER_TOKEN: process.env.BEDROCK_AWS_BEARER_TOKEN, + }; + + process.env.BEDROCK_AWS_ACCESS_KEY_ID = AuthType.USER_PROVIDED; + process.env.BEDROCK_AWS_SECRET_ACCESS_KEY = AuthType.USER_PROVIDED; + process.env.BEDROCK_AWS_SESSION_TOKEN = AuthType.USER_PROVIDED; + process.env.BEDROCK_AWS_BEARER_TOKEN = AuthType.USER_PROVIDED; + + try { + const deps = createMockDeps({ + loadDefaultEndpointsConfig: jest.fn().mockResolvedValue({ + [EModelEndpoint.bedrock]: { userProvide: false, order: 0 }, + }), + }); + const { getEndpointsConfig } = createEndpointsConfigService(deps); + const result = await getEndpointsConfig(fakeReq()); + + expect(result?.[EModelEndpoint.bedrock]).toEqual( + expect.objectContaining({ + userProvideAccessKeyId: true, + userProvideSecretAccessKey: true, + userProvideSessionToken: true, + userProvideBearerToken: true, + }), + ); + } finally { + Object.entries(previousEnv).forEach(([key, value]) => { + if (value == null) { + delete process.env[key]; + } else { + process.env[key] = value; + } + }); + } + }); + it('uses req.config when available instead of calling getAppConfig', async () => { const mockGetAppConfig = jest.fn(); const deps = createMockDeps({ getAppConfig: mockGetAppConfig }); @@ -213,9 +253,7 @@ describe('createEndpointsConfigService', () => { const deps = createMockDeps({ getAppConfig: mockGetAppConfig }); const { getEndpointsConfig } = createEndpointsConfigService(deps); - await getEndpointsConfig( - fakeReq({ user: { id: 'u1', role: 'USER', tenantId: 'tenant-a' } }), - ); + await getEndpointsConfig(fakeReq({ user: { id: 'u1', role: 'USER', tenantId: 'tenant-a' } })); expect(mockGetAppConfig).toHaveBeenCalledWith({ role: 'USER', diff --git a/packages/api/src/endpoints/config/endpoints.ts b/packages/api/src/endpoints/config/endpoints.ts index 0d500219d9..4bdd999a13 100644 --- a/packages/api/src/endpoints/config/endpoints.ts +++ b/packages/api/src/endpoints/config/endpoints.ts @@ -1,4 +1,5 @@ import { + AuthType, EModelEndpoint, isAgentsEndpoint, orderEndpointsConfig, @@ -9,7 +10,7 @@ import type { AgentCapabilities, TEndpointsConfig, TConfig } from 'librechat-dat import type { ServerRequest, TCustomEndpointsConfig } from '~/types'; import { loadCustomEndpointsConfig as defaultLoadCustomEndpoints } from '~/endpoints/custom'; -type PartialEndpointEntry = Partial; +type PartialEndpointEntry = Partial & Record; type DefaultEndpointsResult = Record; type MutableEndpointsConfig = Record; @@ -109,6 +110,17 @@ export function createEndpointsConfigService(deps: EndpointsConfigDeps) { }; } + if (mergedConfig[EModelEndpoint.bedrock]) { + mergedConfig[EModelEndpoint.bedrock] = { + ...mergedConfig[EModelEndpoint.bedrock], + userProvideAccessKeyId: process.env.BEDROCK_AWS_ACCESS_KEY_ID === AuthType.USER_PROVIDED, + userProvideSecretAccessKey: + process.env.BEDROCK_AWS_SECRET_ACCESS_KEY === AuthType.USER_PROVIDED, + userProvideSessionToken: process.env.BEDROCK_AWS_SESSION_TOKEN === AuthType.USER_PROVIDED, + userProvideBearerToken: process.env.BEDROCK_AWS_BEARER_TOKEN === AuthType.USER_PROVIDED, + }; + } + return orderEndpointsConfig(mergedConfig as TEndpointsConfig); } diff --git a/packages/api/src/types/bedrock.ts b/packages/api/src/types/bedrock.ts index e092e6bc04..c60658e33d 100644 --- a/packages/api/src/types/bedrock.ts +++ b/packages/api/src/types/bedrock.ts @@ -8,6 +8,13 @@ import type { BedrockConverseInput } from 'librechat-data-provider'; */ export type BedrockCredentials = Partial; +/** + * User-provided Bedrock credentials can be either AWS credentials or an API key. + */ +export type BedrockUserCredentials = BedrockCredentials & { + bearerToken?: string; +}; + /** * AWS Bedrock Guardrail configuration * @see https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_GuardrailConfiguration.html diff --git a/packages/data-provider/src/types.ts b/packages/data-provider/src/types.ts index 0cfb826819..a6a6be1433 100644 --- a/packages/data-provider/src/types.ts +++ b/packages/data-provider/src/types.ts @@ -388,6 +388,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[];