File size: 4,382 Bytes
b404de3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | from __future__ import annotations
import traceback
from typing import Any, Dict
from openenv.core.client_types import StepResult
from openenv.core.env_client import EnvClient
from cloudenv.models import CloudEnvAction, CloudEnvObservation, CloudEnvState
class ClientEnvironment(
EnvClient[CloudEnvAction, CloudEnvObservation, CloudEnvState]
):
# ---------------------------
# INIT ACTION
# ---------------------------
def initialize_environment(
self, *, task_id: str, episode_id: str
) -> StepResult[CloudEnvObservation]:
try:
action = CloudEnvAction(
task_id=task_id,
episode_id=episode_id,
service="internal", # safer than enum
operation="InitializeEnvironment",
payload={},
)
return self.step(action)
except Exception:
print("[CLIENT ERROR] initialize_environment failed", flush=True)
traceback.print_exc()
return StepResult(
observation=CloudEnvObservation(
episode_id=episode_id,
task_id=task_id,
step_count=0,
operation_result={"error": "init failed"},
reward=0.0,
done=True,
metadata={},
current_pipeline_state={},
),
reward=0.0,
done=True,
)
# ---------------------------
# SERIALIZATION
# ---------------------------
def _step_payload(self, action: CloudEnvAction) -> Dict[str, Any]:
try:
return action.model_dump(mode="json", exclude_none=True)
except Exception:
print(
"[CLIENT ERROR] step payload serialization failed", flush=True
)
traceback.print_exc()
return {}
# ---------------------------
# PARSE RESULT
# ---------------------------
def _parse_result(
self, payload: Dict[str, Any]
) -> StepResult[CloudEnvObservation]:
try:
payload = payload or {}
reward = payload.get("reward", 0.0)
observation = CloudEnvObservation(
episode_id=payload.get("episode_id", ""),
task_id=payload.get("task_id", ""),
step_count=payload.get("step_count", 0),
operation_result=payload.get("operation_result", {}),
reward=float(reward) if reward is not None else 0.0,
done=bool(payload.get("done", False)),
metadata=payload.get("metadata", {}),
current_pipeline_state=payload.get(
"current_pipeline_state", {}
),
)
return StepResult(
observation=observation,
reward=observation.reward,
done=observation.done,
)
except Exception:
print("[CLIENT ERROR] parse_result failed", flush=True)
traceback.print_exc()
return StepResult(
observation=CloudEnvObservation(
episode_id="",
task_id="",
step_count=0,
operation_result={"error": "parse failure"},
reward=0.0,
done=True,
metadata={},
current_pipeline_state={},
),
reward=0.0,
done=True,
)
# ---------------------------
# PARSE STATE
# ---------------------------
def _parse_state(self, payload: Dict[str, Any]) -> CloudEnvState:
try:
payload = payload or {}
return CloudEnvState(
task_id=payload.get("task_id"),
episode_id=payload.get("episode_id"),
step_count=payload.get("step_count", 0),
reward_history=payload.get("reward_history", []),
)
except Exception:
print("[CLIENT ERROR] parse_state failed", flush=True)
traceback.print_exc()
return CloudEnvState(
task_id=None,
episode_id=None,
step_count=0,
reward_history=[],
)
|