Spaces:
Sleeping
Sleeping
| # guardrails/pii_output_guard.py | |
| from typing import Generator, Dict, Any, Tuple | |
| from presidio_analyzer import AnalyzerEngine | |
| from presidio_anonymizer import AnonymizerEngine | |
| class PiiOutputGuard: | |
| """ | |
| A specialized PII guard focused specifically on output processing. | |
| This version includes enhanced features for output testing and monitoring. | |
| """ | |
| def __init__(self, config: Dict[str, Any]): | |
| """Initializes the PiiOutputGuard with a given configuration.""" | |
| self.config = config | |
| self.analyzer = AnalyzerEngine() | |
| self.anonymizer = AnonymizerEngine() | |
| print("✅ PII Output Guard initialized.") | |
| def process_input(self, prompt: str) -> Tuple[str, bool]: | |
| """ | |
| This guard is output-focused, so input processing is minimal. | |
| """ | |
| if not self.config.get("on_input", False): | |
| return prompt, True | |
| # Simple input processing if enabled | |
| analyzer_results = self.analyzer.analyze( | |
| text=prompt, | |
| language="en", | |
| entities=self.config.get("anonymize_entities", []), | |
| ) | |
| if analyzer_results: | |
| pii_types = {res.entity_type for res in analyzer_results} | |
| print(f" ⚠️ Input contains PII: {', '.join(pii_types)}") | |
| return prompt, True # Don't block input in output-focused guard | |
| def process_output_stream( | |
| self, text_stream: Generator[str, None, None] | |
| ) -> Generator[str, None, None]: | |
| """Enhanced PII detection and handling for output streams.""" | |
| if not self.config.get("on_output", True): | |
| yield from text_stream | |
| return | |
| accumulated_text = "" | |
| pii_found = False | |
| for chunk in text_stream: | |
| accumulated_text += chunk | |
| # Analyze the accumulated text for PII | |
| analyzer_results = self.analyzer.analyze( | |
| text=accumulated_text, | |
| language="en", | |
| entities=self.config.get("anonymize_entities", []), | |
| ) | |
| if analyzer_results and not pii_found: | |
| pii_found = True | |
| pii_types = {res.entity_type for res in analyzer_results} | |
| print(f"\n 🔍 PII detected in output: {', '.join(pii_types)}") | |
| # Apply anonymization to the accumulated text | |
| if analyzer_results: | |
| action = self.config.get("output_action", "anonymize") | |
| if action == "block": | |
| pii_types = {res.entity_type for res in analyzer_results} | |
| yield f"\n\n🔒 [OUTPUT BLOCKED: PII detected - {', '.join(pii_types)}]" | |
| return | |
| elif action == "anonymize": | |
| anonymized_result = self.anonymizer.anonymize( | |
| text=accumulated_text, | |
| analyzer_results=analyzer_results, | |
| ) | |
| # Calculate the new chunk based on the difference | |
| anonymized_text = anonymized_result.text | |
| if len(anonymized_text) >= len(accumulated_text) - len(chunk): | |
| new_chunk = anonymized_text[len(accumulated_text) - len(chunk):] | |
| yield new_chunk | |
| accumulated_text = anonymized_text | |
| else: | |
| yield chunk | |
| else: | |
| yield chunk | |
| def process_complete_output(self, text: str) -> Tuple[str, bool]: | |
| """ | |
| Process a complete output text for PII (non-streaming). | |
| This is specifically designed for output testing. | |
| """ | |
| analyzer_results = self.analyzer.analyze( | |
| text=text, | |
| language="en", | |
| entities=self.config.get("anonymize_entities", []), | |
| ) | |
| if not analyzer_results: | |
| return text, True # No PII found | |
| pii_types = {res.entity_type for res in analyzer_results} | |
| print(f"🔍 PII Analysis Results:") | |
| for result in analyzer_results: | |
| print(f" - {result.entity_type}: '{text[result.start:result.end]}' (confidence: {result.score:.2f})") | |
| action = self.config.get("output_action", "anonymize") | |
| if action == "block": | |
| return f"Output blocked: PII detected ({', '.join(pii_types)}).", False | |
| elif action == "anonymize": | |
| anonymized_result = self.anonymizer.anonymize( | |
| text=text, | |
| analyzer_results=analyzer_results, | |
| ) | |
| print(f"✅ PII anonymized in output") | |
| return anonymized_result.text, True | |
| # Default to anonymization for unknown actions | |
| anonymized_result = self.anonymizer.anonymize( | |
| text=text, | |
| analyzer_results=analyzer_results, | |
| ) | |
| return anonymized_result.text, True |