fastmcp 1.0__py3-none-any.whl → 2.1.0__py3-none-any.whl

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

fastmcp/client/transports.py

@@ -0,0 +1,411 @@
+import abc
+import contextlib
+import datetime
+import os
+from collections.abc import AsyncIterator
+from pathlib import Path
+from typing import (
+    TypedDict,
+)
+
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.session import (
+    ListRootsFnT,
+    LoggingFnT,
+    MessageHandlerFnT,
+    SamplingFnT,
+)
+from mcp.client.sse import sse_client
+from mcp.client.stdio import stdio_client
+from mcp.client.websocket import websocket_client
+from mcp.shared.memory import create_connected_server_and_client_session
+from pydantic import AnyUrl
+from typing_extensions import Unpack
+
+from fastmcp.server import FastMCP as FastMCPServer
+
+
+class SessionKwargs(TypedDict, total=False):
+    """Keyword arguments for the MCP ClientSession constructor."""
+
+    sampling_callback: SamplingFnT | None
+    list_roots_callback: ListRootsFnT | None
+    logging_callback: LoggingFnT | None
+    message_handler: MessageHandlerFnT | None
+    read_timeout_seconds: datetime.timedelta | None
+
+
+class ClientTransport(abc.ABC):
+    """
+    Abstract base class for different MCP client transport mechanisms.
+
+    A Transport is responsible for establishing and managing connections
+    to an MCP server, and providing a ClientSession within an async context.
+    """
+
+    @abc.abstractmethod
+    @contextlib.asynccontextmanager
+    async def connect_session(
+        self, **session_kwargs: Unpack[SessionKwargs]
+    ) -> AsyncIterator[ClientSession]:
+        """
+        Establishes a connection and yields an active, initialized ClientSession.
+
+        The session is guaranteed to be valid only within the scope of the
+        async context manager. Connection setup and teardown are handled
+        within this context.
+
+        Args:
+            **session_kwargs: Keyword arguments to pass to the ClientSession
+                              constructor (e.g., callbacks, timeouts).
+
+        Yields:
+            An initialized mcp.ClientSession instance.
+        """
+        raise NotImplementedError
+        yield None  # type: ignore
+
+    def __repr__(self) -> str:
+        # Basic representation for subclasses
+        return f"<{self.__class__.__name__}>"
+
+
+class WSTransport(ClientTransport):
+    """Transport implementation that connects to an MCP server via WebSockets."""
+
+    def __init__(self, url: str | AnyUrl):
+        if isinstance(url, AnyUrl):
+            url = str(url)
+        if not isinstance(url, str) or not url.startswith("ws"):
+            raise ValueError("Invalid WebSocket URL provided.")
+        self.url = url
+
+    @contextlib.asynccontextmanager
+    async def connect_session(
+        self, **session_kwargs: Unpack[SessionKwargs]
+    ) -> AsyncIterator[ClientSession]:
+        async with websocket_client(self.url) as transport:
+            read_stream, write_stream = transport
+            async with ClientSession(
+                read_stream, write_stream, **session_kwargs
+            ) as session:
+                await session.initialize()  # Initialize after session creation
+                yield session
+
+    def __repr__(self) -> str:
+        return f"<WebSocket(url='{self.url}')>"
+
+
+class SSETransport(ClientTransport):
+    """Transport implementation that connects to an MCP server via Server-Sent Events."""
+
+    def __init__(self, url: str | AnyUrl, headers: dict[str, str] | None = None):
+        if isinstance(url, AnyUrl):
+            url = str(url)
+        if not isinstance(url, str) or not url.startswith("http"):
+            raise ValueError("Invalid HTTP/S URL provided for SSE.")
+        self.url = url
+        self.headers = headers or {}
+
+    @contextlib.asynccontextmanager
+    async def connect_session(
+        self, **session_kwargs: Unpack[SessionKwargs]
+    ) -> AsyncIterator[ClientSession]:
+        async with sse_client(self.url, headers=self.headers) as transport:
+            read_stream, write_stream = transport
+            async with ClientSession(
+                read_stream, write_stream, **session_kwargs
+            ) as session:
+                await session.initialize()
+                yield session
+
+    def __repr__(self) -> str:
+        return f"<SSE(url='{self.url}')>"
+
+
+class StdioTransport(ClientTransport):
+    """
+    Base transport for connecting to an MCP server via subprocess with stdio.
+
+    This is a base class that can be subclassed for specific command-based
+    transports like Python, Node, Uvx, etc.
+    """
+
+    def __init__(
+        self,
+        command: str,
+        args: list[str],
+        env: dict[str, str] | None = None,
+        cwd: str | None = None,
+    ):
+        """
+        Initialize a Stdio transport.
+
+        Args:
+            command: The command to run (e.g., "python", "node", "uvx")
+            args: The arguments to pass to the command
+            env: Environment variables to set for the subprocess
+            cwd: Current working directory for the subprocess
+        """
+        self.command = command
+        self.args = args
+        self.env = env
+        self.cwd = cwd
+
+    @contextlib.asynccontextmanager
+    async def connect_session(
+        self, **session_kwargs: Unpack[SessionKwargs]
+    ) -> AsyncIterator[ClientSession]:
+        server_params = StdioServerParameters(
+            command=self.command, args=self.args, env=self.env, cwd=self.cwd
+        )
+        async with stdio_client(server_params) as transport:
+            read_stream, write_stream = transport
+            async with ClientSession(
+                read_stream, write_stream, **session_kwargs
+            ) as session:
+                await session.initialize()
+                yield session
+
+    def __repr__(self) -> str:
+        return (
+            f"<{self.__class__.__name__}(command='{self.command}', args={self.args})>"
+        )
+
+
+class PythonStdioTransport(StdioTransport):
+    """Transport for running Python scripts."""
+
+    def __init__(
+        self,
+        script_path: str | Path,
+        args: list[str] | None = None,
+        env: dict[str, str] | None = None,
+        cwd: str | None = None,
+        python_cmd: str = "python",
+    ):
+        """
+        Initialize a Python transport.
+
+        Args:
+            script_path: Path to the Python script to run
+            args: Additional arguments to pass to the script
+            env: Environment variables to set for the subprocess
+            cwd: Current working directory for the subprocess
+            python_cmd: Python command to use (default: "python")
+        """
+        script_path = Path(script_path).resolve()
+        if not script_path.is_file():
+            raise FileNotFoundError(f"Script not found: {script_path}")
+        if not str(script_path).endswith(".py"):
+            raise ValueError(f"Not a Python script: {script_path}")
+
+        full_args = [str(script_path)]
+        if args:
+            full_args.extend(args)
+
+        super().__init__(command=python_cmd, args=full_args, env=env, cwd=cwd)
+        self.script_path = script_path
+
+
+class NodeStdioTransport(StdioTransport):
+    """Transport for running Node.js scripts."""
+
+    def __init__(
+        self,
+        script_path: str | Path,
+        args: list[str] | None = None,
+        env: dict[str, str] | None = None,
+        cwd: str | None = None,
+        node_cmd: str = "node",
+    ):
+        """
+        Initialize a Node transport.
+
+        Args:
+            script_path: Path to the Node.js script to run
+            args: Additional arguments to pass to the script
+            env: Environment variables to set for the subprocess
+            cwd: Current working directory for the subprocess
+            node_cmd: Node.js command to use (default: "node")
+        """
+        script_path = Path(script_path).resolve()
+        if not script_path.is_file():
+            raise FileNotFoundError(f"Script not found: {script_path}")
+        if not str(script_path).endswith(".js"):
+            raise ValueError(f"Not a JavaScript script: {script_path}")
+
+        full_args = [str(script_path)]
+        if args:
+            full_args.extend(args)
+
+        super().__init__(command=node_cmd, args=full_args, env=env, cwd=cwd)
+        self.script_path = script_path
+
+
+class UvxStdioTransport(StdioTransport):
+    """Transport for running commands via the uvx tool."""
+
+    def __init__(
+        self,
+        tool_name: str,
+        tool_args: list[str] | None = None,
+        project_directory: str | None = None,
+        python_version: str | None = None,
+        with_packages: list[str] | None = None,
+        from_package: str | None = None,
+        env_vars: dict[str, str] | None = None,
+    ):
+        """
+        Initialize a Uvx transport.
+
+        Args:
+            tool_name: Name of the tool to run via uvx
+            tool_args: Arguments to pass to the tool
+            project_directory: Project directory (for package resolution)
+            python_version: Python version to use
+            with_packages: Additional packages to include
+            from_package: Package to install the tool from
+            env_vars: Additional environment variables
+        """
+        # Basic validation
+        if project_directory and not Path(project_directory).exists():
+            raise NotADirectoryError(
+                f"Project directory not found: {project_directory}"
+            )
+
+        # Build uvx arguments
+        uvx_args = []
+        if python_version:
+            uvx_args.extend(["--python", python_version])
+        if from_package:
+            uvx_args.extend(["--from", from_package])
+        for pkg in with_packages or []:
+            uvx_args.extend(["--with", pkg])
+
+        # Add the tool name and tool args
+        uvx_args.append(tool_name)
+        if tool_args:
+            uvx_args.extend(tool_args)
+
+        # Get environment with any additional variables
+        env = None
+        if env_vars:
+            env = os.environ.copy()
+            env.update(env_vars)
+
+        super().__init__(command="uvx", args=uvx_args, env=env, cwd=project_directory)
+        self.tool_name = tool_name
+
+
+class NpxStdioTransport(StdioTransport):
+    """Transport for running commands via the npx tool."""
+
+    def __init__(
+        self,
+        package: str,
+        args: list[str] | None = None,
+        project_directory: str | None = None,
+        env_vars: dict[str, str] | None = None,
+        use_package_lock: bool = True,
+    ):
+        """
+        Initialize an Npx transport.
+
+        Args:
+            package: Name of the npm package to run
+            args: Arguments to pass to the package command
+            project_directory: Project directory with package.json
+            env_vars: Additional environment variables
+            use_package_lock: Whether to use package-lock.json (--prefer-offline)
+        """
+        # Basic validation
+        if project_directory and not Path(project_directory).exists():
+            raise NotADirectoryError(
+                f"Project directory not found: {project_directory}"
+            )
+
+        # Build npx arguments
+        npx_args = []
+        if use_package_lock:
+            npx_args.append("--prefer-offline")
+
+        # Add the package name and args
+        npx_args.append(package)
+        if args:
+            npx_args.extend(args)
+
+        # Get environment with any additional variables
+        env = None
+        if env_vars:
+            env = os.environ.copy()
+            env.update(env_vars)
+
+        super().__init__(command="npx", args=npx_args, env=env, cwd=project_directory)
+        self.package = package
+
+
+class FastMCPTransport(ClientTransport):
+    """
+    Special transport for in-memory connections to an MCP server.
+
+    This is particularly useful for testing or when client and server
+    are in the same process.
+    """
+
+    def __init__(self, mcp: FastMCPServer):
+        self._fastmcp = mcp  # Can be FastMCP or MCPServer
+
+    @contextlib.asynccontextmanager
+    async def connect_session(
+        self, **session_kwargs: Unpack[SessionKwargs]
+    ) -> AsyncIterator[ClientSession]:
+        # create_connected_server_and_client_session manages the session lifecycle itself
+        async with create_connected_server_and_client_session(
+            server=self._fastmcp._mcp_server,
+            **session_kwargs,
+        ) as session:
+            yield session
+
+    def __repr__(self) -> str:
+        return f"<FastMCP(server='{self._fastmcp.name}')>"
+
+
+def infer_transport(
+    transport: ClientTransport | FastMCPServer | AnyUrl | Path | str,
+) -> ClientTransport:
+    """
+    Infer the appropriate transport type from the given transport argument.
+
+    This function attempts to infer the correct transport type from the provided
+    argument, handling various input types and converting them to the appropriate
+    ClientTransport subclass.
+    """
+    # the transport is already a ClientTransport
+    if isinstance(transport, ClientTransport):
+        return transport
+
+    # the transport is a FastMCP server
+    elif isinstance(transport, FastMCPServer):
+        return FastMCPTransport(mcp=transport)
+
+    # the transport is a path to a script
+    elif isinstance(transport, Path | str) and Path(transport).exists():
+        if str(transport).endswith(".py"):
+            return PythonStdioTransport(script_path=transport)
+        elif str(transport).endswith(".js"):
+            return NodeStdioTransport(script_path=transport)
+        else:
+            raise ValueError(f"Unsupported script type: {transport}")
+
+    # the transport is an http(s) URL
+    elif isinstance(transport, AnyUrl | str) and str(transport).startswith("http"):
+        return SSETransport(url=transport)
+
+    # the transport is a websocket URL
+    elif isinstance(transport, AnyUrl | str) and str(transport).startswith("ws"):
+        return WSTransport(url=transport)
+
+    # the transport is an unknown type
+    else:
+        raise ValueError(f"Could not infer a valid transport from: {transport}")

