diff --git a/sentinelguard/pii/__init__.py b/sentinelguard/pii/__init__.py index 7cbcb8a..1383898 100644 --- a/sentinelguard/pii/__init__.py +++ b/sentinelguard/pii/__init__.py @@ -162,7 +162,10 @@ def anonymize(self, text: str, entities: List[PIIEntity]) -> AnonymizedResult: if not entities: return AnonymizedResult(text=text) - sorted_entities = sorted(entities, key=lambda e: e.start, reverse=True) + # Deduplicate overlapping entities — keep the one with higher score + deduped = self._remove_overlaps(entities) + + sorted_entities = sorted(deduped, key=lambda e: e.start, reverse=True) result_text = text items = [] mapping = {} @@ -183,6 +186,25 @@ def anonymize(self, text: str, entities: List[PIIEntity]) -> AnonymizedResult: items.reverse() return AnonymizedResult(text=result_text, items=items, mapping=mapping) + @staticmethod + def _remove_overlaps(entities: List[PIIEntity]) -> List[PIIEntity]: + """Remove overlapping entities, keeping the one with higher score.""" + if not entities: + return [] + # Sort by start position, then by score descending + sorted_ents = sorted(entities, key=lambda e: (e.start, -e.score)) + result = [sorted_ents[0]] + for ent in sorted_ents[1:]: + prev = result[-1] + if ent.start >= prev.end: + # No overlap + result.append(ent) + elif ent.score > prev.score: + # Overlaps but higher score — replace + result[-1] = ent + # else: overlaps with lower score — skip + return result + def _apply_strategy(self, entity: PIIEntity, strategy: str) -> str: if strategy == "replace": return f"<{entity.entity_type}>" diff --git a/sentinelguard/scanners/prompt/pii.py b/sentinelguard/scanners/prompt/pii.py index 29f141d..bc92b1d 100644 --- a/sentinelguard/scanners/prompt/pii.py +++ b/sentinelguard/scanners/prompt/pii.py @@ -9,12 +9,12 @@ from presidio_analyzer import AnalyzerEngine -from sentinelguard.core.scanner import PromptScanner, RiskLevel, ScanResult, register_scanner +from sentinelguard.core.scanner import BaseScanner, ScannerType, RiskLevel, ScanResult, register_scanner @register_scanner -class PIIScanner(PromptScanner): - """Detects personally identifiable information in prompts using Presidio. +class PIIScanner(BaseScanner): + """Detects personally identifiable information using Presidio. Covers 30+ entity types: EMAIL_ADDRESS, PHONE_NUMBER, CREDIT_CARD, US_SSN, IBAN_CODE, US_PASSPORT, IP_ADDRESS, PERSON, LOCATION, CRYPTO, @@ -24,10 +24,11 @@ class PIIScanner(PromptScanner): threshold: Confidence threshold (0.0-1.0). Default 0.5. entities: List of entity types to detect. ``None`` = all. language: Language for Presidio NLP engine. Default "en". - score_threshold: Minimum Presidio confidence score. Default 0.5. + score_threshold: Minimum Presidio confidence score. Default 0.3. """ scanner_name: ClassVar[str] = "pii" + scanner_type: ClassVar[ScannerType] = ScannerType.BOTH # Sensitivity weights per entity type for risk scoring ENTITY_SENSITIVITY: Dict[str, float] = { @@ -45,7 +46,7 @@ def __init__( threshold: float = 0.5, entities: Optional[List[str]] = None, language: str = "en", - score_threshold: float = 0.5, + score_threshold: float = 0.3, **kwargs: Any, ): super().__init__(threshold=threshold, **kwargs) diff --git a/sentinelguard/scanners/prompt/toxicity.py b/sentinelguard/scanners/prompt/toxicity.py index 2db1385..fd1939b 100644 --- a/sentinelguard/scanners/prompt/toxicity.py +++ b/sentinelguard/scanners/prompt/toxicity.py @@ -11,7 +11,7 @@ import re from typing import Any, ClassVar, Dict, List, Optional -from sentinelguard.core.scanner import PromptScanner, RiskLevel, ScanResult, register_scanner +from sentinelguard.core.scanner import BaseScanner, ScannerType, RiskLevel, ScanResult, register_scanner logger = logging.getLogger(__name__) @@ -57,7 +57,7 @@ @register_scanner -class ToxicityScanner(PromptScanner): +class ToxicityScanner(BaseScanner): """Detects toxic, hateful, or offensive content. Combines keyword/pattern matching (hate speech, threats, harassment, @@ -70,6 +70,7 @@ class ToxicityScanner(PromptScanner): """ scanner_name: ClassVar[str] = "toxicity" + scanner_type: ClassVar[ScannerType] = ScannerType.BOTH def __init__( self,