Add locators

pyproject.toml

@@ -39,9 +39,10 @@ path = "src/askui/__init__.py"
 distribution = true
 
 [tool.pdm.scripts]
-test = "pytest"
-"test:unit" = "pytest tests/unit"
-"test:integration" = "pytest tests/integration"
+test = "pytest -n auto"
+"test:e2e" = "pytest -n auto tests/e2e"
+"test:integration" = "pytest -n auto tests/integration"
+"test:unit" = "pytest -n auto tests/unit"
 sort = "isort ."
 format = "black ."
 lint = "ruff check ."
@@ -56,6 +57,7 @@ test = [
     "black>=25.1.0",
     "ruff>=0.9.5",
     "pytest-mock>=3.14.0",
+    "pytest-xdist>=3.6.1",
 ]
 chat = [
     "streamlit>=1.42.0",

src/askui/agent.py

@@ -5,6 +5,7 @@
 from pydantic import Field, validate_call
 
 from askui.container import telemetry
+from askui.locators import Locator
 
 from .tools.askui.askui_controller import (
     AskUiControllerClient,
@@ -15,7 +16,7 @@
 from .models.anthropic.claude import ClaudeHandler
 from .logger import logger, configure_logging
 from .tools.toolbox import AgentToolbox
-from .models.router import ModelRouter
+from .models.router import ModelRouter, Point
 from .reporting.report import SimpleReportGenerator
 import time
 from dotenv import load_dotenv
@@ -59,13 +60,13 @@ def _check_askui_controller_enabled(self) -> None:
                 "AskUI Controller is not initialized. Please, set `enable_askui_controller` to `True` when initializing the `VisionAgent`."
             )
 
-    @telemetry.record_call(exclude={"instruction"})
-    def click(self, instruction: Optional[str] = None, button: Literal['left', 'middle', 'right'] = 'left', repeat: int = 1, model_name: Optional[str] = None) -> None:
+    @telemetry.record_call(exclude={"locator"})
+    def click(self, locator: Optional[str | Locator] = None, button: Literal['left', 'middle', 'right'] = 'left', repeat: int = 1, model_name: Optional[str] = None) -> None:
         """
-        Simulates a mouse click on the user interface element identified by the provided instruction.
+        Simulates a mouse click on the user interface element identified by the provided locator.
 
         Parameters:
-            instruction (str | None): The identifier or description of the element to click.
+            locator (str | Locator | None): The identifier or description of the element to click.
             button ('left' | 'middle' | 'right'): Specifies which mouse button to click. Defaults to 'left'.
             repeat (int): The number of times to click. Must be greater than 0. Defaults to 1.
             model_name (str | None): The model name to be used for element detection. Optional.
@@ -92,29 +93,34 @@ def click(self, instruction: Optional[str] = None, button: Literal['left', 'midd
                 msg = f'{button} ' + msg 
             if repeat > 1:
                 msg += f' {repeat}x times'
-            if instruction is not None:
-                msg += f' on "{instruction}"'
+            if locator is not None:
+                msg += f' on "{locator}"'
             self.report.add_message("User", msg)
-        if instruction is not None:
-            logger.debug("VisionAgent received instruction to click '%s'", instruction)
-            self.__mouse_move(instruction, model_name)
+        if locator is not None:
+            logger.debug("VisionAgent received instruction to click '%s'", locator)
+            self._mouse_move(locator, model_name)
         self.client.click(button, repeat) # type: ignore
-
-    def __mouse_move(self, instruction: str, model_name: Optional[str] = None) -> None:
-        self._check_askui_controller_enabled()
-        screenshot = self.client.screenshot() # type: ignore
-        x, y = self.model_router.locate(screenshot, instruction, model_name)
+
+    def locate(self, locator: str | Locator, screenshot: Optional[Image.Image] = None, model_name: Optional[str] = None) -> Point:
+        if screenshot is None:
+            self._check_askui_controller_enabled()
+            screenshot = self.client.screenshot() # type: ignore
+        point = self.model_router.locate(screenshot, locator, model_name)
         if self.report is not None:
-            self.report.add_message("ModelRouter", f"locate: ({x}, {y})")
-        self.client.mouse(x, y) # type: ignore
+            self.report.add_message("ModelRouter", f"locate: ({point[0]}, {point[1]})")
+        return point
+
+    def _mouse_move(self, locator: str | Locator, model_name: Optional[str] = None) -> None:
+        point = self.locate(locator=locator, model_name=model_name)
+        self.client.mouse(point[0], point[1]) # type: ignore
 
-    @telemetry.record_call(exclude={"instruction"})
-    def mouse_move(self, instruction: str, model_name: Optional[str] = None) -> None:
+    @telemetry.record_call(exclude={"locator"})
+    def mouse_move(self, locator: str | Locator, model_name: Optional[str] = None) -> None:
         """
-        Moves the mouse cursor to the UI element identified by the provided instruction.
+        Moves the mouse cursor to the UI element identified by the provided locator.
 
         Parameters:
-            instruction (str): The identifier or description of the element to move to.
+            locator (str | Locator): The identifier or description of the element to move to.
             model_name (str | None): The model name to be used for element detection. Optional.
 
         Example:
@@ -126,9 +132,9 @@ def mouse_move(self, instruction: str, model_name: Optional[str] = None) -> None
         ```
         """
         if self.report is not None:
-            self.report.add_message("User", f'mouse_move: "{instruction}"')
-        logger.debug("VisionAgent received instruction to mouse_move '%s'", instruction)
-        self.__mouse_move(instruction, model_name)
+            self.report.add_message("User", f'mouse_move: "{locator}"')
+        logger.debug("VisionAgent received instruction to mouse_move to '%s'", locator)
+        self._mouse_move(locator, model_name)
 
     @telemetry.record_call()
     def mouse_scroll(self, x: int, y: int) -> None:

src/askui/chat/__main__.py

@@ -211,7 +211,7 @@ def rerun():
                                 image=screenshot_with_crosshair,
                             )
                             agent.mouse_move(
-                                instruction=element_description.replace('"', ""),
+                                locator=element_description.replace('"', ""),
                                 model_name="anthropic-claude-3-5-sonnet-20241022",
                             )
                         else:

src/askui/locators/__init__.py

@@ -0,0 +1,13 @@
+from .relatable import ReferencePoint
+from .locators import Class, Description, Locator, Text, TextMatchType
+from . import serializers
+
+__all__ = [
+    "Class",
+    "Description",
+    "Locator",
+    "ReferencePoint",
+    "Text",
+    "TextMatchType",
+    "serializers",
+]

src/askui/locators/locators.py

@@ -0,0 +1,77 @@
+from abc import ABC, abstractmethod
+from typing import Generic, Literal, TypeVar
+
+from askui.locators.relatable import Relatable
+
+
+SerializedLocator = TypeVar("SerializedLocator")
+
+
+class LocatorSerializer(Generic[SerializedLocator], ABC):
+    @abstractmethod
+    def serialize(self, locator: "Locator") -> SerializedLocator:
+        raise NotImplementedError()
+
+
+class Locator(Relatable, ABC):
+    def serialize(
+        self, serializer: LocatorSerializer[SerializedLocator]
+    ) -> SerializedLocator:
+        return serializer.serialize(self)
+
+
+class Description(Locator):
+    def __init__(self, description: str):
+        super().__init__()
+        self.description = description
+
+    def __str__(self):
+        result = f'element with description "{self.description}"'
+        return result + super()._relations_str()
+
+
+class Class(Locator):
+    # None is used to indicate that it is an element with a class but not a specific class
+    def __init__(self, class_name: Literal["text", "textfield"] | None = None):
+        super().__init__()
+        self.class_name = class_name
+
+    def __str__(self):
+        result = (
+            f'element with class "{self.class_name}"'
+            if self.class_name
+            else "element that has a class"
+        )
+        return result + super()._relations_str()
+
+
+TextMatchType = Literal["similar", "exact", "contains", "regex"]
+
+
+class Text(Class):
+    def __init__(
+        self,
+        text: str | None = None,
+        match_type: TextMatchType = "similar",
+        similarity_threshold: int = 70,
+    ):
+        super().__init__(class_name="text")
+        self.text = text
+        self.match_type = match_type
+        self.similarity_threshold = similarity_threshold
+
+    def __str__(self):
+        if self.text is None:
+            result = "text"
+        else:
+            result = "text "
+            match self.match_type:
+                case "similar":
+                    result += f'similar to "{self.text}" (similarity >= {self.similarity_threshold}%)'
+                case "exact":
+                    result += f'"{self.text}"'
+                case "contains":
+                    result += f'containing text "{self.text}"'
+                case "regex":
+                    result += f'matching regex "{self.text}"'
+        return result + super()._relations_str()