Integrate PII defense with LiteLLM gateway

Introduction

Engineers often build prompt pipelines that inadvertently expose customer names, Social Security numbers, or credit card numbers to third-party LLM providers — each unscanned request is a potential compliance incident. Adding a scanning layer at the gateway level stops PII before it leaves your network, without requiring every application team to implement their own detection logic. By the end of this lesson, you will be able to implement a LiteLLM CustomLogger hook that integrates Presidio's analyzer and anonymizer engines to enforce per-tenant PII policies on both inbound prompts and outbound model responses.

Key Terminology

• CustomLogger — LiteLLM's base class that PIIGatewayHook subclasses to inject logic into the proxy's request lifecycle, providing the async_pre_call_hook and async_log_success_event extension points used throughout this lesson.
• async_pre_call_hook — A CustomLogger method that fires synchronously before a request is forwarded to the LLM provider; used here to extract the calling tenant's policy, run Presidio analysis on data["messages"], and either raise an HTTPException (block mode) or anonymize detected spans in-place (redact mode).
• async_log_success_event — A CustomLogger method that fires after the LLM provider returns a completion; used here to scan response_obj.choices[].message.content against the tenant's response_entities (falling back to blocked_entities) and redact any PII before the response reaches the caller.
• Per-tenant policy registry — A dictionary (TENANT_POLICIES) mapping tenant IDs to enforcement configuration objects containing blocked_entities, threshold, mode, reversible, and optionally response_entities; the hook resolves the active policy by reading tenant_id from data["metadata"] on each request.
• Asymmetric PII policy — A configuration pattern in which response_entities specifies a distinct set of entity types for response-side scanning, separate from the blocked_entities used on inbound prompts, enabling stricter or entity-different enforcement on model completions than on the prompts that produced them.
• Enforcement mode — The mode field in a tenant policy: "block" causes the hook to raise HTTP 400 and abort the request the moment a detection clears the threshold, while "redact" replaces detected spans with entity-type placeholder tokens (e.g., <US_SSN>) and forwards the sanitized text to the provider.

Concepts

Why Enforce at the Gateway Rather Than in Each Application

When individual application teams each implement their own PII scanning, policies fragment silently: one team sets a loose threshold, another forgets to include SSNs, a third scans prompts but never touches responses. A single missed scan is a compliance incident. Placing the enforcement hook in the LiteLLM gateway centralizes the invariant — every request and every response flows through PIIGatewayHook regardless of which upstream service originated the call. Compliance teams then audit and update one policy registry rather than coordinating across every application. This is the architectural rationale behind subclassing CustomLogger instead of embedding Presidio calls in each service's prompt-building logic.

The LiteLLM Hook Lifecycle: Two Interception Points

LiteLLM's CustomLogger exposes two distinct interception points that together create a bidirectional filter. async_pre_call_hook fires synchronously in the outbound path: the hook can inspect and mutate the messages list before any bytes leave the network, or raise an HTTPException to abort the request entirely. async_log_success_event fires after the provider returns: it can inspect and mutate response_obj.choices[].message.content before the completion reaches the caller. Neither hook can substitute for the other — a pre-call block prevents PII from reaching the provider, but only the success event hook can catch PII the model itself generates in its reply.

Per-Tenant Policies and the Block/Redact Decision

Different tenants carry different risk profiles, which is why the policy registry maps each tenant_id to its own enforcement contract rather than applying a single global rule. The blocked_entities list scopes which Presidio entity types are even checked — types absent from the list are silently ignored even if recognizers would otherwise flag them. The threshold field then filters by confidence, letting high-sensitivity tenants reject borderline detections and low-sensitivity tenants accept them. The mode field determines what happens when a detection clears both gates: "block" aborts the request with HTTP 400, appropriate when no PII should reach the provider at all; "redact" replaces the detected span with an entity-type placeholder and forwards the sanitized text (see Code Walkthrough).

Asymmetric Request and Response Enforcement

A subtlety the lesson surfaces explicitly: the entity types you want to block in a prompt are not necessarily the same ones you need to scrub from a response. A tenant might intentionally include a customer's first name in a prompt for personalization, yet require that any name the model echoes back be removed before the completion reaches the client. The optional response_entities field captures this asymmetry — when present, the response hook (async_log_success_event) uses it instead of blocked_entities. When absent, the response hook falls back to blocked_entities, applying the same entity list in both directions. This fallback is the default behavior for tenants that do not need direction-specific enforcement (see Code Walkthrough).

Code Walkthrough

Now that you understand per-tenant policy configuration — the blocked_entities list, threshold, mode, and reversible flags that define each tenant's enforcement contract — we can translate that structure into a LiteLLM hook that enforces those policies on every request and response.

The PIIGatewayHook class subclasses LiteLLM's CustomLogger and holds shared references to a Presidio AnalyzerEngine and AnonymizerEngine, along with a policy registry keyed by tenant ID. The async_pre_call_hook method extracts the calling tenant's policy, scans the outbound prompt text, and either raises a 400 error (block mode) or anonymizes the detected spans (redact mode) before the request reaches the LLM provider:

