models.py

"""Action and Observation models for the SimLab HR environment."""

from __future__ import annotations

from typing import Literal

from pydantic import Field

from openenv.core.env_server.types import Action, Observation

TOOL_SERVERS = ("hrms", "email", "calendar", "rocketchat")
ToolServerName = Literal["hrms", "email", "calendar", "rocketchat"]


class HRAction(Action):
    """An action in the HR environment — invoke a tool on one of the tool servers."""

    tool_server: ToolServerName = Field(
        ...,
        description="Which tool server to target: 'hrms' (employee records, leave, payroll), "
        "'email' (send/read email), 'calendar' (schedule events), or 'rocketchat' (team chat).",
    )
    tool_name: str = Field(
        ...,
        description="Name of the tool to invoke on the selected server. "
        "Use the tools_available field from the observation to see what's available per server.",
    )
    parameters: dict = Field(
        default_factory=dict,
        description="Tool-specific parameters as a JSON object.",
    )


class HRObservation(Observation):
    """Observation returned by the HR environment after each action."""

    result: str = Field(
        ...,
        description="Result of the tool invocation, or initial environment status on reset.",
    )
    is_error: bool = Field(
        default=False,
        description="Whether the tool invocation resulted in an error.",
    )
    tools_available: dict[str, list[str]] = Field(
        default_factory=dict,
        description="Available tools grouped by server: {'hrms': [...], 'email': [...], ...}.",
    )
    task_instruction: str = Field(
        default="",
        description="The task the agent should complete in this episode.",
    )