Building a PII Data Extraction and Scrubbing System with BAML

1
class PIIData {
2
  index int
3
  dataType string
4
  value string
5
}
6
7
class PIIExtraction {
8
  privateData PIIData[]
9
  containsSensitivePII bool @description("E.g. SSN")
10
}

1
function ExtractPII(document: string) -> PIIExtraction {
2
  client "openai/gpt-4o-mini"
3
  prompt #"
4
    Extract all personally identifiable information (PII) from the given document. Look for items like:
5
    - Names
6
    - Email addresses
7
    - Phone numbers
8
    - Addresses
9
    - Social security numbers
10
    - Dates of birth
11
    - Any other personal data
12
13
    {{ ctx.output_format }}
14
15
    {{ _.role("user") }} 
16
    
17
    {{ document }}
18
  "#
19
}

1
test BasicPIIExtraction {
2
  functions [ExtractPII]
3
  args {
4
    document #"
5
      John Doe was born on 01/02/1980. 
6
      His email is john.doe@email.com and phone is 555-123-4567.
7
      He lives at 123 Main St, Springfield, IL 62704.
8
    "#
9
  }
10
}
11
12
test EmptyDocument {
13
  functions [ExtractPII]
14
  args {
15
    document "This document contains no PII data."
16
  }
17
}

1
from baml_client import b
2
from baml_client.types import PIIExtraction
3
from typing import Dict, Tuple
4
5
def scrub_document(text: str) -> Tuple[str, Dict[str, str]]:
6
    # Extract PII from the document
7
    result = b.ExtractPII(text)
8
    
9
    # Create a mapping of real values to scrubbed placeholders
10
    scrubbed_text = text
11
    pii_mapping = {}
12
    
13
    # Process each PII item and replace with a placeholder
14
    for pii_item in result.privateData:
15
        pii_type = pii_item.dataType.upper()
16
        placeholder = f"[{pii_type}_{pii_item.index}]"
17
        
18
        # Store the mapping for reference
19
        pii_mapping[placeholder] = pii_item.value
20
        
21
        # Replace the PII with the placeholder
22
        scrubbed_text = scrubbed_text.replace(pii_item.value, placeholder)
23
    
24
    return scrubbed_text, pii_mapping
25
26
def restore_document(scrubbed_text: str, pii_mapping: Dict[str, str]) -> str:
27
    """Restore the original text using the PII mapping."""
28
    restored_text = scrubbed_text
29
    for placeholder, original_value in pii_mapping.items():
30
        restored_text = restored_text.replace(placeholder, original_value)
31
    return restored_text
32
33
# Example usage
34
document = """
35
John Smith works at Tech Corp.
36
You can reach him at john.smith@techcorp.com
37
or call 555-0123 during business hours.
38
His employee ID is TC-12345.
39
"""
40
41
# Scrub the document
42
scrubbed_text, pii_mapping = scrub_document(document)
43
44
print("Original Document:")
45
print(document)
46
print("\nScrubbed Document:")
47
print(scrubbed_text)
48
print("\nPII Mapping:")
49
for placeholder, original in pii_mapping.items():
50
    print(f"{placeholder}: {original}")
51
52
# If needed, restore the original document
53
restored_text = restore_document(scrubbed_text, pii_mapping)
54
print("\nRestored Document:")
55
print(restored_text)