Spaces:
Paused
Paused
| """Typed persistent-session client for MATH OpenEnv.""" | |
| from __future__ import annotations | |
| import asyncio | |
| from typing import Any | |
| from urllib.parse import urlsplit | |
| from openenv.core import EnvClient | |
| from openenv.core.client_types import StepResult | |
| from websockets.asyncio.client import connect as ws_connect | |
| from websockets.exceptions import SecurityError | |
| from math_env.models import ( | |
| MathAction, | |
| MathObservation, | |
| MathState, | |
| ) | |
| class _AuthenticatedConnect(ws_connect): | |
| def process_redirect(self, exc: Exception) -> Exception | str: | |
| result = super().process_redirect(exc) | |
| if isinstance(result, str): | |
| return SecurityError("Authenticated WebSocket redirects are disabled") | |
| return result | |
| def _authenticated_ws_connect( | |
| uri: str, | |
| *, | |
| bearer_token: str, | |
| **kwargs: Any, | |
| ): | |
| return _AuthenticatedConnect( | |
| uri, | |
| additional_headers={"Authorization": f"Bearer {bearer_token}"}, | |
| **kwargs, | |
| ) | |
| class MathClient(EnvClient[MathAction, MathObservation, MathState]): | |
| """Typed WebSocket client for one connection-owned episode.""" | |
| def __init__( | |
| self, | |
| *args: Any, | |
| bearer_token: str | None = None, | |
| **kwargs: Any, | |
| ) -> None: | |
| base_url = kwargs.get("base_url", args[0] if args else None) | |
| if bearer_token is not None: | |
| scheme = urlsplit(str(base_url)).scheme.lower() | |
| if scheme not in {"https", "wss"}: | |
| raise ValueError( | |
| "Bearer authentication requires an HTTPS or WSS endpoint" | |
| ) | |
| self._bearer_token = bearer_token | |
| super().__init__(*args, **kwargs) | |
| async def _connect_async(self) -> "MathClient": | |
| if self._bearer_token is None: | |
| await super()._connect_async() | |
| return self | |
| if self._ws is not None: | |
| if self._ws_loop is asyncio.get_running_loop(): | |
| return self | |
| self._ws = None | |
| self._ws_loop = None | |
| try: | |
| self._start_provider_if_needed() | |
| except Exception: | |
| await self.close() | |
| raise | |
| assert self._ws_url is not None | |
| try: | |
| self._ws = await _authenticated_ws_connect( | |
| self._ws_url, | |
| bearer_token=self._bearer_token, | |
| open_timeout=self._connect_timeout, | |
| max_size=self._max_message_size, | |
| ping_interval=self._websocket_ping_interval_s, | |
| ping_timeout=self._websocket_ping_timeout_s, | |
| ) | |
| self._ws_loop = asyncio.get_running_loop() | |
| except Exception as error: | |
| await self.close() | |
| raise ConnectionError( | |
| f"Failed to connect to {self._ws_url}: {error}" | |
| ) from error | |
| return self | |
| def _step_payload(self, action: MathAction) -> dict[str, Any]: | |
| return action.model_dump() | |
| def _parse_result( | |
| self, | |
| payload: dict[str, Any], | |
| ) -> StepResult[MathObservation]: | |
| observation_data = dict(payload.get("observation", {})) | |
| reward_value = payload.get("reward") | |
| reward = None if reward_value is None else float(reward_value) | |
| metadata = payload.get("metadata", observation_data.get("metadata")) | |
| observation_data.update( | |
| reward=reward, | |
| metadata={} if metadata is None else metadata, | |
| ) | |
| observation = MathObservation.model_validate(observation_data) | |
| return StepResult( | |
| observation=observation, | |
| reward=reward, | |
| done=observation.done, | |
| metadata=metadata, | |
| ) | |
| def _parse_state(self, payload: dict[str, Any]) -> MathState: | |
| return MathState.model_validate(payload) |