fastmcp 2.13.3__py3-none-any.whl → 2.14.1__py3-none-any.whl

fastmcp/server/sampling/run.py

@@ -0,0 +1,301 @@
+"""Sampling types and helper functions for FastMCP servers."""
+
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Generic
+
+from mcp.types import (
+    ClientCapabilities,
+    CreateMessageResult,
+    CreateMessageResultWithTools,
+    ModelHint,
+    ModelPreferences,
+    SamplingCapability,
+    SamplingMessage,
+    SamplingToolsCapability,
+    TextContent,
+    ToolChoice,
+    ToolResultContent,
+    ToolUseContent,
+)
+from mcp.types import CreateMessageRequestParams as SamplingParams
+from mcp.types import Tool as SDKTool
+from typing_extensions import TypeVar
+
+from fastmcp.exceptions import ToolError
+from fastmcp.server.sampling.sampling_tool import SamplingTool
+from fastmcp.utilities.logging import get_logger
+
+logger = get_logger(__name__)
+
+if TYPE_CHECKING:
+    from fastmcp.server.context import Context
+
+ResultT = TypeVar("ResultT", default=str)
+
+
+@dataclass
+class SamplingResult(Generic[ResultT]):
+    """Result of a sampling operation.
+
+    Attributes:
+        text: The text representation of the result (raw text or JSON for structured).
+        result: The typed result (str for text, parsed object for structured output).
+        history: All messages exchanged during sampling.
+    """
+
+    text: str | None
+    result: ResultT
+    history: list[SamplingMessage]
+
+
+@dataclass
+class SampleStep:
+    """Result of a single sampling call.
+
+    Represents what the LLM returned in this step plus the message history.
+    """
+
+    response: CreateMessageResult | CreateMessageResultWithTools
+    history: list[SamplingMessage]
+
+    @property
+    def is_tool_use(self) -> bool:
+        """True if the LLM is requesting tool execution."""
+        if isinstance(self.response, CreateMessageResultWithTools):
+            return self.response.stopReason == "toolUse"
+        return False
+
+    @property
+    def text(self) -> str | None:
+        """Extract text from the response, if available."""
+        content = self.response.content
+        if isinstance(content, list):
+            for block in content:
+                if isinstance(block, TextContent):
+                    return block.text
+            return None
+        elif isinstance(content, TextContent):
+            return content.text
+        return None
+
+    @property
+    def tool_calls(self) -> list[ToolUseContent]:
+        """Get the list of tool calls from the response."""
+        content = self.response.content
+        if isinstance(content, list):
+            return [c for c in content if isinstance(c, ToolUseContent)]
+        elif isinstance(content, ToolUseContent):
+            return [content]
+        return []
+
+
+def _parse_model_preferences(
+    model_preferences: ModelPreferences | str | list[str] | None,
+) -> ModelPreferences | None:
+    """Convert model preferences to ModelPreferences object."""
+    if model_preferences is None:
+        return None
+    elif isinstance(model_preferences, ModelPreferences):
+        return model_preferences
+    elif isinstance(model_preferences, str):
+        return ModelPreferences(hints=[ModelHint(name=model_preferences)])
+    elif isinstance(model_preferences, list):
+        if not all(isinstance(h, str) for h in model_preferences):
+            raise ValueError("All elements of model_preferences list must be strings.")
+        return ModelPreferences(hints=[ModelHint(name=h) for h in model_preferences])
+    else:
+        raise ValueError(
+            "model_preferences must be one of: ModelPreferences, str, list[str], or None."
+        )
+
+
+# --- Standalone functions for sample_step() ---
+
+
+def determine_handler_mode(context: Context, needs_tools: bool) -> bool:
+    """Determine whether to use fallback handler or client for sampling.
+
+    Args:
+        context: The MCP context.
+        needs_tools: Whether the sampling request requires tool support.
+
+    Returns:
+        True if fallback handler should be used, False to use client.
+
+    Raises:
+        ValueError: If client lacks required capability and no fallback configured.
+    """
+    fastmcp = context.fastmcp
+    session = context.session
+
+    # Check what capabilities the client has
+    has_sampling = session.check_client_capability(
+        capability=ClientCapabilities(sampling=SamplingCapability())
+    )
+    has_tools_capability = session.check_client_capability(
+        capability=ClientCapabilities(
+            sampling=SamplingCapability(tools=SamplingToolsCapability())
+        )
+    )
+
+    if fastmcp.sampling_handler_behavior == "always":
+        if fastmcp.sampling_handler is None:
+            raise ValueError(
+                "sampling_handler_behavior is 'always' but no handler configured"
+            )
+        return True
+    elif fastmcp.sampling_handler_behavior == "fallback":
+        client_sufficient = has_sampling and (not needs_tools or has_tools_capability)
+        if not client_sufficient:
+            if fastmcp.sampling_handler is None:
+                if needs_tools and has_sampling and not has_tools_capability:
+                    raise ValueError(
+                        "Client does not support sampling with tools. "
+                        "The client must advertise the sampling.tools capability."
+                    )
+                raise ValueError("Client does not support sampling")
+            return True
+    elif fastmcp.sampling_handler_behavior is not None:
+        raise ValueError(
+            f"Invalid sampling_handler_behavior: {fastmcp.sampling_handler_behavior!r}. "
+            "Must be 'always', 'fallback', or None."
+        )
+    elif not has_sampling:
+        raise ValueError("Client does not support sampling")
+    elif needs_tools and not has_tools_capability:
+        raise ValueError(
+            "Client does not support sampling with tools. "
+            "The client must advertise the sampling.tools capability."
+        )
+
+    return False
+
+
+async def call_sampling_handler(
+    context: Context,
+    messages: list[SamplingMessage],
+    *,
+    system_prompt: str | None,
+    temperature: float | None,
+    max_tokens: int,
+    model_preferences: ModelPreferences | str | list[str] | None,
+    sdk_tools: list[SDKTool] | None,
+    tool_choice: ToolChoice | None,
+) -> CreateMessageResult | CreateMessageResultWithTools:
+    """Make LLM call using the fallback handler.
+
+    Note: This function expects the caller (sample_step) to have validated that
+    sampling_handler is set via determine_handler_mode(). The checks below are
+    safeguards against internal misuse.
+    """
+    if context.fastmcp.sampling_handler is None:
+        raise RuntimeError("sampling_handler is None")
+    if context.request_context is None:
+        raise RuntimeError("request_context is None")
+
+    result = context.fastmcp.sampling_handler(
+        messages,
+        SamplingParams(
+            systemPrompt=system_prompt,
+            messages=messages,
+            temperature=temperature,
+            maxTokens=max_tokens,
+            modelPreferences=_parse_model_preferences(model_preferences),
+            tools=sdk_tools,
+            toolChoice=tool_choice,
+        ),
+        context.request_context,
+    )
+
+    if inspect.isawaitable(result):
+        result = await result
+
+    # Convert string to CreateMessageResult
+    if isinstance(result, str):
+        return CreateMessageResult(
+            role="assistant",
+            content=TextContent(type="text", text=result),
+            model="unknown",
+            stopReason="endTurn",
+        )
+
+    return result
+
+
+async def execute_tools(
+    tool_calls: list[ToolUseContent],
+    tool_map: dict[str, SamplingTool],
+    mask_error_details: bool = False,
+) -> list[ToolResultContent]:
+    """Execute tool calls and return results.
+
+    Args:
+        tool_calls: List of tool use requests from the LLM.
+        tool_map: Mapping from tool name to SamplingTool.
+        mask_error_details: If True, mask detailed error messages from tool execution.
+            When masked, only generic error messages are returned to the LLM.
+            Tools can explicitly raise ToolError to bypass masking when they want
+            to provide specific error messages to the LLM.
+
+    Returns:
+        List of tool result content blocks.
+    """
+    tool_results: list[ToolResultContent] = []
+
+    for tool_use in tool_calls:
+        tool = tool_map.get(tool_use.name)
+        if tool is None:
+            tool_results.append(
+                ToolResultContent(
+                    type="tool_result",
+                    toolUseId=tool_use.id,
+                    content=[
+                        TextContent(
+                            type="text",
+                            text=f"Error: Unknown tool '{tool_use.name}'",
+                        )
+                    ],
+                    isError=True,
+                )
+            )
+        else:
+            try:
+                result_value = await tool.run(tool_use.input)
+                tool_results.append(
+                    ToolResultContent(
+                        type="tool_result",
+                        toolUseId=tool_use.id,
+                        content=[TextContent(type="text", text=str(result_value))],
+                    )
+                )
+            except ToolError as e:
+                # ToolError is the escape hatch - always pass message through
+                logger.exception(f"Error calling sampling tool '{tool_use.name}'")
+                tool_results.append(
+                    ToolResultContent(
+                        type="tool_result",
+                        toolUseId=tool_use.id,
+                        content=[TextContent(type="text", text=str(e))],
+                        isError=True,
+                    )
+                )
+            except Exception as e:
+                # Generic exceptions - mask based on setting
+                logger.exception(f"Error calling sampling tool '{tool_use.name}'")
+                if mask_error_details:
+                    error_text = f"Error executing tool '{tool_use.name}'"
+                else:
+                    error_text = f"Error executing tool '{tool_use.name}': {e}"
+                tool_results.append(
+                    ToolResultContent(
+                        type="tool_result",
+                        toolUseId=tool_use.id,
+                        content=[TextContent(type="text", text=error_text)],
+                        isError=True,
+                    )
+                )
+
+    return tool_results

