feat: backwards-compatible create_message overloads for SEP-1577 (#1713)
This commit is contained in:
@@ -948,8 +948,9 @@ async def generate_poem(topic: str, ctx: Context[ServerSession, None]) -> str:
|
||||
max_tokens=100,
|
||||
)
|
||||
|
||||
if all(c.type == "text" for c in result.content_as_list):
|
||||
return "\n".join(c.text for c in result.content_as_list if c.type == "text")
|
||||
# Since we're not passing tools param, result.content is single content
|
||||
if result.content.type == "text":
|
||||
return result.content.text
|
||||
return str(result.content)
|
||||
```
|
||||
|
||||
|
||||
@@ -178,8 +178,9 @@ async def test_sampling(prompt: str, ctx: Context[ServerSession, None]) -> str:
|
||||
max_tokens=100,
|
||||
)
|
||||
|
||||
if any(c.type == "text" for c in result.content_as_list):
|
||||
model_response = "\n".join(c.text for c in result.content_as_list if c.type == "text")
|
||||
# Since we're not passing tools param, result.content is single content
|
||||
if result.content.type == "text":
|
||||
model_response = result.content.text
|
||||
else:
|
||||
model_response = "No response"
|
||||
|
||||
|
||||
@@ -20,6 +20,7 @@ async def generate_poem(topic: str, ctx: Context[ServerSession, None]) -> str:
|
||||
max_tokens=100,
|
||||
)
|
||||
|
||||
if all(c.type == "text" for c in result.content_as_list):
|
||||
return "\n".join(c.text for c in result.content_as_list if c.type == "text")
|
||||
# Since we're not passing tools param, result.content is single content
|
||||
if result.content.type == "text":
|
||||
return result.content.text
|
||||
return str(result.content)
|
||||
|
||||
@@ -13,6 +13,7 @@ from .types import (
|
||||
CompleteRequest,
|
||||
CreateMessageRequest,
|
||||
CreateMessageResult,
|
||||
CreateMessageResultWithTools,
|
||||
ErrorData,
|
||||
GetPromptRequest,
|
||||
GetPromptResult,
|
||||
@@ -42,6 +43,7 @@ from .types import (
|
||||
ResourceUpdatedNotification,
|
||||
RootsCapability,
|
||||
SamplingCapability,
|
||||
SamplingContent,
|
||||
SamplingContextCapability,
|
||||
SamplingMessage,
|
||||
SamplingMessageContentBlock,
|
||||
@@ -75,6 +77,7 @@ __all__ = [
|
||||
"CompleteRequest",
|
||||
"CreateMessageRequest",
|
||||
"CreateMessageResult",
|
||||
"CreateMessageResultWithTools",
|
||||
"ErrorData",
|
||||
"GetPromptRequest",
|
||||
"GetPromptResult",
|
||||
@@ -105,6 +108,7 @@ __all__ = [
|
||||
"ResourceUpdatedNotification",
|
||||
"RootsCapability",
|
||||
"SamplingCapability",
|
||||
"SamplingContent",
|
||||
"SamplingContextCapability",
|
||||
"SamplingMessage",
|
||||
"SamplingMessageContentBlock",
|
||||
|
||||
+67
-21
@@ -38,7 +38,7 @@ be instantiated directly by users of the MCP framework.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
from typing import Any, TypeVar
|
||||
from typing import Any, TypeVar, overload
|
||||
|
||||
import anyio
|
||||
import anyio.lowlevel
|
||||
@@ -233,6 +233,44 @@ class ServerSession(
|
||||
)
|
||||
)
|
||||
|
||||
@overload
|
||||
async def create_message(
|
||||
self,
|
||||
messages: list[types.SamplingMessage],
|
||||
*,
|
||||
max_tokens: int,
|
||||
system_prompt: str | None = None,
|
||||
include_context: types.IncludeContext | None = None,
|
||||
temperature: float | None = None,
|
||||
stop_sequences: list[str] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
model_preferences: types.ModelPreferences | None = None,
|
||||
tools: None = None,
|
||||
tool_choice: types.ToolChoice | None = None,
|
||||
related_request_id: types.RequestId | None = None,
|
||||
) -> types.CreateMessageResult:
|
||||
"""Overload: Without tools, returns single content."""
|
||||
...
|
||||
|
||||
@overload
|
||||
async def create_message(
|
||||
self,
|
||||
messages: list[types.SamplingMessage],
|
||||
*,
|
||||
max_tokens: int,
|
||||
system_prompt: str | None = None,
|
||||
include_context: types.IncludeContext | None = None,
|
||||
temperature: float | None = None,
|
||||
stop_sequences: list[str] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
model_preferences: types.ModelPreferences | None = None,
|
||||
tools: list[types.Tool],
|
||||
tool_choice: types.ToolChoice | None = None,
|
||||
related_request_id: types.RequestId | None = None,
|
||||
) -> types.CreateMessageResultWithTools:
|
||||
"""Overload: With tools, returns array-capable content."""
|
||||
...
|
||||
|
||||
async def create_message(
|
||||
self,
|
||||
messages: list[types.SamplingMessage],
|
||||
@@ -247,7 +285,7 @@ class ServerSession(
|
||||
tools: list[types.Tool] | None = None,
|
||||
tool_choice: types.ToolChoice | None = None,
|
||||
related_request_id: types.RequestId | None = None,
|
||||
) -> types.CreateMessageResult:
|
||||
) -> types.CreateMessageResult | types.CreateMessageResultWithTools:
|
||||
"""Send a sampling/create_message request.
|
||||
|
||||
Args:
|
||||
@@ -278,27 +316,35 @@ class ServerSession(
|
||||
validate_sampling_tools(client_caps, tools, tool_choice)
|
||||
validate_tool_use_result_messages(messages)
|
||||
|
||||
request = types.ServerRequest(
|
||||
types.CreateMessageRequest(
|
||||
params=types.CreateMessageRequestParams(
|
||||
messages=messages,
|
||||
systemPrompt=system_prompt,
|
||||
includeContext=include_context,
|
||||
temperature=temperature,
|
||||
maxTokens=max_tokens,
|
||||
stopSequences=stop_sequences,
|
||||
metadata=metadata,
|
||||
modelPreferences=model_preferences,
|
||||
tools=tools,
|
||||
toolChoice=tool_choice,
|
||||
),
|
||||
)
|
||||
)
|
||||
metadata_obj = ServerMessageMetadata(related_request_id=related_request_id)
|
||||
|
||||
# Use different result types based on whether tools are provided
|
||||
if tools is not None:
|
||||
return await self.send_request(
|
||||
request=request,
|
||||
result_type=types.CreateMessageResultWithTools,
|
||||
metadata=metadata_obj,
|
||||
)
|
||||
return await self.send_request(
|
||||
request=types.ServerRequest(
|
||||
types.CreateMessageRequest(
|
||||
params=types.CreateMessageRequestParams(
|
||||
messages=messages,
|
||||
systemPrompt=system_prompt,
|
||||
includeContext=include_context,
|
||||
temperature=temperature,
|
||||
maxTokens=max_tokens,
|
||||
stopSequences=stop_sequences,
|
||||
metadata=metadata,
|
||||
modelPreferences=model_preferences,
|
||||
tools=tools,
|
||||
toolChoice=tool_choice,
|
||||
),
|
||||
)
|
||||
),
|
||||
request=request,
|
||||
result_type=types.CreateMessageResult,
|
||||
metadata=ServerMessageMetadata(
|
||||
related_request_id=related_request_id,
|
||||
),
|
||||
metadata=metadata_obj,
|
||||
)
|
||||
|
||||
async def list_roots(self) -> types.ListRootsResult:
|
||||
|
||||
+25
-1
@@ -1146,6 +1146,10 @@ class ToolResultContent(BaseModel):
|
||||
SamplingMessageContentBlock: TypeAlias = TextContent | ImageContent | AudioContent | ToolUseContent | ToolResultContent
|
||||
"""Content block types allowed in sampling messages."""
|
||||
|
||||
SamplingContent: TypeAlias = TextContent | ImageContent | AudioContent
|
||||
"""Basic content types for sampling responses (without tool use).
|
||||
Used for backwards-compatible CreateMessageResult when tools are not used."""
|
||||
|
||||
|
||||
class SamplingMessage(BaseModel):
|
||||
"""Describes a message issued to or received from an LLM API."""
|
||||
@@ -1543,7 +1547,27 @@ StopReason = Literal["endTurn", "stopSequence", "maxTokens", "toolUse"] | str
|
||||
|
||||
|
||||
class CreateMessageResult(Result):
|
||||
"""The client's response to a sampling/create_message request from the server."""
|
||||
"""The client's response to a sampling/create_message request from the server.
|
||||
|
||||
This is the backwards-compatible version that returns single content (no arrays).
|
||||
Used when the request does not include tools.
|
||||
"""
|
||||
|
||||
role: Role
|
||||
"""The role of the message sender (typically 'assistant' for LLM responses)."""
|
||||
content: SamplingContent
|
||||
"""Response content. Single content block (text, image, or audio)."""
|
||||
model: str
|
||||
"""The name of the model that generated the message."""
|
||||
stopReason: StopReason | None = None
|
||||
"""The reason why sampling stopped, if known."""
|
||||
|
||||
|
||||
class CreateMessageResultWithTools(Result):
|
||||
"""The client's response to a sampling/create_message request when tools were provided.
|
||||
|
||||
This version supports array content for tool use flows.
|
||||
"""
|
||||
|
||||
role: Role
|
||||
"""The role of the message sender (typically 'assistant' for LLM responses)."""
|
||||
|
||||
@@ -8,8 +8,10 @@ from mcp.shared.memory import (
|
||||
from mcp.types import (
|
||||
CreateMessageRequestParams,
|
||||
CreateMessageResult,
|
||||
CreateMessageResultWithTools,
|
||||
SamplingMessage,
|
||||
TextContent,
|
||||
ToolUseContent,
|
||||
)
|
||||
|
||||
|
||||
@@ -56,3 +58,79 @@ async def test_sampling_callback():
|
||||
assert result.isError is True
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert result.content[0].text == "Error executing tool test_sampling: Sampling not supported"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_message_backwards_compat_single_content():
|
||||
"""Test backwards compatibility: create_message without tools returns single content."""
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
server = FastMCP("test")
|
||||
|
||||
# Callback returns single content (text)
|
||||
callback_return = CreateMessageResult(
|
||||
role="assistant",
|
||||
content=TextContent(type="text", text="Hello from LLM"),
|
||||
model="test-model",
|
||||
stopReason="endTurn",
|
||||
)
|
||||
|
||||
async def sampling_callback(
|
||||
context: RequestContext[ClientSession, None],
|
||||
params: CreateMessageRequestParams,
|
||||
) -> CreateMessageResult:
|
||||
return callback_return
|
||||
|
||||
@server.tool("test_backwards_compat")
|
||||
async def test_tool(message: str):
|
||||
# Call create_message WITHOUT tools
|
||||
result = await server.get_context().session.create_message(
|
||||
messages=[SamplingMessage(role="user", content=TextContent(type="text", text=message))],
|
||||
max_tokens=100,
|
||||
)
|
||||
# Backwards compat: result should be CreateMessageResult
|
||||
assert isinstance(result, CreateMessageResult)
|
||||
# Content should be single (not a list) - this is the key backwards compat check
|
||||
assert isinstance(result.content, TextContent)
|
||||
assert result.content.text == "Hello from LLM"
|
||||
# CreateMessageResult should NOT have content_as_list (that's on WithTools)
|
||||
assert not hasattr(result, "content_as_list") or not callable(getattr(result, "content_as_list", None))
|
||||
return True
|
||||
|
||||
async with create_session(server._mcp_server, sampling_callback=sampling_callback) as client_session:
|
||||
result = await client_session.call_tool("test_backwards_compat", {"message": "Test"})
|
||||
assert result.isError is False
|
||||
assert isinstance(result.content[0], TextContent)
|
||||
assert result.content[0].text == "true"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_message_result_with_tools_type():
|
||||
"""Test that CreateMessageResultWithTools supports content_as_list."""
|
||||
# Test the type itself, not the overload (overload requires client capability setup)
|
||||
result = CreateMessageResultWithTools(
|
||||
role="assistant",
|
||||
content=ToolUseContent(type="tool_use", id="call_123", name="get_weather", input={"city": "SF"}),
|
||||
model="test-model",
|
||||
stopReason="toolUse",
|
||||
)
|
||||
|
||||
# CreateMessageResultWithTools should have content_as_list
|
||||
content_list = result.content_as_list
|
||||
assert len(content_list) == 1
|
||||
assert content_list[0].type == "tool_use"
|
||||
|
||||
# It should also work with array content
|
||||
result_array = CreateMessageResultWithTools(
|
||||
role="assistant",
|
||||
content=[
|
||||
TextContent(type="text", text="Let me check the weather"),
|
||||
ToolUseContent(type="tool_use", id="call_456", name="get_weather", input={"city": "NYC"}),
|
||||
],
|
||||
model="test-model",
|
||||
stopReason="toolUse",
|
||||
)
|
||||
content_list_array = result_array.content_as_list
|
||||
assert len(content_list_array) == 2
|
||||
assert content_list_array[0].type == "text"
|
||||
assert content_list_array[1].type == "tool_use"
|
||||
|
||||
@@ -242,8 +242,9 @@ class ServerTest(Server): # pragma: no cover
|
||||
)
|
||||
|
||||
# Return the sampling result in the tool response
|
||||
if all(c.type == "text" for c in sampling_result.content_as_list):
|
||||
response = "\n".join(c.text for c in sampling_result.content_as_list if c.type == "text")
|
||||
# Since we're not passing tools param, result.content is single content
|
||||
if sampling_result.content.type == "text":
|
||||
response = sampling_result.content.text
|
||||
else:
|
||||
response = str(sampling_result.content)
|
||||
return [
|
||||
|
||||
+23
-2
@@ -8,6 +8,7 @@ from mcp.types import (
|
||||
ClientRequest,
|
||||
CreateMessageRequestParams,
|
||||
CreateMessageResult,
|
||||
CreateMessageResultWithTools,
|
||||
Implementation,
|
||||
InitializeRequest,
|
||||
InitializeRequestParams,
|
||||
@@ -239,7 +240,7 @@ async def test_create_message_request_params_with_tools():
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_message_result_with_tool_use():
|
||||
"""Test CreateMessageResult with tool use content for SEP-1577."""
|
||||
"""Test CreateMessageResultWithTools with tool use content for SEP-1577."""
|
||||
result_data = {
|
||||
"role": "assistant",
|
||||
"content": {"type": "tool_use", "name": "search", "id": "call_123", "input": {"query": "test"}},
|
||||
@@ -247,7 +248,8 @@ async def test_create_message_result_with_tool_use():
|
||||
"stopReason": "toolUse",
|
||||
}
|
||||
|
||||
result = CreateMessageResult.model_validate(result_data)
|
||||
# Tool use content uses CreateMessageResultWithTools
|
||||
result = CreateMessageResultWithTools.model_validate(result_data)
|
||||
assert result.role == "assistant"
|
||||
assert isinstance(result.content, ToolUseContent)
|
||||
assert result.stopReason == "toolUse"
|
||||
@@ -259,6 +261,25 @@ async def test_create_message_result_with_tool_use():
|
||||
assert content_list[0] == result.content
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_message_result_basic():
|
||||
"""Test CreateMessageResult with basic text content (backwards compatible)."""
|
||||
result_data = {
|
||||
"role": "assistant",
|
||||
"content": {"type": "text", "text": "Hello!"},
|
||||
"model": "claude-3",
|
||||
"stopReason": "endTurn",
|
||||
}
|
||||
|
||||
# Basic content uses CreateMessageResult (single content, no arrays)
|
||||
result = CreateMessageResult.model_validate(result_data)
|
||||
assert result.role == "assistant"
|
||||
assert isinstance(result.content, TextContent)
|
||||
assert result.content.text == "Hello!"
|
||||
assert result.stopReason == "endTurn"
|
||||
assert result.model == "claude-3"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_client_capabilities_with_sampling_tools():
|
||||
"""Test ClientCapabilities with nested sampling capabilities for SEP-1577."""
|
||||
|
||||
Reference in New Issue
Block a user