feat(tools): add LoadImageTool and improve file tool flexibility

Author: mlikasam-askui

src/askui/models/shared/tools.py

@@ -163,7 +163,9 @@ def _create_tool_result_block_param_for_playwright_error(
 
 
 class Tool(BaseModel, ABC):
-    model_config = ConfigDict(populate_by_name=True)
+    model_config = ConfigDict(
+        validate_by_alias=True,
+    )
 
     base_name: str = Field(alias="name", description="Name of the tool")
     description: str = Field(description="Description of what the tool does")

src/askui/tools/askui/askui_controller.py

@@ -225,7 +225,7 @@ def _run_recorder_action(
     ) -> controller_v1_pbs.Response_RunRecordedAction:
         time.sleep(self._pre_action_wait)
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
         response: controller_v1_pbs.Response_RunRecordedAction = (
             self._stub.RunRecordedAction(
@@ -241,7 +241,7 @@ def _run_recorder_action(
         num_retries = 0
         for _ in range(self._max_retries):
             assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-                "Stub is not initialized. Call Connect first."
+                "Stub is not initialized. Call `connect()` first."
             )
             poll_response: controller_v1_pbs.Response_Poll = self._stub.Poll(
                 controller_v1_pbs.Request_Poll(
@@ -313,7 +313,7 @@ def __exit__(
 
     def _start_session(self) -> None:
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
         response = self._stub.StartSession(
             controller_v1_pbs.Request_StartSession(
@@ -324,23 +324,23 @@ def _start_session(self) -> None:
 
     def _stop_session(self) -> None:
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
         self._stub.EndSession(
             controller_v1_pbs.Request_EndSession(sessionInfo=self._session_info)
         )
 
     def _start_execution(self) -> None:
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
         self._stub.StartExecution(
             controller_v1_pbs.Request_StartExecution(sessionInfo=self._session_info)
         )
 
     def _stop_execution(self) -> None:
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
         self._stub.StopExecution(
             controller_v1_pbs.Request_StopExecution(sessionInfo=self._session_info)
@@ -361,7 +361,7 @@ def screenshot(self, report: bool = True) -> Image.Image:
 
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
         screenResponse = self._stub.CaptureScreen(
             controller_v1_pbs.Request_CaptureScreen(
@@ -652,7 +652,7 @@ def set_display(self, display: int = 1) -> None:
                 Defaults to `1`.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
         self._stub.SetActiveDisplay(
             controller_v1_pbs.Request_SetActiveDisplay(displayID=display)
@@ -715,7 +715,7 @@ def list_displays(
             DisplaysListResponse
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
 
         self._reporter.add_message("AgentOS", "list_displays()")
@@ -752,7 +752,7 @@ def get_process_list(
                 - processes: List of ProcessInfo objects
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
 
         self._reporter.add_message("AgentOS", f"get_process_list({get_extended_info})")
@@ -781,7 +781,7 @@ def get_window_list(
                 - windows: List of WindowInfo objects with ID and name
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
 
         self._reporter.add_message("AgentOS", f"get_window_list({process_id})")
@@ -809,7 +809,7 @@ def get_automation_target_list(
                 - targets: List of AutomationTarget objects
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
 
         self._reporter.add_message("AgentOS", "get_automation_target_list()")
@@ -832,7 +832,7 @@ def set_mouse_delay(self, delay_ms: int) -> None:
             delay_ms (int): The delay in milliseconds to set for mouse actions.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
 
         self._reporter.add_message("AgentOS", f"set_mouse_delay({delay_ms})")
@@ -852,7 +852,7 @@ def set_keyboard_delay(self, delay_ms: int) -> None:
             delay_ms (int): The delay in milliseconds to set for keyboard actions.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
 
         self._reporter.add_message("AgentOS", f"set_keyboard_delay({delay_ms})")
@@ -881,7 +881,7 @@ def set_active_window(self, process_id: int, window_id: int) -> int:
             If display length is not increased after adding the window.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
 
         self._reporter.add_message(
@@ -914,7 +914,7 @@ def set_active_automation_target(self, target_id: int) -> None:
             target_id (int): The ID of the automation target to set as active.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
 
         self._reporter.add_message(
@@ -945,7 +945,7 @@ def schedule_batched_action(
                 the scheduled action ID.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
 
         self._reporter.add_message(
@@ -971,7 +971,7 @@ def start_batch_run(self) -> None:
         Start executing batched actions.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
 
         self._reporter.add_message("AgentOS", "start_batch_run()")
@@ -986,7 +986,7 @@ def stop_batch_run(self) -> None:
         Stop executing batched actions.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
 
         self._reporter.add_message("AgentOS", "stop_batch_run()")
@@ -1005,7 +1005,7 @@ def get_action_count(self) -> controller_v1_pbs.Response_GetActionCount:
                 containing the action count.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
 
         self._reporter.add_message("AgentOS", "get_action_count()")
@@ -1031,7 +1031,7 @@ def get_action(self, action_index: int) -> controller_v1_pbs.Response_GetAction:
                 - actionParameters: The action parameters
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
 
         self._reporter.add_message("AgentOS", f"get_action({action_index})")
@@ -1053,7 +1053,7 @@ def remove_action(self, action_id: int) -> None:
             action_id (int): The ID of the action to remove.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
 
         self._reporter.add_message("AgentOS", f"remove_action({action_id})")
@@ -1070,7 +1070,7 @@ def remove_all_actions(self) -> None:
         Clear all recorded or batched actions.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
 
         self._reporter.add_message("AgentOS", "remove_all_actions()")
@@ -1095,7 +1095,7 @@ def _send_command(self, command: Command) -> AskUIAgentOSSendResponseSchema:
                 on the server side.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
 
         header = Header(authentication=Guid(root=self._session_guid))
@@ -1126,7 +1126,7 @@ def get_mouse_position(self) -> Coordinate:
             Coordinate: Response containing the result of the mouse position change.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
         self._reporter.add_message("AgentOS", "get_mouse_position()")
         res = self._send_command(GetMousePositionCommand())
@@ -1147,7 +1147,7 @@ def set_mouse_position(self, x: int, y: int) -> None:
             y (int): The vertical coordinate (in pixels) to set the cursor to.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
         location = Location(x=Length(root=x), y=Length(root=y))
         command = SetMousePositionCommand(parameters=[location])
@@ -1166,7 +1166,7 @@ def render_quad(self, style: RenderObjectStyle) -> int:
             int: Object ID.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
         self._reporter.add_message("AgentOS", f"render_quad({style})")
         command = AddRenderObjectCommand(parameters=["Quad", style])
@@ -1186,7 +1186,7 @@ def render_line(self, style: RenderObjectStyle, points: list[Coordinate]) -> int
             int: Object ID.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
         self._reporter.add_message("AgentOS", f"render_line({style}, {points})")
         command = AddRenderObjectCommand(parameters=["Line", style, points])
@@ -1206,7 +1206,7 @@ def render_image(self, style: RenderObjectStyle, image_data: str) -> int:
             int: Object ID.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
         self._reporter.add_message("AgentOS", f"render_image({style}, [image_data])")
         image = RenderImage(root=image_data)
@@ -1228,7 +1228,7 @@ def render_text(self, style: RenderObjectStyle, content: str) -> int:
             int: Object ID.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
         self._reporter.add_message("AgentOS", f"render_text({style}, {content})")
         text = RenderText(root=content)
@@ -1249,7 +1249,7 @@ def update_render_object(self, object_id: int, style: RenderObjectStyle) -> None
             int: Object ID.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
         self._reporter.add_message(
             "AgentOS", f"update_render_object({object_id}, {style})"
@@ -1267,7 +1267,7 @@ def delete_render_object(self, object_id: int) -> None:
             object_id (RenderObjectId): The ID of the render object to delete.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
         self._reporter.add_message("AgentOS", f"delete_render_object({object_id})")
         render_object_id = RenderObjectId(root=object_id)
@@ -1280,7 +1280,7 @@ def clear_render_objects(self) -> None:
         Clear all render objects from the display.
         """
         assert isinstance(self._stub, controller_v1.ControllerAPIStub), (
-            "Stub is not initialized. Call Connect first."
+            "Stub is not initialized. Call `connect()` first."
         )
         self._reporter.add_message("AgentOS", "clear_render_objects()")
         command = ClearRenderObjectsCommand()

src/askui/tools/askui/askui_controller_settings.py

@@ -2,9 +2,8 @@
 import sys
 from functools import cached_property
 
-from pydantic import BaseModel, Field, field_validator, model_validator
+from pydantic import BaseModel, Field, field_validator
 from pydantic_settings import BaseSettings, SettingsConfigDict
-from typing_extensions import Self
 
 
 class RemoteDeviceController(BaseModel):
@@ -93,22 +92,6 @@ def validate_controller_args(cls, value: str) -> str:
 
         return value
 
-    @model_validator(mode="after")
-    def validate_either_component_registry_or_installation_directory_is_set(
-        self,
-    ) -> "Self":
-        if (
-            self.component_registry_file is None
-            and self.installation_directory is None
-            and self.controller_path_setting is None
-        ):
-            error_msg = (
-                "Either ASKUI_COMPONENT_REGISTRY_FILE, ASKUI_INSTALLATION_DIRECTORY, "
-                "or ASKUI_CONTROLLER_PATH environment variable must be set"
-            )
-            raise ValueError(error_msg)
-        return self
-
     def _find_remote_device_controller_by_installation_directory(
         self,
     ) -> pathlib.Path | None:

src/askui/tools/store/universal/__init__.py

@@ -5,6 +5,7 @@
 """
 
 from .list_files_tool import ListFilesTool
+from .load_image_tool import LoadImageTool
 from .print_to_console import PrintToConsoleTool
 from .read_from_file_tool import ReadFromFileTool
 from .write_to_file_tool import WriteToFileTool
@@ -14,4 +15,5 @@
     "PrintToConsoleTool",
     "ReadFromFileTool",
     "WriteToFileTool",
+    "LoadImageTool",
 ]

src/askui/tools/store/universal/list_files_tool.py

@@ -13,7 +13,7 @@ class ListFilesTool(Tool):
     directories during execution.
 
     Args:
-        base_dir (str): The base directory path where file listing will start.
+        base_dir (str | Path): The base directory path where file listing will start.
             All directory paths will be relative to this directory.
 
     Example:
@@ -29,7 +29,10 @@ class ListFilesTool(Tool):
         ```
     """
 
-    def __init__(self, base_dir: str) -> None:
+    def __init__(self, base_dir: str | Path) -> None:
+        if not isinstance(base_dir, Path):
+            base_dir = Path(base_dir)
+        base_dir = base_dir.absolute()
         super().__init__(
             name="list_files_tool",
             description=(
@@ -66,7 +69,7 @@ def __init__(self, base_dir: str) -> None:
                 "required": [],
             },
         )
-        self._base_dir = Path(base_dir)
+        self._base_dir = base_dir
 
     def __call__(self, directory_path: str = "", recursive: bool = False) -> str:
         """