refactor(tools/computer): consistent, more readable coordinate scaling

Author: adi-wan-askui

src/askui/locators/serializers.py

@@ -13,9 +13,7 @@
     Prompt,
     Text,
 )
-from .locators import (
-    AiElement as AiElementLocator,
-)
+from .locators import AiElement as AiElementLocator
 from .relatable import (
     BoundingRelation,
     LogicalRelation,

src/askui/models/anthropic/messages_api.py

@@ -45,8 +45,8 @@
 from askui.utils.image_utils import (
     ImageSource,
     image_to_base64,
-    scale_coordinates_back,
-    scale_image_with_padding,
+    scale_coordinates,
+    scale_image_to_fit,
 )
 
 from .utils import extract_click_coordinates
@@ -156,10 +156,9 @@ def _inference(
         system: str,
         model_choice: str,
     ) -> str:
-        scaled_image = scale_image_with_padding(
+        scaled_image = scale_image_to_fit(
             image.root,
-            self._settings.resolution[0],
-            self._settings.resolution[1],
+            self._settings.resolution,
         )
         message = self.create_message(
             messages=[
@@ -222,16 +221,12 @@ def locate(
                 ),
                 model_choice=model_choice,
             )
-            scaled_x, scaled_y = extract_click_coordinates(content)
-            x, y = scale_coordinates_back(
-                scaled_x,
-                scaled_y,
-                image.root.width,
-                image.root.height,
-                screen_width,
-                screen_height,
+            return scale_coordinates(
+                extract_click_coordinates(content),
+                image.root.size,
+                self._settings.resolution,
+                inverse=True,
             )
-            return int(x), int(y)
         except (
             _UnexpectedResponseError,
             ValueError,

src/askui/models/shared/tools.py

@@ -7,6 +7,7 @@
 from pydantic import BaseModel, Field
 from typing_extensions import Self
 
+from askui.logger import logger
 from askui.models.shared.agent_message_param import (
     Base64ImageSourceParam,
     ContentBlockParam,
@@ -155,6 +156,7 @@ def _run_tool(
         except AgentException:
             raise
         except Exception as e:  # noqa: BLE001
+            logger.error(f"Tool {tool_use_block_param.name} failed: {e}", exc_info=True)
             return ToolResultBlockParam(
                 content=f"Tool {tool_use_block_param.name} failed: {e}",
                 is_error=True,

src/askui/tools/android/agent_os_facade.py

@@ -4,7 +4,7 @@
 
 from askui.reporting import Reporter
 from askui.tools.android.agent_os import ANDROID_KEY, AndroidAgentOs, AndroidDisplay
-from askui.utils.image_utils import scale_coordinates_back, scale_image_with_padding
+from askui.utils.image_utils import scale_coordinates, scale_image_to_fit
 
 
 class AndroidAgentOsFacade(AndroidAgentOs):
@@ -32,10 +32,9 @@ def disconnect(self) -> None:
     def screenshot(self) -> Image.Image:
         screenshot = self._agent_os.screenshot()
         self._real_screen_resolution = screenshot.size
-        scaled_image = scale_image_with_padding(
+        scaled_image = scale_image_to_fit(
             screenshot,
-            self._target_resolution[0],
-            self._target_resolution[1],
+            self._target_resolution,
         )
 
         self._reporter.add_message("AndroidAgentOS", "Screenshot taken", screenshot)
@@ -45,15 +44,12 @@ def _scale_coordinates_back(self, x: int, y: int) -> Tuple[int, int]:
         if self._real_screen_resolution is None:
             self._real_screen_resolution = self._agent_os.screenshot().size
 
-        scaled_x, scaled_y = scale_coordinates_back(
-            x,
-            y,
-            self._real_screen_resolution[0],
-            self._real_screen_resolution[1],
-            self._target_resolution[0],
-            self._target_resolution[1],
+        return scale_coordinates(
+            (x, y),
+            self._real_screen_resolution,
+            self._target_resolution,
+            inverse=True,
         )
-        return int(scaled_x), int(scaled_y)
 
     def tap(self, x: int, y: int) -> None:
         scaled_x, scaled_y = self._scale_coordinates_back(x, y)

src/askui/tools/askui/askui_controller.py

@@ -29,11 +29,11 @@
 from askui.tools.askui.askui_ui_controller_grpc.generated import (
     Controller_V1_pb2_grpc as controller_v1,
 )
-from askui.tools.askui.askui_ui_controller_grpc.generated.AgentOS_Send_Request_2501 import (
-    RenderObjectStyle,  # noqa: E501
+from askui.tools.askui.askui_ui_controller_grpc.generated.AgentOS_Send_Request_2501 import (  # noqa: E501
+    RenderObjectStyle,
 )
-from askui.tools.askui.askui_ui_controller_grpc.generated.AgentOS_Send_Response_2501 import (
-    AskuiAgentosSendResponseSchema,  # noqa: E501
+from askui.tools.askui.askui_ui_controller_grpc.generated.AgentOS_Send_Response_2501 import (  # noqa: E501
+    AskuiAgentosSendResponseSchema,
 )
 from askui.tools.askui.command_helpers import (
     create_clear_render_objects_command,

src/askui/tools/computer.py

@@ -13,11 +13,7 @@
 from typing_extensions import Self, override
 
 from askui.tools.agent_os import AgentOs, Coordinate, ModifierKey, PcKey
-from askui.utils.image_utils import (
-    scale_coordinates_back,
-    scale_coordinates_with_padding,
-    scale_image_with_padding,
-)
+from askui.utils.image_utils import scale_coordinates, scale_image_to_fit
 
 from ..models.shared.tools import InputSchema, Tool
 
@@ -192,20 +188,55 @@ class BetaToolComputerUseParamBase(TypedDict):
     display_height_px: int
 
 
+@dataclass
+class Resolution:
+    width: int
+    height: int
+
+
+# https://github.com/anthropics/anthropic-quickstarts/blob/main/computer-use-demo/README.md
+RESOLUTIONS_RECOMMENDED_BY_ANTHROPIC: dict[str, Resolution] = {
+    "XGA": Resolution(width=1024, height=768),  # 4:3
+    "WXGA": Resolution(width=1280, height=800),  # 16:10
+}
+
+
+def _get_closest_recommended_resolution(resolution: Resolution) -> Resolution:
+    return min(
+        RESOLUTIONS_RECOMMENDED_BY_ANTHROPIC.values(),
+        key=lambda r: abs(r.width - resolution.width)
+        + abs(r.height - resolution.height),
+    )
+
+
 class ComputerToolBase(Tool, ABC):
     def __init__(
         self,
         agent_os: AgentOs,
         input_schema: InputSchema,
+        resolution: Resolution | None = None,
     ) -> None:
         super().__init__(
             name="computer",
             description="A tool for interacting with the computer",
             input_schema=input_schema,
         )
         self._agent_os = agent_os
-        self._width = 1280
-        self._height = 800
+        real_resolution = self._get_real_screen_resolution()
+        self._resolution = resolution or _get_closest_recommended_resolution(
+            Resolution(
+                width=real_resolution[0],
+                height=real_resolution[1],
+            )
+        )
+
+    @property
+    def _width(self) -> int:
+        return self._resolution.width
+
+    @property
+    def _height(self) -> int:
+        return self._resolution.height
 
     @property
     def params_base(
@@ -228,7 +259,7 @@ def __call__(  # noqa: C901
     ) -> Image.Image | None | str:
         match action:
             case "cursor_position":
-                return self._get_mouse_position_scaled()
+                return self._retrieve_cursor_position()
             case "double_click":
                 return self._agent_os.click("left", 2)
             case "key":
@@ -284,17 +315,12 @@ def _scale_coordinates_back(
         self,
         coordinate: tuple[Annotated[int, Field(ge=0)], Annotated[int, Field(ge=0)]],
     ) -> tuple[int, int]:
-        real_screen_width, real_screen_height = self._get_real_screen_resolution()
-        x, y = scale_coordinates_back(
-            coordinate[0],
-            coordinate[1],
-            real_screen_width,
-            real_screen_height,
-            self._width,
-            self._height,
+        return scale_coordinates(
+            coordinate,
+            self._get_real_screen_resolution(),
+            (self._width, self._height),
+            inverse=True,
         )
-        x, y = int(x), int(y)
-        return x, y
 
     @validate_call
     def _mouse_move(
@@ -320,18 +346,15 @@ def _screenshot(self) -> Image.Image:
         Take a screenshot of the current screen, scale it and return it
         """
         screenshot = self._agent_os.screenshot()
-        return scale_image_with_padding(screenshot, self._width, self._height)
+        return scale_image_to_fit(screenshot, (self._width, self._height))
 
-    def _get_mouse_position_scaled(self) -> str:
+    def _retrieve_cursor_position(self) -> str:
         mouse_position: Coordinate = self._agent_os.get_mouse_position()
         real_screen_width, real_screen_height = self._get_real_screen_resolution()
-        x, y = scale_coordinates_with_padding(
-            mouse_position.x,
-            mouse_position.y,
-            real_screen_width,
-            real_screen_height,
-            self._width,
-            self._height,
+        x, y = scale_coordinates(
+            (mouse_position.x, mouse_position.y),
+            (real_screen_width, real_screen_height),
+            (self._width, self._height),
         )
 
         return f"X={x},Y={y}"