Add AWS Comprehend PII redaction

JS/edgechains/arakoodev/package.json

@@ -22,6 +22,7 @@
         "test": "vitest"
     },
     "dependencies": {
+        "@aws-sdk/client-comprehend": "^3.864.0",
         "@babel/core": "^7.24.4",
         "@babel/preset-env": "^7.24.4",
         "@hono/node-server": "^0.6.0",
@@ -48,6 +49,7 @@
         "retell-client-js-sdk": "^2.0.4",
         "retell-sdk": "^4.9.0",
         "retry": "^0.13.1",
+        "rxjs": "^7.8.2",
         "ts-node": "^10.9.2",
         "typeorm": "^0.3.20",
         "vitest": "^2.0.3",

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

@@ -3,3 +3,10 @@ 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,
+    ComprehendClientLike,
+    ComprehendPiiRedactorOptions,
+    RedactableChatOptions,
+} from "./lib/comprehend/comprehendPiiRedactor.js";

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

@@ -0,0 +1,160 @@
+import {
+    ComprehendClient,
+    DetectPiiEntitiesCommand,
+    LanguageCode,
+    type DetectPiiEntitiesCommandOutput,
+    type PiiEntity,
+} from "@aws-sdk/client-comprehend";
+import { Observable, concatMap, type OperatorFunction } from "rxjs";
+
+export interface ComprehendClientLike {
+    send(
+        command: DetectPiiEntitiesCommand,
+        options?: { abortSignal?: AbortSignal }
+    ): Promise<DetectPiiEntitiesCommandOutput>;
+}
+
+export interface ComprehendPiiRedactorOptions {
+    client?: ComprehendClientLike;
+    region?: string;
+    languageCode?: LanguageCode;
+    minScore?: number;
+    replacement?: "label" | "mask";
+    maskCharacter?: string;
+}
+
+export interface RedactableChatOptions {
+    prompt?: string;
+    messages?: Array<{ content?: string; [key: string]: unknown }>;
+    [key: string]: unknown;
+}
+
+export interface ChatEndpoint<Options extends RedactableChatOptions, Result> {
+    chat(options: Options): Promise<Result>;
+}
+
+export class ComprehendPiiRedactor {
+    private readonly client: ComprehendClientLike;
+    private readonly languageCode: LanguageCode;
+    private readonly minScore: number;
+    private readonly replacement: "label" | "mask";
+    private readonly maskCharacter: string;
+
+    constructor(options: ComprehendPiiRedactorOptions = {}) {
+        this.client =
+            options.client ||
+            (new ComprehendClient({
+                region: options.region,
+            }) as ComprehendClientLike);
+        this.languageCode = options.languageCode || LanguageCode.EN;
+        this.minScore = options.minScore ?? 0;
+        this.replacement = options.replacement || "label";
+        this.maskCharacter = options.maskCharacter || "*";
+    }
+
+    async redact(text: string, abortSignal?: AbortSignal): Promise<string> {
+        if (!text) return text;
+        if (new TextEncoder().encode(text).byteLength > 100 * 1024) {
+            throw new Error("Amazon Comprehend accepts at most 100 KiB of UTF-8 text");
+        }
+
+        const response = await this.client.send(
+            new DetectPiiEntitiesCommand({
+                Text: text,
+                LanguageCode: this.languageCode,
+            }),
+            { abortSignal }
+        );
+        const characters = Array.from(text);
+        const entities = selectEntities(response.Entities || [], characters.length, this.minScore);
+
+        for (const entity of entities) {
+            const begin = entity.BeginOffset as number;
+            const end = entity.EndOffset as number;
+            const original = characters.slice(begin, end);
+            const replacement =
+                this.replacement === "mask"
+                    ? Array(original.length).fill(this.maskCharacter)
+                    : Array.from(`[${entity.Type || "PII"}]`);
+            characters.splice(begin, end - begin, ...replacement);
+        }
+        return characters.join("");
+    }
+
+    async redactChatOptions<T extends RedactableChatOptions>(
+        options: T,
+        abortSignal?: AbortSignal
+    ): Promise<T> {
+        return {
+            ...options,
+            ...(typeof options.prompt === "string"
+                ? { prompt: await this.redact(options.prompt, abortSignal) }
+                : {}),
+            ...(options.messages
+                ? {
+                      messages: await Promise.all(
+                          options.messages.map(async (message) => ({
+                              ...message,
+                              ...(typeof message.content === "string"
+                                  ? { content: await this.redact(message.content, abortSignal) }
+                                  : {}),
+                          }))
+                      ),
+                  }
+                : {}),
+        };
+    }
+
+    redact$(text: string): Observable<string> {
+        return new Observable((subscriber) => {
+            const controller = new AbortController();
+            this.redact(text, controller.signal).then(
+                (value) => {
+                    if (!subscriber.closed) {
+                        subscriber.next(value);
+                        subscriber.complete();
+                    }
+                },
+                (error) => {
+                    if (!subscriber.closed) subscriber.error(error);
+                }
+            );
+            return () => controller.abort();
+        });
+    }
+
+    redactOperator(): OperatorFunction<string, string> {
+        return concatMap((text) => this.redact$(text));
+    }
+
+    endpointOperator<Options extends RedactableChatOptions, Result>(
+        endpoint: ChatEndpoint<Options, Result>
+    ): OperatorFunction<Options, Result> {
+        return concatMap(async (options) => endpoint.chat(await this.redactChatOptions(options)));
+    }
+}
+
+function selectEntities(entities: PiiEntity[], length: number, minScore: number): PiiEntity[] {
+    const selected: PiiEntity[] = [];
+    for (const entity of [...entities].sort((a, b) => (b.Score || 0) - (a.Score || 0))) {
+        const begin = entity.BeginOffset;
+        const end = entity.EndOffset;
+        if (
+            !Number.isInteger(begin) ||
+            !Number.isInteger(end) ||
+            (begin as number) < 0 ||
+            (end as number) <= (begin as number) ||
+            (end as number) > length ||
+            (entity.Score || 0) < minScore ||
+            selected.some(
+                (other) =>
+                    (begin as number) < (other.EndOffset as number) &&
+                    (end as number) > (other.BeginOffset as number)
+            )
+        ) {
+            continue;
+        }
+        selected.push(entity);
+    }
+    return selected.sort((a, b) => (b.BeginOffset as number) - (a.BeginOffset as number));
+}

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

