src/auditenv/client.py

from __future__ import annotations

from typing import Any

import requests


class AuditEnvClient:
    """Small HTTP client for OpenEnv-style interactions with the running server."""

    def __init__(self, base_url: str = "http://127.0.0.1:8000") -> None:
        self.base_url = base_url.rstrip("/")
        self.session = requests.Session()
        self.session_id: str | None = None

    def reset(self, task_id: str = "easy", seed: int = 42) -> dict[str, Any]:
        resp = self.session.post(f"{self.base_url}/reset", json={"task_id": task_id, "seed": seed}, timeout=30)
        resp.raise_for_status()
        payload = resp.json()
        self.session_id = payload.get("session_id")
        return payload

    def step(self, action: dict[str, Any]) -> dict[str, Any]:
        payload = dict(action)
        if self.session_id and "session_id" not in payload:
            payload["session_id"] = self.session_id

        resp = self.session.post(f"{self.base_url}/step", json=payload, timeout=30)
        resp.raise_for_status()
        return resp.json()

    def state(self) -> dict[str, Any]:
        resp = self.session.get(f"{self.base_url}/state", timeout=30)
        resp.raise_for_status()
        return resp.json()

    def close(self) -> None:
        self.session.close()


__all__ = ["AuditEnvClient"]