models.py

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