adam/models.py

from __future__ import annotations

from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any
from uuid import uuid4


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


class JobStatus(str, Enum):
    DRAFT = "Draft"
    AWAITING_CONFIRMATION = "Awaiting confirmation"
    QUEUED = "Queued"
    RUNNING = "Running"
    PAUSED = "Paused"
    FINISHED = "Finished"
    FAILED = "Failed"
    CANCELLED = "Cancelled"
    INTERRUPTED = "Interrupted"


class StepStatus(str, Enum):
    PENDING = "Pending"
    RUNNING = "Running"
    FINISHED = "Finished"
    FAILED = "Failed"
    SKIPPED = "Skipped"


@dataclass(slots=True)
class PlanStep:
    tool_id: str
    title: str
    description: str
    arguments: dict[str, Any] = field(default_factory=dict)
    status: StepStatus = StepStatus.PENDING

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "PlanStep":
        payload = dict(data)
        payload["status"] = StepStatus(payload.get("status", StepStatus.PENDING))
        return cls(**payload)


@dataclass(slots=True)
class ExecutionPlan:
    request: str
    summary: str
    steps: list[PlanStep]
    requires_confirmation: bool = False
    confirmation_reason: str = ""
    project_name: str = "Untitled project"
    id: str = field(default_factory=lambda: uuid4().hex[:10])
    created_at: str = field(default_factory=utc_now)
    orion_review: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ExecutionPlan":
        payload = dict(data)
        payload["steps"] = [PlanStep.from_dict(item) for item in payload["steps"]]
        return cls(**payload)


@dataclass(slots=True)
class Job:
    plan: ExecutionPlan
    id: str = field(default_factory=lambda: uuid4().hex[:8].upper())
    status: JobStatus = JobStatus.DRAFT
    progress: int = 0
    current_step: int = -1
    created_at: str = field(default_factory=utc_now)
    started_at: str | None = None
    ended_at: str | None = None
    output_folder: str | None = None
    logs: list[str] = field(default_factory=list)
    error: str | None = None
    preview_path: str | None = None
    preview_epoch: int = 0
    preview_next_epoch: int = 0
    preview_prompt: str = ""
    preview_seed: int | None = None
    preview_steps: int = 0
    preview_kind: str = "training"
    preview_current: int = 0
    preview_total: int = 0
    preview_image_index: int = 0
    preview_image_count: int = 0
    atlas_report: dict[str, Any] = field(default_factory=dict)
    nova_report: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        payload = asdict(self)
        payload["status"] = self.status.value
        for step in payload["plan"]["steps"]:
            step["status"] = step["status"].value
        return payload

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "Job":
        payload = dict(data)
        payload["status"] = JobStatus(payload["status"])
        payload["plan"] = ExecutionPlan.from_dict(payload["plan"])
        return cls(**payload)


@dataclass(slots=True)
class SystemSnapshot:
    cpu_percent: float = 0.0
    memory_percent: float = 0.0
    memory_used_gb: float = 0.0
    memory_total_gb: float = 0.0
    disk_percent: float = 0.0
    disk_used_gb: float = 0.0
    disk_total_gb: float = 0.0
    gpu_name: str = "No NVIDIA GPU detected"
    gpu_percent: float = 0.0
    vram_percent: float = 0.0
    vram_used_gb: float = 0.0
    vram_total_gb: float = 0.0
    gpu_temperature: float | None = None