diff --git a/automation/agent.py b/automation/agent.py new file mode 100644 index 0000000..509c347 --- /dev/null +++ b/automation/agent.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() \ No newline at end of file diff --git a/automation/main.py b/automation/main.py index f5dbc8d..fac48da 100644 --- a/automation/main.py +++ b/automation/main.py @@ -13,7 +13,7 @@ def main(): computer = Computer() model = Model() - goal = input("What would you like Friday to do?\n> ") + goal = input("Open word and write me a paragraph on How AI is good.\n> ") while True: # Take the latest screenshot