diff --git a/src/handlers/handlerUtils.ts b/src/handlers/handlerUtils.ts index 21013dcac..76e447bbf 100644 --- a/src/handlers/handlerUtils.ts +++ b/src/handlers/handlerUtils.ts @@ -197,18 +197,26 @@ export function convertGuardrailsShorthand(guardrailsArr: any, type: string) { }; // if the deny key is present (true or false), add it to hooksObject and remove it from guardrails - ['deny', 'on_fail', 'on_success', 'async', 'onFail', 'onSuccess'].forEach( - (key) => { - if (guardrails.hasOwnProperty(key)) { - hooksObject[key] = guardrails[key]; - delete guardrails[key]; - } + [ + 'deny', + 'on_fail', + 'on_success', + 'async', + 'id', + 'type', + 'guardrail_version_id', + ].forEach((key) => { + if (guardrails.hasOwnProperty(key)) { + hooksObject[key] = guardrails[key]; + delete guardrails[key]; } - ); + }); + + hooksObject = convertKeysToCamelCase(hooksObject); // Now, add all the checks to the checks array hooksObject.checks = Object.keys(guardrails).map((key) => ({ - id: key, + id: key.includes('.') ? key : `default.${key}`, parameters: guardrails[key], })); @@ -983,6 +991,8 @@ export function constructConfigFromRequestHeaders( 'checks', 'vertex_service_account_json', 'conditions', + 'input_guardrails', + 'output_guardrails', ]) as any; } diff --git a/src/middlewares/requestValidator/schema/config.ts b/src/middlewares/requestValidator/schema/config.ts index 3ae25f617..0c7d676de 100644 --- a/src/middlewares/requestValidator/schema/config.ts +++ b/src/middlewares/requestValidator/schema/config.ts @@ -79,8 +79,28 @@ export const configSchema: any = z // Google Vertex AI specific vertex_project_id: z.string().optional(), vertex_region: z.string().optional(), - after_request_hooks: z.any().optional(), - before_request_hooks: z.any().optional(), + after_request_hooks: z + .array(z.object({}).catchall(z.any())) // Allows any object structure + .optional(), + before_request_hooks: z + .array(z.object({}).catchall(z.any())) // Allows any object structure + .optional(), + input_guardrails: z + .union([ + z.array(z.string()), + z.array( + z.object({}).catchall(z.any()) // Allows any object structure + ), + ]) + .optional(), + output_guardrails: z + .union([ + z.array(z.string()), + z.array( + z.object({}).catchall(z.any()) // Allows any object structure + ), + ]) + .optional(), vertex_service_account_json: z.object({}).catchall(z.string()).optional(), // OpenAI specific openai_project: z.string().optional(), @@ -115,7 +135,9 @@ export const configSchema: any = z hasAWSDetails || isVertexAIProvider || value.after_request_hooks || - value.before_request_hooks + value.before_request_hooks || + value.input_guardrails || + value.output_guardrails ); }, {