Sampling

from mcp_use.client import MCPClient
from mcp.client.session import ClientSession
from mcp.types import (
    CreateMessageRequestParams,
    CreateMessageResult,
    ErrorData,
    TextContent
)

async def sampling_callback(
    context: ClientSession,
    params: CreateMessageRequestParams
) -> CreateMessageResult | ErrorData:
    """
    Your sampling callback implementation.
    This function receives a prompt and returns an LLM response.
    """
    # Integrate with your LLM of choice (OpenAI, Anthropic, etc.)
    response = await your_llm.complete(params.messages[-1].content.text)

    return CreateMessageResult(
        content=TextContent(text=response, type="text"),
        model="your-model-name",
        role="assistant"
    )

# Initialize client with sampling support
client = MCPClient(
    config="config.json",
    sampling_callback=sampling_callback
)

# Without sampling callback
client = MCPClient(config="config.json")  # No sampling_callback

# Tool that requires sampling will return an error
result = await session.call_tool("analyze_sentiment", {"text": "Hello"})
# result.isError will be True