Merge pull request #1048 from b4s36t4/feat/azure-guardrail-plugin

feat: azure plugin support for pii and content safety

Author: VisargD

plugins/azure/azure.test.ts

@@ -0,0 +1,173 @@
+import { describe, it, expect, jest, beforeEach } from '@jest/globals';
+import { handler as piiHandler } from './pii';
+import { handler as contentSafetyHandler } from './contentSafety';
+import { HookEventType, PluginContext, PluginParameters } from '../types';
+import { AzureCredentials } from './types';
+import { pii, contentSafety } from './.creds.json';
+
+describe('Azure Plugins', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  describe('PII Plugin', () => {
+    const mockContext: PluginContext = {
+      request: {
+        text: 'My email is [email protected] and SSN is 123-45-6789',
+        json: {
+          messages: [
+            {
+              role: 'user',
+              content: 'My email is [email protected] and SSN is 123-45-6789',
+            },
+          ],
+        },
+      },
+      requestType: 'chatComplete',
+    };
+
+    describe('API Key Authentication', () => {
+      const params: PluginParameters<{ pii: AzureCredentials }> = {
+        credentials: {
+          pii: pii.apiKey as AzureCredentials,
+        },
+        redact: true,
+        apiVersion: '2024-11-01',
+      };
+
+      it('should successfully analyze and redact PII with API key', async () => {
+        const result = await piiHandler(
+          mockContext,
+          params,
+          'beforeRequestHook'
+        );
+
+        expect(result.error).toBeNull();
+        expect(result.verdict).toBe(true);
+        expect(result.transformed).toBe(true);
+      });
+
+      it('should handle API errors gracefully', async () => {
+        const result = await piiHandler(
+          mockContext,
+          {
+            ...params,
+            credentials: {
+              pii: {
+                azureAuthMode: 'apiKey',
+                resourceName: 'wrong-resurce-name',
+                apiKey: 'wrong-api-key',
+              },
+            },
+          },
+          'beforeRequestHook'
+        );
+
+        expect(result.error).toBeDefined();
+        expect(result.verdict).toBe(true);
+        expect(result.data).toBeNull();
+      });
+    });
+
+    describe('Entra ID Authentication', () => {
+      const params: PluginParameters<{ pii: AzureCredentials }> = {
+        credentials: {
+          pii: pii.entra as AzureCredentials,
+        },
+        redact: true,
+      };
+
+      it('should successfully analyze and redact PII with Entra ID', async () => {
+        const result = await piiHandler(
+          mockContext,
+          params,
+          'beforeRequestHook'
+        );
+        expect(result.error).toBeNull();
+        expect(result.verdict).toBe(true);
+        expect(result.transformed).toBe(true);
+      });
+    });
+  });
+
+  describe('Content Safety Plugin', () => {
+    const mockContext = {
+      request: {
+        text: "Fuck you, if you don't answer I'll kill you.",
+        json: {
+          messages: [
+            {
+              role: 'user',
+              content: `Fuck you, if you don't answer I'll kill you.`,
+            },
+          ],
+        },
+      },
+    };
+
+    describe('API Key Authentication', () => {
+      const params: PluginParameters<{ contentSafety: AzureCredentials }> = {
+        credentials: {
+          contentSafety: contentSafety.apiKey as AzureCredentials,
+        },
+        categories: ['Hate', 'Violence'],
+        apiVersion: '2024-09-01',
+      };
+
+      it('should successfully analyze content with API key', async () => {
+        const result = await contentSafetyHandler(
+          mockContext,
+          params,
+          'beforeRequestHook'
+        );
+
+        expect(result.error).toBeNull();
+        expect(result.verdict).toBe(false);
+        expect(result.data).toBeDefined();
+      });
+    });
+
+    describe('Entra ID Authentication', () => {
+      const params: PluginParameters<{ contentSafety: AzureCredentials }> = {
+        credentials: {
+          contentSafety: contentSafety.entra as AzureCredentials,
+        },
+        categories: ['Hate', 'Violence'],
+        apiVersion: '2024-09-01',
+      };
+
+      it('should successfully analyze content with Entra ID', async () => {
+        const result = await contentSafetyHandler(
+          mockContext,
+          params,
+          'beforeRequestHook'
+        );
+
+        expect(result.error).toBeNull();
+        expect(result.verdict).toBe(false);
+        expect(result.data).toBeDefined();
+      });
+
+      it('should detect harmful content correctly', async () => {
+        const harmfulResponse = {
+          categoriesAnalysis: [
+            {
+              category: 'Hate',
+              severity: 2, // High severity
+            },
+          ],
+          blocklistsMatch: [],
+        };
+
+        const result = await contentSafetyHandler(
+          mockContext,
+          params,
+          'beforeRequestHook'
+        );
+        expect(result.error).toBeNull();
+        expect(result.verdict).toBe(false); // Should be false due to high severity
+        expect(result.data).toBeDefined();
+      });
+    });
+  });
+});

