Merge branch 'main' into feat-add-locate-all

Author: adi-wan-askui

README.md

@@ -350,7 +350,7 @@ class MyGetAndLocateModel(GetModel, LocateModel):
     def get(
         self,
         query: str,
-        image: ImageSource,
+        source: Source,
         response_schema: Type[ResponseSchema] | None,
         model_choice: str,
     ) -> ResponseSchema | str:
@@ -640,9 +640,9 @@ else:
     agent.click("Login")
 ```
 
-#### Using custom images
+#### Using custom images and PDFs
 
-Instead of taking a screenshot, you can analyze specific images:
+Instead of taking a screenshot, you can analyze specific images or PDFs:
 
 ```python
 from PIL import Image
@@ -651,10 +651,13 @@ from askui import VisionAgent
 # From PIL Image
 with VisionAgent() as agent:
   image = Image.open("screenshot.png")
-  result = agent.get("What's in this image?", image)
+  result = agent.get("What's in this image?", source=image)
 
   # From file path
-  result = agent.get("What's in this image?", "screenshot.png")
+  result = agent.get("What's in this image?", source="screenshot.png")
+
+  # From PDF
+  result = agent.get("What is this PDF about?", source="document.pdf")
 ```
 
 #### Using response schemas
@@ -696,7 +699,7 @@ with VisionAgent() as agent:
     response = agent.get(
         "What is the current url shown in the url bar?",
         response_schema=UrlResponse,
-        image="screenshot.png",
+        source="screenshot.png",
     )
 
     # Dump whole model
@@ -712,7 +715,7 @@ with VisionAgent() as agent:
     is_login_page = agent.get(
         "Is this a login page?",
         response_schema=bool,
-        image=Image.open("screenshot.png"),
+        source=Image.open("screenshot.png"),
     )
     print(is_login_page)
 
@@ -751,6 +754,7 @@ with VisionAgent() as agent:
 **⚠️ Limitations:**
 - The support for response schemas varies among models. Currently, the `askui` model provides best support for response schemas
   as we try different models under the hood with your schema to see which one works best.
+- PDF processing is only supported for Gemini models hosted on AskUI and for PDFs up to 20MB.
 
 ## What is AskUI Vision Agent?
 

pdm.lock

@@ -25,6 +25,7 @@ dependencies = [
     "jsonref>=1.1.0",
     "protobuf>=6.31.1",
     "google-genai>=1.20.0",
+    "filetype>=1.2.0",
 ]
 requires-python = ">=3.10"
 readme = "README.md"

pyproject.toml

@@ -1,6 +1,6 @@
 """AskUI Vision Agent"""
 
-__version__ = "0.11.0"
+__version__ = "0.12.0"
 
 from .agent import VisionAgent
 from .locators import Locator

src/askui/__init__.py

@@ -48,15 +48,15 @@
 
 _ANTHROPIC__CLAUDE__3_5__SONNET__20241022__ACT_SETTINGS = ActSettings(
     messages=MessageSettings(
-        model=ModelName.ANTHROPIC__CLAUDE__3_5__SONNET__20241022.value,
+        model=ModelName.ANTHROPIC__CLAUDE__3_5__SONNET__20241022,
         system=_SYSTEM_PROMPT,
         betas=[COMPUTER_USE_20241022_BETA_FLAG],
     ),
 )
 
 _CLAUDE__SONNET__4__20250514__ACT_SETTINGS = ActSettings(
     messages=MessageSettings(
-        model=ModelName.CLAUDE__SONNET__4__20250514.value,
+        model=ModelName.CLAUDE__SONNET__4__20250514,
         system=_SYSTEM_PROMPT,
         betas=[COMPUTER_USE_20250124_BETA_FLAG],
         thinking={"type": "enabled", "budget_tokens": 2048},

src/askui/agent.py

@@ -1,6 +1,7 @@
 import time
 import types
 from abc import ABC
+from pathlib import Path
 from typing import Annotated, Optional, Type, overload
 
 from dotenv import load_dotenv
@@ -16,6 +17,8 @@
 from askui.tools.agent_os import AgentOs
 from askui.tools.android.agent_os import AndroidAgentOs
 from askui.utils.image_utils import ImageSource, Img
+from askui.utils.pdf_utils import Pdf
+from askui.utils.source_utils import load_image_source, load_source
 
 from .logger import configure_logging, logger
 from .models import ModelComposition
@@ -190,46 +193,53 @@ def get(
         query: Annotated[str, Field(min_length=1)],
         response_schema: None = None,
         model: str | None = None,
-        image: Optional[Img] = None,
+        source: Optional[Img | Pdf] = None,
     ) -> str: ...
     @overload
     def get(
         self,
         query: Annotated[str, Field(min_length=1)],
         response_schema: Type[ResponseSchema],
         model: str | None = None,
-        image: Optional[Img] = None,
+        source: Optional[Img | Pdf] = None,
     ) -> ResponseSchema: ...
 
-    @telemetry.record_call(exclude={"query", "image", "response_schema"})
+    @telemetry.record_call(exclude={"query", "source", "response_schema"})
     @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
     def get(
         self,
         query: Annotated[str, Field(min_length=1)],
         response_schema: Type[ResponseSchema] | None = None,
         model: str | None = None,
-        image: Optional[Img] = None,
+        source: Optional[Img | Pdf] = None,
     ) -> ResponseSchema | str:
         """
-        Retrieves information from an image (defaults to a screenshot of the current
-        screen) based on the provided `query`.
+        Retrieves information from an image or PDF based on the provided `query`.
+
+        If no `source` is provided, a screenshot of the current screen is taken.
 
         Args:
             query (str): The query describing what information to retrieve.
-            image (Img | None, optional): The image to extract information from.
-                Defaults to a screenshot of the current screen. Can be a path to
-                an image file, a PIL Image object or a data URL.
+            source (Img | Pdf | None, optional): The source to extract information from.
+                Can be a path to a PDF file, a path to an image file, a PIL Image
+                object or a data URL. Defaults to a screenshot of the current screen.
             response_schema (Type[ResponseSchema] | None, optional): A Pydantic model
                 class that defines the response schema. If not provided, returns a
                 string.
             model (str | None, optional): The composition or name of the model(s) to
                 be used for retrieving information from the screen or image using the
                 `query`. Note: `response_schema` is not supported by all models.
+                PDF processing is only supported for Gemini models hosted on AskUI.
 
         Returns:
             ResponseSchema | str: The extracted information, `str` if no
                 `response_schema` is provided.
 
+        Raises:
+            NotImplementedError: If PDF processing is not supported for the selected
+                model.
+            ValueError: If the `source` is not a valid PDF or image.
+
         Example:
             ```python
             from askui import ResponseSchemaBase, VisionAgent
