mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-09 15:19:10 +02:00
f018f68bbf* Expose telemetry attributes on span start 🤖 Generated with Codex * Expose sampling attributes on span start 🤖 Generated with Codex --------- Co-authored-by: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
811 lines
31 KiB
Python
811 lines
31 KiB
Python
"""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.stop_reason == "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, *, client_available: bool = True
|
|
) -> 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.
|
|
client_available: Whether the client back-channel can be reached at all.
|
|
On modern (2026-07-28) connections the server-initiated createMessage
|
|
back-channel was removed (SEP-2577), so the client can never serve a
|
|
sampling request; pass False there to force the server-side handler
|
|
path (``"fallback"`` behaves like ``"always"`` when a handler exists).
|
|
|
|
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. On connections without a
|
|
# back-channel the client can never serve the request regardless of the
|
|
# capabilities it advertised, so treat both as unavailable.
|
|
has_sampling = client_available and session.check_client_capability(
|
|
capability=ClientCapabilities(sampling=SamplingCapability())
|
|
)
|
|
has_tools_capability = client_available and 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(
|
|
system_prompt=system_prompt,
|
|
messages=messages,
|
|
temperature=temperature,
|
|
max_tokens=max_tokens,
|
|
model_preferences=_parse_model_preferences(model_preferences),
|
|
tools=sdk_tools,
|
|
tool_choice=tool_choice,
|
|
),
|
|
# SamplingHandler is typed against the SDK's RequestContext placeholder,
|
|
# but FastMCP hands handlers its own FastMCPRequestContext wrapper at
|
|
# runtime; the two aren't structurally related in the type system.
|
|
context.request_context, # ty: ignore[invalid-argument-type]
|
|
)
|
|
|
|
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",
|
|
stop_reason="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",
|
|
tool_use_id=tool_use.id,
|
|
content=[
|
|
TextContent(
|
|
type="text",
|
|
text=f"Error: Unknown tool '{tool_use.name}'",
|
|
)
|
|
],
|
|
is_error=True,
|
|
)
|
|
|
|
tracer = get_tracer()
|
|
with tracer.start_as_current_span(
|
|
f"sampling tool {tool_use.name}",
|
|
kind=SpanKind.INTERNAL,
|
|
attributes={
|
|
"gen_ai.tool.name": tool_use.name,
|
|
"fastmcp.tool.use_id": tool_use.id,
|
|
},
|
|
) as span:
|
|
try:
|
|
result_value = await tool.run(tool_use.input)
|
|
return ToolResultContent(
|
|
type="tool_result",
|
|
tool_use_id=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",
|
|
tool_use_id=tool_use.id,
|
|
content=[TextContent(type="text", text=str(e))],
|
|
is_error=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",
|
|
tool_use_id=tool_use.id,
|
|
content=[TextContent(type="text", text=error_text)],
|
|
is_error=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": <schema>}
|
|
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,
|
|
client_available: bool = True,
|
|
) -> 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.
|
|
|
|
When ``client_available`` is False (e.g. a modern 2026-07-28 connection with
|
|
no back-channel), the client is never used and a configured sampling handler
|
|
serves the request; the caller is responsible for raising a clear era error
|
|
when no handler can serve it.
|
|
"""
|
|
# 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), client_available=client_available
|
|
)
|
|
|
|
# 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,
|
|
attributes={
|
|
"mcp.method.name": "sampling/createMessage",
|
|
"fastmcp.server.name": context.fastmcp.name,
|
|
},
|
|
record_exception=False,
|
|
set_status_on_exception=False,
|
|
) as span:
|
|
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:
|
|
# Deprecated upstream in SDK v2 but deliberately kept per compat
|
|
# directive; removed with the multi-round-trip follow-up.
|
|
response = await context.session.create_message( # ty: ignore[deprecated]
|
|
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.stop_reason == "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,
|
|
client_available: bool = True,
|
|
) -> 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.
|
|
|
|
When ``client_available`` is False (e.g. a modern 2026-07-28 connection with
|
|
no back-channel), the client is never used and a configured sampling handler
|
|
serves the request; the caller is responsible for raising a clear era error
|
|
when no handler can serve it.
|
|
"""
|
|
# 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,
|
|
client_available=client_available,
|
|
)
|
|
|
|
# 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",
|
|
tool_use_id=tool_call.id,
|
|
content=[
|
|
TextContent(
|
|
type="text",
|
|
text=(
|
|
f"Validation error: {e}. "
|
|
"Please try again with valid data."
|
|
),
|
|
)
|
|
],
|
|
is_error=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})")
|