# src/client.py """ HTTP client for AquaGuard-RL environment. Provides a synchronous client that connects to a running AquaGuard-RL server (local or remote) and provides the standard RL interface: env.reset(task, seed) → AquaGuardObservation env.step(action) → AquaGuardObservation env.state → AquaGuardState env.close() → stops Docker container if started via from_docker_image() Usage: # Connect to a running server env = AquaGuardEnv("http://localhost:8000") obs = env.reset(task="baseline", seed=42) while not obs.done: action = ... obs = env.step(action) env.close() # Start from Docker image env = AquaGuardEnv.from_docker_image("aquaguard-env:latest") obs = env.reset(task="crisis") """ from __future__ import annotations import logging import time from typing import Optional logger = logging.getLogger(__name__) from models import AquaGuardAction, AquaGuardObservation, AquaGuardState def _parse_observation(data: dict) -> AquaGuardObservation: """Parse server response into AquaGuardObservation. Handles both flat responses and openenv_core's nested format: {"observation": {...}, "reward": ..., "done": ...} """ if "observation" in data and isinstance(data["observation"], dict): # Nested format from openenv_core Environment obs_data = data["observation"] obs_data["reward"] = data.get("reward") obs_data["done"] = data.get("done", False) return AquaGuardObservation(**obs_data) return AquaGuardObservation(**data) class AquaGuardEnv: """ Synchronous HTTP client for AquaGuard-RL environment. Implements the standard RL interface (reset/step/state) backed by HTTP requests to a running FastAPI environment server. """ def __init__(self, base_url: str = "http://localhost:8000") -> None: """ Initialize client connected to a running environment server. Args: base_url: Base URL of the environment server. """ try: import httpx self._http = httpx.Client(timeout=120.0) except ImportError: raise ImportError("httpx required: pip install httpx") self._base_url = base_url.rstrip("/") self._container: Optional[str] = None logger.debug(f"AquaGuardEnv client initialized: {self._base_url}") @classmethod def from_docker_image( cls, image: str = "aquaguard-env:latest", port: int = 8000, timeout: int = 60, ) -> "AquaGuardEnv": """ Start AquaGuard-RL in a Docker container and return connected client. Args: image: Docker image name/tag. port: Host port to bind (default 8000). timeout: Seconds to wait for server startup. Returns: Connected AquaGuardEnv client. Raises: RuntimeError: If server fails to start within timeout. """ import subprocess import httpx logger.info(f"Starting Docker container: {image}") result = subprocess.run( ["docker", "run", "-d", "-p", f"{port}:8000", image], capture_output=True, text=True, check=True, ) container_id = result.stdout.strip() logger.info(f"Container started: {container_id[:12]}") base_url = f"http://localhost:{port}" for attempt in range(timeout): try: resp = httpx.get(f"{base_url}/health", timeout=2.0) if resp.status_code == 200: env = cls(base_url) env._container = container_id logger.info(f"Server ready after {attempt + 1}s") return env except Exception: pass time.sleep(1.0) # Cleanup failed container subprocess.run(["docker", "stop", container_id], capture_output=True) raise RuntimeError( f"Environment server failed to start within {timeout}s. " f"Check Docker logs: docker logs {container_id[:12]}" ) def reset( self, task: str = "baseline", seed: Optional[int] = None, episode_id: Optional[str] = None, ) -> AquaGuardObservation: """ Reset environment to start a new episode. Args: task: Task name ('baseline', 'crisis', 'policy_shift', 'climate_shock', 'multi_district'). seed: Optional random seed for reproducibility. episode_id: Optional explicit episode ID. Returns: Initial AquaGuardObservation (step_number=0, reward=None, done=False). """ payload: dict = {"task": task} if seed is not None: payload["seed"] = seed if episode_id is not None: payload["episode_id"] = episode_id resp = self._http.post(f"{self._base_url}/reset", json=payload) resp.raise_for_status() return _parse_observation(resp.json()) def step(self, action: AquaGuardAction) -> AquaGuardObservation: """ Execute one policy step (one growing season ~4 months). Args: action: AquaGuardAction with crop allocation, water quotas, etc. Returns: AquaGuardObservation with updated state, reward, and done flag. """ payload = action.model_dump() if hasattr(action, "model_dump") else action.dict() # Try openenv_core wrapped format first: {"action": {...}}, then flat fallback resp = self._http.post(f"{self._base_url}/step", json={"action": payload}) if resp.status_code == 422: # Server may use standalone (flat) format resp = self._http.post(f"{self._base_url}/step", json=payload) resp.raise_for_status() return _parse_observation(resp.json()) @property def state(self) -> AquaGuardState: """Get current episode state metadata.""" resp = self._http.get(f"{self._base_url}/state") resp.raise_for_status() return AquaGuardState(**resp.json()) def close(self) -> None: """ Close the HTTP client and stop Docker container if started via from_docker_image(). """ try: self._http.close() except Exception: pass if self._container: import subprocess logger.info(f"Stopping container {self._container[:12]}") subprocess.run(["docker", "stop", self._container], capture_output=True) subprocess.run(["docker", "rm", self._container], capture_output=True) def __enter__(self) -> "AquaGuardEnv": return self def __exit__(self, *args) -> None: self.close() def health_check(self) -> bool: """Check if the server is running and healthy.""" try: resp = self._http.get(f"{self._base_url}/health", timeout=5.0) return resp.status_code == 200 except Exception: return False def get_info(self) -> dict: """Get environment metadata from the server.""" resp = self._http.get(f"{self._base_url}/info") resp.raise_for_status() return resp.json()