fastmcp/tests/client/test_sampling.py
Jeremiah Lowin 5c0f05511b 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.
2025-12-04 22:19:27 -05:00

938 lines
34 KiB
Python

import json
from typing import cast
from unittest.mock import AsyncMock
import pytest
from mcp.types import TextContent
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
from fastmcp.utilities.types import Image
@pytest.fixture
def fastmcp_server():
mcp = FastMCP()
@mcp.tool
async def simple_sample(message: str, context: Context) -> str:
result = await context.sample("Hello, world!")
return result.text # type: ignore[attr-defined]
@mcp.tool
async def sample_with_system_prompt(message: str, context: Context) -> str:
result = await context.sample("Hello, world!", system_prompt="You love FastMCP")
return result.text # type: ignore[attr-defined]
@mcp.tool
async def sample_with_messages(message: str, context: Context) -> str:
result = await context.sample(
[
"Hello!",
SamplingMessage(
content=TextContent(
type="text", text="How can I assist you today?"
),
role="assistant",
),
]
)
return result.text # type: ignore[attr-defined]
@mcp.tool
async def sample_with_image(image_bytes: bytes, context: Context) -> str:
image = Image(data=image_bytes)
result = await context.sample(
[
SamplingMessage(
content=TextContent(type="text", text="What's in this image?"),
role="user",
),
SamplingMessage(
content=image.to_image_content(),
role="user",
),
]
)
return result.text # type: ignore[attr-defined]
return mcp
async def test_simple_sampling(fastmcp_server: FastMCP):
def sampling_handler(
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
) -> str:
return "This is the sample message!"
async with Client(fastmcp_server, sampling_handler=sampling_handler) as client:
result = await client.call_tool("simple_sample", {"message": "Hello, world!"})
assert result.data == "This is the sample message!"
async def test_sampling_with_system_prompt(fastmcp_server: FastMCP):
def sampling_handler(
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
) -> str:
assert params.systemPrompt is not None
return params.systemPrompt
async with Client(fastmcp_server, sampling_handler=sampling_handler) as client:
result = await client.call_tool(
"sample_with_system_prompt", {"message": "Hello, world!"}
)
assert result.data == "You love FastMCP"
async def test_sampling_with_messages(fastmcp_server: FastMCP):
def sampling_handler(
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
) -> str:
assert len(messages) == 2
assert isinstance(messages[0].content, TextContent)
assert messages[0].content.type == "text"
assert messages[0].content.text == "Hello!"
assert isinstance(messages[1].content, TextContent)
assert messages[1].content.type == "text"
assert messages[1].content.text == "How can I assist you today?"
return "I need to think."
async with Client(fastmcp_server, sampling_handler=sampling_handler) as client:
result = await client.call_tool(
"sample_with_messages", {"message": "Hello, world!"}
)
assert result.data == "I need to think."
async def test_sampling_with_fallback(fastmcp_server: FastMCP):
openai_sampling_handler = AsyncMock(return_value="But I need to think")
fastmcp_server = FastMCP(
sampling_handler=openai_sampling_handler,
)
@fastmcp_server.tool
async def sample_with_fallback(context: Context) -> str:
sampling_result = await context.sample("Do not think.")
return cast(TextContent, sampling_result).text
client = Client(fastmcp_server)
async with client:
call_tool_result = await client.call_tool("sample_with_fallback")
assert call_tool_result.data == "But I need to think"
async def test_sampling_with_image(fastmcp_server: FastMCP):
def sampling_handler(
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
) -> str:
assert len(messages) == 2
return to_json(messages).decode()
async with Client(fastmcp_server, sampling_handler=sampling_handler) as client:
image_bytes = b"abc123"
result = await client.call_tool(
"sample_with_image", {"image_bytes": image_bytes}
)
assert json.loads(result.data) == [
{
"role": "user",
"content": {
"type": "text",
"text": "What's in this image?",
"annotations": None,
"_meta": None,
},
"_meta": None,
},
{
"role": "user",
"content": {
"type": "image",
"data": "YWJjMTIz",
"mimeType": "image/png",
"annotations": None,
"_meta": None,
},
"_meta": None,
},
]
class TestSamplingWithTools:
"""Tests for sampling with tools functionality."""
async def test_sampling_with_tools_requires_capability(self):
"""Test that sampling with tools raises error when client lacks capability."""
import mcp.types as mcp_types
from fastmcp.exceptions import ToolError
server = FastMCP()
def search(query: str) -> str:
"""Search the web."""
return f"Results for: {query}"
@server.tool
async def sample_with_tool(context: Context) -> str:
# This should fail because the client doesn't advertise tools capability
result = await context.sample(
messages="Search for Python tutorials",
tools=[search],
)
return str(result)
def sampling_handler(
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
) -> str:
return "Response"
# Explicitly disable tools capability by passing SamplingCapability without tools
async with Client(
server,
sampling_handler=sampling_handler,
sampling_capabilities=mcp_types.SamplingCapability(), # No tools
) as client:
with pytest.raises(ToolError, match="sampling.tools capability"):
await client.call_tool("sample_with_tool", {})
async def test_sampling_with_tools_fallback_handler_must_return_correct_type(self):
"""Test that fallback handler must return CreateMessageResultWithTools when tools provided."""
from fastmcp.exceptions import ToolError
# This handler returns a string, which is invalid when tools are provided
invalid_handler = AsyncMock(return_value="Fallback response")
mcp = FastMCP(sampling_handler=invalid_handler)
def search(query: str) -> str:
"""Search the web."""
return f"Results for: {query}"
@mcp.tool
async def sample_with_tool(context: Context) -> str:
result = await context.sample(
messages="Search for Python tutorials",
tools=[search],
)
return str(result)
# Client without sampling handler - will use server's fallback
async with Client(mcp) as client:
# Handler returns string but tools were provided - should error
with pytest.raises(
ToolError, match="must return CreateMessageResultWithTools"
):
await client.call_tool("sample_with_tool", {})
def test_sampling_tool_schema(self):
"""Test that SamplingTool generates correct schema."""
def search(query: str, limit: int = 10) -> str:
"""Search the web for results."""
return f"Results for: {query}"
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."""
def add(a: int, b: int) -> int:
"""Add two numbers."""
return a + b
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."""
async def async_multiply(a: int, b: int) -> int:
"""Multiply two numbers."""
return a * b
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):
"""Test creating SamplingTool from FastMCP Tool."""
from fastmcp.tools.tool import Tool
def original_fn(x: int, y: int) -> int:
"""Add x and y."""
return x + y
mcp_tool = Tool.from_function(original_fn)
sampling = SamplingTool.from_mcp_tool(mcp_tool)
assert sampling.name == "original_fn"
assert sampling.description == "Add x and y."
assert sampling.fn is mcp_tool.fn
def test_tool_choice_parameter(self):
"""Test that tool_choice parameter accepts string literals."""
from fastmcp.server.context import ToolChoiceOption
# Verify ToolChoiceOption type accepts the valid string values
choices: list[ToolChoiceOption] = ["auto", "required", "none"]
assert len(choices) == 3
assert "auto" in choices
assert "required" in choices
assert "none" in choices
class TestAutomaticToolLoop:
"""Tests for automatic tool execution loop in ctx.sample()."""
async def test_automatic_tool_loop_executes_tools(self):
"""Test that ctx.sample() automatically executes tool calls."""
from mcp.types import CreateMessageResultWithTools, ToolUseContent
call_count = 0
tool_was_called = False
def get_weather(city: str) -> str:
"""Get weather for a city."""
nonlocal tool_was_called
tool_was_called = True
return f"Weather in {city}: sunny, 72°F"
def sampling_handler(
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
) -> CreateMessageResultWithTools:
nonlocal call_count
call_count += 1
if call_count == 1:
# First call: return tool use
return CreateMessageResultWithTools(
role="assistant",
content=[
ToolUseContent(
type="tool_use",
id="call_1",
name="get_weather",
input={"city": "Seattle"},
)
],
model="test-model",
stopReason="toolUse",
)
else:
# Second call: return final response
return CreateMessageResultWithTools(
role="assistant",
content=[TextContent(type="text", text="The weather is sunny!")],
model="test-model",
stopReason="endTurn",
)
mcp = FastMCP(sampling_handler=sampling_handler)
@mcp.tool
async def weather_assistant(question: str, context: Context) -> str:
result = await context.sample(
messages=question,
tools=[get_weather],
)
# Get text from SamplingResult
return result.text or ""
async with Client(mcp) as client:
result = await client.call_tool(
"weather_assistant", {"question": "What's the weather?"}
)
assert tool_was_called
assert call_count == 2
assert result.data == "The weather is sunny!"
async def test_automatic_tool_loop_multiple_tools(self):
"""Test that multiple tool calls in one response are all executed."""
from mcp.types import CreateMessageResultWithTools, ToolUseContent
executed_tools: list[str] = []
def tool_a(x: int) -> int:
"""Tool A."""
executed_tools.append(f"tool_a({x})")
return x * 2
def tool_b(y: int) -> int:
"""Tool B."""
executed_tools.append(f"tool_b({y})")
return y + 10
call_count = 0
def sampling_handler(
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
) -> CreateMessageResultWithTools:
nonlocal call_count
call_count += 1
if call_count == 1:
# Return multiple tool calls
return CreateMessageResultWithTools(
role="assistant",
content=[
ToolUseContent(
type="tool_use", id="call_a", name="tool_a", input={"x": 5}
),
ToolUseContent(
type="tool_use", id="call_b", name="tool_b", input={"y": 3}
),
],
model="test-model",
stopReason="toolUse",
)
else:
return CreateMessageResultWithTools(
role="assistant",
content=[TextContent(type="text", text="Done!")],
model="test-model",
stopReason="endTurn",
)
mcp = FastMCP(sampling_handler=sampling_handler)
@mcp.tool
async def multi_tool(context: Context) -> str:
result = await context.sample(messages="Run tools", tools=[tool_a, tool_b])
return result.text or ""
async with Client(mcp) as client:
result = await client.call_tool("multi_tool", {})
assert executed_tools == ["tool_a(5)", "tool_b(3)"]
assert result.data == "Done!"
async def test_automatic_tool_loop_max_iterations(self):
"""Test that max_iterations prevents infinite loops.
On the last iteration, tool_choice is set to 'none' to force text response.
"""
from mcp.types import CreateMessageResultWithTools, ToolUseContent
call_count = 0
received_tool_choices: list = []
def looping_tool() -> str:
"""A tool that always gets called again."""
return "keep going"
def sampling_handler(
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
) -> CreateMessageResultWithTools:
nonlocal call_count
call_count += 1
received_tool_choices.append(params.toolChoice)
# On last iteration (when tool_choice=none), return text
if params.toolChoice and params.toolChoice.mode == "none":
return CreateMessageResultWithTools(
role="assistant",
content=[TextContent(type="text", text="Forced to stop")],
model="test-model",
stopReason="endTurn",
)
# Otherwise keep returning tool use
return CreateMessageResultWithTools(
role="assistant",
content=[
ToolUseContent(
type="tool_use",
id=f"call_{call_count}",
name="looping_tool",
input={},
)
],
model="test-model",
stopReason="toolUse",
)
mcp = FastMCP(sampling_handler=sampling_handler)
@mcp.tool
async def infinite_loop(context: Context) -> str:
result = await context.sample(
messages="Start", tools=[looping_tool], max_iterations=3
)
return result.text or "no text"
async with Client(mcp) as client:
result = await client.call_tool("infinite_loop", {})
# Should complete after 3 iterations with forced text response
assert call_count == 3
assert result.data == "Forced to stop"
# Last call should have tool_choice=none
assert received_tool_choices[-1].mode == "none"
async def test_automatic_tool_loop_handles_unknown_tool(self):
"""Test that unknown tool names result in error being passed to LLM."""
from mcp.types import (
CreateMessageResultWithTools,
ToolResultContent,
ToolUseContent,
)
def known_tool() -> str:
"""A known tool."""
return "known result"
messages_received: list[list[SamplingMessage]] = []
def sampling_handler(
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
) -> CreateMessageResultWithTools:
messages_received.append(list(messages))
if len(messages_received) == 1:
# Request unknown tool
return CreateMessageResultWithTools(
role="assistant",
content=[
ToolUseContent(
type="tool_use",
id="call_1",
name="unknown_tool",
input={},
)
],
model="test-model",
stopReason="toolUse",
)
else:
return CreateMessageResultWithTools(
role="assistant",
content=[TextContent(type="text", text="Handled error")],
model="test-model",
stopReason="endTurn",
)
mcp = FastMCP(sampling_handler=sampling_handler)
@mcp.tool
async def test_unknown(context: Context) -> str:
result = await context.sample(messages="Test", tools=[known_tool])
return result.text or ""
async with Client(mcp) as client:
result = await client.call_tool("test_unknown", {})
# Check that error was passed back in messages
assert len(messages_received) == 2
last_messages = messages_received[1]
# Find the tool result in list content
tool_result = None
for msg in last_messages:
# Tool results are now in a list
if isinstance(msg.content, list):
for item in msg.content:
if isinstance(item, ToolResultContent):
tool_result = item
break
elif isinstance(msg.content, ToolResultContent):
tool_result = msg.content
break
assert tool_result is not None
assert tool_result.isError is True
# Content is list of TextContent objects
error_text = tool_result.content[0].text # type: ignore[union-attr]
assert "Unknown tool" in error_text
assert result.data == "Handled error"
async def test_automatic_tool_loop_handles_tool_exception(self):
"""Test that tool exceptions are caught and passed to LLM as errors."""
from mcp.types import (
CreateMessageResultWithTools,
ToolResultContent,
ToolUseContent,
)
def failing_tool() -> str:
"""A tool that raises an exception."""
raise ValueError("Tool failed intentionally")
messages_received: list[list[SamplingMessage]] = []
def sampling_handler(
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
) -> CreateMessageResultWithTools:
messages_received.append(list(messages))
if len(messages_received) == 1:
return CreateMessageResultWithTools(
role="assistant",
content=[
ToolUseContent(
type="tool_use",
id="call_1",
name="failing_tool",
input={},
)
],
model="test-model",
stopReason="toolUse",
)
else:
return CreateMessageResultWithTools(
role="assistant",
content=[TextContent(type="text", text="Handled error")],
model="test-model",
stopReason="endTurn",
)
mcp = FastMCP(sampling_handler=sampling_handler)
@mcp.tool
async def test_exception(context: Context) -> str:
result = await context.sample(messages="Test", tools=[failing_tool])
return result.text or ""
async with Client(mcp) as client:
result = await client.call_tool("test_exception", {})
# Check that error was passed back
assert len(messages_received) == 2
last_messages = messages_received[1]
# Find the tool result in list content
tool_result = None
for msg in last_messages:
# Tool results are now in a list
if isinstance(msg.content, list):
for item in msg.content:
if isinstance(item, ToolResultContent):
tool_result = item
break
elif isinstance(msg.content, ToolResultContent):
tool_result = msg.content
break
assert tool_result is not None
assert tool_result.isError is True
# Content is list of TextContent objects
error_text = tool_result.content[0].text # type: ignore[union-attr]
assert "Tool failed intentionally" in error_text
assert result.data == "Handled error"
async def test_max_iterations_one_for_manual_loop(self):
"""Test that max_iterations=1 forces text on first call for manual loop building.
With max_iterations=1, tool_choice is set to 'none' on the first (and only) call,
forcing a text response so users can build their own loop using history.
"""
from mcp.types import CreateMessageResultWithTools, ToolUseContent
received_tool_choices: list = []
def my_tool() -> str:
"""A tool."""
return "result"
def sampling_handler(
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
) -> CreateMessageResultWithTools:
received_tool_choices.append(params.toolChoice)
# With tool_choice=none, return text response
if params.toolChoice and params.toolChoice.mode == "none":
return CreateMessageResultWithTools(
role="assistant",
content=[TextContent(type="text", text="Forced text response")],
model="test-model",
stopReason="endTurn",
)
# Otherwise return tool use (shouldn't happen with max_iterations=1)
return CreateMessageResultWithTools(
role="assistant",
content=[
ToolUseContent(
type="tool_use", id="call_1", name="my_tool", input={}
)
],
model="test-model",
stopReason="toolUse",
)
mcp = FastMCP(sampling_handler=sampling_handler)
@mcp.tool
async def single_iteration(context: Context) -> str:
result = await context.sample(
messages="Test", tools=[my_tool], max_iterations=1
)
return result.text or "no text"
async with Client(mcp) as client:
result = await client.call_tool("single_iteration", {})
# Should get forced text response
assert result.data == "Forced text response"
# First and only call should have tool_choice=none
assert len(received_tool_choices) == 1
assert received_tool_choices[0].mode == "none"
class TestSamplingResultType:
"""Tests for result_type parameter (structured output)."""
async def test_result_type_creates_final_response_tool(self):
"""Test that result_type creates a synthetic final_response tool."""
from mcp.types import CreateMessageResultWithTools, ToolUseContent
from pydantic import BaseModel
class MathResult(BaseModel):
answer: int
explanation: str
received_tools: list = []
def sampling_handler(
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
) -> CreateMessageResultWithTools:
received_tools.extend(params.tools or [])
# Return the final_response tool call
return CreateMessageResultWithTools(
role="assistant",
content=[
ToolUseContent(
type="tool_use",
id="call_1",
name="final_response",
input={"answer": 42, "explanation": "The meaning of life"},
)
],
model="test-model",
stopReason="toolUse",
)
mcp = FastMCP(sampling_handler=sampling_handler)
@mcp.tool
async def math_tool(context: Context) -> str:
result = await context.sample(
messages="What is 6 * 7?",
result_type=MathResult,
)
# result.result should be a MathResult object
return f"{result.result.answer}: {result.result.explanation}" # type: ignore[attr-defined]
async with Client(mcp) as client:
result = await client.call_tool("math_tool", {})
# Check that final_response tool was added
tool_names = [t.name for t in received_tools]
assert "final_response" in tool_names
# Check the result
assert result.data == "42: The meaning of life"
async def test_result_type_with_user_tools(self):
"""Test result_type works alongside user-provided tools."""
from mcp.types import CreateMessageResultWithTools, ToolUseContent
from pydantic import BaseModel
class SearchResult(BaseModel):
summary: str
sources: list[str]
def search(query: str) -> str:
"""Search for information."""
return f"Found info about: {query}"
call_count = 0
tool_was_called = False
def sampling_handler(
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
) -> CreateMessageResultWithTools:
nonlocal call_count, tool_was_called
call_count += 1
if call_count == 1:
# First call: use the search tool
return CreateMessageResultWithTools(
role="assistant",
content=[
ToolUseContent(
type="tool_use",
id="call_1",
name="search",
input={"query": "Python tutorials"},
)
],
model="test-model",
stopReason="toolUse",
)
else:
# Second call: call final_response
tool_was_called = True
return CreateMessageResultWithTools(
role="assistant",
content=[
ToolUseContent(
type="tool_use",
id="call_2",
name="final_response",
input={
"summary": "Python is great",
"sources": ["python.org", "docs.python.org"],
},
)
],
model="test-model",
stopReason="toolUse",
)
mcp = FastMCP(sampling_handler=sampling_handler)
@mcp.tool
async def research(context: Context) -> str:
result = await context.sample(
messages="Research Python",
tools=[search],
result_type=SearchResult,
)
return f"{result.result.summary} - {len(result.result.sources)} sources" # type: ignore[attr-defined]
async with Client(mcp) as client:
result = await client.call_tool("research", {})
assert tool_was_called
assert result.data == "Python is great - 2 sources"
async def test_result_type_validation_error_retries(self):
"""Test that validation errors are sent back to LLM for retry."""
from mcp.types import (
CreateMessageResultWithTools,
ToolResultContent,
ToolUseContent,
)
from pydantic import BaseModel
class StrictResult(BaseModel):
value: int # Must be an int
messages_received: list[list[SamplingMessage]] = []
def sampling_handler(
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
) -> CreateMessageResultWithTools:
messages_received.append(list(messages))
if len(messages_received) == 1:
# First call: invalid type
return CreateMessageResultWithTools(
role="assistant",
content=[
ToolUseContent(
type="tool_use",
id="call_1",
name="final_response",
input={"value": "not_an_int"}, # Wrong type
)
],
model="test-model",
stopReason="toolUse",
)
else:
# Second call: valid type after seeing error
return CreateMessageResultWithTools(
role="assistant",
content=[
ToolUseContent(
type="tool_use",
id="call_2",
name="final_response",
input={"value": 42}, # Correct type
)
],
model="test-model",
stopReason="toolUse",
)
mcp = FastMCP(sampling_handler=sampling_handler)
@mcp.tool
async def validate_tool(context: Context) -> str:
result = await context.sample(
messages="Give me a number",
result_type=StrictResult,
)
return str(result.result.value) # type: ignore[attr-defined]
async with Client(mcp) as client:
result = await client.call_tool("validate_tool", {})
# Should have retried after validation error
assert len(messages_received) == 2
# Check that error was passed back
last_messages = messages_received[1]
# Find the tool result in list content
tool_result = None
for msg in last_messages:
# Tool results are now in a list
if isinstance(msg.content, list):
for item in msg.content:
if isinstance(item, ToolResultContent):
tool_result = item
break
elif isinstance(msg.content, ToolResultContent):
tool_result = msg.content
break
assert tool_result is not None
assert tool_result.isError is True
error_text = tool_result.content[0].text # type: ignore[union-attr]
assert "Validation error" in error_text
# Final result should be correct
assert result.data == "42"
async def test_sampling_result_has_text_and_history(self):
"""Test that SamplingResult has text, result, and history attributes."""
from mcp.types import CreateMessageResultWithTools
def sampling_handler(
messages: list[SamplingMessage], params: SamplingParams, ctx: RequestContext
) -> CreateMessageResultWithTools:
return CreateMessageResultWithTools(
role="assistant",
content=[TextContent(type="text", text="Hello world")],
model="test-model",
stopReason="endTurn",
)
mcp = FastMCP(sampling_handler=sampling_handler)
@mcp.tool
async def check_result(context: Context) -> str:
result = await context.sample(messages="Say hello")
# Check all attributes exist
assert result.text == "Hello world"
assert result.result == "Hello world"
assert len(result.history) >= 1
return "ok"
async with Client(mcp) as client:
result = await client.call_tool("check_result", {})
assert result.data == "ok"