"""Pydantic models for LogSentinel v2 — Adaptive Multi-Agent SOC War-Room.""" from __future__ import annotations from enum import Enum from typing import Any, Dict, List, Optional from pydantic import BaseModel, Field, field_validator # --------------------------------------------------------------------------- # Enumerations # --------------------------------------------------------------------------- class AgentRole(str, Enum): """Roles in the SOC war-room. Each role has a distinct view of the world.""" INCIDENT_COMMANDER = "incident_commander" APP_SRE = "app_sre" DB_SRE = "db_sre" SECURITY_ANALYST = "security_analyst" class EpisodePhase(str, Enum): """Ordered lifecycle phases of an incident episode.""" DETECT = "detect" TRIAGE = "triage" MITIGATE = "mitigate" VERIFY = "verify" FINAL_REPORT = "final_report" class DifficultyLevel(str, Enum): EASY = "easy" MEDIUM = "medium" HARD = "hard" # --------------------------------------------------------------------------- # Core log + ground-truth models (backward-compatible) # --------------------------------------------------------------------------- class LogEntry(BaseModel): """A single log entry from a source system.""" timestamp: str source: str level: str message: str metadata: Dict[str, Any] = Field(default_factory=dict) class GroundTruth(BaseModel): """Ground truth for grading agent actions.""" log_classifications: Dict[int, str] # index -> classification incidents: List[Dict[str, Any]] # list of incidents expected_severity: Optional[str] = None # --------------------------------------------------------------------------- # World state (latent + observable) # --------------------------------------------------------------------------- class WorldState(BaseModel): """Full latent world state tracked by the environment.""" service_health: float = Field(default=1.0, ge=0.0, le=1.0) db_replication_lag: float = Field(default=0.0, ge=0.0) # seconds error_rate: float = Field(default=0.0, ge=0.0, le=1.0) # fraction attack_signal: float = Field(default=0.0, ge=0.0, le=1.0) # 0=none, 1=confirmed queue_depth: int = Field(default=0, ge=0) mitigations_applied: List[str] = Field(default_factory=list) containment_status: bool = False false_positive_count: int = 0 time_elapsed_steps: int = 0 phase: EpisodePhase = EpisodePhase.DETECT # --------------------------------------------------------------------------- # Observation (role-filtered) # --------------------------------------------------------------------------- class Observation(BaseModel): """What a single agent sees at each step (role-filtered view).""" agent_role: str log_entries: List[LogEntry] task_description: str time_window: str remaining_steps: int current_phase: str previous_action_result: Optional[str] = None incident_context: Optional[Dict[str, Any]] = None shared_board: Optional[Dict[str, Any]] = None # explicit shared memory visible_metrics: Dict[str, Any] = Field(default_factory=dict) # --------------------------------------------------------------------------- # Action (extended, multi-agent) # --------------------------------------------------------------------------- VALID_ACTION_TYPES = { # Legacy (backward-compatible) "classify_log", "detect_incident", "assign_severity", "correlate_logs", "recommend_action", "submit_report", # New Phase-2 actions "observe_logs", "query_metric", "propose_incident", "vote_severity", "request_handoff", "execute_mitigation", "verify_recovery", "submit_joint_report", } class Action(BaseModel): """What an agent can do at each step (multi-agent extended).""" action_type: str agent_role: Optional[str] = None # required for Phase-2 multi-agent actions target_log_indices: Optional[List[int]] = None classification: Optional[str] = None severity: Optional[str] = None incident_type: Optional[str] = None correlated_indices: Optional[List[int]] = None recommendation: Optional[str] = None report: Optional[Dict[str, Any]] = None # Phase-2 fields metric_name: Optional[str] = None # for query_metric mitigation_id: Optional[str] = None # for execute_mitigation handoff_to: Optional[str] = None # for request_handoff (target role) handoff_note: Optional[str] = None evidence_indices: Optional[List[int]] = None # for propose_incident / vote vote: Optional[str] = None # "confirm" | "reject" | "escalate" @field_validator("action_type") @classmethod def validate_action_type(cls, v: str) -> str: if v not in VALID_ACTION_TYPES: raise ValueError( f"Unknown action_type '{v}'. Valid types: {sorted(VALID_ACTION_TYPES)}" ) return v # --------------------------------------------------------------------------- # Scenario configuration (procedural generation) # --------------------------------------------------------------------------- class ScenarioConfig(BaseModel): """Parameters that drive procedural scenario generation.""" num_incidents: int = Field(default=2, ge=1, le=3) difficulty: DifficultyLevel = DifficultyLevel.MEDIUM attack_subtlety: float = Field(default=0.5, ge=0.0, le=1.0) observability_quality: float = Field(default=0.8, ge=0.0, le=1.0) telemetry_delay_steps: int = Field(default=0, ge=0) confounding_noise_ratio: float = Field(default=0.3, ge=0.0, le=0.8) seed: Optional[int] = None # --------------------------------------------------------------------------- # Task definition (extended) # --------------------------------------------------------------------------- class TaskDefinition(BaseModel): """Definition of a task in the environment.""" name: str description: str difficulty: str max_steps: int num_logs: int num_sources: int num_incidents: int multi_agent: bool = False scenario_config: Optional[ScenarioConfig] = None # --------------------------------------------------------------------------- # Reward breakdown (transparent logging) # --------------------------------------------------------------------------- class RewardBreakdown(BaseModel): """Per-step reward components logged in info/state for plotting.""" r_outcome: float = 0.0 r_detection_f1: float = 0.0 r_severity_accuracy: float = 0.0 r_efficiency: float = 0.0 r_teamwork: float = 0.0 penalty_spam: float = 0.0 penalty_unsupported: float = 0.0 penalty_unsafe: float = 0.0 penalty_noop: float = 0.0 total: float = 0.0 def compute_total(self) -> "RewardBreakdown": """Recompute weighted total from components.""" raw = ( 0.40 * self.r_outcome + 0.20 * self.r_detection_f1 + 0.15 * self.r_severity_accuracy + 0.10 * self.r_efficiency + 0.15 * self.r_teamwork - self.penalty_spam - self.penalty_unsupported - self.penalty_unsafe - self.penalty_noop ) self.total = max(-1.0, min(1.0, raw)) return self # --------------------------------------------------------------------------- # Curriculum tracker # --------------------------------------------------------------------------- class CurriculumState(BaseModel): """Tracks rolling success rate for adaptive difficulty.""" window_size: int = 20 success_history: List[float] = Field(default_factory=list) current_difficulty: DifficultyLevel = DifficultyLevel.EASY up_threshold: float = 0.70 # promote to harder if avg > this down_threshold: float = 0.35 # demote to easier if avg < this def record(self, episode_reward: float) -> None: """Record episode outcome and potentially adjust difficulty.""" success = 1.0 if episode_reward > 0.5 else 0.0 self.success_history.append(success) if len(self.success_history) > self.window_size: self.success_history = self.success_history[-self.window_size:] def adjust_difficulty(self) -> DifficultyLevel: """Check rolling avg and update difficulty level.""" if len(self.success_history) < 5: return self.current_difficulty avg = sum(self.success_history) / len(self.success_history) order = [DifficultyLevel.EASY, DifficultyLevel.MEDIUM, DifficultyLevel.HARD] idx = order.index(self.current_difficulty) if avg > self.up_threshold and idx < len(order) - 1: self.current_difficulty = order[idx + 1] elif avg < self.down_threshold and idx > 0: self.current_difficulty = order[idx - 1] return self.current_difficulty