feat: update engine actions data model

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)

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()

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

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]