refactor: flatten the methods in Client (#1914)
This commit is contained in:
parent
3dc8b72041
commit
bcb07c2b25
4 changed files with 213 additions and 315 deletions
|
|
@ -36,12 +36,7 @@ class InMemoryTransport:
|
|||
result = await client.call_tool("my_tool", {...})
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
server: Server[Any] | FastMCP,
|
||||
*,
|
||||
raise_exceptions: bool = False,
|
||||
) -> None:
|
||||
def __init__(self, server: Server[Any] | FastMCP, *, raise_exceptions: bool = False) -> None:
|
||||
"""Initialize the in-memory transport.
|
||||
|
||||
Args:
|
||||
|
|
|
|||
|
|
@ -58,6 +58,7 @@ class Client:
|
|||
self,
|
||||
server: Server[Any] | FastMCP,
|
||||
*,
|
||||
# TODO(Marcelo): When do `raise_exceptions=True` actually raises?
|
||||
raise_exceptions: bool = False,
|
||||
read_timeout_seconds: float | None = None,
|
||||
sampling_callback: SamplingFnT | None = None,
|
||||
|
|
@ -125,14 +126,9 @@ class Client:
|
|||
self._exit_stack = exit_stack.pop_all()
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_val: BaseException | None,
|
||||
exc_tb: Any,
|
||||
) -> None:
|
||||
async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
|
||||
"""Exit the async context manager."""
|
||||
if self._exit_stack:
|
||||
if self._exit_stack: # pragma: no branch
|
||||
await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)
|
||||
self._session = None
|
||||
|
||||
|
|
@ -177,28 +173,22 @@ class Client:
|
|||
"""Set the logging level on the server."""
|
||||
return await self.session.set_logging_level(level)
|
||||
|
||||
async def list_resources(
|
||||
self,
|
||||
params: types.PaginatedRequestParams | None = None,
|
||||
) -> types.ListResourcesResult:
|
||||
async def list_resources(self, *, cursor: str | None = None) -> types.ListResourcesResult:
|
||||
"""List available resources from the server."""
|
||||
return await self.session.list_resources(params=params)
|
||||
return await self.session.list_resources(params=types.PaginatedRequestParams(cursor=cursor))
|
||||
|
||||
async def list_resource_templates(
|
||||
self,
|
||||
params: types.PaginatedRequestParams | None = None,
|
||||
) -> types.ListResourceTemplatesResult:
|
||||
async def list_resource_templates(self, *, cursor: str | None = None) -> types.ListResourceTemplatesResult:
|
||||
"""List available resource templates from the server."""
|
||||
return await self.session.list_resource_templates(params=params)
|
||||
return await self.session.list_resource_templates(params=types.PaginatedRequestParams(cursor=cursor))
|
||||
|
||||
async def read_resource(self, uri: str | AnyUrl) -> types.ReadResourceResult:
|
||||
"""Read a resource from the server.
|
||||
|
||||
Args:
|
||||
uri: The URI of the resource to read
|
||||
uri: The URI of the resource to read.
|
||||
|
||||
Returns:
|
||||
The resource content
|
||||
The resource content.
|
||||
"""
|
||||
return await self.session.read_resource(uri)
|
||||
|
||||
|
|
@ -239,18 +229,11 @@ class Client:
|
|||
meta=meta,
|
||||
)
|
||||
|
||||
async def list_prompts(
|
||||
self,
|
||||
params: types.PaginatedRequestParams | None = None,
|
||||
) -> types.ListPromptsResult:
|
||||
async def list_prompts(self, *, cursor: str | None = None) -> types.ListPromptsResult:
|
||||
"""List available prompts from the server."""
|
||||
return await self.session.list_prompts(params=params)
|
||||
return await self.session.list_prompts(params=types.PaginatedRequestParams(cursor=cursor))
|
||||
|
||||
async def get_prompt(
|
||||
self,
|
||||
name: str,
|
||||
arguments: dict[str, str] | None = None,
|
||||
) -> types.GetPromptResult:
|
||||
async def get_prompt(self, name: str, arguments: dict[str, str] | None = None) -> types.GetPromptResult:
|
||||
"""Get a prompt from the server.
|
||||
|
||||
Args:
|
||||
|
|
@ -258,7 +241,7 @@ class Client:
|
|||
arguments: Arguments to pass to the prompt
|
||||
|
||||
Returns:
|
||||
The prompt content
|
||||
The prompt content.
|
||||
"""
|
||||
return await self.session.get_prompt(name=name, arguments=arguments)
|
||||
|
||||
|
|
@ -276,21 +259,15 @@ class Client:
|
|||
context_arguments: Additional context arguments
|
||||
|
||||
Returns:
|
||||
Completion suggestions
|
||||
Completion suggestions.
|
||||
"""
|
||||
return await self.session.complete(
|
||||
ref=ref,
|
||||
argument=argument,
|
||||
context_arguments=context_arguments,
|
||||
)
|
||||
return await self.session.complete(ref=ref, argument=argument, context_arguments=context_arguments)
|
||||
|
||||
async def list_tools(
|
||||
self,
|
||||
params: types.PaginatedRequestParams | None = None,
|
||||
) -> types.ListToolsResult:
|
||||
async def list_tools(self, *, cursor: str | None = None) -> types.ListToolsResult:
|
||||
"""List available tools from the server."""
|
||||
return await self.session.list_tools(params=params)
|
||||
return await self.session.list_tools(params=types.PaginatedRequestParams(cursor=cursor))
|
||||
|
||||
async def send_roots_list_changed(self) -> None:
|
||||
"""Send a notification that the roots list has changed."""
|
||||
await self.session.send_roots_list_changed()
|
||||
# TODO(Marcelo): Currently, there is no way for the server to handle this. We should add support.
|
||||
await self.session.send_roots_list_changed() # pragma: no cover
|
||||
|
|
|
|||
|
|
@ -349,10 +349,7 @@ class ClientSession(
|
|||
Args:
|
||||
params: Full pagination parameters including cursor and any future fields
|
||||
"""
|
||||
return await self.send_request(
|
||||
types.ListPromptsRequest(params=params),
|
||||
types.ListPromptsResult,
|
||||
)
|
||||
return await self.send_request(types.ListPromptsRequest(params=params), types.ListPromptsResult)
|
||||
|
||||
async def get_prompt(self, name: str, arguments: dict[str, str] | None = None) -> types.GetPromptResult:
|
||||
"""Send a prompts/get request."""
|
||||
|
|
|
|||
|
|
@ -234,12 +234,12 @@ class BaseSession(
|
|||
metadata: MessageMetadata = None,
|
||||
progress_callback: ProgressFnT | None = None,
|
||||
) -> ReceiveResultT:
|
||||
"""Sends a request and wait for a response. Raises an McpError if the
|
||||
response contains an error. If a request read timeout is provided, it
|
||||
will take precedence over the session read timeout.
|
||||
"""Sends a request and wait for a response.
|
||||
|
||||
Do not use this method to emit notifications! Use send_notification()
|
||||
instead.
|
||||
Raises an McpError if the response contains an error. If a request read timeout is provided, it will take
|
||||
precedence over the session read timeout.
|
||||
|
||||
Do not use this method to emit notifications! Use send_notification() instead.
|
||||
"""
|
||||
request_id = self._request_id
|
||||
self._request_id = request_id + 1
|
||||
|
|
@ -261,15 +261,10 @@ class BaseSession(
|
|||
|
||||
try:
|
||||
jsonrpc_request = JSONRPCRequest(jsonrpc="2.0", id=request_id, **request_data)
|
||||
|
||||
await self._write_stream.send(SessionMessage(message=jsonrpc_request, metadata=metadata))
|
||||
|
||||
# request read timeout takes precedence over session read timeout
|
||||
timeout = None
|
||||
if request_read_timeout_seconds is not None: # pragma: no cover
|
||||
timeout = request_read_timeout_seconds
|
||||
elif self._session_read_timeout_seconds is not None: # pragma: no cover
|
||||
timeout = self._session_read_timeout_seconds
|
||||
timeout = request_read_timeout_seconds or self._session_read_timeout_seconds
|
||||
|
||||
try:
|
||||
with anyio.fail_after(timeout):
|
||||
|
|
@ -279,9 +274,8 @@ class BaseSession(
|
|||
ErrorData(
|
||||
code=httpx.codes.REQUEST_TIMEOUT,
|
||||
message=(
|
||||
f"Timed out while waiting for response to "
|
||||
f"{request.__class__.__name__}. Waited "
|
||||
f"{timeout} seconds."
|
||||
f"Timed out while waiting for response to {request.__class__.__name__}. "
|
||||
f"Waited {timeout} seconds."
|
||||
),
|
||||
)
|
||||
)
|
||||
|
|
@ -302,9 +296,7 @@ class BaseSession(
|
|||
notification: SendNotificationT,
|
||||
related_request_id: RequestId | None = None,
|
||||
) -> None:
|
||||
"""Emits a notification, which is a one-way message that does not expect
|
||||
a response.
|
||||
"""
|
||||
"""Emits a notification, which is a one-way message that does not expect a response."""
|
||||
# Some transport implementations may need to set the related_request_id
|
||||
# to attribute to the notifications to the request that triggered them.
|
||||
jsonrpc_notification = JSONRPCNotification(
|
||||
|
|
@ -373,11 +365,7 @@ class BaseSession(
|
|||
error_response = JSONRPCError(
|
||||
jsonrpc="2.0",
|
||||
id=message.message.id,
|
||||
error=ErrorData(
|
||||
code=INVALID_PARAMS,
|
||||
message="Invalid request parameters",
|
||||
data="",
|
||||
),
|
||||
error=ErrorData(code=INVALID_PARAMS, message="Invalid request parameters", data=""),
|
||||
)
|
||||
session_message = SessionMessage(message=error_response)
|
||||
await self._write_stream.send(session_message)
|
||||
|
|
@ -518,13 +506,9 @@ class BaseSession(
|
|||
total: float | None = None,
|
||||
message: str | None = None,
|
||||
) -> None:
|
||||
"""Sends a progress notification for a request that is currently being
|
||||
processed.
|
||||
"""
|
||||
"""Sends a progress notification for a request that is currently being processed."""
|
||||
|
||||
async def _handle_incoming(
|
||||
self,
|
||||
req: RequestResponder[ReceiveRequestT, SendResultT] | ReceiveNotificationT | Exception,
|
||||
self, req: RequestResponder[ReceiveRequestT, SendResultT] | ReceiveNotificationT | Exception
|
||||
) -> None:
|
||||
"""A generic handler for incoming messages. Overwritten by subclasses."""
|
||||
pass # pragma: no cover
|
||||
|
|
|
|||
Loading…
Reference in a new issue