Merge pull request #870 from Portkey-AI/feat/mutator-hooks

Feat: Mutator Hooks

plugins/index.ts

@@ -34,6 +34,13 @@ import { handler as patronustoxicity } from './patronus/toxicity';
 import { handler as patronuscustom } from './patronus/custom';
 import { mistralGuardrailHandler } from './mistral';
 import { handler as pangeatextGuard } from './pangea/textGuard';
+import { handler as portkeyredactPii } from './portkey/redactPii';
+import { handler as promptfooRedactPii } from './promptfoo/redactPii';
+import { handler as promptfooHarm } from './promptfoo/harm';
+import { handler as promptfooGuard } from './promptfoo/guard';
+import { handler as pangearedactPii } from './pangea/redactPii';
+import { handler as patronusredactPii } from './patronus/redactPii';
+import { handler as patronusredactPhi } from './patronus/redactPhi';
 
 export const plugins = {
   default: {
@@ -58,6 +65,7 @@ export const plugins = {
     language: portkeylanguage,
     pii: portkeypii,
     gibberish: portkeygibberish,
+    redactPii: portkeyredactPii,
   },
   aporia: {
     validateProject: aporiavalidateProject,
@@ -81,11 +89,19 @@ export const plugins = {
     retrievalAnswerRelevance: patronusretrievalAnswerRelevance,
     toxicity: patronustoxicity,
     custom: patronuscustom,
+    redactPii: patronusredactPii,
+    redactPhi: patronusredactPhi,
   },
   mistral: {
     moderateContent: mistralGuardrailHandler,
   },
   pangea: {
     textGuard: pangeatextGuard,
+    redactPii: pangearedactPii,
+  },
+  promptfoo: {
+    redactPii: promptfooRedactPii,
+    harm: promptfooHarm,
+    guard: promptfooGuard,
   },
 };

plugins/pangea/redactPii.ts

@@ -0,0 +1,95 @@
+import {
+  HookEventType,
+  PluginContext,
+  PluginHandler,
+  PluginParameters,
+} from '../types';
+import { getCurrentContentPart, post, setCurrentContentPart } from '../utils';
+import { VERSION } from './version';
+
+export const handler: PluginHandler = async (
+  context: PluginContext,
+  parameters: PluginParameters,
+  eventType: HookEventType
+) => {
+  let transformedData = {
+    request: {
+      json: null,
+    },
+    response: {
+      json: null,
+    },
+  };
+
+  try {
+    if (!parameters.credentials?.domain) {
+      return {
+        error: `'parameters.credentials.domain' must be set`,
+        verdict: true,
+        data: null,
+      };
+    }
+
+    if (!parameters.credentials?.apiKey) {
+      return {
+        error: `'parameters.credentials.apiKey' must be set`,
+        verdict: true,
+        data: null,
+      };
+    }
+
+    const url = `https://redact.${parameters.credentials.domain}/v1/redact_structured`;
+
+    const { content } = getCurrentContentPart(context, eventType);
+
+    if (!content) {
+      return {
+        error: { message: 'request or response json is empty' },
+        verdict: true,
+        data: null,
+        transformedData,
+      };
+    }
+
+    const requestOptions = {
+      headers: {
+        'Content-Type': 'application/json',
+        'User-Agent': 'portkey-ai-plugin/' + VERSION,
+        Authorization: `Bearer ${parameters.credentials.apiKey}`,
+      },
+    };
+    const request = {
+      data: content,
+      ...(Array.isArray(content) && content[0]?.type === 'text'
+        ? { jsonp: ['$[*].text'] }
+        : {}),
+    };
+
+    const response = await post(url, request, requestOptions);
+
+    if (response.result?.count > 0 && response.result.redacted_data) {
+      setCurrentContentPart(
+        context,
+        eventType,
+        transformedData,
+        response.result.redacted_data
+      );
+    }
+
+    return {
+      error: null,
+      verdict: true,
+      data: {
+        summary: response.summary,
+      },
+      transformedData,
+    };
+  } catch (e) {
+    return {
+      error: e as Error,
+      verdict: true,
+      data: null,
+      transformedData,
+    };
+  }
+};

plugins/patronus/globals.ts

@@ -31,3 +31,36 @@ export const postPatronus = async (
 
   return post(BASE_URL, body, options, timeout);
 };
+
+interface Position {
+  positions: [number, number][];
+  extra: any;
+  confidence_interval: any;
+}
+
+// For finding all longest positions of equal length
+interface AllLongestPositionsResult {
+  positions: [number, number][];
+  length: number;
+}
+
+export function findAllLongestPositions(
+  data: Position
+): AllLongestPositionsResult | null {
+  if (!data?.positions?.length) {
+    return null;
+  }
+
+  // Calculate max length
+  const maxLength = Math.max(...data.positions.map((pos) => pos[1] - pos[0]));
+
+  // Find all positions with max length
+  const longestPositions = data.positions.filter(
+    (pos) => pos[1] - pos[0] === maxLength
+  );
+
+  return {
+    positions: longestPositions,
+    length: maxLength,
+  };
+}

plugins/patronus/redactPhi.ts

@@ -0,0 +1,94 @@
+import {
+  HookEventType,
+  PluginContext,
+  PluginHandler,
+  PluginParameters,
+} from '../types';
+import { getCurrentContentPart, setCurrentContentPart } from '../utils';
+import { findAllLongestPositions, postPatronus } from './globals';
+import { maskEntities } from './redactPii';
+
+const redactPhi = async (text: string, credentials: any) => {
+  const evaluator = 'phi';
+
+  const evaluationBody: any = {
+    output: text,
+  };
+
+  const result: any = await postPatronus(
+    evaluator,
+    credentials,
+    evaluationBody
+  );
+  const evalResult = result.results[0];
+
+  const positionsData = evalResult.evaluation_result.additional_info;
+  if (
+    Array.isArray(positionsData?.positions) &&
+    positionsData.positions.length > 0
+  ) {
+    const longestPosition = findAllLongestPositions(positionsData);
+    if (longestPosition?.positions && longestPosition.positions.length > 0) {
+      const maskedText = maskEntities(text, longestPosition.positions);
+      return maskedText;
+    }
+  }
+  return text;
+};
+
+export const handler: PluginHandler = async (
+  context: PluginContext,
+  parameters: PluginParameters,
+  eventType: HookEventType
+) => {
+  const transformedData: Record<string, any> = {
+    request: {
+      json: null,
+    },
+    response: {
+      json: null,
+    },
+  };
+
+  try {
+    const { content, textArray } = getCurrentContentPart(context, eventType);
+
+    if (!content) {
+      return {
+        error: { message: 'request or response json is empty' },
+        verdict: true,
+        data: null,
+        transformedData,
+      };
+    }
+
+    const transformedTextPromise = textArray.map((text) =>
+      redactPhi(text, parameters.credentials)
+    );
+
+    const transformedText = await Promise.all(transformedTextPromise);
+
+    setCurrentContentPart(
+      context,
+      eventType,
+      transformedData,
+      null,
+      transformedText
+    );
+
+    return {
+      error: null,
+      verdict: true,
+      data: null,
+      transformedData,
+    };
+  } catch (e: any) {
+    delete e.stack;
+    return {
+      error: e as Error,
+      verdict: true,
+      data: null,
+      transformedData,
+    };
+  }
+};