Add support for Elicitation (#625)
Co-authored-by: David Soria Parra <davidsp@anthropic.com>
This commit is contained in:
parent
a3bcabdce2
commit
69e6572d48
9 changed files with 643 additions and 8 deletions
48
README.md
48
README.md
|
|
@ -31,6 +31,8 @@
|
|||
- [Images](#images)
|
||||
- [Context](#context)
|
||||
- [Completions](#completions)
|
||||
- [Elicitation](#elicitation)
|
||||
- [Authentication](#authentication)
|
||||
- [Running Your Server](#running-your-server)
|
||||
- [Development Mode](#development-mode)
|
||||
- [Claude Desktop Integration](#claude-desktop-integration)
|
||||
|
|
@ -74,7 +76,7 @@ The Model Context Protocol allows applications to provide context for LLMs in a
|
|||
|
||||
### Adding MCP to your python project
|
||||
|
||||
We recommend using [uv](https://docs.astral.sh/uv/) to manage your Python projects.
|
||||
We recommend using [uv](https://docs.astral.sh/uv/) to manage your Python projects.
|
||||
|
||||
If you haven't created a uv-managed project yet, create one:
|
||||
|
||||
|
|
@ -372,6 +374,50 @@ async def handle_completion(
|
|||
return Completion(values=filtered)
|
||||
return None
|
||||
```
|
||||
### Elicitation
|
||||
|
||||
Request additional information from users during tool execution:
|
||||
|
||||
```python
|
||||
from mcp.server.fastmcp import FastMCP, Context
|
||||
from mcp.server.elicitation import (
|
||||
AcceptedElicitation,
|
||||
DeclinedElicitation,
|
||||
CancelledElicitation,
|
||||
)
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
mcp = FastMCP("Booking System")
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def book_table(date: str, party_size: int, ctx: Context) -> str:
|
||||
"""Book a table with confirmation"""
|
||||
|
||||
# Schema must only contain primitive types (str, int, float, bool)
|
||||
class ConfirmBooking(BaseModel):
|
||||
confirm: bool = Field(description="Confirm booking?")
|
||||
notes: str = Field(default="", description="Special requests")
|
||||
|
||||
result = await ctx.elicit(
|
||||
message=f"Confirm booking for {party_size} on {date}?", schema=ConfirmBooking
|
||||
)
|
||||
|
||||
match result:
|
||||
case AcceptedElicitation(data=data):
|
||||
if data.confirm:
|
||||
return f"Booked! Notes: {data.notes or 'None'}"
|
||||
return "Booking cancelled"
|
||||
case DeclinedElicitation():
|
||||
return "Booking declined"
|
||||
case CancelledElicitation():
|
||||
return "Booking cancelled"
|
||||
```
|
||||
|
||||
The `elicit()` method returns an `ElicitationResult` with:
|
||||
- `action`: "accept", "decline", or "cancel"
|
||||
- `data`: The validated response (only when accepted)
|
||||
- `validation_error`: Any validation error message
|
||||
|
||||
### Authentication
|
||||
|
||||
|
|
|
|||
|
|
@ -22,6 +22,14 @@ class SamplingFnT(Protocol):
|
|||
) -> types.CreateMessageResult | types.ErrorData: ...
|
||||
|
||||
|
||||
class ElicitationFnT(Protocol):
|
||||
async def __call__(
|
||||
self,
|
||||
context: RequestContext["ClientSession", Any],
|
||||
params: types.ElicitRequestParams,
|
||||
) -> types.ElicitResult | types.ErrorData: ...
|
||||
|
||||
|
||||
class ListRootsFnT(Protocol):
|
||||
async def __call__(
|
||||
self, context: RequestContext["ClientSession", Any]
|
||||
|
|
@ -58,6 +66,16 @@ async def _default_sampling_callback(
|
|||
)
|
||||
|
||||
|
||||
async def _default_elicitation_callback(
|
||||
context: RequestContext["ClientSession", Any],
|
||||
params: types.ElicitRequestParams,
|
||||
) -> types.ElicitResult | types.ErrorData:
|
||||
return types.ErrorData(
|
||||
code=types.INVALID_REQUEST,
|
||||
message="Elicitation not supported",
|
||||
)
|
||||
|
||||
|
||||
async def _default_list_roots_callback(
|
||||
context: RequestContext["ClientSession", Any],
|
||||
) -> types.ListRootsResult | types.ErrorData:
|
||||
|
|
@ -91,6 +109,7 @@ class ClientSession(
|
|||
write_stream: MemoryObjectSendStream[SessionMessage],
|
||||
read_timeout_seconds: timedelta | None = None,
|
||||
sampling_callback: SamplingFnT | None = None,
|
||||
elicitation_callback: ElicitationFnT | None = None,
|
||||
list_roots_callback: ListRootsFnT | None = None,
|
||||
logging_callback: LoggingFnT | None = None,
|
||||
message_handler: MessageHandlerFnT | None = None,
|
||||
|
|
@ -105,12 +124,16 @@ class ClientSession(
|
|||
)
|
||||
self._client_info = client_info or DEFAULT_CLIENT_INFO
|
||||
self._sampling_callback = sampling_callback or _default_sampling_callback
|
||||
self._elicitation_callback = elicitation_callback or _default_elicitation_callback
|
||||
self._list_roots_callback = list_roots_callback or _default_list_roots_callback
|
||||
self._logging_callback = logging_callback or _default_logging_callback
|
||||
self._message_handler = message_handler or _default_message_handler
|
||||
|
||||
async def initialize(self) -> types.InitializeResult:
|
||||
sampling = types.SamplingCapability() if self._sampling_callback is not _default_sampling_callback else None
|
||||
elicitation = (
|
||||
types.ElicitationCapability() if self._elicitation_callback is not _default_elicitation_callback else None
|
||||
)
|
||||
roots = (
|
||||
# TODO: Should this be based on whether we
|
||||
# _will_ send notifications, or only whether
|
||||
|
|
@ -128,6 +151,7 @@ class ClientSession(
|
|||
protocolVersion=types.LATEST_PROTOCOL_VERSION,
|
||||
capabilities=types.ClientCapabilities(
|
||||
sampling=sampling,
|
||||
elicitation=elicitation,
|
||||
experimental=None,
|
||||
roots=roots,
|
||||
),
|
||||
|
|
@ -362,6 +386,12 @@ class ClientSession(
|
|||
client_response = ClientResponse.validate_python(response)
|
||||
await responder.respond(client_response)
|
||||
|
||||
case types.ElicitRequest(params=params):
|
||||
with responder:
|
||||
response = await self._elicitation_callback(ctx, params)
|
||||
client_response = ClientResponse.validate_python(response)
|
||||
await responder.respond(client_response)
|
||||
|
||||
case types.ListRootsRequest():
|
||||
with responder:
|
||||
response = await self._list_roots_callback(ctx)
|
||||
|
|
|
|||
111
src/mcp/server/elicitation.py
Normal file
111
src/mcp/server/elicitation.py
Normal file
|
|
@ -0,0 +1,111 @@
|
|||
"""Elicitation utilities for MCP servers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import types
|
||||
from typing import Generic, Literal, TypeVar, Union, get_args, get_origin
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from mcp.server.session import ServerSession
|
||||
from mcp.types import RequestId
|
||||
|
||||
ElicitSchemaModelT = TypeVar("ElicitSchemaModelT", bound=BaseModel)
|
||||
|
||||
|
||||
class AcceptedElicitation(BaseModel, Generic[ElicitSchemaModelT]):
|
||||
"""Result when user accepts the elicitation."""
|
||||
|
||||
action: Literal["accept"] = "accept"
|
||||
data: ElicitSchemaModelT
|
||||
|
||||
|
||||
class DeclinedElicitation(BaseModel):
|
||||
"""Result when user declines the elicitation."""
|
||||
|
||||
action: Literal["decline"] = "decline"
|
||||
|
||||
|
||||
class CancelledElicitation(BaseModel):
|
||||
"""Result when user cancels the elicitation."""
|
||||
|
||||
action: Literal["cancel"] = "cancel"
|
||||
|
||||
|
||||
ElicitationResult = AcceptedElicitation[ElicitSchemaModelT] | DeclinedElicitation | CancelledElicitation
|
||||
|
||||
|
||||
# Primitive types allowed in elicitation schemas
|
||||
_ELICITATION_PRIMITIVE_TYPES = (str, int, float, bool)
|
||||
|
||||
|
||||
def _validate_elicitation_schema(schema: type[BaseModel]) -> None:
|
||||
"""Validate that a Pydantic model only contains primitive field types."""
|
||||
for field_name, field_info in schema.model_fields.items():
|
||||
if not _is_primitive_field(field_info):
|
||||
raise TypeError(
|
||||
f"Elicitation schema field '{field_name}' must be a primitive type "
|
||||
f"{_ELICITATION_PRIMITIVE_TYPES} or Optional of these types. "
|
||||
f"Complex types like lists, dicts, or nested models are not allowed."
|
||||
)
|
||||
|
||||
|
||||
def _is_primitive_field(field_info: FieldInfo) -> bool:
|
||||
"""Check if a field is a primitive type allowed in elicitation schemas."""
|
||||
annotation = field_info.annotation
|
||||
|
||||
# Handle None type
|
||||
if annotation is types.NoneType:
|
||||
return True
|
||||
|
||||
# Handle basic primitive types
|
||||
if annotation in _ELICITATION_PRIMITIVE_TYPES:
|
||||
return True
|
||||
|
||||
# Handle Union types
|
||||
origin = get_origin(annotation)
|
||||
if origin is Union or origin is types.UnionType:
|
||||
args = get_args(annotation)
|
||||
# All args must be primitive types or None
|
||||
return all(arg is types.NoneType or arg in _ELICITATION_PRIMITIVE_TYPES for arg in args)
|
||||
|
||||
return False
|
||||
|
||||
|
||||
async def elicit_with_validation(
|
||||
session: ServerSession,
|
||||
message: str,
|
||||
schema: type[ElicitSchemaModelT],
|
||||
related_request_id: RequestId | None = None,
|
||||
) -> ElicitationResult[ElicitSchemaModelT]:
|
||||
"""Elicit information from the client/user with schema validation.
|
||||
|
||||
This method can be used to interactively ask for additional information from the
|
||||
client within a tool's execution. The client might display the message to the
|
||||
user and collect a response according to the provided schema. Or in case a
|
||||
client is an agent, it might decide how to handle the elicitation -- either by asking
|
||||
the user or automatically generating a response.
|
||||
"""
|
||||
# Validate that schema only contains primitive types and fail loudly if not
|
||||
_validate_elicitation_schema(schema)
|
||||
|
||||
json_schema = schema.model_json_schema()
|
||||
|
||||
result = await session.elicit(
|
||||
message=message,
|
||||
requestedSchema=json_schema,
|
||||
related_request_id=related_request_id,
|
||||
)
|
||||
|
||||
if result.action == "accept" and result.content:
|
||||
# Validate and parse the content using the schema
|
||||
validated_data = schema.model_validate(result.content)
|
||||
return AcceptedElicitation(data=validated_data)
|
||||
elif result.action == "decline":
|
||||
return DeclinedElicitation()
|
||||
elif result.action == "cancel":
|
||||
return CancelledElicitation()
|
||||
else:
|
||||
# This should never happen, but handle it just in case
|
||||
raise ValueError(f"Unexpected elicitation action: {result.action}")
|
||||
|
|
@ -34,6 +34,7 @@ from mcp.server.auth.provider import OAuthAuthorizationServerProvider
|
|||
from mcp.server.auth.settings import (
|
||||
AuthSettings,
|
||||
)
|
||||
from mcp.server.elicitation import ElicitationResult, ElicitSchemaModelT, elicit_with_validation
|
||||
from mcp.server.fastmcp.exceptions import ResourceError
|
||||
from mcp.server.fastmcp.prompts import Prompt, PromptManager
|
||||
from mcp.server.fastmcp.resources import FunctionResource, Resource, ResourceManager
|
||||
|
|
@ -972,6 +973,37 @@ class Context(BaseModel, Generic[ServerSessionT, LifespanContextT, RequestT]):
|
|||
assert self._fastmcp is not None, "Context is not available outside of a request"
|
||||
return await self._fastmcp.read_resource(uri)
|
||||
|
||||
async def elicit(
|
||||
self,
|
||||
message: str,
|
||||
schema: type[ElicitSchemaModelT],
|
||||
) -> ElicitationResult[ElicitSchemaModelT]:
|
||||
"""Elicit information from the client/user.
|
||||
|
||||
This method can be used to interactively ask for additional information from the
|
||||
client within a tool's execution. The client might display the message to the
|
||||
user and collect a response according to the provided schema. Or in case a
|
||||
client is an agent, it might decide how to handle the elicitation -- either by asking
|
||||
the user or automatically generating a response.
|
||||
|
||||
Args:
|
||||
schema: A Pydantic model class defining the expected response structure, according to the specification,
|
||||
only primive types are allowed.
|
||||
message: Optional message to present to the user. If not provided, will use
|
||||
a default message based on the schema
|
||||
|
||||
Returns:
|
||||
An ElicitationResult containing the action taken and the data if accepted
|
||||
|
||||
Note:
|
||||
Check the result.action to determine if the user accepted, declined, or cancelled.
|
||||
The result.data will only be populated if action is "accept" and validation succeeded.
|
||||
"""
|
||||
|
||||
return await elicit_with_validation(
|
||||
session=self.request_context.session, message=message, schema=schema, related_request_id=self.request_id
|
||||
)
|
||||
|
||||
async def log(
|
||||
self,
|
||||
level: Literal["debug", "info", "warning", "error"],
|
||||
|
|
|
|||
|
|
@ -121,6 +121,10 @@ class ServerSession(
|
|||
if client_caps.sampling is None:
|
||||
return False
|
||||
|
||||
if capability.elicitation is not None:
|
||||
if client_caps.elicitation is None:
|
||||
return False
|
||||
|
||||
if capability.experimental is not None:
|
||||
if client_caps.experimental is None:
|
||||
return False
|
||||
|
|
@ -251,6 +255,35 @@ class ServerSession(
|
|||
types.ListRootsResult,
|
||||
)
|
||||
|
||||
async def elicit(
|
||||
self,
|
||||
message: str,
|
||||
requestedSchema: types.ElicitRequestedSchema,
|
||||
related_request_id: types.RequestId | None = None,
|
||||
) -> types.ElicitResult:
|
||||
"""Send an elicitation/create request.
|
||||
|
||||
Args:
|
||||
message: The message to present to the user
|
||||
requestedSchema: Schema defining the expected response structure
|
||||
|
||||
Returns:
|
||||
The client's response
|
||||
"""
|
||||
return await self.send_request(
|
||||
types.ServerRequest(
|
||||
types.ElicitRequest(
|
||||
method="elicitation/create",
|
||||
params=types.ElicitRequestParams(
|
||||
message=message,
|
||||
requestedSchema=requestedSchema,
|
||||
),
|
||||
)
|
||||
),
|
||||
types.ElicitResult,
|
||||
metadata=ServerMessageMetadata(related_request_id=related_request_id),
|
||||
)
|
||||
|
||||
async def send_ping(self) -> types.EmptyResult:
|
||||
"""Send a ping request."""
|
||||
return await self.send_request(
|
||||
|
|
|
|||
|
|
@ -11,7 +11,14 @@ import anyio
|
|||
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
||||
|
||||
import mcp.types as types
|
||||
from mcp.client.session import ClientSession, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
|
||||
from mcp.client.session import (
|
||||
ClientSession,
|
||||
ElicitationFnT,
|
||||
ListRootsFnT,
|
||||
LoggingFnT,
|
||||
MessageHandlerFnT,
|
||||
SamplingFnT,
|
||||
)
|
||||
from mcp.server import Server
|
||||
from mcp.shared.message import SessionMessage
|
||||
|
||||
|
|
@ -53,6 +60,7 @@ async def create_connected_server_and_client_session(
|
|||
message_handler: MessageHandlerFnT | None = None,
|
||||
client_info: types.Implementation | None = None,
|
||||
raise_exceptions: bool = False,
|
||||
elicitation_callback: ElicitationFnT | None = None,
|
||||
) -> AsyncGenerator[ClientSession, None]:
|
||||
"""Creates a ClientSession that is connected to a running MCP server."""
|
||||
async with create_client_server_memory_streams() as (
|
||||
|
|
@ -83,6 +91,7 @@ async def create_connected_server_and_client_session(
|
|||
logging_callback=logging_callback,
|
||||
message_handler=message_handler,
|
||||
client_info=client_info,
|
||||
elicitation_callback=elicitation_callback,
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
yield client_session
|
||||
|
|
|
|||
|
|
@ -216,6 +216,12 @@ class SamplingCapability(BaseModel):
|
|||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
|
||||
class ElicitationCapability(BaseModel):
|
||||
"""Capability for elicitation operations."""
|
||||
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
|
||||
class ClientCapabilities(BaseModel):
|
||||
"""Capabilities a client may support."""
|
||||
|
||||
|
|
@ -223,6 +229,8 @@ class ClientCapabilities(BaseModel):
|
|||
"""Experimental, non-standard capabilities that the client supports."""
|
||||
sampling: SamplingCapability | None = None
|
||||
"""Present if the client supports sampling from an LLM."""
|
||||
elicitation: ElicitationCapability | None = None
|
||||
"""Present if the client supports elicitation from the user."""
|
||||
roots: RootsCapability | None = None
|
||||
"""Present if the client supports listing roots."""
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
|
@ -1186,11 +1194,49 @@ class ClientNotification(
|
|||
pass
|
||||
|
||||
|
||||
class ClientResult(RootModel[EmptyResult | CreateMessageResult | ListRootsResult]):
|
||||
# Type for elicitation schema - a JSON Schema dict
|
||||
ElicitRequestedSchema: TypeAlias = dict[str, Any]
|
||||
"""Schema for elicitation requests."""
|
||||
|
||||
|
||||
class ElicitRequestParams(RequestParams):
|
||||
"""Parameters for elicitation requests."""
|
||||
|
||||
message: str
|
||||
requestedSchema: ElicitRequestedSchema
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
|
||||
class ElicitRequest(Request[ElicitRequestParams, Literal["elicitation/create"]]):
|
||||
"""A request from the server to elicit information from the client."""
|
||||
|
||||
method: Literal["elicitation/create"]
|
||||
params: ElicitRequestParams
|
||||
|
||||
|
||||
class ElicitResult(Result):
|
||||
"""The client's response to an elicitation request."""
|
||||
|
||||
action: Literal["accept", "decline", "cancel"]
|
||||
"""
|
||||
The user action in response to the elicitation.
|
||||
- "accept": User submitted the form/confirmed the action
|
||||
- "decline": User explicitly declined the action
|
||||
- "cancel": User dismissed without making an explicit choice
|
||||
"""
|
||||
|
||||
content: dict[str, str | int | float | bool | None] | None = None
|
||||
"""
|
||||
The submitted form data, only present when action is "accept".
|
||||
Contains values matching the requested schema.
|
||||
"""
|
||||
|
||||
|
||||
class ClientResult(RootModel[EmptyResult | CreateMessageResult | ListRootsResult | ElicitResult]):
|
||||
pass
|
||||
|
||||
|
||||
class ServerRequest(RootModel[PingRequest | CreateMessageRequest | ListRootsRequest]):
|
||||
class ServerRequest(RootModel[PingRequest | CreateMessageRequest | ListRootsRequest | ElicitRequest]):
|
||||
pass
|
||||
|
||||
|
||||
|
|
|
|||
210
tests/server/fastmcp/test_elicitation.py
Normal file
210
tests/server/fastmcp/test_elicitation.py
Normal file
|
|
@ -0,0 +1,210 @@
|
|||
"""
|
||||
Test the elicitation feature using stdio transport.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from mcp.server.fastmcp import Context, FastMCP
|
||||
from mcp.shared.memory import create_connected_server_and_client_session
|
||||
from mcp.types import ElicitResult, TextContent
|
||||
|
||||
|
||||
# Shared schema for basic tests
|
||||
class AnswerSchema(BaseModel):
|
||||
answer: str = Field(description="The user's answer to the question")
|
||||
|
||||
|
||||
def create_ask_user_tool(mcp: FastMCP):
|
||||
"""Create a standard ask_user tool that handles all elicitation responses."""
|
||||
|
||||
@mcp.tool(description="A tool that uses elicitation")
|
||||
async def ask_user(prompt: str, ctx: Context) -> str:
|
||||
result = await ctx.elicit(
|
||||
message=f"Tool wants to ask: {prompt}",
|
||||
schema=AnswerSchema,
|
||||
)
|
||||
|
||||
if result.action == "accept" and result.data:
|
||||
return f"User answered: {result.data.answer}"
|
||||
elif result.action == "decline":
|
||||
return "User declined to answer"
|
||||
else:
|
||||
return "User cancelled"
|
||||
|
||||
return ask_user
|
||||
|
||||
|
||||
async def call_tool_and_assert(
|
||||
mcp: FastMCP,
|
||||
elicitation_callback,
|
||||
tool_name: str,
|
||||
args: dict,
|
||||
expected_text: str | None = None,
|
||||
text_contains: list[str] | None = None,
|
||||
):
|
||||
"""Helper to create session, call tool, and assert result."""
|
||||
async with create_connected_server_and_client_session(
|
||||
mcp._mcp_server, elicitation_callback=elicitation_callback
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
result = await client_session.call_tool(tool_name, args)
|
||||
assert len(result.content) == 1
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
|
||||
if expected_text is not None:
|
||||
assert result.content[0].text == expected_text
|
||||
elif text_contains is not None:
|
||||
for substring in text_contains:
|
||||
assert substring in result.content[0].text
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_stdio_elicitation():
|
||||
"""Test the elicitation feature using stdio transport."""
|
||||
mcp = FastMCP(name="StdioElicitationServer")
|
||||
create_ask_user_tool(mcp)
|
||||
|
||||
# Create a custom handler for elicitation requests
|
||||
async def elicitation_callback(context, params):
|
||||
if params.message == "Tool wants to ask: What is your name?":
|
||||
return ElicitResult(action="accept", content={"answer": "Test User"})
|
||||
else:
|
||||
raise ValueError(f"Unexpected elicitation message: {params.message}")
|
||||
|
||||
await call_tool_and_assert(
|
||||
mcp, elicitation_callback, "ask_user", {"prompt": "What is your name?"}, "User answered: Test User"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_stdio_elicitation_decline():
|
||||
"""Test elicitation with user declining."""
|
||||
mcp = FastMCP(name="StdioElicitationDeclineServer")
|
||||
create_ask_user_tool(mcp)
|
||||
|
||||
async def elicitation_callback(context, params):
|
||||
return ElicitResult(action="decline")
|
||||
|
||||
await call_tool_and_assert(
|
||||
mcp, elicitation_callback, "ask_user", {"prompt": "What is your name?"}, "User declined to answer"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_elicitation_schema_validation():
|
||||
"""Test that elicitation schemas must only contain primitive types."""
|
||||
mcp = FastMCP(name="ValidationTestServer")
|
||||
|
||||
def create_validation_tool(name: str, schema_class: type[BaseModel]):
|
||||
@mcp.tool(name=name, description=f"Tool testing {name}")
|
||||
async def tool(ctx: Context) -> str:
|
||||
try:
|
||||
await ctx.elicit(message="This should fail validation", schema=schema_class)
|
||||
return "Should not reach here"
|
||||
except TypeError as e:
|
||||
return f"Validation failed as expected: {str(e)}"
|
||||
|
||||
return tool
|
||||
|
||||
# Test cases for invalid schemas
|
||||
class InvalidListSchema(BaseModel):
|
||||
names: list[str] = Field(description="List of names")
|
||||
|
||||
class NestedModel(BaseModel):
|
||||
value: str
|
||||
|
||||
class InvalidNestedSchema(BaseModel):
|
||||
nested: NestedModel = Field(description="Nested model")
|
||||
|
||||
create_validation_tool("invalid_list", InvalidListSchema)
|
||||
create_validation_tool("nested_model", InvalidNestedSchema)
|
||||
|
||||
# Dummy callback (won't be called due to validation failure)
|
||||
async def elicitation_callback(context, params):
|
||||
return ElicitResult(action="accept", content={})
|
||||
|
||||
async with create_connected_server_and_client_session(
|
||||
mcp._mcp_server, elicitation_callback=elicitation_callback
|
||||
) as client_session:
|
||||
await client_session.initialize()
|
||||
|
||||
# Test both invalid schemas
|
||||
for tool_name, field_name in [("invalid_list", "names"), ("nested_model", "nested")]:
|
||||
result = await client_session.call_tool(tool_name, {})
|
||||
assert len(result.content) == 1
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert "Validation failed as expected" in result.content[0].text
|
||||
assert field_name in result.content[0].text
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_elicitation_with_optional_fields():
|
||||
"""Test that Optional fields work correctly in elicitation schemas."""
|
||||
mcp = FastMCP(name="OptionalFieldServer")
|
||||
|
||||
class OptionalSchema(BaseModel):
|
||||
required_name: str = Field(description="Your name (required)")
|
||||
optional_age: int | None = Field(default=None, description="Your age (optional)")
|
||||
optional_email: str | None = Field(default=None, description="Your email (optional)")
|
||||
subscribe: bool | None = Field(default=False, description="Subscribe to newsletter?")
|
||||
|
||||
@mcp.tool(description="Tool with optional fields")
|
||||
async def optional_tool(ctx: Context) -> str:
|
||||
result = await ctx.elicit(message="Please provide your information", schema=OptionalSchema)
|
||||
|
||||
if result.action == "accept" and result.data:
|
||||
info = [f"Name: {result.data.required_name}"]
|
||||
if result.data.optional_age is not None:
|
||||
info.append(f"Age: {result.data.optional_age}")
|
||||
if result.data.optional_email is not None:
|
||||
info.append(f"Email: {result.data.optional_email}")
|
||||
info.append(f"Subscribe: {result.data.subscribe}")
|
||||
return ", ".join(info)
|
||||
else:
|
||||
return f"User {result.action}"
|
||||
|
||||
# Test cases with different field combinations
|
||||
test_cases = [
|
||||
(
|
||||
# All fields provided
|
||||
{"required_name": "John Doe", "optional_age": 30, "optional_email": "john@example.com", "subscribe": True},
|
||||
"Name: John Doe, Age: 30, Email: john@example.com, Subscribe: True",
|
||||
),
|
||||
(
|
||||
# Only required fields
|
||||
{"required_name": "Jane Smith"},
|
||||
"Name: Jane Smith, Subscribe: False",
|
||||
),
|
||||
]
|
||||
|
||||
for content, expected in test_cases:
|
||||
|
||||
async def callback(context, params):
|
||||
return ElicitResult(action="accept", content=content)
|
||||
|
||||
await call_tool_and_assert(mcp, callback, "optional_tool", {}, expected)
|
||||
|
||||
# Test invalid optional field
|
||||
class InvalidOptionalSchema(BaseModel):
|
||||
name: str = Field(description="Name")
|
||||
optional_list: list[str] | None = Field(default=None, description="Invalid optional list")
|
||||
|
||||
@mcp.tool(description="Tool with invalid optional field")
|
||||
async def invalid_optional_tool(ctx: Context) -> str:
|
||||
try:
|
||||
await ctx.elicit(message="This should fail", schema=InvalidOptionalSchema)
|
||||
return "Should not reach here"
|
||||
except TypeError as e:
|
||||
return f"Validation failed: {str(e)}"
|
||||
|
||||
await call_tool_and_assert(
|
||||
mcp,
|
||||
lambda c, p: ElicitResult(action="accept", content={}),
|
||||
"invalid_optional_tool",
|
||||
{},
|
||||
text_contains=["Validation failed:", "optional_list"],
|
||||
)
|
||||
|
|
@ -14,14 +14,14 @@ from typing import Any
|
|||
|
||||
import pytest
|
||||
import uvicorn
|
||||
from pydantic import AnyUrl
|
||||
from pydantic import AnyUrl, BaseModel, Field
|
||||
from starlette.applications import Starlette
|
||||
from starlette.requests import Request
|
||||
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.client.sse import sse_client
|
||||
from mcp.client.streamable_http import streamablehttp_client
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from mcp.server.fastmcp import Context, FastMCP
|
||||
from mcp.server.fastmcp.resources import FunctionResource
|
||||
from mcp.shared.context import RequestContext
|
||||
from mcp.types import (
|
||||
|
|
@ -30,6 +30,7 @@ from mcp.types import (
|
|||
CompletionContext,
|
||||
CreateMessageRequestParams,
|
||||
CreateMessageResult,
|
||||
ElicitResult,
|
||||
GetPromptResult,
|
||||
InitializeResult,
|
||||
LoggingMessageNotification,
|
||||
|
|
@ -98,6 +99,20 @@ def make_fastmcp_app():
|
|||
def echo(message: str) -> str:
|
||||
return f"Echo: {message}"
|
||||
|
||||
# Add a tool that uses elicitation
|
||||
@mcp.tool(description="A tool that uses elicitation")
|
||||
async def ask_user(prompt: str, ctx: Context) -> str:
|
||||
class AnswerSchema(BaseModel):
|
||||
answer: str = Field(description="The user's answer to the question")
|
||||
|
||||
result = await ctx.elicit(message=f"Tool wants to ask: {prompt}", schema=AnswerSchema)
|
||||
|
||||
if result.action == "accept" and result.data:
|
||||
return f"User answered: {result.data.answer}"
|
||||
else:
|
||||
# Handle cancellation or decline
|
||||
return f"User cancelled or declined: {result.action}"
|
||||
|
||||
# Create the SSE app
|
||||
app = mcp.sse_app()
|
||||
|
||||
|
|
@ -261,6 +276,49 @@ def make_everything_fastmcp() -> FastMCP:
|
|||
context_data["path"] = request.url.path
|
||||
return json.dumps(context_data)
|
||||
|
||||
# Restaurant booking tool with elicitation
|
||||
@mcp.tool(description="Book a table at a restaurant with elicitation")
|
||||
async def book_restaurant(
|
||||
date: str,
|
||||
time: str,
|
||||
party_size: int,
|
||||
ctx: Context,
|
||||
) -> str:
|
||||
"""Book a table - uses elicitation if requested date is unavailable."""
|
||||
|
||||
class AlternativeDateSchema(BaseModel):
|
||||
checkAlternative: bool = Field(description="Would you like to try another date?")
|
||||
alternativeDate: str = Field(
|
||||
default="2024-12-26",
|
||||
description="What date would you prefer? (YYYY-MM-DD)",
|
||||
)
|
||||
|
||||
# For testing: assume dates starting with "2024-12-25" are unavailable
|
||||
if date.startswith("2024-12-25"):
|
||||
# Use elicitation to ask about alternatives
|
||||
result = await ctx.elicit(
|
||||
message=(
|
||||
f"No tables available for {party_size} people on {date} "
|
||||
f"at {time}. Would you like to check another date?"
|
||||
),
|
||||
schema=AlternativeDateSchema,
|
||||
)
|
||||
|
||||
if result.action == "accept" and result.data:
|
||||
if result.data.checkAlternative:
|
||||
alt_date = result.data.alternativeDate
|
||||
return f"✅ Booked table for {party_size} on {alt_date} at {time}"
|
||||
else:
|
||||
return "❌ No booking made"
|
||||
elif result.action in ("decline", "cancel"):
|
||||
return "❌ Booking cancelled"
|
||||
else:
|
||||
# Handle case where action is "accept" but data is None
|
||||
return "❌ No booking data received"
|
||||
else:
|
||||
# Available - book directly
|
||||
return f"✅ Booked table for {party_size} on {date} at {time}"
|
||||
|
||||
return mcp
|
||||
|
||||
|
||||
|
|
@ -652,6 +710,22 @@ class NotificationCollector:
|
|||
await self.handle_tool_list_changed(message.root.params)
|
||||
|
||||
|
||||
async def create_test_elicitation_callback(context, params):
|
||||
"""Shared elicitation callback for tests.
|
||||
|
||||
Handles elicitation requests for restaurant booking tests.
|
||||
"""
|
||||
# For restaurant booking test
|
||||
if "No tables available" in params.message:
|
||||
return ElicitResult(
|
||||
action="accept",
|
||||
content={"checkAlternative": True, "alternativeDate": "2024-12-26"},
|
||||
)
|
||||
else:
|
||||
# Default response
|
||||
return ElicitResult(action="decline")
|
||||
|
||||
|
||||
async def call_all_mcp_features(session: ClientSession, collector: NotificationCollector) -> None:
|
||||
"""
|
||||
Test all MCP features using the provided session.
|
||||
|
|
@ -747,6 +821,21 @@ async def call_all_mcp_features(session: ClientSession, collector: NotificationC
|
|||
assert "info" in log_levels
|
||||
assert "warning" in log_levels
|
||||
|
||||
# 5. Test elicitation tool
|
||||
# Test restaurant booking with unavailable date (triggers elicitation)
|
||||
booking_result = await session.call_tool(
|
||||
"book_restaurant",
|
||||
{
|
||||
"date": "2024-12-25", # Unavailable date to trigger elicitation
|
||||
"time": "19:00",
|
||||
"party_size": 4,
|
||||
},
|
||||
)
|
||||
assert len(booking_result.content) == 1
|
||||
assert isinstance(booking_result.content[0], TextContent)
|
||||
# Should have booked the alternative date from elicitation callback
|
||||
assert "✅ Booked table for 4 on 2024-12-26" in booking_result.content[0].text
|
||||
|
||||
# Test resources
|
||||
# 1. Static resource
|
||||
resources = await session.list_resources()
|
||||
|
|
@ -887,8 +976,6 @@ async def test_fastmcp_all_features_sse(everything_server: None, everything_serv
|
|||
# Create notification collector
|
||||
collector = NotificationCollector()
|
||||
|
||||
# Create a sampling callback that simulates an LLM
|
||||
|
||||
# Connect to the server with callbacks
|
||||
async with sse_client(everything_server_url + "/sse") as streams:
|
||||
# Set up message handler to capture notifications
|
||||
|
|
@ -901,6 +988,7 @@ async def test_fastmcp_all_features_sse(everything_server: None, everything_serv
|
|||
async with ClientSession(
|
||||
*streams,
|
||||
sampling_callback=sampling_callback,
|
||||
elicitation_callback=create_test_elicitation_callback,
|
||||
message_handler=message_handler,
|
||||
) as session:
|
||||
# Run the common test suite
|
||||
|
|
@ -933,7 +1021,37 @@ async def test_fastmcp_all_features_streamable_http(
|
|||
read_stream,
|
||||
write_stream,
|
||||
sampling_callback=sampling_callback,
|
||||
elicitation_callback=create_test_elicitation_callback,
|
||||
message_handler=message_handler,
|
||||
) as session:
|
||||
# Run the common test suite with HTTP-specific test suffix
|
||||
await call_all_mcp_features(session, collector)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_elicitation_feature(server: None, server_url: str) -> None:
|
||||
"""Test the elicitation feature."""
|
||||
|
||||
# Create a custom handler for elicitation requests
|
||||
async def elicitation_callback(context, params):
|
||||
# Verify the elicitation parameters
|
||||
if params.message == "Tool wants to ask: What is your name?":
|
||||
return ElicitResult(content={"answer": "Test User"}, action="accept")
|
||||
else:
|
||||
raise ValueError("Unexpected elicitation message")
|
||||
|
||||
# Connect to the server with our custom elicitation handler
|
||||
async with sse_client(server_url + "/sse") as streams:
|
||||
async with ClientSession(*streams, elicitation_callback=elicitation_callback) as session:
|
||||
# First initialize the session
|
||||
result = await session.initialize()
|
||||
assert isinstance(result, InitializeResult)
|
||||
assert result.serverInfo.name == "NoAuthServer"
|
||||
|
||||
# Call the tool that uses elicitation
|
||||
tool_result = await session.call_tool("ask_user", {"prompt": "What is your name?"})
|
||||
# Verify the result
|
||||
assert len(tool_result.content) == 1
|
||||
assert isinstance(tool_result.content[0], TextContent)
|
||||
# # The test should only succeed with the successful elicitation response
|
||||
assert tool_result.content[0].text == "User answered: Test User"
|
||||
|
|
|
|||
Loading…
Reference in a new issue