Merge pull request #22 from google/ericpts/v2

Implement the v2 actions.

Author: ericpts

agent.py

@@ -39,8 +39,8 @@ def __init__(
         browser_computer: Computer,
         query: str,
         model_name: Literal[
-            "computer-use-exp"
-        ] = "computer-use-exp",
+            "computer-use-exp-6-11"
+        ] = "computer-use-exp-6-11",
     ):
         self._browser_computer = browser_computer
         self._query = query
@@ -50,6 +50,10 @@ def __init__(
             vertexai=os.environ.get("USE_VERTEXAI", "0").lower() in ["true", "1"],
             project=os.environ.get("VERTEXAI_PROJECT"),
             location=os.environ.get("VERTEXAI_LOCATION"),
+            http_options=types.HttpOptions(
+                api_version="v1alpha",
+                base_url="https://generativelanguage.googleapis.com",
+                )
         )
         self._contents: list[Content] = [
             Content(
@@ -94,13 +98,32 @@ def handle_action(self, action: types.FunctionCall) -> EnvState:
         elif action.name == "type_text_at":
             x = self.normalize_x(action.args["x"])
             y = self.normalize_y(action.args["y"])
+            press_enter = action.args.get("press_enter", True)
+            clear_before_typing = action.args.get("clear_before_typing", True)
             return self._browser_computer.type_text_at(
                 x=x,
                 y=y,
                 text=action.args["text"],
+                press_enter=press_enter,
+                clear_before_typing=clear_before_typing,
             )
         elif action.name == "scroll_document":
             return self._browser_computer.scroll_document(action.args["direction"])
+        elif action.name == "scroll_at":
+            x = self.normalize_x(action.args["x"])
+            y = self.normalize_y(action.args["y"])
+            magnitude = action.args.get("magnitude", 200)
+            direction = action.args["direction"]
+
+            if direction in ("up", "down"):
+                magnitude = self.normalize_y(magnitude)
+            elif direction in ("left", "right"):
+                magnitude = self.normalize_x(magnitude)
+            else:
+                raise ValueError("Unknown direction: ", direction)
+            return self._browser_computer.scroll_at(
+                x=x, y=y, direction=direction, magnitude=magnitude
+            )
         elif action.name == "wait_5_seconds":
             return self._browser_computer.wait_5_seconds()
         elif action.name == "go_back":
@@ -115,6 +138,17 @@ def handle_action(self, action: types.FunctionCall) -> EnvState:
             return self._browser_computer.key_combination(
                 action.args["keys"].split("+")
             )
+        elif action.name == "drag_and_drop":
+            x = self.normalize_x(action.args["x"])
+            y = self.normalize_y(action.args["y"])
+            destination_x = self.normalize_x(action.args["destination_x"])
+            destination_y = self.normalize_y(action.args["destination_y"])
+            return self._browser_computer.drag_and_drop(
+                x=x,
+                y=y,
+                destination_x=destination_x,
+                destination_y=destination_y,
+            )
         else:
             raise ValueError(f"Unsupported function: {action}")
 

apiserver/commands.py

@@ -24,23 +24,39 @@ class NavigateArgs(BaseModel):
 
 
 class ClickAtArgs(BaseModel):
-    y: int
     x: int
+    y: int
 
 
 class HoverAtArgs(BaseModel):
-    y: int
     x: int
+    y: int
 
 
 class TypeTextAtArgs(BaseModel):
-    y: int
     x: int
+    y: int
     text: str
+    press_enter: bool
+    clear_before_typing: bool
 
 
 class ScrollDocumentArgs(BaseModel):
-    direction: Literal["up"] | Literal["down"]
+    direction: Literal["up", "down", "left", "right"]
+
+
+class ScrollAtArgs(BaseModel):
+    x: int
+    y: int
+    direction: Literal["up", "down", "left", "right"]
+    magnitude: int
+
+
+class DragAndDropArgs(BaseModel):
+    x: int
+    y: int
+    destination_x: int
+    destination_y: int
 
 
 class KeyCombinationArgs(BaseModel):
@@ -72,6 +88,16 @@ class ScrollDocument(BaseModel):
     args: ScrollDocumentArgs
 
 
+class ScrollAt(BaseModel):
+    name: Literal["scroll_at"]
+    args: ScrollAtArgs
+
+
+class DragAndDrop(BaseModel):
+    name: Literal["drag_and_drop"]
+    args: DragAndDropArgs
+
+
 class Wait5Seconds(BaseModel):
     name: Literal["wait_5_seconds"]
     args: Optional[EmptyJson] = Field(None)
@@ -120,6 +146,8 @@ class Shutdown(BaseModel):
         HoverAt,
         TypeTextAt,
         ScrollDocument,
+        ScrollAt,
+        DragAndDrop,
         GoBack,
         GoForward,
         Search,

apiserver/test_app.py

@@ -114,7 +114,12 @@ def test_get_command_timeout_standard_command():
 
 def test_get_command_timeout_type_text_at_command():
     text = "hello"
-    command = TypeTextAt(name="type_text_at", args=TypeTextAtArgs(y=0, x=0, text=text))
+    command = TypeTextAt(
+        name="type_text_at",
+        args=TypeTextAtArgs(
+            y=0, x=0, text=text, press_enter=False, clear_before_typing=True
+        ),
+    )
     base_timeout = 10
     key_delay = 0.1
     expected_timeout = base_timeout + (len(text) * key_delay)
@@ -124,7 +129,12 @@ def test_get_command_timeout_type_text_at_command():
 
 def test_get_command_timeout_type_text_at_command_zero_delay():
     text = "world"
-    command = TypeTextAt(name="type_text_at", args=TypeTextAtArgs(y=0, x=0, text=text))
+    command = TypeTextAt(
+        name="type_text_at",
+        args=TypeTextAtArgs(
+            y=0, x=0, text=text, press_enter=False, clear_before_typing=True
+        ),
+    )
     base_timeout = 20
     key_delay = 0.0
     expected_timeout = base_timeout

apiserver/test_commands.py

@@ -19,6 +19,8 @@
     HoverAt,
     TypeTextAt,
     ScrollDocument,
+    ScrollAt,
+    DragAndDrop,
     GoBack,
     GoForward,
     Search,
@@ -38,21 +40,27 @@ def test_navigate(self):
         self.assertEqual(command.model_dump(), function_call)
 
     def test_click_at(self):
-        function_call = {"name": "click_at", "args": {"y": 1, "x": 2}}
+        function_call = {"name": "click_at", "args": {"x": 2, "y": 1}}
         command = CommandModel.model_validate(function_call)
         self.assertIsInstance(command.root, ClickAt)
         self.assertEqual(command.model_dump(), function_call)
 
     def test_hover_at(self):
-        function_call = {"name": "hover_at", "args": {"y": 1, "x": 2}}
+        function_call = {"name": "hover_at", "args": {"x": 2, "y": 1}}
         command = CommandModel.model_validate(function_call)
         self.assertIsInstance(command.root, HoverAt)
         self.assertEqual(command.model_dump(), function_call)
 
     def test_type_text_at(self):
         function_call = {
             "name": "type_text_at",
-            "args": {"y": 1, "x": 2, "text": "one"},
+            "args": {
+                "x": 2,
+                "y": 1,
+                "text": "one",
+                "press_enter": True,
+                "clear_before_typing": False,
+            },
         }
         command = CommandModel.model_validate(
             function_call,
@@ -62,7 +70,13 @@ def test_type_text_at(self):
             command.model_dump(),
             {
                 "name": "type_text_at",
-                "args": {"y": 1, "x": 2, "text": "one"},
+                "args": {
+                    "x": 2,
+                    "y": 1,
+                    "text": "one",
+                    "press_enter": True,
+                    "clear_before_typing": False,
+                },
             },
         )
 
@@ -77,6 +91,58 @@ def test_scroll_document(self):
         self.assertIsInstance(command.root, ScrollDocument)
         self.assertEqual(command.model_dump(), function_call)
 
+        function_call = {
+            "name": "scroll_document",
+            "args": {"direction": "down"},
+        }
+        command = CommandModel.model_validate(
+            function_call,
+        )
+        self.assertIsInstance(command.root, ScrollDocument)
+        self.assertEqual(command.model_dump(), function_call)
+
+        function_call = {
+            "name": "scroll_document",
+            "args": {"direction": "left"},
+        }
+        command = CommandModel.model_validate(
+            function_call,
+        )
+        self.assertIsInstance(command.root, ScrollDocument)
+        self.assertEqual(command.model_dump(), function_call)
+
+        function_call = {
+            "name": "scroll_document",
+            "args": {"direction": "right"},
+        }
+        command = CommandModel.model_validate(
+            function_call,
+        )
+        self.assertIsInstance(command.root, ScrollDocument)
+        self.assertEqual(command.model_dump(), function_call)
+
+    def test_scroll_at(self):
+        function_call = {
+            "name": "scroll_at",
+            "args": {"x": 1, "y": 2, "direction": "up", "magnitude": 10},
+        }
+        command = CommandModel.model_validate(
+            function_call,
+        )
+        self.assertIsInstance(command.root, ScrollAt)
+        self.assertEqual(command.model_dump(), function_call)
+
+    def test_drag_and_drop(self):
+        function_call = {
+            "name": "drag_and_drop",
+            "args": {"x": 1, "y": 2, "destination_x": 3, "destination_y": 4},
+        }
+        command = CommandModel.model_validate(
+            function_call,
+        )
+        self.assertIsInstance(command.root, DragAndDrop)
+        self.assertEqual(command.model_dump(), function_call)
+
     def test_go_back(self):
         function_call = {"name": "go_back", "args": {}}
         command = CommandModel.model_validate(function_call)

computers/cloud_run/cloud_run.py

@@ -14,7 +14,7 @@
 import base64
 import termcolor
 import time
-from typing import Any, Optional
+from typing import Any, Optional, Literal
 import requests
 from urllib.parse import urljoin
 from ..computer import (
@@ -110,25 +110,46 @@ def _run_command(
     def open_web_browser(self) -> EnvState:
         return self._run_command("open_web_browser")
 
-    def click_at(self, y, x):
+    def click_at(self, x: int, y: int) -> EnvState:
         return self._run_command("click_at", args={"x": x, "y": y})
 
-    def hover_at(self, y, x):
+    def hover_at(self, x: int, y: int) -> EnvState:
         return self._run_command("hover_at", args={"x": x, "y": y})
 
-    def type_text_at(self, x: int, y: int, text: str) -> EnvState:
+    def type_text_at(
+        self,
+        x: int,
+        y: int,
+        text: str,
+        press_enter: bool,
+        clear_before_typing: bool,
+    ) -> EnvState:
         return self._run_command(
             "type_text_at",
             args={
                 "x": x,
                 "y": y,
                 "text": text,
+                "press_enter": press_enter,
+                "clear_before_typing": clear_before_typing,
             },
         )
 
     def scroll_document(self, direction: str) -> EnvState:
         return self._run_command("scroll_document", args={"direction": direction})
 
+    def scroll_at(
+        self,
+        x: int,
+        y: int,
+        direction: Literal["up", "down", "left", "right"],
+        magnitude: int,
+    ) -> EnvState:
+        return self._run_command(
+            "scroll_at",
+            args={"x": x, "y": y, "direction": direction, "magnitude": magnitude},
+        )
+
     def wait_5_seconds(self) -> EnvState:
         return self._run_command("wait_5_seconds")
 
@@ -147,6 +168,19 @@ def navigate(self, url: str) -> EnvState:
     def key_combination(self, keys: list[str]) -> EnvState:
         return self._run_command("key_combination", args={"keys": "+".join(keys)})
 
+    def drag_and_drop(
+        self, x: int, y: int, destination_x: int, destination_y: int
+    ) -> EnvState:
+        return self._run_command(
+            "drag_and_drop",
+            args={
+                "x": x,
+                "y": y,
+                "destination_x": destination_x,
+                "destination_y": destination_y,
+            },
+        )
+
     def current_state(self) -> EnvState:
         return self._run_command("screenshot")