Merge pull request #118 from askui/feat-add-locate-all

Feat: Add locate all function

Author: mlikasam-askui

README.md

@@ -367,12 +367,12 @@ class MyGetAndLocateModel(GetModel, LocateModel):
         locator: str | Locator,
         image: ImageSource,
         model_choice: ModelComposition | str,
-    ) -> Point:
+    ) -> PointList:
         # Implement custom locate logic, e.g.:
         # - Use a different object detection model
         # - Implement custom element finding
         # - Call external vision services
-        return (100, 100)  # Example coordinates
+        return [(100, 100)]  # Example coordinates
 
 
 # Create model registry

src/askui/__init__.py

@@ -25,6 +25,7 @@
     OnMessageCb,
     OnMessageCbParam,
     Point,
+    PointList,
     TextBlockParam,
     TextCitationParam,
     ToolResultBlockParam,
@@ -82,6 +83,7 @@
     "OnMessageCbParam",
     "PcKey",
     "Point",
+    "PointList",
     "ResponseSchema",
     "ResponseSchemaBase",
     "Retry",

src/askui/agent.py

@@ -182,7 +182,7 @@ def _click(
     def _mouse_move(
         self, locator: str | Locator, model: ModelComposition | str | None = None
     ) -> None:
-        point = self._locate(locator=locator, model=model)
+        point = self._locate(locator=locator, model=model)[0]
         self.tools.os.mouse_move(point[0], point[1])
 
     @telemetry.record_call(exclude={"locator"})

src/askui/agent_base.py

@@ -29,6 +29,7 @@
     ModelName,
     ModelRegistry,
     Point,
+    PointList,
     TotalModelChoice,
 )
 from .models.types.response_schemas import ResponseSchema
@@ -352,13 +353,14 @@ class LinkedListNode(ResponseSchemaBase):
         self._reporter.add_message("Agent", message_content)
         return response
 
+    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
     def _locate(
         self,
         locator: str | Locator,
         screenshot: Optional[Img] = None,
         model: ModelComposition | str | None = None,
-    ) -> Point:
-        def locate_with_screenshot() -> Point:
+    ) -> PointList:
+        def locate_with_screenshot() -> PointList:
             _screenshot = load_image_source(
                 self._agent_os.screenshot() if screenshot is None else screenshot
             )
@@ -368,10 +370,10 @@ def locate_with_screenshot() -> Point:
                 model_choice=model or self._model_choice["locate"],
             )
 
-        point = self._retry.attempt(locate_with_screenshot)
-        self._reporter.add_message("ModelRouter", f"locate: ({point[0]}, {point[1]})")
-        logger.debug("ModelRouter locate: (%d, %d)", point[0], point[1])
-        return point
+        points = self._retry.attempt(locate_with_screenshot)
+        self._reporter.add_message("ModelRouter", f"locate {len(points)} elements")
+        logger.debug("ModelRouter locate: %d elements", len(points))
+        return points
 
     @telemetry.record_call(exclude={"locator", "screenshot"})
     @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
@@ -382,7 +384,7 @@ def locate(
         model: ModelComposition | str | None = None,
     ) -> Point:
         """
-        Locates the UI element identified by the provided locator.
+        Locates the first matching UI element identified by the provided locator.
 
         Args:
             locator (str | Locator): The identifier or description of the element to
@@ -405,8 +407,53 @@ def locate(
                 print(f"Element found at coordinates: {point}")
             ```
         """
