[Multi-Agent Privacy] Detection tools implementation

pyproject.toml

@@ -136,6 +136,17 @@ privacy = [
     "mmore[llm]",
     "langgraph>=1.0",
     "langgraph-checkpoint-sqlite>=2.0",
+    "gliner",
+    "presidio-analyzer",
+    "presidio-anonymizer",
+    "spacy",
+    "dspy>=2.6",
+]
+
+privacy-openai-filter = [
+    "mmore[privacy]",
+    "transformers>=5",
+    "peft>=0.18",
 ]
 
 # --- Composite + variant extras ---
@@ -173,6 +184,16 @@ conflicts = [
     { extra = "cpu" },
     { extra = "cu126" },
   ],
+  [
+    # privacy-openai-filter needs transformers>=5 but has conflicts with 
+    # marker-pdf in `process`/`all` (transformers<5)
+    { extra = "privacy-openai-filter" },
+    { extra = "process" },
+  ],
+  [
+    { extra = "privacy-openai-filter" },
+    { extra = "all" },
+  ],
 ]
 
 override-dependencies = ["anthropic==0.73.0"]

src/mmore/privacy/detection/__init__.py

@@ -0,0 +1,22 @@
+"""PII detection engines exposed as agent tools."""
+
+from .base import DetectionEngine, PIISpan
+from .config import DetectionConfig
+from .gliner_engine import GLiNEREngine, detect_pii_gliner
+from .llm_engine import LLMDetectionEngine, detect_pii_llm
+from .openai_filter_engine import OpenAIFilterEngine, detect_pii_openai_filter
+from .presidio_engine import PresidioEngine, detect_pii_presidio
+
+__all__ = [
+    "DetectionConfig",
+    "DetectionEngine",
+    "GLiNEREngine",
+    "LLMDetectionEngine",
+    "OpenAIFilterEngine",
+    "PIISpan",
+    "PresidioEngine",
+    "detect_pii_gliner",
+    "detect_pii_llm",
+    "detect_pii_openai_filter",
+    "detect_pii_presidio",
+]

src/mmore/privacy/detection/base.py

