From 368782c15543a655d06e6098afe81e3926275712 Mon Sep 17 00:00:00 2001 From: TerminalMan <84923604+SecretiveShell@users.noreply.github.com> Date: Sat, 4 Jan 2025 16:29:37 +0000 Subject: [PATCH] add request handler --- src/mcp/client/session.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/src/mcp/client/session.py b/src/mcp/client/session.py index 254b5bc..87df780 100644 --- a/src/mcp/client/session.py +++ b/src/mcp/client/session.py @@ -6,7 +6,7 @@ from pydantic import AnyUrl import mcp.types as types -from mcp.shared.session import BaseSession +from mcp.shared.session import BaseSession, RequestResponder from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS sampling_function_signature = Callable[ @@ -255,3 +255,17 @@ async def send_roots_list_changed(self) -> None: ) ) ) + + async def _received_request( + self, responder: RequestResponder["types.ServerRequest", "types.ClientResult"] + ) -> None: + if isinstance(responder.request.root, types.CreateMessageRequest): + print("Received create message request") + if self.sampling_callback is None: + raise RuntimeError("Sampling callback is not set") + response = await self.sampling_callback(responder.request.root.params) + + client_response = types.ClientResult(**response.model_dump()) + + print(f"Response: {response.dict()}") + await responder.respond(client_response)