cace_env/client.py

"""
cace_env/client.py
CACEEnvClient — typed async/sync client for the CACE environment.
Inherits from openenv.core.EnvClient.

Usage (async):
    async with CACEEnvClient(base_url="ws://localhost:7860") as env:
        result = await env.reset()
        result = await env.step(CACEAction(action_int=0))

Usage (sync):
    with CACEEnvClient(base_url="ws://localhost:7860").sync() as env:
        result = env.reset()
        result = env.step(CACEAction(action_int=0))

Usage (from Docker):
    env = await CACEEnvClient.from_docker_image("cace-env:latest")
    async with env:
        result = await env.reset()

Usage (from HF Space):
    env = await CACEEnvClient.from_env("YOUR_USERNAME/cace-env")
    async with env:
        result = await env.reset()
"""

from typing import Any, Dict, Optional
from openenv.core.env_client import EnvClient
from cace_env.models import CACEAction, CACEObservation, CACEState


class CACEEnvClient(EnvClient[CACEAction, CACEObservation, CACEState]):
    """
    Typed async client for the Cultural Context Arbitration Environment.

    Implements the two abstract parse methods required by EnvClient:
        _parse_result(payload) → StepResult[CACEObservation]
        _parse_state(payload)  → CACEState
    """

    def _parse_result(self, payload: Dict[str, Any]):
        """Parse server step/reset response into StepResult[CACEObservation]."""
        from openenv.core.env_client import StepResult
        obs_data = payload.get("observation", payload)
        if isinstance(obs_data, str):
            obs = CACEObservation(
                observation=obs_data,
                case_id=payload.get("case_id", "unknown"),
                language=payload.get("language", "Unknown"),
                region=payload.get("region", "Unknown"),
                complexity=payload.get("complexity", "medium"),
                done=payload.get("done", False),
                reward=payload.get("reward", None),
            )
        else:
            obs = CACEObservation(**obs_data)

        return StepResult(
            observation=obs,
            done=payload.get("done", obs.done),
            reward=payload.get("reward", obs.reward),
            info=payload.get("info", {}),
        )

    def _parse_state(self, payload: Dict[str, Any]) -> CACEState:
        """Parse server state response into CACEState."""
        return CACEState(**{k: v for k, v in payload.items() if k in CACEState.model_fields})

    # ── Convenience methods ───────────────────────────────────────────────────

    async def reset_v1(self, seed: Optional[int] = None):
        """Reset in V1 mode (single case)."""
        return await self.reset(mode="v1", seed=seed)

    async def reset_v2(self, seed: Optional[int] = None):
        """Reset in V2 mode (network batch of 20 posts)."""
        return await self.reset(mode="v2", seed=seed)

    async def decide(self, action_int: int, selected_indices: list[int] | None = None):
        """Shorthand for step with a moderation decision."""
        action = CACEAction(action_int=action_int, selected_indices=selected_indices)
        return await self.step(action)

    # ── Class methods for easy instantiation ──────────────────────────────────

    @classmethod
    async def from_local(cls, port: int = 7860, **kwargs):
        """Connect to a locally running server."""
        client = cls(base_url=f"ws://localhost:{port}", **kwargs)
        await client.connect()
        return client

    # Inherited from EnvClient:
    #   CACEEnvClient.from_docker_image("cace-env:latest")
    #   CACEEnvClient.from_env("username/cace-env")
    #   .sync()  → SyncEnvClient wrapper for synchronous use