Detection Without Hallucination

{
  "$schema": "http://json-schema.org/draft-07/schema#",
  "type": "object",
  "required": ["is_pii", "pii_type", "confidence", "reason"],
  "additionalProperties": false,
  "properties": {
    "is_pii": {
      "type": "boolean",
      "description": "True if the span is actual PII/CUI in context, not a placeholder or field name"
    },
    "pii_type": {
      "type": "string",
      "enum": ["SSN", "EMAIL", "PHONE", "CREDIT_CARD", "PASSPORT",
               "IP_ADDRESS", "NAME", "ADDRESS", "DATE_OF_BIRTH",
               "MEDICAL_RECORD", "CUI_EXPORT", "CUI_LAW_ENFORCEMENT",
               "CUI_GOVERNMENT_CONTRACT", "OTHER", "NOT_PII"],
      "description": "Most specific matching type, or NOT_PII if rejected"
    },
    "confidence": {
      "type": "number",
      "minimum": 0.0,
      "maximum": 1.0,
      "description": "Model's confidence that this is real PII/CUI in context"
    },
    "reason": {
      "type": "string",
      "maxLength": 200,
      "description": "One sentence explaining the classification decision"
    }
  }
}

from __future__ import annotations
import json
from dataclasses import dataclass, field
from typing import Sequence

import anthropic
from presidio_analyzer import AnalyzerEngine, RecognizerResult

# ── Data model ──────────────────────────────────────────────────────────────

@dataclass
class DetectedEntity:
    text: str
    start: int
    end: int
    pii_type: str
    confidence: float
    reason: str

# ── Pass 1: deterministic sweep via Presidio ─────────────────────────────────

_analyzer = AnalyzerEngine()

def _pass1_candidates(text: str) -> list[RecognizerResult]:
    """Fast regex + NER sweep. High recall, lower precision."""
    return _analyzer.analyze(text=text, language="en")

# ── Pass 2: model-grounded confirmation ──────────────────────────────────────

_client = anthropic.Anthropic()  # reads ANTHROPIC_API_KEY from env

_SCHEMA = {
    "name": "pii_classification",
    "description": "Classify whether a text span is real PII/CUI in context",
    "input_schema": {
        "type": "object",
        "required": ["is_pii", "pii_type", "confidence", "reason"],
        "properties": {
            "is_pii": {"type": "boolean"},
            "pii_type": {"type": "string", "enum": [
                "SSN", "EMAIL", "PHONE", "CREDIT_CARD", "PASSPORT",
                "IP_ADDRESS", "NAME", "ADDRESS", "DATE_OF_BIRTH",
                "MEDICAL_RECORD", "CUI_EXPORT", "CUI_LAW_ENFORCEMENT",
                "CUI_GOVERNMENT_CONTRACT", "OTHER", "NOT_PII"
            ]},
            "confidence": {"type": "number", "minimum": 0.0, "maximum": 1.0},
            "reason": {"type": "string"}
        }
    }
}

def _pass2_confirm(
    full_text: str,
    span: str,
    start: int,
    end: int,
    candidate_type: str,
) -> dict:
    """Send candidate to LLM for context-aware confirmation."""
    # Include 200 chars of context on each side
    ctx_start = max(0, start - 200)
    ctx_end = min(len(full_text), end + 200)
    context_snippet = full_text[ctx_start:ctx_end]

    prompt = (
        f"You are a PII/CUI classifier. Given the following text context, "
        f"determine whether the highlighted span is real PII or CUI — "
        f"not a placeholder, field name, or example value.\n\n"
        f"Context:\n---\n{context_snippet}\n---\n\n"
        f"Highlighted span: '{span}'\n"
        f"Candidate type from regex: {candidate_type}\n\n"
        f"Classify the span using the pii_classification tool."
    )

    response = _client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=256,
        tools=[_SCHEMA],
        tool_choice={"type": "tool", "name": "pii_classification"},
        messages=[{"role": "user", "content": prompt}],
    )

    for block in response.content:
        if block.type == "tool_use":
            return block.input
    return {"is_pii": False, "pii_type": "NOT_PII", "confidence": 0.0, "reason": "No tool call"}

# ── Main detector ─────────────────────────────────────────────────────────────

def detect(
    text: str,
    confidence_threshold: float = 0.75,
) -> list[DetectedEntity]:
    """
    Run two-pass PII detection.
    Returns confirmed entities above the confidence threshold.
    """
    candidates = _pass1_candidates(text)
    confirmed: list[DetectedEntity] = []

    for cand in candidates:
        span_text = text[cand.start : cand.end]
        result = _pass2_confirm(
            full_text=text,
            span=span_text,
            start=cand.start,
            end=cand.end,
            candidate_type=cand.entity_type,
        )

        if result["is_pii"] and result["confidence"] >= confidence_threshold:
            confirmed.append(DetectedEntity(
                text=span_text,
                start=cand.start,
                end=cand.end,
                pii_type=result["pii_type"],
                confidence=result["confidence"],
                reason=result["reason"],
            ))

    # Deduplicate overlapping spans: keep highest-confidence
    confirmed.sort(key=lambda e: e.confidence, reverse=True)
    deduped: list[DetectedEntity] = []
    for entity in confirmed:
        overlaps = any(
            not (entity.end <= kept.start or entity.start >= kept.end)
            for kept in deduped
        )
        if not overlaps:
            deduped.append(entity)

    return deduped