Merge pull request #851 from Portkey-AI/chore/shorthand-guardrails-up…

…dates Chore: shorthand guardrails updates
Portkey-AI · Dec 27, 2024 · 24ead1b · 24ead1b
2 parents 66e9213 + 45d0f42
commit 24ead1b
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 11 deletions.
diff --git 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
@@ -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
       );
     },
     {