""" PostgreSQL-compatible Memory Store for Learning Engine This replaces the SQLite-based memory store with PostgreSQL to maintain database consistency in production. """ import logging from typing import List, Dict, Any, Optional, Tuple, Union from datetime import datetime, timedelta from dataclasses import dataclass, asdict from pathlib import Path import json logger = logging.getLogger(__name__) @dataclass class AttackRecord: """Single attack execution record.""" attack_id: str attack_type: str attack_category: str target_model: str dataset: str prompt: str success: bool safety_score: float risk_level: str response_text: Optional[str] = None response_length: Optional[int] = None inference_time_ms: Optional[float] = None timestamp: datetime = None metadata: Dict[str, Any] = None def __post_init__(self): if self.timestamp is None: self.timestamp = datetime.utcnow() if self.metadata is None: self.metadata = {} @dataclass class PatternMetrics: """Pattern analysis metrics.""" attack_type: str success_rate: float avg_safety_score: float total_attempts: int successful_attempts: int failure_patterns: List[str] success_patterns: List[str] risk_distribution: Dict[str, int] last_updated: datetime = None def __post_init__(self): if self.last_updated is None: self.last_updated = datetime.utcnow() class PostgreSQLMemoryStore: """PostgreSQL-based memory store for learning data.""" def __init__(self, db_session=None, enable_persistence: bool = True): """ Initialize PostgreSQL memory store. Args: db_session: Async database session enable_persistence: Whether to enable database persistence """ self.db_session = db_session self.enable_persistence = enable_persistence # In-memory cache for frequent access self._attack_cache: List[AttackRecord] = [] self._pattern_cache: Dict[str, PatternMetrics] = {} logger.info("PostgreSQL Memory store initialized") async def store_attack(self, attack: AttackRecord) -> bool: """Store attack record in PostgreSQL.""" if not self.enable_persistence or not self.db_session: # Store in memory only self._attack_cache.append(attack) return True try: from sqlalchemy import text # Insert into PostgreSQL query = text(""" INSERT INTO learning_attacks ( attack_id, attack_type, attack_category, target_model, dataset, prompt, success, safety_score, risk_level, response_text, response_length, inference_time_ms, timestamp, metadata ) VALUES ( :attack_id, :attack_type, :attack_category, :target_model, :dataset, :prompt, :success, :safety_score, :risk_level, :response_text, :response_length, :inference_time_ms, :timestamp, :metadata ) ON CONFLICT (attack_id) DO UPDATE SET success = EXCLUDED.success, safety_score = EXCLUDED.safety_score, response_text = EXCLUDED.response_text, timestamp = EXCLUDED.timestamp """) await self.db_session.execute(query, { 'attack_id': attack.attack_id, 'attack_type': attack.attack_type, 'attack_category': attack.attack_category, 'target_model': attack.target_model, 'dataset': attack.dataset, 'prompt': attack.prompt, 'success': attack.success, 'safety_score': attack.safety_score, 'risk_level': attack.risk_level, 'response_text': attack.response_text, 'response_length': attack.response_length, 'inference_time_ms': attack.inference_time_ms, 'timestamp': attack.timestamp, 'metadata': json.dumps(attack.metadata) }) # Also store in cache self._attack_cache.append(attack) return True except Exception as e: logger.error(f"Failed to store attack in PostgreSQL: {e}") # Fallback to memory cache self._attack_cache.append(attack) return False async def get_recent_attacks(self, limit: int = 100, attack_type: str = None) -> List[AttackRecord]: """Get recent attacks from PostgreSQL or cache.""" if not self.enable_persistence or not self.db_session: # Return from cache attacks = self._attack_cache if attack_type: attacks = [a for a in attacks if a.attack_type == attack_type] return attacks[-limit:] try: from sqlalchemy import text query = text(""" SELECT attack_id, attack_type, attack_category, target_model, dataset, prompt, success, safety_score, risk_level, response_text, response_length, response_time_ms, timestamp, metadata FROM learning_attacks WHERE (:attack_type IS NULL OR attack_type = :attack_type) ORDER BY timestamp DESC LIMIT :limit """) result = await self.db_session.execute(query, { 'attack_type': attack_type, 'limit': limit }) attacks = [] for row in result: attack = AttackRecord( attack_id=row[0], attack_type=row[1], attack_category=row[2], target_model=row[3], dataset=row[4], prompt=row[5], success=row[6], safety_score=row[7], risk_level=row[8], response_text=row[9], response_length=row[10], response_time_ms=row[11], timestamp=row[12], metadata=json.loads(row[13]) if row[13] else {} ) attacks.append(attack) return attacks except Exception as e: logger.error(f"Failed to get attacks from PostgreSQL: {e}") # Fallback to cache attacks = self._attack_cache if attack_type: attacks = [a for a in attacks if a.attack_type == attack_type] return attacks[-limit:] async def get_pattern_metrics(self, attack_type: Optional[str] = None) -> Union[PatternMetrics, Dict[str, PatternMetrics]]: """Get pattern metrics from PostgreSQL or cache.""" if not self.enable_persistence or not self.db_session: # Return from cache if attack_type: return self._pattern_cache.get(attack_type, PatternMetrics( attack_type=attack_type, success_rate=0.0, avg_safety_score=0.0, total_attempts=0, successful_attempts=0, failure_patterns=[], success_patterns=[], risk_distribution={} )) return self._pattern_cache try: from sqlalchemy import text query = text(""" SELECT attack_type, COUNT(*) as total_attempts, SUM(CASE WHEN success THEN 1 ELSE 0 END) as successful_attempts, AVG(safety_score) as avg_safety_score FROM learning_attacks WHERE (:attack_type IS NULL OR attack_type = :attack_type) GROUP BY attack_type """) result = await self.db_session.execute(query, {'attack_type': attack_type}) if attack_type: row = result.fetchone() if row: return PatternMetrics( attack_type=row[0], success_rate=row[2] / row[1] if row[1] > 0 else 0.0, avg_safety_score=row[3] or 0.0, total_attempts=row[1], successful_attempts=row[2], failure_patterns=[], success_patterns=[], risk_distribution={} ) else: return PatternMetrics( attack_type=attack_type, success_rate=0.0, avg_safety_score=0.0, total_attempts=0, successful_attempts=0, failure_patterns=[], success_patterns=[], risk_distribution={} ) else: metrics = {} for row in result: metrics[row[0]] = PatternMetrics( attack_type=row[0], success_rate=row[2] / row[1] if row[1] > 0 else 0.0, avg_safety_score=row[3] or 0.0, total_attempts=row[1], successful_attempts=row[2], failure_patterns=[], success_patterns=[], risk_distribution={} ) return metrics except Exception as e: logger.error(f"Failed to get pattern metrics from PostgreSQL: {e}") # Fallback to cache if attack_type: return self._pattern_cache.get(attack_type, PatternMetrics( attack_type=attack_type, success_rate=0.0, avg_safety_score=0.0, total_attempts=0, successful_attempts=0, failure_patterns=[], success_patterns=[], risk_distribution={} )) return self._pattern_cache # Global instance _postgresql_memory_store_instance = None def get_postgresql_memory_store(db_session=None) -> PostgreSQLMemoryStore: """Get PostgreSQL memory store instance.""" global _postgresql_memory_store_instance if _postgresql_memory_store_instance is None: _postgresql_memory_store_instance = PostgreSQLMemoryStore(db_session) return _postgresql_memory_store_instance