diff --git a/JS/edgechains/arakoodev/package.json b/JS/edgechains/arakoodev/package.json index 0b0bd3784..58d66467e 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.1076.0", "@babel/core": "^7.24.4", "@babel/preset-env": "^7.24.4", "@hono/node-server": "^0.6.0", diff --git a/JS/edgechains/arakoodev/src/ai/src/index.ts b/JS/edgechains/arakoodev/src/ai/src/index.ts index 2c98f37dc..cf37da488 100644 --- a/JS/edgechains/arakoodev/src/ai/src/index.ts +++ b/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"; 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..1cadbcaf8 --- /dev/null +++ b/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; +}; + +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; +} + +export class ComprehendPiiRedactor { + private readonly client: ComprehendLikeClient; + private readonly languageCode: LanguageCode; + private readonly minScore: number; + private readonly entityTypes?: Set; + 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 { + 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 { + 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( + chatOptions: T, + ): Promise { + 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( + endpoint: ChatEndpoint, + ): ChatEndpoint { + return { + chat: async (chatOptions: TOptions) => { + return endpoint.chat(await this.redactChatOptions(chatOptions)); + }, + }; + } + + asOperator(): ( + value: T, + ) => Promise { + return async (value: T) => { + if (typeof value === "string") { + return (await this.redact(value)) as T; + } + + return this.redactChatOptions(value) as Promise; + }; + } + + 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; + } +} 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..77aba1fc8 --- /dev/null +++ b/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(); + + await expect(redact("Hi Jane")).resolves.toBe("Hi [NAME]"); + }); +}); diff --git a/JS/edgechains/examples/comprehend-pii-redaction/README.md b/JS/edgechains/examples/comprehend-pii-redaction/README.md new file mode 100644 index 000000000..ebf0cb5c6 --- /dev/null +++ b/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." +``` diff --git a/JS/edgechains/examples/comprehend-pii-redaction/jsonnet/main.jsonnet b/JS/edgechains/examples/comprehend-pii-redaction/jsonnet/main.jsonnet new file mode 100644 index 000000000..8304445bd --- /dev/null +++ b/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, +} diff --git a/JS/edgechains/examples/comprehend-pii-redaction/package.json b/JS/edgechains/examples/comprehend-pii-redaction/package.json new file mode 100644 index 000000000..44b668af4 --- /dev/null +++ b/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": {} +} diff --git a/JS/edgechains/examples/comprehend-pii-redaction/src/edgechains-ai.d.ts b/JS/edgechains/examples/comprehend-pii-redaction/src/edgechains-ai.d.ts new file mode 100644 index 000000000..3bda74009 --- /dev/null +++ b/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; + }; + + export class ComprehendPiiRedactor { + constructor(options?: { + client?: ComprehendLikeClient; + region?: string; + languageCode?: string; + minScore?: number; + }); + + redact(text: string): Promise; + } +} diff --git a/JS/edgechains/examples/comprehend-pii-redaction/src/index.ts b/JS/edgechains/examples/comprehend-pii-redaction/src/index.ts new file mode 100644 index 000000000..2b86e87e0 --- /dev/null +++ b/JS/edgechains/examples/comprehend-pii-redaction/src/index.ts @@ -0,0 +1,68 @@ +import { + ComprehendPiiRedactor, + type ComprehendLikeClient, +} from "@arakoodev/edgechains.js/ai"; +import Jsonnet from "@arakoodev/jsonnet"; +import path from "path"; +import { fileURLToPath } from "url"; + +const __dirname = path.dirname(fileURLToPath(import.meta.url)); +const jsonnet = new Jsonnet(); + +const request = + process.argv.slice(2).join(" ") || + "Jane Doe emailed jane@example.com about invoice 555-010-9999."; + +jsonnet.extString("request", request); + +const config = JSON.parse( + jsonnet.evaluateFile(path.join(__dirname, "../jsonnet/main.jsonnet")), +); + +const mockClient: ComprehendLikeClient = { + async send(command: any) { + const text = command.input.Text || ""; + const nameStart = text.indexOf("Jane Doe"); + const emailStart = text.indexOf("jane@example.com"); + const phoneStart = text.indexOf("555-010-9999"); + + return { + Entities: [ + nameStart >= 0 + ? { + Type: "NAME", + BeginOffset: nameStart, + EndOffset: nameStart + 8, + Score: 0.99, + } + : undefined, + emailStart >= 0 + ? { + Type: "EMAIL", + BeginOffset: emailStart, + EndOffset: emailStart + 16, + Score: 0.99, + } + : undefined, + phoneStart >= 0 + ? { + Type: "PHONE", + BeginOffset: phoneStart, + EndOffset: phoneStart + 12, + Score: 0.99, + } + : undefined, + ].filter(Boolean) as any, + }; + }, +}; + +const redactor = new ComprehendPiiRedactor({ + client: process.env.USE_REAL_AWS === "true" ? undefined : mockClient, + region: process.env.AWS_REGION || "us-east-1", + languageCode: config.languageCode, + minScore: config.minScore, +}); + +const safePrompt = await redactor.redact(config.prompt); +console.log(safePrompt); diff --git a/JS/edgechains/examples/comprehend-pii-redaction/tsconfig.json b/JS/edgechains/examples/comprehend-pii-redaction/tsconfig.json new file mode 100644 index 000000000..e38187bac --- /dev/null +++ b/JS/edgechains/examples/comprehend-pii-redaction/tsconfig.json @@ -0,0 +1,13 @@ +{ + "compilerOptions": { + "target": "ES2022", + "module": "NodeNext", + "moduleResolution": "NodeNext", + "esModuleInterop": true, + "strict": true, + "skipLibCheck": true, + "outDir": "dist", + "rootDir": "src" + }, + "include": ["src"] +}