ui-locator-cli 0.1.0__py3-none-any.whl

ui_locator/__init__.py

@@ -0,0 +1,54 @@
+"""ui-locator — precise click coordinates for UI elements in screenshots.
+
+Given a screenshot and a natural-language query, return absolute-pixel
+coordinates of the target element, ready for a computer-use agent.
+
+>>> from ui_locator import Locator
+>>> loc = Locator()
+>>> res = loc.locate("screen.png", "where is the 'Submit' button?")
+>>> res.click_point
+Coordinates(x=520, y=430)
+
+Default backend on Apple Silicon is UI-TARS-1.5-7B via MLX; any
+OpenAI-compatible endpoint (mlx-vlm server, vLLM, HuggingFace Inference
+Endpoint, Volcengine/Doubao Ark for UI-TARS-2) works via the ``openai`` provider.
+"""
+
+from __future__ import annotations
+
+from .api import Locator, locate
+from .config import ProviderName, Settings
+from .errors import (
+    CoordinateParseError,
+    ElementNotFoundError,
+    ProviderError,
+    ProviderNotAvailableError,
+    UILocatorError,
+)
+from .image import load_image
+from .types import BoundingBox, Coordinates, Element, LocateResult, ProviderKind
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "__version__",
+    # API
+    "Locator",
+    "locate",
+    "load_image",
+    # Types
+    "BoundingBox",
+    "Coordinates",
+    "Element",
+    "LocateResult",
+    "ProviderKind",
+    # Config
+    "Settings",
+    "ProviderName",
+    # Errors
+    "UILocatorError",
+    "ProviderError",
+    "ProviderNotAvailableError",
+    "ElementNotFoundError",
+    "CoordinateParseError",
+]

ui_locator/__main__.py

@@ -0,0 +1,6 @@
+"""Allow ``python -m ui_locator``."""
+
+from .cli import entrypoint
+
+if __name__ == "__main__":  # pragma: no cover
+    entrypoint()

ui_locator/api.py

