From 8d074699514737e72865bec3f8054f432d60382b Mon Sep 17 00:00:00 2001 From: Alexis Deprez Date: Mon, 16 Sep 2024 22:50:55 +0200 Subject: [PATCH] feat: update engine actions data model --- lavague-core/lavague/core/action/__init__.py | 10 ++- lavague-core/lavague/core/action/base.py | 24 +++++-- .../lavague/core/action/navigation.py | 56 ++++++++++------- .../lavague/core/action/navigation_control.py | 62 ------------------- lavague-core/lavague/core/trajectory.py | 8 ++- 5 files changed, 65 insertions(+), 95 deletions(-) delete mode 100644 lavague-core/lavague/core/action/navigation_control.py diff --git a/lavague-core/lavague/core/action/__init__.py b/lavague-core/lavague/core/action/__init__.py index a22320e4..c38e5c8f 100644 --- a/lavague-core/lavague/core/action/__init__.py +++ b/lavague-core/lavague/core/action/__init__.py @@ -1,7 +1,11 @@ -from lavague.core.action.base import Action, ActionParser, DEFAULT_PARSER +from lavague.core.action.base import ( + Action, + ActionStatus, + ActionParser, + DEFAULT_PARSER, + UnhandledTypeException, +) from lavague.core.action.navigation import NavigationAction -from lavague.core.action.navigation_control import ControlsAction DEFAULT_PARSER.register("navigation", NavigationAction) -DEFAULT_PARSER.register("controls", ControlsAction) diff --git a/lavague-core/lavague/core/action/base.py b/lavague-core/lavague/core/action/base.py index e78e4535..95b08ecb 100644 --- a/lavague-core/lavague/core/action/base.py +++ b/lavague-core/lavague/core/action/base.py @@ -1,18 +1,23 @@ from PIL import Image import base64 import os -from typing import TypeVar, Generic, Dict, Type, Optional +from typing import Dict, Type, Optional from pydantic import BaseModel, validate_call import time +from enum import Enum -T = TypeVar("T") +class ActionStatus(Enum): + COMPLETED = "completed" + FAILED = "failed" -class Action(BaseModel, Generic[T]): + +class Action(BaseModel): """Action performed by the agent.""" engine: str - args: T + action: str + status: ActionStatus preaction_screenshot: Optional[str] = None postaction_screenshot: Optional[str] = None @@ -65,8 +70,11 @@ def parse(self, action_dict: Dict) -> Action: action_dict.get("postaction_screenshot") ) - target_type = self.engine_action_builders.get(engine, Action) - return target_type.parse(action_dict) + target_type: Type[Action] = self.engine_action_builders.get(engine, Action) + try: + return target_type.parse(action_dict) + except UnhandledTypeException: + return Action.parse(action_dict) def _store_image(self, image: str) -> str: """Store image on disk and return absolute path""" @@ -89,4 +97,8 @@ def _store_image(self, image: str) -> str: return os.path.abspath(image_path) +class UnhandledTypeException(Exception): + pass + + DEFAULT_PARSER = ActionParser() diff --git a/lavague-core/lavague/core/action/navigation.py b/lavague-core/lavague/core/action/navigation.py index 08e87d42..5ee34c03 100644 --- a/lavague-core/lavague/core/action/navigation.py +++ b/lavague-core/lavague/core/action/navigation.py @@ -1,7 +1,6 @@ -from lavague.core.action import Action +from lavague.core.action import Action, UnhandledTypeException from enum import Enum -from typing import Any, ClassVar, Dict, Type, Optional -from pydantic import BaseModel +from typing import ClassVar, Dict, Type, Optional class NavigationActionType(Enum): @@ -12,27 +11,29 @@ class NavigationActionType(Enum): SET_VALUE_AND_ENTER = "setValueAndEnter" DROPDOWN_SELECT = "dropdownSelect" HOVER = "hover" + SCROLL_DOWN = "scroll_down" + SCROLL_UP = "scroll_up" + BACK = "back" -class NavigationActionArgs(BaseModel): - """Arguments for navigation action.""" - - xpath: str - name: NavigationActionType - value: Optional[str] = None - - -class NavigationAction(Action[NavigationActionArgs]): +class NavigationAction(Action): """Navigation action performed by the agent.""" subtypes: ClassVar[Dict[str, Type["NavigationAction"]]] = {} - args: NavigationActionArgs + xpath: str + action: NavigationActionType + value: Optional[str] = None @classmethod - def parse(cls, action_dict: Any) -> "NavigationAction": - subtype = action_dict.get("args", {}).get("name") - target_type = cls.subtypes.get(subtype, NavigationAction) + def parse(cls, action_dict: Dict) -> "NavigationAction": + action_name = action_dict.get("action") + try: + NavigationActionType(action_name) + except ValueError: + raise UnhandledTypeException(f"Unhandled action type: {action_name}") + + target_type = cls.subtypes.get(action_name, NavigationAction) return target_type(**action_dict) @classmethod @@ -45,16 +46,10 @@ def register_navigation(name: str): return lambda cls: NavigationAction.register_subtype(name, cls) -class NavigationWithValueActionArgs(NavigationActionArgs): - """Arguments for navigation action with a value.""" - - value: str - - class NavigationWithValueAction(NavigationAction): """Navigation action performed by the agent with a value.""" - args: NavigationWithValueActionArgs + value: str @register_navigation(NavigationActionType.CLICK.value) @@ -80,3 +75,18 @@ class SetValueAndEnterAction(SetValueAction): @register_navigation(NavigationActionType.DROPDOWN_SELECT.value) class DropdownSelectAction(NavigationWithValueAction): pass + + +@register_navigation(NavigationActionType.SCROLL_DOWN.value) +class ScrollDownAction(NavigationAction): + pass + + +@register_navigation(NavigationActionType.SCROLL_UP.value) +class ScrollUpAction(NavigationAction): + pass + + +@register_navigation(NavigationActionType.BACK.value) +class BackAction(NavigationAction): + pass \ No newline at end of file diff --git a/lavague-core/lavague/core/action/navigation_control.py b/lavague-core/lavague/core/action/navigation_control.py deleted file mode 100644 index 8e1460b1..00000000 --- a/lavague-core/lavague/core/action/navigation_control.py +++ /dev/null @@ -1,62 +0,0 @@ -from lavague.core.action import Action -from enum import Enum -from pydantic import BaseModel -from typing import Any, ClassVar, Dict, Type - - -class ControlsActionType(Enum): - """Types of controls actions.""" - - SCROLL_DOWN = "scroll_down" - SCROLL_UP = "scroll_up" - BACK = "back" - SCAN = "scan" - - -class ControlsActionArgs(BaseModel): - """Arguments for controls action.""" - - name: ControlsActionType - - -class ControlsAction(Action[ControlsActionArgs]): - """Controls action performed by the agent.""" - - subtypes: ClassVar[Dict[str, Type["ControlsActionArgs"]]] = {} - - args: ControlsActionArgs - - @classmethod - def parse(cls, action_dict: Any) -> "ControlsAction": - subtype = action_dict.get("args", {}).get("name") - target_type = cls.subtypes.get(subtype, ControlsAction) - return target_type(**action_dict) - - @classmethod - def register_subtype(cls, subtype: str, action: Type["ControlsAction"]): - cls.subtypes[subtype] = action - return cls - - -def register_control(name: str): - return lambda cls: ControlsAction.register_subtype(name, cls) - - -@register_control(ControlsActionType.SCROLL_UP.value) -class ScrollUpAction(ControlsAction): - pass - - -@register_control(ControlsActionType.SCROLL_DOWN.value) -class ScrollDownAction(ControlsAction): - pass - - -@register_control(ControlsActionType.BACK.value) -class BackAction(ControlsAction): - pass - - -@register_control(ControlsActionType.SCAN.value) -class ScanAction(ControlsAction): - pass diff --git a/lavague-core/lavague/core/trajectory.py b/lavague-core/lavague/core/trajectory.py index e014d0fa..acb50982 100644 --- a/lavague-core/lavague/core/trajectory.py +++ b/lavague-core/lavague/core/trajectory.py @@ -1,6 +1,12 @@ from typing import Optional, Generator, List from pydantic import BaseModel from lavague.core.action.base import Action +from enum import Enum + + +class TrajectoryStatus(Enum): + COMPLETED = "completed" + FAILED = "failed" class Trajectory(BaseModel): @@ -9,7 +15,7 @@ class Trajectory(BaseModel): url: str objective: str actions: List[Action] - success: bool + status: TrajectoryStatus final_html: Optional[str] output: Optional[str]