feat: integrate AWS Comprehend as a utility to redact data

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",

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";

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<string> {
+        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;
+        }
+    }
+}

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("");
+        });
+    });
+});
+