DEV Community

Tracepilot
Tracepilot

Posted on

Build a PII Redaction Pipeline with AWS Comprehend & EdgeChains

#ai #debugging #python

Build a PII Redaction Pipeline with AWS Comprehend & EdgeChains

What we're building: A chainable PII redaction utility that automatically strips sensitive information (names, emails, SSNs, etc.) from prompts before they reach your LLM — using AWS Comprehend and EdgeChains' observable architecture.

Prerequisites

  • Node.js 18+
  • An AWS account with Comprehend access
  • The EdgeChains JS SDK installed: npm install edgechains
  • AWS SDK v3: npm install @aws-sdk/client-comprehend
  • An OpenAI API key (for the demo)

Step 1: Set up your project 

mkdir pii-redactor
cd pii-redactor
npm init -y
npm install edgechains @aws-sdk/client-comprehend openai dotenv
Enter fullscreen mode Exit fullscreen mode

Create .env:

AWS_ACCESS_KEY_ID=your_key_here
AWS_SECRET_ACCESS_KEY=your_secret_here
AWS_REGION=us-east-1
OPENAI_API_KEY=sk-...
Enter fullscreen mode Exit fullscreen mode

Step 2: Create the Comprehend PII Redactor class

This is the core — a class that wraps AWS Comprehend's detectPiiEntities and redacts matches. It's designed to be chained with EdgeChains' Endpoint classes.

// redactor.ts
import {
  ComprehendClient,
  DetectPiiEntitiesCommand,
  type Entity,
} from "@aws-sdk/client-comprehend";
import { Endpoint } from "edgechains";

type RedactConfig = {
  /** Which PII types to redact. Default: all */
  entityTypes?: string[];
  /** Replacement character. Default: '*' */
  maskChar?: string;
};

export class ComprehendPIIRedactor extends Endpoint {
  private client: ComprehendClient;
  private config: Required<RedactConfig>;

  constructor(config: RedactConfig = {}) {
    super();
    this.client = new ComprehendClient({
      region: process.env.AWS_REGION || "us-east-1",
    });
    this.config = {
      entityTypes: config.entityTypes || [],
      maskChar: config.maskChar || "*",
    };
  }

  /**
   * Detect PII entities in the input text
   */
  private async detectPII(text: string): Promise<Entity[]> {
    const command = new DetectPiiEntitiesCommand({
      Text: text,
      LanguageCode: "en",
    });
    const response = await this.client.send(command);
    return response.Entities || [];
  }

  /**
   * Redact detected PII entities from the text
   */
  private redactPII(text: string, entities: Entity[]): string {
    // Sort entities by offset in reverse to avoid index shifting
    const sorted = [...entities]
      .filter((e) => {
        if (this.config.entityTypes.length === 0) return true;
        return this.config.entityTypes.includes(e.Type!);
      })
      .sort((a, b) => (b.BeginOffset || 0) - (a.BeginOffset || 0));

    let result = text;
    for (const entity of sorted) {
      const start = entity.BeginOffset || 0;
      const end = entity.EndOffset || 0;
      const original = result.slice(start, end);
      const masked = this.config.maskChar.repeat(original.length);
      result = result.slice(0, start) + masked + result.slice(end);
    }
    return result;
  }

  /**
   * The main redact method — chainable with EdgeChains
   * Returns an observable that emits the redacted text
   */
  async redact(input: string): Promise<string> {
    const entities = await this.detectPII(input);
    if (entities.length === 0) {
      console.log("✅ No PII detected");
      return input;
    }

    console.log(`🔍 Found ${entities.length} PII entities:`, 
      entities.map((e) => `${e.Type} at [${e.BeginOffset}-${e.EndOffset}]`).join(", ")
    );

    return this.redactPII(input, entities);
  }
}
Enter fullscreen mode Exit fullscreen mode

Step 3: Chain it with an LLM endpoint

Now the magic — chain the redactor with an OpenAI endpoint so every prompt gets sanitized before the LLM sees it.

// chain.ts
import { Endpoint } from "edgechains";
import OpenAI from "openai";
import { ComprehendPIIRedactor } from "./redactor";

class OpenAIEndpoint extends Endpoint {
  private client: OpenAI;

  constructor() {
    super();
    this.client = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
  }

  async generate(prompt: string): Promise<string> {
    const response = await this.client.chat.completions.create({
      model: "gpt-4o-mini",
      messages: [{ role: "user", content: prompt }],
    });
    return response.choices[0]?.message?.content || "";
  }
}

async function main() {
  const redactor = new ComprehendPIIRedactor({
    entityTypes: ["NAME", "EMAIL", "PHONE", "SSN"], // only redact these types
    maskChar: "█",
  });

  const llm = new OpenAIEndpoint();

  // The chain: prompt → redact PII → send to LLM
  const rawPrompt = `Hi, I'm John Smith. My email is john.smith@example.com and my SSN is 123-45-6789. Can you help me reset my password?`;

  console.log("📝 Original prompt:", rawPrompt);

  // Chain the endpoints
  const redactedPrompt = await redactor.redact(rawPrompt);
  console.log("🛡️  Redacted prompt:", redactedPrompt);

  const response = await llm.generate(redactedPrompt);
  console.log("🤖 LLM response:", response);
}

main().catch(console.error);
Enter fullscreen mode Exit fullscreen mode

Step 4: Add observability with TracePilot

This is where debugging becomes obvious. Add TracePilot to see exactly what PII was redacted, when, and how the LLM responded.

npm install tracepilot-sdk
Enter fullscreen mode Exit fullscreen mode 

typescript
// tracepilot-chain.ts
import { TracePilot } from "tracepilot-sdk";
import { ComprehendPIIRedactor } from "./redactor";
import { Endpoint } from "edgechains";
import OpenAI from "openai";

const tp = new TracePilot(process.env.TRACEPILOT_API_KEY!);

class ObservableRedactor extends ComprehendPIIRedactor {
  async redact(input: string): Promise<string> {
    return await tp.wrapToolCall(
      "comprehend-pii-redaction",
      () => super.redact(input),
      undefined, // parent span — set if chaining
      1,         // step order
      true       // destructive? Yes — we're modifying data
    );
  }
}

class ObservableLLM extends Endpoint {
  private client: OpenAI;

  constructor() {
    super();
    this.client = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
  }

  async generate(prompt: string): Promise<string> {
    const { result } = await tp.wrapOpenAI(
      () =>
        this.client.chat.completions.create({
          model: "gpt-4o-mini",
          messages: [{ role: "user", content: prompt }],
        }),
      [{ role: "user", content: prompt }]
    );
    return result.choices[0]?.message?.content || "";
  }
}

async function main() {
  await tp.startTrace("pii-redaction-pipeline");

  const redactor = new ObservableRedactor({
    entityTypes: ["NAME", "EMAIL", "PHONE", "SSN"],
    maskChar: "█",
  });

  const llm = new ObservableLLM();

  const rawPrompt = `Call me at 555-123-4567 or email support@bank.com. My name is Jane Doe.`;
  const redactedPrompt = await redactor.redact(rawPrompt);
  const response = await llm.generate(redactedPrompt);

  console.log("Final response:",

---

**Debugging AI agents shouldn't feel like reading The Matrix.** 
Join other engineers who are building reliable autonomous workflows in our community: [TracePilot Discord](https://discord.gg/KzXRAXFM8)
Enter fullscreen mode Exit fullscreen mode

Top comments (0)