Add AWS Comprehend PII redaction utility

JS/edgechains/arakoodev/package.json

@@ -24,6 +24,7 @@
     "dependencies": {
         "@babel/core": "^7.24.4",
         "@babel/preset-env": "^7.24.4",
+        "@aws-sdk/client-comprehend": "^3.716.0",
         "@hono/node-server": "^0.6.0",
         "@lifeomic/attempt": "^3.1.0",
         "@playwright/test": "^1.45.3",

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

@@ -3,3 +3,4 @@ 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 { AwsComprehendRedactor } from "./lib/aws-comprehend/awsComprehendRedactor.js";

JS/edgechains/arakoodev/src/ai/src/lib/aws-comprehend/awsComprehendRedactor.ts

@@ -0,0 +1,121 @@
+import {
+    ComprehendClient,
+    DetectPiiEntitiesCommand,
+    type DetectPiiEntitiesCommandOutput,
+    type LanguageCode,
+    type PiiEntity,
+} from "@aws-sdk/client-comprehend";
+
+type ChatMessage = {
+    content: string;
+    [key: string]: unknown;
+};
+
+type PromptInput = {
+    prompt?: string;
+    messages?: ChatMessage[];
+    [key: string]: unknown;
+};
+
+type ComprehendLikeClient = {
+    send(command: DetectPiiEntitiesCommand): Promise<DetectPiiEntitiesCommandOutput>;
+};
+
+type RedactionResult = {
+    content: string;
+    entities: PiiEntity[];
+};
+
+type AwsComprehendRedactorOptions = {
+    client?: ComprehendLikeClient;
+    region?: string;
+    languageCode?: LanguageCode;
+    replacement?: string;
+    replacementForType?: (entity: PiiEntity) => string;
+};
+
+export class AwsComprehendRedactor {
+    private client: ComprehendLikeClient;
+    private languageCode: LanguageCode;
+    private replacement: string;
+    private replacementForType?: (entity: PiiEntity) => string;
+
+    constructor(options: AwsComprehendRedactorOptions = {}) {
+        this.client =
+            options.client ||
+            new ComprehendClient({
+                region: options.region || process.env.AWS_REGION || "us-east-1",
+            });
+        this.languageCode = options.languageCode || "en";
+        this.replacement = options.replacement || "[REDACTED]";
+        this.replacementForType = options.replacementForType;
+    }
+
+    async redactText(text: string): Promise<RedactionResult> {
+        if (!text) {
+            return { content: text, entities: [] };
+        }
+
+        const response = await this.client.send(
+            new DetectPiiEntitiesCommand({
+                Text: text,
+                LanguageCode: this.languageCode,
+            })
+        );
+        const entities = (response.Entities || []).filter(hasOffsets);
+
+        return {
+            content: redactByOffsets(text, entities, (entity) => this.maskFor(entity)),
+            entities,
+        };
+    }
+
+    async redactPrompt<T extends PromptInput>(input: T): Promise<T> {
+        const nextInput = { ...input };
+        if (typeof input.prompt === "string") {
+            nextInput.prompt = (await this.redactText(input.prompt)).content;
+        }
+        if (Array.isArray(input.messages)) {
+            nextInput.messages = await Promise.all(
+                input.messages.map(async (message) => ({
+                    ...message,
+                    content:
+                        typeof message.content === "string"
+                            ? (await this.redactText(message.content)).content
+                            : message.content,
+                }))
+            );
+        }
+        return nextInput;
+    }
+
+    async chain<T extends PromptInput, R>(input: T, next: (redactedInput: T) => Promise<R>): Promise<R> {
+        return next(await this.redactPrompt(input));
+    }
+
+    private maskFor(entity: PiiEntity): string {
+        if (this.replacementForType) {
+            return this.replacementForType(entity);
+        }
+        return entity.Type ? `[REDACTED_${entity.Type}]` : this.replacement;
+    }
+}
+
+function hasOffsets(entity: PiiEntity): entity is PiiEntity & { BeginOffset: number; EndOffset: number } {
+    return typeof entity.BeginOffset === "number" && typeof entity.EndOffset === "number";
+}
+
+function redactByOffsets(
+    text: string,
+    entities: Array<PiiEntity & { BeginOffset: number; EndOffset: number }>,
+    maskFor: (entity: PiiEntity) => string
+): string {
+    return [...entities]
+        .sort((a, b) => b.BeginOffset - a.BeginOffset)
+        .reduce((content, entity) => {
+            if (entity.BeginOffset < 0 || entity.EndOffset > content.length || entity.BeginOffset >= entity.EndOffset) {
+                return content;
+            }
+            return content.slice(0, entity.BeginOffset) + maskFor(entity) + content.slice(entity.EndOffset);
+        }, text);
+}

JS/edgechains/arakoodev/src/ai/src/tests/awsComprehendRedactor.test.ts

@@ -0,0 +1,68 @@
+import { AwsComprehendRedactor } from "../lib/aws-comprehend/awsComprehendRedactor";
+
+describe("AwsComprehendRedactor", () => {
+    test("redacts PII returned by AWS Comprehend", async () => {
+        const client = {
+            send: jest.fn().mockResolvedValue({
+                Entities: [
+                    {
+                        BeginOffset: 8,
+                        EndOffset: 24,
+                        Type: "EMAIL",
+                        Score: 0.99,
+                    },
+                ],
+            }),
+        };
+        const redactor = new AwsComprehendRedactor({ client });
+
+        const result = await redactor.redactText("Contact jane@example.com");
+
+        expect(result.content).toBe("Contact [REDACTED_EMAIL]");
+        expect(result.entities).toHaveLength(1);
+        expect(client.send).toHaveBeenCalledTimes(1);
+    });
+
+    test("redacts prompts before chaining to another endpoint", async () => {
+        const client = {
+            send: jest.fn().mockResolvedValue({
+                Entities: [
+                    {
+                        BeginOffset: 11,
+                        EndOffset: 27,
+                        Type: "EMAIL",
+                        Score: 0.99,
+                    },
+                ],
+            }),
+        };
+        const redactor = new AwsComprehendRedactor({ client });
+        const next = jest.fn().mockResolvedValue({ content: "ok" });
+
+        await redactor.chain({ prompt: "Summarize jane@example.com" }, next);
+
+        expect(next).toHaveBeenCalledWith({ prompt: "Summarize [REDACTED_EMAIL]" });
+    });
+
+    test("redacts chat message content", async () => {
+        const client = {
+            send: jest.fn().mockResolvedValue({
+                Entities: [
+                    {
+                        BeginOffset: 5,
+                        EndOffset: 17,
+                        Type: "PHONE",
+                        Score: 0.99,
+                    },
+                ],
+            }),
+        };
+        const redactor = new AwsComprehendRedactor({ client });
+
+        const result = await redactor.redactPrompt({
+            messages: [{ role: "user", content: "Call 555-123-4567" }],
+        });
+
+        expect(result.messages?.[0].content).toBe("Call [REDACTED_PHONE]");
+    });
+});

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

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

JS/edgechains/examples/aws-comprehend-redaction/readme.md

@@ -0,0 +1,17 @@
+# AWS Comprehend Redaction
+
+This example redacts PII from a prompt with AWS Comprehend before forwarding the prompt to an EdgeChains AI endpoint.
+
+Required environment variables:
+
+- `AWS_REGION`
+- AWS credentials supported by the AWS SDK
+- `OPENAI_API_KEY`
+- `OPENAI_ORG_ID`
+
+Run:
+
+```bash
+npm install
+npm start
+```

JS/edgechains/examples/aws-comprehend-redaction/src/index.ts

@@ -0,0 +1,21 @@
+import { OpenAI, AwsComprehendRedactor } from "@arakoodev/edgechains.js/ai";
+
+const redactor = new AwsComprehendRedactor({
+    region: process.env.AWS_REGION || "us-east-1",
+});
+
+const openai = new OpenAI({
+    apiKey: process.env.OPENAI_API_KEY,
+    orgId: process.env.OPENAI_ORG_ID,
+});
+
+const prompt = "Summarize this support request: My email is jane@example.com and my phone is 555-123-4567.";
+
+const response = await redactor.chain({ prompt }, (redactedPrompt) =>
+    openai.chat({
+        ...redactedPrompt,
+        max_tokens: 128,
+    })
+);
+
+console.log(response);

JS/edgechains/examples/aws-comprehend-redaction/tsconfig.json

@@ -0,0 +1,11 @@
+{
+    "compilerOptions": {
+        "target": "ES2022",
+        "module": "NodeNext",
+        "moduleResolution": "NodeNext",
+        "strict": true,
+        "esModuleInterop": true,
+        "outDir": "./dist"
+    },
+    "include": ["src/**/*.ts"]
+}