diff --git a/src/fastmcp/models.py b/src/fastmcp/models.py deleted file mode 100644 index dce1168e5..000000000 --- a/src/fastmcp/models.py +++ /dev/null @@ -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 diff --git a/src/fastmcp/server.py b/src/fastmcp/server.py index 333272378..b135eafa6 100644 --- a/src/fastmcp/server.py +++ b/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 ] diff --git a/src/fastmcp/tools.py b/src/fastmcp/tools.py index d8dc2f0f3..e37830fab 100644 --- a/src/fastmcp/tools.py +++ b/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) diff --git a/tests/test_tools.py b/tests/test_tools.py new file mode 100644 index 000000000..16f244231 --- /dev/null +++ b/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})