@@ -0,0 +1,82 @@
+import { describe, expect, it, vi } from "vitest";
+import { firstValueFrom, of, toArray } from "rxjs";
+import { ComprehendPiiRedactor } from "../../lib/comprehend/comprehendPiiRedactor.js";
+
+function entity(text: string, value: string, type: string, score = 0.99) {
+    const offset = text.indexOf(value);
+    const begin = Array.from(text.slice(0, offset)).length;
+    return {
+        Type: type,
+        Score: score,
+        BeginOffset: begin,
+        EndOffset: begin + Array.from(value).length,
+    };
+}
+
+describe("ComprehendPiiRedactor", () => {
+    it("redacts Unicode-safe PII spans", async () => {
+        const text = "🔒 Contact José at jose@example.com";
+        const client = {
+            send: vi.fn().mockResolvedValue({
+                Entities: [entity(text, "José", "NAME"), entity(text, "jose@example.com", "EMAIL")],
+            }),
+        };
+
+        await expect(new ComprehendPiiRedactor({ client }).redact(text)).resolves.toBe(
+            "🔒 Contact [NAME] at [EMAIL]"
+        );
+    });
+
+    it("redacts endpoint options without mutating them", async () => {
+        const prompt = "Email jane@example.com";
+        const client = {
+            send: vi.fn().mockResolvedValue({
+                Entities: [entity(prompt, "jane@example.com", "EMAIL")],
+            }),
+        };
+        const endpoint = { chat: vi.fn(async (options) => options.prompt) };
+        const input = { prompt };
+        const redactor = new ComprehendPiiRedactor({ client });
+
+        await expect(
+            firstValueFrom(of(input).pipe(redactor.endpointOperator(endpoint)))
+        ).resolves.toBe("Email [EMAIL]");
+        expect(input.prompt).toBe(prompt);
+    });
+
+    it("preserves source order while awaiting Comprehend", async () => {
+        let release: (() => void) | undefined;
+        const first = new Promise<{ Entities: never[] }>((resolve) => {
+            release = () => resolve({ Entities: [] });
+        });
+        const client = {
+            send: vi
+                .fn()
+                .mockImplementationOnce(() => first)
+                .mockResolvedValue({ Entities: [] }),
+        };
+        const redactor = new ComprehendPiiRedactor({ client });
+        const result = firstValueFrom(
+            of("first", "second").pipe(redactor.redactOperator(), toArray())
+        );
+
+        await Promise.resolve();
+        expect(client.send).toHaveBeenCalledTimes(1);
+        release?.();
+        await expect(result).resolves.toEqual(["first", "second"]);
+    });
+
+    it("aborts an in-flight request when unsubscribed", async () => {
+        let signal: AbortSignal | undefined;
+        const client = {
+            send: vi.fn((_command, options) => {
+                signal = options?.abortSignal;
+                return new Promise(() => undefined);
+            }),
+        };
+        const subscription = new ComprehendPiiRedactor({ client }).redact$("Jane").subscribe();
+
+        subscription.unsubscribe();
+        expect(signal?.aborted).toBe(true);
+    });
+});

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