plugins/azure/contentSafety.ts

@@ -0,0 +1,131 @@
+import {
+  HookEventType,
+  PluginContext,
+  PluginHandler,
+  PluginParameters,
+} from '../types';
+import { post, getText } from '../utils';
+import { AzureCredentials } from './types';
+import { getAccessToken } from './utils';
+
+const defaultCategories = ['Hate', 'SelfHarm', 'Sexual', 'Violence'];
+
+export const handler: PluginHandler<{
+  contentSafety: AzureCredentials;
+}> = async (
+  context: PluginContext,
+  parameters: PluginParameters<{ contentSafety: AzureCredentials }>,
+  eventType: HookEventType,
+  options
+) => {
+  let error = null;
+  let verdict = true;
+  let data = null;
+
+  const credentials = parameters.credentials?.contentSafety;
+
+  if (!credentials) {
+    return {
+      error: new Error('parameters.credentials.contentSafety must be set'),
+      verdict: true,
+      data,
+    };
+  }
+
+  // Validate required credentials
+  if (!credentials?.resourceName) {
+    return {
+      error: new Error('Content Safety credentials must include resourceName'),
+      verdict: true,
+      data,
+    };
+  }
+
+  // prefer api key over auth mode
+  if (!credentials?.azureAuthMode && !credentials?.apiKey) {
+    return {
+      error: new Error(
+        'Content Safety credentials must include either apiKey or azureAuthMode'
+      ),
+      verdict: true,
+      data,
+    };
+  }
+
+  const text = getText(context, eventType);
+  if (!text) {
+    return {
+      error: new Error('request or response text is empty'),
+      verdict: true,
+      data,
+    };
+  }
+
+  const apiVersion = parameters.apiVersion || '2024-11-01';
+
+  const url = `https://${credentials.resourceName}.cognitiveservices.azure.com/contentsafety/text:analyze?api-version=${apiVersion}`;
+
+  const { token, error: tokenError } = await getAccessToken(
+    credentials as any,
+    'contentSafety',
+    options,
+    options?.env
+  );
+
+  if (tokenError) {
+    return {
+      error: tokenError,
+      verdict: true,
+      data,
+    };
+  }
+
+  const headers: Record<string, string> = {
+    'Content-Type': 'application/json',
+    'User-Agent': 'portkey-ai-plugin/',
+    'Ocp-Apim-Subscription-Key': token,
+  };
+
+  if (credentials?.azureAuthMode && credentials?.azureAuthMode !== 'apiKey') {
+    headers['Authorization'] = `Bearer ${token}`;
+    delete headers['Ocp-Apim-Subscription-Key'];
+  }
+
+  const request = {
+    text: text,
+    categories: parameters.categories || defaultCategories,
+    blocklistNames: parameters.blocklistNames || [],
+  };
+
+  const timeout = parameters.timeout || 5000;
+  let response;
+  try {
+    response = await post(url, request, { headers }, timeout);
+  } catch (e) {
+    return { error: e, verdict: true, data };
+  }
+
+  if (response) {
+    data = response;
+
+    // Check if any category exceeds the threshold (default: 2 - Medium)
+    const hasHarmfulContent = response.categoriesAnalysis?.some(
+      (category: any) => {
+        return category.severity >= (parameters.severity || 2);
+      }
+    );
+
+    // Check if any blocklist items were hit
+    const hasBlocklistHit = response.blocklistsMatch?.some((match: any) => {
+      return match.matchResults.length > 0;
+    });
+
+    verdict = !(hasHarmfulContent || hasBlocklistHit);
+  }
+
+  return {
+    error,
+    verdict,
+    data,
+  };
+};