safety_wrapper.py

"""
Safety wrapper for Helion-V2.0-Thinking
Implements content filtering, rate limiting, and safety checks
"""

import torch
from transformers import AutoModelForCausalLM, AutoProcessor
from typing import Optional, List, Dict, Any, Union
from PIL import Image
import json
import re
import time
from collections import defaultdict, deque
from datetime import datetime, timedelta
import hashlib


class SafetyViolation(Exception):
    """Exception raised when safety policies are violated"""
    pass


class ContentFilter:
    """Content filtering for inputs and outputs"""
    
    # Harmful patterns to detect
    HARMFUL_PATTERNS = [
        r'(?i)how\s+to\s+(hack|crack|break\s+into)',
        r'(?i)make\s+(explosive|bomb|weapon)',
        r'(?i)(kill|murder|harm)\s+(myself|someone|people)',
        r'(?i)credit\s+card\s+number',
        r'(?i)social\s+security\s+number',
        r'(?i)(steal|fraud|scam)\s+',
        r'(?i)illegal\s+(drugs|substances)',
        r'(?i)child\s+(abuse|exploitation)',
    ]
    
    # PII patterns
    PII_PATTERNS = {
        'email': r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b',
        'phone': r'\b\d{3}[-.]?\d{3}[-.]?\d{4}\b',
        'ssn': r'\b\d{3}-\d{2}-\d{4}\b',
        'credit_card': r'\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b',
    }
    
    # Toxic keywords
    TOXIC_KEYWORDS = [
        'hate', 'violence', 'threat', 'abuse', 'harass',
        'discriminate', 'racist', 'sexist', 'offensive'
    ]
    
    @staticmethod
    def check_harmful_content(text: str) -> tuple[bool, Optional[str]]:
        """
        Check if text contains harmful content
        
        Returns:
            (is_harmful, reason)
        """
        for pattern in ContentFilter.HARMFUL_PATTERNS:
            if re.search(pattern, text):
                return True, f"Matched harmful pattern: {pattern}"
        
        return False, None
    
    @staticmethod
    def check_pii(text: str) -> tuple[bool, List[str]]:
        """
        Check for personally identifiable information
        
        Returns:
            (contains_pii, pii_types)
        """
        found_pii = []
        for pii_type, pattern in ContentFilter.PII_PATTERNS.items():
            if re.search(pattern, text):
                found_pii.append(pii_type)
        
        return len(found_pii) > 0, found_pii
    
    @staticmethod
    def check_toxicity(text: str) -> tuple[float, List[str]]:
        """
        Simple toxicity check based on keywords
        
        Returns:
            (toxicity_score, matched_keywords)
        """
        text_lower = text.lower()
        matched = [kw for kw in ContentFilter.TOXIC_KEYWORDS if kw in text_lower]
        score = len(matched) / len(ContentFilter.TOXIC_KEYWORDS)
        
        return score, matched
    
    @staticmethod
    def redact_pii(text: str) -> str:
        """Redact PII from text"""
        redacted = text
        
        for pii_type, pattern in ContentFilter.PII_PATTERNS.items():
            redacted = re.sub(pattern, f'[REDACTED_{pii_type.upper()}]', redacted)
        
        return redacted


class RateLimiter:
    """Rate limiting for API usage"""
    
    def __init__(
        self,
        requests_per_minute: int = 60,
        tokens_per_minute: int = 90000,
        concurrent_limit: int = 10
    ):
        self.requests_per_minute = requests_per_minute
        self.tokens_per_minute = tokens_per_minute
        self.concurrent_limit = concurrent_limit
        
        self.request_times = defaultdict(deque)
        self.token_counts = defaultdict(deque)
        self.active_requests = defaultdict(int)
    
    def check_rate_limit(self, user_id: str, estimated_tokens: int = 0) -> tuple[bool, Optional[str]]:
        """
        Check if request is within rate limits
        
        Returns:
            (allowed, reason_if_denied)
        """
        now = datetime.now()
        minute_ago = now - timedelta(minutes=1)
        
        # Clean old entries
        while self.request_times[user_id] and self.request_times[user_id][0] < minute_ago:
            self.request_times[user_id].popleft()
        
        while self.token_counts[user_id] and self.token_counts[user_id][0][0] < minute_ago:
            self.token_counts[user_id].popleft()
        
        # Check requests per minute
        if len(self.request_times[user_id]) >= self.requests_per_minute:
            return False, f"Rate limit exceeded: {self.requests_per_minute} requests per minute"
        
        # Check tokens per minute
        total_tokens = sum(t[1] for t in self.token_counts[user_id])
        if total_tokens + estimated_tokens > self.tokens_per_minute:
            return False, f"Token limit exceeded: {self.tokens_per_minute} tokens per minute"
        
        # Check concurrent requests
        if self.active_requests[user_id] >= self.concurrent_limit:
            return False, f"Concurrent request limit exceeded: {self.concurrent_limit}"
        
        return True, None
    
    def record_request(self, user_id: str, tokens: int = 0):
        """Record a request"""
        now = datetime.now()
        self.request_times[user_id].append(now)
        self.token_counts[user_id].append((now, tokens))
        self.active_requests[user_id] += 1
    
    def release_request(self, user_id: str):
        """Release an active request slot"""
        if self.active_requests[user_id] > 0:
            self.active_requests[user_id] -= 1


