mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-22 21:44:18 +02:00
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:
parent
ead4730423
commit
5c0f05511b
7 changed files with 60 additions and 145 deletions
|
|
@ -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}"
|
||||
|
|
|
|||
|
|
@ -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()))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue