Skip to content

Commit d69b290

Browse files
chore: create union for working with message content (#939)
1 parent 185fa49 commit d69b290

File tree

9 files changed

+25
-51
lines changed

9 files changed

+25
-51
lines changed

examples/servers/simple-streamablehttp-stateless/mcp_simple_streamablehttp_stateless/server.py

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -41,14 +41,7 @@ def main(
4141
app = Server("mcp-streamable-http-stateless-demo")
4242

4343
@app.call_tool()
44-
async def call_tool(
45-
name: str, arguments: dict
46-
) -> list[
47-
types.TextContent
48-
| types.ImageContent
49-
| types.AudioContent
50-
| types.EmbeddedResource
51-
]:
44+
async def call_tool(name: str, arguments: dict) -> list[types.Content]:
5245
ctx = app.request_context
5346
interval = arguments.get("interval", 1.0)
5447
count = arguments.get("count", 5)

examples/servers/simple-streamablehttp/mcp_simple_streamablehttp/server.py

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -45,14 +45,7 @@ def main(
4545
app = Server("mcp-streamable-http-demo")
4646

4747
@app.call_tool()
48-
async def call_tool(
49-
name: str, arguments: dict
50-
) -> list[
51-
types.TextContent
52-
| types.ImageContent
53-
| types.AudioContent
54-
| types.EmbeddedResource
55-
]:
48+
async def call_tool(name: str, arguments: dict) -> list[types.Content]:
5649
ctx = app.request_context
5750
interval = arguments.get("interval", 1.0)
5851
count = arguments.get("count", 5)

examples/servers/simple-tool/mcp_simple_tool/server.py

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,7 @@
77

88
async def fetch_website(
99
url: str,
10-
) -> list[
11-
types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource
12-
]:
10+
) -> list[types.Content]:
1311
headers = {
1412
"User-Agent": "MCP Test Server (github.com/modelcontextprotocol/python-sdk)"
1513
}
@@ -31,14 +29,7 @@ def main(port: int, transport: str) -> int:
3129
app = Server("mcp-website-fetcher")
3230

3331
@app.call_tool()
34-
async def fetch_tool(
35-
name: str, arguments: dict
36-
) -> list[
37-
types.TextContent
38-
| types.ImageContent
39-
| types.AudioContent
40-
| types.EmbeddedResource
41-
]:
32+
async def fetch_tool(name: str, arguments: dict) -> list[types.Content]:
4233
if name != "fetch":
4334
raise ValueError(f"Unknown tool: {name}")
4435
if "url" not in arguments:

src/mcp/server/fastmcp/prompts/base.py

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,18 +7,16 @@
77
import pydantic_core
88
from pydantic import BaseModel, Field, TypeAdapter, validate_call
99

10-
from mcp.types import AudioContent, EmbeddedResource, ImageContent, TextContent
11-
12-
CONTENT_TYPES = TextContent | ImageContent | AudioContent | EmbeddedResource
10+
from mcp.types import Content, TextContent
1311

1412

1513
class Message(BaseModel):
1614
"""Base class for all prompt messages."""
1715

1816
role: Literal["user", "assistant"]
19-
content: CONTENT_TYPES
17+
content: Content
2018

21-
def __init__(self, content: str | CONTENT_TYPES, **kwargs: Any):
19+
def __init__(self, content: str | Content, **kwargs: Any):
2220
if isinstance(content, str):
2321
content = TextContent(type="text", text=content)
2422
super().__init__(content=content, **kwargs)
@@ -29,7 +27,7 @@ class UserMessage(Message):
2927

3028
role: Literal["user", "assistant"] = "user"
3129

32-
def __init__(self, content: str | CONTENT_TYPES, **kwargs: Any):
30+
def __init__(self, content: str | Content, **kwargs: Any):
3331
super().__init__(content=content, **kwargs)
3432

3533

@@ -38,7 +36,7 @@ class AssistantMessage(Message):
3836

3937
role: Literal["user", "assistant"] = "assistant"
4038

41-
def __init__(self, content: str | CONTENT_TYPES, **kwargs: Any):
39+
def __init__(self, content: str | Content, **kwargs: Any):
4240
super().__init__(content=content, **kwargs)
4341

4442

src/mcp/server/fastmcp/server.py

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -52,10 +52,8 @@
5252
from mcp.shared.context import LifespanContextT, RequestContext, RequestT
5353
from mcp.types import (
5454
AnyFunction,
55-
AudioContent,
56-
EmbeddedResource,
55+
Content,
5756
GetPromptResult,
58-
ImageContent,
5957
TextContent,
6058
ToolAnnotations,
6159
)
@@ -256,9 +254,7 @@ def get_context(self) -> Context[ServerSession, object, Request]:
256254
request_context = None
257255
return Context(request_context=request_context, fastmcp=self)
258256

259-
async def call_tool(
260-
self, name: str, arguments: dict[str, Any]
261-
) -> Sequence[TextContent | ImageContent | AudioContent | EmbeddedResource]:
257+
async def call_tool(self, name: str, arguments: dict[str, Any]) -> Sequence[Content]:
262258
"""Call a tool by name with arguments."""
263259
context = self.get_context()
264260
result = await self._tool_manager.call_tool(name, arguments, context=context)
@@ -842,12 +838,12 @@ async def get_prompt(self, name: str, arguments: dict[str, Any] | None = None) -
842838

843839
def _convert_to_content(
844840
result: Any,
845-
) -> Sequence[TextContent | ImageContent | AudioContent | EmbeddedResource]:
841+
) -> Sequence[Content]:
846842
"""Convert a result to a sequence of content objects."""
847843
if result is None:
848844
return []
849845

850-
if isinstance(result, TextContent | ImageContent | AudioContent | EmbeddedResource):
846+
if isinstance(result, Content):
851847
return [result]
852848

853849
if isinstance(result, Image):

src/mcp/server/lowlevel/server.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -384,9 +384,7 @@ def call_tool(self):
384384
def decorator(
385385
func: Callable[
386386
...,
387-
Awaitable[
388-
Iterable[types.TextContent | types.ImageContent | types.AudioContent | types.EmbeddedResource]
389-
],
387+
Awaitable[Iterable[types.Content]],
390388
],
391389
):
392390
logger.debug("Registering handler for CallToolRequest")

src/mcp/types.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -667,11 +667,14 @@ class EmbeddedResource(BaseModel):
667667
model_config = ConfigDict(extra="allow")
668668

669669

670+
Content = TextContent | ImageContent | AudioContent | EmbeddedResource
671+
672+
670673
class PromptMessage(BaseModel):
671674
"""Describes a message returned as part of a prompt."""
672675

673676
role: Role
674-
content: TextContent | ImageContent | AudioContent | EmbeddedResource
677+
content: Content
675678
model_config = ConfigDict(extra="allow")
676679

677680

@@ -787,7 +790,7 @@ class CallToolRequest(Request[CallToolRequestParams, Literal["tools/call"]]):
787790
class CallToolResult(Result):
788791
"""The server's response to a tool call."""
789792

790-
content: list[TextContent | ImageContent | AudioContent | EmbeddedResource]
793+
content: list[Content]
791794
isError: bool = False
792795

793796

tests/issues/test_88_random_error.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
from mcp.client.session import ClientSession
1212
from mcp.server.lowlevel import Server
1313
from mcp.shared.exceptions import McpError
14-
from mcp.types import AudioContent, EmbeddedResource, ImageContent, TextContent
14+
from mcp.types import Content, TextContent
1515

1616

1717
@pytest.mark.anyio
@@ -31,7 +31,7 @@ async def test_notification_validation_error(tmp_path: Path):
3131
slow_request_complete = anyio.Event()
3232

3333
@server.call_tool()
34-
async def slow_tool(name: str, arg) -> Sequence[TextContent | ImageContent | AudioContent | EmbeddedResource]:
34+
async def slow_tool(name: str, arg) -> Sequence[Content]:
3535
nonlocal request_count
3636
request_count += 1
3737

tests/server/fastmcp/test_server.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from starlette.routing import Mount, Route
99

1010
from mcp.server.fastmcp import Context, FastMCP
11-
from mcp.server.fastmcp.prompts.base import EmbeddedResource, Message, UserMessage
11+
from mcp.server.fastmcp.prompts.base import Message, UserMessage
1212
from mcp.server.fastmcp.resources import FileResource, FunctionResource
1313
from mcp.server.fastmcp.utilities.types import Image
1414
from mcp.shared.exceptions import McpError
@@ -18,6 +18,8 @@
1818
from mcp.types import (
1919
AudioContent,
2020
BlobResourceContents,
21+
Content,
22+
EmbeddedResource,
2123
ImageContent,
2224
TextContent,
2325
TextResourceContents,
@@ -192,7 +194,7 @@ def image_tool_fn(path: str) -> Image:
192194
return Image(path)
193195

194196

195-
def mixed_content_tool_fn() -> list[TextContent | ImageContent | AudioContent]:
197+
def mixed_content_tool_fn() -> list[Content]:
196198
return [
197199
TextContent(type="text", text="Hello"),
198200
ImageContent(type="image", data="abc", mimeType="image/png"),

0 commit comments

Comments
 (0)