Free lesson · GenAI Security Engineering

Build bidirectional PII redaction pipeline

Implement request-side PII redaction before LLM calls and response-side detection. Build reversible tokenization for authorized de-anonymization.

Course: AI Security Engineering · Chapter 8 · PII Leakage Engineering

Free to read — no subscription required.

Introduction

Engineers often focus PII defenses on what enters the model while ignoring what comes back out — a blind spot that lets the LLM echo a phone number or social security number directly into the caller's response. Protecting LLM traffic in both directions requires separate analysis passes: one that sanitizes the request before the model sees it, and one that enforces data-loss-prevention policy on the response before it reaches the user. By the end of this lesson you will have built a ReversibleRedactionPipeline and a BidirectionalPIIGuard that use Presidio to detect, tokenize, and policy-gate PII across both request and response flows, with configurable DLP tiers wired to LiteLLM hooks.

Key Terminology

  • Reversible Redaction — replacing each detected PII span with a structured UUID-backed token (e.g., <EMAIL_ADDRESS_3f2a1b4c>) so the original value can be restored after the model responds, unlike permanent anonymization that discards the value entirely.
  • Vault — the per-request dictionary inside ReversibleRedactionPipeline._vaults that maps each token back to the original PII string it replaced; identified by a vault_id returned from redact_request and destroyed on the first call to restore_response.
  • DLP Policy Tier — one of three configurable entity-type sets (BLOCK_ENTITIES, REDACT_ENTITIES, LOG_ONLY_ENTITIES) that determine which DLPAction BidirectionalPIIGuard.process_response applies when Presidio detects PII in an LLM response.
  • DLPAction — the enum (ALLOW, REDACT, BLOCK) returned by process_response to signal whether the response should be delivered unchanged, overwritten with anonymized text, or halted before reaching the caller.
  • AnalyzerEngine — Presidio's entity-detection component, invoked in both redact_request and process_response to locate PII spans and their confidence scores; results are filtered by the configurable score_threshold before any anonymization step.
  • async_post_call_success_hook — LiteLLM's post-completion callback where BidirectionalPIIGuard enforces response-side DLP: a BLOCK outcome raises an exception that halts delivery, while a REDACT outcome overwrites response.choices[0].message.content before the caller receives it.

Concepts

The Blind Spot in One-Directional Guards

Engineers commonly wire PII detection only at the input boundary — scrubbing prompts before sending them to the model — and assume that eliminates the exposure. It does not. A model can echo values it inferred from surrounding context, reproduce entities embedded in few-shot examples, or reflect structured data back verbatim in a formatted response. Redacting a phone number in the prompt does not stop the model from generating one in its reply. Real data-loss prevention requires an independent analysis pass on the response before it reaches the caller, which is the founding design principle of BidirectionalPIIGuard.

Vault-Based Tokenization: Redaction That Remembers

Simple anonymization permanently destroys the original value, which breaks any use-case where the application legitimately needs to present PII to the user after the model has processed the sanitized text. Reversible redaction avoids that trade-off by replacing each detected span with a structured UUID-backed token and storing the token-to-original mapping in a per-request vault. The vault is keyed by a vault_id returned alongside the sanitized text; the model receives only tokens like <EMAIL_ADDRESS_3f2a1b4c> and cannot see the originals. When restore_response is called with the matching vault_id, it pops the vault entry and swaps every token back to its original value.

The pop-on-read design is deliberate: once the vault entry is consumed, the mapping is gone, making post-call exposure impossible without re-running the pipeline. The threading.Lock wrapping all vault reads and writes ensures that concurrent requests sharing the same ReversibleRedactionPipeline instance cannot corrupt each other's mappings or bleed one caller's PII into another vault (see Code Walkthrough).

DLP Policy Tiers: Classifying Response Risk

