Add tool tests

src/fastmcp/models.py

@@ -1,18 +0,0 @@
-"""Pydantic models for FastMCP."""
-
-from typing import Callable, Optional, Type
-
-from pydantic import BaseModel
-
-
-class Tool(BaseModel):
-    """Internal tool registration info."""
-
-    model_config: dict = dict(arbitrary_types_allowed=True)
-
-    func: Callable
-    name: str
-    description: str
-    input_schema: dict
-    is_async: bool
-    pydantic_model: Optional[Type[BaseModel]] = None

src/fastmcp/server.py

@@ -35,7 +35,7 @@ class FastMCPServer:
                 Tool(
                     name=info.name,
                     description=info.description,
-                    inputSchema=info.input_schema,
+                    inputSchema=info.parameters,
                 )
                 for info in tools
             ]

src/fastmcp/tools.py

@@ -1,12 +1,54 @@
 """Tool management for FastMCP."""
 
 import inspect
-from typing import Any, Callable, Dict, Optional, get_type_hints
+from typing import Any, Callable, Dict, Optional
 
-from pydantic import BaseModel, create_model
+from pydantic import BaseModel, Field, TypeAdapter
 
 from .exceptions import ToolError
-from .models import Tool
+
+
+class Tool(BaseModel):
+    """Internal tool registration info."""
+
+    func: Callable = Field(exclude=True)
+    name: str = Field(description="Name of the tool")
+    description: str = Field(description="Description of what the tool does")
+    parameters: dict = Field(description="JSON schema for tool parameters")
+    is_async: bool = Field(description="Whether the tool is async")
+
+    @classmethod
+    def from_function(
+        cls,
+        func: Callable,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+    ) -> "Tool":
+        """Create a Tool from a function."""
+        func_name = name or func.__name__
+        func_doc = description or func.__doc__ or ""
+        is_async = inspect.iscoroutinefunction(func)
+
+        # Get schema from TypeAdapter - will fail if function isn't properly typed
+        schema = TypeAdapter(func).json_schema()
+
+        return cls(
+            func=func,
+            name=func_name,
+            description=func_doc,
+            parameters=schema,
+            is_async=is_async,
+        )
+
+    async def run(self, arguments: dict) -> Any:
+        """Run the tool with arguments."""
+        try:
+            # Call function with proper async handling
+            if self.is_async:
+                return await self.func(**arguments)
+            return self.func(**arguments)
+        except Exception as e:
+            raise ToolError(f"Error executing tool {self.name}: {e}") from e
 
 
 class ToolManager:
@@ -30,61 +72,12 @@ class ToolManager:
         description: Optional[str] = None,
     ) -> None:
         """Add a tool to the server."""
-        func_name = name or func.__name__
-        func_doc = description or func.__doc__ or ""
-        is_async = inspect.iscoroutinefunction(func)
-
-        # Get type hints for parameters
-        hints = get_type_hints(func)
-        if "return" in hints:
-            del hints["return"]
-
-        # Check for Pydantic model parameter
-        if len(hints) == 1 and issubclass(next(iter(hints.values())), BaseModel):
-            model = next(iter(hints.values()))
-            schema = model.model_json_schema()
-            pydantic_model = model
-        else:
-            # Create parameter schema from type hints
-            fields = {}
-            sig = inspect.signature(func)
-            for param_name, param in sig.parameters.items():
-                param_type = hints.get(param_name, Any)
-                default = (
-                    ... if param.default is inspect.Parameter.empty else param.default
-                )
-                fields[param_name] = (param_type, default)
-
-            model = create_model(f"{func_name}Args", **fields)
-            schema = model.model_json_schema()
-            pydantic_model = model
-
-        self._tools[func_name] = Tool(
-            func=func,
-            name=func_name,
-            description=func_doc,
-            input_schema=schema,
-            is_async=is_async,
-            pydantic_model=pydantic_model,
-        )
+        tool = Tool.from_function(func, name=name, description=description)
+        self._tools[tool.name] = tool
 
     async def call_tool(self, name: str, arguments: dict) -> Any:
         """Call a tool by name with arguments."""
         tool = self.get_tool(name)
         if not tool:
             raise ToolError(f"Unknown tool: {name}")
