Spaces:
Running
Running
File size: 4,688 Bytes
51336fb 09d1b8e 51336fb 09d1b8e 51336fb 09d1b8e 09438a8 09d1b8e 51336fb 09d1b8e 09438a8 09d1b8e 51336fb 09d1b8e 09438a8 09d1b8e 51336fb 09d1b8e aa46c09 09d1b8e aa46c09 09d1b8e aa46c09 51336fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | import json
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.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 messages[0].content.type == "text"
assert messages[0].content.text == "Hello!"
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_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,
},
},
{
"role": "user",
"content": {
"type": "image",
"data": "YWJjMTIz",
"mimeType": "image/png",
"annotations": None,
"_meta": None,
},
},
]
|