workflow_twin/models.py

from typing import Any, Literal

from pydantic import BaseModel, Field


class Observation(BaseModel):
    ticket_id: str | None = None
    ticket_status: str | None = None
    attempts_used: int = 0
    attempts_remaining: int = 0
    severity: Literal["low", "medium", "high"] = "low"
    summary: str = ""
    embedding: list[float] | None = None

    queue_length: int = 0
    waiting_tickets: list[dict[str, Any]] = Field(default_factory=list)

    agents_available: int = 0
    current_time: float = 0.0
    level: int = 1
    memory_used: int = 0
    memory_budget: int = 0


class Action(BaseModel):
    action_type: Literal[
        "triage",
        "respond",
        "resolve",
        "escalate",
        "assign",
        "reprioritize",
        "skip",
        "compress_memory",
    ]
    ticket_id: str | None = None
    agent_id: str | None = None
    priority: int | None = None
    note: str = ""


class RewardSignal(BaseModel):
    value: float
    reason: str


class StepResult(BaseModel):
    observation: Observation
    reward: RewardSignal
    done: bool
    info: dict[str, Any] = Field(default_factory=dict)