diff --git a/examples/sampling.py b/examples/sampling.py new file mode 100644 index 000000000..385f9d576 --- /dev/null +++ b/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()) diff --git a/src/fastmcp/client/client.py b/src/fastmcp/client/client.py index c6d2a7cdd..ea4f81cd4 100644 --- a/src/fastmcp/client/client.py +++ b/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 diff --git a/src/fastmcp/client/roots.py b/src/fastmcp/client/roots.py new file mode 100644 index 000000000..a04cefd8d --- /dev/null +++ b/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 diff --git a/src/fastmcp/client/sampling.py b/src/fastmcp/client/sampling.py new file mode 100644 index 000000000..2945ef24a --- /dev/null +++ b/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 diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index 7c099cc2a..830bb6ed4 100644 --- a/src/fastmcp/server/context.py +++ b/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, diff --git a/tests/client/test_roots.py b/tests/client/test_roots.py new file mode 100644 index 000000000..0ede20f01 --- /dev/null +++ b/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", + ] diff --git a/tests/client/test_sampling.py b/tests/client/test_sampling.py new file mode 100644 index 000000000..0707f3a44 --- /dev/null +++ b/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." diff --git a/tests/server.py b/tests/server.py deleted file mode 100644 index e69de29bb..000000000