@@ -0,0 +1,26 @@
+# AWS Comprehend PII redaction
+
+This example redacts PII before an endpoint receives a prompt. It runs offline by default and uses
+the real Amazon Comprehend client when `USE_REAL_AWS=true`.
+
+```sh
+npm install
+npm run demo
+```
+
+Expected output:
+
+```text
+Original: 🔒 Contact jane@example.com before sharing this prompt.
+Endpoint received: 🔒 Contact [EMAIL] before sharing this prompt.
+```
+
+For a live run, provide AWS credentials through the standard AWS credential provider chain and set
+`USE_REAL_AWS=true AWS_REGION=us-east-1`. Never put AWS keys in this repository.
+
+## Loom checklist
+
+1. Show issue #290 and its completion criteria.
+2. Run the offline demo and focused test file.
+3. Show the RxJS endpoint chain and cancellation test.
+4. Upload the recording to Loom and add its URL to the pull request.

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

@@ -0,0 +1,18 @@
+{
+    "name": "aws-comprehend-redaction-example",
+    "version": "1.0.0",
+    "private": true,
+    "type": "module",
+    "scripts": {
+        "build": "tsc",
+        "demo": "npm run build && node dist/index.js"
+    },
+    "dependencies": {
+        "@arakoodev/edgechains.js": "file:../../arakoodev",
+        "rxjs": "^7.8.2"
+    },
+    "devDependencies": {
+        "@types/node": "^20.17.2",
+        "typescript": "^5.6.3"
+    }
+}

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

@@ -0,0 +1,34 @@
+import { ComprehendPiiRedactor, type ComprehendClientLike } from "@arakoodev/edgechains.js/ai";
+import { firstValueFrom, of } from "rxjs";
+
+const prompt = "🔒 Contact jane@example.com before sharing this prompt.";
+const email = "jane@example.com";
+const start = Array.from(prompt.slice(0, prompt.indexOf(email))).length;
+const offlineClient: ComprehendClientLike = {
+    async send() {
+        return {
+            $metadata: {},
+            Entities: [
+                {
+                    Type: "EMAIL",
+                    Score: 0.999,
+                    BeginOffset: start,
+                    EndOffset: start + Array.from(email).length,
+                },
+            ],
+        };
+    },
+};
+
+const redactor = new ComprehendPiiRedactor({
+    client: process.env.USE_REAL_AWS === "true" ? undefined : offlineClient,
+    region: process.env.AWS_REGION || "us-east-1",
+});
+const endpoint = {
+    async chat(options: { prompt: string }) {
+        return `Endpoint received: ${options.prompt}`;
+    },
+};
+
+console.log("Original:", prompt);
+console.log(await firstValueFrom(of({ prompt }).pipe(redactor.endpointOperator(endpoint))));

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

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