feat: add comprehend pii redactor

JS/edgechains/arakoodev/package.json

@@ -22,6 +22,7 @@
         "test": "vitest"
     },
     "dependencies": {
+        "@aws-sdk/client-comprehend": "^3.1076.0",
         "@babel/core": "^7.24.4",
         "@babel/preset-env": "^7.24.4",
         "@hono/node-server": "^0.6.0",

JS/edgechains/arakoodev/src/ai/src/index.ts

@@ -3,3 +3,12 @@ export { GeminiAI } from "./lib/gemini/gemini.js";
 export { LlamaAI } from "./lib/llama/llama.js";
 export { RetellAI } from "./lib/retell-ai/retell.js";
 export { RetellWebClient } from "./lib/retell-ai/retellWebClient.js";
+export { ComprehendPiiRedactor } from "./lib/comprehend/comprehendPiiRedactor.js";
+export type {
+  ChatEndpoint,
+  ComprehendLikeClient,
+  ComprehendPiiRedactorOptions,
+  RedactableChatOptions,
+  RedactableMessage,
+  RedactionReplacement,
+} from "./lib/comprehend/comprehendPiiRedactor.js";

JS/edgechains/arakoodev/src/ai/src/lib/comprehend/comprehendPiiRedactor.ts

@@ -0,0 +1,182 @@
+import {
+  ComprehendClient,
+  DetectPiiEntitiesCommand,
+  type DetectPiiEntitiesCommandOutput,
+  type LanguageCode,
+  type PiiEntity,
+} from "@aws-sdk/client-comprehend";
+
+export type ComprehendLikeClient = {
+  send(
+    command: DetectPiiEntitiesCommand,
+  ): Promise<DetectPiiEntitiesCommandOutput>;
+};
+
+export type RedactionReplacement = "entityType" | "mask";
+
+export interface ComprehendPiiRedactorOptions {
+  client?: ComprehendLikeClient;
+  region?: string;
+  languageCode?: LanguageCode;
+  minScore?: number;
+  entityTypes?: string[];
+  replacement?: RedactionReplacement;
+  maskCharacter?: string;
+}
+
+export interface RedactableMessage {
+  content?: string;
+  [key: string]: unknown;
+}
+
+export interface RedactableChatOptions {
+  prompt?: string;
+  messages?: RedactableMessage[];
+  [key: string]: unknown;
+}
+
+export interface ChatEndpoint<
+  TOptions extends RedactableChatOptions = RedactableChatOptions,
+  TResult = unknown,
+> {
+  chat(chatOptions: TOptions): Promise<TResult>;
+}
+
+export class ComprehendPiiRedactor {
+  private readonly client: ComprehendLikeClient;
+  private readonly languageCode: LanguageCode;
+  private readonly minScore: number;
+  private readonly entityTypes?: Set<string>;
+  private readonly replacement: RedactionReplacement;
+  private readonly maskCharacter: string;
+
+  constructor(options: ComprehendPiiRedactorOptions = {}) {
+    this.client =
+      options.client || new ComprehendClient({ region: options.region });
+    this.languageCode = options.languageCode || "en";
+    this.minScore = options.minScore ?? 0;
+    this.entityTypes = options.entityTypes
+      ? new Set(options.entityTypes)
+      : undefined;
+    this.replacement = options.replacement || "entityType";
+    this.maskCharacter = options.maskCharacter || "*";
+  }
+
+  async detect(text: string): Promise<PiiEntity[]> {
+    if (!text) return [];
+
+    const response = await this.client.send(
+      new DetectPiiEntitiesCommand({
+        Text: text,
+        LanguageCode: this.languageCode,
+      }),
+    );
+
+    return (response.Entities || []).filter((entity) =>
+      this.shouldRedact(entity),
+    );
+  }
+
+  async redact(text: string): Promise<string> {
+    if (!text) return text;
+
+    const entities = this.selectNonOverlappingEntities(await this.detect(text));
+    let redacted = text;
+
+    for (const entity of entities.sort(
+      (a, b) => (b.BeginOffset || 0) - (a.BeginOffset || 0),
+    )) {
+      const begin = entity.BeginOffset || 0;
+      const end = entity.EndOffset || begin;
+      const original = redacted.slice(begin, end);
+      const replacement =
+        this.replacement === "mask"
+          ? this.maskCharacter.repeat(original.length)
+          : `[${entity.Type || "PII"}]`;
+
+      redacted = redacted.slice(0, begin) + replacement + redacted.slice(end);
+    }
+
+    return redacted;
+  }
+
+  async redactChatOptions<T extends RedactableChatOptions>(
+    chatOptions: T,
+  ): Promise<T> {
+    const redactedOptions = { ...chatOptions };
+
+    if (typeof chatOptions.prompt === "string") {
+      redactedOptions.prompt = await this.redact(chatOptions.prompt);
+    }
+
+    if (Array.isArray(chatOptions.messages)) {
+      redactedOptions.messages = await Promise.all(
+        chatOptions.messages.map(async (message) => {
+          if (typeof message.content !== "string") return message;
+
+          return {
+            ...message,
+            content: await this.redact(message.content),
+          };
+        }),
+      );
+    }
+
+    return redactedOptions;
+  }
+
+  wrapEndpoint<TOptions extends RedactableChatOptions, TResult>(
+    endpoint: ChatEndpoint<TOptions, TResult>,
+  ): ChatEndpoint<TOptions, TResult> {
+    return {
+      chat: async (chatOptions: TOptions) => {
+        return endpoint.chat(await this.redactChatOptions(chatOptions));
+      },
+    };
+  }
+
+  asOperator<T extends string | RedactableChatOptions>(): (
+    value: T,
+  ) => Promise<T> {
+    return async (value: T) => {
+      if (typeof value === "string") {
+        return (await this.redact(value)) as T;
+      }
+
+      return this.redactChatOptions(value) as Promise<T>;
+    };
+  }
+
+  private shouldRedact(entity: PiiEntity): boolean {
+    if (entity.BeginOffset === undefined || entity.EndOffset === undefined)
+      return false;
+    if ((entity.Score ?? 0) < this.minScore) return false;
+    if (
+      this.entityTypes &&
+      (!entity.Type || !this.entityTypes.has(entity.Type))
+    )
+      return false;
+    return true;
+  }
+
+  private selectNonOverlappingEntities(entities: PiiEntity[]): PiiEntity[] {
+    const selected: PiiEntity[] = [];
+    const sorted = [...entities].sort(
+      (a, b) => (a.BeginOffset || 0) - (b.BeginOffset || 0),
+    );
+
+    for (const entity of sorted) {
+      const begin = entity.BeginOffset || 0;
+      const end = entity.EndOffset || begin;
+      const overlaps = selected.some((selectedEntity) => {
+        const selectedBegin = selectedEntity.BeginOffset || 0;
+        const selectedEnd = selectedEntity.EndOffset || selectedBegin;
+        return begin < selectedEnd && end > selectedBegin;
+      });
+
+      if (!overlaps) selected.push(entity);
+    }
+
+    return selected;
+  }
+}

JS/edgechains/arakoodev/src/ai/src/tests/comprehend/comprehendPiiRedactor.test.ts

@@ -0,0 +1,101 @@
+import { describe, expect, it, vi } from "vitest";
+import { ComprehendPiiRedactor } from "../../lib/comprehend/comprehendPiiRedactor.js";
+
+describe("ComprehendPiiRedactor", () => {
+  it("redacts PII entities with entity labels", async () => {
+    const client = {
+      send: vi.fn().mockResolvedValue({
+        Entities: [
+          {
+            Type: "EMAIL",
+            BeginOffset: 14,
+            EndOffset: 27,
+            Score: 0.99,
+          },
+          {
+            Type: "PHONE",
+            BeginOffset: 31,
+            EndOffset: 43,
+            Score: 0.98,
+          },
+        ],
+      }),
+    };
+    const redactor = new ComprehendPiiRedactor({ client });
+
+    const result = await redactor.redact(
+      "Contact me at a@example.com or 555-010-9999.",
+    );
+
+    expect(result).toBe("Contact me at [EMAIL] or [PHONE].");
+    expect(client.send).toHaveBeenCalledOnce();
+    expect(client.send.mock.calls[0][0].input).toMatchObject({
+      Text: "Contact me at a@example.com or 555-010-9999.",
+      LanguageCode: "en",
+    });
+  });
+
+  it("redacts prompt and chat message content", async () => {
+    const client = {
+      send: vi
+        .fn()
+        .mockResolvedValueOnce({
+          Entities: [
+            { Type: "NAME", BeginOffset: 6, EndOffset: 10, Score: 0.98 },
+          ],
+        })
+        .mockResolvedValueOnce({
+          Entities: [
+            { Type: "EMAIL", BeginOffset: 6, EndOffset: 19, Score: 0.99 },
+          ],
+        }),
+    };
+    const redactor = new ComprehendPiiRedactor({ client });
+
+    const result = await redactor.redactChatOptions({
+      prompt: "Hello Jane",
+      messages: [{ role: "user", content: "Email a@example.com" }],
+      temperature: 0.2,
+    });
+
+    expect(result).toEqual({
+      prompt: "Hello [NAME]",
+      messages: [{ role: "user", content: "Email [EMAIL]" }],
+      temperature: 0.2,
+    });
+  });
+
+  it("wraps endpoint-style chat calls", async () => {
+    const client = {
+      send: vi.fn().mockResolvedValue({
+        Entities: [
+          { Type: "EMAIL", BeginOffset: 6, EndOffset: 19, Score: 0.99 },
+        ],
+      }),
+    };
+    const endpoint = {
+      chat: vi.fn().mockResolvedValue({ content: "done" }),
+    };
+    const redactor = new ComprehendPiiRedactor({ client });
+    const wrappedEndpoint = redactor.wrapEndpoint(endpoint);
+
+    const result = await wrappedEndpoint.chat({
+      prompt: "Email a@example.com",
+    });
+
+    expect(result).toEqual({ content: "done" });
+    expect(endpoint.chat).toHaveBeenCalledWith({ prompt: "Email [EMAIL]" });
+  });
+
+  it("supports a chainable async operator", async () => {
+    const client = {
+      send: vi.fn().mockResolvedValue({
+        Entities: [{ Type: "NAME", BeginOffset: 3, EndOffset: 7, Score: 0.99 }],
+      }),
+    };
+    const redactor = new ComprehendPiiRedactor({ client });
+    const redact = redactor.asOperator<string>();
+
+    await expect(redact("Hi Jane")).resolves.toBe("Hi [NAME]");
+  });
+});

JS/edgechains/examples/comprehend-pii-redaction/README.md

@@ -0,0 +1,19 @@
+# AWS Comprehend PII Redaction
+
+This example loads a prompt from Jsonnet, redacts sensitive values with
+`ComprehendPiiRedactor`, and prints the safe prompt.
+
+```bash
+npm install
+npm run build
+npm start
+```
+
+By default it uses a mock Comprehend client so the example can run without AWS
+credentials. To call AWS Comprehend directly:
+
+```bash
+export AWS_REGION="us-east-1"
+export USE_REAL_AWS="true"
+npm start -- "Jane Doe emailed jane@example.com about invoice 555-010-9999."
+```

JS/edgechains/examples/comprehend-pii-redaction/jsonnet/main.jsonnet

@@ -0,0 +1,13 @@
+local promptTemplate = |||
+  Please summarize this support request without exposing private information:
+
+  {request}
+|||;
+
+local request = std.extVar('request');
+
+{
+  prompt: std.strReplace(promptTemplate, '{request}', request),
+  languageCode: 'en',
+  minScore: 0.5,
+}

JS/edgechains/examples/comprehend-pii-redaction/package.json

@@ -0,0 +1,17 @@
+{
+  "name": "comprehend-pii-redaction",
+  "version": "1.0.0",
+  "type": "module",
+  "scripts": {
+    "build": "tsc",
+    "start": "node dist/index.js"
+  },
+  "dependencies": {
+    "@arakoodev/edgechains.js": "file:../../arakoodev",
+    "@arakoodev/jsonnet": "^0.1.0",
+    "@aws-sdk/client-comprehend": "^3.1076.0",
+    "@types/node": "^20.17.2",
+    "typescript": "^5.6.3"
+  },
+  "devDependencies": {}
+}

JS/edgechains/examples/comprehend-pii-redaction/src/edgechains-ai.d.ts

@@ -0,0 +1,16 @@
+declare module "@arakoodev/edgechains.js/ai" {
+  export type ComprehendLikeClient = {
+    send(command: any): Promise<any>;
+  };
+
+  export class ComprehendPiiRedactor {
+    constructor(options?: {
+      client?: ComprehendLikeClient;
+      region?: string;
+      languageCode?: string;
+      minScore?: number;
+    });
+
+    redact(text: string): Promise<string>;
+  }
+}