feat: add description, class and text locators

Author: adi-wan-askui

pyproject.toml

@@ -40,8 +40,9 @@ distribution = true
 
 [tool.pdm.scripts]
 test = "pytest"
-"test:unit" = "pytest tests/unit"
+"test:e2e" = "pytest tests/e2e"
 "test:integration" = "pytest tests/integration"
+"test:unit" = "pytest tests/unit"
 sort = "isort ."
 format = "black ."
 lint = "ruff check ."

src/askui/agent.py

@@ -5,6 +5,7 @@
 from pydantic import Field, validate_call
 
 from askui.container import telemetry
+from askui.models.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/models/__init__.py

@@ -87,8 +87,8 @@ def __init__(self, additional_ai_element_locations: Optional[List[pathlib.Path]]
 
         logger.debug("AI Element locations: %s", self.ai_element_locations)
 
-    def find(self, name: str):
-        ai_elements = []
+    def find(self, name: str) -> list[AiElement]:
+        ai_elements: list[AiElement] = []
 
         for location in self.ai_element_locations:
             path = pathlib.Path(location)
@@ -105,4 +105,4 @@ def find(self, name: str):
                 if ai_element.metadata.name == name:
                     ai_elements.append(ai_element)
 
-        return ai_elements
+        return ai_elements

src/askui/models/askui/ai_element_utils.py

@@ -4,8 +4,9 @@
 import requests
 
 from PIL import Image
-from typing import List, Union
+from typing import Any, List, Union
 from askui.models.askui.ai_element_utils import AiElement, AiElementCollection, AiElementNotFound
+from askui.models.locators import AskUiLocatorSerializer, Locator
 from askui.utils import image_to_base64
 from askui.logger import logger
 
@@ -23,6 +24,7 @@ def __init__(self):
             self.authenticated = False
 
         self.ai_element_collection = AiElementCollection()
+        self._locator_serializer = AskUiLocatorSerializer()
 
 
 
@@ -32,7 +34,7 @@ def _build_askui_token_auth_header(self, bearer_token: str | None = None) -> dic
         token_base64 = base64.b64encode(self.token.encode("utf-8")).decode("utf-8")
         return {"Authorization": f"Basic {token_base64}"}
 
-    def _build_custom_elements(self, ai_elements: List[AiElement] | None):
+    def _build_custom_elements(self, ai_elements: List[AiElement] | None) -> list[dict[str, str]]:
         """
         Converts AiElements to the CustomElementDto format expected by the backend.
         
@@ -43,9 +45,9 @@ def _build_custom_elements(self, ai_elements: List[AiElement] | None):
             dict: Custom elements in the format expected by the backend
         """
         if not ai_elements:
-            return {}
+            return []
 
-        custom_elements = []
+        custom_elements: list[dict[str, str]] = []
         for element in ai_elements:
             custom_element = {
                 "customImage": "," + image_to_base64(element.image),            
@@ -54,24 +56,22 @@ def _build_custom_elements(self, ai_elements: List[AiElement] | None):
             }
             custom_elements.append(custom_element)
 
-        return {
-            "customElements": custom_elements
-        }  
-    def __build_model_composition(self):
-        return {}
+        return custom_elements
 
     def __build_base_url(self, endpoint: str = "inference") -> str:
         return f"{self.inference_endpoint}/api/v3/workspaces/{self.workspace_id}/{endpoint}"
 
-    def predict(self, image: Union[pathlib.Path, Image.Image], locator: str, ai_elements: List[pathlib.Path] = None) -> tuple[int | None, int | None]:
+    def predict(self, image: Union[pathlib.Path, Image.Image], locator: str | Locator, ai_elements: List[AiElement] | None = None) -> tuple[int | None, int | None]:
+        json: dict[str, Any] = {
+            "image": f",{image_to_base64(image)}",
+        }
+        if locator is not None:
+            json["instruction"] = locator if isinstance(locator, str) else locator.serialize(serializer=self._locator_serializer)
+        if ai_elements is not None:
+            json["customElements"] = self._build_custom_elements(ai_elements)
         response = requests.post(
             self.__build_base_url(),
-            json={
-                "image": f",{image_to_base64(image)}",
-                **({"instruction": locator} if locator is not None else {}),
-                **self.__build_model_composition(),
-                **self._build_custom_elements(ai_elements)
-            },
+            json=json,
             headers={"Content-Type": "application/json", **self._build_askui_token_auth_header()},
             timeout=30,
         )
@@ -83,23 +83,23 @@ def predict(self, image: Union[pathlib.Path, Image.Image], locator: str, ai_elem
         actions = [el for el in content["data"]["actions"] if el["inputEvent"] == "MOUSE_MOVE"]
         if len(actions) == 0:
             return None, None
-        position = actions[0]["position"]
 
+        position = actions[0]["position"]
         return int(position["x"]), int(position["y"])
 
-    def locate_pta_prediction(self, image: Union[pathlib.Path, Image.Image], locator: str) -> tuple[int | None, int | None]:
-        askui_locator = f'Click on pta "{locator}"'
-        return self.predict(image, askui_locator)
+    def locate_pta_prediction(self, image: Union[pathlib.Path, Image.Image], locator: str | Locator) -> tuple[int | None, int | None]:
+        _locator = f'Click on pta "{locator}"' if isinstance(locator, str) else locator
+        return self.predict(image, _locator)
 
-    def locate_ocr_prediction(self, image: Union[pathlib.Path, Image.Image], locator: str) -> tuple[int | None, int | None]:
-        askui_locator = f'Click on with text "{locator}"'
-        return self.predict(image, askui_locator)
+    def locate_ocr_prediction(self, image: Union[pathlib.Path, Image.Image], locator: str | Locator) -> tuple[int | None, int | None]:
+        _locator = f'Click on with text "{locator}"' if isinstance(locator, str) else locator
+        return self.predict(image, _locator)
 
     def locate_ai_element_prediction(self, image: Union[pathlib.Path, Image.Image], name: str) -> tuple[int | None, int | None]:
         ai_elements = self.ai_element_collection.find(name)
 
         if len(ai_elements) == 0:
             raise AiElementNotFound(f"Could not locate AI element with name '{name}'")
 
-        askui_instruction = f'Click on custom element with text "{name}"'
-        return self.predict(image, askui_instruction, ai_elements=ai_elements)
+        _locator = f'Click on custom element with text "{name}"'
+        return self.predict(image, _locator, ai_elements=ai_elements)

src/askui/models/askui/api.py

@@ -0,0 +1,115 @@
+from abc import ABC, abstractmethod
+from typing import Literal, TypeVar, Generic
+
+
+SerializedLocator = TypeVar('SerializedLocator')
+
+
+class LocatorSerializer(Generic[SerializedLocator], ABC):
+    @abstractmethod
+    def serialize(self, locator: "Locator") -> SerializedLocator:
+        raise NotImplementedError()
+
+
+class Locator:
+    def serialize(self, serializer: LocatorSerializer[SerializedLocator]) -> SerializedLocator:
+        return serializer.serialize(self)
+
+
+class Description(Locator):
+    def __init__(self, description: str):
+        self.description = description
+
+    def __str__(self):
+        return f'element with description "{self.description}"'
+
+
+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):
+        self.class_name = class_name
+
+    def __str__(self):
+        return f'element with class "{self.class_name}"' if self.class_name else "element that has a class"
+
+
+class Text(Class):
+    def __init__(
+        self,
+        text: str | None = None,
+        match_type: Literal["similar", "exact", "contains", "regex"] = "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):
+        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
+
+
+class AskUiLocatorSerializer(LocatorSerializer[str]):
+    _TEXT_DELIMITER = "<|string|>"
+    
+    def serialize(self, locator: Locator) -> str:
+        prefix = "Click on "
+        if isinstance(locator, Text):
+            return prefix + self._serialize_text(locator)
+        elif isinstance(locator, Class):
+            return prefix + self._serialize_class(locator)
+        elif isinstance(locator, Description):
+            return prefix + self._serialize_description(locator)
+        else:
+            raise ValueError(f"Unsupported locator type: {type(locator)}")
+
+    def _serialize_class(self, class_: Class) -> str:
+        return class_.class_name or "element"
+        
+    def _serialize_description(self, description: Description) -> str:
+        return f'pta {self._TEXT_DELIMITER}{description.description}{self._TEXT_DELIMITER}'
+
+    def _serialize_text(self, text: Text) -> str:
+        match text.match_type:
+            case "similar":
+                return f'with text {self._TEXT_DELIMITER}{text.text}{self._TEXT_DELIMITER} that matches to {text.similarity_threshold} %'
+            case "exact":
+                return f'equals text {self._TEXT_DELIMITER}{text.text}{self._TEXT_DELIMITER}'
+            case "contains":
+                return f'contain text {self._TEXT_DELIMITER}{text.text}{self._TEXT_DELIMITER}'
+            case "regex":
+                return f'match regex pattern {self._TEXT_DELIMITER}{text.text}{self._TEXT_DELIMITER}'
+
+
+class VlmLocatorSerializer(LocatorSerializer[str]):
+    def serialize(self, locator: Locator) -> str:
+        if isinstance(locator, Text):
+            return self._serialize_text(locator)
+        elif isinstance(locator, Class):
+            return self._serialize_class(locator)
+        elif isinstance(locator, Description):
+            return self._serialize_description(locator)
+        else:
+            raise ValueError(f"Unsupported locator type: {type(locator)}")
+
+    def _serialize_class(self, class_: Class) -> str:
+        return class_.class_name or "ui element"
+        
+    def _serialize_description(self, description: Description) -> str:
+        return description.description
+
+    def _serialize_text(self, text: Text) -> str:
+        if text.match_type == "similar":
+            return f'text similar to "{text.text}"'
+
+        return str(text)