Merge branch 'main' into feat/traces-for-plugins

Author: VisargD

.github/pull_request_template.md

@@ -1,12 +1,35 @@
-**Title:** 
-- Brief Description of Changes
+## Description
+<!-- Provide a brief description of the changes in this PR -->
 
-**Description:** (optional)
-- Detailed change 1
-- Detailed change 2
+## Motivation
+<!-- Provide a brief motivation of why the changes in this PR are needed -->
 
-**Motivation:** (optional)
-- Why this change is necessary or beneficial
+## Type of Change
+<!-- Put an 'x' in the boxes that apply -->
+- [ ] Bug fix (non-breaking change which fixes an issue)
+- [ ] New feature (non-breaking change which adds functionality)
+- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
+- [ ] Documentation update
+- [ ] Refactoring (no functional changes)
 
-**Related Issues:** (optional)
-- #issue-number
+## How Has This Been Tested?
+<!-- Describe the tests you ran to verify your changes -->
+- [ ] Unit Tests
+- [ ] Integration Tests
+- [ ] Manual Testing
+
+## Screenshots (if applicable)
+<!-- Add screenshots to help explain your changes -->
+
+## Checklist
+<!-- Put an 'x' in the boxes that apply -->
+- [ ] My code follows the style guidelines of this project
+- [ ] I have performed a self-review of my own code
+- [ ] I have commented my code, particularly in hard-to-understand areas
+- [ ] I have made corresponding changes to the documentation
+- [ ] My changes generate no new warnings
+- [ ] I have added tests that prove my fix is effective or that my feature works
+- [ ] New and existing unit tests pass locally with my changes
+
+## Related Issues
+<!-- Link related issues below. Insert the issue link or reference -->

conf.json

@@ -6,7 +6,8 @@
     "sydelabs",
     "pillar",
     "patronus",
-    "pangea"
+    "pangea",
+    "promptsecurity"
   ],
   "credentials": {
     "portkey": {

cookbook/guardrails/Langchain Chatbot with PII Guardrails.ipynb

@@ -29,7 +29,7 @@
         "- **Custom guardrail**s: Integrate your existing guardrail systems custom guardrail integration\n",
         "-**LLM-based guardrails** (e.g., gibberish detection, prompt injection scanning)\n",
         "\n",
-        "With **20+ deterministic guardrail**s and integrations with platforms like **Aporia, Pillar and Patronus AI,** Portkey provides comprehensive AI safety solutions. Guardrails can be configured for inputs, outputs, or both, with actions ranging from request denial to alternative LLM fallbacks.\n",
+        "With **20+ deterministic guardrail**s and integrations with platforms like **Aporia, Pillar, Patronus AI and Prompt Security,** Portkey provides comprehensive AI safety solutions. Guardrails can be configured for inputs, outputs, or both, with actions ranging from request denial to alternative LLM fallbacks.\n",
         "\n",
         "For more details, visit the [Portkey Guardrails documentation](https://docs.portkey.ai/docs/product/guardrails/list-of-guardrail-checks/patronus-ai).\n",
         "\n",

package-lock.json

@@ -1,6 +1,6 @@
 {
   "name": "@portkey-ai/gateway",
-  "version": "1.9.15",
+  "version": "1.9.17",
   "description": "A fast AI gateway by Portkey",
   "repository": {
     "type": "git",

package.json

@@ -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 abc@xyz.com and SSN is 123-45-6789',
+        json: {
+          messages: [
+            {
+              role: 'user',
+              content: 'My email is abc@xyz.com 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/azure.test.ts

@@ -0,0 +1,129 @@
+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?.length > 0;
+
+    verdict = !(hasHarmfulContent || hasBlocklistHit);
+  }
+
+  return {
+    error,
+    verdict,
+    data,
+  };
+};