feat: backwards-compatible create_message overloads for SEP-1577 (#1713)
This commit is contained in:
parent
f82b0c9371
commit
fa851d93a2
9 changed files with 209 additions and 32 deletions
|
|
@ -948,8 +948,9 @@ async def generate_poem(topic: str, ctx: Context[ServerSession, None]) -> str:
|
||||||
max_tokens=100,
|
max_tokens=100,
|
||||||
)
|
)
|
||||||
|
|
||||||
if all(c.type == "text" for c in result.content_as_list):
|
# Since we're not passing tools param, result.content is single content
|
||||||
return "\n".join(c.text for c in result.content_as_list if c.type == "text")
|
if result.content.type == "text":
|
||||||
|
return result.content.text
|
||||||
return str(result.content)
|
return str(result.content)
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -178,8 +178,9 @@ async def test_sampling(prompt: str, ctx: Context[ServerSession, None]) -> str:
|
||||||
max_tokens=100,
|
max_tokens=100,
|
||||||
)
|
)
|
||||||
|
|
||||||
if any(c.type == "text" for c in result.content_as_list):
|
# Since we're not passing tools param, result.content is single content
|
||||||
model_response = "\n".join(c.text for c in result.content_as_list if c.type == "text")
|
if result.content.type == "text":
|
||||||
|
model_response = result.content.text
|
||||||
else:
|
else:
|
||||||
model_response = "No response"
|
model_response = "No response"
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -20,6 +20,7 @@ async def generate_poem(topic: str, ctx: Context[ServerSession, None]) -> str:
|
||||||
max_tokens=100,
|
max_tokens=100,
|
||||||
)
|
)
|
||||||
|
|
||||||
if all(c.type == "text" for c in result.content_as_list):
|
# Since we're not passing tools param, result.content is single content
|
||||||
return "\n".join(c.text for c in result.content_as_list if c.type == "text")
|
if result.content.type == "text":
|
||||||
|
return result.content.text
|
||||||
return str(result.content)
|
return str(result.content)
|
||||||
|
|
|
||||||
|
|
@ -13,6 +13,7 @@ from .types import (
|
||||||
CompleteRequest,
|
CompleteRequest,
|
||||||
CreateMessageRequest,
|
CreateMessageRequest,
|
||||||
CreateMessageResult,
|
CreateMessageResult,
|
||||||
|
CreateMessageResultWithTools,
|
||||||
ErrorData,
|
ErrorData,
|
||||||
GetPromptRequest,
|
GetPromptRequest,
|
||||||
GetPromptResult,
|
GetPromptResult,
|
||||||
|
|
@ -42,6 +43,7 @@ from .types import (
|
||||||
ResourceUpdatedNotification,
|
ResourceUpdatedNotification,
|
||||||
RootsCapability,
|
RootsCapability,
|
||||||
SamplingCapability,
|
SamplingCapability,
|
||||||
|
SamplingContent,
|
||||||
SamplingContextCapability,
|
SamplingContextCapability,
|
||||||
SamplingMessage,
|
SamplingMessage,
|
||||||
SamplingMessageContentBlock,
|
SamplingMessageContentBlock,
|
||||||
|
|
@ -75,6 +77,7 @@ __all__ = [
|
||||||
"CompleteRequest",
|
"CompleteRequest",
|
||||||
"CreateMessageRequest",
|
"CreateMessageRequest",
|
||||||
"CreateMessageResult",
|
"CreateMessageResult",
|
||||||
|
"CreateMessageResultWithTools",
|
||||||
"ErrorData",
|
"ErrorData",
|
||||||
"GetPromptRequest",
|
"GetPromptRequest",
|
||||||
"GetPromptResult",
|
"GetPromptResult",
|
||||||
|
|
@ -105,6 +108,7 @@ __all__ = [
|
||||||
"ResourceUpdatedNotification",
|
"ResourceUpdatedNotification",
|
||||||
"RootsCapability",
|
"RootsCapability",
|
||||||
"SamplingCapability",
|
"SamplingCapability",
|
||||||
|
"SamplingContent",
|
||||||
"SamplingContextCapability",
|
"SamplingContextCapability",
|
||||||
"SamplingMessage",
|
"SamplingMessage",
|
||||||
"SamplingMessageContentBlock",
|
"SamplingMessageContentBlock",
|
||||||
|
|
|
||||||
|
|
@ -38,7 +38,7 @@ be instantiated directly by users of the MCP framework.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from enum import Enum
|
from enum import Enum
|
||||||
from typing import Any, TypeVar
|
from typing import Any, TypeVar, overload
|
||||||
|
|
||||||
import anyio
|
import anyio
|
||||||
import anyio.lowlevel
|
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(
|
async def create_message(
|
||||||
self,
|
self,
|
||||||
messages: list[types.SamplingMessage],
|
messages: list[types.SamplingMessage],
|
||||||
|
|
@ -247,7 +285,7 @@ class ServerSession(
|
||||||
tools: list[types.Tool] | None = None,
|
tools: list[types.Tool] | None = None,
|
||||||
tool_choice: types.ToolChoice | None = None,
|
tool_choice: types.ToolChoice | None = None,
|
||||||
related_request_id: types.RequestId | None = None,
|
related_request_id: types.RequestId | None = None,
|
||||||
) -> types.CreateMessageResult:
|
) -> types.CreateMessageResult | types.CreateMessageResultWithTools:
|
||||||
"""Send a sampling/create_message request.
|
"""Send a sampling/create_message request.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
|
|
@ -278,27 +316,35 @@ class ServerSession(
|
||||||
validate_sampling_tools(client_caps, tools, tool_choice)
|
validate_sampling_tools(client_caps, tools, tool_choice)
|
||||||
validate_tool_use_result_messages(messages)
|
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(
|
return await self.send_request(
|
||||||
request=types.ServerRequest(
|
request=request,
|
||||||
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,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
),
|
|
||||||
result_type=types.CreateMessageResult,
|
result_type=types.CreateMessageResult,
|
||||||
metadata=ServerMessageMetadata(
|
metadata=metadata_obj,
|
||||||
related_request_id=related_request_id,
|
|
||||||
),
|
|
||||||
)
|
)
|
||||||
|
|
||||||
async def list_roots(self) -> types.ListRootsResult:
|
async def list_roots(self) -> types.ListRootsResult:
|
||||||
|
|
|
||||||
|
|
@ -1146,6 +1146,10 @@ class ToolResultContent(BaseModel):
|
||||||
SamplingMessageContentBlock: TypeAlias = TextContent | ImageContent | AudioContent | ToolUseContent | ToolResultContent
|
SamplingMessageContentBlock: TypeAlias = TextContent | ImageContent | AudioContent | ToolUseContent | ToolResultContent
|
||||||
"""Content block types allowed in sampling messages."""
|
"""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):
|
class SamplingMessage(BaseModel):
|
||||||
"""Describes a message issued to or received from an LLM API."""
|
"""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):
|
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
|
role: Role
|
||||||
"""The role of the message sender (typically 'assistant' for LLM responses)."""
|
"""The role of the message sender (typically 'assistant' for LLM responses)."""
|
||||||
|
|
|
||||||
|
|
@ -8,8 +8,10 @@ from mcp.shared.memory import (
|
||||||
from mcp.types import (
|
from mcp.types import (
|
||||||
CreateMessageRequestParams,
|
CreateMessageRequestParams,
|
||||||
CreateMessageResult,
|
CreateMessageResult,
|
||||||
|
CreateMessageResultWithTools,
|
||||||
SamplingMessage,
|
SamplingMessage,
|
||||||
TextContent,
|
TextContent,
|
||||||
|
ToolUseContent,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -56,3 +58,79 @@ async def test_sampling_callback():
|
||||||
assert result.isError is True
|
assert result.isError is True
|
||||||
assert isinstance(result.content[0], TextContent)
|
assert isinstance(result.content[0], TextContent)
|
||||||
assert result.content[0].text == "Error executing tool test_sampling: Sampling not supported"
|
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
|
# Return the sampling result in the tool response
|
||||||
if all(c.type == "text" for c in sampling_result.content_as_list):
|
# Since we're not passing tools param, result.content is single content
|
||||||
response = "\n".join(c.text for c in sampling_result.content_as_list if c.type == "text")
|
if sampling_result.content.type == "text":
|
||||||
|
response = sampling_result.content.text
|
||||||
else:
|
else:
|
||||||
response = str(sampling_result.content)
|
response = str(sampling_result.content)
|
||||||
return [
|
return [
|
||||||
|
|
|
||||||
|
|
@ -8,6 +8,7 @@ from mcp.types import (
|
||||||
ClientRequest,
|
ClientRequest,
|
||||||
CreateMessageRequestParams,
|
CreateMessageRequestParams,
|
||||||
CreateMessageResult,
|
CreateMessageResult,
|
||||||
|
CreateMessageResultWithTools,
|
||||||
Implementation,
|
Implementation,
|
||||||
InitializeRequest,
|
InitializeRequest,
|
||||||
InitializeRequestParams,
|
InitializeRequestParams,
|
||||||
|
|
@ -239,7 +240,7 @@ async def test_create_message_request_params_with_tools():
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_create_message_result_with_tool_use():
|
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 = {
|
result_data = {
|
||||||
"role": "assistant",
|
"role": "assistant",
|
||||||
"content": {"type": "tool_use", "name": "search", "id": "call_123", "input": {"query": "test"}},
|
"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",
|
"stopReason": "toolUse",
|
||||||
}
|
}
|
||||||
|
|
||||||
result = CreateMessageResult.model_validate(result_data)
|
# Tool use content uses CreateMessageResultWithTools
|
||||||
|
result = CreateMessageResultWithTools.model_validate(result_data)
|
||||||
assert result.role == "assistant"
|
assert result.role == "assistant"
|
||||||
assert isinstance(result.content, ToolUseContent)
|
assert isinstance(result.content, ToolUseContent)
|
||||||
assert result.stopReason == "toolUse"
|
assert result.stopReason == "toolUse"
|
||||||
|
|
@ -259,6 +261,25 @@ async def test_create_message_result_with_tool_use():
|
||||||
assert content_list[0] == result.content
|
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
|
@pytest.mark.anyio
|
||||||
async def test_client_capabilities_with_sampling_tools():
|
async def test_client_capabilities_with_sampling_tools():
|
||||||
"""Test ClientCapabilities with nested sampling capabilities for SEP-1577."""
|
"""Test ClientCapabilities with nested sampling capabilities for SEP-1577."""
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue