diff --git a/JS/edgechains/arakoodev/package.json b/JS/edgechains/arakoodev/package.json index 0b0bd3784..dda664c35 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.1083.0", "@babel/core": "^7.24.4", "@babel/preset-env": "^7.24.4", "@hono/node-server": "^0.6.0", @@ -65,6 +66,7 @@ "@types/node": "^20.17.2", "@types/pdf-parse": "^1.1.4", "@types/ws": "^8.5.12", + "aws-sdk-client-mock": "^4.1.0", "buffer": "^6.0.3", "crypto-browserify": "^3.12.1", "jest": "^29.7.0", diff --git a/JS/edgechains/arakoodev/src/ai/src/index.ts b/JS/edgechains/arakoodev/src/ai/src/index.ts index 2c98f37dc..7522c252a 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 { AwsComprehend } from "./lib/aws/comprehend.js"; diff --git a/JS/edgechains/arakoodev/src/ai/src/lib/aws/comprehend.ts b/JS/edgechains/arakoodev/src/ai/src/lib/aws/comprehend.ts new file mode 100644 index 000000000..4e94f2ca8 --- /dev/null +++ b/JS/edgechains/arakoodev/src/ai/src/lib/aws/comprehend.ts @@ -0,0 +1,39 @@ +import { ComprehendClient, DetectPiiEntitiesCommand, ComprehendClientConfig } from "@aws-sdk/client-comprehend"; + +export class AwsComprehend { + private client: ComprehendClient; + + constructor(options?: ComprehendClientConfig) { + this.client = new ComprehendClient(options || { region: "us-east-1" }); + } + + async redact(text: string, languageCode: string = "en"): Promise { + if (!text) return text; + const command = new DetectPiiEntitiesCommand({ + Text: text, + LanguageCode: languageCode as any + }); + + try { + const response = await this.client.send(command); + const entities = response.Entities || []; + + // Sort entities by BeginOffset in descending order + // This ensures that string replacement from end to start doesn't affect previous offsets + entities.sort((a, b) => (b.BeginOffset || 0) - (a.BeginOffset || 0)); + + let redactedText = text; + for (const entity of entities) { + const begin = entity.BeginOffset; + const end = entity.EndOffset; + if (begin !== undefined && end !== undefined) { + redactedText = redactedText.substring(0, begin) + `[${entity.Type || "PII"}]` + redactedText.substring(end); + } + } + return redactedText; + } catch (error) { + console.error("Error redacting text using AWS Comprehend:", error); + throw error; + } + } +} diff --git a/JS/edgechains/arakoodev/src/ai/src/tests/awsComprehend.test.ts b/JS/edgechains/arakoodev/src/ai/src/tests/awsComprehend.test.ts new file mode 100644 index 000000000..83b17e1d5 --- /dev/null +++ b/JS/edgechains/arakoodev/src/ai/src/tests/awsComprehend.test.ts @@ -0,0 +1,79 @@ +import { AwsComprehend } from "../lib/aws/comprehend.ts"; +import { describe, test, expect, vi, beforeEach } from "vitest"; + +const sendMock = vi.fn(); + +vi.mock("@aws-sdk/client-comprehend", () => { + return { + ComprehendClient: vi.fn().mockImplementation(() => { + return { + send: sendMock + }; + }), + DetectPiiEntitiesCommand: vi.fn().mockImplementation((args) => args) + }; +}); + +describe("AwsComprehend", () => { + beforeEach(() => { + sendMock.mockReset(); + }); + + describe("redact", () => { + test("should redact PII data from prompt", async () => { + const mockResponse = { + Entities: [ + { + Type: "NAME", + BeginOffset: 18, + EndOffset: 26, + Score: 0.99, + }, + { + Type: "EMAIL", + BeginOffset: 43, + EndOffset: 63, + Score: 0.99, + } + ] + }; + + sendMock.mockResolvedValue(mockResponse); + + const awsComprehend = new AwsComprehend({ region: "us-east-1", credentials: { accessKeyId: "test", secretAccessKey: "test" } }); + + const prompt = "Hello, my name is John Doe and my email is john.doe@example.com."; + const redactedPrompt = await awsComprehend.redact(prompt); + + expect(sendMock).toHaveBeenCalledTimes(1); + expect(redactedPrompt).toEqual("Hello, my name is [NAME] and my email is [EMAIL]."); + }); + + test("should return original text if no PII is found", async () => { + const mockResponse = { + Entities: [] + }; + + sendMock.mockResolvedValue(mockResponse); + + const awsComprehend = new AwsComprehend({ region: "us-east-1", credentials: { accessKeyId: "test", secretAccessKey: "test" } }); + + const prompt = "Hello, this text is safe."; + const redactedPrompt = await awsComprehend.redact(prompt); + + expect(sendMock).toHaveBeenCalledTimes(1); + expect(redactedPrompt).toEqual("Hello, this text is safe."); + }); + + test("should return empty string if input is empty", async () => { + const awsComprehend = new AwsComprehend(); + + const prompt = ""; + const redactedPrompt = await awsComprehend.redact(prompt); + + expect(sendMock).not.toHaveBeenCalled(); + expect(redactedPrompt).toEqual(""); + }); + }); +}); +