Merge pull request #220 from askui/feat/refactor-tools-auto-inject-agent-os

refactor(tools): auto-inject agent OS and split computer tool

- Split monolithic Computer20250124Tool into 16 individual tool classes
- Add ComputerBaseTool/AndroidBaseTool with automatic agent_os injection
- Enhance ToolCollection to auto-inject agent_os via required_tags
- Introduce ComputerAgentOsFacade for coordinate scaling
- Simplify agent initialization by removing beta flag logic
- Reorganize geometry types to askui.models.types.geometry
- Introduce tool store (askui.tools.store) with optional tools organized
  by category (computer, android, universal)

Improves maintainability and type safety while maintaining backward compatibility.

Author: mlikasam-askui

README.md

@@ -9,16 +9,27 @@ Join the [AskUI Discord](https://discord.gg/Gu35zMGxbx).
 
 ## Table of Contents
 
-- [📖 Introduction](#-introduction)
-- [📦 Installation](#-installation)
-  - [AskUI Python Package](#askui-python-package)
-  - [AskUI Agent OS](#askui-agent-os)
-- [🚀 Quickstart](#-quickstart)
+- [🤖 AskUI Vision Agent](#-askui-vision-agent)
+  - [Table of Contents](#table-of-contents)
+  - [📖 Introduction](#-introduction)
+  - [📦 Installation](#-installation)
+    - [AskUI Python Package](#askui-python-package)
+    - [AskUI Agent OS](#askui-agent-os)
+      - [AMD64](#amd64)
+      - [ARM64](#arm64)
+      - [AMD64](#amd64-1)
+      - [ARM64](#arm64-1)
+      - [ARM64](#arm64-2)
+  - [🚀 Quickstart](#-quickstart)
     - [🧑 Control your devices](#-control-your-devices)
     - [🤖 Let AI agents control your devices](#-let-ai-agents-control-your-devices)
-- [📚 Further Documentation](#-further-documentation)
-- [🤝 Contributing](#-contributing)
-- [📜 License](#-license)
+      - [🔐 Sign up with AskUI](#-sign-up-with-askui)
+      - [⚙️ Configure environment variables](#️-configure-environment-variables)
+      - [💻 Example](#-example)
+    - [🛠️ Extending Agents with Tool Store](#️-extending-agents-with-tool-store)
+  - [📚 Further Documentation](#-further-documentation)
+  - [🤝 Contributing](#-contributing)
+  - [📜 License](#-license)
 
 ## 📖 Introduction
 
@@ -184,6 +195,41 @@ Run the script with `python <file path>`, e.g `python test.py`.
 
 If you see a lot of logs and the first paragraph of the introduction in the console, congratulations! You've successfully let AI agents control your device to automate a task! If you have any issues, please check the [documentation](https://docs.askui.com/01-tutorials/01-your-first-agent#common-issues-and-solutions) or join our [Discord](https://discord.gg/Gu35zMGxbx) for support.
 
+### 🛠️ Extending Agents with Tool Store
+
+The Tool Store provides optional tools to extend your agents' capabilities. Import tools from `askui.tools.store` and pass them to `agent.act()` or pass them to the agent constructor as `act_tools`.
+
+**Example passing tools to `agent.act()`:**
+```python
+from askui import VisionAgent
+from askui.tools.store.computer import ComputerSaveScreenshotTool
+from askui.tools.store.universal import PrintToConsoleTool
+
+with VisionAgent() as agent:
+    agent.act(
+        "Take a screenshot and save it as demo/demo.png, then print a status message",
+        tools=[
+            ComputerSaveScreenshotTool(base_dir="./screenshots"),
+            PrintToConsoleTool()
+        ]
+    )
+```
+
+**Example passing tools to the agent constructor:**
+```python
+from askui import VisionAgent
+from askui.tools.store.computer import ComputerSaveScreenshotTool
+from askui.tools.store.universal import PrintToConsoleTool
+
+with VisionAgent(act_tools=[
+    ComputerSaveScreenshotTool(base_dir="./screenshots"),
+    PrintToConsoleTool()
+]) as agent:
+    agent.act("Take a screenshot and save it as demo/demo.png, then print a status message")
+```
+
+Tools are organized by category: `universal/` (work with any agent), `computer/` (require `AgentOs`) works  only with VisionAgent and `android/` (require `AndroidAgentOs`) works only with AndroidVisionAgent.
+
 ## 📚 Further Documentation
 
 Aside from our [official documentation](https://docs.askui.com), we also have some additional guides and examples under the [docs](docs) folder that you may find useful, for example:

src/askui/agent.py

@@ -1,37 +1,39 @@
 import logging
-from typing import TYPE_CHECKING, Annotated, Literal, Optional
+from typing import Annotated, Literal, Optional
 
-from anthropic import Omit, omit
 from pydantic import ConfigDict, Field, validate_call
-from typing_extensions import override
 
 from askui.agent_base import AgentBase
 from askui.container import telemetry
 from askui.locators.locators import Locator
-from askui.models.shared.settings import (
-    COMPUTER_USE_20250124_BETA_FLAG,
-    COMPUTER_USE_20251124_BETA_FLAG,
-    ActSettings,
-    MessageSettings,
-)
+from askui.models.shared.settings import ActSettings, MessageSettings
 from askui.models.shared.tools import Tool
 from askui.prompts.system import COMPUTER_AGENT_SYSTEM_PROMPT
-from askui.tools.computer import Computer20250124Tool
+from askui.tools.computer import (
+    ComputerGetMousePositionTool,
+    ComputerKeyboardPressedTool,
+    ComputerKeyboardReleaseTool,
+    ComputerKeyboardTapTool,
+    ComputerListDisplaysTool,
+    ComputerMouseClickTool,
+    ComputerMouseHoldDownTool,
+    ComputerMouseReleaseTool,
+    ComputerMouseScrollTool,
+    ComputerMoveMouseTool,
+    ComputerRetrieveActiveDisplayTool,
+    ComputerScreenshotTool,
+    ComputerSetActiveDisplayTool,
+    ComputerTypeTool,
+)
 from askui.tools.exception_tool import ExceptionTool
-from askui.tools.list_displays_tool import ListDisplaysTool
-from askui.tools.retrieve_active_display_tool import RetrieveActiveDisplayTool
-from askui.tools.set_active_display_tool import SetActiveDisplayTool
 
 from .models import ModelComposition
 from .models.models import ModelChoice, ModelRegistry, Point
 from .reporting import CompositeReporter, Reporter
 from .retry import Retry
-from .tools import AgentToolbox, ModifierKey, PcKey
+from .tools import AgentToolbox, ComputerAgentOsFacade, ModifierKey, PcKey
 from .tools.askui import AskUiControllerClient
 
-if TYPE_CHECKING:
-    from anthropic.types import AnthropicBetaParam
-
 logger = logging.getLogger(__name__)
 
 
@@ -88,14 +90,35 @@ def __init__(
             models=models,
             tools=[
                 ExceptionTool(),
-                SetActiveDisplayTool(agent_os=self.tools.os),
-                RetrieveActiveDisplayTool(agent_os=self.tools.os),
-                ListDisplaysTool(agent_os=self.tools.os),
+                ComputerGetMousePositionTool(),
+                ComputerKeyboardPressedTool(),
+                ComputerKeyboardReleaseTool(),
+                ComputerKeyboardTapTool(),
+                ComputerMouseClickTool(),
+                ComputerMouseHoldDownTool(),
+                ComputerMouseReleaseTool(),
+                ComputerMouseScrollTool(),
+                ComputerMoveMouseTool(),
+                ComputerScreenshotTool(),
+                ComputerTypeTool(),
+                ComputerListDisplaysTool(),
+                ComputerRetrieveActiveDisplayTool(),
+                ComputerSetActiveDisplayTool(),
             ]
             + (act_tools or []),
             agent_os=self.tools.os,
             model_provider=model_provider,
         )
+        self.act_agent_os_facade: ComputerAgentOsFacade = ComputerAgentOsFacade(
+            self.tools.os
+        )
+        self.act_tool_collection.add_agent_os(self.act_agent_os_facade)
+        self.act_settings = ActSettings(
+            messages=MessageSettings(
+                system=COMPUTER_AGENT_SYSTEM_PROMPT,
+                thinking={"type": "enabled", "budget_tokens": 2048},
+            )
+        )
 
     @telemetry.record_call(exclude={"locator"})
     @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
@@ -396,34 +419,6 @@ def mouse_down(
         logger.debug("VisionAgent received instruction to mouse_down '%s'", button)
         self.tools.os.mouse_down(button)
 
-    @override
-    def _get_default_settings_for_act(self, model: str) -> ActSettings:
-        computer_use_beta_flag: list[AnthropicBetaParam] | Omit
-        if "claude-opus-4-5-20251101" in model:
-            computer_use_beta_flag = [COMPUTER_USE_20251124_BETA_FLAG]
-        elif (
-            "claude-sonnet-4-5-20250929" in model
-            or "claude-haiku-4-5-20251001" in model
-            or "claude-opus-4-1-20250805" in model
-            or "claude-opus-4-20250514" in model
-            or "claude-sonnet-4-20250514" in model
-            or "claude-3-7-sonnet-20250219" in model
-        ):
-            computer_use_beta_flag = [COMPUTER_USE_20250124_BETA_FLAG]
-        else:
-            computer_use_beta_flag = omit
-        return ActSettings(
-            messages=MessageSettings(
-                system=COMPUTER_AGENT_SYSTEM_PROMPT,
-                betas=computer_use_beta_flag,
-                thinking={"type": "enabled", "budget_tokens": 2048},
-            ),
-        )
-
-    @override
-    def _get_default_tools_for_act(self, model: str) -> list[Tool]:
-        return self._tools + [Computer20250124Tool(agent_os=self.tools.os)]
-
     @telemetry.record_call()
     @validate_call
     def keyboard(

src/askui/agent_base.py

@@ -37,10 +37,9 @@
     ModelChoice,
     ModelName,
     ModelRegistry,
-    Point,
-    PointList,
     TotalModelChoice,
 )
+from .models.types.geometry import Point, PointList
 from .models.types.response_schemas import ResponseSchema
 from .reporting import Reporter
 from .retry import ConfigurableRetry, Retry
@@ -106,6 +105,11 @@ def __init__(
             reporter=self._reporter, models=models or {}
         )
 
+        self.act_tool_collection = ToolCollection(tools=tools)
+
+        self.act_settings = ActSettings()
+        self.caching_settings = CachingSettings()
+
     def _init_model_router(
         self,
         reporter: Reporter,
@@ -300,11 +304,9 @@ def act(
             [MessageParam(role="user", content=goal)] if isinstance(goal, str) else goal
         )
         _model = self._get_model(model, "act")
-        _settings = settings or self._get_default_settings_for_act(_model)
+        _settings = settings or self.act_settings
 
-        _caching_settings: CachingSettings = (
-            caching_settings or self._get_default_caching_settings_for_act(_model)
-        )
+        _caching_settings: CachingSettings = caching_settings or self.caching_settings
 
         tools, on_message, cached_execution_tool = self._patch_act_with_cache(
             _caching_settings, _settings, tools, on_message
@@ -323,14 +325,14 @@ def act(
         )
 
     def _build_tools(
-        self, tools: list[Tool] | ToolCollection | None, model: str
+        self, tools: list[Tool] | ToolCollection | None, _model: str
     ) -> ToolCollection:
-        default_tools = self._get_default_tools_for_act(model)
+        tool_collection = self.act_tool_collection
         if isinstance(tools, list):
-            return ToolCollection(tools=default_tools + tools)
+            tool_collection.append_tool(*tools)
         if isinstance(tools, ToolCollection):
-            return ToolCollection(default_tools) + tools
-        return ToolCollection(tools=default_tools)
+            tool_collection += tools
+        return tool_collection
 
     def _patch_act_with_cache(
         self,
@@ -399,15 +401,6 @@ def _patch_act_with_cache(
 
         return tools, on_message, cached_execution_tool
 
-    def _get_default_settings_for_act(self, model: str) -> ActSettings:  # noqa: ARG002
-        return ActSettings()
-
-    def _get_default_caching_settings_for_act(self, model: str) -> CachingSettings:  # noqa: ARG002
-        return CachingSettings()
-
-    def _get_default_tools_for_act(self, model: str) -> list[Tool]:  # noqa: ARG002
-        return self._tools
-
     @overload
     def get(
         self,

src/askui/android_agent.py

@@ -2,7 +2,6 @@
 from typing import Annotated, overload
 
 from pydantic import ConfigDict, Field, validate_call
-from typing_extensions import override
 
 from askui.agent_base import AgentBase
 from askui.container import telemetry
@@ -64,7 +63,7 @@ class AndroidVisionAgent(AgentBase):
         ```
     """
 
-    @telemetry.record_call(exclude={"model_router", "reporters", "tools"})
+    @telemetry.record_call(exclude={"model_router", "reporters", "tools", "act_tools"})
     @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
     def __init__(
         self,
@@ -85,25 +84,33 @@ def __init__(
             retry=retry,
             models=models,
             tools=[
-                AndroidScreenshotTool(self.act_agent_os_facade),
-                AndroidTapTool(self.act_agent_os_facade),
-                AndroidTypeTool(self.act_agent_os_facade),
-                AndroidDragAndDropTool(self.act_agent_os_facade),
-                AndroidKeyTapEventTool(self.act_agent_os_facade),
-                AndroidSwipeTool(self.act_agent_os_facade),
-                AndroidKeyCombinationTool(self.act_agent_os_facade),
-                AndroidShellTool(self.act_agent_os_facade),
-                AndroidSelectDeviceBySerialNumberTool(self.act_agent_os_facade),
-                AndroidSelectDisplayByUniqueIDTool(self.act_agent_os_facade),
-                AndroidGetConnectedDevicesSerialNumbersTool(self.act_agent_os_facade),
-                AndroidGetConnectedDisplaysInfosTool(self.act_agent_os_facade),
-                AndroidGetCurrentConnectedDeviceInfosTool(self.act_agent_os_facade),
+                AndroidScreenshotTool(),
+                AndroidTapTool(),
+                AndroidTypeTool(),
+                AndroidDragAndDropTool(),
+                AndroidKeyTapEventTool(),
+                AndroidSwipeTool(),
+                AndroidKeyCombinationTool(),
+                AndroidShellTool(),
+                AndroidSelectDeviceBySerialNumberTool(),
+                AndroidSelectDisplayByUniqueIDTool(),
+                AndroidGetConnectedDevicesSerialNumbersTool(),
+                AndroidGetConnectedDisplaysInfosTool(),
+                AndroidGetCurrentConnectedDeviceInfosTool(),
                 ExceptionTool(),
             ]
             + (act_tools or []),
             agent_os=self.os,
             model_provider=model_provider,
         )
+        self.act_tool_collection.add_agent_os(self.act_agent_os_facade)
+        self.act_settings = ActSettings(
+            messages=MessageSettings(
+                system=ANDROID_AGENT_SYSTEM_PROMPT,
+                thinking={"type": "disabled"},
+                temperature=0.0,
+            ),
+        )
 
     @overload
     def tap(
@@ -353,13 +360,3 @@ def set_device_by_serial_number(
             f"set_device_by_serial_number(device_sn='{device_sn}')",
         )
         self.os.set_device_by_serial_number(device_sn)
-
-    @override
-    def _get_default_settings_for_act(self, model: str) -> ActSettings:
-        return ActSettings(
-            messages=MessageSettings(
-                system=ANDROID_AGENT_SYSTEM_PROMPT,
-                thinking={"type": "disabled"},
-                temperature=0.0,
-            ),
-        )

src/askui/chat/__main__.py

@@ -18,4 +18,5 @@
         reload=False,
         workers=1,
         log_config=None,
+        timeout_graceful_shutdown=5,
     )