Skip to content

Commit 0ec8847

Browse files
mlikasam-askuimlikasam-askui
committed
fix tools
1 parent f9393df commit 0ec8847

File tree

21 files changed

+174
-98
lines changed

21 files changed

+174
-98
lines changed

‎src/askui/agent.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
from askui.prompts.system import COMPUTER_AGENT_SYSTEM_PROMPT
1212
from askui.tools.computer import (
1313
ComputerGetMousePositionTool,
14+
ComputerGetSystemInfoTool,
1415
ComputerKeyboardPressedTool,
1516
ComputerKeyboardReleaseTool,
1617
ComputerKeyboardTapTool,
@@ -90,6 +91,7 @@ def __init__(
9091
models=models,
9192
tools=[
9293
ExceptionTool(),
94+
ComputerGetSystemInfoTool(),
9395
ComputerGetMousePositionTool(),
9496
ComputerKeyboardPressedTool(),
9597
ComputerKeyboardReleaseTool(),

‎src/askui/agent_base.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@ def __init__(
106106
)
107107

108108
self.act_tool_collection = ToolCollection(tools=tools)
109+
self.act_tool_collection.add_agent_os(agent_os)
109110

110111
self.act_settings = ActSettings()
111112
self.caching_settings = CachingSettings()

‎src/askui/models/askui/askui_computer_base_tool.py‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
from typing import Any
2+
3+
from askui.models.shared import ComputerBaseTool
4+
from askui.models.shared.tool_tags import ToolTags
5+
from askui.tools.agent_os import AgentOs
6+
from askui.tools.agent_os_type_error import AgentOsTypeError
7+
from askui.tools.android.agent_os import AndroidAgentOs
8+
from askui.tools.askui.askui_controller import AskUiControllerClient
9+
10+
11+
class AskUiComputerBaseTool(ComputerBaseTool):
12+
"""
13+
Base tool for AskUI computer tools.
14+
Tools that can only operate with the AskUiControllerClient as agent_os.
15+
"""
16+
17+
def __init__(
18+
self,
19+
agent_os: AgentOs | None = None,
20+
required_tags: list[str] | None = None,
21+
**kwargs: Any,
22+
) -> None:
23+
super().__init__(
24+
required_tags=[ToolTags.COMPUTER.value, ToolTags.ASKUI_CONTROLLER.value]
25+
+ (required_tags or []),
26+
agent_os=agent_os,
27+
**kwargs,
28+
)
29+
30+
@property
31+
def agent_os(self) -> AskUiControllerClient:
32+
"""Get the agent OS.
33+
34+
Returns:
35+
AgentOs: The agent OS instance.
36+
"""
37+
agent_os = super().agent_os
38+
if not isinstance(agent_os, AskUiControllerClient):
39+
raise AgentOsTypeError(
40+
expected_type=AskUiControllerClient,
41+
actual_type=type(agent_os),
42+
)
43+
return agent_os
44+
45+
@agent_os.setter
46+
def agent_os(self, agent_os: AgentOs | AndroidAgentOs) -> None:
47+
"""Set the agent OS facade.
48+
49+
Args:
50+
agent_os (AgentOs | AndroidAgentOs): The agent OS facade instance to set.
51+
52+
Raises:
53+
TypeError: If the agent OS is not an AgentOs instance.
54+
"""
55+
if not isinstance(agent_os, AskUiControllerClient):
56+
raise AgentOsTypeError(
57+
expected_type=AskUiControllerClient,
58+
actual_type=type(agent_os),
59+
)
60+
self._agent_os = agent_os

‎src/askui/models/shared/android_base_tool.py‎

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from askui.models.shared.tool_tags import ToolTags
44
from askui.models.shared.tools import ToolWithAgentOS
55
from askui.tools import AgentOs
6+
from askui.tools.agent_os_type_error import AgentOsTypeError
67
from askui.tools.android.agent_os import AndroidAgentOs
78

89

@@ -33,12 +34,10 @@ def agent_os(self) -> AndroidAgentOs:
3334
"""
3435
agent_os = super().agent_os
3536
if not isinstance(agent_os, AndroidAgentOs):
36-
msg = (
37-
"Agent OS is not an AndroidAgentOs. "
38-
"Call `agent_os = ...` or initialize the tool with an "
39-
"AndroidAgentOs."
37+
raise AgentOsTypeError(
38+
expected_type=AndroidAgentOs,
39+
actual_type=type(agent_os),
4040
)
41-
raise TypeError(msg)
4241
return agent_os
4342

4443
@agent_os.setter
@@ -52,9 +51,8 @@ def agent_os(self, agent_os: AgentOs | AndroidAgentOs) -> None:
5251
TypeError: If the agent OS is not an AndroidAgentOs instance.
5352
"""
5453
if not isinstance(agent_os, AndroidAgentOs):
55-
msg = (
56-
"Agent OS must be an AndroidAgentOs instance. "
57-
f"Got {type(agent_os).__name__} instead."
54+
raise AgentOsTypeError(
55+
expected_type=AndroidAgentOs,
56+
actual_type=type(agent_os),
5857
)
59-
raise TypeError(msg)
6058
self._agent_os = agent_os

‎src/askui/models/shared/computer_base_tool.py‎

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from askui.models.shared.tool_tags import ToolTags
44
from askui.models.shared.tools import ToolWithAgentOS
55
from askui.tools.agent_os import AgentOs
6+
from askui.tools.agent_os_type_error import AgentOsTypeError
67
from askui.tools.android.agent_os import AndroidAgentOs
78

89

@@ -30,12 +31,10 @@ def agent_os(self) -> AgentOs:
3031
"""
3132
agent_os = super().agent_os
3233
if not isinstance(agent_os, AgentOs):
33-
msg = (
34-
"Agent OS is not a ComputerAgentOs. "
35-
"Call `agent_os = ...` or initialize the tool with a "
36-
"ComputerAgentOs."
34+
raise AgentOsTypeError(
35+
expected_type=AgentOs,
36+
actual_type=type(agent_os),
3737
)
38-
raise TypeError(msg)
3938
return agent_os
4039

4140
@agent_os.setter
@@ -49,9 +48,8 @@ def agent_os(self, agent_os: AgentOs | AndroidAgentOs) -> None:
4948
TypeError: If the agent OS is not an AgentOs instance.
5049
"""
5150
if not isinstance(agent_os, AgentOs):
52-
msg = (
53-
"Agent OS must be an AgentOs instance. "
54-
f"Got {type(agent_os).__name__} instead."
51+
raise AgentOsTypeError(
52+
expected_type=AgentOs,
53+
actual_type=type(agent_os),
5554
)
56-
raise TypeError(msg)
5755
self._agent_os = agent_os

‎src/askui/models/shared/tool_tags.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,5 +5,6 @@ class ToolTags(str, Enum):
55
"""Default tool tags."""
66

77
ANDROID = "android"
8+
ASKUI_CONTROLLER = "askui_controller"
89
COMPUTER = "computer"
910
SCALED_AGENT_OS = "scaled_agent_os"

‎src/askui/models/shared/tools.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -473,12 +473,13 @@ def _run_regular_tool(
473473
except AgentException:
474474
raise
475475
except Exception as e: # noqa: BLE001
476+
error_message = getattr(e, "message", str(e))
476477
logger.warning(
477478
"Tool failed",
478-
extra={"tool_name": tool_use_block_param.name, "error": str(e)},
479+
extra={"tool_name": tool_use_block_param.name, "error": error_message},
479480
)
480481
return ToolResultBlockParam(
481-
content=str(e),
482+
content=f"Tool raised an unexpected error: {error_message}",
482483
is_error=True,
483484
tool_use_id=tool_use_block_param.id,
484485
)

‎src/askui/tools/agent_os_type_error.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
class AgentOsTypeError(TypeError):
2+
"""Exception raised when the agent OS is not of the expected type."""
3+
4+
def __init__(self, expected_type: type, actual_type: type):
5+
self.expected_type = expected_type
6+
self.actual_type = actual_type
7+
super().__init__(
8+
f"Agent OS must be an {expected_type.__name__} instance. "
9+
f"Got {actual_type.__name__} instead."
10+
)

‎src/askui/tools/askui/__init__.py‎

Lines changed: 0 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,6 @@
11
from .askui_controller import AskUiControllerClient, AskUiControllerServer
2-
from .get_system_info_tool import GetSystemInfoTool
3-
from .window_managment import (
4-
AddWindowAsVirtualDisplayTool,
5-
ListProcessTool,
6-
ListProcessWindowsTool,
7-
SetProcessInFocusTool,
8-
SetWindowInFocusTool,
9-
)
102

113
__all__ = [
124
"AskUiControllerClient",
135
"AskUiControllerServer",
14-
"AddWindowAsVirtualDisplayTool",
15-
"GetSystemInfoTool",
16-
"ListProcessTool",
17-
"ListProcessWindowsTool",
18-
"SetWindowInFocusTool",
19-
"SetProcessInFocusTool",
206
]

‎src/askui/tools/askui/askui_controller.py‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from typing_extensions import Self, override
1414

1515
from askui.container import telemetry
16+
from askui.models.shared.tool_tags import ToolTags
1617
from askui.reporting import NULL_REPORTER, Reporter
1718
from askui.tools.agent_os import (
1819
AgentOs,
@@ -193,6 +194,7 @@ def __init__(
193194
self._controller_server = controller_server or AskUiControllerServer()
194195
self._session_guid = "{" + str(uuid.uuid4()) + "}"
195196
self._settings = settings or AskUiControllerClientSettings()
197+
self.tags.append(ToolTags.ASKUI_CONTROLLER.value)
196198

197199
@telemetry.record_call()
198200
@override
@@ -269,11 +271,19 @@ def disconnect(self) -> None:
269271
This method stops the execution, ends the session, closes the gRPC channel,
270272
and stops the controller server.
271273
"""
272-
self._stop_execution()
273-
self._stop_session()
274-
if self._channel is not None:
275-
self._channel.close()
276-
self._controller_server.stop()
274+
try:
275+
self._stop_execution()
276+
self._stop_session()
277+
if self._channel is not None:
278+
self._channel.close()
279+
self._controller_server.stop()
280+
except Exception as e: # noqa: BLE001
281+
# We want to catch all other exceptions here and not re-raise them
282+
msg = (
283+
"Error while disconnecting from the AskUI Remote Device Controller"
284+
f"Error: {e}"
285+
)
286+
logger.exception(msg)
277287

278288
@telemetry.record_call()
279289
def __enter__(self) -> Self:

0 commit comments

Comments
 (0)