client.py

"""
PermitPathfinder Environment Client.

Thin typed wrapper over openenv-core's EnvClient. Connect to a running
PermitPathfinder server via base_url, or spin one up from a Docker image.

Example:
    >>> with PermitEnv(base_url="http://localhost:8000") as env:
    ...     result = env.reset()
    ...     print(result.observation.available_actions)
    ...     result = env.step(PermitAction(action_type="list"))

    >>> # Or launch a container automatically:
    >>> env = PermitEnv.from_docker_image("permit_env-env:latest")
"""

from typing import Dict

from openenv.core import EnvClient
from openenv.core.client_types import StepResult
from openenv.core.env_server.types import State

try:
    from .models import PermitAction, PermitObservation
except ImportError:  # running as flat script / container
    from models import PermitAction, PermitObservation


class PermitEnv(EnvClient[PermitAction, PermitObservation, State]):
    """Client for the PermitPathfinder environment."""

    def _step_payload(self, action: PermitAction) -> Dict:
        return {
            "action_type": action.action_type,
            "permit_id": action.permit_id,
        }

    def _parse_result(self, payload: Dict) -> StepResult[PermitObservation]:
        obs_data = payload.get("observation", {}) or {}
        observation = PermitObservation(
            message=obs_data.get("message", ""),
            permits=obs_data.get("permits", {}) or {},
            budget_remaining=obs_data.get("budget_remaining", 0.0) or 0.0,
            wasted_submissions=obs_data.get("wasted_submissions", 0) or 0,
            last_action_error=obs_data.get("last_action_error"),
            available_actions=obs_data.get("available_actions", []) or [],
            task_name=obs_data.get("task_name", "easy_foodtruck"),
            done=payload.get("done", False),
            reward=payload.get("reward"),
            metadata=obs_data.get("metadata", {}) or {},
        )
        return StepResult(
            observation=observation,
            reward=payload.get("reward"),
            done=payload.get("done", False),
        )

    def _parse_state(self, payload: Dict) -> State:
        return State(
            episode_id=payload.get("episode_id"),
            step_count=payload.get("step_count", 0),
        )