client.py

"""Shopify Store Audit Environment Client."""

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 ShopifyStoreAuditAction, ShopifyStoreAuditObservation
except ImportError:
    from models import ShopifyStoreAuditAction, ShopifyStoreAuditObservation


class ShopifyStoreAuditEnv(
    EnvClient[ShopifyStoreAuditAction, ShopifyStoreAuditObservation, State]
):
    """
    Client for the Shopify Store Audit & Remediation Environment.

    Example:
        >>> async with ShopifyStoreAuditEnv(base_url="http://localhost:8000") as client:
        ...     result = await client.reset(task="product_listing_qa")
        ...     print(result.observation.message)
        ...
        ...     result = await client.step(ShopifyStoreAuditAction(
        ...         command="query_store_health", params={}
        ...     ))

    Example with Docker:
        >>> env = await ShopifyStoreAuditEnv.from_docker_image(
        ...     "shopify-store-audit:latest"
        ... )
    """

    def _step_payload(self, action: ShopifyStoreAuditAction) -> Dict:
        return {
            "command": action.command,
            "params": action.params,
        }

    def _parse_result(
        self, payload: Dict
    ) -> StepResult[ShopifyStoreAuditObservation]:
        obs_data = payload.get("observation", {})
        observation = ShopifyStoreAuditObservation(
            message=obs_data.get("message", ""),
            data=obs_data.get("data", {}),
            issues_remaining=obs_data.get("issues_remaining", 0),
            issues_fixed=obs_data.get("issues_fixed", 0),
            total_issues=obs_data.get("total_issues", 0),
            store_health_score=obs_data.get("store_health_score", 0.0),
            available_commands=obs_data.get("available_commands", []),
            task_name=obs_data.get("task_name", ""),
            done=payload.get("done", False),
            reward=payload.get("reward"),
            metadata=obs_data.get("metadata", {}),
        )
        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),
        )