-
-        try:
-            # Validate arguments using schema
-            if tool.pydantic_model:
-                validated_args = tool.pydantic_model(**arguments)
-                args_dict = validated_args.model_dump()
-            else:
-                args_dict = arguments
-
-            # Call function with proper async handling
-            if tool.is_async:
-                return await tool.func(**args_dict)
-            return tool.func(**args_dict)
-        except Exception as e:
-            raise ToolError(f"Error executing tool {name}: {e}") from e
+        return await tool.run(arguments)

tests/test_tools.py

@@ -0,0 +1,119 @@
+"""Test tool registration and execution."""
+
+import pytest
+from pydantic import BaseModel
+
+from fastmcp.exceptions import ToolError
+from fastmcp.tools import ToolManager
+
+
+class TestAddTools:
+    def test_basic_function(self):
+        """Test registering and running a basic function."""
+
+        def add(a: int, b: int) -> int:
+            """Add two numbers."""
+            return a + b
+
+        manager = ToolManager()
+        manager.add_tool(add)
+
+        tool = manager.get_tool("add")
+        assert tool is not None
+        assert tool.name == "add"
+        assert tool.description == "Add two numbers."
+        assert tool.is_async is False
+        assert tool.parameters["properties"]["a"]["type"] == "integer"
+        assert tool.parameters["properties"]["b"]["type"] == "integer"
+
+    async def test_async_function(self):
+        """Test registering and running an async function."""
+
+        async def fetch_data(url: str) -> str:
+            """Fetch data from URL."""
+            return f"Data from {url}"
+
+        manager = ToolManager()
+        manager.add_tool(fetch_data)
+
+        tool = manager.get_tool("fetch_data")
+        assert tool is not None
+        assert tool.name == "fetch_data"
+        assert tool.description == "Fetch data from URL."
+        assert tool.is_async is True
+        assert tool.parameters["properties"]["url"]["type"] == "string"
+
+    def test_pydantic_model_function(self):
+        """Test registering a function that takes a Pydantic model."""
+
+        class UserInput(BaseModel):
+            name: str
+            age: int
+
+        def create_user(user: UserInput, flag: bool) -> dict:
+            """Create a new user."""
+            return {"id": 1, **user.model_dump()}
+
+        manager = ToolManager()
+        manager.add_tool(create_user)
+
+        tool = manager.get_tool("create_user")
+        assert tool is not None
+        assert tool.name == "create_user"
+        assert tool.description == "Create a new user."
+        assert tool.is_async is False
+        assert "name" in tool.parameters["$defs"]["UserInput"]["properties"]
+        assert "age" in tool.parameters["$defs"]["UserInput"]["properties"]
+        assert "flag" in tool.parameters["properties"]
+
+    def test_add_invalid_tool(self):
+        manager = ToolManager()
+        with pytest.raises(AttributeError):
+            manager.add_tool(1)
+
+
+class TestCallTools:
+    async def test_call_tool(self):
+        def add(a: int, b: int) -> int:
+            """Add two numbers."""
+            return a + b
+
+        manager = ToolManager()
+        manager.add_tool(add)
+        result = await manager.call_tool("add", {"a": 1, "b": 2})
+        assert result == 3
+
+    async def test_call_async_tool(self):
+        async def double(n: int) -> int:
+            """Double a number."""
+            return n * 2
+
+        manager = ToolManager()
+        manager.add_tool(double)
+        result = await manager.call_tool("double", {"n": 5})
+        assert result == 10
+
+    async def test_call_tool_with_default_args(self):
+        def add(a: int, b: int = 1) -> int:
+            """Add two numbers."""
+            return a + b
+
+        manager = ToolManager()
+        manager.add_tool(add)
+        result = await manager.call_tool("add", {"a": 1})
+        assert result == 2
+
+    async def test_call_tool_with_missing_args(self):
+        def add(a: int, b: int) -> int:
+            """Add two numbers."""
+            return a + b
+
+        manager = ToolManager()
+        manager.add_tool(add)
+        with pytest.raises(ToolError):
+            await manager.call_tool("add", {"a": 1})
+
+    async def test_call_unknown_tool(self):
+        manager = ToolManager()
+        with pytest.raises(ToolError):
+            await manager.call_tool("unknown", {"a": 1})