@@ -0,0 +1,329 @@
+"""High-level API: the :class:`Locator` and the convenience :func:`locate`.
+
+Example
+-------
+>>> from ui_locator import Locator
+>>> loc = Locator()                       # auto-selects MLX on Apple Silicon
+>>> res = loc.locate("screen.png", "dove si trova il bottone 'Invia'?")
+>>> res.click_point
+Coordinates(x=742, y=511)
+"""
+
+from __future__ import annotations
+
+import os
+import platform
+
+from PIL import Image
+
+from .config import OpenAIPromptStyle, ProviderName, Settings, get_settings
+from .errors import (
+    ElementNotFoundError,
+    ProviderError,
+    ProviderNotAvailableError,
+    UILocatorError,
+)
+from .image import ImageLike, load_image, peek_size
+from .mapstore import CROP_MAX_DIFF, ElementMap, MapEntry, crop_diff, make_crop
+from .providers.base import GroundingProvider
+from .types import Element, LocateResult, MapMeta
+
+
+class Locator:
+    """Locate UI elements in screenshots via natural-language queries.
+
+    Parameters
+    ----------
+    settings:
+        Override the auto-loaded :class:`Settings`. Useful to force a provider
+        or model programmatically.
+    provider:
+        Force a specific provider instance (mostly for testing).
+    """
+
+    def __init__(
+        self,
+        settings: Settings | None = None,
+        *,
+        provider: GroundingProvider | None = None,
+    ) -> None:
+        self._settings = settings or get_settings()
+        self._provider: GroundingProvider | None = provider
+        self._map: ElementMap | None = None
+
+    @property
+    def settings(self) -> Settings:
+        return self._settings
+
+    @property
+    def element_map(self) -> ElementMap:
+        """The on-disk element map (lazily opened)."""
+        if self._map is None:
+            self._map = ElementMap(self._settings.map_dir)
+        return self._map
+
+    @property
+    def provider(self) -> GroundingProvider:
+        if self._provider is None:
+            self._provider = self._select_provider()
+        return self._provider
+
+    def _select_provider(self) -> GroundingProvider:
+        """Pick a provider according to settings + environment."""
+        choice = self._settings.provider
+
+        if choice == ProviderName.MLX:
+            return self._make_mlx()
+        if choice == ProviderName.QWEN:
+            # 'qwen' means Qwen grounding wherever it is available: the cloud
+            # endpoint when one is configured, the local checkpoint otherwise.
+            # Either way the Qwen prompt style is the only one whose <box>
+            # output the parser understands.
+            if self._cloud_configured():
+                return self._make_openai(prompt=OpenAIPromptStyle.QWEN)
+            return self._make_qwen()
+        if choice == ProviderName.QWEN_LOCAL:
+            return self._make_qwen()
+        if choice == ProviderName.OPENAI:
+            return self._make_openai()
+        if choice == ProviderName.ANTHROPIC:
+            return self._make_anthropic()
+
+        # AUTO: prefer MLX on Apple Silicon, fall back to OpenAI if configured.
+        if platform.system() == "Darwin" and platform.machine() == "arm64":
+            try:
+                return self._make_mlx()
+            except ProviderNotAvailableError:
+                pass
+        # OpenAI-compatible cloud (if configured).
+        try:
+            return self._make_openai()
+        except ProviderNotAvailableError:
+            pass
+
+        raise UILocatorError(
+            "No grounding provider available. Either:\n"
+            "  - install the MLX extra on Apple Silicon: "
+            "uv pip install -e '.[mlx,uitars]', or\n"
+            "  - configure a cloud endpoint with "
+            "UI_LOCATOR_OPENAI_BASE_URL / UI_LOCATOR_OPENAI_API_KEY / "
+            "UI_LOCATOR_OPENAI_MODEL (uv pip install -e '.[cloud]')."
+        )
+
+    def _make_mlx(self) -> GroundingProvider:
+        from .providers.mlx_uitars import MlxUITarsProvider
+
+        return MlxUITarsProvider(self._settings)
+
+    def _make_qwen(self) -> GroundingProvider:
+        from .providers.mlx_qwen import MlxQwenProvider
+
+        return MlxQwenProvider(self._settings)
+
+    def _cloud_configured(self) -> bool:
+        return bool(self._settings.openai_base_url or os.environ.get("OPENAI_BASE_URL"))
+
+    def _make_openai(self, *, prompt: OpenAIPromptStyle | None = None) -> GroundingProvider:
+        from .providers.openai_compat import OpenAICompatProvider
+
+        settings = self._settings
+        if prompt is not None and settings.openai_prompt != prompt:
+            settings = settings.model_copy(update={"openai_prompt": prompt})
+        return OpenAICompatProvider(settings)
+
+    def _make_anthropic(self) -> GroundingProvider:
+        from .providers.cloud_cua import AnthropicProvider
+
+        return AnthropicProvider(self._settings)
+
+    def _engine_id(self, provider_name: str) -> str:
+        """A short provenance string for the backend that found an element."""
+        s = self._settings
+        model = {
+            "mlx-uitars": s.model,
+            "mlx-qwen": s.qwen_model,
+            "openai": s.openai_model or "",
+            "anthropic": s.anthropic_model,
+        }.get(provider_name, "")
+        return f"{provider_name}:{model}" if model else provider_name
+
+    def locate(
+        self,
+        image: ImageLike,
+        query: str | None = None,
+        *,
+        element_id: str | None = None,
+        session: str | None = None,
+        use_map: bool | None = None,
+        verify: bool | None = None,
+    ) -> LocateResult:
+        """Locate the UI element described by ``query`` in ``image``.
+
+        Args:
+            image: A path, ``PIL.Image``, or bytes.
+            query: Natural-language description / intent, e.g.
+                ``"where is the 'Invia' button?"`` or ``"chiudi il terminale"``.
+                Optional when ``element_id`` already has a map entry.
+            element_id: Dotted, caller-declared element id (e.g.
+                ``"sap.screen.submit"``). Asserts that this element sits at a
+                stable position on that screen: a stored hit is returned
+                without decoding the screenshot or loading a model, and a
+                model-discovered position is written back for next time.
+            session: Tag recorded on learned entries, so a whole run can be
+                dropped later with ``ElementMap.forget(session=...)``.
+            use_map: Set ``False`` to skip the lookup and re-ask the model;
+                the fresh result still replaces the stored entry, which is the
+                recovery path after a stored point turned out to be wrong.
+            verify: On a map hit, compare a small crop around the stored point
+                against the screenshot; on mismatch fall back to the model.
+
+        Returns:
+            A :class:`LocateResult` whose ``best`` element has absolute-pixel
+            ``click_point`` (and ``bbox`` when available). ``result.map``
+            carries the provenance when ``element_id`` was given.
+
+        Raises:
+            ElementNotFoundError: if the model produced no usable coordinates,
+                or if no query was given and the map has no entry.
+            ProviderError: if the backend call or parsing failed.
+            ProviderNotAvailableError: if dependencies are missing.
+        """
+        settings = self._settings
+        session = session if session is not None else settings.session
+        verify = settings.verify if verify is None else verify
+        read_map = bool(element_id) and (settings.use_map if use_map is None else use_map)
+
+        if read_map and element_id:
+            hit = self._map_lookup(element_id, image, verify=verify)
+            if hit is not None:
+                return hit.model_copy(update={"query": query or hit.query})
+
+        if query is None or not query.strip():
+            if element_id:
+                raise ElementNotFoundError(
+                    f"No map entry for id {element_id!r} at this screenshot size. "
+                    "Pass a query so the model can locate it."
+                )
+            raise ValueError("query must be a non-empty string.")
+
+        img = load_image(image)
+        provider = self.provider
+
+        candidates = provider.ground(img, query)
+        if not candidates:
+            raise ElementNotFoundError(
+                f"No element found for query {query!r} via provider '{provider.name}'."
+            )
+
+        candidates.sort(key=lambda e: e.confidence, reverse=True)
+        best = candidates[0]
+
+        meta: MapMeta | None = None
+        if element_id:
+            meta = self._map_store(element_id, img, best, query, session)
+
+        return LocateResult(
+            query=query,
+            image_size=img.size,
+            best=best,
+            candidates=candidates,
+            map=meta,
+        )
+
+    def _map_lookup(
+        self, element_id: str, image: ImageLike, *, verify: bool
+    ) -> LocateResult | None:
+        """Return a stored result for ``element_id``, or ``None`` to ask the model."""
+        size = peek_size(image)
+        if size is None:
+            return None
+        entry = self.element_map.lookup(element_id, size)
+        if entry is None:
+            return None
+
+        verified: bool | None = None
+        diff: float | None = None
+        if verify:
+            if entry.crop:
+                diff = crop_diff(load_image(image), entry.click_point, entry.crop)
+            verified = diff is not None and diff <= CROP_MAX_DIFF
+            if not verified:
+                return None
+
+        self.element_map.record_hit(element_id, entry)
+        return LocateResult(
+            query=entry.query,
+            image_size=size,
+            best=entry.to_element(),
+            candidates=[],
+            map=MapMeta(
+                element_id=element_id,
+                hit=True,
+                source=entry.source,
+                hits=entry.hits + 1,
+                verified=verified,
+                diff=diff,
+            ),
+        )
+
+    def _map_store(
+        self,
+        element_id: str,
+        img: Image.Image,
+        best: Element,
+        query: str,
+        session: str | None,
+    ) -> MapMeta:
+        """Write a model-discovered element back to the map."""
+        entry = MapEntry(
+            click_point=(best.click_point.x, best.click_point.y),
+            bbox=(
+                (best.bbox.x1, best.bbox.y1, best.bbox.x2, best.bbox.y2) if best.bbox else None
+            ),
+            size=img.size,
+            engine=self._engine_id(self.provider.name),
+            source="model",
+            query=query,
+            label=best.label,
+            confidence=best.confidence,
+            session=session,
+            crop=make_crop(img, (best.click_point.x, best.click_point.y)),
+        )
+        saved = self.element_map.save(element_id, entry)
+        return MapMeta(
+            element_id=element_id,
+            hit=False,
+            source="model",
+            stored=saved is entry,
+        )
+
+
+def locate(
+    image: ImageLike,
+    query: str | None = None,
+    *,
+    element_id: str | None = None,
+    **kwargs: object,
+) -> LocateResult:
+    """One-shot convenience wrapper around :meth:`Locator.locate`.
+
+    Extra keyword arguments are forwarded to :class:`Settings` (e.g.
+    ``model=..., think=True, session="123"``).
+    """
+    settings = Settings(**kwargs)  # type: ignore[arg-type]
+    return Locator(settings=settings).locate(image, query, element_id=element_id)
+
+
+__all__ = [
+    "Locator",
+    "locate",
+    "Element",
+    "ElementMap",
+    "MapEntry",
+    "Image",
+    "LocateResult",
+    "ProviderError",
+    "ProviderNotAvailableError",
+    "ElementNotFoundError",
+    "UILocatorError",
+]