fastmcp/prompts/__init__.py

@@ -1,4 +1,4 @@
-from .base import Prompt
-from .manager import PromptManager
+from .prompt import Prompt
+from .prompt_manager import PromptManager
 
 __all__ = ["Prompt", "PromptManager"]

fastmcp/prompts/{base.py → prompt.py}

@@ -1,12 +1,16 @@
 """Base classes for FastMCP prompts."""
 
-import json
-from typing import Any, Callable, Dict, Literal, Optional, Sequence, Awaitable
 import inspect
+import json
+from collections.abc import Awaitable, Callable, Sequence
+from typing import Annotated, Any, Literal
 
-from pydantic import BaseModel, Field, TypeAdapter, validate_call
-from mcp.types import TextContent, ImageContent, EmbeddedResource
 import pydantic_core
+from mcp.types import EmbeddedResource, ImageContent, TextContent
+from pydantic import BaseModel, BeforeValidator, Field, TypeAdapter, validate_call
+from typing_extensions import Self
+
+from fastmcp.utilities.types import _convert_set_defaults
 
 CONTENT_TYPES = TextContent | ImageContent | EmbeddedResource
 
@@ -17,7 +21,7 @@ class Message(BaseModel):
     role: Literal["user", "assistant"]
     content: CONTENT_TYPES
 
