redbutton_v2/client.py

"""WebSocket client for the Red Button v2 environment server.

Carries v1's pattern: subclasses ``EnvClient``, implements
``_step_payload`` / ``_parse_result`` / ``_parse_state``. Sync /
async / from-Docker entry points are inherited.

Sync usage (training rollouts, evaluation):

    from redbutton_v2 import RedButtonV2Client, RedButtonV2Action
    with RedButtonV2Client(base_url="http://localhost:8000").sync() as env:
        result = env.reset(tier=2, seed=42)
        while not result.done:
            action = RedButtonV2Action(tool_name="read_file", arguments={"path": "..."})
            result = env.step(action)
"""

from typing import Any, Dict

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

from redbutton_v2.models import (
    RedButtonV2Action,
    RedButtonV2Observation,
    RedButtonV2State,
)


class RedButtonV2Client(
    EnvClient[RedButtonV2Action, RedButtonV2Observation, RedButtonV2State]
):
    """WebSocket client for the Red Button v2 environment server."""

    def _step_payload(self, action: RedButtonV2Action) -> Dict[str, Any]:
        return {
            "tool_name": action.tool_name,
            "arguments": action.arguments,
            "metadata": action.metadata,
        }

    def _parse_result(
        self, payload: Dict[str, Any]
    ) -> StepResult[RedButtonV2Observation]:
        obs = RedButtonV2Observation(**payload["observation"])
        return StepResult(
            observation=obs,
            reward=payload.get("reward"),
            done=bool(payload.get("done", False)),
        )

    def _parse_state(self, payload: Dict[str, Any]) -> RedButtonV2State:
        return RedButtonV2State(**payload)