client.py

"""Python client for the Fake Gang Detection OpenEnv environment."""

from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Any, Dict, Optional

try:
    import requests
except ImportError:
    requests = None  # type: ignore

from models import (
    AccountProfile,
    FakeGangAction,
    FakeGangObservation,
    FakeGangState,
    ActionType,
)


# ---------------------------------------------------------------------------
# Result container
# ---------------------------------------------------------------------------

@dataclass
class StepResult:
    observation: FakeGangObservation
    done: bool
    reward: Optional[float]
    message: str


# ---------------------------------------------------------------------------
# Sync HTTP client
# ---------------------------------------------------------------------------

class FakeGangEnvClient:
    """Synchronous HTTP client for the Fake Gang Detection environment."""

    def __init__(self, base_url: str = "http://localhost:8000") -> None:
        if requests is None:
            raise ImportError("Install 'requests' to use FakeGangEnvClient.")
        self.base_url = base_url.rstrip("/")
        self._session = requests.Session()

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def reset(
        self,
        task: str = "easy",
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
    ) -> StepResult:
        payload = {"task": task}
        if seed is not None:
            payload["seed"] = seed
        if episode_id is not None:
            payload["episode_id"] = episode_id
        resp = self._post("/reset", payload)
        return self._parse_result(resp)

    def step(self, action: FakeGangAction) -> StepResult:
        resp = self._post("/step", action.model_dump())
        return self._parse_result(resp)

    def state(self) -> FakeGangState:
        resp = self._session.get(f"{self.base_url}/state")
        resp.raise_for_status()
        return FakeGangState(**resp.json())

    def health(self) -> Dict[str, str]:
        resp = self._session.get(f"{self.base_url}/health")
        resp.raise_for_status()
        return resp.json()

    # ------------------------------------------------------------------
    # Convenience shortcuts
    # ------------------------------------------------------------------

    def inspect(self, account_id: str) -> StepResult:
        return self.step(FakeGangAction(action_type=ActionType.INSPECT, account_id=account_id))

    def investigate_network(self, account_id: str) -> StepResult:
        return self.step(FakeGangAction(action_type=ActionType.INVESTIGATE_NETWORK, account_id=account_id))

    def flag(self, account_id: str) -> StepResult:
        return self.step(FakeGangAction(action_type=ActionType.FLAG, account_id=account_id))

    def unflag(self, account_id: str) -> StepResult:
        return self.step(FakeGangAction(action_type=ActionType.UNFLAG, account_id=account_id))

    def submit(self) -> StepResult:
        return self.step(FakeGangAction(action_type=ActionType.SUBMIT))

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _post(self, path: str, payload: Dict[str, Any]) -> Dict[str, Any]:
        resp = self._session.post(f"{self.base_url}{path}", json=payload)
        resp.raise_for_status()
        return resp.json()

    def _parse_result(self, payload: Dict[str, Any]) -> StepResult:
        obs_data = payload["observation"]
        profiles = [AccountProfile(**p) for p in obs_data.get("visible_accounts", [])]
        obs = FakeGangObservation(
            done=obs_data.get("done", False),
            reward=obs_data.get("reward"),
            visible_accounts=profiles,
            visible_account_ids=obs_data.get("visible_account_ids", []),
            flagged_ids=obs_data.get("flagged_ids", []),
            inspected_ids=obs_data.get("inspected_ids", []),
            graph_edges=obs_data.get("graph_edges", {}),
            steps_remaining=obs_data.get("steps_remaining", 0),
            evasion_triggered=obs_data.get("evasion_triggered", False),
            evasion_count=obs_data.get("evasion_count", 0),
            task=obs_data.get("task", "easy"),
            message=obs_data.get("message", ""),
            suspect_ids=obs_data.get("suspect_ids", []),
        )
        return StepResult(
            observation=obs,
            done=payload.get("done", False),
            reward=payload.get("reward"),
            message=payload.get("message", ""),
        )

    # ------------------------------------------------------------------
    # Context manager support
    # ------------------------------------------------------------------

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

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