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]]