Coverage for me2ai_mcp\base.py: 11%

94 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-04-13 11:30 +0200

1""" 

2Base classes for ME2AI MCP servers. 

3 

4This module provides base classes for creating standardized ME2AI MCP servers 

5that extend the functionality of the official MCP package. 

6""" 

7from typing import Dict, List, Any, Optional, Union, Callable, Type, TypeVar 

8import logging 

9import asyncio 

10import inspect 

11import json 

12from dataclasses import dataclass, field 

13from functools import wraps 

14import os 

15from pathlib import Path 

16 

17# Import the official MCP package 

18from mcp import MCPServer as OfficialMCPServer 

19from mcp import register_tool as official_register_tool 

20from mcp import MCPToolInput 

21 

22# Configure logging 

23logging.basicConfig( 

24 level=logging.INFO, 

25 format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" 

26) 

27 

28T = TypeVar("T", bound="ME2AIMCPServer") 

29 

30 

31class ME2AIMCPServer(OfficialMCPServer): 

32 """Base class for all ME2AI MCP servers with enhanced functionality.""" 

33 

34 def __init__( 

35 self, 

36 server_name: str, 

37 description: str = "", 

38 version: str = "0.1.0", 

39 debug: bool = False 

40 ) -> None: 

41 """Initialize a ME2AI MCP server. 

42  

43 Args: 

44 server_name: Unique name of the MCP server 

45 description: Human-readable description of the server 

46 version: Server version 

47 debug: Whether to enable debug logging 

48 """ 

49 super().__init__(server_name) 

50 

51 # Additional properties 

52 self.description = description 

53 self.version = version 

54 self.metadata = { 

55 "server_name": server_name, 

56 "description": description, 

57 "version": version, 

58 "framework": "ME2AI MCP" 

59 } 

60 

61 # Configure logging 

62 self.logger = logging.getLogger(f"me2ai-mcp-{server_name}") 

63 

64 if debug: 

65 self.logger.setLevel(logging.DEBUG) 

66 else: 

67 self.logger.setLevel(logging.INFO) 

68 

69 self.logger.info(f"Initializing ME2AI MCP Server: {server_name} v{version}") 

70 

71 # Statistics tracking 

72 self.stats = { 

73 "requests": 0, 

74 "errors": 0, 

75 "tool_calls": {} 

76 } 

77 

78 # Register built-in tools 

79 self._register_builtin_tools() 

80 

81 def _register_builtin_tools(self) -> None: 

82 """Register built-in tools for all ME2AI MCP servers.""" 

83 # Register server info tool 

84 @official_register_tool 

85 async def server_info(self) -> Dict[str, Any]: 

86 """Get information about this MCP server.""" 

87 self.stats["requests"] += 1 

88 self.stats["tool_calls"]["server_info"] = self.stats["tool_calls"].get("server_info", 0) + 1 

89 

90 return { 

91 "success": True, 

92 "server": self.metadata, 

93 "stats": self.stats, 

94 "tools": list(self.tools.keys()) 

95 } 

96 

97 @classmethod 

98 def from_config(cls: Type[T], config_path: Union[str, Path]) -> T: 

99 """Create an MCP server from a configuration file. 

100  

101 Args: 

102 config_path: Path to the configuration file 

103  

104 Returns: 

105 Configured MCP server instance 

106 """ 

107 config_path = Path(config_path) 

108 

109 if not config_path.exists(): 

110 raise FileNotFoundError(f"Configuration file not found: {config_path}") 

111 

112 with open(config_path, "r") as f: 

113 config = json.load(f) 

114 

115 server_name = config.get("server_name", "unnamed-server") 

116 description = config.get("description", "") 

117 version = config.get("version", "0.1.0") 

118 debug = config.get("debug", False) 

119 

120 server = cls( 

121 server_name=server_name, 

122 description=description, 

123 version=version, 

124 debug=debug 

125 ) 

126 

127 return server 

128 

129 async def start(self) -> None: 

130 """Start the MCP server with enhanced logging and error handling.""" 

131 try: 

132 self.logger.info(f"Starting {self.server_name} MCP server (v{self.version})") 

133 self.logger.info(f"Available tools: {', '.join(self.tools.keys())}") 

134 

135 # Start the official MCP server 

136 await super().start() 

137 

138 except Exception as e: 

139 self.logger.error(f"Error starting MCP server: {str(e)}") 

140 self.stats["errors"] += 1 

141 raise 

142 

143 

144@dataclass 

145class BaseTool: 

146 """Base class for ME2AI MCP tools with enhanced functionality.""" 

147 

148 name: str 

149 description: str = "" 

150 enabled: bool = True 

151 stats: Dict[str, Any] = field(default_factory=lambda: { 

152 "calls": 0, 

153 "errors": 0, 

154 "last_call": None 

155 }) 

156 

157 def __post_init__(self) -> None: 

158 """Initialize the tool after all fields are set.""" 

159 self.logger = logging.getLogger(f"me2ai-mcp-tool-{self.name}") 

160 

161 async def execute(self, params: Dict[str, Any]) -> Dict[str, Any]: 

162 """Execute the tool with the given parameters. 

163  

164 This method should be overridden by subclasses. 

165  

166 Args: 

167 params: Tool parameters 

168  

169 Returns: 

170 Tool execution result 

171 """ 

172 raise NotImplementedError("Tool execution not implemented") 

173 

174 

175def register_tool(func: Optional[Callable] = None, *, name: Optional[str] = None) -> Callable: 

176 """Enhanced decorator for registering tools with ME2AI MCP servers. 

177  

178 This extends the official register_tool decorator with additional 

179 functionality like automatic error handling and logging. 

180  

181 Args: 

182 func: Function to register as a tool 

183 name: Custom name for the tool (optional) 

184  

185 Returns: 

186 Decorated function 

187 """ 

188 def decorator(func: Callable) -> Callable: 

189 # Get the function's signature for better error messages 

190 sig = inspect.signature(func) 

191 

192 @wraps(func) 

193 async def wrapper(self, *args, **kwargs) -> Dict[str, Any]: 

194 tool_name = name or func.__name__ 

195 

196 # Update stats 

197 if hasattr(self, "stats") and isinstance(self.stats, dict): 

198 self.stats["requests"] += 1 

199 self.stats["tool_calls"][tool_name] = self.stats["tool_calls"].get(tool_name, 0) + 1 

200 

201 # Log the call 

202 if hasattr(self, "logger"): 

203 params_str = str(kwargs) if kwargs else str(args) 

204 self.logger.info(f"Tool call: {tool_name}({params_str[:100]}{'...' if len(params_str) > 100 else ''})") 

205 

206 try: 

207 # Execute the tool function 

208 result = await func(self, *args, **kwargs) 

209 

210 # Ensure the result is a dictionary with success status 

211 if isinstance(result, dict) and "success" not in result: 

212 result["success"] = True 

213 

214 return result 

215 

216 except Exception as e: 

217 # Handle errors consistently 

218 error_message = f"Error executing {tool_name}: {str(e)}" 

219 

220 if hasattr(self, "logger"): 

221 self.logger.error(error_message, exc_info=True) 

222 

223 if hasattr(self, "stats") and isinstance(self.stats, dict): 

224 self.stats["errors"] += 1 

225 

226 return { 

227 "success": False, 

228 "error": error_message, 

229 "exception_type": type(e).__name__ 

230 } 

231 

232 # Register with the official decorator 

233 official_register_tool(wrapper) 

234 

235 return wrapper 

236 

237 # Handle both @register_tool and @register_tool(name="custom_name") forms 

238 if func is None: 

239 return decorator 

240 else: 

241 return decorator(func)