From cc68dc868835557e2fa0ec20594fbbca8ef97ef2 Mon Sep 17 00:00:00 2001 From: Luvi-1 <115382358+Luvi-1@users.noreply.github.com> Date: Sun, 2 Aug 2026 23:33:05 -0700 Subject: [PATCH] Add AWS Comprehend PII redaction --- JS/edgechains/arakoodev/package.json | 2 + JS/edgechains/arakoodev/src/ai/src/index.ts | 7 + .../lib/comprehend/comprehendPiiRedactor.ts | 160 ++++++++++++++++++ .../comprehend/comprehendPiiRedactor.test.ts | 82 +++++++++ .../aws-comprehend-redaction/README.md | 26 +++ .../aws-comprehend-redaction/package.json | 18 ++ .../aws-comprehend-redaction/src/index.ts | 34 ++++ .../aws-comprehend-redaction/tsconfig.json | 12 ++ 8 files changed, 341 insertions(+) create mode 100644 JS/edgechains/arakoodev/src/ai/src/lib/comprehend/comprehendPiiRedactor.ts create mode 100644 JS/edgechains/arakoodev/src/ai/src/tests/comprehend/comprehendPiiRedactor.test.ts create mode 100644 JS/edgechains/examples/aws-comprehend-redaction/README.md create mode 100644 JS/edgechains/examples/aws-comprehend-redaction/package.json create mode 100644 JS/edgechains/examples/aws-comprehend-redaction/src/index.ts create mode 100644 JS/edgechains/examples/aws-comprehend-redaction/tsconfig.json diff --git a/JS/edgechains/arakoodev/package.json b/JS/edgechains/arakoodev/package.json index 0b0bd3784..4d48cb8bf 100644 --- a/JS/edgechains/arakoodev/package.json +++ b/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", diff --git a/JS/edgechains/arakoodev/src/ai/src/index.ts b/JS/edgechains/arakoodev/src/ai/src/index.ts index 2c98f37dc..7157f1357 100644 --- a/JS/edgechains/arakoodev/src/ai/src/index.ts +++ b/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"; diff --git a/JS/edgechains/arakoodev/src/ai/src/lib/comprehend/comprehendPiiRedactor.ts b/JS/edgechains/arakoodev/src/ai/src/lib/comprehend/comprehendPiiRedactor.ts new file mode 100644 index 000000000..dcdc61d97 --- /dev/null +++ b/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; +} + +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 { + chat(options: Options): Promise; +} + +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 { + 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( + options: T, + abortSignal?: AbortSignal + ): Promise { + 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 { + 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 { + return concatMap((text) => this.redact$(text)); + } + + endpointOperator( + endpoint: ChatEndpoint + ): OperatorFunction { + 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)); +} diff --git a/JS/edgechains/arakoodev/src/ai/src/tests/comprehend/comprehendPiiRedactor.test.ts b/JS/edgechains/arakoodev/src/ai/src/tests/comprehend/comprehendPiiRedactor.test.ts new file mode 100644 index 000000000..e4a52eb7a --- /dev/null +++ b/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); + }); +}); diff --git a/JS/edgechains/examples/aws-comprehend-redaction/README.md b/JS/edgechains/examples/aws-comprehend-redaction/README.md new file mode 100644 index 000000000..dd74e5499 --- /dev/null +++ b/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. diff --git a/JS/edgechains/examples/aws-comprehend-redaction/package.json b/JS/edgechains/examples/aws-comprehend-redaction/package.json new file mode 100644 index 000000000..2d899cc3b --- /dev/null +++ b/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" + } +} diff --git a/JS/edgechains/examples/aws-comprehend-redaction/src/index.ts b/JS/edgechains/examples/aws-comprehend-redaction/src/index.ts new file mode 100644 index 000000000..69cbb6808 --- /dev/null +++ b/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)))); diff --git a/JS/edgechains/examples/aws-comprehend-redaction/tsconfig.json b/JS/edgechains/examples/aws-comprehend-redaction/tsconfig.json new file mode 100644 index 000000000..007c59350 --- /dev/null +++ b/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"] +}