The AI Firewall: Using Local Small Language Models (SLMs) to Scrub PII Before Cloud Processing

# macOS
curl -fsSL https://ollama.ai/install.sh | sh

# Linux
curl -fsSL https://ollama.ai/install.sh | sh

# Windows (PowerShell, requires WSL2)
irm https://ollama.ai/install.ps1 | iex

# Option 1: Phi-3 Mini (lightweight, fast)
ollama pull phi3

# Option 2: Llama 3.2 (more accurate, slightly larger)
ollama pull llama3.2:3b

# Option 3: Mistral (good balance)
ollama pull mistral:7b-instruct

import ollama
import json
import re
from dataclasses import dataclass
from typing import List, Dict, Tuple

@dataclass
class PIIEntity:
    """Represents a detected PII entity."""
    text: str
    category: str
    start_pos: int
    end_pos: int
    confidence: float

class LocalPIIDetector:
    """
    Uses a local SLM to detect PII in text.
    Acts as an AI Firewall before cloud processing.
    """
    
    PII_CATEGORIES = [
        "PERSON_NAME", "EMAIL", "PHONE_NUMBER", 
        "SSN", "CREDIT_CARD", "ADDRESS",
        "DATE_OF_BIRTH", "BANK_ACCOUNT", "IP_ADDRESS"
    ]
    
    def __init__(self, model_name: str = "phi3"):
        self.model = model_name
        self.detection_prompt = self._build_detection_prompt()
    
    def _build_detection_prompt(self) -> str:
        return """You are a PII (Personal Identifiable Information) detection system.
Analyze the following text and identify ALL instances of PII.

For each PII found, output in this exact JSON format:
{
    "entities": [
        {"text": "exact text", "category": "CATEGORY", "confidence": 0.95}
    ]
}

Categories to detect: PERSON_NAME, EMAIL, PHONE_NUMBER, SSN, CREDIT_CARD, 
ADDRESS, DATE_OF_BIRTH, BANK_ACCOUNT, IP_ADDRESS

If no PII is found, return: {"entities": []}

TEXT TO ANALYZE:
"""
    
    def detect(self, text: str) -> List[PIIEntity]:
        """Detect PII entities in the given text."""
        response = ollama.chat(
            model=self.model,
            messages=[
                {
                    "role": "user",
                    "content": f"{self.detection_prompt}\n{text}"
                }
            ],
            options={"temperature": 0.1}  # Low temp for consistency
        )
        
        return self._parse_response(response["message"]["content"], text)
    
    def _parse_response(self, response: str, original_text: str) -> List[PIIEntity]:
        """Parse the LLM response into PIIEntity objects."""
        entities = []
        try:
            # Extract JSON from response
            json_match = re.search(r'\{[^{}]*"entities"[^{}]*\[.*?\][^{}]*\}', 
                                   response, re.DOTALL)
            if json_match:
                data = json.loads(json_match.group())
                for item in data.get("entities", []):
                    # Find position in original text
                    start = original_text.find(item["text"])
                    if start != -1:
                        entities.append(PIIEntity(
                            text=item["text"],
                            category=item["category"],
                            start_pos=start,
                            end_pos=start + len(item["text"]),
                            confidence=item.get("confidence", 0.8)
                        ))
        except (json.JSONDecodeError, KeyError) as e:
            print(f"Warning: Could not parse LLM response: {e}")
        
        return entities

import hashlib
from typing import Optional
import uuid