Code snippetpython
1from presidio_analyzer import AnalyzerEngine
2from presidio_anonymizer import AnonymizerEngine
3from litellm.integrations.custom_logger import CustomLogger
4from fastapi import HTTPException
5
6TENANT_POLICIES = {
7    "tenant_a": {
8        "blocked_entities": ["PERSON", "CREDIT_CARD", "US_SSN"],
9        "threshold": 0.85,
10        "mode": "redact",
11        "reversible": False,
12    },
13    "tenant_b": {
14        "blocked_entities": ["PERSON", "EMAIL_ADDRESS", "PHONE_NUMBER"],
15        "response_entities": ["PERSON", "EMAIL_ADDRESS"],
16        "threshold": 0.75,
17        "mode": "block",
18        "reversible": False,
19    },
20}
21
22class PIIGatewayHook(CustomLogger):
23    def __init__(self):
24        self.analyzer = AnalyzerEngine()
25        self.anonymizer = AnonymizerEngine()
26        self.policies = TENANT_POLICIES
27
28    async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
29        tenant_id = (data.get("metadata") or {}).get("tenant_id", "default")
30        policy = self.policies.get(tenant_id)
31        if not policy:
32            return data
33
34        text = " ".join(
35            m.get("content", "")
36            for m in data.get("messages", [])
37            if isinstance(m.get("content"), str)
38        )
39        entities = policy["blocked_entities"]
40        threshold = policy.get("threshold", 0.8)
41        results = self.analyzer.analyze(text=text, entities=entities, language="en")
42        hits = [r for r in results if r.score >= threshold]
43
44        if not hits:
45            return data
46
47        if policy.get("mode") == "block":
48            raise HTTPException(
49                status_code=400, detail="PII Policy Violation: request blocked"
50            )
51
52        for msg in data.get("messages", []):
53            if isinstance(msg.get("content"), str):
54                anonymized = self.anonymizer.anonymize(
55                    text=msg["content"], analyzer_results=hits
56                )
57                msg["content"] = anonymized.text
58        return data

The response-side hook mirrors this pattern. It reads response_entities from the policy — falling back to blocked_entities when the field is absent, which is the asymmetric-policy behavior described in the Concepts section — then scans the model's completion and anonymizes any hits before they reach the client:

Code snippetpython
1    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
2        tenant_id = (kwargs.get("metadata") or {}).get("tenant_id", "default")
3        policy = self.policies.get(tenant_id)
4        if not policy:
5            return
6
7        entities = policy.get("response_entities", policy["blocked_entities"])
8        threshold = policy.get("threshold", 0.8)
9        for choice in getattr(response_obj, "choices", []):
10            content = getattr(getattr(choice, "message", None), "content", None)
11            if isinstance(content, str):
12                results = self.analyzer.analyze(
13                    text=content, entities=entities, language="en"
14                )
15                hits = [r for r in results if r.score >= threshold]
16                if hits:
17                    anonymized = self.anonymizer.anonymize(
18                        text=content, analyzer_results=hits
19                    )
20                    choice.message.content = anonymized.text

You'll know it works when a prompt containing a synthetic SSN is either rejected with HTTP 400 (block mode) or returned with the SSN replaced by a <US_SSN> token (redact mode), and when the same tenant's response text has any model-generated PII scrubbed before the completion reaches the client.

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

1. ✓Do instantiate AnalyzerEngine and AnonymizerEngine once in PIIGatewayHook.__init__ — both engines load spaCy NLP models on construction; sharing them as self.analyzer and self.anonymizer across all hook invocations avoids reloading those models on every gateway request, which would add hundreds of milliseconds of per-call latency at scale.
2. ✓Do filter analyzer.analyze() results to hits = [r for r in results if r.score >= threshold] before passing them to anonymizer.anonymize(analyzer_results=hits) — the per-tenant threshold field is each tenant's confidence cutoff; passing the raw results list instead anonymizes low-confidence detections that are likely false positives, over-redacting benign content against the tenant's stated policy.
3. ✓Do implement both async_pre_call_hook and async_log_success_event in the same PIIGatewayHook class — the pre-call hook blocks or redacts PII in outbound prompts before they reach the LLM provider, but without async_log_success_event, model-generated PII in completions flows back to the client entirely unscanned.

Don'ts

1. ✗Don't pass raw analyzer.analyze() results directly to anonymizer.anonymize() without threshold filtering — the unfiltered list includes every detection regardless of confidence score, causing tenant_a's 0.85 threshold to be silently ignored and redacting false-positive spans that should have passed through.
2. ✗Don't hardcode the response-scan entity list as identical to blocked_entities — the asymmetric-policy design uses policy.get("response_entities", policy["blocked_entities"]) so tenants like tenant_b can enforce a different entity set on model completions than on inbound prompts; using blocked_entities unconditionally in async_log_success_event collapses both sides of the policy into one and breaks any tenant that defines response_entities.
3. ✗Don't extract tenant_id from anywhere other than (data.get("metadata") or {}).get("tenant_id") — LiteLLM forwards caller-supplied metadata through the data["metadata"] dict; reading it from headers, query params, or a direct data["tenant_id"] key misses the actual routing path and falls through to the "default" policy, silently applying no enforcement for tenants whose IDs never match.