diff --git a/CHANGELOG.MD b/CHANGELOG.MD index a2655567..13383256 100644 --- a/CHANGELOG.MD +++ b/CHANGELOG.MD @@ -2,6 +2,24 @@ ## [Unreleased] +### Added + +- **Batch redaction API**: `datafog.redact_many()` scans text fragments + independently while sharing token numbering and pseudonym assignments. + `ScanResult`, `RedactResult`, and `BatchRedactResult` expose computed + `entity_counts` dictionaries so integrations no longer need to recount + result entities. + +### Fixed + +- **LiteLLM request and response coverage**: the compatibility guardrail now + scans Responses API input/output, top-level instructions and system prompts, + text-completion prompts and choices, tool/function arguments and schemas, + structured-output schemas, Anthropic content, and buffered streaming output. + It supports `during_call` blocking, preserves exact and regex allowlists, + adds locale configuration, keeps legacy option names compatible, and never + yields a streaming fragment before the complete text has been checked. + ## [2026-07-05] ### `datafog-python` [4.8.0] diff --git a/README.md b/README.md index b974e48a..3371566e 100644 --- a/README.md +++ b/README.md @@ -31,12 +31,18 @@ values never echoed into logs or transcripts: Manual hook setup and limitations: [examples/claude_code_hook/](examples/claude_code_hook/). - **LiteLLM guardrail** (`DataFogGuardrail`): redacts or blocks PII in - requests and responses at the gateway, for any LiteLLM-proxied provider. - In-process (~40µs per message scanned; a request clears the guardrail in - well under a millisecond), no sidecar service. Setup: + chat, Responses API, text-completion, tool, schema, and streaming traffic at + the gateway, for any LiteLLM-proxied provider. In-process (~40µs per message + scanned; a request clears the guardrail in well under a millisecond), no + sidecar service. Setup: [examples/litellm_guardrail/](examples/litellm_guardrail/). -Both default to the high-precision entity set (`EMAIL`, `PHONE`, +- **Batch redaction** (`redact_many`): scans separate text fragments without + joining them, while preserving token numbering and pseudonyms across the + batch. Scan and redaction results expose privacy-safe `entity_counts` for + policy decisions and observability. + +Both adapters default to the high-precision entity set (`EMAIL`, `PHONE`, `CREDIT_CARD`, `SSN`); noisier types are opt-in. Known-safe values can be exempted with an allowlist: `scan(text, allowlist=[...])` for exact values, `allowlist_patterns=[...]` for full-match regexes (e.g. `^\d{10}$` to stop diff --git a/datafog/__init__.py b/datafog/__init__.py index 7236ac10..07e229c7 100644 --- a/datafog/__init__.py +++ b/datafog/__init__.py @@ -16,10 +16,12 @@ # Core API functions - always available (lightweight) from .core import anonymize_text, detect_pii, get_supported_entities, scan_text -from .engine import Entity, RedactResult, ScanResult +from .engine import BatchRedactResult, Entity, RedactResult, ScanResult from .engine import redact as _redact_entities +from .engine import redact_many as _redact_many_entities from .engine import scan as _scan from .engine import scan_and_redact as _scan_and_redact +from .engine import scan_and_redact_many as _scan_and_redact_many # Essential models - always available from .models.common import EntityTypes @@ -224,6 +226,48 @@ def redact( ) +def redact_many( + texts: list[str], + entities: list[list[Entity]] | None = None, + engine: str = "regex", + entity_types: list[str] | None = None, + strategy: str = "token", + preset: str | None = None, + locales: list[str] | None = None, + allowlist: list[str] | None = None, + allowlist_patterns: list[str] | None = None, +) -> BatchRedactResult: + """Redact multiple fragments with stable numbering across the batch. + + Detection still runs independently per fragment, so entities cannot span + fragment boundaries. Token counters and pseudonym assignments are shared. + """ + if preset is not None: + try: + strategy = _REDACT_PRESETS[preset] + except KeyError as exc: + allowed = ", ".join(sorted(_REDACT_PRESETS)) + raise ValueError(f"preset must be one of: {allowed}") from exc + + if entities is not None: + if allowlist or allowlist_patterns: + raise ValueError( + "allowlist/allowlist_patterns cannot be combined with explicit " + "entities; filter the entities before calling redact_many" + ) + return _redact_many_entities(texts=texts, entities=entities, strategy=strategy) + + return _scan_and_redact_many( + texts=texts, + engine=engine, + entity_types=entity_types, + strategy=strategy, + locales=locales, + allowlist=allowlist, + allowlist_patterns=allowlist_patterns, + ) + + def protect( entity_types: list[str] | None = None, engine: str = "regex", @@ -388,8 +432,10 @@ def process(text: str, anonymize: bool = False, method: str = "redact") -> dict: "Entity", "ScanResult", "RedactResult", + "BatchRedactResult", "scan", "redact", + "redact_many", "protect", "detect", "process", diff --git a/datafog/engine.py b/datafog/engine.py index 419558e9..7b07e739 100644 --- a/datafog/engine.py +++ b/datafog/engine.py @@ -5,6 +5,7 @@ import hashlib import re import warnings +from collections.abc import Iterable from dataclasses import dataclass from functools import lru_cache from typing import Optional @@ -94,6 +95,11 @@ class ScanResult: text: str engine_used: str + @property + def entity_counts(self) -> dict[str, int]: + """Return detected entity totals keyed by canonical entity type.""" + return _entity_counts(self.entities) + @dataclass class RedactResult: @@ -103,6 +109,36 @@ class RedactResult: mapping: dict[str, str] entities: list[Entity] + @property + def entity_counts(self) -> dict[str, int]: + """Return redacted entity totals keyed by canonical entity type.""" + return _entity_counts(self.entities) + + +@dataclass +class BatchRedactResult: + """Result of redacting multiple text fragments as one logical batch.""" + + redacted_texts: list[str] + mapping: dict[str, str] + entities: list[list[Entity]] + + @property + def entity_counts(self) -> dict[str, int]: + """Return redacted entity totals across every batch fragment.""" + return _entity_counts( + entity + for fragment_entities in self.entities + for entity in fragment_entities + ) + + +def _entity_counts(entities: Iterable[Entity]) -> dict[str, int]: + counts: dict[str, int] = {} + for entity in entities: + counts[entity.type] = counts.get(entity.type, 0) + 1 + return counts + def _canonical_type(entity_type: str) -> str: normalized = entity_type.upper().strip() @@ -483,10 +519,26 @@ def redact( if strategy not in {"token", "mask", "hash", "pseudonymize"}: raise ValueError("strategy must be one of: token, mask, hash, pseudonymize") + return _redact_with_state( + text=text, + entities=entities, + strategy=strategy, + counters={}, + pseudonym_by_value={}, + ) + + +def _redact_with_state( + text: str, + entities: list[Entity], + strategy: str, + counters: dict[str, int], + pseudonym_by_value: dict[tuple[str, str], str], +) -> RedactResult: + """Redact one fragment while sharing numbering state with its batch.""" + redacted_text = text mapping: dict[str, str] = {} - counters: dict[str, int] = {} - pseudonym_by_value: dict[tuple[str, str], str] = {} valid_entities = [ entity @@ -539,6 +591,50 @@ def redact( ) +def redact_many( + texts: list[str], + entities: list[list[Entity]], + strategy: str = "token", +) -> BatchRedactResult: + """Redact text fragments with stable numbering across the whole batch.""" + if not isinstance(texts, list) or not all(isinstance(text, str) for text in texts): + raise TypeError("texts must be a list of strings") + if ( + not isinstance(entities, list) + or len(entities) != len(texts) + or not all( + isinstance(fragment_entities, list) for fragment_entities in entities + ) + ): + raise ValueError("entities must contain one entity list per text") + if strategy not in {"token", "mask", "hash", "pseudonymize"}: + raise ValueError("strategy must be one of: token, mask, hash, pseudonymize") + + counters: dict[str, int] = {} + pseudonym_by_value: dict[tuple[str, str], str] = {} + redacted_texts: list[str] = [] + mapping: dict[str, str] = {} + redacted_entities: list[list[Entity]] = [] + + for text, fragment_entities in zip(texts, entities, strict=True): + result = _redact_with_state( + text=text, + entities=fragment_entities, + strategy=strategy, + counters=counters, + pseudonym_by_value=pseudonym_by_value, + ) + redacted_texts.append(result.redacted_text) + mapping.update(result.mapping) + redacted_entities.append(result.entities) + + return BatchRedactResult( + redacted_texts=redacted_texts, + mapping=mapping, + entities=redacted_entities, + ) + + def scan_and_redact( text: str, engine: str = "smart", @@ -558,3 +654,30 @@ def scan_and_redact( allowlist_patterns=allowlist_patterns, ) return redact(text=text, entities=scan_result.entities, strategy=strategy) + + +def scan_and_redact_many( + texts: list[str], + engine: str = "smart", + entity_types: Optional[list[str]] = None, + strategy: str = "token", + locales: Optional[list[str]] = None, + allowlist: Optional[list[str]] = None, + allowlist_patterns: Optional[list[str]] = None, +) -> BatchRedactResult: + """Scan and redact text fragments with batch-stable replacements.""" + if not isinstance(texts, list) or not all(isinstance(text, str) for text in texts): + raise TypeError("texts must be a list of strings") + + entity_groups = [ + scan( + text=text, + engine=engine, + entity_types=entity_types, + locales=locales, + allowlist=allowlist, + allowlist_patterns=allowlist_patterns, + ).entities + for text in texts + ] + return redact_many(texts=texts, entities=entity_groups, strategy=strategy) diff --git a/datafog/integrations/litellm_guardrail.py b/datafog/integrations/litellm_guardrail.py index 7e33ed55..4c4350e5 100644 --- a/datafog/integrations/litellm_guardrail.py +++ b/datafog/integrations/litellm_guardrail.py @@ -1,43 +1,45 @@ -"""LiteLLM guardrail adapter: redact or block PII at the gateway. +"""DataFog PII guardrail: offline, in-process PII redaction and blocking. -Usage (LiteLLM proxy ``config.yaml``):: +Uses the `datafog` library (https://github.com/DataFog/datafog-python) to +detect and redact PII locally — no external service, no sidecar, no network +calls. Scans run in microseconds per request, so the guardrail can sit on +every call without a latency budget. - guardrails: - - guardrail_name: "datafog-pii" - litellm_params: - guardrail: datafog.integrations.litellm_guardrail.DataFogGuardrail - mode: "pre_call" - action: "redact" # or "block" - fail_policy: "open" # or "closed" - # entity_types: ["EMAIL", "PHONE", "CREDIT_CARD", "SSN"] - -Behavior: - -- ``pre_call`` — scans request messages. ``redact`` (default) replaces +- ``pre_call``: scans request messages. ``redact`` (default) replaces findings with ``[TYPE_N]`` tokens before the request leaves the gateway; - ``block`` rejects the request outright. -- ``post_call`` — redacts findings from model responses before they reach + ``block`` rejects the request with HTTP 400. +- ``during_call``: same as pre_call, but runs in parallel with the LLM call + (block only — content cannot be modified mid-flight). +- ``post_call``: redacts findings from model responses before they reach the client. -- ``fail_policy`` — ``open`` (default) lets traffic through unscanned if - the engine errors, so a guardrail bug never takes down the gateway; - ``closed`` rejects traffic instead, for compliance deployments where - unscanned egress is worse than downtime. -Errors and block messages report entity type counts only — matched PII is -never echoed into logs, exceptions, or proxy responses. +Defaults to the high-precision entity types (EMAIL, PHONE, CREDIT_CARD, +SSN). Noisier types (IP_ADDRESS, DOB, ZIP, DE_* locale entities) are opt-in +via ``datafog_entity_types`` because version strings, dates, and short +numbers saturate technical text. + +Errors and block messages report entity type counts only; matched PII is +never echoed into logs, exceptions, or proxy responses. Engine exceptions +are re-raised with ``from None`` and logged by type name only, since their +messages can embed the text being scanned. -Requires ``litellm`` and ``fastapi`` (this module is not imported by -``datafog`` core; the LiteLLM proxy, where this runs, always ships fastapi). +This compatibility class accepts DataFog's original option names (``action``, +``entity_types``, ``fail_policy``, ``allowlist``, ``allowlist_patterns``, and +``locales``) plus the native LiteLLM provider's ``datafog_*`` aliases. """ +import json import logging -from typing import Any, Optional +from collections.abc import AsyncGenerator, Callable +from typing import TYPE_CHECKING, Any, Literal from fastapi import HTTPException from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.types.guardrails import GuardrailEventHooks + +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth -# High-precision defaults, matching the Claude Code hook adapter: noisy-in- -# practice types (IP_ADDRESS, DOB, ZIP) must be opted into explicitly. DEFAULT_ENTITY_TYPES = ["EMAIL", "PHONE", "CREDIT_CARD", "SSN"] VALID_ACTIONS = {"redact", "block"} @@ -49,8 +51,9 @@ def _redact_text( text: str, entity_types: list[str], - allowlist: list[str] | None = None, - allowlist_patterns: list[str] | None = None, + locales: list[str] | None, + allowlist: list[str] | None, + allowlist_patterns: list[str] | None, ) -> tuple[str, dict[str, int]]: """Redact ``text``; return (redacted_text, counts per entity type).""" import datafog @@ -59,69 +62,112 @@ def _redact_text( text, engine="regex", entity_types=entity_types, + locales=locales, allowlist=allowlist, allowlist_patterns=allowlist_patterns, ) - counts: dict[str, int] = {} - for entity in result.entities: - counts[entity.type] = counts.get(entity.type, 0) + 1 - return result.redacted_text, counts + return result.redacted_text, result.entity_counts def _summary(counts: dict[str, int]) -> str: return ", ".join(f"{etype} x{n}" for etype, n in sorted(counts.items())) +def _merge_counts(total_counts: dict[str, int], counts: dict[str, int]) -> None: + for etype, n in counts.items(): + total_counts[etype] = total_counts.get(etype, 0) + n + + +def _get_field(container: Any, field_name: str) -> Any: + if isinstance(container, dict): + return container.get(field_name) + return getattr(container, field_name, None) + + +def _set_field(container: Any, field_name: str, value: Any) -> None: + if isinstance(container, dict): + container[field_name] = value + else: + setattr(container, field_name, value) + + class DataFogGuardrail(CustomGuardrail): - """Offline PII guardrail for the LiteLLM proxy, powered by datafog.""" + """Offline PII guardrail powered by the datafog library.""" def __init__( self, - action: str = "redact", - entity_types: Optional[list[str]] = None, - fail_policy: str = "open", - allowlist: Optional[list[str]] = None, - allowlist_patterns: Optional[list[str]] = None, + action: Literal["redact", "block"] | None = "redact", + entity_types: list[str] | None = None, + fail_policy: Literal["open", "closed"] | None = "open", + allowlist: list[str] | None = None, + allowlist_patterns: list[str] | None = None, + locales: list[str] | None = None, + datafog_action: Literal["redact", "block"] | None = None, + datafog_entity_types: list[str] | None = None, + datafog_locales: list[str] | None = None, + datafog_fail_policy: Literal["open", "closed"] | None = None, + datafog_allowlist: list[str] | None = None, + datafog_allowlist_patterns: list[str] | None = None, **kwargs: Any, ) -> None: - if action not in VALID_ACTIONS: + resolved_action = ( + datafog_action if datafog_action is not None else action or "redact" + ) + resolved_fail_policy = ( + datafog_fail_policy + if datafog_fail_policy is not None + else fail_policy or "open" + ) + if resolved_action not in VALID_ACTIONS: raise ValueError(f"action must be one of: {sorted(VALID_ACTIONS)}") - if fail_policy not in VALID_FAIL_POLICIES: + if resolved_fail_policy not in VALID_FAIL_POLICIES: raise ValueError( f"fail_policy must be one of: {sorted(VALID_FAIL_POLICIES)}" ) - self.action = action - self.entity_types = entity_types or DEFAULT_ENTITY_TYPES - self.fail_policy = fail_policy - self.allowlist = allowlist - self.allowlist_patterns = allowlist_patterns - super().__init__(**kwargs) + + self.action = resolved_action + self.entity_types = ( + datafog_entity_types + if datafog_entity_types is not None + else entity_types or DEFAULT_ENTITY_TYPES + ) + self.locales = datafog_locales if datafog_locales is not None else locales + self.fail_policy = resolved_fail_policy + self.allowlist = ( + datafog_allowlist if datafog_allowlist is not None else allowlist + ) + self.allowlist_patterns = ( + datafog_allowlist_patterns + if datafog_allowlist_patterns is not None + else allowlist_patterns + ) + default_on = kwargs.pop("default_on", True) + super().__init__(default_on=default_on, **kwargs) + + def _redact(self, text: str) -> tuple[str, dict[str, int]]: + return _redact_text( + text, + self.entity_types, + self.locales, + self.allowlist, + self.allowlist_patterns, + ) def _process_content(self, content: Any) -> tuple[Any, dict[str, int]]: """Redact a message content value (str or list of content parts).""" counts: dict[str, int] = {} if isinstance(content, str): - redacted, counts = _redact_text( - content, self.entity_types, self.allowlist, self.allowlist_patterns - ) - return redacted, counts + return self._redact(content) if isinstance(content, list): new_parts = [] skipped_parts = 0 for part in content: if isinstance(part, dict) and isinstance(part.get("text"), str): - redacted, part_counts = _redact_text( - part["text"], - self.entity_types, - self.allowlist, - self.allowlist_patterns, - ) + redacted, part_counts = self._redact(part["text"]) new_parts.append({**part, "text": redacted}) for etype, n in part_counts.items(): counts[etype] = counts.get(etype, 0) + n else: - # Images and other non-text parts are not scanned — - # count them so the blind spot is auditable. new_parts.append(part) skipped_parts += 1 if skipped_parts: @@ -132,79 +178,608 @@ def _process_content(self, content: Any) -> tuple[Any, dict[str, int]]: return new_parts, counts return content, counts + def _process_tool_payload(self, payload: Any) -> tuple[Any, dict[str, int]]: + """Redact text inside supported tool/function argument payloads.""" + return self._process_tool_payload_inner(payload, seen=set()) + + def _process_tool_payload_inner( + self, payload: Any, seen: set[int] + ) -> tuple[Any, dict[str, int]]: + if isinstance(payload, str): + return self._redact(payload) + if not isinstance(payload, (list, dict)) or id(payload) in seen: + return payload, {} + + seen.add(id(payload)) + counts: dict[str, int] = {} + changed = False + + if isinstance(payload, list): + new_items: list[Any] = [] + for value in payload: + new_value, value_counts = self._process_tool_payload_inner(value, seen) + new_items.append(new_value) + if value_counts: + changed = True + _merge_counts(counts, value_counts) + seen.remove(id(payload)) + return (new_items if changed else payload), counts + + new_mapping: dict[Any, Any] = {} + for key, value in payload.items(): + new_value, value_counts = self._process_tool_payload_inner(value, seen) + new_mapping[key] = new_value + if value_counts: + changed = True + _merge_counts(counts, value_counts) + seen.remove(id(payload)) + return (new_mapping if changed else payload), counts + + def _scan_argument_mapping( + self, mapping: dict[str, Any] + ) -> tuple[dict[str, Any], dict[str, int]]: + counts: dict[str, int] = {} + new_mapping = dict(mapping) + changed = False + for field_name in ("arguments", "input"): + payload = mapping.get(field_name) + if not isinstance(payload, (str, list, dict)): + continue + new_payload, payload_counts = self._process_tool_payload(payload) + if payload_counts: + new_mapping[field_name] = new_payload + changed = True + _merge_counts(counts, payload_counts) + return (new_mapping if changed else mapping), counts + + def _scan_argument_container( + self, container: Any, *, redact: bool + ) -> dict[str, int]: + """Scan supported argument fields on a dict/object container.""" + counts: dict[str, int] = {} + for field_name in ("arguments", "input"): + payload = _get_field(container, field_name) + if not isinstance(payload, (str, list, dict)): + continue + new_payload, payload_counts = self._process_tool_payload(payload) + if payload_counts: + if redact: + _set_field(container, field_name, new_payload) + _merge_counts(counts, payload_counts) + return counts + + def _scan_request_tool_calls(self, tool_calls: Any) -> tuple[Any, dict[str, int]]: + if not isinstance(tool_calls, list): + return tool_calls, {} + counts: dict[str, int] = {} + new_tool_calls = [] + changed = False + for tool_call in tool_calls: + if isinstance(tool_call, dict) and isinstance( + tool_call.get("function"), dict + ): + new_function, function_counts = self._scan_argument_mapping( + tool_call["function"] + ) + if function_counts: + new_tool_calls.append({**tool_call, "function": new_function}) + changed = True + _merge_counts(counts, function_counts) + continue + new_tool_calls.append(tool_call) + return (new_tool_calls if changed else tool_calls), counts + + def _scan_message( + self, message: dict[str, Any] + ) -> tuple[dict[str, Any], dict[str, int]]: + counts: dict[str, int] = {} + new_message = dict(message) + changed = False + + if "content" in message: + new_content, content_counts = self._process_content(message["content"]) + if content_counts: + new_message["content"] = new_content + changed = True + _merge_counts(counts, content_counts) + + function_call = message.get("function_call") + if isinstance(function_call, dict): + new_function_call, function_counts = self._scan_argument_mapping( + function_call + ) + if function_counts: + new_message["function_call"] = new_function_call + changed = True + _merge_counts(counts, function_counts) + + new_tool_calls, tool_counts = self._scan_request_tool_calls( + message.get("tool_calls") + ) + if tool_counts: + new_message["tool_calls"] = new_tool_calls + changed = True + _merge_counts(counts, tool_counts) + + return (new_message if changed else message), counts + + def _scan_responses_input_item(self, item: Any) -> tuple[Any, dict[str, int]]: + if isinstance(item, str): + return self._redact(item) + if not isinstance(item, dict): + return item, {} + + counts: dict[str, int] = {} + new_item, message_counts = self._scan_message(item) + changed = bool(message_counts) + _merge_counts(counts, message_counts) + + argument_item, argument_counts = self._scan_argument_mapping(new_item) + if argument_counts: + new_item = argument_item + changed = True + _merge_counts(counts, argument_counts) + + for field_name in ("text", "output"): + payload = new_item.get(field_name) + if not isinstance(payload, (str, list, dict)): + continue + new_payload, payload_counts = self._process_tool_payload(payload) + if payload_counts: + if not changed: + new_item = dict(new_item) + new_item[field_name] = new_payload + changed = True + _merge_counts(counts, payload_counts) + + return (new_item if changed else item), counts + + def _scan_responses_api_input(self, input_value: Any) -> tuple[Any, dict[str, int]]: + if isinstance(input_value, str): + return self._redact(input_value) + if not isinstance(input_value, list): + return input_value, {} + + counts: dict[str, int] = {} + new_items = [] + changed = False + for item in input_value: + new_item, item_counts = self._scan_responses_input_item(item) + new_items.append(new_item) + if item_counts: + changed = True + _merge_counts(counts, item_counts) + + return (new_items if changed else input_value), counts + + def _scan_top_level_prompt_fields(self, data: dict) -> tuple[dict, dict[str, int]]: + counts: dict[str, int] = {} + new_data = data + for field_name in ("instructions", "system"): + if field_name not in data: + continue + new_value, field_counts = self._process_content(data[field_name]) + if field_counts: + if new_data is data: + new_data = dict(data) + new_data[field_name] = new_value + _merge_counts(counts, field_counts) + return new_data, counts + + def _scan_schema_fields(self, data: dict) -> tuple[dict, dict[str, int]]: + counts: dict[str, int] = {} + new_data = data + for field_name in ("tools", "functions", "response_format"): + schema = data.get(field_name) + if not isinstance(schema, (list, dict)): + continue + new_schema, field_counts = self._process_tool_payload(schema) + if field_counts: + if new_data is data: + new_data = dict(data) + new_data[field_name] = new_schema + _merge_counts(counts, field_counts) + return new_data, counts + + def _process_completion_prompt(self, prompt: Any) -> tuple[Any, dict[str, int]]: + if isinstance(prompt, str): + return self._redact(prompt) + if not isinstance(prompt, list) or not all( + isinstance(item, str) for item in prompt + ): + return prompt, {} + + counts: dict[str, int] = {} + new_prompt = [] + changed = False + for item in prompt: + new_item, item_counts = self._redact(item) + new_prompt.append(new_item) + if item_counts: + changed = True + _merge_counts(counts, item_counts) + return (new_prompt if changed else prompt), counts + def _handle_engine_error(self, exc: Exception) -> None: - # Only the exception *type* is ever logged or re-raised. Engine - # exception messages can embed the text being scanned, so chaining - # (`from exc`) or interpolating str(exc) would leak matched PII into - # tracebacks and logs — the exact thing this guardrail exists to - # prevent. `from None` suppresses both __cause__ and __context__. + """Apply the fail policy without leaking scanned text. + + Only the exception type is logged or re-raised: engine exception + messages can embed the text being scanned, so chaining via ``from + exc`` or interpolating ``str(exc)`` would leak matched PII into + tracebacks. A missing dependency is a config error and is surfaced + regardless of fail policy. + """ + if isinstance(exc, ImportError): + raise exc if self.fail_policy == "closed": - # RuntimeError (no status_code attr) intentionally surfaces as - # HTTP 500: an engine failure is a server fault, distinct from - # the policy block below, which is a 400. raise RuntimeError( - "DataFog guardrail failed and fail_policy is 'closed'; " - f"rejecting unscanned traffic ({type(exc).__name__})." + "DataFog guardrail failed and fail_policy is " + f"'closed'; rejecting unscanned traffic ({type(exc).__name__})." ) from None logger.warning( "DataFog guardrail error (fail-open, traffic unscanned): %s", type(exc).__name__, ) - async def async_pre_call_hook( - self, - user_api_key_dict: Any, - cache: Any, - data: dict, - call_type: str, - ) -> dict: + def _scan_messages(self, data: dict) -> tuple[dict, dict[str, int]]: + """Scan/redact all message contents; return (new_data, counts).""" messages = data.get("messages") if not isinstance(messages, list): - return data + return data, {} total_counts: dict[str, int] = {} new_messages = [] - try: - for message in messages: - if isinstance(message, dict) and "content" in message: - new_content, counts = self._process_content(message["content"]) - new_messages.append({**message, "content": new_content}) - for etype, n in counts.items(): - total_counts[etype] = total_counts.get(etype, 0) + n - else: - new_messages.append(message) - except Exception as exc: # noqa: BLE001 — fail policy decides - self._handle_engine_error(exc) - return data + for message in messages: + if isinstance(message, dict): + new_message, counts = self._scan_message(message) + new_messages.append(new_message) + _merge_counts(total_counts, counts) + else: + new_messages.append(message) if not total_counts: - return data + return data, {} + return {**data, "messages": new_messages}, total_counts - if self.action == "block": - self._record_guardrail_logging(data, total_counts) - # HTTPException(400) is one of the exception types litellm's - # _is_guardrail_intervention recognizes, so the block is - # classified as a policy intervention (not a backend failure) - # and reaches the client as a 400, not a 500. - # Counts only — never the matched values. - raise HTTPException( - status_code=400, - detail={ - "error": ( - f"DataFog PII guardrail: request blocked, messages " - f"contain {_summary(total_counts)}." + def _scan_request_data(self, data: dict) -> tuple[dict, dict[str, int]]: + """Scan/redact supported request payload fields.""" + new_data, total_counts = self._scan_messages(data) + + new_input, input_counts = self._scan_responses_api_input(new_data.get("input")) + if input_counts: + new_data = {**new_data, "input": new_input} + _merge_counts(total_counts, input_counts) + + new_prompt_data, prompt_counts = self._scan_top_level_prompt_fields(new_data) + if prompt_counts: + new_data = new_prompt_data + _merge_counts(total_counts, prompt_counts) + + if "prompt" in new_data: + new_prompt, completion_prompt_counts = self._process_completion_prompt( + new_data["prompt"] + ) + if completion_prompt_counts: + new_data = {**new_data, "prompt": new_prompt} + _merge_counts(total_counts, completion_prompt_counts) + + new_schema_data, schema_counts = self._scan_schema_fields(new_data) + if schema_counts: + new_data = new_schema_data + _merge_counts(total_counts, schema_counts) + + if not total_counts: + return data, {} + return new_data, total_counts + + def _scan_text_field( + self, container: Any, field_name: str, *, redact: bool + ) -> dict[str, int]: + text = _get_field(container, field_name) + if not isinstance(text, str): + return {} + new_text, counts = self._redact(text) + if counts and redact: + _set_field(container, field_name, new_text) + return counts + + def _scan_response_content(self, container: Any, *, redact: bool) -> dict[str, int]: + content = _get_field(container, "content") + if content is None: + return {} + new_content, counts = self._process_content(content) + if counts and redact: + _set_field(container, "content", new_content) + if isinstance(content, list): + for part in content: + if isinstance(part, dict) and part.get("type") == "tool_use": + input_counts = self._scan_argument_container(part, redact=redact) + _merge_counts(counts, input_counts) + return counts + + def _scan_response_tool_calls( + self, message: Any, *, redact: bool + ) -> dict[str, int]: + counts: dict[str, int] = {} + function_call = _get_field(message, "function_call") + if function_call is not None: + _merge_counts( + counts, self._scan_argument_container(function_call, redact=redact) + ) + + tool_calls = _get_field(message, "tool_calls") + if isinstance(tool_calls, list): + for tool_call in tool_calls: + function = _get_field(tool_call, "function") + if function is not None: + _merge_counts( + counts, self._scan_argument_container(function, redact=redact) + ) + return counts + + def _scan_response_message(self, message: Any, *, redact: bool) -> dict[str, int]: + counts = self._scan_response_content(message, redact=redact) + _merge_counts(counts, self._scan_response_tool_calls(message, redact=redact)) + return counts + + def _scan_responses_api_output( + self, response: Any, *, redact: bool + ) -> dict[str, int]: + counts: dict[str, int] = {} + output = _get_field(response, "output") + if not isinstance(output, list): + return counts + for item in output: + item_type = _get_field(item, "type") + if item_type == "message": + _merge_counts(counts, self._scan_response_content(item, redact=redact)) + elif item_type in {"function_call", "custom_tool_call"}: + _merge_counts( + counts, self._scan_argument_container(item, redact=redact) + ) + return counts + + @staticmethod + def _stream_group_key(prefix: str, *parts: Any) -> tuple[Any, ...]: + return (prefix, *parts) + + @staticmethod + def _field_setter(container: Any, field_name: str) -> Callable[[str], None]: + def set_value(new_text: str) -> None: + _set_field(container, field_name, new_text) + + return set_value + + def _sse_chunk_setter( + self, chunks: list[Any], chunk_index: int + ) -> Callable[[str], None]: + def set_value(new_text: str) -> None: + chunks[chunk_index] = self._replace_sse_text_delta( + chunks[chunk_index], new_text + ) + + return set_value + + @staticmethod + def _decode_sse_text_delta(chunk: bytes | bytearray) -> tuple[Any, str] | None: + raw = bytes(chunk).decode("utf-8", errors="replace") + for line in raw.splitlines(): + stripped = line.strip() + if not stripped.startswith("data: "): + continue + try: + data = json.loads(stripped[6:]) + except (json.JSONDecodeError, ValueError): + continue + if not isinstance(data, dict) or data.get("type") != "content_block_delta": + continue + delta = data.get("delta") + if ( + isinstance(delta, dict) + and delta.get("type") == "text_delta" + and isinstance(delta.get("text"), str) + ): + return data.get("index"), delta["text"] + return None + + @staticmethod + def _replace_sse_text_delta(chunk: bytes | bytearray, new_text: str) -> bytes: + raw = bytes(chunk).decode("utf-8", errors="replace") + lines = raw.splitlines(keepends=True) + for index, line in enumerate(lines): + line_ending = "" + line_body = line + if line.endswith("\r\n"): + line_ending = "\r\n" + line_body = line[:-2] + elif line.endswith("\n"): + line_ending = "\n" + line_body = line[:-1] + + leading = line_body[: len(line_body) - len(line_body.lstrip())] + stripped = line_body.lstrip() + if not stripped.startswith("data: "): + continue + try: + data = json.loads(stripped[6:]) + except (json.JSONDecodeError, ValueError): + continue + if not isinstance(data, dict) or data.get("type") != "content_block_delta": + continue + delta = data.get("delta") + if not isinstance(delta, dict) or delta.get("type") != "text_delta": + continue + delta["text"] = new_text + lines[index] = ( + f"{leading}data: {json.dumps(data, separators=(',', ':'))}{line_ending}" + ) + break + return "".join(lines).encode("utf-8") + + def _append_chat_stream_text_refs( + self, + chunk: Any, + refs: list[tuple[tuple[Any, ...], str, Callable[[str], None]]], + ) -> None: + choices = _get_field(chunk, "choices") + if not isinstance(choices, list): + return + + for choice_index, choice in enumerate(choices): + choice_text = _get_field(choice, "text") + if isinstance(choice_text, str) and choice_text: + key = self._stream_group_key( + "completion", _get_field(choice, "index") or choice_index + ) + refs.append( + ( + key, + choice_text, + self._field_setter(choice, "text"), ) - }, + ) + + delta = _get_field(choice, "delta") + if delta is None: + continue + content = _get_field(delta, "content") + if not isinstance(content, str) or not content: + continue + key = self._stream_group_key( + "chat", _get_field(choice, "index") or choice_index + ) + refs.append( + ( + key, + content, + self._field_setter(delta, "content"), + ) ) - new_data = {**data, "messages": new_messages} - self._record_guardrail_logging(new_data, total_counts) - return new_data + def _append_responses_stream_text_ref( + self, + chunk: Any, + refs: list[tuple[tuple[Any, ...], str, Callable[[str], None]]], + ) -> None: + if _get_field(chunk, "type") != "response.output_text.delta": + return + delta = _get_field(chunk, "delta") + if not isinstance(delta, str) or not delta: + return + key = self._stream_group_key( + "responses", + _get_field(chunk, "item_id"), + _get_field(chunk, "output_index"), + _get_field(chunk, "content_index"), + ) + refs.append( + ( + key, + delta, + self._field_setter(chunk, "delta"), + ) + ) + + def _append_anthropic_stream_text_ref( + self, + chunk: Any, + refs: list[tuple[tuple[Any, ...], str, Callable[[str], None]]], + ) -> None: + if _get_field(chunk, "type") != "content_block_delta": + return + delta = _get_field(chunk, "delta") + if _get_field(delta, "type") != "text_delta": + return + text = _get_field(delta, "text") + if not isinstance(text, str) or not text: + return + key = self._stream_group_key("anthropic", _get_field(chunk, "index")) + refs.append( + ( + key, + text, + self._field_setter(delta, "text"), + ) + ) + + def _collect_stream_text_refs( + self, chunks: list[Any] + ) -> list[tuple[tuple[Any, ...], str, Callable[[str], None]]]: + refs: list[tuple[tuple[Any, ...], str, Callable[[str], None]]] = [] + for chunk_index, chunk in enumerate(chunks): + if isinstance(chunk, (bytes, bytearray)): + sse_delta = self._decode_sse_text_delta(chunk) + if sse_delta: + content_index, text = sse_delta + key = self._stream_group_key("anthropic_sse", content_index) + refs.append( + ( + key, + text, + self._sse_chunk_setter(chunks, chunk_index), + ) + ) + continue + + self._append_chat_stream_text_refs(chunk, refs) + self._append_responses_stream_text_ref(chunk, refs) + self._append_anthropic_stream_text_ref(chunk, refs) + + return refs + + def _scan_stream_text_refs( + self, + refs: list[tuple[tuple[Any, ...], str, Callable[[str], None]]], + *, + redact: bool, + ) -> dict[str, int]: + counts: dict[str, int] = {} + refs_by_group: dict[ + tuple[Any, ...], list[tuple[str, Callable[[str], None]]] + ] = {} + for key, text, setter in refs: + refs_by_group.setdefault(key, []).append((text, setter)) + + for grouped_refs in refs_by_group.values(): + stream_text = "".join(text for text, _ in grouped_refs) + redacted_text, group_counts = self._redact(stream_text) + if not group_counts: + continue + _merge_counts(counts, group_counts) + if not redact: + continue + first_ref = True + for _, setter in grouped_refs: + setter(redacted_text if first_ref else "") + first_ref = False + + return counts + + def _scan_streaming_chunks( + self, chunks: list[Any], *, redact: bool + ) -> dict[str, int]: + refs = self._collect_stream_text_refs(chunks) + if not refs: + return {} + return self._scan_stream_text_refs(refs, redact=redact) + + def _raise_block(self, total_counts: dict[str, int]) -> None: + """Reject with HTTP 400 so the block is classified as a guardrail + intervention by ``_is_guardrail_intervention`` rather than a + backend failure, and reaches the client as 400 instead of 500.""" + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"Violated DataFog PII guardrail policy: request contains {_summary(total_counts)}." + ) + }, + ) def _record_guardrail_logging( self, data: dict, total_counts: dict[str, int] ) -> None: - """Record the decision into litellm's standard guardrail logging.""" + """Record the decision into standard guardrail logging.""" try: self.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response=_summary(total_counts), @@ -212,45 +787,167 @@ def _record_guardrail_logging( guardrail_status="guardrail_intervened", masked_entity_count=dict(total_counts), ) - except Exception: # noqa: BLE001 — observability must never break traffic - logger.debug("could not record guardrail logging information") + except ( + Exception + ): # noqa: BLE001 # logging failures must not break guardrail enforcement + logger.debug("DataFog guardrail: could not record logging information") + + async def async_pre_call_hook( + self, + user_api_key_dict: "UserAPIKeyAuth", + cache: Any, + data: dict, + call_type: str, + ) -> Exception | str | dict | None: + if ( + self.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.pre_call + ) + is not True + ): + return data + try: + new_data, total_counts = self._scan_request_data(data) + except ( + Exception + ) as exc: # noqa: BLE001 # engine failures are handled by the configured fail policy + self._handle_engine_error(exc) + return data + + if not total_counts: + return data + + if self.action == "block": + self._record_guardrail_logging(data, total_counts) + self._raise_block(total_counts) + self._record_guardrail_logging(new_data, total_counts) + return new_data + + async def async_moderation_hook( + self, + data: dict, + user_api_key_dict: "UserAPIKeyAuth", + call_type: str, + ) -> Any: + """Block on PII during the parallel call window. + + Content cannot be rewritten mid-flight, so redact mode is a no-op + here; only block has an effect. + """ + if self.action != "block": + return data + if ( + self.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.during_call + ) + is not True + ): + return data + try: + _, total_counts = self._scan_request_data(data) + except ( + Exception + ) as exc: # noqa: BLE001 # engine failures are handled by the configured fail policy + self._handle_engine_error(exc) + return data + + if total_counts: + self._record_guardrail_logging(data, total_counts) + self._raise_block(total_counts) + return data async def async_post_call_success_hook( self, data: dict, - user_api_key_dict: Any, + user_api_key_dict: "UserAPIKeyAuth", response: Any, ) -> Any: - """Redact PII from model responses. + """Redact PII from model responses, or block when action is block. - Mutates ``response`` in place — deliberate: litellm post_call - guardrails share the response object rather than cloning it, and - an unredacted clone escaping through another callback would defeat - the purpose. + Redaction mutates ``response`` in place, deliberately: post_call + guardrails share the response object, and an unredacted copy + escaping through another callback would defeat the purpose. """ - choices = getattr(response, "choices", None) - if not choices: + if ( + self.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.post_call + ) + is not True + ): return response + response_counts: dict[str, int] = {} try: - skipped_parts = 0 - for choice in choices: - message = getattr(choice, "message", None) - if message is not None and isinstance(message.content, str): - redacted, counts = _redact_text( - message.content, - self.entity_types, - self.allowlist, - self.allowlist_patterns, + redact = self.action != "block" + choices = getattr(response, "choices", None) + if choices: + for choice in choices: + message = getattr(choice, "message", None) + if message is not None: + _merge_counts( + response_counts, + self._scan_response_message(message, redact=redact), + ) + _merge_counts( + response_counts, + self._scan_text_field(choice, "text", redact=redact), ) - if counts: - message.content = redacted - elif message is not None and message.content is not None: - skipped_parts += 1 - if skipped_parts: - logger.debug( - "DataFog guardrail: %d non-text response parts not scanned", - skipped_parts, + _merge_counts( + response_counts, + self._scan_responses_api_output(response, redact=redact), + ) + if isinstance(response, dict): + _merge_counts( + response_counts, + self._scan_response_content(response, redact=redact), ) - except Exception as exc: # noqa: BLE001 — fail policy decides + except ( + Exception + ) as exc: # noqa: BLE001 # engine failures are handled by the configured fail policy self._handle_engine_error(exc) + return response + if response_counts: + self._record_guardrail_logging(data, response_counts) + if self.action == "block": + self._raise_block(response_counts) return response + + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: "UserAPIKeyAuth", + response: Any, + request_data: dict, + ) -> AsyncGenerator[Any, None]: + if ( + self.should_run_guardrail( + data=request_data, event_type=GuardrailEventHooks.post_call + ) + is not True + ): + async for chunk in response: + yield chunk + return + + chunks = [] + async for chunk in response: + chunks.append(chunk) + + stream_counts: dict[str, int] = {} + try: + stream_counts = self._scan_streaming_chunks( + chunks, redact=self.action != "block" + ) + except ( + Exception + ) as exc: # noqa: BLE001 # engine failures are handled by the configured fail policy + self._handle_engine_error(exc) + for chunk in chunks: + yield chunk + return + + if stream_counts: + self._record_guardrail_logging(request_data, stream_counts) + if self.action == "block": + self._raise_block(stream_counts) + + for chunk in chunks: + yield chunk diff --git a/docs/python-sdk.rst b/docs/python-sdk.rst index ce70f577..39da94e8 100644 --- a/docs/python-sdk.rst +++ b/docs/python-sdk.rst @@ -19,9 +19,21 @@ and do not require OCR, Spark, model downloads, or distributed dependencies. redact_result = datafog.redact(text, engine="regex") print(redact_result.redacted_text) + print(redact_result.entity_counts) + + batch_result = datafog.redact_many( + ["Email jane@example.com", "Call 415-555-1212"], + engine="regex", + ) + print(batch_result.redacted_texts) print(datafog.sanitize(text)) +``redact_many()`` scans each fragment independently while sharing token +numbering and pseudonym assignments across the batch. This is useful for +messages, tool payloads, log fields, and other structured records where text +boundaries must remain intact. + The backward-compatible ``DataFog`` and ``TextService`` classes remain available for existing users. ``TextService(engine="regex")`` is the dependency-light service path; ``spacy``, ``gliner``, ``smart``, OCR, and Spark diff --git a/examples/litellm_guardrail/README.md b/examples/litellm_guardrail/README.md index 7d10b44a..e5882511 100644 --- a/examples/litellm_guardrail/README.md +++ b/examples/litellm_guardrail/README.md @@ -19,6 +19,22 @@ pip install datafog litellm litellm --config config.yaml # see config.yaml in this directory ``` +LiteLLM releases with the native DataFog provider can use the shorter +registration below: + +```yaml +guardrails: + - guardrail_name: "datafog-pii" + litellm_params: + guardrail: datafog + mode: ["pre_call", "post_call"] + default_on: true + datafog_action: "redact" +``` + +The bundled `config.yaml` uses DataFog's compatibility class so it also works +with older LiteLLM releases and supports DataFog-specific allowlists. + With `action: redact` (default), a request containing > email the report to jane.doe@example.invalid @@ -27,9 +43,16 @@ reaches your model provider as > email the report to [EMAIL_1] -Response-side redaction only runs when `post_call` is included in `mode` -(the example config registers both: `mode: ["pre_call", "post_call"]`). -With it, PII in model _responses_ is redacted before reaching the client. +Request scanning covers chat messages, Responses API input and instructions, +text-completion prompts, tool arguments, tool definitions, legacy functions, +and structured-output schemas. Response scanning covers chat and text +completions, Responses API output, Anthropic content, tool-call arguments, and +streaming output. Streaming responses are buffered until the complete text can +be checked, so no partial PII is yielded before detection. + +Response-side redaction only runs when `post_call` is included in `mode` (the +example config registers both: `mode: ["pre_call", "post_call"]`). With it, +PII in model _responses_ is redacted before reaching the client. In `block` mode, rejected requests return **HTTP 400** with an entity-type summary — litellm classifies them as guardrail interventions, not backend @@ -44,3 +67,11 @@ errors, so monitoring stays accurate. - `fail_policy`: `open` (engine error → traffic passes unscanned, gateway stays up) or `closed` (engine error → traffic rejected; for compliance deployments where unscanned egress is worse than downtime) +- `allowlist`: exact entity values that should remain unchanged +- `allowlist_patterns`: full-match regular expressions for known-safe values; + treat these as trusted operator configuration, not user input +- `locales`: locale packs to enable, such as `["de"]` for German identifiers + +The compatibility class also accepts the native provider's prefixed forms: +`datafog_action`, `datafog_entity_types`, `datafog_fail_policy`, +`datafog_allowlist`, `datafog_allowlist_patterns`, and `datafog_locales`. diff --git a/examples/litellm_guardrail/config.yaml b/examples/litellm_guardrail/config.yaml index f9be59ed..b1629728 100644 --- a/examples/litellm_guardrail/config.yaml +++ b/examples/litellm_guardrail/config.yaml @@ -15,3 +15,6 @@ guardrails: action: "redact" # "redact" replaces PII with [TYPE_N] tokens; "block" rejects fail_policy: "open" # "closed" rejects traffic if the scan engine errors # entity_types: ["EMAIL", "PHONE", "CREDIT_CARD", "SSN"] # defaults shown + # allowlist: ["support@example.com"] + # allowlist_patterns: ["^555-0100$"] # trusted operator config only + # locales: ["de"] diff --git a/tests/test_engine_api.py b/tests/test_engine_api.py index 2a43e147..d6caabbf 100644 --- a/tests/test_engine_api.py +++ b/tests/test_engine_api.py @@ -4,7 +4,14 @@ import pytest -from datafog.engine import Entity, redact, scan, scan_and_redact +from datafog.engine import ( + Entity, + redact, + redact_many, + scan, + scan_and_redact, + scan_and_redact_many, +) from datafog.exceptions import EngineNotAvailable @@ -15,6 +22,7 @@ def test_scan_regex_detects_structured_entities() -> None: assert "EMAIL" in entity_types assert "SSN" in entity_types assert result.engine_used == "regex" + assert result.entity_counts == {"EMAIL": 1, "SSN": 1} def test_scan_filters_entity_types() -> None: @@ -82,6 +90,56 @@ def test_redact_token_numbering_follows_document_order() -> None: "[EMAIL_1]": "alpha@example.com", "[EMAIL_2]": "beta@example.com", } + assert result.entity_counts == {"EMAIL": 2} + + +def test_redact_many_preserves_numbering_across_fragments() -> None: + texts = ["first alpha@example.com", "then beta@example.com"] + entities = [ + [_entity_at(texts[0], "alpha@example.com")], + [_entity_at(texts[1], "beta@example.com")], + ] + + result = redact_many(texts=texts, entities=entities) + + assert result.redacted_texts == ["first [EMAIL_1]", "then [EMAIL_2]"] + assert result.mapping == { + "[EMAIL_1]": "alpha@example.com", + "[EMAIL_2]": "beta@example.com", + } + assert result.entity_counts == {"EMAIL": 2} + + +def test_redact_many_reuses_pseudonym_across_fragments() -> None: + texts = ["first alpha@example.com", "again alpha@example.com"] + entities = [ + [_entity_at(texts[0], "alpha@example.com")], + [_entity_at(texts[1], "alpha@example.com")], + ] + + result = redact_many(texts=texts, entities=entities, strategy="pseudonymize") + + assert result.redacted_texts == [ + "first [EMAIL_PSEUDO_1]", + "again [EMAIL_PSEUDO_1]", + ] + assert result.mapping == {"[EMAIL_PSEUDO_1]": "alpha@example.com"} + + +def test_scan_and_redact_many_scans_each_fragment() -> None: + result = scan_and_redact_many( + ["first alpha@example.com", "then beta@example.com"], + engine="regex", + ) + + assert result.redacted_texts == ["first [EMAIL_1]", "then [EMAIL_2]"] + + +def test_redact_many_validates_shape() -> None: + with pytest.raises(TypeError, match="list of strings"): + redact_many(["valid", None], [[], []]) # type: ignore[list-item] + with pytest.raises(ValueError, match="one entity list per text"): + redact_many(["one", "two"], [[]]) def test_redact_pseudonymize_numbering_follows_document_order() -> None: diff --git a/tests/test_litellm_guardrail.py b/tests/test_litellm_guardrail.py index 1d7656a4..dfeaf158 100644 --- a/tests/test_litellm_guardrail.py +++ b/tests/test_litellm_guardrail.py @@ -1,21 +1,26 @@ -"""Tests for the LiteLLM guardrail adapter (DataFogGuardrail). +""" +Tests for the DataFog PII Guardrail. -PII literals below are split ("jane.doe@" "acme.com") so this source file -itself never contains a contiguous match — the values only assemble at -runtime. This keeps write-time PII scanners (including our own Claude Code -hook) quiet while the tests exercise real detections. +All PII values below are synthetic test fixtures (documentation-reserved +domains, test card numbers, invalid SSN ranges). """ +import json +from types import SimpleNamespace + import pytest -litellm = pytest.importorskip("litellm") -pytest.importorskip("fastapi") # adapter raises fastapi.HTTPException on block +pytest.importorskip("litellm") +pytest.importorskip("fastapi") + +from fastapi import HTTPException # noqa: E402 from datafog.integrations.litellm_guardrail import DataFogGuardrail # noqa: E402 -EMAIL = "jane.doe@" "acme.com" +EMAIL = "jane.doe@" "example.com" CARD = "4242 4242 " "4242 4242" SSN = "856-45-" "6789" +DE_TAX_ID = "12345678901" def _chat_data(content) -> dict: @@ -23,306 +28,1120 @@ def _chat_data(content) -> dict: def _model_response(text: str): + import litellm + resp = litellm.ModelResponse() resp.choices[0].message.content = text return resp +async def _async_iter(items): + for item in items: + yield item + + +async def _collect_async(async_iterable): + return [item async for item in async_iterable] + + @pytest.mark.asyncio -class TestPreCallRedact: - async def test_redacts_email_in_message(self): - guardrail = DataFogGuardrail(guardrail_name="datafog-pii") - data = await guardrail.async_pre_call_hook( +async def test_pre_call_redacts_email(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=_chat_data(f"email the report to {EMAIL} please"), + call_type="completion", + ) + content = data["messages"][0]["content"] + assert EMAIL not in content + assert "[EMAIL_1]" in content + + +@pytest.mark.asyncio +async def test_pre_call_clean_message_unchanged(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=_chat_data("summarize this design doc"), + call_type="completion", + ) + assert data["messages"][0]["content"] == "summarize this design doc" + + +@pytest.mark.asyncio +async def test_pre_call_redacts_content_parts_and_skips_images(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=_chat_data( + [ + {"type": "image_url", "image_url": {"url": "data:image/png;base64,x"}}, + {"type": "text", "text": f"ssn is {SSN}"}, + ] + ), + call_type="completion", + ) + parts = data["messages"][0]["content"] + assert parts[0]["type"] == "image_url" + assert SSN not in parts[1]["text"] + + +@pytest.mark.asyncio +async def test_pre_call_redacts_tool_and_function_call_arguments(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + original = { + "messages": [ + { + "role": "assistant", + "content": None, + "function_call": { + "name": "legacy_lookup", + "arguments": json.dumps({"ssn": SSN}), + }, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "send_email", + "arguments": json.dumps({"recipient": EMAIL}), + }, + } + ], + } + ] + } + data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=original, + call_type="completion", + ) + + message = data["messages"][0] + assert message["content"] is None + assert SSN not in message["function_call"]["arguments"] + assert EMAIL not in message["tool_calls"][0]["function"]["arguments"] + assert SSN in original["messages"][0]["function_call"]["arguments"] + assert EMAIL in original["messages"][0]["tool_calls"][0]["function"]["arguments"] + + +@pytest.mark.asyncio +async def test_pre_call_redacts_tool_and_function_definitions(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + original = { + "messages": [{"role": "user", "content": "look up the customer"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup_customer", + "description": f"Send the result to {EMAIL}", + "parameters": { + "type": "object", + "properties": { + "card": { + "type": "string", + "description": f"Customer card number, such as {CARD}", + } + }, + }, + }, + }, + { + "name": "anthropic_lookup", + "description": f"Look up customer {SSN}", + "input_schema": { + "type": "object", + "properties": { + "email": {"type": "string", "examples": [EMAIL]}, + }, + }, + }, + ], + "functions": [ + { + "name": "legacy_lookup", + "description": f"Look up card {CARD}", + "parameters": { + "type": "object", + "properties": { + "ssn": {"type": "string", "default": SSN}, + }, + }, + } + ], + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "customer_record", + "schema": { + "type": "object", + "properties": { + "email": { + "type": "string", + "description": f"For example, {EMAIL}", + }, + }, + }, + }, + }, + } + + data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=original, + call_type="completion", + ) + + assert EMAIL not in data["tools"][0]["function"]["description"] + assert ( + CARD + not in data["tools"][0]["function"]["parameters"]["properties"]["card"][ + "description" + ] + ) + assert SSN not in data["tools"][1]["description"] + assert ( + EMAIL not in data["tools"][1]["input_schema"]["properties"]["email"]["examples"] + ) + assert CARD not in data["functions"][0]["description"] + assert SSN not in data["functions"][0]["parameters"]["properties"]["ssn"]["default"] + assert ( + EMAIL + not in data["response_format"]["json_schema"]["schema"]["properties"]["email"][ + "description" + ] + ) + assert EMAIL in original["tools"][0]["function"]["description"] + assert SSN in original["functions"][0]["parameters"]["properties"]["ssn"]["default"] + assert ( + EMAIL + in original["response_format"]["json_schema"]["schema"]["properties"]["email"][ + "description" + ] + ) + + +@pytest.mark.asyncio +async def test_pre_call_redacts_responses_api_input_string(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + original = {"input": f"Please contact {EMAIL}"} + + data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=original, + call_type="responses", + ) + + assert EMAIL not in data["input"] + assert "[EMAIL_1]" in data["input"] + assert original["input"] == f"Please contact {EMAIL}" + + +@pytest.mark.asyncio +async def test_pre_call_redacts_responses_api_input_items(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + original = { + "input": [ + { + "role": "user", + "content": [{"type": "input_text", "text": f"email {EMAIL}"}], + }, + { + "type": "function_call_output", + "call_id": "call_1", + "output": {"ssn": SSN}, + }, + ] + } + + data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=original, + call_type="responses", + ) + + assert EMAIL not in data["input"][0]["content"][0]["text"] + assert SSN not in data["input"][1]["output"]["ssn"] + assert EMAIL in original["input"][0]["content"][0]["text"] + assert SSN in original["input"][1]["output"]["ssn"] + + +@pytest.mark.asyncio +async def test_pre_call_redacts_top_level_instructions_and_system_string(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + original = { + "instructions": f"Route customer follow-up to {EMAIL}", + "system": f"Billing card on file is {CARD}", + "input": "summarize the customer account", + } + + data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=original, + call_type="responses", + ) + + assert EMAIL not in data["instructions"] + assert "[EMAIL_1]" in data["instructions"] + assert CARD not in data["system"] + assert "[CREDIT_CARD_1]" in data["system"] + assert original["instructions"] == f"Route customer follow-up to {EMAIL}" + assert original["system"] == f"Billing card on file is {CARD}" + + +@pytest.mark.asyncio +async def test_pre_call_redacts_top_level_anthropic_system_content_blocks(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + original = { + "system": [ + {"type": "text", "text": f"Use tax profile {SSN} for context"}, + {"type": "cache_control", "cache_control": {"type": "ephemeral"}}, + ], + "messages": [{"role": "user", "content": "hello"}], + } + + data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=original, + call_type="completion", + ) + + assert SSN not in data["system"][0]["text"] + assert "[SSN_1]" in data["system"][0]["text"] + assert data["system"][1] == original["system"][1] + assert SSN in original["system"][0]["text"] + + +@pytest.mark.asyncio +async def test_pre_call_redacts_text_completion_prompt_string_and_list(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + original_string = {"prompt": f"Complete the account note for {EMAIL}"} + original_list = {"prompt": [f"Email {EMAIL}", f"SSN {SSN}"]} + + string_data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=original_string, + call_type="completion", + ) + list_data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=original_list, + call_type="completion", + ) + + assert EMAIL not in string_data["prompt"] + assert "[EMAIL_1]" in string_data["prompt"] + assert EMAIL not in list_data["prompt"][0] + assert SSN not in list_data["prompt"][1] + assert original_string["prompt"] == f"Complete the account note for {EMAIL}" + assert original_list["prompt"] == [f"Email {EMAIL}", f"SSN {SSN}"] + + +def test_process_content_returns_unmodified_non_text_content(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + content = {"structured": "payload"} + assert guardrail._process_content(content) == (content, {}) + + +def test_process_tool_payload_handles_nested_lists_and_non_text_values(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + payload = ["plain", {"email": EMAIL}, [SSN], 42] + redacted, counts = guardrail._process_tool_payload(payload) + + assert EMAIL not in redacted[1]["email"] + assert SSN not in redacted[2][0] + assert redacted[3] == 42 + assert counts == {"EMAIL": 1, "SSN": 1} + assert guardrail._process_tool_payload(42) == (42, {}) + + +def test_process_tool_payload_skips_seen_containers(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + payload = {"self": None, "email": EMAIL} + payload["self"] = payload + list_payload = [EMAIL] + list_payload.append(list_payload) + + redacted, counts = guardrail._process_tool_payload(payload) + redacted_list, list_counts = guardrail._process_tool_payload(list_payload) + + assert EMAIL not in redacted["email"] + assert counts == {"EMAIL": 1} + assert EMAIL not in redacted_list[0] + assert list_counts == {"EMAIL": 1} + + +def test_scan_request_tool_calls_preserves_unsupported_entries(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + tool_calls = ["not-a-tool", {"id": "call_1", "type": "function"}] + assert guardrail._scan_request_tool_calls(tool_calls) == (tool_calls, {}) + + +def test_scan_messages_ignores_non_message_payloads(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + data = {"messages": "not-a-list"} + assert guardrail._scan_messages(data) == (data, {}) + + data = {"messages": ["not-a-dict", {"role": "system"}]} + assert guardrail._scan_messages(data) == (data, {}) + + +@pytest.mark.asyncio +async def test_block_raises_http_400_without_echoing_pii(): + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", default_on=True, datafog_action="block" + ) + with pytest.raises(HTTPException) as exc: + await guardrail.async_pre_call_hook( user_api_key_dict=None, cache=None, - data=_chat_data(f"email the report to {EMAIL} please"), + data=_chat_data(f"send {CARD} to billing"), call_type="completion", ) - content = data["messages"][0]["content"] - assert EMAIL not in content - assert "[EMAIL_1]" in content - - async def test_clean_message_unchanged(self): - guardrail = DataFogGuardrail(guardrail_name="datafog-pii") - original = _chat_data("summarize this design doc") - data = await guardrail.async_pre_call_hook( - user_api_key_dict=None, cache=None, data=original, call_type="completion" - ) - assert data["messages"][0]["content"] == "summarize this design doc" + assert exc.value.status_code == 400 + detail = str(exc.value.detail) + assert "CREDIT_CARD" in detail + assert CARD not in detail - async def test_redacts_content_parts_form(self): - guardrail = DataFogGuardrail(guardrail_name="datafog-pii") - data = await guardrail.async_pre_call_hook( + +@pytest.mark.asyncio +async def test_during_call_blocks_when_action_is_block(): + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", default_on=True, datafog_action="block" + ) + with pytest.raises(HTTPException): + await guardrail.async_moderation_hook( + data=_chat_data(f"reach me at {EMAIL}"), user_api_key_dict=None, - cache=None, - data=_chat_data([{"type": "text", "text": f"ssn is {SSN}"}]), call_type="completion", ) - part = data["messages"][0]["content"][0]["text"] - assert SSN not in part - async def test_redacts_multiple_messages_and_roles(self): - guardrail = DataFogGuardrail(guardrail_name="datafog-pii") - data = await guardrail.async_pre_call_hook( - user_api_key_dict=None, - cache=None, + +@pytest.mark.asyncio +async def test_during_call_blocks_tool_call_arguments_when_action_is_block(): + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", default_on=True, datafog_action="block" + ) + with pytest.raises(HTTPException): + await guardrail.async_moderation_hook( data={ "messages": [ - {"role": "system", "content": f"support contact: {EMAIL}"}, - {"role": "user", "content": f"card on file {CARD}"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "send_email", + "arguments": json.dumps({"recipient": EMAIL}), + }, + } + ], + } ] }, + user_api_key_dict=None, call_type="completion", ) - assert EMAIL not in data["messages"][0]["content"] - assert CARD not in data["messages"][1]["content"] @pytest.mark.asyncio -class TestPreCallBlock: - async def test_block_raises_http_400_without_echoing_pii(self): - # HTTPException(400) is what litellm's _is_guardrail_intervention - # recognizes as a policy block; a bare exception would surface as - # HTTP 500 and be misclassified as a backend failure. - from fastapi import HTTPException +async def test_during_call_blocks_tool_and_function_definitions_when_action_is_block(): + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", default_on=True, datafog_action="block" + ) - guardrail = DataFogGuardrail(guardrail_name="datafog-pii", action="block") - with pytest.raises(HTTPException) as exc: - await guardrail.async_pre_call_hook( - user_api_key_dict=None, - cache=None, - data=_chat_data(f"send {CARD} to billing"), - call_type="completion", - ) - assert exc.value.status_code == 400 - detail = str(exc.value.detail) - assert "CREDIT_CARD" in detail - assert CARD not in detail - - async def test_block_action_allows_clean_request(self): - guardrail = DataFogGuardrail(guardrail_name="datafog-pii", action="block") - data = await guardrail.async_pre_call_hook( + with pytest.raises(HTTPException): + await guardrail.async_moderation_hook( + data={ + "tools": [ + { + "type": "function", + "function": { + "name": "notify_customer", + "description": f"Notify {EMAIL}", + }, + } + ] + }, + user_api_key_dict=None, + call_type="completion", + ) + + with pytest.raises(HTTPException): + await guardrail.async_moderation_hook( + data={ + "functions": [ + { + "name": "legacy_lookup", + "parameters": { + "type": "object", + "description": f"Use card {CARD}", + }, + } + ] + }, user_api_key_dict=None, - cache=None, - data=_chat_data("hello"), call_type="completion", ) - assert data["messages"][0]["content"] == "hello" @pytest.mark.asyncio -class TestPostCall: - async def test_redacts_model_response(self): - guardrail = DataFogGuardrail(guardrail_name="datafog-pii") - response = _model_response(f"the customer is reachable at {EMAIL}") - await guardrail.async_post_call_success_hook( - data={}, user_api_key_dict=None, response=response +async def test_during_call_blocks_responses_api_input_when_action_is_block(): + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", default_on=True, datafog_action="block" + ) + with pytest.raises(HTTPException): + await guardrail.async_moderation_hook( + data={ + "input": [ + { + "role": "user", + "content": [{"type": "input_text", "text": f"email {EMAIL}"}], + } + ] + }, + user_api_key_dict=None, + call_type="responses", ) - assert EMAIL not in response.choices[0].message.content - async def test_response_without_choices_is_returned_untouched(self): - guardrail = DataFogGuardrail(guardrail_name="datafog-pii") - opaque = object() - result = await guardrail.async_post_call_success_hook( - data={}, user_api_key_dict=None, response=opaque - ) - assert result is opaque - async def test_non_text_response_content_is_skipped_not_crashed(self): - guardrail = DataFogGuardrail(guardrail_name="datafog-pii") - response = _model_response("placeholder") - response.choices[0].message.content = [{"type": "tool_use"}] - result = await guardrail.async_post_call_success_hook( - data={}, user_api_key_dict=None, response=response - ) - assert result.choices[0].message.content == [{"type": "tool_use"}] +@pytest.mark.asyncio +async def test_during_call_blocks_top_level_prompt_fields_when_action_is_block(): + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", default_on=True, datafog_action="block" + ) - async def test_post_call_fail_open_returns_unredacted_response(self, monkeypatch): - guardrail = DataFogGuardrail(guardrail_name="datafog-pii", fail_policy="open") - monkeypatch.setattr( - "datafog.integrations.litellm_guardrail._redact_text", - lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom")), + with pytest.raises(HTTPException): + await guardrail.async_moderation_hook( + data={"instructions": f"Contact billing at {EMAIL}"}, + user_api_key_dict=None, + call_type="responses", ) - response = _model_response(f"reach me at {EMAIL}") - result = await guardrail.async_post_call_success_hook( - data={}, user_api_key_dict=None, response=response + + with pytest.raises(HTTPException): + await guardrail.async_moderation_hook( + data={"system": [{"type": "text", "text": f"Use card {CARD} for lookup"}]}, + user_api_key_dict=None, + call_type="completion", ) - assert result.choices[0].message.content == f"reach me at {EMAIL}" @pytest.mark.asyncio -class TestEdgeShapes: - async def test_data_without_messages_passes_through(self): - guardrail = DataFogGuardrail(guardrail_name="datafog-pii") - data = {"input": f"embed {EMAIL}"} - result = await guardrail.async_pre_call_hook( - user_api_key_dict=None, cache=None, data=data, call_type="aembedding" - ) - assert result == data +async def test_during_call_blocks_text_completion_prompt_when_action_is_block(): + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", default_on=True, datafog_action="block" + ) - async def test_message_without_content_key_passes_through(self): - guardrail = DataFogGuardrail(guardrail_name="datafog-pii") - data = {"messages": [{"role": "assistant", "tool_calls": []}]} - result = await guardrail.async_pre_call_hook( - user_api_key_dict=None, cache=None, data=data, call_type="completion" + with pytest.raises(HTTPException): + await guardrail.async_moderation_hook( + data={"prompt": [f"Complete customer note for {EMAIL}"]}, + user_api_key_dict=None, + call_type="completion", ) - assert result == data - async def test_mixed_content_parts_skips_non_text_and_redacts_text(self): - guardrail = DataFogGuardrail(guardrail_name="datafog-pii") - data = _chat_data( - [ - {"type": "image_url", "image_url": {"url": "data:image/png;base64,xx"}}, - {"type": "text", "text": f"card {CARD}"}, + +@pytest.mark.asyncio +async def test_during_call_noop_when_action_is_redact(): + # during_call cannot modify content mid-flight; redact mode is a no-op. + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", default_on=True, datafog_action="redact" + ) + data = _chat_data(f"reach me at {EMAIL}") + result = await guardrail.async_moderation_hook( + data=data, user_api_key_dict=None, call_type="completion" + ) + assert result == data + + +@pytest.mark.asyncio +async def test_post_call_redacts_model_response(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + response = _model_response(f"the customer is reachable at {EMAIL}") + await guardrail.async_post_call_success_hook( + data={}, user_api_key_dict=None, response=response + ) + assert EMAIL not in response.choices[0].message.content + + +@pytest.mark.asyncio +async def test_post_call_redacts_text_completion_choice_text(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + response = SimpleNamespace( + choices=[SimpleNamespace(text=f"the customer email is {EMAIL}")] + ) + + await guardrail.async_post_call_success_hook( + data={}, user_api_key_dict=None, response=response + ) + + assert EMAIL not in response.choices[0].text + assert "[EMAIL_1]" in response.choices[0].text + + +@pytest.mark.asyncio +async def test_post_call_redacts_model_response_tool_and_function_arguments(): + import litellm + + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + response = litellm.ModelResponse( + choices=[ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "content": None, + "role": "assistant", + "function_call": { + "name": "legacy_lookup", + "arguments": json.dumps({"ssn": SSN}), + }, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "send_email", + "arguments": json.dumps({"recipient": EMAIL}), + }, + } + ], + }, + } + ] + ) + + await guardrail.async_post_call_success_hook( + data={}, user_api_key_dict=None, response=response + ) + + message = response.choices[0].message + assert message.content is None + assert SSN not in message.function_call.arguments + assert EMAIL not in message.tool_calls[0].function.arguments + + +@pytest.mark.asyncio +async def test_post_call_redacts_responses_api_output_text_and_arguments(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + response = SimpleNamespace( + output=[ + { + "type": "message", + "content": [{"type": "output_text", "text": f"contact {EMAIL}"}], + }, + { + "type": "function_call", + "arguments": json.dumps({"card": CARD}), + }, + ] + ) + + await guardrail.async_post_call_success_hook( + data={}, user_api_key_dict=None, response=response + ) + + assert EMAIL not in response.output[0]["content"][0]["text"] + assert CARD not in response.output[1]["arguments"] + + +@pytest.mark.asyncio +async def test_post_call_redacts_anthropic_messages_content_and_tool_use_input(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + response = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [ + {"type": "text", "text": f"email {EMAIL}"}, + { + "type": "tool_use", + "id": "toolu_1", + "name": "lookup_customer", + "input": {"ssn": SSN}, + }, + ], + } + + await guardrail.async_post_call_success_hook( + data={}, user_api_key_dict=None, response=response + ) + + assert EMAIL not in response["content"][0]["text"] + assert SSN not in response["content"][1]["input"]["ssn"] + + +@pytest.mark.asyncio +async def test_post_call_returns_original_when_event_hook_does_not_match(): + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", default_on=True, event_hook="pre_call" + ) + response = _model_response(f"the customer is reachable at {EMAIL}") + result = await guardrail.async_post_call_success_hook( + data={}, user_api_key_dict=None, response=response + ) + assert result is response + assert ( + response.choices[0].message.content == f"the customer is reachable at {EMAIL}" + ) + + +@pytest.mark.asyncio +async def test_post_call_returns_response_without_choices(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + response = SimpleNamespace(choices=[]) + result = await guardrail.async_post_call_success_hook( + data={}, user_api_key_dict=None, response=response + ) + assert result is response + + +@pytest.mark.asyncio +async def test_post_call_skips_non_text_response_parts(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + response = SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content=[{"type": "image"}]))] + ) + result = await guardrail.async_post_call_success_hook( + data={}, user_api_key_dict=None, response=response + ) + assert result is response + + +@pytest.mark.asyncio +async def test_post_call_fail_open_returns_response_on_engine_error(monkeypatch): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + monkeypatch.setattr( + "datafog.integrations.litellm_guardrail._redact_text", + lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom")), + ) + response = _model_response(f"the customer is reachable at {EMAIL}") + result = await guardrail.async_post_call_success_hook( + data={}, user_api_key_dict=None, response=response + ) + assert result is response + assert ( + response.choices[0].message.content == f"the customer is reachable at {EMAIL}" + ) + + +@pytest.mark.asyncio +async def test_streaming_redacts_model_response_chunks_before_yielding(): + from litellm.types.utils import ModelResponseStream + + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + chunks = [ + ModelResponseStream( + choices=[{"index": 0, "delta": {"content": "email jane.doe@"}}] + ), + ModelResponseStream( + choices=[ + { + "index": 0, + "delta": {"content": "example.com"}, + "finish_reason": "stop", + } ] + ), + ] + + result = await _collect_async( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, + response=_async_iter(chunks), + request_data={}, ) - result = await guardrail.async_pre_call_hook( - user_api_key_dict=None, cache=None, data=data, call_type="completion" - ) - parts = result["messages"][0]["content"] - assert parts[0]["type"] == "image_url" # untouched - assert CARD not in parts[1]["text"] - - async def test_non_string_non_list_content_passes_through(self): - guardrail = DataFogGuardrail(guardrail_name="datafog-pii") - data = _chat_data(None) - result = await guardrail.async_pre_call_hook( - user_api_key_dict=None, cache=None, data=data, call_type="completion" - ) - assert result["messages"][0]["content"] is None - - async def test_logging_helper_failure_never_breaks_traffic(self, monkeypatch): - guardrail = DataFogGuardrail(guardrail_name="datafog-pii") - monkeypatch.setattr( - DataFogGuardrail, - "add_standard_logging_guardrail_information_to_request_data", - lambda self, **kw: (_ for _ in ()).throw(RuntimeError("obs down")), - ) - data = await guardrail.async_pre_call_hook( + ) + + text = "".join(chunk.choices[0].delta.content or "" for chunk in result) + assert EMAIL not in text + assert "[EMAIL_1]" in text + + +@pytest.mark.asyncio +async def test_streaming_redacts_responses_api_output_text_delta_events(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + chunks = [ + SimpleNamespace( + type="response.output_text.delta", + item_id="msg_1", + output_index=0, + content_index=0, + delta="email jane.doe@", + ), + SimpleNamespace( + type="response.output_text.delta", + item_id="msg_1", + output_index=0, + content_index=0, + delta="example.com", + ), + ] + + result = await _collect_async( + guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=None, - cache=None, - data=_chat_data(f"reach me at {EMAIL}"), - call_type="completion", + response=_async_iter(chunks), + request_data={}, ) - assert EMAIL not in data["messages"][0]["content"] # redaction still happened + ) + + text = "".join(chunk.delta for chunk in result) + assert EMAIL not in text + assert "[EMAIL_1]" in text @pytest.mark.asyncio -class TestConfig: - async def test_noisy_entities_off_by_default(self): - guardrail = DataFogGuardrail(guardrail_name="datafog-pii") - data = await guardrail.async_pre_call_hook( +async def test_streaming_redacts_anthropic_sse_text_delta_bytes(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + chunks = [ + b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,' + b'"delta":{"type":"text_delta","text":"email jane.doe@"}}\n\n', + b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,' + b'"delta":{"type":"text_delta","text":"example.com"}}\n\n', + ] + + result = await _collect_async( + guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=None, - cache=None, - data=_chat_data("ping 192.168.1.1 about build 2020-01-02"), - call_type="completion", + response=_async_iter(chunks), + request_data={}, ) - assert ( - data["messages"][0]["content"] == "ping 192.168.1.1 about build 2020-01-02" + ) + + raw_text = b"".join(result).decode("utf-8") + assert EMAIL not in raw_text + assert "[EMAIL_1]" in raw_text + + +@pytest.mark.asyncio +async def test_streaming_redacts_text_completion_choice_text_before_yielding(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + chunks = [ + SimpleNamespace(choices=[SimpleNamespace(index=0, text="email jane.doe@")]), + SimpleNamespace(choices=[SimpleNamespace(index=0, text="example.com")]), + ] + + result = await _collect_async( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, + response=_async_iter(chunks), + request_data={}, ) + ) + + text = "".join(chunk.choices[0].text for chunk in result) + assert EMAIL not in text + assert "[EMAIL_1]" in text + + +@pytest.mark.asyncio +async def test_streaming_block_raises_before_mutating_chunks(): + from litellm.types.utils import ModelResponseStream + + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", default_on=True, datafog_action="block" + ) + chunks = [ + ModelResponseStream( + choices=[{"index": 0, "delta": {"content": "email jane.doe@"}}] + ), + ModelResponseStream( + choices=[ + { + "index": 0, + "delta": {"content": "example.com"}, + "finish_reason": "stop", + } + ] + ), + ] - async def test_entity_types_override(self): - guardrail = DataFogGuardrail( - guardrail_name="datafog-pii", entity_types=["IP_ADDRESS"] + with pytest.raises(HTTPException) as exc: + await _collect_async( + guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=None, + response=_async_iter(chunks), + request_data={}, + ) ) - data = await guardrail.async_pre_call_hook( + + assert exc.value.status_code == 400 + assert EMAIL not in str(exc.value.detail) + assert chunks[0].choices[0].delta.content == "email jane.doe@" + assert chunks[1].choices[0].delta.content == "example.com" + + +@pytest.mark.asyncio +async def test_noisy_entities_off_by_default(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=_chat_data("ping 192.168.1.1 about build 2020-01-02"), + call_type="completion", + ) + assert data["messages"][0]["content"] == "ping 192.168.1.1 about build 2020-01-02" + + +@pytest.mark.asyncio +async def test_entity_types_override(): + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", + default_on=True, + datafog_entity_types=["IP_ADDRESS"], + ) + data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=_chat_data("ping 192.168.1.1"), + call_type="completion", + ) + assert "192.168.1.1" not in data["messages"][0]["content"] + + +@pytest.mark.asyncio +async def test_german_locale_entities(): + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", + default_on=True, + datafog_entity_types=["DE_TAX_ID"], + datafog_locales=["de"], + ) + data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=_chat_data(f"Steuer-ID {DE_TAX_ID} liegt vor."), + call_type="completion", + ) + assert DE_TAX_ID not in data["messages"][0]["content"] + + +@pytest.mark.asyncio +async def test_fail_open_passes_data_through_on_engine_error(monkeypatch): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + monkeypatch.setattr( + "datafog.integrations.litellm_guardrail._redact_text", + lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom")), + ) + original = _chat_data(f"reach me at {EMAIL}") + data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, cache=None, data=original, call_type="completion" + ) + assert data["messages"][0]["content"] == f"reach me at {EMAIL}" + + +@pytest.mark.asyncio +async def test_fail_closed_raises_without_pii_or_cause_chain(monkeypatch): + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", default_on=True, datafog_fail_policy="closed" + ) + monkeypatch.setattr( + "datafog.integrations.litellm_guardrail._redact_text", + lambda *a, **k: (_ for _ in ()).throw( + RuntimeError(f"parser choked on: reach me at {EMAIL}") + ), + ) + with pytest.raises(RuntimeError, match="fail_policy") as exc: + await guardrail.async_pre_call_hook( user_api_key_dict=None, cache=None, - data=_chat_data("ping 192.168.1.1"), + data=_chat_data(f"reach me at {EMAIL}"), call_type="completion", ) - assert "192.168.1.1" not in data["messages"][0]["content"] + assert exc.value.__cause__ is None + assert EMAIL not in str(exc.value) @pytest.mark.asyncio -class TestFailPolicy: - async def test_fail_open_passes_data_through_on_engine_error(self, monkeypatch): - guardrail = DataFogGuardrail(guardrail_name="datafog-pii", fail_policy="open") - monkeypatch.setattr( - "datafog.integrations.litellm_guardrail._redact_text", - lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom")), - ) - original = _chat_data(f"reach me at {EMAIL}") - data = await guardrail.async_pre_call_hook( - user_api_key_dict=None, cache=None, data=original, call_type="completion" +async def test_redact_logging_metadata_survives_on_returned_data(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii", default_on=True) + data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=_chat_data(f"email {EMAIL}"), + call_type="completion", + ) + assert "metadata" in data + + +@pytest.mark.asyncio +async def test_post_call_block_raises_on_response_pii(): + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", default_on=True, datafog_action="block" + ) + response = _model_response(f"the customer is reachable at {EMAIL}") + with pytest.raises(HTTPException) as exc: + await guardrail.async_post_call_success_hook( + data={}, user_api_key_dict=None, response=response ) - assert data["messages"][0]["content"] == f"reach me at {EMAIL}" + assert exc.value.status_code == 400 + assert EMAIL not in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_post_call_block_raises_on_response_tool_call_arguments(): + import litellm + + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", default_on=True, datafog_action="block" + ) + response = litellm.ModelResponse( + choices=[ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "content": None, + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "send_email", + "arguments": json.dumps({"recipient": EMAIL}), + }, + } + ], + }, + } + ] + ) - async def test_fail_closed_raises_on_engine_error(self, monkeypatch): - guardrail = DataFogGuardrail(guardrail_name="datafog-pii", fail_policy="closed") - monkeypatch.setattr( - "datafog.integrations.litellm_guardrail._redact_text", - lambda *a, **k: (_ for _ in ()).throw(RuntimeError("boom")), + with pytest.raises(HTTPException) as exc: + await guardrail.async_post_call_success_hook( + data={}, user_api_key_dict=None, response=response ) - with pytest.raises(RuntimeError, match="fail_policy is 'closed'"): - await guardrail.async_pre_call_hook( - user_api_key_dict=None, - cache=None, - data=_chat_data(f"reach me at {EMAIL}"), - call_type="completion", - ) - async def test_allowlist_exempts_configured_values(self): - own = "sid@" "example.com" - other = "jane.doe@" "example.com" - guardrail = DataFogGuardrail(guardrail_name="datafog-pii", allowlist=[own]) - data = await guardrail.async_pre_call_hook( - user_api_key_dict=None, - cache=None, - data=_chat_data(f"contact {own} or {other}"), - call_type="completion", + assert exc.value.status_code == 400 + assert EMAIL not in str(exc.value.detail) + assert response.choices[0].message.tool_calls[0].function.arguments == json.dumps( + {"recipient": EMAIL} + ) + + +@pytest.mark.asyncio +async def test_post_call_block_raises_on_text_completion_choice_text(): + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", default_on=True, datafog_action="block" + ) + response = SimpleNamespace( + choices=[SimpleNamespace(text=f"the customer email is {EMAIL}")] + ) + + with pytest.raises(HTTPException) as exc: + await guardrail.async_post_call_success_hook( + data={}, user_api_key_dict=None, response=response ) - content = data["messages"][0]["content"] - assert own in content - assert other not in content - - async def test_invalid_config_rejected(self): - with pytest.raises(ValueError): - DataFogGuardrail(guardrail_name="datafog-pii", action="explode") - with pytest.raises(ValueError): - DataFogGuardrail(guardrail_name="datafog-pii", fail_policy="maybe") - - async def test_fail_closed_error_carries_no_pii_and_no_cause_chain( - self, monkeypatch - ): - # Engine exceptions can embed the text being scanned. The re-raise - # must not chain them (`from None`): a chained __cause__ is printed - # by traceback.format_exc(), which litellm calls for logging. - guardrail = DataFogGuardrail(guardrail_name="datafog-pii", fail_policy="closed") - monkeypatch.setattr( - "datafog.integrations.litellm_guardrail._redact_text", - lambda *a, **k: (_ for _ in ()).throw( - RuntimeError(f"parser choked on: reach me at {EMAIL}") - ), + + assert exc.value.status_code == 400 + assert EMAIL not in str(exc.value.detail) + assert response.choices[0].text == f"the customer email is {EMAIL}" + + +def test_invalid_config_rejected(): + with pytest.raises(ValueError): + DataFogGuardrail( + guardrail_name="datafog-pii", default_on=True, datafog_action="explode" ) - with pytest.raises(RuntimeError) as exc: - await guardrail.async_pre_call_hook( - user_api_key_dict=None, - cache=None, - data=_chat_data(f"reach me at {EMAIL}"), - call_type="completion", - ) - assert exc.value.__cause__ is None - assert exc.value.__suppress_context__ is True - assert EMAIL not in str(exc.value) - assert not hasattr(exc.value, "status_code") # engine fault -> 500, by design - - async def test_fail_open_log_carries_no_pii(self, monkeypatch, caplog): - import logging - - guardrail = DataFogGuardrail(guardrail_name="datafog-pii", fail_policy="open") - monkeypatch.setattr( - "datafog.integrations.litellm_guardrail._redact_text", - lambda *a, **k: (_ for _ in ()).throw( - RuntimeError(f"parser choked on: reach me at {EMAIL}") - ), + with pytest.raises(ValueError): + DataFogGuardrail( + guardrail_name="datafog-pii", default_on=True, datafog_fail_policy="maybe" ) - with caplog.at_level(logging.WARNING): - await guardrail.async_pre_call_hook( - user_api_key_dict=None, - cache=None, - data=_chat_data(f"reach me at {EMAIL}"), - call_type="completion", - ) - assert EMAIL not in caplog.text - assert "RuntimeError" in caplog.text + + +@pytest.mark.asyncio +async def test_legacy_config_names_and_allowlists_remain_supported(): + own_email = "sid@" "example.com" + other_email = "jane.doe@" "example.com" + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", + action="redact", + entity_types=["EMAIL"], + fail_policy="open", + allowlist=[own_email], + ) + + data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=_chat_data(f"contact {own_email} or {other_email}"), + call_type="completion", + ) + + assert own_email in data["messages"][0]["content"] + assert other_email not in data["messages"][0]["content"] + + +@pytest.mark.asyncio +async def test_prefixed_allowlist_patterns_remain_supported(): + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", + datafog_entity_types=["PHONE"], + datafog_allowlist_patterns=[r"^555-0100$"], + ) + + data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=_chat_data("call 555-0100"), + call_type="completion", + ) + + assert data["messages"][0]["content"] == "call 555-0100" + + +@pytest.mark.asyncio +async def test_direct_adapter_usage_runs_by_default(): + guardrail = DataFogGuardrail(guardrail_name="datafog-pii") + + data = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data={"input": f"embed {EMAIL}"}, + call_type="aembedding", + ) + + assert EMAIL not in data["input"] + + +@pytest.mark.asyncio +async def test_explicit_default_off_requires_request_selection(): + guardrail = DataFogGuardrail( + guardrail_name="datafog-pii", + default_on=False, + event_hook="pre_call", + ) + original = _chat_data(f"email {EMAIL}") + + skipped = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data=original, + call_type="completion", + ) + selected = await guardrail.async_pre_call_hook( + user_api_key_dict=None, + cache=None, + data={**original, "guardrails": ["datafog-pii"]}, + call_type="completion", + ) + + assert skipped is original + assert EMAIL in skipped["messages"][0]["content"] + assert EMAIL not in selected["messages"][0]["content"] diff --git a/tests/test_no_network_core.py b/tests/test_no_network_core.py index f06e8360..d35936b2 100644 --- a/tests/test_no_network_core.py +++ b/tests/test_no_network_core.py @@ -39,6 +39,9 @@ def blocked(*_args, **_kwargs): redact_result = datafog.redact("Email jane@example.com or call 415-555-1212") assert "jane@example.com" not in redact_result.redacted_text assert "415-555-1212" not in redact_result.redacted_text + +batch_result = datafog.redact_many(["Email jane@example.com", "Call 415-555-1212"]) +assert batch_result.redacted_texts == ["Email [EMAIL_1]", "Call [PHONE_1]"] """ ) @@ -59,6 +62,11 @@ def fail_optional_engine_probe(): redact_result = datafog.redact("Email jane@example.com") assert redact_result.redacted_text == "Email [EMAIL_1]" + batch_result = datafog.redact_many( + ["Email jane@example.com", "Email support@example.com"] + ) + assert batch_result.redacted_texts == ["Email [EMAIL_1]", "Email [EMAIL_2]"] + guardrail = datafog.protect() guarded = guardrail.filter("Email jane@example.com") assert guarded.redacted_text == "Email [EMAIL_1]" @@ -139,6 +147,7 @@ def find_spec(self, fullname, path=None, target=None): assert datafog.scan("Email jane@example.com").entities assert datafog.redact("Email jane@example.com").redacted_text == "Email [EMAIL_1]" +assert datafog.redact_many(["Email jane@example.com"]).redacted_texts == ["Email [EMAIL_1]"] assert datafog.protect().filter("Email jane@example.com").redacted_text == "Email [EMAIL_1]" assert datafog.sanitize("Email jane@example.com") == "Email [EMAIL_1]" assert datafog.scan_prompt("Email jane@example.com").entities diff --git a/tests/test_v44_bridge_api.py b/tests/test_v44_bridge_api.py index a26099a9..6158b140 100644 --- a/tests/test_v44_bridge_api.py +++ b/tests/test_v44_bridge_api.py @@ -47,6 +47,16 @@ def test_top_level_redact_supports_preview_presets() -> None: assert "[EMAIL_1]" in result.redacted_text +def test_top_level_redact_many_uses_batch_stable_tokens() -> None: + result = datafog.redact_many( + ["Email alpha@example.com", "Email beta@example.com"], + engine="regex", + ) + + assert result.redacted_texts == ["Email [EMAIL_1]", "Email [EMAIL_2]"] + assert result.entity_counts == {"EMAIL": 2} + + def test_top_level_protect_returns_guardrail() -> None: guardrail = datafog.protect(on_detect="redact")