@@ -254,7 +264,7 @@ class LinkedListNode(ResponseSchemaBase):
                 response = agent.get(
                     "What is the current url shown in the url bar?",
                     response_schema=UrlResponse,
-                    image="screenshot.png",
+                    source="screenshot.png",
                 )
                 # Dump whole model
                 print(response.model_dump_json(indent=2))
@@ -269,7 +279,7 @@ class LinkedListNode(ResponseSchemaBase):
                 is_login_page = agent.get(
                     "Is this a login page?",
                     response_schema=bool,
-                    image=Image.open("screenshot.png"),
+                    source=Image.open("screenshot.png"),
                 )
                 print(is_login_page)
 
@@ -303,13 +313,34 @@ class LinkedListNode(ResponseSchemaBase):
                 while current:
                     print(current.value)
                     current = current.next
+
+                # Get text from PDF
+                text = agent.get(
+                    "Extract all text from the PDF",
+                    source="document.pdf",
+                )
+                print(text)
             ```
         """
         logger.debug("VisionAgent received instruction to get '%s'", query)
-        _image = ImageSource(self._agent_os.screenshot() if image is None else image)
-        self._reporter.add_message("User", f'get: "{query}"', image=_image.root)
+        _source = (
+            ImageSource(self._agent_os.screenshot())
+            if source is None
+            else load_source(source)
+        )
+
+        # Prepare message content with file path if available
+        user_message_content = f'get: "{query}"' + (
+            f" from '{source}'" if isinstance(source, (str, Path)) else ""
+        )
+
+        self._reporter.add_message(
+            "User",
+            user_message_content,
+            image=_source.root if isinstance(_source, ImageSource) else None,
+        )
         response = self._model_router.get(
-            image=_image,
+            source=_source,
             query=query,
             response_schema=response_schema,
             model_choice=model or self._model_choice["get"],
@@ -329,8 +360,8 @@ def _locate(
         screenshot: Optional[Img] = None,
         model: ModelComposition | str | None = None,
     ) -> PointList:
-        def locate_with_screenshot() -> PointList:
-            _screenshot = ImageSource(
+        def locate_with_screenshot() -> Point:
+            _screenshot = load_image_source(
                 self._agent_os.screenshot() if screenshot is None else screenshot
             )
             return self._model_router.locate(

src/askui/agent_base.py

@@ -103,15 +103,15 @@
 
 _ANTHROPIC__CLAUDE__3_5__SONNET__20241022__ACT_SETTINGS = ActSettings(
     messages=MessageSettings(
-        model=ModelName.ANTHROPIC__CLAUDE__3_5__SONNET__20241022.value,
+        model=ModelName.ANTHROPIC__CLAUDE__3_5__SONNET__20241022,
         system=_SYSTEM_PROMPT,
         betas=[],
     ),
 )
 
 _CLAUDE__SONNET__4__20250514__ACT_SETTINGS = ActSettings(
     messages=MessageSettings(
-        model=ModelName.CLAUDE__SONNET__4__20250514.value,
+        model=ModelName.CLAUDE__SONNET__4__20250514,
         system=_SYSTEM_PROMPT,
         thinking={"type": "enabled", "budget_tokens": 2048},
         betas=[],

src/askui/android_agent.py

@@ -64,7 +64,7 @@ def list_(self, thread_id: ThreadId, query: ListQuery) -> list[Message]:
             raise FileNotFoundError(error_msg)
 
         messages: list[Message] = []
-        with thread_file.open("r") as f:
+        with thread_file.open("r", encoding="utf-8") as f:
             for line in f:
                 msg = Message.model_validate_json(line)
                 messages.append(msg)

src/askui/chat/api/messages/service.py

@@ -7,7 +7,7 @@
 from pydantic import ConfigDict, Field, validate_call
 
 from askui.locators.relatable import Relatable
-from askui.utils.image_utils import ImageSource
+from askui.utils.source_utils import load_image_source
 
 TextMatchType = Literal["similar", "exact", "contains", "regex"]
 """The type of match to use.
@@ -303,7 +303,7 @@ def __init__(
             image_compare_format=image_compare_format,
             name=_generate_name() if name is None else name,
         )
-        self._image = ImageSource(image)
+        self._image = load_image_source(image)
 
 
 class AiElement(ImageBase):

src/askui/locators/locators.py

@@ -48,6 +48,8 @@
     scale_coordinates,
     scale_image_to_fit,
 )
+from askui.utils.pdf_utils import PdfSource
+from askui.utils.source_utils import Source
 
 from .utils import extract_click_coordinates
 
@@ -63,10 +65,6 @@ def build_system_prompt_locate(screen_width: str, screen_height: str) -> str:
                 "claude-3-5-sonnet-20241022"
             ),
             ModelName.CLAUDE__SONNET__4__20250514: "claude-sonnet-4-20250514",
-            ModelName.ANTHROPIC__CLAUDE__3_5__SONNET__20241022.value: (
-                "claude-3-5-sonnet-20241022"
-            ),
-            ModelName.CLAUDE__SONNET__4__20250514.value: "claude-sonnet-4-20250514",
         }
     )
 )
@@ -240,16 +238,19 @@ def locate(
     def get(
         self,
         query: str,
-        image: ImageSource,
+        source: Source,
         response_schema: Type[ResponseSchema] | None,
         model_choice: str,
     ) -> ResponseSchema | str:
+        if isinstance(source, PdfSource):
+            err_msg = f"PDF processing is not supported for the model {model_choice}"
+            raise NotImplementedError(err_msg)
         try:
             if response_schema is not None:
                 error_msg = "Response schema is not yet supported for Anthropic"
                 raise NotImplementedError(error_msg)
             return self._inference(
-                image=image,
+                image=source,
                 prompt=query,
                 system=SYSTEM_PROMPT_GET,
                 model_choice=model_choice,