Not all PII in an LLM response carries the same compliance weight. A social security number echoed verbatim is a regulatory incident; a person's first name in a summary may be acceptable. BidirectionalPIIGuard models this gradient with three entity-type sets that map directly to enforcement actions through the DLPAction enum:

  • BLOCK: entities in BLOCK_ENTITIES (e.g., US_SSN, CREDIT_CARD) suppress the response entirely. In the async_post_call_success_hook, this surfaces as a raised exception that halts delivery before response.choices[0].message.content is returned to the caller.
  • REDACT: entities in REDACT_ENTITIES (e.g., PERSON, EMAIL_ADDRESS, PHONE_NUMBER) are anonymized in the response text, and the result overwrites response.choices[0].message.content in-place before the caller receives it.
  • ALLOW: entities limited to LOG_ONLY_ENTITIES (e.g., LOCATION, DATE_TIME) pass through without mutation; the vault restores the original request-side values and the response is delivered normally.

The evaluation order is strict — BLOCK is checked before REDACT before ALLOW — so a response containing both a credit card number and an email address is blocked outright rather than partially redacted (see Code Walkthrough).

Loading diagram...

Code Walkthrough

Now that you understand the vault management strategy, DLP policy tiers, and thread-safety requirements from the Concepts section, the code below assembles those ideas into two production-ready classes.

The first class, ReversibleRedactionPipeline, handles the request side. Its redact_request method runs Presidio's AnalyzerEngine to detect entities above a configurable confidence threshold, replaces each span with a UUID-backed token stored in an internal vault protected by a threading.Lock, and returns both the sanitized text and a vault identifier. restore_response uses that identifier to pop the vault entry and swap tokens back to original values.

Code snippetpython
1import uuid 2import threading 3from typing import Dict, Tuple 4from presidio_analyzer import AnalyzerEngine 5from presidio_anonymizer import AnonymizerEngine 6from presidio_anonymizer.entities import OperatorConfig 7 8class ReversibleRedactionPipeline: 9 def __init__(self, score_threshold: float = 0.7): 10 self._analyzer = AnalyzerEngine() 11 self._anonymizer = AnonymizerEngine() 12 self._vaults: Dict[str, Dict[str, str]] = {} 13 self._lock = threading.Lock() 14 self.score_threshold = score_threshold 15 16 def redact_request(self, text: str, language: str = "en") -> Tuple[str, str]: 17 results = self._analyzer.analyze( 18 text=text, language=language, score_threshold=self.score_threshold 19 ) 20 vault: Dict[str, str] = {} 21 operator_config: Dict[str, OperatorConfig] = {} 22 for result in results: 23 token = f"<{result.entity_type}_{uuid.uuid4().hex[:8]}>" 24 vault[token] = text[result.start:result.end] 25 operator_config[result.entity_type] = OperatorConfig( 26 "replace", {"new_value": token} 27 ) 28 anonymized = self._anonymizer.anonymize( 29 text=text, analyzer_results=results, operators=operator_config 30 ) 31 vault_id = uuid.uuid4().hex 32 with self._lock: 33 self._vaults[vault_id] = vault 34 return anonymized.text, vault_id 35 36 def restore_response(self, text: str, vault_id: str) -> str: 37 with self._lock: 38 vault = self._vaults.pop(vault_id, {}) 39 for token, original in vault.items(): 40 text = text.replace(token, original) 41 return text

The second class, BidirectionalPIIGuard, wraps the pipeline and enforces DLP policy on the response side. It evaluates entity types found in the LLM's output against three configurable sets — BLOCK_ENTITIES, REDACT_ENTITIES, and LOG_ONLY_ENTITIES. When wired to a LiteLLM async_post_call_success_hook, a BLOCK result raises an exception halting delivery, while a REDACT result overwrites response.choices[0].message.content before the caller receives it.

