From f5d82bd22966af44c4603104147ba98f09e67c3c Mon Sep 17 00:00:00 2001 From: David Soria Parra Date: Mon, 11 Nov 2024 20:05:51 +0000 Subject: [PATCH] Remove helper types The helper types in mcp.server.types got really confusioning during implementation as they overlapped with mcp.types. I now believe it is better if we stay more low level to the spec types. To do this, we now only use mcp.types everywhere. We renamed mcp.server.types to mcp.server.models and removed it to the absolute minimum. --- src/mcp/server/__init__.py | 58 +++++++++--------------------------- src/mcp/server/__main__.py | 2 +- src/mcp/server/models.py | 19 ++++++++++++ src/mcp/server/session.py | 2 +- src/mcp/server/types.py | 46 ---------------------------- tests/conftest.py | 2 +- tests/server/test_session.py | 2 +- uv.lock | 7 ++--- 8 files changed, 40 insertions(+), 98 deletions(-) create mode 100644 src/mcp/server/models.py delete mode 100644 src/mcp/server/types.py diff --git a/src/mcp/server/__init__.py b/src/mcp/server/__init__.py index f7e66a3..29133e0 100644 --- a/src/mcp/server/__init__.py +++ b/src/mcp/server/__init__.py @@ -7,7 +7,7 @@ from typing import Any, Sequence from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream from pydantic import AnyUrl -from mcp.server import types +from mcp.server.models import InitializationOptions from mcp.server.session import ServerSession from mcp.server.stdio import stdio_server as stdio_server from mcp.shared.context import RequestContext @@ -15,6 +15,10 @@ from mcp.shared.session import RequestResponder from mcp.types import ( METHOD_NOT_FOUND, CallToolRequest, + GetPromptResult, + GetPromptRequest, + GetPromptResult, + ImageContent, ClientNotification, ClientRequest, CompleteRequest, @@ -84,7 +88,7 @@ class Server: self, notification_options: NotificationOptions | None = None, experimental_capabilities: dict[str, dict[str, Any]] | None = None, - ) -> types.InitializationOptions: + ) -> InitializationOptions: """Create initialization options from this server instance.""" def pkg_version(package: str) -> str: @@ -99,7 +103,7 @@ class Server: return "unknown" - return types.InitializationOptions( + return InitializationOptions( server_name=self.name, server_version=pkg_version("mcp"), capabilities=self.get_capabilities( @@ -168,50 +172,16 @@ class Server: return decorator def get_prompt(self): - from mcp.types import ( - GetPromptRequest, - GetPromptResult, - ImageContent, - ) - from mcp.types import ( - Role as Role, - ) - def decorator( func: Callable[ - [str, dict[str, str] | None], Awaitable[types.PromptResponse] + [str, dict[str, str] | None], Awaitable[GetPromptResult] ], ): logger.debug("Registering handler for GetPromptRequest") async def handler(req: GetPromptRequest): prompt_get = await func(req.params.name, req.params.arguments) - messages: list[PromptMessage] = [] - for message in prompt_get.messages: - match message.content: - case str() as text_content: - content = TextContent(type="text", text=text_content) - case types.ImageContent() as img_content: - content = ImageContent( - type="image", - 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)}" - ) - - prompt_message = PromptMessage(role=message.role, content=content) - messages.append(prompt_message) - - return ServerResult( - GetPromptResult(description=prompt_get.desc, messages=messages) - ) + return ServerResult(prompt_get) self.request_handlers[GetPromptRequest] = handler return func @@ -338,7 +308,7 @@ class Server: def decorator( func: Callable[ ..., - Awaitable[Sequence[str | types.ImageContent | types.EmbeddedResource]], + Awaitable[Sequence[TextContent | ImageContent | EmbeddedResource]], ], ): logger.debug("Registering handler for CallToolRequest") @@ -351,15 +321,15 @@ class Server: match result: case str() as text: content.append(TextContent(type="text", text=text)) - case types.ImageContent() as img: + case ImageContent() as img: content.append( ImageContent( type="image", data=img.data, - mimeType=img.mime_type, + mimeType=img.mimeType, ) ) - case types.EmbeddedResource() as resource: + case EmbeddedResource() as resource: content.append( EmbeddedResource( type="resource", resource=resource.resource @@ -427,7 +397,7 @@ class Server: self, read_stream: MemoryObjectReceiveStream[JSONRPCMessage | Exception], write_stream: MemoryObjectSendStream[JSONRPCMessage], - initialization_options: types.InitializationOptions, + initialization_options: InitializationOptions, # When True, exceptions are returned as messages to the client. # When False, exceptions are raised, which will cause the server to shut down # but also make tracing exceptions much easier during testing and when using diff --git a/src/mcp/server/__main__.py b/src/mcp/server/__main__.py index 2313f46..417380d 100644 --- a/src/mcp/server/__main__.py +++ b/src/mcp/server/__main__.py @@ -6,7 +6,7 @@ import anyio from mcp.server.session import ServerSession from mcp.server.stdio import stdio_server -from mcp.server.types import InitializationOptions +from mcp.server.models import InitializationOptions from mcp.types import ServerCapabilities if not sys.warnoptions: diff --git a/src/mcp/server/models.py b/src/mcp/server/models.py new file mode 100644 index 0000000..8920597 --- /dev/null +++ b/src/mcp/server/models.py @@ -0,0 +1,19 @@ +""" +This module provides simpler types to use with the server for managing prompts +and tools. +""" + +from dataclasses import dataclass +from typing import Literal + +from pydantic import BaseModel + +from mcp.types import ( + ServerCapabilities, +) + + +class InitializationOptions(BaseModel): + server_name: str + server_version: str + capabilities: ServerCapabilities diff --git a/src/mcp/server/session.py b/src/mcp/server/session.py index f6ed1b3..03e6882 100644 --- a/src/mcp/server/session.py +++ b/src/mcp/server/session.py @@ -6,7 +6,7 @@ import anyio.lowlevel from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream from pydantic import AnyUrl -from mcp.server.types import InitializationOptions +from mcp.server.models import InitializationOptions from mcp.shared.session import ( BaseSession, RequestResponder, diff --git a/src/mcp/server/types.py b/src/mcp/server/types.py deleted file mode 100644 index 7946a4b..0000000 --- a/src/mcp/server/types.py +++ /dev/null @@ -1,46 +0,0 @@ -""" -This module provides simpler types to use with the server for managing prompts -and tools. -""" - -from dataclasses import dataclass -from typing import Literal - -from pydantic import BaseModel - -from mcp.types import ( - BlobResourceContents, - Role, - ServerCapabilities, - TextResourceContents, -) - - -@dataclass -class ImageContent: - type: Literal["image"] - data: str - mime_type: str - - -@dataclass -class EmbeddedResource: - resource: TextResourceContents | BlobResourceContents - - -@dataclass -class Message: - role: Role - content: str | ImageContent | EmbeddedResource - - -@dataclass -class PromptResponse: - messages: list[Message] - desc: str | None = None - - -class InitializationOptions(BaseModel): - server_name: str - server_version: str - capabilities: ServerCapabilities diff --git a/tests/conftest.py b/tests/conftest.py index 10f3ed6..8d792aa 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,7 +2,7 @@ import pytest from pydantic import AnyUrl from mcp.server import Server -from mcp.server.types import InitializationOptions +from mcp.server.models import InitializationOptions from mcp.types import Resource, ServerCapabilities TEST_INITIALIZATION_OPTIONS = InitializationOptions( diff --git a/tests/server/test_session.py b/tests/server/test_session.py index 66ac58d..728d276 100644 --- a/tests/server/test_session.py +++ b/tests/server/test_session.py @@ -4,7 +4,7 @@ import pytest from mcp.client.session import ClientSession from mcp.server import NotificationOptions, Server from mcp.server.session import ServerSession -from mcp.server.types import InitializationOptions +from mcp.server.models import InitializationOptions from mcp.types import ( ClientNotification, InitializedNotification, diff --git a/uv.lock b/uv.lock index 35e52f4..49d4e13 100644 --- a/uv.lock +++ b/uv.lock @@ -331,15 +331,14 @@ wheels = [ [[package]] name = "pyright" -version = "1.1.388" +version = "1.1.378" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nodeenv" }, - { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/9c/83/e9867538a794638d2d20ac3ab3106a31aca1d9cfea530c9b2921809dae03/pyright-1.1.388.tar.gz", hash = "sha256:0166d19b716b77fd2d9055de29f71d844874dbc6b9d3472ccd22df91db3dfa34", size = 21939 } +sdist = { url = "https://files.pythonhosted.org/packages/3d/f0/e8aa5555410d88f898bef04da2102b0a9bf144658c98d34872e91621ced2/pyright-1.1.378.tar.gz", hash = "sha256:78a043be2876d12d0af101d667e92c7734f3ebb9db71dccc2c220e7e7eb89ca2", size = 17486 } wheels = [ - { url = "https://files.pythonhosted.org/packages/03/57/7fb00363b7f267a398c5bdf4f55f3e64f7c2076b2e7d2901b3373d52b6ff/pyright-1.1.388-py3-none-any.whl", hash = "sha256:c7068e9f2c23539c6ac35fc9efac6c6c1b9aa5a0ce97a9a8a6cf0090d7cbf84c", size = 18579 }, + { url = "https://files.pythonhosted.org/packages/38/c6/f0d4bc20c13b20cecfbf13c699477c825e45767f1dc5068137323f86e495/pyright-1.1.378-py3-none-any.whl", hash = "sha256:8853776138b01bc284da07ac481235be7cc89d3176b073d2dba73636cb95be79", size = 18222 }, ] [[package]]