diff --git a/mcp_python/server/__init__.py b/mcp_python/server/__init__.py index ff8900f..5c85eea 100644 --- a/mcp_python/server/__init__.py +++ b/mcp_python/server/__init__.py @@ -42,6 +42,9 @@ from mcp_python.types import ( SubscribeRequest, Tool, UnsubscribeRequest, + TextContent, + EmbeddedResource, + PromptMessage, ) logger = logging.getLogger(__name__) @@ -117,8 +120,6 @@ class Server: GetPromptRequest, GetPromptResult, ImageContent, - SamplingMessage, - TextContent, ) from mcp_python.types import ( Role as Role, @@ -133,7 +134,7 @@ class Server: async def handler(req: GetPromptRequest): prompt_get = await func(req.params.name, req.params.arguments) - messages: list[SamplingMessage] = [] + messages: list[PromptMessage] = [] for message in prompt_get.messages: match message.content: case str() as text_content: @@ -144,15 +145,20 @@ class Server: data=img_content.data, mimeType=img_content.mime_type, ) + case types.EmbeddedResource() as resource: + content = EmbeddedResource( + type="resource", + resource=resource.resource + ) case _: raise ValueError( f"Unexpected content type: {type(message.content)}" ) - sampling_message = SamplingMessage( + prompt_message = PromptMessage( role=message.role, content=content ) - messages.append(sampling_message) + messages.append(prompt_message) return ServerResult( GetPromptResult(description=prompt_get.desc, messages=messages) @@ -276,14 +282,46 @@ class Server: return decorator def call_tool(self): - from mcp_python.types import CallToolResult + from mcp_python.types import CallToolResult, TextContent, ImageContent, EmbeddedResource - def decorator(func: Callable[..., Awaitable[Any]]): + def decorator( + func: Callable[..., Awaitable[list[str | types.ImageContent | types.EmbeddedResource]]] + ): logger.debug("Registering handler for CallToolRequest") async def handler(req: CallToolRequest): - result = await func(req.params.name, (req.params.arguments or {})) - return ServerResult(CallToolResult(toolResult=result)) + try: + results = await func(req.params.name, (req.params.arguments or {})) + content = [] + for result in results: + match result: + case str() as text: + content.append(TextContent(type="text", text=text)) + case types.ImageContent() as img: + content.append(ImageContent( + type="image", + data=img.data, + mimeType=img.mime_type + )) + case types.EmbeddedResource() as resource: + content.append(EmbeddedResource( + type="resource", + resource=resource.resource + )) + + return ServerResult( + CallToolResult( + content=content, + isError=False + ) + ) + except Exception as e: + return ServerResult( + CallToolResult( + content=[TextContent(type="text", text=str(e))], + isError=True + ) + ) self.request_handlers[CallToolRequest] = handler return func diff --git a/mcp_python/server/types.py b/mcp_python/server/types.py index 7632406..437bc29 100644 --- a/mcp_python/server/types.py +++ b/mcp_python/server/types.py @@ -1,5 +1,5 @@ """ -This module provides simpler types to use with the server for managing prompts. +This module provides simpler types to use with the server for managing prompts and tools. """ from dataclasses import dataclass @@ -7,7 +7,7 @@ from typing import Literal from pydantic import BaseModel -from mcp_python.types import Role, ServerCapabilities +from mcp_python.types import Role, ServerCapabilities, TextResourceContents, BlobResourceContents @dataclass @@ -17,10 +17,15 @@ class ImageContent: mime_type: str +@dataclass +class EmbeddedResource: + resource: TextResourceContents | BlobResourceContents + + @dataclass class Message: role: Role - content: str | ImageContent + content: str | ImageContent | EmbeddedResource @dataclass