@@ -0,0 +1,28 @@
+"""PII detection interface.
+
+Each engine implements ``DetectionEngine.detect`` and returns a list of
+``PIISpan`` records. Engines are independently registered as agent tools so a
+sanitizer agent can resolve them by name from YAML.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import List
+
+
+@dataclass
+class PIISpan:
+    """A single detected PII occurrence in some text."""
+
+    start: int
+    end: int
+    label: str
+    score: float
+
+
+class DetectionEngine(ABC):
+    """Abstract base for PII detection backends."""
+
+    @abstractmethod
+    def detect(self, text: str) -> List[PIISpan]:
+        """Return all PII spans found in ``text``."""

src/mmore/privacy/detection/config.py

@@ -0,0 +1,17 @@
+"""Configuration dataclass for the PII detection toolkit."""
+
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+from ...rag.llm import LLMConfig
+from .defaults import DEFAULT_CONFIDENCE_THRESHOLD
+
+
+@dataclass
+class DetectionConfig:
+    """Schema for the ``privacy.detection`` block of a YAML config."""
+
+    engine: str
+    entity_types: List[str] = field(default_factory=list)
+    confidence_threshold: float = DEFAULT_CONFIDENCE_THRESHOLD
+    llm: Optional[LLMConfig] = None

src/mmore/privacy/detection/defaults.py

@@ -0,0 +1,55 @@
+"""Shared defaults for the PII detection engines."""
+
+from ...rag.llm import LLMConfig
+
+DEFAULT_LANGUAGE = "en"
+
+DEFAULT_GLINER_MODEL = "nvidia/gliner-PII"
+DEFAULT_OPENAI_FILTER_MODEL = "openai/privacy-filter"
+DEFAULT_PRESIDIO_SPACY_MODEL = "en_core_web_lg"
+
+DEFAULT_LLM_CONFIG = LLMConfig(
+    llm_name="Qwen/Qwen2.5-3B-Instruct",
+    max_new_tokens=512,
+)
+
+DEFAULT_CONFIDENCE_THRESHOLD = 0.7
+
+# TODO: Later add new labels to the list
+DEFAULT_LABELS = [
+    "PERSON",
+    "PHONE",
+    "EMAIL",
+    "MRN",
+    "DATE",
+    "LOCATION",
+    "SSN",
+    "INSURANCE_ID",
+]
+
+# TODO: Later add new patterns to the list
+PRESIDIO_CLINICAL_PATTERNS = [
+    {
+        "entity": "MRN",
+        "patterns": [
+            ("mrn_with_prefix", r"\bMRN[\s:#]*\d{6,10}\b", 0.9),
+            ("mrn_bare_8_digits", r"\b\d{8}\b", 0.4),
+        ],
+        "context": ["mrn", "medical record", "record number", "patient id"],
+    },
+    {
+        "entity": "HOSPITAL_DATE",
+        "patterns": [
+            ("iso_date", r"\b\d{4}-\d{2}-\d{2}\b", 0.6),
+            ("us_date", r"\b\d{1,2}/\d{1,2}/\d{4}\b", 0.6),
+        ],
+        "context": ["admission", "discharge", "appointment", "hospital", "clinic"],
+    },
+    {
+        "entity": "INSURANCE_ID",
+        "patterns": [
+            ("insurance_alnum", r"\b[A-Z]{2,3}\d{6,12}\b", 0.7),
+        ],
+        "context": ["insurance", "policy", "member id", "subscriber"],
+    },
+]

src/mmore/privacy/detection/gliner_engine.py

@@ -0,0 +1,109 @@
+"""GLiNER-based PII detection engine."""
+
+import logging
+import threading
+from typing import Any, Dict, List, Optional, Sequence
+
+from typing_extensions import Self
+
+from ..agents.registry import register_tool
+from .base import DetectionEngine, PIISpan
+from .config import DetectionConfig
+from .defaults import (
+    DEFAULT_CONFIDENCE_THRESHOLD,
+    DEFAULT_GLINER_MODEL,
+    DEFAULT_LABELS,
+)
+
+logger = logging.getLogger(__name__)
+
+_model_cache: Dict[str, Any] = {}
+_model_cache_lock = threading.Lock()
+
+
+def _load_gliner_model(model_name: str) -> Any:
+    from gliner import GLiNER
+
+    return GLiNER.from_pretrained(model_name)
+
+
+def _get_or_load_model(model_name: str) -> Any:
+    cached = _model_cache.get(model_name)
+    if cached is not None:
+        return cached
+    with _model_cache_lock:
+        cached = _model_cache.get(model_name)
+        if cached is None:
+            cached = _load_gliner_model(model_name)
+            _model_cache[model_name] = cached
+        return cached
+
+
+def clear_gliner_cache() -> None:
+    """Drop all cached GLiNER models."""
+    with _model_cache_lock:
+        _model_cache.clear()
+
+
+class GLiNEREngine(DetectionEngine):
+    """Detect PII spans with a GLiNER model.
+
+    Each instance carries its own ``entity_types`` and ``confidence_threshold``,
+    models with the same ``model_name`` are shared via ``_models_cache``.
+    """
+
+    def __init__(
+        self,
+        model_name: str = DEFAULT_GLINER_MODEL,
+        entity_types: Optional[Sequence[str]] = None,
+        confidence_threshold: float = DEFAULT_CONFIDENCE_THRESHOLD,
+    ):
+        self._model_name = model_name
+        self._entity_types: List[str] = (
+            list(entity_types) if entity_types else list(DEFAULT_LABELS)
+        )
+        self._confidence_threshold = confidence_threshold
+
+    @classmethod
+    def from_config(cls, config: DetectionConfig) -> Self:
+        """Build an engine from a ``DetectionConfig``."""
+        return cls(
+            entity_types=config.entity_types or None,
+            confidence_threshold=config.confidence_threshold,
+        )
+
+    @property
+    def model(self) -> Any:
+        """Lazy-load and cache the LLM on first access."""
+        return _get_or_load_model(self._model_name)
+
+    def detect(self, text: str) -> List[PIISpan]:
+        raw = self.model.predict_entities(
+            text=text,
+            labels=self._entity_types,
+            threshold=self._confidence_threshold,
+            multi_label=False,
+        )
+        return [
+            PIISpan(
+                start=int(r["start"]),
+                end=int(r["end"]),
+                label=str(r["label"]),
+                score=float(r["score"]),
+            )
+            for r in raw
+        ]
+
+
+@register_tool("detect_pii_gliner")
+def detect_pii_gliner(text: str) -> List[PIISpan]:
+    """Detect PII spans in ``text`` using a default-configured GLiNER engine.
+
+    Agents needing per-config behavior should be wired by setup code that
+    builds a ``GLiNEREngine.from_config(detection_cfg)`` and registers its
+    ``detect()`` function under a distinct tool name, e.g.::
+
+        engine = GLiNEREngine.from_config(detection_cfg)
+        register_tool("detect_pii_gliner_custom", engine.detect)
+    """
+    return GLiNEREngine().detect(text)