agent-studio api

Author: mbrunel

lavague-integrations/drivers/lavague-drivers-selenium/lavague/drivers/selenium/base.py

@@ -25,6 +25,8 @@
     NoPageException,
 )
 
+from lavague.sdk.action.navigation import WebNavigationAction, NavigationCommand
+from lavague.sdk.action import ActionStatus
 from selenium.common.exceptions import (
     NoSuchElementException,
     TimeoutException,
@@ -39,7 +41,6 @@
 from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.chrome.service import Service
 
-
 class SeleniumDriver(BaseDriver[SeleniumNode]):
     driver: WebDriver
 

lavague-sdk/lavague/sdk/action/__init__.py

@@ -1,5 +1,7 @@
 from lavague.sdk.action.base import (
     Action,
+    Instruction,
+    EngineType,
     ActionType,
     ActionStatus,
     ActionParser,

lavague-sdk/lavague/sdk/action/base.py

@@ -14,6 +14,13 @@ class ActionType(str, Enum):
     EXTRACTION = "web_extraction"
 
 
+class EngineType(str, Enum):
+    NAVIGATION = "Navigation Engine"
+    EXTRACTION = "Element Extraction Engine"
+    CONTROLS = "Navigation Controls"
+    COMPLETE = "COMPLETE"
+
+
 T = TypeVar("T")
 
 
@@ -55,6 +62,12 @@ def parse(self, action_dict: Dict) -> Action:
             return Action.parse(action_dict)
 
 
+class Instruction(BaseModel):
+    chain_of_toughts: str
+    engine: EngineType
+    engine_instruction: str
+
+
 class UnhandledTypeException(Exception):
     pass
 

lavague-sdk/lavague/sdk/base_driver/base.py

@@ -67,6 +67,7 @@ def execute(self, action: NavigationOutput) -> None:
                     raise NotImplementedError(
                         f"Action {action.navigation_command} not implemented"
                     )
+        self.wait_for_idle()
 
     @abstractmethod
     def destroy(self) -> None:

lavague-sdk/lavague/sdk/client.py

@@ -2,7 +2,7 @@
 from typing import Any, Optional, Tuple
 
 import requests
-from lavague.sdk.action import DEFAULT_PARSER, ActionParser
+from lavague.sdk.action import DEFAULT_PARSER, ActionParser, Instruction, Action
 from lavague.sdk.trajectory import Trajectory
 from lavague.sdk.trajectory.controller import TrajectoryController
 from lavague.sdk.trajectory.model import StepCompletion
@@ -86,7 +86,30 @@ def next_step(self, run_id: str) -> StepCompletion:
             f"/runs/{run_id}/step",
             "POST",
         )
-        return StepCompletion.model_validate_json(content)
+        return StepCompletion.from_data(content)
+
+    def generate_instruction(self, run_id: str) -> Instruction:
+        content = self.request_api(
+            f"/runs/{run_id}/step/instruction",
+            "POST",
+        )
+        return Instruction.model_validate_json(content)
+
+    def generate_action(self, run_id: str, instruction: Instruction) -> StepCompletion:
+        content = self.request_api(
+            f"/runs/{run_id}/step/action",
+            "POST",
+            instruction.model_dump(),
+        )
+        return StepCompletion.from_data(content)
+    
+    def execute_action(self, run_id: str, action: StepCompletion) -> StepCompletion:
+        content = self.request_api(
+            f"/runs/{run_id}/step/execution",
+            "POST",
+            action.model_dump(),
+        )
+        return StepCompletion.from_data(content)
 
     def stop(self, run_id: str) -> None:
         self.request_api(

lavague-sdk/lavague/sdk/trajectory/model.py

@@ -1,9 +1,10 @@
 from enum import Enum
 from typing import Any, Dict, List, Tuple, Optional
 from pydantic import BaseModel, SerializeAsAny
-from lavague.sdk.action import Action
+from lavague.sdk.action import Action, ActionParser
 from lavague.sdk.action.base import DEFAULT_PARSER
 from pydantic import model_validator
+from pydantic_core import from_json
 
 class RunStatus(str, Enum):
     STARTING = "starting"
@@ -59,4 +60,26 @@ def write_to_file(self, file_path: str):
 class StepCompletion(BaseModel):
     run_status: RunStatus
     action: Optional[Action]
-    run_mode: RunMode
+    run_mode: RunMode
+
+    @classmethod
+    def from_data(
+        cls,
+        data: str | bytes | bytearray,
+        parser: ActionParser = DEFAULT_PARSER,
+    ):
+        obj = from_json(data)
+        return cls.from_dict(obj, parser)
+
+    @classmethod
+    def from_dict(
+        cls,
+        data: Dict,
+        parser: ActionParser = DEFAULT_PARSER,
+    ):
+        action = data.get("action")
+        action = parser.parse(action) if action else None
+        return cls.model_validate({
+            **data,
+            "action": action
+        })