client.py

from __future__ import annotations

from typing import Any

import requests

from models import SepsisAction, SepsisObservation, SepsisState
from openenv_compat import EnvClient, OPENENV_AVAILABLE, StepResult
from server.sepsis_environment import SepsisTreatmentEnvironment


class SepsisTreatmentEnv(EnvClient):
    def __init__(self, base_url: str | None = None, task_id: str = "easy"):
        if base_url is not None:
            super().__init__(base_url=base_url)
        else:
            self.base_url = None
            if OPENENV_AVAILABLE:
                self._provider = None
                self._ws = None
        self.task_id = task_id
        self._local_env = None if base_url else SepsisTreatmentEnvironment(task_id=task_id)

    def _step_payload(self, action: SepsisAction) -> dict[str, Any]:
        return action.model_dump()

    def _parse_result(self, payload: dict[str, Any]) -> StepResult[SepsisObservation]:
        return StepResult(
            observation=SepsisObservation(**payload["observation"]),
            reward=payload.get("reward"),
            done=payload.get("done", False),
            info=payload.get("info", {}),
        )

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

    def reset(self) -> StepResult[SepsisObservation]:
        if self._local_env is not None:
            observation = self._local_env.reset(task_id=self.task_id)
            return StepResult(
                observation=observation,
                reward=0.0,
                done=False,
                info={"tasks": self._local_env.available_tasks()},
            )

        response = requests.post(f"{self.base_url.rstrip('/')}/reset", json={"task_id": self.task_id}, timeout=30)
        response.raise_for_status()
        return self._parse_result(response.json())

    def step(self, action: SepsisAction) -> StepResult[SepsisObservation]:
        if self._local_env is not None:
            observation = self._local_env.step(action)
            return StepResult(
                observation=observation,
                reward=observation.reward,
                done=observation.done,
                info={"metrics": self._local_env.current_metrics()},
            )

        action_payload = self._step_payload(action)
        response = requests.post(f"{self.base_url.rstrip('/')}/step", json=action_payload, timeout=30)
        if response.status_code == 422:
            response = requests.post(
                f"{self.base_url.rstrip('/')}/step",
                json={"action": action_payload},
                timeout=30,
            )
        response.raise_for_status()
        return self._parse_result(response.json())

    def state(self) -> SepsisState:
        if self._local_env is not None:
            return self._local_env.state
        response = requests.get(f"{self.base_url.rstrip('/')}/state", timeout=30)
        response.raise_for_status()
        return self._parse_state(response.json())

    def metadata(self) -> dict[str, Any]:
        if self._local_env is not None:
            return self._local_env.metadata()
        response = requests.get(f"{self.base_url.rstrip('/')}/metadata", timeout=30)
        response.raise_for_status()
        return response.json()

    def close(self) -> None:
        self._local_env = None