pii-shield-llm 1.0.0

Safe LLM usage by automatically masking PII/PFI before it leaves your system.

pii-shield 🔒

Safe LLM usage by automatically masking PII/PFI before it leaves your system.

PII Shield intercepts sensitive data, replaces it with typed UUID tokens (<<EMAIL_3f2a1b>>), sends the clean text to your LLM, then restores original values in the response — all transparently.

---

Installation

pip install pii-shield                    # core (regex detection only)
pip install "pii-shield[nlp]"             # + spaCy NER for names & orgs
pip install "pii-shield[openai]"          # + OpenAI wrapper
pip install "pii-shield[anthropic]"       # + Anthropic wrapper
pip install "pii-shield[all]"             # everything

---

Quick Start

Basic masking & restoration

from pii_shield import PIISession

with PIISession() as session:
    masked, mapping = session.mask(
        "Hi, I'm Alice Smith. Reach me at [e-mail] or 415-555-0199."
    )
    print(masked)
    # → "Hi, I'm Alice Smith. Reach me at <<EMAIL_d3a1f9>> or <<PHONE_7b2c04>>."
    #   (Names need spaCy NER; emails & phones detected by regex)

    llm_response = your_llm_call(masked)          # send safe text to LLM
    clean = session.restore(llm_response)          # PII restored in response
    print(clean)
# Session auto-saved to encrypted vault on context manager exit

With the built-in LLM wrapper (OpenAI)

import openai
from pii_shield import PIISession
from pii_shield.llm_client import LLMClient

client = LLMClient(
    session=PIISession(),
    backend="openai",
    openai_client=openai.OpenAI(),
    model="gpt-4o",
    system_prompt="You are a helpful customer support agent.",
)

result = client.chat(
    "Please look up account for [e-mail], SSN 123-45-6789."
)
print(result.restored_text)    # PII restored
print(result.masked_prompt)    # What was actually sent to OpenAI

With Anthropic Claude

import anthropic
from pii_shield import PIISession
from pii_shield.llm_client import LLMClient

client = LLMClient(
    session=PIISession(),
    backend="anthropic",
    anthropic_client=anthropic.Anthropic(),
    model="claude-sonnet-4-20250514",
)

result = client.chat("Summarise the risk profile for card 4111111111111111.")
print(result.restored_text)

Multi-turn conversation

client = LLMClient(session=PIISession(), backend="openai", openai_client=..., model="gpt-4o")

r1 = client.chat("My name is Bob and my email is [e-mail].")
r2 = client.chat("What was the email I just gave you?")
# History is maintained; PII masked/restored across all turns

---

Entity Types Detected

Entity	Method	Example
EMAIL	Regex	[e-mail]
PHONE	Regex	[phone]
SSN	Regex	123-45-6789
CREDIT_CARD	Regex + Luhn	4111 1111 1111 1111
ADDRESS	Regex	123 Main St, NY 10001
IP_ADDRESS	Regex	192.168.1.100
PASSPORT	Regex	A12345678
DOB	Regex	01/15/1990
NAME	spaCy NER*	Alice Smith
ORG	spaCy NER*	Goldman Sachs

*Requires pip install "pii-shield[nlp]" + python -m spacy download en_core_web_sm

---

Token Format

<<ENTITY_TYPE_shortid>>

Examples:
  <<EMAIL_d3a1f9>>
  <<PHONE_7b2c04>>
  <<SSN_1a9e3c>>
  <<CREDIT_CARD_f00b12>>

• Deterministic within a session: same raw value → same token (LLM output stays coherent)
• Cross-session isolated: different sessions never share tokens
• Type-preserving: LLM sees the entity category without the actual data

---

Vault & Key Management

Session mappings are persisted in an AES-encrypted file (Fernet / AES-128-CBC + HMAC-SHA256).

On first use, a key is auto-generated and written to .env:

PII_SHIELD_KEY="base64url-encoded-32-byte-key"

Keep this key safe — losing it means mappings in the vault cannot be decrypted.

To use a custom vault path or bring your own key:

from pii_shield import PIIVault, PIISession

vault = PIIVault(
    vault_path="/secure/path/pii.vault.enc",
    env_path="/secure/path/.env",
)
session = PIISession(vault=vault)

---

Advanced Usage

Restrict entity types

from pii_shield import PIISession, EntityType

session = PIISession(entity_types=[EntityType.EMAIL, EntityType.CREDIT_CARD])

Resume a named session

# Session 1 — mask and save
s = PIISession(session_id="order-12345")
masked, _ = s.mask("Card: 4111111111111111")
s.save()

# Later — restore using the same session
s2 = PIISession(session_id="order-12345")
print(s2.restore(masked))   # → "Card: 4111111111111111"

Custom LLM backend

def my_llm(prompt: str) -> str:
    # Call your own model
    return my_model.generate(prompt)

client = LLMClient(session=PIISession(), backend="custom", custom_fn=my_llm)

---

Architecture

Input Text
    │
    ▼
[PIIDetector]   regex patterns + optional spaCy NER
    │
    ▼
[PIIReplacer]   value → <<TYPE_shortid>> (deterministic, collision-safe)
    │
    ▼
Masked Text ──→ LLM API  (OpenAI / Anthropic / custom)
    │
    ▼
LLM Response (may echo tokens back)
    │
    ▼
[PIIRestorer]   <<TYPE_shortid>> → original value
    │
    ▼
Restored Output

[PIIVault]  ←→  AES-encrypted session mappings on disk

---

Running Tests

pip install pytest
pytest tests/ -v

---

License

MIT