hr_env/client.py

"""EnvClient subclass for the HR Productivity Environment."""

from __future__ import annotations

from typing import Any, Dict

from openenv.core import EnvClient
from openenv.core.client_types import StepResult

from hr_env.models import HRAction, HRObservation, HRState


class HRProductivityEnv(EnvClient[HRAction, HRObservation, HRState]):
    """Typed client for connecting to the HR Productivity Environment server."""

    def _step_payload(self, action: HRAction) -> Dict[str, Any]:
        """Serialize an HRAction for the WebSocket protocol."""
        return action.model_dump(exclude_none=True)

    def _parse_result(self, payload: Dict[str, Any]) -> StepResult[HRObservation]:
        """Parse a step/reset response into a StepResult."""
        obs = HRObservation.model_validate(payload)
        return StepResult(
            observation=obs,
            reward=obs.reward,
            done=obs.done,
        )

    def _parse_state(self, payload: Dict[str, Any]) -> HRState:
        """Parse a state response into HRState."""
        return HRState.model_validate(payload)