models.py

from typing import List, Dict, Literal, Optional, Any
from pydantic import BaseModel

class Action(BaseModel):
    action_type: Literal["detect", "classify", "mitigate", "request_more_logs"]
    value: Optional[str] = None  # e.g., "anomaly", "cpu_spike", "scale_up", or None if requesting logs

class Observation(BaseModel):
    logs: List[str]
    metrics: Dict[str, float]
    timestep: int
    max_steps: int
    info: str = ""  # General feedback/info from the previous step

class State(BaseModel):
    """Full internal state — contains ground_truth. Never expose this directly."""
    episode_id: str
    step_count: int
    scenario_name: str
    total_reward: float
    is_done: bool
    ground_truth: Dict[str, str]  # Internal only — NOT exposed via /state endpoint

class PublicState(BaseModel):
    """Safe state returned by the /state endpoint — ground_truth is stripped."""
    episode_id: str
    step_count: int
    scenario_name: str
    total_reward: float
    is_done: bool