client.py

"""Client for the Data Extraction Environment."""

from dataclasses import dataclass
from typing import Any
from openenv.core.env_client import EnvClient
from openenv.core.env_server.types import State

try:
    from .models import ExtractionAction, ExtractionObservation
except ImportError:
    from models import ExtractionAction, ExtractionObservation


@dataclass
class StepResult:
    """Simple container for the result of a reset() or step() call."""
    observation: Any
    reward: float = 0.0
    done: bool = False


class DataExtractionEnv(EnvClient):
    """Client for the Data Extraction environment."""
    action_type = ExtractionAction
    observation_type = ExtractionObservation

    def _step_payload(self, action: ExtractionAction) -> dict:
        return action.model_dump()

    def _parse_result(self, data: dict) -> StepResult:
        obs_data = data.get("observation", {})
        obs = ExtractionObservation(**obs_data)
        return StepResult(
            observation=obs,
            reward=data.get("reward", 0.0),
            done=data.get("done", False),
        )

    def _parse_state(self, data: dict) -> State:
        return State(**data)