Spaces:
Sleeping
Sleeping
| """Typed OpenEnv client for the Hospital ED environment. | |
| Maintains a persistent WebSocket connection to a running Hospital ED | |
| server and exposes typed ``reset()`` / ``step()`` / ``state()`` methods | |
| that return :class:`HospitalObservation` and :class:`HospitalState`. | |
| Use this from the baseline ``inference.py`` or any other LLM agent | |
| driver, either pointing it at a locally-running server or spinning up | |
| a fresh Docker container via :meth:`HospitalEdEnv.from_docker_image`:: | |
| # Connect to a running server | |
| with HospitalEdEnv(base_url="http://localhost:8000") as client: | |
| result = client.reset(task="surge", seed=0) | |
| while not result.done: | |
| result = client.step(HospitalAction(action=0)) | |
| # Or launch a fresh container (hackathon inference pattern) | |
| client = HospitalEdEnv.from_docker_image("hospital-ed-env:latest") | |
| """ | |
| from __future__ import annotations | |
| from typing import Any, Dict | |
| from openenv.core import EnvClient | |
| from openenv.core.client_types import StepResult | |
| from models import HospitalAction, HospitalObservation, HospitalState | |
| class HospitalEdEnv( | |
| EnvClient[HospitalAction, HospitalObservation, HospitalState] | |
| ): | |
| """Client for the Hospital Emergency Department OpenEnv environment. | |
| This client maintains a persistent WebSocket connection to the | |
| environment server, enabling efficient multi-step interactions with | |
| lower latency. Each instance holds a dedicated session on the server. | |
| """ | |
| # ------------------------------------------------------------------ | |
| # OpenEnv EnvClient hooks | |
| # ------------------------------------------------------------------ | |
| def _step_payload(self, action: HospitalAction) -> Dict[str, Any]: | |
| """Serialize ``action`` to the JSON payload expected by ``/step``.""" | |
| return {"action": int(action.action)} | |
| def _parse_result( | |
| self, payload: Dict[str, Any] | |
| ) -> StepResult[HospitalObservation]: | |
| """Parse a server response into a typed :class:`StepResult`.""" | |
| obs_data = payload.get("observation", payload) or {} | |
| observation = HospitalObservation( | |
| done=bool(payload.get("done", obs_data.get("done", False))), | |
| reward=payload.get("reward", obs_data.get("reward")), | |
| metadata=obs_data.get("metadata", {}) or {}, | |
| bed_occupancy=list(obs_data.get("bed_occupancy", [])), | |
| icu_occupancy=list(obs_data.get("icu_occupancy", [])), | |
| ventilator_status=list(obs_data.get("ventilator_status", [])), | |
| waiting_queue=[ | |
| list(row) for row in obs_data.get("waiting_queue", []) | |
| ], | |
| time_step=int(obs_data.get("time_step", 0) or 0), | |
| stats=list(obs_data.get("stats", [])), | |
| action_mask=[bool(x) for x in obs_data.get("action_mask", [])], | |
| survival_rate=float(obs_data.get("survival_rate", 0.0) or 0.0), | |
| critical_survival_rate=float( | |
| obs_data.get("critical_survival_rate", 1.0) or 1.0 | |
| ), | |
| task_score=float(obs_data.get("task_score", 0.0) or 0.0), | |
| task=obs_data.get("task"), | |
| ) | |
| return StepResult( | |
| observation=observation, | |
| reward=payload.get("reward"), | |
| done=bool(payload.get("done", False)), | |
| ) | |
| def _parse_state(self, payload: Dict[str, Any]) -> HospitalState: | |
| """Parse a ``/state`` payload into a typed :class:`HospitalState`.""" | |
| return HospitalState( | |
| episode_id=payload.get("episode_id"), | |
| step_count=int(payload.get("step_count", 0) or 0), | |
| total_treated=int(payload.get("total_treated", 0) or 0), | |
| total_deaths=int(payload.get("total_deaths", 0) or 0), | |
| total_admitted=int(payload.get("total_admitted", 0) or 0), | |
| critical_total=int(payload.get("critical_total", 0) or 0), | |
| critical_saved=int(payload.get("critical_saved", 0) or 0), | |
| queue_len=int(payload.get("queue_len", 0) or 0), | |
| general_occupancy=float(payload.get("general_occupancy", 0.0) or 0.0), | |
| icu_occupancy=float(payload.get("icu_occupancy", 0.0) or 0.0), | |
| ventilator_utilization=float( | |
| payload.get("ventilator_utilization", 0.0) or 0.0 | |
| ), | |
| invalid_action_rate=float( | |
| payload.get("invalid_action_rate", 0.0) or 0.0 | |
| ), | |
| episode_reward=float(payload.get("episode_reward", 0.0) or 0.0), | |
| scenario=payload.get("scenario"), | |
| ) | |