From 5c0f05511bf8dfa5539a5322fcc817a3704b11aa Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Thu, 4 Dec 2025 22:19:27 -0500 Subject: [PATCH] Remove @sampling_tool decorator - pass functions directly to sample() Functions passed to ctx.sample(tools=[...]) are now auto-converted via SamplingTool.from_function(). Users can still use that method directly for custom name/description overrides. --- docs/servers/sampling.mdx | 29 +++++---- examples/advanced_sampling/tool_use.py | 5 +- src/fastmcp/server/context.py | 16 +++-- src/fastmcp/server/sampling/__init__.py | 3 +- src/fastmcp/server/sampling/sampling_tool.py | 68 ++++---------------- tests/client/test_sampling.py | 30 +++------ tests/server/sampling/test_sampling_tool.py | 54 ++++------------ 7 files changed, 60 insertions(+), 145 deletions(-) diff --git a/docs/servers/sampling.mdx b/docs/servers/sampling.mdx index c52bf7ea5..f4573f0ff 100644 --- a/docs/servers/sampling.mdx +++ b/docs/servers/sampling.mdx @@ -248,24 +248,37 @@ Sampling with tools enables agentic workflows where the LLM can request tool cal ### Creating Sampling Tools -Use the `@sampling_tool` decorator to create tools for use during sampling: +Define regular Python functions and pass them directly to `ctx.sample()`: ```python from fastmcp import FastMCP, Context -from fastmcp.server.sampling import sampling_tool -@sampling_tool def search(query: str) -> str: """Search the web for information.""" return f"Search results for: {query}" -@sampling_tool def get_current_time() -> str: """Get the current time.""" from datetime import datetime return datetime.now().strftime("%H:%M:%S") ``` + +To create a tool with a custom name or description, use `SamplingTool.from_function()`: + +```python +from fastmcp.server.sampling import SamplingTool + +tool = SamplingTool.from_function( + my_search_function, + name="web_search", + description="Search the web for information" +) + +result = await ctx.sample(messages="...", tools=[tool]) +``` + + You can also pass existing FastMCP tools directly to `ctx.sample()`: ```python @@ -292,14 +305,11 @@ Pass tools to `ctx.sample()` to enable agentic tool use. FastMCP automatically h ```python from fastmcp import FastMCP, Context -from fastmcp.server.sampling import sampling_tool -@sampling_tool def search(query: str) -> str: """Search the web.""" return f"Results for: {query}" -@sampling_tool def get_time() -> str: """Get the current time.""" from datetime import datetime @@ -370,19 +380,16 @@ Combine `result_type` with tools for agentic workflows that return structured da ```python from pydantic import BaseModel from fastmcp import FastMCP, Context -from fastmcp.server.sampling import sampling_tool class ResearchResult(BaseModel): summary: str sources: list[str] confidence: float -@sampling_tool def search(query: str) -> str: """Search for information.""" return f"Found information about: {query}" -@sampling_tool def fetch_url(url: str) -> str: """Fetch content from a URL.""" return f"Content from {url}" @@ -538,12 +545,10 @@ The fallback handler fully supports sampling with tools: ```python from fastmcp import FastMCP, Context -from fastmcp.server.sampling import sampling_tool from fastmcp.experimental.sampling.handlers.openai import OpenAISamplingHandler from openai import OpenAI import os -@sampling_tool def search(query: str) -> str: """Search for information.""" return f"Results for: {query}" diff --git a/examples/advanced_sampling/tool_use.py b/examples/advanced_sampling/tool_use.py index 8629f0fb3..a3302e47a 100644 --- a/examples/advanced_sampling/tool_use.py +++ b/examples/advanced_sampling/tool_use.py @@ -19,11 +19,9 @@ from pydantic import BaseModel from fastmcp import Client, Context, FastMCP from fastmcp.experimental.sampling.handlers.openai import OpenAISamplingHandler -from fastmcp.server.sampling import sampling_tool -# Define sampling tools (available to the LLM during sampling) -@sampling_tool +# Define tools (available to the LLM during sampling) def search_web(query: str) -> str: """Search the web for information.""" # Simulated search results @@ -38,7 +36,6 @@ def search_web(query: str) -> str: return f"No results found for: {query}" -@sampling_tool def get_word_count(text: str) -> str: """Count words in text.""" return str(len(text.split())) diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index 7650b00b6..2c8de0e22 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -5,7 +5,7 @@ import inspect import json import logging import weakref -from collections.abc import Generator, Mapping, Sequence +from collections.abc import Callable, Generator, Mapping, Sequence from contextlib import contextmanager from contextvars import ContextVar, Token from dataclasses import dataclass @@ -519,7 +519,7 @@ class Context: temperature: float | None = None, max_tokens: int | None = None, model_preferences: ModelPreferences | str | list[str] | None = None, - tools: Sequence[SamplingTool | FastMCPTool] | None = None, + tools: Sequence[SamplingTool | FastMCPTool | Callable[..., Any]] | None = None, tool_choice: ToolChoiceOption | None = None, max_iterations: int = 10, result_type: type[ResultT], @@ -536,7 +536,7 @@ class Context: temperature: float | None = None, max_tokens: int | None = None, model_preferences: ModelPreferences | str | list[str] | None = None, - tools: Sequence[SamplingTool | FastMCPTool] | None = None, + tools: Sequence[SamplingTool | FastMCPTool | Callable[..., Any]] | None = None, tool_choice: ToolChoiceOption | None = None, max_iterations: int = 10, result_type: None = None, @@ -552,7 +552,7 @@ class Context: temperature: float | None = None, max_tokens: int | None = None, model_preferences: ModelPreferences | str | list[str] | None = None, - tools: Sequence[SamplingTool | FastMCPTool] | None = None, + tools: Sequence[SamplingTool | FastMCPTool | Callable[..., Any]] | None = None, tool_choice: ToolChoiceOption | None = None, max_iterations: int = 10, result_type: type[ResultT] | None = None, @@ -581,8 +581,8 @@ class Context: temperature: Optional sampling temperature. max_tokens: Maximum tokens to generate. Defaults to 512. model_preferences: Optional model preferences. - tools: Optional list of tools the LLM can use. Accepts both - SamplingTools and FastMCP Tools (which are auto-converted). + tools: Optional list of tools the LLM can use. Accepts plain + functions, SamplingTools, or FastMCP Tools (all are auto-converted). When provided, the method automatically handles tool execution and returns the final response after all tool calls complete. tool_choice: Optional control over tool usage behavior. Only valid @@ -626,9 +626,11 @@ class Context: sampling_tools.append(t) elif isinstance(t, FastMCPTool): sampling_tools.append(SamplingTool.from_mcp_tool(t)) + elif callable(t): + sampling_tools.append(SamplingTool.from_function(t)) else: raise TypeError( - f"Expected SamplingTool or FastMCP Tool, got {type(t)}" + f"Expected SamplingTool, FastMCP Tool, or callable, got {type(t)}" ) # Create synthetic final_response tool for structured output diff --git a/src/fastmcp/server/sampling/__init__.py b/src/fastmcp/server/sampling/__init__.py index f5238d230..6485e658e 100644 --- a/src/fastmcp/server/sampling/__init__.py +++ b/src/fastmcp/server/sampling/__init__.py @@ -1,10 +1,9 @@ """Sampling module for FastMCP servers.""" from fastmcp.server.sampling.handler import ServerSamplingHandler -from fastmcp.server.sampling.sampling_tool import SamplingTool, sampling_tool +from fastmcp.server.sampling.sampling_tool import SamplingTool __all__ = [ "SamplingTool", "ServerSamplingHandler", - "sampling_tool", ] diff --git a/src/fastmcp/server/sampling/sampling_tool.py b/src/fastmcp/server/sampling/sampling_tool.py index 85ce62427..fad0b706b 100644 --- a/src/fastmcp/server/sampling/sampling_tool.py +++ b/src/fastmcp/server/sampling/sampling_tool.py @@ -4,7 +4,7 @@ from __future__ import annotations import inspect from collections.abc import Callable -from typing import TYPE_CHECKING, Any, overload +from typing import TYPE_CHECKING, Any from mcp.types import Tool as SDKTool from pydantic import BaseModel, ConfigDict @@ -22,21 +22,24 @@ class SamplingTool(BaseModel): an executor function, enabling servers to execute agentic workflows where the LLM can request tool calls during sampling. - Create a SamplingTool using the @sampling_tool decorator or class methods: + In most cases, pass functions directly to ctx.sample(): - @sampling_tool def search(query: str) -> str: '''Search the web.''' return web_search(query) - # Or from an existing FastMCP Tool - sampling_tool = SamplingTool.from_mcp_tool(existing_tool) - - # Then use in sampling result = await context.sample( messages="Find info about Python", - tools=[search], + tools=[search], # Plain functions work directly ) + + Create a SamplingTool explicitly when you need custom name/description: + + tool = SamplingTool.from_function(search, name="web_search") + + Or from an existing FastMCP Tool: + + sampling_tool = SamplingTool.from_mcp_tool(existing_tool) """ name: str @@ -136,52 +139,3 @@ class SamplingTool(BaseModel): parameters=parsed.input_schema, fn=parsed.fn, ) - - -@overload -def sampling_tool(fn: Callable[..., Any]) -> SamplingTool: ... - - -@overload -def sampling_tool( - *, - name: str | None = None, - description: str | None = None, -) -> Callable[[Callable[..., Any]], SamplingTool]: ... - - -def sampling_tool( - fn: Callable[..., Any] | None = None, - *, - name: str | None = None, - description: str | None = None, -) -> SamplingTool | Callable[[Callable[..., Any]], SamplingTool]: - """Decorator to create a SamplingTool from a function. - - Can be used with or without arguments: - - @sampling_tool - def search(query: str) -> str: - '''Search the web.''' - return web_search(query) - - @sampling_tool(name="web_search", description="Search the internet") - def search(query: str) -> str: - return web_search(query) - - Args: - fn: The function to wrap (when used without parentheses). - name: Optional name override for the tool. - description: Optional description override for the tool. - - Returns: - A SamplingTool if called directly on a function, or a decorator - if called with arguments. - """ - if fn is not None: - return SamplingTool.from_function(fn) - - def decorator(fn: Callable[..., Any]) -> SamplingTool: - return SamplingTool.from_function(fn, name=name, description=description) - - return decorator diff --git a/tests/client/test_sampling.py b/tests/client/test_sampling.py index 6c0bf43e6..915c1ce90 100644 --- a/tests/client/test_sampling.py +++ b/tests/client/test_sampling.py @@ -8,7 +8,7 @@ from pydantic_core import to_json from fastmcp import Client, Context, FastMCP from fastmcp.client.sampling import RequestContext, SamplingMessage, SamplingParams -from fastmcp.server.sampling import SamplingTool, sampling_tool +from fastmcp.server.sampling import SamplingTool from fastmcp.utilities.types import Image @@ -177,7 +177,6 @@ class TestSamplingWithTools: server = FastMCP() - @sampling_tool def search(query: str) -> str: """Search the web.""" return f"Results for: {query}" @@ -214,7 +213,6 @@ class TestSamplingWithTools: mcp = FastMCP(sampling_handler=invalid_handler) - @sampling_tool def search(query: str) -> str: """Search the web.""" return f"Results for: {query}" @@ -238,36 +236,36 @@ class TestSamplingWithTools: def test_sampling_tool_schema(self): """Test that SamplingTool generates correct schema.""" - @sampling_tool def search(query: str, limit: int = 10) -> str: """Search the web for results.""" return f"Results for: {query}" - assert search.name == "search" - assert search.description == "Search the web for results." - assert "query" in search.parameters.get("properties", {}) - assert "limit" in search.parameters.get("properties", {}) + tool = SamplingTool.from_function(search) + assert tool.name == "search" + assert tool.description == "Search the web for results." + assert "query" in tool.parameters.get("properties", {}) + assert "limit" in tool.parameters.get("properties", {}) async def test_sampling_tool_run(self): """Test that SamplingTool.run() executes correctly.""" - @sampling_tool def add(a: int, b: int) -> int: """Add two numbers.""" return a + b - result = await add.run({"a": 5, "b": 3}) + tool = SamplingTool.from_function(add) + result = await tool.run({"a": 5, "b": 3}) assert result == 8 async def test_sampling_tool_run_async(self): """Test that SamplingTool.run() works with async functions.""" - @sampling_tool async def async_multiply(a: int, b: int) -> int: """Multiply two numbers.""" return a * b - result = await async_multiply.run({"a": 4, "b": 7}) + tool = SamplingTool.from_function(async_multiply) + result = await tool.run({"a": 4, "b": 7}) assert result == 28 def test_sampling_tool_from_mcp_tool(self): @@ -307,7 +305,6 @@ class TestAutomaticToolLoop: call_count = 0 tool_was_called = False - @sampling_tool def get_weather(city: str) -> str: """Get weather for a city.""" nonlocal tool_was_called @@ -370,13 +367,11 @@ class TestAutomaticToolLoop: executed_tools: list[str] = [] - @sampling_tool def tool_a(x: int) -> int: """Tool A.""" executed_tools.append(f"tool_a({x})") return x * 2 - @sampling_tool def tool_b(y: int) -> int: """Tool B.""" executed_tools.append(f"tool_b({y})") @@ -436,7 +431,6 @@ class TestAutomaticToolLoop: call_count = 0 received_tool_choices: list = [] - @sampling_tool def looping_tool() -> str: """A tool that always gets called again.""" return "keep going" @@ -498,7 +492,6 @@ class TestAutomaticToolLoop: ToolUseContent, ) - @sampling_tool def known_tool() -> str: """A known tool.""" return "known result" @@ -573,7 +566,6 @@ class TestAutomaticToolLoop: ToolUseContent, ) - @sampling_tool def failing_tool() -> str: """A tool that raises an exception.""" raise ValueError("Tool failed intentionally") @@ -649,7 +641,6 @@ class TestAutomaticToolLoop: received_tool_choices: list = [] - @sampling_tool def my_tool() -> str: """A tool.""" return "result" @@ -763,7 +754,6 @@ class TestSamplingResultType: summary: str sources: list[str] - @sampling_tool def search(query: str) -> str: """Search for information.""" return f"Found info about: {query}" diff --git a/tests/server/sampling/test_sampling_tool.py b/tests/server/sampling/test_sampling_tool.py index 19808ecd2..1025084a6 100644 --- a/tests/server/sampling/test_sampling_tool.py +++ b/tests/server/sampling/test_sampling_tool.py @@ -2,7 +2,7 @@ import pytest -from fastmcp.server.sampling import SamplingTool, sampling_tool +from fastmcp.server.sampling import SamplingTool from fastmcp.tools.tool import Tool @@ -94,75 +94,43 @@ class TestSamplingToolFromMCPTool: SamplingTool.from_mcp_tool(tool) -class TestSamplingToolDecorator: - """Tests for the @sampling_tool decorator.""" - - def test_decorator_without_args(self): - @sampling_tool - def search(query: str) -> str: - """Search the web.""" - return f"Results for: {query}" - - assert isinstance(search, SamplingTool) - assert search.name == "search" - assert search.description == "Search the web." - - def test_decorator_with_args(self): - @sampling_tool(name="web_search", description="Custom description") - def search(query: str) -> str: - return f"Results for: {query}" - - assert isinstance(search, SamplingTool) - assert search.name == "web_search" - assert search.description == "Custom description" - - def test_decorator_with_partial_args(self): - @sampling_tool(name="custom_name") - def search(query: str) -> str: - """Original docstring.""" - return f"Results for: {query}" - - assert search.name == "custom_name" - assert search.description == "Original docstring." - - class TestSamplingToolRun: """Tests for SamplingTool.run().""" async def test_run_sync_function(self): - @sampling_tool def add(a: int, b: int) -> int: """Add two numbers.""" return a + b - result = await add.run({"a": 2, "b": 3}) + tool = SamplingTool.from_function(add) + result = await tool.run({"a": 2, "b": 3}) assert result == 5 async def test_run_async_function(self): - @sampling_tool async def async_add(a: int, b: int) -> int: """Add two numbers asynchronously.""" return a + b - result = await async_add.run({"a": 2, "b": 3}) + tool = SamplingTool.from_function(async_add) + result = await tool.run({"a": 2, "b": 3}) assert result == 5 async def test_run_with_no_arguments(self): - @sampling_tool def get_value() -> str: """Return a fixed value.""" return "hello" - result = await get_value.run() + tool = SamplingTool.from_function(get_value) + result = await tool.run() assert result == "hello" async def test_run_with_none_arguments(self): - @sampling_tool def get_value() -> str: """Return a fixed value.""" return "hello" - result = await get_value.run(None) + tool = SamplingTool.from_function(get_value) + result = await tool.run(None) assert result == "hello" @@ -170,12 +138,12 @@ class TestSamplingToolSDKConversion: """Tests for SamplingTool._to_sdk_tool() internal method.""" def test_to_sdk_tool(self): - @sampling_tool def search(query: str) -> str: """Search the web.""" return f"Results for: {query}" - sdk_tool = search._to_sdk_tool() + tool = SamplingTool.from_function(search) + sdk_tool = tool._to_sdk_tool() assert sdk_tool.name == "search" assert sdk_tool.description == "Search the web."