diff --git a/librechat.example.yaml b/librechat.example.yaml index 983a5b611d..a7ab9ee054 100644 --- a/librechat.example.yaml +++ b/librechat.example.yaml @@ -427,6 +427,22 @@ endpoints: forcePrompt: false modelDisplayLabel: 'Portkey' iconURL: https://images.crunchbase.com/image/upload/c_pad,f_auto,q_auto:eco,dpr_1/rjqy7ghvjoiu4cd1xjbf + + # AWS Bedrock Example + # Note: Bedrock endpoint is configured via environment variables + # bedrock: + # # Guardrail Configuration + # guardrailConfig: + # guardrailIdentifier: "your-guardrail-id" + # guardrailVersion: "1" + # + # # Trace behavior for debugging (optional) + # # - "enabled": Include basic trace information about guardrail assessments + # # - "enabled_full": Include comprehensive trace details (recommended for debugging) + # # - "disabled": No trace information (default) + # # Trace output is logged to application log files for compliance auditing + # trace: "enabled" + # Example modelSpecs configuration showing grouping options # The 'group' field organizes model specs in the UI selector: # - If 'group' matches an endpoint name (e.g., "openAI", "groq"), the spec appears nested under that endpoint diff --git a/packages/api/src/endpoints/bedrock/initialize.spec.ts b/packages/api/src/endpoints/bedrock/initialize.spec.ts new file mode 100644 index 0000000000..9b0ba152d7 --- /dev/null +++ b/packages/api/src/endpoints/bedrock/initialize.spec.ts @@ -0,0 +1,316 @@ +import { AuthType, EModelEndpoint } from 'librechat-data-provider'; +import { initializeBedrock } from './initialize'; +import type { BaseInitializeParams, BedrockLLMConfigResult } from '~/types'; +import { checkUserKeyExpiry } from '~/utils'; + +jest.mock('https-proxy-agent', () => ({ + HttpsProxyAgent: jest.fn().mockImplementation((proxy) => ({ proxy })), +})); + +jest.mock('@smithy/node-http-handler', () => ({ + NodeHttpHandler: jest.fn().mockImplementation((options) => ({ ...options })), +})); + +jest.mock('@aws-sdk/client-bedrock-runtime', () => ({ + BedrockRuntimeClient: jest.fn().mockImplementation((config) => ({ + ...config, + _isBedrockClient: true, + })), +})); + +jest.mock('~/utils', () => ({ + checkUserKeyExpiry: jest.fn(), +})); + +const mockedCheckUserKeyExpiry = jest.mocked(checkUserKeyExpiry); + +const createMockParams = ( + overrides: Partial<{ + config: Record; + body: Record; + user: { id: string }; + model_parameters: Record; + env: Record; + }> = {}, +): BaseInitializeParams => { + const mockDb = { + getUserKey: jest.fn().mockResolvedValue( + JSON.stringify({ + accessKeyId: 'user-access-key', + secretAccessKey: 'user-secret-key', + }), + ), + }; + + return { + req: { + config: overrides.config ?? {}, + body: overrides.body ?? {}, + user: overrides.user ?? { id: 'test-user-id' }, + }, + endpoint: EModelEndpoint.bedrock, + model_parameters: overrides.model_parameters ?? { model: 'anthropic.claude-3-sonnet' }, + db: mockDb, + } as unknown as BaseInitializeParams; +}; + +describe('initializeBedrock', () => { + const originalEnv = process.env; + + beforeEach(() => { + jest.clearAllMocks(); + process.env = { ...originalEnv }; + 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'; + }); + + afterEach(() => { + process.env = originalEnv; + }); + + describe('Basic Configuration', () => { + it('should create a basic configuration with credentials from environment', async () => { + const params = createMockParams(); + const result = (await initializeBedrock(params)) as BedrockLLMConfigResult; + + expect(result.llmConfig).toHaveProperty('credentials'); + expect(result.llmConfig.credentials).toEqual({ + accessKeyId: 'test-access-key', + secretAccessKey: 'test-secret-key', + }); + expect(result.llmConfig).toHaveProperty('model', 'anthropic.claude-3-sonnet'); + }); + + it('should include region from environment', async () => { + const params = createMockParams(); + const result = await initializeBedrock(params); + + expect(result.llmConfig).toHaveProperty('region', 'us-east-1'); + }); + + it('should handle model_parameters', async () => { + const params = createMockParams({ + model_parameters: { + model: 'anthropic.claude-3-opus', + temperature: 0.7, + maxTokens: 4096, + }, + }); + const result = await initializeBedrock(params); + + expect(result.llmConfig).toHaveProperty('model', 'anthropic.claude-3-opus'); + expect(result.llmConfig).toHaveProperty('temperature', 0.7); + expect(result.llmConfig).toHaveProperty('maxTokens', 4096); + }); + + it('should handle session token when provided', async () => { + process.env.BEDROCK_AWS_SESSION_TOKEN = 'test-session-token'; + const params = createMockParams(); + const result = (await initializeBedrock(params)) as BedrockLLMConfigResult; + + expect(result.llmConfig.credentials).toEqual({ + accessKeyId: 'test-access-key', + secretAccessKey: 'test-secret-key', + sessionToken: 'test-session-token', + }); + }); + }); + + describe('GuardrailConfig', () => { + it('should apply guardrailConfig from backend config', async () => { + const guardrailConfig = { + guardrailIdentifier: 'test-guardrail-id', + guardrailVersion: '1', + trace: 'enabled' as const, + }; + + const params = createMockParams({ + config: { + endpoints: { + [EModelEndpoint.bedrock]: { + guardrailConfig, + }, + }, + }, + }); + + const result = (await initializeBedrock(params)) as BedrockLLMConfigResult; + + expect(result.llmConfig).toHaveProperty('guardrailConfig'); + expect(result.llmConfig.guardrailConfig).toEqual(guardrailConfig); + }); + + it('should NOT include guardrailConfig when not configured', async () => { + const params = createMockParams({ + config: { + endpoints: { + [EModelEndpoint.bedrock]: {}, + }, + }, + }); + + const result = await initializeBedrock(params); + + expect(result.llmConfig).not.toHaveProperty('guardrailConfig'); + }); + + it('should apply guardrailConfig regardless of model_parameters', async () => { + const guardrailConfig = { + guardrailIdentifier: 'admin-guardrail', + guardrailVersion: 'DRAFT', + }; + + const params = createMockParams({ + config: { + endpoints: { + [EModelEndpoint.bedrock]: { + guardrailConfig, + }, + }, + }, + model_parameters: { + model: 'anthropic.claude-3-sonnet', + temperature: 0.5, + }, + }); + + const result = (await initializeBedrock(params)) as BedrockLLMConfigResult; + + expect(result.llmConfig.guardrailConfig).toEqual(guardrailConfig); + expect(result.llmConfig).toHaveProperty('temperature', 0.5); + }); + + it('should handle guardrailConfig with enabled_full trace', async () => { + const guardrailConfig = { + guardrailIdentifier: 'compliance-guardrail', + guardrailVersion: '2', + trace: 'enabled_full' as const, + }; + + const params = createMockParams({ + config: { + endpoints: { + [EModelEndpoint.bedrock]: { + guardrailConfig, + }, + }, + }, + }); + + const result = (await initializeBedrock(params)) as BedrockLLMConfigResult; + + expect(result.llmConfig.guardrailConfig).toEqual(guardrailConfig); + expect(result.llmConfig.guardrailConfig?.trace).toBe('enabled_full'); + }); + }); + + describe('Proxy Configuration', () => { + it('should create BedrockRuntimeClient with proxy when PROXY is set', async () => { + process.env.PROXY = 'http://proxy:8080'; + 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).not.toHaveProperty('credentials'); + }); + + it('should include reverse proxy endpoint when BEDROCK_REVERSE_PROXY is set with PROXY', async () => { + process.env.PROXY = 'http://proxy:8080'; + process.env.BEDROCK_REVERSE_PROXY = 'custom-bedrock-endpoint.com'; + const params = createMockParams(); + const result = (await initializeBedrock(params)) as BedrockLLMConfigResult; + + expect(result.llmConfig).toHaveProperty('client'); + expect(result.llmConfig.client).toHaveProperty( + 'endpoint', + 'https://custom-bedrock-endpoint.com', + ); + }); + }); + + describe('Reverse Proxy Configuration', () => { + it('should set endpointHost when BEDROCK_REVERSE_PROXY is set without PROXY', async () => { + process.env.BEDROCK_REVERSE_PROXY = 'reverse-proxy.example.com'; + const params = createMockParams(); + const result = await initializeBedrock(params); + + expect(result.llmConfig).toHaveProperty('endpointHost', 'reverse-proxy.example.com'); + expect(result.llmConfig).not.toHaveProperty('client'); + }); + }); + + describe('User-Provided Credentials', () => { + it('should fetch credentials from database when user-provided', async () => { + process.env.BEDROCK_AWS_SECRET_ACCESS_KEY = AuthType.USER_PROVIDED; + const params = createMockParams({ + body: { key: '2024-12-31T23:59:59Z' }, + }); + + const result = (await initializeBedrock(params)) as BedrockLLMConfigResult; + + expect(params.db.getUserKey).toHaveBeenCalledWith({ + userId: 'test-user-id', + name: EModelEndpoint.bedrock, + }); + expect(result.llmConfig.credentials).toEqual({ + accessKeyId: 'user-access-key', + secretAccessKey: 'user-secret-key', + }); + }); + + it('should check key expiry for user-provided credentials', async () => { + process.env.BEDROCK_AWS_SECRET_ACCESS_KEY = AuthType.USER_PROVIDED; + const expiresAt = '2024-12-31T23:59:59Z'; + const params = createMockParams({ + body: { key: expiresAt }, + }); + + await initializeBedrock(params); + + expect(mockedCheckUserKeyExpiry).toHaveBeenCalledWith(expiresAt, EModelEndpoint.bedrock); + }); + }); + + describe('Credentials Edge Cases', () => { + it('should set credentials to undefined when access key and secret are empty', async () => { + process.env.BEDROCK_AWS_ACCESS_KEY_ID = ''; + process.env.BEDROCK_AWS_SECRET_ACCESS_KEY = ''; + const params = createMockParams(); + const result = (await initializeBedrock(params)) as BedrockLLMConfigResult; + + expect(result.llmConfig.credentials).toBeUndefined(); + }); + + it('should set credentials to undefined when access key and secret are undefined', async () => { + delete process.env.BEDROCK_AWS_ACCESS_KEY_ID; + delete process.env.BEDROCK_AWS_SECRET_ACCESS_KEY; + const params = createMockParams(); + const result = (await initializeBedrock(params)) as BedrockLLMConfigResult; + + expect(result.llmConfig.credentials).toBeUndefined(); + }); + + it('should throw error when user-provided credentials are not found', async () => { + process.env.BEDROCK_AWS_SECRET_ACCESS_KEY = AuthType.USER_PROVIDED; + const params = createMockParams(); + (params.db.getUserKey as jest.Mock).mockResolvedValue(null); + + await expect(initializeBedrock(params)).rejects.toThrow( + 'Bedrock credentials not provided. Please provide them again.', + ); + }); + }); + + describe('Return Structure', () => { + it('should return llmConfig and configOptions', async () => { + const params = createMockParams(); + const result = await initializeBedrock(params); + + expect(result).toHaveProperty('llmConfig'); + expect(result).toHaveProperty('configOptions'); + expect(typeof result.configOptions).toBe('object'); + }); + }); +}); diff --git a/packages/api/src/endpoints/bedrock/initialize.ts b/packages/api/src/endpoints/bedrock/initialize.ts index 500285e02f..9f5db47bf7 100644 --- a/packages/api/src/endpoints/bedrock/initialize.ts +++ b/packages/api/src/endpoints/bedrock/initialize.ts @@ -8,7 +8,12 @@ import { bedrockOutputParser, removeNullishValues, } from 'librechat-data-provider'; -import type { BaseInitializeParams, InitializeResultBase, BedrockCredentials } from '~/types'; +import type { + BaseInitializeParams, + InitializeResultBase, + BedrockCredentials, + GuardrailConfiguration, +} from '~/types'; import { checkUserKeyExpiry } from '~/utils'; /** @@ -42,6 +47,11 @@ export async function initializeBedrock({ db, }: BaseInitializeParams): Promise { void endpoint; + const appConfig = req.config; + const bedrockConfig = appConfig?.endpoints?.[EModelEndpoint.bedrock] as + | ({ guardrailConfig?: GuardrailConfiguration } & Record) + | undefined; + const { BEDROCK_AWS_SECRET_ACCESS_KEY, BEDROCK_AWS_ACCESS_KEY_ID, @@ -89,15 +99,23 @@ export async function initializeBedrock({ const llmConfig = bedrockOutputParser( bedrockInputParser.parse( - removeNullishValues({ ...requestOptions, ...(model_parameters ?? {}) }), + removeNullishValues({ + ...requestOptions, + ...(model_parameters ?? {}), + }), ), ) as InitializeResultBase['llmConfig'] & { region?: string; client?: BedrockRuntimeClient; credentials?: BedrockCredentials; endpointHost?: string; + guardrailConfig?: GuardrailConfiguration; }; + if (bedrockConfig?.guardrailConfig) { + llmConfig.guardrailConfig = bedrockConfig.guardrailConfig; + } + /** Only include credentials if they're complete (accessKeyId and secretAccessKey are both set) */ const hasCompleteCredentials = credentials && diff --git a/packages/api/src/types/bedrock.ts b/packages/api/src/types/bedrock.ts index a2e68cf9ac..22c8464619 100644 --- a/packages/api/src/types/bedrock.ts +++ b/packages/api/src/types/bedrock.ts @@ -8,6 +8,19 @@ import type { BedrockConverseInput } from 'librechat-data-provider'; */ export type BedrockCredentials = Partial; +/** + * AWS Bedrock Guardrail configuration + * @see https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_GuardrailConfiguration.html + */ +export interface GuardrailConfiguration { + /** The identifier for the guardrail (ID or ARN) */ + guardrailIdentifier: string; + /** The version of the guardrail (version number or "DRAFT") */ + guardrailVersion: string; + /** The trace behavior for the guardrail */ + trace?: 'enabled' | 'disabled' | 'enabled_full'; +} + /** * Configuration options for Bedrock LLM */ @@ -21,6 +34,8 @@ export interface BedrockConfigOptions { credentials?: BedrockCredentials; /** Custom endpoint host for reverse proxy */ endpointHost?: string; + /** Guardrail configuration for content filtering */ + guardrailConfig?: GuardrailConfiguration; } /** @@ -32,6 +47,7 @@ export interface BedrockLLMConfigResult { client?: BedrockRuntimeClient; credentials?: BedrockCredentials; endpointHost?: string; + guardrailConfig?: GuardrailConfiguration; }; configOptions: Record; }