This commit is contained in:
David Soria Parra 2024-11-11 20:17:39 +00:00
commit ec8c85edea
No known key found for this signature in database
8 changed files with 37 additions and 21 deletions

View file

@ -3,9 +3,9 @@ from datetime import timedelta
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from pydantic import AnyUrl from pydantic import AnyUrl
import mcp.types as types
from mcp.shared.session import BaseSession from mcp.shared.session import BaseSession
from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
import mcp.types as types
class ClientSession( class ClientSession(

View file

@ -84,10 +84,8 @@ async def sse_client(
case "message": case "message":
try: try:
message = ( message = types.JSONRPCMessage.model_validate_json( # noqa: E501
types.JSONRPCMessage.model_validate_json( sse.data
sse.data
)
) )
logger.debug( logger.debug(
f"Received server message: {message}" f"Received server message: {message}"

View file

@ -7,12 +7,12 @@ from typing import Any, Sequence
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from pydantic import AnyUrl from pydantic import AnyUrl
import mcp.types as types
from mcp.server.models import InitializationOptions from mcp.server.models import InitializationOptions
from mcp.server.session import ServerSession from mcp.server.session import ServerSession
from mcp.server.stdio import stdio_server as stdio_server from mcp.server.stdio import stdio_server as stdio_server
from mcp.shared.context import RequestContext from mcp.shared.context import RequestContext
from mcp.shared.session import RequestResponder from mcp.shared.session import RequestResponder
import mcp.types as types
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -36,7 +36,9 @@ class NotificationOptions:
class Server: class Server:
def __init__(self, name: str): def __init__(self, name: str):
self.name = name self.name = name
self.request_handlers: dict[type, Callable[..., Awaitable[types.ServerResult]]] = { self.request_handlers: dict[
type, Callable[..., Awaitable[types.ServerResult]]
] = {
types.PingRequest: _ping_handler, types.PingRequest: _ping_handler,
} }
self.notification_handlers: dict[type, Callable[..., Awaitable[None]]] = {} self.notification_handlers: dict[type, Callable[..., Awaitable[None]]] = {}
@ -153,7 +155,9 @@ class Server:
async def handler(_: Any): async def handler(_: Any):
resources = await func() resources = await func()
return types.ServerResult(types.ListResourcesResult(resources=resources)) return types.ServerResult(
types.ListResourcesResult(resources=resources)
)
self.request_handlers[types.ListResourcesRequest] = handler self.request_handlers[types.ListResourcesRequest] = handler
return func return func
@ -249,7 +253,11 @@ class Server:
def decorator( def decorator(
func: Callable[ func: Callable[
..., ...,
Awaitable[Sequence[types.TextContent | types.ImageContent | types.EmbeddedResource]], Awaitable[
Sequence[
types.TextContent | types.ImageContent | types.EmbeddedResource
]
],
], ],
): ):
logger.debug("Registering handler for CallToolRequest") logger.debug("Registering handler for CallToolRequest")
@ -261,7 +269,9 @@ class Server:
for result in results: for result in results:
match result: match result:
case str() as text: case str() as text:
content.append(types.TextContent(type="text", text=text)) content.append(
types.TextContent(type="text", text=text)
)
case types.ImageContent() as img: case types.ImageContent() as img:
content.append( content.append(
types.ImageContent( types.ImageContent(
@ -277,7 +287,9 @@ class Server:
) )
) )
return types.ServerResult(types.CallToolResult(content=content, isError=False)) return types.ServerResult(
types.CallToolResult(content=content, isError=False)
)
except Exception as e: except Exception as e:
return types.ServerResult( return types.ServerResult(
types.CallToolResult( types.CallToolResult(
@ -312,7 +324,10 @@ class Server:
def decorator( def decorator(
func: Callable[ func: Callable[
[types.PromptReference | types.ResourceReference, types.CompletionArgument], [
types.PromptReference | types.ResourceReference,
types.CompletionArgument,
],
Awaitable[types.Completion | None], Awaitable[types.Completion | None],
], ],
): ):

View file

@ -4,9 +4,9 @@ import sys
import anyio import anyio
from mcp.server.models import InitializationOptions
from mcp.server.session import ServerSession from mcp.server.session import ServerSession
from mcp.server.stdio import stdio_server from mcp.server.stdio import stdio_server
from mcp.server.models import InitializationOptions
from mcp.types import ServerCapabilities from mcp.types import ServerCapabilities
if not sys.warnoptions: if not sys.warnoptions:

View file

@ -3,9 +3,6 @@ This module provides simpler types to use with the server for managing prompts
and tools. and tools.
""" """
from dataclasses import dataclass
from typing import Literal
from pydantic import BaseModel from pydantic import BaseModel
from mcp.types import ( from mcp.types import (

View file

@ -6,12 +6,12 @@ import anyio.lowlevel
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from pydantic import AnyUrl from pydantic import AnyUrl
import mcp.types as types
from mcp.server.models import InitializationOptions from mcp.server.models import InitializationOptions
from mcp.shared.session import ( from mcp.shared.session import (
BaseSession, BaseSession,
RequestResponder, RequestResponder,
) )
import mcp.types as types
class InitializationState(Enum): class InitializationState(Enum):
@ -37,7 +37,9 @@ class ServerSession(
write_stream: MemoryObjectSendStream[types.JSONRPCMessage], write_stream: MemoryObjectSendStream[types.JSONRPCMessage],
init_options: InitializationOptions, init_options: InitializationOptions,
) -> None: ) -> None:
super().__init__(read_stream, write_stream, types.ClientRequest, types.ClientNotification) super().__init__(
read_stream, write_stream, types.ClientRequest, types.ClientNotification
)
self._initialization_state = InitializationState.NotInitialized self._initialization_state = InitializationState.NotInitialized
self._init_options = init_options self._init_options = init_options
@ -65,7 +67,9 @@ class ServerSession(
"Received request before initialization was complete" "Received request before initialization was complete"
) )
async def _received_notification(self, notification: types.ClientNotification) -> None: async def _received_notification(
self, notification: types.ClientNotification
) -> None:
# Need this to avoid ASYNC910 # Need this to avoid ASYNC910
await anyio.lowlevel.checkpoint() await anyio.lowlevel.checkpoint()
match notification.root: match notification.root:

View file

@ -30,7 +30,9 @@ class SseServerTransport:
""" """
_endpoint: str _endpoint: str
_read_stream_writers: dict[UUID, MemoryObjectSendStream[types.JSONRPCMessage | Exception]] _read_stream_writers: dict[
UUID, MemoryObjectSendStream[types.JSONRPCMessage | Exception]
]
def __init__(self, endpoint: str) -> None: def __init__(self, endpoint: str) -> None:
""" """

View file

@ -3,8 +3,8 @@ import pytest
from mcp.client.session import ClientSession from mcp.client.session import ClientSession
from mcp.server import NotificationOptions, Server from mcp.server import NotificationOptions, Server
from mcp.server.session import ServerSession
from mcp.server.models import InitializationOptions from mcp.server.models import InitializationOptions
from mcp.server.session import ServerSession
from mcp.types import ( from mcp.types import (
ClientNotification, ClientNotification,
InitializedNotification, InitializedNotification,