server/client.py

"""Minimal Python client for the Tetris OpenEnv FastAPI server."""

from __future__ import annotations

from typing import Any, Dict, Tuple

import requests


class TetrisClient:
    """Thin HTTP client wrapping the /reset, /step, and /state endpoints."""

    def __init__(self, base_url: str = "http://localhost:8000", timeout: float = 30.0) -> None:
        self.base_url: str = base_url.rstrip("/")
        self.timeout: float = timeout

    def reset(self) -> Dict[str, Any]:
        """Reset the remote environment and return the observation dict."""
        r = requests.post(f"{self.base_url}/reset", timeout=self.timeout)
        r.raise_for_status()
        return r.json()["observation"]

    def step(self, action: str) -> Tuple[Dict[str, Any], Dict[str, float], bool, Dict[str, Any]]:
        """Apply an action and return (observation, reward, done, info)."""
        r = requests.post(
            f"{self.base_url}/step",
            json={"action": action},
            timeout=self.timeout,
        )
        r.raise_for_status()
        data = r.json()
        return data["observation"], data["reward"], data["done"], data["info"]

    def state(self) -> Dict[str, Any]:
        """Fetch the current observation without stepping."""
        r = requests.get(f"{self.base_url}/state", timeout=self.timeout)
        r.raise_for_status()
        return r.json()