supportdesk_env/client.py

"""HTTP client for interacting with a deployed SupportDesk environment."""

from __future__ import annotations

from supportdesk_env.models import SupportDeskAction, SupportDeskObservation, SupportDeskState
from supportdesk_env.openenv_compat import EnvClient, StepResult


def _validate(model_cls, payload):
    if hasattr(model_cls, "model_validate"):
        return model_cls.model_validate(payload)
    return model_cls(**payload)  # pragma: no cover - pydantic v1 fallback


class SupportDeskEnv(EnvClient[SupportDeskAction, SupportDeskObservation, SupportDeskState]):
    """Typed client for a locally running or deployed OpenEnv server."""

    def _parse_state(self, payload) -> SupportDeskState:
        return _validate(SupportDeskState, payload)

    def _parse_reset(self, payload) -> SupportDeskObservation:
        return _validate(SupportDeskObservation, payload)

    def _parse_result(self, payload) -> StepResult[SupportDeskObservation]:
        observation = _validate(SupportDeskObservation, payload["observation"])
        return StepResult(
            observation=observation,
            reward=payload["reward"],
            done=payload["done"],
            info=payload.get("info", {}),
        )