openadapt-ml 0.1.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

openadapt_ml/schema/episode.py

@@ -0,0 +1,470 @@
+"""
+Episode Schema for GUI Trajectory Data
+
+Canonical contract for episode/demonstration data in GUI automation. Designed for
+interoperability across training pipelines, benchmarks, and human demonstrations.
+
+Features:
+- Pydantic models with runtime validation
+- JSON Schema export for language-agnostic tooling
+- Supports pixel coordinates AND normalized (0-1) coordinates
+- Extensible via `raw` and `metadata` fields
+- Converters for common formats (WAA, WebArena, etc.)
+
+Quick Start:
+    from openadapt_ml.schema import Episode, Step, Action, Observation, ActionType
+
+    episode = Episode(
+        episode_id="demo_001",
+        instruction="Open the Settings app and enable Dark Mode",
+        steps=[
+            Step(
+                step_index=0,
+                observation=Observation(screenshot_path="step_0.png"),
+                action=Action(
+                    type=ActionType.CLICK,
+                    coordinates={"x": 512, "y": 384},
+                    # Or use normalized coords for resolution independence:
+                    # normalized_coordinates=(0.5, 0.375),
+                ),
+                reasoning="Click on Settings icon",
+            ),
+        ],
+        success=True,
+    )
+
+    # Validate any dict against the schema
+    from openadapt_ml.schema import validate_episode
+    is_valid, error = validate_episode(data)
+
+    # Export JSON Schema for external tools
+    from openadapt_ml.schema import export_json_schema
+    export_json_schema("episode.schema.json")
+
+Schema Version: 1.0.0
+- Core models: Episode, Step, Action, Observation
+- 24 action types covering mouse, keyboard, navigation, and system actions
+- Support for both pixel and normalized coordinates
+- Extension points: raw, metadata fields
+
+Evolution Policy (SemVer):
+- PATCH (1.0.x): Documentation, bug fixes (no schema changes)
+- MINOR (1.x.0): New optional fields with defaults (backward compatible)
+- MAJOR (x.0.0): Breaking changes (field removal, type changes, new required fields)
+
+Migration Guide:
+- MINOR bumps: No action needed, old data validates
+- MAJOR bumps: Use converters or migration scripts (provided in release notes)
+"""
+
+from __future__ import annotations
+
+import json
+from datetime import datetime
+from enum import Enum
+from pathlib import Path
+from typing import Any, Literal, Optional, Union
+
+from pydantic import BaseModel, Field, field_validator, model_validator
+
+
+# Schema version - follows semver
+SCHEMA_VERSION = "1.0.0"
+
+
+class ActionType(str, Enum):
+    """Supported action types for GUI automation."""
+
+    # Mouse actions
+    CLICK = "click"
+    DOUBLE_CLICK = "double_click"
+    RIGHT_CLICK = "right_click"
+    DRAG = "drag"
+    SCROLL = "scroll"
+    HOVER = "hover"
+
+    # Keyboard actions
+    TYPE = "type"
+    KEY = "key"
+    HOTKEY = "hotkey"
+
+    # Combined/special actions
+    CLICK_AND_TYPE = "click_and_type"
+    WAIT = "wait"
+    SCREENSHOT = "screenshot"
+
+    # Navigation (for web)
+    GOTO = "goto"
+    BACK = "back"
+    FORWARD = "forward"
+    REFRESH = "refresh"
+
+    # System actions
+    OPEN_APP = "open_app"
+    CLOSE_APP = "close_app"
+    SELECT_MONITOR = "select_monitor"  # Multi-monitor: focus a specific display
+    WINDOW_FOCUS = "window_focus"  # Focus a specific window
+    WINDOW_RESIZE = "window_resize"  # Resize window
+    WINDOW_MOVE = "window_move"  # Move window
+
+    # Meta actions
+    DONE = "done"
+    FAIL = "fail"
+
+
+class BenchmarkSource(str, Enum):
+    """Source benchmark/dataset for the episode."""
+
+    WAA = "waa"  # Windows Agent Arena
+    WEBARENA = "webarena"
+    OSWORLD = "osworld"
+    MINIWOB = "miniwob"
+    HUMAN = "human"  # Human demonstration
+    SYNTHETIC = "synthetic"  # Generated/augmented
+
+
+class Coordinates(BaseModel):
+    """Screen coordinates for mouse actions."""
+
+    x: int = Field(..., description="X coordinate (pixels from left)")
+    y: int = Field(..., description="Y coordinate (pixels from top)")
+
+    @field_validator("x", "y")
+    @classmethod
+    def validate_non_negative(cls, v: int) -> int:
+        if v < 0:
+            raise ValueError("Coordinates must be non-negative")
+        return v
+
+
+class BoundingBox(BaseModel):
+    """Bounding box for UI elements."""
+
+    x: int = Field(..., description="Left edge X coordinate")
+    y: int = Field(..., description="Top edge Y coordinate")
+    width: int = Field(..., ge=0, description="Width in pixels")
+    height: int = Field(..., ge=0, description="Height in pixels")
+
+    @property
+    def center(self) -> Coordinates:
+        """Get center point of bounding box."""
+        return Coordinates(x=self.x + self.width // 2, y=self.y + self.height // 2)
+
+
+class UIElement(BaseModel):
+    """UI element information from accessibility tree or DOM."""
+
+    role: Optional[str] = Field(
+        None, description="Element role (button, textbox, etc.)"
+    )
+    name: Optional[str] = Field(None, description="Element accessible name")
+    value: Optional[str] = Field(None, description="Element value (for inputs)")
+    bounds: Optional[BoundingBox] = Field(None, description="Element bounding box")
+    element_id: Optional[str] = Field(None, description="Unique element identifier")
+    xpath: Optional[str] = Field(None, description="XPath selector (web)")
+    selector: Optional[str] = Field(None, description="CSS selector (web)")
+    automation_id: Optional[str] = Field(None, description="Automation ID (Windows)")
+
+
+class Action(BaseModel):
+    """An action taken by the agent."""
+
+    type: ActionType = Field(..., description="Type of action")
+
+    # Mouse action parameters
+    coordinates: Optional[Coordinates] = Field(
+        None, description="Target coordinates for mouse actions"
+    )
+    start_coordinates: Optional[Coordinates] = Field(
+        None, description="Start coordinates for drag actions"
+    )
+    end_coordinates: Optional[Coordinates] = Field(
+        None, description="End coordinates for drag actions"
+    )
+    scroll_direction: Optional[Literal["up", "down", "left", "right"]] = Field(
+        None, description="Scroll direction"
+    )
+    scroll_amount: Optional[int] = Field(None, description="Scroll amount in pixels")
+
+    # Keyboard action parameters
+    text: Optional[str] = Field(None, description="Text to type")
+    key: Optional[str] = Field(None, description="Key to press (e.g., 'enter', 'tab')")
+    modifiers: Optional[list[str]] = Field(
+        None, description="Modifier keys (ctrl, alt, shift, meta)"
+    )
+
+    # Element targeting (alternative to coordinates)
+    element: Optional[UIElement] = Field(
+        None, description="Target element (for element-based actions)"
+    )
+
+    # Additional parameters
+    url: Optional[str] = Field(None, description="URL for goto action")
+    app_name: Optional[str] = Field(None, description="Application name for open/close")
+    duration: Optional[float] = Field(
+        None, description="Duration in seconds (for wait)"
+    )
+    monitor_id: Optional[int] = Field(
+        None, description="Monitor ID for select_monitor action"
+    )
+    window_title: Optional[str] = Field(
+        None, description="Window title for window_focus action"
+    )
+
+    # Normalized coordinates (0.0-1.0) - alternative to pixel coordinates
+    # Useful for resolution-independent recordings
+    normalized_coordinates: Optional[tuple[float, float]] = Field(
+        None, description="Normalized (x, y) coordinates (0.0-1.0 range)"
+    )
+    normalized_start: Optional[tuple[float, float]] = Field(
+        None, description="Normalized start coordinates for drag (0.0-1.0 range)"
+    )
+    normalized_end: Optional[tuple[float, float]] = Field(
+        None, description="Normalized end coordinates for drag (0.0-1.0 range)"
+    )
+
+    # Raw/original action data
+    raw: Optional[dict[str, Any]] = Field(
+        None, description="Original action data from source format"
+    )
+
+    @model_validator(mode="after")
+    def validate_action_params(self) -> "Action":
+        """Validate that required parameters are present for action type."""
+        if self.type in {
+            ActionType.CLICK,
+            ActionType.DOUBLE_CLICK,
+            ActionType.RIGHT_CLICK,
+        }:
+            if self.coordinates is None and self.element is None:
+                # Allow missing coordinates - can be inferred from context
+                pass
+
+        if self.type == ActionType.TYPE and self.text is None:
+            raise ValueError("TYPE action requires 'text' parameter")
+
+        if self.type == ActionType.KEY and self.key is None:
+            raise ValueError("KEY action requires 'key' parameter")
+
+        if self.type == ActionType.GOTO and self.url is None:
+            raise ValueError("GOTO action requires 'url' parameter")
+
+        return self
+
+
+class Observation(BaseModel):
+    """An observation of the environment state."""
+
+    # Visual observation
+    screenshot_path: Optional[str] = Field(
+        None, description="Path to screenshot image file"
+    )
+    screenshot_base64: Optional[str] = Field(
+        None, description="Base64-encoded screenshot (for inline storage)"
+    )
+
+    # Structured observations
+    a11y_tree: Optional[dict[str, Any]] = Field(
+        None, description="Accessibility tree snapshot"
+    )
+    dom: Optional[str] = Field(None, description="DOM HTML snapshot (web)")
+
+    # Window/screen info
+    window_title: Optional[str] = Field(None, description="Active window title")
+    app_name: Optional[str] = Field(
+        None, description="Application name (e.g., 'Chrome', 'System Settings')"
+    )
+    url: Optional[str] = Field(None, description="Current URL (for web apps)")
+    screen_size: Optional[tuple[int, int]] = Field(
+        None, description="Screen dimensions (width, height)"
+    )
+
+    # Focused element
+    focused_element: Optional[UIElement] = Field(
+        None, description="Currently focused UI element"
+    )
+
+    # Additional metadata
+    timestamp: Optional[float] = Field(None, description="Unix timestamp")
+    raw: Optional[dict[str, Any]] = Field(
+        None, description="Original observation data from source format"
+    )
+
+
+class Step(BaseModel):
+    """A single step in an episode (observation -> action pair)."""
+
+    step_index: int = Field(..., ge=0, description="Step number (0-indexed)")
+
+    # Core data
+    observation: Observation = Field(..., description="State observation before action")
+    action: Action = Field(..., description="Action taken")
+
+    # Agent reasoning (for demos/training)
+    reasoning: Optional[str] = Field(
+        None, description="Agent's reasoning for the action (chain-of-thought)"
+    )
+
+    # Outcome
+    reward: Optional[float] = Field(None, description="Reward signal (if available)")
+    done: Optional[bool] = Field(
+        None, description="Whether episode ended after this step"
+    )
+
+    # Timing
+    timestamp: Optional[float] = Field(None, description="Unix timestamp of action")
+    duration_ms: Optional[int] = Field(
+        None, description="Time taken for this step in milliseconds"
+    )
+
+
+class Episode(BaseModel):
+    """A complete episode/demonstration for GUI automation.
+
+    This is the canonical format for storing and exchanging GUI trajectory data.
+    All benchmark-specific formats should be converted to/from this format.
+    """
+
+    # Schema metadata
+    schema_version: str = Field(
+        default=SCHEMA_VERSION, description="Schema version for compatibility checking"
+    )
+
+    # Episode identification
+    episode_id: str = Field(..., description="Unique episode identifier")
+    task_id: Optional[str] = Field(None, description="Task identifier (from benchmark)")
+
+    # Task specification
+    instruction: str = Field(..., description="Natural language task instruction")
+    goal: Optional[str] = Field(
+        None, description="Detailed goal description (if different from instruction)"
+    )
+
+    # Episode data
+    steps: list[Step] = Field(..., description="Sequence of steps in the episode")
+
+    # Outcome
+    success: Optional[bool] = Field(
+        None, description="Whether task was completed successfully"
+    )
+    final_reward: Optional[float] = Field(None, description="Final reward/score")
+
+    # Provenance
+    source: Optional[BenchmarkSource] = Field(
+        None, description="Source benchmark/dataset"
+    )
+    source_file: Optional[str] = Field(None, description="Original source file path")
+
+    # Metadata
+    created_at: Optional[datetime] = Field(
+        default_factory=datetime.utcnow, description="When episode was created/recorded"
+    )
+    agent_model: Optional[str] = Field(
+        None, description="Model that generated this episode (e.g., 'gpt-4o')"
+    )
+    environment: Optional[str] = Field(
+        None, description="Environment info (OS, browser, etc.)"
+    )
+    tags: Optional[list[str]] = Field(None, description="Tags for categorization")
+
+    # Extension point for benchmark-specific data
+    metadata: Optional[dict[str, Any]] = Field(
+        None, description="Additional metadata from source"
+    )
+
+    @property
+    def num_steps(self) -> int:
+        """Number of steps in the episode."""
+        return len(self.steps)
+
+    @property
+    def action_types(self) -> list[ActionType]:
+        """List of action types in this episode."""
+        return [step.action.type for step in self.steps]
+
+    def to_json(self, indent: int = 2) -> str:
+        """Serialize to JSON string."""
+        return self.model_dump_json(indent=indent)
+
+    @classmethod
+    def from_json(cls, json_str: str) -> "Episode":
+        """Deserialize from JSON string."""
+        return cls.model_validate_json(json_str)
+
+    @classmethod
+    def json_schema(cls) -> dict[str, Any]:
+        """Get JSON Schema for Episode format."""
+        return cls.model_json_schema()
+
+
+# ============================================================================
+# Utility Functions
+# ============================================================================
+
+
+def validate_episode(data: dict[str, Any]) -> tuple[bool, Optional[str]]:
+    """Validate episode data against schema.
+
+    Args:
+        data: Episode data as dictionary
+
+    Returns:
+        Tuple of (is_valid, error_message)
+    """
+    try:
+        Episode.model_validate(data)
+        return True, None
+    except Exception as e:
+        return False, str(e)
+
+
+def load_episode(path: Union[str, Path]) -> Episode:
+    """Load episode from JSON file.
+
+    Args:
+        path: Path to JSON file
+
+    Returns:
+        Episode instance
+    """
+    path = Path(path)
+    with open(path, "r") as f:
+        data = json.load(f)
+
+    episode = Episode.model_validate(data)
+
+    # Set source_file if not already set
+    if episode.source_file is None:
+        episode = episode.model_copy(update={"source_file": str(path)})
+
+    return episode
+
+
+def save_episode(episode: Episode, path: Union[str, Path], indent: int = 2) -> None:
+    """Save episode to JSON file.
+
+    Args:
+        episode: Episode to save
+        path: Output path
+        indent: JSON indentation
+    """
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    with open(path, "w") as f:
+        f.write(episode.to_json(indent=indent))
+
+
+def export_json_schema(path: Union[str, Path]) -> None:
+    """Export JSON Schema to file for documentation/tooling.
+
+    Args:
+        path: Output path for schema file
+    """
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    schema = Episode.json_schema()
+
+    with open(path, "w") as f:
+        json.dump(schema, f, indent=2)