diff --git a/docs/documentation/privacy-guard/architecture/entity-processing-engines.md b/docs/documentation/privacy-guard/architecture/entity-processing-engines.md index d92d8c9..da4cb82 100644 --- a/docs/documentation/privacy-guard/architecture/entity-processing-engines.md +++ b/docs/documentation/privacy-guard/architecture/entity-processing-engines.md @@ -294,15 +294,64 @@ Important properties: The third-party `regex` backend is used because it can interrupt an individual search when the timeout expires. +## NEREngine + +`NEREngine` is a resource-backed built-in for general named-entity recognition. +Its policy configuration owns: + +- an ordered non-empty list of printable labels +- a required finite threshold from zero through one +- `nested` or `flat` overlap behavior +- an optional template replacement allowing only literal text and `{entity}` + +The deployment registers one immutable `NERResources` bundle containing a +provider-neutral `NERModel`. The facade has one synchronous +`predict_entities` operation. Endpoint URLs, model identifiers, chunk +configuration, device placement, and loaded model objects never enter policy. +The default built-in registry remains Regex-only; it registers NER only when +the operator explicitly supplies `ner_resources`. + +The package provides two focused facades: + +- `NERExtractEndpointModel` sends the exact `POST /v1/extract` JSON contract, + bounds response bytes before decoding, performs no retries, and translates + expected transport or schema failures without exposing content, endpoint + details, or raw collaborator errors. +- `LocalNERModel` receives an already-loaded object without importing GLiNER, + serializes access, processes long text in overlapping code-point windows, + rebases offsets, and removes exact chunk-overlap duplicates. + +The local facade's lock acquisition is bounded by the shared timeout. Once a +loaded model call starts, Python's worker-thread architecture cannot preempt +it; expiration is observed after that call returns. + +The engine accepts only returned labels declared by policy, canonicalizes +case-only differences, rejects invalid spans and scores, retains the +highest-scoring exact duplicate, and sorts results deterministically. Raw +scores remain private to duplicate and replacement selection and do not become +finding confidence or metadata. + +Nested detection may return overlapping findings. Replacement keeps every +finding but selects non-overlapping winners by score, span length, earlier +start, and configured label order. It projects UTF-8 output size before +allocating the replacement. + +The extract endpoint contract is not the self-hosted Anonymizer GLiNER +`/v1/chat/completions` contract. Supporting that deployment requires a future +explicit facade with its own response validation and must not rely on protocol +autodetection. A future full Anonymizer engine remains separately configured +because it may own validation, augmentation, or rewriting beyond general NER. + ## Tool-specific custom engines Tool integrations belong in custom engines until they have a complete, production-backed implementation. Privacy Guard does not ship placeholder engine types or runtime protocols that merely resemble a third-party tool. -The first NeMo Anonymizer integration will be a custom engine. Its configuration -and replacement types should preserve Anonymizer's native concepts, while the -engine itself owns all translation to and from the actual Anonymizer SDK. +A future full NeMo Anonymizer integration will be a custom engine. Its +configuration and replacement types should preserve Anonymizer's native +concepts, while the engine itself owns translation to and from the actual +Anonymizer SDK. ## State diff --git a/docs/documentation/privacy-guard/architecture/index.md b/docs/documentation/privacy-guard/architecture/index.md index 2e8d932..322d367 100644 --- a/docs/documentation/privacy-guard/architecture/index.md +++ b/docs/documentation/privacy-guard/architecture/index.md @@ -64,8 +64,8 @@ Source paths on these pages are relative to policy action. It does not import gRPC or implement an engine's algorithms. - `engines/` defines the custom-engine contract, registers engine implementations and operator-owned resources, builds the exact Pydantic - discriminated union, and contains the built-in Regex implementation. Each - engine owns its detection and replacement algorithms. + discriminated union, and contains the built-in Regex and resource-backed NER + implementations. Each engine owns its detection and replacement algorithms. - `config.py` defines ordered stages and the required policy action. - top-level `base.py` defines the package-wide strict immutable domain-model base. diff --git a/docs/documentation/privacy-guard/architecture/safety-and-limits.md b/docs/documentation/privacy-guard/architecture/safety-and-limits.md index 3b22fec..8c70949 100644 --- a/docs/documentation/privacy-guard/architecture/safety-and-limits.md +++ b/docs/documentation/privacy-guard/architecture/safety-and-limits.md @@ -44,6 +44,7 @@ reimplement either. | Maximum configured timeout | `MAX_TIMEOUT_SECONDS` | 30 seconds | Server, processor, and `Timeout` | | Active processor workers | `MAX_CONCURRENT_PROCESSING` | 4 | Service | | Concurrent gRPC calls | `MAX_CONCURRENT_RPCS` | 16 | gRPC server | +| NER endpoint response | `MAX_NER_ENDPOINT_RESPONSE_BYTES` | 1 MiB | NER endpoint facade | The gRPC receive allowance is 1 MiB larger than the request body limit for the protobuf envelope. The service still enforces the advertised 4 MiB body bound. @@ -100,6 +101,32 @@ The per-evaluation config `Struct` is limited to 64 KiB. A catalog may satisfy Privacy Guard's engine limits yet remain too large for the current upstream protocol. Internal caching does not raise that transport ceiling. +## NER configuration and execution limits + +| Limit | Constant | Value | +| --- | --- | ---: | +| Labels per NER stage | `MAX_NER_LABELS` | 128 | +| UTF-8 bytes per label | `MAX_NER_LABEL_BYTES` | 256 | +| Total UTF-8 label bytes | `MAX_NER_LABELS_BYTES` | 16 KiB | +| Endpoint response bytes | `MAX_NER_ENDPOINT_RESPONSE_BYTES` | 1 MiB | + +NER labels must be non-empty printable Unicode without leading or trailing +whitespace and must be unique after Unicode case folding. Model identity, +endpoint, and execution chunking are deployment resources rather than policy +fields. + +The endpoint facade bounds response bytes before JSON decoding and performs one +request without retry or fallback. The local facade processes every supported +input in overlapping code-point windows and rebases returned offsets. Its +overlap must be smaller than its chunk length. Operators should choose overlap +using the longest entity spans expected in the deployment corpus. + +Both facades return only normalized label, start, end, and score values. The +engine rejects unknown labels, invalid or out-of-bounds spans, non-finite or +out-of-range scores, and excessive result cardinality. Request text, returned +matched text, raw response bodies, endpoint URLs, credentials, and collaborator +exception messages are not logged or included in errors. + ## Regex execution safety `RegexEngine` validates and compiles the complete catalog before accepting diff --git a/docs/documentation/privacy-guard/architecture/service-boundary.md b/docs/documentation/privacy-guard/architecture/service-boundary.md index 01ca64f..961ab22 100644 --- a/docs/documentation/privacy-guard/architecture/service-boundary.md +++ b/docs/documentation/privacy-guard/architecture/service-boundary.md @@ -216,9 +216,10 @@ EngineRegistry -> gRPC server ``` -The built-in registry includes `RegexEngine`. Operators register custom engines -and resource-backed tool integrations before registry finalization, then pass -that registry to `PrivacyGuardServer`. +The built-in registry always includes `RegexEngine` and includes `NEREngine` +only when the operator supplies explicit NER resources. Operators register +custom engines and other resource-backed tool integrations before registry +finalization, then pass that registry to `PrivacyGuardServer`. The registry is an explicit application-scoped dependency, not a global singleton. `PrivacyGuardServer` and `PrivacyGuardMiddleware` reject unfinalized diff --git a/projects/privacy-guard/README.md b/projects/privacy-guard/README.md index efb98dd..f2c6e41 100644 --- a/projects/privacy-guard/README.md +++ b/projects/privacy-guard/README.md @@ -103,6 +103,48 @@ contract. Privacy Guard owns the catalog schema but maintains no authoritative patterns. +### NEREngine + +`NEREngine` is a resource-backed general named-entity recognizer. Policy +declares an ordered non-empty label list, a required score threshold, nested or +flat overlap behavior, and an optional constrained replacement template. +Deployment supplies one provider-neutral model facade through `NERResources`. +The package includes facades for the explicit GLiNER-compatible +`POST /v1/extract` contract and for an already-loaded local model; neither +model identity nor endpoint appears in policy. + +The built-in registry remains Regex-only unless NER resources are supplied: + +```python +from privacy_guard.engines import ( + NERExtractEndpointModel, + NERResources, +) +from privacy_guard.engines.registry import create_builtin_registry + +registry = create_builtin_registry( + ner_resources=NERResources( + model=NERExtractEndpointModel( + endpoint="http://model-host:8002/v1/extract", + model="operator-selected-model", + ) + ) +) +``` + +NER model scores remain internal and are not mapped to Privacy Guard's +categorical confidence because model scores are not universally calibrated. +Case-only returned-label differences are canonicalized to the configured +label. Exact duplicate spans retain the highest score. Replacement keeps every +finding while choosing non-overlapping winners by score, span length, start, +and configured label order. + +See the +[NER engine end-to-end example](https://github.com/NVIDIA/OpenShell-Research/blob/main/projects/privacy-guard/examples/ner-engine/README.md) +for endpoint and local registry factories. The self-hosted Anonymizer GLiNER +chat-completions contract is a future dedicated facade rather than an +auto-detected variant of the extract endpoint. + ## Custom engines Custom engines are a first-class extension point. Authors declare one typed @@ -120,8 +162,8 @@ must contain no policy behavior or per-request state and must be safe for concurrent use. Resource-free engines omit the second `EntityProcessingEngine` generic argument entirely. -The first NeMo Anonymizer integration will be implemented as a custom engine, -not as a built-in or placeholder abstraction in Privacy Guard. +A future full NeMo Anonymizer integration remains a separate custom engine, +not an expansion of the general built-in NER facade. Application startup registers engines and operator-owned resources, then returns one finalized registry: diff --git a/projects/privacy-guard/examples/ner-engine/README.md b/projects/privacy-guard/examples/ner-engine/README.md new file mode 100644 index 0000000..919f788 --- /dev/null +++ b/projects/privacy-guard/examples/ner-engine/README.md @@ -0,0 +1,98 @@ +# NER engine example + +This example registers Privacy Guard's general `ner` engine with either the +explicit `POST /v1/extract` endpoint contract or an already-downloaded local +GLiNER model. Policy owns labels, threshold, overlap behavior, replacement, +and the final action. The registry factory owns the model and its execution +location. + +The initial `nvidia/gliner-PII` checkpoint is optimized for English PII/PHI. +It is neither a required model nor a built-in entity catalog. Validate the +chosen model, labels, threshold, and representative inputs before production. + +## Endpoint-backed registry + +Set the exact endpoint and model identifier. Request text leaves the Privacy +Guard process, so use only an operator-approved endpoint and protected network. +The endpoint URL must be an `http` or `https` URL whose path is exactly +`/v1/extract`; embedded credentials, query strings, and fragments are rejected. + +```bash +cd projects/privacy-guard +export PYTHONPATH="$PWD/examples/ner-engine" +export PRIVACY_GUARD_NER_ENDPOINT="http://spark-9107.local:8002/v1/extract" +export PRIVACY_GUARD_NER_MODEL="nvidia/gliner-PII" + +uv run privacy-guard \ + --registry-factory endpoint_registry:create_registry engines +uv run privacy-guard \ + --registry-factory endpoint_registry:create_registry configuration-schema +uv run privacy-guard \ + --registry-factory endpoint_registry:create_registry configure-gateway \ + --host-ip YOUR_HOST_IPV4 +uv run privacy-guard \ + --registry-factory endpoint_registry:create_registry serve \ + --listen 0.0.0.0:50051 \ + --timeout-seconds 30 +``` + +The endpoint receives `text`, `labels`, `model`, `threshold`, `chunk_length`, +`overlap`, and `flat_ner`. Privacy Guard performs no retry, fallback, health +check, or protocol detection. + +## Direct local registry + +The direct example deliberately keeps GLiNER and PyTorch out of Privacy Guard's +base dependencies. Install the characterized library version in the deployment +environment and download the model through an operator-controlled workflow: + +```bash +uv pip install "gliner==0.2.27" +export PYTHONPATH="$PWD/examples/ner-engine" +export PRIVACY_GUARD_NER_MODEL_PATH="/absolute/path/to/downloaded/model" + +uv run privacy-guard \ + --registry-factory local_registry:create_registry engines +uv run privacy-guard \ + --registry-factory local_registry:create_registry serve \ + --listen 0.0.0.0:50051 \ + --timeout-seconds 30 +``` + +`local_files_only=True` prevents an implicit model download. Local calls are +serialized because loaded-model thread safety and GPU memory behavior are not +assumed. Long input is processed in overlapping Unicode code-point windows, +with global offset rebasing and exact duplicate removal. A running model call +cannot be preempted by Privacy Guard's worker thread; timeout is checked after +the call returns. + +`PRIVACY_GUARD_NER_CHUNK_LENGTH` and +`PRIVACY_GUARD_NER_CHUNK_OVERLAP` optionally override the defaults of 384 and +128. The endpoint interprets these according to its contract. The local facade +uses them as code-point window sizes and requires overlap to be smaller than +chunk length. + +## Exercise policy actions + +Use [`policy.yaml`](policy.yaml) as the OpenShell policy source. The extracted +[`privacy-guard-config.yaml`](privacy-guard-config.yaml) is convenient for +direct middleware tests. Its configured replacement may stay present for every +action: + +- `detect` allows the original request and reports findings. +- `block` denies a request when a configured entity is found. +- `replace` sends the body after deterministic, non-overlapping template + replacement. + +Change only `on_detection.action` to exercise these modes. Nested detection +returns all valid model spans. Replacement keeps all findings but selects +non-overlapping winners by score, span length, start offset, and configured +label order. + +Scores are not exposed as categorical confidence because model scores are not +universally calibrated. Tune the required `threshold` using use-case-specific +evaluation. + +The Anonymizer self-hosted GLiNER `/v1/chat/completions` contract is not +accepted by this example. It remains a dedicated future adapter so Privacy +Guard never guesses a provider protocol from a URL. diff --git a/projects/privacy-guard/examples/ner-engine/endpoint_registry.py b/projects/privacy-guard/examples/ner-engine/endpoint_registry.py new file mode 100644 index 0000000..e87de60 --- /dev/null +++ b/projects/privacy-guard/examples/ner-engine/endpoint_registry.py @@ -0,0 +1,49 @@ +"""Registry factory for the explicit ``POST /v1/extract`` NER endpoint.""" + +from __future__ import annotations + +import os + +from privacy_guard.engines import NERExtractEndpointModel, NERResources +from privacy_guard.engines.registry import EngineRegistry + + +def create_registry() -> EngineRegistry: + """Create a built-in registry backed by the configured NER endpoint.""" + endpoint = _required_environment("PRIVACY_GUARD_NER_ENDPOINT") + model = _required_environment("PRIVACY_GUARD_NER_MODEL") + resources = NERResources( + model=NERExtractEndpointModel( + endpoint=endpoint, + model=model, + chunk_length=_environment_integer( + "PRIVACY_GUARD_NER_CHUNK_LENGTH", + default=384, + ), + overlap=_environment_integer( + "PRIVACY_GUARD_NER_CHUNK_OVERLAP", + default=128, + ), + ) + ) + return EngineRegistry( + include_builtin_engines=True, + ner_resources=resources, + ).finalize() + + +def _required_environment(name: str) -> str: + value = os.environ.get(name) + if value is None or not value: + raise RuntimeError(f"Required environment variable {name} is not set") + return value + + +def _environment_integer(name: str, *, default: int) -> int: + value = os.environ.get(name) + if value is None: + return default + try: + return int(value) + except ValueError: + raise RuntimeError(f"Environment variable {name} must be an integer") from None diff --git a/projects/privacy-guard/examples/ner-engine/local_registry.py b/projects/privacy-guard/examples/ner-engine/local_registry.py new file mode 100644 index 0000000..fcfb81b --- /dev/null +++ b/projects/privacy-guard/examples/ner-engine/local_registry.py @@ -0,0 +1,50 @@ +"""Registry factory for an already-downloaded local GLiNER model.""" + +from __future__ import annotations + +import os +from importlib import import_module + +from privacy_guard.engines import LocalNERModel, NERResources +from privacy_guard.engines.registry import EngineRegistry + + +def create_registry() -> EngineRegistry: + """Load one local model at startup and create the built-in NER registry.""" + model_path = _required_environment("PRIVACY_GUARD_NER_MODEL_PATH") + gliner_type = getattr(import_module("gliner"), "GLiNER") + loaded_model = gliner_type.from_pretrained(model_path, local_files_only=True) + resources = NERResources( + model=LocalNERModel( + model=loaded_model, + chunk_length=_environment_integer( + "PRIVACY_GUARD_NER_CHUNK_LENGTH", + default=384, + ), + overlap=_environment_integer( + "PRIVACY_GUARD_NER_CHUNK_OVERLAP", + default=128, + ), + ) + ) + return EngineRegistry( + include_builtin_engines=True, + ner_resources=resources, + ).finalize() + + +def _required_environment(name: str) -> str: + value = os.environ.get(name) + if value is None or not value: + raise RuntimeError(f"Required environment variable {name} is not set") + return value + + +def _environment_integer(name: str, *, default: int) -> int: + value = os.environ.get(name) + if value is None: + return default + try: + return int(value) + except ValueError: + raise RuntimeError(f"Environment variable {name} must be an integer") from None diff --git a/projects/privacy-guard/examples/ner-engine/policy.yaml b/projects/privacy-guard/examples/ner-engine/policy.yaml new file mode 100644 index 0000000..c0eba9c --- /dev/null +++ b/projects/privacy-guard/examples/ner-engine/policy.yaml @@ -0,0 +1,60 @@ +version: 1 + +filesystem_policy: + include_workdir: true + read_only: [/usr, /lib, /proc, /dev/urandom, /app, /etc, /var/log] + read_write: [/sandbox, /tmp, /dev/null] +landlock: + compatibility: best_effort +process: + run_as_user: sandbox + run_as_group: sandbox + +network_policies: + claude_code: + name: Claude Code subscription access + endpoints: + - host: api.anthropic.com + port: 443 + protocol: rest + enforcement: enforce + access: full + - host: platform.claude.com + port: 443 + protocol: rest + enforcement: enforce + access: full + - host: claude.ai + port: 443 + - { host: statsig.anthropic.com, port: 443 } + - { host: sentry.io, port: 443 } + binaries: + - { path: /usr/local/bin/claude } + - { path: /usr/bin/node } + +network_middlewares: + privacy_guard_ner: + name: Detect named entities + middleware: privacy-guard-ner + order: 0 + config: + entity_processing: + stages: + - name: general-entities + config: + engine: ner + labels: + - person + - email + - phone_number + threshold: 0.5 + overlap_mode: nested + replacement: + strategy: template + template: "[{entity}]" + on_detection: + action: detect + on_error: fail_closed + endpoints: + include: + - api.anthropic.com diff --git a/projects/privacy-guard/examples/ner-engine/privacy-guard-config.yaml b/projects/privacy-guard/examples/ner-engine/privacy-guard-config.yaml new file mode 100644 index 0000000..359ba3b --- /dev/null +++ b/projects/privacy-guard/examples/ner-engine/privacy-guard-config.yaml @@ -0,0 +1,16 @@ +entity_processing: + stages: + - name: general-entities + config: + engine: ner + labels: + - person + - email + - phone_number + threshold: 0.5 + overlap_mode: nested + replacement: + strategy: template + template: "[{entity}]" +on_detection: + action: detect diff --git a/projects/privacy-guard/src/privacy_guard/constants.py b/projects/privacy-guard/src/privacy_guard/constants.py index 4fbda0c..de29e58 100644 --- a/projects/privacy-guard/src/privacy_guard/constants.py +++ b/projects/privacy-guard/src/privacy_guard/constants.py @@ -42,6 +42,10 @@ # Engine configuration and regex execution limits. MAX_ENTITY_PROCESSING_STAGES = 10 +MAX_NER_LABELS = 128 +MAX_NER_LABEL_BYTES = 256 +MAX_NER_LABELS_BYTES = 16 * 1024 +MAX_NER_ENDPOINT_RESPONSE_BYTES = 1024 * 1024 MAX_REGEX_NAME_BYTES = 128 MAX_REGEX_ENTITIES_PER_CATALOG = 2_000 MAX_REGEX_RULES_PER_CATALOG = 10_000 diff --git a/projects/privacy-guard/src/privacy_guard/engines/__init__.py b/projects/privacy-guard/src/privacy_guard/engines/__init__.py index 59a843d..b060fa0 100644 --- a/projects/privacy-guard/src/privacy_guard/engines/__init__.py +++ b/projects/privacy-guard/src/privacy_guard/engines/__init__.py @@ -13,6 +13,20 @@ EntityProcessingStrategy, TextProcessingResult, ) +from privacy_guard.engines.ner import ( + NEREngine, + NEREngineConfig, + NEROverlapMode, + NERReplacement, + NERResources, +) +from privacy_guard.engines.ner_model import ( + LocalNERModel, + LocalNERPredictor, + NERExtractEndpointModel, + NERModel, + NERModelEntity, +) from privacy_guard.engines.regex import ( RegexEngine, RegexEngineConfig, @@ -43,6 +57,16 @@ "EntityProcessingEngine", "EntityProcessingError", "EntityProcessingStrategy", + "LocalNERModel", + "LocalNERPredictor", + "NEREngine", + "NEREngineConfig", + "NERExtractEndpointModel", + "NERModel", + "NERModelEntity", + "NEROverlapMode", + "NERReplacement", + "NERResources", "RegexEngine", "RegexEngineConfig", "RegexEntity", diff --git a/projects/privacy-guard/src/privacy_guard/engines/_replacement.py b/projects/privacy-guard/src/privacy_guard/engines/_replacement.py new file mode 100644 index 0000000..0ab6d4a --- /dev/null +++ b/projects/privacy-guard/src/privacy_guard/engines/_replacement.py @@ -0,0 +1,77 @@ +"""Private constrained-template validation and bounded rendering.""" + +from __future__ import annotations + +from string import Formatter +from typing import Protocol + +from privacy_guard.constants import MAX_BODY_BYTES, MAX_DIAGNOSTIC_TEXT_BYTES +from privacy_guard.errors import EngineLimitExceededError +from privacy_guard.string_validators import validate_scalar_string + + +class _ReplacementSpan(Protocol): + @property + def entity(self) -> str: ... + + @property + def start(self) -> int: ... + + @property + def end(self) -> int: ... + + +def validate_replacement_template(value: object) -> str: + """Return a bounded template containing only literals and ``{entity}``.""" + template = validate_scalar_string(value) + if len(template.encode("utf-8")) > MAX_DIAGNOSTIC_TEXT_BYTES: + raise ValueError("replacement template exceeds the size limit") + try: + for _, field_name, format_spec, conversion in Formatter().parse(template): + if field_name is not None and field_name != "entity": + raise ValueError + if format_spec or conversion is not None: + raise ValueError + except ValueError: + raise ValueError("replacement template syntax is invalid") from None + return template + + +def render_bounded_replacement( + text: str, + spans: tuple[_ReplacementSpan, ...], + template: str, + *, + limit_message: str, +) -> str: + """Project UTF-8 size, then render ordered non-overlapping spans.""" + projected_size = 0 + cursor = 0 + for span in spans: + projected_size += len(text[cursor : span.start].encode("utf-8")) + projected_size += _rendered_template_size(template, span.entity) + if projected_size > MAX_BODY_BYTES: + raise EngineLimitExceededError(limit_message) + cursor = span.end + projected_size += len(text[cursor:].encode("utf-8")) + if projected_size > MAX_BODY_BYTES: + raise EngineLimitExceededError(limit_message) + + parts: list[str] = [] + cursor = 0 + for span in spans: + parts.append(text[cursor : span.start]) + parts.append(template.format(entity=span.entity)) + cursor = span.end + parts.append(text[cursor:]) + return "".join(parts) + + +def _rendered_template_size(template: str, entity: str) -> int: + size = 0 + entity_size = len(entity.encode("utf-8")) + for literal, field_name, _, _ in Formatter().parse(template): + size += len(literal.encode("utf-8")) + if field_name is not None: + size += entity_size + return size diff --git a/projects/privacy-guard/src/privacy_guard/engines/ner.py b/projects/privacy-guard/src/privacy_guard/engines/ner.py new file mode 100644 index 0000000..1bd34d0 --- /dev/null +++ b/projects/privacy-guard/src/privacy_guard/engines/ner.py @@ -0,0 +1,274 @@ +"""Resource-backed general named-entity recognition engine.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import StrEnum +from typing import Literal + +from pydantic import Field, field_validator + +from privacy_guard.base import StrictDomainModel +from privacy_guard.constants import ( + MAX_DETECTIONS_PER_STAGE, + MAX_NER_LABEL_BYTES, + MAX_NER_LABELS, + MAX_NER_LABELS_BYTES, +) +from privacy_guard.engines._replacement import ( + render_bounded_replacement, + validate_replacement_template, +) +from privacy_guard.engines.base import ( + EngineConfig, + EngineResources, + EntityDetection, + EntityProcessingEngine, + EntityProcessingStrategy, + TextProcessingResult, +) +from privacy_guard.engines.ner_model import NERModel, NERModelEntity +from privacy_guard.errors import ( + EngineConfigurationError, + EngineContractError, + EngineLimitExceededError, +) +from privacy_guard.string_validators import validate_scalar_string +from privacy_guard.timeout import Timeout + + +class NEROverlapMode(StrEnum): + """Select whether the model may return overlapping nested entities.""" + + NESTED = "nested" + FLAT = "flat" + + +class NERReplacement(StrictDomainModel): + """A constrained entity-label template replacement.""" + + strategy: Literal["template"] = "template" + template: str = Field(default="[{entity}]", repr=False) + + @field_validator("template", mode="before") + @classmethod + def _template_is_safe_and_bounded(cls, value: object) -> str: + return validate_replacement_template(value) + + +class NEREngineConfig(EngineConfig): + """Exact policy configuration for general named-entity recognition.""" + + engine: Literal["ner"] = "ner" + labels: tuple[str, ...] + threshold: float = Field(ge=0, le=1, allow_inf_nan=False) + overlap_mode: NEROverlapMode = NEROverlapMode.NESTED + replacement: NERReplacement | None = None + + @field_validator("labels", mode="before") + @classmethod + def _labels_are_an_ordered_tuple(cls, value: object) -> object: + if not isinstance(value, list | tuple) or not value: + raise ValueError("labels must be a non-empty list") + return tuple(value) + + @field_validator("labels") + @classmethod + def _labels_are_bounded_and_unambiguous( + cls, + value: tuple[str, ...], + ) -> tuple[str, ...]: + if len(value) > MAX_NER_LABELS: + raise ValueError("labels exceed the count limit") + total_bytes = 0 + normalized: set[str] = set() + for raw_label in value: + label = validate_scalar_string(raw_label) + label_bytes = len(label.encode("utf-8")) + if ( + not label + or not label.isprintable() + or label != label.strip() + or label_bytes > MAX_NER_LABEL_BYTES + ): + raise ValueError("NER label is invalid") + casefolded = label.casefold() + if casefolded in normalized: + raise ValueError("NER labels must be unique ignoring case") + normalized.add(casefolded) + total_bytes += label_bytes + if total_bytes > MAX_NER_LABELS_BYTES: + raise ValueError("labels exceed the total size limit") + return value + + @field_validator("overlap_mode", mode="before") + @classmethod + def _parse_overlap_mode(cls, value: object) -> NEROverlapMode: + return NEROverlapMode(validate_scalar_string(value)) + + +@dataclass(frozen=True) +class NERResources(EngineResources): + """Operator-owned NER model facade shared by engine instances.""" + + model: NERModel + + +class NEREngine(EntityProcessingEngine[NEREngineConfig, NERResources]): + """Detect named entities with an operator-supplied model facade.""" + + supported_strategies = frozenset( + { + EntityProcessingStrategy.DETECT, + EntityProcessingStrategy.REPLACE, + } + ) + + @classmethod + def _validate_config( + cls, + config: NEREngineConfig, + resources: NERResources, + ) -> None: + if not isinstance(resources.model, NERModel): + raise EngineConfigurationError("NER model resources are invalid") + + @classmethod + def _validate_run_config( + cls, + config: NEREngineConfig, + resources: NERResources, + *, + strategy: EntityProcessingStrategy, + ) -> None: + if strategy is EntityProcessingStrategy.REPLACE and config.replacement is None: + raise EngineConfigurationError("NER replacement configuration is required") + + def _initialize(self) -> None: + self._model = self.resources.model + self._canonical_labels = { + label.casefold(): (index, label) + for index, label in enumerate(self.config.labels) + } + + def _run( + self, + text: str, + *, + strategy: EntityProcessingStrategy, + timeout: Timeout, + ) -> TextProcessingResult: + raw_entities = self._model.predict_entities( + text, + labels=self.config.labels, + threshold=self.config.threshold, + flat_ner=self.config.overlap_mode is NEROverlapMode.FLAT, + timeout=timeout, + ) + if not isinstance(raw_entities, tuple): + raise EngineContractError("NER model output is invalid") + unique: dict[tuple[int, int, str], NERModelEntity] = {} + for index, entity in enumerate(raw_entities): + if index >= MAX_DETECTIONS_PER_STAGE: + raise EngineLimitExceededError("NER detection count exceeds the limit") + normalized = self._validate_entity(entity, text) + key = (normalized.start, normalized.end, normalized.label) + current = unique.get(key) + if current is None or normalized.score > current.score: + unique[key] = normalized + if len(unique) > MAX_DETECTIONS_PER_STAGE: + raise EngineLimitExceededError("NER detection count exceeds the limit") + entities = tuple( + sorted( + unique.values(), + key=lambda item: ( + item.start, + item.end, + self._canonical_labels[item.label.casefold()][0], + ), + ) + ) + detections = tuple( + EntityDetection( + entity=entity.label, + start=entity.start, + end=entity.end, + ) + for entity in entities + ) + output_text = text + if strategy is EntityProcessingStrategy.REPLACE and entities: + replacement = self.config.replacement + if replacement is None: + raise EngineConfigurationError( + "NER replacement configuration is required" + ) + winners = self._resolve_replacement_winners(entities) + output_text = render_bounded_replacement( + text, + tuple( + EntityDetection( + entity=winner.label, + start=winner.start, + end=winner.end, + ) + for winner in winners + ), + replacement.template, + limit_message="NER replacement exceeds the size limit", + ) + return TextProcessingResult(text=output_text, detections=detections) + + def _validate_entity(self, value: object, text: str) -> NERModelEntity: + if not isinstance(value, NERModelEntity): + raise EngineContractError("NER model entity is invalid") + canonical = self._canonical_labels.get(value.label.casefold()) + if canonical is None or value.end > len(text): + raise EngineContractError("NER model entity is invalid") + return NERModelEntity( + label=canonical[1], + start=value.start, + end=value.end, + score=value.score, + ) + + def _resolve_replacement_winners( + self, + entities: tuple[NERModelEntity, ...], + ) -> tuple[NERModelEntity, ...]: + winners: list[NERModelEntity] = [] + ranked = sorted( + entities, + key=lambda item: ( + -item.score, + -(item.end - item.start), + item.start, + self._canonical_labels[item.label.casefold()][0], + item.end, + ), + ) + for candidate in ranked: + if all( + candidate.end <= winner.start or candidate.start >= winner.end + for winner in winners + ): + winners.append(candidate) + return tuple( + sorted( + winners, + key=lambda item: ( + item.start, + item.end, + self._canonical_labels[item.label.casefold()][0], + ), + ) + ) + + +__all__ = [ + "NEREngine", + "NEREngineConfig", + "NEROverlapMode", + "NERReplacement", + "NERResources", +] diff --git a/projects/privacy-guard/src/privacy_guard/engines/ner_model.py b/projects/privacy-guard/src/privacy_guard/engines/ner_model.py new file mode 100644 index 0000000..ec0dde3 --- /dev/null +++ b/projects/privacy-guard/src/privacy_guard/engines/ner_model.py @@ -0,0 +1,340 @@ +"""Provider-neutral model facade for named-entity recognition.""" + +from __future__ import annotations + +import json +from collections.abc import Iterator, Mapping +from dataclasses import dataclass, field +from threading import Lock +from typing import Protocol, Self, runtime_checkable +from urllib.error import HTTPError, URLError +from urllib.parse import urlsplit +from urllib.request import Request, urlopen + +from pydantic import Field, model_validator + +from privacy_guard.base import StrictDomainModel +from privacy_guard.constants import ( + MAX_DETECTIONS_PER_STAGE, + MAX_NER_ENDPOINT_RESPONSE_BYTES, +) +from privacy_guard.errors import EngineExecutionError, EngineLimitExceededError +from privacy_guard.string_validators import ScalarString, validate_scalar_string +from privacy_guard.timeout import Timeout + + +class NERModelEntity(StrictDomainModel): + """One normalized model detection using Unicode code-point offsets.""" + + label: ScalarString + start: int = Field(ge=0) + end: int + score: float = Field(ge=0, le=1, allow_inf_nan=False) + + @model_validator(mode="after") + def _span_is_non_empty(self) -> Self: + if self.end <= self.start: + raise ValueError("NER model entity span must be non-empty") + return self + + +@runtime_checkable +class NERModel(Protocol): + """Complete execution facade required by the built-in NER engine.""" + + def predict_entities( + self, + text: str, + *, + labels: tuple[str, ...], + threshold: float, + flat_ner: bool, + timeout: Timeout, + ) -> tuple[NERModelEntity, ...]: + """Return normalized entities for the complete input text.""" + ... + + +@dataclass(frozen=True) +class NERExtractEndpointModel: + """Call the explicit GLiNER-compatible ``POST /v1/extract`` contract.""" + + endpoint: str = field(repr=False) + model: str = field(repr=False) + chunk_length: int = 384 + overlap: int = 128 + max_response_bytes: int = MAX_NER_ENDPOINT_RESPONSE_BYTES + + def __post_init__(self) -> None: + try: + endpoint = validate_scalar_string(self.endpoint) + model = validate_scalar_string(self.model) + except ValueError: + raise ValueError("NER endpoint configuration is invalid") from None + parsed = urlsplit(endpoint) + if ( + parsed.scheme not in {"http", "https"} + or not parsed.hostname + or parsed.username is not None + or parsed.password is not None + or parsed.query + or parsed.fragment + or parsed.path != "/v1/extract" + ): + raise ValueError("NER extract endpoint is invalid") + if not model or not model.isprintable(): + raise ValueError("NER model identifier is invalid") + if ( + isinstance(self.chunk_length, bool) + or not isinstance(self.chunk_length, int) + or self.chunk_length <= 0 + or isinstance(self.overlap, bool) + or not isinstance(self.overlap, int) + or self.overlap < 0 + or self.overlap >= self.chunk_length + ): + raise ValueError("NER chunk configuration is invalid") + if ( + isinstance(self.max_response_bytes, bool) + or not isinstance(self.max_response_bytes, int) + or self.max_response_bytes <= 0 + or self.max_response_bytes > MAX_NER_ENDPOINT_RESPONSE_BYTES + ): + raise ValueError("NER endpoint response limit is invalid") + + def predict_entities( + self, + text: str, + *, + labels: tuple[str, ...], + threshold: float, + flat_ner: bool, + timeout: Timeout, + ) -> tuple[NERModelEntity, ...]: + payload = json.dumps( + { + "text": text, + "labels": labels, + "model": self.model, + "threshold": threshold, + "chunk_length": self.chunk_length, + "overlap": self.overlap, + "flat_ner": flat_ner, + }, + allow_nan=False, + ensure_ascii=False, + separators=(",", ":"), + ).encode("utf-8") + request = Request( + self.endpoint, + data=payload, + headers={ + "Accept": "application/json", + "Content-Type": "application/json; charset=utf-8", + }, + method="POST", + ) + try: + with timeout.enforce(): + with urlopen( + request, + timeout=timeout.remaining_seconds(), + ) as response: + response_bytes = response.read(self.max_response_bytes + 1) + except (HTTPError, URLError, OSError, TimeoutError): + raise EngineExecutionError("NER endpoint request failed") from None + if len(response_bytes) > self.max_response_bytes: + raise EngineLimitExceededError( + "NER endpoint response exceeds the size limit" + ) + try: + decoded = json.loads(response_bytes) + return _normalize_model_output(decoded) + except (UnicodeError, json.JSONDecodeError, TypeError, ValueError): + raise EngineExecutionError("NER endpoint response is invalid") from None + + +class LocalNERPredictor(Protocol): + """Narrow structural contract implemented by a loaded GLiNER model.""" + + def predict_entities( + self, + text: str, + labels: list[str], + *, + threshold: float, + flat_ner: bool, + ) -> object: + """Return the loaded model's entity payload.""" + ... + + +class LocalNERModel: + """Serialize calls to an already-loaded model that processes complete text.""" + + def __init__( + self, + *, + model: LocalNERPredictor, + chunk_length: int = 384, + overlap: int = 128, + ) -> None: + if ( + isinstance(chunk_length, bool) + or not isinstance(chunk_length, int) + or chunk_length <= 0 + or isinstance(overlap, bool) + or not isinstance(overlap, int) + or overlap < 0 + or overlap >= chunk_length + ): + raise ValueError("NER chunk configuration is invalid") + self._model = model + self._chunk_length = chunk_length + self._overlap = overlap + self._lock = Lock() + + @property + def chunk_length(self) -> int: + """Return the operator-selected model chunk length.""" + return self._chunk_length + + @property + def overlap(self) -> int: + """Return the operator-selected model chunk overlap.""" + return self._overlap + + def predict_entities( + self, + text: str, + *, + labels: tuple[str, ...], + threshold: float, + flat_ner: bool, + timeout: Timeout, + ) -> tuple[NERModelEntity, ...]: + acquired = self._lock.acquire(timeout=timeout.remaining_seconds()) + if not acquired: + timeout.raise_if_expired() + raise EngineExecutionError("NER local model is unavailable") + try: + unique: dict[tuple[int, int, str], NERModelEntity] = {} + for start, chunk in _iter_text_chunks( + text, + chunk_length=self._chunk_length, + overlap=self._overlap, + ): + with timeout.enforce(): + output = self._model.predict_entities( + chunk, + list(labels), + threshold=threshold, + flat_ner=flat_ner, + ) + for entity in _normalize_model_output(output): + if ( + entity.start < 0 + or entity.end <= entity.start + or entity.end > len(chunk) + ): + raise ValueError("local entity span is invalid") + rebased = NERModelEntity( + label=entity.label, + start=start + entity.start, + end=start + entity.end, + score=entity.score, + ) + key = (rebased.start, rebased.end, rebased.label) + current = unique.get(key) + if current is None or rebased.score > current.score: + unique[key] = rebased + if len(unique) > MAX_DETECTIONS_PER_STAGE: + raise EngineLimitExceededError( + "NER model returned too many detections" + ) + entities = tuple(unique.values()) + if flat_ner: + entities = _select_flat_entities(entities) + return tuple( + sorted( + entities, + key=lambda item: (item.start, item.end, item.label), + ) + ) + except (TypeError, ValueError, RuntimeError, OSError): + raise EngineExecutionError("NER local model inference failed") from None + finally: + self._lock.release() + + +def _normalize_model_output(value: object) -> tuple[NERModelEntity, ...]: + if isinstance(value, Mapping): + value = value.get("entities") + if not isinstance(value, list | tuple): + raise ValueError("entity output must be a list") + if len(value) > MAX_DETECTIONS_PER_STAGE: + raise EngineLimitExceededError("NER model returned too many detections") + normalized: list[NERModelEntity] = [] + for item in value: + if not isinstance(item, Mapping): + raise ValueError("entity item must be an object") + normalized.append( + NERModelEntity.model_validate( + { + "label": item.get("label"), + "start": item.get("start"), + "end": item.get("end"), + "score": item.get("score"), + } + ) + ) + return tuple(normalized) + + +def _iter_text_chunks( + text: str, + *, + chunk_length: int, + overlap: int, +) -> Iterator[tuple[int, str]]: + if len(text) <= chunk_length: + yield 0, text + return + start = 0 + while start < len(text): + end = min(start + chunk_length, len(text)) + yield start, text[start:end] + if end == len(text): + break + start = end - overlap + + +def _select_flat_entities( + entities: tuple[NERModelEntity, ...], +) -> tuple[NERModelEntity, ...]: + winners: list[NERModelEntity] = [] + for candidate in sorted( + entities, + key=lambda item: ( + -item.score, + -(item.end - item.start), + item.start, + item.end, + item.label, + ), + ): + if all( + candidate.end <= winner.start or candidate.start >= winner.end + for winner in winners + ): + winners.append(candidate) + return tuple(winners) + + +__all__ = [ + "LocalNERModel", + "LocalNERPredictor", + "NERExtractEndpointModel", + "NERModel", + "NERModelEntity", +] diff --git a/projects/privacy-guard/src/privacy_guard/engines/regex.py b/projects/privacy-guard/src/privacy_guard/engines/regex.py index 4b0323b..70e4e51 100644 --- a/projects/privacy-guard/src/privacy_guard/engines/regex.py +++ b/projects/privacy-guard/src/privacy_guard/engines/regex.py @@ -9,7 +9,6 @@ from dataclasses import dataclass from pathlib import Path from stat import S_ISREG -from string import Formatter from threading import RLock from typing import Literal, Protocol, Self @@ -23,9 +22,7 @@ from privacy_guard.base import StrictDomainModel from privacy_guard.constants import ( - MAX_BODY_BYTES, MAX_DETECTIONS_PER_STAGE, - MAX_DIAGNOSTIC_TEXT_BYTES, MAX_REGEX_CATALOG_FILE_BYTES, MAX_REGEX_CATALOG_PATH_BYTES, MAX_REGEX_COMPILED_CACHE_WEIGHT_BYTES, @@ -35,6 +32,10 @@ MAX_REGEX_RULES_PER_CATALOG, REGEX_COMPILED_RULE_WEIGHT_BYTES, ) +from privacy_guard.engines._replacement import ( + render_bounded_replacement, + validate_replacement_template, +) from privacy_guard.engines.base import ( ConfidenceLevel, EngineConfig, @@ -148,17 +149,7 @@ class RegexReplacement(StrictDomainModel): @field_validator("template") @classmethod def _template_is_safe_and_bounded(cls, value: str) -> str: - if len(value.encode("utf-8")) > MAX_DIAGNOSTIC_TEXT_BYTES: - raise ValueError("replacement template exceeds the size limit") - try: - for _, field_name, format_spec, conversion in Formatter().parse(value): - if field_name is not None and field_name != "entity": - raise ValueError - if format_spec or conversion is not None: - raise ValueError - except ValueError: - raise ValueError("replacement template syntax is invalid") from None - return value + return validate_replacement_template(value) class RegexEngineConfig(EngineConfig): @@ -183,7 +174,6 @@ def _load_pattern_catalog( cls, value: object, ) -> object: - del cls if isinstance(value, str): return _load_pattern_catalog_file(value) return value @@ -215,7 +205,6 @@ def _validate_run_config( *, strategy: EntityProcessingStrategy, ) -> None: - del cls, resources if strategy is EntityProcessingStrategy.REPLACE and config.replacement is None: raise EngineConfigurationError( "regex replacement configuration is required" @@ -288,10 +277,11 @@ def _run( "regex replacement configuration is required" ) winners = _resolve_overlaps(detections_with_identity) - output_text = _render_bounded_replacement( + output_text = render_bounded_replacement( text, winners, replacement.template, + limit_message="regex replacement exceeds the size limit", ) return TextProcessingResult(text=output_text, detections=detections) @@ -659,43 +649,6 @@ def _categorical_confidence_rank(confidence: object) -> int: return _CONFIDENCE_RANK[confidence] -def _render_bounded_replacement( - text: str, - detections: tuple[EntityDetection, ...], - template: str, -) -> str: - projected_size = 0 - cursor = 0 - for detection in detections: - projected_size += len(text[cursor : detection.start].encode("utf-8")) - projected_size += _rendered_template_size(template, detection.entity) - if projected_size > MAX_BODY_BYTES: - raise EngineLimitExceededError("regex replacement exceeds the size limit") - cursor = detection.end - projected_size += len(text[cursor:].encode("utf-8")) - if projected_size > MAX_BODY_BYTES: - raise EngineLimitExceededError("regex replacement exceeds the size limit") - - parts: list[str] = [] - cursor = 0 - for detection in detections: - parts.append(text[cursor : detection.start]) - parts.append(template.format(entity=detection.entity)) - cursor = detection.end - parts.append(text[cursor:]) - return "".join(parts) - - -def _rendered_template_size(template: str, entity: str) -> int: - size = 0 - entity_size = len(entity.encode("utf-8")) - for literal, field_name, _, _ in Formatter().parse(template): - size += len(literal.encode("utf-8")) - if field_name is not None: - size += entity_size - return size - - _NAME_PATTERN = regex.compile(r"[A-Za-z_][A-Za-z0-9_-]*\Z") _INLINE_FLAG_PATTERN = regex.compile(r"[A-Za-z0-9-]+(?=[:)])") _CONFIDENCE_RANK = { diff --git a/projects/privacy-guard/src/privacy_guard/engines/registry.py b/projects/privacy-guard/src/privacy_guard/engines/registry.py index bee9deb..64989be 100644 --- a/projects/privacy-guard/src/privacy_guard/engines/registry.py +++ b/projects/privacy-guard/src/privacy_guard/engines/registry.py @@ -24,6 +24,7 @@ EntityProcessingEngine, EntityProcessingStrategy, ) +from privacy_guard.engines.ner import NEREngine, NERResources from privacy_guard.engines.regex import ( RegexEngine, ) @@ -47,13 +48,24 @@ class EngineDescription: class EngineRegistry: """Register engine implementations and finalize their exact policy union.""" - def __init__(self, *, include_builtin_engines: bool = False) -> None: + def __init__( + self, + *, + include_builtin_engines: bool = False, + ner_resources: NERResources | None = None, + ) -> None: self._registrations: dict[str, _Registration] = {} self._config_adapter: TypeAdapter[PrivacyGuardConfig[EngineConfig]] | None = ( None ) if include_builtin_engines: self.register(RegexEngine) + if ner_resources is not None: + self.register(NEREngine, resources=ner_resources) + elif ner_resources is not None: + raise EngineRegistryError( + "NER resources require built-in engines to be enabled" + ) @property def is_finalized(self) -> bool: @@ -231,9 +243,15 @@ def _require_config_adapter( return self._config_adapter -def create_builtin_registry() -> EngineRegistry: +def create_builtin_registry( + *, + ner_resources: NERResources | None = None, +) -> EngineRegistry: """Build the finalized registry shipped by the base package.""" - return EngineRegistry(include_builtin_engines=True).finalize() + return EngineRegistry( + include_builtin_engines=True, + ner_resources=ner_resources, + ).finalize() @dataclass(frozen=True) diff --git a/projects/privacy-guard/tests/engines/test_ner.py b/projects/privacy-guard/tests/engines/test_ner.py new file mode 100644 index 0000000..6c70268 --- /dev/null +++ b/projects/privacy-guard/tests/engines/test_ner.py @@ -0,0 +1,285 @@ +"""Tests for the built-in general NER engine.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import pytest +from pydantic import ValidationError + +import privacy_guard.engines._replacement as replacement_module +import privacy_guard.engines.ner as ner_module +from privacy_guard.engines import ( + EngineConfigurationError, + EngineContractError, + EngineLimitExceededError, + EntityProcessingStrategy, + NEREngine, + NEREngineConfig, + NERModelEntity, + NERResources, +) +from privacy_guard.engines.registry import create_builtin_registry +from privacy_guard.timeout import Timeout + + +@dataclass +class FakeNERModel: + entities: tuple[NERModelEntity, ...] + on_call: Callable[..., None] | None = None + + def predict_entities( + self, + text: str, + *, + labels: tuple[str, ...], + threshold: float, + flat_ner: bool, + timeout: Timeout, + ) -> tuple[NERModelEntity, ...]: + if self.on_call is not None: + self.on_call(text, labels, threshold, flat_ner, timeout) + return self.entities + + +def _config( + *, + labels: list[str] | None = None, + threshold: object = 0.5, + overlap_mode: str = "nested", + replacement: dict[str, object] | None = None, +) -> NEREngineConfig: + values: dict[str, object] = { + "engine": "ner", + "labels": ["person", "email"] if labels is None else labels, + "threshold": threshold, + "overlap_mode": overlap_mode, + } + if replacement is not None: + values["replacement"] = replacement + return NEREngineConfig.model_validate(values) + + +def _run( + model: FakeNERModel, + config: NEREngineConfig, + text: str, + *, + strategy: EntityProcessingStrategy = EntityProcessingStrategy.DETECT, +) -> tuple[str, list[tuple[str, int, int]]]: + result = NEREngine(config, NERResources(model=model)).run( + text, + strategy=strategy, + timeout=Timeout.from_seconds(1), + ) + return result.text, [ + (detection.entity, detection.start, detection.end) + for detection in result.detections + ] + + +def test_forwards_policy_behavior_and_returns_deterministic_detections() -> None: + calls: list[tuple[object, ...]] = [] + + def record_call(*values: object) -> None: + calls.append(values) + + model = FakeNERModel( + entities=( + NERModelEntity(label="EMAIL", start=16, end=19, score=0.8), + NERModelEntity(label="PERSON", start=0, end=5, score=0.7), + NERModelEntity(label="person", start=0, end=5, score=0.9), + ), + on_call=record_call, + ) + config = _config(threshold=0.3, overlap_mode="flat") + + output, detections = _run(model, config, "Alice contacted x@y") + + assert output == "Alice contacted x@y" + assert detections == [("person", 0, 5), ("email", 16, 19)] + assert calls + assert calls[0][0:4] == ( + "Alice contacted x@y", + ("person", "email"), + 0.3, + True, + ) + + +def test_nested_mode_maps_to_flat_ner_false() -> None: + flat_values: list[bool] = [] + + def record_call( + text: str, + labels: tuple[str, ...], + threshold: float, + flat_ner: bool, + timeout: Timeout, + ) -> None: + del text, labels, threshold, timeout + flat_values.append(flat_ner) + + _run(FakeNERModel(entities=(), on_call=record_call), _config(), "text") + + assert flat_values == [False] + + +def test_replacement_keeps_all_detections_and_ranks_overlap_winners() -> None: + model = FakeNERModel( + entities=( + NERModelEntity(label="person", start=0, end=5, score=0.6), + NERModelEntity(label="email", start=1, end=4, score=0.9), + NERModelEntity(label="person", start=6, end=10, score=0.8), + ) + ) + config = _config(replacement={"strategy": "template", "template": "<{entity}>"}) + + output, detections = _run( + model, + config, + "abcdefghij", + strategy=EntityProcessingStrategy.REPLACE, + ) + + assert output == "aef" + assert len(detections) == 3 + + +def test_replace_requires_replacement_configuration() -> None: + with pytest.raises(EngineConfigurationError): + NEREngine.validate_run_config( + _config(), + NERResources(model=FakeNERModel(entities=())), + strategy=EntityProcessingStrategy.REPLACE, + ) + + +def test_registry_validation_is_pure_and_does_not_run_model() -> None: + calls = 0 + + def record_call(*values: object) -> None: + nonlocal calls + del values + calls += 1 + + registry = create_builtin_registry( + ner_resources=NERResources(model=FakeNERModel(entities=(), on_call=record_call)) + ) + + registry.validate_config( + { + "entity_processing": { + "stages": [ + { + "config": { + "engine": "ner", + "labels": ["person"], + "threshold": 0.5, + } + } + ] + }, + "on_detection": {"action": "detect"}, + } + ) + + assert calls == 0 + + +@pytest.mark.parametrize( + "labels", + [ + [], + ["person", "PERSON"], + [""], + [" person"], + ["line\nbreak"], + ["x" * 257], + ], +) +def test_invalid_labels_are_rejected(labels: list[str]) -> None: + with pytest.raises(ValidationError): + _config(labels=labels) + + +@pytest.mark.parametrize( + "threshold", + [-0.1, 1.1, float("nan"), float("inf"), True, "0.5"], +) +def test_invalid_thresholds_are_rejected(threshold: object) -> None: + with pytest.raises(ValidationError): + _config(threshold=threshold) + + +@pytest.mark.parametrize( + "replacement", + [ + {"strategy": "template", "template": "{unknown}"}, + {"strategy": "template", "template": "{entity!r}"}, + {"strategy": "template", "template": "{entity:>10}"}, + {"strategy": "template", "template": "{"}, + ], +) +def test_replacement_template_is_constrained( + replacement: dict[str, object], +) -> None: + with pytest.raises(ValidationError): + _config(replacement=replacement) + + +@pytest.mark.parametrize( + "entity", + [ + NERModelEntity(label="unknown", start=0, end=1, score=0.5), + NERModelEntity(label="person", start=0, end=5, score=0.5), + ], +) +def test_engine_rejects_entities_invalid_for_policy_or_input( + entity: NERModelEntity, +) -> None: + with pytest.raises(EngineContractError): + _run(FakeNERModel(entities=(entity,)), _config(), "abc") + + +@pytest.mark.parametrize( + "values", + [ + {"label": "person", "start": -1, "end": 1, "score": 0.5}, + {"label": "person", "start": 1, "end": 1, "score": 0.5}, + {"label": "person", "start": 0, "end": 1, "score": -0.1}, + {"label": "person", "start": 0, "end": 1, "score": 1.1}, + {"label": "person", "start": 0, "end": 1, "score": float("nan")}, + {"label": "person", "start": True, "end": 1, "score": 0.5}, + ], +) +def test_normalized_entity_invariants_are_pydantic_validated( + values: dict[str, object], +) -> None: + with pytest.raises(ValidationError): + NERModelEntity.model_validate(values) + + +def test_detection_count_is_bounded(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(ner_module, "MAX_DETECTIONS_PER_STAGE", 1) + entities = ( + NERModelEntity(label="person", start=0, end=1, score=0.5), + NERModelEntity(label="person", start=1, end=2, score=0.5), + ) + + with pytest.raises(EngineLimitExceededError): + _run(FakeNERModel(entities=entities), _config(), "ab") + + +def test_replacement_size_is_projected_before_rendering( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(replacement_module, "MAX_BODY_BYTES", 4) + config = _config(replacement={"strategy": "template", "template": "[{entity}]"}) + model = FakeNERModel( + entities=(NERModelEntity(label="person", start=0, end=1, score=0.5),) + ) + + with pytest.raises(EngineLimitExceededError): + _run(model, config, "x", strategy=EntityProcessingStrategy.REPLACE) diff --git a/projects/privacy-guard/tests/engines/test_ner_model.py b/projects/privacy-guard/tests/engines/test_ner_model.py new file mode 100644 index 0000000..310cd41 --- /dev/null +++ b/projects/privacy-guard/tests/engines/test_ner_model.py @@ -0,0 +1,417 @@ +"""Contract tests for NER model facade implementations.""" + +from __future__ import annotations + +import json +import os +from collections.abc import Iterator +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from importlib import import_module +from threading import Event, Thread + +import pytest + +from privacy_guard.engines import LocalNERModel, NERExtractEndpointModel +from privacy_guard.errors import ( + EngineExecutionError, + EngineLimitExceededError, + TimeoutExpiredError, +) +from privacy_guard.timeout import Timeout + + +class _EndpointHandler(BaseHTTPRequestHandler): + response_body = b'{"entities":[]}' + received_body = b"" + received_path = "" + + def do_POST(self) -> None: + type(self).received_path = self.path + content_length = int(self.headers.get("Content-Length", "0")) + type(self).received_body = self.rfile.read(content_length) + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(type(self).response_body))) + self.end_headers() + self.wfile.write(type(self).response_body) + + def log_message(self, format: str, *args: object) -> None: + del format, args + + +@contextmanager +def _endpoint( + response_body: bytes, +) -> Iterator[tuple[str, type[_EndpointHandler]]]: + handler = type("EndpointHandler", (_EndpointHandler,), {}) + handler.response_body = response_body + server = ThreadingHTTPServer(("127.0.0.1", 0), handler) + thread = Thread(target=server.serve_forever) + thread.start() + try: + host = server.server_address[0] + port = server.server_address[1] + yield f"http://{host}:{port}/v1/extract", handler + finally: + server.shutdown() + thread.join() + server.server_close() + + +def test_extract_endpoint_sends_exact_contract_and_normalizes_entities() -> None: + response = json.dumps( + { + "entities": [ + { + "text": "Alice", + "label": "person", + "start": 0, + "end": 5, + "score": 0.75, + "unknown": "ignored", + } + ], + "tagged_text": "ignored", + "total_entities": 1, + } + ).encode() + with _endpoint(response) as (endpoint, handler): + model = NERExtractEndpointModel( + endpoint=endpoint, + model="nvidia/gliner-PII", + chunk_length=384, + overlap=128, + ) + + entities = model.predict_entities( + "Alice", + labels=("person",), + threshold=0.3, + flat_ner=False, + timeout=Timeout.from_seconds(2), + ) + + assert handler.received_path == "/v1/extract" + assert json.loads(handler.received_body) == { + "text": "Alice", + "labels": ["person"], + "model": "nvidia/gliner-PII", + "threshold": 0.3, + "chunk_length": 384, + "overlap": 128, + "flat_ner": False, + } + assert [ + (entity.label, entity.start, entity.end, entity.score) for entity in entities + ] == [("person", 0, 5, 0.75)] + + +@pytest.mark.parametrize( + "body", + [ + b"not-json", + b"{}", + b'{"entities":{}}', + b'{"entities":[{"label":"person","start":0,"end":1}]}', + b'{"entities":[{"label":"person","start":false,"end":1,"score":0.5}]}', + b'{"entities":[{"label":"person","start":1,"end":1,"score":0.5}]}', + b'{"entities":[{"label":"person","start":0,"end":1,"score":2}]}', + ], +) +def test_extract_endpoint_rejects_malformed_responses(body: bytes) -> None: + with _endpoint(body) as (endpoint, _): + model = NERExtractEndpointModel(endpoint=endpoint, model="model") + + with pytest.raises(EngineExecutionError): + model.predict_entities( + "x", + labels=("person",), + threshold=0.5, + flat_ner=True, + timeout=Timeout.from_seconds(2), + ) + + +def test_extract_endpoint_bounds_response_before_json_decoding() -> None: + with _endpoint(b'{"entities":[]}' + b" " * 100) as (endpoint, _): + model = NERExtractEndpointModel( + endpoint=endpoint, + model="model", + max_response_bytes=16, + ) + + with pytest.raises(EngineLimitExceededError): + model.predict_entities( + "x", + labels=("person",), + threshold=0.5, + flat_ner=True, + timeout=Timeout.from_seconds(2), + ) + + +def test_extract_endpoint_translates_connection_failure_content_safely() -> None: + model = NERExtractEndpointModel( + endpoint="http://127.0.0.1:1/v1/extract", + model="secret-model-name", + ) + + with pytest.raises(EngineExecutionError) as exception_info: + model.predict_entities( + "secret request text", + labels=("person",), + threshold=0.5, + flat_ner=True, + timeout=Timeout.from_seconds(1), + ) + + message = str(exception_info.value) + assert "secret request text" not in message + assert "secret-model-name" not in message + assert "127.0.0.1" not in message + assert "127.0.0.1" not in repr(model) + assert "secret-model-name" not in repr(model) + + +@pytest.mark.parametrize( + "endpoint", + [ + "ftp://example.com/v1/extract", + "http://user:password@example.com/v1/extract", + "http://example.com/other", + "http://example.com/v1/extract?token=secret", + ], +) +def test_extract_endpoint_requires_an_explicit_safe_contract_url( + endpoint: str, +) -> None: + with pytest.raises(ValueError): + NERExtractEndpointModel(endpoint=endpoint, model="model") + + +class _LoadedModel: + def __init__(self) -> None: + self.calls: list[tuple[object, ...]] = [] + + def predict_entities( + self, + text: str, + labels: list[str], + *, + threshold: float, + flat_ner: bool, + ) -> object: + self.calls.append((text, labels, threshold, flat_ner)) + return [ + { + "text": "ignored", + "label": "person", + "start": 0, + "end": 1, + "score": 0.9, + } + ] + + +def test_local_model_passes_explicit_arguments_and_normalizes_output() -> None: + loaded = _LoadedModel() + model = LocalNERModel(model=loaded, chunk_length=512, overlap=64) + + entities = model.predict_entities( + "x", + labels=("person",), + threshold=0.4, + flat_ner=True, + timeout=Timeout.from_seconds(1), + ) + + assert loaded.calls == [("x", ["person"], 0.4, True)] + assert model.chunk_length == 512 + assert model.overlap == 64 + assert [ + (entity.label, entity.start, entity.end, entity.score) for entity in entities + ] == [("person", 0, 1, 0.9)] + + +class _ChunkingLoadedModel: + def __init__(self) -> None: + self.calls: list[str] = [] + + def predict_entities( + self, + text: str, + labels: list[str], + *, + threshold: float, + flat_ner: bool, + ) -> object: + del labels, threshold, flat_ner + self.calls.append(text) + entities: list[dict[str, object]] = [] + position = text.find("éx") + if position >= 0: + entities.append( + { + "label": "token", + "start": position, + "end": position + 2, + "score": 0.8, + } + ) + return entities + + +def test_local_model_chunks_complete_input_rebases_unicode_and_deduplicates() -> None: + loaded = _ChunkingLoadedModel() + model = LocalNERModel(model=loaded, chunk_length=5, overlap=2) + + entities = model.predict_entities( + "abcéxyz", + labels=("token",), + threshold=0.5, + flat_ner=False, + timeout=Timeout.from_seconds(1), + ) + + assert loaded.calls == ["abcéx", "éxyz"] + assert [ + (entity.label, entity.start, entity.end, entity.score) for entity in entities + ] == [("token", 3, 5, 0.8)] + + +class _BlockingLoadedModel: + def __init__(self) -> None: + self.started = Event() + self.release = Event() + + def predict_entities( + self, + text: str, + labels: list[str], + *, + threshold: float, + flat_ner: bool, + ) -> object: + del text, labels, threshold, flat_ner + self.started.set() + self.release.wait(timeout=2) + return [] + + +def test_local_model_bounds_serialized_lock_acquisition() -> None: + loaded = _BlockingLoadedModel() + model = LocalNERModel(model=loaded) + first_error: list[BaseException] = [] + + def first_call() -> None: + try: + model.predict_entities( + "first", + labels=("person",), + threshold=0.5, + flat_ner=True, + timeout=Timeout.from_seconds(2), + ) + except BaseException as error: + first_error.append(error) + + thread = Thread(target=first_call) + thread.start() + assert loaded.started.wait(timeout=1) + try: + with pytest.raises(TimeoutExpiredError): + model.predict_entities( + "second", + labels=("person",), + threshold=0.5, + flat_ner=True, + timeout=Timeout.from_seconds(0.01), + ) + finally: + loaded.release.set() + thread.join() + + assert first_error == [] + + +class _FailingLoadedModel: + def predict_entities( + self, + text: str, + labels: list[str], + *, + threshold: float, + flat_ner: bool, + ) -> object: + del labels, threshold, flat_ner + raise RuntimeError(text) + + +def test_local_model_translates_runtime_failure_without_content() -> None: + model = LocalNERModel(model=_FailingLoadedModel()) + + with pytest.raises(EngineExecutionError) as exception_info: + model.predict_entities( + "secret text", + labels=("person",), + threshold=0.5, + flat_ner=False, + timeout=Timeout.from_seconds(1), + ) + + assert "secret text" not in str(exception_info.value) + + +@pytest.mark.skipif( + "PRIVACY_GUARD_DGX_NER_ENDPOINT" not in os.environ, + reason="DGX NER smoke test is explicitly opt-in", +) +def test_opt_in_dgx_endpoint_detects_sample_entities() -> None: + endpoint = os.environ["PRIVACY_GUARD_DGX_NER_ENDPOINT"] + model_name = os.environ.get( + "PRIVACY_GUARD_DGX_NER_MODEL", + "nvidia/gliner-PII", + ) + model = NERExtractEndpointModel(endpoint=endpoint, model=model_name) + + entities = model.predict_entities( + "Contact Alice Example at alice@example.com or +1 202-555-0147.", + labels=("person", "email", "phone_number"), + threshold=0.3, + flat_ner=False, + timeout=Timeout.from_seconds(30), + ) + + assert {entity.label.casefold() for entity in entities} >= { + "person", + "email", + "phone_number", + } + + +@pytest.mark.skipif( + "PRIVACY_GUARD_LOCAL_NER_MODEL_PATH" not in os.environ, + reason="local GLiNER smoke test is explicitly opt-in", +) +def test_opt_in_local_gliner_detects_sample_entities_without_download() -> None: + model_path = os.environ["PRIVACY_GUARD_LOCAL_NER_MODEL_PATH"] + gliner_type = getattr(import_module("gliner"), "GLiNER") + loaded_model = gliner_type.from_pretrained( + model_path, + local_files_only=True, + ) + model = LocalNERModel(model=loaded_model) + + entities = model.predict_entities( + "Contact Alice Example at alice@example.com.", + labels=("person", "email"), + threshold=0.3, + flat_ner=False, + timeout=Timeout.from_seconds(30), + ) + + assert {entity.label.casefold() for entity in entities} >= { + "person", + "email", + } diff --git a/projects/privacy-guard/tests/engines/test_regex.py b/projects/privacy-guard/tests/engines/test_regex.py index 54ddd3f..554a38f 100644 --- a/projects/privacy-guard/tests/engines/test_regex.py +++ b/projects/privacy-guard/tests/engines/test_regex.py @@ -7,6 +7,7 @@ import pytest from pydantic import ValidationError +import privacy_guard.engines._replacement as replacement_module import privacy_guard.engines.regex as regex_module from privacy_guard.engines import ( EngineConfigurationError, @@ -282,7 +283,7 @@ def test_template_language_allows_only_literal_text_and_entity( def test_replacement_size_is_projected_before_rendering( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setattr(regex_module, "MAX_BODY_BYTES", 4) + monkeypatch.setattr(replacement_module, "MAX_BODY_BYTES", 4) config = _config( [{"pattern": "x", "confidence": "high"}], replacement={"strategy": "template", "template": "[{entity}]"}, diff --git a/projects/privacy-guard/tests/engines/test_registry.py b/projects/privacy-guard/tests/engines/test_registry.py index fc60047..e4d9c95 100644 --- a/projects/privacy-guard/tests/engines/test_registry.py +++ b/projects/privacy-guard/tests/engines/test_registry.py @@ -15,6 +15,8 @@ EngineResources, EntityProcessingEngine, EntityProcessingStrategy, + NERModelEntity, + NERResources, RegexEngine, TextProcessingResult, ) @@ -61,7 +63,6 @@ def _validate_run_config( *, strategy: EntityProcessingStrategy, ) -> None: - del cls, resources if strategy is EntityProcessingStrategy.REPLACE and config.replacement is None: raise EngineConfigurationError("acme replacement configuration is required") @@ -127,6 +128,40 @@ def test_builtin_registry_contains_the_builtin_regex_engine() -> None: ) +class _EmptyNERModel: + def predict_entities( + self, + text: str, + *, + labels: tuple[str, ...], + threshold: float, + flat_ner: bool, + timeout: Timeout, + ) -> tuple[NERModelEntity, ...]: + del text, labels, threshold, flat_ner, timeout + return () + + +def test_builtin_registry_adds_ner_only_with_explicit_resources() -> None: + resources = NERResources(model=_EmptyNERModel()) + + registry = create_builtin_registry(ner_resources=resources) + + assert tuple(item.engine_name for item in registry.describe_engines()) == ( + "regex", + "ner", + ) + schema = registry.configuration_json_schema() + definitions = schema["$defs"] + assert isinstance(definitions, dict) + assert "NEREngineConfig" in definitions + + +def test_ner_resources_require_builtins_to_be_enabled() -> None: + with pytest.raises(EngineRegistryError): + EngineRegistry(ner_resources=NERResources(model=_EmptyNERModel())) + + def test_registry_can_include_builtin_engines_before_custom_registration() -> None: registry = EngineRegistry(include_builtin_engines=True) registry.register(AcmeEngine, resources=AcmeResources(prefix="token")) diff --git a/projects/privacy-guard/tests/examples/test_ner_engine.py b/projects/privacy-guard/tests/examples/test_ner_engine.py new file mode 100644 index 0000000..7514313 --- /dev/null +++ b/projects/privacy-guard/tests/examples/test_ner_engine.py @@ -0,0 +1,139 @@ +"""End-to-end checks for the built-in NER engine example.""" + +from __future__ import annotations + +import asyncio +import json +import os +import subprocess +import sys +from pathlib import Path + +import yaml +from google.protobuf import json_format + +from privacy_guard.bindings import supervisor_middleware_pb2 as pb2 +from privacy_guard.engines import NERModelEntity, NERResources +from privacy_guard.engines.registry import create_builtin_registry +from privacy_guard.service.servicer import PrivacyGuardMiddleware +from privacy_guard.timeout import Timeout + +EXAMPLE_DIRECTORY = Path(__file__).parents[2] / "examples" / "ner-engine" + + +class _ExampleNERModel: + def predict_entities( + self, + text: str, + *, + labels: tuple[str, ...], + threshold: float, + flat_ner: bool, + timeout: Timeout, + ) -> tuple[NERModelEntity, ...]: + del text, labels, threshold, flat_ner, timeout + return ( + NERModelEntity( + label="EMAIL", + start=8, + end=24, + score=0.9, + ), + ) + + +def test_ner_example_runs_through_the_middleware_boundary() -> None: + values = yaml.safe_load( + (EXAMPLE_DIRECTORY / "privacy-guard-config.yaml").read_text() + ) + assert isinstance(values, dict) + config = pb2.HttpRequestEvaluation().config + json_format.ParseDict(values, config) + + async def evaluate() -> None: + middleware = PrivacyGuardMiddleware( + create_builtin_registry( + ner_resources=NERResources(model=_ExampleNERModel()) + ) + ) + try: + result = await middleware._evaluate_http_request( + pb2.HttpRequestEvaluation( + phase=pb2.SUPERVISOR_MIDDLEWARE_PHASE_PRE_CREDENTIALS, + config=config, + body=b"Contact user@example.com", + ) + ) + finally: + await middleware.close() + + assert result.decision == pb2.DECISION_ALLOW + assert result.has_body is False + assert {finding.label for finding in result.findings} == { + "email (general-entities)" + } + assert all(finding.confidence == "" for finding in result.findings) + + asyncio.run(evaluate()) + + +def test_endpoint_registry_drives_cli_discovery_and_schema() -> None: + command = str(Path(sys.executable).with_name("privacy-guard")) + environment = os.environ.copy() + environment.update( + { + "PYTHONPATH": str(EXAMPLE_DIRECTORY), + "PRIVACY_GUARD_NER_ENDPOINT": "http://model-host:8002/v1/extract", + "PRIVACY_GUARD_NER_MODEL": "operator-model", + } + ) + common = [ + command, + "--registry-factory", + "endpoint_registry:create_registry", + ] + + engines = subprocess.run( + [*common, "engines"], + cwd=EXAMPLE_DIRECTORY, + env=environment, + check=True, + capture_output=True, + text=True, + ) + schema = subprocess.run( + [*common, "configuration-schema"], + cwd=EXAMPLE_DIRECTORY, + env=environment, + check=True, + capture_output=True, + text=True, + ) + + assert "regex\tdetect,replace\t" in engines.stdout + assert "ner\tdetect,replace\t" in engines.stdout + serialized_schema = json.loads(schema.stdout) + assert "NEREngineConfig" in serialized_schema["$defs"] + assert "NERReplacement" in serialized_schema["$defs"] + + +def test_ner_walkthrough_uses_current_policy_and_registry_contract() -> None: + policy = yaml.safe_load((EXAMPLE_DIRECTORY / "policy.yaml").read_text()) + config = yaml.safe_load( + (EXAMPLE_DIRECTORY / "privacy-guard-config.yaml").read_text() + ) + readme = (EXAMPLE_DIRECTORY / "README.md").read_text() + + assert isinstance(policy, dict) + assert isinstance(config, dict) + middleware_config = policy["network_middlewares"]["privacy_guard_ner"] + assert middleware_config["middleware"] == "privacy-guard-ner" + assert middleware_config["config"] == config + stage_config = config["entity_processing"]["stages"][0]["config"] + assert stage_config["engine"] == "ner" + assert stage_config["threshold"] == 0.5 + assert stage_config["labels"] == ["person", "email", "phone_number"] + assert "endpoint_registry:create_registry configuration-schema" in readme + assert "local_registry:create_registry serve" in readme + assert 'uv pip install "gliner==0.2.27"' in readme + assert "/v1/chat/completions" in readme diff --git a/projects/privacy-guard/tests/test_request_processor.py b/projects/privacy-guard/tests/test_request_processor.py index be5a22c..34c25d1 100644 --- a/projects/privacy-guard/tests/test_request_processor.py +++ b/projects/privacy-guard/tests/test_request_processor.py @@ -10,7 +10,11 @@ from privacy_guard.config import PolicyAction from privacy_guard.constants import MAX_BODY_BYTES -from privacy_guard.engines import RegexEngine +from privacy_guard.engines import ( + NERModelEntity, + NERResources, + RegexEngine, +) from privacy_guard.engines.registry import EngineRegistry from privacy_guard.errors import ( EngineConfigurationError, @@ -23,6 +27,30 @@ from privacy_guard.timeout import Timeout +class _MarkerNERModel: + def predict_entities( + self, + text: str, + *, + labels: tuple[str, ...], + threshold: float, + flat_ner: bool, + timeout: Timeout, + ) -> tuple[NERModelEntity, ...]: + del labels, threshold, flat_ner, timeout + start = text.find("[person]") + if start < 0: + return () + return ( + NERModelEntity( + label="marker", + start=start, + end=start + len("[person]"), + score=0.9, + ), + ) + + def _values( action: PolicyAction, *, @@ -119,6 +147,74 @@ def test_replace_runs_stages_sequentially_over_the_current_text() -> None: ) +def test_replace_runs_mixed_regex_and_ner_stages_sequentially() -> None: + registry = EngineRegistry( + include_builtin_engines=True, + ner_resources=NERResources(model=_MarkerNERModel()), + ).finalize() + values: dict[str, object] = { + "entity_processing": { + "stages": [ + { + "name": "people", + "config": { + "engine": "regex", + "pattern_catalog": { + "entities": [ + { + "name": "person", + "rules": [ + { + "pattern": "Alice", + "confidence": "high", + } + ], + } + ] + }, + "replacement": { + "strategy": "template", + "template": "[{entity}]", + }, + }, + }, + { + "name": "model-marker", + "config": { + "engine": "ner", + "labels": ["marker"], + "threshold": 0.5, + "overlap_mode": "nested", + "replacement": { + "strategy": "template", + "template": "<{entity}>", + }, + }, + }, + ] + }, + "on_detection": {"action": "replace"}, + } + config = registry.validate_config(values) + prepared = tuple( + ( + stage.diagnostic_name(index), + registry.create_engine(stage.config), + ) + for index, stage in enumerate(config.entity_processing.stages, start=1) + ) + + result = RequestProcessor(config, prepared).process("Hello Alice") + + assert result.replacement_text == "Hello " + assert tuple( + (item.entity, item.source_stage) for item in result.detection_summaries + ) == ( + ("person", "people"), + ("marker", "model-marker"), + ) + + def test_detect_reports_without_returning_replacement_text() -> None: result = _processor(PolicyAction.DETECT).process("Hello Alice")