diff --git a/src/guardian_rag/pii/__init__.py b/src/guardian_rag/pii/__init__.py index dbb80ef..61f533d 100644 --- a/src/guardian_rag/pii/__init__.py +++ b/src/guardian_rag/pii/__init__.py @@ -1 +1,41 @@ """PII detection and redaction pipeline.""" + +from __future__ import annotations + +from guardian_rag.pii.audit import FAIL_CLOSED_REPLACEMENT, PostGenerationAuditor +from guardian_rag.pii.deberta_validator import DebertaValidator, DebertaValidatorBackend +from guardian_rag.pii.detector import PiiDetector +from guardian_rag.pii.gliner_recognizer import ( + DEFAULT_PII_LABELS, + GlinerBackend, + GlinerRecognizer, + GlinerSpan, +) +from guardian_rag.pii.models import ( + AuditResult, + PiiEntity, + PiiEntityType, + PiiStage, + RedactionMap, +) +from guardian_rag.pii.presidio_engine import PresidioEngine +from guardian_rag.pii.redactor import Redactor + +__all__ = [ + "DEFAULT_PII_LABELS", + "FAIL_CLOSED_REPLACEMENT", + "AuditResult", + "DebertaValidator", + "DebertaValidatorBackend", + "GlinerBackend", + "GlinerRecognizer", + "GlinerSpan", + "PiiDetector", + "PiiEntity", + "PiiEntityType", + "PiiStage", + "PostGenerationAuditor", + "PresidioEngine", + "RedactionMap", + "Redactor", +] diff --git a/src/guardian_rag/pii/_real_backends.py b/src/guardian_rag/pii/_real_backends.py new file mode 100644 index 0000000..11c7c5c --- /dev/null +++ b/src/guardian_rag/pii/_real_backends.py @@ -0,0 +1,66 @@ +"""Production backends that wrap real HuggingFace models. + +Imports of ``gliner`` and ``transformers`` are performed lazily inside the +constructors so the test suite, which substitutes deterministic fakes, +never has to install or download the heavy ML stack. + +This module is excluded from pyright and coverage. Keep it small: every +line here is a line the type-checker cannot verify. +""" + +from __future__ import annotations + +from typing import Any, cast + +from guardian_rag.pii.gliner_recognizer import GlinerSpan + + +class RealGlinerBackend: + """Stage-2 backend backed by the official ``gliner`` package.""" + + def __init__(self, model_name: str = "urchade/gliner_multi_pii-v1") -> None: + from gliner import GLiNER + + self._model = GLiNER.from_pretrained(model_name) + + def predict_entities( + self, + text: str, + labels: list[str], + threshold: float, + ) -> list[GlinerSpan]: + spans = cast( + "list[GlinerSpan]", + self._model.predict_entities(text, labels, threshold=threshold), + ) + return spans + + +class RealDebertaValidatorBackend: + """Stage-3 backend backed by a HuggingFace zero-shot NLI pipeline.""" + + HYPOTHESIS_TEMPLATE = "This text is {}." + + def __init__( + self, + model_name: str = "MoritzLaurer/deberta-v3-xsmall-zeroshot-v1.1-all-33", + ) -> None: + from transformers import pipeline + + self._pipeline = pipeline("zero-shot-classification", model=model_name) + + def score(self, text: str, hypothesis: str) -> float: + result = cast( + "dict[str, Any]", + self._pipeline( + text, + candidate_labels=[hypothesis, "something else"], + hypothesis_template=self.HYPOTHESIS_TEMPLATE, + ), + ) + labels = cast("list[str]", result["labels"]) + scores = cast("list[float]", result["scores"]) + for label, score in zip(labels, scores, strict=True): + if label == hypothesis: + return float(score) + return 0.0 diff --git a/src/guardian_rag/pii/audit.py b/src/guardian_rag/pii/audit.py new file mode 100644 index 0000000..016c581 --- /dev/null +++ b/src/guardian_rag/pii/audit.py @@ -0,0 +1,104 @@ +"""Post-generation PII audit. + +Most RAG privacy projects stop at ingestion. The LLM can still paraphrase +or infer personal data from clean chunks. This module re-runs the same +three-stage detector over every generated answer before it reaches the +user. + +Failure semantics matter here: if the audit raises, we MUST NOT return +the raw answer. The default is fail-closed: an audit exception turns into +a redacted answer with `fail_closed=True`. Setting `fail_closed=False` +re-raises instead, which only makes sense in offline tooling like the +eval harness. +""" + +from __future__ import annotations + +import structlog + +from guardian_rag.pii.detector import PiiDetector +from guardian_rag.pii.models import AuditResult +from guardian_rag.pii.redactor import Redactor + +logger = structlog.get_logger(__name__) + + +# Returned in place of the LLM answer when the audit itself crashed. Kept +# deliberately short so the API layer can surface it to the user without +# leaking implementation details. +FAIL_CLOSED_REPLACEMENT = "[Audit failed: response redacted for safety]" + + +class PostGenerationAuditor: + """Run the PII detector over an LLM answer and redact any leak. + + Parameters + ---------- + detector: + The same three-stage detector ingestion uses. Sharing one instance + means the audit catches exactly what would have been redacted at + ingest time. + redactor: + Typed-placeholder redactor. Same instance as ingestion for + consistent placeholder labels. + fail_closed: + When True (the default), any exception inside `audit()` is caught, + the answer is replaced with `FAIL_CLOSED_REPLACEMENT`, and the + result is returned with `fail_closed=True`. When False the + exception is re-raised, which is useful for offline tooling. + """ + + def __init__( + self, + detector: PiiDetector, + redactor: Redactor, + fail_closed: bool = True, + ) -> None: + self._detector = detector + self._redactor = redactor + self._fail_closed = fail_closed + + def audit(self, answer: str, document_id: str = "answer") -> AuditResult: + """Return an `AuditResult` for one LLM answer. + + On `clean=True` the returned `answer` is identical to the input. + On `clean=False` it is the redacted version. On `fail_closed=True` + the input is replaced with a safe placeholder. + """ + try: + entities = self._detector.detect(answer) + except Exception: + if not self._fail_closed: + raise + return _fail_closed(answer, "detect") + + if not entities: + return AuditResult(clean=True, answer=answer) + + try: + redacted, _map = self._redactor.redact(answer, entities, document_id) + except Exception: + if not self._fail_closed: + raise + return _fail_closed(answer, "redact") + + logger.warning( + "audit.leak_detected", + n_entities=len(entities), + entity_types=sorted({e.entity_type.value for e in entities}), + ) + return AuditResult(clean=False, answer=redacted, entities=entities) + + +def _fail_closed(answer: str, phase: str) -> AuditResult: + logger.exception( + "audit.failed_closed", + phase=phase, + answer_preview=answer[:120], + ) + return AuditResult( + clean=False, + answer=FAIL_CLOSED_REPLACEMENT, + entities=[], + fail_closed=True, + ) diff --git a/src/guardian_rag/pii/deberta_validator.py b/src/guardian_rag/pii/deberta_validator.py new file mode 100644 index 0000000..5ab9864 --- /dev/null +++ b/src/guardian_rag/pii/deberta_validator.py @@ -0,0 +1,119 @@ +"""Stage 3 detector: conditional DeBERTa zero-shot validation. + +When GLiNER produces a span whose confidence is below the configured +threshold, we ask a DeBERTa zero-shot NLI model "is this span really a +{label}?". This is the safety net for borderline detections without paying +DeBERTa inference cost on every span. + +The backend is exposed as a `Protocol` so tests can drop in a deterministic +fake. The real backend lives in `_real_backends.py`. +""" + +from __future__ import annotations + +from typing import Protocol, runtime_checkable + +import structlog + +from guardian_rag.pii.models import PiiEntity, PiiStage + +logger = structlog.get_logger(__name__) + + +# Each domain PII type gets a short, human-readable label that DeBERTa NLI +# can score against. The wording matters: "personal name" works better than +# "PERSON" because the model was trained on natural language hypotheses. +ENTITY_TO_HYPOTHESIS: dict[str, str] = { + "PERSON": "a personal name", + "EMAIL": "an email address", + "PHONE": "a phone number", + "ADDRESS": "a postal address", + "DATE_OF_BIRTH": "a date of birth", + "ID_NUMBER": "an identification number", + "BANK_ACCOUNT": "a bank account number", + "CREDIT_CARD": "a credit card number", + "SOCIAL_SECURITY": "a social security number", + "IP_ADDRESS": "an IP address", + "URL": "a URL", + "ORGANIZATION": "an organisation name", + "LOCATION": "a geographic location", +} + + +@runtime_checkable +class DebertaValidatorBackend(Protocol): + """Pluggable DeBERTa backend so tests can avoid loading weights.""" + + def score(self, text: str, hypothesis: str) -> float: + """Probability in [0, 1] that ``hypothesis`` entails ``text``.""" + ... + + +class DebertaValidator: + """Validate low-confidence GLiNER spans against a zero-shot NLI model. + + Parameters + ---------- + backend: + Anything implementing `DebertaValidatorBackend`. + score_threshold: + GLiNER spans with score below this are sent to DeBERTa for + validation. Spans with score at or above this are accepted as-is. + accept_threshold: + DeBERTa probability above which a validated span is kept. Below + this the span is dropped (treated as a GLiNER false positive). + """ + + def __init__( + self, + backend: DebertaValidatorBackend, + score_threshold: float = 0.85, + accept_threshold: float = 0.5, + ) -> None: + if not 0.0 <= score_threshold <= 1.0: + raise ValueError("score_threshold must be in [0, 1]") + if not 0.0 <= accept_threshold <= 1.0: + raise ValueError("accept_threshold must be in [0, 1]") + self._backend = backend + self._score_threshold = score_threshold + self._accept_threshold = accept_threshold + + def validate(self, entities: list[PiiEntity]) -> list[PiiEntity]: + """Return the subset of ``entities`` that survive validation. + + High-confidence spans (>= score_threshold) pass through unchanged. + Low-confidence spans are scored by DeBERTa. Spans whose DeBERTa + probability is below `accept_threshold` are dropped. Surviving + validated spans have their stage marked GLINER_VALIDATED and their + score replaced with the DeBERTa probability. + """ + survivors: list[PiiEntity] = [] + for entity in entities: + if entity.score >= self._score_threshold: + survivors.append(entity) + continue + hypothesis = ENTITY_TO_HYPOTHESIS.get(entity.entity_type.value) + if hypothesis is None: + # No hypothesis means we cannot validate; keep the span to + # err on the privacy-safe side. + survivors.append(entity) + continue + probability = float(self._backend.score(entity.text, hypothesis)) + if probability < self._accept_threshold: + logger.debug( + "deberta.validator.dropped", + text=entity.text, + entity_type=entity.entity_type.value, + gliner_score=entity.score, + deberta_score=probability, + ) + continue + survivors.append( + entity.model_copy( + update={ + "score": probability, + "stage": PiiStage.GLINER_VALIDATED, + } + ) + ) + return survivors diff --git a/src/guardian_rag/pii/detector.py b/src/guardian_rag/pii/detector.py new file mode 100644 index 0000000..1b3aa3a --- /dev/null +++ b/src/guardian_rag/pii/detector.py @@ -0,0 +1,68 @@ +"""Top-level PII detector orchestrating Presidio + GLiNER + DeBERTa. + +This is the single entry point ingestion, the query path, and the audit all +share. Each layer (Presidio, GLiNER, optional DeBERTa) is injected so the +test suite can swap any of them for deterministic fakes. + +The detector returns a flat `list[PiiEntity]` in source-document order. Use +the redactor to produce text + RedactionMap. +""" + +from __future__ import annotations + +import structlog + +from guardian_rag.pii.deberta_validator import DebertaValidator +from guardian_rag.pii.gliner_recognizer import GlinerRecognizer +from guardian_rag.pii.models import PiiEntity +from guardian_rag.pii.presidio_engine import PresidioEngine + +logger = structlog.get_logger(__name__) + + +class PiiDetector: + """Run Presidio, GLiNER, and the optional DeBERTa validator in sequence. + + Parameters + ---------- + presidio: + Stage 1 engine for structured PII (regex). Required. + gliner: + Stage 2 recogniser for unstructured PII (NER). Optional. If absent, + only Presidio entities are returned. + validator: + Stage 3 conditional DeBERTa validator. Optional. Only fires on + low-confidence GLiNER spans. + """ + + def __init__( + self, + presidio: PresidioEngine, + gliner: GlinerRecognizer | None = None, + validator: DebertaValidator | None = None, + ) -> None: + self._presidio = presidio + self._gliner = gliner + self._validator = validator + + def detect(self, text: str) -> list[PiiEntity]: + """Return all PII entities found in ``text``, in source order.""" + if not text: + return [] + + presidio_entities = self._presidio.detect(text) + gliner_entities: list[PiiEntity] = [] + if self._gliner is not None: + gliner_entities = self._gliner.detect(text) + if self._validator is not None: + gliner_entities = self._validator.validate(gliner_entities) + + merged = [*presidio_entities, *gliner_entities] + merged.sort(key=lambda e: (e.start, e.end)) + logger.debug( + "pii.detect.complete", + n_presidio=len(presidio_entities), + n_gliner=len(gliner_entities), + n_total=len(merged), + ) + return merged diff --git a/src/guardian_rag/pii/gliner_recognizer.py b/src/guardian_rag/pii/gliner_recognizer.py new file mode 100644 index 0000000..4ca1c56 --- /dev/null +++ b/src/guardian_rag/pii/gliner_recognizer.py @@ -0,0 +1,122 @@ +"""Stage 2 detector: GLiNER NER for unstructured PII. + +GLiNER is a span-based NER model that beats regex on contextual entities +like names, organisations, and locations. We expose the backend as a +`Protocol` so tests can plug in a deterministic fake without loading the +real model. The real backend lives in `_real_backends.py`. + +The recogniser converts GLiNER spans into `PiiEntity` records using the +domain enum, dropping any spans whose label does not map onto a Guardian +PII type. +""" + +from __future__ import annotations + +from typing import Protocol, TypedDict, runtime_checkable + +from guardian_rag.pii.models import PiiEntity, PiiEntityType, PiiStage + + +class GlinerSpan(TypedDict): + """Wire shape of a single GLiNER prediction. + + Matches the dict the official ``gliner`` package returns from + `predict_entities`. We use a `TypedDict` here so backends do not have + to depend on pydantic. + """ + + start: int + end: int + text: str + label: str + score: float + + +# GLiNER label -> domain PII type. Labels are lower case by convention. +# Anything not in this map is ignored, which lets us reuse the same backend +# for non-PII GLiNER use cases without leaking unrelated entities. +GLINER_LABEL_TO_DOMAIN: dict[str, PiiEntityType] = { + "person": PiiEntityType.PERSON, + "email": PiiEntityType.EMAIL, + "phone number": PiiEntityType.PHONE, + "address": PiiEntityType.ADDRESS, + "date of birth": PiiEntityType.DATE_OF_BIRTH, + "id number": PiiEntityType.ID_NUMBER, + "bank account": PiiEntityType.BANK_ACCOUNT, + "credit card": PiiEntityType.CREDIT_CARD, + "social security": PiiEntityType.SOCIAL_SECURITY, + "ip address": PiiEntityType.IP_ADDRESS, + "url": PiiEntityType.URL, + "organization": PiiEntityType.ORGANIZATION, + "location": PiiEntityType.LOCATION, +} + +DEFAULT_PII_LABELS: tuple[str, ...] = tuple(GLINER_LABEL_TO_DOMAIN.keys()) + + +@runtime_checkable +class GlinerBackend(Protocol): + """Pluggable GLiNER backend so tests can avoid downloading weights.""" + + def predict_entities( + self, + text: str, + labels: list[str], + threshold: float, + ) -> list[GlinerSpan]: + """Run GLiNER over ``text`` for the given candidate ``labels``.""" + ... + + +class GlinerRecognizer: + """Convert GLiNER predictions into `PiiEntity` records. + + Parameters + ---------- + backend: + Anything implementing the `GlinerBackend` protocol. + labels: + Candidate PII labels. Defaults to the full Guardian set. + threshold: + Minimum score for GLiNER to surface a span. Lower this to be more + permissive (more entities, more false positives) and raise it to + be stricter. + """ + + def __init__( + self, + backend: GlinerBackend, + labels: tuple[str, ...] = DEFAULT_PII_LABELS, + threshold: float = 0.5, + ) -> None: + if not 0.0 <= threshold <= 1.0: + raise ValueError("threshold must be in [0, 1]") + self._backend = backend + self._labels = list(labels) + self._threshold = threshold + + def detect(self, text: str) -> list[PiiEntity]: + if not text: + return [] + spans = self._backend.predict_entities(text, self._labels, self._threshold) + entities: list[PiiEntity] = [] + for span in spans: + label = span["label"].lower() + if label not in GLINER_LABEL_TO_DOMAIN: + continue + start = int(span["start"]) + end = int(span["end"]) + if end <= start: + continue + entities.append( + PiiEntity( + entity_type=GLINER_LABEL_TO_DOMAIN[label], + start=start, + end=end, + score=float(span["score"]), + text=text[start:end], + stage=PiiStage.GLINER, + ) + ) + entities.sort(key=lambda e: (e.start, e.end)) + return entities diff --git a/src/guardian_rag/pii/models.py b/src/guardian_rag/pii/models.py new file mode 100644 index 0000000..03a6655 --- /dev/null +++ b/src/guardian_rag/pii/models.py @@ -0,0 +1,117 @@ +"""Domain models for the PII detection pipeline. + +Every cross-module boundary uses frozen pydantic models. Detection stages +attach their identity to each entity so downstream audit reports can show +which stage flagged what. +""" + +from __future__ import annotations + +from enum import StrEnum + +from pydantic import BaseModel, ConfigDict, Field, model_validator + + +class PiiStage(StrEnum): + """Which detection stage surfaced a candidate entity.""" + + PRESIDIO = "presidio" + GLINER = "gliner" + GLINER_VALIDATED = "gliner_validated" + + +class PiiEntityType(StrEnum): + """Canonical PII entity types Guardian RAG recognises. + + The string values double as the placeholder label, so a `PERSON` entity + becomes `[PERSON_1]` after redaction. Adding a new value automatically + extends placeholder generation. + """ + + PERSON = "PERSON" + EMAIL = "EMAIL" + PHONE = "PHONE" + ADDRESS = "ADDRESS" + DATE_OF_BIRTH = "DATE_OF_BIRTH" + ID_NUMBER = "ID_NUMBER" + BANK_ACCOUNT = "BANK_ACCOUNT" + CREDIT_CARD = "CREDIT_CARD" + SOCIAL_SECURITY = "SOCIAL_SECURITY" + IP_ADDRESS = "IP_ADDRESS" + URL = "URL" + ORGANIZATION = "ORGANIZATION" + LOCATION = "LOCATION" + + +class PiiEntity(BaseModel): + """A single PII span detected in source text. + + `start` and `end` are character offsets into the original (unredacted) text. + `score` is the detector's confidence in [0, 1]. `stage` records which + detection stage surfaced the entity, which lets the audit report attribute + detections back to Presidio, GLiNER, or the validator. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + entity_type: PiiEntityType + start: int = Field(ge=0) + end: int = Field(ge=0) + score: float = Field(ge=0.0, le=1.0) + text: str + stage: PiiStage + + @model_validator(mode="after") + def _check_offsets(self) -> PiiEntity: + if self.end <= self.start: + raise ValueError(f"end ({self.end}) must be greater than start ({self.start})") + return self + + @property + def length(self) -> int: + return self.end - self.start + + def overlaps(self, other: PiiEntity) -> bool: + """True if this entity shares any character with ``other``.""" + return self.start < other.end and other.start < self.end + + +class RedactionMap(BaseModel): + """Mapping between placeholder tokens and the original PII spans. + + Stored separately from the redacted text. Never embedded into the vector + store. Used to render before/after views in the dashboard and to support + audit replay. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + document_id: str + placeholders: dict[str, str] = Field(default_factory=lambda: {}) + """Mapping from placeholder token (e.g. ``[PERSON_1]``) to original text.""" + + entities: list[PiiEntity] = Field(default_factory=lambda: []) + """All entities that were redacted, in source-document order.""" + + def reverse(self) -> dict[str, str]: + """Original text -> placeholder. Used for consistency checks.""" + return {v: k for k, v in self.placeholders.items()} + + +class AuditResult(BaseModel): + """Outcome of the post-generation PII audit on a single answer. + + `clean=True` means the audit ran successfully and found nothing. A + `clean=False` answer is always returned redacted, never raw. `entities` + lists what was found so the dashboard can show the leak. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + clean: bool + answer: str + """The user-facing answer. Identical to the input if clean, otherwise redacted.""" + + entities: list[PiiEntity] = Field(default_factory=lambda: []) + fail_closed: bool = False + """True if the audit raised and the answer was redacted as a safety default.""" diff --git a/src/guardian_rag/pii/presidio_engine.py b/src/guardian_rag/pii/presidio_engine.py new file mode 100644 index 0000000..8740cbb --- /dev/null +++ b/src/guardian_rag/pii/presidio_engine.py @@ -0,0 +1,118 @@ +"""Stage 1 detector: Presidio regex recognisers for structured PII. + +Presidio ships a battery of `EntityRecognizer` subclasses (email, IBAN, +credit card, IP, SSN, phone number). They do not need an NLP engine, so we +call them directly without loading spaCy or Stanza. This keeps Stage 1 in +the ~5ms latency budget the architecture decision committed to. + +The `PresidioEngine` exposes a single `detect()` method that returns a list +of `PiiEntity` records in source-document order, ready to merge with the +GLiNER output. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from guardian_rag.pii.models import PiiEntity, PiiEntityType, PiiStage + +if TYPE_CHECKING: # pragma: no cover + from presidio_analyzer import EntityRecognizer + + +# Map Presidio's canonical entity strings onto our domain enum. Anything we +# do not have a slot for is dropped silently. Presidio supports more entity +# types than this list, but we restrict ourselves to entities we can credibly +# redact and re-use as RAG context. +_PRESIDIO_TO_DOMAIN: dict[str, PiiEntityType] = { + "EMAIL_ADDRESS": PiiEntityType.EMAIL, + "PHONE_NUMBER": PiiEntityType.PHONE, + "IBAN_CODE": PiiEntityType.BANK_ACCOUNT, + "CREDIT_CARD": PiiEntityType.CREDIT_CARD, + "IP_ADDRESS": PiiEntityType.IP_ADDRESS, + "US_SSN": PiiEntityType.SOCIAL_SECURITY, + "URL": PiiEntityType.URL, +} + + +def _build_default_recognizers() -> list[EntityRecognizer]: + """Instantiate the Presidio recognisers Guardian RAG enables by default. + + Import is lazy so that `from guardian_rag.pii import ...` stays cheap; + the module only pays for Presidio when a `PresidioEngine` is built. + """ + from presidio_analyzer.predefined_recognizers import ( + CreditCardRecognizer, + EmailRecognizer, + IbanRecognizer, + IpRecognizer, + PhoneRecognizer, + UrlRecognizer, + UsSsnRecognizer, + ) + + return [ + EmailRecognizer(), + PhoneRecognizer(), + IbanRecognizer(), + CreditCardRecognizer(), + IpRecognizer(), + UsSsnRecognizer(), + UrlRecognizer(), + ] + + +class PresidioEngine: + """Thin wrapper around a set of Presidio `EntityRecognizer`s. + + Parameters + ---------- + recognizers: + Iterable of Presidio recognisers. If omitted, the default set + registers EmailRecognizer, PhoneRecognizer, IbanRecognizer, + CreditCardRecognizer, IpRecognizer, UsSsnRecognizer, UrlRecognizer. + language: + Presidio language code. Defaults to English. Other languages need + their own recogniser configuration. + """ + + def __init__( + self, + recognizers: list[EntityRecognizer] | None = None, + language: str = "en", + ) -> None: + self._recognizers = recognizers if recognizers is not None else _build_default_recognizers() + self._language = language + + def detect(self, text: str) -> list[PiiEntity]: + """Run every recogniser against ``text`` and return merged entities.""" + if not text: + return [] + + results: list[PiiEntity] = [] + for recognizer in self._recognizers: + entities = recognizer.supported_entities + for entity_name in entities: + if entity_name not in _PRESIDIO_TO_DOMAIN: + continue + # `nlp_artifacts` is optional for EntityRecognizers; passing + # None skips the spaCy code path entirely. + matches = recognizer.analyze( + text=text, + entities=[entity_name], + nlp_artifacts=None, # type: ignore[arg-type] + ) + for match in matches or []: + results.append( + PiiEntity( + entity_type=_PRESIDIO_TO_DOMAIN[entity_name], + start=int(match.start), + end=int(match.end), + score=float(match.score), + text=text[int(match.start) : int(match.end)], + stage=PiiStage.PRESIDIO, + ) + ) + + results.sort(key=lambda e: (e.start, e.end)) + return results diff --git a/src/guardian_rag/pii/redactor.py b/src/guardian_rag/pii/redactor.py new file mode 100644 index 0000000..3fd020f --- /dev/null +++ b/src/guardian_rag/pii/redactor.py @@ -0,0 +1,97 @@ +"""Typed placeholder redaction. + +Blank redaction destroys document structure and breaks retrieval. We replace +each PII span with a typed placeholder of the form `[PERSON_1]`, +`[EMAIL_2]`, etc. The same source span resolves to the same placeholder +within a document, so a chunk that mentions the same person twice remains +internally consistent. + +`Redactor.redact()` returns the redacted text plus the `RedactionMap` that +records every placeholder. The map is stored separately and never embedded. +""" + +from __future__ import annotations + +from guardian_rag.pii.models import PiiEntity, PiiEntityType, RedactionMap + + +def _resolve_overlaps(entities: list[PiiEntity]) -> list[PiiEntity]: + """Drop overlapping entities, keeping the higher-scoring one. + + Presidio and GLiNER can flag the same span twice when they agree. They + can also disagree on boundaries. We resolve by sorting on (start, -score) + and dropping anything that overlaps an entity already kept. + """ + if not entities: + return [] + ordered = sorted(entities, key=lambda e: (e.start, -e.score, e.end - e.start)) + kept: list[PiiEntity] = [] + for entity in ordered: + if any(entity.overlaps(k) for k in kept): + continue + kept.append(entity) + kept.sort(key=lambda e: e.start) + return kept + + +class Redactor: + """Replace PII spans in source text with typed placeholders. + + Parameters + ---------- + placeholder_format: + Format string with two placeholders: entity type and index. Defaults + to ``"[{type}_{index}]"`` which produces e.g. ``[PERSON_1]``. + """ + + def __init__(self, placeholder_format: str = "[{type}_{index}]") -> None: + self._format = placeholder_format + + def redact( + self, + text: str, + entities: list[PiiEntity], + document_id: str, + ) -> tuple[str, RedactionMap]: + """Redact ``entities`` from ``text`` and return the redacted text and map.""" + if not entities: + return text, RedactionMap(document_id=document_id) + + deduped = _resolve_overlaps(entities) + + # Allocate stable placeholders per entity type per unique span text. + # `[PERSON_1]` should refer to the same person every time it appears + # in the document so the redacted chunk is still readable. + index_by_type: dict[PiiEntityType, int] = {} + token_by_span: dict[tuple[PiiEntityType, str], str] = {} + placeholders: dict[str, str] = {} + + redacted_entities: list[PiiEntity] = [] + out: list[str] = [] + cursor = 0 + for entity in deduped: + if entity.start < cursor: + # Should not happen after `_resolve_overlaps`, but defend + # the loop so a bad input cannot corrupt the output. + continue + key = (entity.entity_type, entity.text) + token = token_by_span.get(key) + if token is None: + index_by_type[entity.entity_type] = index_by_type.get(entity.entity_type, 0) + 1 + token = self._format.format( + type=entity.entity_type.value, + index=index_by_type[entity.entity_type], + ) + token_by_span[key] = token + placeholders[token] = entity.text + out.append(text[cursor : entity.start]) + out.append(token) + cursor = entity.end + redacted_entities.append(entity) + out.append(text[cursor:]) + + return "".join(out), RedactionMap( + document_id=document_id, + placeholders=placeholders, + entities=redacted_entities, + ) diff --git a/tests/test_pii_audit.py b/tests/test_pii_audit.py new file mode 100644 index 0000000..569423d --- /dev/null +++ b/tests/test_pii_audit.py @@ -0,0 +1,122 @@ +"""Post-generation audit tests.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest + +from guardian_rag.pii.audit import FAIL_CLOSED_REPLACEMENT, PostGenerationAuditor +from guardian_rag.pii.detector import PiiDetector +from guardian_rag.pii.gliner_recognizer import GlinerRecognizer, GlinerSpan +from guardian_rag.pii.presidio_engine import PresidioEngine +from guardian_rag.pii.redactor import Redactor + + +@dataclass +class FakeGliner: + spans: list[GlinerSpan] + + def predict_entities( + self, + text: str, + labels: list[str], + threshold: float, + ) -> list[GlinerSpan]: + del text, labels, threshold + return list(self.spans) + + +def _auditor( + gliner_spans: list[GlinerSpan] | None = None, + fail_closed: bool = True, +) -> PostGenerationAuditor: + gliner = GlinerRecognizer(FakeGliner(spans=gliner_spans or [])) + detector = PiiDetector(presidio=PresidioEngine(), gliner=gliner) + return PostGenerationAuditor(detector=detector, redactor=Redactor(), fail_closed=fail_closed) + + +def test_clean_answer_returns_unchanged(): + auditor = _auditor() + result = auditor.audit("This is a perfectly safe answer.") + assert result.clean is True + assert result.answer == "This is a perfectly safe answer." + assert result.entities == [] + assert result.fail_closed is False + + +def test_email_leak_caught_and_redacted(): + auditor = _auditor() + answer = "Please contact alice@example.com for next steps." + result = auditor.audit(answer) + assert result.clean is False + assert "alice@example.com" not in result.answer + assert "[EMAIL_1]" in result.answer + assert any(e.text == "alice@example.com" for e in result.entities) + + +def test_gliner_only_leak_caught(): + auditor = _auditor( + gliner_spans=[ + {"start": 0, "end": 10, "text": "John Smith", "label": "person", "score": 0.95}, + ] + ) + result = auditor.audit("John Smith approved the request.") + assert result.clean is False + assert "John Smith" not in result.answer + assert "[PERSON_1]" in result.answer + + +def test_detector_exception_fails_closed_by_default(): + class BrokenDetector: + def detect(self, text: str) -> list[object]: + del text + raise RuntimeError("backend down") + + auditor = PostGenerationAuditor( + detector=BrokenDetector(), # type: ignore[arg-type] + redactor=Redactor(), + ) + result = auditor.audit("anything") + assert result.clean is False + assert result.fail_closed is True + assert result.answer == FAIL_CLOSED_REPLACEMENT + assert result.entities == [] + + +def test_detector_exception_re_raised_when_fail_closed_off(): + class BrokenDetector: + def detect(self, text: str) -> list[object]: + del text + raise RuntimeError("backend down") + + auditor = PostGenerationAuditor( + detector=BrokenDetector(), # type: ignore[arg-type] + redactor=Redactor(), + fail_closed=False, + ) + with pytest.raises(RuntimeError, match="backend down"): + auditor.audit("anything") + + +def test_redactor_exception_also_fails_closed(): + class BrokenRedactor: + def redact(self, *args: object, **kwargs: object) -> tuple[str, object]: + del args, kwargs + raise RuntimeError("redactor broken") + + auditor = PostGenerationAuditor( + detector=PiiDetector(presidio=PresidioEngine()), + redactor=BrokenRedactor(), # type: ignore[arg-type] + ) + result = auditor.audit("Contact alice@example.com please.") + assert result.clean is False + assert result.fail_closed is True + assert result.answer == FAIL_CLOSED_REPLACEMENT + + +def test_empty_answer_is_clean(): + auditor = _auditor() + result = auditor.audit("") + assert result.clean is True + assert result.answer == "" diff --git a/tests/test_pii_deberta.py b/tests/test_pii_deberta.py new file mode 100644 index 0000000..f24834f --- /dev/null +++ b/tests/test_pii_deberta.py @@ -0,0 +1,72 @@ +"""DeBERTa validator tests, using a deterministic fake backend.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from guardian_rag.pii.deberta_validator import DebertaValidator +from guardian_rag.pii.models import PiiEntity, PiiEntityType, PiiStage + + +@dataclass +class FakeBackend: + """Score backend that returns a canned score for every call.""" + + score_value: float = 0.9 + calls: list[tuple[str, str]] = field(default_factory=list) + + def score(self, text: str, hypothesis: str) -> float: + self.calls.append((text, hypothesis)) + return self.score_value + + +def _entity(score: float, etype: PiiEntityType = PiiEntityType.PERSON) -> PiiEntity: + return PiiEntity( + entity_type=etype, + start=0, + end=10, + score=score, + text="John Smith", + stage=PiiStage.GLINER, + ) + + +def test_high_confidence_entities_pass_through_without_validation(): + backend = FakeBackend(score_value=0.0) + validator = DebertaValidator(backend, score_threshold=0.85) + survivors = validator.validate([_entity(score=0.95)]) + assert len(survivors) == 1 + assert survivors[0].stage is PiiStage.GLINER + assert backend.calls == [] + + +def test_low_confidence_entity_accepted_when_deberta_agrees(): + backend = FakeBackend(score_value=0.8) + validator = DebertaValidator(backend, score_threshold=0.85, accept_threshold=0.5) + survivors = validator.validate([_entity(score=0.6)]) + assert len(survivors) == 1 + assert survivors[0].stage is PiiStage.GLINER_VALIDATED + assert survivors[0].score == 0.8 + + +def test_low_confidence_entity_dropped_when_deberta_disagrees(): + backend = FakeBackend(score_value=0.2) + validator = DebertaValidator(backend, score_threshold=0.85, accept_threshold=0.5) + assert validator.validate([_entity(score=0.6)]) == [] + + +def test_threshold_bounds(): + import pytest + + backend = FakeBackend() + with pytest.raises(ValueError, match="score_threshold"): + DebertaValidator(backend, score_threshold=1.5) + with pytest.raises(ValueError, match="accept_threshold"): + DebertaValidator(backend, accept_threshold=-0.1) + + +def test_every_entity_type_has_a_hypothesis(): + from guardian_rag.pii.deberta_validator import ENTITY_TO_HYPOTHESIS + + for etype in PiiEntityType: + assert etype.value in ENTITY_TO_HYPOTHESIS diff --git a/tests/test_pii_detector.py b/tests/test_pii_detector.py new file mode 100644 index 0000000..80ce891 --- /dev/null +++ b/tests/test_pii_detector.py @@ -0,0 +1,84 @@ +"""Orchestrator tests for the three-stage PII detector.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from guardian_rag.pii.deberta_validator import DebertaValidator +from guardian_rag.pii.detector import PiiDetector +from guardian_rag.pii.gliner_recognizer import GlinerRecognizer, GlinerSpan +from guardian_rag.pii.models import PiiEntityType, PiiStage +from guardian_rag.pii.presidio_engine import PresidioEngine + + +@dataclass +class FakeGliner: + spans: list[GlinerSpan] = field(default_factory=list) + + def predict_entities( + self, + text: str, + labels: list[str], + threshold: float, + ) -> list[GlinerSpan]: + del text, labels, threshold + return list(self.spans) + + +@dataclass +class FakeDeberta: + score_value: float = 0.9 + + def score(self, text: str, hypothesis: str) -> float: + del text, hypothesis + return self.score_value + + +def test_detector_returns_presidio_and_gliner_entities_sorted(): + text = "John Smith emailed alice@example.com about the project." + presidio = PresidioEngine() + gliner = GlinerRecognizer( + FakeGliner( + spans=[ + {"start": 0, "end": 10, "text": "John Smith", "label": "person", "score": 0.95}, + ] + ) + ) + detector = PiiDetector(presidio=presidio, gliner=gliner) + entities = detector.detect(text) + + types = {e.entity_type for e in entities} + assert PiiEntityType.PERSON in types + assert PiiEntityType.EMAIL in types + starts = [e.start for e in entities] + assert starts == sorted(starts) + + +def test_detector_without_gliner_returns_only_presidio(): + text = "Email: jane@example.com" + detector = PiiDetector(presidio=PresidioEngine()) + entities = detector.detect(text) + assert all(e.stage is PiiStage.PRESIDIO for e in entities) + + +def test_detector_with_validator_filters_low_confidence_gliner_spans(): + text = "John works at Acme." + gliner = GlinerRecognizer( + FakeGliner( + spans=[ + # Low GLiNER confidence: will be sent to DeBERTa. + {"start": 0, "end": 4, "text": "John", "label": "person", "score": 0.40}, + ] + ) + ) + # DeBERTa returns 0.1 -> below accept_threshold of 0.5, so the entity + # is dropped. + validator = DebertaValidator(FakeDeberta(score_value=0.1), score_threshold=0.85) + detector = PiiDetector(presidio=PresidioEngine(), gliner=gliner, validator=validator) + entities = detector.detect(text) + assert entities == [] + + +def test_detector_empty_text_short_circuits(): + detector = PiiDetector(presidio=PresidioEngine()) + assert detector.detect("") == [] diff --git a/tests/test_pii_gliner.py b/tests/test_pii_gliner.py new file mode 100644 index 0000000..044fcd3 --- /dev/null +++ b/tests/test_pii_gliner.py @@ -0,0 +1,93 @@ +"""GLiNER recogniser tests, using a deterministic fake backend.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from guardian_rag.pii.gliner_recognizer import ( + DEFAULT_PII_LABELS, + GlinerRecognizer, + GlinerSpan, +) +from guardian_rag.pii.models import PiiEntityType, PiiStage + + +@dataclass +class FakeBackend: + """Deterministic backend that replays canned spans regardless of input.""" + + spans: list[GlinerSpan] = field(default_factory=list) + calls: list[tuple[str, tuple[str, ...], float]] = field(default_factory=list) + + def predict_entities( + self, + text: str, + labels: list[str], + threshold: float, + ) -> list[GlinerSpan]: + self.calls.append((text, tuple(labels), threshold)) + return list(self.spans) + + +def test_maps_person_label_to_domain_entity(): + text = "John Smith reported to HR." + backend = FakeBackend( + spans=[ + {"start": 0, "end": 10, "text": "John Smith", "label": "person", "score": 0.92}, + ] + ) + rec = GlinerRecognizer(backend) + entities = rec.detect(text) + assert len(entities) == 1 + assert entities[0].entity_type is PiiEntityType.PERSON + assert entities[0].stage is PiiStage.GLINER + assert entities[0].text == "John Smith" + + +def test_drops_unknown_labels(): + text = "Some text" + backend = FakeBackend( + spans=[ + {"start": 0, "end": 4, "text": "Some", "label": "unknown_label", "score": 0.9}, + ] + ) + rec = GlinerRecognizer(backend) + assert rec.detect(text) == [] + + +def test_default_labels_match_domain_map(): + backend = FakeBackend() + rec = GlinerRecognizer(backend) + rec.detect("text") + assert backend.calls[0][1] == DEFAULT_PII_LABELS + + +def test_empty_text_short_circuits(): + backend = FakeBackend( + spans=[ + {"start": 0, "end": 4, "text": "Some", "label": "person", "score": 0.9}, + ] + ) + rec = GlinerRecognizer(backend) + assert rec.detect("") == [] + assert backend.calls == [] + + +def test_threshold_propagates_to_backend(): + backend = FakeBackend() + rec = GlinerRecognizer(backend, threshold=0.42) + rec.detect("anything") + assert backend.calls[0][2] == 0.42 + + +def test_entities_are_sorted_by_start(): + text = "A B C D E F" + backend = FakeBackend( + spans=[ + {"start": 6, "end": 7, "text": "D", "label": "person", "score": 0.9}, + {"start": 0, "end": 1, "text": "A", "label": "person", "score": 0.9}, + ] + ) + rec = GlinerRecognizer(backend) + entities = rec.detect(text) + assert [e.start for e in entities] == [0, 6] diff --git a/tests/test_pii_models.py b/tests/test_pii_models.py new file mode 100644 index 0000000..062930e --- /dev/null +++ b/tests/test_pii_models.py @@ -0,0 +1,66 @@ +"""Domain-model invariants for the PII pipeline.""" + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from guardian_rag.pii.models import ( + AuditResult, + PiiEntity, + PiiEntityType, + PiiStage, + RedactionMap, +) + + +def _entity(start: int, end: int, **overrides: object) -> PiiEntity: + defaults: dict[str, object] = { + "entity_type": PiiEntityType.PERSON, + "start": start, + "end": end, + "score": 0.9, + "text": "John", + "stage": PiiStage.GLINER, + } + defaults.update(overrides) + return PiiEntity(**defaults) # type: ignore[arg-type] + + +def test_entity_rejects_inverted_offsets(): + with pytest.raises(ValidationError): + _entity(10, 5) + + +def test_entity_rejects_zero_length(): + with pytest.raises(ValidationError): + _entity(5, 5) + + +def test_entity_overlap_detection(): + a = _entity(0, 10) + b = _entity(5, 15) + c = _entity(20, 30) + assert a.overlaps(b) + assert b.overlaps(a) + assert not a.overlaps(c) + + +def test_entity_length(): + assert _entity(3, 11).length == 8 + + +def test_redaction_map_reverse_round_trips(): + m = RedactionMap( + document_id="doc-1", + placeholders={"[PERSON_1]": "John Smith", "[EMAIL_1]": "john@example.com"}, + ) + rev = m.reverse() + assert rev["John Smith"] == "[PERSON_1]" + assert rev["john@example.com"] == "[EMAIL_1]" + + +def test_audit_result_clean_default(): + res = AuditResult(clean=True, answer="hello world") + assert res.entities == [] + assert res.fail_closed is False diff --git a/tests/test_pii_presidio.py b/tests/test_pii_presidio.py new file mode 100644 index 0000000..9753442 --- /dev/null +++ b/tests/test_pii_presidio.py @@ -0,0 +1,60 @@ +"""Presidio Stage-1 detector tests. + +These tests do exercise the real Presidio recognisers (pattern matchers, +no NLP engine). They run offline and never touch the network. If Presidio's +default thresholds change upstream, the assertions stay valid because we +check entity types rather than raw scores. +""" + +from __future__ import annotations + +from guardian_rag.pii.models import PiiEntityType, PiiStage +from guardian_rag.pii.presidio_engine import PresidioEngine + + +def test_email_detection(): + engine = PresidioEngine() + text = "Please email john.smith@example.com for details." + entities = engine.detect(text) + emails = [e for e in entities if e.entity_type is PiiEntityType.EMAIL] + assert len(emails) == 1 + assert emails[0].text == "john.smith@example.com" + assert emails[0].stage is PiiStage.PRESIDIO + + +def test_iban_detection(): + engine = PresidioEngine() + text = "Transfer to BE71096123456769 by Friday." + entities = engine.detect(text) + ibans = [e for e in entities if e.entity_type is PiiEntityType.BANK_ACCOUNT] + assert len(ibans) == 1 + assert "BE71096123456769" in ibans[0].text + + +def test_credit_card_detection(): + engine = PresidioEngine() + text = "Card 4111-1111-1111-1111 expires soon." + entities = engine.detect(text) + cards = [e for e in entities if e.entity_type is PiiEntityType.CREDIT_CARD] + assert len(cards) >= 1 + + +def test_ip_address_detection(): + engine = PresidioEngine() + text = "Origin IP 192.168.1.42 reported at 09:00." + entities = engine.detect(text) + ips = [e for e in entities if e.entity_type is PiiEntityType.IP_ADDRESS] + assert len(ips) >= 1 + assert ips[0].text == "192.168.1.42" + + +def test_empty_text_returns_empty(): + assert PresidioEngine().detect("") == [] + + +def test_entities_returned_in_source_order(): + engine = PresidioEngine() + text = "Send to alice@ex.com from 10.0.0.1 by Friday." + entities = engine.detect(text) + starts = [e.start for e in entities] + assert starts == sorted(starts) diff --git a/tests/test_pii_redactor.py b/tests/test_pii_redactor.py new file mode 100644 index 0000000..7537da5 --- /dev/null +++ b/tests/test_pii_redactor.py @@ -0,0 +1,88 @@ +"""Typed-placeholder redactor tests.""" + +from __future__ import annotations + +from guardian_rag.pii.models import PiiEntity, PiiEntityType, PiiStage +from guardian_rag.pii.redactor import Redactor + + +def _entity(start: int, end: int, text: str, etype: PiiEntityType, score: float = 0.9) -> PiiEntity: + return PiiEntity( + entity_type=etype, + start=start, + end=end, + score=score, + text=text, + stage=PiiStage.GLINER, + ) + + +def test_empty_entities_round_trip_unchanged(): + text = "Nothing to redact here." + redacted, mapping = Redactor().redact(text, [], "doc-1") + assert redacted == text + assert mapping.placeholders == {} + assert mapping.entities == [] + + +def test_typed_placeholder_format(): + text = "John Smith works here." + entities = [_entity(0, 10, "John Smith", PiiEntityType.PERSON)] + redacted, mapping = Redactor().redact(text, entities, "doc-1") + assert redacted == "[PERSON_1] works here." + assert mapping.placeholders == {"[PERSON_1]": "John Smith"} + + +def test_same_span_resolves_to_same_placeholder(): + text = "John Smith met John Smith again." + entities = [ + _entity(0, 10, "John Smith", PiiEntityType.PERSON), + _entity(15, 25, "John Smith", PiiEntityType.PERSON), + ] + redacted, mapping = Redactor().redact(text, entities, "doc-1") + assert redacted == "[PERSON_1] met [PERSON_1] again." + assert mapping.placeholders == {"[PERSON_1]": "John Smith"} + + +def test_different_spans_get_distinct_indices(): + text = "Alice met Bob today." + entities = [ + _entity(0, 5, "Alice", PiiEntityType.PERSON), + _entity(10, 13, "Bob", PiiEntityType.PERSON), + ] + redacted, _mapping = Redactor().redact(text, entities, "doc-1") + assert redacted == "[PERSON_1] met [PERSON_2] today." + + +def test_different_types_use_separate_index_space(): + text = "Alice emailed alice@ex.com." + entities = [ + _entity(0, 5, "Alice", PiiEntityType.PERSON), + _entity(14, 26, "alice@ex.com", PiiEntityType.EMAIL), + ] + redacted, _mapping = Redactor().redact(text, entities, "doc-1") + assert redacted == "[PERSON_1] emailed [EMAIL_1]." + + +def test_overlapping_entities_dedupe_keeping_higher_score(): + text = "John Smith works here." + entities = [ + _entity(0, 10, "John Smith", PiiEntityType.PERSON, score=0.6), + _entity(0, 4, "John", PiiEntityType.PERSON, score=0.9), + ] + redacted, _ = Redactor().redact(text, entities, "doc-1") + # The higher-scoring entity wins, which is the 4-char "John" span. + assert redacted.startswith("[PERSON_1] Smith") + + +def test_redaction_preserves_structure_outside_entities(): + text = "Contact: John (HR) john@ex.com, Friday only." + entities = [ + _entity(9, 13, "John", PiiEntityType.PERSON), + _entity(19, 30, "john@ex.com", PiiEntityType.EMAIL), + ] + redacted, _ = Redactor().redact(text, entities, "doc-1") + assert "(HR)" in redacted + assert "Friday only." in redacted + assert "[PERSON_1]" in redacted + assert "[EMAIL_1]" in redacted