-        self._reporter.add_message("User", f"locate {locator}")
-        logger.debug("VisionAgent received instruction to locate %s", locator)
+        self._reporter.add_message("User", f"locate first matching element {locator}")
+        logger.debug(
+            "VisionAgent received instruction to locate first matching element %s",
+            locator,
+        )
+        return self._locate(locator, screenshot, model)[0]
+
+    @telemetry.record_call(exclude={"locator", "screenshot"})
+    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
+    def locate_all(
+        self,
+        locator: str | Locator,
+        screenshot: Optional[Img] = None,
+        model: ModelComposition | str | None = None,
+    ) -> PointList:
+        """
+        Locates all matching UI elements identified by the provided locator.
+
+        Note: Some LocateModels can only locate a single element. In this case, the
+        returned list will have a length of 1.
+
+        Args:
+            locator (str | Locator): The identifier or description of the element to
+                locate.
+            screenshot (Img | None, optional): The screenshot to use for locating the
+                element. Can be a path to an image file, a PIL Image object or a data
+                URL. If `None`, takes a screenshot of the currently selected display.
+            model (ModelComposition | str | None, optional): The composition or name
+                of the model(s) to be used for locating the element using the `locator`.
+
+        Returns:
+            PointList: The coordinates of the elements as a list of tuples (x, y).
+
+        Example:
+            ```python
+            from askui import VisionAgent
+
+            with VisionAgent() as agent:
+                points = agent.locate_all("Submit button")
+                print(f"Found {len(points)} elements at coordinates: {points}")
+            ```
+        """
+        self._reporter.add_message("User", f"locate all matching UI elements {locator}")
+        logger.debug(
+            "VisionAgent received instruction to locate all matching UI elements %s",
+            locator,
+        )
         return self._locate(locator, screenshot, model)
 
     @telemetry.record_call()

src/askui/android_agent.py

@@ -198,7 +198,7 @@ def tap(
             msg += f" on {target}"
             self._reporter.add_message("User", msg)
             logger.debug("VisionAgent received instruction to click on %s", target)
-            point = self._locate(locator=target, model=model)
+            point = self._locate(locator=target, model=model)[0]
             self.os.tap(point[0], point[1])
 
     @telemetry.record_call(exclude={"text"})

src/askui/models/__init__.py

@@ -9,6 +9,7 @@
     ModelName,
     ModelRegistry,
     Point,
+    PointList,
 )
 from .openrouter.model import OpenRouterModel
 from .openrouter.settings import ChatCompletionsCreateSettings, OpenRouterSettings
@@ -53,6 +54,7 @@
     "OpenRouterModel",
     "OpenRouterSettings",
     "Point",
+    "PointList",
     "TextBlockParam",
     "TextCitationParam",
     "ToolResultBlockParam",

src/askui/models/anthropic/messages_api.py

@@ -27,7 +27,7 @@
     LocateModel,
     ModelComposition,
     ModelName,
-    Point,
+    PointList,
 )
 from askui.models.shared.agent_message_param import (
     Base64ImageSourceParam,
@@ -198,7 +198,7 @@ def locate(
         locator: str | Locator,
         image: ImageSource,
         model_choice: ModelComposition | str,
-    ) -> Point:
+    ) -> PointList:
         if not isinstance(model_choice, str):
             error_msg = "Model composition is not supported for Claude"
             raise NotImplementedError(error_msg)
@@ -219,12 +219,14 @@ def locate(
                 ),
                 model_choice=model_choice,
             )