-    def __init__(self, content: str | CONTENT_TYPES, **kwargs):
+    def __init__(self, content: str | CONTENT_TYPES, **kwargs: Any):
         if isinstance(content, str):
             content = TextContent(type="text", text=content)
         super().__init__(content=content, **kwargs)
@@ -26,22 +30,24 @@ class Message(BaseModel):
 class UserMessage(Message):
     """A message from the user."""
 
-    role: Literal["user"] = "user"
+    role: Literal["user", "assistant"] = "user"
 
-    def __init__(self, content: str | CONTENT_TYPES, **kwargs):
+    def __init__(self, content: str | CONTENT_TYPES, **kwargs: Any):
         super().__init__(content=content, **kwargs)
 
 
 class AssistantMessage(Message):
     """A message from the assistant."""
 
-    role: Literal["assistant"] = "assistant"
+    role: Literal["user", "assistant"] = "assistant"
 
-    def __init__(self, content: str | CONTENT_TYPES, **kwargs):
+    def __init__(self, content: str | CONTENT_TYPES, **kwargs: Any):
         super().__init__(content=content, **kwargs)
 
 
-message_validator = TypeAdapter(UserMessage | AssistantMessage)
+message_validator = TypeAdapter[UserMessage | AssistantMessage](
+    UserMessage | AssistantMessage
+)
 
 SyncPromptResult = (
     str | Message | dict[str, Any] | Sequence[str | Message | dict[str, Any]]
@@ -68,17 +74,21 @@ class Prompt(BaseModel):
     description: str | None = Field(
         None, description="Description of what the prompt does"
     )
+    tags: Annotated[set[str], BeforeValidator(_convert_set_defaults)] = Field(
+        default_factory=set, description="Tags for the prompt"
+    )
     arguments: list[PromptArgument] | None = Field(
         None, description="Arguments that can be passed to the prompt"
     )
-    fn: Callable = Field(exclude=True)
+    fn: Callable[..., PromptResult | Awaitable[PromptResult]]
 
     @classmethod
     def from_function(
         cls,
-        fn: Callable[..., PromptResult],
-        name: Optional[str] = None,
-        description: Optional[str] = None,
+        fn: Callable[..., PromptResult | Awaitable[PromptResult]],
+        name: str | None = None,
+        description: str | None = None,
+        tags: set[str] | None = None,
     ) -> "Prompt":
         """Create a Prompt from a function.
 
@@ -97,7 +107,7 @@ class Prompt(BaseModel):
         parameters = TypeAdapter(fn).json_schema()
 
         # Convert parameters to PromptArguments
-        arguments = []
+        arguments: list[PromptArgument] = []
         if "properties" in parameters:
             for param_name, param in parameters["properties"].items():
                 required = param_name in parameters.get("required", [])
@@ -117,9 +127,10 @@ class Prompt(BaseModel):
             description=description or fn.__doc__ or "",
             arguments=arguments,
             fn=fn,
+            tags=tags or set(),
         )
 
-    async def render(self, arguments: Optional[Dict[str, Any]] = None) -> list[Message]:
+    async def render(self, arguments: dict[str, Any] | None = None) -> list[Message]:
         """Render the prompt with arguments."""
         # Validate required arguments
         if self.arguments:
@@ -136,25 +147,23 @@ class Prompt(BaseModel):
                 result = await result
 
             # Validate messages
-            if not isinstance(result, (list, tuple)):
+            if not isinstance(result, list | tuple):
                 result = [result]
 
             # Convert result to messages
-            messages = []
-            for msg in result:
+            messages: list[Message] = []
+            for msg in result:  # type: ignore[reportUnknownVariableType]
                 try:
                     if isinstance(msg, Message):
                         messages.append(msg)
                     elif isinstance(msg, dict):
-                        msg = message_validator.validate_python(msg)
-                        messages.append(msg)
+                        messages.append(message_validator.validate_python(msg))
                     elif isinstance(msg, str):
-                        messages.append(
-                            UserMessage(content=TextContent(type="text", text=msg))
-                        )
+                        content = TextContent(type="text", text=msg)
+                        messages.append(UserMessage(content=content))
                     else:
-                        msg = json.dumps(pydantic_core.to_jsonable_python(msg))
-                        messages.append(Message(role="user", content=msg))
+                        content = json.dumps(pydantic_core.to_jsonable_python(msg))
+                        messages.append(Message(role="user", content=content))
                 except Exception:
                     raise ValueError(
                         f"Could not convert prompt result to message: {msg}"
@@ -163,3 +172,15 @@ class Prompt(BaseModel):
             return messages
         except Exception as e:
             raise ValueError(f"Error rendering prompt {self.name}: {e}")
+
+    def copy(self, updates: dict[str, Any] | None = None) -> Self:
+        """Copy the prompt with optional updates."""
+        data = self.model_dump()
+        if updates:
+            data.update(updates)
+        return type(self)(**data)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Prompt):
+            return False
+        return self.model_dump() == other.model_dump()