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._vaultsthat maps each token back to the original PII string it replaced; identified by avault_idreturned fromredact_requestand destroyed on the first call torestore_response. - DLP Policy Tier — one of three configurable entity-type sets (
BLOCK_ENTITIES,REDACT_ENTITIES,LOG_ONLY_ENTITIES) that determine whichDLPActionBidirectionalPIIGuard.process_responseapplies when Presidio detects PII in an LLM response. - DLPAction — the enum (
ALLOW,REDACT,BLOCK) returned byprocess_responseto 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_requestandprocess_responseto locate PII spans and their confidence scores; results are filtered by the configurablescore_thresholdbefore any anonymization step. - async_post_call_success_hook — LiteLLM's post-completion callback where
BidirectionalPIIGuardenforces response-side DLP: aBLOCKoutcome raises an exception that halts delivery, while aREDACToutcome overwritesresponse.choices[0].message.contentbefore 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 theasync_post_call_success_hook, this surfaces as a raised exception that halts delivery beforeresponse.choices[0].message.contentis 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 overwritesresponse.choices[0].message.contentin-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).
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
- ✓Do build per-entity
OperatorConfigentries with UUID-backed tokens inredact_request— each detected span must get its own hex token (e.g.,<EMAIL_ADDRESS_a3f7b2c1>) so thatrestore_responsecan 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. - ✓Do protect every read and write of
_vaultswithself._lock— LiteLLM hooks execute concurrently across requests, so two threads can interleave vault inserts or pops without athreading.Lock, silently corrupting the token-to-original mapping or dropping vault entries mid-flight, which causes tokens to survive unexpanded into the final response. - ✓Do run a dedicated
AnalyzerEnginepass on the LLM response insideBidirectionalPIIGuard.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 theBLOCK_ENTITIESandREDACT_ENTITIEStiers.
Don'ts
- ✗Don't call
restore_responsemore than once with the samevault_id—restore_responseusesdict.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. - ✗Don't place the same entity type in both
BLOCK_ENTITIESandREDACT_ENTITIES—process_responseevaluatesBLOCK_ENTITIESfirst; 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. - ✗Don't omit or set
score_thresholdto0.0— Presidio will surface low-confidence partial matches on common words and short numeric strings, inserting spurious vault tokens that either bloat_vaultsmemory across concurrent requests or corrupt non-PII text whenrestore_responsesubstitutes them back; the default of0.7is 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
- Ch 7Enforce supply chain policies in GKE with Binary Authorization
- Ch 8Build bidirectional PII redaction pipelineYou are here
- Ch 8Integrate PII defense with LiteLLM gateway
- Ch 8Deploy PII defense pipeline on GKE
- Ch 9Implement row-level access control for vector stores
- Ch 9Implement embedding integrity and provenance tracking
- Ch 9Encrypt embeddings at rest and in transit