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.
This commit is contained in:
Jeremiah Lowin 2025-12-04 22:19:27 -05:00
commit 5c0f05511b
7 changed files with 60 additions and 145 deletions

View file

@ -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")
```
<Note>
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])
```
</Note>
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}"

View file

@ -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()))

View file

@ -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

View file

@ -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",
]

View file

@ -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

View file

@ -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}"

View file

@ -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."