from __future__ import annotations import re from enum import Enum from typing import Any, Dict, List, Optional from openenv.core.env_server.types import Action, Observation, State from pydantic import Field, model_validator class ActionType(str, Enum): EXECUTE = "execute" DEFER = "defer" DELEGATE = "delegate" OPTIMIZE = "optimize" class TaskStatus(str, Enum): PENDING = "pending" IN_PROGRESS = "in_progress" COMPLETED = "completed" DEFERRED = "deferred" FAILED = "failed" class DifficultyLevel(str, Enum): EASY = "easy" MEDIUM = "medium" HARD = "hard" def _coerce_action_type_from_text(text: str) -> ActionType: """ Convert loose text from the web UI into the closest valid action type. The default web interface validates form fields against the Action model before our environment can run `message_to_action()`, so we accept friendly free-form input here and normalize it. """ normalized = (text or "").strip().lower() if not normalized: return ActionType.EXECUTE keyword_map = ( (ActionType.OPTIMIZE, ("optimize", "optimise", "tune", "analyze", "analyse")), (ActionType.DELEGATE, ("delegate", "assign", "handoff", "hand off", "offload")), (ActionType.DEFER, ("defer", "later", "wait", "skip", "postpone")), (ActionType.EXECUTE, ("execute", "run", "do", "complete", "process", "start")), ) for action_type, keywords in keyword_map: if any(keyword in normalized for keyword in keywords): return action_type return ActionType.EXECUTE def _extract_task_id_from_text(text: str) -> Optional[int]: if not text: return None explicit_match = re.search(r"(?:task|id|#)\s*(\d+)", text, flags=re.IGNORECASE) if explicit_match: return int(explicit_match.group(1)) loose_match = re.search(r"\b(\d+)\b", text) if loose_match: return int(loose_match.group(1)) return None class TaskInfo(object): """Lightweight task descriptor (not a BaseModel to avoid nesting issues).""" def __init__( self, task_id: int, name: str, priority: float, deadline: int, uncertainty: float, value: float, required_energy: float, required_budget: float, category: str, status: str = "pending", ): self.task_id = task_id self.name = name self.priority = priority self.deadline = deadline self.uncertainty = uncertainty self.value = value self.required_energy = required_energy self.required_budget = required_budget self.category = category self.status = status def to_dict(self) -> Dict[str, Any]: return { "task_id": self.task_id, "name": self.name, "priority": round(self.priority, 3), "deadline": self.deadline, "uncertainty": round(self.uncertainty, 3), "value": round(self.value, 3), "required_energy": round(self.required_energy, 3), "required_budget": round(self.required_budget, 3), "category": self.category, "status": self.status, } class AetherTaskFlowAction(Action): """ Action for the AETHER-TaskFlow environment. The agent selects a task by ID and decides how to act on it. """ @model_validator(mode="before") @classmethod def normalize_web_input(cls, data: Any) -> Any: """ Make the default web form resilient to casual text input. Examples that should validate cleanly: {"action_type": "hi"} {"action_type": "execute task 3"} {"message": "delegate 2"} """ if isinstance(data, str): data = {"action_type": data} if not isinstance(data, dict): return data payload = dict(data) raw_message = payload.get("message") raw_action_type = payload.get("action_type") # When the UI sends a free-form message, treat it as action text. if isinstance(raw_message, str) and not raw_action_type: raw_action_type = raw_message payload.pop("message", None) if isinstance(raw_action_type, str): parsed_task_id = _extract_task_id_from_text(raw_action_type) payload["action_type"] = _coerce_action_type_from_text(raw_action_type).value if payload.get("task_id") in (None, "") and parsed_task_id is not None: payload["task_id"] = parsed_task_id if payload.get("reasoning") in (None, "") and raw_action_type.strip(): payload["reasoning"] = f"parsed from '{raw_action_type.strip()[:80]}'" if payload.get("task_id") in (None, ""): payload["task_id"] = 0 return payload action_type: ActionType = Field( ..., description=( "How to act on the selected task. " "'execute': consume resources and complete the task; " "'defer': postpone to a later step (low penalty); " "'delegate': offload at reduced reward but no resource cost; " "'optimize': reduce task uncertainty before execution." ), ) task_id: int = Field( ..., ge=0, description="ID of the task to act on (from the current task list).", ) reasoning: Optional[str] = Field( default=None, max_length=500, description="Optional agent reasoning for this action (logged but not scored).", ) class AetherTaskFlowObservation(Observation): """ Observation returned after each step in the AETHER-TaskFlow environment. Contains all information the agent needs to make the next decision. """ # Task queue tasks: List[Dict[str, Any]] = Field( default_factory=list, description="Current list of pending/deferred tasks as dicts.", ) # Resource pool time_remaining: int = Field( default=10, ge=0, description="Time steps remaining in this episode.", ) energy_remaining: float = Field( default=10.0, ge=0.0, description="Energy units remaining.", ) budget_remaining: float = Field( default=50.0, ge=0.0, description="Budget units remaining.", ) # System health system_health: float = Field( default=1.0, ge=0.0, le=1.0, description="Overall system health [0,1]. Drops on overload or missed deadlines.", ) # Episode progress step_number: int = Field( default=0, ge=0, description="Current step number in this episode.", ) tasks_completed: int = Field( default=0, ge=0, description="Total tasks completed so far.", ) tasks_failed: int = Field( default=0, ge=0, description="Total tasks that missed their deadline.", ) cumulative_value: float = Field( default=0.0, description="Total value accumulated so far.", ) # Last action feedback last_action_type: Optional[str] = Field( default=None, description="Action type taken in the previous step.", ) last_action_task_id: Optional[int] = Field( default=None, description="Task ID acted on in the previous step.", ) last_action_outcome: Optional[str] = Field( default=None, description="Human-readable outcome of the last action.", ) # Episode info difficulty: str = Field( default="easy", description="Current task difficulty level.", ) episode_id: Optional[str] = Field( default=None, description="Unique identifier for this episode.", ) class AetherTaskFlowState(State): """ Internal state of the AETHER-TaskFlow environment. This is the ground truth state used by the grader. """ difficulty: str = Field(default="easy") tasks: List[Dict[str, Any]] = Field(default_factory=list) completed_tasks: List[Dict[str, Any]] = Field(default_factory=list) failed_tasks: List[Dict[str, Any]] = Field(default_factory=list) deferred_tasks: List[Dict[str, Any]] = Field(default_factory=list) resources: Dict[str, float] = Field( default_factory=lambda: {"time": 10.0, "energy": 10.0, "budget": 50.0} ) initial_resources: Dict[str, float] = Field( default_factory=lambda: {"time": 10.0, "energy": 10.0, "budget": 50.0} ) system_health: float = Field(default=1.0) cumulative_value: float = Field(default=0.0) cumulative_reward: float = Field(default=0.0) tasks_completed: int = Field(default=0) tasks_failed: int = Field(default=0) episode_done: bool = Field(default=False) seed: Optional[int] = Field(default=None)