-            return scale_coordinates(
-                extract_click_coordinates(content),
-                image.root.size,
-                self._settings.resolution,
-                inverse=True,
-            )
+            return [
+                scale_coordinates(
+                    extract_click_coordinates(content),
+                    image.root.size,
+                    self._settings.resolution,
+                    inverse=True,
+                )
+            ]
         except (
             _UnexpectedResponseError,
             ValueError,

src/askui/models/askui/inference_api.py

@@ -20,7 +20,7 @@
 from askui.locators.serializers import AskUiLocatorSerializer, AskUiSerializedLocator
 from askui.logger import logger
 from askui.models.exceptions import ElementNotFoundError
-from askui.models.models import GetModel, LocateModel, ModelComposition, Point
+from askui.models.models import GetModel, LocateModel, ModelComposition, PointList
 from askui.models.shared.agent_message_param import MessageParam
 from askui.models.shared.messages_api import MessagesApi
 from askui.models.shared.settings import MessageSettings
@@ -162,7 +162,7 @@ def locate(
         locator: str | Locator,
         image: ImageSource,
         model_choice: ModelComposition | str,
-    ) -> Point:
+    ) -> PointList:
         serialized_locator = (
             self._locator_serializer.serialize(locator=locator)
             if isinstance(locator, Locator)
@@ -171,7 +171,7 @@ def locate(
         logger.debug(f"serialized_locator:\n{json_lib.dumps(serialized_locator)}")
         json: dict[str, Any] = {
             "image": image.to_data_url(),
-            "instruction": f"Click on {serialized_locator['instruction']}",
+            "instruction": f"get element {serialized_locator['instruction']}",
         }
         if "customElements" in serialized_locator:
             json["customElements"] = serialized_locator["customElements"]
@@ -182,17 +182,20 @@ def locate(
             )
         response = self._post(path="/inference", json=json)
         content = response.json()
-        assert content["type"] == "COMMANDS", (
+        assert content["type"] == "DETECTED_ELEMENTS", (
             f"Received unknown content type {content['type']}"
         )
-        actions = [
-            el for el in content["data"]["actions"] if el["inputEvent"] == "MOUSE_MOVE"
-        ]
-        if len(actions) == 0:
+        detected_elements = content["data"]["detected_elements"]
+        if len(detected_elements) == 0:
             raise ElementNotFoundError(locator, serialized_locator)
 
-        position = actions[0]["position"]
-        return int(position["x"]), int(position["y"])
+        return [
+            (
+                int((element["bndbox"]["xmax"] + element["bndbox"]["xmin"]) / 2),
+                int((element["bndbox"]["ymax"] + element["bndbox"]["ymin"]) / 2),
+            )
+            for element in detected_elements
+        ]
 
     @override
     def get(

src/askui/models/askui/model_router.py

@@ -8,7 +8,7 @@
     ElementNotFoundError,
     ModelNotFoundError,
 )
-from askui.models.models import LocateModel, ModelComposition, ModelName, Point
+from askui.models.models import LocateModel, ModelComposition, ModelName, PointList
 from askui.utils.image_utils import ImageSource
 
 
@@ -18,7 +18,7 @@ def __init__(self, inference_api: AskUiInferenceApi):
 
     def _locate_with_askui_ocr(
         self, screenshot: ImageSource, locator: str | Text
-    ) -> Point:
+    ) -> PointList:
         locator = Text(locator) if isinstance(locator, str) else locator
         return self._inference_api.locate(
             locator, screenshot, model_choice=ModelName.ASKUI__OCR
@@ -30,7 +30,7 @@ def locate(
         locator: str | Locator,
         image: ImageSource,
         model_choice: ModelComposition | str,
-    ) -> Point:
+    ) -> PointList:
         if (
             isinstance(model_choice, ModelComposition)
             or model_choice == ModelName.ASKUI

src/askui/models/huggingface/spaces_api.py

@@ -10,7 +10,7 @@
 from askui.exceptions import AutomationError
 from askui.locators.locators import Locator
 from askui.locators.serializers import VlmLocatorSerializer
-from askui.models.models import LocateModel, ModelComposition, ModelName, Point
+from askui.models.models import LocateModel, ModelComposition, ModelName, PointList
 from askui.utils.image_utils import ImageSource
 
 
@@ -65,7 +65,7 @@ def locate(
         locator: str | Locator,
         image: ImageSource,
         model_choice: ModelComposition | str,
-    ) -> Point:
+    ) -> PointList:
         """Predict element location using Hugging Face Spaces."""
         if not isinstance(model_choice, str):
             error_msg = "Model composition is not supported for Hugging Face Spaces"
@@ -76,9 +76,9 @@ def locate(
                 if isinstance(locator, Locator)
                 else locator
             )
-            return self._spaces[model_choice](
-                image.root, serialized_locator, model_choice
-            )
+            return [
+                self._spaces[model_choice](image.root, serialized_locator, model_choice)
+            ]
         except (ValueError, json.JSONDecodeError, httpx.HTTPError) as e:
             error_msg = f"Hugging Face Spaces Exception: {e}"
             raise AutomationError(error_msg) from e