Refactor ToolDefinition architecture to use subclass pattern for all tools (#971)

Co-authored-by: openhands <[email protected]>
Co-authored-by: simonrosenberg <[email protected]>

Author: 3 people

examples/01_standalone_sdk/02_custom_tools.py

@@ -26,8 +26,8 @@
 )
 from openhands.tools.execute_bash import (
     BashExecutor,
+    BashTool,
     ExecuteBashAction,
-    execute_bash_tool,
 )
 from openhands.tools.file_editor import FileEditorTool
 
@@ -115,6 +115,42 @@ def __call__(self, action: GrepAction, conversation=None) -> GrepObservation:  #
 * When you are doing an open ended search that may require multiple rounds of globbing and grepping, use the Agent tool instead
 """  # noqa: E501
 
+
+# --- Tool Definition ---
+
+
+class GrepTool(ToolDefinition[GrepAction, GrepObservation]):
+    """A custom grep tool that searches file contents using regular expressions."""
+
+    @classmethod
+    def create(
+        cls, conv_state, bash_executor: BashExecutor | None = None
+    ) -> Sequence[ToolDefinition]:
+        """Create GrepTool instance with a GrepExecutor.
+
+        Args:
+            conv_state: Conversation state to get working directory from.
+            bash_executor: Optional bash executor to reuse. If not provided,
+                         a new one will be created.
+
+        Returns:
+            A sequence containing a single GrepTool instance.
+        """
+        if bash_executor is None:
+            bash_executor = BashExecutor(working_dir=conv_state.workspace.working_dir)
+        grep_executor = GrepExecutor(bash_executor)
+
+        return [
+            cls(
+                name="grep",
+                description=_GREP_DESCRIPTION,
+                action_type=GrepAction,
+                observation_type=GrepObservation,
+                executor=grep_executor,
+            )
+        ]
+
+
 # Configure LLM
 api_key = os.getenv("LLM_API_KEY")
 assert api_key is not None, "LLM_API_KEY environment variable is not set."
@@ -135,16 +171,11 @@ def _make_bash_and_grep_tools(conv_state) -> list[ToolDefinition]:
     """Create execute_bash and custom grep tools sharing one executor."""
 
     bash_executor = BashExecutor(working_dir=conv_state.workspace.working_dir)
-    bash_tool = execute_bash_tool.set_executor(executor=bash_executor)
-
-    grep_executor = GrepExecutor(bash_executor)
-    grep_tool = ToolDefinition(
-        name="grep",
-        description=_GREP_DESCRIPTION,
-        action_type=GrepAction,
-        observation_type=GrepObservation,
-        executor=grep_executor,
-    )
+    # bash_tool = execute_bash_tool.set_executor(executor=bash_executor)
+    bash_tool = BashTool.create(conv_state, executor=bash_executor)[0]
+
+    # Use the GrepTool.create() method with shared bash_executor
+    grep_tool = GrepTool.create(conv_state, bash_executor=bash_executor)[0]
 
     return [bash_tool, grep_tool]
 

openhands-sdk/openhands/sdk/__init__.py

@@ -37,7 +37,6 @@
     Action,
     Observation,
     Tool,
-    ToolBase,
     ToolDefinition,
     list_registered_tools,
     register_tool,
@@ -67,7 +66,6 @@
     "RedactedThinkingBlock",
     "Tool",
     "ToolDefinition",
-    "ToolBase",
     "AgentBase",
     "Agent",
     "Action",

openhands-sdk/openhands/sdk/agent/agent.py

@@ -37,7 +37,6 @@
 from openhands.sdk.security.llm_analyzer import LLMSecurityAnalyzer
 from openhands.sdk.tool import (
     Action,
-    FinishTool,
     Observation,
 )
 from openhands.sdk.tool.builtins import FinishAction, ThinkAction
@@ -431,6 +430,6 @@ def _execute_action_event(
         on_event(obs_event)
 
         # Set conversation state
-        if tool.name == FinishTool.name:
+        if tool.name == "finish":
             state.agent_status = AgentExecutionStatus.FINISHED
         return obs_event

openhands-sdk/openhands/sdk/agent/base.py

@@ -223,7 +223,9 @@ def _initialize(self, state: "ConversationState"):
             )
 
         # Always include built-in tools; not subject to filtering
-        tools.extend(BUILT_IN_TOOLS)
+        # Instantiate built-in tools using their .create() method
+        for tool_class in BUILT_IN_TOOLS:
+            tools.extend(tool_class.create(state))
 
         # Check tool types
         for tool in tools:

openhands-sdk/openhands/sdk/llm/llm.py

@@ -26,7 +26,7 @@
 
 
 if TYPE_CHECKING:  # type hints only, avoid runtime import cycle
-    from openhands.sdk.tool.tool import ToolBase
+    from openhands.sdk.tool.tool import ToolDefinition
 
 from openhands.sdk.utils.pydantic_diff import pretty_pydantic_diff
 
@@ -425,7 +425,7 @@ def restore_metrics(self, metrics: Metrics) -> None:
     def completion(
         self,
         messages: list[Message],
-        tools: Sequence[ToolBase] | None = None,
+        tools: Sequence[ToolDefinition] | None = None,
         _return_metrics: bool = False,
         add_security_risk_prediction: bool = False,
         **kwargs,
@@ -562,7 +562,7 @@ def _one_attempt(**retry_kwargs) -> ModelResponse:
     def responses(
         self,
         messages: list[Message],
-        tools: Sequence[ToolBase] | None = None,
+        tools: Sequence[ToolDefinition] | None = None,
         include: list[str] | None = None,
         store: bool | None = None,
         _return_metrics: bool = False,

openhands-sdk/openhands/sdk/llm/router/base.py

@@ -11,7 +11,7 @@
 from openhands.sdk.llm.llm_response import LLMResponse
 from openhands.sdk.llm.message import Message
 from openhands.sdk.logger import get_logger
-from openhands.sdk.tool.tool import ToolBase
+from openhands.sdk.tool.tool import ToolDefinition
 
 
 logger = get_logger(__name__)
@@ -49,7 +49,7 @@ def validate_llms_not_empty(cls, v):
     def completion(
         self,
         messages: list[Message],
-        tools: Sequence[ToolBase] | None = None,
+        tools: Sequence[ToolDefinition] | None = None,
         return_metrics: bool = False,
         add_security_risk_prediction: bool = False,
         **kwargs,

openhands-sdk/openhands/sdk/mcp/utils.py

@@ -8,7 +8,7 @@
 
 from openhands.sdk.logger import get_logger
 from openhands.sdk.mcp import MCPClient, MCPToolDefinition
-from openhands.sdk.tool.tool import ToolBase
+from openhands.sdk.tool.tool import ToolDefinition
 
 
 logger = get_logger(__name__)
@@ -30,9 +30,9 @@ async def log_handler(message: LogMessage):
     logger.log(level, msg, extra=extra)
 
 
-async def _list_tools(client: MCPClient) -> list[ToolBase]:
+async def _list_tools(client: MCPClient) -> list[ToolDefinition]:
     """List tools from an MCP client."""
-    tools: list[ToolBase] = []
+    tools: list[ToolDefinition] = []
 
     async with client:
         assert client.is_connected(), "MCP client is not connected."

openhands-sdk/openhands/sdk/tool/__init__.py

@@ -14,7 +14,6 @@
 from openhands.sdk.tool.tool import (
     ExecutableTool,
     ToolAnnotations,
-    ToolBase,
     ToolDefinition,
     ToolExecutor,
 )
@@ -23,7 +22,6 @@
 __all__ = [
     "Tool",
     "ToolDefinition",
-    "ToolBase",
     "ToolAnnotations",
     "ToolExecutor",
     "ExecutableTool",

openhands-sdk/openhands/sdk/tool/builtins/finish.py

@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Self
 
 from pydantic import Field
 from rich.text import Text
@@ -16,6 +16,7 @@
 
 if TYPE_CHECKING:
     from openhands.sdk.conversation.base import BaseConversation
+    from openhands.sdk.conversation.state import ConversationState
 
 
 class FinishAction(Action):
@@ -67,17 +68,42 @@ def __call__(
         return FinishObservation(message=action.message)
 
 
-FinishTool = ToolDefinition(
-    name="finish",
-    action_type=FinishAction,
-    observation_type=FinishObservation,
-    description=TOOL_DESCRIPTION,
-    executor=FinishExecutor(),
-    annotations=ToolAnnotations(
-        title="finish",
-        readOnlyHint=True,
-        destructiveHint=False,
-        idempotentHint=True,
-        openWorldHint=False,
-    ),
-)
+class FinishTool(ToolDefinition[FinishAction, FinishObservation]):
+    """Tool for signaling the completion of a task or conversation."""
+
+    @classmethod
+    def create(
+        cls,
+        conv_state: "ConversationState | None" = None,  # noqa: ARG003
+        **params,
+    ) -> Sequence[Self]:
+        """Create FinishTool instance.
+
+        Args:
+            conv_state: Optional conversation state (not used by FinishTool).
+            **params: Additional parameters (none supported).
+
+        Returns:
+            A sequence containing a single FinishTool instance.
+
+        Raises:
+            ValueError: If any parameters are provided.
+        """
+        if params:
+            raise ValueError("FinishTool doesn't accept parameters")
+        return [
+            cls(
+                name="finish",
+                action_type=FinishAction,
+                observation_type=FinishObservation,
+                description=TOOL_DESCRIPTION,
+                executor=FinishExecutor(),
+                annotations=ToolAnnotations(
+                    title="finish",
+                    readOnlyHint=True,
+                    destructiveHint=False,
+                    idempotentHint=True,
+                    openWorldHint=False,
+                ),
+            )
+        ]

openhands-sdk/openhands/sdk/tool/builtins/think.py

@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Self
 
 from pydantic import Field
 from rich.text import Text
@@ -16,6 +16,7 @@
 
 if TYPE_CHECKING:
     from openhands.sdk.conversation.base import BaseConversation
+    from openhands.sdk.conversation.state import ConversationState
 
 
 class ThinkAction(Action):
@@ -83,16 +84,41 @@ def __call__(
         return ThinkObservation()
 
 
-ThinkTool = ToolDefinition(
-    name="think",
-    description=THINK_DESCRIPTION,
-    action_type=ThinkAction,
-    observation_type=ThinkObservation,
-    executor=ThinkExecutor(),
-    annotations=ToolAnnotations(
-        readOnlyHint=True,
-        destructiveHint=False,
-        idempotentHint=True,
-        openWorldHint=False,
-    ),
-)
+class ThinkTool(ToolDefinition[ThinkAction, ThinkObservation]):
+    """Tool for logging thoughts without making changes."""
+
+    @classmethod
+    def create(
+        cls,
+        conv_state: "ConversationState | None" = None,  # noqa: ARG003
+        **params,
+    ) -> Sequence[Self]:
+        """Create ThinkTool instance.
+
+        Args:
+            conv_state: Optional conversation state (not used by ThinkTool).
+            **params: Additional parameters (none supported).
+
+        Returns:
+            A sequence containing a single ThinkTool instance.
+
+        Raises:
+            ValueError: If any parameters are provided.
+        """
+        if params:
+            raise ValueError("ThinkTool doesn't accept parameters")
+        return [
+            cls(
+                name="think",
+                description=THINK_DESCRIPTION,
+                action_type=ThinkAction,
+                observation_type=ThinkObservation,
+                executor=ThinkExecutor(),
+                annotations=ToolAnnotations(
+                    readOnlyHint=True,
+                    destructiveHint=False,
+                    idempotentHint=True,
+                    openWorldHint=False,
+                ),
+            )
+        ]