| |
| |
| |
|
|
| """ |
| OrgSim Environment Client. |
| |
| Connects to OrgSimEnvironment server via WebSocket for persistent sessions. |
| """ |
|
|
| from typing import Any, Dict |
| from openenv.core.client_types import StepResult |
| from openenv.core.env_client import EnvClient |
|
|
| from .models import OrgAction, OrgObservation, OrgState |
|
|
|
|
| class OrgSimEnv(EnvClient[OrgAction, OrgObservation, OrgState]): |
| """Client for the Organization Simulation Environment.""" |
|
|
| def _step_payload(self, action: OrgAction) -> Dict[str, Any]: |
| """Convert OrgAction to JSON payload.""" |
| return { |
| "action_type": action.action_type, |
| "target_id": action.target_id, |
| "payload": action.payload, |
| "metadata": action.metadata, |
| } |
|
|
| def _parse_result(self, payload: Dict[str, Any]) -> StepResult[OrgObservation]: |
| """Parse server response into StepResult.""" |
| obs_data = payload.get("observation", {}) |
|
|
| observation = OrgObservation( |
| my_agent_id=obs_data.get("my_agent_id", ""), |
| my_team=obs_data.get("my_team", ""), |
| my_role=obs_data.get("my_role", "member"), |
| available_tasks=obs_data.get("available_tasks", []), |
| active_task=obs_data.get("active_task"), |
| inbox=obs_data.get("inbox", []), |
| team_status=obs_data.get("team_status", {}), |
| resources=obs_data.get("resources", {}), |
| metrics=obs_data.get("metrics", {}), |
| 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[str, Any]) -> OrgState: |
| """Parse server response into OrgState.""" |
| return OrgState( |
| episode_id=payload.get("episode_id"), |
| step_count=payload.get("step_count", 0), |
| tasks=payload.get("tasks", {}), |
| team_members=payload.get("team_members", {}), |
| resource_status=payload.get("resource_status", {}), |
| ) |
|
|
