From f4636e54c4f53355548a47dbb0cb45b16b1007eb Mon Sep 17 00:00:00 2001 From: Albert2026py Date: Wed, 1 Jul 2026 15:12:53 +0800 Subject: [PATCH] Add AWS Comprehend PII redaction utility --- JS/edgechains/arakoodev/package.json | 1 + JS/edgechains/arakoodev/src/ai/src/index.ts | 1 + .../aws-comprehend/awsComprehendRedactor.ts | 121 ++++++++++++++++++ .../src/tests/awsComprehendRedactor.test.ts | 68 ++++++++++ .../aws-comprehend-redaction/package.json | 16 +++ .../aws-comprehend-redaction/readme.md | 17 +++ .../aws-comprehend-redaction/src/index.ts | 21 +++ .../aws-comprehend-redaction/tsconfig.json | 11 ++ 8 files changed, 256 insertions(+) create mode 100644 JS/edgechains/arakoodev/src/ai/src/lib/aws-comprehend/awsComprehendRedactor.ts create mode 100644 JS/edgechains/arakoodev/src/ai/src/tests/awsComprehendRedactor.test.ts create mode 100644 JS/edgechains/examples/aws-comprehend-redaction/package.json create mode 100644 JS/edgechains/examples/aws-comprehend-redaction/readme.md 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..f3f510b0a 100644 --- a/JS/edgechains/arakoodev/package.json +++ b/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", diff --git a/JS/edgechains/arakoodev/src/ai/src/index.ts b/JS/edgechains/arakoodev/src/ai/src/index.ts index 2c98f37dc..7ffb20b27 100644 --- a/JS/edgechains/arakoodev/src/ai/src/index.ts +++ b/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"; diff --git a/JS/edgechains/arakoodev/src/ai/src/lib/aws-comprehend/awsComprehendRedactor.ts b/JS/edgechains/arakoodev/src/ai/src/lib/aws-comprehend/awsComprehendRedactor.ts new file mode 100644 index 000000000..17a5693ff --- /dev/null +++ b/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; +}; + +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 { + 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(input: T): Promise { + 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(input: T, next: (redactedInput: T) => Promise): Promise { + 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, + 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); +} diff --git a/JS/edgechains/arakoodev/src/ai/src/tests/awsComprehendRedactor.test.ts b/JS/edgechains/arakoodev/src/ai/src/tests/awsComprehendRedactor.test.ts new file mode 100644 index 000000000..cc77ac48d --- /dev/null +++ b/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]"); + }); +}); 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..fa02ed9a5 --- /dev/null +++ b/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" + } +} 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..4c48f4fce --- /dev/null +++ b/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 +``` 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..7368148f9 --- /dev/null +++ b/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); 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..e6d24ebfe --- /dev/null +++ b/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"] +}