"""Sampling types and helper functions for FastMCP servers.""" from __future__ import annotations import inspect import json from collections.abc import Callable, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Generic, Literal, cast import anyio from mcp.types import ( ClientCapabilities, CreateMessageResult, CreateMessageResultWithTools, ModelHint, ModelPreferences, SamplingCapability, SamplingMessage, SamplingMessageContentBlock, SamplingToolsCapability, TextContent, ToolChoice, ToolResultContent, ToolUseContent, ) from mcp.types import CreateMessageRequestParams as SamplingParams from mcp.types import Tool as SDKTool from opentelemetry.trace import SpanKind, Status, StatusCode from pydantic import ValidationError from typing_extensions import TypeVar from fastmcp import settings from fastmcp.exceptions import ToolError from fastmcp.server.sampling.sampling_tool import SamplingTool from fastmcp.telemetry import get_tracer from fastmcp.tools.function_tool import FunctionTool from fastmcp.tools.tool_transform import TransformedTool from fastmcp.utilities.async_utils import gather from fastmcp.utilities.json_schema import compress_schema from fastmcp.utilities.logging import get_logger from fastmcp.utilities.types import get_cached_typeadapter logger = get_logger(__name__) if TYPE_CHECKING: from fastmcp.server.context import Context ResultT = TypeVar("ResultT") # Maximum number of consecutive final_response validation retries (not # counting the initial attempt) before aborting. Total attempts = N + 1. _MAX_VALIDATION_RETRIES = 3 # Simplified tool choice type - just the mode string instead of the full MCP object ToolChoiceOption = Literal["auto", "required", "none"] # How many times we retry when the LLM returns text instead of calling final_response _MAX_TEXT_RESPONSE_RETRIES = 3 @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 result = cast("str | CreateMessageResult | CreateMessageResultWithTools", 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, tool_concurrency: int | None = None, ) -> 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. tool_concurrency: Controls parallel execution of tools: - None (default): Sequential execution (one at a time) - 0: Unlimited parallel execution - N > 0: Execute at most N tools concurrently If any tool has sequential=True, all tools execute sequentially regardless of this setting. Returns: List of tool result content blocks in the same order as tool_calls. """ if tool_concurrency is not None and tool_concurrency < 0: raise ValueError( f"tool_concurrency must be None, 0 (unlimited), or a positive integer, " f"got {tool_concurrency}" ) async def _execute_single_tool(tool_use: ToolUseContent) -> ToolResultContent: """Execute a single tool and return its result.""" tool = tool_map.get(tool_use.name) if tool is None: return ToolResultContent( type="tool_result", toolUseId=tool_use.id, content=[ TextContent( type="text", text=f"Error: Unknown tool '{tool_use.name}'", ) ], isError=True, ) tracer = get_tracer() with tracer.start_as_current_span( f"sampling tool {tool_use.name}", kind=SpanKind.INTERNAL, ) as span: if span.is_recording(): span.set_attribute("gen_ai.tool.name", tool_use.name) span.set_attribute("fastmcp.tool.use_id", tool_use.id) try: result_value = await tool.run(tool_use.input) return ToolResultContent( type="tool_result", toolUseId=tool_use.id, content=[TextContent(type="text", text=str(result_value))], ) except ToolError as e: if span.is_recording(): span.set_attribute("error.type", "tool_error") span.record_exception(e) span.set_status(Status(StatusCode.ERROR, str(e))) logger.log( e.log_level, f"Error calling sampling tool '{tool_use.name}'", exc_info=True, ) return ToolResultContent( type="tool_result", toolUseId=tool_use.id, content=[TextContent(type="text", text=str(e))], isError=True, ) except Exception as e: if span.is_recording(): span.set_attribute("error.type", type(e).__qualname__) span.record_exception(e) span.set_status(Status(StatusCode.ERROR, str(e))) 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}" return ToolResultContent( type="tool_result", toolUseId=tool_use.id, content=[TextContent(type="text", text=error_text)], isError=True, ) # Check if any tool requires sequential execution requires_sequential = any( tool.sequential for tool_use in tool_calls if (tool := tool_map.get(tool_use.name)) is not None ) # Execute sequentially if required or if concurrency is None (default) if tool_concurrency is None or requires_sequential: tool_results: list[ToolResultContent] = [] for tool_use in tool_calls: result = await _execute_single_tool(tool_use) tool_results.append(result) return tool_results # Execute in parallel if tool_concurrency == 0: # Unlimited parallel execution return await gather(*[_execute_single_tool(tc) for tc in tool_calls]) else: # Bounded parallel execution with semaphore semaphore = anyio.Semaphore(tool_concurrency) async def bounded_execute(tool_use: ToolUseContent) -> ToolResultContent: async with semaphore: return await _execute_single_tool(tool_use) return await gather(*[bounded_execute(tc) for tc in tool_calls]) # --- Helper functions for sampling --- def prepare_messages( messages: str | Sequence[str | SamplingMessage], ) -> list[SamplingMessage]: """Convert various message formats to a list of SamplingMessage objects.""" if isinstance(messages, str): return [ SamplingMessage( content=TextContent(text=messages, type="text"), role="user" ) ] else: return [ SamplingMessage(content=TextContent(text=m, type="text"), role="user") if isinstance(m, str) else m for m in messages ] def prepare_tools( tools: Sequence[SamplingTool | FunctionTool | TransformedTool | Callable[..., Any]] | None, ) -> list[SamplingTool] | None: """Convert tools to SamplingTool objects. Accepts SamplingTool instances, FunctionTool instances, TransformedTool instances, or plain callable functions. FunctionTool and TransformedTool are converted using from_callable_tool(), while plain functions use from_function(). Args: tools: Sequence of tools to prepare. Can be SamplingTool, FunctionTool, TransformedTool, or plain callable functions. Returns: List of SamplingTool instances, or None if tools is None. """ if tools is None: return None sampling_tools: list[SamplingTool] = [] for t in tools: if isinstance(t, SamplingTool): sampling_tools.append(t) elif isinstance(t, (FunctionTool, TransformedTool)): sampling_tools.append(SamplingTool.from_callable_tool(t)) elif callable(t): sampling_tools.append(SamplingTool.from_function(t)) else: raise TypeError( f"Expected SamplingTool, FunctionTool, TransformedTool, or callable, got {type(t)}" ) return sampling_tools if sampling_tools else None def extract_tool_calls( response: CreateMessageResult | CreateMessageResultWithTools, ) -> list[ToolUseContent]: """Extract tool calls from a response.""" content = response.content if isinstance(content, list): return [c for c in content if isinstance(c, ToolUseContent)] elif isinstance(content, ToolUseContent): return [content] return [] def create_final_response_tool(result_type: type) -> SamplingTool: """Create a synthetic 'final_response' tool for structured output. This tool is used to capture structured responses from the LLM. The tool's schema is derived from the result_type. """ type_adapter = get_cached_typeadapter(result_type) schema = type_adapter.json_schema() schema = compress_schema(schema, prune_titles=True) # Tool parameters must be object-shaped. Wrap primitives in {"value": } if schema.get("type") != "object": schema = { "type": "object", "properties": {"value": schema}, "required": ["value"], } # The fn just returns the input as-is (validation happens in the loop) def final_response(**kwargs: Any) -> dict[str, Any]: return kwargs return SamplingTool( name="final_response", description=( "Call this tool to provide your final response. " "Use this when you have completed the task and are ready to return the result." ), parameters=schema, fn=final_response, ) # --- Implementation functions for Context methods --- async def sample_step_impl( context: Context, messages: str | Sequence[str | SamplingMessage], *, system_prompt: str | None = None, temperature: float | None = None, max_tokens: int | None = None, model_preferences: ModelPreferences | str | list[str] | None = None, tools: Sequence[SamplingTool | FunctionTool | TransformedTool | Callable[..., Any]] | None = None, tool_choice: ToolChoiceOption | str | None = None, auto_execute_tools: bool = True, mask_error_details: bool | None = None, tool_concurrency: int | None = None, ) -> SampleStep: """Implementation of Context.sample_step(). Make a single LLM sampling call. This is a stateless function that makes exactly one LLM call and optionally executes any requested tools. """ # Convert messages to SamplingMessage objects current_messages = prepare_messages(messages) # Convert tools to SamplingTools sampling_tools = prepare_tools(tools) sdk_tools: list[SDKTool] | None = ( [t._to_sdk_tool() for t in sampling_tools] if sampling_tools else None ) tool_map: dict[str, SamplingTool] = ( {t.name: t for t in sampling_tools} if sampling_tools else {} ) # Determine whether to use fallback handler or client use_fallback = determine_handler_mode(context, bool(sampling_tools)) # Build tool choice effective_tool_choice: ToolChoice | None = None if tool_choice is not None: if tool_choice not in ("auto", "required", "none"): raise ValueError( f"Invalid tool_choice: {tool_choice!r}. " "Must be 'auto', 'required', or 'none'." ) effective_tool_choice = ToolChoice( mode=cast(Literal["auto", "required", "none"], tool_choice) ) # Effective max_tokens effective_max_tokens = max_tokens if max_tokens is not None else 512 # Make the LLM call tracer = get_tracer() with tracer.start_as_current_span( "sampling create_message", kind=SpanKind.CLIENT, record_exception=False, set_status_on_exception=False, ) as span: if span.is_recording(): span.set_attribute("mcp.method.name", "sampling/createMessage") span.set_attribute("fastmcp.server.name", context.fastmcp.name) try: if use_fallback: response = await call_sampling_handler( context, current_messages, system_prompt=system_prompt, temperature=temperature, max_tokens=effective_max_tokens, model_preferences=model_preferences, sdk_tools=sdk_tools, tool_choice=effective_tool_choice, ) else: response = await context.session.create_message( messages=current_messages, system_prompt=system_prompt, temperature=temperature, max_tokens=effective_max_tokens, model_preferences=_parse_model_preferences(model_preferences), tools=sdk_tools, tool_choice=effective_tool_choice, related_request_id=context.origin_request_id, ) except Exception as e: if span.is_recording(): span.set_attribute("error.type", type(e).__qualname__) span.record_exception(e) span.set_status(Status(StatusCode.ERROR, str(e))) raise # Check if this is a tool use response is_tool_use_response = ( isinstance(response, CreateMessageResultWithTools) and response.stopReason == "toolUse" ) # Always include the assistant response in history current_messages.append(SamplingMessage(role="assistant", content=response.content)) # If not a tool use, return immediately if not is_tool_use_response: return SampleStep(response=response, history=current_messages) # If not executing tools, return with assistant message but no tool results if not auto_execute_tools: return SampleStep(response=response, history=current_messages) # Execute tools and add results to history step_tool_calls = extract_tool_calls(response) if step_tool_calls: effective_mask = ( mask_error_details if mask_error_details is not None else settings.mask_error_details ) tool_results: list[ToolResultContent] = await execute_tools( step_tool_calls, tool_map, mask_error_details=effective_mask, tool_concurrency=tool_concurrency, ) if tool_results: current_messages.append( SamplingMessage( role="user", content=cast(list[SamplingMessageContentBlock], tool_results), ) ) return SampleStep(response=response, history=current_messages) async def sample_impl( context: Context, messages: str | Sequence[str | SamplingMessage], *, system_prompt: str | None = None, temperature: float | None = None, max_tokens: int | None = None, model_preferences: ModelPreferences | str | list[str] | None = None, tools: Sequence[SamplingTool | FunctionTool | TransformedTool | Callable[..., Any]] | None = None, result_type: type[ResultT] | None = None, mask_error_details: bool | None = None, tool_concurrency: int | None = None, ) -> SamplingResult[ResultT]: """Implementation of Context.sample(). Send a sampling request to the client and await the response. This method runs to completion automatically, executing a tool loop until the LLM provides a final text response. """ # Safety limit to prevent infinite loops max_iterations = 100 # Convert tools to SamplingTools sampling_tools = prepare_tools(tools) # Handle structured output with result_type tool_choice: str | None = None if result_type is not None and result_type is not str: final_response_tool = create_final_response_tool(result_type) sampling_tools = list(sampling_tools) if sampling_tools else [] sampling_tools.append(final_response_tool) # Always require tool calls when result_type is set - the LLM must # eventually call final_response (text responses are not accepted) tool_choice = "required" # Convert messages for the loop current_messages: str | Sequence[str | SamplingMessage] = messages text_response_retries = 0 consecutive_validation_failures = 0 for _iteration in range(max_iterations): step = await sample_step_impl( context, messages=current_messages, system_prompt=system_prompt, temperature=temperature, max_tokens=max_tokens, model_preferences=model_preferences, tools=sampling_tools, tool_choice=tool_choice, mask_error_details=mask_error_details, tool_concurrency=tool_concurrency, ) # Check for final_response tool call for structured output had_final_response = False if result_type is not None and result_type is not str and step.is_tool_use: for tool_call in step.tool_calls: if tool_call.name == "final_response": had_final_response = True # Validate and return the structured result type_adapter = get_cached_typeadapter(result_type) # Unwrap if we wrapped primitives (non-object schemas) input_data = tool_call.input original_schema = compress_schema( type_adapter.json_schema(), prune_titles=True ) if ( original_schema.get("type") != "object" and isinstance(input_data, dict) and "value" in input_data ): input_data = input_data["value"] try: validated_result = type_adapter.validate_python(input_data) text = json.dumps( type_adapter.dump_python(validated_result, mode="json") ) return SamplingResult( text=text, result=validated_result, history=step.history, ) except ValidationError as e: consecutive_validation_failures += 1 if consecutive_validation_failures > _MAX_VALIDATION_RETRIES: raise RuntimeError( f"Structured output validation failed " f"{consecutive_validation_failures} consecutive " f"times for type {result_type.__name__}: {e}" ) from e # Validation failed - add error as tool result step.history.append( SamplingMessage( role="user", content=[ ToolResultContent( type="tool_result", toolUseId=tool_call.id, content=[ TextContent( type="text", text=( f"Validation error: {e}. " "Please try again with valid data." ), ) ], isError=True, ) ], ) ) # The LLM called tools but not final_response — reset validation counter if not had_final_response: consecutive_validation_failures = 0 # If not a tool use response, we're done if not step.is_tool_use: # For structured output, the LLM must use the final_response tool if result_type is not None and result_type is not str: text_response_retries += 1 if text_response_retries > _MAX_TEXT_RESPONSE_RETRIES: raise RuntimeError( f"Expected structured output of type {result_type.__name__}, " "but the LLM returned a text response instead of calling " f"the final_response tool ({text_response_retries} attempts)." ) # Nudge the LLM to use the tool step.history.append( SamplingMessage( role="user", content=TextContent( type="text", text=( "You must call the `final_response` tool to provide " "your answer. Do not respond with text — use the tool." ), ), ) ) current_messages = step.history continue return SamplingResult( text=step.text, result=cast(ResultT, step.text if step.text else ""), history=step.history, ) # Continue with the updated history current_messages = step.history # After first iteration, reset tool_choice to auto (unless structured output is required) if result_type is None or result_type is str: tool_choice = None raise RuntimeError(f"Sampling exceeded maximum iterations ({max_iterations})")