Code snippetpython
1from enum import Enum 2 3class DLPAction(Enum): 4 ALLOW = "allow" 5 REDACT = "redact" 6 BLOCK = "block" 7 8BLOCK_ENTITIES = {"US_SSN", "CREDIT_CARD"} 9REDACT_ENTITIES = {"PERSON", "EMAIL_ADDRESS", "PHONE_NUMBER"} 10LOG_ONLY_ENTITIES = {"LOCATION", "DATE_TIME"} 11 12class BidirectionalPIIGuard: 13 def __init__(self): 14 self._pipeline = ReversibleRedactionPipeline() 15 self._response_analyzer = AnalyzerEngine() 16 self._response_anonymizer = AnonymizerEngine() 17 18 def process_request(self, prompt: str) -> Tuple[str, str]: 19 return self._pipeline.redact_request(prompt) 20 21 def process_response( 22 self, response_text: str, vault_id: str 23 ) -> Tuple[str, DLPAction]: 24 results = self._response_analyzer.analyze( 25 text=response_text, language="en" 26 ) 27 found = {r.entity_type for r in results} 28 if found & BLOCK_ENTITIES: 29 return "", DLPAction.BLOCK 30 if found & REDACT_ENTITIES: 31 redacted = self._response_anonymizer.anonymize( 32 text=response_text, analyzer_results=results 33 ) 34 restored = self._pipeline.restore_response(redacted.text, vault_id) 35 return restored, DLPAction.REDACT 36 restored = self._pipeline.restore_response(response_text, vault_id) 37 return restored, DLPAction.ALLOW

Confirm that calling guard.process_request("Email me at alice@example.com") returns a sanitized prompt with the address replaced by a token, and that a subsequent guard.process_response(...) using the same vault ID restores alice@example.com in the final output when the DLP action resolves to ALLOW.

Do's and Don'ts

Now that you have worked through the implementation, the practices below separate a durable approach from a fragile one.

Do's

  1. Do build per-entity OperatorConfig entries with UUID-backed tokens in redact_request — each detected span must get its own hex token (e.g., <EMAIL_ADDRESS_a3f7b2c1>) so that restore_response can map every token back to its unique original value; reusing a shared placeholder like <EMAIL_ADDRESS> collapses all instances of a type into one vault key and breaks reversal whenever a prompt contains more than one address.
  2. Do protect every read and write of _vaults with self._lock — LiteLLM hooks execute concurrently across requests, so two threads can interleave vault inserts or pops without a threading.Lock, silently corrupting the token-to-original mapping or dropping vault entries mid-flight, which causes tokens to survive unexpanded into the final response.
  3. Do run a dedicated AnalyzerEngine pass on the LLM response inside BidirectionalPIIGuard.process_response — the model can echo, synthesize, or infer PII that was never present in the original prompt; a request-only analysis pass leaves SSNs, credit card numbers, and phone numbers that appear only in the model's output completely ungated by the BLOCK_ENTITIES and REDACT_ENTITIES tiers.

Don'ts

  1. Don't call restore_response more than once with the same vault_idrestore_response uses dict.pop(vault_id, {}) to atomically remove the entry; a second call receives an empty dict and silently leaves tokens like <PHONE_NUMBER_d9e2a1f4> unexpanded in the caller's output, leaking the token string instead of the original value with no error raised.
  2. Don't place the same entity type in both BLOCK_ENTITIES and REDACT_ENTITIESprocess_response evaluates BLOCK_ENTITIES first; any type listed in both sets always triggers the block branch, making the redact branch dead code for that type and obscuring whether the intent is delivery prevention or in-line masking; keep the three DLP tiers mutually exclusive.
  3. Don't omit or set score_threshold to 0.0 — Presidio will surface low-confidence partial matches on common words and short numeric strings, inserting spurious vault tokens that either bloat _vaults memory across concurrent requests or corrupt non-PII text when restore_response substitutes them back; the default of 0.7 is the floor that keeps false-positive token injection manageable.

This lesson is free to read. Its 3 hands-on labs — real code, in a cloud IDE — are part of the GenAI Security Engineering subscription.

From · cancel anytime · Already a subscriber? Sign in →

More free lessons in AI Security Engineering

All free lessons in GenAI Security Engineering