Build the first version of the computer-use agent

automation/actions.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Annotated, Literal, Union
+
+from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
+
+
+class ActionType(str, Enum):
+    """Actions the model is allowed to request."""
+
+    CLICK = "click"
+    DOUBLE_CLICK = "double_click"
+    RIGHT_CLICK = "right_click"
+    MOVE = "move"
+    TYPE = "type"
+    PRESS = "press"
+    HOTKEY = "hotkey"
+    SCROLL = "scroll"
+    WAIT = "wait"
+    FINISH = "finish"
+    FAIL = "fail"
+
+
+class BaseAction(BaseModel):
+    """Shared validation for every model action."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    action: ActionType
+    reason: str = Field(
+        min_length=1,
+        max_length=300,
+        description="Short explanation of why this action is needed.",
+    )
+
+
+class CoordinateAction(BaseAction):
+    """Base class for actions that target a screen coordinate."""
+
+    x: int = Field(ge=0)
+    y: int = Field(ge=0)
+
+
+class ClickAction(CoordinateAction):
+    action: Literal[ActionType.CLICK] = ActionType.CLICK
+
+
+class DoubleClickAction(CoordinateAction):
+    action: Literal[ActionType.DOUBLE_CLICK] = ActionType.DOUBLE_CLICK
+
+
+class RightClickAction(CoordinateAction):
+    action: Literal[ActionType.RIGHT_CLICK] = ActionType.RIGHT_CLICK
+
+
+class MoveAction(CoordinateAction):
+    action: Literal[ActionType.MOVE] = ActionType.MOVE
+    duration: float = Field(
+        default=0.2,
+        ge=0,
+        le=2,
+        description="Seconds used to move the cursor.",
+    )
+
+
+class TypeAction(BaseAction):
+    action: Literal[ActionType.TYPE] = ActionType.TYPE
+    text: str = Field(
+        min_length=1,
+        max_length=5000,
+        description="Text to type into the focused application.",
+    )
+    interval: float = Field(
+        default=0.01,
+        ge=0,
+        le=0.25,
+        description="Delay between keystrokes.",
+    )
+
+
+class PressAction(BaseAction):
+    action: Literal[ActionType.PRESS] = ActionType.PRESS
+    key: str = Field(
+        min_length=1,
+        max_length=30,
+        description="One keyboard key, such as enter, tab, or esc.",
+    )
+    presses: int = Field(default=1, ge=1, le=20)
+    interval: float = Field(default=0.05, ge=0, le=1)
+
+
+class HotkeyAction(BaseAction):
+    action: Literal[ActionType.HOTKEY] = ActionType.HOTKEY
+    keys: list[str] = Field(
+        min_length=2,
+        max_length=5,
+        description="Keys pressed together, such as ['ctrl', 'l'].",
+    )
+
+
+class ScrollAction(BaseAction):
+    action: Literal[ActionType.SCROLL] = ActionType.SCROLL
+    amount: int = Field(
+        ge=-20,
+        le=20,
+        description="Positive scrolls up and negative scrolls down.",
+    )
+    x: int | None = Field(
+        default=None,
+        ge=0,
+        description="Optional horizontal position before scrolling.",
+    )
+    y: int | None = Field(
+        default=None,
+        ge=0,
+        description="Optional vertical position before scrolling.",
+    )
+
+
+class WaitAction(BaseAction):
+    action: Literal[ActionType.WAIT] = ActionType.WAIT
+    seconds: float = Field(
+        ge=0.1,
+        le=10,
+        description="How long to wait for the interface to update.",
+    )
+
+
+class FinishAction(BaseAction):
+    action: Literal[ActionType.FINISH] = ActionType.FINISH
+    summary: str = Field(
+        min_length=1,
+        max_length=500,
+        description="What was completed.",
+    )
+
+
+class FailAction(BaseAction):
+    action: Literal[ActionType.FAIL] = ActionType.FAIL
+    error: str = Field(
+        min_length=1,
+        max_length=500,
+        description="Why the task cannot continue.",
+    )
+
+
+# The action field tells Pydantic which schema to use.
+ComputerAction = Annotated[
+    Union[
+        ClickAction,
+        DoubleClickAction,
+        RightClickAction,
+        MoveAction,
+        TypeAction,
+        PressAction,
+        HotkeyAction,
+        ScrollAction,
+        WaitAction,
+        FinishAction,
+        FailAction,
+    ],
+    Field(discriminator="action"),
+]
+
+ACTION_ADAPTER = TypeAdapter(ComputerAction)
+
+
+def parse_action(data: str | bytes | dict) -> ComputerAction:
+    """Validate a model response and return a typed action."""
+
+    if isinstance(data, dict):
+        return ACTION_ADAPTER.validate_python(data)
+
+    return ACTION_ADAPTER.validate_json(data)
+
+
+def action_json_schema() -> dict:
+    """Return the schema sent to the local model."""
+
+    return ACTION_ADAPTER.json_schema()