Merge pull request #117 from jlowin/sampling

Add sampling and roots

README.md

@@ -42,10 +42,11 @@ FastMCP handles the complex protocol details and server management, letting you
 ## Key Features:
 
 *   **Simple Server Creation:** Build MCP servers with minimal boilerplate using intuitive decorators (`@tool`, `@resource`, `@prompt`).
-*   **Powerful Clients:** Programmatically interact with *any* MCP server, regardless of how it was built.
-*   **Proxy MCP Servers:** Create proxy servers to expose existing MCP servers or clients with modifications, or **convert between transport protocols** (e.g., expose a Stdio server via SSE for web access).
+*   **Proxy MCP Servers:** Create proxy servers to expose existing MCP servers or clients with modifications, or convert between transport protocols (e.g., expose a Stdio server via SSE for web access).
 *   **Compose MCP Servers:** Compose complex applications by mounting multiple FastMCP servers together.
 *   **API Generation:** Automatically create MCP servers from existing **OpenAPI specifications** or **FastAPI applications**.
+*   **Powerful Clients:** Programmatically interact with *any* MCP server, regardless of how it was built.
+*   **LLM Sampling:** Request completions from client LLMs directly within your MCP tools.
 *   **Pythonic Interface:** Designed with familiar Python patterns like decorators and type hints.
 *   **Context Injection:** Easily access core MCP capabilities like sampling, logging, and progress reporting within your functions.
 