fastmcp/server/sampling/sampling_tool.py

@@ -0,0 +1,108 @@
+"""SamplingTool for use during LLM sampling requests."""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import Callable
+from typing import Any
+
+from mcp.types import Tool as SDKTool
+from pydantic import BaseModel, ConfigDict
+
+from fastmcp.tools.tool import ParsedFunction
+
+
+class SamplingTool(BaseModel):
+    """A tool that can be used during LLM sampling.
+
+    SamplingTools bundle a tool's schema (name, description, parameters) with
+    an executor function, enabling servers to execute agentic workflows where
+    the LLM can request tool calls during sampling.
+
+    In most cases, pass functions directly to ctx.sample():
+
+        def search(query: str) -> str:
+            '''Search the web.'''
+            return web_search(query)
+
+        result = await context.sample(
+            messages="Find info about Python",
+            tools=[search],  # Plain functions work directly
+        )
+
+    Create a SamplingTool explicitly when you need custom name/description:
+
+        tool = SamplingTool.from_function(search, name="web_search")
+    """
+
+    name: str
+    description: str | None = None
+    parameters: dict[str, Any]
+    fn: Callable[..., Any]
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    async def run(self, arguments: dict[str, Any] | None = None) -> Any:
+        """Execute the tool with the given arguments.
+
+        Args:
+            arguments: Dictionary of arguments to pass to the tool function.
+
+        Returns:
+            The result of executing the tool function.
+        """
+        if arguments is None:
+            arguments = {}
+
+        result = self.fn(**arguments)
+        if inspect.isawaitable(result):
+            result = await result
+        return result
+
+    def _to_sdk_tool(self) -> SDKTool:
+        """Convert to an mcp.types.Tool for SDK compatibility.
+
+        This is used internally when passing tools to the MCP SDK's
+        create_message() method.
+        """
+        return SDKTool(
+            name=self.name,
+            description=self.description,
+            inputSchema=self.parameters,
+        )
+
+    @classmethod
+    def from_function(
+        cls,
+        fn: Callable[..., Any],
+        *,
+        name: str | None = None,
+        description: str | None = None,
+    ) -> SamplingTool:
+        """Create a SamplingTool from a function.
+
+        The function's signature is analyzed to generate a JSON schema for
+        the tool's parameters. Type hints are used to determine parameter types.
+
+        Args:
+            fn: The function to create a tool from.
+            name: Optional name override. Defaults to the function's name.
+            description: Optional description override. Defaults to the function's docstring.
+
+        Returns:
+            A SamplingTool wrapping the function.
+
+        Raises:
+            ValueError: If the function is a lambda without a name override.
+        """
+        parsed = ParsedFunction.from_function(fn, validate=True)
+
+        if name is None and parsed.name == "<lambda>":
+            raise ValueError("You must provide a name for lambda functions")
+
+        return cls(
+            name=name or parsed.name,
+            description=description or parsed.description,
+            parameters=parsed.input_schema,
+            fn=parsed.fn,
+        )