From 766dce6e7c662dbebe661785be81885d16723e92 Mon Sep 17 00:00:00 2001
From: Tianqi Xu <tianqi.xu@kaust.edu.sa>
Date: Sat, 7 Sep 2024 23:51:23 +0300
Subject: [PATCH] Add MacOS support

---
 crab-benchmark-v0/main.py  | 23 +++++++++++++++++++++++
 crab/environments/macos.py | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)
 create mode 100644 crab/environments/macos.py

diff --git a/crab-benchmark-v0/main.py b/crab-benchmark-v0/main.py
index adefbc8..5a622cb 100644
--- a/crab-benchmark-v0/main.py
+++ b/crab-benchmark-v0/main.py
@@ -20,6 +20,7 @@
     BenchmarkConfig,
     Experiment,
     MessageType,
+    Task,
     TaskGenerator,
     create_benchmark,
 )
@@ -32,6 +33,8 @@
 )
 from crab.core.agent_policy import AgentPolicy
 from crab.core.benchmark import Benchmark
+from crab.core.decorators import evaluator
+from crab.environments.macos import mac_env
 
 from .android_env import ANDROID_ENV
 from .dataset.android_subtasks import android_subtasks
@@ -78,6 +81,11 @@ def get_prompt(self):
         return result_prompt
 
 
+@evaluator(env_name="macos")
+def empty_evaluator() -> bool:
+    return False
+
+
 def get_benchmark(env: str, ubuntu_url: str):
     ubuntu_env = UBUNTU_ENV.model_copy()
     ubuntu_env.remote_url = ubuntu_url
@@ -87,6 +95,9 @@ def get_benchmark(env: str, ubuntu_url: str):
     android_tool = {
         "screenshot": groundingdino_easyocr(font_size=40) >> get_elements_prompt
     }
+    mac_tool = {
+        "screenshot": groundingdino_easyocr(font_size=24) >> get_elements_prompt
+    }
 
     if env == "ubuntu":
         prompting_tools = {"ubuntu": ubuntu_tool}
@@ -121,6 +132,18 @@ def get_benchmark(env: str, ubuntu_url: str):
             root_action_space=[complete],
             multienv=True,
         )
+    elif env == "mac":
+        task = Task(description="Open firefox in both macos and android.", id="0",evaluator=empty_evaluator)
+        prompting_tools = {"macos": mac_tool, "android": android_tool}
+        mac_env.remote_url = "http://10.85.170.240:8000"
+        benchmark_config = BenchmarkConfig(
+            name="mac_benchmark",
+            tasks=[task],
+            environments=[mac_env, ANDROID_ENV],
+            prompting_tools=prompting_tools,
+            root_action_space=[complete],
+            multienv=True,
+        )
     else:
         raise ValueError("Env not support")
 
diff --git a/crab/environments/macos.py b/crab/environments/macos.py
new file mode 100644
index 0000000..09734c2
--- /dev/null
+++ b/crab/environments/macos.py
@@ -0,0 +1,32 @@
+from crab import action
+from crab.actions.crab_actions import complete, get_element_position
+from crab.actions.desktop_actions import (
+    click_position,
+    key_press,
+    press_hotkey,
+    right_click,
+    screenshot,
+    write_text,
+)
+from crab.core import EnvironmentConfig
+
+
+@action(local=True)
+def click(element: int, env) -> None:
+    """
+    Click an UI element shown on the desktop screen. A simple use case can be
+    click(5), which clicks the UI element labeled with the number 5.
+
+    Args:
+        element: A numeric tag assigned to an UI element shown on the screenshot.
+    """
+    x, y = get_element_position(element, env)
+    env._action_endpoint(click_position, {"x": round(x / 2), "y": round(y / 2)})
+
+
+mac_env = EnvironmentConfig(
+    name="macos",
+    action_space=[click, key_press, write_text, press_hotkey, right_click, complete],
+    observation_space=[screenshot],
+    description="A Macbook laptop environment with a single display.",
+)