@@ -80,10 +81,12 @@ FastMCP v1's core approach of using the `@tool`, `@resource`, `@prompt` decorato
   - [Context](#context)
   - [Images](#images)
 - [Advanced Features](#advanced-features)
-  - [MCP Client](#mcp-client)
   - [Proxy Servers](#proxy-servers)
   - [Composing MCP Servers](#composing-mcp-servers)
   - [OpenAPI \& FastAPI Generation](#openapi--fastapi-generation)
+  - [MCP Client](#mcp-client)
+    - [LLM Sampling](#llm-sampling)
+    - [Roots Access](#roots-access)
 - [Running Your Server](#running-your-server)
   - [Development Mode (Recommended for Building \& Testing)](#development-mode-recommended-for-building--testing)
   - [Claude Desktop Integration (For Regular Use)](#claude-desktop-integration-for-regular-use)
@@ -340,38 +343,6 @@ FastMCP handles the conversion to/from the base64-encoded format required by the
 
 Building on the core concepts, FastMCP v2 introduces powerful features for more complex scenarios:
 
-### MCP Client
-
-The client allows your Python code to interact with *any* MCP server, whether it's built with FastMCP, the official SDK, or another implementation. This is essential for testing, building meta-tools, or integrating MCP servers.
-
-```python
-import asyncio
-from fastmcp import Client
-from fastmcp.client.transports import StdioTransport # Example transport
-
-async def main():
-    # Connect to a server running via standard I/O
-    # Replace with the actual command to start your target server
-    client = Client(StdioTransport(command="python", args=["path/to/target_server.py"]))
-
-    async with client:
-        # Discover tools
-        tools_result = await client.list_tools()
-        print(f"Available Tools: {[t.name for t in tools_result.tools]}")
-
-        # Call a tool
-        add_result = await client.call_tool("add", {"a": 10, "b": 5})
-        print(f"Result of add(10, 5): {add_result.content[0].text}") # Output: 15
-
-        # Read a resource
-        greeting = await client.read_resource("greeting://Client")
-        print(f"Resource Content: {greeting.contents[0].text}") # Output: Hello, Client!
-
-if __name__ == "__main__":
-    asyncio.run(main())
-```
-
-The client supports various transports (`WSTransport`, `SSETransport`, `StdioTransport`, `FastMCPTransport`) and intelligently infers the correct one based on the connection information provided (URL, `FastMCP` instance, command arguments, etc.).
 
 ### Proxy Servers
 
@@ -509,6 +480,80 @@ if __name__ == "__main__":
     mcp_server.run()
 ```
 
+### MCP Client
+
+The `Client` class lets you interact with any MCP server (not just FastMCP ones) from Python code:
+
+```python
+from fastmcp import Client
+
+async with Client("path/to/server") as client:
+    # Call a tool
+    result = await client.call_tool("weather", {"location": "San Francisco"})
+    print(result)
+    
+    # Read a resource
+    res = await client.read_resource("db://users/123/profile")
+    print(res)
+```
+
+You can connect to servers using any supported transport protocol (Stdio, SSE, FastMCP, etc.). If you don't specify a transport, the `Client` class automatically attempts to detect an appropriate one from your connection string or server object.
+
+#### LLM Sampling
+
+Sampling is an MCP feature that allows a server to request a completion from the client LLM, enabling sophisticated use cases while maintaining security and privacy on the server.
+
+```python
+import marvin  # Or any other LLM client
+from fastmcp import Client, Context, FastMCP
+from fastmcp.client.sampling import RequestContext, SamplingMessage, SamplingParams
+
+# -- Create a server that requests LLM completions from the client
+
+mcp = FastMCP("Sampling Example")
+
+@mcp.tool()
+async def generate_poem(topic: str, context: Context) -> str:
+    """Generate a short poem about the given topic."""
+    response = await context.sample(
+        f"Write a short poem about {topic}",
+        system_prompt="You are a talented poet who writes concise, evocative verses."
+    )
+    return response.text
+
+# -- Create a client that handles the sampling requests
+
+async def sampling_handler(
+    messages: list[SamplingMessage],
+    params: SamplingParams,
+    ctx: RequestContext,
+) -> str:
+    # Use your preferred LLM client to generate completions
+    return await marvin.say_async(
+        message=[m.content.text for m in messages if m.content.type == "text"],
+        instructions=params.systemPrompt,
+    )
+
+# Connect them together
+async with Client(mcp, sampling_handler=sampling_handler) as client:
+    result = await client.call_tool("generate_poem", {"topic": "autumn leaves"})
+    print(result.content[0].text)
+```
+
+#### Roots Access
+
+FastMCP exposes the MCP roots functionality, allowing clients to specify which file system roots they can access. This creates a secure boundary for tools that need to work with files. Note that the server must account for client roots explicitly.
+
+```python
+from fastmcp import Client, RootsList
+
+# Specify file roots that the client can access
+roots = ["file:///path/to/allowed/directory"]
+
+async with Client(mcp_server, roots=roots) as client:
+    # Now tools in the MCP server can access files in the specified roots
+    await client.call_tool("process_file", {"filename": "data.csv"})
+```
 ## Running Your Server
 
 Choose the method that best suits your needs:
@@ -569,6 +614,7 @@ Explore the `examples/` directory for code samples demonstrating various feature
 *   `simple_echo.py`: Basic tool, resource, and prompt.
 *   `complex_inputs.py`: Using Pydantic models for tool inputs.
 *   `mount_example.py`: Mounting multiple FastMCP servers.
+*   `sampling.py`: Using LLM completions within your MCP server.
 *   `screenshot.py`: Tool returning an Image object.
 *   `text_me.py`: Tool interacting with an external API.
 *   `memory.py`: More complex example with database interaction.

examples/sampling.py

@@ -0,0 +1,52 @@
+"""
+Example of using sampling to request an LLM completion via Marvin
+"""
+
+import asyncio
+
+import marvin
+from mcp.types import TextContent
+
+from fastmcp import Client, Context, FastMCP
+from fastmcp.client.sampling import RequestContext, SamplingMessage, SamplingParams
+
+# -- Create a server that sends a sampling request to the LLM
+
+mcp = FastMCP("Sampling Example")
+
+
+@mcp.tool()
+async def example_tool(prompt: str, context: Context) -> str:
+    """Sample a completion from the LLM."""
+    response = await context.sample(
+        "What is your favorite programming language?",
+        system_prompt="You love languages named after snakes.",
+    )
+    assert isinstance(response, TextContent)
+    return response.text
+
+
+# -- Create a client that can handle the sampling request
+
+
+async def sampling_fn(
+    messages: list[SamplingMessage],
+    params: SamplingParams,
+    ctx: RequestContext,
+) -> str:
+    return await marvin.say_async(
+        message=[m.content.text for m in messages],
+        instructions=params.systemPrompt,
+    )
+
+
+async def run():
+    async with Client(mcp, sampling_handler=sampling_fn) as client:
+        result = await client.call_tool(
+            "example_tool", {"prompt": "What is the best programming language?"}
+        )
+        print(result)
+
+
+if __name__ == "__main__":
+    asyncio.run(run())

src/fastmcp/client/client.py

@@ -6,26 +6,22 @@ from typing import Any
 import mcp.types
 from mcp import ClientSession
 from mcp.client.session import (
-    ListRootsFnT,
     LoggingFnT,
     MessageHandlerFnT,
-    SamplingFnT,
 )
-from mcp.shared.context import LifespanContextT, RequestContext
 from pydantic import AnyUrl
 
+from fastmcp.client.roots import (
+    RootsHandler,
+    RootsList,
+    create_roots_callback,
+)
+from fastmcp.client.sampling import SamplingHandler, create_sampling_callback
 from fastmcp.server import FastMCP
 
 from .transports import ClientTransport, SessionKwargs, infer_transport
 
-
-def _get_roots_callback(roots: list[mcp.types.Root]) -> ListRootsFnT | None:
-    async def _roots_callback(
-        context: RequestContext[ClientSession, LifespanContextT],
-    ) -> mcp.types.ListRootsResult:
-        return mcp.types.ListRootsResult(roots=roots)
-
-    return _roots_callback
+__all__ = ["Client", "RootsHandler", "RootsList"]
 
 
 class Client:
@@ -40,10 +36,9 @@ class Client:
         self,
         transport: ClientTransport | FastMCP | AnyUrl | Path | str,
         # Common args
-        roots: list[mcp.types.Root] | None = None,
-        sampling_callback: SamplingFnT | None = None,
-        list_roots_callback: ListRootsFnT | None = None,
-        logging_callback: LoggingFnT | None = None,
+        roots: RootsList | RootsHandler | None = None,
+        sampling_handler: SamplingHandler | None = None,
+        log_handler: LoggingFnT | None = None,
         message_handler: MessageHandlerFnT | None = None,
         read_timeout_seconds: datetime.timedelta | None = None,
     ):
@@ -51,21 +46,20 @@ class Client:
         self._session: ClientSession | None = None
         self._session_cm: AbstractAsyncContextManager[ClientSession] | None = None
 
-        # Store common kwargs to pass to transport.connect_session
-        if roots is not None and list_roots_callback is not None:
-            raise ValueError("Cannot provide both `roots` and `list_roots_callback`.")
-        resolved_list_roots_callback = list_roots_callback or (
-            _get_roots_callback(roots) if roots else None
-        )
-
         self._session_kwargs: SessionKwargs = {
-            "sampling_callback": sampling_callback,
-            "list_roots_callback": resolved_list_roots_callback,
-            "logging_callback": logging_callback,
+            "sampling_callback": None,
+            "list_roots_callback": None,
+            "logging_callback": log_handler,
             "message_handler": message_handler,
             "read_timeout_seconds": read_timeout_seconds,
         }
 
+        if roots is not None:
+            self.set_roots(roots)
+
+        if sampling_handler is not None:
+            self.set_sampling_callback(sampling_handler)
+
     @property
     def session(self) -> ClientSession:
         """Get the current active session. Raises RuntimeError if not connected."""
@@ -75,6 +69,16 @@ class Client:
             )
         return self._session
 
+    def set_roots(self, roots: RootsList | RootsHandler) -> None:
+        """Set the roots for the client. This does not automatically call `send_roots_list_changed`."""
+        self._session_kwargs["list_roots_callback"] = create_roots_callback(roots)
+
+    def set_sampling_callback(self, sampling_callback: SamplingHandler) -> None:
+        """Set the sampling callback for the client."""
+        self._session_kwargs["sampling_callback"] = create_sampling_callback(
+            sampling_callback
+        )
+
     def is_connected(self) -> bool:
         """Check if the client is currently connected."""
         return self._session is not None

src/fastmcp/client/roots.py

@@ -0,0 +1,75 @@
+import inspect
+from collections.abc import Awaitable, Callable
+from typing import TypeAlias
+
+import mcp.types
+import pydantic
+from mcp import ClientSession
+from mcp.client.session import ListRootsFnT
+from mcp.shared.context import LifespanContextT, RequestContext
+
+RootsList: TypeAlias = list[str] | list[mcp.types.Root] | list[str | mcp.types.Root]
+
+RootsHandler: TypeAlias = (
+    Callable[[RequestContext[ClientSession, LifespanContextT]], RootsList]
+    | Callable[[RequestContext[ClientSession, LifespanContextT]], Awaitable[RootsList]]
+)
+
+
+def convert_roots_list(roots: RootsList) -> list[mcp.types.Root]:
+    roots_list = []
+    for r in roots:
+        if isinstance(r, mcp.types.Root):
+            roots_list.append(r)
+        elif isinstance(r, pydantic.FileUrl):
+            roots_list.append(mcp.types.Root(uri=r))
+        elif isinstance(r, str):
+            roots_list.append(mcp.types.Root(uri=pydantic.FileUrl(r)))
+        else:
+            raise ValueError(f"Invalid root: {r}")
+    return roots_list
+
+
+def create_roots_callback(
+    handler: RootsList | RootsHandler,
+) -> ListRootsFnT:
+    if isinstance(handler, list):
+        return _create_roots_callback_from_roots(handler)
+    elif inspect.isfunction(handler):
+        return _create_roots_callback_from_fn(handler)
+    else:
+        raise ValueError(f"Invalid roots handler: {handler}")
+
+
+def _create_roots_callback_from_roots(
+    roots: RootsList,
+) -> ListRootsFnT:
+    roots = convert_roots_list(roots)
+
+    async def _roots_callback(
+        context: RequestContext[ClientSession, LifespanContextT],
+    ) -> mcp.types.ListRootsResult:
+        return mcp.types.ListRootsResult(roots=roots)
+
+    return _roots_callback
+
+
+def _create_roots_callback_from_fn(
+    fn: Callable[[RequestContext[ClientSession, LifespanContextT]], RootsList]
+    | Callable[[RequestContext[ClientSession, LifespanContextT]], Awaitable[RootsList]],
+) -> ListRootsFnT:
+    async def _roots_callback(
+        context: RequestContext[ClientSession, LifespanContextT],
+    ) -> mcp.types.ListRootsResult | mcp.types.ErrorData:
+        try:
+            roots = fn(context)
+            if inspect.isawaitable(roots):
+                roots = await roots
+            return mcp.types.ListRootsResult(roots=convert_roots_list(roots))
+        except Exception as e:
+            return mcp.types.ErrorData(
+                code=mcp.types.INTERNAL_ERROR,
+                message=str(e),
+            )
+
+    return _roots_callback

src/fastmcp/client/sampling.py

@@ -0,0 +1,50 @@
+import inspect
+from collections.abc import Awaitable, Callable
+from typing import TypeAlias
+
+import mcp.types
+from mcp import ClientSession, CreateMessageResult
+from mcp.client.session import SamplingFnT
+from mcp.shared.context import LifespanContextT, RequestContext
+from mcp.types import CreateMessageRequestParams as SamplingParams
+from mcp.types import SamplingMessage
+
+
+class MessageResult(CreateMessageResult):
+    role: mcp.types.Role = "assistant"
+    content: mcp.types.TextContent | mcp.types.ImageContent
+    model: str = "client-model"
+
+
+SamplingHandler: TypeAlias = Callable[
+    [
+        list[SamplingMessage],
+        SamplingParams,
+        RequestContext[ClientSession, LifespanContextT],
+    ],
+    str | CreateMessageResult | Awaitable[str | CreateMessageResult],
+]
+
+
+def create_sampling_callback(sampling_handler: SamplingHandler) -> SamplingFnT:
+    async def _sampling_handler(
+        context: RequestContext[ClientSession, LifespanContextT],
+        params: SamplingParams,
+    ) -> CreateMessageResult | mcp.types.ErrorData:
+        try:
+            result = sampling_handler(params.messages, params, context)
+            if inspect.isawaitable(result):
+                result = await result
+
+            if isinstance(result, str):
+                result = MessageResult(
+                    content=mcp.types.TextContent(type="text", text=result)
+                )
+            return result
+        except Exception as e:
+            return mcp.types.ErrorData(
+                code=mcp.types.INTERNAL_ERROR,
+                message=str(e),
+            )
+
+    return _sampling_handler

src/fastmcp/server/context.py

@@ -1,6 +1,5 @@
 from __future__ import annotations as _annotations
 
-from collections.abc import Iterable
 from typing import Any, Generic, Literal
 
 from mcp.server.lowlevel.helper_types import ReadResourceContents
@@ -9,6 +8,7 @@ from mcp.shared.context import LifespanContextT, RequestContext
 from mcp.types import (
     CreateMessageResult,
     ImageContent,
+    Root,
     SamplingMessage,
     TextContent,
 )
@@ -106,7 +106,7 @@ class Context(BaseModel, Generic[ServerSessionT, LifespanContextT]):
             progress_token=progress_token, progress=progress, total=total
         )
 
-    async def read_resource(self, uri: str | AnyUrl) -> Iterable[ReadResourceContents]:
+    async def read_resource(self, uri: str | AnyUrl) -> list[ReadResourceContents]:
         """Read a resource by URI.
 
         Args:
@@ -175,9 +175,14 @@ class Context(BaseModel, Generic[ServerSessionT, LifespanContextT]):
         """Send an error log message."""
         await self.log("error", message, **extra)
 
+    async def list_roots(self) -> list[Root]:
+        """List the roots available to the server, as indicated by the client."""
+        result = await self.request_context.session.list_roots()
+        return result.roots
+
     async def sample(
         self,
-        message: str,
+        messages: str | list[str | SamplingMessage],
         system_prompt: str | None = None,
         temperature: float | None = None,
         max_tokens: int | None = None,
@@ -193,16 +198,22 @@ class Context(BaseModel, Generic[ServerSessionT, LifespanContextT]):
         if max_tokens is None:
             max_tokens = 512
 
-        assert self._request_context is not None
-        assert self._request_context.session is not None
-
-        sampling_message = SamplingMessage(
-            content=TextContent(text=message, type="text"),
-            role="user",
-        )
+        if isinstance(messages, str):
+            sampling_messages = [
+                SamplingMessage(
+                    content=TextContent(text=messages, type="text"), role="user"
+                )
+            ]
+        elif isinstance(messages, list):
+            sampling_messages = [
+                SamplingMessage(content=TextContent(text=m, type="text"), role="user")
+                if isinstance(m, str)
+                else m
+                for m in messages
+            ]
 
         result: CreateMessageResult = await self.request_context.session.create_message(
-            messages=[sampling_message],
+            messages=sampling_messages,
             system_prompt=system_prompt,
             temperature=temperature,
             max_tokens=max_tokens,

tests/client/test_roots.py

@@ -0,0 +1,48 @@
+import json
+
+import pytest
+from mcp.types import TextContent
+
+from fastmcp import Client, Context, FastMCP
+
+
+@pytest.fixture
+def fastmcp_server():
+    mcp = FastMCP()
+
+    @mcp.tool()
+    async def list_roots(context: Context) -> list[str]:
+        roots = await context.list_roots()
+        return [str(r.uri) for r in roots]
+
+    return mcp
+
+
+class TestClientRoots:
+    @pytest.mark.parametrize("roots", [["x"], ["x", "y"]])
+    async def test_invalid_roots(self, fastmcp_server: FastMCP, roots: list[str]):
+        """
+        Roots must be URIs
+        """
+        with pytest.raises(ValueError, match="Input should be a valid URL"):
+            async with Client(fastmcp_server, roots=roots):
+                pass
+
+    @pytest.mark.parametrize("roots", [["https://x.com"]])
+    async def test_invalid_urls(self, fastmcp_server: FastMCP, roots: list[str]):
+        """
+        At this time, root URIs must start with file://
+        """
+        with pytest.raises(ValueError, match="URL scheme should be 'file'"):
+            async with Client(fastmcp_server, roots=roots):
+                pass
+
+    @pytest.mark.parametrize("roots", [["file://x/y/z", "file://x/y/z"]])
+    async def test_valid_roots(self, fastmcp_server: FastMCP, roots: list[str]):
+        async with Client(fastmcp_server, roots=roots) as client:
+            result = await client.call_tool("list_roots", {})
+            assert isinstance(result.content[0], TextContent)
+            assert json.loads(result.content[0].text) == [
+                "file://x/y/z",
+                "file://x/y/z",
+            ]

tests/client/test_sampling.py

@@ -0,0 +1,85 @@
+from typing import cast
+
+import pytest
+from mcp.types import TextContent
+
+from fastmcp import Client, Context, FastMCP
+from fastmcp.client.sampling import RequestContext, SamplingMessage, SamplingParams
+
+
+@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 cast(TextContent, result).text
+
+    @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 cast(TextContent, result).text
+
+    @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 cast(TextContent, result).text
+
+    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!"})
+        reply = cast(TextContent, result.content[0])
+        assert reply.text == "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!"}
+        )
+        reply = cast(TextContent, result.content[0])
+        assert reply.text == "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!"}
+        )
+        reply = cast(TextContent, result.content[0])
+        assert reply.text == "I need to think."