Merge pull request #306 from CasperGN/chore/stronger-tool-return-types

fix: improve type checking for the mcp and toolbox functions

Author: sicoyle

dapr_agents/tool/base.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import inspect
 import logging
 from typing import TYPE_CHECKING, Callable, Type, Optional, Any, Dict
@@ -80,10 +82,10 @@ def from_func(cls, func: Callable) -> "AgentTool":
     @classmethod
     def from_mcp(
         cls,
-        mcp_tool: "MCPTool",
-        session: "Optional[ClientSession]" = None,
+        mcp_tool: MCPTool,
+        session: Optional[ClientSession] = None,
         connection: Any = None,
-    ) -> "AgentTool":
+    ) -> AgentTool:
         """
         Create an AgentTool from an MCPTool and a session or connection.
 
@@ -157,10 +159,10 @@ async def executor(**kwargs: Any) -> Any:
     @classmethod
     def from_mcp_many(
         cls,
-        mcp_tools: list,
-        session: "ClientSession" = None,
+        mcp_tools: list[MCPTool],
+        session: Optional[ClientSession] = None,
         connection: Any = None,
-    ) -> list:
+    ) -> list[AgentTool]:
         """
         Batch-create AgentTool objects from a list of MCPTool objects.
 
@@ -182,7 +184,7 @@ def from_mcp_many(
         ]
 
     @classmethod
-    async def from_mcp_session(cls, session: "ClientSession") -> list:
+    async def from_mcp_session(cls, session: ClientSession) -> list[AgentTool]:
         """
         Fetch all tools and wrap them as AgentTool objects.
 
@@ -201,8 +203,8 @@ async def from_mcp_session(cls, session: "ClientSession") -> list:
     @classmethod
     def from_toolbox(
         cls,
-        toolbox_tool: "ToolboxSyncTool",
-    ) -> "AgentTool":
+        toolbox_tool: ToolboxSyncTool,
+    ) -> AgentTool:
         """
         Create an AgentTool from a ToolboxSyncTool.
 
@@ -275,8 +277,8 @@ def executor(**kwargs: Any) -> Any:
     @classmethod
     def from_toolbox_many(
         cls,
-        toolbox_tools: list,
-    ) -> list:
+        toolbox_tools: list[ToolboxSyncTool],
+    ) -> list[AgentTool]:
         """
         Batch-create AgentTool objects from a list of ToolboxSyncTool objects.