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

openadapt_ml/schema/converters.py

@@ -0,0 +1,588 @@
+"""
+Converters for benchmark-specific episode formats.
+
+Supported formats:
+- WAA (Windows Agent Arena)
+- WebArena (coming soon)
+- OSWorld (coming soon)
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from pathlib import Path
+from typing import Any, Optional, Union
+
+from openadapt_ml.schema.episode import (
+    Action,
+    ActionType,
+    BenchmarkSource,
+    Coordinates,
+    Episode,
+    Observation,
+    Step,
+    UIElement,
+)
+
+
+# ============================================================================
+# WAA (Windows Agent Arena) Converter
+# ============================================================================
+
+
+def _parse_waa_action(action_str: str) -> tuple[ActionType, dict[str, Any]]:
+    """Parse WAA action string into ActionType and parameters.
+
+    WAA action format examples:
+    - pyautogui.click(100, 200)
+    - pyautogui.write('hello')
+    - pyautogui.press('enter')
+    - pyautogui.hotkey('ctrl', 'c')
+    - pyautogui.scroll(3)
+    - DONE
+    - FAIL
+    """
+    action_str = action_str.strip()
+
+    # Meta actions
+    if action_str == "DONE":
+        return ActionType.DONE, {}
+    if action_str == "FAIL":
+        return ActionType.FAIL, {}
+
+    # Parse pyautogui calls
+    match = re.match(r"pyautogui\.(\w+)\((.*)\)", action_str)
+    if not match:
+        # Try without pyautogui prefix
+        match = re.match(r"(\w+)\((.*)\)", action_str)
+
+    if match:
+        func_name = match.group(1).lower()
+        args_str = match.group(2)
+
+        # Parse arguments (handle strings with commas inside)
+        args = []
+        current_arg = ""
+        in_string = False
+        string_char = None
+
+        for char in args_str:
+            if char in "'\"" and not in_string:
+                in_string = True
+                string_char = char
+            elif char == string_char and in_string:
+                in_string = False
+                string_char = None
+            elif char == "," and not in_string:
+                if current_arg.strip():
+                    args.append(current_arg.strip())
+                current_arg = ""
+                continue
+            current_arg += char
+
+        if current_arg.strip():
+            args.append(current_arg.strip())
+
+        # Clean up string arguments
+        cleaned_args = []
+        for arg in args:
+            arg = arg.strip()
+            if (arg.startswith("'") and arg.endswith("'")) or (
+                arg.startswith('"') and arg.endswith('"')
+            ):
+                cleaned_args.append(arg[1:-1])
+            else:
+                try:
+                    cleaned_args.append(int(arg))
+                except ValueError:
+                    try:
+                        cleaned_args.append(float(arg))
+                    except ValueError:
+                        cleaned_args.append(arg)
+
+        # Map function to action type
+        if func_name == "click":
+            params = {}
+            if len(cleaned_args) >= 2:
+                params["coordinates"] = Coordinates(
+                    x=int(cleaned_args[0]), y=int(cleaned_args[1])
+                )
+            return ActionType.CLICK, params
+
+        elif func_name == "doubleclick":
+            params = {}
+            if len(cleaned_args) >= 2:
+                params["coordinates"] = Coordinates(
+                    x=int(cleaned_args[0]), y=int(cleaned_args[1])
+                )
+            return ActionType.DOUBLE_CLICK, params
+
+        elif func_name == "rightclick":
+            params = {}
+            if len(cleaned_args) >= 2:
+                params["coordinates"] = Coordinates(
+                    x=int(cleaned_args[0]), y=int(cleaned_args[1])
+                )
+            return ActionType.RIGHT_CLICK, params
+
+        elif func_name in ("write", "typewrite"):
+            return ActionType.TYPE, {"text": cleaned_args[0] if cleaned_args else ""}
+
+        elif func_name == "press":
+            return ActionType.KEY, {"key": cleaned_args[0] if cleaned_args else ""}
+
+        elif func_name == "hotkey":
+            if len(cleaned_args) >= 2:
+                return ActionType.HOTKEY, {
+                    "key": cleaned_args[-1],
+                    "modifiers": list(cleaned_args[:-1]),
+                }
+            return ActionType.KEY, {"key": cleaned_args[0] if cleaned_args else ""}
+
+        elif func_name == "scroll":
+            amount = cleaned_args[0] if cleaned_args else 0
+            direction = "up" if amount > 0 else "down"
+            return ActionType.SCROLL, {
+                "scroll_direction": direction,
+                "scroll_amount": abs(int(amount)) * 100,  # Convert to pixels
+            }
+
+        elif func_name == "moveto":
+            params = {}
+            if len(cleaned_args) >= 2:
+                params["coordinates"] = Coordinates(
+                    x=int(cleaned_args[0]), y=int(cleaned_args[1])
+                )
+            return ActionType.HOVER, params
+
+        elif func_name == "drag" or func_name == "dragto":
+            params = {}
+            if len(cleaned_args) >= 2:
+                params["end_coordinates"] = Coordinates(
+                    x=int(cleaned_args[0]), y=int(cleaned_args[1])
+                )
+            return ActionType.DRAG, params
+
+    # Fallback - treat as raw text if nothing matched
+    return ActionType.TYPE, {"text": action_str, "raw": {"original": action_str}}
+
+
+def from_waa_trajectory(
+    trajectory: list[dict[str, Any]],
+    task_info: dict[str, Any],
+    episode_id: Optional[str] = None,
+) -> Episode:
+    """Convert WAA trajectory format to Episode.
+
+    Args:
+        trajectory: List of WAA step dictionaries with keys like:
+            - screenshot_path: Path to screenshot
+            - action: Action string (pyautogui format)
+            - a11y_tree: Accessibility tree (optional)
+            - thought: Agent reasoning (optional)
+        task_info: Task metadata with keys like:
+            - id: Task ID
+            - instruction: Task instruction
+            - domain: Task domain (file_explorer, etc.)
+
+    Returns:
+        Episode instance
+    """
+    steps = []
+
+    for i, step_data in enumerate(trajectory):
+        # Parse observation
+        observation = Observation(
+            screenshot_path=step_data.get("screenshot_path"),
+            a11y_tree=step_data.get("a11y_tree"),
+            window_title=step_data.get("window_title"),
+            raw=step_data.get("observation_raw"),
+        )
+
+        # Parse action
+        action_str = step_data.get("action", "")
+        action_type, action_params = _parse_waa_action(action_str)
+
+        action = Action(
+            type=action_type,
+            raw={"original": action_str},
+            **action_params,
+        )
+
+        # Create step
+        step = Step(
+            step_index=i,
+            observation=observation,
+            action=action,
+            reasoning=step_data.get("thought") or step_data.get("reasoning"),
+            reward=step_data.get("reward"),
+            done=step_data.get("done"),
+        )
+        steps.append(step)
+
+    # Extract task info
+    task_id = task_info.get("id") or task_info.get("task_id")
+    instruction = task_info.get("instruction") or task_info.get("goal", "")
+
+    if episode_id is None:
+        episode_id = f"waa_{task_id}" if task_id else f"waa_episode_{id(trajectory)}"
+
+    return Episode(
+        episode_id=episode_id,
+        task_id=task_id,
+        instruction=instruction,
+        steps=steps,
+        success=task_info.get("success"),
+        source=BenchmarkSource.WAA,
+        metadata={
+            "domain": task_info.get("domain"),
+            "difficulty": task_info.get("difficulty"),
+            **{
+                k: v
+                for k, v in task_info.items()
+                if k
+                not in [
+                    "id",
+                    "task_id",
+                    "instruction",
+                    "goal",
+                    "success",
+                    "domain",
+                    "difficulty",
+                ]
+            },
+        },
+    )
+
+
+def to_waa_trajectory(episode: Episode) -> tuple[list[dict[str, Any]], dict[str, Any]]:
+    """Convert Episode to WAA trajectory format.
+
+    Args:
+        episode: Episode instance
+
+    Returns:
+        Tuple of (trajectory, task_info)
+    """
+    trajectory = []
+
+    for step in episode.steps:
+        step_data = {
+            "screenshot_path": step.observation.screenshot_path,
+            "a11y_tree": step.observation.a11y_tree,
+            "window_title": step.observation.window_title,
+        }
+
+        # Convert action back to pyautogui format
+        action = step.action
+        if action.raw and "original" in action.raw:
+            step_data["action"] = action.raw["original"]
+        else:
+            step_data["action"] = _action_to_pyautogui(action)
+
+        if step.reasoning:
+            step_data["thought"] = step.reasoning
+
+        if step.reward is not None:
+            step_data["reward"] = step.reward
+
+        if step.done is not None:
+            step_data["done"] = step.done
+
+        trajectory.append(step_data)
+
+    task_info = {
+        "id": episode.task_id,
+        "instruction": episode.instruction,
+        "success": episode.success,
+    }
+
+    if episode.metadata:
+        task_info.update(episode.metadata)
+
+    return trajectory, task_info
+
+
+def _action_to_pyautogui(action: Action) -> str:
+    """Convert Action to pyautogui string format."""
+    if action.type == ActionType.DONE:
+        return "DONE"
+    if action.type == ActionType.FAIL:
+        return "FAIL"
+
+    if action.type == ActionType.CLICK:
+        if action.coordinates:
+            return f"pyautogui.click({action.coordinates.x}, {action.coordinates.y})"
+        return "pyautogui.click()"
+
+    if action.type == ActionType.DOUBLE_CLICK:
+        if action.coordinates:
+            return (
+                f"pyautogui.doubleClick({action.coordinates.x}, {action.coordinates.y})"
+            )
+        return "pyautogui.doubleClick()"
+
+    if action.type == ActionType.RIGHT_CLICK:
+        if action.coordinates:
+            return (
+                f"pyautogui.rightClick({action.coordinates.x}, {action.coordinates.y})"
+            )
+        return "pyautogui.rightClick()"
+
+    if action.type == ActionType.TYPE:
+        text = action.text or ""
+        # Escape single quotes
+        text = text.replace("'", "\\'")
+        return f"pyautogui.write('{text}')"
+
+    if action.type == ActionType.KEY:
+        return f"pyautogui.press('{action.key}')"
+
+    if action.type == ActionType.HOTKEY:
+        modifiers = action.modifiers or []
+        keys = modifiers + [action.key]
+        keys_str = ", ".join(f"'{k}'" for k in keys)
+        return f"pyautogui.hotkey({keys_str})"
+
+    if action.type == ActionType.SCROLL:
+        amount = action.scroll_amount or 100
+        if action.scroll_direction in ("down", "right"):
+            amount = -amount
+        return f"pyautogui.scroll({amount // 100})"
+
+    if action.type == ActionType.HOVER:
+        if action.coordinates:
+            return f"pyautogui.moveTo({action.coordinates.x}, {action.coordinates.y})"
+        return "pyautogui.moveTo()"
+
+    if action.type == ActionType.DRAG:
+        if action.end_coordinates:
+            return f"pyautogui.dragTo({action.end_coordinates.x}, {action.end_coordinates.y})"
+        return "pyautogui.drag()"
+
+    return f"# Unknown action: {action.type}"
+
+
+# ============================================================================
+# Internal Format Converter (openadapt_ml.schemas.sessions)
+# ============================================================================
+
+
+def from_internal_episode(
+    internal_episode: Any,
+    episode_id: Optional[str] = None,
+) -> Episode:
+    """Convert from internal training format (openadapt_ml.schemas.sessions.Episode).
+
+    This converts from the dataclass-based format used by the training pipeline
+    to the Pydantic-based Episode format used for external interoperability.
+
+    Args:
+        internal_episode: An openadapt_ml.schemas.sessions.Episode instance
+        episode_id: Override episode ID (defaults to internal_episode.id)
+
+    Returns:
+        Episode instance in the new format
+    """
+    steps = []
+    for i, step in enumerate(internal_episode.steps):
+        # Convert observation
+        obs = Observation(
+            screenshot_path=step.observation.image_path,
+            a11y_tree=step.observation.accessibility_tree,
+            dom=step.observation.dom_html,
+            window_title=step.observation.window_title,
+            raw=step.observation.meta,
+        )
+
+        # Convert action - note: internal format uses normalized coords in x/y
+        action_type_map = {
+            "click": ActionType.CLICK,
+            "double_click": ActionType.DOUBLE_CLICK,
+            "right_click": ActionType.RIGHT_CLICK,
+            "drag": ActionType.DRAG,
+            "scroll": ActionType.SCROLL,
+            "type": ActionType.TYPE,
+            "key": ActionType.KEY,
+            "wait": ActionType.WAIT,
+            "done": ActionType.DONE,
+            "failed": ActionType.FAIL,
+            "answer": ActionType.DONE,  # Map answer to done
+        }
+        action_type = action_type_map.get(step.action.type, ActionType.CLICK)
+
+        action = Action(
+            type=action_type,
+            # Store normalized coords from internal format
+            normalized_coordinates=(step.action.x, step.action.y)
+            if step.action.x is not None and step.action.y is not None
+            else None,
+            text=step.action.text,
+            key=step.action.key,
+            modifiers=step.action.modifiers,
+            scroll_direction=step.action.scroll_direction,
+            scroll_amount=int(step.action.scroll_amount)
+            if step.action.scroll_amount
+            else None,
+            normalized_end=(step.action.end_x, step.action.end_y)
+            if step.action.end_x is not None and step.action.end_y is not None
+            else None,
+            element=UIElement(
+                element_id=step.action.target_node_id,
+                role=step.action.target_role,
+                name=step.action.target_name,
+            )
+            if step.action.target_node_id
+            else None,
+            raw=step.action.raw,
+        )
+
+        steps.append(
+            Step(
+                step_index=i,
+                observation=obs,
+                action=action,
+                reasoning=step.thought,
+                timestamp=step.t,
+            )
+        )
+
+    return Episode(
+        episode_id=episode_id or internal_episode.id,
+        instruction=internal_episode.goal,
+        steps=steps,
+        success=internal_episode.success,
+        metadata={
+            "workflow_id": internal_episode.workflow_id,
+            "summary": internal_episode.summary,
+        }
+        if internal_episode.workflow_id or internal_episode.summary
+        else None,
+    )
+
+
+def to_internal_episode(episode: Episode) -> dict:
+    """Convert Episode to internal training format (as dict).
+
+    Returns a dict matching openadapt_ml.schemas.sessions.Episode structure.
+    The caller can construct the dataclass from this dict.
+
+    Args:
+        episode: Episode in new format
+
+    Returns:
+        Dict matching internal Episode structure
+    """
+    steps = []
+    for step in episode.steps:
+        # Get normalized coordinates
+        norm_x, norm_y = None, None
+        if step.action.normalized_coordinates:
+            norm_x, norm_y = step.action.normalized_coordinates
+        elif step.action.coordinates:
+            # Can't convert pixel to normalized without screen size
+            # Store in raw for reference
+            pass
+
+        step_dict = {
+            "t": step.timestamp or float(step.step_index),
+            "observation": {
+                "image_path": step.observation.screenshot_path,
+                "accessibility_tree": step.observation.a11y_tree,
+                "dom_html": step.observation.dom,
+                "window_title": step.observation.window_title,
+                "meta": step.observation.raw,
+            },
+            "action": {
+                "type": step.action.type.value,
+                "x": norm_x,
+                "y": norm_y,
+                "text": step.action.text,
+                "key": step.action.key,
+                "modifiers": step.action.modifiers,
+                "scroll_direction": step.action.scroll_direction,
+                "scroll_amount": step.action.scroll_amount,
+                "end_x": step.action.normalized_end[0]
+                if step.action.normalized_end
+                else None,
+                "end_y": step.action.normalized_end[1]
+                if step.action.normalized_end
+                else None,
+                "target_node_id": step.action.element.element_id
+                if step.action.element
+                else None,
+                "target_role": step.action.element.role
+                if step.action.element
+                else None,
+                "target_name": step.action.element.name
+                if step.action.element
+                else None,
+                "raw": step.action.raw,
+            },
+            "thought": step.reasoning,
+        }
+        steps.append(step_dict)
+
+    return {
+        "id": episode.episode_id,
+        "goal": episode.instruction,
+        "steps": steps,
+        "success": episode.success,
+        "workflow_id": episode.metadata.get("workflow_id")
+        if episode.metadata
+        else None,
+        "summary": episode.metadata.get("summary") if episode.metadata else None,
+    }
+
+
+def load_waa_result(result_dir: Union[str, Path]) -> Episode:
+    """Load episode from WAA result directory.
+
+    WAA result directories contain:
+    - result.txt: Final score
+    - trajectory.json or similar: Step-by-step data
+
+    Args:
+        result_dir: Path to WAA result directory
+
+    Returns:
+        Episode instance
+    """
+    result_dir = Path(result_dir)
+
+    # Try to find trajectory file
+    trajectory_files = list(result_dir.glob("*trajectory*.json")) + list(
+        result_dir.glob("*steps*.json")
+    )
+
+    trajectory = []
+    task_info = {}
+
+    if trajectory_files:
+        with open(trajectory_files[0]) as f:
+            data = json.load(f)
+            if isinstance(data, list):
+                trajectory = data
+            elif isinstance(data, dict):
+                trajectory = data.get("steps", data.get("trajectory", []))
+                task_info = {
+                    k: v for k, v in data.items() if k not in ["steps", "trajectory"]
+                }
+
+    # Try to read result
+    result_file = result_dir / "result.txt"
+    if result_file.exists():
+        with open(result_file) as f:
+            result_str = f.read().strip()
+            try:
+                task_info["success"] = float(result_str) > 0
+            except ValueError:
+                pass
+
+    # Try to get task info from parent directory name
+    task_id = result_dir.name
+    if task_id and "task_id" not in task_info:
+        task_info["task_id"] = task_id
+
+    return from_waa_trajectory(trajectory, task_info, episode_id=f"waa_{task_id}")