Merge pull request #1277 from Portkey-AI/chore/update-azure-guardrails

Allow failures when PII guardrail fails, manage blank agents in requests

Author: VisargD

plugins/azure/azure.test.ts

@@ -45,6 +45,29 @@ describe('Azure Plugins', () => {
         expect(result.error).toBeNull();
         expect(result.verdict).toBe(true);
         expect(result.transformed).toBe(true);
+      }, 10000);
+
+      it('should not redact anything if text has no PII', async () => {
+        const context = structuredClone(mockContext);
+        context.request.text = "hello, I'm a harmless string";
+        context.request.json = {
+          messages: [{ role: 'user', content: "hello, I'm a harmless string" }],
+        };
+        const result = await piiHandler(context, params, 'beforeRequestHook');
+        expect(result.error).toBeNull();
+        expect(result.verdict).toBe(true);
+        expect(result.transformed).toBe(false);
+      });
+
+      it('should not redact anything if redact is false', async () => {
+        const result = await piiHandler(
+          mockContext,
+          { ...params, redact: false },
+          'beforeRequestHook'
+        );
+        expect(result.error).toBeNull();
+        expect(result.verdict).toBe(false);
+        expect(result.transformed).toBe(false);
       });
 
       it('should handle API errors gracefully', async () => {

plugins/azure/contentSafety.ts

@@ -17,7 +17,7 @@ export const handler: PluginHandler<{
   context: PluginContext,
   parameters: PluginParameters<{ contentSafety: AzureCredentials }>,
   eventType: HookEventType,
-  options
+  pluginOptions?: Record<string, any>
 ) => {
   let error = null;
   let verdict = true;
@@ -69,8 +69,8 @@ export const handler: PluginHandler<{
   const { token, error: tokenError } = await getAccessToken(
     credentials as any,
     'contentSafety',
-    options,
-    options?.env
+    pluginOptions,
+    pluginOptions?.env
   );
 
   if (tokenError) {
@@ -110,14 +110,13 @@ export const handler: PluginHandler<{
   };
 
   const timeout = parameters.timeout || 5000;
+  const requestOptions: Record<string, any> = { headers };
+  if (agent) {
+    requestOptions.dispatcher = agent;
+  }
   let response;
   try {
-    response = await post(
-      url,
-      request,
-      { headers, dispatcher: agent },
-      timeout
-    );
+    response = await post(url, request, requestOptions, timeout);
   } catch (e) {
     return { error: e, verdict: true, data };
   }

plugins/azure/pii.ts

@@ -12,7 +12,7 @@ import { getAccessToken } from './utils';
 const redact = async (
   documents: any[],
   parameters: PluginParameters<{ pii: AzureCredentials }>,
-  options?: Record<string, any>
+  pluginOptions?: Record<string, any>
 ) => {
   const body = {
     kind: 'PiiEntityRecognition',
@@ -35,8 +35,8 @@ const redact = async (
   const { token, error: tokenError } = await getAccessToken(
     credentials as any,
     'pii',
-    options,
-    options?.env
+    pluginOptions,
+    pluginOptions?.env
   );
 
   const headers: Record<string, string> = {
@@ -66,20 +66,19 @@ const redact = async (
   }
 
   const timeout = parameters.timeout || 5000;
-  const response = await post(
-    url,
-    body,
-    { headers, dispatcher: agent },
-    timeout
-  );
+  const requestOptions: Record<string, any> = { headers };
+  if (agent) {
+    requestOptions.dispatcher = agent;
+  }
+  const response = await post(url, body, requestOptions, timeout);
   return response;
 };
 
 export const handler: PluginHandler<{ pii: AzureCredentials }> = async (
   context: PluginContext,
   parameters: PluginParameters<{ pii: AzureCredentials }>,
   eventType: HookEventType,
-  options?: Record<string, any>
+  pluginOptions?: Record<string, any>
 ) => {
   let error = null;
   let verdict = true;
@@ -134,9 +133,18 @@ export const handler: PluginHandler<{ pii: AzureCredentials }> = async (
   }));
 
   try {
-    const response = await redact(documents, parameters, options);
+    const response = await redact(documents, parameters, pluginOptions);
+    if (!response?.results?.documents) {
+      throw new Error('Invalid response from Azure PII API');
+    }
     data = response.results.documents;
-    if (parameters.redact) {
+    const containsPII =
+      data.length > 0 && data.some((doc: any) => doc.entities.length > 0);
+    if (containsPII) {
+      verdict = false;
+    }
+    if (parameters.redact && containsPII) {
+      verdict = true;
       const redactedData = (response.results.documents ?? []).map(
         (doc: any) => doc.redactedText
       );