client.py

"""Client for OversightArena environment."""
from __future__ import annotations

import requests
from models import OversightAction, OversightObservation


class OversightArenaClient:
    def __init__(self, base_url: str = "http://localhost:8000"):
        self.base_url = base_url.rstrip("/")

    def health(self) -> dict:
        return requests.get(f"{self.base_url}/health").json()

    def tasks(self) -> dict:
        return requests.get(f"{self.base_url}/tasks").json()

    def reset(self, task_id: str = "easy") -> OversightObservation:
        r = requests.post(
            f"{self.base_url}/reset",
            json={"task_id": task_id}
        )
        r.raise_for_status()
        return OversightObservation(**r.json()["observation"])

    def step(self, action: OversightAction) -> tuple[OversightObservation, float, bool]:
        r = requests.post(
            f"{self.base_url}/step",
            json={"action": action.model_dump()}
        )
        r.raise_for_status()
        data = r.json()
        obs = OversightObservation(**data["observation"])
        return obs, data["reward"], data["done"]

    def baseline(self) -> dict:
        return requests.get(f"{self.base_url}/baseline").json()


if __name__ == "__main__":
    client = OversightArenaClient()
    print("Health:", client.health())
    print("Tasks:", client.tasks())
    
    obs = client.reset("easy")
    print(f"Reset: step={obs.step_number}, flags={obs.flags_remaining}")
    
    action = OversightAction(
        action_type="flag",
        question_id=0,
        error_type="wrong_value",
        reasoning="The value appears inconsistent with source data",
        confidence=0.8
    )
    obs, reward, done = client.step(action)
    print(f"Step: reward={reward}, done={done}")