backend/ai/learning/postgresql_memory_store.py

"""
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