client.py

"""
CleanRL Python client — a thin HTTP wrapper around the CleanRL OpenEnv server.

Usage:
    from client import CleanRLClient

    client = CleanRLClient("http://localhost:7860")
    obs = client.reset("basic_tabular_cleaning")

    while True:
        action = {"action_type": "remove_duplicates"}
        result = client.step(action)
        print(result["reward"], result["done"])
        if result["done"]:
            break
"""
from __future__ import annotations

import json
from typing import Any

import requests


class CleanRLClient:
    """HTTP client for the CleanRL OpenEnv server."""

    def __init__(self, base_url: str = "http://localhost:7860", timeout: int = 30) -> None:
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout

    # ── Core OpenEnv methods ───────────────────────────────────────────────

    def reset(self, task_id: str = "basic_tabular_cleaning") -> dict[str, Any]:
        """Start a new episode. Returns the initial observation."""
        r = requests.post(
            f"{self.base_url}/reset",
            params={"task_id": task_id},
            timeout=self.timeout,
        )
        r.raise_for_status()
        return r.json()

    def step(self, action: dict[str, Any]) -> dict[str, Any]:
        """
        Submit one action. Returns::

            {
                "observation": {...},
                "reward": float,
                "done": bool,
                "info": {...}
            }
        """
        r = requests.post(
            f"{self.base_url}/step",
            json=action,
            timeout=self.timeout,
        )
        r.raise_for_status()
        return r.json()

    def state(self) -> dict[str, Any]:
        """Return full environment state for the current session."""
        r = requests.get(f"{self.base_url}/state", timeout=self.timeout)
        r.raise_for_status()
        return r.json()

    def close(self) -> None:
        """No-op — HTTP env has no persistent connection to close."""

    # ── Convenience helpers ────────────────────────────────────────────────

    def auto_step(self) -> dict[str, Any]:
        """Trigger one heuristic auto-step (no LLM required)."""
        r = requests.post(f"{self.base_url}/auto", timeout=self.timeout)
        r.raise_for_status()
        return r.json()

    def health(self) -> dict[str, Any]:
        r = requests.get(f"{self.base_url}/health", timeout=self.timeout)
        r.raise_for_status()
        return r.json()

    def tasks(self) -> list[dict[str, Any]]:
        """List all available tasks."""
        r = requests.get(f"{self.base_url}/tasks", timeout=self.timeout)
        r.raise_for_status()
        return r.json()["tasks"]

    def ground_truth(self) -> dict[str, Any]:
        """Return the clean (target) dataset for the current session."""
        r = requests.get(f"{self.base_url}/ground_truth", timeout=self.timeout)
        r.raise_for_status()
        return r.json()

    def validate(self) -> dict[str, Any]:
        """Run heuristic agent on all tasks and return scores."""
        r = requests.get(f"{self.base_url}/validate", timeout=120)
        r.raise_for_status()
        return r.json()

    # ── Context manager support ────────────────────────────────────────────

    def __enter__(self) -> "CleanRLClient":
        return self

    def __exit__(self, *_: Any) -> None:
        self.close()

    def __repr__(self) -> str:
        return f"CleanRLClient(base_url={self.base_url!r})"