class PIIScrubber:
    """
    Scrubs detected PII from text using various strategies.
    Maintains a vault for reversible operations.
    """
    
    def __init__(self, strategy: str = "tokenize"):
        self.strategy = strategy
        self.vault: Dict[str, str] = {}  # token -> original value
        self.token_counter = 0
    
    def scrub(self, text: str, entities: List[PIIEntity]) -> Tuple[str, Dict]:
        """
        Scrub PII from text and return scrubbed version with mapping.
        """
        # Sort entities by position (reverse) to replace from end
        sorted_entities = sorted(entities, key=lambda x: x.start_pos, reverse=True)
        
        scrubbed = text
        mappings = {}
        
        for entity in sorted_entities:
            replacement = self._get_replacement(entity)
            mappings[replacement] = entity.text
            self.vault[replacement] = entity.text
            
            scrubbed = (
                scrubbed[:entity.start_pos] + 
                replacement + 
                scrubbed[entity.end_pos:]
            )
        
        return scrubbed, mappings
    
    def _get_replacement(self, entity: PIIEntity) -> str:
        """Generate replacement based on strategy."""
        if self.strategy == "redact":
            return f"[{entity.category}]"
        
        elif self.strategy == "tokenize":
            self.token_counter += 1
            return f"<>"
        
        elif self.strategy == "hash":
            hash_val = hashlib.sha256(entity.text.encode()).hexdigest()[:8]
            return f"[{entity.category}:{hash_val}]"
        
        elif self.strategy == "pseudonymize":
            return self._generate_fake(entity.category)
        
        return "[REDACTED]"
    
    def _generate_fake(self, category: str) -> str:
        """Generate fake but realistic replacement data."""
        fakes = {
            "PERSON_NAME": "John Smith",
            "EMAIL": "[email protected]",
            "PHONE_NUMBER": "(555) 123-4567",
            "SSN": "XXX-XX-XXXX",
            "ADDRESS": "123 Main St, Anytown, ST 12345",
        }
        return fakes.get(category, "[REDACTED]")
    
    def restore(self, scrubbed_text: str) -> str:
        """Restore original PII from vault (for authorized use only)."""
        restored = scrubbed_text
        for token, original in self.vault.items():
            restored = restored.replace(token, original)
        return restored

import openai  # For cloud AI calls
from datetime import datetime
import logging

class AIFirewall:
    """
    Main AI Firewall class that orchestrates PII detection,
    scrubbing, cloud processing, and response handling.
    """
    
    def __init__(
        self,
        local_model: str = "phi3",
        scrub_strategy: str = "tokenize",
        cloud_provider: str = "openai"
    ):
        self.detector = LocalPIIDetector(model_name=local_model)
        self.scrubber = PIIScrubber(strategy=scrub_strategy)
        self.cloud_provider = cloud_provider
        self.audit_log = []
        
        logging.basicConfig(level=logging.INFO)
        self.logger = logging.getLogger("AIFirewall")
    
    def process(
        self, 
        text: str, 
        cloud_prompt: str,
        restore_response: bool = False
    ) -> dict:
        """
        Main processing pipeline:
        1. Detect PII locally
        2. Scrub PII from text
        3. Send safe text to cloud AI
        4. Optionally restore PII in response
        5. Return results with audit trail
        """
        timestamp = datetime.utcnow().isoformat()
        
        # Step 1: Detect PII using local SLM
        self.logger.info("Detecting PII with local model...")
        entities = self.detector.detect(text)
        self.logger.info(f"Found {len(entities)} PII entities")
        
        # Step 2: Scrub PII
        scrubbed_text, mappings = self.scrubber.scrub(text, entities)
        self.logger.info("PII scrubbed from text")
        
        # Step 3: Send to cloud AI
        self.logger.info("Sending sanitized text to cloud AI...")
        cloud_response = self._call_cloud_ai(scrubbed_text, cloud_prompt)
        
        # Step 4: Optionally restore PII in response
        if restore_response and mappings:
            final_response = self.scrubber.restore(cloud_response)
        else:
            final_response = cloud_response
        
        # Step 5: Create audit record
        audit_record = {
            "timestamp": timestamp,
            "pii_detected": len(entities),
            "pii_categories": list(set(e.category for e in entities)),
            "scrub_strategy": self.scrubber.strategy,
            "text_length_original": len(text),
            "text_length_scrubbed": len(scrubbed_text),
        }
        self.audit_log.append(audit_record)
        
        return {
            "original_text": text,
            "scrubbed_text": scrubbed_text,
            "cloud_response": cloud_response,
            "final_response": final_response,
            "pii_entities": [
                {"text": e.text, "category": e.category, "confidence": e.confidence}
                for e in entities
            ],
            "audit": audit_record
        }
    
    def _call_cloud_ai(self, safe_text: str, prompt: str) -> str:
        """Send sanitized text to cloud AI service."""
        try:
            response = openai.chat.completions.create(
                model="gpt-4",
                messages=[
                    {"role": "system", "content": prompt},
                    {"role": "user", "content": safe_text}
                ]
            )
            return response.choices[0].message.content
        except Exception as e:
            self.logger.error(f"Cloud AI error: {e}")
            return f"Error: {str(e)}"

