wehub-resource-sync
105 lines
3.9 KiB
Python
105 lines
3.9 KiB
Python
from contextlib import AsyncExitStack
|
|
from enum import Enum
|
|
|
|
from mcp import ClientSession, StdioServerParameters
|
|
from mcp.client.stdio import stdio_client
|
|
|
|
from ..tools.mcp_tool import MCPTool
|
|
from .config import MCPServerConfig
|
|
|
|
|
|
class MCPServerStatus(Enum):
|
|
DISCONNECTED = "disconnected" # Server is disconnected or experiencing errors
|
|
CONNECTING = "connecting" # Server is in the process of connecting
|
|
CONNECTED = "connected" # Server is connected and ready to use
|
|
|
|
|
|
class MCPDiscoveryState(Enum):
|
|
"""State of MCP discovery process."""
|
|
|
|
NOT_STARTED = "not_started" # Discovery has not started yet
|
|
IN_PROGRESS = "in_progress" # Discovery is currently in progress
|
|
# Discovery has completed (with or without errors)
|
|
COMPLETED = "completed"
|
|
|
|
|
|
class MCPClient:
|
|
def __init__(self):
|
|
# Initialize session and client objects
|
|
self.session: ClientSession | None = None
|
|
self.exit_stack = AsyncExitStack()
|
|
self.mcp_servers_status: dict[str, MCPServerStatus] = {}
|
|
|
|
def get_mcp_server_status(self, mcp_server_name: str) -> MCPServerStatus:
|
|
return self.mcp_servers_status.get(mcp_server_name, MCPServerStatus.DISCONNECTED)
|
|
|
|
def update_mcp_server_status(self, mcp_server_name, status: MCPServerStatus):
|
|
self.mcp_servers_status[mcp_server_name] = status
|
|
|
|
async def connect_and_discover(
|
|
self,
|
|
mcp_server_name: str,
|
|
mcp_server_config: MCPServerConfig,
|
|
mcp_tools_container: list,
|
|
model_provider,
|
|
):
|
|
transport = None
|
|
if mcp_server_config.http_url:
|
|
raise NotImplementedError("HTTP transport is not implemented yet")
|
|
elif mcp_server_config.url:
|
|
raise NotImplementedError("WebSocket transport is not implemented yet")
|
|
elif mcp_server_config.command:
|
|
params = StdioServerParameters(
|
|
command=mcp_server_config.command,
|
|
args=mcp_server_config.args,
|
|
env=mcp_server_config.env,
|
|
cwd=mcp_server_config.cwd,
|
|
)
|
|
transport = await self.exit_stack.enter_async_context(stdio_client(params))
|
|
else:
|
|
# error
|
|
raise ValueError(
|
|
f"Invalid MCP server configuration for {mcp_server_name}. "
|
|
"Please provide either a command or a URL."
|
|
)
|
|
try:
|
|
await self.connect_to_server(mcp_server_name, transport)
|
|
mcp_tools = await self.list_tools()
|
|
for tool in mcp_tools.tools:
|
|
mcp_tool = MCPTool(self, tool, model_provider)
|
|
mcp_tools_container.append(mcp_tool)
|
|
except Exception as e:
|
|
raise e
|
|
|
|
async def connect_to_server(self, mcp_server_name, transport):
|
|
"""Connect to an MCP server
|
|
|
|
Args:
|
|
server_params: Parameters for connecting to the MCP server.
|
|
"""
|
|
if self.get_mcp_server_status(mcp_server_name) != MCPServerStatus.CONNECTED:
|
|
self.update_mcp_server_status(mcp_server_name, MCPServerStatus.CONNECTING)
|
|
try:
|
|
stdio, write = transport
|
|
self.session = await self.exit_stack.enter_async_context(
|
|
ClientSession(stdio, write)
|
|
)
|
|
await self.session.initialize()
|
|
self.update_mcp_server_status(mcp_server_name, MCPServerStatus.CONNECTED)
|
|
except Exception as e:
|
|
self.update_mcp_server_status(mcp_server_name, MCPServerStatus.DISCONNECTED)
|
|
raise e
|
|
|
|
async def call_tool(self, name, args):
|
|
output = await self.session.call_tool(name, args)
|
|
return output
|
|
|
|
async def list_tools(self):
|
|
tools = await self.session.list_tools()
|
|
return tools
|
|
|
|
async def cleanup(self, mcp_server_name):
|
|
"""Clean up resources"""
|
|
await self.exit_stack.aclose()
|
|
self.update_mcp_server_status(mcp_server_name, MCPServerStatus.DISCONNECTED)
|