""" Online Learning System for Cyber-LLM Enables real-time model updates from operational feedback and new threat intelligence. Author: Muzan Sano """ import asyncio import json import logging from datetime import datetime, timedelta from typing import Dict, List, Optional, Tuple, Any from dataclasses import dataclass, asdict from enum import Enum import numpy as np import torch from torch.utils.data import DataLoader, Dataset from transformers import AutoTokenizer, AutoModelForCausalLM import redis from pydantic import BaseModel, Field from ..utils.logging_system import CyberLLMLogger from ..utils.secrets_manager import SecretsManager # Configure logging logger = CyberLLMLogger(__name__).get_logger() class LearningEventType(Enum): """Types of learning events that can trigger model updates""" FEEDBACK_POSITIVE = "feedback_positive" FEEDBACK_NEGATIVE = "feedback_negative" NEW_THREAT_INTELLIGENCE = "new_threat_intel" SECURITY_INCIDENT = "security_incident" AGENT_SUCCESS = "agent_success" AGENT_FAILURE = "agent_failure" OPSEC_VIOLATION = "opsec_violation" FALSE_POSITIVE = "false_positive" @dataclass class LearningEvent: """Structure for learning events""" event_id: str event_type: LearningEventType timestamp: datetime source: str # Which agent or system generated this event context: Dict[str, Any] # Relevant context for learning feedback_score: Optional[float] = None # Human feedback score (0-1) confidence: float = 1.0 # Confidence in this event priority: int = 1 # Priority level (1=low, 5=critical) def to_dict(self) -> Dict[str, Any]: data = asdict(self) data['timestamp'] = self.timestamp.isoformat() data['event_type'] = self.event_type.value return data class OnlineDataset(Dataset): """Dataset for online learning from streaming events""" def __init__(self, events: List[LearningEvent], tokenizer, max_length: int = 512): self.events = events self.tokenizer = tokenizer self.max_length = max_length def __len__(self) -> int: return len(self.events) def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: event = self.events[idx] # Convert learning event to training sample context_text = self._event_to_text(event) # Tokenize encoding = self.tokenizer( context_text, truncation=True, padding='max_length', max_length=self.max_length, return_tensors='pt' ) return { 'input_ids': encoding['input_ids'].squeeze(), 'attention_mask': encoding['attention_mask'].squeeze(), 'labels': encoding['input_ids'].squeeze(), 'event_weight': torch.tensor(event.confidence * event.priority, dtype=torch.float32) } def _event_to_text(self, event: LearningEvent) -> str: """Convert learning event to training text""" if event.event_type == LearningEventType.FEEDBACK_POSITIVE: return f"[POSITIVE_FEEDBACK] Context: {event.context.get('query', '')} Response: {event.context.get('response', '')} Score: {event.feedback_score}" elif event.event_type == LearningEventType.FEEDBACK_NEGATIVE: return f"[NEGATIVE_FEEDBACK] Context: {event.context.get('query', '')} Response: {event.context.get('response', '')} Score: {event.feedback_score}" elif event.event_type == LearningEventType.NEW_THREAT_INTELLIGENCE: return f"[THREAT_INTEL] {event.context.get('threat_description', '')} TTPs: {event.context.get('ttps', [])}" elif event.event_type == LearningEventType.SECURITY_INCIDENT: return f"[INCIDENT] {event.context.get('incident_description', '')} Response: {event.context.get('response_actions', [])}" else: return f"[{event.event_type.value.upper()}] {json.dumps(event.context)}" class OnlineLearningManager: """Manages online learning process for Cyber-LLM""" def __init__(self, model_name: str = "microsoft/DialoGPT-medium", redis_host: str = "localhost", redis_port: int = 6379, learning_rate: float = 1e-5, batch_size: int = 4, update_frequency: int = 100, # Update after N events max_events_memory: int = 10000): self.model_name = model_name self.learning_rate = learning_rate self.batch_size = batch_size self.update_frequency = update_frequency self.max_events_memory = max_events_memory # Initialize components self.secrets_manager = SecretsManager() self.redis_client = redis.Redis(host=redis_host, port=redis_port, decode_responses=True) # Load model and tokenizer self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.model = AutoModelForCausalLM.from_pretrained(model_name) if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token # Learning state self.learning_events: List[LearningEvent] = [] self.total_events_processed = 0 self.last_update_time = datetime.now() # Performance tracking self.learning_metrics = { 'total_updates': 0, 'successful_updates': 0, 'failed_updates': 0, 'average_loss': 0.0, 'learning_rate_history': [] } logger.info(f"OnlineLearningManager initialized with model: {model_name}") async def add_learning_event(self, event: LearningEvent) -> None: """Add a new learning event to the queue""" try: # Store event in memory self.learning_events.append(event) # Store event in Redis for persistence event_key = f"learning_event:{event.event_id}" await self._store_event_redis(event_key, event) # Maintain memory limit if len(self.learning_events) > self.max_events_memory: self.learning_events.pop(0) self.total_events_processed += 1 logger.info(f"Added learning event: {event.event_type.value} from {event.source}") # Trigger update if threshold reached if len(self.learning_events) >= self.update_frequency: await self.trigger_model_update() except Exception as e: logger.error(f"Error adding learning event: {str(e)}") async def trigger_model_update(self) -> Dict[str, Any]: """Trigger an online model update based on accumulated events""" if not self.learning_events: logger.warning("No learning events available for model update") return {'success': False, 'reason': 'no_events'} try: logger.info(f"Starting online model update with {len(self.learning_events)} events") # Prepare dataset dataset = OnlineDataset(self.learning_events, self.tokenizer) dataloader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True) # Configure optimizer optimizer = torch.optim.AdamW(self.model.parameters(), lr=self.learning_rate) # Training loop self.model.train() total_loss = 0.0 num_batches = 0 for batch in dataloader: optimizer.zero_grad() outputs = self.model( input_ids=batch['input_ids'], attention_mask=batch['attention_mask'], labels=batch['labels'] ) # Apply event weights to loss loss = outputs.loss * batch['event_weight'].mean() loss.backward() torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) optimizer.step() total_loss += loss.item() num_batches += 1 avg_loss = total_loss / num_batches if num_batches > 0 else 0.0 # Update metrics self.learning_metrics['total_updates'] += 1 self.learning_metrics['successful_updates'] += 1 self.learning_metrics['average_loss'] = avg_loss self.learning_metrics['learning_rate_history'].append(self.learning_rate) # Clear processed events self.learning_events.clear() self.last_update_time = datetime.now() logger.info(f"Online model update completed. Average loss: {avg_loss:.4f}") # Store updated model (in production, would save to model registry) await self._save_model_checkpoint() return { 'success': True, 'average_loss': avg_loss, 'events_processed': num_batches * self.batch_size, 'timestamp': self.last_update_time.isoformat() } except Exception as e: self.learning_metrics['failed_updates'] += 1 logger.error(f"Online model update failed: {str(e)}") return {'success': False, 'error': str(e)} async def process_feedback(self, query: str, response: str, feedback_score: float, source: str = "human_feedback") -> None: """Process human feedback for online learning""" event_type = LearningEventType.FEEDBACK_POSITIVE if feedback_score > 0.5 else LearningEventType.FEEDBACK_NEGATIVE event = LearningEvent( event_id=f"feedback_{datetime.now().timestamp()}", event_type=event_type, timestamp=datetime.now(), source=source, context={ 'query': query, 'response': response, 'feedback_score': feedback_score }, feedback_score=feedback_score, priority=3 if abs(feedback_score - 0.5) > 0.3 else 2 # Higher priority for strong feedback ) await self.add_learning_event(event) async def process_threat_intelligence(self, threat_data: Dict[str, Any], source: str = "threat_intel") -> None: """Process new threat intelligence for online learning""" event = LearningEvent( event_id=f"threat_{datetime.now().timestamp()}", event_type=LearningEventType.NEW_THREAT_INTELLIGENCE, timestamp=datetime.now(), source=source, context=threat_data, priority=4, # High priority for new threats confidence=threat_data.get('confidence', 0.8) ) await self.add_learning_event(event) async def process_agent_performance(self, agent_name: str, task: str, success: bool, performance_data: Dict[str, Any]) -> None: """Process agent performance data for online learning""" event_type = LearningEventType.AGENT_SUCCESS if success else LearningEventType.AGENT_FAILURE event = LearningEvent( event_id=f"agent_{agent_name}_{datetime.now().timestamp()}", event_type=event_type, timestamp=datetime.now(), source=agent_name, context={ 'task': task, 'performance_data': performance_data, 'success': success }, priority=2 if success else 3, # Higher priority for failures to learn from confidence=performance_data.get('confidence', 0.9) ) await self.add_learning_event(event) async def get_learning_statistics(self) -> Dict[str, Any]: """Get comprehensive learning statistics""" return { 'total_events_processed': self.total_events_processed, 'current_events_in_memory': len(self.learning_events), 'last_update_time': self.last_update_time.isoformat(), 'metrics': self.learning_metrics, 'event_type_distribution': self._get_event_type_distribution(), 'learning_rate': self.learning_rate, 'update_frequency': self.update_frequency } def _get_event_type_distribution(self) -> Dict[str, int]: """Get distribution of event types in current memory""" distribution = {} for event in self.learning_events: event_type = event.event_type.value distribution[event_type] = distribution.get(event_type, 0) + 1 return distribution async def _store_event_redis(self, key: str, event: LearningEvent) -> None: """Store learning event in Redis for persistence""" try: event_data = json.dumps(event.to_dict()) self.redis_client.setex(key, timedelta(days=7), event_data) except Exception as e: logger.warning(f"Failed to store event in Redis: {str(e)}") async def _save_model_checkpoint(self) -> None: """Save model checkpoint after online learning update""" try: checkpoint_path = f"models/online_learning_checkpoint_{datetime.now().strftime('%Y%m%d_%H%M%S')}" self.model.save_pretrained(checkpoint_path) self.tokenizer.save_pretrained(checkpoint_path) logger.info(f"Model checkpoint saved to {checkpoint_path}") except Exception as e: logger.error(f"Failed to save model checkpoint: {str(e)}") # Factory function for easy instantiation def create_online_learning_manager(**kwargs) -> OnlineLearningManager: """Factory function to create OnlineLearningManager with default configuration""" return OnlineLearningManager(**kwargs)