diff --git a/packages/data-provider/src/config.spec.ts b/packages/data-provider/src/config.spec.ts index f5c107c58f..ae7e7ca526 100644 --- a/packages/data-provider/src/config.spec.ts +++ b/packages/data-provider/src/config.spec.ts @@ -29,6 +29,34 @@ describe('excludedKeys', () => { }); }); +describe('bedrockEndpointSchema', () => { + it('preserves guardrailConfig from configSchema parsing', () => { + const guardrailConfig = { + guardrailIdentifier: '${BEDROCK_GUARDRAIL_ID}', + guardrailVersion: '${BEDROCK_GUARDRAIL_VERSION}', + trace: 'enabled_full', + streamProcessingMode: 'sync', + }; + + const result = configSchema.safeParse({ + version: '1.0', + endpoints: { + bedrock: { + streamRate: 25, + availableRegions: ['us-west-2'], + guardrailConfig, + }, + }, + }); + + expect(result.success).toBe(true); + if (!result.success) { + return; + } + expect(result.data.endpoints?.bedrock?.guardrailConfig).toEqual(guardrailConfig); + }); +}); + describe('resolveEndpointType', () => { describe('non-agents endpoints', () => { it('returns the config type for a custom endpoint', () => { diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index b0d2217968..cb48a15a64 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -403,10 +403,18 @@ export const baseEndpointSchema = z.object({ export type TBaseEndpoint = z.infer; +export const bedrockGuardrailConfigSchema = z.object({ + guardrailIdentifier: z.string(), + guardrailVersion: z.string(), + trace: z.enum(['enabled', 'disabled', 'enabled_full']).optional(), + streamProcessingMode: z.enum(['sync', 'async']).optional(), +}); + export const bedrockEndpointSchema = baseEndpointSchema.merge( z.object({ availableRegions: z.array(z.string()).optional(), models: z.array(z.string()).optional(), + guardrailConfig: bedrockGuardrailConfigSchema.optional(), inferenceProfiles: z.record(z.string(), z.string()).optional(), }), );