class SafeHelionWrapper:
    """Safety wrapper for Helion-V2.0-Thinking"""
    
    def __init__(
        self,
        model_name: str = "DeepXR/Helion-V2.0-Thinking",
        safety_config_path: Optional[str] = None,
        enable_safety: bool = True,
        enable_rate_limiting: bool = True,
        device: str = "auto"
    ):
        """
        Initialize safe wrapper
        
        Args:
            model_name: Model name or path
            safety_config_path: Path to safety_config.json
            enable_safety: Enable safety checks
            enable_rate_limiting: Enable rate limiting
            device: Device for model
        """
        print(f"Loading model with safety wrapper: {model_name}")
        
        # Load safety config
        if safety_config_path:
            with open(safety_config_path, 'r') as f:
                self.safety_config = json.load(f)
        else:
            self.safety_config = self._default_safety_config()
        
        self.enable_safety = enable_safety
        self.enable_rate_limiting = enable_rate_limiting
        
        # Initialize components
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.bfloat16,
            device_map=device,
            trust_remote_code=True
        )
        self.processor = AutoProcessor.from_pretrained(model_name)
        self.model.eval()
        
        self.content_filter = ContentFilter()
        self.rate_limiter = RateLimiter(
            requests_per_minute=self.safety_config['safety_settings']['rate_limiting']['requests_per_minute'],
            tokens_per_minute=self.safety_config['safety_settings']['rate_limiting']['tokens_per_minute'],
            concurrent_requests=self.safety_config['safety_settings']['rate_limiting']['concurrent_requests']
        )
        
        # Violation tracking
        self.violation_log = []
        
        print("Safety wrapper initialized successfully")
    
    def _default_safety_config(self) -> Dict[str, Any]:
        """Default safety configuration"""
        return {
            "safety_settings": {
                "rate_limiting": {
                    "requests_per_minute": 60,
                    "tokens_per_minute": 90000,
                    "concurrent_requests": 10
                },
                "content_filtering": {
                    "profanity_filter": {"enabled": True},
                    "pii_detection": {"enabled": True},
                    "toxicity_detection": {"enabled": True, "threshold": 0.7}
                }
            }
        }
    
    def _validate_input(
        self,
        prompt: str,
        images: Optional[List[Image.Image]] = None,
        user_id: str = "default"
    ):
        """Validate input against safety policies"""
        if not self.enable_safety:
            return
        
        # Check harmful content
        is_harmful, reason = self.content_filter.check_harmful_content(prompt)
        if is_harmful:
            self._log_violation(user_id, "harmful_content", reason)
            raise SafetyViolation(f"Input rejected: {reason}")
        
        # Check PII
        if self.safety_config['safety_settings']['content_filtering']['pii_detection']['enabled']:
            has_pii, pii_types = self.content_filter.check_pii(prompt)
            if has_pii:
                self._log_violation(user_id, "pii_detected", f"Types: {pii_types}")
                print(f"Warning: PII detected in input: {pii_types}")
        
        # Check toxicity
        if self.safety_config['safety_settings']['content_filtering']['toxicity_detection']['enabled']:
            toxicity_score, keywords = self.content_filter.check_toxicity(prompt)
            threshold = self.safety_config['safety_settings']['content_filtering']['toxicity_detection']['threshold']
            
            if toxicity_score > threshold:
                self._log_violation(user_id, "high_toxicity", f"Score: {toxicity_score}")
                raise SafetyViolation(f"Input rejected: High toxicity score ({toxicity_score:.2f})")
        
        # Validate images
        if images:
            max_images = self.safety_config.get('guardrails', {}).get('input_validation', {}).get('max_images_per_request', 10)
            if len(images) > max_images:
                raise SafetyViolation(f"Too many images: {len(images)} (max: {max_images})")
    
    def _validate_output(self, output: str, user_id: str = "default"):
        """Validate output against safety policies"""
        if not self.enable_safety:
            return output
        
        # Check for harmful content in output
        is_harmful, reason = self.content_filter.check_harmful_content(output)
        if is_harmful:
            self._log_violation(user_id, "harmful_output", reason)
            return "I cannot provide that information as it may be harmful."
        
        # Redact PII if found
        if self.safety_config['safety_settings']['content_filtering']['pii_detection']['enabled']:
            output = self.content_filter.redact_pii(output)
        
        return output
    
    def _log_violation(self, user_id: str, violation_type: str, details: str):
        """Log safety violation"""
        self.violation_log.append({
            "timestamp": datetime.now().isoformat(),
            "user_id": hashlib.sha256(user_id.encode()).hexdigest()[:16],
            "type": violation_type,
            "details": details
        })
        
        # Keep only last 1000 violations
        if len(self.violation_log) > 1000:
            self.violation_log = self.violation_log[-1000:]
    
    def generate(
        self,
        prompt: str,
        images: Optional[List[Image.Image]] = None,
        user_id: str = "default",
        max_new_tokens: int = 512,
        temperature: float = 0.7,
        **kwargs
    ) -> str:
        """
        Safe generation with input/output filtering
        
        Args:
            prompt: Input prompt
            images: Optional list of images
            user_id: User identifier for rate limiting
            max_new_tokens: Maximum tokens to generate
            temperature: Sampling temperature
            **kwargs: Additional generation parameters
        
        Returns:
            Generated text (filtered)
        """
        # Check rate limit
        if self.enable_rate_limiting:
            allowed, reason = self.rate_limiter.check_rate_limit(user_id, max_new_tokens)
            if not allowed:
                raise SafetyViolation(reason)
            
            self.rate_limiter.record_request(user_id, max_new_tokens)
        
        try:
            # Validate input
            self._validate_input(prompt, images, user_id)
            
            # Generate
            if images:
                inputs = self.processor(text=prompt, images=images, return_tensors="pt").to(self.model.device)
            else:
                inputs = self.processor(text=prompt, return_tensors="pt").to(self.model.device)
            
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=max_new_tokens,
                    temperature=temperature,
                    **kwargs
                )
            
            response = self.processor.decode(outputs[0], skip_special_tokens=True)
            
            # Remove prompt from response
            if response.startswith(prompt):
                response = response[len(prompt):].strip()
            
            # Validate output
            response = self._validate_output(response, user_id)
            
            return response
        
        finally:
            if self.enable_rate_limiting:
                self.rate_limiter.release_request(user_id)
    
    def get_violation_stats(self) -> Dict[str, Any]:
        """Get violation statistics"""
        if not self.violation_log:
            return {"total_violations": 0}
        
        violation_types = defaultdict(int)
        for log in self.violation_log:
            violation_types[log['type']] += 1
        
        return {
            "total_violations": len(self.violation_log),
            "by_type": dict(violation_types),
            "recent_violations": self.violation_log[-10:]
        }
    
    def export_violation_log(self, filename: str = "violations.json"):
        """Export violation log to file"""
        with open(filename, 'w') as f:
            json.dump(self.violation_log, f, indent=2)
        print(f"Violation log exported to {filename}")


def main():
    """Example usage of safe wrapper"""
    # Initialize safe wrapper
    wrapper = SafeHelionWrapper(
        model_name="DeepXR/Helion-V2.0-Thinking",
        enable_safety=True,
        enable_rate_limiting=True
    )
    
    # Test cases
    test_prompts = [
        "Explain how photosynthesis works.",  # Safe
        "Write a poem about nature.",  # Safe
        "How do I hack into an email account?",  # Should be blocked
    ]
    
    print("\n" + "="*60)
    print("Testing Safety Wrapper")
    print("="*60 + "\n")
    
    for prompt in test_prompts:
        print(f"Prompt: {prompt}")
        try:
            response = wrapper.generate(prompt, max_new_tokens=128)
            print(f"Response: {response}\n")
        except SafetyViolation as e:
            print(f"BLOCKED: {e}\n")
        except Exception as e:
            print(f"ERROR: {e}\n")
    
    # Print violation stats
    print("="*60)
    print("Violation Statistics")
    print("="*60)
    stats = wrapper.get_violation_stats()
    print(json.dumps(stats, indent=2))


if __name__ == "__main__":
    main()