def main():
    # Initialize the firewall
    firewall = AIFirewall(
        local_model="phi3",
        scrub_strategy="tokenize"
    )
    
    # Sample text with PII
    customer_feedback = """
    Hi, my name is Sarah Johnson and I'm writing about my order #12345.
    You can reach me at [email protected] or call me at (415) 555-0123.
    My billing address is 742 Evergreen Terrace, Springfield, IL 62704.
    I paid with my card ending in 4532 and my SSN is 123-45-6789 which
    was required for the credit check.
    """
    
    # Process through firewall
    result = firewall.process(
        text=customer_feedback,
        cloud_prompt="Summarize this customer feedback and identify the main concern.",
        restore_response=False  # Keep PII scrubbed in response
    )
    
    print("=== ORIGINAL TEXT ===")
    print(result["original_text"])
    
    print("\n=== SCRUBBED TEXT (Sent to Cloud) ===")
    print(result["scrubbed_text"])
    
    print("\n=== PII DETECTED ===")
    for entity in result["pii_entities"]:
        print(f"  {entity['category']}: {entity['text']} ({entity['confidence']:.0%})")
    
    print("\n=== CLOUD AI RESPONSE ===")
    print(result["final_response"])
    
    print("\n=== AUDIT LOG ===")
    print(result["audit"])

if __name__ == "__main__":
    main()

=== SCRUBBED TEXT (Sent to Cloud) ===
Hi, my name is <> and I'm writing about my order #12345.
You can reach me at <> or call me at <>.
My billing address is <>.
I paid with my card ending in <
       > and my SSN is < 
       
        > which was required for the credit check. === PII DETECTED === PERSON_NAME: Sarah Johnson (95%) EMAIL: [email protected] (98%) PHONE_NUMBER: (415) 555-0123 (97%) ADDRESS: 742 Evergreen Terrace, Springfield, IL 62704 (92%) CREDIT_CARD: 4532 (85%) SSN: 123-45-6789 (99%)

class HybridPIIDetector:
    """
    Combines regex patterns (fast, high-precision) with
    SLM detection (catches context-dependent PII).
    """
    
    PATTERNS = {
        "EMAIL": r'[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}',
        "PHONE_NUMBER": r'\(?\d{3}\)?[-.s]?\d{3}[-.s]?\d{4}',
        "SSN": r'\d{3}-\d{2}-\d{4}',
        "CREDIT_CARD": r'\b\d{4}[-s]?\d{4}[-s]?\d{4}[-s]?\d{4}\b',
        "IP_ADDRESS": r'\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b',
    }
    
    def __init__(self, model_name: str = "phi3"):
        self.slm_detector = LocalPIIDetector(model_name)
        self.compiled_patterns = {
            k: re.compile(v) for k, v in self.PATTERNS.items()
        }
    
    def detect(self, text: str) -> List[PIIEntity]:
        """Detect PII using both regex and SLM."""
        entities = []
        
        # Fast regex pass
        for category, pattern in self.compiled_patterns.items():
            for match in pattern.finditer(text):
                entities.append(PIIEntity(
                    text=match.group(),
                    category=category,
                    start_pos=match.start(),
                    end_pos=match.end(),
                    confidence=0.99  # Regex matches are high confidence
                ))
        
        # SLM pass for context-dependent PII (names, addresses)
        slm_entities = self.slm_detector.detect(text)
        
        # Merge, avoiding duplicates
        existing_positions = {(e.start_pos, e.end_pos) for e in entities}
        for e in slm_entities:
            if (e.start_pos, e.end_pos) not in existing_positions:
                entities.append(e)
        
        return entities

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional

app = FastAPI(title="AI Firewall API", version="1.0.0")

firewall = AIFirewall(local_model="phi3", scrub_strategy="tokenize")

class ProcessRequest(BaseModel):
    text: str
    cloud_prompt: str
    restore_response: Optional[bool] = False

@app.post("/process")
async def process_text(request: ProcessRequest):
    """Process text through the AI Firewall."""
    try:
        result = firewall.process(
            text=request.text,
            cloud_prompt=request.cloud_prompt,
            restore_response=request.restore_response
        )
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health_check():
    """Health check endpoint."""
    return {"status": "healthy", "model": firewall.detector.model}

# Run with: uvicorn app:app --host 0.0.0.0 --port 8000