34df96bf2c
Co-authored-by: Matt Zhou <mattzh1314@gmail.com> Co-authored-by: Kian Jones <11655409+kianjones9@users.noreply.github.com>
157 lines
6.2 KiB
Python
157 lines
6.2 KiB
Python
from letta.log import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
# class BaseMCPClient:
|
|
# def __init__(self, server_config: BaseServerConfig):
|
|
# self.server_config = server_config
|
|
# self.session: Optional[ClientSession] = None
|
|
# self.stdio = None
|
|
# self.write = None
|
|
# self.initialized = False
|
|
# self.loop = asyncio.new_event_loop()
|
|
# self.cleanup_funcs = []
|
|
#
|
|
# def connect_to_server(self):
|
|
# asyncio.set_event_loop(self.loop)
|
|
# success = self._initialize_connection(self.server_config, timeout=tool_settings.mcp_connect_to_server_timeout)
|
|
#
|
|
# if success:
|
|
# try:
|
|
# self.loop.run_until_complete(
|
|
# asyncio.wait_for(self.session.initialize(), timeout=tool_settings.mcp_connect_to_server_timeout)
|
|
# )
|
|
# self.initialized = True
|
|
# except asyncio.TimeoutError:
|
|
# raise MCPTimeoutError("initializing session", self.server_config.server_name, tool_settings.mcp_connect_to_server_timeout)
|
|
# else:
|
|
# raise RuntimeError(
|
|
# f"Connecting to MCP server failed. Please review your server config: {self.server_config.model_dump_json(indent=4)}"
|
|
# )
|
|
#
|
|
# def _initialize_connection(self, server_config: BaseServerConfig, timeout: float) -> bool:
|
|
# raise NotImplementedError("Subclasses must implement _initialize_connection")
|
|
#
|
|
# def list_tools(self) -> List[MCPTool]:
|
|
# self._check_initialized()
|
|
# try:
|
|
# response = self.loop.run_until_complete(
|
|
# asyncio.wait_for(self.session.list_tools(), timeout=tool_settings.mcp_list_tools_timeout)
|
|
# )
|
|
# return response.tools
|
|
# except asyncio.TimeoutError:
|
|
# logger.error(
|
|
# f"Timed out while listing tools for MCP server {self.server_config.server_name} (timeout={tool_settings.mcp_list_tools_timeout}s)."
|
|
# )
|
|
# raise MCPTimeoutError("listing tools", self.server_config.server_name, tool_settings.mcp_list_tools_timeout)
|
|
#
|
|
# def execute_tool(self, tool_name: str, tool_args: dict) -> Tuple[str, bool]:
|
|
# self._check_initialized()
|
|
# try:
|
|
# result = self.loop.run_until_complete(
|
|
# asyncio.wait_for(self.session.call_tool(tool_name, tool_args), timeout=tool_settings.mcp_execute_tool_timeout)
|
|
# )
|
|
#
|
|
# parsed_content = []
|
|
# for content_piece in result.content:
|
|
# if isinstance(content_piece, TextContent):
|
|
# parsed_content.append(content_piece.text)
|
|
# print("parsed_content (text)", parsed_content)
|
|
# else:
|
|
# parsed_content.append(str(content_piece))
|
|
# print("parsed_content (other)", parsed_content)
|
|
#
|
|
# if len(parsed_content) > 0:
|
|
# final_content = " ".join(parsed_content)
|
|
# else:
|
|
# # TODO move hardcoding to constants
|
|
# final_content = "Empty response from tool"
|
|
#
|
|
# return final_content, result.isError
|
|
# except asyncio.TimeoutError:
|
|
# logger.error(
|
|
# f"Timed out while executing tool '{tool_name}' for MCP server {self.server_config.server_name} (timeout={tool_settings.mcp_execute_tool_timeout}s)."
|
|
# )
|
|
# raise MCPTimeoutError(f"executing tool '{tool_name}'", self.server_config.server_name, tool_settings.mcp_execute_tool_timeout)
|
|
#
|
|
# def _check_initialized(self):
|
|
# if not self.initialized:
|
|
# logger.error("MCPClient has not been initialized")
|
|
# raise RuntimeError("MCPClient has not been initialized")
|
|
#
|
|
# def cleanup(self):
|
|
# try:
|
|
# for cleanup_func in self.cleanup_funcs:
|
|
# cleanup_func()
|
|
# self.initialized = False
|
|
# if not self.loop.is_closed():
|
|
# self.loop.close()
|
|
# except Exception as e:
|
|
# logger.warning(e)
|
|
# finally:
|
|
# logger.info("Cleaned up MCP clients on shutdown.")
|
|
#
|
|
#
|
|
# class BaseAsyncMCPClient:
|
|
# def __init__(self, server_config: BaseServerConfig):
|
|
# self.server_config = server_config
|
|
# self.session: Optional[ClientSession] = None
|
|
# self.stdio = None
|
|
# self.write = None
|
|
# self.initialized = False
|
|
# self.cleanup_funcs = []
|
|
#
|
|
# async def connect_to_server(self):
|
|
#
|
|
# success = await self._initialize_connection(self.server_config, timeout=tool_settings.mcp_connect_to_server_timeout)
|
|
#
|
|
# if success:
|
|
# self.initialized = True
|
|
# else:
|
|
# raise RuntimeError(
|
|
# f"Connecting to MCP server failed. Please review your server config: {self.server_config.model_dump_json(indent=4)}"
|
|
# )
|
|
#
|
|
# async def list_tools(self) -> List[MCPTool]:
|
|
# self._check_initialized()
|
|
# response = await self.session.list_tools()
|
|
# return response.tools
|
|
#
|
|
# async def execute_tool(self, tool_name: str, tool_args: dict) -> Tuple[str, bool]:
|
|
# self._check_initialized()
|
|
# result = await self.session.call_tool(tool_name, tool_args)
|
|
#
|
|
# parsed_content = []
|
|
# for content_piece in result.content:
|
|
# if isinstance(content_piece, TextContent):
|
|
# parsed_content.append(content_piece.text)
|
|
# else:
|
|
# parsed_content.append(str(content_piece))
|
|
#
|
|
# if len(parsed_content) > 0:
|
|
# final_content = " ".join(parsed_content)
|
|
# else:
|
|
# # TODO move hardcoding to constants
|
|
# final_content = "Empty response from tool"
|
|
#
|
|
# return final_content, result.isError
|
|
#
|
|
# def _check_initialized(self):
|
|
# if not self.initialized:
|
|
# logger.error("MCPClient has not been initialized")
|
|
# raise RuntimeError("MCPClient has not been initialized")
|
|
#
|
|
# async def cleanup(self):
|
|
# try:
|
|
# for cleanup_func in self.cleanup_funcs:
|
|
# cleanup_func()
|
|
# self.initialized = False
|
|
# if not self.loop.is_closed():
|
|
# self.loop.close()
|
|
# except Exception as e:
|
|
# logger.warning(e)
|
|
# finally:
|
|
# logger.info("Cleaned up MCP clients on shutdown.")
|
|
#
|