hr_env/models.py

"""Pydantic models for the HR Productivity Environment OpenEnv interface."""

from __future__ import annotations

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

from openenv.core.env_server.types import Action, Observation, State
from pydantic import Field


class HRAction(Action):
    """Action space for the HR Productivity Environment.

    The agent selects an action_type and provides relevant parameters.
    Actions are gated by the current HCM:21 phase.
    """

    action_type: Literal[
        # Scanning phase
        "query_department",
        "query_employees",
        "calculate_metric",
        "review_financials",
        # Planning phase
        "set_hiring_target",
        "set_training_budget",
        "set_compensation_policy",
        "set_retention_program",
        # Producing phase
        "execute_hiring",
        "execute_promotion",
        "execute_transfer",
        "execute_training",
        "execute_termination",
        # Phase/quarter control
        "advance_phase",
        "advance_quarter",
        "submit_report",
    ]
    department: Optional[str] = Field(default=None, description="Target department name")
    employee_ids: Optional[List[str]] = Field(default=None, description="Target employee IDs")
    metric_name: Optional[str] = Field(
        default=None,
        description="Metric to calculate: hcva, hcroi, qips, five_indexes, employee_value",
    )
    amount: Optional[float] = Field(default=None, description="Dollar amount or percentage")
    count: Optional[int] = Field(default=None, description="Headcount or quantity")
    parameters: Optional[Dict[str, Any]] = Field(default=None, description="Additional action parameters")
    rationale: Optional[str] = Field(default=None, description="Agent's reasoning for this action")


class HRObservation(Observation):
    """Observation returned after each step."""

    message: str = Field(description="Human-readable result of the action")
    data: Optional[Dict[str, Any]] = Field(default=None, description="Structured response data")
    available_actions: List[str] = Field(default_factory=list, description="Valid action_types in current phase")
    current_phase: str = Field(default="scanning", description="Current HCM:21 phase")
    current_quarter: int = Field(default=1, description="Current quarter (1-6)")
    step_in_quarter: int = Field(default=0, description="Steps taken in current quarter")
    warnings: List[str] = Field(default_factory=list, description="Budget alerts, attrition spikes, etc.")


class HRState(State):
    """Full environment state exposed via the state endpoint."""

    current_quarter: int = Field(default=1, description="Current quarter (1-6)")
    current_phase: str = Field(default="scanning", description="Current HCM:21 phase")
    total_steps: int = Field(default=0, description="Total steps across all quarters")
    company_snapshot: Dict[str, Any] = Field(default_factory=dict, description="Aggregated company data")
    metric_history: List[Dict[str, Any]] = Field(
        default_factory=list, description="Per-quarter metric snapshots"
    )
    budget_remaining: Dict[str, float] = Field(default_factory=dict, description="Remaining budget allocations")
    active_initiatives: List[Dict[str, Any]] = Field(default_factory=list, description="In-progress programs")
    recent_events: List[str] = Field(default_factory=list, description="Recent stochastic events")