mirror of
https://github.com/sstent/AICycling_mcp.git
synced 2026-03-06 13:06:12 +00:00
added custom MCP
This commit is contained in:
260
mcp_client.py
260
mcp_client.py
@@ -1,47 +1,103 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
MCP Client - Handles MCP server connections and tool management
|
||||
Updated MCP Client - Uses our custom Garth implementation
|
||||
Fallback to standard MCP if needed, but prefer custom implementation
|
||||
"""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import logging
|
||||
import asyncio
|
||||
from typing import List, Dict, Any, Optional
|
||||
|
||||
from config import Config
|
||||
|
||||
# Try to import both implementations
|
||||
try:
|
||||
from custom_garth_mcp import CustomGarthMCP, GarthTool
|
||||
CUSTOM_GARTH_AVAILABLE = True
|
||||
except ImportError:
|
||||
CUSTOM_GARTH_AVAILABLE = False
|
||||
CustomGarthMCP = None
|
||||
GarthTool = None
|
||||
|
||||
try:
|
||||
from pydantic_ai.mcp import MCPServerStdio
|
||||
MCP_AVAILABLE = True
|
||||
import shutil
|
||||
import os
|
||||
STANDARD_MCP_AVAILABLE = True
|
||||
except ImportError:
|
||||
MCP_AVAILABLE = False
|
||||
STANDARD_MCP_AVAILABLE = False
|
||||
MCPServerStdio = None
|
||||
|
||||
from config import Config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class MCPClient:
|
||||
"""Manages MCP server connection and tool interactions"""
|
||||
"""
|
||||
Enhanced MCP Client that prefers custom Garth implementation
|
||||
Falls back to standard MCP if needed
|
||||
"""
|
||||
|
||||
def __init__(self, config: Config):
|
||||
self.config = config
|
||||
self.mcp_server = None
|
||||
self.garth_mcp = None
|
||||
self.standard_mcp = None
|
||||
self.available_tools = []
|
||||
self._initialized = False
|
||||
self._use_custom = True
|
||||
|
||||
if not MCP_AVAILABLE:
|
||||
logger.warning("MCP not available. Tool functionality will be limited.")
|
||||
return
|
||||
|
||||
# Set up MCP server
|
||||
self._setup_mcp_server()
|
||||
# Decide which implementation to use
|
||||
if CUSTOM_GARTH_AVAILABLE and config.garth_token:
|
||||
logger.info("Using custom Garth MCP implementation")
|
||||
self._use_custom = True
|
||||
elif STANDARD_MCP_AVAILABLE and config.garth_token:
|
||||
logger.info("Falling back to standard MCP implementation")
|
||||
self._use_custom = False
|
||||
else:
|
||||
logger.warning("No MCP implementation available")
|
||||
|
||||
def _setup_mcp_server(self):
|
||||
"""Set up MCP server connection"""
|
||||
async def initialize(self):
|
||||
"""Initialize the preferred MCP implementation"""
|
||||
if not self.config.garth_token:
|
||||
logger.warning("No GARTH_TOKEN provided. MCP tools will be unavailable.")
|
||||
return
|
||||
|
||||
try:
|
||||
if self._use_custom and CUSTOM_GARTH_AVAILABLE:
|
||||
await self._initialize_custom_garth()
|
||||
elif STANDARD_MCP_AVAILABLE:
|
||||
await self._initialize_standard_mcp()
|
||||
else:
|
||||
logger.error("No MCP implementation available")
|
||||
return
|
||||
|
||||
self._initialized = True
|
||||
logger.info("MCP client initialized successfully")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"MCP initialization failed: {e}")
|
||||
# Try fallback if custom failed
|
||||
if self._use_custom and STANDARD_MCP_AVAILABLE:
|
||||
logger.info("Trying fallback to standard MCP")
|
||||
try:
|
||||
self._use_custom = False
|
||||
await self._initialize_standard_mcp()
|
||||
self._initialized = True
|
||||
logger.info("Fallback MCP initialization successful")
|
||||
except Exception as fallback_error:
|
||||
logger.error(f"Fallback MCP initialization also failed: {fallback_error}")
|
||||
|
||||
async def _initialize_custom_garth(self):
|
||||
"""Initialize custom Garth MCP"""
|
||||
self.garth_mcp = CustomGarthMCP(
|
||||
garth_token=self.config.garth_token,
|
||||
cache_ttl=self.config.cache_ttl
|
||||
)
|
||||
await self.garth_mcp.initialize()
|
||||
logger.info("Custom Garth MCP initialized")
|
||||
|
||||
async def _initialize_standard_mcp(self):
|
||||
"""Initialize standard MCP (fallback)"""
|
||||
if not self.config.garth_token:
|
||||
raise ValueError("GARTH_TOKEN required for standard MCP")
|
||||
|
||||
# Set up environment
|
||||
os.environ["GARTH_TOKEN"] = self.config.garth_token
|
||||
env = os.environ.copy()
|
||||
@@ -50,85 +106,116 @@ class MCPClient:
|
||||
# Find server executable
|
||||
server_executable = shutil.which(self.config.garth_mcp_server_path)
|
||||
if not server_executable:
|
||||
logger.error(f"'{self.config.garth_mcp_server_path}' not found in PATH")
|
||||
return
|
||||
raise FileNotFoundError(f"'{self.config.garth_mcp_server_path}' not found in PATH")
|
||||
|
||||
self.mcp_server = MCPServerStdio(
|
||||
self.standard_mcp = MCPServerStdio(
|
||||
command=server_executable,
|
||||
args=["garth-mcp-server"],
|
||||
env=env,
|
||||
)
|
||||
|
||||
async def initialize(self):
|
||||
"""Initialize MCP server connection"""
|
||||
if not self.mcp_server:
|
||||
logger.warning("MCP server not configured")
|
||||
return
|
||||
|
||||
try:
|
||||
logger.info("Initializing MCP server connection...")
|
||||
|
||||
# The MCP server will be initialized when used by the agent
|
||||
# For now, we'll try to list tools to verify connection
|
||||
await asyncio.sleep(0.1) # Give it a moment
|
||||
|
||||
logger.info("MCP server connection established")
|
||||
self._initialized = True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"MCP server initialization failed: {e}")
|
||||
self.mcp_server = None
|
||||
logger.info("Standard MCP initialized")
|
||||
|
||||
async def cleanup(self):
|
||||
"""Cleanup MCP server connection"""
|
||||
if self.mcp_server:
|
||||
# MCP server cleanup is handled by the agent
|
||||
pass
|
||||
"""Cleanup MCP resources"""
|
||||
if self.garth_mcp:
|
||||
await self.garth_mcp.cleanup()
|
||||
# Standard MCP cleanup is handled by the agent
|
||||
|
||||
async def list_tools(self) -> List[Any]:
|
||||
"""List available MCP tools"""
|
||||
if not self.mcp_server:
|
||||
if not self._initialized:
|
||||
return []
|
||||
|
||||
try:
|
||||
if not self.available_tools:
|
||||
self.available_tools = await self.mcp_server.list_tools()
|
||||
return self.available_tools
|
||||
if self._use_custom and self.garth_mcp:
|
||||
if not self.available_tools:
|
||||
self.available_tools = self.garth_mcp.list_tools()
|
||||
return self.available_tools
|
||||
elif self.standard_mcp:
|
||||
if not self.available_tools:
|
||||
self.available_tools = await self.standard_mcp.list_tools()
|
||||
return self.available_tools
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing tools: {e}")
|
||||
return []
|
||||
|
||||
return []
|
||||
|
||||
def list_tools_sync(self) -> List[Any]:
|
||||
"""Synchronous version for compatibility"""
|
||||
if self._use_custom and self.garth_mcp:
|
||||
return self.garth_mcp.list_tools()
|
||||
return []
|
||||
|
||||
async def has_tool(self, tool_name: str) -> bool:
|
||||
"""Check if a specific tool is available"""
|
||||
tools = await self.list_tools()
|
||||
return any(tool.name == tool_name for tool in tools)
|
||||
|
||||
async def call_tool(self, tool_name: str, parameters: Dict[str, Any]) -> Any:
|
||||
"""Call a specific MCP tool directly"""
|
||||
if not self.mcp_server:
|
||||
raise RuntimeError("MCP server not available")
|
||||
if not self._initialized:
|
||||
return False
|
||||
|
||||
try:
|
||||
result = await self.mcp_server.direct_call_tool(tool_name, parameters)
|
||||
return result.output if hasattr(result, 'output') else result
|
||||
if self._use_custom and self.garth_mcp:
|
||||
return await self.garth_mcp.has_tool(tool_name)
|
||||
elif self.standard_mcp:
|
||||
tools = await self.list_tools()
|
||||
return any(tool.name == tool_name for tool in tools)
|
||||
except Exception as e:
|
||||
logger.error(f"Error checking tool {tool_name}: {e}")
|
||||
return False
|
||||
|
||||
return False
|
||||
|
||||
async def call_tool(self, tool_name: str, parameters: Dict[str, Any]) -> Any:
|
||||
"""Call a specific MCP tool"""
|
||||
if not self._initialized:
|
||||
raise RuntimeError("MCP client not initialized")
|
||||
|
||||
try:
|
||||
if self._use_custom and self.garth_mcp:
|
||||
result = await self.garth_mcp.call_tool(tool_name, parameters)
|
||||
logger.debug(f"Custom MCP tool {tool_name} called successfully")
|
||||
return result
|
||||
elif self.standard_mcp:
|
||||
result = await self.standard_mcp.direct_call_tool(tool_name, parameters)
|
||||
# Handle different result formats
|
||||
if hasattr(result, 'output'):
|
||||
return result.output
|
||||
elif hasattr(result, 'content'):
|
||||
return result.content
|
||||
else:
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Error calling tool {tool_name}: {e}")
|
||||
raise
|
||||
|
||||
def get_tool_info(self, tool_name: str) -> Optional[Dict[str, Any]]:
|
||||
"""Get information about a specific tool"""
|
||||
for tool in self.available_tools:
|
||||
tools = self.list_tools_sync() if self._use_custom else []
|
||||
|
||||
for tool in tools:
|
||||
if tool.name == tool_name:
|
||||
return {
|
||||
"name": tool.name,
|
||||
"description": getattr(tool, 'description', ''),
|
||||
"parameters": getattr(tool, 'inputSchema', {}).get('properties', {})
|
||||
}
|
||||
if self._use_custom and isinstance(tool, GarthTool):
|
||||
return {
|
||||
"name": tool.name,
|
||||
"description": tool.description,
|
||||
"parameters": tool.parameters
|
||||
}
|
||||
else:
|
||||
# Standard MCP tool
|
||||
return {
|
||||
"name": tool.name,
|
||||
"description": getattr(tool, 'description', ''),
|
||||
"parameters": getattr(tool, 'inputSchema', {}).get('properties', {})
|
||||
}
|
||||
return None
|
||||
|
||||
def print_tools(self):
|
||||
"""Pretty print available tools"""
|
||||
if not self.available_tools:
|
||||
if self._use_custom and self.garth_mcp:
|
||||
self.garth_mcp.print_tools()
|
||||
return
|
||||
|
||||
tools = self.list_tools_sync()
|
||||
if not tools:
|
||||
print("No MCP tools available")
|
||||
return
|
||||
|
||||
@@ -136,7 +223,7 @@ class MCPClient:
|
||||
print("AVAILABLE MCP TOOLS")
|
||||
print(f"{'='*60}")
|
||||
|
||||
for i, tool in enumerate(self.available_tools, 1):
|
||||
for i, tool in enumerate(tools, 1):
|
||||
print(f"\n{i}. {tool.name}")
|
||||
if hasattr(tool, 'description') and tool.description:
|
||||
print(f" Description: {tool.description}")
|
||||
@@ -156,5 +243,42 @@ class MCPClient:
|
||||
|
||||
@property
|
||||
def is_available(self) -> bool:
|
||||
"""Check if MCP server is available"""
|
||||
return self.mcp_server is not None and self._initialized
|
||||
"""Check if MCP is available and initialized"""
|
||||
return self._initialized and (
|
||||
(self._use_custom and self.garth_mcp is not None) or
|
||||
(not self._use_custom and self.standard_mcp is not None)
|
||||
)
|
||||
|
||||
@property
|
||||
def implementation_type(self) -> str:
|
||||
"""Get the type of MCP implementation being used"""
|
||||
if self._use_custom:
|
||||
return "custom_garth"
|
||||
else:
|
||||
return "standard_mcp"
|
||||
|
||||
def get_cache_stats(self) -> Dict[str, Any]:
|
||||
"""Get cache statistics (only available with custom implementation)"""
|
||||
if self._use_custom and self.garth_mcp:
|
||||
return self.garth_mcp.get_cache_stats()
|
||||
else:
|
||||
return {"message": "Cache stats only available with custom implementation"}
|
||||
|
||||
def clear_cache(self):
|
||||
"""Clear cache (only available with custom implementation)"""
|
||||
if self._use_custom and self.garth_mcp:
|
||||
self.garth_mcp.clear_cache()
|
||||
logger.info("Cache cleared")
|
||||
else:
|
||||
logger.warning("Cache clearing only available with custom implementation")
|
||||
|
||||
# Compatibility methods for existing code
|
||||
@property
|
||||
def mcp_server(self):
|
||||
"""Compatibility property for existing code using agent integration"""
|
||||
if self._use_custom:
|
||||
# For custom implementation, we can't provide direct agent integration
|
||||
# Return None to indicate tools should be called directly
|
||||
return None
|
||||
else:
|
||||
return self.standard_mcp
|
||||
Reference in New Issue
Block a user