pii-shield-llm 🔒

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

pii-shield-llm 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-llm                    # core (regex detection only)
pip install "pii-shield-llm[nlp]"             # + spaCy NER for names & orgs
pip install "pii-shield-llm[openai]"          # + OpenAI wrapper
pip install "pii-shield-llm[anthropic]"       # + Anthropic wrapper
pip install "pii-shield-llm[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 alice@acme.com 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 john@corp.com, 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 bob@corp.com.")
r2 = client.chat("What was the email I just gave you?")
# History is maintained; PII masked/restored across all turns

Entity Types Detected

*Requires pip install "pii-shield-llm[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