env_models.py

from pydantic import BaseModel
from typing import Optional, List
from models import Category, Priority, Status


class TicketObservation(BaseModel):
    task: str
    ticket_id: str
    subject: str
    description: str
    current_status: Status
    current_category: Optional[Category] = None
    current_priority: Optional[Priority] = None
    ai_resolution: Optional[str] = None
    step: int
    done: bool


class TicketAction(BaseModel):
    category: Optional[Category] = None
    priority: Optional[Priority] = None
    resolution: Optional[str] = None
    status: Optional[Status] = None


class TicketReward(BaseModel):
    reward: float
    done: bool
    info: dict


class StepResult(BaseModel):
    observation: TicketObservation
    reward: float
    done: bool
    info: dict


class EnvState(BaseModel):
    task: str
    step: int
    done: bool
    tickets: List[dict]
    score: float