Merge pull request #904 from Portkey-AI/chore/return-transformed-flag…

…-in-redaction-check-result Chore: return transformed flag in redaction check result
Portkey-AI · Jan 30, 2025 · a10d0ca · a10d0ca
2 parents bb5859d + 3cbcad2
commit a10d0ca
Show file tree

Hide file tree

Showing 15 changed files with 101 additions and 29 deletions.
diff --git a/.gitignore b/.gitignore
@@ -141,3 +141,4 @@ build
 .idea
 plugins/**/.creds.json
 plugins/**/creds.json
+plugins/**/.parameters.json
diff --git a/plugins/bedrock/bedrock.test.ts b/plugins/bedrock/bedrock.test.ts
@@ -2,15 +2,16 @@ import { HookEventType, PluginContext, PluginParameters } from '../types';
 import { pluginHandler } from './index';
 import testCreds from './.creds.json';
 import { BedrockParameters } from './type';
+import parametersCreds from './.parameters.json';
 
 /**
  * @example Parameters object
  * 
  * {
     "credentials": {
-      "accessKeyId": "keyId",
-      "accessKeySecret": "keysecret",
-      "region": "us-east-1"
+      "awsAccessKeyId": "keyId",
+      "awsSecretAccessKey": "keysecret",
+      "awsRegion": "us-east-1"
     },
     "guardrailId": "xyxyxyx",
     "guardrailVersion": "1"
@@ -53,6 +54,7 @@ describe('Credentials check', () => {
     expect(result.verdict).toBe(true);
     expect(result.error).toBeDefined();
     expect(result.data).toBeNull();
+    expect(result.transformed).toBe(false);
   });
 
   test('Should fail with wrong creds', async () => {
@@ -91,6 +93,7 @@ describe('Credentials check', () => {
     expect(result.verdict).toBe(true);
     expect(result.error).toBeDefined();
     expect(result.data).toBeNull();
+    expect(result.transformed).toBe(false);
   });
 
   it('should only detect PII', async () => {
@@ -110,9 +113,9 @@ describe('Credentials check', () => {
       requestType: 'chatComplete',
     };
     const parameters = {
-      credentials: testCreds,
-      guardrailId: testCreds.guardrailId,
-      guardrailVersion: testCreds.guardrailVersion,
+      credentials: testCreds as BedrockParameters['credentials'],
+      guardrailId: parametersCreds.guardrailId,
+      guardrailVersion: parametersCreds.guardrailVersion,
     };
 
     const result = await pluginHandler(
@@ -128,6 +131,7 @@ describe('Credentials check', () => {
     expect(result.error).toBeNull();
     expect(result.data).toBeDefined();
     expect(result.transformedData?.request?.json).toBeNull();
+    expect(result.transformed).toBe(false);
   });
 
   it('should detect and redact PII in request text', async () => {
@@ -146,10 +150,10 @@ describe('Credentials check', () => {
       requestType: 'chatComplete',
     };
     const parameters = {
-      credentials: testCreds,
+      credentials: testCreds as BedrockParameters['credentials'],
       redact: true,
-      guardrailId: testCreds.guardrailId,
-      guardrailVersion: testCreds.guardrailVersion,
+      guardrailId: parametersCreds.guardrailId,
+      guardrailVersion: parametersCreds.guardrailVersion,
     };
 
     const result = await pluginHandler(
@@ -166,6 +170,7 @@ describe('Credentials check', () => {
     expect(result.transformedData?.request?.json?.messages?.[0]?.content).toBe(
       'My SSN is {US_SOCIAL_SECURITY_NUMBER} and some random text'
     );
+    expect(result.transformed).toBe(true);
   });
 
   it('should detect and redact PII in request text with multiple content parts', async () => {
@@ -193,10 +198,10 @@ describe('Credentials check', () => {
       requestType: 'chatComplete',
     };
     const parameters = {
-      credentials: testCreds,
+      credentials: testCreds as BedrockParameters['credentials'],
       redact: true,
-      guardrailId: testCreds.guardrailId,
-      guardrailVersion: testCreds.guardrailVersion,
+      guardrailId: parametersCreds.guardrailId,
+      guardrailVersion: parametersCreds.guardrailVersion,
     };
 
     const result = await pluginHandler(
@@ -217,6 +222,7 @@ describe('Credentials check', () => {
     expect(
       result.transformedData?.request?.json?.messages?.[0]?.content?.[1]?.text
     ).toBe('My SSN is {US_SOCIAL_SECURITY_NUMBER} and some random text');
+    expect(result.transformed).toBe(true);
   });
 
   it('should detect and redact PII in response text', async () => {
@@ -237,10 +243,10 @@ describe('Credentials check', () => {
       requestType: 'chatComplete',
     };
     const parameters = {
-      credentials: testCreds,
+      credentials: testCreds as BedrockParameters['credentials'],
       redact: true,
-      guardrailId: testCreds.guardrailId,
-      guardrailVersion: testCreds.guardrailVersion,
+      guardrailId: parametersCreds.guardrailId,
+      guardrailVersion: parametersCreds.guardrailVersion,
     };
 
     const result = await pluginHandler(
@@ -258,6 +264,7 @@ describe('Credentials check', () => {
     expect(
       result.transformedData?.response?.json?.choices?.[0]?.message?.content
     ).toBe('My SSN is {US_SOCIAL_SECURITY_NUMBER} and some random text');
+    expect(result.transformed).toBe(true);
   });
 
   it('should pass text without PII', async () => {
@@ -279,9 +286,9 @@ describe('Credentials check', () => {
       requestType: 'chatComplete',
     };
     const parameters = {
-      credentials: testCreds,
-      guardrailId: testCreds.guardrailId,
-      guardrailVersion: testCreds.guardrailVersion,
+      credentials: testCreds as BedrockParameters['credentials'],
+      guardrailId: parametersCreds.guardrailId,
+      guardrailVersion: parametersCreds.guardrailVersion,
     };
 
     const result = await pluginHandler(
@@ -297,5 +304,6 @@ describe('Credentials check', () => {
     expect(result.error).toBeNull();
     expect(result.data).toBeDefined();
     expect(result.transformedData?.response?.json).toBeNull();
+    expect(result.transformed).toBe(false);
   });
 });
diff --git a/plugins/bedrock/index.ts b/plugins/bedrock/index.ts
@@ -32,6 +32,7 @@ export const pluginHandler: PluginHandler<
       json: null,
     },
   };
+  let transformed = false;
   const credentials = parameters.credentials;
 
   const validate = validateCreds(credentials);
@@ -48,6 +49,8 @@ export const pluginHandler: PluginHandler<
       verdict,
       error: { message: 'Missing required credentials' },
       data,
+      transformed,
+      transformedData,
     };
   }
 
@@ -68,6 +71,7 @@ export const pluginHandler: PluginHandler<
         verdict: true,
         data: null,
         transformedData,
+        transformed,
       };
     }
 
@@ -118,6 +122,7 @@ export const pluginHandler: PluginHandler<
       );
 
       setCurrentContentPart(context, eventType, transformedData, maskedTexts);
+      transformed = true;
     }
 
     if (hasPii && flaggedCategories.size === 1 && redact) {
@@ -138,5 +143,6 @@ export const pluginHandler: PluginHandler<
     error,
     data,
     transformedData,
+    transformed,
   };
 };
diff --git a/plugins/pangea/pangea.test.ts b/plugins/pangea/pangea.test.ts
@@ -235,6 +235,7 @@ describe('pii handler', () => {
     expect(result.error).toBeNull();
     expect(result.data).toBeDefined();
     expect(result.transformedData?.request?.json).toBeNull();
+    expect(result.transformed).toBe(false);
   });
 
   it('should detect and redact PII in request text', async () => {
@@ -271,6 +272,7 @@ describe('pii handler', () => {
     expect(result.transformedData?.request?.json?.messages?.[0]?.content).toBe(
       'My email is <EMAIL_ADDRESS> and some random text'
     );
+    expect(result.transformed).toBe(true);
   });
 
   it('should detect and redact PII in request text with multiple content parts', async () => {
@@ -320,6 +322,7 @@ describe('pii handler', () => {
     expect(
       result.transformedData?.request?.json?.messages?.[0]?.content?.[1]?.text
     ).toBe('My email is <EMAIL_ADDRESS> and some random text');
+    expect(result.transformed).toBe(true);
   });
 
   it('should detect and redact PII in response text', async () => {
@@ -359,6 +362,7 @@ describe('pii handler', () => {
     expect(
       result.transformedData?.response?.json?.choices?.[0]?.message?.content
     ).toBe('My email is <EMAIL_ADDRESS> and some random text');
+    expect(result.transformed).toBe(true);
   });
 
   it('should pass text without PII', async () => {
@@ -393,5 +397,6 @@ describe('pii handler', () => {
     expect(result.verdict).toBe(true);
     expect(result.error).toBeNull();
     expect(result.data).toBeDefined();
+    expect(result.transformed).toBe(false);
   });
 });
diff --git a/plugins/pangea/pii.ts b/plugins/pangea/pii.ts
@@ -20,6 +20,7 @@ export const handler: PluginHandler = async (
       json: null,
     },
   };
+  let transformed = false;
   const redact = parameters.redact || false;
 
   try {
@@ -29,6 +30,7 @@ export const handler: PluginHandler = async (
         verdict: true,
         data: null,
         transformedData,
+        transformed,
       };
     }
 
@@ -37,6 +39,8 @@ export const handler: PluginHandler = async (
         error: `'parameters.credentials.domain' must be set`,
         verdict: true,
         data: null,
+        transformedData,
+        transformed,
       };
     }
 
@@ -45,6 +49,8 @@ export const handler: PluginHandler = async (
         error: `'parameters.credentials.apiKey' must be set`,
         verdict: true,
         data: null,
+        transformedData,
+        transformed,
       };
     }
 
@@ -58,6 +64,7 @@ export const handler: PluginHandler = async (
         verdict: true,
         data: null,
         transformedData,
+        transformed,
       };
     }
 
@@ -87,6 +94,7 @@ export const handler: PluginHandler = async (
         response.result.redacted_data
       );
       shouldBlock = false;
+      transformed = true;
     }
 
     return {
@@ -96,13 +104,15 @@ export const handler: PluginHandler = async (
         summary: response.summary,
       },
       transformedData,
+      transformed,
     };
   } catch (e) {
     return {
       error: e as Error,
       verdict: true,
       data: null,
       transformedData,
+      transformed,
     };
   }
 };
diff --git a/plugins/patronus/patronus.test.ts b/plugins/patronus/patronus.test.ts
@@ -4,7 +4,7 @@ import { handler as piiHandler } from './pii';
 import { handler as toxicityHandler } from './toxicity';
 import { handler as retrievalAnswerRelevanceHandler } from './retrievalAnswerRelevance';
 import { handler as customHandler } from './custom';
-import { HookEventType, PluginContext } from '../types';
+import { PluginContext } from '../types';
 
 describe('phi handler', () => {
   it('should pass when text is clean', async () => {
@@ -35,6 +35,7 @@ describe('phi handler', () => {
     expect(result.verdict).toBe(true);
     expect(result.error).toBeNull();
     expect(result.data).toBeDefined();
+    expect(result.transformed).toBe(false);
   });
 
   it('should fail when text contains PHI', async () => {
@@ -67,6 +68,7 @@ describe('phi handler', () => {
     expect(result.data).toBeDefined();
     expect(result.transformedData?.response?.json).toBeNull();
     expect(result.transformedData?.request?.json).toBeNull();
+    expect(result.transformed).toBe(false);
   });
 
   it('should detect and redact PII in request text', async () => {
@@ -104,6 +106,7 @@ describe('phi handler', () => {
     expect(result.transformedData?.request?.json?.messages?.[0]?.content).toBe(
       'J******e has a history of heart disease'
     );
+    expect(result.transformed).toBe(true);
   });
 
   it('should detect and redact PII in request text with multiple content parts', async () => {
@@ -153,6 +156,7 @@ describe('phi handler', () => {
     expect(
       result.transformedData?.request?.json?.messages?.[0]?.content?.[1]?.text
     ).toBe('J******e has a history of heart disease and some random text');
+    expect(result.transformed).toBe(true);
   });
 
   it('should detect and redact PHI in response text', async () => {
@@ -193,6 +197,7 @@ describe('phi handler', () => {
     expect(
       result.transformedData?.response?.json?.choices?.[0]?.message?.content
     ).toBe('J******e has a history of heart disease and some random text');
+    expect(result.transformed).toBe(true);
   });
 });
 
@@ -225,6 +230,7 @@ describe('pii handler', () => {
     expect(result.verdict).toBe(true);
     expect(result.error).toBeNull();
     expect(result.data).toBeDefined();
+    expect(result.transformed).toBe(false);
   });
 
   it('should fail when text contains PII', async () => {
@@ -257,6 +263,7 @@ describe('pii handler', () => {
     expect(result.data).toBeDefined();
     expect(result.transformedData?.response?.json).toBeNull();
     expect(result.transformedData?.request?.json).toBeNull();
+    expect(result.transformed).toBe(false);
   });
 
   it('should detect and redact PII in request text', async () => {
@@ -294,6 +301,7 @@ describe('pii handler', () => {
     expect(result.transformedData?.request?.json?.messages?.[0]?.content).toBe(
       'My email is a*********m and some random text'
     );
+    expect(result.transformed).toBe(true);
   });
 
   it('should detect and redact PII in request text with multiple content parts', async () => {
@@ -343,6 +351,7 @@ describe('pii handler', () => {
     expect(
       result.transformedData?.request?.json?.messages?.[0]?.content?.[1]?.text
     ).toBe('My email is a*********m and some random text');
+    expect(result.transformed).toBe(true);
   });
 
   it('should detect and redact PII in response text', async () => {
@@ -382,6 +391,7 @@ describe('pii handler', () => {
     expect(
       result.transformedData?.response?.json?.choices?.[0]?.message?.content
     ).toBe('My email is a*********m and some random text');
+    expect(result.transformed).toBe(true);
   });
 });