feat: Sampling (#170)

* Removed unused _setup_client

* Implement sampling in client and connectors

* Tests

* Nicer names in action

* Update previous tests

* Docstrings

* Fixed naming and created other primitives tests

* Supported features table

* docs

Author: pietrozullo

.github/workflows/testprimitives.yml

@@ -0,0 +1,41 @@
+name: Primitives
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  primitives:
+    name: ${{ matrix.primitive }}
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        primitive: [sampling, tools, resources, prompts]
+        python-version: ["3.11"]
+
+    steps:
+      - uses: actions/checkout@v3
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install uv
+        run: |
+          pip install uv
+
+      - name: Install dependencies
+        run: |
+          # Install project with dev and optional extras similar to unit tests environment
+          uv pip install --system .[dev,anthropic,openai,search,e2b]
+
+      - name: Lint with ruff
+        run: |
+          ruff check .
+
+      - name: Run integration tests for ${{ matrix.primitive }} primitive
+        run: |
+          pytest tests/integration/primitives/test_${{ matrix.primitive }}.py

README.md

@@ -44,7 +44,12 @@
 
 💬 Get started quickly - chat with your servers on our <b>hosted version</b>! [Try mcp-use chat (beta)](https://chat.mcp-use.com).
 
-# Features
+| Supports  | |
+| :--- | :--- |
+| **Primitives** | [![Tools](https://img.shields.io/github/actions/workflow/status/pietrozullo/mcp-use/testprimitives.yml?job=tools&label=Tools&style=flat)](https://github.com/pietrozullo/mcp-use/actions/workflows/testprimitives.yml) [![Resources](https://img.shields.io/github/actions/workflow/status/pietrozullo/mcp-use/testprimitives.yml?job=resources&label=Resources&style=flat)](https://github.com/pietrozullo/mcp-use/actions/workflows/testprimitives.yml) [![Prompts](https://img.shields.io/github/actions/workflow/status/pietrozullo/mcp-use/testprimitives.yml?job=prompts&label=Prompts&style=flat)](https://github.com/pietrozullo/mcp-use/actions/workflows/testprimitives.yml) [![Sampling](https://img.shields.io/github/actions/workflow/status/pietrozullo/mcp-use/testprimitives.yml?job=sampling&label=Sampling&style=flat)](https://github.com/pietrozullo/mcp-use/actions/workflows/testprimitives.yml) |
+| **Transports** | [![Stdio](https://img.shields.io/github/actions/workflow/status/pietrozullo/mcp-use/transportstests.yml?job=stdio&label=Stdio&style=flat)](https://github.com/pietrozullo/mcp-use/actions/workflows/transportstests.yml) [![SSE](https://img.shields.io/github/actions/workflow/status/pietrozullo/mcp-use/transportstests.yml?job=sse&label=SSE&style=flat)](https://github.com/pietrozullo/mcp-use/actions/workflows/transportstests.yml) [![Streamable HTTP](https://img.shields.io/github/actions/workflow/status/pietrozullo/mcp-use/transportstests.yml?job=streamableHttp&label=Streamable%20HTTP&style=flat)](https://github.com/pietrozullo/mcp-use/actions/workflows/transportstests.yml) |
+
+## Features
 
 <table>
   <tr>

docs/docs.json

@@ -30,7 +30,14 @@
             "pages": [
               "essentials/configuration",
               "essentials/client-configuration",
-              "essentials/connection-types"
+              {
+                "group": "Client Features",
+                "icon": "list-plus",
+                "pages": [
+                  "essentials/connection-types",
+                  "essentials/sampling"
+                ]
+              }
             ]
           },
           {

docs/essentials/sampling.mdx

@@ -0,0 +1,86 @@
+---
+title: "Sampling"
+description: "Enable LLM sampling capabilities for MCP tools"
+icon: "pipette"
+---
+
+# Sampling
+
+<Info>
+Sampling allows MCP tools to request LLM completions during their execution.
+</Info>
+
+## Configuration
+
+To enable sampling, provide a `sampling_callback` function when initializing the MCPClient:
+
+```python
+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
+)
+```
+
+
+## Creating Sampling-Enabled Tools
+
+When building MCP servers, tools can request sampling using the context parameter:
+
+```python
+from fastmcp import Context, FastMCP
+
+mcp = FastMCP(name="MyServer")
+
+@mcp.tool
+async def analyze_sentiment(text: str, ctx: Context) -> str:
+    """Analyze the sentiment of text using the client's LLM."""
+    prompt = f"""Analyze the sentiment of the following text as positive, negative, or neutral.
+    Just output a single word - 'positive', 'negative', or 'neutral'.
+
+    Text to analyze: {text}"""
+
+    # Request LLM analysis through sampling
+    response = await ctx.sample(prompt)
+
+    return response.text.strip()
+```
+
+## Error Handling
+
+If no sampling callback is provided but a tool requests sampling:
+
+```python
+# 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
+```

docs/essentials/server-manager.mdx

@@ -1,8 +1,7 @@
-
 ---
 title: 'Server Manager'
 description: 'Intelligent management of multiple MCP servers and tool discovery'
-icon: 'server'
+icon: 'server-cog'
 ---
 
 # Server Manager

mcp_use/client.py

@@ -9,6 +9,8 @@
 import warnings
 from typing import Any
 
+from mcp.client.session import SamplingFnT
+
 from mcp_use.types.sandbox import SandboxOptions
 
 from .config import create_connector_from_config, load_config_file
@@ -28,6 +30,7 @@ def __init__(
         config: str | dict[str, Any] | None = None,
         sandbox: bool = False,
         sandbox_options: SandboxOptions | None = None,
+        sampling_callback: SamplingFnT | None = None,
     ) -> None:
         """Initialize a new MCP client.
 
@@ -36,13 +39,14 @@ def __init__(
                    If None, an empty configuration is used.
             sandbox: Whether to use sandboxed execution mode for running MCP servers.
             sandbox_options: Optional sandbox configuration options.
+            sampling_callback: Optional sampling callback function.
         """
         self.config: dict[str, Any] = {}
         self.sandbox = sandbox
         self.sandbox_options = sandbox_options
         self.sessions: dict[str, MCPSession] = {}
         self.active_sessions: list[str] = []
-
+        self.sampling_callback = sampling_callback
         # Load configuration if provided
         if config is not None:
             if isinstance(config, str):
@@ -151,7 +155,10 @@ async def create_session(self, server_name: str, auto_initialize: bool = True) -
 
         # Create connector with options
         connector = create_connector_from_config(
-            server_config, sandbox=self.sandbox, sandbox_options=self.sandbox_options
+            server_config,
+            sandbox=self.sandbox,
+            sandbox_options=self.sandbox_options,
+            sampling_callback=self.sampling_callback,
         )
 
         # Create the session

mcp_use/config.py

@@ -7,6 +7,8 @@
 import json
 from typing import Any
 
+from mcp.client.session import SamplingFnT
+
 from mcp_use.types.sandbox import SandboxOptions
 
 from .connectors import (
@@ -36,6 +38,7 @@ def create_connector_from_config(
     server_config: dict[str, Any],
     sandbox: bool = False,
     sandbox_options: SandboxOptions | None = None,
+    sampling_callback: SamplingFnT | None = None,
 ) -> BaseConnector:
     """Create a connector based on server configuration.
     This function can be called with just the server_config parameter:
@@ -44,7 +47,7 @@ def create_connector_from_config(
         server_config: The server configuration section
         sandbox: Whether to use sandboxed execution mode for running MCP servers.
         sandbox_options: Optional sandbox configuration options.
-
+        sampling_callback: Optional sampling callback function.
     Returns:
         A configured connector instance
     """
@@ -55,6 +58,7 @@ def create_connector_from_config(
             command=server_config["command"],
             args=server_config["args"],
             env=server_config.get("env", None),
+            sampling_callback=sampling_callback,
         )
 
     # Sandboxed connector
@@ -64,6 +68,7 @@ def create_connector_from_config(
             args=server_config["args"],
             env=server_config.get("env", None),
             e2b_options=sandbox_options,
+            sampling_callback=sampling_callback,
         )
 
     # HTTP connector
@@ -72,6 +77,7 @@ def create_connector_from_config(
             base_url=server_config["url"],
             headers=server_config.get("headers", None),
             auth_token=server_config.get("auth_token", None),
+            sampling_callback=sampling_callback,
         )
 
     # WebSocket connector

mcp_use/connectors/base.py

@@ -8,11 +8,14 @@
 from abc import ABC, abstractmethod
 from typing import Any
 
-from mcp import ClientSession
+from mcp import ClientSession, Implementation
+from mcp.client.session import SamplingFnT
 from mcp.shared.exceptions import McpError
 from mcp.types import CallToolResult, GetPromptResult, Prompt, ReadResourceResult, Resource, Tool
 from pydantic import AnyUrl
 
+import mcp_use
+
 from ..logging import logger
 from ..task_managers import ConnectionManager
 
@@ -23,7 +26,7 @@ class BaseConnector(ABC):
     This class defines the interface that all MCP connectors must implement.
     """
 
-    def __init__(self):
+    def __init__(self, sampling_callback: SamplingFnT | None = None):
         """Initialize base connector with common attributes."""
         self.client_session: ClientSession | None = None
         self._connection_manager: ConnectionManager | None = None
@@ -33,6 +36,16 @@ def __init__(self):
         self._connected = False
         self._initialized = False  # Track if client_session.initialize() has been called
         self.auto_reconnect = True  # Whether to automatically reconnect on connection loss (not configurable for now)
+        self.sampling_callback = sampling_callback
+
+    @property
+    def client_info(self) -> Implementation:
+        """Get the client info for the connector."""
+        return Implementation(
+            name="mcp-use",
+            version=mcp_use.__version__,
+            url="https://github.com/mcp-use/mcp-use",
+        )
 
     @abstractmethod
     async def connect(self) -> None:
@@ -143,7 +156,8 @@ async def initialize(self) -> dict[str, Any]:
 
         logger.debug(
             f"MCP session initialized with {len(self._tools)} tools, "
-            "{len(self._resources)} resources, and {len(self._prompts)} prompts"
+            f"{len(self._resources)} resources, "
+            f"and {len(self._prompts)} prompts"
         )
 
         return result

mcp_use/connectors/http.py

@@ -7,9 +7,10 @@
 
 import httpx
 from mcp import ClientSession
+from mcp.client.session import SamplingFnT
 
 from ..logging import logger
-from ..task_managers import ConnectionManager, SseConnectionManager, StreamableHttpConnectionManager
+from ..task_managers import SseConnectionManager, StreamableHttpConnectionManager
 from .base import BaseConnector
 
 
@@ -27,6 +28,7 @@ def __init__(
         headers: dict[str, str] | None = None,
         timeout: float = 5,
         sse_read_timeout: float = 60 * 5,
+        sampling_callback: SamplingFnT | None = None,
     ):
         """Initialize a new HTTP connector.
 
@@ -36,8 +38,9 @@ def __init__(
             headers: Optional additional headers.
             timeout: Timeout for HTTP operations in seconds.
             sse_read_timeout: Timeout for SSE read operations in seconds.
+            sampling_callback: Optional sampling callback.
         """
-        super().__init__()
+        super().__init__(sampling_callback=sampling_callback)
         self.base_url = base_url.rstrip("/")
         self.auth_token = auth_token
         self.headers = headers or {}
@@ -46,14 +49,6 @@ def __init__(
         self.timeout = timeout
         self.sse_read_timeout = sse_read_timeout
 
-    async def _setup_client(self, connection_manager: ConnectionManager) -> None:
-        """Set up the client session with the provided connection manager."""
-
-        self._connection_manager = connection_manager
-        read_stream, write_stream = await self._connection_manager.start()
-        self.client_session = ClientSession(read_stream, write_stream, sampling_callback=None)
-        await self.client_session.__aenter__()
-
     async def connect(self) -> None:
         """Establish a connection to the MCP implementation."""
         if self._connected:
@@ -76,7 +71,9 @@ async def connect(self) -> None:
             read_stream, write_stream = await connection_manager.start()
 
             # Test if this actually works by trying to create a client session and initialize it
-            test_client = ClientSession(read_stream, write_stream, sampling_callback=None)
+            test_client = ClientSession(
+                read_stream, write_stream, sampling_callback=self.sampling_callback, client_info=self.client_info
+            )
             await test_client.__aenter__()
 
             try:
@@ -154,7 +151,12 @@ async def connect(self) -> None:
                     read_stream, write_stream = await connection_manager.start()
 
                     # Create the client session for SSE
-                    self.client_session = ClientSession(read_stream, write_stream, sampling_callback=None)
+                    self.client_session = ClientSession(
+                        read_stream,
+                        write_stream,
+                        sampling_callback=self.sampling_callback,
+                        client_info=self.client_info,
+                    )
                     await self.client_session.__aenter__()
                     self.transport_type = "SSE"